SMAT: Simple and Efficient Merge-Aware Training

SMAT. Simple and Efficient Merge-Aware Training. By Yanggan Gu and seven others, from the Hong Kong Polytechnic University, the Hong Kong University of Science and Technology in Guangzhou, and three other institutions. Published on the twenty-seventh of September. Its code repository is worked out: about fifty-five source files, with a readme and a licence.

Model merging combines the capabilities of several experts without retraining them jointly. The catch sits upstream. Each expert is normally trained on its own task loss and nothing else, so there is no guarantee it will still perform well once it has been merged with the others. Merge-aware training, or MAT, tries to fix that by targeting merged performance directly. Existing MAT methods, though, do not fully account for the merging operations people commonly use, and they add training cost.

The starting observation is a change of viewpoint. Stand where a single expert stands, and the common merging methods reduce to three operations. Scale reweights the expert's own update. Mask removes selected coordinates. Perturb adds in updates from the other experts. If merging looks like that from the inside, an expert can be trained to expect it. SMAT, short for Simple MAT, optimizes two things jointly: the expert's ordinary loss, and its expected loss at simulated merged parameters. Those simulated parameters come from sampling scaling coefficients, masks, and additive noise, one sample standing in for each of the three operations.

The remaining work is making this cheap. SMAT adds periodic scheduling, kernel fusion, and switching of parameter storage, and with those in place each training step costs one forward pass and one backward pass.

The evaluation spans four backbones, covering both language and vision-language models, and scores each trained expert under five different merging methods. Averaged over those five, SMAT beats the strongest baseline for each backbone by between one point zero seven and two point one six points. The training-time overhead over standard fine-tuning stays below two percent.

Put plainly, the paper replaces a vague goal, be good after merging, with a concrete one: be good under random scaling, random masking and random interference from neighbours. Because those three draws are cheap to simulate, merge-awareness comes at almost the price of ordinary fine-tuning, and it holds across several merging methods rather than being tuned to one.