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

TL;DR

AutoMixAlign (AMA) is a theoretically-grounded adaptive data mixing algorithm that improves multi-task preference optimization in large language models by balancing task performance during training.

Contribution

AMA introduces a novel minimax optimization framework with two algorithms, AMA-R and AMA-S, for adaptive data mixing in multi-task LLM alignment, backed by convergence guarantees.

Findings

AMA outperforms standard total loss optimization in multi-task alignment.

AMA surpasses model merging methods in multi-task preference performance.

Both AMA-R and AMA-S achieve $O(1/\sqrt{T})$ convergence rate in convex settings.

Abstract

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.