Improving Model Fusion by Training-time Neuron Alignment with Fixed Neuron Anchors

TL;DR

This work addresses the barrier to model fusion caused by permutation invariance in neural networks by introducing training-time neuron alignment. The proposed method, TNA-PFN, fixes a subset of neuron weights as anchors and trains multiple models in a shared permutation subspace, reducing linear mode connectivity barriers and improving fusion outcomes. The authors provide theoretical guarantees (Theorem III.4) and extensive empirical evidence across model soup, ColD fusion, and federated learning, along with practical federated variants FedPFN and FedPNU. The results show significant gains in interpolated accuracies and reduced barriers, highlighting a scalable approach to fuse pretrained models and improve global generalization in heterogeneous settings.

Abstract

Model fusion aims to integrate several deep neural network (DNN) models' knowledge into one by fusing parameters, and it has promising applications, such as improving the generalization of foundation models and parameter averaging in federated learning. However, models under different settings (data, hyperparameter, etc.) have diverse neuron permutations; in other words, from the perspective of loss landscape, they reside in different loss basins, thus hindering model fusion performances. To alleviate this issue, previous studies highlighted the role of permutation invariance and have developed methods to find correct network permutations for neuron alignment after training. Orthogonal to previous attempts, this paper studies training-time neuron alignment, improving model fusion without the need for post-matching. Training-time alignment is cheaper than post-alignment and is applicable in various model fusion scenarios. Starting from fundamental hypotheses and theorems, a simple yet lossless algorithm called TNA-PFN is introduced. TNA-PFN utilizes partially fixed neuron weights as anchors to reduce the potential of training-time permutations, and it is empirically validated in reducing the barriers of linear mode connectivity and multi-model fusion. It is also validated that TNA-PFN can improve the fusion of pretrained models under the setting of model soup (vision transformers) and ColD fusion (pretrained language models). Based on TNA-PFN, two federated learning methods, FedPFN and FedPNU, are proposed, showing the prospects of training-time neuron alignment. FedPFN and FedPNU reach state-of-the-art performances in federated learning under heterogeneous settings and can be compatible with the server-side algorithm.

Figures (19)

• Figure 1: Organization overview of the paper.
• Figure 2: LMC of random pruning at initialization and TNA-PFN under different mask ratios. For pruning, the mask ratio is the pruning ratio. "Avg. Acc." means averaged accuracies of individual models, and "Interp. Acc." means the accuracy of the interpolated model ($\alpha=0.5$) of two modes. The shadow areas mean the accuracy barriers in LMC, the smaller the better.
• Figure 3: A simple demonstration of reducing permutation symmetries via fixed neuron anchors. There are 6 networks (1 original + 5 counterparts) which are functionally identical but with different permutations. The number of permutation symmetries can be reduced by asymmetrically fixing some weights (in red). Though this demonstration presents a simple static network permutation, we will show in the main paper that this kind of method can realize better neuron alignment under training dynamics.
• Figure 4: Left two: Accuracy barriers of MLP under different hidden layers ($h$) and widths ($w$). Right two: Loss landscapes of MLP. For MLPs, if the barriers exist, TNA-PFN can reduce them. The shadow areas refer to the standard deviations.
• Figure 5: Accuracy barriers under different model architecture. WRN56 abbreviates for WideResNet56. CIFAR-10.

Theorems & Definitions (5)

• Definition 3.1
• Definition 3.2
• Theorem 3.4
• Remark 3.5
• Theorem B.1