Scaling with Collapse: Efficient and Predictable Training of LLM Families
Authors: Shane Bergsma
Bin Claire Zhang
Nolan Simran Dey
Shaheer Muhammad
Gurpreet Gosal
Joel Hestness
Venue: ICLR 2026 Poster (2026)
Keywords: Training loss curve collapse
Compute-efficient LLM pre-training
Tokens-per-parameter (TPP)
AdamW EMA timescale
Learning-rate schedules
Scale-stable dynamics (μP)
Early stopping for hyperparameter tuning
Primary area: foundation or frontier models, including LLMs
TL;DR: We show that loss curves *collapse* across LLM scales when training at fixed TPP and with AdamW timescale set optimally for that TPP, making collapse a marker of compute-efficient training and a tool for tuning, diagnostics, and early stopping.

Abstract: Effective LLM training depends on predictable scaling of key quantities—such as final loss and optimal hyperparameters—with model and dataset size. Qiu et al. (2025) recently showed that this predictability can extend beyond scalars: whole training loss curves can *collapse* onto a universal trajectory after a simple normalization. What remains unclear is whether this phenomenon persists for LLM families trained under *practical scaling recipes*, where width, depth, learning rate, batch size, and weight decay are scaled jointly. We show that it does: loss curves collapse across scales precisely when optimization hyperparameters are set optimally for the given data budget, in accordance with recent empirical scaling laws. Collapse therefore emerges as a signature of compute-efficient training. We demonstrate two applications at scale: (1) deviation-from-collapse provides a sensitive, early diagnostic of training pathologies, and (2) predictability of collapsed curves enables early stopping in large-scale hyperparameter tuning. Finally, we train a competitive LLM family, *Celerity*, using these insights, establishing collapse as an effective tool for developing efficient LLMs.

