---
title: 'AutoMixAlign (AMA): Multi-Task LLM Alignment'
url: https://www.emergentmind.com/topics/automixalign-ama
type: topic
---

# AutoMixAlign (AMA): Multi-Task LLM Alignment

AutoMixAlign (AMA) is a theoretically-grounded algorithmic framework for multi-task preference optimization in aligning large language models (LLMs) by adaptive data mixing. AMA introduces task-adaptive training workflows for optimizing LLM alignment across multiple objectives, such as helpfulness, harmlessness, and coding ability, by systematically minimizing worst-case excess loss relative to task-specific specialist models. The framework consists of two main variants: adaptive reweighting (AMA-R) and adaptive resampling (AMA-S), both capable of minimizing the maximum per-task clipped excess loss using online learning techniques with provable convergence rates [2506.00569].

## 1. Multi-Task Preference Alignment Problem Formulation

Given $k$ preference-optimization tasks with corresponding datasets $D_1, ..., D_k$ where each $D_i = \{z = (x, y_+, y_-)\}$, the goal is to train a generalist LLM $\pi_\theta$ using Direct Preference Optimization (DPO) over a reference model $\pi_\mathrm{ref}$. The DPO per-example loss is defined as:
\[
\ell(\theta; z) = -\log \sigma\left(\beta \cdot \left(\frac{\pi_\theta(y_+|x)}{\pi_\mathrm{ref}(y_+|x)} - \frac{\pi_\theta(y_-|x)}{\pi_\mathrm{ref}(y_-|x)}\right)\right)
\]
Specialist models $\theta^*_i$ are obtained by training independently on each $D_i$. For a generalist model $\theta$, the average loss on $D_i$ is $L_i(\theta) = |D_i|^{-1} \sum_{z\in D_i}\ell(\theta;z)$, and the reference specialist loss is $L_i^* = L_i(\theta^*_i)$. The clipped excess loss is:
\[
L^{\text{excess}}_i(\theta) = |D_i|^{-1} \sum_{z\in D_i}\max\{\ell(\theta;z) - \ell(\theta^*_i;z), 0\}
\]
AMA seeks $\theta$ that minimizes the worst-case excess loss across all tasks:
\[
\min_\theta \max_{i\in[k]} L^{\text{excess}}_i(\theta)
\]
Introducing $\alpha\in\Delta^k$ (probability simplex), this is equivalently
\[
\min_\theta \max_{\alpha\in\Delta^k} \sum_{i=1}^k \alpha_i L^{\text{excess}}_i(\theta)
\]

## 2. AMA-R: Adaptive Reweighting via Minimax Optimization

AMA-R casts the objective as a two-player minimax game. The $\alpha$-player (task weights) performs exponentiated-gradient (EG) ascent to emphasize tasks where current generalist losses most exceed the specialist baseline, while the $\theta$-player updates parameters to minimize the weighted excess loss:
\[
\min_\theta \max_{\alpha\in\Delta^k} f(\theta,\alpha), \quad f(\theta,\alpha) = \sum_{i=1}^k \alpha_i L^{\text{excess}}_i(\theta)
\]
Algorithmic steps per iteration:
- $\alpha$-player: EG update $\alpha_i^{(t+1)} \propto \alpha_i^{(t)} \exp(\eta L^{\text{excess}}_i(\theta^{(t)}))$
- $\theta$-player: Stochastic gradient descent on $\sum_i \alpha_i^{(t)} L^{\text{excess}}_i(\theta)$

In practice, $\alpha$ is smoothed: $\alpha_i^{(t)} \gets (1-c)q_i^{(t)} + c/k$, where $q_i$ is the EG internal weight and $c$ is a smoothing parameter. Under convexity, converges at rate $O(1/\sqrt{T})$ as implied by Sagawa et al. (2019) [2506.00569].

## 3. AMA-S: Adaptive Resampling via Bandit Algorithms

AMA-S adaptively adjusts the sampling distribution over tasks using the bandit algorithm EXP3, rather than reweighting objective components. At each iteration:
- A minibatch is formed by first sampling task counts $(n_1, ..., n_k)\sim $ Multinomial$(b; p_1,...,p_k)$ where $p$ is the smoothed distribution.
- The loss gradient is computed on the batch using the clipped excess loss.
- The internal distribution $q$ is updated via $q_i^{(t+1)} \propto q_i^{(t)}\exp(\eta \hat\ell_i^{(t)}/p_i^{(t)})$ where $\hat\ell_i^{(t)}$ is the empirical average excess loss for task $i$ in the minibatch.

