Model Fusion via Neuron Transplantation

Abstract

Ensemble learning is a widespread technique to improve the prediction performance of neural networks. However, it comes at the price of increased memory and inference time. In this work we propose a novel model fusion technique called \emph{Neuron Transplantation (NT)} in which we fuse an ensemble of models by transplanting important neurons from all ensemble members into the vacant space obtained by pruning insignificant neurons. An initial loss in performance post-transplantation can be quickly recovered via fine-tuning, consistently outperforming individual ensemble members of the same model capacity and architecture. Furthermore, NT enables all the ensemble members to be jointly pruned and jointly trained in a combined model. Comparing it to alignment-based averaging (like Optimal-Transport-fusion), it requires less fine-tuning than the corresponding OT-fused model, the fusion itself is faster and requires less memory, while the resulting model performance is comparable or better. The code is available under the following link: https://github.com/masterbaer/neuron-transplantation.

18 pages, 7 figures, conference: ECML-PKDD 2024

References in corpus (13)

• PyTorch: An Imperative Style, High-Performance Deep Learning Library
• Distilling the Knowledge in a Neural Network
• PaLM: Scaling Language Modeling with Pathways
• A Survey on Ensemble Learning under the Era of Deep Learning
• Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time
• Distributed Deep Learning Using Synchronous Stochastic Gradient Descent
• A Survey on Deep Neural Network Compression: Challenges, Overview, and Solutions
• Git Re-Basin: Merging Models modulo Permutation Symmetries
• Lost in Pruning: The Effects of Pruning Neural Networks beyond Test Accuracy
• Deep Model Fusion: A Survey
• Non-Iterative Knowledge Fusion in Deep Convolutional Neural Networks
• DepGraph: Towards Any Structural Pruning
• Ensemble-Compression: A New Method for Parallel Training of Deep Neural Networks