Predicting Training Re-evaluation Curves Enables Effective Data Curriculums for LLMs
Authors: Shane Bergsma
Nolan Simran Dey
Joel Hestness
Venue: ICLR 2026 Poster (2026)
Keywords: training re-evaluation curve
data curriculum / data placement
large language model (LLM) pre-training
AdamW EMA timescale
learning-rate schedules
tokens-per-parameter ratio
Primary area: foundation or frontier models, including LLMs
TL;DR: We evaluate fully-trained LLMs on their original training data, measuring retention across steps; a predictive model of the resulting "re-evaluation curve" identifies optimal spots for high-quality data, surpassing default end-of-training placement.

Abstract: Data curriculums have become central to successful LLM training, yet principles governing optimal data placement remain unclear. We introduce the *training re-evaluation curve (TREC)*, a diagnostic that retrospectively evaluates training batches *using the final model weights*. The TREC characterizes how well a trained model retains training data as a function of *when* the data was encountered during training. Analyzing TRECs for models from 111M to 3.9B parameters, we show that placing high-quality data at low points on the TREC significantly improves performance. Importantly, while a TREC is initially observable only after training, we demonstrate it can be *predicted in advance* from AdamW’s implicit EMA coefficients, enabling proactive curriculum design. By predicting TRECs for published training recipes, we explain prior ablations and reveal suboptimal data placements. We also align high-quality data with TREC minima in order to improve continual pre-training of a 3.9B-parameter LLM trained on 900B tokens.

Published as a conference paper at ICLR 2026
PREDICTING TRAINING RE-EVALUATION CURVES EN- ABLES EFFECTIVE DATA CURRICULUMS FOR LLMs
Shane Bergsma, Nolan Dey & Joel Hestness Cerebras Systems
(shane.bergsme bergsma, nolan, Joel Jecerebras..et
ABSTRACT
Data curriculums have become central to successful LLM training, yet princi- ples governing optimal data placement remain unclear. We introduce the training re-evaluation curve (TREC), diagnostic that retrospectively evaluates training batches using the final model weights. The TREC characterizes how well trained model retains training data as function of when the data was encountered dur- ing training. Analyzing TRECs for models from 11 IM to 3.9B parameters, we show that placing high-quality data at low points on the TREC significantly im- proves performance. Importantly, while TREC is initially observable only after training, we demonstrate it can be predicted in advance from AdamW's implicit EMA coefficients, enabling proactive curriculum design. By predicting TRECs for published training recipes, we explain prior ablations and reveal suboptimal data placements. We also align high-quality data with TREC minima in order improve continual pre-training of a 3.9B- L.9B-parameter LLM trained on 900B tokens.
TREC predicts best data place
TRECs weight decay. varies
TREC's align
scale (same
LR drops 100x
1.0
3.4
Train
3.0
TREC
0.8
265M
data pos
Normalized TREC
2.8
610M
3.0
0.6
1.78
2.8
X
2.6
3.38
0.4
2.6
+0.0000
0.1
2.4
2.4
0.2
0.0125
0.4
2.2
2.2
0.0
0.05
1.0
0.00 0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
Ordered training tokens x10
training tokens,
Fraction training tokens,
Figure Left: (610M parameters): train loss steadily falls after learning rate drops to 0.01 (after 70% of steps), but the optimal position for placing high-quality (HQ) data is before the LR drop, the TREC valley. Middle: (610M, linear LR decay): TREC shape varies with AdamW timescale (varied via weight decay A). Right: (size varies, linear LR decay-to-zero, all 20 tokens- per-parameter):: TRECs align across 1000 scaling of training compute, when 7 matches.
INTRODUCTION
LLM training now often includes mid-training or annealing phases where special data is upsampled in the final stages of pre-training. These data can be high-quality (Shen et al., 2024), recent (Dubey math (OLMo et al.. 2024), code (Zhang et al., 2024a) or
et al., 2024) or domain-specific-e instructions (Hu et al., 2024). It is assumed that placing such data at the end of training, when learning rate is near zero, maximizes its effectiveness. While some work explores when this phase should begin (Feng et al., 2024; Parmar et al., 2024; Liu et al., 2025), "many interesting questions remain around finding the optimal dataset distribution for pre-training" (Anil al., 2023).
We introduce the training re-evaluation curve (TREC), a diagnostic for understanding LLM training and data placement. A TREC measures how well fully trained model performs training batches

---

Published as a conference paper at ICLR 2026
as function of when those batches were seen. Defined over homogeneous data, construct it by evaluating the final model on the i..d. training sequence, in order. If data were retained equally, the TREC would be flat; in practice, models perform better on tokens from specific points in training.
Under Step-decay learning rate (LR) schedules, c.g., as used (with multi-step-decay) in DeepSeck LLM (Bi et al., 2024), the TREC may bottom-out well before the end of training (Fig. left), while the standard training curve (on unseen batches) continues to decline. We hypothesize that high- quality (HQ) data should be placed where the TREC is lowest-that is, where the final model would achieve its lowest retrospective loss. Indeed, when retraining different models, each with the same HQ data inserted into different 10% segment of the training trajectory, placing HQ data where the TREC is lowest yields the best validation loss (loss-by by-semeent plotted as blue points in Fig. 1, left).
Of course, it is impractical train a very large model, measure its TREC, and then re-train with a TREC-inforned curriculum. Fortunately, we show the TREC is predictable, enabling proactive cur- riculum design. For Adam W the dominant optimizer in LLM pre-training TRECss are governed by the EMA timescale (Sec. 3): AdamW parameters can be viewed as an EMA over weight updates (over data), so timescale T controls data influence across training, and thus where TREC perfor- mance peaks (Fig. 1, middle). Sweeping learning rate, weight decay, or batch size with matching timescales yields identical TREC shapes (Fig. 3). Shape also persists across scale: despite x more training FLOPs, 3.3B model matches a 11 IM model's shape (Fig. 1, right). In Sec. 4, we formalize a predictive model of TRECs based on an expanded view of the Adam timescale-one that handles arbitrary LR schedules, including Step drops.
Based on our work, practitioners can use TRECs to guide data ordering, avoiding flawed heuristics and costly ablations. More specifically, our main contributions are:
Introducing train re-evaluation curves, showing they predict optimal data placement for a given optimizer configuration (Sec. 2).
A large-scale study of 600 TRECs in models from 111M to 3.9B parameters, and datasets from 20 to 1280 (Sec. 3), and enables an analytical model for predicting curves in advance of training (Sec. 4). Explaining findings in sparse mixture-of-expert
(TPP). The study connects TRECs to the AdamW timescale
(Sec. 5.1) and prior data curriculum work (Sec. 5.2). For example, TREC prediction can explain why Llama-3 405B did not benefit, on GSM8k validation, from annealing on the GSM8k training set.
Leveraging TRECs to improve a 3.9B-parameter LLM trained on 900B tokens (Sec. 5.3).
2 TRECs PREDICT EFFECTIVE DATA PLACEMENT
We now formally define TRECs, state our key hypothesis, and describe its evaluation.
Definition 1 (TREC). Let B By denote the sequence of batches used during training, drawn independently and identically from data distribution DD and let 0r represent the fully-trained model parameters. The training re-evaluation curve (TREC) is the sequence of scalar loss values: (..(() == C(B; Or), for t = 1,.... T. where C(B; 0r) denotes the loss (c.g., cross-entropy) evaluated on training batch B3 from step t using final parameters Or. Intuitively, lower TREC loss suggests greater alignment with 0r. and may indicate B, contributed more significantly to the final model.
It is worth clarifying why re-evaluation loss may depend on order, even when training batches are drawn i.i.d. In any online optimization process, each batch is encountered at different point along the parameter trajectory, with a different amount of subsequent parameter evolution. Early batches are learned when the parameters are far from their final values, and updates derived from these batches lose their effect as the parameters subsequently evolve. The result is higher re-evaluation loss on these earlier batches: the model performs on these forgotten data similarly to how it performs on unseen validation samples. This is related to catastrophic forgetting in continual learning (Mc- Closkey & Cohen, 1989; Kirkpatrick et al., 2017): earlier examples can be forgotten, even when the training distribution does not change. In Appendix J, we formalize this intuition and show that TREC loss can be analytically predicted as a function of parameter evolution under a simplified quadratic model.
Data), Given training distribution Dorig a distribution Dıo is high- quality for task if replacing Dorig samples with Dnq samples improves task performance.
Definition 2
2

---

Published as a conference paper at ICLR 2026
Schedule
Post- Post- warmup decay v7(1)
Used in
Step drop at 70% training (Step decay)
v
Mmaxx
t t< 0.77
Fig. (lef). Fig.
0.01 T/maxi
12 0.71
(left)
Linear decay 0. D.Jax (10x decay)
n(t) O. 0.17max +0.97mmax
T-
Fig. (right)
Linear decay to zero (D2Z decay)
(7()) Tmax
1-w
Fig. 11 (left)
7-27
Table 1: Learning-rate schedules used in placement tests. All schedules share linear warmup over
0.17 steps. Appendix Fig. 11 (right) visualizes the LR curves.
The canonical TREC is defined under homogeneous i..d, sampling from a base distribution D. We average across batches when plotting TRECs in order isolate the effect of training position from any variation in intrinsic batch difficulty. We hypothesize the resulting (t) primarily reflects optimization dynamics, rather than the specific DD used to compute it. That is, given fixed optimizer configuration, TREC characterizes how strongly data encountered at each training position is retained by the final regeleregardless of what data actually inserted at those positions, or even if different data (e.g., high-quality data) is used heterogeneously across positions. Our belief in the effectiveness of this positional ranking leads to our main hypothesis:
Hypothesis 1: Placing a fixed number of high-quality samples at steps where Lre(!) lowest maximizes performance on target lask.
Experimental setup. Experiments use 610M parameter GPT2-style LLMs trained with Adam W, μP (Yang et al., 2021), and learning rate warmup over 10% of steps. After warmup, LR follows one of the decay schedules in Table 1. Full architecture and tuning details are in Appendix C. Models are trained on SlimPajama splits (Soboleva et al., 2023) (blend weights in Table 2), using the general blend (GB) DOrig and code blend (CB) as DHQ, with CB validation loss the target task. We use 45B (90%) GB tokens and 5B (10%) CB tokens. We also construct an aggregate blend (AB) with uniform 90/10 GB/CB mix. All models train 50B tokens total (82 TPP) with no repetition. Train and TREC loss curves are computed on homogeneous AB batches.
Table 2: SlimPajama mixes used in Fig. 2 placement tests: Gen- eral Blend is the original distri- bution, Code Blend is the HQ data.
Step decay (step at 70% of training)
10x decay
loss
3.4
TREC
3.4
TREC
3.2
data pos
3.2
loss ve pos
-- Val loss using agx blend
3.0
2.8
2.8
SlimPJ Subset
General Blend
Code Blend
Loes for this pos
Commcrawl
55.1
37.1
General bland
Code
C4
28.2
19.0
GitHub
0.0
30.0
Ordered training tokens
x10
Ordered training tokens X10
Books
4.4
3.0
ArXiv
4.9
3.3
Figure 2: TRECs predict best placement. Left: Example place-
Wikipedia
4.0
2.7
ment curriculum 5-of-10 10 for the Step decay LR schedule. Right:
StackExch.
3.5
5.0
Results for linear decay to 0.17maxx
To evaluate Hypothesis 1, we conduct data placement sweep: for each LR schedule, we train 10 models, each with the 5B CB tokens inserted into a different 10% segment of the training data. example curriculum shown in Fig. 2, left. This figure also plots the standard training curve (loss on unseen batches during training). While the train loss shows familiar decrease with LR annealing (Tissue et al., 2024; Schaipp et al., 2025), note the TREC increases after the LR drops.
Results. Results in Fig. 1 (left, Step decay schedule). Fig. (right, linear decay to 0.17maxx) and appendix Fig. 11 (linear decay to zero) show, in all cases, placing CB data at the step with lowest (t)(t) yields the best validation loss. Data placement also outperforms uniform training on AB
3

---

Published as a conference paper at ICLR 2026
(appendix Fig. 12). TRECs do not merely rediscover the put HQ data at the end heuristic used in prior work (with varying effectiveness, see 5.2). Looking only at the standard loss curve for Step decay (Fig. 1, left), might think it best "anneal" on high-quality data during the low-LR region where training loss drops rapidly. Our placement experiments show this intuition is incorrect: placing data at the TREC minimum yields better validation performance.
Under the decay-to-zero schedule (Fig. 11), the TREC minimum is likewise not at the literal end of training; the curve turns upward the LR approaches zero. If arbitrary placement were allowed (rather than fixed 10% segments), the optimal window would occur slightly before the final steps. In practice, HQ budgets are typically far smaller than 10% of pre-training tokens, making this dis- tinction even more consequential: sub-segments earlier within the final decile can outperform end placement. Consistent with this observation, Sec. 5.3 shows that in CPT under linear decay-to-7ent mid-CPT placement outperforms end placement across three peak learning rates.
Key takeaway 1: TRECs reliably indicate the best point for high-quality data insertion within a given optimizer configuration (1₁(t,, X, B3.
Interestingly, we find that absolute (t) values only weakly generalize across different optimizer consigurations for placement prediction. Appendix D.2 investigates the predictiveness of TREC loss for placement across configurations, and discusses a grokking-ss-memorize ve-memorization view that may explain why configurations with similarly low TREC loss can differ in validation performance.
3 TREC SHAPE IS GOVERNED BY THE ADAMW TIMESCALE
To apply the insights of Sec. (without having to train model twice), we need to predict TRECs in advance. Before presenting predictive model (Sec. 4), we build intuition for what controls TRECs (under a fixed LR schedule), finding they are mainly governed by the AdamW timescale.
Background: the AdamW EMA and its timescale. AdamW (Loshchilov & Hutter, 2017) updates at step can be expressed in terms of learning rate 1) and weight decay À as: O, moving aver-
(1 n)))6,-1 1) Vorte where my and 0, are bias-corrected exponentially-weigen ages (EMAs) of gradients and squared gradients, respectively (Kingma & Ba, 2014).
Wang & Aitchison (2024) observed that Adam parameters 0, can also be viewed an EMA- of weight updates. Specifically, the standard EMA form y/+ = (1 - 0)) + am; matches Adam when y/ = 8,, a = ndλ, and xe = - me Aitchison, represents the approximate number of iterations over which updates are averaged.
Vo,+e* The timescale 1/a = 1/1>, denoted Titrr by Wang &
When expressed in epochs as Tepoch = Wang & Aitchison found the optimal Tepach remained stable under model and dataset scaling in image tasks. Maintaining a constant Tepoch requires decreasing À proportionally when M increases.
Tite/M, where M is the number of iterations per epoch,
Since LLM pre-training typically uses single epoch, we follow Dey et al. (2025) and Bergsma et al. (2025a) in defining normalized timescale 7 = Titer/ /T, where is the total number of optimization steps. As T D/ B (total tokens D) divided by batch size B):
1
B
MAD
Hypothesis 2:
For a given learning-rate rate decay schedule (Linear, Cosine, Step, etc.), the TREC is controlled by the AdamW W timescale
In other words, because parameters in AdamW are implicitly weighted averages over updates (de- rived from training data), the EMA timescale 7 governs the scope of data influence the final model, and thus the TREC. Yet higher EMA weight alone may not lower TREC loss: earlier updates can lose influence our position on the loss surface shifts (discussed more in Sec. 4).
Training fraction. Viewing training in terms of discrete optimizer steps becomes limiting when batch sizes and sequence lengths vary. Instead, we view optimization continuous stochastic pro- cess over a fixed dataset, with different batch sizes yielding different discretizations. In this perspec- tive, plot TRECs against training fraction , 1/T t.B D; this naturally facilitates comparing
4

---

