Model training method improves combined expert model performance efficiently
SMAT: Simple and Efficient Merge-Aware Training
Machine LearningComputation and Language
Summary
Combining different expert models can help machines do many tasks better, but training each model separately doesn’t guarantee that mixing them works well. The authors see common combining steps as scaling, masking, and adding changes, then train the models by simulating these steps during learning. Their new method, SMAT, trains experts to work better when merged, with little extra time needed. They tested SMAT on language and vision models and found it improves results across several merging techniques.
What this means in practice
- •For machine learning engineers: Train expert models that merge more effectively without expensive joint retraining steps.
- •For ai product developers: Improve combined AI model performance in language and vision tasks by reducing merging accuracy loss.$Commercial implications: Enables better merged AI components for commercial applications in natural language and vision AI products with minimal extra training cost.
Authors
Yanggan Gu, Yuanyi Wang, Zhen Li, Shuo Cai, Yuhang Liu, Junzhuo Li, Zihao Wang, Hongxia Yang
Abstract
Model merging integrates the capabilities of multiple experts without joint retraining, but standard expert training optimizes task loss alone and does not guarantee good performance after merging. Merge-aware training (MAT) aims to improve merged performance, but existing methods do not fully account for common merging operations and add training cost. We observe that, from an expert's perspective, common merging methods can be described by three operations: Scale reweights its own update, Mask removes selected coordinates, and Perturb adds updates from other experts. Based on this view, we introduce SMAT (Simple MAT), which jointly optimizes expert loss and expected loss at simulated merged parameters generated by sampling scaling coefficients, masks, and additive noise. We further introduce periodic scheduling, kernel fusion, and parameter storage switching to make SMAT efficient, with one forward and one backward pass per step. Across four language and vision-language backbones, SMAT improves the mean score across five merging methods by 1.07-2.16 points over the strongest baseline for each backbone, with less than 2% training-time overhead over standard fine-tuning.