Predicting Training Re-evaluation Curves Enables Effective Data Curriculums for LLMs

1. Problem Statement

Large Language Model (LLM) training increasingly relies on data curriculums, where high-quality (HQ) or domain-specific data (e.g., math, code) is upsampled during specific phases of pre-training (often called "annealing" or "mid-training").

• The Challenge: The optimal timing for introducing this HQ data remains unclear. The prevailing heuristic is to place HQ data at the very end of training when the learning rate (LR) is lowest, assuming this maximizes retention.
• The Gap: There is no principled, predictive framework to determine when a model is most receptive to specific data types. Existing methods rely on costly ablation studies (training multiple models with different data placements) or flawed heuristics that may lead to suboptimal performance or catastrophic forgetting.

2. Methodology: Training Re-evaluation Curves (TREC)

The authors introduce the Training Re-evaluation Curve (TREC) as a diagnostic tool to measure data retention and guide curriculum design.

Definition of TREC

A TREC measures how well a fully trained model ($\theta_T$) performs on training batches ($B_t$) as a function of when those batches were encountered during training ($t$).

• Metric: $L_{re}(t) = L(B_t; \theta_T)$, where $L$ is the cross-entropy loss.
• Interpretation: A lower TREC loss at a specific training fraction $\hat{t}$ indicates that the model retains the information from that phase better. Conversely, a high TREC loss implies the model has "forgotten" or drifted away from the optimal parameters for that specific data.
• Observation: Under standard learning rate schedules (e.g., step decay), the TREC often bottoms out before the end of training, contradicting the heuristic that "end-of-training" is always best.

Predictive Framework

Since calculating a TREC requires training a model to completion (which is expensive), the authors propose a method to predict TREC shapes in advance without retraining.

1. AdamW EMA Perspective: They leverage the insight that AdamW parameters can be viewed as an Exponential Moving Average (EMA) of weight updates. The influence of a past update on the final model is governed by an EMA timescale ($\tau$).
   • $\tau = \frac{1}{\eta \lambda T} = \frac{B}{\eta \lambda D}$
   • Where $\eta$ is learning rate, $\lambda$ is weight decay, $B$ is batch size, $D$ is total tokens, and $T$ is total steps.
2. Minimizer Drift: The authors argue that early gradients lose effectiveness not just because of EMA decay, but because the "local minimizer" for a specific batch drifts over time as the model learns new features.
3. Predictive Model: They formulate the normalized TREC shape ($\hat{L}_{re}$) as:
   $$\hat{L}_{re}(\hat{t}) = 1 - c(\hat{t})^p \cdot \hat{t}^m$$
   • $c(\hat{t})$: The EMA coefficient derived from the LR and weight decay schedule.
   • $\hat{t}$: The training fraction ($0$ to $1$).
   • $m$: A training-fraction exponent that captures the rate of minimizer drift.
   • $p$: A scaling exponent (fixed to 0.5 in experiments).

The exponent $m^*$ is predicted via a power law based on Tokens-Per-Parameter (TPP) and the timescale $\tau$:
$$m^* = C \cdot (\text{TPP})^{\mu_1} \cdot (\tau)^{\mu_2}$$

3. Key Contributions

1. Introduction of TREC: A diagnostic metric that retrospectively evaluates training data retention, revealing that optimal data placement often occurs before the final training steps (in the "TREC valley").
2. Predictive Theory: Demonstrated that TREC shapes are governed primarily by the AdamW timescale ($\tau$) and can be predicted analytically using EMA coefficients and a training-fraction drift term.
3. Large-Scale Empirical Validation: Conducted experiments on models ranging from 111M to 3.9B parameters and datasets from 20 to 1280 TPP, showing that TREC shapes align across scales when $\tau$ and TPP are matched.
4. Curriculum Optimization: Proved that placing HQ data at the TREC minimum yields significantly better validation performance than placing it at the end of training or using uniform mixing.
5. Explanation of Prior Work: Used TREC analysis to explain why certain existing recipes (e.g., Llama-3 405B's GSM8k annealing) failed (the LR was too low/short, causing the TREC to rise before the data was introduced) and why others succeeded.

4. Key Results

• Optimal Placement: In experiments with a 610M model, placing high-quality code data at the TREC minimum (often ~70-80% through training, depending on the LR schedule) resulted in the lowest validation loss. Placing it at the very end (standard practice) was suboptimal.
• Predictive Accuracy: The analytical model successfully predicted TREC shapes across different model sizes (111M to 3.3B) and hyperparameter settings with high correlation ($r_p \approx 98\%$), even when the exact exponent $m^*$ prediction had slight errors.
• Scaling Laws:
  • Timescale ($\tau$): The primary driver of TREC shape. Higher $\tau$ (slower decay) shifts the TREC valley earlier.
  • TPP: Higher TPP (more training data per parameter) slightly shifts the TREC valley later and reduces the magnitude of the loss drop (indicating reduced memorization/plasticity).
  • Batch Size: TRECs remain stable up to the critical batch size but diverge significantly beyond it.
• Continual Pre-Training (CPT): In a 3.9B model CPT experiment, placing math data at the predicted TREC minimum (mid-CPT) outperformed placing it at the end, across various learning rates.
• MoE Architectures: Sparse Mixture-of-Experts (MoE) models behave as if they have a lower effective TPP per expert, shifting their TREC valleys earlier, which aligns with empirical observations of increased memorization in MoEs.

5. Significance and Impact

• Paradigm Shift: Moves data curriculum design from "heuristic trial-and-error" to a principled, predictive science. Practitioners can now determine the optimal data insertion point before training a massive model.
• Cost Efficiency: Eliminates the need for expensive ablation studies (training multiple full-scale models with different data mixes) to find the best curriculum.
• Explaining Failures: Provides a theoretical explanation for why some state-of-the-art models (like Llama-3 405B) failed to benefit from specific annealing phases (the TREC was rising, meaning the model had already "forgotten" the context needed to retain that data).
• Generalizability: The framework is applicable to various optimizers (with implicit EMA forms) and scales, offering a robust tool for optimizing frontier-scale LLM training.

In summary, the paper establishes that data retention is a function of training dynamics (specifically the AdamW timescale and minimizer drift) rather than just data quality. By predicting where the model is most "receptive" (the TREC valley), researchers can strategically place high-value data to maximize downstream performance.