PaperlayerView exact revision on arXiv

Source-linked AI summary

SMAT: Simple and Efficient Merge-Aware Training

Yanggan Gu, Yuanyi Wang, Zhen Li, Shuo Cai, Yuhang Liu, Junzhuo Li, Zihao Wang, Hongxia Yang

arXiv:2609.33437v1cs.LGcs.CL

TL;DR

Standard expert training optimizes task loss alone and does not guarantee good performance after merging, while existing merge-aware methods do not fully account for common merging operations. SMAT simulates Scale, Mask, and Perturb while jointly optimizing expert and expected merged loss. Across language and vision-language models, it improves merged performance across multiple merging methods with less than 2% training-time overhead.

  • Problem

    Standard fine-tuning optimizes each expert’s task loss alone, which does not guarantee good performance after merging, and existing MAT methods incompletely account for common merging operations while adding training cost.

    Source: §1 · p. 2

  • Method

    SMAT jointly optimizes expert loss and expected loss at simulated merged parameters by sampling scaling coefficients, masks, and additive noise.

    Source: §1 · p. 2

  • Results

    Across language and vision-language models, SMAT consistently improves merged performance across multiple merging methods with less than 2% training-time overhead.

    Source: §6 · p. 10

  • Takeaways & Limitations

    The results identify merge-induced parameter variation as a useful training signal, with Scale and Mask providing complementary gains beyond Perturb.

    Source: §6 · p. 10

    How closely the objective represents actual merging depends on the sampled states, with loss differences related to differences in state means and second moments.

    Source: §4.2 · p. 5

Abstract

from arXiv · show

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.

1 INTRODUCTION

SMAT trains experts for merging by simulating common changes to their updates while retaining expert-loss optimization. It reports improved merged performance across language and vision-language models with less than 2% training-time overhead over standard fine-tuning.

1 INTRODUCTION

  • Prior MAT methods do not fully account for common merging operations, especially scaling and masking an expert’s own update, and add training cost.SMAT is motivated by standard fine-tuning’s task-loss-only objective, which does not guarantee good post-merging performance.Source: §1 · p. 2
  • Common merging methods can be described from an expert’s perspective as Scale, which reweights its update; Mask, which removes coordinates; and Perturb, which adds other experts’ updates.The formulation is linked to Task Arithmetic [Ilharco et al., 2023], TIES-Merging [Yadav et al., 2023], DARE [Yu et al., 2024], and DELLA [Deep et al., 2024].Source: §1 · p. 2
  • SMAT jointly optimizes expert loss and expected loss at simulated merged parameters by sampling scaling coefficients, masks, and noise.Periodic scheduling, kernel fusion, and parameter storage switching reduce the cost of simulating merged states; training uses one forward and one backward pass per step.Source: §1 · p. 2
  • SMAT improves merged performance by 1.12 points over OrthoReg on Llama-3.2-1B-Instruct, with less than 2% training-time overhead over standard fine-tuning.Across four backbones, it improves merged performance and has shorter training time than compared MAT methods.Source: §1 · p. 2

2 RELATED WORK

The related work covers weight- and task-vector merging methods, then contrasts existing merge-aware training approaches with SMAT’s simulation of common merging operations.

2 RELATED WORK

  • Other merging approaches average weights, reweight parameters, preserve layer outputs with regression, or use subspaces defined by task-update singular vectors.The passage contrasts these methods with SMAT’s simulation of scaling, masking, and additive changes during expert training.Source: §2 · pp. 2–3
  • Merging methods alter task vectors through weighting, learned coefficients, magnitude or sign selection, random deletion, masks, or rescalers.Examples include Task Arithmetic [Ilharco et al., 2023], AdaMerging [Yang et al., 2024], TIES-Merging [Yadav et al., 2023], DARE [Yu et al., 2024], and DELLA [Deep et al., 2024].Source: §2 · p. 3
  • Existing MAT methods use linearized fine-tuning, sharpness-aware training, sampled merge offsets, or update-matrix orthogonality, but do not fully consider scaling and masking.They incur extra computation through gradient evaluations, perturbation processing, or matrix regularization; SMAT simulates Scale, Mask, and Perturb and uses periodic scheduling and fused operations to reduce overhead.Source: §2 · p. 3

3 PRELIMINARIES

The preliminaries establish independent experts with a shared initialization, describe how task vectors are merged, and define merge-aware training as balancing expert and merged performance.

3.1 PROBLEM SETUP

  • The setup considers N tasks with a shared pretrained initialization, where each expert trains on its own data without access to other experts’ data or checkpoints.Source: §3.1 · p. 3

3.2 TASK VECTORS AND MODEL MERGING

  • An expert’s task vector is its fine-tuned parameters minus the shared initialization; merging weights and masks task vectors before adding them to that initialization.The formulation includes parameter-varying weights and methods such as Task Arithmetic [Ilharco et al., 2023], TIES [Yadav et al., 2023], DARE [Yu et al., 2024], and DELLA [Deep et al., 2024].Source: §3.2 · p. 3

