多任务偏好对齐（helpfulness、harmlessness、coding 等）面临一个工程痛点：数据配比怎么选？均匀混合容易被大数据集主导，按任务等权又可能在简单任务上浪费容量。常见做法是跑大量消融实验或凭经验设定比例，成本高且不一定最优。

AutoMixAlign（AutoMixAlign: Adaptive Data Mixing for Multi-Task Preference Optimization in LLMs，ACL 2025）的做法是：先为每个任务单独训练 specialist 模型作为参照基线，再在训练 generalist 模型时基于 minimax 优化动态调整各任务的权重或采样比例，优先关注那些 generalist 与 specialist 之间 loss 差距最大的任务。在 Zephyr-7B 上的实验中，多任务平均性能最高提升约 9.42%，且避免了标准 DPO 混合训练中常见的"某个任务退化"问题。

问题设定与标准方法的不足

假设有 $k$ 个偏好数据集 $D_1, \ldots, D_k$，每个对应一个对齐任务（如 helpfulness、safety、coding）。标准 DPO 训练把所有数据合并后最小化加权 loss：

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

常见的权重选择有两种。第一种是 uniform weighting，$\alpha_i = 1$，即把所有样本的 loss 直接求和。问题在于大数据集对总 loss 的贡献更大，梯度方向会偏向大数据集对应的任务。第二种是 task-normalized weighting，$\alpha_i = 1/|D_i|$，使各任务贡献它们各自的平均 loss，任务间等权。但等权同样不一定最优：如果某些任务天然更容易学习，固定等权会导致模型在简单任务上过度优化，忽略更需要关注的难任务。

Specialist 基线与 Minimax 优化

上图展示 AMA 的三步流程：①为每个任务单独跑 DPO 训练得到 specialist 模型；②用 specialist 在各自任务数据上计算 loss 并存入数据集；③训练 generalist 时，根据各任务的 excess loss（generalist loss 减去 specialist loss）动态调整任务权重或采样概率，excess loss 大的任务获得更多关注。

AMA 的核心设计是引入 specialist 模型作为"各任务能做到多好"的参照。对每个任务 $i$，先单独用 $D_i$ 做 DPO 训练得到 specialist 参数 $\theta_i$。Specialist 在本任务上表现好但跨任务能力有限，它的 loss 提供了"单任务训练在该任务上能达到什么水平"的参照值。

接下来 AMA 要找一组 generalist 参数 $\theta$，使其在每个任务上都尽可能接近对应 specialist 的水平。形式化为一个可行性问题：找 $\theta$ 满足对所有 $i$，generalist loss 不超过 specialist loss。由于精确求解不可行，将其松弛为最小化最大 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\}$$

这里的 $\max\{\cdot, 0\}$ 称为 clipping：当 generalist 在某个样本上已经比 specialist 做得好时，不再继续压低这个样本的 loss，避免过拟合。这是 AMA 与标准 Group DRO 方法的主要区别。

两种优化算法

AMA-R（Reweighting）

将上述 minimax 问题改写为 $\min_\theta \max_{\alpha \in \Delta_k} \sum_i \alpha_i E_i(\theta)$，其中 $E_i(\theta)$ 是任务 $i$ 的 clipped excess loss。训练过程中交替更新：用 exponentiated gradient ascent 增大 excess loss 较大的任务的权重 $\alpha_i$，然后用更新后的加权 loss 对 $\theta$ 做梯度下降。

AMA-R 每步的 batch 包含所有任务的样本（每个任务约 b/k 个），用各任务的加权 loss 做一次参数更新。

AMA-S（Resampling）

AMA-R 的一个效率问题：当某些任务权重很小时，从这些任务采的样本对梯度几乎没有贡献，计算浪费。AMA-S 用 EXP3（多臂老虎机算法）动态调整各任务的采样概率，excess loss 大的任务被采到的概率更高，计算资源集中在需要关注的任务上。

两种算法在凸情况下均有 $O(1/\sqrt{T})$ 的收敛率保证。AMA-R 的收敛性直接继承自 Sagawa et al. (2019) 的 Group DRO 分析，AMA-S 的证明则把问题建模为零和博弈，利用 EXP3 的 regret bound 推导。

Clipping 的作用

消融实验中，去掉 clipping（AMA-S Unclipped、AMA-R Unclipped）后 MBPP 等 coding 指标明显退化。原因在于：当 generalist 在某任务上的 loss 已经低于 specialist 时，继续优化会让模型在该任务的训练数据上过拟合。论文在实验中观察到，不做 clipping 的版本在 0.5 个 epoch 后 excess loss 变为负值，表明模型已经越过了 specialist 水平，此后的优化是有害的。

