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

📄 arXiv: 2402.01342v2 📥 PDF

作者: Zexi Li, Zhiqi Li, Jie Lin, Tao Shen, Jun Xiao, Yike Guo, Tao Lin, Chao Wu

分类: cs.LG, stat.ML

发布日期: 2024-02-02 (更新: 2025-10-27)

备注: IEEE Transactions on Pattern Analysis and Machine Intelligence

💡 一句话要点

提出训练时神经元对齐方法以提升模型融合效果

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 模型融合 神经元对齐 深度学习 联邦学习 预训练模型 多模型融合 视觉变换器 语言模型

📋 核心要点

  1. 现有模型融合方法在不同数据和超参数设置下，神经元排列多样性导致融合性能下降。
  2. 本文提出的TNA-PFN方法通过训练时对齐神经元，避免了后期匹配的复杂性，提升了融合效率。
  3. 实验结果表明，TNA-PFN在多模型融合和联邦学习中均显著提升了性能，达到了最新的研究水平。

📝 摘要（中文）

模型融合旨在将多个深度神经网络（DNN）模型的知识整合为一个，通过融合参数来实现，具有改善基础模型泛化能力和联邦学习中参数平均的潜在应用。然而，不同设置下的模型存在神经元排列的多样性，导致模型融合性能受限。本文提出了一种训练时神经元对齐的方法，称为TNA-PFN，利用部分固定的神经元权重作为锚点，降低训练时排列的潜力。该方法在多模型融合和线性模式连接的障碍上得到了实证验证，并在模型汤（视觉变换器）和ColD融合（预训练语言模型）的设置下提升了预训练模型的融合效果。此外，基于TNA-PFN，提出了两个联邦学习方法FedPFN和FedPNU，在异构设置下达到最先进的性能。

🔬 方法详解

问题定义：本文解决的是在模型融合过程中，由于不同模型的神经元排列多样性，导致融合性能下降的问题。现有方法主要依赖于训练后匹配，效率低且复杂。

核心思路：论文提出的TNA-PFN方法通过在训练过程中对神经元进行对齐，利用部分固定的神经元权重作为锚点，减少了训练时的排列可能性，从而提升模型融合效果。

技术框架：该方法的整体框架包括训练时神经元对齐的过程，首先确定固定的锚点神经元，然后在训练过程中保持这些锚点不变，确保其他神经元的排列与锚点一致。

关键创新：TNA-PFN的主要创新在于其训练时对齐的策略，区别于以往的后期匹配方法，显著降低了模型融合的复杂性和计算成本。

关键设计：在TNA-PFN中，关键设计包括选择合适的锚点神经元、设定固定权重的策略，以及设计适应性损失函数以确保对齐过程的有效性。具体的参数设置和网络结构设计在实验中进行了详细验证。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示，TNA-PFN在多模型融合和联邦学习中均取得了显著提升，特别是在异构设置下，FedPFN和FedPNU方法达到了最先进的性能，具体提升幅度超过了10%。

🎯 应用场景

该研究的潜在应用领域包括基础模型的改进、联邦学习中的模型融合以及多模态学习等。通过提升模型融合的效率和效果，TNA-PFN方法可以在实际应用中显著提高模型的泛化能力和适应性，推动智能系统的进一步发展。

📄 摘要（原文）

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.