Published as a conference paper at ICLR 2026
SCALING WITH COLLAPSE: EFFICIENT AND PRE- DICTABLE TRAINING OF LLM FAMILIES
Shane Bergsma Bin Claire Zhang Nolan Dey Shaheer Muhammad Gurpreet Gosal
Joel Hestness
Cerebras Systems
(shane eccane.bergeee bergsma, joel
ceeeeeeerrss...e. joeraaceebbassssss.o.o net
ABSTRACT
Effective LLM training depends predictable scaling of key quantities-sue final loss and optimal hyperparameters- -with model and dataset size. Qiu al. (2025) recently showed that this predictability can extend beyond scalars: whole training loss curves can collapse onto a universal trajectory after a simple nor- malization. What remains unclear is whether this phenomenon persists for LLM families trained under practical scaling recipes, where width, depth, learning rate, batch size, and weight decay are scaled jointly. We show that it does: loss curves collapse across scales precisely when optimization hyperparameters are set op- timally for the given data budget, in accordance with recent empirical scaling laws. Collapse therefore emerges as a signature of compute-efficient training. We demonstrate two applications scale: (1) crriaationfrom-colel
provides a sen- sitive, early diagnostic of training pathologies, and (2) predictability of collapsed curves enables early stopping in large-scale hyperparameter tuning. Finally, we train competitive LLM family, Celerity, using these insights, establishing col- lapse as an effective tool for developing efficient LLMs.
Llama-2 train loss; TPP &
Celerity curves collapse: TPP & fixed
Collapse residuals can detect issues
0.000
1.4
1ama-22 Models
14
Calerity Models
with normalized 500M curve
Calerity Models
M,
286 r=0.07
234 TPP, 0.0.5
0.025
500M
Normalized train loss
13B: 154 TPP, r=0.07
Normalized train lass
500M: 234 r
900M
1.3
34B: TPP, r-0.13
1.3
900M: 234
000.5
0.020
1.8 (initial)
70R: 29 TPP. 1-1.33
1.8P: 234 TPP, 0-0.55
0.015
1,88 (repaired)
3,98: 234 TPP. 0=0.05
1.2
1.1
1.1
0.005
0.000
1.0
1.0
-0.005
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.75
1.00
0.6
0.8
1.0
Fraction
training tokens
Fraction training tokens,
Fraction training tokens,
Figure I: Left: Prior LLM families like Llama-2 train at varying tokens-per-parameter (TPP; D/N) and AdamW timescale т; training loss curves do not collapse. Middle: Fixing TPP and setting T optimally for that TPP, Celerity loss curves do collapse. Right: Deviations from collapse allow precise identification (and earlier repair) of numerics issues in large-scale training runs.
INTRODUCTION
Scaling up pre-training has emerged as the primary route to improving LLM performance (Brown al., 2020; Achiam et al., 2023). Yet once we reach frontier scales, opportunities for direct exper- imentation disappear (Xiao, 2024). How then can we train effectively at those scales-whi size of model should we use, and how should we set hyperparameters? Encouragingly, recent work has revealed that several quantities remarkably predictable we scale deep learning. These include model performance as function of model and dataset size (Hestness et al., 2017; Kaplan et al.. 2020), well as hyperparameters under maximal update parameterization (nP), which en- ables optimal base learning rates and initializations to approximately transfer across widths (Yang al., 2021). In this paper build on this trajectory of predictability: show that, LLM scale,

---

Published as a conference paper at ICLR 2026
training loss curves (TLCs) from different model sizes collapse onto single universal curve after simple normalization provided models are trained with a particular typerparametter-sallinn recipe.
Qiu et al. (2025) only recently demonstrated this striking regularity in TLCs, showing collapse when training with μ&P on small-scale autoregressive tasks. As their validation was limited to small models trained with vanilla Adam (Kingma & Ba, 2014), without weight decay, they explicitly call for tests at larger scales with practical scaling ladders that co-scale width, depth, batch size, and weight decay. Our work addresses this gap, showing that collapse persists in full-scale LLM families.
While modeling LLM loss is an active research topic (Sec. 6), the ability to predict TLCs has great practical value. For example, human judgment is now required to decide whether training has recovered from a loss spike-or whether rewinding/restariing is needed (Chowdhery et al., 2022; Zhang al., 2022a). Other subjective signals, such as a gradual upward trend (Zhang et al., 2022b), can also trigger interventions. Yet criteria remain vague: Touvron et al. (2023b) report Llama- TLCs "did not show any sign of saturation," " but how to recognize saturation unclear. If TLCs collapse across sizes, practitioners can compare in-progress training universal reference, monitor residuals, and extrapolate final loss from partial trajectories. Teams already rely on TLCs in this way, often without principled account of what governs TLC shape; for example, Falcon's 's final LR was chosen by simply continuing the rum performing best after warmup (Almazrouei et al., 2023).
In this paper show that the essential condition for collapse under μP is that the LR schedule, tokens-per tokens-per-paannen per-parameter ratio (TPP), and Adam W timescale T (Wang & Aitchison, 2024) are held fixed across model sizes. This reflects deeper regularity: prior work showed that optimal depends only on TPP (Bergsma et al., 2025a). Thus, scaling across fixed TPP with 7 chosen optimally guarantees collapse, and collapse emerges robust marker of compute-efficiet and stable pre-training. When T is mis-scaled-an as in the Llama-2 family (Fig. 1, normalzzzd curves fail to align.
We introduce Celerity as the first LLM family trained with both optimal T scaling and demonstrable TLC collapse (Fig.
70.0
better
Gemmaz-28
Llama2
Llama-78 78
middle). Effective parameterization, in- cluding tuning and transferring т, helped Celerity land on the compute-efficieney frontier for open models of its scale (Fig. 2). Meanwhile, deviations from col- lapse provided a sensitive diagnostic of training issues: in our 1.8B run, merical instability became evident from collapse residuals (Fig. 1, right) well be- fore the raw TLC showed an upward trend (Fig. 6, right). Celerity exemplifies scal- ing with collapse: efficient, predictable training. across scales and throughout the run.
Average Accuracy 1%
67.5
Celenty-3.9e
Smo1LM2-1.78
OLMa 7B
65.0
Zarbba2 1.28
SmolLM 1.78
62.5
Celerity-1.aa 1.88
OLMo2
BTLM-38
60.0
Gemma3 1B
57.5
Mm 360M
OLMo-18
CerebrasGPT-6.77 CerebrasGPT-138
55.0
SmojLM-360M
Celerity-90MM
Celerity Fit: 100
(resew 0-0057
52.6oa
10
10
Training Compute FLOPs (C)
Figure 2: Celerity is at the compute-efficiency fron- tier. (Average accuracy tasks arc-c, arc-e, boolg, hellaswag, piqa, siqa, winogrande; see Table 10.)
In summary, our main contributions are:
Identifying the key factors influencing loss curve shape under jPP: the LR schedule, the TPP ratio, and the AdamW timescale т, and explaining shape dependence on these quantities (Sec. 3).
Demonstrating that when 5 is set optimally for given TPP, TLCs collapse across model scales, providing a signature of compute-efficient training (Sec. 3).
Introducing the Celerity family, the first large-scale LLMs trained in a collapse regime (Sec. 4). Proposing simple functional form for normalized TLCs, and showing that fitting this form on small-scale training runs enables early stopping in large-scale hyperparameter tuning (Sec. 5).
BACKGROUND
Training loss curves (TLCs) for different model sizes typically differ in scale, duration and final loss. Yet after a simple normalization, they can align closely, or collapse, onto a common trajectory. Collapse emerges only when three controls are matched across scales. The tokens-per-parameter
2

---

Published as a conference paper at ICLR 2026
ratio (TPP) determines how much data each parameter sees and thus affects the normalized pace of improvement; the AdamW timescale r governs how long the optimizer "remembers" past gradients, shaping the bias-variance trade-off over training; and the learning-rate schedule phases early bias reduction against late variance suppression. When TPP and r are chosen consistently across model sizes and the LR schedule is fixed, the resulting normalized TL.Cs become approximately scale- invariant. We now formalize these quantities; Table summarizes key symbols.
TPP. The TPP ratio is equal to number of training tokens D divided by the model size N. This sim- ple quantity plays surprisingly profound role in compute-efficient LLM training and TLC shape. Hoffmann et al. (2022) investigated, for a given compute budget C. how to allocate D and N in order to minimize loss. They found optimal D and N scale roughly equally as C' increases, with the optimal D/N ratio relatively constant at around 20 TPP (Appendix C. Replication studies have found similar results (Besiroglu et al., 2024; Porian et al., 2024), and 20 TPP has emerged as a fule-of-thumb for compute-optimal training (Dey et al., 2023a; Zhang et al., 2024b).
uPP nP (Yang& & Hu, 2020) and related parameterizations for depth (Bordelon et al., 2023; Yang et al., 2023; Dey et al., 2025) seek to achieve consistent, stable training dynamics as networks scale up. Moreover, with μP, base hyperparameters can be tuned transferred to larger scales. Given the width of the proxy model, dp. and target, d, uP prescribes scaling factors to apply to the base LR, initial weight variance, and other base HPs.
small prexy model and then
uPP is increasingly used in LLM training (Dey et al., 2023a;b; Sengupta et al., 2023; Shen et al., 2024: Hu et al., 2024). Moreover, recent work has shown that, when using μP, other important aspects of training may decouple from model size, including optimal batch size (scaling primarily in the total number of tokens (Zhang et al., 2024b; Bergsma et al., 2025a)), and optimal AdamW timescale/weight decay (scaling primarily in TPP (Bergsma et al., 2025a)).
Supercollapse. Using pP. Qiu et al. (2025) observed that TLCs for different model sizes, despite varying widely over compute and absolute loss, appear to follow a consistent shape. This motivated them to affinely rescale the curves to the normalized loss l given by:
(1) where w is the random sced, i is the fraction of training completed (what Qiu et al. (2025) refer to as normalized compute), N is the number of model parameters, and T* (N) is the corresponding compute-optimal number of training steps, estimated from a power law fit. L is an offset, which they subsequently set to the estimated irreducible loss of their power law.
(i) = (L(E (NT) - ())(L((""()
Training compute-optimally under μP (on small-scale autoregressive tasks, e.g., predicting chess moves), Qiu et al. (2025) showed TLCs collapse under this normalization-indedd, they su- percollapse, meaning they differ by less than the noise from inter-run variation. They further show that collapse arises naturally in leanstant--ernnn--raa rate models where loss obeys typical neural power laws, while extending the theory to arbitrary LR schedules via theoretical model of quadratic loss.
The Adam EMA and its timescale. AdamW updates at step can be expressed in terms of learnnng rate 7) and weight decay as: A: (1 7 - 7 Vette where my and êt are bias- corrected EMAs of gradients and squared gradients, respectively (Kingma & Ba, 2014). Wang & Aitchison (2024) observed that Adam parameters 0g also be viewed an EMA-of weight updates. That is, the standard EMA form ÿt = a a) + O.2;, matches Adam when yr- = 6,, Aie = =,,,, and x; = - Vo The timescale Titer = 1/a = 1/nA represents the approximate number of iterations over which updates are averaged. When expressed in epochs as Tepoch /iter where M is the number of iterations per epoch, Wang & Aitchison (2024) found the optimal Topoch (swept by varying A) remains stable under model and dataset scaling on image tasks.
Since LLM pre-training typically uses a single epoch, we follow Bergsma et al. (2025a) in defining a normalized timescale Titer /T, where T is the total number of optimization steps. As T D/B (total tokens/batch size):
(T) B/ (/())))
(2)
In contrast with the results in Wang & Aitchison (2024), Bergsma et al. (2025a) did not find optimal to remain stable in L.M training, but instead to decrease a (scale-invariant) power law in TPP.
3

---

Published as a conference paper at ICLR 2026
Table 1: Core quantities used throughout the paper.
Symbol
Meaning
N
Number of model parameters
D
Total number of training tokens
Batch size (tokens per optimization step)
T D/B
Total number of optimization steps
1=
Fraction of training completed Tokens-per perapaameteee ratio (TPP)
L(()
Training loss at step (
((1
Normalized training loss curve (TLC)
3
Learning rate
Weight decay coefficient
Titer
1/(TA)) Tiuer T 1/(1AT)= B/(A)
Adam W timescale (in steps)
Normalized Adam timescale
3 WHAT FACTORS MODULATE TRAINING CURVE SHAPE?
Experimental setup. We use GPT2-like LLM (Radford et al., 2019), with ALiBi embed- dings (Press et al., 2022) and SwiGLU (Shazeer, 2020). We train on SlimPajama (Soboleva et al., 2023). Models are trained with Adam and μP. We use linear decay-to-zero LR schedule, context length of 2048, and the GPT2 vocabulary. Full architecture and other details are in Appendix B.
We plot vs. training fraction f t/T tB/D, with step count f, total steps T, batch size B, and dataset size D. To reduce noise in small-B settings, we post hoc aggregate losses {(f) using a moving-average filter over a window of 100 steps, smoothing curves without altering the underlying trajectory. We also consistently found simply dividing by the final training loss (i.e., L 0 in Eq. (1)) resulted in optimal alignment across scales, so use this for all curves.
varies y= 1-622-22222 02 A=0.1
B=504
A=0.1
B=504
p=1.62e-02 varies
B (r)
1) (+)
1.4
A
126 (0.025)
4.05e 03 (0.421)
0.025 (0.421)
Normalized train lazs
Normalized train loss
252 (0.053)
8.09e-03 03 (0.210)
0.05 (0:210)
1.3
504 (0.100 105)
Normalized train
1.5
1.82e-02 (0.105)
1.3
0.1 (0. 106)
(0.210)
(0.053)
(0.053)
1.2
2016 (0.421)
5.48e-02 (0.026)
1.2
0.4 (0.026)
1.1
1.1
1.0
1.0
1.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 tokenns,
Fraction of training tokens,
Figure 3: AdamW timescale т modulates TLC shape (610M, S0TPP): Sweeping n (left), (middle), B (right) produces matching variations in normalized TLCs when varies identically.
modulates TLC shape. Fig. 3 shows normalized TLCs for 610M models trained to 80 TPP, sweeping either learning rate weight decay X, or batch size B in each subplot. Across hyperparameters, TL.Cs with matching r exhibit very similar shapes, reflecting consistent timescale control. Similar patterns hold across other scales and dataset sizes. Generally, as 7 increases, TLCs drop more early and less later. This is also a function of the LR schedule: when we switch to using Constant LR, there is no final drop, lower-7 TLCs are lower throughout (appendix Fig. 10).
Finding:
Finding: TPP modulates TLC shape. We now fix 7 and test increasing TPP, finding TLCs drop earlier and flatten for longer (Fig. 4, left, middle; see also Llama-2 for 0.07 and - 0.13 in Fig. 1, left). Intuitively, relative to length of training. higher TPP drops loss more at the beginning and then obtains diminishing returns later model scales at the same TPP (when 7 is roughly equal), showing TPP's shaping effect is scale- invariant (scaling from 11 1M to 3.3B at fixed TPP represents x000 increase in training FL.PPs.
In Fig. 4, right, TLC shape is quite similar across
4

---

Published as a conference paper at ICLR 2026
7-0.021
7-0.105
Curves align when {TPP. r} const.
TPP
TPP
- 20
20
111M, 7-0.33
Normalized train lazs
Normalzzed train loss
266M, 0.2.27
1.3
320
1.3
520
1.3
610M, 7-0.21
1.78. r=0.30
1.2
1.2
3.38. r=0.31
1.1
1.1
1.0
1.0
1.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 takens,
Fraction of training tokens,
Figure 4: TPP modulates TLC shape. Fixing T for 111M (left) & 610M (middle) while increasing TPP, curves shift down. When const. and TPP also fixed (at 20), curves roughly collapse (right).
Using bias and variance to explain the effect of Classical analyses of stochastic optimization (e.g., Bottou et al. (2018); D' Angelo et al. (2024)) decompose loss into two components: bias term, reflecting dependence on initial weights, and a variance term, reflecting noise from stochastic gradients. Early in training bias dominates; later, variance determines the attainable loss floor.
In our setting, the AdamW timescale controls this trade-off through the EMA over weight updates (Sec. 2). A smaller corresponds to shorter memory: updates depend mainly on recent gradients, yielding rapid bias reduction but a higher variance floor. larger T averages over more past gradi- ents, reducing variance more effectively but slowing early progress. Empirically, TLCs reflect this pacing, with smaller 7 producing faster early descent (e.g., left panel of appendix Fig. 10).
This intuition is formalized in Appendix B.3 via a noisy quadratic model: h
h ee-2/// E(((0)*)
B[
e1-e -2///
21/7
(3)
4
where the first term approaches variance floor x /T and the second is an exponentially decaying bias term. Smaller thus yields faster initial decay but a higher floor, while larger yields slower initial decay but lower floor, matching the observed fast-ther-flatten behavior under constant LR. With LR decay, } decreases and the instantaneous timescale т₃ = AT) increases, enhancing late-stage variance suppression and steepening final drop (e.g.. Fig. 10: more LR decay, more drop).
Scale invariance. After normalizing by the final loss, the curvature factor cancels (Appendix B.3). Provided residual bias end-of- trainirriiing is negligible relative the variance floor, the normalized TLC depends only on 7 and f. Thus, at matched т, normalized TLCs collapse across scales.
Explaining effect of TPP. TPP affects TLCs via power laws. With constant LR, every step is the endpoint of a shorter run, trained to i. TPP. Qiu et al. (2025) note L(i) therefore follows the same power law as final loss of a run fully trained to that effective budget. Normalizing by the projected loss at total TPP removes dependence on model and dataset size, so f(t) depends only on f and total TPP (Appendix B.2). Higher-TPP curves analytically decay faster and level off sooner. LR schedules deform the curves, but deformation is also scale invariant given consistent curvature of the loss landscape across model sizes under pP (Noci et al., 2024).
Key takeaway 1: TLC shape is governed by three scale-invariant controls: (I) AdamW timescale (bias-variance trade-off), (2) tokens-per-Pe per-paamtrrrr ratio (TPP), which sets the relative pace of improvement, and (3) learning- rattengg-reeeeeee schedule, which phases early and late loss reduction. When matched across model sizes, normalized TLCs follow the same trajectory-that that is, they collapse.
4 CELERITY: A COMPUTE- MODEL FAMILY WITH COLLAPSE
We have established that collapse arises when 7 and TPP are held fixed across model sizes. Mean- while, prior work has shown optimal to depend only TPP (Bergsma et al., 2025a). Here, introduce a model family, Celerity trained at fixed TPP and with т chosen optimally for that TPP, i.e. regime where collapse emerges naturally as consequence of good training.
5

---

Published as a conference paper at ICLR 2026
Compute vs. parameter efficiency. A key question for Celerity is which TPP to use: ~ 20 is compute-optimal (Sec. 2). while higher TPP means greater parameter efficiency (fewer parameters to obtain same loss). For small, inference-ready models, parameter efficiency is paramount, but such models are usually distilled (Tunstall et al., 2023; Wang et al., 2025) rather than pre-trained at high-TPP. Our main interest is developing pre-training strategies for very large models. As models scale, the relative importance of compute-efficiency increases-indedd
public families often have declining TPP as size increases (Touvron al., 2023a; Biderman al., 2023).
Yet even for the largest models, parameter efficiency re- mains valuable, e.g., when generating distillation logits or synthetic data. To choose Celerity's 's TPP, we analyze this trade-off: Appendix C. derives an expression for the extra compute required to compress a model to fraction of the size when training compute optimally- while maintaining the same loss. This expression leverages power law fits
234 TPP: high compression /low cost
TPP: D/N
Cost: Extra compute (C:Can)
6 9k
better
10
TFP-20 20
(ky=10 C/Cgm110
10a
from prior work. Fig. 5 plots the trade-off, where a TPP ratio of 234 is estimated achieve 62% reduction in pa- rameters with only a 67% increase in total FLOPs (relative to 20 TPP). This is responsible balance point, near what has been called the critical model size-the point where further increasing compute obtains mussively-diminihhing returns in parameter efficiency (De Vries, 2023) (c.g., dou- bling our FLOPs, to 3.34 compute-optimal, reduces N by only further = 11%).
10
Compression: Fraction Nog (kw)
Figure Expected iso-loss compute vs. compress trade-off as TPP varies.
Even if the ultimate goal is a "herd" of models varying TPP, such as Llama-2 2 in Fig. 1, left, there are advantages to training different "bands" within the herd, e.g., 7B, 13B, 34B, 70B all at 29 TPP:
Tuning: you can fine-tune T at a smaller scale and zero-shot transfer to larger models.
Diagnostics: Because TLCs collapse. deviations provide an early warning of training issues.
Cost: Fixed-TPe TPP bands are cheap (e.g., 10x lower N + 10x smaller D + 100 less compute).
Philosophy. Celerity aims to advance general LLM capabilities using public pre-training corpora and fully-open, consistent methods-rathe than targeting specifie benchmarks. In contrast, the majority of LLMs now anneal on training subsets of downstream benchmarks (Dubey et al., 2024; Achiam et al., 2023), or inject special high-quality math (OLMo al., 2024), code (Zhang et al., 2024a), or instruction (Hu et al., 2024) data during a late-stage mid-training process. Since these practices make evaluation problematic (Dominguez-OIil Olmeedo al., 2024), Celerity can serve as a comparison for models trained without (or prior to) applying such techniques.
Experimental details. Celerity pre- trains in bands of 20, 80, and 234 TPP, each spanning 300M-3.9B models (Ta- ble 2); see Appendix C.2 for further de- tails. Key enablers of Celerity's reliable, efficient training include:
Table 2: Architecture of Celerity
Celerity:
300M
500M
900M
1.8B
3.9B
Hidden Dim
640
896
1152
1536
2048
Num Heads
10
14
9.
12
16
Head Size
64
64
128
128
128
Data: emphasizing (open) educa- tional, math, and coding data through- out training (appendix Table 6); this outperformed training on the general SlimPajama dataset (Table 7).
Layers
13
17
23
30
40
Batch Size
176
240
336
464
672
Vocabulary Embeddings Seq Length
Llama-3 (size 128256)
ALiBi, Untied
8192
Using CompleteP, which enables hyperparameter trans- fer over width and depth, was more efficienn//eliable than μP (Fig. 15).
Parameterization:
Non-linearity FFN Mult
Squared ReLU 8x Pre- Normalization, = 10
Norm Type
LR Schedule
Peak: 0.15, linear decay-to-zero
Optimization: LR, T, batch size tuned small, transferred via scaling rules.
LR warmup
min( 10% of total tokens, 375M tokens)

---

Published as a conference paper at ICLR 2026
Evaluation results. Appendix Table 10 provides full downstream evaluation results for Celerity and other public models tested on seven common downstream tasks. Fig. 2 shows that Celerity mod- els form the accuracy/compute Pareto frontier up our largest training budget. Against BTLM (Dey et al., 20233) before task-specific data annealing became standard- standard-Celerity achieves com- parable accuracy with 75% fewer training FLOPs. Extrapolation via fitted power law in compute (dashed line in plot) suggests smooth scaling and continued competitiveness. For comparison with distilled models, we count only student FLOPs in Fig. 2. Including teacher FLOPs (forward passes), or the cost of teacher training, strengthens Celerity further (appendix Fig. 16).
In terms of parameter efficiency, Celerity is weaker than high- TPP families (Figs. 19 and 20), mean- ing such models save FLOPs at inference. However, beyond the importance of studying compute efficiency for hyper-scale training, there is strong motivation to train and study compute-efficien smaller models: growing evidence suggests some models may be counter-productively (even catas- trophically) overtrained, making them harder to fine-tune (Springer et al., 2025) and quantize (Ku- mar et al., 2024). Compute-efficient alternatives therefore serve both as principled baseline for understanding scaling and as a practical fallback when high-TPP models prove brittle.
Collapse results. In Sec. 3, we normalized training loss curves by dividing by the final loss value, L(T) (Eq. (I)). To use collapse as a diagnostic during training, we need a way to normalize when L(T) is still unknown. We explored two strategies and use early-align in our experiments:
Estimate: extrapolate L(T) from power law fit at lower scales.
Early-align: choose L(T) so E(f) best aligns with the smallest-scale curve over 25-50% portion.
Celerity collapse: TPP
Colerity train collapse: TPP
1.88 (initial) curce with apparent blip
3.0
N (r=0.175)
(r=07))7)
Training loss (unrormaized)
(234 TPP, r=0.061)
Normalized train loss
300M
1.88
Noralize train loss
300M
900M
2.8
(initial)
500M
3.98
500M
-.88
1.8B (initial. unsmoothed)
1.3
900M
1.3
2.6
1.38 (repaireed)
1.2
1.2
2.4
1.1
1.1
2.2
2.0
1.0
1.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 training tokens,
Figure 6: Collapse in Celerity models. Celerity 20 TPP (left) and 80 TPP (middle) models exhibit collapse. Right: smoothing helps detect blip in loss near the end of initial 1.8B run (red curve), but divergence can be detected much earlier using collapse residuals (Fig. 1, right).
Fig. 6 shows normalized curves. Collapse is tight at 80 TPP (middle). At 20 TPP (left), we see small early deviations, which we attribute to differing LR warmup proportions (Table 2). At 234 TPP, divergences appear late in training for larger models (Fig. 1, middle). Investigating. we find loss improves disproportionately on training data, while held-out data remains aligned with projections.
Collapse for monitoring. Fig. (right) shows the unnormalized TLC for our original 1.8B, 234 TPP run. Smoothing helps reveal sudden rise in training loss, but only after 90% of training. Without collapse reference, would be impossible to see that problems began much earlier. By comparing against the 500M TLC reference (Fig. right), we pinpoint divergence starting near 60%. Knowing this timing was crucial: we did not waste effort investigating late-stage data redun- dancy, and instead realized the problem coincided with a job restart under new compute allocation.
The collapse reference was also essential for debugging: by running ablations with different batch sizes and measuring divergence from the reference, we confirmed the anomaly arose from numer- ical issue in a loss kernel triggered only at specific microbatch sizes. After fixing the kernel and restarting from before the divergence, training tracked the reference TLC closely (Fig. 1).
7

---

Published as a conference paper at ICLR 2026
Key takeaway 2: Celerity trains models in fixed-TPP bands the optimal T for each band. This produces collapse across scales, enabling hyperparameter transfer, early detection of training issues, and reliable cross-scale comparisons. The 234-TPP band lies on the compute- frontier while using =62% fewer parameters than compute-optimal training at equal loss.
COLLAPSE ENABLES EARLY STOPPING IN HYPERPARAMETER TUNING
Training to completion is expensive. If normalized TLCs behave predictably, can we stop earlier and still recover the final loss? We show collapse enables principled early stopping in tuning, and introduce a predictive model-fit at small scales, and re-used to extrapolate large-scale TL.Cs.
1.7B. 20 TPP: sweeping DJ A- 0.1
1.7B.
20 TPP: sweeping 13 T-0.15
3.6
3.6
B
B
Training lass (unnormalize lizee)
3.4
126
1008
Training lass (unnormalizee lizee))
3.4
126
1008
252
2016
252
2015
3-2
504
4032
3.2
504
4032
3.0
3.0
2.8
2.8
2.6
2.6
2.4
2.4
0.00
0.25
0.50
0.75
1.00
0.00
01.25
0.50
0.77
1.00
Fraction of training tokens,
Fraction of training tokens,
Figure 7: Predictability of batch size sweeps. Left: When fixing weight decay in batch size sweeps (standard practice), normalized TLCs cross, making final loss hard predict from partial results. Right: Fixing instead (by adjusting A), TLCs maintain ordering, enabling early stopping.
Role of r in tuning. Recent work tunes learning rate r and batch size B at smaller scales and extrapolates via power laws (Hu al., 2024; Bi al., 2024; Porian et al., 2024). These studies typically fix weight decay X, which unintentionally varies r-and hence TLC shape. As Fig. (left) shows, when varies, mid-training loss is a poor predictor of final outcomes during batch size sweep. In contrast, when r is fixed during tuning (by adjusting A), the ordering of B-specific curves preserved throughout training (right); runs can be stopped early (e.g.. at 25%) while still reliably identifying the best batch size.
There are, however, cases where 7 must vary. For example, Bergsma et al. (2025a) found optimal was no longer constant once B > Berit potentially requiring retuning of X.
Exploiting collapse for early stopping. Suppose we wish to find the optimal setting of hyperpa- rameter (HP) in large-scale training. Naively, we could sweep across settings of the HP and train large-scale model to completion at each setting; the lowest final loss would identify the best choice. Alternatively, rather than training to completion, we propose a procedure to infer the final loss values from partial training runs. We do this by exploiting the collapse phenomenon as follows:
1. For each large-scale setting in the sweep, first identify the corresponding TLC controls, i.e, the LR schedule, т, and TPP.
2. Train smaller (e.g.. (00Mpaaramete) model for each unique combination of TLC controls.
3. Normalize these small-run loss curves to obtain the normalized universal TLCs (E(I) corre-
sponding to each set of controls.
4. Perform partial training runs (e.g., to 30% of tokens) at each HP setting.
5. For each partial large-scale loss curve, use its corresponding universal TLC to predict the final loss value. Do this by determining the divisor L(T') that maximally aligns the partial curve with the same segment of the corresponding universal TLC. These divisors thus act both as normalizing constants and, by construction, calibrated extrapolations of the final loss.
6. Select the optimal hyperparameter setting corresponding to the lowest predicted final loss.
8

---

Published as a conference paper at ICLR 2026
Predicting the normalized universal TLCs. In some cases, may already have many small- scale runs, but not loss curves for each of the specific TLC controls corresponding to our large-scale sweep. In this situation, we hypothesize that a parametric surrogate model can generate high-fideling normalized universal TLCss as a function of and TPP. This allows us to obtain the normalized TL.C shapes required for Step 5 of the above procedure, without training small-models at exactly- matching controls. Fitting such a surrogate lets us leverage our broader TLC dataset, and enables estimation of normalized TLC shapes that go beyond trained regimes (see Appendix D.3 for an example).
For the surrogate ((t) model, we experimented with several functional forms and ablations on our 11 data, focusing on:
(f)) = ((1 + a) + (n(() + c2)
(4)
The first term captures power-law improvement in training fraction (Appendix B.2) while the second term modulates this by the LR schedule rf(f), reflecting how variance suppression phased over training (Appendix B.3). m, b, q. 61 and 62 are fit parameters. We divide ε(t) by its final value so that {11) = 1.0. Fixing 61 = 0.001 and 62 0.1 avoids large swings at {(O) and ((1).
In practice, we find m can be fixed (we use 0.05). Parameters and q then vary systematically with T and TPP, respectively, which we capture with power laws:
(5) Because b and q interact, jointly fitting their parameters would require a Ogg²) grid search (with g the grid resolution). Instead, alternate: fit (beonst: desp) with fixed q, then fit (goosst; Tesp) with fixed b, iterating to convergence. This reduces cost to O(g²) while yielding stable fits.
b
boonst (r) q = Iconst (TPP) Sop
Results: prediction. We fit the b and g power laws on 11 11 data and evaluate using mean absolute eπror (MAE) between é and true f, computed over & E [0.2, 1] (ignoring error around LR warmup, when initial curves are noisy). We report unweighted mean MAE across all curves.
3,3R, 30 TPP: swneping X,
2018
1.78, 20 TPP: sweeping A, 8081
3.3B, 30 TPP: sweeping A, 2016
1.4
Loss dift predicted us. trie best (%)
Choose randomly
Training loss (unnormalized)
0.05
0.1 0.2 0.4
1.2
Loss dift predicted vs. troe best
1.2
Choose current best
Choose predicted
1.3
1.0
1.0
+ b (())"
pm
randomly Choose current best
1.2
0.6
predicted
0.6
1.1
0.4
0.4
0.0
0.0
0.00
0.25
0.50
0.75
1.00
0.2
0.4
0.6
0.8
1.0
0.2
0.4
0.6
0.8
1.0
Fraction of training tokens, i
Stop point training
Stop point training
Figure 8: 3.3B-scale predictions and true normalized TLCs.
Figure 9: Early stopping works best with predicted loss:
Tuning À in 1.7B (left) and 3.3B (right) models.
Results show that predictions are good: MAE is low and actually improves with scale (appendix Table 11), likely because (1) larger datasets yield smoother TLCs, and (2) fewer extreme hyperpa- rameter tests at larger scales. Fig. 8 (left) shows example: predictions trained on 11 M--scale TLCs (1000x fewer FLOPs) closely match observed curves for a 3.3B model.
Estimating b and y as power laws reduces MAE by two-thirds compared to using fixed values (Ta- ble 12), though error remains =2x higher than an oracle fit of and g per curve. Adjusting for both and TPP is vital; however, fitting 6 and q jointly on both did not improve further.
Results: tuning. We now test whether optimal LLM settings can be predicted from partial training runs. At different stopping points in training, we choose setting as the best, and evaluate the gap between the chosen setting's final loss and the true best setting. We compare the following choices:
1. Random baseline: randomly choose one setting the best.

---

Published as a conference paper at ICLR 2026
2. Current best: choose the setting giving the best result at the stop point.
3. Predicted best: Align partial TLCs with predicted é, choose lowest fitted normalizer L(T).
Fig. shows results for sweeps at 1.7B/20TPP (left) and 3.3B/30TPP (right). Predicted best achieves negligible loss gaps when stopping after just 30% and 10% of training, respectively. In contrast, current hest-used in Almazrouei et al. (2023) for LR tuning-Succeeds succeeds initially at 3.3B but fails at 1.7B, showing it is not general solution. Further experiments are in Appendix D.2.
Key takeaway 3: Collapse enables reliable early stopping. By aligning partial training curves to a small-scale reference, we can predict final loss L(T) and select hyperparumeters after only 10-30% of training, substantially reducing tuning compute.
6 RELATED WORK
Scaling laws and scale-stable dynamics. Neural scaling laws relate loss (generally obtained from separate training runs) growth in model, data, and compute sizes, via power laws (Hestness al., 2017: Kaplan et al., 2020: Henighan et al., 2020; Hoffmann et al., 2022; Caballero et al., 2022; Alabdulmohsin al., 2022). To ensure stable training as models scale, parameterizations such as IP transfer base hyperparameters across sizes, and yield early dynamics that are scale-stable (Yang et al., 2021; Vyas al., 2023; Kalra et al., 2023), even super-consistent (in curvature) (Noci al., 2024). Observing suboptimal LRs under μP as data scales, recent work has proposed decreasing the LR as a function of D (Shen et al., 2024; Bjorck et al., 2024); Bergsma et al. (2025a) unify these techniques forms of 7 adjustment. Qiu et al. (2025) show that, for compute-optimal ladders, TLCs collapse after normalization. We build on these threads at LLM scale while co-scaling width, depth, batch size, and weight decay, identifying new controls that govern TLC collapse.
LLM loss-curve prediction. While Kaplan et al. (2020) fit a simple power law to TLCs, recent papers make loss prediction explicitly LR- Leddependent (Tissue et al., 2024; Luo et al., 2025; Schaipp et al., 2025; Qiu al., 2025; Hong & Wang. 2025). Complementary to these, we take timescale- centric view: Adam W implements an EMA over updates, and the normalized timescale 7 (jointly set by LR, weight decay, and batch size) acts to control an implicit batch size, one that trades bias reduc- tion vs. variance suppression and thereby shapes TLCs. In noisy quadratic model ( (Appendix B.3), derive expression for training loss under constant LR, and explain why decaying schedules invert the ordering of TLCSs across т, with deformations remaining scale-invariant once normalized.
Early stopping, HPO, and monitoring. Early-termination and HPO methods extrapolate TLCs or prune trials (Swersky et al., 2014: Domhan et al., 2015; Jaderberg et al., 2017: Zela et al., 2018; Li et al., 2018; Choi et al., 2018; Akiba et al., 2019), but typically require many short runs and are not tailored LLM pre-training regimes. Our approach leverages collapse itself: fit small-scale predictor of normalized TLCs, align in-progress curves to infer L(T), and select winners by 10-30% of training. Operationally, large-scale reports document spikes and divergences (Chowdhery al., 2022; Zhang et al., 2022a; Wortsman et al., 2023; Molybog et al., 2023); we show collapse residuals provide quantitative, scale-normalized early-warning signal and practical aid for debugging.
CONCLUSION
At LLM scale, normalized training loss curves collapse across model sizes when three controls align: the AdamW timescale r, the tokens--er-paramette ratio (TPP), and the learning-rate schedule. Empirically, T (bias-variance smoothing) and TPP (power-law improvement rate) set TL.C shape, while the schedule phases these effects. Fixing TPP and setting optimally for that TPP yiclds alignment across 100M-3.9B parameters in our experiments.
We instantiate this in Celerity: fixed TPP with optimal produces tight collapse and competitive accuracy. Collapse residuals surface issues early, localize their onset, and enable safer restarts. A simple predictor for normalized TLCs (fit at small scale) supports early stopping in HPO: by 10- 30% of training we can select winners and estimate L(T), saving tuning compute. For SIB runs, collapse provides valuable reference trajectory: keeping training on track, every step of the way.
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Martín Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Gre- gory S. Corrado, Andy Davis, Jeffrey Dean, Matthieu Devin, et TensorFlow: Large-scale machine learning heterogeneous distributed systems, 2016.
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, 1ige Akkaya, Florencia Leoni Ale- man, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. GPT-4 technical report. arXiv preprint arXiv:2303.08772 2023.
Takuya Akiba, Shotaro Sano, Toshihiko Yanase, Takeru Ohta, and Masanori Koyama. Optuna: A next-generation hyperparameter optimization framework. In Proceedings of the 25th ACM SIGKDD international conference on knowledge discovery & data mining, pp. 2623-222222222 2631, 2019.
Ibrahim M Alabdulmohsin, Behnam Neyshabur, and Xiaohua Zhai. Revisiting neural scaling laws in language and vision. Advances in Neural Information Processing Systems, 35:22300-2223 2022.
Loubna Ben Allal, Anton Lozhkov, Elie Bakouch, Leandro von Werra, and Thomas Wolf. SmolLM- blazingly fast and remarkably powerful. Hugging Face Blog, 16, 2024.
Loubna Ben Allal, Anton Lozhkov, Elie Bakouch, Gabriel Martín Blázquez, Guilherme Penedo, Lewis Tunstall, Andrés Marafioti, Hynek Kydlíček, Agustín Piqueres Lajarín, Vaibhav Srivastav, al. SmoILM2: When Smol goes big-data-centric training of small language model. arXiv preprint arXiv:2502.02737. 2025.
Ebtesam Almazrouei, Hamza Alobeidli, Abdulaziz Alshamsi, Alessandro Cappelli, Ruxandra Co- jocaru, Mérouane Debbah, étienne Goffinet, Daniel Hesslow, Julien Launay, Quentin Malartic, et al. The Falcon series of open language models. arXiv preprint arXiv:2311.1680 2023.
Loubna Ben Allal, Anton Lozhkov, Guilherme Penedo, Thomas Wolf, and Leandro von Werra. Cosmopedia. Hugging Face, 2024.
Shane Bergsma, Nolan Dey, Gurpreet 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 zero works best for LLMs. arXiv preprint arXiv:2502. 1533::2500..5938, 2025b.
Tamay Besiroglu, Ege Erdil, Matthew Barnett, and Josh You. Chinchilla scaling: A replication attempt. arXiv preprint 24Xv:2444.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 urXiv:2401.02954 2024.
Stella Biderman, Hailey Schoelkopf, Quentin Anthony, Herbie Bradley, Kyle O'Brien, Eric Hal- lahan, 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.
Johan Bjorck, Alon Benhaim, Vishrav Chaudhary, Furu Wei, and Xia Song. Scaling optimal LR across token horizons. arXiv preprint arXiv:2409. 2024.
Blake Bordelon, Lorenzo Noci, Mufan Bill Li, Boris Hanin, and Cengiz Pehlevan. Depthwise hyperparameter transfer in residual networks: Dynamics and scaling limit. arXiv preprint arXiv:2309. 16620, 2023.
Léon Bottou, Frank Curtis, and Jorge Nocedal. Optimization methods for large-scale machine learning. SIAM review, 60(2):223-311, 2018

---

Published as a conference paper at ICLR 2026
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 leamners Advances in Neural Information Processing Systems, 33:1877-1901, 2020.
Dan Busbridge, Amitis Shidani, Floris Weers, Jason Ramapuram, Etai Littwin, and Russ Webb. Distillation scaling laws, arXiv preprint arXiv:2502.08606 2025.
Ethan Caballero, Kshitij Gupta, Irina Rish, and David Krueger. Broken neural scaling laws. arXiv preprint arXiv:2210. 148.999, 2022.
Daeyoung Choi, Hyunghun Cho, and Wonjong Rhee. On the difficulty of DNN hyperparameter op- timization using learning curve prediction. In TENCON 2018-2018 IEEE Region 10 Conference, pp. 0651-0656. IEEE, 2018
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, al. PaLM: Scaling language modeling with pathways, 2022.
Francesco D' Angelo, Maksym Andriushchenkoo, Aditya Vardhan Varre, and Nicolas Flammarion. Why do we need weight decay in modern deep learning? Advances in Neural Information Pro- cessing Systems, 37:23191 11-2223, 2024.
Harm De Vries. Go smol go home. Blog post, 2023.
Aaron Defazio, Ashok Cutkosky, Harsh Mehta, and Konstantin Mishchenko. Optimal linear decay learning rate schedules and further refinements. arXiv preprint arXiv:2310.07831, 2023.
Aaron Defazio, Xingyu (Alice) Yang. Harsh Mehta, Konstantin Mishchenko, Ahmed Khaled, and Ashok Cutkosky. The road less scheduled. arXiv preprint arXiv:2405. 2024.
Nolan Dey, Gurpreet Gosal, Hemant Khachane, William Marshall, Ribhu Pathria, Marvin Tom, and Joel Hestness. Strebras-GTT. Open compute-opümal language models trained on the Cerebras wafer-scale cluster. arXiv preprint arXiv:2304.0322 2023a.
Nolan Dey, Daria Soboleva, Faisal Al-Khateeb, Bowen Yang, Ribhu Pathria, Hemant Khachane, Shaheer Muhammad, Zhiming, Chen, Robert Myers, Jacob Robert Steeves, Natalia Vassilieva, Marvin Tom, and Joel Hestness. BTLM-3B-8K. 7B parameter performance in a 3B parameter model, 2023b.
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.0161. 2025.
Tobias Domhan, Jost Tobias Springenberg, Frank Hutter, et al. Speeding up automatic hyperparame- ter optimization of deep neural networks by extrapolation of learning curves. In IJCAI, volume 15, pp. 3460-8, 2015.
Ricardo Dominguez-OIinn Olmeeo, Florian Dorner, and Moritz Hardt. Training on the test task con- founds evaluation and emergence. arXiv preprint arXiv:2407.07890, 2024.
Abhimanyu Dubey, Abhinav Jauhri, Abhinay Pandey, Abhishek Kadian, Ahmad AI-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et al. The Llama 3 herd of models. arXiv preprint arXiv:2407.21788 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(120):1-33) 2022.
Steven Feng, Shrimai Prabhumoye, Kezhi Kong, Dan Su, Mostofa Patwary, Mohammad Shoeybi, and Bryan Catanzaro. Maximize your data's 's potential: Enhancing LLM accuracy with two-phase pretraining. arXiv preprint arXiv:2412.15285 2024.
12

---

Published as a conference paper at ICLR 2026
Sebastian Gabarain. Ultratextbooks-.00 Hugging Face, 2024.
Leo Gao, Jonathan Tow, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Kyle McDonell, Niklas Muennighoff, Jason Phang, Laria Reynolds, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. A framework for few-shot language model evaluation. Zenodo, 2021.
Paolo Glorioso, Quentin Anthony, Yury Tokpanov, Anna Golubeva, Vasudev Shyam, James Whit- tington, Jonathan Pilault, and Beren Millidge. The Zamba2 suite: Technical report. arXiv preprint arXiv:2411.15242, 2024.
Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor opti- mization. In International Conference on Machine Learning, pp. 1842-1850. PMLR, 2018.
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.
Tom Henighan, Jared Kaplan, Mor Katz, Mark Chen, Christopher Hesse, Jacob Jackson, Heewoo Jun, Tom B Brown, Prafulla Dhariwal, Scott Gray, et al. Scaling laws for autoregressive generative modeling. arXiv preprint arXiv:2010. 1001::20100..701,, 2020.
Joel Hestmess, Sharan Narang, Newsha Ardalani, Gregory Diamos, Heewoo Jun, Hassan Kianinejad, Md. Mostofa Ali Patwary, Yang Yang, and Yanqi Zhou. Deep learning scaling is predictable, empirically, 2017.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, al. An empirical analysis of compute-optimal large language model training. Advances in Neural Infor- mation Processing Systems, 35, 2022.
Letong Hong and Zhangyang Wang. On the provable separation scales in maximal update param- eterization. In Forry-second International Conference on Machine Learning, 2025.
Shengding Hu, Yuge Tu, Xu Han, Chaoqun He, Ganqu Cui, Xiang Long, Zhi Zheng, Yewei Fang, Yuxiang Huang, Weilin Zhao, et MiniCPM: Unveiling the potential of small language models with scalable training strategies. arXiv preprint arXiv:2404.06395. 2024.
Max Jaderberg, Valentin Dalibard, Simon Osindero, Wojeiech M Czarnecki, Jeff Donahue, Ali Razavi, Oriol Vinyals, Tim Green, lain Dunning. Karen Simonyan, et al. Population based train- ing of neural networks. arXiv preprint arXiv:1711.09844 2017.
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.02070 2023.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.08360 2020.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.69880 2014.
Atli Kosson, Bettina Messmer, and Martin Jaggi. Analyzing & reducing the need for learning rate warmup in GPT training. arXiv preprint arXiv:2410.23992. 2024.
Jakub Krajewski, Jan Ludziejewski, Kamil Adamczewski, Maciej Pióro, Michal Krutul, Szymon Antoniak, Kamil Ciebiera, Krystian Król, Tomasz Odrzygó2dz, Piotr Sankowski, Marek Cy- gan, and Sebastian Jaszczur. Scaling laws for ane-grained mixture of experts. arXiv preprint arXiv.2402.07871, 2024.
Tanishq Kumar, Zachary Ankner, Benjamin F Spector, Blake Bordelon, Niklas Muennighoff, Man- sheej Paul, Cengiz Pehlevan, Christopher Ré, and Aditi Raghunathan. Scaling laws for precision. arXiv preprint arxiv:241 1.04300, 2024.
13

---

Published as a conference paper at ICLR 2026
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 preprini arXiv:2006. 16668, 2020.
Lisha Li, Kevin Jamieson, Giulia DeSalvo, Afshin Rostamizadeh, and Ameet Talwalkar. Hyperband: A novel bandit-based approach to hyperparameter optimization. Journal of Machine Learning Research, 18(18)):1 2018.
Raymond Li, Loubna Ben Allal, Yangtian Zi, Niklas Muennighoff, Denis Kocetkov, et al. StarCoder: may the source be with you!. 2023.
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. /4342, 2023.
Iyya Loshchilov and Frank Hutter. SGDR: Stochastic gradient descent with warm restarts. arXiv preprint 1608 00w.::000800088,,, 2016.
Anton Lozhkov, Loubna Ben Allal, Leandro von Werra, and Thomas Wolf. FineWeb-EEdu the finest collection of educational content. Hugging Face, 2024.
Jan Ludzicjewski, Maciej Pióro, Jakub Krajewski, Maciej Stefaniak. Michal Krutul, Jan Malasnicki, Marek Cygan, Piotr Sankowski, Kamil Adamczewski, Piotr Miloš, et al. Joint MoE scaling laws: Mixture of experts can be memory efficient. arXiv preprint arXiv:2502.05172. 2025.
Kairong Luo, Haodong Wen, Shengding Hu, Zhenbo Sun, Zhiyuan I Liu, Maosong Sun, Kaifeng Lyu, and Wenguang Chen. A multi-power law for loss curve prediction across learning rate schedules. arXiv preprint urXiv:2503.12811 2025.
Sam McCandlish, Jared Kaplan, Dario Amodei, et al. An empirical model of large-batch training. arXiv preprint arXiv: 10:::81222.0 /2260//66 2018.
Alexandru Meterez, Depen Morwani, Jingfeng Wu, Costin-Andrei Oncescu, Cengiz Pehlevan, and Sham Kakade. Seesaw: Accelerating training by balancing learning rate and batch size schedul- ing. arXiv preprint arXtv:2510.14711. 2025.
Igor Molybog, Peter Albert, Moya Chen, Zachary Veto,o, David Esiobu, Naman Goyal, Punit Singh Koura, Sharan Narang, Andrew Poulton, Ruan Silva, al. A theory on Adam instability in large- scale machine learning. arXiv preprint arXiv:2304.099:: 009771 2023.
Niklas Muennighoff, Alexander Rush, Boaz Barak, Teven Le Scao, Nouamane Tazi, Aleksandra Piktus, Sampo Pyysalo, Thomas Wolf, and Colin A Raffel. Scaling cata-constrained language models. Advances in Neural Information Processing Systems, 36, 2023.
Niklas Muennighoff, Luca Soldaini, Dirk Groeneveld, Kyle Lo, Jacob Morrison, Sewon Min, Wei- jia Shi, Pete Walsh, Oyvind Tafjord, Nathan Lambert, et al. OLMoE: Open mixture-of-experts language models. arXiv preprint arXiv:2409.02066 2024.
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, 37: 172106666-127433 2024.
Team OLMo, Pete Walsh, Luca Soldaini, Dirk Groeneveld, Kyle Lo, Shane Arora, Akshita Bha- gia, Yuling Gu, Shengyi Huang, Matt Jordan, et al. 2 OLMo 2 Furious. arXiv preprint arXiv:2501.00656. 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.
14

---

Published as a conference paper at ICLR 2026
Ofir Press, Noah Smith, and Mike Lewis. Train short, test long: input length extrapolation. In International Conference on Learning Representations, 2022.
Attention with linear biases enables
Shikai Qiu, Lechao Xiao, Andrew Gordon Wilson, Jeffrey Pennington, and Atish Agarwala. Scaling collapse reveals universal dynamics in compute-optimally trained neural networks. arXiv preprint arXiv:2507.02119, 2025.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners, 2019.
Jack W. Rae, Sebastian Borgeaud, Trevor Cai, Katie Millican, Jordan Hoffmann, Francis Song, John Aslanides, Sarah Henderson, Roman Ring. Susamnah Young, et al. Scaling language models: Methods, analysis & insights from training Gopher, 2022.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of Machine Learning Research, 2020.
Stephen Roller, Sainbayar Sukhbaatar, Arthur Szlam, and Jason Weston. Hash layers for large sparse models. Advances in Neural Information Processing Systems, 34: 14:15555-77666 2021.
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 preprint arXiv:2501.18965, 2025.
David Sculley, Gary Holt, Daniel Golovin, Eugene Davydov, Todd Phillips, Dietmar Ebner, Vinay Chaudhary, Michael Young, Jean-Francois Crespo, and Dan Dennison. Hidden technical debt in machine learning systems. Advances in neural information processing systems, 28, 2015.
Neha Sengupta, Sunil Kumar Sahu, Bokang Jia, Satheesh Katipomu, Haonan Li, Fajri Koto, William Marshall, Gurpreet Gosal, Cynthia Liu, Zhiming Chen, et al. Jais and Jais-chat: Arabic- centric foundation and instruction-tuned open generative large language models. arXiv preprint arXiv:2308. 16149, 2023.
Noam Shazeer. GLU variants improve transformer. arXiv preprint arXiv:2002.05200 2020.
Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In International Conference on Machine Learning, pp. 4596-4604. PMLR, 2018.
Yikang Shen, Matthew Stallone, Mayank Mishra, Gaoyuan Zhang, Shawn Tan, Aditya Prasad, Adri- ana Meza Soria, David D Cox, and Rameswar Panda. Power scheduler: batch size and token number agnostic learning rate scheduler. arXiv preprint arXiv:2408. 13359, 2024.
Daria Soboleva, Faisal Al-Khateed Robert Myers, Jacob R Steeves. Joel Hestness, and Nolan Dey. SlimPajama: 627B token cleaned and deduplicated version of RedPajama. Web page, 2023.
Minhak Song, Beomhan Baek, Kwangjun Ahn, and Chulhee Yun. Through the river: Under- standing the benefit of schedule-free methods for language model training. arXiv preprint 25KV::2007.09866. 2025.
Jacob Mitchell Springer, Sachin Goyal, Kaiyue Wen, Tanishq Kumar, Xiang Yue, Sadhika Malladi, Graham Neubig, and Aditi Raghunathan. Overtrained language models are harder to fine-tune. arXiv preprint arXiv:2503. 2025.
Kevin Swersky, Jasper Snoek, and Ryan Prescott Adams. Freeze-thaw Bayesian optimization. arXiv preprint arXiv:1406.3896 2014.
Gemma Team, Morgane Riviere, Shreya Pathak, Pier Giuseppe Sessa, Cassidy Hardin, Surya Bhupatiraju, Léonard Hussenot, Thomas Mesnard, Bobak Shahriari, Alexandre Ramé, et al. Gemma Improving open language models at a practical size. arXiv preprint arXiv:2408.00 18, 2024.
Gemma Team, Aishwarya Kamath, Johan Ferret, Shreya Pathak, Nino Vieillard, Ramona Merhej, Sarah Perrin, Tatiana Matejovicova, Alexandre Ramé, Morgane Rivière, et al. Gemma 3 technical report. arXiv preprint a5Xiv:2503. 19786, 2025.
15

---

Published as a conference paper at ICLR 2026
Kimi Team. Kimi K2: Open agentic intelligence, 2025.
URL https://gittubbccm
MoonshotAI CatettAAAKKiiiiAii X2 wttp/////maa////tttttt.././///aa rew.rtcc.eppp.pp......
Howe Tissue, Venus Wang, and Lu Wang. Scaling law with learning rate annealing, arXiv preprint arXiv:2408. 11029, 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, 2023a.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Niko- lay Bashlykov, Soumya Batra, Prajiwal Bhargava, Shruti Bhosale, et al. LLaMA 2: Open foun- dation and fine-tuned chat models. arXiv preprint arXiv:2307.09288 2023b.
Lewis Tunstall, Edward Beeching, Nathan Lambert, Nazneen Rajani, Kashif Rasul, Younes Belkada, Shengyi Huang, Leandro Von Werra, Clémentine Fourrier, Nathan Habib, et al. Zephyr: Direct distillation of LM alignment. arXiv preprint 11111223100.6444,, 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Infor- mation Processing Systems, 2017.
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: 1036-000000, 2023.
Nikhil Vyas, Depen Morwani, Rosie Zhao, Mujin Kwun, Itai Shapira, David Brandfonbrener, Lucas Janson, and Sham Kakade. SOAP: Improving and stabilizing shampoo using adam. arXiv preprint arXiv:2409. 13321, 2024.
Chengyu Wang, Junbing Yan, Yuanhao Yue, and Jun Huang. DistilQwen2.5: Industrial practices of training distilled open lightweight language models. arXiv preprint arXiv:2504. /5027, 2025.
Xi Wang and Laurence Aitchison. How to set AdamW's weight decay as you scale model and dataset size. arXiv preprint arXiv:2405.13669 2024.
Kaiyue Wen, Zhiyuan Li, Jason Wang, David Hall, Percy Liang, and Tengyu Ma. Understanding warmup-stable-des decay learning rates: A river valley loss landscape perspective. arXiv preprint arXiv:2410.05192 2024.
Mitchell Wortsman, Peter Liu, Lechao Xiao, Katie Everett, Alex Alemi, Ben Adlam, John Co- Reyes, Izzeddin Gur, Abhishek Kumar, Roman Novak, et al. Small-scale proxies for large-scale transformer training instabilities. arXiv preprint arXiv:2309. 14322, 2023.
Lechao Xiao. Rethinking conventional wisdom in machine learning: From generalization scaling. arXiv preprint urXiv:2409. 1510:::::00.11166.. 2024.
An Yang. Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chengyuan Li, Dayiheng Liu, Fei Huang, Haoran Wei, et Qwen2.5 technical report. arXiv preprint 2m11v12412.15115. 2024.
An Yang, Anfeng Li, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao. Chengen Huang, Chenxu Lv, et al. Qwen3 technical report. arXiv preprint arXiv:2505.09388, 2025.
Greg Yang and Edward J Hu. Feature learning in infinite-width neural networks. arXiv preprint arXiv:2011.14522 2020.
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 Newral Information Processing Systems, 2021.
16

---

Published as a conference paper at ICLR 2026
Greg Yang, Dingli Yu, Chen Zhu, and Soufiane Hayou. Tensor programs VI: Feature learning in infinite-depth neural networks. arXiv preprint arXiv:2310.02244, 2023.
Arber Zela, Aaron Klein, Stefan Falkner, and Frank Hutter. Towards automated deep learning: Efficient joint neural architecture and hyperparameter search. arXiv preprint arXiv:1807.06 2018.
Ge Zhang, Scott Qu, Jiaheng Liu, Chenchen Zhang. Chenghua Lin, Chou Leuang Yu, Danny Pan, Esther Cheng, Jie Liu, Qunshu Lin, al. MAP-Neo: Highly capable and transparent bilingual large language model series. arXiv preprint arXiv:2405.1932 2024a.
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.
Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christo- pher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, Todor Mihaylov, Myle Ott, Sam Shleifer, Kurt Shuster, Daniel Simig, Punit Singh Koura, Anjali Sridhar, Tianlu Wang, and Luke Zettlemoyer. OPT: Open pre-trained transformer language models, 2022a.
Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christo- pher Dewan, Mona Diab, Xian Li, et al. Opt-175 logbook. PDF, 2022b.
17

---

Published as a conference paper at ICLR 2026
A LIMITATIONS AND FUTURE DIRECTIONS
Across our initial (Secs. and 5) and Celerity (Sec. 4) setups, we have tested collapse across two distinct settings of architecture (and context length), dataset (and vocabulary size), and parameteri- zation. We directly compare TLCs from these two settings in Appendix B.4, while also describing further experiments in learning-rate schedule (Constant vs. 10x vs. D2Z), Adam /1/ parameters, and dense vs. sparse mixture-of of-urr-----eeerrr (MoE) architectures. However, in all cases our results are established under single-epoch pre-training with AdamW. The observed patterns may change un- der extreme TPP, multi-epoch training, alternate optimizens/schedules, or heavy mid-training data annealing/curricua..
Optimizers. We hypothesize the optimizer timescale will remain primary control of TLC shape for other optimizers with decoupled weight decay (e.g.. Sophia (Liu et al., 2023), MuonClip (Team, 2025)) whose update rules can be expressed in EMA form analogous AdamW (Sec. 2). Likewise, the perspective should also hold when Adam W is applied in alternate weight bases, c.g., as in SOAP (Vyas et al., 2024), where AdamW is applied in Shampoo's eigenbasis (Gupta et al., 2018). Extending timescale analysis to optimizers without a natural EMA form (e.g., Adagrad (Duchi al., 2011), Adafactor (Shazeer Stern, 2018), SGD variants) is an important future direction.
Data curricula. Given the growing use of data curricula and ltte- data annealing in LLM pre-training, it is valuable to study how shifts in data affect TLC shape across scales. Collapse may also inform curriculum design by serving as transfer marker. For example, observing limited opportunities for experimentation at large scale, Feng et al. (2024) experiment at smaller scales with downsampled datasets that simulate the repetition occurring larger sizes (due to limited high- quality tokens), Consistency in TLC shape could serve as an indicator of whether the downsample proportions reflect a consistent overfiting/geneeralzatiin trade-off across scales. Collapse can thus serve to confirm that smaller-scal settings provide suitable proxies for optimizing data mixes and other settings.
Celerity extensions. Beyond choosing TPP (controlling placement on the cost/compression curve; Fig. 5), we aim to understand which factors or training strategies shift the curve itself. For faster inference, we especially interested in the location of the parameter wall- the minimal capacity achieving a target loss-and and how architecture (dense vs. MoE), routing, and depth/width changes affect collapse and efficiency.
We intentionally chose dense models for our initial Celerity series because dense models have fewer confounding factors (e.g., routing strategy, number of experts), making them simpler study and build upon. In our own practice, algorithmic innovations are typically validated on dense models first. However, we are interested in scaling MoE-variants of our Celerity series due to their docu- mented savings in training compute (Krajewski et al., 2024; Ludziejewski et al., 2025).
Train loss vs. generalization. We focus on training loss because (i) it is FLOPs-free to monitor, (ii) in LLM pre-training it typically tracks validation under stationary data, and (iii) it surfaces is- sues earlier (e.g., duplicated segments), enabling targeted intervention before held-out degradation. Late-stage annealing and domain shift decouple train-loss collapse from downstream behav- ior. We study validation collapse for Celerity Appendix C.5; future work should also consider downstreum-collappe. and measure train--vvll-doowstrean correlations across schedules and data mixtures.
Predictive model and schedules. Both collapse itself, and our predictive model's ability to accu- rately forecast normalized TLCs, is impaired by loss spikes and divergences, which move the nor- malized curve away from the universal trajectory (sometimes temporarily, sometimes for extended periods). From one perspective, this is a feature not bug, the resulting collapse anomalies provide a useful mechanism for detecting training issues (discussed further below).
Empirically, dividing by the final training loss (i=0) aligned curves best; future work will study why irreducible-loss offsets, in Qiu et al. (2025), were not beneficial. In terms of our predictive model. next steps include factoring LR envelope vs. anncal-phase effects (cf. (Tissue et al., 2024; Luo et al., 2025)), adding uncertainty (e.g., seed bootstraps) and uncertainty- early-stopping
18

---

Published as a conference paper at ICLR 2026
Table 3: Model architectures used in Sec. 3 and Sec.
Model
dmodel
Tlayers
dit
dhead
111M
768
10
2048
64
266M
768
32
2048
64
610M
2048
10
5461
64
1.7B
2048
32
5461
64
3.3B
2048
64
5461
64
policies. Given that in our experiments, the parametric predictor was fit for one specific LR schedule (D2Z), we should also revisit whether b(T), q(TP)) and possibly m vary systematically across cosine (Loshchilov Hutter, 2016), inverse square-root (Vaswani al., 2017; Raffel et al., 2020; Shen et al., 2024), and warmup-stable-decay (WSD) (Hu et al., 2024; Hägele et al., 2024; Wen et al., 2024; Song et al., 2025) schedules, and schedule-free schemes (Defazio et al., 2024).
It would also be interesting to measure collapse when batch size schedules are used. Theoretically, we could maintain collapse by adjusting weight decay whenever batch size changes, maintaining the r invariant. We could compare such adjustments to adjustments of LR, e.g., as in Meterez et al. (2025).
Systems effects and making collapse a practice. Collapse residuals are sensitive to systems wwm.rmbacchinn/cccuuuuaaaa
precision, kernel changes, restarts -whahh can create arti- facts or reveal true pathologies. To pay down "hidden technical debt" (Sculley et al., 2015), we advocate lightweight collapse monitor: log fraction-of-data in addition to raw step count (easy to add in TensorBoard (Abadi et al., 2016)), as well as microbatch statistics and restart boundaries; nor- malize online and alert when residuals exceed policy thresholds. Treating collapse as operational invariant reduces configuration fragility and surfaces data/numerics issues early.
B EXPLAINING TLC SHAPE: FURTHER DETAILS
B. FULL EXPERIMENTAL DETAILS
In this section, we provide details on the model architecture (Table 3) and training data (Table 4) for models used in experiments in Sec. Sec. 5, and elsewhere in the appendix. Experimental details for the Celerity model series are in Appendix C.2.
In total, =600 TLCs were analyzed for these experiments. All such models were GPT2-style LLMs (Radford et al., 2019) with ALiBi (Press et al., 2022) embeddings and SwiGLU (Shazeer, 2020) non-lincarity. We use the AdamW optimizer. Following standard practice, we do not apply weight decay or bias to LayerNorm layers. Default AdamW settings are /31 0.9, B2 = 0.95, and 1e-8. We report cross-entropy loss. We parameterize with maximal update parameterization, uP (Yang et al., 2021), with hyperparameters set via proxy tuning, as described below. For 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, warmup 10% of the total steps. Learning rate warmup is stan- dard practice in LLM training (Brown et al., 2020; Rae et al., 2022; Biderman et al., 2023; Dubey et al., 2024; Kosson et al., 2024).
Note Fig. (and later Table 2, Fig. 7, Fig. 8, and Fig. 9) report batch size in sequences rather than tokens.
All models in the main experiments were trained on a Cerebras CS-3 system. 610M parammtr 20TPP models take roughly 6 hours each to train single CS-3.
Proxy model hyperparameter tuning. To find optimal μP hyperparameters (HPs), we trained 39M proxy model using a width dproxy of 256, with 24 layers and head size of 64. We trained this model on 800M tokens with B=256 sequences and context length 2048. We randomly sampled 350 configurations of base learning rates, base initialization standard deviation, and embedding and output logits scaling factors, and used the top-perforning values tuned HPs (Table 5).
19

---

Published as a conference paper at ICLR 2026
Table 4: Models, sasens-per-paaameee (TPP) and corresponding dataset sizes (in tokens), number of model variants trained (over LR schedule type, X, B) for models used in Sec. 3 and Sec. 5. In total, =:00 TLCs were analyzed.
Model
TPP
D
Variants trained
111M
20
2.19B
74
111
80
8.76B
50
111M
200
21.9B
28
111M
320
35.0B
40
111M
1280
140.1B
11
266M
20
5,31B
25
266M
80
21.2B
19
266M
320
85.0B
19
266M
1280
339.8B
3
610M
20
12.1B
205
610M
80
48.5B
53
610M
200
121.3B
610M
320
194.1B
5
1.7B
20
34.3B
31
1.7B
80
137.2B
11
1.7B
160
274.3B
1.7B
320
548.6B
3.3B
20
66.5B
2
3.3B
23
76.5B
3.3B
30
99.8B
5
Table 5: Tuned hyperparameters for jP proxy model for models used in Sec. and Sec. 5.
Wase T)
8.67e-02 1.62e-02
Olinput 9.17
Cooutput
1.095
20

---

Published as a conference paper at ICLR 2026
It is worth noting that the LR values reported in this paper and shown in figures base μPP LRs before uP-adjjstment. Calculation of т (Sec. 2) requires the adjusted LR (i.e., multiplying by aprocy model).. Also, when LR decay is used, reported LR values always refer to the peak/max LR of the LR schedule.
B.2 EXPLAINING TLC DEPENDENCE ON TPP
Schedules with decaying LR reach their minimum value only at the final step (after D tokens). However, for a constant LR schedule, every step of training is equivalent to a complete training run ending at that step. Qiu al. (2025) make the observation that therefore the loss at every training fraction == L/T E [0, 1] should respect the same fitted scaling law, but for training budget of f. D tokens.
Starting from the Chinchilla functional form L(N, D) = with a constant LR schedule, training until a certain final every fraction of training I, we will have trained for an intermediate TPP of i. i.e., using N- total tokens. To arrive at scale invariance, and dataset exponents a and / were roughly equal: this rough equality has also been repeatedly validated in replication studies (Besiroglu al., 2024; Porian et al., 2024). Using = B, and focusing on the reducible loss, we obtain a final training loss of:
assume that we train
ratio A = D/N. At
note that Hoffmann et al. (2022) found their fitted model
L(N, k N) AN- + B(k N
AN- + "N
(6)
Meanwhile, training for training fraction i, the predicted loss is L(N, L l N) = AN + Bi "k "N"NN"
(7)
We now normalize by the final loss expose the shape of the training loss curve. The resulting
normalized loss L(N, f k NN/L(N, k N) is independent of model (and dataset) size, depending only on the training fraction f and the target TPP ratio k:
Bi-ak-
E(f, k) = A+ Bk-
(8)
In other words, for TLCs using a constant LR schedule, collapse approximately holds under this normalization. Scaling of l in t also motivates own TLC predictive form (Eq. (4)).
As shown in Appendix C.I. a a = 3 implics there a single optimal TPP ratio r, and morcover, that the Chinchilla coefficients obey B = Ard. For given training run, suppose that the TPP at which we train is a multiple of the optimal TPP by the ratio v, e.g., k = u. Thus, e = corresponds to overtraining. We can reparameterize the E
corresponds to optimal TPP, while equation in terms of v as:
1 1 0
(ff,) v)
1 1+ vー
This simple equation clarifies how the overtraining factor e influences the shape of the TLCs. When is small (undertraining), the power law term dominates, and the TLC gradually decays in f. When v is large (overtraining). the power law only plays a role for smaller f, the curve drops quickly and 1. Intuitively, for overtrained models, we make gains quickly the beginning
then flattens of training and then obtain diminishing returns as training progresses.
Qiu al. (2025) further show that for non-uniform LR schedules, the loss curve deformed by (0). but, given consistent curvature of the loss landscape across model scales under juP (Noci et al., 2024), the noise-induced deformation is invariant to model size, and thus collapse still holds.
B.3 EXPLAINING TLC DEPENDENCE ON
As noted in Sec. 3, the AdamW timescale (n(nXT) controls the effective memory length of the parameter updates: smaller 7 emphasizes recent updates (bias reduction), while larger 7 averages more broadly (variance reduction). In this sense, acts an implicit batch size.
21

---

Published as a conference paper at ICLR 2026
To provide further insight into the role of 7 in shaping TLCs, now derive analytical expression for training loss under a constant learning rate, using a simple noisy quadratic model (NQM). While LLM training minimizes cross-entropy loss, it is common to perform a local quadratic approxima- tion, i.e., a second-order Taylor expansion in the parameters, with the constant Hessian replaced by the instantaneous Hessian along the training trajectory (LeCun et al., 1989). Thus conclusions drawn from quadratic models often generalize to large, realistic networks (Zhang et al., 2019).
Setup. Following Zhang et al. (2019), we assume the optimizer dynamics are invariant to rotation and translation, allowing model, without loss of generality, locally quadratic loss, separable across dimensions, and having an optimum at zero. Specifically, we consider single quadratic mode with curvature h O, optimum at 0 0, and parameters 0,, where the step index:
((t) 1 h
(10)
With AdamW optimization, Og evolves as an exponential moving average (EMA) of stochastic up- dates 2, with constant smoothing a y7 (Sec. 3):
02 = (1 a) 0 + 0
(11)
Unrolling the recurrence gives the general form
0. = a:) 0o
(1-a)-----
az.
(12)
4=0
The first term is the (decaying) contribution of the initialization, while the second term reflects the accumulation of stochastic updates.
Continuous (fraining-fraction) limit. We now switch to fractional time i= L/T and define the AdamW timescale - = 1/(oT)) Approximating (1 o)--1- x e-((t-1-i) and interpreting the sum as Riemann approximation T co, we obtain
1 о f.e e- e(-fss)/T (ss)ds
e(f) 1/7 0(0)
+
(13)
That is, 0 consists of contributions: decaying memory the initialization, and convolution of the update signal r(s) with an exponential kernel of timescale (an EMA filter over updates).
Noise model. Following Zhang al. (2019), model the update signal s=f)) preconditioned white noise: a zero-mean process with constant variance of and no temporal correlation,
[f)( 0,
E [[=(f)((s)) of 5(t -).
This idcalized assumption isolates the effeect of by removing structure in gradient noise beyond its overall scale.
Mean and variance.
The EMA filter preserves initialization, which decays exponentially:
B(((O) e -//
i/7 0(0).
The variance from stochastic updates is
Var[O(f)] = 02 (1- 11-c-21/ -21//)
(14)
27
Thus the total second moment is
a2 (1-e -2/T7)
B(0(i)²) = -21/ 0(0)2
2
+ 2T
In words, the initialization bias decays away on timescale T, while variance from noisy updates accumulates toward floor proportional to 1/T.
22

---

Published as a conference paper at ICLR 2026
Expected loss.
The per-mode loss is
L(i) hh((1)²
Taking expectation, and using the decomposition into bias and variance,
E[L(O] = 1/hh ( ( 0(0)2 (1)) oe
2//7
2T (1- 26/7)
(15)
The first term reflects exponentially decaying initialization bias, while the second reflects variance accumulation to floor proportional to 1/7.
If intiiialiiztoonn is zero-mean in expectation (E] = 0), the bias term vanishes and the expression simplifies to
h ho
B[I(O)]
4T - (1-e-2//)
(16)
which captures the characteristic fast- then-flattenn TLC shape under constant learning rate.
Interpretation. Equation 15 decomposes the expected loss into an exponentially decaying bias term (xx e 26/r 1022////000 ((0))) and variance term that rises to floor (xx 1/T). This yields two opposing effects of 7 on TLCs:
Smaller 7 suppresses initialization bias more rapidly (via the e 26/r decay), but accumulates higher variance, yielding higher asymptotic loss floor (x 1/T).
Larger 7 reduces variance more effectively, lowering the final loss, but is slower eliminate bias from initialization.
When initialization is zero-mean in expectation, the bias term vanishes and the expression reduces to Eq. 16.
Constant LR: + slower drop
10x LR decay: T + faster drop
decay-to-zero: 1 faster drop
1.4
(x)
1.4
7 (+)
1.4
7 (x)
1.0ee (8.)
(0.811)
(0.84L)
2.02e (0.841)
Nomalized train loss
4.05e (0.421)
Normalized train loos
4.05e (0.421)
Training lazs
1.3
4.05e-03 (2.)
2.3
8 0be-00 (0.210)
1.3
8.09e-03 (0.210)
8.09e (0.210)
1 62e (0.105)
1.6e 02 (0.105)
1.2
1.62e (0.105)
1.2
(0.053)
(0.053)
3.24e-02 (0.053)
6. 6.48e-02 (0.026)
1.1
1.1
1.1
1.0
1.0
1.0
0.22
0.50
1.00
0.25
1.00
0.00
0.75
Fraction training tokens,
Fraction
training takens,
Fraction training tokenss,
Figure 10: Effect of LR schedule TLC shape (610M, 80TPP). Left: Constant LR, Middle: Linear 10x decay, Right: Linear decay-to-zero. Different schedules deform the TLCs distinct ways, yet in all cases the AdamW timescale governs the bias-variance trade-off. With a Constant LR, smaller 7 accelerates early loss reduction. With D2Z, the effect inverts: smaller yields larger late-stage drop. Although here is varied by changing LR, equivalent effects arise when varying weight decay or batch size (Fig. 3), confirming as unifying control knob for TLC shape.
This interpretation matches our empirical findings for normalized Constant-LR TLCs (Fig. 10, left). The situation is different, however, for LR decay schedules, which we discuss next.
Finally, 7 a normalized timescale and thus invariant the absolute number of steps. In the NQM, after normalizing by the final loss L(1) the curvature factor h cancels exactly. The normalized curve takes the form
L(t) L(1) (1 (1- 21TT tee 2/T
e1-e 28/T) + tke 2//T
2T E
Thus, when the initialization contribution is negligible by the end of training (or when the ratio is approximately scale-invariant), the normalized TLC depends only (7, f), and curves matched
23

---

Published as a conference paper at ICLR 2026
collapse across model sizes. If varies across scales, small early deviations can appear (bias- dominated regime) but typically diminish as e e-2f/r decays.
Remark. Qiu et al. (2025) observed collapse without AdamW. Empirically, as À 0 TLCs approach a limiting shape: vanilla Adam behaves like AdamW with x= (effectively =o0)
Extension to decaying LR schedules. The constant-LR analysis in Eq. 16 shows that sets the trade-off: smaller 7 accelerates early bias reduction but saturates at a higher variance-driven floor, while larger reduces variance more slowly but to a lower asymptote. With decaying LR schedule, the smoothing at = THt À decreases after warmup, so the instantaneous timescale т; = 1/(mxAT')) in- creases training progresses. In this setting, small-7 runs still make rapid early progress (fast bias reduction), but during the decay phase they gain additional variance suppression as T lengthens, often producing noticeable late-stage drop in loss. By contrast, large-7 runs emphasize variance reduction throughout, yielding steadier curves without the same end-of-training acceleration. Equiv- alently, in the EMA view, decay flattens the contribution coefficients Cr, averaging over more (car- lier) updates near the end. Thus LR decay effectively combines the early bias-reducing dynamics of small with the late variance-reducing dynamics of large т, inverting the TLC ordering observed under constant LR (Fig. 10).
This analysis aligns with Bergsma al. (2025b), who attribute the effectiveness of D2Z sched- ules to balancing early bias reduction with later variance suppression (building on D' Angelo et al., 2024). Their treatment primarily conceptual; here we show how the same virianariance dynamics manifest directly in TLC shapes and provide simple analytical form under the NQM.
B.4 ADDITIONAL TLC EXPERIMENTS
Constant LR: (TPP, +] const.
10x-Decay LR: (TPP, r) const.
Decay to zero LR: (TPP, 7) const.
1.4
N
1.4
N
1.4
NN
T-0.33
0-3.3
0-3.33
Normalized train loas
610M,
r-0.22
Normalized train loss
-610M --0.2
Normalized train loas
610M, 7-0.2
1.3
0.330
1.3
1.78, r-0.30
1.3
1.78, --0.30
1.2
1.2
1.1
1.1
1.1
1.0
1.0
1.0
1.00
0.25
1.00
0.25
0.75
1.00
Fraction training tokens, 2
Fraction training tokens,
Fraction training tokens,
Figure 11: Collapse in differrnt LR schedules, Left: Constant LR, Middle: Linear 10x decay, Right: Linear decay-to-zero. In contrast to Fig. 10, where 7 varies, here TPP-20 and t 0.3: curves collapse across scales.
Collapse under alternative LR schedules. Fig. 11 shows that normalized TLCs also collapse under a Constant schedule, a 10x decay schedule, and our decay-to-zero schedule (all with 10% warmup). At corresponding model sizes, we the same batch size, peak LR, and weight decay, so same-size results across schedules differ only in their final LR. Collapse is slightly looser than in the Celerity runs because the resulting not matched exactly across schedules (see plot an- notations), but the qualitative agreement is strong. These results are consistent with our analysis in Appendix B.3 and echo the cross-schedull findings of Qiu al. (2025).
Collapse across datasets and architectures. TLC shape can in principle depend task, data, and architecture (c.g., multi-epoch training on a small corpus can yield faster apparent improvement than single-epoch pre-training). We therefore ask: how much does normalized TLC shape change as we vary parameterization, vocabulary size, architecture, context length, and dataset mix?
As first probe, compare Celerity TLCs to earlier non-Celerity runs, at the same TPP (20) and similar т s 0.2, while varying all items above: Celerity uses CompleteP (vs. vanilla pP), a larger vocabulary, different nonlinearity and FFN multiplier, 4x longer context, and different data
24

---

Published as a conference paper at ICLR 2026
Collapse across dataset/ /architecctre
Sparse mixture-of-experts (MoE)
Varying Adam and A
Model
experts
2.0
8.(%)
Normalized training loss
610M), 0
111M", T-1.666
Normalized training loss
8
16
Normalized troining loss
1.8
0.0 (0.95)
0.5 (0.95)
1.3
1.78I, r-0.149
1.3
32
0.9 (0.95)
900M". T=0.175
1.6
(0.99)
1.2
1.2
0.95 (0.99)
1.4
0.99 (0.9900)
1.1
L.1
0.999 (0 (0.9999)
1.2
1.0
1.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 takens, 2
Fraction of training tokens,
Fraction of training tokens,
Figure 12: Collapse across 'original, 'Celerity setups. 20 TPP.
Figure 13: Collapse as E varies in a sparse MoE. 111M, 0.33, 20 TPP.
Collapse as Adam 31 and 32 vary.
Figure 14:
610M, 0.21, 20 TPP.
mixture (Appendices B.1 and C.2). Despite these differences, the TLCs loosely collapse (Fig. 12). The Celerity 900M model tracks closer to the 610M model than the 1.7B model, although its intermediate between these two. Overall, we view this as evidence that the normalized TLC shape surprisingly robust when LR schedule, TPP, and r are held (approximately) fixed.
Collapse in sparse mixture-of-expertt (MoE). We next analyze sparse MoE architectures. where only a subset of parameters are active per token (Lepikhin et al., 2020; Fedus et al., 2022). Starting from our 111M dense model (Appendix B.1), we replace each FFN with a sparse MoE layer and vary the number of experts E E {1, 2, 4, 8, 16, 32}. Tokens are routed to one expert via hash routing (Roller et al., 2021) so each expert processes a similar token count. Global training tokens and datasets are identical across E, hence the effective TPP per expert decreases from 20 (dense) to 20/E E as E grows. Note also that as the number of experts E' increases, and the effective tokens per expert decrease proportionally, both the expert's ''s effective batch size B and effective dataset size D are reduced by a factor of E. Since T = B/ (n)D)), these reductions cancel, leaving the overall timescale unchanged (for fixed 11, A)-
Fig. 13 shows that lower E (higher effective TPP per expert) yields slightly earlier drops and slightly flatter tails, broadly obeying the TPP effect characterized in Sec. 3. Thus, the observed deformation is explained by effective TPP rather than differing training dynamics per se.
Collapse across Adam 3, and 32. Finally, we vary (3,, at fixed LR, batch size, and weight decay (r 0.21, 610M model, 20 TPP). The default (0.9, 0.95) gives the lowest absolute loss in this experiment, but several "standard" settings- (0.9, 0.95). (0.95, 0.99), and even (0.99, 0.9999)- produce normalized TLCs that collapse (Fig. 14). In contrast, runs with (0.0, 0.95). (0.5, 0.95), and a noisy instance of (0.9, 0.99) exhibit early loss spikes; when the loss fails to recover promptly, the curves remain elevated and do not rejoin the main trajectory, breaking collapse (early loss spikes also distort early collapse for noisy, large-batch-size runs, e.g., Fig. 7). We also observe that (0.999, 0.9999), which aggregates gradients over a much longer horizon, follows a systematically slower (but eventually convergent) trajectory-ccosiiterr with enlarged momentum timescale prioritizing variance reduction over bias, akin to increasing T.
Overall, aside from extreme momentum settings or instability-induced spikes, setting of (31, has limited effect on the shape of normalized TLCs. The AdamW timescale 7 remains the dominant optimization polmiization-based control for TLC trajectories.
Key takeaway 4: Normalized TLCs are strikingly robust: they largely collapse across diverse datasets and architectures, remain predictable under sparse MoE routing (scaling in effective TPP theory suggests), and are insensitive to typical Adam Bi, B2 settings. Apart from pathological loss spikes, the dominant factor shaping TLCs is still the Adam W timescale T.
25

---

Published as a conference paper at ICLR 2026
C CELERITY MODELS: FURTHER DETAILS
C.I COMPUTE COST AS A FUNCTION OF MODEL COMPRESSION
Starting from compute-optimal model size, required (C/Copt) to compress a model to a smaller (less efficient) size, while maintaining the same loss. We use the resulting equation to plot the compression vs. cost trade-off in Fig. 5, This analysis motivated the selection of max TPP in the Celerity model series.
derive expression for the extra compute
We begin again with the Chinchilla functional form from Hoffmann al. (2022), giving loss as function of model size N and data size D:
L(N, = + AN + BD-
(17)
where F,, A, B, and 3 are parameters to be fit on observed training runs.
Hoffmann et al. (2022) asked, for fixed training compute budget C' (in FLOPs), how should we allocate model size N versus number of training tokens D in order to minimize loss? From Eq. (17), they derived functions for loss-optimal (C)) and Dope(C) (constraining L(N, D) by the com- mon approximation C' ~ 6N
Nopt (C) =C)= ( G
a+d
a+d
and Dopt = G-1
(18)
Results indicated that Nop: and Dopt scale roughly equally as C increases. This analysis agreed with their other methods for estimating compute-optimal scaling, and guided their N and D allocation for training their large-scale Chinchilla model.
where G
AA 35
Let r be the optimal (C) (C) ratio. If r is roughly independent of C, this implies a = B.
Using a a B, we obtain:
B
(19)
or equivalently B = Ar
Replication studies have found 3 = 0.35, and optimal TPP of around 20 (Besiroglu et al., 2024; Porian et al., 2024) (as noted Sec. 2).
Now, suppose a = a = 3 and we obtain loss of L at the optimal TPP ratio (where Dopt = r:Nopt):
i AN
opt BD-
upt + (Ar") (r.opt)
= E + AN- opt
E + 2-N- opt
(20)
We now wish to train compressed model with fraction ky of parameters compared to Nopt but obtaining the same loss. Let kx Nopt- If to Dopt in order to reach the loss target. Let D kn Dnpt Rather than training at r TPP, we will (
Nopt. we will need kn extra tokens compared
train at a higher ratio derivation to De Vries (2023), the estimated loss will be:
= (kD/KN)r From Eq. (17), and following a similar
L(N,D) E A(kx B(k
(21)
Again substituting a = a B and B = Ar to obtain the target loss L, we set the loss in Eq. (21)
to equal f. in Eq. (20), and solve for kD, finding:
kp (2 - kN =
(22)
The compute cost C' of the compressed training will be derive the extra compute ratio compared Copt 6.Nope Dopt:
from which we can
C/Copt kNNDD
k2 - kN -a)-
(23)
26

---

Published as a conference paper at ICLR 2026
Eq. (23) allows to vary ky and obtain the corresponding compute overhead. When planning the Celerity training runs, we assumed r = 20 corresponded to the compute-optimg model size (following the Chinchilla rule-of-thumb) and we tested different values of reported in prior work, using a = 0.35 in Fig. 5.
C.2 CELERITY RECIPE DETAILS
In this section, we provide further details for the techniques that most impacted Celerity's 's perfor- mance and compute efficiency, including parameterization, learning rate and weight decay schedul- ing, architecture, and dataset construction.
3.0
Celerity ((C)=(cTTEEE) 0.04
2.9
HP L(C) (())-(yvies? 0.54
2
2.8
CompleteP ((C)=/pmee -0.031
Va 1on 9 9
2.7
Validation
2.6
2.5
2.4
2.3
1000
10
10a2
Compute FLOPs (C)
Figure 15: Scaling law comparison between CompleteP and μP. CompleteP scales more predictably: the power law is a better fit to the observed points. CompleteP also exhibits steeper slope, improv- ing loss faster in compute FLOPs, likely due to both better HP transfer (across model width and depth), and better compute efficiency.
Parameterization. We compare the effect of different parameterizations and their influence compute efficiency in Fig. 15. Specifically, we compare μP (Yang et al., 2021), which accounts for scaling in width, and CompleteP (Dey et al., 2025), which accounts for scaling in both width and depth. For each parameterization, the hyperparameters (HP), such as learning rate, weight initialization, and multipliers, tuned at depth 32 and then directly applied the target model training.
Here note observations: the first observation is that μP points do not align well on standard scaling law. We attribute this to HP de-tuning when transferring HPs from proxy model depth to target model depth. Such de-tuning is not seen the scaling laws for CompleteP and Celerity (which uses CompleteP but a different dataset), where the points align with minimal error on the scaling law. Dey et al. (2023a) showed poor scaling law fits when scaling width in the standard parameterization vs. μPP we believe a similar phenomenon is now happening when scaling in depth.
The second observation is that CompleteP is more compute efficient than jaP, which prior work has explained through the lens of feature learning (Dey et al., 2025). Indeed, Fig. 4 in Dey et al. (2025) suggest that CompleteP models exhibit better scaling behavior than j.P even when both are comprehensively tuned.
Based on these observations, we used CompleteP for training the Celerity series, using the proxy model's tuned HPs across all scales. Sample code for implementing CompleteP is available at https: / /ttptbb..tttt///////////tt...h com/
httpr//nuuuuptteee//mooeeeeagg completep.
Learning Rate and Weight Decay. We chose the linear decay-to-zero (D2Z) learning rate sched- ule based its empirical success and conceptual motivations in Bergsma et al. (2025b). In par- ticular, Bergsma et al. (2025b) showed that as TPP increases beyond compute-optimal 20 TPP, the relative benefit of D2Z also increases, in scale-invariant manner. This makes D2Z particularly appropriate for parameter-efficins train with linear warmup to the peak LR, over the minimum of 10%-of-total-tooer
training (c.g.. Celerity's s 234 TPP model band). All models also
375M-tokens.
27

---

Published as a conference paper at ICLR 2026
We tuned smaller scale and smaller TPP, and transferred across TPP using the power law fit from Bergsma et al. (2025a). Given learning rate is determined by CompleteP, and batch size is optimized according to separate scaling rule (described below), adjusted weight decay in order to obtain the desired 7 setting at each scale.
Batch Size. In early experimentation, the batch sizes were chosen such that they around the critical batch size (McCandlish et al., 2018). Later we used the insights from Bergsma et al. (2025a) and started following the rule Bop: x D0.5, tuning B at small scale and then inferring optimal batch sizes on larger datasets via the power law.
Table Composition of the Celerity pre-training dataset.
Data Subset
Percentage (%)
Fine FineWeb-Eod (Lozhkov et al., 2024)
64.75
StarCoder (Li et al., 2023)
10.8
Cosmopedia (Ben Allal et al., 2024)
4.66
SlimPajama (Soboleva al., 2023)
17.49
Open WebMath (Paster et al., 2023)
1.88
TextBookss (Gabarain, 2024)
0.42
Data Selection. Over the course of experiments, we found that adding more refined data, partic- ularly educational, math, and coding datasets, generally helps the models score higher common benchmarks. In Table 6, we break down the datasets used for Celerity model training, including the proportion assigned each subset. A large portion of the datasets focused on educational materials, math, and coding. We use only the non-web-crawled subsets of SlimPajama (Soboleva al., 2023), i.e., excluding C4 and CommonCrawl, which are effectively replaced by Fine Web-Edibl As noted in Sec. 4, we do not schedule the data sources, i.c., we do not employ data curriculum in the training of Celerity, nor do we include (benchmark) task-specific data Celerity training.
Table 7: Comparison of Celerity models trained on different datasets.
Name
Downstream Accuracy (Num Shots)
arc-c
boolq
hellaswag
piga
siqa
winogrande
Avg.
(25)
(0)
(0)
(10)
(0)
(0)
(5)
Celerity 300M
27.82
50.63
52.75
37.57
66.21
37.77
52.25
46.43
Celerity 300M SlimPJ
24.32
42.17
61.53
36.04
65.56
37.97
50.99
45.51
Celerity 900M
39.68
64.52
47.92
55.02
72.03
41.97
58.48
54.23
Celerity 900M SlimPJ
30.89
54.67
55.47
53.74
71.00
40.89
57.46
52.02
Comparison in Table shows that the same model configurations trained on general dataset like SlimPajama result in worse downstream performance compared to the Celerity data mix. While dataset optimization was not focus of Celerity, these results do underscore the importance of dataset composition in pre-training. This also makes clear why hyperscalers invest tremendous amount of work into data preparation, synthesis, filtering, and refinement.
Table 8 summarizes the dataset sizes for all models the Celerity model series.
Model Architecture. Celerity models use a decoder-only GPT2-style transformer architecture. Table summarizes the architecture dimensions, hyperparameters, and other details of the Celerity model family. We trained five Celerity model sizes from scratch with parameter counts roughly 300M, 500M, 900M, 1.8B, and 3.9B. All models are trained under consistent data and optimization methods, public datasets, in order to foster open science and fair comparison.
C.3 CELERITY FURTHER RESULTS
In our empirical evaluation of Celerity, we necessarily only compare to model families with suffi- ciently precise and complete training details, in particular the total training tokens. E.g., Llama-3.2
28

---

Published as a conference paper at ICLR 2026
Table 8: Models, pprrmpareettrr and corresponding dataset sizes (in tokens) for Celerity.
Model
TPP
D
300M
20
5.4B
300M
80
21.7B
300M
234
63.4B
500M
20
10.1B
500M
80
40.2B
500M
234
117.8B
900M
20
18.1B
900M
SO
72.5B
900M
234
212.3B
1.8B
20
36.2B
1.8B
80
144,8B
1.8B
234
424.0B
3.9B
20
77.6B
3.9B
80
310.4B
3.9B
234
909.2B
reports using "up to" 9T tokens (Llama-3.2 Overview) while Llama-3.1 and Qwen-3 sizes ported "approximately." Moreover, it is unclear from the papers whether the full pre-training datasets (i.e., used train the flagship models) were also used for the smaller models. For approximate sense of how these models compare, we include Llama-3 and Qwen models in Table 10 based on the assumption the models do the full (approximatell--epporeel) datasets. Note that if these sumptions hold, the smaller Llama-3 and Qwen-3 models would be rather compute-inefficien e.g., they would all be beyond the plotting range of Fig. 2 (i.e., 1024 FLOPs) and score well below the Celerity extrapolation (and Gemma results) in Fig. 16.
Semma3-278
80
Celerity
Gemma2-27B
Cx Cssudent, train
Gemma2-99
Average Accuracy (%%
75
Ca,1/= Cat Creacher, forward
Gerhima 3-128
Cu. Cat, Creacher, traio
70
Gemmaz-28
GBBma3.4B
Celerity-3.98
65
Celerity-1.88
60
Gemma3-18
Cel Celerity-90MM
10a1
10
10a3 Compute FLOPS (C)
10
Figure 16: Celerity compute efliciency vs. distilled models: Downstream accuracy. Celerity mod- els perform similarly to distilled distillation student FLOPS. When considering teacher forward pass FLOPs, Gemma curves shift away from Pareto frontier (worse), with a further shift if we account for FLOPs to train the teacher.
models, when generously only accounting for
Compute efficiency. Fig. 16 provides further downstream results for Celerity models (and their fitted extrapolation), in comparison to larger Gemma-2 and Gemma-3 models. The plot shows how the accuracy FLOPs comparison depends whether account for teaching FLOPs (e.g.. generating logits for student training), or the initial cost of educating the teacher.
Token efficiency. Fig. 17 compares the token efficiency of Celerity other model families. Using the Celerity models trained in the fixed 234 TPP band, we fit a power law D and extrapolate token efficiency to larger scales.
29

---

Published as a conference paper at ICLR 2026
70.0
Gemmia2-28
Liama-78
Llama2-78
67.5
SmolLM2-1.78
average Accuracy 18)
Celerily-3.99
OLMo-78
65.0
Zamba2-1.20
SmoILM-1.78
OLMo2-18
62.5
Celerity-1.88
BTLM-BB
60.0
Gemma3-18
SmoiLM2-360M
57.5
OLMp-18
CerebrasGPT-138
55.0
CembraaGPT.
SmolLM 360M
Celerity-900M
Celerity Fit: 100
-4.100
52 1gu1
10t2
10
Training (D)
Figure 17: Celerity token efficiency: Downstream accuracy. Celerity models are at the Pareto frontier compared other model families.
Celerity Fit
Gemma3.275
Remmaa-277
Datution, tran
Gemma2
Datusomn tra Dassher, tram
Gemma3-120
Average Accuracy (%)
70
2emmma2-2
Gemma3-4B
Celerity-3.98
65
Celerity-1.88
Gemma3-18
Celerity-900M
1041
1022
1043
Tokens (D)
Figure 18: Celerity token efficiency vs. distilled models: Downstream accuracy. Celerity models are on par with distilled models.
Generally, larger models should more token-efficient for the same token budget. Theoretically, distillation should also offer greater token efficiency-aaa given TPP (Busbridge et al., 2025)- but by training small models to very-high TPP, the distilled models in Fig. 18 train mainly in a regime of diminishing returns, and ultimately end up without advantage over Celerity's standard next- token-prediction training.
There many interesting questions around token efficiency scale, and indeed token efficiency may become more critical as frontier models reach the limits of high-quality data (Muennighoff al., 2023).
Parameter efficiency. Celerity. Celerity models are less parameter efficient than models specifically designed for parame- ter efficiency.
Finally, Figs. 19 and 20 provide the parameter efficiency comparisons for
C.4 OPEN MODEL EVALUATION AND FLOP CALCULATION METHODS
All models are obtained from HuggingFace and evaluated using the Eleuther Eval Harness frame- work (Gao et al., 2021). The downstream tasks with number of shots are arc-challenge (25), arc-easy (0), boolq (0), hellaswag (10), piqa (0), siqa (0) and winogrande (5). These tasks are chosen as they the most commonly reported downstream benchmarks for pre-trained base models, and are ap- propriate for Celerity models the scale that we compare (i.c., tasks where small models perform above random chance).
30

---

Published as a conference paper at ICLR 2026
70.0
Gemma2-28
67.5
average Accuracy 18)
SmolLM2-1.78
OLMo-78
65.0
Ce
SmoILM-1.7B
62.5
OLMo2-1B
BTLM-38
Celerity-1188
60.0
Gemma3-18
57.5
SmoILM2-360M
CerebrasGPT-138
OLMo-1B
55.0
SmolLM-360M
CerebrasGPT-6.72
Celerity: 100
(5mirr)) -4.100
Celerity-900M
52.5
10a
10a0
Parameters (N)
Figure 19: Celerity parameter efficiency: Downstream accuracy. Celerity models are less param- eter efficient than models trained at much higher TPP, while better than prior models also aiming for compute efficiency (Cerebras GPT).
Gemma2-278
Gemma3-278
75
Go memmaa--12
Average Accuracy (%)
70
Gemma2-28 Gemma3-4B
Mm.1LM2-1.78
65
Celerity-3.99
Smoll M-1.7
Celerity: 100
53311177
Celerity-1.88
SmolLM2: 100
-0155
60
G Gemma3-16
SMoILM2-360M
SmolLM: 100
SmolLM M-30M
Gemma2: 100
0016
Celerity-900M
100
-..19
50
10a
1020
10t2
(N)
Figure 20: Celerity parameter efficiency scaling comparison: Downstream accuracy. Preliminary accuracy vs. model size power law comparison between Gemma (distillation). SmolLM (refined data), and Celerity models (standard pre-training). Distilled model families have the largest scaling exponent, suggesting distillation may scale better in parameters.
For full transparency, method for counting FLOPs across the different models families is given in Table 9, while table all the raw downstream evaluation scores are in Table 10.
C.5 VALIDATION LOSS COLLAPSE
In Fig. we show normalized training loss curves where we evaluate the validation loss of model checkpoints during training. For each training run, we evaluate checkpoints at 5% intervals, com- puting loss on the same 493M-token held-out portion of SlimPajama. Similar to training, collapse occurs except for the initial few checkpoints; we attribute the initial differences to differing LR warmup proportions (Table 2). Nevertheless, validation collapse is sufficient for deviations to pro- vide useful diagnostic of any training issues. Validation collapse, measured on fixed datasets, could be a particularly valuable diagnostic if late-stage annealing or curriculum strategies distort training curves due data differences.
31

---

Published as a conference paper at ICLR 2026
Table 9: Forward FLOPs calculation for self-attention block and Mamba-2 block. This table only lists operations that are not covered 6 * Apaaram * Titokens. which should take care of all operations that involves matmul with a weight matrix. For Zamba2, the training FLOPs can be calculated as 6 Plparams Nttoeens 2+ V D Dottn* MTokeens 3 Lattn (L* mambee + Cattn), while the rest of the models analyzed are variations of decodeer-transfommers whose training FLOPs can be estimated as 6 * peramm tokens 2* V * Dattn Rttkeens + 3 Lattn Cattn Here 3 represents I flop per forward op and flops per backward op.
Operation
FLOPs, given input:
B x S D (or Data))
Attention: QK
2BS Dottn
Self-Attentio
Attention:
softmax, scaling, mask
3BS²
Attention: V matmul
2BS Dattn
Cutta
Attention: O projection
2BSD atin
Feedforward: activation
BSD
dt softplus
3BSH
xBC conv 1d, silu
BS(ED 2
Mamba-2
sampling x,
BSED+
SSD, A prefix sum
BHS
Cmambo2
SSD, compute output for each intra-chunk
4BHSC BSEDNC
SSD, compute state for cach intra-chunk
2BHS+ BSEDN
SSD, compute inter-chunk recurrence
4BH(Z 1)² BN(Z 1)
SSD, compute output from state per chunk
BHS BSEDN BSED
"D
2BSED
silu, norm
6BSED
Params
B: batch size, S: sequence length, V: vocabulary size
D Dattn: attention hidden dim, Lattn num attention layers
D: mamba2 hidden dim, L: mamba2 layers, E: expansion factor
N: mamba2 state dim, H: mamba2 num heads, P: mamba2 head dim
C: mamba2 chunk size, Z: mamba2 num chunks, K: mamba2 conv dim
loss collapse Calerity
1.4
NV
Normalized valication loss
300M
500M
1.3
1.80
1.2
1.1
1.0
0.25
0.50
0.75
1.00
Fraction
training tokens,
Figure 21: Training loss curves also collapse in validation loss: Normalized validation loss across 4 Celerity model sizes, all trained to 80 TPP. Validation loss collected at 5% intervals of training.
D COLLAPSE ENABLES EARLY STOPPING: FURTHER DETAILS
D. PREDICTING NORMALIZED TRAINING LOSS CURVES
This section provides further details regarding the development of the functional form in Eq. (4), which we use to predict normalized TLCs and, through these, extrapolate in in-progress TLCs. Based on Sec. 3, we know that TLC shape modulated by LR schedule, TPP, and т. Prior theoretical and empirical work has mostly focused how loss proceeds schedule (Defazio et al., 2023; Tissue et al., 2024; Schaipp et al., 2025; Luo et al., 2025: Qiu et al., 2025). To incorporate these factors into single functional form, take the following approach:
function of training steps and LR
32

---

Published as a conference paper at ICLR 2026
Table 10: Evaluations, params, tokens, and FLOPs for all models evaluated.
Name
Params
Tokens
FLOPs
aro-c (25)
arc-e (0)
buoly (0)
Downstream Aceuracy (Num Shots)
bellacwag
plua
siqg.
winugrande
Avg.
(10)
(0)
(0)
(5)
BTLM-16-68 St- base (Dey et 20236)
2,60E+099
6.27E+11
1.53B+22
40.70
66.79
69.72
70.92
77.20
43.50
65.90
62.10
Cerehns-GPPT-I.3 (Dey et al., 2023a)
1.308:09
2.60
2.458+20
26.79
45.85
59.33
38.55
66.76
38.59
51.70
46.79
Cerebra--GPT-2211 GPT-2.78
2:70E:09
5. 5.408=100
1.048+21
29.52
52.57
59.24
49.74
70.78
40.25
54.85
50.99
Corehma-GiP--
6.70F4499
1.335+11
6. 11684-1
36.01
57.91
62.81
59.45
73.99
41.50
59.98
55.95
Gemma 2-26 (Jeam ol., 2024)
2,61E+099
2,00E+12
4.251+222
53.41
80.22
73.58
74.62
79.11
51.23
71.51
69.10
Gemma
9,24E:09
8.00E+12
5.736+23
68.34
82.76
55.08
Gemma-2-27b
2.72E:10
13
2.44
69.62
88.30
94,85
97.00
54.55
83:03
Gemma-2:22 2-2beforward
2.612:099
2.008+122
2.
53.41
80.22
73.58
74.62
79.11
51.23
71.51
69.10
(emmo-222b44are 20om0-222hffowwueeeeeerr
2.614409
3.31
$3.41
80.22
73.58
74.62
79.11
51.23
71.51
69 10
Semma-2-9h4fannee
9.24E409
8.00
24
08.34
87.88
94.22
82.76
82.97
55.48
80.35
77
Gemma (oauma-2-99ree 2-br-------foowartteeeeee
9,24E:09
1.10
3.84E+24
68.34
87.88
84.22
82.76
82.97
55.48
80.35
77.43
Gemma-3- 1-t.-----TTeemmmmeme al., 2025)
L00E+099
12
3-4865233
42.78
Gemma-3-4a gt
4.90
+.002:12
2.
58.28
81.69
78.96
77.78
79.87
49.13
Gemma-3 pt
1,22E+10
1.341+24
67.69
87.75
85.41
84.12
81.88
52.15
80.03
76.98
Gemma-3-27bgge
2.74E: 10
+1
3.3
70.51
88.17
87.25
86.14
83.95
53.99
82.95
78.97
Gemma-3 1omma--3-11 3-1b-ge-forward
L000
2. 10
24
39.16
71.93
66.67
62.98
74.54
42.78
62.19
60.04
Gemma-3- w-w.---g----eeeaadeeeeeee
L.00F4099
60F
4 44004-44
59.16
71.95
66.67
62.99
74.54
42.78
62.19
60.04
(emmma-2-dh-g--fewaaa
4.00F:09
400
1.308-24
58.28
78.96
77.78
Gemman- hatmmng---------fwwaaattttteeee
4.00E:09
80E++3
4.636-24
58.28
78.96
77.78
72.22
Gemma-3-1 2-1-ht----n--tt
1.20F+10
1.20F+13
2.4681-74
67.49
87.75
85,41
84.17
81.88
52.15
80.03
76-98
Gemma:3 3- 12b pl-fowauc-teeccerr
1,20E+10
2,60E+13
5.808++2
67.49
87.75
85.41
84.12
81.88
52.15
80.03
76.95
Liana 70.777//ooo (Touvron et al., 2023a)
7,00E+09
L000++1
2.82
50.77
72,90
75.05
77.84
79.00
45.91
71.11
67.51
1,30
1. 1.00
8 8.909422
55.55
74.54
77.98
81.18
80.36
46.62
76.95
70.45
Lhama 2: hf (Touvzon al., 2023b)
7,00E+09
2.00E+12
1.03E-23
$2.65
74.54
77.71
78.96
79.11
16.11
74.19
69.01
Lhama-2- 13b-h/
1,30+100
2.00E+12
1.886-23
59.47
80.56
82.23
80.52
47.34
LJana-3-8B (Dubey et ol., 2024)
8.00E:09
1.50
23
77.61
80,95
47.08
77.51
Lhama-3 1-8B
8.00E:09
50E+13
3.858-24
57.85
81.19
87.05
81.91
81.01
77.19
72.60
1lama-1.2.1 12-18
1,23E409
9. 9.00++12
23
39.59
60.61
63399
65.51
74.27
42.99
62.27
45.45
Liama-3.2 3,2-30
3.21E409
9.00E+12
1.405+21
50.58
71.84
72.75
76.42
77.37
47.39
71,82
66.90
DLMM-1B--M MM-II-aaaM at. 2024)
L,00F409
2. 2000+12
1. 2656422
34.47
57.28
61.74
63.81
75.14
42.12
60.46
56.43
OLMo-78-hf
7,00E409
2.46
1.268:-23
45.14
72.45
77.13
44.52
70.96
0LMo-2-0025-1 OLMo-2-1124.7B
(OLMo 2021)
L00E+09
4,00E+12
3.041-22
68.71
13.76
7,00E:09
4.00
2.038423
64.51
82.87
90.00
91.93
81.01
51.33
77.03
2112-2-1124-138
1.30E+10
5.00E+12
4.636+23
66.13
81.31
73.91
84.99
82.15
52.05
83.03
34.30
DI 01-e--20025-522
3,20F410
6. 6004412
69.45
85.94
82.81
87.33
82.97
54.25
83,90
78.00
Qwend- O.6B-Base YYang et al., 2025)
6,00E+088
3. 3.60E+13
5.318+23
44.80
58.00
69.82
$3.46
69.80
43.30
60.46
57.09
Qwen3-1.70 1.70-BBBe
1,70E+09
3. 3.60++13
1.18--21
55.20
68.60
79,24
67.19
75.52
18.62
65.27
65.66
Qwen3-48-BB
4.00E409
60
196-24
87.91
75.64
Qwen3-98
8.00E:09
3.60
3.90E+24
67.24
79.88
93.09
79.55
54.76
Qwenš-140 14B 4BB-Buse
1,40E+10
3.60E+13
6.091+22
69.97
81.86
86.76
82.69
82.10
$5.89
79:48
76.96
Qwen2 5-0.55 (Yang et al., 2014)
4.00F4088
1. 80F+13
2.04.223
35.24
$8.55 54
61.47
$1.89
69.80
44.17
$6.50
$3.95
Qwen2.5-1..
LS00400
80E+13
1.386-24
$1.86
72.10
72.118
67.86
75,90
19.08
65.27
65.36
Qwen2.5-33 5-38
3,00E409
806+13
8.5113233
56.31
73.02
77.43
74.54
78.67
49.80
71.67
68.78
Qwen2.5-7B
7,00E+09
1800+11
3.62E-24
63.65
77.48
84.66
80.19
79.82
54.61
76.40
73.83
Qwen2.5: 5-14
1,40++10
L80E+13
5.58-221
67.58
55.48
Q.e52.5.328
3.20E:10
1808413
2.298+25
70.65
77.99
87.49
95.16
82.45
56.29
82.08
M-355 (Allal et al., 2024)
L3SE+08
6.00++11
1.518-21
32,00
56.14
60.09
12.92
68.01
39.56
$2.25
50.14
Smoll M-560M
3.60E408
6. 600++1
3. 1666-1
38.65
63.59
$5.05
54.24
71.44
40.99
$7.
54.44
Smoll M-1.7B
L7704409
100
2.54P4
49.40
73.57
66.15
67.33
75.95
43.35
61.72
62.50
SmolLMM 135M (Allal al., 2025)
L35E+08
2,00E+12
2.48E+21
33:02
58.38
60.06
43.64
68.12
39.25
53.12
50.80
Smoll M2-360M
3.60F:08
4.00
1.2084222
40.78
57.46
40.89
58.41
Smoll M2-1.7B
1.70E409
1.10
1.21B+23
53,50
73.27
72.32
73.16
27.55
44.52
68.35
SmoolLMB333 Base
3,00E+09
1.12E+13
8.56B+23
59.81
76.85
80.09
77.18
79.11
16.78
73:40
70.52
Zamba2-1.28 (Glorisso et al. 2024)
4,20
3 3.00
3.868:23
53.92
66.71
70.18
12.21
77.20
46.42
68.98
65 09
Zambu2-2.70
2,70E409
3 3.00E+12
4.7715-23
60.67
73.82
78.07
77.72
79.69
45,50
76.01
70.18
Zambu2-78
7,40F:499
2.00F+12
7.688+23
68:34
80.39
83,70
83.53
80.69
49.90
79.72
75.18 18
Colerity-0000 300M
2.71E:088
6.34E+10
1.47Ex200
27.82
50.65
52.75
37.57
66.21
37.77
52.25
46.43
Cakrity-S00M
5.03F::08
1.18F+11
5.155-30
34.39
61.22
45.96
40.23
52.64
Coliitt--00MM
9.06F:088
2112
1.688.21
39.68
64.5
47.92
55.02
72.03
41.97
Colerity .SB
1811+09
4,24E+11
6.54E+21
48.55
70.29
65.17
64.34
75.46
42.99
60.22
61.00
Cokrity-3.99
3.88F409
9.08F+11
2.80F+23
$4.01
75.55
66.61
72.19
77.97
44,73
65.90
65:28
Use functional form that accounts for training fraction and LR schedule
Make the parameters of this functional form depend on TPP and r
This led to Eq. (4). Our initial aim here not to develop the best possible TL.C predictor, but to obtain a simple, effective, and interpretable method for extrapolating TLCs, allowing us to test the value of this extrapolation for early stopping in hyperparameter tuning.
We conducted variety of preliminary experiments at 11 MM-scale, using the same data as in Sec. 3 (with details in Appendix B. 1). As an input to Eq. (4), the LR schedule is normalized to be at 1.0 at its peak. It's also interpreted over training fraction, so from 0.0 to 1.0. Experiments in this data only to-zooo use fits for linear decay-to- zero schedules. To get an initial sense of how the parameters in Eq. (4)
vary, we did dilei--imeesiooal grid search determine optimal parameters for each individual
33

---

Published as a conference paper at ICLR 2026
curve, measuring total macroaveraged MAE loss over all 11 TLCs. Over the course of these experiments, we found total MAE did not change substantially when we fixed m = 0.05, and we subsequently tuned and E2 small constants in order avoid boundary effects at f and (when v(() goes to zero). Prior to fitting, training curves were smoothed using a moving average filter covering 12288 sequences (equal to the largest batch size in the dataset), and we ignore error on the first 20% of each curve (around LR warmup when curves are noisy).
Optimal 7, TPP
Optimal
TPP. not
10"
TPP
TPP
10"
20
Optimal
V 80
Optimal
80
200
200
320
320
1280
1280
10-1
10
10-2
10
10-2
10-1
Figure 22: Trends in fits for training curve prediction. Optimal per-ourve fits (from per-curve grid searches) for Eq. (4): {(t) = 1/m0005 + b rn(f) and parameters. Left: Optimal b varies strongly in T (Pearson's -0.59), weakly in TPP 0.17). Right: Optimal q varies somewhat in TPP (r -0.30), while overall stronger in (r = 0.55), but 7 trends reverse at higher TPP.
Fig. 22 shows the optimal fits for b and when each curve is independently. For optimal b, found that correlation in was much stronger than correlation in TPP (Pearson's -0.59 for
т, r = 0.17 for TPP). On the other hand, while optimal seems to increase with т for TPP 20, the relationship with other TPP appears random. Furthermore, note larger TPP values do correspond to lower optimal (r = -0.30). Based on these fits, we hypothesize we could obtain reasonable predictions by fitting as a power law in T, and g as power law TPP:
beonst
Alve q
Aconst TPP$0
(24)
As noted in Sec. 5. Also, as reported in that section, we developed an alternating greedy optimization procedure fit these four parameters, exponentially reducing the cost of the grid search space.
Results. We first note that the fits improve over the iterations of our alternating grid search pro- cedure, demonstrating that optimal parameters of the power laws do depend on each other, and can reach stable fits through iterative alternating fitting.
Table 11: Predictions improve with scale: fit at 11 IM scale, evaluated at larger scales.
Evaluation scale
MAE
Number of evaluation curves
1 (fitting points)
1.37%
112
266M
0.75%
40
610M
1.07%
102
1.7B
0.66%
3.3B
0.54%
Table 11 and Table 12 are the tables discussed in the main paper, showing how fits obtained at 11 perform at other scales (Table 11), and how different fitting procedures perform the 610M-scale evaluation data (Table 12). Fitting and g with the optimum values per-curve (i.e., oracle fits) achieves MAE of 0.504%, roughly half that of the dual power law extrapolations.
34

---

Published as a conference paper at ICLR 2026
Table 12: Separate power laws for b and q work well: fit at 111M scale (112 TLCs), evaluation at
610M (102 TLCs).
Method for estimating b
Method for estimating
MAE
Global fixed optimum
Global fixed optimum
3.03%
Global fixed optimum
= PowerLaw(TPP) (TPP)
3.35%
b PowerLaw()
Global fixed optimum
2.08%
PowerLaw
LawerLaw (TPP)
1.07%
b PowerLaw (T. TPP)
q PowerLaw(r. TPP)
1.07%
D.2 EARLY STOPPING IN TUNING: FURTHER RESULTS
In this section we describe some further early stopping experiments, and present additional evalua- tion metries.
1.78.
20 TPP: sweeping À
617M
TPP: sweeping
Loss diff prodicted vs. trie best (7)
1.4
Loss diff predicted vs. true best (76)
1.4
Choose randomly
1:2
Choose current best
Choose predicted best
1.0
1.0
0.8
0.8
Choose randomly
Choese current best
0.6
0.6
Choose predicted best
0.4
0.4
0.2
0.2
0.0
0.0
0.4
0.6
1.0
0.4
0.6
1.0
Stop point training
Stop point training
Figure 23: Early stopping comparison: further setups. Companion to Fig. 9, now comparing early stopping accuracy (final loss of predicted vs. actual best) for sweeps at 1.7B (left) and 617M (right) (both 20 TPP). Current best works well very early, but is worse for most of training.
Fig. 23 evaluates early stopping strategies in batch-size sweeps a fixed À value. Fig. 23, left. uses the same data as in Fig. 7, left. While we do not advocate keeping À fixed during B sweeps in practice, this data can nevertheless serve to evaluate prediction of early winners in tuning. Both of these plots exhibit the phenomenon also observed in Fig. 9, right: choosing the current best setting after LR warmup, as was done in Falcon (Almazrouei et al., 2023), better than selecting the best during the middle of training. However, as seen in Fig. 9, left, this method is not always successful. In Fig. 23, left, choosing the extrapolated best setting outperforms choosing the current best from 40% of training, while it picks the correct winner from the beginning in Fig. 23, right.
1.78, 20 TPP: sweeping ), B 8061
3.38. 30 TPP: sweeping A, B = 2016
Loss all prods vz. true loses (56)
1.4
Loss dif all preds vs. true losses (36
1.4
Output current
Output current loss Predict final loss
1.2
Predic: final loss
1.2
1.0
1.0
0.8
0.8
0.6
0.6
0.4
0.4
0.2
0.2
0.0
0.0
0.4
0.6
1.0
0.6
1.0
Stop point training
Stop point training
Figure 24: Early stopping comparison: MAE at 1.7B, 3.3B: À sweeps. Mean absolute error of all predicted final losses, comparing taking current loss vs. extrapolating final loss.
In many cases, rather than caring purely about which setting is best, we care about the actual pro- jected final loss. This may be useful for fitting scaling laws, for helping practitioners reason about
35

---

Published as a conference paper at ICLR 2026
1.78,
TPP: sweeping
617M,
sweeping
Loss dit. all prods vs. thin looss (%)
1.4
Loss def pedds vs. true losve (%)
1.4
Output current loss
Output current lcss
1:2
Predict final
Predict final
1.0
1.0
0.8
088
0.G
0.6
0.4
0.4
0.2
0.2
0.0
0.0
0.4
0.6
0.8
1.0
0.2
0.4
0.8
1.0
Stop point training
Stop point training
Figure 25: Early stopping comparison: MAE at 1.7B, 617M: B sweeps. Mean absolute error of all predicted final losses, comparing taking current loss vs. extrapolating final loss.
the trade-offs of, for example, greater throughput from larger B vs. suffering higher final loss. We therefore evaluated the same four hyperparameter sweeps above, but now evaluating the average loss difference between the predicted final loss and the true final loss for all curves. The baseline chooses the current loss for each curve at the given training fraction, which will overestimate the final loss. Results in Figs. 24 and 25 show that in three of four cases, extrapolating the final loss using our predictive form results in much smaller average error than using the current value.
The only instance where predicting the final loss incurred significant error was the 1.7B, 20 TPP model with / = 8064. We note that the TLCs are very noisy at this high batch size across al- most all À settings and therefore to the predicted TLC. Increasing smoothing reduces the predicted error somewhat, but the primary issue is that the noise affects the TLC mainly in the first 60% of training, thus distorting even the smoothed loss from the universal trajectory. Accurate prediction in the presence of loss spikes is an acknowledged limitation of our methodology (Appendix A).
evidently challenging to align the in in-progress training runs
D.3 OPTIMAL AND SUBOPTIMAL TLCS AS TPP SCALES
Suboptimal curves: fixed A, B
Suboptimal curves: ==0.22
Optimal curves: TPP
10
101
10
1.20
1.20
1.20
Narmalized train loss
Nomalizd train loss
Normalized train loss
1.15
10P
1.15
10a
1.15
10
TPP
TPP
TPP
1.10
I.10
1.10
11P
10a
1.00
0.0
0.5
1.0
0.0
0.5
1.0
0.0
0.5 Training fraction
1.0
Training fraction
Training fraction
Figure 26: Evolution of train curve shape. Left: When TPP is scaled but 7, À and B are held con- stant, curve shape varies significantly. Middle: When is instead held constant, shape evolves more gradually. Right: When T scales with TPP according to established power laws, curves maintain their concave structure.
Given a fitted predictive form (Eq. (4)), it is natural to ask how TLC shape varies as TPP increases, under various hyperparameter (HP) scaling strategies. In this section, we consider three scenarios:
1. No adjustment to any HPs: basically standard practice under jaP until very recently.
2. Maintain constant т: i.e., following the prescription of Wang & Aitchison (2024).
3. Optimize 7: adjust for each TPP setting following the power law of Bergsma et al. (2025a).
36

---

Published as a conference paper at ICLR 2026
Results in Fig. 26 demonstrate that, with no HP adjustments, curve shape changes substantially across TPP (left). Fixing T results in more consistent shapes (middle). but only when T is scaled for TPP do curves maintain their characteristic concave shape, with noticeable drop near the end of training (right). One may view this final period as the annealing phase of training, or the phase where variance is reduced and we descend the valley into the river (Wen et al., 2024). As TPP increases, we must reduce correspondingly to prioritize exploration for the majority of training, enabling this final descent only the final phases.
37


Reviews:
[{"review_id": "FTyQuM64IQ", "reviewer": "Reviewer_ak8p", "summary": "This paper demonstrates that training loss curves (what they call TLCs) of LLM training collapse onto a universal trajectory when normalized, provided that three key factors match: the AdamW timescale τ, the tokens-per-parameter (\"TPP\") ratio, and the learning rate schedule. The authors show that this collapse phenomenon emerges naturally when optimization hyperparameters are set optimally for a given data budget. They then train Celerity, a family of LLMs trained at fixed TPP with optimal τ, demonstrating the collapse across scales, collapse residuals as early diagnostics for training issues, and a predictive model enabling early stopping in hyperparameter tuning.", "questions": "* Maybe I am missing something, or why are the Llama3 and Qwen3 models not included in any figure comparison? The models are in the table, yet their points miss in plots.\n\n* I would highly encourage the authors to include the list of benchmarks in the Fig. 2 caption, link to them, or make it clearer what is being measured. It is a good overview but I would initially suspicious because I could not immediately find what accuracy was actually measured.\n\n* More broadly, figure captions are not self-contained, which makes it hard to grasp details or what exact setup was used; instead I often needed to jump back and forth between text and figures.\n\n* The bias-variance tradeoff is somewhat glossed over, and without spending more extensive time studying the paper, I could not grasp this discussion exactly. Could you potentially clarify directly?\n\n* Figure 15 shows CompleteP outperform μP quite substantially in terms of not only offset but slope. You provide some explanations, but have you verified experimentally how the μP transfer is suboptimal in your LR setup?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 2, "contribution": 3, "strengths": "Quality: The biggest strength of the paper is that the experimental work is thorough and well-executed. Every major ablation you would need is provided, with lots of comparisons (e.g., varying batchsizes, LR sweeps, weight decay, ...). The appendix provides even more. \n\nOriginality: The work follows a continuous line of research about stable and predictable dynamics, and specifically extends the findings of Qiu et al. to LLM training. As such, the contribution is clear: to me, it provides sound evidence of a phenomena we previously believed to be true. The application to early stopping/hyperparameter testing is interesting and well motivated, and the case study with the 1.8B run recovery is cool.\n\nSignificance: Predictability of large scale LLM training runs is extremely important, e.g., being able to reason about optimality, avoid wasting computational resources, or restart", "weaknesses": "The authors are already well aware of the main limitations: the results are done with a fixed (but certainly standard) way of pretraining, i.e., single epoch, AdamW, moderate dataset sizes (\"TPP\"). The conclusions might not hold or look different for setups with much longer training, data curriculums or optimizers and schedules. Especially longer training (higher \"TPP\") would be relevant (is there a point where i breaks down?) and extrapolations beyond just training loss. That being said, the paper is already extensive enough and cannot cover too many topics at once.\n\nMinor: Celerity uses a much more refined datamix than just Slimpajama, which makes some comparison (e.g. to BTLM) somewhat hard. But that is discussed (and represents a general improvement of pretraining strategies including data+optimization).\n\nBesides, I believe the biggest weaknesses revolve around accessibility, which I would cluster into two points: readability and availability of resources. First, as mentioned above, the paper is the newest in a serious of publications around LLM training predictability and scaling, which all build upon each other. I understand the authors have an in-depth understanding of this area, but the paper is extremely dense, with lots of notation and terms that make it daunting and complex to grasp for new readers. I think this presentation can be improved. Second, and very related, the authors have not published the code or any of the models as far as I can tell. Since the exact setup is quite unique (e.g., CompleteP), releasing those artifacts would be crucial in enabling better study of the entirety of methods that are included in this line of work.", "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": "IiLQraVlmE", "reviewer": "Reviewer_J2cj", "summary": "The authors use recent work on learning curve universality in compute-optimally trained networks, combined with other observations on the timescales of weight decay to come up with a training/scaling procedure for language models that takes advantage of these effects. They show that correct selection of the weight decay timescale, via either the (coupled) weight decay parameter, learning rate, or batch size/total number of tokens can reshape the loss curve in a consistent manner. They also show that the collapse principle can be used to predict early training failures, and to develop a model that can use detailed data from small scale training runs to predict the results of large scale training runs by 10-30% of the length of the larger scale training run.", "questions": "Can you explain the details of the point raised in weaknesses? This is key to the underpinning of the model/fitting procedure proposed in Section 5; right now it is still unclear to me how the proposed model allows for optimization of the unnormalized loss at larger scale.\n\nWith regards to the learning curve predictions: can this model be used to predict the results of varied learning rate schedules, or is it mainly suited to predicting the correct overall timescales?", "limitations": "", "rating": 8, "confidence": 5, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "This paper well-integrates key recent results on scaling behavior and uses it to provide a compelling angle on how to scale large models. The experiments are generally thorough and the paper is overall well-presented. The authors provide strong evidence that the universality of learning curves, combined with properly chosen scaling procedures, can be used to train compute-optimally with less risk than with previous methods.", "weaknesses": "One thing that is not immediately clear to me is the utility of the model presented in section 5. It seems like the best procedure would be:\n\n* Search over hyperparameters at small scale (e.g. 100M).\n* Match parameters at larger scale.\n* Carry out partial training runs at large scale.\n* Use the partial training runs to select the best run and finish it.\n\nIt seems to me that even with the model, one may have to run many training runs at the largest scale in order to obtain the normalization factor, aka the finalized loss.", "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": "9vOjMuJNys", "reviewer": "Reviewer_SS9H", "summary": "The paper studies how the training of large language models can be made more predictable and efficient. The authors observe that when optimization hyperparameters are set correctly, the normalized training loss curves for models of different sizes collapse onto a single curve. This loss collapse suggests a consistent and predictable scaling behavior across model sizes.\n\nThe paper uses this finding to propose practical tools for large-scale training, e.g., using loss collapse for early diagnostics and for tuning hyperparameters more efficiently. The authors also introduce a new model family, Celerity, to demonstrate that these principles can lead to faster and more compute-efficient training.\n\nThe results are supported by experiments across multiple model sizes, and the figures clearly show the observed alignment of loss curves.", "questions": "- Have you tested the method with more recent state-of-the-art optimizers, e.g. Muon [1]?\n\n- Does the same loss collapse happen for validation loss, or only for training loss?\n\n[1] Keller Jordan, Yuchen Jin, Vlado Boza, You Jiacheng, Franz Cesista, Laker Newhouse, and Jeremy Bernstein. Muon: An optimizer for hidden layers in neural networks, 2024.", "limitations": "", "rating": 8, "confidence": 2, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "(I am not very familiar with this specific literature. My assessment is therefore based on the content and presentation of this paper itself, rather than detailed knowledge of prior work. Judging solely by the paper, it appears to be a solid piece of work.)\n\nStrengths:\n\n- The paper addresses an important and very active research area (making large-scale model training more predictable and compute-efficient through analysis of loss-curve regularities).\n\n- The authors present convincing experimental evidence that normalized training loss curves collapse across model scales when key hyperparameters are properly aligned. The results span a range of model sizes.\n\n- The idea of using loss-curve collapse as a diagnostic tool and a guide for efficient hyperparameter tuning is valuable in large-scale training settings.\n\n- The paper is generally well-written and easy to follow. The figures are informative and communicate the main ideas well.\n\n- The introduced Celerity model family shows that the proposed ideas can work well in practice.", "weaknesses": "- The paper's main ideas are based mostly on experiments, and the theoretical explanations are limited.\n\n- The experiments focus on GPT-like architectures trained with AdamW. It is not clear whether the same collapse behavior holds for other architectures or optimizers.", "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)