Skip to main content
Back to timeline
arXivSource publication:

SMAT decomposes merging into Scale, Mask and Perturb simulations, lifting the five-merger average by 1.07–2.16 points across four backbones with under 2% training overhead

Synopsis

The work introduces SMAT (Simple MAT), which from a single expert's perspective describes common model-merging methods through three operations—Scale (reweighting its own update), Mask (removing selected coordinates) and Perturb (adding other experts' updates)—and during training samples scaling coefficients, masks and additive noise to simulate merged parameters while jointly optimizing the expert loss and the expected loss at simulated merged parameters, made efficient by periodic scheduling, kernel fusion and parameter storage switching with one forward and one backward pass per step; across four backbones (Llama-3.2-1B-Instruct, Llama-3.1-8B-Instruct, CLIP ViT-B/32 and ViT-L/14), SMAT improves the mean score over five merging methods by 1.07–2.

AI-generated editorial illustration: SMAT: Simple and Efficient Merge-Aware Training

Interpretation

The paper describes common merging methods such as Task Arithmetic, TIES-Merging, DARE and DELLA from a single expert's view as three operations: Scale reweights that expert's own update, Mask removes selected coordinates, and Perturb adds updates from other experts; SMAT therefore samples scaling coefficients, masks and additive noise during independent training to build simulated merged parameters and jointly optimize the expert loss and the simulated merged loss. Earlier merge-aware training methods (SAFT, MergOPT, OrthoReg) do not fully account for common merging operations, especially scaling and masking an expert's own update; SMAT brings all three operations explicitly into the training objective. The paper gives a unified form in Eq. (1) and states it covers Model Soups, Fisher Merging, Task Arithmetic and the masking of TIES, DARE and DELLA; ablations show removing Scale, Mask or Perturb lowers the five-merger average by 1.25, 1.07 and 2.10 points respectively, so all three contribute, with Perturb having the largest effect.

SMAT consistently improves merged performance on four backbones: 45.77 average on Llama-3.2-1B-Instruct and 58.49 on Llama-3.1-8B-Instruct, exceeding the strongest baseline OrthoReg by 1.12 and 1.07 points; 74.46 on CLIP ViT-B/32 (1.89 points above MergOPT) and 87.78 on ViT-L/14 (2.16 points above OrthoReg). Relative to each backbone's strongest baseline, SMAT gains 1.07–2.16 points on the average over five mergers (WA, TA, TIES, DARE, DELLA), leading three of five merging columns on 1B and all five columns on 8B and both CLIP backbones. Experiments span language and vision-language models from 151M to 8B parameters, with each model trained and evaluated three times using different random seeds and averaged; language tasks come from the seven TRACE tasks and vision from Cars, DTD, EuroSAT, GTSRB, MNIST, RESISC45, SUN397 and SVHN, with all methods sharing data, initialization, learning rate, batch size and training budget.

Through periodic scheduling, kernel fusion and parameter storage switching, SMAT needs only one forward and one backward pass per step; across the four backbones training-time overhead is about 0.2–1.9% over standard fine-tuning, with peak GPU memory increasing by 1.7–24.3%. Proposition 1 shows mixed-loss updates can be replaced to first order by single-loss steps with counts proportional to their weights, motivating a schedule of one simulated-loss update every four steps; compared with the extra gradient evaluations, perturbation processing or matrix regularization of ASAM, MergOPT and OrthoReg, SMAT keeps overhead close to standard fine-tuning. Tables 1 and 2 report time and memory changes relative to FT; Appendix E gives absolute times, e.g. mean time per step on Llama-1B is 0.0687 s for FT versus 0.0700 s for SMAT and 0.3543 s versus 0.3594 s on Llama-8B; on Llama-1B SMAT improves over OrthoReg by 1.12 points while using 82% less training time.

Mechanistic analysis shows SMAT broadens low-loss regions along merge-relevant directions: along rescaling an expert's own update and adding updates from other tasks, SMAT's low-loss regions are broader than FT and MergOPT across six tasks, consistent with Perturb smoothing the conditional training loss; SMAT also reaches 44.15 average under the Muon optimizer, 0.74 points above OrthoReg and 5.90 points above FT, at training time comparable to FT. The paper attributes merge robustness to loss smoothing and provides a theoretical account via moments of the simulated state and a local loss expansion; it also shows the gains extend beyond AdamW and that partially replacing FT experts can improve an existing expert set (normalized score rising from 77.55% for all-FT to 86.08% for all-SMAT). Loss curves use the Table 1 FT, MergOPT and SMAT expert sets with identical FT partner updates and evaluation samples; Muon results are reported on Llama-1B where SMAT's expert score of 55.07 is close to FT's 55.31; the expert-count study shows SMAT's normalized score falling from 93.20% with two experts to 86.08% with all seven, with its advantage clearer as expert count grows.

Perspective

The result targets sets of independently trained experts that share a pretrained initialization and shared trainable parameters, in multi-task merging settings on language and vision-language backbones; for practitioners who want to improve post-merge average performance without joint retraining and who are sensitive to training time and memory increments, SMAT offers a training approach that plugs into existing fine-tuning pipelines and works with mergers such as WA, TA, TIES, DARE and DELLA. The paper also shows an existing FT expert set can be upgraded by partial replacement with SMAT experts, and that SMAT can be used with the Muon optimizer.

The paper assumes shared expert architectures and initialization, so behavior under heterogeneous models, different initializations and joint compute, memory and checkpoint-access budgets remains to be evaluated; the gap between the simulated merged loss and the true merge distribution is characterized by differences in state means and second moments, and the paper explicitly notes that this expected-loss comparison does not guarantee higher task scores; Perturb approximates other experts' updates with random noise, and uniform, Gaussian and Laplace noise behave similarly at matched RMS, leaving the role of precise distribution matching unclear; on the 8B backbone SMAT's independent expert score is below FT, indicating that better merged performance need not coincide with better single-expert performance, a trade-off still worth watching in deployment decisions.

Sources