Published as a conference paper at ICLR 2026
across different batch sizes, step counts, and dataset sizes. To reduce noise in small-batch settings (and align with large-batch models), we also smooth TRECs using a noving-averagg filter (typically over window of 100 steps), smoothing curves without altering the underlying trajectory.
Results. We follow the architecture and setup of Sec. 2, but apply: (a) linear decay-to-zero LR schedule, (b) a context length of 2048, and (c) the original GPT2 vocabulary (50257 tokens). We train on standard SlimPajama splits with original source weightings. Plot axes indicate whether the TREC is shown in absolute loss or normalized (min-max scaled) form.
varies y=1.62e02 0=0.1
B=504 y varies A=0.1
B=504 ny=1.22-02
varies
1.0
1.0
1.0
Normaizzd TREC
0.8
Nomalzzd TREC
0.8
Normalizdd TREC
0.8
0.6
B (v)
# (+)
A (v)
0.4
(0.026)
0.4
-4.05e-03 (0.421)
0.4
0.025 (0-421)
252 (0.053)
8 05.000 (0.210)
0.05 (0.210)
0.2
504 (0. 105]
0.2
62e-00 02 (0.105)
0.2
0.1 (0 106)
1008 (0 (0:210)
3.24e 02 (0.053)
0.2 (0 (0.053)
0.0
(0:421)
0.0
(0.026)
0.9
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
Fraction of training tokens,
2
Fraction training tokens,
Fraction of training tokens,
Figure 3: Timescale determines TREC shape (610M, 80 TPP). Sweeping B (left), 1 (middle),
or À (right) produces matching variations in TRECs when T (Eq. (1)) varies identically.
Fig. 3 shows normalized TRECs for 610M models trained to 80 TPP, sweeping 7) X, or B in each subplot. Across hyperparameters, curves with matching exhibit very similar shapes, reflecting consistent timescale control. This aligmment is especially clear here since all models use the same ordered training data. Similar patterns hold across other scales and dataset sizes. Generally, as timescale expands (7 increases), the TREC minimum (valley) shifts earlier!
We next examine TREC alignment as models and datasets scale. When models share the same and TPP, their TRECs largely align (Fig. 4, left), with small differennces causing corresponding shifts. The size of absolute TREC drops is also similar across scales when training to the same TPP (Fig. 1, right). As TPP increases, the TRECs move slightly right (Fig. 4, middle, right); e.g., at 111M, increasing TPP by 64x moves the TREC valley slightly later in training. Interestingly, as TPP increases, the size of the TREC drop diminishes (appendix Fig. 15): at lower TPP, training appears to emphasize memorization. This finding aligns with a recent study of LLM memorization from an information theoretic perspective (Morris al., 2025).
Finally, Appendix E.3 shows that large changes in AdamW's B1 and B2 parameters do not signif- icantly alter TREC shape (even when 31 0, i.c., no momentum). This suggests the timescale imposed by weight decay is far more influential than the timescales of momentum and velocity.
Key takeaway 2: Overall, the data broadly supports Hypothesis 2:
the AdamW timescale (T)
predominantly controls TREC shape, with TPP playing secondary, smaller role.
However, seen in Fig. 2, the LR schedule itself significantly affects TREC shape. To integrate LR schedules into our analysis, we next adopt an expanded view of the Adam timescale.
4 PREDICTING TRECs: ADJUSTING FOR TRAINING FRACTION
Following Sec. one could in principle predict a TREC for any model by re-evaluating a smaller model with matching т, TPP, and LR schedule. More generally, we propose functional form that extends the Adam W EMA perspective to time-varying learning rates, while also incorporating a training-fractin fraction term to capture minimizer drifi, described below.
'Tii shift has limit: as X + 0 (r + oo), the curves converge (Fig. 1, middle). This occurs because the shape of the EMA coefficients (Sec. 4) converges the shape of the I.R schedule (Appendix F.3).
5

---

Published as a conference paper at ICLR 2026
111M-3.38 3B, 20TPP, v9:0.3
111M,
021
610M 7=0.105
1.0
1.0
1.0
Normalized TREC
0.8
Normaizzed TREC
0.8
Normalized TREC
0.8
0.6
N (r)
0.6
0.6
0.4
(0.333)
0.4
0.4
266M
20
TPP
0.2
610M (0.211)
0.
0.2
1.78 (0.298)
520
80
0.0
3.38 (0.307)
0.0
1280
0.0
320
0.00
0.25
0.30
0.75
1.00
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
Fraction of training tokens, 2
Fraction training tokens.
Fraction of training tokens,
Figure 4: Timescale and TREC shape across model/dataset scales. Left: Similar yields similar TREC shapes across model scales when training to 20 TPP (T z 0.3). At 111M (middle, 0.021) and 610M (right, 0.105), increasing TPP shifts TRECs slightly right.
Background: the extended Adam EMA perspective. Bergsma et al. (2025b) extend Wang & Aitchison (2024) by considering EMAs with time-varying smoothing parameters a E [0, 1). Setting a1 = 1 (so 21), they show the recursion = (1 a2 a112 + 02222, and in general:
y/2 =
Σ i=1
II (1 ag)) j=i+1
0,
Σ Cri
(2)
where Ce2 quantifies the contribution of input æ; the output Vr-
With AdamW, W, the LR/weight decay schedule defines the time-varying smoothing a, = 7J (Sec. 3). rh. The EMA operates over weight updates 2, = - Vortr* with larger or indicating the ith update contributes more to model weights Ye 0, at step t. We focus on coefficients for final model weights Ør, dropping subscript T' for clarity: G T7 11,-++14 n;A). To connect with our continuous-time framing, we interpret over 7 steps continuous function, reparameterizing as e(i) via = i/T, enabling direct comparison with TREC (i).
Cyclic EMA coefficients
Cyclic predicted TRECs
Cyclic normalized TRECs
Normalized coefficient values
1.0
n (r)
1.0
1.0
[0 (0.421)
Predicted TREC /REEC(scall) (scaledd
0.8
8.09e-03 (0.210)
0.8
Narmalized TREC
0.8
1.62e 02 (0.105)
0.6
3.24e-02 (0.053)
0.6
0.4
0.4
(=))
0.4
1
-.005ee (0.421)
4.05e 03 (0.421)
0.2
8 1.000 03 (0.210)
8.09e-03 (0:210)
6.e
(0.105)
1.62e 02 (0.105)
(0.053)
(0.063)
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
0.00
0.20
0.50
0.75
1.00
Fraction of training tokens,
Fraction of training tokens,
Fraction of training tokenss,
Figure 5: Predicting TRECs (610M, 80 TPP, Cyclic LR). Left: Normalized EMA coefficients c(f) in our model. Middle: Predicted TRECs (t) from Eq. (3). Right: True TRECs (t). Predictions match TREC dips, with early damping and late alignment with (inverted) EMA.
Fig. 5 (left) shows c(f) for eyclic LR schedule (schedule shown in Fig. 19). When e(f) drops to zero (i.e., LR is zero), TRECs (right) return to baseline (1.0); when c(t) is higher, TRECs dip lower. Yet c() influence fades earlier in training, suggesting e({) alone does not fully explain TREC shape.
Predicting TREC shape. EMA coefficients e(f) quantify each update's contribution to the fi- nal weights, but an update's s effectiveness can fade if the batch-specific loss surface shifts after the gradient computed, i.c., due to minimizer drift. Motivated by simplified quadratic analysis (Appendix J), we model drift on a training fraction clock that schedule-dependent, similar the of normalized compute in Qiu et al. (2025).
under jP yet LR-

---

Published as a conference paper at ICLR 2026
Hypothesis 3: TREC shape can be predicted using the AdamW EMA coefficients combined with an adjustment for training fraction.
Let (f) denote the TREC and c(f) the EMA coefficients of the final weights 6y, both indexed by training fraction 8 [0, 1]. We model the normalized TREC shape using the simple form:
()) c(ipp imm
where p and m are exponents to be fit. Exponent p controls the strength of the EMA contribution, while m (the training-racttinn exponent) determines when the predicted (f) begins to reflect c(t). For example, with m = 1, fluctuations in c(f) appear immediately; with a larger m, (t) remains near for most of training and only incorporates c(t) near the end.
Fig. illustrates this formulation: EMA fluctuations (leff) are dampened early in training in both predicted (middle, using fits of Eq. (3)) and true TRECs (right). We focus on predicting shape, not absolute values, and use Eq. (3) as a normalized functional form. We find tuning p has minor impact on shape prediction compared to m, so we fix p = 0.5 across experiments and focus on fitting m.
Predicting the training-fractionn exponent. For given set of TRECs, we define the optimal m as the one that maximizes shape agreement between the predicted and true TRECs. We use the Pearson correlation Tp as a scale- and shift-invariant measure of this agreement, finding it aligns well with visual assessments (Appendix F.I). Empirically, we find that optimal exponent m closely follows a power-law relationship with tokens celss--eepppe per-maareeterr (TPP) and the Adam timescale т, which we express
m = C . (TPP)* (7)
(4)
Fitting C, μl, and μl2 at small scale enables predict m² systematically across dodel/dataset sizes and hyperparameter settings, completing the components needed for TREC prediction.
Power law preds& & observed 111 data
Power law prods observed 610M data
Power law preds & abserved 1. 1.78 data
Optimal Train fic. exponent, m"²
TPP
TPP
TPP
60
Optinal Trainffcc. expoment,
20
Optimal Train exponent,
20
80
v 80
v 80
200
20
320
320
10
10
5
2
1.4
TPP6.1
Tppl:164
2
TPPl.:
10
10
10
30
10a
10-2
10
10
Figure 6: Fitted power law for m² aligns with observed optima. Left: 111M fit: middle: 610M eval; right: 1.7B eval. Accuracy of 111M fit holds across scale, degrading slightly for larger models.
Results. We follow the setup of Sec. 3, using a Linear decay-to-zen zero LR schedule unless noted. We evaluate prediction accuracy across scales and datasets using: (i) R for predicted vs. true m² values, and (ii) Pearson Tp for predicted vs. true TREC shape. We fit the power law Eq. (4) using m= values from small-scale 11 M models, trained across varied TPPs and timescales (Fig. 6, left).
Fits at 111M generalize to larger scales (Fig. 6, appendix Fig. 21). As shown in appendix Table 6, while fit R² declines from 99% at 111M 77% at 3.3B, TREC prediction accuracy remains high across scales (rp ~ 98%), confirming robust predictive performance even when vn* fits are imperfect. Ablation results (appendix Table 7, fitting 111M, evaluating at 610M scale) show that both TPP and especially 7 are important for accurate m prediction and TREC shape matching.
The observed dependence of m² on and TPP, and also LR schedule (Appendix G.4), matches the minimizer drift account in Appendix J. In the quadratic view, drift accumulates with the schedule- weighted curvature [ (7(s) h(s) ds; shorter EMA timescales shorten memory and increase drift (larger m), higher TPP increases the extent of curvature evolution and hence drift (larger m), and different LR schedules directly change the cumulative integral, yielding schedule-specific m.

---

Published as a conference paper at ICLR 2026
"CPT" prodicted TRECs
"CPT" normalized TRECc
"CPT" Training Position
"CPT" Fraction
1.0
After
1.0
After PT
From scratch
From scratch
TPP after PT
Pediccdd TREC
0.8
Normalized TREC
0.8
234 TPP
0.98-1
0.4
0.4
TPP from scratch
0.2
0.2
0-1
0.0
0.0
0.00
0.25
0.50
0.75
1.00
0.00
0.25
01.00
0.77
1.00
Fraction of "CPT" training takens, F
Fraction of "CPT" training tokens, F
Figure 7: Impact of training fraction TRECs (3.9B). Left: Two 3.9B models trained with identical dataset, LR, weight decay, and batch size, differing only in init (scratch vs. checkpoint after 234 TPP) and thus f. Middle: Predicted TRECs using framework. Right: True TRECs closely match predictions, showing training fraction determines shape under controlled conditions.
Fig. 7 illustrates the importance of the training fraction term in Eq. (3): two 3.9B models trained with identical dataset, learning rate, weight decay, and batch size have the same c(f). However, differences in initialization (from scratch vs, pre-trained checkpoint) mean training evolves over different fr TREC predictions that differ solely on the basis of erd closely match actual TRECs.
Key takeaway 3: TREC shape can be accurately predicted from EMA coefficients and a training- fraction adjustment, enabling proactive curriculum design.
5 APPLICATIONS
5.1
APPLICATION TO SPARSE MIXTURE-OF-EXPERT (MoE)
We now apply TREC analysis to sparse MoE architectures, where only subset of parameters acti- vate per input (Lepikhin et al., 2020; Fedus et al., 2022). We replace each FFN block our 111M model with sparse MoE layer, varying the number of experts E from 1 (dense) to 32. Tokens are routed to experts via hash routing (Roller et al., 2021), ensuring balanced usage. All models train with identical total tokens and datasets. Yet each expert receives only 1/E of tokens, reducing an expert's 's effective tokens-pers peee--ee--ppaaaaee Timescale T, however, is E--ivariint: both batch size B and total tokens D scale identically with 1/E; since B/ (1/A),, these reductions cancel.2
Larger E remembers train more
Larger E - earlier TREC drap
Power law precis & observed MoE data
3.05
1.0
Optimal Train frac. exponent, Ite"
TPP (E)
(8)
3.00
Noralize TREC
1.25
(4)
0.625 (32)
0.6
2.90
0.4
Num. experts
Num. experts
2.85
0.2
2.80
32
0.0
1.4 TTPP)
0.25
0.25
1.00
10-2
10-1
1P
Fraction of training tokiens,
Fraction training takens,
Figure 8: Sparse MoE TRECs reflect reduced effective TPP. Left: Absolute TRECs for 111M models with increasing expert count E (more sparsity) and increasing TREC drop. Middle: Nor- malized curves show larger E shifts the valley earlier. Right: m² predictions from Eq. (4) match true optima, confirming MoEs with more experts behave as if trained at reduced TPP.
Timescale invariance may explain why correcting IR for the effectively 1 FE-smaller per-expert batches is unnecessary (Wei al., 2024, Appendix C.1); MoEs might simply have won the parameterization lottery.
8

---

Published as a conference paper at ICLR 2026
Fig. 8 shows larger E produces greater and earlier TREC drops, indicating stronger memorization: MoE layers behave as trained at their effective TPP. Indeed, empirically-optimal m values for these curves align well with predictions from dense-data power-law model (Eq. (4)), when the power law uses the effective TPP (R² R 83%, Fig. 8-right). These results complement prior obser- vations by Jelassi al. (2024), who found that increasing experts boosts memorization reasoning; our analysis suggests lower effective TPP may partly drive this effect.
than
5.2 APPLICATION TO EVALUATING LLM RECIPES
While prior LLM recipes place their high-quality data at the end of training, our TREC plots suggest this strategy is suboptimal. Still. given this common practice, we ask whether the onser of the HQ phase-ane its reported success or failure consistent with the predicted TREC.
Llama (Dubey al., 2024) evaluated annealing on GSM8k and MATH training sets and reported strong gains for Llama 8B but none for their flagship 405B model. TRECs explain this outcome: the 405B model annealed its LR from 8 to over the final 40M tokens. Since the batch size 16M, this phase spans only ~3 optimizer steps. With such a short window and vanishing LR, EMA coefficients are essentially zero-the model retains little from this final data.
Feng el. Best* 95% aligns w. TREC
OLMc2:
phase aligns w.
Re-eval-based place" best for all
1.0
1.0
Pro-training
OpanWeeaMathh validation loas
0.0015
0.8
1.702
Prediceed TREC
Predicted TREC
0.015
0.6
0.6
1.790
0.4
50
0.4
L,788
30
1,786
0.0
10
0.0
Mid-training stage
1.784
Fraction training takens,
7
Fraction training tokens, f
CPT placement location
Figure 9: Prior work agrees with predicted TRECs. Left: The optimal onset point for the HQ data phase in Feng et al. (2024) aligns with the TREC bottom. Right: Placement of HQ data in OLMo et al. (2024) aligns with the much-narrower TREC valley.
Figure 10: 3.9B CPT results.
Placing data halfway (at TREC bottom) is most effective across all v settings.
Fig. shows two other examples. Feng et al. (2024) tested different onset points for their HQ blend, finding best results when it used for the final 40% region aligning with the predicted TREC valley (leff). In contrast, OLMo-2 13B (OLMo et al., 2024) uses its HQ blend for only the final 5.7% of training, which again aligns with the (narrower) predicted TREC dip (right.. TREC analysis enables finding these optimal placement locations without costly trial-and-error testing.
5.3 CONTINUAL PRE-TRAINING OF A 3.9B LLM USING TREC INSIGHTS
We now test whether TREC-guided data placement improves outcomes in continual pre-training (CPT) setting. Prior placement experiments in Sec. inserted high-quality data partway through base-mode training. In contrast, CPT typically refers to additional training performed after the base model trained. To isolate this setting, we define CPT strictly as continued training after learning rate decay to zero- zero-excluding works such as Parmar et al. (2024), where the CPT LR continues from a 10% decay value, thus arguably reflecting a mid-training strategy (Appendix H.3). This definition matches intuition: CPT can be performed multiple times, mid-trae mid-training is only performed once.
Under this stricter definition, we take a 3.9B-paramett model trained to 234 TPP (900B tokens) with learning-rate decay-to-zero, and continue training for an additional 18B tokens (~5 TPP; illustrated earlier in Fig. 7, left). During this CPT phase, we continue training on the same data blend, but insert 1.3B-token segment of HQ data up-weighting math, targeting improved performance validation (Paster et al., 2023). We compare placing this HQ data halfway through CPT, where observed the empirical TREC for the vanilla CPT be low (orange line in
Open WebMath OpenWebaath

---

Published as a conference paper at ICLR 2026
Fig. 7, right), versus the very end (as is standard), where the TREC returns to baseline. Full model, dataset, and experimental details are in Appendix I.
We test three CPT learning rates, the optimal value is unknown. While effects are necessarily small (HQ data represents only 0.14% of total PT+CPT steps), the trend is consistent: TREC-guided placement remains effective during CPT. Placing HQ data at the TREC minimum outperforms end placement across all LRs (Fig. 10). However, when LR is too high (0.015), even correctly placed data fails to match the performance of the original base model (dashed red line on plot).'
RELATED WORK
Here we highlight the most relevant prior directions, with Appendix B providing full details.
Data curriculums, influence, and attribution. Curriculum learning explores strategies for effec- tively ordering training data (Bengio et al., 2009). Recent LLM curriculums often emphasize high- quality or domain-specific data in later training phases (OLMo et al., 2024; Dubey et al., 2024). While effective, crafting these approaches typically relies on heuristics expensive experiments.
Meanwhile, quantifying influence of training points aids interpretability, auditing, and compensa- tion (Koh & Liang. 2017; Grosse et al., 2023). Recent scalable methods, such data value embed- dings (Wang et al., 2024), recognize the importance of ordering but typically focus on retrospective attribution rather than guiding training. Memorization research examines how models retain training data, motivated by copyright or privacy concerns (Carlini et al., 2022; Schwarzschild et al., 2024).
Re-evaluation. Several works have analyzed loss on training data in order to probe retention dy- namics. Pagliardini et al. (2024) examined memory of specific training batches function of optimizer type and LR schedule. Lesci et al. (2024) estimate memorization profiles across train- ing and checkpoint steps for Pythia models, showing that retention depends data order and LR, and exhibits stable cross-scale trends. Bergsma et al. (2025b) first linked retrospective losses to AdamW's EMA dynamics. In contrast, we systematically vary hyperparameters and LR schedules in order to understand, predict, and exploit data retention in LLM training.
Scaling collapse. Qiu et al. (2025) show that when training loss is normalized appropriately (and training progress is normalized similarly to i), training loss curves collapse onto a universal trajec- tory across model scales. Deformations due to LR schedules are explained by a noisy-quadratic anal- ysis. Our theoretical analysis (Appendix J) and empirical results are complementary: we study ret- rospective re-evaluation rather than initial training loss, account for both TPP and EMA timescale, and translate re-eval structure into actionable curriculums. In later work (Bergsma al., 2025c), we found the same controls that govern re-evaluation loss (TPP, T, LR schedule) govern training loss curve shape.
7 CONCLUSION
We introduced the training re-evaluation curve, simple diagnostic that evaluates how well a trained model retains individual training batches as function of when they appeared in training. Align- ing high-quality data with TREC minima improves final validation loss across models and training regimes (in both PT and CPT), surpassing the default end-of tead----rrainngggggg placement. Crucially, TRECs are largely determined by AdamW EMA coefficients and can be predicted in advance, enabling proactive data placement. We provide a theoretical account in which the training- term 1 captures, in phenomenological form, the view, the fitted m is directly linked to the cumulative influence of schedule and curvature dynamics. Our insights explain existing curriculum strategies, identify suboptimal ones, and yield improved performance in large-scale continual pre-training. Taken together, our results position TREC-based placement principled alternative suboptimal heuristics and costly data-onset ablations.
drift of batch-specific minimizers. In this
'e also tested inserting HQ data at the predicted TREC minimum during pre-training (~-97%). This mid-training strategy matched the best CPT result, suggesting rationale for why mid-training is arguably sup- planting CPT in practice: it achieves comparable gains without requiring extra LR tuning or training compute.
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Marah Abdin, Jyoti Aneja, Hany Awadalla, Ahmed Awadallah, Ammar Ahmad Awan, Nguyen Bach, Amit Bahree, Arash Bakhtiari, Jianmin Bao, Harkirat Behl, et al. Phi-3 technical report: A highly capable language model locally on your phone. arXiv preprint arXiv:2404. 14219, 2024.
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, lige Akkaya, Florencia Leoni Ale- man, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. GPT-4 technical report. arXiv preprint 2rXiv::303.087777 2023.
Rohan Anil, Sebastian Borgeaud, Jean-Baptiin Alayrac, Jiahui Yu, Radu Soricut, Johan Schalkwyk, Andrew M Dai, Anja Hauth, Katie Millican, et al. Gemini: a family of highly capable multimodal models. arXiv preprint arXiv: arXiv:2312.11805 2023.
Jordan Ash and Ryan P Adams. On warm-starting neural network training. Advances in Neural Information Processing Systems, 33:3884-333 -3994, 2020.
Yoshua Bengio, Jérôme Louradour, Ronan Collobert, and Jason Weston. Curriculum learning. In Proceedings of the 26th annual international conference on machine learning, pp. 41-48, 2009.
Shane Bergsma, Nolan Dey, Gurprect Gosal, Gavia Gray, Daria Soboleva, and Joel Hestness. Power lines: Scaling laws for weight decay and batch size in LLM pre-training. arXiv preprint arXiv:2505. 13738, 2025a.
Shane Bergsma, Nolan Dey, Gurpreet Gosal, Gavia Gray, Daria Soboleva, and Joel Hestness. Straight to zero: Why linearly decaying the learning rate to zero works best for LLMs. arXiv preprint arXiv:2502.15938 2025b.
Shane Bergsma, Bin Claire Zhang, Nolan Dey, Shaheer Muhammad, Gurpreet Gosal, and Joel Hes- tness. Scaling with collapse: Efficient and predictable training of LLM families. arXiv preprint UrXiv:2509.250000 2025c.
Tamay Besiroglu, Ege Erdil, Matthew Barnett, and Josh You. Chinchilla scaling: A replication attempt. arXiv preprint arXiv:2404. 10102, 2024.
Xiao Bi, Deli Chen, Guanting Chen, Shanhuang Chen, Damai Dai, Chengqi Deng, Honghui Ding, Kai Dong, Qiushi Du Zhe Fu, et al. DeepSeek LLM: Scaling open-source language models with longtermism. arXiv preprint arXiv:2401.02954 2024.
Stella Biderman, Hailey Schoelkopf, Quentin Anthony, Herbie Bradley, Kyle O'Brien, Eric Hal- Jahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, Aviya Skowron, Lintang Sutawika, and Oskar van der Wal. Pythia: A suite for analyzing large language models across training and scaling, 2023.
Cody Blakeney, Mansheej Paul, Brett W Larsen, Sean Owen, and Jonathan Frankle. Does your data spark joy? performance gains from domain upsampling at the end of training. arXiv preprint arXiv:2406.03476 2024.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in Neural Information Processing Systems, 33:1877-1901, 2020.
Nicholas Carlini, Daphne Ippolito, Matthew Jagielski, Katherine Lee, Florian Tramer, and Chiyuan Zhang. Quantifying memorization across neural language models. In International Conference on Learning Representations, 2022.
Sang Keun Choe, Hwijeen Ahn, Juhan Bae, Kewen Zhao, Minsoo Kang, Youngseog Chung, Adithya Pratapa, Willie Neiswanger, Emma Strubell, Teruko Mitamura, et al. What is your data worth to GPT? LLM-scale data valuation with influence functions. arXiv preprint Att::2005.19544 2024.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Rejichiro Nakano, et al. Training verifiers to solve math word problems. arXiv preprint arXiv:21 10.14:210 2021.
11

---

Published as a conference paper at ICLR 2026
Aaron Defazio, Ashok Cutkosky, Harsh Mehta, and Konstantin Mishchenko. When, why and how much? Adaptive learning rate scheduling by refinement. arXiv preprint arXiv:2100.08811, 2023.
Nolan Dey, Gurpreet Gosal, Hemant Khachane, William Marshall, Ribhu Pathria, Marvin Tom, and Joel Hestmess. Cerebras-GrP GPT: Open compute-optimal language models trained on the Cerebras wafer-scale cluster. arXiv preprint wwt.:::3:..3320000 2023.
Nolan Dey, Bin Claire Zhang, Lorenzo Noci, Mufan Li, Blake Bordelon, Shane Bergsma, Cengiz Pehlevan, Boris Hanin, and Joel Hestness. Don't 't be lazy: CompleteP enables compute-efficient deep transformers. arXiv preprint arXiv:2505.01618, 2025.
Shibhansh Dohare, Fernando Hernandez-Garcia. Qingfeng Lan, Parash Rahman, A Rupam Mah- mood, and Richard Sutton. Loss of plasticity in deep continual learning. Nature, 682(8026): 768-774, 2024.
Ricardo Dominguez-Olmedo, Florian Dorner, and Moritz Hardt. Training on the test task con- founds evaluation and emergence. arXiv preprint arXiv:240.007890, 2024.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et al. The Llama 3 herd of models. arXiv preprint arXiv:2407.21188, 2024.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research, 12(7), 2011.
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 23(12)):1-39 2022.
Vitaly Feldman and Chiyuan Zhang. What neural networks memorize and why: Discovering the long tail via influence estimation. Advances in Neural Information Processing Systems, 33:2881- 2020.
Steven Feng, Shrimai Prabhumoye, Kezhi Kong, Dan Su, Mostofa Patwary, Mohammad Shoeybi, and Bryan Catanzaro. Maximize your data's potential: Enhancing LLM accuracy with two-phase pretraining. arXiv preprint arXiv:2412.15285. 2024.
Petko Georgiev, Ving Ian Lei, Ryan Burnell, Libin Bai, Anmol Gulati, Garrett Tanzer, Damien Vincent, Zhufeng Pan, Shibo Wang, et al. Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context. arXiv preprint arXiv:2403.05530. 2024.
Jacopo Graldi, Alessandro Breccia, Giulia Lanzillotta, Thomas Hofmann, and Lorenzo Noci. The importance of being lazy: Scaling limits of continual learning. arXiv preprint arXiv:2506. 16884, 2025.
Alex Graves. Generating sequences with recurrent neural networks UrXiv:1308, 0850, 2013.
arXiv preprint
Roger Grosse, Juhan Bae, Cem Anil, Nelson Elhage. Alex Tamkin, Amirhossein Tajdini, Benoit Steiner, Dustin Li, Esin Durmus, Ethan Perez, et al. Studying large language model generalization with influence functions. arXiv preprint arXiv:2308.03222 2023.
Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor opti- mization. In International Conference Machine Learning. pp. 1850-1850. PMLR, 2018.
Suchin Gururangan, Ana Marasović, Swabha Swayamdipta, Kyle Lo, Iz Beltagy, Doug Downey, and Noah A Smith. Don't stop pretraining: Adapt language models to domains and tasks. arXiv preprint arXiv:2004. 10V:44004..094,, 2020.
Alexander Hägele, Elie Bakouch, Atli Kosson, Loubna Ben Allal, Leandro Von Werra, and Martin Jaggi. Scaling laws and compute-optimal training beyond fixed training durations. arXiv preprint arXiv:2405. 18392, 2024.
12

---

Published as a conference paper at ICLR 2026
Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the MATH dataset. arXiv preprint ar arXiv:2103.038// 2021.
Danny Hernandez, Tom Brown, Tom Conerly, Nova DasSarma, Dawn Drain, Sheer El-Showk, Nel- son Elhage, Zac Hatfeel--Dodds, Tom Henighan, Tristan Hume, al. Scaling laws and inter- pretability of learning from repeated data. arXiv preprint arXiv:2205. 20l.:2205.00877 10487, 2022.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. An empirical analysis of compute-optimal large language model training. Advances in Neural Infor- mation Processing Systems, 35, 2022.
Shengding Hu, Yuge Tu, Xu Han, Chaoqun He, Ganqu Cui, Xiang Long, Zhi Zheng, Yewei Fang. Yuxiang Huang, Weilin Zhao, et al. MiniCPM: Unveiling the potential of small language models with scalable training strategies. arXiv preprint arXiv:2404.00 06395, 2024.
Kai Hua, Steven Wu, Ge Zhang, and Ke Shen. AttentionIafluence: Adopting attention head influence for weak-to-strong pretraining data selection. arXiv preprint arXiv:2505.0723 072993 2025.
Adam Ibrahim, Benjamin Thérien, Kshitij Gupta, Mats L Richter, Quentin Anthony, Timothée Lesort, Eugene Belilovsky, and Irina Rish. Simple and scalable strategies to continually pre-train large language models. arXiv preprint arXiv:2403.0876. 2024.
Samy Jelassi, Clara Mohri, David Brandfonbrener. Alex Gu, Nikhil Vyas, Nikhil Anand, David Alvarez-Melis, Yuanzhi Li, Sham M Kakade, and Eran Malach. Mixture of parrots: Experts improve memorization more than reasoning. arXiv preprint arXiv:2410. 19034, 2024.
Dayal Singh Kalra, Tianyu He, and Maissam Barkeshli. Universal sharpness dynamics in neural network training: Fixed point analysis, edge of stability, and route to chaos. arXiv preprint arXiv:2311.02076, 2023.
Feiyang Kang, Yifan Sun, Bingbing Wen, Si Chen, Dawn Song, Rañid Mahmood, and Ruoxi Jia. AutoScale: Automatic prediction of compute-optimal data compositions for training 1lms. 2024.
Diederik Kingma and Jimmy Ba. Adam: method for stochastic optimization, arXiv preprint arXiv:1412.6980, 2014.
James Kirkpatrick, Razvan Pascanu, Neil Rabinowitz, Joel Veness, Guillaume Desjardins, Andrei A Rusu, Kieran Milan, John Quan, Tiago Ramalho, Agnieszka Trabbka-BBrwiinka. et al. Overcom- ing catastrophic forgetting neural networks. Proceedings of the National Acudemy of Sciences.
114(13):3521-3526 2017.
Pang Wei Koh and Percy Liang. Understanding black-box predictions influence functions. In International conference machine learning, pp. 1894. PMLR, 2017.
Atli Kosson, Bettina Messmer, and Martin Jaggi. Analyzing & reducing the need for learning rate warmup in GPT training. arXiv preprint axxiv:2410.2332222222 2024.
Tanishq Kumar, Zachary Ankner, Benjamin Spector, Blake Bordelon, Niklas Muennighoff, Man- sheej Paul, Cengiz Pehlevan, Christopher Ré, and Aditi Raghunathan. Scaling laws for precision. arXiv preprint arXiv:2411.04333 2024.
Yongchan Kwon, Eric Wu, Kevin Wu, and James Zou. DataInf: Efficiently estimating data influence in lora-tuned LLMs and diffusion models. arXiv preprint arXiv:2310.00900 2023.
Yann LeCun, John Denker, and Sara Solla. Optimal brain damage. Advances in Neural Information Processing Systems, 2, 1989.
Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu. Dehao Chen. Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. Gshard: Scaling giant models with conditional computation and automatic sharding. arXiv preprint arXiv:2006. 16668, 2020.
13

---

Published as a conference paper at ICLR 2026
Pietro Lesci, Clara Meister, Thomas Hofmann, Andreas Vlachos, and Tiago Pimentel. Causal esti- mation of memorisation profiles. arXiv preprint 2xxi::2406604327,, 2024.
Jeffrey Li, Alex Fang, Georgios Smyrnis, Maor Ivgi, Matt Jordan, Samir Yitzhak Gadre, Hritik Bansal, Etash Guha, Sedrick Scott Keh, Kushal Arora, et al. DataComp-LM: In search of the next generation of training sets for language models. Advances in Neural Information Processing Systems, 12110-14022222 2024a.
Ruihang Li, Yixuan Wei, Miaosen Zhang, Nenghai Yu, Han Hu, and Houwen Peng. Scal- ingFilter: Assessing data quality through inverse utilization of scaling laws. arXiv preprint Sttv:240080831100 2024b.
Emmy Liu, Graham Neubig, and Chenyan Xiong. Midtraining bridges pretraining and posttraining distributions. arXiv preprint arXiv:2510.14865, 2025.
Hong Liu, Zhiyuan Li, David Hall, Percy Liang, and Tengyu Ma. Sophia: A scalable stochastic second-order optimizer for language model pre-training. arXiv preprint arXiv:2305. 14342, 2023.
Qian Liu, Xiaosen Zheng, Niklas Muennighoff, Guangtao Zeng, Longxu Dou, Tianyu Pang, Jing Jiang, and Min Lin. RegMix: Data mixture regression for language model pre-training. arXiv preprint ar 2rXiv:2407.01492, 2024.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Confer- ence on Learning Representations, 2017.
Clare Lyle, Zeyu Zheng, Evgenii Nikishin, Bernardo Avila Pires, Razvan Pascanu, and Will Dabney. Understanding plasticity in neural networks. arXiv preprint wrXiv:2303.01466 2023.
Sam McCandlish, Jared Kaplan, Dario Amodei, et An empirical model of large-batch training. arXiv preprint arXiv:1812.06100 2018.
Michael McCloskey and Neal J Cohen. Catastrophic interference in comnectionist networks: The sequential learning problem. In Psychology of learning and motivation, volume 24, pp. 109-165.
1989.
William Merrill, Shane Arora, Dirk Groeneveld, and Hannanch Hajishirzi. Critical batch size re- visited: A simple empirical approach to large-batch language model training. arXiv preprint arXiv:2505.23977 2025.
Thomas Mesnard, Cassidy Hardin, Robert Dadashi, Surya Bhupatiraju, Shreya Pathak, Laurent Sifre, Morgane Rivière, Mihir Sanjay Kale, Juliette Love, et al. Gemma: Open models based gemini research and technology. arXiv preprint arXiv:2403.08295, 2024.
Sören Mindermann, Jan M Brauner, Muhammed T Razzak, Mrinank Sharma, Andreas Kirsch, Win- nie Xu, Benedikt Höltgen, Aidan N Gomez, Adrien Morisot, Sebastian Farquhar, et al. Prioritized training on points that are learnable, worth learning, and not yet learnt. In International Confer- ence on Machine Learning, pp. 11630-15649, PMLR, 2022.
John X Morris, Chawin Sitawarin, Chuan Guo, Narine Kokhlikyan, G Edward Suh, Alexander M Rush, Kamalika Chaudhuri, and Saeed Mahloujifar. How much do language models memorize? arXiv preprint arXiv:2505.2488 24822. 2025.
Lorenzo Noci, Alexandru Meterez, Thomas Hofmann, and Antonio Orvieto. Super consistency of neural network landscapes and learning rate transfer. Advances in Neural Information Processing Systems, 17:102996-002433, 2024.
Team OLMo, Pete Walsh, Luca Soldaini, Dirk Groeneveld, Kyle Lo, Shane Arora, Akshita Bha- gia, Yuling Gu, Shengyi Huang. Matt Jordan, al. OLMo 2 Furious. arXiv preprint arXiv:2501.00656, 2024.
Firat Öncel, Matthias Bethge, Beyza Ermis, Mirco Ravanelli, Cem Subakan, and Çağatay Yıldı. Adaptation odyssey in LLMs: Why does additional pretraining sometimes fail to improve? arXiv preprint 10.55:2400.05581, 2024.
14

---

Published as a conference paper at ICLR 2026
Matteo Pagliardini, Pierre Ablin, and David Grangier. The AdEMAMix optimizer: Better, faster, older. arXiv preprint arXiv:2409.03137, 2024.
Jupinder Parmar, Sanjev Satheesh, Mostofa Patwary, Mohammad Shocybi, and Bryan Catanzaro. Reuse, don't 't retrain: A recipe for continued pretraining of language models. arXiv preprint arXiv:2407.07263, 2024.
Keiran Paster, Marco Dos Santos, Zhangir Azerbayev, and Jimmy Ba. OpenWebMath: An open dataset of high-quality mathematical web text, 2023.
Tomer Porian, Mitchell Wortsman, Jenia Jitsev, Ludwig Schmidt, and Yair Carmon. Resolving discrepancies in compute-optimal scaling of language models. arXiv preprint arXiv:2406.19146 2024.
Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra. Grokking: Gen- eralization beyond overfitting on small algorithmic datasets. arXiv preprint arXiv:2201.02177, 2022.
Ofir Press, Noah Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. In International Conference on Learning Representations, 2022.
Yujia Qin, Jiajic Zhang, Yankai Lin, Zhiyuan Liu, Peng Li, Maosong Sun, and Jie Zhou. Elle: Efficient lifelong pre-training for emerging data. arXiv preprint arXiv:2203.063: 2022.
Shikai Qiu, Lechao Xiao, Andrew Gordon Wilson, Jeffrey Pennington, and Atish Agarwala. Scaling collapse reveals universal dynamics in compute- compute-optimally trained neural networks. arXiv preprint arXiv:2507.0211 19, 2025.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners, 2019.
Jack W. Rac, Scbastian Borgeaud, Trevor Cai, Katic Millican, Jordan Hoffmaann, Francis Song, John Aslanides, Sarah Henderson, Roman Ring, Susannah Young, et al. Scaling language models: Methods, analysis & insights from training Gopher, 2022.
Stephen Roller, Sainbayar Sukhbaatar, Arthur Szlam, and Jason Weston. Hash layers for large sparse models. Advances in Neural Information Processing Systems, 34:1 1451155517566, 2021.
Noveen Sachdeva, Benjamin Coleman, Bang-Cheng Kang, Jianmo Ni, Lichan Hong, Ed H Chi, James Caverlee, Julian McAuley, and Derek Zhiyuan Cheng. How to train data-efficient LLMs. arXiv preprint arXiv:2402.09666 2024.
Fabian Schaipp, Alexander Hägele, Adrien Taylor, Umut Simsekli, and Francis Bach. The surpris- ing agreement between convex optimization theory and learning-rate scheduling for large model training. arXiv preprini 2rXiv:2501.18965 2025.
Avi Schwarzschild, Zhili Feng, Pratyush Maini, Zachary Lipton, and J Zico Kolter. Rethinking LLM memorization through the lens of adversarial compression. Advances Neural Information Processing Systems, 37:56244-56267, 2024.
Christopher Shallue, Jaehoon Lee, Joseph Antognini, Jascha Sohl-Dickstein, Roy Frostig, and George E. Dahl. Measuring the effects of data parallelism on neural network training. Journal of Machine Learning Research, 20( 12):1449, 2019.
Noam Shazeer. GLU variants improve transformer. arXiv preprini arXiv:2002.052000 2020.
Noam Shazeer and Mitchell Stem. Adafactor: Adaptive learning rates with sublinear memory cost. In International Conference on Machine Learning, pp. 4606. PMLR, 2018.
Yikang Shen, Zhen Guo, Tianle Cai, and Zengyi Qin. JetMoE: Reaching Llama2 performance with 0.1M dollars. arXiv preprint arXiv:2404.07413, 2024.
David So, Wojeiech Manke, Hanxiao Liu, Zihang Dai. Noam Shazeer, and Quoc V Le. Searching for efficient transformers for language modeling. Advances in neural information processing systems, 34:6010-600 6022, 2021.
15

---

Published as a conference paper at ICLR 2026
Daria Soboleva, Faisal Al-Khateels Robert Myers, Jacob R Steeves, Joel Hestness, and Nolan Dey. SlimPajama: A 627B token cleaned and deduplicated version of RedPajama. Web page, 2023.
Jacob Mitchell Springer, Sachin Goyal, Kaiyue Wen, Tanishq Kumar, Xiang Yue, Sadhika Malladi, Graham Neubig, and Aditi Raghunathan. Overtrained language models harder to fine-tune. arXiv preprini arXiv:2503. 19206, 2025.
Falcon-LLM Team. Falcon-H1: A family of hybrid-head language models redefining efficiency and performance, May 2025a. URL https: /ffalpo 1m. github. httpb/log/aalcooooo/
Kimi Team. Kimi K2: Open agentic intelligence, 2025b.
URL httpa://githhb.coo
Howe Tissue, Venus Wang, and Lu Wang. Scaling law with learning rate annealing. arXiv preprint
arXiv:2408. 1029, 2024.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurelien Rodriguez, Ar- mand Joulin, Edouard Grave, and Guillaume Lample. LLaMA: Open and efficient foundation language models, 2023.
Nikhil Vyas, Alexander Atanasov, Blake Bordelon, Depen Morwani, Sabarish Sainathan, and Cen- giz Pehlevan. Feature-learning networks are consistent across widths at realistic scales, Advances in Neural Information Processing Systems, 36: 10:6006-000, 2023.
Nikhil Vyas, Depen Morwani, Rosie Zhao, Mujin Kwun, Itai Shapira, David Brandfonbrener, Lueas
Janson, and Sham Kakade. SOAP: Improving and stabilizing shampoo using adam. arXiv preprint arXiv:2409. 11311, 2024.
Alan Wake, Bei Chen, CX Lv, Chao Li, Chengen Huang, Chenglin Cai, Chujic Zheng, Daniel Cooper, Fan Zhou, Feng Hu, al. Yi-Lightning technical report. arXiv preprint arXiv:2412.01253, 2024.
Jiachen T Wang, Dawn Song, James Zou, Prateek Mittal, and Ruoxi Jia. Capturing the temporal dependence of training data influence. arXiv preprint arXiv:2412.09538 2024.
Xi Wang and Laurence Aitchison. How to set AdamW's weight decay as you scale model and dataset size. arXiv preprint arXiv:2405.13660 2024.
Zengzhi Wang, Fan Zhou, Xuefeng Li, and Pengfei Liu. OctoThinker: Mid-training incentivizes reinforcement learning scaling. arXiv preprint 1.XXv:2506.20//2 2025.
Tianwen Wei, Bo Zhu, Liang Zhao, Cheng Cheng, Biye Li, Weiwei Lü, Peng Cheng, Jianhao Zhang, Xiaoyu Zhang, Liang Zeng, et al. Skywork-MoE: A deep dive into training techniques for mixture-of-experts language models. arXiv preprint 24t.::240000065633 2024.
Mengzhou Xia, Sadhika Malladi, Suchin Gururangan, Sanjeev Arora, and Danqi Chen. LESS: Selecting influential data for targeted instruction tuning. arXiv preprint 2xXiv:2402.04333 2024.
Greg Yang, Edward Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tuning large neural networks via zero-shot hyperparameter transfer. In Advances in Neural Information Processing Systems, 2021.
Jiasheng Yc, Peiju Liu, Tianxiang Sun, Jun Zhan, Yunhua Zhou, and Xipeng Qiu. Data mixing laws: Optimizing data mixtures by predicting language modeling performance. arXiv preprint arXiv:2403.16952, 2024.
Yichun Yin, Wenyong Huang, Kaikai Song, Yehui Tang, Xueyu Wu, Wei Guo, Peng Guo, Yaoyuan Wang, Xiaojun Meng. Yasheng Wang, al. Pangu Ultra: Pushing the limits of dense large language models ascend NPUs. arXiv preprint arXiv:2504.07866, 2025.
Ge Zhang, Scott Qu, Jiaheng Liu, Chenchen Zhang. Chenghua Lin, Chou Leuang Yu, Danny Pan, Esther Cheng, Jie Liu. Qunshu Lin, et al. MAP-Neo: Highly capable and transparent bilingual large language model series. arXiv preprint arXiv: 19327, 2024a.
16

---

Published as a conference paper at ICLR 2026
Guodong Zhang, Lala Li, Zachary Nado, James Martens, Sushant Sachdeva, George Dahl, Chris Shallue, and Roger Grosse. Which algorithmic choices matter at which batch sizes? insights from noisy quadratic model. Advances in neural information processing systems, 32, 2019.
Hanlin Zhang, Depen Morwani, Nikhil Vyas, Jingfeng Wu, Difan Zou, Udaya Ghai, Dean Fos- ter, and Sham Kakade. How does critical batch size scale in pre-training? arXiv preprint arXiv:2410.21676, 2024b.
Jingwei Zuo, Maksim Velikanov, Dhia Eddine Rhaiem, Ilyas Chahed, Younes Belkada, Guillaume Kunsch, and Hakim Hacid. Falcon Mamba: The first competitive attention-free 7B language model. XiXiv preprint arX::22410.05355, 2024.
17

---

Published as a conference paper at ICLR 2026
A LIMITATIONS AND FUTURE WORK
While TRECs offer practical and predictive diagnostic for guiding data placement in LLM training, several limitations and opportunities for future work remain.
Optimizer scope. The TREC itself is not optimizer-specffi:: TRECs can be computed for any training run regardless of the optimization algorithm. However, our predictive analysis of TRECs is tailored to the Adam W optimizer, drawing on its implicit EMA formulation and corresponding timescale. Update rules for other optimizers using weight decay (such as Sophia (Liu et al., 2023, Algorithm 3) and MuonClip (Team, 2025b, Algorithm 1)) can be directly converted to an extended EMA form, exactly as was done with AdamW (Sec. 4). Moreover, the EMA perspective should also hold when AdamW is applied a different weight basis, e.g., as SOAP (Vyas et al., 2024). where AdamW W is applied in Shampoo's 's eigenbasis (Gupta et al., 2018).
Extending predictive TREC models to optimizers without an implicit EMA formulation-including Adagrad (Duchi et al., 2011), Adafactor (Shazeer & Stern, 2018), or SGD variants-remain an important avenue for future exploration. For optimizers that do not use weight decay, such as Adam (Kingma& & Ba, 2014), we can potentially view them as the limit of senggh--eca--ehhnneed versions, as weight decay goes to zero. As shown in Appendix F.3, such optimizers may approach a specific EMA shape, even when the timescale is undefined.
Training setups and model scales. Our experiments focus primarily on compute-optimal and overtrained regimes, using training runs or beyond 20 likees-per-parameter (TPP). While this aligns with common practices in large-scale LLM development, further work is needed to under- stand how TREC dynamics behave in undertrained or data-scarce settings. We also primarily study models in the 100M-4B range; exploring scaling trends for smaller or larger models could refine our conclusions.
Data types, quality, and curriculums. Our work focuses on the placement of high-quality data presumed to be limited in quantity, consistent with prior observations that "truly high-quality tokens are still scarce at this moment" (Wang et al., 2025). We test placing discrete segments of HQ data, rather than continuously evolving data distributions. Moreover, although our findings generalize across blends involving code, math, and web text, we do not explicitly analyze how TRECs (or the reliability of TREC-guided placements) vary for distinct data types-suue as factual vs. reasoning, or instruction vs. narrative content. A natural extension would be to combine our placement frame- work with recent approaches for data selection and weighting, including AutoScale (Kang et al., 2024) and RegMix (Liu et al., 2024).
Optimizing placement to maximize retention and downstream task performance raises important questions. For example, increasing memorization may come at the expense of general reasoning ability (Jelassi et al., 2024), may confound evaluation by overly tailoring training to benchmark tasks Soominguez-Olmenn has been shown be essential for models benefit from downstream reinforcement learning (Wang et al., 2025), or SFT datasets that differ substantially from the PT distribution (Liu et al., 2025). The TREC framework provides tool for maximizing the effect of limited training data-but whether learning from such data is ultimately beneficial remains an open, complex, and important question in language model research.
et al., 2024). At the same time, high-quality domain-specific pre-training
Context length, vocabulary, and data diversity. Our study is limited to two vocabulary configu- rations and context lengths of 2048 and 8192 tokens. The interaction between TREC dynamics and architectural choices such tokenizer design or sequence length remains underexplored. Similarly, we evaluate a modest range of dataset blends, and further validation is needed to assess generality across diverse languages or modalities.
Evaluating memorization and generalization. TRECs measure how well the final model retains or forgets data presented at different points in training. However, we do not attempt to quantify exact memorization defined in prior work (e.g., substring continuation (Carlini et al., 2022: Georgiev et al., 2024; Schwarzschild et al., 2024)). Future studies could examine how TREC loss relates to sequence-level memorization and whether schedules that maximize retention may inadvertently
18

---

Published as a conference paper at ICLR 2026
encourage overfitting. In similar vein, Biderman et al. (2023) offered advice where to "place sequences that are undesirable to memorize." " It is worth studying whether deliberately placing biased undesirable data away from the TREC minimum may offer new tool for mitigating unwanted retention. See also our notes on the "memorization window" in Appendix B.
Predictive scope and validation. Our predictive framework successfully anticipates TREC shape across optimizer settings and learning rate schedules. However, predictions of optimal placement for prior work are not verified via end-to-end rettriinnn Additionally, our CPT experiments suggest that predictive placement is most reliable within particular optimizer configuration (ry(1, X, B), and may fail to generalize across configurations. Finding effective techniques that jointly choose both placement and optimizer configuration an open goal and an important next step. Doing so will likely require predicting when and why a model transitions from memorization into grokking, as discussed further in Appendix D.2.
Also, our predictive form focused on the (normalized) shape of the TREC loss (Eq. (3)), rather than its absolute magnitude. would be valuable to explore other forms that can predict both shape and magnitude. It would also be valuable to expand our theoretical model to encompass these extended forms (Appendix J).
Designing LR schedules that do not forget. Given our predictive form for TREC loss (Eq. (3)), is possible to design a learning rate schedule such that the EMA coefficients "cancel out" the effects of the training fraction term. That is, in theory, we may design schedule order to obtain flat TREC.
In practice, such no-forgetting schedule would have to rapidly decrease the LR as a power law in training steps, in order to offset the increasing power-law of the training fraction term. In prior work, such rapidly decreasing LR schedules do not perform as well as more-gradual decline (Defazio et al., 2023; Bergsma et al., 2025b). Based on the discussion in Appendix J, we may understand why: since gradients lose their effectiveness, some forgetting is actually desirable in LLM pre-training. Indeed, the optimal EMA timescale has been shown decrease 2025a), meaning that when training longer, relatively more of the data should be forgotten.
power law in TPP (Bergsma et al.,
However, there contexts where avoiding forgetting may be important (e.g., when performing CPT or SFT). would also be interesting to investigate and mitigate sources of local optimizer drift, as means to reducing the need for forgetting. Can we transition learning to regime where new knowledge can be added indefinitely, without fundamentally changing the representation of such knowledge? Recent work using lazy learning to avoid catastrophic forgetting may provide valuable insights here for the TREC perspective (Graldi et al., 2025).
Toward practical deployment. Finally, while we offer actionable guidance for curriculum design (e.g.. predicting TRECs in advance, avoiding late placement under step-drop or D27. LR schedules, or leveraging homogeneous CPT phases to measure TREC dips), wider adoption will depend on usability. We will explore mechanisms to make TREC tools more performant and accessible the community. For example, for cases where TRECs are constructed through explicit re-evaluation rather than prediction, compute can be saved by sampling portion of the training batches re- evaluate on, rather than evaluating on every batch (as we do). Determining acceptable fraction depends on loss variance, and should be investigated systematically future work.
B ADDITIONAL RELATED WORK
Optimizing data mixtures and quality. Recent work has explored methods for improving LLM training via better data selection or mixture strategies. One line of work focuses on identifying what constitutes high-quality data, including weak-to-strong selection using attention mechanisms (Hua et al., 2025), filtering based on scaling law deviations (Li et al., 2024b), or assessments from trained models (Sachdeva et al., 2024; Li et al., 2024a). Another line of work addresses how to mix datasets drawn from multiple domains. Ye et al. (2024) introduce data mixing laws, showing that validation loss be predicted to generalize to larger model/data regimes. Liu et al. (2024) similarly use small-scale training runs to regress mixture efficacy and extrapolate to larger models, outperforming heuristic and prior-
function of mixture proportions and proposing nested use of scaling laws
19

---

Published as a conference paper at ICLR 2026
optimized mixtures. AutoScale (Kang al., 2024) takes more theoretical approach, modeling how optimal domain weights vary with training scale, and deriving recipes that converge faster and perform better than baselines. While these efforts offer valuable tools for which data to include and how to weight it, our work focuses on when to introduce data (possibly very limited in size) during training.
Data curriculums and annealing phases. Our work shares conceptual motivation with currieu- lum learning (Bengio et al., 2009; Mindermann et al., 2022), which aims to improve generalization by presenting examples in the most effective order. In modern LLM training, this principle manifests as staged or "annealed" data mixtures, where high-quality or domain-specific datasets are introduced later in training. In Llama-3 (Dubey et al., 2024), the authors explicitly report following OpenAl's strategy (Achiam et al., 2023) of annealing in-domain datasets like GSM8K (Cobbe et al., 2021) and MATH (Hendrycks et al., 2021).* Other state-of of oftte-of-thearrr the models using such practices include OLMo-2 (OLMo al., 2024), JetMoR (Shen al., 2024), Phi-3 (Abdin al., 2024), Gemini (Anil et al., 2023), Gemma (Mesnard et al., 2024), MAP-Neo (Zhang et al., 2024a), Falcon-Mamba (Zuo et al., 2024) and Yi-Lightning (Wake et al., 2024). These strategies are premised on the assumption that exposing key data late in training improves downstream task performance. Recent work has also used annealing phases to assess data quality efficiently (Blakeney et al., 2024; Dubey et al., 2024: OLMo et al., 2024), enabling comparisons without full-scale pre-training. However, placing the highest-quality data at the very end is flawed heuristic, and using ablations to determine the optimal onset for the high-quality phase is expensive. Our TREC diagnostic addresses these issues by offering a scalable method for identifying optimal data placement locations.
Mcasuring loss on training data post-hoc. As far back as Graves (2013), it was observed that generated text from neural language model tends to inordinately reflect the final training batches; from a TREC- -errppees,ee, this manifests as lower TREC losses on the final training samples. While not the primary focus of these works, previous studies have also explicitly measured loss on previ- ously seen training examples in order to understand model retention dynamics. In a comparison of different LR schedules, Bergsma et al. (2025b) qualitatively connected training re-evaluation losses to the exponential moving average (EMA) coefficients of AdamW updates. Pagliardini et al. (2024) propose AdEMAMix, an optimizer designed to forget more slowly. They visualize loss trajectories for individual batches across training and conclude that learning rate decay is the dominant factor controlling forgetting.
Lesci al. (2024) introduce difference-in- -naie-ffriifffeeessss--- estimator for counterfactual memorisation, constructing a two-dimensional memorisation profile Tge that measures the causal effect of training on batch g model performance at checkpoint c. Their analysis shows that memorisation strength in the Pythia model suite (Biderman et al., 2023) depends on learning rate, data order, and model scale. Conceptually, our TREC can be viewed as specific slice of such memorisation profile: we fix c = T (the final checkpoint) and measure the retrospective loss as a function of training position g- By not relying on public checkpoints, but instead systematically varying optimizer hyperparame- ters during pre-training, we develop a deeper understanding of the factors that govern retention. We also differ from all these works in that we evaluate TRECs as a tool to guide curriculum design.
Influence estimation and data attribution. Training data influence has long been studied for in- terpretability and accountability. Influence function approaches (Koh & Liang, 2017) and retraining- based approximations (Feldman & Zhang, 2020; Grosse et al., 2023) estimate data value by measur- ing its effect on final model behavior. Recent scalable methods such as DataInf (Kwon al., 2023), LESS (Xia et al., 2024) and LoGRA (Choe et al., 2024) leverage gradient-based approximations (e.g., via LoRA) to approximate influence scale. Wang et al. (2024) break from permutation- invariant assumptions and introduce trajectory-specfiic data value embeddings that explicitly model training data order. Their method uncovers distinct training phases in LLMs: a "high-impact warmup phase," followed by "low-impact basin," and then a "gradual ascending" region. While insightful, this pattern differs from our findings: both EMA analysis and TREC diagnosties suggest
'While the placement details are often not reported (i.e., when exactly such data is introduced during pre- training), Dominguez Olmeeo et al. (2024) specifically identify November 2023 as turning point, after which technical reports "start referencing certain pre-training practices that may amount to training the test task," such as using instruction-tuming data or QA templates.
20

---

Published as a conference paper at ICLR 2026
that specific batches in early training data have minimal influence final weights. In this sense, our work serves as a valuable cross-chect for attribution Tallyyes-TREE offer a simple, forward-only sanity-check that can validate or challenge more complex influence models.
Memorization and forgetting dynamics. A growing literature studies memorization in LLMs, often motivated by copyright or privacy risks. Typical methods identify memorized sequences by testing whether prompting with part of a training example elicits exact or near-exact continu- ations (Carlini et al., 2022; Georgiev et al., 2024; Schwarzschild et al., 2024). Morris et al. (2025) propose an information-theoree show how it varies across dataset size, with conclusions that align with our own TREC results (Sec. 3).
definition of memorization based on Kolmogorov complexity, and
In large-scale analysis, Biderman al. (2023) found that training order had little impact mem- orization, with memorized sequences distributed approximately as a Poisson process across train- ing. This contrasts sharply with our results: TRECs consistently show strong order effects on loss, indicating meaningful differences in retention and learning across training time. Exact sequence reproduction may be too coarse a signal capture the subtler, gradient- adaptations revealed by TREC analysis.
Seeking to avoid the known deleterious effects of data repetition (Hernandez et al.. 2022), yet eager to make full use of their high-quality data, Team (2025a) report successful re-use of such data, so long as one "carefully estimates" and avoids the model's "memorization window." " Though no further details are offered, their approach suggests another promising application of TREC analysis: it could help identify when high-quality data can be safely repeated.
While the prior works aim to avoid memorization, separate line of work focuses on the opposite failure mode: catastrophic forgetting. This is a central concern in continual learning (Kirkpatrick et al,, 2017), where models forget previously-leanan
knowledge when exposed new data. Closely related is the issue of loss of plasticity, in which extensive pre-training reduces a model's 's capacity to acquire new information (Ash & Adams, 2020; Lyle et al., 2023; Dohare et al., 2024; Kumar et al., 2024; Springer et al., 2025). Our TREC framework helps disentangle these effects. For example, training to higher TPP reduces the magnitude of the TREC drop (reflecting reduced plasticity), while mainly preserving the shape of the curve (indicating which segments are most forgotten).
Continual pre-training (CPT) dynamics. Continual pre-training (CPT), or lifelong learning, in- volves adapting models to new data distributions beyond the initial training set. Earlier work in this area emphasized domain adaptation for classification tasks (Gururangan et al., 2020: Qin al., 2022), while recent LLM recipes apply CPT to full model continuation. However, such practices often suffer from performance degradation, even when continuing on the original domain, due to optimization challenges (Ibrahim et al., 2024; Óncel et al., 2024). CPT strategies commonly repur- pose pre-training heuristics such as annealing phases or late-stage data swaps- principled guidance. Our work introduces a TREC framework that directly informs when and how to incorpo- rate new data during CPT.
Scale-stable dynamics under juP and normalized compute. Feature-learning parameterizations such as μP can transfer hyperparameters across scale and yield early-time consistency of dynam- ies across widths (Vyas et al., 2023; Kalra et al., 2023), though finite-width deviations grow on harder tasks and later epochs. Complementing this, Noci al. (2024) provide evidence of super- consistency in curvature (e.g., Hessian eigenvalues) along the training trajectory, supporting training fraction as a natural coordinate (Appendix J.5). Building on these observations, Qiu al. (2025) show that when loss is indexed by normalized compute x = (p)" where / is the current step and (p)) is the compute-optimal step count for model size p, then training-loss curves collapse across scales, with LR wwwhhduleddependeeeee
deformations explained by moisy-quadrattc analysis; collapse also holds for fixed multiples of t" (p). Their experiments, however, are limited to proof-of-concept models at relatively small scale, whereas we validate TREC dynamics in LLMs up to 3.9B param- eters. Our setting also differs in target (re-evaluation vs. training loss) and mechanism (Adam EMA timescale + drift).
In later work (Bergsma et al.. 2025c), we found the same controls that govern re-evaluation loss (TPP, LR schedule) govern training loss curve shape.
21

---

Published as a conference paper at ICLR 2026
Table 3: Model architectures used in main experiments
Model
mmodel
Players
d+i dhead
Experiments
111M
768
10
2048
64
Sec. 3 and Sec. 5.1
266M
768
32
2048
64
Sec. 3
610M
2048
10
5461
64
Secs. and 3
1.7B
2048
32
5461
64
Sec. 3
3.3B
2048
64
5461
64
Sec.
3.9B
2048
40
16384
128
Sec. 4 (Fig. 7), Sec. 5.3
Table 4: Models, pekens-per-parameter (TPP) and corresponding dataset sizes (in tokens), number of model variants trained (LR schedule type, 7), A, B, or data placement strategy), and purpose of trained models. total, 41 models were trained with different data placements, and 578 TRECs were computed over different optimizer hyperparameters.
Model
TPP
D
Variants trained
Purpose
111M
20
2.19B
61
Fitting/evaluating TREC prediction
111M
80
8.76B
50
Fittingevvaluating
TREC prediction
111M
200
21.9B
21
Fitting/evaluating
TREC prediction
111M
320
35.0B
40
Fitting/vvaluating
TREC prediction
111M
1280
140.1B
11
Fittingovalluating
TREC prediction
266M
20
5.31B
25
Fitting/vvaluating TREC prediction
266M
80
21.2B
19
Fitting/vvaluating TREC prediction
266M
320
85.0B
19
Fitting/evaluating TREC prediction
266M
1280
339.8B
Fitting/vaaluating TREC prediction
610M
20
12.1B
205
Fitting/evaluating TREC prediction
610M
80
48.5B
53
Fittingevalluaiing TREC prediction
610M
82
50.0B
30
Mid- Mid-training data placement
placpeemmm tests (code blend)
610M
200
121.3B
14
Fitring/evaluating TREC prediction
610M
320
194.1B
5
Fitting/vvaluating TREC prediction
L.7B
20
34.3B
Fitting/vaauuating TREC prediction
L.7B
80
137.2B
11
Fitting/vvaluating TREC prediction
L.7B
160
274.3B
Fitting/evaluating TREC prediction
1.7B
320
548.6B
Fitting/evaluating TREC prediction
3.3B
20
66.5B
Fitting/evaluating TREC prediction
3.3B
23
76.5B
Fitting/evaluating TREC prediction
3.9B
234
909.2B
2
Mid-training data placement tests (math blend)
3.9B
239
923.4B
Continual PT data placement tests (math blend)
C EXPERIMENTAL DETAILS
Table provides details on the architecture for models used in the main experiments, while Ta- ble 4 provides, for each model scale and TPP, the dataset sizes used in training, and the number of training variations explored that scale (varying data placement strategy, or LR schedule and hyperparameters n, X, B). In total, 578 TRECs were computed for the main experiments.
Further details of the 3.9B model and settings for continual pre-training experiments are in Ap- pendix I. In the remainder of this section, discuss the main pre-training settings.
All trained models were GPT2-style LLMs (Radford et al., 2019) with ALiBi (Press et al., 2022) embeddings and SwiGLU (Shazeer, 2020) non-linearity. We use the AdamW optimizer. Following standard practice, we do not apply weight decay or bias to LayerNorm layers. AdamW settings are 31 0.9, 32 0.95, and 1. We report cross-entropy loss. By default parameterize with maximal update parameterization, μP (Yang et al., 2021), with hyperparameters set via proxy tuning, described below.
22

---

Published as a conference paper at ICLR 2026
Table Tuned hyperparameters for μPP proxy model
Wasse 8.67e-02 1.62c-02
Ckinput 9.17
Ocoulpul 1.095
For a given TPP, all models have the exact same warmup phase: linear warmup of the learning rate from 0 to the maximum value. In all runs (aside from training of the 3.9B model), warmup was 10% of the total steps. Learning rate warmup is standard practice in LLM training (Brown et al., 2020; Rae al., 2022; Biderman et al., 2023; Dubey al., 2024: Kosson al., 2024).
All models in the experiments were trained on Cerebras CS-3 systems. 610M-parameter 20 TPP models take roughly hours each to train on single CS-3.
Proxy model hyperparameter tuning. We now describe how we tuned the jP base hyperparam- eters (HPs). Our proxy model is 39M-parameter LLM with width dproxy of 256, depth of 24 layers, and a head size of 64. Tuning runs were conducted on 800M tokens with B = 256 sequences and a context length of 2048 tokens. Tuning was performed by randomly sampling 350 configu- rations base LRs, initialization standard deviations, and embeeding/outpuut-oon scaling factors. Table 5 gives the resulting perpefforming values, which used
tuned HPs.
is also worth noting that the LR values shown in Fig. 3, Fig. and the appendix figures are the base yP LRs before j.P-adjustment. Calculation of (Sec. 3) and the EMA coefficients (Eq. (2)) requires the adjusted IR (i.e., multiplying by dprecyy / ddmode)). Also, when LR decay used, reported LR values always refer to the peak/max LR of the LR schedule.
C.I PLACEMENT TESTS: EXPERIMENTAL DETAILS
For the placement tests in See. 2, we used the standard training and validation splits of the SlimPa- jama dataset (Soboleva et al., 2023), but with different weighting of subsets given in Table 2. For these experiments, we used a context length 8192 tokens, batch sizes of 126, base peak LR of n 1.62e-2 (the MUP proxy-model-tume LR), and weight decay of X 0.1. We used the Llama-3 (Dubey et al., 2024) vocabulary size of 128256.
C.2 TREC FITTING AND PREDICTION: EXPERIMENTAL DETAILS
For the experiments and analysis in Sec. 3, Sec. 4 and Sec. 5.1, we use context length of 2048 tokens and the GPT2 (Radford et al., 2019) vocabulary of size 50257. For these experiments, we use the default source weightings for the SlimPajama dataset.
D FURTHER DATA PLACEMENT RESULTS
We present additional experimental results supporting our finding that optimal data placement cor- responds to the lowest point the TREC. These results cover additional learning rate schedules, summarize placement effectiveness across blends, and extend validation to alternative metrics.
D.1 WITHIN-SCHEDULEE PLACEMENT OUTCOMES
Fig. 11 provides the decay-to-zero (D2Z) data placement results, along with the LR schedules for D2Z, 10x, and Step drop (after 70% training); these are the three schedules tested in the Sec. experiments. For D2Z, although the TRECs bend back up baseline the end expected, note the tenth data placement position still obtains the lowest average TREC loss, and consequently is the optimal data placement location.
Fig. 12 isolates the results of the placement tests. Here, the y-axis gives the code blend (CB) val- idation loss. In each case, the aggregate blend, corresponding uniform mix of the code and
23

---

Published as a conference paper at ICLR 2026
Docay-to-zend training. TREC, CB
schods curriculum tests
X10
Train lass
Leaming rate (after pP adjustment)
3.4
TREC
2.0
3.2
Val HQ data pos
1.5
3.0
2.8
LO
LR decay function
0.5
Docay-to-zere
10x
0.0
Jump to-1%
Ordered training tokens x10"
Ordered training tokens x10'
Figure 11: TREC-guided data placement: further details. Left: TRECs predict best data place- ment in terms of resulting CB validation loss, for a linear decay-to-zens (D2Z) LR schedule. Right: Plots of the D2Z, 10x, and Step LR schedules used in Sec. 2 (formulas in Table 1).
Step docay
decay
Decay-to-
Codl-blled (CB) validaaioo loss
2,38
2.34
2.36
blend (CB) validation
2:52
Codieblend (C9) (CB) validation
2.25
2,34
2.24
Ry position
2:28
By position
position general blend
2.32
All general blend
general blend
2.22
All aggregate bland
2.25
aggregate blend
aggregate blend
2,30
2.24
2.28
2.22
2.26
2.16
Ordered training takens x102
Ordered training tokens x10
Ordered training tokens x10
Code blend (CB) validation loss across different LR schedules and CB training- data placements. Same result data as in Fig. 1 (left), Fig. 2, and Fig. 11. Left: Step decay schedule, Middle: 10x decay, Right: D27. All 30 models that trained with data placement improve no placement at all ("All general blend" line), while placing in the optimal TREC-guided position always improves over the aggregate blend ("All aggregate blend" line, i.e., mixing the code blend uniformly across all training steps).
Figure 12:
general data, strong baseline: it only bested by placing the code data at the optimal placement position.
Fig. 13 presents results when all the above trained models evaluated on the general blend (GB) validation data, rather than the code blend. This experiment can be interpreted as form of ablation: if we replace GB data with non-GB (i.e., CB) data at given position during training, which position leads to the greatest degradation in performance on GB validation (i.e., where not to place)? In ef- fect, we are testing the necessity GB data at each position by observing the impact of its omission. The greatest degradation consistently occurs at the position of minimum TREC loss, thus validating the placement hypothesis (Hypothesis 1) through omission rather than commission of task-relevant data.
Key takeaway 4: The TREC minimum also the most important placement position for GB data.
D.2
TREC PLACEMENT HYPOTHESIS
Hypothesis 4: Across different LR schedules, position-wiss TREC loss predicts the effective- ness of high-quality data placement at corresponding positions.
This hypothesis extends Hypothesis 1, proposing that TRECs are not only useful within a single training run but that their absolute values are meaningful and comparable across optimization sched-
24

---

Published as a conference paper at ICLR 2026
Stop decay
decay
Docay-to-2ero
Genera-blend (GB) validation loss
Glenral-bieend (GB) validation lost
Gineral-blend (GB) validation los
2,82
2.72
2.76
2.80
2.70
By position
2.74
position general blend
position
2.78
All general blend
2,68
general blend
All aggregate blend
aggregate blend
aggregate blend
2.70
2.74
2.64
2.68
2.72
*g
2.02
2.66
2.70
Ordered training takens x10"
Ordered training tokens x10
Ordered training takens x10
Figure 13: General blend (GB) validation loss across different LR schedules and code-blend training-data placements. Counterpoint to Fig. 12 (same trained models), except now the y-axis gives loss on the general blend. Left: Step decay schedule, Middle: 10x decay, Right: D22.. Placing code blend (CB) data at the TREC minimum significantly impairs performance on the general blend.
ules. For example, in the context of Sec. the question is: does absolute TREC loss predict CB validation loss across Step, 10x, and decay-to-zero LR schedules?
If true, one could use predicted TRECs from multiple learning rate schedules (or weight decay/batch size configurations) to identify the globally best HQ data placement and optimization settings-i.e. the position, LR schedule, and other hyperparameters yielding the lowest TREC loss. Consider 5B HQ token budget alongside 45B baseline tokens. One could scan different schedules, find the decile with the lowest predicted (t), and insert HQ data there. Similarly, if HQ data arrives late (e.g., recent web data to "advance the model's knowledge cut-off" (Dubey et al., 2024)), one might select the schedule with the lowest TREC loss in the final decile, knowing this will maximize task performance.
Absolute TREC not perfectly predict HQ benefit
2:40
Decay-to-zero
deciyy
C8 val. loss when placing
2.35
Step decay
2.30
2.25
2.50
00
2.15
2.54
2.56
2.58
2,00
2.62
2.64
2.66
2.68
Aug. TREC loss train part
Figure 14: Correlation between CB validation loss and TREC loss, by placement position. CB validation loss from placement in each decile of training. versus the actual absolute TREC loss measured in that decile. Markers labeled with their deciles (1-10), dashed lines to show linear trends. While CB validation varies monotonically with TREC loss within one LR schedule, the correlation does not hold well across schedules.
Fig. 14 illustrates this idea. Each marker plots CB validation loss (y-axis) from placing HQ data in a training decile against the TREC loss (x-axis) in that decile. Markers are grouped by LR schedule, with dashed lines showing idealized linear trends.
Within each schedule, TREC loss monotonically predicts validation performance, consistent with Hypothesis 1. But across schedules, the alignment breaks down. For example, placing HQ data in the final decile of the 10x schedule yields lower validation loss than many decay-to-zero place- ments, despite higher TREC loss. Furthermore, large TREC differences sometimes translate to minor validation gains, and vice versa, indicating TREC loss alone does not capture all relevant dynamics.
25

---

Published as a conference paper at ICLR 2026
We revisit this in Appendix I, where our 3.9B experiments show similar disconnect: CPT segments with the lowest TREC loss (typically under high LR) do not always yield the best validation loss when HQ data is inserted. Interestingly, while TREC loss the placed data itself matches TREC predictions (from the general blend), the gains do not translate to the held-out validation sets. is worth re-iterating, however, that this generalization gap only occurs across optimizer configurations (e.g., different LRs or LR schedules), not within a given configuration.
Key takeaway 5: While TRECs reliably guide IQ placement within given optimizer configu- ration, their absolute values do not consistently predict optimal placement across configurations. Further work is needed to guide schedule selection for data placement.
Toward an explanation of the generalization gap. We can summarize the gap as follows: two schedules A and B may attain the same TREC loss at certain positions, yet A may generalize better than B. One possible explanation is that A and B can achieve similar TREC loss via qualitatively different mechanisms. Recent grokking studies (Power et al., 2022; Morris et al., 2025) show that models initially reduce training loss through memorization, and only later transition into grokking, where they learn structure that generalizes. Both mechanisms can yield equally low TREC loss, but grokking will produce stronger validation performance.
This provides natural interpretation of why placement at position 10 of the 10x schedule gener- alizes better than early decay-to-zero placement in Fig. 14, even though both achieve similar TREC losses: in the tenth decile of training, the 10x schedule evidently leverages memorization. Likewise, in the CPT LR sweeps, high LR appears to push the model back into a memorization regime, while low LR preserves grokked structure and yields better validation loss.
grokking and less
These observations connect to classical factors known to influence generalization (such as curva- ture, SGD noise, and basin geometry), which may vary systematically across LR schedules. If we can identify signals that distinguish memorization-driven vs. grokking-driven TREC improvements (e.g., the width of the TREC valley, or deviations between the TREC and the standard training-loss curve), we could potentially correct for these generalization differences and enable cross-schedule prediction.
E FURTHER TREC RESULTS
This section provides supporting analysis for the TREC behavior discussed in Sec. 3 of the main paper. We expand on trends observed across model scale, dataset size, and peeens-per-prrmmett (TPP), and include additional plots that quantify absolute loss drops and timescale effects. Our findings reinforce the role of the AdamW timescale 7 in shaping TRECs and offer deeper insight into how training dynamics evolve across compute regimes.
E.1 FURTHER SCALING RESULTS
The main paper mainly focused on the shape and position of TREC valleys, as these are most pertinent for optimal data placement. As part of those findings, we found that т and TPP both modulate the shape of the TRECs. We now examine the absolute magnitude of the TREC loss drops. As shown in Fig. 15, when models are trained at constant TPP (e.g., 20), the overall TREC trajectories exhibit similar total drops across scales (left). However, when we increase TPP while holding model scale fixed (middle and right panels), the absolute TREC drop shrinks.
This behavior is notable because it raises the hypothesis that at high TPP, models become more inertial or pisiibyy due to saturation or reduced plasticity less signal per training fraction. This aligns with recent findings compute budgets (e.g., Kumar et al. (2024); Springer et al. (2025)). From another perspective, it could also mean that lower-TPP training focuses more memorization, and fitting particular training examples, consistent with other recent findings (Morris et al., 2025; Jelassi et al., 2024), as discussed in Sec. 3 and Sec. 5.1.
overtraining and reduced update effectiveness high
Fig. 16 shows that even though the shape of the TREC shifts with r. the total (summed) TREC drop across training steps is fairly consistent (for a given TPP regime). Together with the above findings,
26

---

Published as a conference paper at ICLR 2026
Same TPP (=20)
same
TREC drop
STREC-drop
-TREC-drpp
TPP
Awanywa W
TPP
111M
3.1
20
20
266M
2.8
610M
320
2.55
320
1.78
3.0
TREC
TREC
2.0
3.3B
2.9
2,45
2.4
2.40
2.2
2.7
2.35
0.00
0.25
0..00
0.75
1.00
0.00
01.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
Fraction of training tokens, 2
Fraction of training tokens,
Fraction of training tokens, F
Figure 15: Absolute magnitude of TREC drops decrease with TPP. Plots show absolute umnor- malized TRECs, same data as in Fig. 4: Left: 0.3; TPP is 20 for all model scales, and absolute magnitudes of drops similar. Middle: т 0.021, all models 111M, and magnitude of drop decreases with TPP. Right: T = 0.105, all models are 610M, and magnitude of drop again decreases with TPP.
B varies 1=2.22e02 A=0.1
B=504
varies
A=0
B=504 y= 1.6.2-022 varies
Drop from maximum TREC loss
Drop from maximum TREC loss
0.00
Drap from maximum TREC loos
-0.01
-0.01
-0.02
-0.02
-0.02
-0.03
-0.03
(?)
-0.03
(F)
126 (0:026)
1305e-08 (0:421)
0.025 (0.421)
-0.04
252 (0.053)
-0.04
3.09 03 (0.210)
-0.04
0.05 (0.210)
504 (0.105)
60e-02 (0.105)
0.1 (0.105)
1008 (0.210)
-0.05
(0.053)
-0.05
(0.053)
-0.06
2015 (0:421)
-0.06
6.48c 02 (0.026)
-0.6
0.4 (0.025)
0.00
0.25
0.50
0.25
1.00
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
Fraction of training tokens,
Fraction of training tokens,
Fraction training tokens,
Figure 16: Total summed TREC drop invariant T (610M, 80 TPP). Same data as in Fig. 3, but showing absolute differences between the per-step and everall-maximun TREC loss. Within and across sweeps of B (left), (middle), or À (right) the total summed TREC drop is fairly consistent. Thus the more narrow the timescale, the larger the drop.
this suggests that v governs the width of the valley (how long data influences the model), while TPP sets its average depth (how strongly the model responds).
The upshot is that compute-efficient training regimes-like 20 TPP, where TREC valleys are sharp- may benefit most from intelligent data placement, Note such regimes are often used for frontier- scale training due to their compute-efficiencyy MoEs are also now often used for frontier-seale efforts, and in MoEs, expert parameters also relatively-low effective TPP (Sec. 5.1). TREC- guided data curriculums are therefore likely to offer benefits to frontier-scale training going forward.
E.2 BATCH SIZE
We now investigate how TRECs behave across a wide range of batch sizes. particular, we study regimes well above the critical batch size Bexit (McCandlish et al., 2018; Shallue et al., 2019; Merrill et al., 2025)-the the point beyond which increasing batch size significantly degrades loss as a function of total tokens trained.
'Hoffmann al. (2022) found the optimal model size Nog and dataset size Dap to scale roughly equally as compute increases, with the optimal P/N ratio around 20 TPP. Further studies have found similar re- sults (Besiroglu et al., 2024; Porian et al., 2024), and 20 TPP has become synonymous with conpute-optina training (Dey et al., 2023; Zhang et al., 2024b). Starting with Llama (Touvron et al., 2023), released models often trained for more than 20 TPP hecause smaller, overtrained models more efficient for inference.
27

---

Published as a conference paper at ICLR 2026
Following Bergsma et al. (2025a). estimate Bera ~ 2150 tokens for 610M models trained to 20 TPP (via a fitted power law in training tokens), and we sweep batch sizes from 63 to 8064 (over two orders of magnitude). In each case, adjust weight decay to maintain constant AdamW timescale of r = 0.421, allowing us to isolate batch size effects from timescale variation.
Similar TREC shape across 23
Very large B forgets training
B
Drop from maximum TREC loss
2.9
1008
-0.01
126
2016
2.8
252
4032
-0.02
504
3064
TREC
-0.03
2.7
12
-0.0
126
2016
-0.06
252
4092
2.3
8064
-0.07
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
Fraction of training tokens,
Fraction training tokens,
Figure 17: TRECs and batch size (610M, 20 TPP, r 0.421). Left: Absolute TREC losses across batch sizes; very large batches achieve worse loss, but patterns are somewhat similar. Right: Absolute TREC losses, but normalized so they have same maximum value; B = 8064 behaves quite differently than all other settings.
Fig. 17 presents the resulting TRECs. The left panel shows the absolute loss curves: as batch size increases, the total TREC drop seems to become somewhat shallower, especially once batch size exceeds Beris- The right panel aligns curves by their maximum values order to compare shapes directly. For B < 2016, the curves remain remarkably similar in shape and position. However, for 8064 (well beyond Beat), the curve diverges notably, indicating a qualitatively different training dynamic.
This divergence aligns with theoretical expectations: as exceeds Berit gradient estimates become increasingly redundant, reducing the marginal utility of each new example. From TREC perspec- tive, gradients from individual samples have less impact on the overall update, which focuses more on common features than idiosyncrasies of individual batches.
Viewed through the lens of TRECs, these results offer a novel diagnostic perspective on the hatch size scaling frontier. These findings also suggest that models trained far beyond Berit may benefit less from careful data placement, as individual batches contribute less distinct signal to model up- dates. On the other hand, in situations where we might wish to avoid memorization (Appendix A), using larger batches could help accomplish this objective.
Key takeaway 6: TRECs remain stable across batch sizes up to Berit but diverge significantly beyond it, highlighting diminishing marginal data influence large- batch regimes.
E.3 ADAM MOMENTUM B1 AND VELOCITY 32
To assess whether TREC shape is driven by other timescales in the optimizer, we sweep the AdamW momentum (81) and velocity (32) parameters while holding other hyperparameters fixed (r 0.210, 610M model, 20 TPP). We tried variety of settings, and report those that completed training successfully (i.c., without failure due to numerical instabilities).
Fig. 18 shows that, although absolute TREC losses vary with 31 and 32 (especially extremes such as 81 = 0 or 31 = 0.999), TREC shape remains remarkably consistent across settings. This especially evident in the right panel, where curves normalized. Only very large momentum begins to shift the TREC minimum to slightly earlier in training.
These results reinforce the conclusion from the main text: while standard settings of B1 and 8n affect training dynamics and absolute loss, they do not significantly alter TREC shape. The AdamW timescale (via 7) appears to be the dominant factor shaping TREC trajectories.
28

---

Published as a conference paper at ICLR 2026
shape
&. B
Large
slightly-carlier re-eval drop
2.70
Normalized re-vvalation loss
1.0
Re-evaluation loss
(1(3))
0.0 (0.95)
0.0 (0.95)
2.00
0.5 (0.95)
0.5 (0.95)
(0.95)
(0.95)
2.55
(0.99)
0.9 (0.99)
0.95 (0.99)
0.2
0.95 (0.99)
0.99 (0.9999)
0.99 (0.9999)
(0.9999)
(0.9999)
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
Fraction of training takens,
Fraction of training takens,
Figure 18: TRECs and Adam 31 and B2 (610M, 20 TPP, 0.210). Left: Absolute TREC losses for range 31 and 3g settings. Both no momentum (3₂ = 0.0) and too much momentum (31 = 0.999) have highest absolute loss. Right: Normalized TREC losses; despite wide variation in absolute loss, normalized TREC shape is remarkably similar across all settings, with only very large momentum beginning slightly shift the TREC minimum earlier training.
Key takeaway 7: TREC shape largely invariant to changes in 81 and B2, underscoring the dominant role of the Adam W timescale in shaping learning dynamics.
F FURTHER PREDICTION RESULTS
In this section, we provide additional details on our predictive framework for TRECs (Sec. 4), in- cluding the specific evaluation metrics used, derivations underlying key equations, and further anal- ysis across model scales and I.R schedules.
F.1 FURTHER PREDICTION DETAILS
When computing predictions from our analytical framework (Eq. (3)), we discard the initial EMA coefficient co, corresponds to the influence of the model's random initialization rather than any datapoint observed during training. We focus on the remaining coefficients 6 for i 2 1, which quantify the contribution of training updates the final model weights.
Pearson correlation (rp). Since we aim to predict the shape of TRECs (t) rather than absolute scale, we evaluate prediction quality using a scale- and shift-invarian metric: the Pearson correlation Tp between the predicted curve (f) and the true curve (f). All plotted curves are similarly normalized to emphasize shape agreement.
Empirically, we found that Pearson correlation better aligned with human judgments of prediction quality than alternatives like l2 loss, R2, MSE. is computed follows:
Tp (Loo(( - (4...) (Ero 122..)
(5)
-t) / (Cro - 12_
where pe., and Mč., are the means of the true and predicted TREC values, respectively.
This measure ranges from -1 (perfect inverse correlation) to +1 (perfect match), with 0 indicating no correlation.
Illustrative example. The middle panel of Fig. 5 iillustrates the full prediction process for Cyclic LR schedule. (For reference, Fig. 19, bottom left, shows the full LR trajectory.) We compute EMA coefficients c(f) from the LR and weight decay schedule (Fig. 5, left), apply Eq. (3) with a fitted exponent obtain (f) (Fig. 5, middle), and compare it to the actual TREC (Fig, 5, right). The predicted curve tracks the shape of the true TREC, particularly in later training stages where the EMA contribution becomes more pronounced.
29

---

Published as a conference paper at ICLR 2026
F.2 SCHEDULE HISTORY AND TRECs: Cyclic vs. WSD
In this section, we study the question: does the full LR schedule history matter, or is late-stage alignment of EMA and sufficient to determine TREC shape?
Motivation. Our predictive framework assumes that TREC shape is primarily determined by the EMA coefficients and the training fraction (Eq. (3)). However, real-world schedules can differ substantially in their early phases while converging later. To test the extent to which schedule history leaves residual effects on TRECs, we directly compare Cyclic and WSD schedules under controlled conditions.
Experimental setup. We train 6 100Mpparametrr models to 80 TPP on SlimPajama using identical peak LR, batch size, weight decay, and dataset. The only difference is the LR schedule:
The WSD schedule warms up, then maintains long flat LR phase, before decaying to zero.
The Cyclic schedule oscillates but aligns with WSD in the final 20% of training.
Because the batch size and dataset are fixed, the training fraction f at each step matches across schedules.
x10-s
WSD LR schedule
WSD EMA coefficients
WSD normalized TREC:
Learning rate (afer yP adjustment)
-005-03 (0.421)
7 (T)
Normalized coefficient values
1.0
17 (+)
(0.421)
1.0
8.09 (0.210)
0000-03 (0.210)
Normalized TREC
1.62e-02 (0.105)
62e-02 (0.105)
3.24e-02 (0.063)
0.6
24e-02 (0.053)
0.6
0.4
0.4
1) (r)
(0.421)
0.2
8.09e-03 (0.210)
1 1.622-22 (0.105)
0.0
0.0
3.24e-02 (0.063)
1.00
Fraction of training takens,
Fraction training takens,
Fraction training toleons,
-10-3
Cyclic LR schedule
Cyclic EMA coeflicients
Cyclic normalized TRECs
Learning rate (after jP adjustment)
Normalized coefiicent values
1.0
1.0
Normalized TREC
0.6
0.4
0.4
0:2
0.2
0.25
0..50
1.00
00.25
0.75
1.00
0.00
0.75
Fraction of training tokens,
Fraction
training tokens,
Fraction of training tokens,
Figure 19: Comparison of Cyclic and WSD schedules (610M, 80 TPP). Top Using WSD schedules: bottom row: Using Cyclic schedules. Left: LR schedule. Training shares the same batch size, weight decay, and dataset, differing only in LR function (cyclic vs. WSD). Middle: Corresponding EMA coefficients e(t). In the final 20% of training, schedules align in LR decay, resulting in identical EMA coefficients. Right: TRECs. In the final portion of training where EMA and training fraction align, TRECs also align closely, indicating that EMA coefficients and training fraction, rather than prior LR schedule history, predominantly determine TREC shape.
Results. Fig. 19 shows LR schedules (leff), EMA coefficients (middle)) and true TRECs (right) for both schedules (top: WSD, bottom: Cyclic). In the final 20% of training-wwii where LR schedules align-the the EMA coefficients converge exactly. Correspondingly, after brief transient period, the TRECs also align, despite the differences in earlier schedule history.
In summary:
30

---

Published as a conference paper at ICLR 2026
When EMA and align, TRECs align across different LR histories.
Prior LR fluctuations leave negligible residual effect on TREC shape once EMA and é match.
Key takeaway 8: These results reinforce that given similar amounts of pre-training, EMA coeffi- cients and training fraction are sufficient to determine TREC shape, supporting the generality of predictive framework across schedules.
F.3 ADAM AS THE LIMIT OF ADAMW WHEN - 0
We now explain why TRECs converge as À + 0, or equivalently, as the AdamW timescale 1/A + co. In this regime, the effect of weight decay vanishes, and Adam behavior approaches that of standard Adam.
Fig. (middle) illustrates this convergence at 20 TPP: the TREC for X = 0 (vanilla Adam) closely resembles that for = 0.001. Since both TPP and T are fixed, the training- fraction term im in Eq. (3) remains unchanged across settings, convergence must arise from changes in the EMA coefficients c(f).
Decay-to-zero EMA coefficients
10x decay EMA coefficients
10x decay mirrors Adam as A
Noralized EMA coefiicents
1.0
A
Normlizzd EMA confficients
1.0
0.8
0.0125
0.8
(r)
Normalized TREC
0.8
0.6
0.05
0.6
0.001 (42.1)
0.6
0.01 (4.21)
A (r)
0.4
0.2
0-4
0.1 (0-421)
0.4
0.0 (NaN)
1.0 (0.042)
0.001 (42.1)
(4.21)
0.1 (0 421)
0.0
6.0
0.0
1.0 (0.042)
training
Fraction training tokens,
Fraction training
Figure 20: Effect weight decay on EMA coefficients and TRECs. Left: Normalized EMA coefficients for the decay-to-zero LR schedule across decreasing À values. As X + 0. the EMA coefficients flatten toward the shape of the LR schedule (coefficients undefined at = 0). Cor- responding TRECs schedule. As weight decay decreases, the curve approaches the À = 0 (Adam) case. Right: TRECs under the 10x-decay schedule. As À - TRECs converge the 0 baseline, confirming that TREC shape becomes increasingly determined by the LR schedule alone.
shown in Fig. 1, middle. Middle: EMA coefficients for the 10x-decay
Fig. 20 provides direct evidence. As shown in the left, decreasing causes the EMA coefficients to flatten toward the shape of the LR schedule. We observe this convergence behavior both for the decay-to-zer schedule (left) and the 10x decay schedule (middle), with corresponding TRECs for 10x shown in the right panel.
fact EMA convergence can be derived mathematically.
Derivation: C -> 17s À O. To clarify the behavior of the EMA coefficients as À 0 (and thus т > oo), is helpful to separate the scale of the learning rate from its shape. We write the learning-rate schedule as
T) = i
where y captures the overall scale (e.g., , the peak learning rate) and f E [0, 1] denotes the normalized schedule shape.
Under this parameterization, the EMA coefficients can be written (up to an overall scale factor) as
T
= II (1-7\î)),
j=++1
where é denotes the coefficients with scale factors divided out.
31

---

Published as a conference paper at ICLR 2026
This form makes the limiting behavior transparent: as X O, the product term approaches and therefore
n-
Thus, in the zero-weight-decay limit, the temporal structure of the coefficients is governed entirely by the learning-rate schedule shape.
An analogous limiting behavior occurs as + 0 (i.e.. the overall learning-rate scale 7mma van- ishes), since the multiplicative term again approaches 1, yielding & + 1).
Moreover, the same structural behavior also emerges in the large-batch regime B + oo, with and À held fixed. Increasing B reduces the total number of optimizer steps T, so the product spans fewer terms and has less opportunity to compound decay. In the limit, the cumulative effect becomes negligible, and the coefficients inherit the learning-rate schedule shape.
We thank Fabian Schaipp for suggesting this scale-shape reparameterization, which clarifies these limiting regimes.
G FITTING THE TRAINING-FRACTION EXPONENT m
This section provides additional details about methodology for fitting and evaluating the power- law functional form for the training-fracee fraction exponent m" (Eq. (4)). We describe how we select fitting data, evaluate prediction quality, validate generalization across model scales, and assess how the fit transfers across learning rate schedules.
G.1 DATA FILTERING AND FIT CRITERIA
To ensure robust and meaningful fits, we restrict fitting dataset regimes exhibiting well- behaved training and stable TRECs:
Effective range: We include only runs with values between 0.001 and 1.0, corresponding to training runs where learning was stable and meaningful signal is present.
Excluding unstable hyperparameters: We discard configurations with extremely high or low learning rates weight decay values, which frequently result in divergence, loss spikes, poor convergence.
Batch size filtering: We remove runs where batch size exceeds the estimated critical batch size Berit, beyond which training enters distinct large-batch regime (see Appendix E.2).
These filters ensure that the resulting fit captures relationships in the regime of effective optimiza- tion, avoiding pathological outliers.
G.2 EVALUATION METRIC: R2 IN LOG SPACE
We quantify the accuracy of our predicted m values using the coefficient determination, R?, computed in log space.
Let }m²} denote the true optimal training-fracee fraction exponents, i.e., those where the prediction has the highest Pearson rp agreement with the true TRECs. Let {m,) be the values predicted by pawer-law fit. Then:
Σ, (log m, log mi;
Ra =1
(6)
E, (log m; log m
This log-space evaluation accounts for multiplicative relationships and downweights the influence of very large or small outliers.
32

---

Published as a conference paper at ICLR 2026
Table 6: TREC Tp rises with scale: R with true m² declines.
Table 7: Both TPP and 7 improve estimate of m²: : fit at
111M scale, evaluation (TRECs & m fits) at 610M.
Eval scale
m²: R
TREC: Yp
Fit of m = C. (TPP)) . (T) uses 7
m R²
TREC: Fp
+ uses TPP
111M
98.9% %
96.6%
266M
97.2%
97.5%
-5.9%
84.1%
610M
98.7%
98.4%
15.1%
90.6%
1.7B
89.0%
98.7%
X
90.9%
97.7%
3.3B
76.7%
98.6%
98.7%
98.4%
Power law preds & observed 266M data
Power law preds & observed 3. 3.38 data
Optinal Train-frac. exponent, n1*
30
TPP
TPP
Optimal Train-facc. expeeet,
20
20
30
20
+ 320
10
6
L.4 TPP
TPP
10
10 10"
10
10
10"
Figure 21: Generalization of m² power-law fits to larger models. We evaluate the power-law fit made 111M scale: Left: Fit applied to 266M model; right: Fit applied 3.3B model. While fit accuracy declines slightly at larger scale, TREC predictions remain strong overall.
G.3 FITS AT OTHER MODEL SCALES
Fig. 21 extends the evaluation of our m fit additional model scales not included in the primary figure (Fig. 6). Fit quality slightly degrades at the 3.3B scale.
However, even though the optimal m values exhibit growing prediction error (Table 6, middle column R results), the predicted TRECs still closely match the true TREC shapes (Table 6, right column T results)-whichl
what we ultimately care about. We attribute this resilience to (1) TREC predictions being robust to exact training fraction exponents, and (2) reduced noise in the reference TRECs at larger model scales, which compensates for slightly-worse m. prediction.
Table provides the ablation results mentioned in the main paper, showing how the power law performs when only function of T, only function of TPP, or when function of neither- i.e. using a constant tuned over all the fitting data.
G.4 FIT GENERALIZATION ACROSS LEARNING RATE SCHEDULES
All m power-law fits in the main paper are derived from training runs using a Linear decay-to-zero LR schedule. We assess whether these coefficients generalize other schedules by applying them to TREC predictions models trained with a Cosine decay schedule.
Using the Linear-based fit, we achieve TREC prediction accuracy of rp 94.2% on Cosine-
schedule runs.
Fitting a dedicated power-law using Cosine runs improves the prediction score to rp = 97.8%.
These results indicate that while the predictive framework retains strong performance across related schedules, schedule-specific fit yields better accuracy. Accordingly, when evaluating prior work that uses Cosine pre-training schedules (e.g.. Sec. 5.2 and Appendix H), we employ our Cosine- specific m fit in order to optimize our TREC predictions.
In summary, if practitioner wishes to use a more complex schedule (c.g., WSD, or a custom multi- stage schedule), the Linear-based fit provide strong out-of-the-box guidance. In particular,
33

---

Published as a conference paper at ICLR 2026
this guidance is far stronger than the default end-of-training heuristic. For large-scale production training, hypersealers could re-fit schedule-specific TREC predictors at small scale (analogous to how they re-fit scaling laws), and use these instead. Or, more directly, obtain TREC prediction by simply re-evaluating smaller model with matching т, TPP, and LR schedule.
Key takeaway 9: Power-law prediction of m" remains robust across model scales and related schedules, though schedule -eseuuee--ciiicccccce fits can improve accuracy. Even when m² prediction slightly degrades, the resulting TREC predictions remain highly accurate.
H EVALUATING PRIOR LLM RECIPES: FURTHER DETAILS
This section provides methodological details for how we predicted TRECs for prior LLM training recipes, including Llama-3 (405B), OLMo-2 (13B), Feng et al. (8B), Pangu-Ultra (135B), and Nemotron-4 (15B).
H.1 METHODOLOGG FOR PREDICTING TRECs OF PRIOR LLMS
Pangu-Ultra: Learning rate
Pangu-Ultra: coefficients
Pangu-Ultra: HQ phases predicted
-10
Reasoning
Normalized EMA coefficients
1.0
1.0
Annealing phase
0.5
0.8
Predicted TREC
Learning rate
0.6
0.6
0.4
0.4
0.4
0.2
0.2
0.2
Reasoning phase Anmealing phase
Reasoning phase Annealing phase
0.0
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
Fraction of training tokens,
Fraction of training tokens,
Fraction of training tokens,
Figure 22: Predicted TRECs for the Pangu-Ultra training recipe. Left: LR schedule used in training the Pangu-Ultra model (Yin et al., 2025), including distinct "reasoning" and "annealing" phases. Middle: EMA coefficients computed from the LR schedule, batch size, and weight de- cay. Right: Predicted TREC. Both the reasoning and annealing phases-where high-quality data is introduced- align with the predicted TREC valley, indicating the model is well-positioned to retain this curated data.
To enable TREC prediction for existing models, we reconstructed the learning rate (LR) schedules from published training recipes. For example, Fig. 22 (left) shows the full LR schedule for Pangu- Ultra (Yin et al., 2025). including the designated "reasoning" and "annealing" phases. From these schedules, and given all these recipes used the AdamW optimizer, we computed the corresponding Adam W EMA coefficients (e.g., Fig. 22, middle).
We also extracted dataset size, batch size, sequence length, and weight decay values from the training documentation. If batch size or sequence length varied over the course of training (as is sometimes the case), we used the values applied for the majority of training. Early-stage variations typically occur well before the predicted TREC valley and thus have minimal impact on our prediction. More- over, in continual tattinull-time interpretation, such variations do not materially affect the TREC prediction once progress is normalized to training fraction i..
To determine the appropriate training-fraction exponent m, we applied the power-law fit described in Eq. (4), using models trained with Cosine decay schedule (Appendix G.4). This was appro- priate because the prior work considered here-Llama-3, Feng al., OLMo-2, Pangu-Ultra, and Nemotron-4-all used Cosine LR schedules during most of pre-training. However, special mid- training phases that hold the LR constant or rapidly decay it to zero (as Fig. 22) means that not exactly Cosine schedule that is ultimately used over all of pre-training. Yet, while TREC prediction accuracy can benefit from a schedule-matched m recall that our earlier analysis has shown:
34

---

Published as a conference paper at ICLR 2026
m" is relatively robust to schedule type,
TREC prediction remains accurate even when m² is slightly off, and
our goal here is to assess general placement trends, rather than fit precise minima.
H.2 PANGU-ULTRA RESULTS
Fig. 22 (right) illustrates the predicted TREC for the Pangu-Ultr training recipe (Yin al., 2025), based on its published learning rate schedule and training phases.
The Pangu-Ultra recipe includes two late-stage phases specifically intended to improve model rea- soning and inslroction-oolloiinng capabilities. As described in the paper:
"In the second reasoning phase, we increase the proportion of high-quality and diverse mathematical and coding data-raising it over 60% of the corpus enhance the reasoning capabilities of Pangu Ultra. [..]] Moreover, LLM-generated synthetic data is widely incorporated enrich the corpus.
The third annealing phase designed to help the model consolidate and effec- tively apply the knowledge and reasoning skills acquired in the previous stages. Therefore, we place greater emphasis on instruction data, [...] [including] care- fully refined ... short and long ofuin-of-thuught responses."
As shown in Fig. 22 (right), both the reasoning and annealing phases occur during the TREC valley predicted by our framework. In some sense, this alignment between real-world data placement and our predicted optimal locations supports the validity of our framework. However, note the TRECs rise quickly during the reasoning phase, suggesting that shifting this HQ phase slightly earlier may have been beneficial. That is, small timing adjustment could have helped avoid placing expensively-curated data in the region where the model has already begun to stabilize, minimizing the risk of diminished impact.
H.3 ANALYSIS OF NEMOTRON-CPT (PARMAR AL., 2024) STRATEGY
A particularly informative paper to interpret through our framework is the recent work by Parmar et al. (2024). While their study includes number of thoughtful ablations, it lacks a unifying theory of data placement. As such, their findings can benefit from reinterpretation in light of our insights into TRECs and training curriculums.
The starting point of their study is the Nemotron-4 15B base model, pre-trained on 8T tokens (533 TPP). Aiming to enhance this model without restarting training from scratch, they explore continual pre-training (CPT) strateg--adding 300B further tokens, including 2.8B high-quality QA tokens mixed into a "high-quality" blend. The HQ blend is placed at the end of the CPT phase. Their main research questions concern: (1) the optimal onset point for introducing the HQ blend during CPT, and (2) the learning rate schedule to apply during CPT.
Although this work is framed CPT, we first note that the original model was trained using 10x LR decay schedule, which is suboptimal relative to decay-to-zero (D2Z. (Bergsma et al., 2025b). Indeed, their own results show D2Z is a better schedule: continuing with the same pre-training data distribution but using D2Z during the 300B token extension improves downstream accuracy from 48.9% 51.5%- 5.3% relative gain over the base model, using less than 4% extra data. In other words, decaying to zero raises accuracy much more than would be expected by instead simply extending the 10x schedule over the same number of tokens. While placing HQ data boosts performance further, the total improvement naturally reflects the combined effects of both D27. and HQ data placement. Crucially, if Parmar et al. were to train further on an additional 300B tokens, they would not benefit again from applying This is precisely why we define methods that place data before the LR has fully decayed to be mid-training rather than true CPT (Sec. 5.3): true CPT methods can be applied repeatedly (e.g., repeatedly re- re-warming and re-decaying the LR), while mid-training benefits, like those in Parmar et al. (2024), can only be obtained once.
benefit would already have been consumed.
Second, by placing HQ data at the very end of training, their ablations of LR schedules are intrin- sically confounded. For example, they find that 100x decay performs marginally better than D27.
35

---

Published as a conference paper at ICLR 2026
during the "CPT" phase. But have shown, D2Z TRECs rise sharply near the end, so late-stage data is unlikely to be retained. By not decaying fully, they may make better use of the (sub-optimally) placed data. Similarly, they observe that the WSD schedule (Hu et al., 2024; Hägele et al., 2024) underperforms compared to Cosine decay, and "hypothesize that in continued pre-training, switch- ing the decay schedule from the one used during pre-training is harmful." However, this explanation is doubtful: different schedules yield different EMA coefficients, and thus different TRECs, Their effectiveness therefore depends not on matching schedule shapes, but on whether the HQ data is placed near the TREC minimum.
Third, their experiments are affected by data repetition, which further complicates interpretation. Repetition is known to degrade pre-training quality (Hermandez et al., 2022), and this likely applies even high-quality data (Team, 2025a). While the subsequent work by the same team (Feng al., 2024) controlled for repetition more carefully across training scales, the original Nemotron-CPP study in Parmar et al. (2024) allows repetition to grow significantly as the size of the CPT phase increases from 100B to 300B to IT tokens. Specifically, the 2.8B QA tokens represent 10% of the HQ blend, implying roughly 10.7x repetition if used for the full the 300B-token phase and 35.7x in the IT-token case. Their results confirm the downside of this: training with HQ data for the full CPT phase yields 53.6% accuracy at 100B tokens, but only 52.8% at 300B i.e. longer training hurly when it induces excessive repetition.
For all these reasons, it is unlikely that some of the paper's 's specific conclusions will generalize. For instance, their recommendation that "the switch [to the high-quality] data distribution should occur at Thwa//5 in the learning rate schedule" is not robust across LR schedules or training regimes. More generally, our framework offers principled alternative for determining optimal data placement that accounts for optimizer dynamics.
I FURTHER DETAILS ON 3.9B CPT EXPERIMENTS
I.1 EXPERIMENTAL DETAILS
Table 3 and Table 4 provide high-level architectural and dataset details for the 3.9B model experi- ments. This model generally follows the experimental setup of Appendix C, with a few noteworthy exceptions detailed here. The 3.9B model is GPT2-style L.M (Radford et al., 2019) using AL- iBi (Press et al., 2022) position embeddings and a squared- squared-ReLU (So et al., 2021) activation func- tion. As shown in Table 3, it employs an unusually wide FFN dimension (8 (8xdmodel), which we found to be effective in early experiments.
Both pre-tnaining and CPT phases use the AdamW optimizer with 81 0.9, B2 0.95, and E = 10 10-16, 16, The pre- pre-training warmup spans 375M tokens (not 10% of total tokens), followed by linear decay zero learning rate-also used in the CPT phase. Training is parameterized via the CompleteP variant of jP (Dey et al., 2025), with hyperparameters tuned using a depth-32 proxy model. The model uses context length of 8192 tokens and batch size of 672 sequences. Weight decay is set to À = 7.9 x 10-4 based on a projection for optimal T following Bergsma et al. (2025a). We also apply the layerwise weight decay correction from Dey al. (2025, Table 1). The peak learning rate during pre- pre-training is 17 = 0.15 (selected via proxy tuning). and we compare three different CPT learning rates (10% 3%, and 1% of the peak PT learning rate); the specific results for each are noted in the corresponding figures.
The data blends used in these experiments are listed Table 8. The general blend (GB) is used ex- clusively during pre-training and as background data in CPT. The math blend (MB), which heavily up-weights Open WebMath, serves as our designated high-quality (HQ) data and is inserted at dif- ferent positions during CPT to assess TREC-guides placement strategies, noted in Sec. 5.3. The MB phase is 234 steps, which comprise the final 234 steps of the 3303-step CPT phase when placing at the End, while placement at Half starts step 1980. Placement results were given previously in Fig. 10.
I.2 EFFECT OF LEARNING RATE ON TREC LOSS DURING CPT
Fig. 23 (left) shows the predicted TRECs for the 3.9B model (PT+CPT), zoomed-in to show the final portion of pre-training and including the CPT phase (shown in blue shading). The depth of the
36

---

Published as a conference paper at ICLR 2026
Table 8: SlimPJ and other sources used in training the 3.9B model. General Blend (GB) is used in pre-training and continual pre-training phases, while Math Blend (MB), serving high-quality data via up- weightngg Open WebMath, is evaluated after placing particular locations during CPT.
Subset
General Blend (GB)
Math Blend (MB)
GitHub
3.83%
3.07%
Books
3.62%
2.89%
ArXiv
4.21%
3.37%
Wikipedia
3.16%
2.53%
StackExchange
2.67%
2.14%
Fineweb-
64.75%
51.80%
Cosmopedia
4.66%
3.73%
OpenWebMaid
1.88%
28. 18%
UltraTex(Books2.00
0.42%
2.29%
StarCoder 10.79%
0.0%
controls emphasis of PT vs. CPT
Larger LRs - lower TREC loss
1.78
Absolute TREC
1.77
W
1.78
0.0015
1.75
0.0015
0.015
CPT phase
1.74
0.015
0.80
0.95
1.00
0.00
0.75
training tokens,
Fraction CPT training tokens,
Figure 23: 3.9B CPT: higher y in CPT leads to lower TREC loss. Left: Predicted TRECs for combined PT+CPT (zoomed into final 20%). Right: Actual TRECs in CPT phase: 17 0.015 has largest drop.
TREC valley varies with the CPT learning rate: 7) increases, the minimum TREC loss becomes deeper. For both SFT and CPT contexts, we hypothesize that TREC predictions can help suggest the optimal post-training learning rate (or least, the optimal range over which LR should be swept). We are exploring this further in ongoing work.
The observed TRECs from the actual training runs (Fig. 23, right) match the predicted trajectories (in blue shaded area) closely.
Interestingly, although 1) 0.015 produces the deepest TREC valley, results in the worst validation performance (Fig. 10). This illustrates a key failure mode for Hypothesis 4: while TREC-guided placement robust within a given training schedule, does not reliably generalize across learning rate schedules.
We investigate this issue further in Fig. 24, where we now also assess the TREC losses on the placed data segments-ing the HQ data that was positioned at specific CPT points (either halfway through CPT or the end). Fig. 24 (left) shows that the TREC loss on the general blend (in the same segment) correlates nearly perfectly with the loss on the placed math blend, regardless of the LR or where the data was inserted. This confirms that TREC behavior generalizes from homogeneous to heterogeneous data schedules: the curve retains its predictive structure when the inserted data differs in content.
Given this strong correlation, the critical question becomes whether performance on the training HQ data (math blend) generalizes to the validation HQ data. This is assessed in Fig. 24, right. Here, we observe two distinct regimes. For CPT runs using lower peak learning rates (n < 0.0045), there is a strong linear relationship between TREC loss and validation loss on the math blend. However, at 1 0.015, validation performance deviates from this trend entirely.
37

---

Published as a conference paper at ICLR 2026
placed
regimes
1.835
1.791
TREC loss placed data
0.0015
0.0045
1.750
1.825
0.015
Validation loss on math
1.780
1,820
1.815
1.810
1,787
0.00015
1.806
L.786
0.015
1,800
1.78
TREC loss general blend
TREC loss on general blend
Figure 24: 3.9B CPT: Predictive power of TRECs across LR schedules. Left: TREC loss positions in general blend (training without placement) correlates very well with TREC loss on the inserted/training math blend (i.e., when this math blend data is inserted those positions, i.e., Half and End placement). Right: TREC loss on positions general blend correlates with validation loss on math blend data (i.e., when separate training math blend data is inserted those positions) when 1j <= 0.0045; large 1) seems to induce another TREC loss regime, with different correlation with validation performance.
We speculate that high learning rates may push the model into a different region of the optimization possibbly one that emphasizes memorization shallower features- the cor- respondence between training and validation loss. This phenomenon is intriguing and merits further study in future work.
J THEORETICAL ANALYSIS OF TRECS
J.1 MOTIVATION AND KEY CONCEPTS
Sec. 4 showed that TREC shapes are predicted well by the AdamW W EMA coefficients, but only after adjusting for training fraction. Why is this adjustment needed?
Our key hypothesis is that the effectiveness of gradient update depends on where in parameter space the update is implicitly applied. An update that was useful when computed may be less effective later in training if the local minimum for its batch has shifted. We argue that the pace of this minimizer drift largely scale-invarian and governed by the training fraction.
More formally, each training batch B,, as an empirical sample, defines corresponding loss surface ((B,, 0) over model parameters This surface fixed by the batch itself (the data, the model architecture, and the parameter space do not change), but our position on the surface, 0,, changes over training. Fig. 25 illustrates: at t = 100, a gradient on B, moves toward that batch's 's minimum. At T 1000, reapplying the same update from different position increase loss, because the optimal direction has changed.
This perspective follows naturally from the EMA view of AdamW (and SGD): the final parameters 0 are a weighted sum of earlier updates, effectively "applying" each update to the final model state, even if the update gradients more aligned with 0r, and hence the gradients on these batches retain more of their original usefulness. That is, when applied at step T. these gradients are more effective in terms of lowering loss on their original batches, which manifests as lower loss upon re-evaluation (using parameters 0r).
computed long before. Batches nearer the end of training tend to have
In the remainder of this section, we formalize this idea using simple quadratic model that allows us isolate the factors EMM weighting and minimizer drift that govem TREC loss. We proceed in three steps:
1. Begin with SGD, showing how it accumulates local minimizers via EMA-like dynamics.
2. Derive TREC loss function of EMA weights and minimizer drift.
38

---

Published as a conference paper at ICLR 2026
{(B,,))
((B,,)) 0)
Batch B,, t=100
((Br,0) 8)
Params 0,, i varies
e(B,, 0)
Pr
0r-1
((B2, 0)
01 O2
0r
0
0
(a) Sctup: Each training batch defines distinct em- pirical loss surface over model parameters 0. We fo- cus on batch & 100 (blue curve).
(b) Fixed loss surface, changing 0;: loss on batch B, changes parameters updated (via gradients on other batches).
{(Br,)) 0)
Batch B,, 1=100 Params Be, 1=100
E(Br,
Batch B, t=100
Params 0r, T=1000
Vol(Be, 0.)
Or
(". 0.)
0' (B,,, 0T)
a Vo. 0.)
0
(c) Training Bs: at step t of training, following the gradient Vol(B,, 0) g==. specific and 8 -pppciificc local minimum.
(d) Re-evaluation: applying the earlier gradient to the final model parameters Or now hurts batch-specific loss.
0 toward batch-
Illustration of how parameters and gradients affect batch-specific loss: (a) batch B, defines its own unique loss surface E),,,0), with two local minima. (b) Each step of training defines a new set of parameters 0,, which determines the batch-specific loss. Now, the effectiveness of a gradient update to parameters 9 depends on O's current position: (c) during training, an update in the gradient direction (red arrow) is effective-it moves 6g toward the local minimum 0 (B,, 0z); (d) after training, when we re-evaluate B,, incorporating that same gradient update is ineffective now moves us away from the (new) local optimum, 8* (B,, 0r). Updates from batches seen later in training, when is closer to T, are more likely to improve loss on batch Br.
Figure 25:
3. Extend AdamW, W, where the EMA is over preconditioned gradient updates.
We then explain why the training-fraction exponent in our predictive form should be scale-invariant, and conclude by discussing how this viewpoint aligns with the empirical evidence presented else- where in the paper.
J.2 SETUP AND PRELIMINARIES
To gain insight into TREC loss, we adopt simplified analytical model similar in spirit to the work of Zhang et al. (2019), who derived closed-form convergence and loss expressions for various opti- mizers under different batch sizes. As their setup, we assume the optimizer dynamics are invariant to rotation and translation, allowing us to model loss as locally quadratic and separable across di- mensions.
Specifically, we assume that each training batch B, defines an empirical loss surface L(B,, 0) over model parameters We approximate this surface as locally quadratic the current point in train- ing." Let 0* (B,) denote the local minimum. In D dimensions, the loss at step t, given model
While LLM training minimizes cross-entie coss-entroppy loss, is common to perform local quadratic approxima- tion, i.c., a second-ord order Taylor expansion in the parameters, with the constant Hessian replaced by the instan- taneous Hessian along the training trajectory (LeCun al., 1989). Thus conclusions drawn from quadratic models often generalize large, realistic networks (Zhang al., 2019).
39

