Unified language models improve by training on mixed data blends
Mixture-Trained Merging for Unified Multi-Objective Models
Machine Learning
Summary
Large language models often try to combine different skills like math, coding, or following instructions, but training for all of these separately and then joining them can make models unstable or lose some skills. The authors found that training models on mixtures of task types rather than one at a time makes their components more compatible when merged. They developed a method called Mixture-Trained Merging (MTM) that picks the best blends of data for each skill part, improving how well the final combined model works without mixing up its behaviors. This approach helps build stronger models that can effectively handle multiple tasks without interference.
What this means in practice
- •For machine learning engineers: Combine multiple specialized models into one unified model that retains each skill’s strength without degradation.
- •For ai system integrators: Build more reliable multi-capability AI systems by using mixture-trained merging to maintain separation between different task behaviors.
Authors
SeongHyeon Kim, Chaeyun Jang, Seungyoo Lee, Jiyeon Ham, Yunju Bak, Boseop Kim, Juho Lee
Abstract
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.