Model Fusion via Neuron Transplantation

TL;DR

This paper introduces Neuron Transplantation, a novel model fusion method that combines ensemble models by transplanting important neurons, reducing memory and inference costs while maintaining or improving performance.

Contribution

The paper presents a new neuron transplantation technique for model fusion that outperforms traditional ensemble methods in efficiency and often in accuracy.

Findings

Neuron Transplantation outperforms individual models after fine-tuning.

NT requires less fine-tuning and memory than OT-fusion.

The fused models achieve comparable or better performance.

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.