Multi-task preference alignment (helpfulness, harmlessness, coding, etc.) faces a practical engineering problem: how to choose the data mixture. Uniform mixing lets large datasets dominate the gradient, while equal task weighting may waste capacity on easy tasks. The common practice of running extensive ablation studies or manually tuning ratios is expensive and often suboptimal.

AutoMixAlign (AutoMixAlign: Adaptive Data Mixing for Multi-Task Preference Optimization in LLMs, ACL 2025) takes a different approach: first train a specialist model for each task as a reference baseline, then dynamically adjust task weights or sampling probabilities during generalist training via minimax optimization, prioritizing tasks where the generalist-specialist loss gap is largest. On Zephyr-7B, multi-task average performance improves by up to 9.42% while avoiding the “one task degrades” failure mode common in standard DPO mixed training.

Problem Setup and Limitations of Standard Approaches

Given $k$ preference datasets $D_1, \ldots, D_k$, each corresponding to an alignment task (e.g., helpfulness, safety, coding), standard DPO training merges all data and minimizes a weighted loss:

$$L(\theta) = \sum_{i=1}^{k} \alpha_i \sum_{z \in D_i} L(\theta, z)$$

Two common weight choices exist. The first is uniform weighting, $\alpha_i = 1$, which sums all sample losses directly. The problem is that larger datasets contribute more to the total loss and the gradient direction is biased toward tasks with more data. The second is task-normalized weighting, $\alpha_i = 1/|D_i|$, which makes each task contribute its average loss equally. But equal weighting is not necessarily optimal either: if some tasks are inherently easier to learn, a fixed equal weight causes the model to over-optimize on easy tasks while neglecting harder ones.

Specialist Baselines and Minimax Optimization

The figure above shows AMA’s three-step pipeline: ① train a specialist model for each task via DPO individually; ② compute specialist losses on each task’s data and store them; ③ during generalist training, dynamically adjust task weights or sampling probabilities based on excess loss (generalist loss minus specialist loss), giving more attention to tasks with larger gaps.

The core design of AMA is introducing specialist models as a reference for “how well each task can be done.” For each task $i$, a specialist $\theta_i$ is trained via DPO on $D_i$ alone. The specialist performs well on its own task but has limited cross-task capabilities; its loss provides a reference point for what single-task training can achieve.

AMA then seeks generalist parameters $\theta$ that match each specialist’s performance as closely as possible. This is formalized as a feasibility problem: find $\theta$ such that generalist loss does not exceed specialist loss for all $i$. Since exact solutions are intractable, the problem is relaxed to minimizing the maximum excess loss:

$$\min_\theta \max_{i \in [k]} \frac{1}{|D_i|} \sum_{z \in D_i} \max\{L(\theta, z) - L(\theta_i, z),\ 0\}$$

The $\max\{\cdot, 0\}$ term is called clipping: when the generalist already outperforms the specialist on a sample, no further loss reduction is pursued for that sample, preventing overfitting. This is the primary distinction between AMA and standard Group DRO methods.

Two Optimization Algorithms

AMA-R (Reweighting)

The minimax problem is rewritten as $\min_\theta \max_{\alpha \in \Delta_k} \sum_i \alpha_i E_i(\theta)$, where $E_i(\theta)$ is task $i$’s clipped excess loss. Training alternates between: using exponentiated gradient ascent to increase weights $\alpha_i$ for tasks with larger excess loss, then performing gradient descent on $\theta$ with the reweighted loss.

Each AMA-R batch contains samples from all tasks (approximately b/k per task), and one parameter update is made using the weighted loss across all tasks.

AMA-S (Resampling)

An efficiency issue with AMA-R: when some task weights are very small, samples from those tasks contribute almost nothing to the gradient, wasting computation. AMA-S uses EXP3 (a multi-armed bandit algorithm) to dynamically adjust sampling probabilities across tasks, concentrating compute on tasks with higher excess loss.

Both algorithms achieve $O(1/\sqrt{T})$ convergence rates in the convex case. AMA-R’s convergence follows directly from Sagawa et al. (2019)’s Group DRO analysis, while AMA-S’s proof models the problem as a zero-sum game and leverages EXP3’s regret bound.

The Role of Clipping

In ablation experiments, removing clipping (AMA-S Unclipped, AMA-R Unclipped) causes significant degradation on coding benchmarks like MBPP. The reason: when the generalist’s loss on a task already drops below the specialist’s, continued optimization causes overfitting on that task’s training data. The paper observes that the unclipped variant reaches negative excess loss after 0.5 epochs, indicating the model has optimized past the specialist level, and further optimization is harmful.

