AutoMixAlign: Adaptive Data Mixing for Multi-Task Preference Optimization in LLMs

📄 arXiv: 2506.00569v1 📥 PDF

作者: Nicholas E. Corrado, Julian Katz-Samuels, Adithya Devraj, Hyokun Yun, Chao Zhang, Yi Xu, Yi Pan, Bing Yin, Trishul Chilimbi

分类: cs.LG

发布日期: 2025-05-31

备注: ACL 2025, Main Conference

💡 一句话要点

提出AutoMixAlign以解决多任务偏好优化问题

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 多任务学习 偏好优化 自适应数据混合 大型语言模型 极小极大优化

📋 核心要点

  1. 现有方法在选择训练数据混合时面临高成本和效果不佳的问题，难以在多任务上实现平衡表现。
  2. 论文提出的AutoMixAlign算法通过自适应混合数据集，利用专门模型和通用模型的损失差异来优化任务表现。
  3. 实验结果表明，AMA在多任务对齐设置中优于传统的总损失优化方法和模型合并方法，展现出显著的性能提升。

📝 摘要（中文）

在对齐大型语言模型（LLMs）时，其在多项任务（如有用性、无害性和诚实性）上的表现高度依赖于训练数据的组成。然而，选择一种能够在所有任务上都表现良好的数据混合方式是具有挑战性的。现有方法依赖于大规模消融研究、启发式方法或人类直觉，但这些方法成本高昂且效果不佳。本文研究了通过DPO进行偏好优化的问题，提出了AutoMixAlign（AMA），这是一种理论基础的算法，能够在训练过程中自适应地混合数据集，以平衡各任务的表现。AMA首先为每个任务训练“专门模型”，以确定与强任务表现相对应的损失，然后使用一种新颖的极小极大优化方法训练通用模型，优先考虑通用模型损失与专门模型损失偏差最大的任务。我们提出了两种算法：AMA-R和AMA-S，均在凸情况下实现了$O(1/ extsqrt{T})$的收敛速率。

🔬 方法详解

问题定义：本文旨在解决大型语言模型在多任务偏好优化中，如何选择训练数据混合以实现各任务表现平衡的问题。现有方法依赖于消融研究和启发式方法，成本高且效果不理想。

核心思路：AutoMixAlign（AMA）通过训练专门模型来确定各任务的损失，并利用这些损失的偏差来指导通用模型的训练，从而实现自适应的数据混合。

技术框架：AMA的整体流程包括两个主要阶段：首先为每个任务训练专门模型以获取损失信息，然后通过极小极大优化方法训练通用模型，优先考虑损失偏差较大的任务。

关键创新：AMA的核心创新在于其自适应数据混合策略，通过动态调整任务的优先级和数据采样量，显著提升了多任务学习的效果，与传统方法相比具有本质区别。

关键设计：AMA-R算法通过自适应重加权目标函数来优先考虑任务，而AMA-S算法则通过调整每个任务的数据采样量来实现优先级调整。两者在凸情况下均实现了$O(1/ extsqrt{T})$的收敛速率。具体损失函数和参数设置在实验中进行了详细验证。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示，AutoMixAlign在多个多任务对齐设置中均优于标准对齐方法，优化总损失的传统方法，以及模型合并方法，展现出显著的性能提升，具体提升幅度在实验中得到了验证。

🎯 应用场景

该研究的潜在应用领域包括多任务学习、自然语言处理和智能助手等场景。通过优化大型语言模型的训练数据组合，能够提升模型在多项任务上的综合表现，具有重要的实际价值和广泛的应用前景。

📄 摘要（原文）

When aligning large language models (LLMs), their performance on various tasks (such as being helpful, harmless, and honest) depends heavily on the composition of their training data. However, selecting a data mixture that achieves strong performance across all tasks is challenging. Existing approaches rely on large ablation studies, heuristics, or human intuition, but these can be prohibitively expensive and suboptimal. We study this problem in the setting of preference optimization via DPO and introduce AutoMixAlign (AMA), a theoretically-grounded algorithm that adaptively mixes datasets during training to balance performance across tasks. AMA first trains \textit{specialist models} for each task to determine losses that correspond to strong task performance. Then, it trains a generalist model using a novel minimax optimization that prioritizes tasks for which generalist model losses deviate most from specialist model losses. To optimize this problem, we propose two algorithms: (1) AMA-R, which adaptively reweights the objective to prioritize tasks, and (2) AMA-S, which adaptively adjusts how much data is sampled from each task to prioritize tasks. Both algorithms achieve a convergence rate of $O(1/\sqrt{T})$ in the convex case. AMA-R's convergence result follows from Sagawa et al. (2019), and we provide a convergence proof for AMA-S using online learning techniques such as EXP3. We evaluate AMA on several multitask alignment setups and find that AMA outperforms the standard alignment approach -- which simply optimizes the total loss across all tasks -- and also outperforms model merging methods.