Model Fusion via Neuron Transplantation
This work addresses the efficiency issue in ensemble learning for neural network practitioners, offering a novel fusion method that reduces computational overhead while maintaining or improving accuracy.
The paper tackles the problem of ensemble learning's high memory and inference costs by proposing Neuron Transplantation, a model fusion technique that transplants important neurons from ensemble members into a pruned model, achieving performance comparable to or better than individual ensemble members after fine-tuning, with faster fusion and less memory usage than alignment-based methods.
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.