3.3 THE MAT OBJECTIVE

  • The MAT objective balances preserving each expert’s performance with improving the merged model, with λ = 0 and λ = 1 representing expert-only and merged-only training, respectively.SMAT simulates Scale, Mask, and Perturb when optimizing for merged performance.Source: §3.3 · pp. 3–4, Figure 2 · p. 4

4 METHOD

SMAT simulates common merging operations during expert training by sampling scale, mask, and perturbation states, then optimizes loss at those states alongside expert loss. Periodic scheduling and implementation changes reduce the added training cost.

4.1 SIMULATING MERGED PARAMETER STATES

  • SMAT samples scaling coefficients, masks, and additive noise to simulate merged parameters, covering coordinate removal and reweighting as well as other experts’ updates.The scale is shared across coordinates; masks are sampled independently across coordinates, and the perturbation represents other experts’ combined updates.Source: §4.1 · p. 4
  • Uniform, Gaussian, and Laplace perturbations show broadly similar trends at matched RMS, with uniform giving the highest merged performance at RMS = 2 × 10−3.The expected perturbed loss averages task loss over nearby parameters, encouraging tolerance to additive merge offsets; the method samples zero-mean uniform noise with fixed RMS.Source: §4.1 · p. 5, Figure 3 · p. 5

4.2 A PRACTICAL TRAINING OBJECTIVE

  • SMAT’s practical objective combines expert loss with expected loss over simulated parameters, using one sampled state per simulated-loss update as a stochastic estimate.How closely the objective represents actual merging depends on the sampled states; the paper relates loss differences to differences in state means and second moments.Source: §4.2 · p. 5
  • The simulated-loss gradient is weighted by the coordinatewise participation coefficient α_i m_i, so excluded task-vector coordinates receive zero gradient from that loss.The other gradient factor is the simulated merged model’s gradient on the expert’s task.Source: §4.2 · p. 5

4.3 SCHEDULING AND EFFICIENT IMPLEMENTATION

  • Mixed-loss updates can be approximated to first order by separate single-loss steps, motivating periodic scheduling.The proposition assumes constant step size and loss gradients that are Lipschitz and bounded in a neighborhood containing both trajectories.Source: §4.3 · pp. 5–6
  • With t = 4, periodic scheduling uses three expert-loss updates and one simulated-loss update, with one forward and one backward pass per step.Each cycle matches λ = 1/t in the objective.Source: §4.3 · p. 6, Figure 4 · p. 6
  • Kernel fusion combines weight simulation and gradient rescaling, while storage switching reuses a simulated-weight buffer and restores original storage for optimizer updates.The fused kernels reduce memory traffic and kernel launches; storage switching requires one additional GPU buffer the size of the trainable weights.Source: §4.3 · p. 6

5 EXPERIMENTS

The experiments compare SMAT with fine-tuning and merge-aware baselines across language and vision-language models, then examine efficiency and merging behavior. SMAT improves merged scores across settings, with analyses testing its components, robustness, and loss-smoothing interpretation.

5.1 EXPERIMENTAL SETUP

  • Across four backbones, methods share data, initialization, learning rate, batch size, and training budget; SMAT simulates every four steps and results average three random seeds.Evaluation uses TRACE language tasks and CLIP vision tasks, with TA, TIES, DARE, and DELLA configurations selected on held-out development data.Source: §5.1 · pp. 6–7

5.2 MAIN RESULTS

  • SMAT exceeds the strongest baseline in Avg. across five mergers by 1.12 and 1.07 points on Llama-1B and Llama-8B, and by 1.89 and 2.16 points on ViT-B/32 and ViT-L/14.The CLIP baselines are MergOPT (72.57) for ViT-B/32 and OrthoReg (85.62) for ViT-L/14.Source: §5.2 · p. 7, Table 1 · p. 7, Table 2 · p. 8
  • With frozen-base offload, SMAT adds 0.2–1.9% training-time overhead across four backbones, while peak GPU memory increases by 1.7–24.3%.Timing measures expert training loops.Source: §5.2 · p. 8

