Summary of REGMIX-D: Dynamic Data Mixing via Proxy Training Trajectories

Summary (Overview)

• Proposes REGMIX-D, a simple extension of the RegMix framework that enables dynamic data mixture scheduling during LLM pretraining, rather than using a single static mixture.
• Key insight: Proxy training runs produce full loss trajectories, not just endpoint losses; these trajectories can train a regression model to predict future losses conditioned on current training state (step, mixture, loss).
• Two deployment variants: Offline (generates complete schedule before target training via recursive predictions) and Online (adapts mixture during training using observed target losses with a cross-scale correction).
• Results: Outperforms RegMix and DoReMi across 13 downstream tasks with a 1B-parameter target model on 25B Pile tokens; achieves gains even with only 128 proxy models (25% of RegMix's compute budget).
• Robustness: Gains persist across ablations on switch-point count (N), scaling factor (β), and proxy model size.

Introduction and Theoretical Foundation

• Problem: Data mixture selection is critical for LLM pretraining performance, but existing methods (RegMix, DoReMi) assume a single optimal static mixture for the entire training run.
• Motivation: The optimal domain composition likely shifts during training—early stages may benefit from diverse general data, while later stages may need domain-specific data. Static mixtures therefore compromise across training phases.
• Existing dynamic methods (Aioli, TiKMiX) require additional machinery (online optimization, gradient influence, RL), adding overhead to target training.
• REGMIX-D's approach: Reuses the regression framework from RegMix but extends it to predict local transitions $(t_\ell, m^{(i)}, \ell_\ell^{(i)}) \to \ell_{\ell+1}^{(i)}$, enabling dynamic scheduling without any extra machinery during target training.

Methodology

Trajectory-Conditioned Regression

• For each proxy run $i$ with mixture $m^{(i)}$ and validation losses $\{\ell_1^{(i)}, \ldots, \ell_{N+1}^{(i)}\}$ at intervals $\{t_1, \ldots, t_{N+1}\}$, form training pairs: $$\underbrace{\left(t_\ell, m^{(i)}, \ell_\ell^{(i)}\right)}_{\text{input}} \longrightarrow \underbrace{\ell_{\ell+1}^{(i)}}_{\text{target}} \tag{1}$$
• This yields $M \times N$ training pairs; the model $f_N$ learns to predict the next-interval loss given current step, mixture, and observed loss.

Offline Mixture Schedule (Algorithm 1)

• Starts with human-curated Pile prior $m_0^{\text{human}}$ for the first segment.
• At each switch point $s_j$, searches for mixture $m_j^*$ minimizing $f_N(s_j, m, \ell_{j-1})$ over Dirichlet-sampled candidates.
• Recursive prediction: predicted loss $\ell_j \gets f_N(s_j, m_j^*, \ell_{j-1})$ feeds into the next query.
• Maps proxy steps to target steps via $s_j^{\text{target}} = s_j \cdot (T_{\text{target}}/T_p)$.

Online Dynamic Mixture (Algorithm 2)

• Uses the same frozen regression model $f_N$ but replaces predicted losses with observed target-model losses at each switch point.
• Cross-scale loss correction: Since $f_N$ is trained on proxy-scale losses (model size $P_p$) but target losses come from a larger model (size $P_t$), applies: $$\tilde{\ell}_{S_j} = \hat{\ell}_{S_j} \cdot (P_t/P_p)^{\beta} \tag{2}$$ with $\beta = 0.05$ (robust across sweep $\beta \in \{0.01, 0.02, 0.05, 0.10\}$; performance varies < 0.25).

Empirical Validation / Results

Main Results (Table 1)

• Setup: 1B TinyLlama target on 25B Pile tokens; 1M-parameter proxy; 5 switch points (N=5); compared against Human, DoReMi, RegMix.

• Key findings:
  • REGMIX-D (128) with only 25% proxy compute surpasses RegMix (512).
  • Online > Offline in both proxy budgets (grounding in target dynamics improves accuracy).
  • Gains are broad (top-1 or top-2 on 11/13 tasks) rather than concentrated in one task category.
  • Compute scaling (128→512 proxies) yields modest gains, suggesting the dynamic-vs-static gap dominates.

Ablation Study (Table 3, Appendix)

• All 26 configurations outperform RegMix (47.41), showing robustness.
• Switch points N: N=5 is best; N=3, 7, 9 also reliably surpass RegMix.
• Scaling factor β: Max Avg variation 0.25, justifying treating it as a fixed constant.
• Proxy size: 1M vs. 120M proxies yield nearly identical performance (within 0.05 Avg).

Theoretical and Practical Implications

• Theoretical insight: Loss trajectories contain richer information than endpoint losses; exploiting this enables dynamic scheduling without additional optimization machinery.
• Practical significance:
  • Cost efficiency: Dynamic mixing gains achievable at 25% of RegMix's proxy compute.
  • Deployment flexibility: Offline mode requires no changes during target training; Online mode adds minimal overhead (just loss measurement and mixture lookup at switch points).
  • Robustness: Performance is insensitive to hyperparameter choices (N, β, proxy size), lowering tuning burden.

Conclusion

• REGMIX-D demonstrates that effective dynamic data mixing can be achieved via simple regression on proxy loss trajectories, with both offline and online variants.
• It outperforms static baselines (RegMix, DoReMi) across 13 downstream tasks and is data-efficient (128 proxies suffice to beat RegMix with 512).
• Limitations: Single-domain optimization signal (pile-cc); fair comparison with other dynamic methods (Aioli, TiKMiX) is challenging due to different corpora/setups; future work may extend to multi-target objectives and broader benchmarks.