IDEAL: Data Equilibrium Adaptation for Multi-Capability Language Model Alignment

TL;DR

IDEAL tackles the data distribution problem in multi-domain SFT by learning a domain-weight vector $\bm{\beta}$ that rebalances data volumes using a second-order, Hessian-informed objective. Through a bi-level optimization and efficient Hessian handling via K-FAC-based approximations, IDEAL iteratively adjusts domain data proportions to minimize the reference loss $\mathcal{L}(\mathcal{D}_{ref}, \bm{\theta}^*)$, achieving robust multi-task improvements. Key contributions include a principled gradient-based update rule, an efficient computation strategy for large models, and empirical evidence showing ~7% average gains across diverse capabilities with stable two-iteration convergence. The approach reduces reliance on costly hyperparameter sweeps and demonstrates data distribution as a critical lever for aligning LLM capabilities across domains, albeit with limitations from Hessian approximations and dataset quality requirements.

Abstract

Large Language Models (LLMs) have achieved impressive performance through Supervised Fine-tuning (SFT) on diverse instructional datasets. When training on multiple capabilities simultaneously, the mixture training dataset, governed by volumes of data from different domains, is a critical factor that directly impacts the final model's performance. Unlike many studies that focus on enhancing the quality of training datasets through data selection methods, few works explore the intricate relationship between the compositional quantity of mixture training datasets and the emergent capabilities of LLMs. Given the availability of a high-quality multi-domain training dataset, understanding the impact of data from each domain on the model's overall capabilities is crucial for preparing SFT data and training a well-balanced model that performs effectively across diverse domains. In this work, we introduce IDEAL, an innovative data equilibrium adaptation framework designed to effectively optimize volumes of data from different domains within mixture SFT datasets, thereby enhancing the model's alignment and performance across multiple capabilities. IDEAL employs a gradient-based approach to iteratively refine the training data distribution, dynamically adjusting the volumes of domain-specific data based on their impact on downstream task performance. By leveraging this adaptive mechanism, IDEAL ensures a balanced dataset composition, enabling the model to achieve robust generalization and consistent proficiency across diverse tasks. Experiments across different capabilities demonstrate that IDEAL outperforms conventional uniform data allocation strategies, achieving a comprehensive improvement of approximately 7% in multi-task evaluation scores.

Figures (3)

• Figure 1: IDEAL adjusts data mixture proportions to optimize model performance, leading to a decrease in loss.
• Figure 2: Models' performance on different domains with varying $m$ values.
• Figure : IDEAL Algorithm

Theorems & Definitions (2)

• Lemma 1
• Proof 1