5.3 ANALYSIS

  • With Muon on Llama-1B, SMAT’s average across five mergers is 44.15, exceeding OrthoReg by 0.74 points and FT by 5.90, with training time at 1.00× FT.SMAT leads four of five mergers, while independent-expert performance remains close to FT (55.07 vs. 55.31).Source: §5.3 · p. 8, Table 3 · p. 8
  • As expert count rises, SMAT’s advantage over three baselines becomes clearer; replacing FT experts raises normalized score from 77.55% to 86.08% for all-SMAT.SMAT leads the baselines at K = 5, 6, 7, and its replacement curve stays above the other methods for q > 0.Source: §5.3 · pp. 8–9, Figure 5 · p. 9
  • SMAT achieves the best mean merged performance by balancing expert quality and retention after merging, rather than by minimizing merging drop alone.Compared with OrthoReg, SMAT has a higher expert score but a larger drop; compared with FT and MergOPT, it has both a higher expert score and a smaller drop.Source: §5.3 · p. 9
  • Across six tasks, SMAT has broader low-loss regions than FT and MergOPT along tested merge directions, consistent with Perturb smoothing conditional training loss.The analysis probes rescaling an expert’s update and adding updates from other tasks; loss slices use a shared NLL scale within each task, with scales differing across tasks.Source: §5.3 · pp. 9–10, Figure 6 · p. 9
  • Removing Scale, Mask, or Perturb lowers Avg. across five mergers by 1.25, 1.07, or 2.10 points, respectively, with Perturb having the largest effect.The ablation keeps training budget and merger settings fixed.Source: §5.3 · p. 10

6 CONCLUSION

The conclusion presents SMAT as an efficient way to prepare independently trained experts for merging by simulating common merge operations. Its analyses support loss smoothing as a contributing mechanism and identify complementary gains from Scale and Mask.

6 CONCLUSION

  • SMAT prepares independently trained experts for merging by simulating Scale, Mask, and Perturb with efficient periodic training.The method jointly accounts for merge-induced parameter variation as a training signal.Source: §6 · p. 10
  • Across language and vision-language models, SMAT improves merged performance across multiple merging methods with less than 2% training-time overhead.The conclusion attributes part of its merge robustness to broader low-loss regions along merge-relevant directions.Source: §6 · p. 10
  • Scale and Mask provide complementary gains beyond Perturb, while loss smoothing is supported as a contributing mechanism for merge robustness.Source: §6 · p. 10

AI USE STATEMENT

The authors report using generative AI for language editing and drafting code for functional testing, with human review of the assisted text and code.

AI USE STATEMENT

  • The authors used generative AI for language editing and functional-testing code drafts, then reviewed the text and manually inspected and tested the code.The authors take responsibility for the final content and code correctness.Source: AI USE STATEMENT · p. 10

REPRODUCIBILITY STATEMENT

The paper identifies where its code, experimental settings, implementation procedures, proofs, and assumptions are documented.

REPRODUCIBILITY STATEMENT

  • Code is available on GitHub; experimental settings, implementation and measurement procedures, proofs, and assumptions are described in specified sections and appendices.The cited locations are Section 4, Section 5.1, Appendices A, E, C, and D.Source: REPRODUCIBILITY STATEMENT · p. 10

A ADDITIONAL EXPERIMENTAL SETTINGS

The appendix records training, merger-calibration, decoding, and split settings for language and vision experiments. It also describes the expert-count evaluation protocol.

A ADDITIONAL EXPERIMENTAL SETTINGS

  • Language training uses AdamW with batch size 8, learning rate 2 × 10^-5, weight decay 10^-3, gradient clipping at 1, and a maximum length of 512.Epoch counts for C-STANCE, FOMC, MeetingBank, ScienceQA, NumGLUE-cm, NumGLUE-ds, and 20Minuten are 5, 3, 7, 3, 5, 5, and 7.Source: A ADDITIONAL EXPERIMENTAL SETTINGS · p. 14
  • Vision training uses Adam with learning rate 10^-5, cosine decay without warmup, zero weight decay, batch size 128, and 4,000 updates per task.The text encoder, projection, and classification heads remain frozen.Source: A ADDITIONAL EXPERIMENTAL SETTINGS · p. 14
  • Splits are separated by normalized prompt or RGB-image identity, and evaluation uses released language eval splits or vision test splits.Language dev sizes range from 50 to 200 across the listed tasks; vision uses 200 examples per task, and NumGLUE-cm has 41 evaluation examples.Source: A ADDITIONAL EXPERIMENTAL SETTINGS · p. 14
  • Sparse-merger calibration selects TIES retention from {0.1, 0.2, 0.3, 0.5}, DARE/DELLA drop probability from {0.1, 0.3, 0.5, 0.7, 0.9}, and fixes DELLA’s window at 0.14.The final settings are TIES retention 0.5 and DARE drop probability 0.3 for both families; DELLA uses 0.5 for language and 0.1 for vision.Source: A ADDITIONAL EXPERIMENTAL SETTINGS · p. 14
  • Selected sparse settings transfer across methods and model sizes within each family, while each expert set selects its own merging coefficient.Reference coefficients are 0.3 for language TIES/DARE, 0.2 for vision TIES/DARE, and 1.0 for DELLA; ties favor less pruning.Source: A ADDITIONAL EXPERIMENTAL SETTINGS · p. 14
  • TRACE dev selection gives each candidate one decoding run, then gives the top two two additional runs on three tasks and selects by the three-run mean.The selection mean also applies to sparse calibration; table differences are computed before rounding.Source: A ADDITIONAL EXPERIMENTAL SETTINGS · p. 15

