Improving Model Fusion by Training-Time Neuron Alignment With Fixed Neuron Anchors.

Area of Science:

• Artificial Intelligence
• Machine Learning
• Deep Learning

Background:

• Model fusion integrates knowledge from multiple deep neural network (DNN) models for improved generalization and parameter averaging.
• Diverse neuron permutations across models trained under different settings hinder effective model fusion by causing models to reside in different loss basins.
• Existing methods for neuron alignment focus on post-training permutation matching, which can be computationally expensive and less versatile.

Purpose of the Study:

• To develop a novel approach for training-time neuron alignment to improve model fusion performance.
• To introduce a method that enhances model fusion without requiring post-training permutation matching.
• To demonstrate the applicability and effectiveness of training-time alignment in various model fusion scenarios, including federated learning.

Main Methods:

• Introduced TNA-PFN, a simple, lossless algorithm for training-time neuron alignment using partially fixed neuron weights as anchors.
• Reduced the potential for training-time permutations by utilizing fixed neuron weights.
• Developed FedPFN and FedPNU, federated learning methods based on TNA-PFN for heterogeneous settings.

Main Results:

• Empirically validated TNA-PFN's effectiveness in reducing barriers to linear mode connectivity and multi-model fusion.
• Demonstrated improved fusion of pre-trained models in settings like model soup (vision transformers) and ColD fusion (language models).
• Achieved state-of-the-art performance with FedPFN and FedPNU in federated learning under heterogeneous conditions, compatible with server-side algorithms.

Conclusions:

• Training-time neuron alignment offers a more efficient and versatile alternative to post-training alignment for model fusion.
• TNA-PFN provides a foundational method for achieving training-time alignment, enhancing both direct model fusion and federated learning.
• The proposed FedPFN and FedPNU methods represent significant advancements in federated learning, particularly in handling data heterogeneity.