Technical Summary: DYNAMIXSFT

Problem Statement

As the landscape of instruction-tuning datasets expands, the challenge of dynamically balancing and optimizing their mixtures has become critical. While diverse dataset mixtures are pivotal for enhancing the generalization of Large Language Models (LLMs) across domains like reasoning, coding, and knowledge, most existing approaches rely on static, manually curated heuristics or fixed rules. Alternative dynamic methods often require additional models (e.g., actor networks), external validation sets, or significant computational resources, limiting their scalability and practicality for large-scale, uncurated collections. There is a need for a lightweight, self-evolving method that can adapt dataset mixtures during training without external supervision or heavy computational overhead.

Methodology: DYNAMIXSFT

The authors propose DYNAMIXSFT, a dynamic and automated framework for optimizing instruction-tuning dataset mixtures. The approach is built on three core components:

1. Multi-Armed Bandit (MAB) Formulation

The problem of allocating training samples across $K$ heterogeneous data sources is formulated as a Multi-Armed Bandit problem, where each dataset is treated as an "arm." The goal is to learn a sampling policy that maximizes model performance based on current training dynamics.

2. Prior-scaled Boltzmann Exploration

To address the non-stationary nature of rewards during fine-tuning while preserving the inherent diversity and coverage of the dataset collection, the authors introduce a Prior-scaled Boltzmann Exploration strategy.

• Mechanism: Unlike standard Boltzmann exploration which treats all rewards as equally comparable, this method incorporates the original dataset proportions $p^{(0)}$ (often based on dataset size) as a prior.
• Formula: The sampling probability $p_k$ for dataset $k$ is calculated as:
  $$p_k = (1-\gamma) \cdot \frac{\exp(\beta \cdot Q_k) \cdot p^{(0)}_k}{\sum_j \exp(\beta \cdot Q_j) \cdot p^{(0)}_j} + \frac{\gamma}{K}$$
  Where:
  • $Q_k$ is the estimated reward (value) of the dataset.
  • $\beta$ is a sharpness factor controlling exploitation.
  • $\gamma$ is a uniformity factor ensuring a minimum floor probability ($\gamma/K$) for all datasets to prevent the "never-sampling" problem and guarantee exploration.
• Effect: This "softly anchors" the learned sampling distribution to the original proportions, allowing for adaptive updates while maintaining broad domain coverage.

3. 1-Step Look-ahead Reward

To update the sampling probabilities, DYNAMIXSFT employs a lightweight 1-Step Look-ahead Reward.

• Mechanism: At fixed update intervals, the model performs a virtual one-step gradient update on a mini-batch from each dataset. The reward $r_k$ is computed as the relative reduction in loss:
  $$r_k = \frac{1}{|B_k|} \sum_{x \in B_k} \frac{L_{pre}(x) - L_{post}(x)}{L_{pre}(x) + \epsilon}$$
• Advantage: This provides an immediate signal of a dataset's utility based on the model's current state without requiring additional supervision, validation sets, or permanent parameter updates.
• Smoothing: To handle noise and non-stationarity, rewards are maintained via an Exponential Moving Average (EMA): $Q_k \leftarrow \alpha \cdot Q_k + (1-\alpha) \cdot r_k$.

Key Contributions

1. Framework Proposal: The introduction of DYNAMIXSFT, a dynamic dataset mixture optimization method that formulates the problem as an MAB with a Prior-scaled Boltzmann Exploration policy.
2. Lightweight Adaptation: The development of a 1-Step Look-ahead Reward mechanism that enables adaptive sampling based on the model's own learning progress, requiring no external models or validation data.
3. Empirical Validation: Application of the method to large-scale collections (TÜLU-2 with 16 datasets and TÜLU-3 with 19 datasets) across 10 benchmarks, demonstrating consistent improvements on LLaMA3.2 1B and LLaMA3.1 8B models.
4. Comprehensive Analysis: Detailed visualizations and ablation studies revealing how mixture proportions evolve over time and the critical roles of prior scaling, exploration/exploitation balance, and reward smoothing.

Experimental Results

The authors evaluated DYNAMIXSFT on LLaMA3.2 1B and LLaMA3.1 8B models using the TÜLU-2 and TÜLU-3 mixture collections.

• Performance Gains: DYNAMIXSFT consistently outperformed static baselines (Full Coverage, Uniform Sampling) and existing dynamic methods (MultiDDS, MultiUAT, HBO).
  • On TÜLU-2, it achieved relative improvements of up to +5.1% (1B) and +5.3% (8B) averaged across 10 benchmarks compared to naive proportional sampling.
  • It surpassed HBO (which requires an additional actor network) despite DYNAMIXSFT's lighter architecture.
• Efficiency: The method introduced only ~12.7% additional computational overhead over naive sampling. In contrast, prior dynamic methods incurred overheads ranging from +139% to +760%. This efficiency is attributed to the 1-Step Look-ahead Reward, which relies solely on forward passes.
• Ablation Insights:
  • Prior Scaling: Removing the prior $p^{(0)}$ led to a substantial performance drop, confirming that anchoring to original proportions is crucial for stability and coverage.
  • Hyperparameters: Optimal performance required a balance between exploitation sharpness ($\beta$) and exploration uniformity ($\gamma$), as well as moderate EMA smoothing ($\alpha$) to prevent volatile allocation.
  • Update Interval: Updating the policy every 50 steps yielded the best balance between responsiveness and stability.

Significance and Claims

The paper claims that DYNAMIXSFT enables LLMs to optimize their own dataset mixtures during post-training without the need for hand-crafted tuning or external supervision. By treating dataset mixture optimization as a self-evolving process, the method makes it feasible to leverage large-scale, heterogeneous dataset collections in a more principled and automated manner.

The authors emphasize that their approach offers a lightweight, self-contained solution that avoids the computational costs of proxy models or actor networks while maintaining the diversity and coverage inherent in the original dataset collections. The work provides practical insights into the adaptive dynamics of training, showing that optimal data mixtures are state-dependent rather than static, and that dynamic adjustment based on immediate learning signals can lead to superior generalization across knowledge, reasoning, coding, and mathematical tasks.