B.1 EXPERT-COUNT EVALUATION

  • For each expert count K = 2–6, seven cyclic subsets balance task inclusion; K = 7 uses all tasks, and scores are normalized by independent fine-tuning scores.Methods share subsets and three decoding repeats, and each merged model is evaluated only on its included tasks.Source: §B.1 · p. 15

B.2 MUON SETTINGS AND COMPLETE RESULTS

The section gives Muon experiment settings and theoretical analyses of SMAT’s simulated states and loss. These analyses relate expected-loss behavior to state moments, curvature, and noise assumptions.

B.2 MUON SETTINGS AND COMPLETE RESULTS

  • Muon experiments use shared matrix and auxiliary-parameter learning rates, seven-task splits and budgets, batch size 8, and maximum length 512; Table 6 reports Llama-3.2-1B-Instruct results.SMAT uses scalar Scale with αmin = 0.2, Mask probability p = 0.5, Perturb RMS 0.002, and interval t = 4; embeddings are excluded from masking.Source: §B.2 · p. 15, Table 6 · p. 15

C THEORY OF SIMULATED MERGING

  • The simulated-state covariance reflects shared scalar scaling, independent coordinate masking, and centered independent Perturb noise.The stated moments also include the simulated-state mean; no Muon result values are supplied in the table caption.Source: C THEORY OF SIMULATED MERGING · p. 15
  • With a twice-differentiable loss and Lipschitz Hessian, the centered displacement removes the linear term, while the trace term penalizes curvature when the Hessian is positive semidefinite.The expansion includes a third-order Taylor remainder.Source: C THEORY OF SIMULATED MERGING · p. 16
  • For small σ, independent uniform, Gaussian, and Laplace noise with equal coordinate variance share the leading loss-expansion term, but the expansion does not rank them at finite scales.The comparison assumes the stated smoothness conditions.Source: C THEORY OF SIMULATED MERGING · p. 16
  • The expected-loss difference from a target merge distribution depends on displacement moments and a third-order remainder, without requiring matching moments.The result assumes finite third absolute moments and a Lipschitz Hessian on the relevant convex hull; expected-loss comparison does not guarantee higher task scores.Source: C THEORY OF SIMULATED MERGING · p. 16

D TRAINING AND SCHEDULING DETAILS

The section analyzes periodic scheduling, implementation and measurement details, and future extensions. Its scheduling guarantees are local or expectation-based, while the stated scope assumes shared expert architectures and initialization.

D TRAINING AND SCHEDULING DETAILS

  • The scheduling comparison is local to one cycle: for fixed cycle length, endpoint difference is second order in the learning rate under ordinary gradient descent, not a guarantee of matching full trajectories.The analysis does not cover AdamW’s adaptive moments or Muon’s momentum and matrix transformation.Source: D TRAINING AND SCHEDULING DETAILS · pp. 16–17
  • One simulated-loss update every t steps matches mixed-objective weight λ = 1/t in cycle-averaged population gradients, with a drift bound under bounded steps and Lipschitz fields.Changing parameters and sampled gradients add variability; the result does not imply identical updates for stochastic or stateful optimizers.Source: D TRAINING AND SCHEDULING DETAILS · p. 17

E.1 HARDWARE, SOFTWARE, AND IMPLEMENTATION

  • The fused implementation uses simulated weights for forward and backward passes, then restores original storage for optimizer updates; frozen base weights remain in pinned CPU memory.Training runs use H800 GPUs, with Llama-8B using two GPUs per expert and other backbones using one.Source: D TRAINING AND SCHEDULING DETAILS · p. 17, §E.1 · p. 18

E.2 TRAINING-TIME MEASUREMENT

  • Table 7 timing totals sum training-loop wall-clock seconds across a complete expert set, rather than reporting parallel elapsed time.Timing includes loading, warmup, and logging but excludes initialization, checkpoint saving, hyperparameter search, merging, and evaluation.Source: §E.2 · p. 18, Table 7 · p. 18

E.3 MEMORY MEASUREMENT

  • Memory is the largest per-GPU PyTorch allocated peak across experts, and the reported peaks cover training and final checkpoint serialization.For Muon, the maximum is measured per device across the seven experts.Source: §E.3 · p. 18

F FUTURE WORK

  • Future work identifies shared expert architectures and initialization as current assumptions, and proposes studying scaling, budgets, heterogeneous fusion, compression, calibration, and reliability.The proposed reliability evaluations include per-task regressions, uncertainty, and constraint violations alongside mean accuracy.Source: F FUTURE WORK · p. 19
Loading 2609.33437v1…