---

Published as a conference paper at ICLR 2026
parameters 0,, is:
1 L(B,, 01) = -B"())) Hr (0, (B.)) Σ (B, 0((0),
(B,,
d=1
where H, = (ha ..... specific loss curvature), and
1 is diagonal positive semi-definite Hessian (reflecting the batch-
(B, 0!
(B,
n
(8!0) - (0_)) -B)(((B)))
(8)
is the per-dimension contribution to the loss.
The key idea: "optimal" parameters can change. To streamline the analysis. we now focus on a single dimension and drop the superscript notation. Unlike Zhang et al. (2019), we explicitly consider that the locally-optimal parameter 0" for a batch B, may depend on the model state at a given training step, i.e., what's optimal depends on where we are on the loss surface. We therefore denote this local optimum as (B,, 0.): the local minimizer for batch B, as computed at step 8; think of B, as defining the loss surface, while 0, defines our current position on it-and the local minimizer depends on our position. This is illustrated in Fig. 25, where the loss surface is fixed, but the local minimizer is different step 100 (Fig. 25c) and step 1000 (Fig. 25d), as our position on the loss surface changes. Furthermore, for simplicity, we assume the batch-specific curvature h remains fixed across steps.
At each step of training, the loss is evaluated at the current model parameters 0g (our current position on the loss surface), the standard training loss curve is:
1 E(B,, 0z) 2 h4 10.000. - 00 (B.,
However, for TREC, compute loss using the final model parameters 0r (our final position on the loss surface), i.e., using the same empirical loss surface, but positioned near a (potentially) different local minimizer. Thus the TREC is:
(t) = E(B,, 1 he (0r 0" (B,,0T))
(10)
We have endeavored to make this key concept clear because this distinction-berwerr what's opti- mal the time batch was seen (Eq. (9)) and what's optimal when re-evaluating it later using the final model parameters (Eq. (10))is central to understanding TREC dynamics. We will now ex- amine how minibatch SGD accumulates these resition-specilee minimizers and how temporal shifts between 0 (B,, 0,) and 6* (B,, 0r) affect the TREC loss.
J.3 SGD: TREC LOSS AND EMA COFFFICIENTS
We now derive the TREC loss under vanilla SGD and show how it relates to the EMA coefficients and shifts in local optima over time.
The gradient of the quadratic loss surface at step f is:
Vol(Bi, h4 (0. - 0 (B,, 0.)).
(11)
and the parameter update under SGD is:
4+1 - T/ VoL(B, 0,) - (0. 0r 0,)),
(12)
where 7/t is the learning rate at step l.
This can be rearranged as:
P111 (1 - mhhh) O2 + nche- 0" (B., 0t),
(13)
which is equivalent to an exponential moving average (EMA):
y/ (1 0z 1 + 04
(14)
40

---

Published as a conference paper at ICLR 2026
where 8,, a nche and 2e 8" (Bi, 0e) (assuming the LR 7) is chosen that the sum is stable, i.e., 0 I 04 $ 1).
Just in the AdamW EMA case (Sec. 4), unroll the recursion and explicitly compute the contribution of each step to the final parameter value:
T Os
TT
II (1 (1-a)))
01(1) = Σ (0)
(15)
/=+11
where c; a; IT)=411( -a)).
As noted above, we are primarily interested in the TREC loss of batch Br, evaluated at the final model parameters 0r, on a loss surface with minimizer 0" (B,, 0r) (Eq. (10)). Substituting our EMA expression for Or into Eq. (10):
- (
1 2
(BB.
{BB 0r)
h4
(16)
To isolate the contribution from i
E(B,, 0T) = -.200 l4 (
Ea0((0) (1)
+ c.0 (B, -B) 0r)
(17)
igft
In the special case where 8* (B, 0() H" (Be, Br) (i.e., the local minimizer has not changed):
1 (BB,, 0T) h4
+ (c - (1,-1)0.BBB.0
(18)
The loss in Eq. (18) is fully minimized when c; = 1 and all other C; O. More generally, the steps with the highest EMA coefficient will obtain the lowest loss: with static local minimizer, the EMA coefficients fully define the TREC trajectory.
However, in practice, the local optimum for batch may drift over the course of training. We model this by assuming:
0" (B.,, 0.) rr 0" (B., 0r):
(19)
with F2 R scaling factor (possibly < 0). The TREC loss becomes:
E(B,, 0r) = 2 he (I Essr( (0) cift
+ -crt-1100
- (B,, 0r)
(20)
This expression emphasizes that the TREC loss is minimized when Cer2 1 (and the sum over j vanishes), i.., when the model fully incorporates the local minimizer for batch t via a high ct, and this minimizer well-aligned with the one the final step T. This analytical result mirrors our empirical finding (Sec. 4) that FMA coefficients alone are insufficient predict TREC loss: early gradients may have high coefficients c,, but the local minimizers of their corresponding batches can drift substantially over training (changing rt), leading poor retention.
J.4 ADAMW: TREC LOSS AND PRECONDITIONED UPDATES
We obtain a similar result when analyzing AdamW with the same locally-quadratic loss surface. As derived in Sec. AdamW parameters the final step T can be expressed as convex combination of weight updates:
T
f,
Σ
Ca
(21)
41

---

Published as a conference paper at ICLR 2026
where each update zy has the form:
1 my
(22)
X
V0l +
and each coefficient is defined as:
T
C =
II (1 (1-1171 1z211 (1--)))) ndd. j=++1
(23)
Substituting this decomposition into the TREC loss expression, we obtain:
e(B,, 02) - hes (1x
C Eylh
+ +Ca 8" (Be, (r)
(24)
TREC loss is minimized when Ct is large and the update ia points the same direction as final-step local minimizer 0* (B,, 0r). In other words, it is primarily the sign aligument of xa and 8" (B, 0r) that determines whether the update helps or hurts TREC loss.
In the quadratic model (and optimization invariant to translation), can assume without loss of generality that 0g = at the time of the update. The gradient becomes:
VgE 0z) = h1 (0: - 0* (B., 0.)) = h₂0 (B.,0.)
(25)
and assuming momentum for simplicity, the update is:
1 h (B..
vur + k2 (Be, 0z),
(26)
where 2 0 absorbs the preconditioning, curvature, and scaling terms.
To model drift in the loss surface, we again assume that the local minimizer for batch changes over time and satisfies:
(B,, 0₂) = r2 -0" (B,, Or).
(27)
Substituting the change in local minimizer (Eq. (27)) into our simplified expression for the update (Eq. (26)), and using this update in the TREC. loss equation (Eq. (24)) yields:
E(B,, 0T) = he (I c: + -ammr 110 0T)
2
1
(28)
This loss is reduced when the AdamW update contributes value that matches the current minimizer direction, where r, controls the update direction (since c, 2 and & 2 0). As before, loss is max- imally reduced when e and the weight on all other updates vanishes. Loss also maximally reduced when the product CeKer Re 1. In practice, if the local minimizer has shifted substantially especially if r, 0-then the early update may push in the wrong direction, increasing TREC loss despite having high EMA coefficient.
This again echoes the central insight: an update's value determined not only by its EMA weight ce but by its alignment with 0y 's position in the loss landscape.
J.5 LOCAL MINIMIZER DRIFT AND THE TRAINING-FRACTION CLOCK
In the preceding analysis, the term 7t captures the effect of local minimizer drifi-the the change in the batch-specific optimum 8* (B,, between when the gradient was computed (at step f) and when it is effectively applied (in the final model 0r). We now explain why the natural coordinate for this drift is the training fraction f t/T.
In the SGD case, the final parameters fr can be written as an EMA with sime-varying smoothing Of u,h,, where Ne is the learning rate and ht the curvature along the current parameter direction
42

---

Published as a conference paper at ICLR 2026
(Eqs. (13) and (15)). When a values are larger, more weight is put on recent minimizers (shrinking the EMA timescale), and parameters evolve faster from step to step. In particular, the pace at which the final parameters 0T move away from earlier parameters 0, is directly governed by the sequence (ns ) We naturally assume that larger movement in parameters coincides with larger move- ment in local optimizers, i.e., that the distance between 0* (B,, 0.) and 0* (B₁, 0r) is also governed by this sequence.
In AdamW, the same qualitative picture holds: the learning rate enters directly into the EMA co- efficients (Bq. (23)), while the curvature appears as a scaling factor in the updates (Eq, (26)). Although preconditioning partly normalizes curvature variation, the product of learning rate and ef- fective curvature still governs the rate at which parameters update, and hence the rate at which local minimizers drift.
Recent work by Noci et al. (2024) shows that, under maximal update parameterization (JPP), curva- ture statistics such as the largest Hessian eigenvalues evolve in nearly ceale-independent way, a phenomenon they term super-consislency. Qiu et al. (2025) build on this observation, hypothesizing, and providing empirical evidence, that a related curvature proxy depends only on training fraction (what they term "normalized compute") and is largely scale-independent under μP, That is, if we align two models of different sizes by their training fraction t/T, their curvature trajectories h(t) are nearly identical.
If both the curvature h(f) and the learning rate schedule (f) are indeeinndeenneent functions of f, then their product (ff)h(f)- which governs the pace parameter drift in the quadratic model-is also scale-independent. Consequently, the sequence of curvature/leaming-raa conditions experi- enced by a small model
0.3 is essentially the same as that experienced by a large model at = 0.3. Based on the prior work in μP and our own empirical results (discussed below, Ap- pendix J.7), we speculate that in more complex settings, other aspects of training dynamics also evolve with training fraction. For example, in μP neural networks trained with cross-entropy loss, the extent of feature learning and the evolution of representations earlier layers should also evolve invariant to model size. This yields tuaiing--ruction clock for minimizer drift: relative training progress, not absolute tokens, is the natural coordinate for r in re-evaluation dynamics. In this view, r, should be written rff).
J.6 FROM CURVATURE TO FUNCTIONAL FORM
Our quadratic-mont analysis suggests that the product leeaaiii) rate times curvature sets the magnitude of parameter updates, and hence the rate at which local minimizers move. A gradient computed at training fraction will remain well aligned with the final model parameters Or only if the subsequent minimizer motion between { and the end of training is small.
Conceptually, we can think of a drift rate at each point in training, proportional to rf))ii(), which measures how fast the optimizer "forgets" earlier minimizers. The cumulative drift that erodes the utility a gradient from fraction is then the integral of this rate over the remaining training.
r(f x [. (())((s))) ds.
Empirically, we find that r() is well fit by a power law im as used in our predictive form for the TREC (Eq. (3)):
-0 =1 - pm
where e(t) is the AdamW EMA coefficient and m captures the cumulative effect of curvature- and LR-driven drift along the training-fraction
fraction clock.
J.7 EMPIRICAL EVIDENCE FOR VARIATION IN
Interpreting the fm term cumulative drift means that larger m values (higher drift) correspond to gradients losing effectiveness more quickly; higher drift affects the TRECs by increasing the span of "forgotten" iata-i.e., the data with baseline higher TREC loss (regardless of e(l) coefficient).
43

---

Published as a conference paper at ICLR 2026
From analysis, m should change under interventions that modify either the eurvature/leaningg rate product n)(t)((t) or the optimizer's 's effective memory. Consistent with this prediction, we ob- serve systematic variation in m across multiple settings presented earlier in the paper:
Lower EMA timescale. Shortening the AdamW EMA timescale decreases the optimizer's 's ef- fective memory, and therefore increases the rate-of-change of parameters. This should lead to faster reductions in gradient effectiveness, with an effect TRECs independent of the separate, analyti- cal effect of a smaller T on the EMA coefficients themselves. Indeed, we observe higher optimal m values for lower timescales in our experiments (c.g., Fig. 6). Moreover, our m power law fitted at 111M scale predicts well across model sizes: the r-governed drift is scale-invariant.
Higher tokens-per-parant parameerrr (TPP). We also see small but consistent increase in optimal m when increasing the TPP ratio (also shown in Fig. 6). Note the normalized LR schedule n))) does not change with TPP. We therefore hypothesize that the degree of overtraining (as defined by TPP) slightly affects the pace of change curvature (as a function training fraction), which drives the pace that earlier gradients become ineffective. Yet it remains notable that these changes in m are solely TPP-dependent and do not depend on the absolute number training steps (which scales with model size for given TPP).
Fewer MoE experts. In sparse Mixture-of -fe---f--E--ee Ef-extsts (MoE) models, reducing the number of experts increases the number of tokens processed by each expert's parameters, raising the expert's 's effective TPP. The impact on closely matches the shift observed when directly changing TPP in dense models, suggesting that MoE layers inherit the same drift scaling in TPP as non-expert dense layers.
CPT vs. from-scratch training. Recall Fig. 7 comparing two 3.9B models trained with identical configurations- same data, batch size, weight decay, and an identical learning-rate schedule ry(()= but different starts: one from random initialization and one from a near-optimal checkpoint (i.e., this variant undergoes continual pre-training CPT). Although rf() and EMA coefficients c(f) were iden- tical in both runs, the two trajectories occurred different points in training fraction and thus likely under different curvature statistics h(t). Beyond curvature variation, changes in the extent of feature learning and evolving representations in earlier layers may also play a role. In deep network, early updates are large compared to weight magnitudes (despite preconditioning) (Kosson et al., 2024), and thus optimization steps can produce disproportionately large shifts the parameters, and thus the optimal batch minimizers. Empirically, the CPT model, starting in a region already well aligned with the final parameters, exhibited little minimizer drift: its tdh term remains close to I, producing a TREC prediction closely following the EMA coefficients. In contrast, the scratch model traversed a much larger region of the loss landscape, with substantial drift that reduced aligmment to early updates; its fr term rises gradually from 0 to 1, with corresponding TREC valley only aligning to the EMA coefficients near the very end of training.
WSD vs. Cyclic LR schedules. In Appendix F.2, we compared the WSD and the Cyclic schedules. Training runs with the two schedules completely align in step-wise LR, weight decay, and batch size-in the final 20% of training. In a way, this is a similar test to the CPT vs. from-scratch comparison above: two similar training runs, but initializing from different start points. In this case, however, the two models have undergone similar total amount of training during their initial phases, and consequently the initial gradients in the final phase have similar alignment with the final model parameters. Consequently, the final period spans the same range of training fraction, and their TRECs align very well over this final period. For the same peak LR, the WSD schedule does appear to drop slightly lower at around, e.g.. i = 0.8, which may reflect more mature model at that stage (perhaps because this schedule, unlike Cyclic, does not have the periods where LR returns to zero).
Different m for different learning rate schedules. While the minor differences between WSD and Cyclic are unlikely to be consequential in terms of data placement strategy, differences in the final stages of training have a greater impact.
As we noted in Appendix G.4, optimal tends to change when the learning rate decay pattern changes sechedulesppccfiii fits yield the best predictive performance). We also separately observed that Constant learning rate schedules produce the largest optimal m values. This is consistent with
44

---

Published as a conference paper at ICLR 2026
drift interpretation: without decay, rf() remains high even late in training, inducing substantial movement in fy and its associated local minimizer. Earlier gradients are therefore misaligned more quickly, requiring larger m to match the observed TRECs.
Across these cases, the observed shifts in m are consistent with changes in the effective 7(0)/() product. Together, these results suggest that functional form may exist for local minimizer drift that not only normalizes across scale (as does the training fraction term), but also across LR schedules. That is, the functional form for this drift could directly incorporate the LR schedule itself.
45


Reviews:
[{"review_id": "09GffhCPSw", "reviewer": "Reviewer_wgGG", "summary": "This paper introduces the Training Re-Evaluation Curve (TREC), a diagnostic tool that measures a trained LLM's retention of training batches as a function of their position, revealing optimal placements for high-quality data that outperform end-of-training heuristics. By modeling TRECs analytically via AdamW's EMA coefficients adjusted for training fraction drift, the authors enable proactive curriculum design, with predictions generalizing across scales from 111M to 3.9B parameters and 20-1280 TPP. Extensive experiments explain suboptimal strategies in prior works like Llama-3 and OLMo-2, while demonstrating TREC-guided placement improves continual pre-training in validation loss.\n\nOverall, I really like this paper. Meanwhile, as a practicer in this field, I have many questions and want to discuss with authors. I am looking forward to replies from the authors.", "questions": "All the questions are recorded in weaknesses sections.\n\nI greatly enjoyed this paper. If the authors can thoughtfully address and resolve my concerns, I would be delighted to raise my score to a 10. Conversely, given the fundamental nature of my questions—particularly the first one—if they remain inadequately resolved, I may lower my score to 6 or even below. I look forward to further discussion with the authors during the rebuttal phase.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper addresses a critical challenge in LLM pre-training from an innovative and unanticipated perspective, demonstrating strong novelty that sets it apart in the field.\n2. The authors conduct a comprehensive suite of rigorous experiments, involving substantial computational resources, which will provide valuable insights to researchers in the LLM pre-training community.\n3. I particularly appreciate the paper's clear and logical structure: it begins with well-defined hypotheses, builds through theoretical derivations, validates them empirically, and systematically distills a series of key takeaways.", "weaknesses": "1. **Difficulty Understanding Loss Variations in TREC**: What I find challenging to comprehend is the section defining TREC. For a given data distribution, different training batches should be nearly identically distributed (i.i.d.). Why, then, does the fully trained model exhibit such substantial loss differences across these batches—particularly in Figure 1 (left) and Figure 2 (left)? In my view, the losses should approximate a flat line parallel to the x-axis, with values closely aligning with the final train loss. This discrepancy warrants further clarification.\n\n2. **Scope of TREC's Applicability**: It appears that this principle applies only to validation sets identical to the training distribution, since TREC is computed on training batches. Is my interpretation correct? If so, does this imply that the optimization ultimately targets only validation sets matching the training distribution, rather than arbitrary domain-specific tasks (e.g., code)? Your statement on line 097—\"we hypothesize the TREC remains predictive even when data distributions vary\"—seems to address this concern. However, this hypothesis comes with no more explanation. Is this hypothesis overly ambitious, given the potential risks of distribution shift?\n\n3. **Inconsistent and Unclear Naming of LR Schedules**: I feel that your selection and nomenclature for LR schedules lack rigor and clarity, manifesting in several ways:  \n   a. Terms like \"10x\" and \"step decay\" are non-standard for LR schedules, and I have rarely encountered similar naming conventions in other literature. While I eventually deciphered \"10x\" via the appendix, this remains confusing and detracts from accessibility.  \n\n   b. \"Step-decay learning rate schedule\" is also uncommon; I am still unclear on its precise form. Does it entail a sudden drop in LR (e.g., at 70% of training) to 10% of its prior value? You reference it \"as used in DeepSeek LLM,\" but DeepSeek employs a \"multi-step cosine scheduler,\" which introduces ambiguity.\n\n   c. I strongly recommend adopting more precise terminology for your experimental LR schedules, ideally with explicit mathematical functions describing LR evolution at each step.\n\n4. **Summary of Predicted TREC Positions Across Schedules**: Allow me to summarize the optimal high-quality data placement positions predicted by your TREC across different LR schedules:  \n   a. Linear decay (commonly used): End of training (Figure 2 right).  \n   b. Decay to zero (commonly used): End of training (Figure 11 left).  \n   c. Decay to 0.1x LR (your \"10x,\" commonly used): Not explicitly described, but I infer it aligns with the end, given its similarity to D2Z.  \n   d. Step drop (uncommonly used): Pre-decay position (Figure 2 left), with the final val loss showing minimal gap compared to end placement.  \n   **Thus, can I conclude that, for the vast majority of common LR schedules, TREC predicts the end of training as the optimal placement?** This suggests some degree of over-claiming on your part, while also raising concerns about TREC's practical utility in experiments. I acknowledge that TREC offers mechanistic insights, but its impact on guiding real-world practice seems limited.\n\n5. **Alignment in Figure 9 and Practical Implications**: Following the previous point, in Figure 9, your TREC predictions align with prior findings, yet these still involve placing high-quality data in the final stages. Notably, in Figure 9 left, your prediction actually points to roughly the 0.5–0.9 fraction for inserting that 40% high-quality data, rather than purely the end (0.6–1.0 fraction). I am skeptical: under a common LR schedule, would shifting high-quality data away from the very end—per TREC guidance—truly yield superior performance **in practice**? This merits empirical validation beyond the current analysis.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "BQji5VRdoL", "reviewer": "Reviewer_MKNN", "summary": "This paper introduces the Training Re-evaluation Curve (TREC), a novel diagnostic tool for understanding data retention in Large Language Model (LLM) training. The TREC is computed retrospectively by evaluating the final trained model on the sequence of training batches it encountered. The authors hypothesize that the optimal placement for high-quality (HQ) data in a curriculum is at the minimum of this curve (the \"TREC valley\"), a point that often precedes the end of training.\n\nThe central challenge is that TREC is a retrospective measure. The paper's main contribution is demonstrating that the TREC's shape is predictable in advance. The authors show that TREC shape is governed by the AdamW optimizer's implicit EMA timescale and a \"training-fraction\" component. They develop a predictive model for TREC and show that its parameters can be fitted on small models (e.g., 111M) and generalized to larger ones (up to 3.3B).\n\nThe authors validate this framework by showing that:\n\n1. Directly placing HQ data at the TREC minimum in 610M model runs yields the best validation performance.\n    \n2. The framework successfully explains findings in prior work, such as the (ineffective) late-stage annealing in Llama 3 405B and the (effective) data phase timing in OLMo-2.\n    \n3. It provides a compelling explanation for memorization in Mixture-of-Experts (MoE) models, linking it to a reduced \"effective TPP\" for each expert.\n    \n4. TREC-guided placement improves performance in a 3.9B model continual pre-training (CPT) task.", "questions": "1. **On the 3.9B CPT Experiment:** In the 3.9B CPT experiment (Sec. 5.3), was the \"Halfway\" placement point proactively calculated using the power-law predictor (Eq. 4) that was fit on the 111M model? Or was this point chosen based on a retrospective analysis of a 3.9B run, or simply as an intuitive \"non-end\" baseline? Clarifying this is critical for understanding the full practical utility of the predictive framework at scale.\n    \n2. **On Cross-Schedule Comparison:** The finding that absolute TREC loss does not predict performance *across* LR schedules (Fig. 14, Key takeaway 5) is an important limitation. Do the authors have any intuition on what additional signal might be required to make such a cross-schedule prediction? For example, does the absolute value of the *training loss* (the orange curve) at the TREC valley, or the *width* of the valley, play a role in this disconnect?\n    \n3. **On Predictor Robustness:** The $m^\\*$ power-law was fit using a Linear decay schedule, and the paper notes a Cosine-specific fit is more accurate. How sensitive is the framework to the LR schedule shape? If a practitioner wishes to use a more complex schedule (e.g., WSD, or a custom multi-stage schedule), must they re-fit the entire power-law (which requires many small-scale runs)? Or is the provided Linear-fit predictor \"good enough\" for most practical purposes to identify the approximate valley?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. **Originality and Significance:** The paper's primary strength is its originality and practical significance. It defines a novel diagnostic (TREC) and, crucially, provides a complete, principled framework to *predict* it in advance. This transforms curriculum design from a heuristic, high-cost ablation study into a predictable, low-cost design choice. This is a significant step forward for the field.\n    \n2. **High-Quality, Thorough Experimentation:** The empirical work is extensive and of high quality. The authors support their claims with a large study of models (111M to 3.9B) , diverse experimental settings (pre-training, CPT, MoE) , and a large-scale analysis (over 600 TRECs). The \"data placement sweep\" (Fig. 2) is a direct, clear, and compelling validation of the core hypothesis.\n    \n3. **Presentation:** As noted above, the paper is exceptionally clear. The logical flow (Here's a tool -> Here's the problem -> Here's the solution -> Here's why it works -> Here's how to use it) is perfectly executed. The figures are exemplary.\n    \n4. **Actionable Insights:** The framework is not just theoretical; it provides immediate, actionable insights. The explanations for MoE memorization (linked to effective TPP) and the failure of Llama 3's late-stage annealing (linked to vanishing EMA coefficients) are powerful demonstrations of the framework's explanatory power. This gives confidence that the model is capturing a true, underlying dynamic of the training process.", "weaknesses": "1. **Limited Scale of Core Validation:** The primary weakness is that the main *predictive loop* (fit on small model, predict for large model, train large model based on prediction) is not fully closed at the largest scales.\n    \n    - The core *validation* of the hypothesis (the data placement sweep in Fig. 2) was done on a 610M model.\n        \n    - The *prediction* framework ($m^\\*$) was fit on 111M and *evaluated* on runs up to 3.3B.\n        \n    - The 3.9B CPT experiment is the closest to a full-loop validation. However, it's a CPT setting, not pre-training from scratch. Furthermore, it's not explicitly stated whether the \"Halfway\" point was *proactively chosen* by the 111M-derived predictor, or if it was just a reasonable choice *inspired* by the general shape of TRECs.\n        \n    - To be fully convincing, a stronger validation would be to use the 111M-derived $m^\\*$ predictor to find the optimal placement for a 7B+ *pre-training* run and show it outperforms the standard heuristic.\n        \n2. **Cross-Schedule Generalization Limitation:** The paper is commendably honest about this limitation, but it is a significant one. The framework can reliably find the best data placement *within* a given LR schedule. However, it *cannot* be used to compare *across* schedules.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "4DI6EO5ie7", "reviewer": "Reviewer_79pF", "summary": "This paper introduces the Training Re-Evaluation Curve (TREC) — a diagnostic that evaluates how well a trained model retains training batches as a function of when they appeared during training. TREC measures retrospective loss on training data using the final model weights and reveals that models tend to retain information from certain phases of training more effectively.\nThe authors show that placing high-quality data at the low points of the TREC (where retrospective loss is smallest) improves performance, contradicting the common practice of placing such data at the end of training.\nThey further demonstrate that TRECs can be predicted in advance from AdamW’s exponential moving average (EMA) timescale, enabling proactive curriculum design without retraining.\nThe paper supports its claims with both theoretical derivations and large-scale empirical results (from 111M to 3.9B parameters). It applies TREC analysis to settings such as sparse mixture-of-experts (MoE) models and continual pre-training (CPT), showing predictive consistency and practical gains.", "questions": "1. The setting for each experiments may need further clarification or better organized. The pretraining loss has a sudden drop at 70% (i.e. figure 1 most left sub-figure), does that correspond to the experiment with Jump-to-1% learning rate scheduler?\n2. Is there a specific kind of data (i.e. domains or other patterns) tend to have lower TREC? Or it is just related to the order across the whole training?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- The paper provides a great amount of experiments on models from 111M to 3.9B with several hundred billions of tokens.\n- The theoretical framing is well-grounded, extending prior work on AdamW dynamics and EMA analysis And the predictive model of TRECs is supported by strong correlations.", "weaknesses": "Since the experimental setup's setting defines HQ(high quality data) as a data with close distribution to the evaluation target ( e.g. code data vs general data and evaluating on code data, not general HQ data). Intuitively, place these data near the end of training will improve the performance as they are less likely to be forgotten by LLM, which is similar to annealing, and the results of the paper also show that place the data at the end of learning yields better results in linear decay cases.\nAlso for step scheduler, place the data at the end of learning (the last part before the lr step decay, as the remaining part may not change much on the parameters due to very low lr) yields better results. And TREC can be interpreted as a reflect how model remember/forget the data.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "OXmTD0lyaU", "reviewer": "Reviewer_YJr3", "summary": "This paper introduces the Training Re-evaluation Curve (TREC), a novel diagnostic tool for designing effective data curriculums for Large Language Models (LLMs). The TREC measures the retrospective loss of the final trained model on individual training batches, plotted as a function of when those batches were seen during training. The core hypothesis is that placing high-quality (HQ) data at the minimum of this curve—the \"TREC valley\"—maximizes its positive impact on the model's final performance.\n\nThe central contribution is a framework for predicting the TREC in advance, obviating the need for a costly initial training run. This prediction is grounded in the dynamics of the AdamW optimizer, modeling its parameters as an Exponential Moving Average (EMA) of weight updates. The authors extend this by incorporating a \"training-fraction\" term to phenomenologically account for \"minimizer drift,\" where the loss landscape shifts during training, reducing the effectiveness of earlier updates. The paper provides extensive empirical validation for models ranging from 111M to 3.9B parameters, showing that the predictive model accurately captures TREC shapes and that TREC-guided data placement outperforms baseline strategies. Finally, the framework is applied to explain and critique existing training recipes and to improve performance in a large-scale continual pre-training setting.", "questions": "see weakness", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1.  The TREC itself is a simple yet useful concept, and the ability to predict it proactively is a contribution.\n\n2.  The claims are supported by a large-scale study of over many TRECs across a wide range of model sizes and training configurations. \n\n3.  The framework successfully explains previously ambiguous empirical results from the community. For instance, it provides a clear rationale for why the late-stage annealing in Llama-3 405B failed to yield gains (the placement occurred on a rising part of the TREC) and why the mid-training interventions in models like OLMo-2 were effective.", "weaknesses": "1.  A significant scope limitation, explicitly noted by the authors (Appendix D.2, Figure 14), is that the framework can only identify the optimal data placement within a given learning rate schedule. \n\n2.  How to explain the schedule with the highest learning rate (η=0.015) produced the deepest TREC valley, which according to the framework's core premise, should have yielded the best results. Instead, it produced the worst validation performance. This directly contradicts the assumption that minimizing re-evaluation loss on training data (what TREC measures) always translates to improved generalization.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)