Weight Scope Alignment: A Frustratingly Easy Method for Model Merging

TL;DR

This paper introduces Weight Scope Alignment (WSA), a simple regularization method that improves model merging by aligning weight scopes, addressing challenges caused by training randomness and data heterogeneity.

Contribution

The paper proposes a novel, easy-to-implement regularization technique called WSA that aligns weight scopes to enhance model merging across various training scenarios.

Findings

WSA improves model merging effectiveness in experiments.

WSA is effective in Mode Connectivity and Federated Learning.

Weight scope variations significantly impact model fusion success.

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.