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

• URL: http://arxiv.org/abs/2506.00569v1
• Date: Sat, 31 May 2025 13:57:10 GMT
• Title: AutoMixAlign: Adaptive Data Mixing for Multi-Task Preference Optimization in LLMs
• Authors: Nicholas E. Corrado, Julian Katz-Samuels, Adithya Devraj, Hyokun Yun, Chao Zhang, Yi Xu, Yi Pan, Bing Yin, Trishul Chilimbi,
• Abstract summary: We study a theoretically-grounded algorithm that adaptively mixes during training to balance performance across tasks.<n>AMA-R adaptively reweights the objective to prioritize tasks, and AMA-S adaptively adjusts how much data is sampled from each task to prioritize tasks.<n>We evaluate on several generalist alignment setups and find that AMA outperforms the standard alignment approach.
• Score: 33.87805544053889
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• DriveLMM-o1: A Step-by-Step Reasoning Dataset and Large Multimodal Model for Driving Scenario Understanding [76.3876070043663]
  We propose DriveLMM-o1, a dataset and benchmark designed to advance step-wise visual reasoning for autonomous driving.<n>Our benchmark features over 18k VQA examples in the training set and more than 4k in the test set, covering diverse questions on perception, prediction, and planning.<n>Our model achieves a +7.49% gain in final answer accuracy, along with a 3.62% improvement in reasoning score over the previous best open-source model.
  arXiv  Detail & Related papers  (2025-03-13T17:59:01Z)
• DUET: Optimizing Training Data Mixtures via Feedback from Unseen Evaluation Tasks [40.91931801667421]
  Our paper presents a novel global-to-local algorithm that interleaves influence function as a data selection method with Bayesian optimization to optimize data mixture via feedback from a specific unseen evaluation task.<n>By analyzing DUET's cumulative regret, we theoretically show that DUET converges to the optimal training data mixture for an unseen task even without any data knowledge of the task.
  arXiv  Detail & Related papers  (2025-02-01T01:52:32Z)
• Ensembles of Low-Rank Expert Adapters [9.599957499802446]
  We propose the Ensembles of Low-Rank Expert Adapters (ELREA) framework to improve the model's capability to handle diverse tasks.<n>ELREA clusters the training instructions based on their gradient directions, representing different areas of expertise.<n>During inference, ELREA combines predictions from the most relevant expert adapters based on the input data's gradient similarity to the training clusters.
  arXiv  Detail & Related papers  (2025-01-31T18:07:21Z)
• Modeling Multi-Task Model Merging as Adaptive Projective Gradient Descent [72.10987117380584]
  Merging multiple expert models offers a promising approach for performing multi-task learning without accessing their original data.<n>We find existing methods discard task-specific information that, while causing conflicts, is crucial for performance.<n>Our approach consistently outperforms previous methods, achieving state-of-the-art results across diverse architectures and tasks in both vision and NLP domains.
  arXiv  Detail & Related papers  (2025-01-02T12:45:21Z)
• A Systematic Examination of Preference Learning through the Lens of Instruction-Following [83.71180850955679]
  We use a novel synthetic data generation pipeline to generate 48,000 instruction unique-following prompts.<n>With our synthetic prompts, we use two preference dataset curation methods - rejection sampling (RS) and Monte Carlo Tree Search (MCTS)<n>Experiments reveal that shared prefixes in preference pairs, as generated by MCTS, provide marginal but consistent improvements.<n>High-contrast preference pairs generally outperform low-contrast pairs; however, combining both often yields the best performance.
  arXiv  Detail & Related papers  (2024-12-18T15:38:39Z)
• Scalable Fine-tuning from Multiple Data Sources: A First-Order Approximation Approach [17.79010397902909]
  We study the problem of fine-tuning a language model (LM) for a target task by optimally using the information from $n$ auxiliary tasks.<n>This problem has broad applications in NLP, such as targeted instruction tuning and data selection in chain-of-thought fine-tuning.<n>We introduce a new algorithm for estimating model fine-tuning performance without requiring repeated training.
  arXiv  Detail & Related papers  (2024-09-28T21:26:50Z)
• Robust Multi-Task Learning with Excess Risks [24.695243608197835]
  Multi-task learning (MTL) considers learning a joint model for multiple tasks by optimizing a convex combination of all task losses. Existing methods use an adaptive weight updating scheme, where task weights are dynamically adjusted based on their respective losses to prioritize difficult tasks. We propose Multi-Task Learning with Excess Risks (ExcessMTL), an excess risk-based task balancing method that updates the task weights by their distances to convergence.
  arXiv  Detail & Related papers  (2024-02-03T03:46:14Z)
• STG-MTL: Scalable Task Grouping for Multi-Task Learning Using Data Map [4.263847576433289]
  Multi-Task Learning (MTL) is a powerful technique that has gained popularity due to its performance improvement over traditional Single-Task Learning (STL) However, MTL is often challenging because there is an exponential number of possible task groupings. We propose a new data-driven method that addresses these challenges and provides a scalable and modular solution for classification task grouping.
  arXiv  Detail & Related papers  (2023-07-07T03:54:26Z)
• FAMO: Fast Adaptive Multitask Optimization [48.59232177073481]
  We introduce Fast Adaptive Multitask Optimization FAMO, a dynamic weighting method that decreases task losses in a balanced way. Our results indicate that FAMO achieves comparable or superior performance to state-of-the-art gradient manipulation techniques.
  arXiv  Detail & Related papers  (2023-06-06T15:39:54Z)
• Improving Generalization in Meta-learning via Task Augmentation [69.83677015207527]
  We propose two task augmentation methods, including MetaMix and Channel Shuffle. Both MetaMix and Channel Shuffle outperform state-of-the-art results by a large margin across many datasets.
  arXiv  Detail & Related papers  (2020-07-26T01:50:42Z)
• Non-Adaptive Adaptive Sampling on Turnstile Streams [57.619901304728366]
  We give the first relative-error algorithms for column subset selection, subspace approximation, projective clustering, and volume on turnstile streams that use space sublinear in $n$. Our adaptive sampling procedure has a number of applications to various data summarization problems that either improve state-of-the-art or have only been previously studied in the more relaxed row-arrival model.
  arXiv  Detail & Related papers  (2020-04-23T05:00:21Z)
• Improving Multi-Turn Response Selection Models with Complementary Last-Utterance Selection by Instance Weighting [84.9716460244444]
  We consider utilizing the underlying correlation in the data resource itself to derive different kinds of supervision signals. We conduct extensive experiments in two public datasets and obtain significant improvement in both datasets.
  arXiv  Detail & Related papers  (2020-02-18T06:29:01Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.