This process is a bandit-style solver for the minimax: $\min_\theta \max_{p\in\Delta^k} \sum_i p_i L^{\text{excess}}_i(\theta)$. With convex loss and boundedness assumptions, $O(1/\sqrt{T})$ convergence of the worst-case per-task excess loss is guaranteed:
\[
\max_{i} \frac{1}{T}\sum_t L_i^{\text{excess}}(\theta^{(t)}) \leq \min_\theta\max_i L^{\text{excess}}_i(\theta) + O\left(\sqrt{\frac{k \log k + C^2}{T}}\right)
\]
where $C$ bounds $\theta$-player regret.

## 4. Algorithmic Procedures

Summary of the two core AMA variants:

| Variant          | Adaptation Mechanism       | Update Rule for Task Distribution |
|------------------|---------------------------|-----------------------------------|
| AMA-R            | Objective reweighting      | EG: $\alpha_i^{(t+1)} \propto \alpha_i^{(t)} \exp(\eta L_i^{\text{excess}})$ |
| AMA-S            | Data resampling            | EXP3: $q_i^{(t+1)} \propto q_i^{(t)}\exp(\eta \hat\ell_i^{(t)}/p_i^{(t)})$ |

Both algorithms return the average model parameters $\bar\theta = (1/T) \sum_{t} \theta^{(t)}$. Smoothing and hyperparameters, such as $c \approx 0.1$, learning rates $\eta \approx 1.0$ (for AMA-R) and $5\times 10^{-7}$ for $\theta$-updates, are empirically recommended.

Practical guidelines include:
- Precomputing specialist losses $\ell^*_i(z)$ for all $z\in D_i$
- Using clipped excess loss to prevent overfitting to easier tasks
- Regular checkpointing, model selection using confidence interval overlap in multitask accuracy

## 5. Experimental Setup and Empirical Results

Experiments use Zephyr-7B SFT Full as the base LLM, with DPO and AMA for generalist tuning over 1–3 epochs (batch size 256, gradient accumulation 4, per-device batch 8, AdamW, $\beta_1=0.9$, $\beta_2=0.999$, warmup 10%). Specialist models are trained for each task with identical hyperparameters.

Task domains:
- Helpfulness: Chatbot Arena 2024 (LLM-as-judge), UltraFeedback
- Coding: CodeUltraFeedback, MBPP, HumanEval
- Harmlessness: SafeRLHF, Toxigen

Baselines:
- Standard: Uniform sampling over all data, minimize total loss
- Standard-Uniform: Uniform task-level sampling
- Model Averaging: Uniform parameter-wise average of all specialists

Key empirical findings:

| Setup                                 | Standard   | Standard-Uniform | Model Averaging | AMA-R   | AMA-S   |
|----------------------------------------|------------|------------------|-----------------|---------|---------|
| Helpfulness (Arena) + Coding (CodeUF)  | 35.14%     | 35.47%           | 39.43%          | 41.75%  | 40.19%  |
| Helpfulness (UltraFB) + Harmlessness   | 49.37%     | 53.21%           | 50.51%          | 53.50%  | 53.81%  |
| Help + Code + Harmlessness             | 50.95%     | 44.96%           | 51.50%          | 54.38%  | 53.18%  |

AMA improves the average metric by up to 9.42% over standard training. Notably, Figure X from the source demonstrates that AMA-S rapidly rebalances sampling to the task with greater excess loss, then stabilizes task allocation as losses converge. This suggests that AMA-S effectively tracks and mitigates task-level underperformance during multitask alignment [2506.00569].

## 6. Theoretical Properties and Extensions

AMA variants offer $O(1/\sqrt{T})$ convergence under convexity and bounded loss, and are robust to the number of tasks $k$. Specialist training is parallelizable, and total compute is approximately double the standard DPO regime, independent of $k$.

AMA methodologies extend directly to any DPO-style training, and are potentially applicable to RLHF, PPO, or supervised fine-tuning by substituting appropriate loss functions. Use of unclipped excess loss is discouraged as it leads to overfitting trivial tasks. Regular checkpointing, smoothing, and careful model selection using confidence intervals are recommended for practical deployment.

## 7. Broader Implications and Future Work

AMA addresses the critical challenge of data mixture selection in LLM alignment by algorithmically optimizing mixture weights or samplings to directly target minimax excess loss, sidestepping reliance on large-scale ablation or subjective heuristics. A plausible implication is that this approach could generalize to related domains in safe or robust multi-objective optimization. Future work may further explore extensions to non-convex LLM landscapes, broader RLHF settings, or adaptive mixture optimization in large-scale distributed fine-tuning [2506.00569].

Source: https://www.emergentmind.com/topics/automixalign-ama