RegMean++: Enhancing Effectiveness and Generalization of Regression Mean for Model Merging

2025-08-05

AI Summary

RegMean independently optimizes individual linear layers during model merging, neglecting inter-layer feature propagation and dependencies, thereby inadequately modeling merged-model behavior. To address this, we propose RegMean++, the first method within the regression-mean framework to explicitly model intra-layer consistency and cross-layer dependencies. RegMean++ integrates feature propagation analysis with linear regression and employs a closed-form solution for efficient, interpretable weight optimization. Experimental results demonstrate that RegMean++ significantly improves prediction consistency across diverse evaluation settings: it outperforms RegMean in both in-distribution and out-of-distribution generalization, scales effectively to large-scale tasks, and exhibits robustness under distributional shift. Notably, RegMean++ achieves state-of-the-art performance while retaining computational efficiency and analytical transparency.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Natural Language Processing: Interpretability, Analysis, and Evaluation of NLP Models
• Reasoning under Uncertainty: Relational Probabilistic Models

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

Regression Mean (RegMean), an approach that formulates model merging as a linear regression problem, aims to find the optimal weights for each linear layer in the merge model by minimizing the discrepancy in predictions between the merge and candidate models. RegMean provides a precise closed-form solution for the merging problem; therefore, it offers explainability and computational efficiency. However, RegMean merges each linear layer independently, overlooking how the features and information in the earlier layers propagate through the layers and influence the final prediction in the merge model. In this paper, we introduce RegMean++, a simple yet effective alternative to RegMean, that explicitly incorporates both intra- and cross-layer dependencies between merge models' layers into RegMean's objective. By accounting for these dependencies, RegMean++ better captures the behaviors of the merge model. Extensive experiments demonstrate that RegMean++ consistently outperforms RegMean across diverse settings, including in-domain (ID) and out-of-domain (OOD) generalization, sequential merging, large-scale tasks, and robustness under several types of distribution shifts. Furthermore, RegMean++ achieves competitive or state-of-the-art performance compared to various recent advanced model merging methods. Our code is available at https://github.com/nthehai01/RegMean-plusplus.

Problem

• Improving model merging by addressing independent layer merging limitations
• Incorporating intra- and cross-layer dependencies for better merge model behavior
• Enhancing generalization and robustness across diverse task settings

Innovation

• Linear regression for optimal model merging
• Incorporates intra- and cross-layer dependencies
• Enhances generalization and robustness performance