# RegMix-D: Dynamic Data Mixing via Proxy Training Trajectories

> REGMIX-D learns dynamic data mixture schedules from proxy loss trajectories, outperforming static baselines on 13 tasks while using only 25% of RegMix's compute budget.

- **Source:** [arXiv](https://arxiv.org/abs/2606.18663)
- **Published:** 2026-08-29
- **Permalink:** https://picx.dev/p/D7rWpX
- **Whiteboard:** https://picx.dev/p/D7rWpX/image

## Summary

# 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.

| **Method** | **Avg Score (13 tasks)** |
|------------|--------------------------|
| Human      | 44.80                    |
| DoReMi     | 47.10                    |
| RegMix     | 47.41                    |
| RegMix-D (128) Offline | **47.93** |
| RegMix-D (128) Online  | **48.15** |
| RegMix-D (512) Offline | **48.09** |
| RegMix-D (512) Online  | **48.44** (best) |

- **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.

---

_Markdown view of https://picx.dev/p/D7rWpX, served by PicX — AI-generated visual whiteboard summaries of research papers._