Source-linked AI summary
Mixture-Trained Merging for Unified Multi-Objective Models
SeongHyeon Kim, Chaeyun Jang, Seungyoo Lee, Jiyeon Ham, Yunju Bak, Boseop Kim, Juho Lee
TL;DR
Unified models need to combine capabilities across objectives, but global post-training recipes are sensitive to design choices and hard-expert merges can lose domain performance or mode control. MTM trains objective-biased branches and iteratively selects mixtures and merges them. Across the evaluated tasks and models, it reports stronger average performance while keeping objectives near expert performance and preserving think/non-think separation.
Problem
Single-trajectory training is sensitive to training order, data ratios, schedules, and stopping criteria, while merging single-objective experts can degrade domains or collapse response modes.
Source: §1 · p. 1, §3.1 · p. 4Method
MTM iteratively trains objective-biased branches and merges them, using merged-model evaluations to guide branch-specific mixture selection.
Source: §4.1 · p. 5, §4.2 · p. 5Results
Takeaways & Limitations
The results support training branches to be mergeable so that a unified model can retain performance across objectives and preserve distinct response modes.
Source: §1 · p. 2, §7 · p. 9MTM-Single requires a second full training pass to produce the deployed model, at roughly twice a single-trajectory baseline's budget.
Source: §4.2 · p. 5
Abstract
from arXiv · showhide
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.
1 Introduction
Unified models combine heterogeneous capabilities, but single-trajectory training is sensitive to recipe choices and merging single-objective experts can degrade domains or collapse response modes. MTM trains objective-biased branches and iteratively selects mixtures and merges to keep branches specialized yet compatible.
1 Introduction
- Unified models combine reasoning, direct responses, coding, and instruction following in one model, potentially replacing separate model stacks with a simpler deployment interface.Source: §1 · p. 1
- Single-trajectory post-training is sensitive to training order, data ratios, and schedules, making robust unified-model recipes costly to tune.The text also identifies stopping criteria and near-full-budget runs as sources of difficulty.Source: §1 · pp. 1–2
- Post-hoc merging of single-objective experts can collapse think/non-think behavior toward the dominant branch or degrade code and math performance.Source: §1 · p. 2
- MTM trains branches on controlled mixtures of objectives, aiming to reduce weight-space divergence while preserving specialization and making branches mutually mergeable.The method integrates branch-mixture selection, specialization, and merging in an iterative loop.Source: §1 · p. 2
- MTM uses merged-model performance as a proxy for mixture training and optimizes branch-specific mixtures with multi-objective Bayesian optimization.Each round selects mixtures, continues branch training, and merges branches into a common reference model.Source: §1 · p. 2
- The paper reports consistent gains across math, code, and instruction-following benchmarks on Qwen3-4B and OLMo-7B.Source: §1 · p. 2
2 Related Works
Prior work studies global data-mixture selection and parameter-space merging as separate routes to unified models. MTM differs by training mixture-biased branches to shape compatibility before merging.
2.1 Building Unified Models via Data Mixtures
- Earlier mixture-selection approaches use proxy-model reweighting, small-run regression, scaling laws, or weight-space approximations, but finding one global mixture remains difficult as objectives grow.Source: §2.1 · p. 2
- The figures contrast single-domain SFT with round-wise mixture-trained branches and promoted merges, and show hard versus soft Math/Code and think/non-think branch setups.Source: Figure 1 · p. 3, Figure 2 · p. 3, Figure 3 · p. 3
2.2 Weight-Space Model Merging
- Model merging combines objective-specific models without increasing inference cost; prior approaches include checkpoint averaging, task-vector arithmetic, and branching from a shared point.Source: §2.2 · p. 3
- MTM trains each branch on a domain-biased mixture to make task vectors compatible before merging, rather than treating independently trained experts as fixed.Source: §2.2 · p. 3
- Post-hoc interference methods adjust updates or fusion coefficients, while application-specific methods constrain subspaces for reasoning and instruction models.Source: §2.2 · p. 3
3 Background and Motivation
The section frames unified-model construction as combining performance across multiple objectives, either through one global training mixture or by merging separately trained branches. It identifies domain trade-offs and mode collapse as failures of hard-expert merging.
3 Background and Motivation
- The problem setting seeks one parameter vector that performs well across K objectives, each with its own training corpus and normalized evaluation metric.Source: §3 · p. 3
- Single-trajectory training must reconcile different task saturation steps and metric values through a shared mixture ratio, stage order, learning-rate schedule, and stopping criterion.Source: §3 · p. 4
- A branch-based alternative fine-tunes one hard expert per objective from the base model, then combines the experts with simplex merge weights.This makes construction a set of independent fine-tunes followed by a mixing step.Source: §3 · p. 4
3.1 Failure Modes of Hard-Expert Merging
- Across the examined settings, hard-expert merging produces sharp cross-domain trade-offs or a single-mode hybrid, while objective-mixture branches alleviate both failures.Source: §3.1 · p. 4
- Hard-expert merging trades off code and math: as the code-expert ratio rises, MBPP+ accuracy improves while MATH-500 accuracy falls from about 85 to below 50.No interior merge ratio reaches expert-level performance on both tasks; mixture-trained branches remain strong across a wider range of ratios.Source: §3.1 · p. 4
- Hard-expert merges also collapse think/non-think behavior: generation lengths move together across interior merge ratios, while mixture-trained branches retain short no-think and long think outputs.Qwen3-4B-Base also shows a length blow-up at one merge ratio.Source: §3.1 · p. 4
4 MTM: Mixture-Trained Merging
MTM trains objective-biased branches and merges them, with an iterative variant that reuses each round’s merged model to guide subsequent training.
4 MTM: Mixture-Trained Merging
- MTM offers MTM-Single, which trains mixture-biased branches once, and MTM-Multi, which reuses successive merged models as training references.Both variants select merge weights using multi-objective Bayesian optimization over the simplex.Source: §4 · pp. 4–5
4.1 Merge-Based Mixture Search
- Each soft-expert favors its target domain while including controlled amounts of other objectives, and the branches are merged into a unified model.MTM uses merge evaluations as a proxy for data-mixture search, approximating mixture behavior under local fine-tuning dynamics.Source: §4.1 · p. 5
4.2 Iterative Merging with Branch Reuse
- MTM-Single requires a second full training pass because its search branches differ from the final deployed model.Its total budget is roughly twice that of a single-trajectory baseline, and the final model does not benefit from search merge evaluations.Source: §4.2 · p. 5
- MTM-Multi uses each round’s merge evaluations both to promote the next reference model and to update branch mixtures, reusing branches to keep training cost aligned with a single-trajectory baseline.The selected merged model also realigns branches between rounds, while mixture updates sharpen each target domain subject to preserving other objectives.Source: §4.2 · pp. 5–6
4.3 Multi-Objective Bayesian Optimization over the Merge Simplex
- MOBO adaptively proposes simplex merge weights to expand the Pareto frontier, avoiding dense grids that grow costly or coarse grids that miss interior trade-offs.qNEHVI handles noisy batched evaluations and targets hypervolume improvement; the resulting candidate set supports global model promotion and constrained branch-mixture updates.Source: §4.3 · p. 6
5 Main Experiments
Across multi-domain and hybrid-mode tests, MTM improves on merging-only alternatives while retaining cross-domain performance and think/non-think separation.
5.1 Experimental Setup
- Experiments use Qwen3-4B and OLMo-7B across math, code, instruction-following, and think/non-think benchmarks, comparing MTM with merging and mixture baselines.The multi-domain data include 10K examples per objective; hybrid branches use 10K examples per mode.Source: §5 · p. 7, §5.1 · p. 7
5.2 Multi-Domain
- MTM-Multi achieves the highest average score on Qwen3-4B-Base and OLMo-7B, and alone keeps math, code, and instruction-following performance close to the best single-objective experts.It also consistently improves over MergeMix, particularly on instruction-following metrics.Source: §5.2 · pp. 7–8
- MTM-Multi’s branches train independently in parallel and reuse initial experts, keeping wall-clock cost comparable to same-stage mixture training.Table 2 compares total training budget and effective optimization path length.Source: Table 2 · p. 8, §5.2 · p. 8
- MTM-Multi consistently outperforms MTM-Single, with the largest gains on interference-sensitive metrics such as IFEval.Each round realigns branches around the promoted reference model and adapts subsequent mixtures using specialist ratios.Source: §5.2 · p. 8
5.3 Hybrid Mode
- MTM preserves think/non-think separation and shifts the length-accuracy frontier toward fewer no-think tokens while matching or exceeding Mixture Training accuracy under think prompts.Linear Merge collapses the modes, while Mixture Training recovers their separation.Source: §5.3 · p. 8, Figure 4 · p. 8
- Soft branches yield a flatter, higher-valued merge-simplex interior than Hard branches on MATH-500 and IFEval.The comparison is shown as ternary heatmaps over merge coefficients on Qwen3-4B-Base.Source: Figure 5 · p. 9
6 Additional Experiments
MTM’s mixture-trained branches produce a more favorable merge-simplex interior, and recipes selected on Qwen3-4B-Base transfer to Qwen2.5-0.5B without additional search.
6 Additional Experiments
- MTM Transfer achieves a higher average score than both size-matched baselines on Qwen2.5-0.5B using Qwen3-4B-Base recipes without rerunning the search.All methods use a fixed three-epoch budget; the baselines use uniform merge or uniform mixture rather than the source recipe.Source: Table 3 · p. 9, §6 · p. 9
- Hard branches concentrate accuracy near simplex vertices and decay along interior rays, whereas Soft branches yield a flatter, higher-valued interior with non-degenerate solutions satisfying cross-objective constraints.This is geometric evidence that mergeability is a property of the branches, not the merger.Source: §6 · p. 9
7 Conclusion
The conclusion presents Mixture-Trained Merging (MTM) as an iterative way to train branches that remain compatible when merged, reporting strong multi-objective performance and preserved think/non-think separation. The study is limited to SFT-stage merging, with larger objective sets and post-RL training left for future work.
7 Conclusion
- MTM-Multi achieves the best average performance across Qwen3-4B-Base and OLMo-7B, keeps objectives close to single-domain experts, and preserves think/non-think separation.The conclusion also reports a flatter, higher-valued simplex for mixture-trained branches.Source: §7 · p. 9
- The authors attribute failures of single-objective merging to experts drifting away from a shared high-performing basin, causing cross-domain trade-offs or think/non-think mode collapse.MTM trains each branch on an objective-biased mixture so task vectors remain compatible.Source: §7 · p. 9
- MTM iterates between MOBO-based mixture search, branch specialization, and merging, using merged-model evaluations as a proxy for data-mixture training.The conclusion states that mergeability depends on the branches, not only on the merging rule.Source: §7 · p. 9
- The study covers SFT-stage merging; larger objective sets and post-RL training remain future directions.Source: §7 · p. 9
A.1 Detail for Qwen3-4B-Base
The Qwen3-4B-Base setup fine-tunes on domain-specific instruction trajectories and uses sampled domain proportions for mixture-data SFT. Training uses assistant-token loss with validation-based early stopping, plus a geometric step-budget schedule for BO branch retraining.
A.1 Detail for Qwen3-4B-Base
- Qwen training uses 10,000 examples per run from separate math, coding, and precise-IF corpora, with NOTHINK trajectories generated from Qwen3-4B-Instruct-25079.The setup filters samples over 15,000 total tokens and caps generated continuations at 8,192 tokens.Source: §A.1 · p. 12
- For mixture-data SFT in BO branch retrains, domain proportions are sampled from the math, code, and instruction-following corpora.Source: §A.1 · p. 12
- Training uses 4× H200 GPUs with DeepSpeed ZeRO-2, bf16 weights, TF32 matmul, and loss computed only over assistant tokens.Source: §A.1 · p. 12
- Training uses plateau-based early stopping, while BO branch retraining also applies a geometric step budget with decay factor r=0.5.Early stopping monitors an EMA of validation token-level accuracy after a 100-step warm-up.Source: §A.1 · p. 12
A.2 Detail for OLMo-7B
The OLMo-7B setup covers multi-domain and think/non-think training, while Bayesian optimization searches merge coefficients and selects global and specialist ratios for subsequent iterations.
A.2 Detail for OLMo-7B
- OLMo-7B is fine-tuned in multi-domain and hybrid-mode settings, with separate math, coding, and precise-IF corpora and additional THINK trajectories for hybrid mode.THINK trajectories include a reasoning block before the final answer, unlike NOTHINK trajectories.Source: §A.2 · p. 12
- OLMo sampling matches the Qwen setup, and training computes loss over assistant tokens, including the reasoning block in THINK mode.Unlike Qwen, OLMo uses a fixed epoch budget without early stopping; Table 5 reports identical hyperparameters across THINK and NOTHINK runs, with different input corpora and trajectory-length caps.Source: §A.2 · p. 12, Table 5 · p. 13
A.3 Bayesian Optimization Details
- At each MTM iteration, Bayesian optimization evaluates candidate merged checkpoints on all objectives, beginning with simplex vertices, a central or previously best point, and space-filling random candidates.Source: §A.3 · p. 13
- The optimization models objectives with independent Gaussian-process surrogates and uses sequential qNEHVI to seek hypervolume gains near the expert frontier while balancing objectives.Source: §A.3 · p. 13
- After the BO budget is exhausted, a balanced score selects the global merge coefficient, while specialist coefficients keep the target near its expert score and constrain degradation on other objectives.The selected coefficients guide base promotion and mixture propagation in the next iteration.Source: §A.3 · p. 13
A.4 Evaluation Settings
The evaluation setup averages four generations per prompt and uses benchmark-specific settings. The accompanying algorithm iterates through branch training, Bayesian-optimization-based merge-ratio search, base promotion, and mixture updates.
A.4 Evaluation Settings
- Models are evaluated with NeMo-Skills and official benchmark evaluators, averaging one greedy and three sampled responses per prompt.Generation limits vary by benchmark from 8,196 to 30,000 tokens; Table 6 summarizes benchmark settings and generation counts.Source: §A.4 · p. 13, Table 6 · p. 13
B Algorithm
- The algorithm initializes the candidate set with simplex vertices, a uniform weight, and random Dirichlet samples before proposing and evaluating merge ratios.Source: §A.4 · p. 14
- The algorithm selects a global ratio and specialist ratios based on constrained target-domain performance, then updates branch mixtures for the next iteration.Source: §A.4 · p. 14
- Algorithm 1 describes the full MTM iteration, including branch training, BO-based merge-ratio search, base promotion, and mixture updates.Source: B Algorithm · p. 14