Terminal Shrinkage Averaging Reveals a Schedule-Estimator Interaction in LLM Pretraining
cs.LG · cs.AI
Quick verdict
TSA: shrinkage blend of final iterate and checkpoint average. NanoChat experiments support schedule-estimator interaction; effects small, scale preliminary.
Summary
LLM pretraining conventionally couples two decisions: the terminal learning-rate (cooldown) schedule and the estimator that constructs the deployed model. This paper proposes Terminal Shrinkage Averaging (TSA), a one-parameter interpolation between the raw final iterate and the uniform average (LAWA) of the last K checkpoints. Under a local quadratic loss approximation, the excess risk of TSA is an exact convex quadratic in the shrinkage coefficient, formalizing a lag-versus-noise-reduction tradeoff, and a one-step corollary predicts that an averaged output prefers a more active terminal schedule. Controlled depth-12 NanoChat experiments (286M parameters; five paired data-order streams per condition) cross terminal LR floors of 5/10/15% of peak with the two estimators. Raising the floor worsens validation BPB for the raw iterate (+0.003757±0.000769 at 15%) but improves it under TSA (−0.000844±0.000737); the paired interaction is positive in every stream and replicates under pure AdamW, and the best α shifts from 0.40 to 0.70 as the floor rises. At depth-22, the combined recipe improves validation BPB in all three seeds (0.719580 vs 0.720235) and yields one DCLM CORE-qualifying time-to-GPT-2 run (only 1/3 seeds surpass the 0.256525 threshold). Limitations: effect sizes are tiny (~0.0007–0.005 BPB), evidence is confined to NanoChat, there is no empirical comparison against Schedule-Free, WSM, or PACE, and the 'finishes faster' timing claim is not horizon-matched to the public baseline.
Review complete
Assessment at a glance
5
risk flags identified
● high 0
● medium 3
● low 2
5 strengths · 5 weaknesses · 5 author questions
AI text signal = Binoculars zero-shot detector score over 20 sampled passages (Fast-DetectGPT co-reported) — absence does not prove human authorship, and a flag does not imply misconduct.
Levels reflect verification complexity — not a quality score. Evolve never rates your paper.
Key claims under review
- Under a local quadratic loss, the risk of TSA is an exact convex quadratic in the shrinkage coefficient, with an interior optimum whenever noise removal outweighs checkpoint lag (Theorem 1).
- The preferred terminal learning-rate activity is higher for an averaged output than for the raw final iterate (§3.2; one-step corollary η*α = η*raw/(1−α/2)).
- In depth-12 NanoChat, raising the terminal LR floor from 5% to 15% worsens the raw final iterate (+0.003757±0.000769 BPB) but improves TSA (−0.000844±0.000737), with the paired interaction positive in all five streams and replicating under pure AdamW.
- At depth 22, a 15% floor plus TSA improves validation BPB in all three seeds (0.719580 vs 0.720235) and produces one DCLM CORE-qualifying time-to-GPT-2 run, which the authors themselves treat as preliminary.
Strengths
Causal isolation of the interaction
A predeclared factorial crosses terminal floor (5/10/15%) with estimator over five paired streams (§4.3, Table 1); within-stream pairing (schedule-effect correlation r=0.995, A.3) shrinks interaction CIs to ±0.00009 — roughly 10× tighter than the simple effects — and the interaction is positive in every stream.
Development/confirmation hygiene
Scalar α=0.55 was frozen on separate development streams before thirty confirmation trajectories were evaluated on an untouched 96-batch holdout (A.2, A.3); post-freeze α grids are explicitly labeled descriptive and not used for selection.
Optimizer robustness
Theory validated at its stated scope
Ablation depth and honest scale reporting
Weaknesses
Tiny effect sizes, unclear practical significance
Thin depth-22 evidence
Missing head-to-head baselines
Gap between theoretical condition and measurement
No code release stated
Full Review Report
$1.99 · one-timeUnlock the complete preliminary analysis Evolve produced for this paper.
Report preview
Scope & Claims
The paper separates two pretraining decisions usually conflated: the terminal learning-rate schedule (the cooldown) and the output estimator mapping the parameter trajectory to the deployed model. It introduces Terminal Shrinkage Averaging (TSA, Eq. 2): θ̂ = (1−α)θ_T + α·θ̄ over K=8 checkpoints spaced s=32 steps. Three claims are made: (1) under a local quadratic loss, TSA risk is an exact convex quadratic in α with an interior optimum (Theorem 1); (2) an averaged output prefers a more active terminal schedule (§3.2; one-step corollary η*α = η*raw/(1−α/2)); (3) controlled depth-12 experiments confirm this schedule–estimator interaction and it transfers to depth 22 (§4.5). Extraction confidence is high: the text includes appendices A–E with numeric tables.
The remaining sections cover Methodology Assessment, Statistical Evidence, Results & Reproducibility, Limitations & Risks, and Verdict & Recommendations.
- All 5 strengths & 5 weaknesses, analyzed in detail
- 5 risk flags with severity, section & reviewer questions
- Methods & statistical evidence assessment
- 5 questions a peer reviewer would ask
- Full six-dimension evaluation report & recommendations
One-time purchase · this paper only
One-time purchase, this paper only. AI-generated preliminary analysis — not peer review.
Frequently asked questions
What is Terminal Shrinkage Averaging (TSA)?
TSA is a rule for constructing the final model at the end of LLM pretraining: it returns (1−α) times the raw final weights plus α times the uniform average of the last K checkpoints. α=0 gives the standard final model; α=1 gives LAWA-style full checkpoint averaging. The paper finds α around 0.4–0.7 works best, depending on how active the terminal learning-rate schedule is.
Can weight averaging replace the learning-rate cooldown in LLM training?
Partially, according to this paper: with TSA averaging (α=0.55), holding the terminal learning rate at 15% of peak instead of decaying to 5% improved validation BPB in depth-12 NanoChat, whereas the raw final iterate got worse. Averaging does not remove the need for a schedule — it changes which schedule is best.
Did this paper set a NanoChat speedrun record?
No official record. The 15%-floor + TSA configuration reached the DCLM CORE qualification threshold (0.256525) in only one of three seeds (mean 0.2579±0.0159), and its 77.9-minute training time covers 10,172 steps versus the public 81.8-minute baseline's 10,624 steps, so the times are not horizon-matched. The authors call the result preliminary.
Does the method depend on the Muon optimizer?
No. The schedule–estimator interaction replicates under pure AdamW at roughly half the magnitude (interaction +0.002536±0.000204 BPB at the 15% floor vs +0.004601±0.000091 under Muon+AdamW; five paired streams each, positive in all repetitions).
About this review
This is an AI-generated preliminary review by Evolve.Pub, produced from the openly available text of the paper. It is not peer review and does not represent the authors' views. The canonical record of the paper is at the publisher link above.
- Figures are available only as captions and axis labels in the extraction; point-level values were taken from embedded tables and caption text, not visually verified.
- Math notation is partially garbled by extraction (hats, summation signs); formulas were reconstructed from context and the appendix proofs.
- Review covers arXiv v1 (2609.25482, 21 Sep 2026); no external repository or supplementary code was inspected.
Go deeper on this paper
Wrote a paper? Stress-test it before reviewers do.
Upload your manuscript for the same structured scrutiny — reasoning, statistics, evidence, figures, and reproducibility.