Clipping 的语义很清晰：specialist loss 是"这个任务做到这个水平就够了"的信号，超过这个水平的 loss 减少不计入优化目标。

实验结果

实验在 Zephyr-7B-SFT-Full 上进行，评测涵盖 helpfulness（AlpacaEval 2.0、IFEval）、coding（MBPP、HumanEval）、harmlessness（Toxigen）。考察了三种数据集组合：

Setup 1（Helpfulness + Coding）：Standard DPO 在 coding 上退化，MBPP 从基线的约 50% 降到约 38%。AMA-R 保持 helpfulness 并把 MBPP 提升到约 51%，平均分从约 35% 提升到约 42%。

Setup 2（Helpfulness + Harmlessness）：Standard DPO 在 AlpacaEval 上退化严重（从约 9% 降到约 7%），AMA-S 将 AlpacaEval 提升到约 17% 并保持接近 100% 的 Toxigen 得分。不过 Standard Uniform（各任务等概率采样）在这个设置下平均分约 53.21%，已经接近 AMA-S 的 53.81%，说明当任务数少且数据集比例恰好合适时，简单的等概率采样也能获得不错的平衡。

Setup 3（三任务联合）：AMA-R 达到约 54% 的平均分，比 Standard 高约 3.4 个百分点，比 Model Averaging 高约 3 个百分点。五个基准上均无退化。

Model Averaging（均匀平均 specialist 参数）在所有设置中都不如 AMA，但在 Setup 1 和 Setup 3 中仍优于 Standard DPO，说明 merging 有一定效果但不如端到端的自适应训练。

Resampling 与 Reweighting 的差异

上图是 Coding（CodeUltraFeedback）+ Harmlessness（SafeRLHF）设置下 AMA-S 与 AMA-R 的对比。上子图：Toxigen 无毒率随训练 epoch 的变化，AMA-S 在 0.4 epoch 后接近 100%，AMA-R 到 1.0 epoch 仍停留在约 90%。中子图：coding 指标（HumanEval + MBPP 均值），AMA-S 持续高于 AMA-R。下子图：SafeRLHF 的采样权重，AMA-S（蓝线）快速提升到 0.5 以上，而 AMA-R（橙线）始终固定在约 1/11。

在数据集大小极度不均衡时，AMA-S 的优势更明显。论文用一个 Coding + Harmlessness 的实验说明：CodeUltraFeedback 数据集远大于 SafeRLHF（约 10:1），AMA-R 的均匀采样导致每个 batch 中只有约 1/11 的样本来自 SafeRLHF，权重再大也无法补偿样本不足。AMA-S 能把 SafeRLHF 的采样比例快速提升，在同样的训练步数内达到接近 100% 的无毒率，而 AMA-R 最高只到约 90%。

直觉上，AMA-R 用权重调整 loss 在目标函数里的比例，但所有任务的 batch 大小是均匀的；AMA-S 直接调整看到多少该任务的样本。当某任务需要更多数据才能学好时，AMA-S 的调节粒度更细。

计算开销与局限

AMA 的额外开销集中在 specialist 训练阶段。训练 $k$ 个 specialist 各跑 $m$ 个 epoch，共 $m|D|$ 次更新；加上 generalist 阶段的 $m|D|$ 次更新，总计约为标准 DPO 的 2 倍。这个比例不随任务数增长。与跑多组消融实验确定配比相比，2 倍的固定开销通常更划算。

论文也提到了降低 specialist 开销的可能：用小模型估计 specialist loss（类似 DoReMi 的做法），或用 Hessian 近似（ExcessMTL）来替代完整的 specialist 训练。但这些变体在本文中未实验。

另一个局限是实验规模：所有实验只在 Zephyr-7B 上验证，任务数最多为 3。在更大模型或更多任务（如 10+ 个维度同时对齐）时的表现尚未验证。

小结

AMA 解决的问题清晰：多任务偏好优化中，手动选数据配比费时且容易偏向某些任务。解法是先用 specialist 定义各任务的 loss 目标，再用 minimax 优化在训练中自动把注意力分配给 loss 差距大的任务。两个算法变体分别适合不同场景：数据集大小均衡时 AMA-R 和 AMA-S 性能接近；数据集极度不均衡时 AMA-S 的动态采样更有效。Clipping 机制防止模型在已优于 specialist 的任务上继续过拟合，是方法稳健性的核心保障。