The semantics of clipping are straightforward: specialist loss serves as a “this task is good enough at this level” signal, and loss reductions beyond that level are excluded from the optimization objective.

Experimental Results

Experiments use Zephyr-7B-SFT-Full as the base model, evaluating helpfulness (AlpacaEval 2.0, IFEval), coding (MBPP, HumanEval), and harmlessness (Toxigen). Three dataset combinations are tested:

Setup 1 (Helpfulness + Coding): Standard DPO degrades coding, with MBPP dropping from ~50% to ~38%. AMA-R maintains helpfulness while boosting MBPP to ~51%, with average scores rising from ~35% to ~42%.

Setup 2 (Helpfulness + Harmlessness): Standard DPO severely degrades AlpacaEval (from ~9% to ~7%). AMA-S raises AlpacaEval to ~17% while maintaining near-100% Toxigen scores. However, Standard Uniform (equal task sampling probability) achieves ~53.21% average in this setup, close to AMA-S’s 53.81%, suggesting that when task counts are small and dataset proportions happen to be suitable, simple equal-probability sampling can also achieve reasonable balance.

Setup 3 (all three tasks): AMA-R reaches ~54% average, approximately 3.4 points above Standard and 3 points above Model Averaging. No degradation on any of the five benchmarks.

Model Averaging (uniform averaging of specialist parameters) underperforms AMA in all setups, though it still outperforms Standard DPO in Setups 1 and 3, indicating that merging has some effect but is inferior to end-to-end adaptive training.

Resampling vs. Reweighting

The figure above compares AMA-S and AMA-R on a Coding (CodeUltraFeedback) + Harmlessness (SafeRLHF) setup. Top panel: Toxigen non-toxicity rate over training epochs; AMA-S reaches near 100% after 0.4 epochs while AMA-R plateaus around 90% at 1.0 epoch. Middle panel: coding metrics (HumanEval + MBPP average); AMA-S consistently outperforms AMA-R. Bottom panel: SafeRLHF sampling weight; AMA-S (blue) quickly ramps up above 0.5, while AMA-R (orange) remains fixed at ~1/11.

When dataset sizes are highly imbalanced, AMA-S has a clear advantage. CodeUltraFeedback is roughly 10x larger than SafeRLHF, so AMA-R’s uniform sampling only allocates about 1/11 of each batch to SafeRLHF; no matter how large the weight becomes, it cannot compensate for insufficient samples. AMA-S can rapidly increase SafeRLHF’s sampling proportion, reaching near-100% non-toxicity within the same training steps, while AMA-R maxes out at ~90%.

Intuitively, AMA-R adjusts the proportion of each task’s loss in the objective function, but batch composition remains uniform across tasks; AMA-S directly controls how many samples from each task the model sees. When a task requires more data exposure to learn well, AMA-S provides finer-grained control.

Computational Overhead and Limitations

AMA’s additional overhead is concentrated in the specialist training phase. Training $k$ specialists for $m$ epochs requires $m|D|$ total updates; combined with $m|D|$ updates for generalist training, the total is approximately 2x standard DPO. This ratio does not grow with the number of tasks. Compared to running multiple ablation studies to determine mixture ratios, a fixed 2x overhead is usually a better deal.

The paper mentions possible approaches to reduce specialist overhead: estimating specialist losses with smaller models (similar to DoReMi), or using Hessian approximations (ExcessMTL) to replace full specialist training. These variants are not experimentally evaluated in this work.

Another limitation is experimental scale: all experiments are conducted on Zephyr-7B only, with at most 3 tasks. Performance on larger models or more tasks (e.g., 10+ alignment dimensions simultaneously) remains unverified.

Summary

AMA addresses a clear problem: manually selecting data mixtures for multi-task preference optimization is time-consuming and prone to favoring certain tasks. The solution is to use specialists to define per-task loss targets, then let minimax optimization automatically allocate attention to tasks with the largest loss gaps during training. The two algorithm variants suit different scenarios: when dataset sizes are balanced, AMA-R and AMA-S perform similarly; when datasets are highly imbalanced, AMA-S’s dynamic sampling is more effective. The clipping mechanism prevents the model from continuing to overfit on tasks where it already surpasses the specialist, serving as the core safeguard for robustness.