Mixture-Trained Merging for Unified Multi-Objective Models
Unified language models must handle heterogeneous capabilities such as mathematics, coding, instruction following, and controllable reasoning, yet standard single-trajectory training is notoriously sensitive to ordering and data ratios. Naive weight-space merging of single-objective experts offers a modular alternative but often collapses distinct behaviors or degrades cross-domain performance. This paper proposes Mixture-Trained Merging (MTM), which trains each branch on an objective-biased data mixture rather than a pure objective so that the branches stay in mutually compatible regions of parameter space and can be merged at deployment time.
The paper presents a well-motivated and empirically supported method for building unified models via weight-space merging. Its central diagnosis---that hard experts drift into incompatible basins and their interpolations land outside shared low-loss regions---is compelling and supported by merge-simplex landscape visualizations. However, the evaluation is limited to small-scale supervised fine-tuning (10K examples per domain, 4B--7B parameters), and the authors explicitly acknowledge that SFT-stage merging is a restricted setting.
The failure-mode analysis is convincing. The authors clearly demonstrate mode collapse in think/non-think merging and cross-domain degradation in math/code merges, and they connect these empirically to accuracy barriers along linear interpolation paths. The iterative MOBO-driven mixture search is principled, and the ablation showing flatter merge-simplex interiors for soft branches directly supports the claim that mergeability is determined by the branches, not only by the merging rule. Cross-architecture recipe transfer and comparisons against multiple post-hoc correction methods further strengthen the case.
The experimental scale is modest. With only 10K examples per branch and models up to 7B parameters, it remains unclear whether the findings transfer to larger models, longer training horizons, or more than three objectives.
The evaluation of think/non-think mode preservation relies heavily on generation length as a proxy metric, which is crude and does not establish format adherence or reasoning correctness under mode-conditioned prompts. Additionally, while MTM-Multi amortizes cost by reusing branches, the absolute overhead of parallel branch training plus Bayesian optimization is not carefully benchmarked against end-to-end mixture search methods such as RegMix or DoReMi.
Comparisons to merging-only baselines and to MergeMix are fair and controlled: MTM consistently achieves the highest average scores while keeping individual objectives near their expert levels. The direct comparison to MergeMix is especially informative, because MTM surpasses it and thereby shows that deployment-time merging of branch-specific mixtures dominates training a single model on a globally selected mixture. What is missing is a rigorous single-trajectory multi-task baseline with a comparably tuned global mixture; without this, one cannot rule out that the gains stem from having more tunable degrees of freedom (per-branch mixtures plus merge weights) rather than from the merging framework itself.
The appendix provides substantial training detail, including hyperparameter tables, early-stopping criteria, BO configuration, and evaluation protocols using standard frameworks (DeepSpeed, vLLM, NeMo-Skills). The main barriers to independent reproduction are the lack of released code or checkpoints, the unreported stochastic details of the GP surrogate and Dirichlet sampling in BO, and the run-to-run variability introduced by plateau-based early stopping. Furthermore, each MTM-Multi round requires training $K$ parallel branches plus many merged-model evaluations, making exact replication costly even with the documented hyperparameters.
Unified language models are increasingly expected to combine heterogeneous capabilities, such as mathematics, code, instruction following, and controllable thinking behavior, within a single set of parameters. A common solution is sequential post-training on multiple objectives, but this entangles all objectives along one optimization trajectory and makes the final model highly sensitive to training order, data ratios, schedules, and stopping criteria. Weight-space merging offers a modular alternative, but naive merging of single-objective experts often fails: domain capabilities degrade sharply, or think/non-think modes collapse into one dominant behavior. We attribute both failures to incompatible weight-space geometry: experts trained on single objectives drift to distant regions of parameter space, placing their interpolations outside any shared low-loss basin. We propose Mixture-Trained Merging (MTM), which trains each branch on an objective-biased data mixture rather than a single objective, exposing it to cross-objective interactions and making branches compatible at merge time. MTM uses merged-model evaluations as a low-cost signal for selecting branch mixtures, avoiding expensive data-mixture ablations. The procedure is iterative: each round promotes the base model using globally selected merge coefficients and refines each branch mixture using domain-preferred coefficients under constraints that preserve other objectives. To scale beyond simplex grid search, MTM uses qNEHVI-based multi-objective Bayesian optimization. Across code, mathematics, instruction following, and think/non-think control, MTM outperforms naive merging and preserves behavioral separation where single-objective merging collapses, suggesting that effective unified models require branches trained to be mergeable.
Pick a starting point or write your own. Challenges run in the background, so you can keep reading while the AI investigates.
No challenges yet. Disagree with the review? Ask the AI to revisit a specific claim.