Weight Scope Alignment: A Frustratingly Easy Method for Model Merging

TL;DR

This work identifies weight-scope mismatches as a key hidden factor limiting model merging, especially under non-I.I.D. training and diverse optimization settings. It proposes Weight Scope Alignment (WSA), comprising Weight Scope Regularization via KL divergence to a target distribution and Weight Scope Fusion to merge multiple models’ scopes into a unified one, with formal Gaussian weight-model assumptions. The framework is extended to two pivotal fusion contexts: mode connectivity and federated learning, where WSA demonstrably reduces interpolation barriers and improves aggregation performance, even when compared with strong baselines and permutation-based methods. Empirically, weight scopes tend to follow Gaussian distributions regardless of initialization, and dynamic, adaptive scope fusion outperforms fixed distributions, highlighting WSA’s practical value for robust, scalable model merging in real-world settings.

Abstract

Merging models becomes a fundamental procedure in some applications that consider model efficiency and robustness. The training randomness or Non-I.I.D. data poses a huge challenge for averaging-based model fusion. Previous research efforts focus on element-wise regularization or neural permutations to enhance model averaging while overlooking weight scope variations among models, which can significantly affect merging effectiveness. In this paper, we reveal variations in weight scope under different training conditions, shedding light on its influence on model merging. Fortunately, the parameters in each layer basically follow the Gaussian distribution, which inspires a novel and simple regularization approach named Weight Scope Alignment (WSA). It contains two key components: 1) leveraging a target weight scope to guide the model training process for ensuring weight scope matching in the subsequent model merging. 2) fusing the weight scope of two or more models into a unified one for multi-stage model fusion. We extend the WSA regularization to two different scenarios, including Mode Connectivity and Federated Learning. Abundant experimental studies validate the effectiveness of our approach.

Figures (17)

• Figure 1: The Gaussian distribution of parameters in VGGBN16 trained on CIFAR-100. Both uniform and normal initialization lead to Gaussian parameter distributions.
• Figure 2: The linear interpolation curves between a model with another model (Scale=1.0) and its scaled versions (Scale $\neq$ 1.0).
• Figure 3: Model interpolation on SVHN and CIFAR-10. WSA enhances OTFusion singh2020ModelFusion performance.
• Figure 4: WSA facilitates Git Re-Basin ainsworth2023GitRebasin under different epochs, widths and depths, ensuring a smaller loss barrier.
• Figure 5: Loss landscape w/o WSA and w/ WSA.