PaperlayerView exact revision on arXiv

Source-linked AI summary

Terminal Shrinkage Averaging Reveals a Schedule-Estimator Interaction in LLM Pretraining

Adam Ousherovitch, Yixin Wang

arXiv:2609.25482v1cs.LGcs.AI

TL;DR

LLM pretraining typically deploys the raw final iterate, coupling the learning-rate trajectory with the model estimator even though averaging and cooldown strength may be jointly relevant. The paper proposes TSA and analyzes this schedule-estimator interaction, then reports improved validation quality in depth-22 NanoChat and one qualifying time-to-GPT-2 run.

  • Problem

    Pretraining typically deploys the raw final iterate, coupling the choice of terminal trajectory with the choice of model estimator, while it remains open whether an averaged output makes a weaker cooldown preferable.

    Source: §1 · p. 1

  • Method

    TSA interpolates between the raw final iterate and a recent-checkpoint average, with a local quadratic analysis and controlled NanoChat experiments studying its interaction with terminal learning-rate activity.

    Source: §1 · pp. 1–2

  • Results

    At depth 22, the combined intervention improves validation BPB over a matched calibrated baseline in all three paired repetitions and yields one qualifying NanoChat time-to-GPT-2 run.

    Source: §5 · p. 9

  • Takeaways & Limitations

    The results suggest treating the optimization schedule and the rule for constructing the deployed model as a joint design problem rather than optimizing them independently.

    Source: §5 · p. 9

    TSA explores only a small part of possible estimators and training trajectories, using one global shrinkage coefficient and a simple learning-rate floor for terminal optimization.

    Source: §5 · p. 9

Abstract

from arXiv · show

Large language model (LLM) pretraining conventionally returns the raw final iterate. This couples two design choices: the learning-rate schedule that generates the parameter trajectory and the estimator that constructs the deployed model (e.g. the raw final iterate or a checkpoint average). A schedule that promotes optimization progress may differ from one that minimizes variation in the raw final iterate. Separating these choices creates an opportunity to maintain progress late in training while reducing variation in the returned model. To this end, we propose \emph{Terminal Shrinkage Averaging (TSA)}, which interpolates between the raw final iterate and the average of recent checkpoints to balance recent progress against terminal variation. We analyze how TSA changes the preferred terminal learning-rate schedule under a local quadratic approximation and test this interaction through a sequence of controlled NanoChat experiments. Finally, we demonstrate that the resulting gains transfer to depth-22 NanoChat, where the combined schedule and estimator improve validation quality. A qualifying time-to-GPT-2 run also finishes faster than the public baseline used in our experiments, providing preliminary evidence of benchmark acceleration.

1 INTRODUCTION

The paper separates the terminal learning-rate schedule from the rule used to construct the deployed model, proposing Terminal Shrinkage Averaging (TSA) to interpolate between the raw final iterate and recent-checkpoint averaging. Controlled experiments show that increasing terminal learning-rate activity can worsen validation BPB for the raw final iterate while improving it for TSA.

1 INTRODUCTION

  • Pretraining conventionally uses one choice to set the late-training trajectory and another to determine which function of that trajectory is deployed.A small terminal learning rate reduces raw-final-iterate variation but also limits how far parameters can still move.Source: §1 · p. 1
  • The motivation for TSA is that checkpoint averaging may become more effective at higher learning rates, raising whether averaging and cooldown strength should be chosen together.The introduction situates this question alongside prior work on weight averaging and LAWA.Source: §1 · p. 1
  • TSA interpolates between the raw final iterate and the average of the last K checkpoints, with one shrinkage coefficient controlling the averaging strength.The raw final iterate and uniform LAWA are the interpolation endpoints.Source: §1 · p. 1
  • Under a local quadratic approximation, TSA has a quadratic surrogate risk that separates mean-output displacement from stochastic variation and can have an interior optimum.The analysis also shows that an averaged output can prefer a more active terminal trajectory than the raw final iterate.Source: §1 · p. 2
  • Across five repetitions, raising the terminal learning-rate floor worsens validation BPB for the raw final iterate but improves it for TSA, with both effects growing as the floor rises.The paired design tests floors of 5%, 10%, and 15% of peak; pure AdamW reproduces the interaction's direction.Source: §1 · p. 2
  • At a matched training budget, the full schedule-estimator recipe improves validation quality over the underlying public depth-22 recipe and produces a qualifying time-to-GPT-2 run.The paper describes this as preliminary evidence that the interaction transfers to scale.Source: §1 · p. 2

2 RELATED WORK

Prior work studies averaging as a way to reduce variation or accelerate training, and connects averaging with learning-rate schedules and cooldowns. This paper focuses on whether choosing an averaged output changes which terminal trajectory should be trained.

2 RELATED WORK

  • Classical and finite-sample analyses study averaging's guarantees and variance reduction, while SWA and LAWA apply late-trajectory or recent-checkpoint averaging in deep learning.LAWA is framed as a way to accelerate training (Kaddour, 2022).Source: §2 · pp. 2–3
  • For language-model pretraining, Sanyal et al. (2024) find checkpoint averaging more effective at higher learning rates; this paper instead asks whether an averaged output changes the preferred terminal trajectory.The paper characterizes its focus as complementary to work studying averaging under a given trajectory.Source: §2 · p. 3
  • Unlike Lookahead, which repeatedly interpolates fast and slow states during optimization, TSA performs one terminal interpolation without changing optimizer updates.The terminal interpolation combines the raw final iterate with a finite late-checkpoint average.Source: §2 · p. 3
  • Prior schedule-and-averaging work includes replacing cooldowns with SWA and merging checkpoints to emulate decay, whereas TSA studies their interpolation and interaction.TSA leaves live optimizer updates unchanged and separately intervenes on terminal learning-rate activity.Source: §2 · p. 3
  • TSA leaves the optimizer update rule unchanged, making it complementary to optimizer design; the primary experiments use NanoChat's Muon+AdamW optimizer.The paper also examines how averaging gains are distributed across parameter groups.Source: §2 · p. 3
  • The depth-22 comparisons start from the exact PR #830 recipe available during the experiments and use it as the matched-endpoint scale-transfer baseline.NanoChat provides the time-to-GPT-2 benchmark harness, and qualification uses the 22-task DCLM CORE metric.Source: §2 · p. 3

3 METHOD

The method first defines TSA as an interpolation between the raw final iterate and a recent-checkpoint average, then analyzes its risk for a fixed trajectory and how the estimator can shift the preferred terminal schedule. The analysis identifies a bias-variance tradeoff and a condition under which averaging favors more terminal activity.

3.1 TERMINAL SHRINKAGE AVERAGING

  • TSA returns a convex interpolation between the raw final iterate and the checkpoint average; α=0 gives the raw iterate and α=1 gives uniform LAWA.The checkpoint average uses K saved checkpoints separated by s optimizer steps, with the final checkpoint receiving extra weight under shrinkage.Source: §3.1 · p. 4
  • Under a positive-semidefinite local quadratic model for the short terminal window, each checkpoint is decomposed into a mean and a stochastic residual to analyze TSA risk.The approximation is intended to describe the short late-training window used by TSA, not the full optimization trajectory.Source: §3.1 · p. 4
  • The risk decomposition separates the mean displacement from stochastic variation; averaging can reduce stochastic error but incur lag when optimization is still making progress.An interior optimum occurs when the stochastic error removed outweighs the cost of moving toward older checkpoints, so full averaging need not be optimal.Source: §3.1 · pp. 4–5
  • TSA's practical aim is to retain enough weight on the raw iterate to limit checkpoint lag while averaging enough to reduce terminal variation.The fixed-trajectory result is followed by analysis of whether the tradeoff changes which trajectory should be trained.Source: §3.1 · p. 5

3.2 A SUFFICIENT CONDITION FOR A MORE ACTIVE TERMINAL SCHEDULE

  • If averaging's stochastic benefit grows with terminal activity faster than its additional deterministic cost, the risk-minimizing averaged trajectory has greater terminal activity.Here, terminal activity is parameterized by ρ and corresponds experimentally to the terminal learning-rate floor.Source: §3.2 · pp. 5–6
  • In a one-dimensional noisy terminal-update example, every α>0 strictly increases the preferred terminal step size, and full averaging doubles it.The result applies when the risk-minimizing step sizes are interior.Source: §3.2 · p. 6
  • The condition does not claim that increased terminal activity is always beneficial under averaging or predict the size of the optimal shift.It identifies when changing the returned estimator changes the preferred terminal schedule.Source: §3.2 · p. 6

4 EXPERIMENTS

The experiments use controlled NanoChat interventions to separate terminal schedule effects from output-estimator effects, reporting validation BPB and, at depth 22, DCLM CORE. They test partial averaging, terminal activity, their interaction, optimizer robustness, and transfer to a larger model.

4 EXPERIMENTS

  • The experiments use depth-12 NanoChat for controlled ablations and depth-22 for a time-to-GPT-2 experiment, pairing depth-12 comparisons on initialization, optimizer, batches, and validation data.The reported metrics are validation BPB, where lower is better, and DCLM CORE for the depth-22 run.Source: §4 · p. 6

4.1 PARTIAL SHRINKAGE HAS A REPLICATED INTERIOR OPTIMUM

  • At the 5% terminal floor, α = 0.40 is the best tested point for mean paired BPB improvement over the raw final iterate across five trajectories; uniform LAWA worsens BPB by 0.000734 ± 0.000172.Source: Figure 2 · p. 7, §4.1 · p. 7

4.2 MORE TERMINAL ACTIVITY WORSENS THE RAW FINAL ITERATE

  • For the raw final iterate, raising the floor from 5% to 10% worsens BPB by 0.000915 ± 0.000836, and raising it to 15% worsens BPB by 0.003757 ± 0.000769.The descriptive sweep shows monotone, increasingly steep degradation over the tested range.Source: §4.2 · p. 7, Figure 3 · p. 7

4.3 THE OUTPUT ESTIMATOR CHANGES THE EFFECT OF TERMINAL ACTIVITY

  • With TSA fixed at α = 0.55, the 15% floor worsens the raw final iterate but improves TSA, and the paired interaction is positive for every stream at both tested floors.The 10% intervention also yields a clear interaction, although its TSA simple effect is not individually resolved.Source: §4.3 · p. 8, Table 1 · p. 8, Figure 4 · p. 8

4.4 THE INTERACTION PERSISTS UNDER PURE ADAMW

  • Under pure AdamW, the interaction is 0.000782 ± 0.000089 BPB at 10% and 0.002536 ± 0.000204 BPB at 15%, positive in all five repetitions.It is smaller than under the native optimizer but persists without Muon updates.Source: §4.4 · p. 8

4.5 SCALE TRANSFER TO THE NANOCHAT TIME-TO-GPT-2 BENCHMARK

  • At depth 22, TSA with a 15% terminal floor improves validation BPB from 0.720235 ± 0.000735 to 0.719580 ± 0.000562 versus the calibrated baseline.The paired improvement is 0.000654 ± 0.000184 BPB, favorable in all three seeds.Source: §4.5 · p. 9, Table 2 · p. 9

5 CONCLUSION

The paper argues that partial averaging can balance optimization progress against terminal variation, while treating the schedule and returned estimator as joint design choices. The appendix details the depth-12 setup and its development–confirmation process.

5 CONCLUSION

  • Partial, rather than uniform, averaging can balance optimization progress against stochastic variation under the local quadratic analysis.The conclusion frames TSA as a way to keep optimization more active near training’s end while constructing the returned model from the terminal trajectory.Source: §5 · p. 9
  • The paper limits its evidence to NanoChat experiments and a local quadratic explanation; cross-architecture, optimizer, dataset, and larger-scale transfer remains future work.Training-time differences mostly reflect allocation throughput and are not interpreted as algorithmic speed differences.Source: §5 · p. 9

A EXPERIMENTAL DETAILS

  • Appendix A specifies the trajectories, estimator selection, evaluation protocol, and scale-transfer comparison underlying the experiments.Source: A EXPERIMENTAL DETAILS · p. 13

A.1 DEPTH-12 MODEL AND COMMON TRAINING SETUP

  • Depth-12 NanoChat uses 12 transformer layers, 286,261,730 trainable parameters, and 3,000 optimizer steps; TSA uses eight checkpoints spaced 32 steps apart unless stated otherwise.The native optimizer assigns matrix parameters to Muon and remaining groups to AdamW; a pure-AdamW control is used for within-optimizer estimator comparisons.Source: §A.1 · p. 13

A.2 DEVELOPMENT AND CONFIRMATION SPLIT

  • Estimator rules are developed on five streams and confirmed on five new streams, with paired optimizer and floor arms sharing initialization and data order within each stream.The native development rule selects α = 0.55 and is used unchanged across both optimizer regimes in the main text.Source: §A.2 · p. 13

A.3 EVALUATION PROTOCOL AND UNCERTAINTY

The appendix reports robustness checks for shrinkage, checkpoint windows, and alternative estimators, alongside paired-uncertainty and depth-22 protocol details. Across the tested settings, partial shrinkage remains useful, while adaptive tensorwise coefficients are presented as a heuristic.

A.3 EVALUATION PROTOCOL AND UNCERTAINTY

  • Replicated intervals are 95% Student-t intervals across five independently ordered streams, and schedule–estimator interactions are computed within stream before averaging.The paired schedule effects are strongly correlated across streams (r = 0.995 for 5% → 10% and r = 0.994 for 5% → 15% under the native optimizer).Source: §A.3 · p. 14

A.4 DEPTH-22 SCALE-TRANSFER PROTOCOL

  • The depth-22 comparison uses three configurations at the same ratio-9.0 endpoint, with the treatment combining a 15% terminal floor and TSA.The checkpoint window is held fixed between the direct-9.0 control and treatment; training time is used to establish matched compute, not an algorithmic speed difference.Source: §A.4 · p. 14

B ADDITIONAL DEPTH-12 RESULTS

  • The appendix examines coefficient response, checkpoint selection, alternative estimators, and parameter-specific rules beyond the fixed scalar TSA used in the main text.Source: B ADDITIONAL DEPTH-12 RESULTS · p. 15

B.1 REPLICATED SHRINKAGE RESPONSES ACROSS FLOORS AND OPTIMIZERS

  • The best tested α shifts upward as the terminal floor rises: native optimizer 0.40, 0.55, 0.70 and pure AdamW 0.30, 0.50, 0.65 at 5%, 10%, and 15%.The fixed native-optimizer α = 0.55 retains approximately 90%, 100%, and 96% of the best observed scalar-TSA gain at those respective floors.Source: Figure 5 · p. 15, §B.1 · p. 15

B.2 SENSITIVITY TO CHECKPOINT COUNT AND SPACING

  • 0.002070 ± 0.000205 BPB for K = 8, s = 32 and 0.002363 ± 0.000207 BPB for K = 16, s = 32 show TSA gains across all nine tested windows.Uniform LAWA is more sensitive to window choice; at K = 16, s = 64 it is worse than the raw final iterate by 0.005222 ± 0.000626 BPB.Source: Figure 6 · p. 16, §B.2 · p. 16

B.3 ALTERNATIVE OUTPUT ESTIMATORS

  • 0.002070 ± 0.000205 BPB is the TSA gain, versus 0.001910 ± 0.000209 BPB for the strongest matched EWA; TSA’s paired advantage is 0.000159 ± 0.000101.Checkpoint EMA gains 0.001113 ± 0.000250 BPB, while the long SWA-style average is worse than the raw final iterate by 0.000775 ± 0.000375 BPB.Source: Figure 7 · p. 17, §B.3 · p. 17

B.4 EXPLORATORY TENSORWISE ADAPTIVE TSA

  • Adaptive TSA gains 0.002002 ± 0.000170 BPB across five trajectories; applying it only to Muon-managed tensors gains 0.001951 ± 0.000140 BPB.The rule is a trajectory-statistic heuristic, not an estimator of the exact optimum in Theorem 1.Source: §B.4 · p. 17

B.5 INITIAL OPTIMIZER-ALIGNED STRUCTURED TSA AND PARAMETER LOCALIZATION

The optimizer-aligned two-group TSA reveals different shrinkage gains across parameter groups, but adds little over scalar TSA on native-optimizer confirmation runs. The group asymmetry also persists when all tensors are trained with AdamW.

B.5 INITIAL OPTIMIZER-ALIGNED STRUCTURED TSA AND PARAMETER LOCALIZATION

  • Development selects the same off-diagonal pair, (αA, αB) = (0.625, 0.250), under both optimizers for the predeclared matrix/complement split.Group A is the 29.67% of parameters assigned to Muon in the native optimizer; the identical coordinate partition is retained in pure-AdamW runs.Source: Figure 8 · p. 18, §B.5 · p. 18
  • Structured TSA adds only 0.000040 ± 0.000050 BPB beyond the best frozen scalar rule at the native 10% decision gate, with an interval that includes zero.Additional gains are similarly small at 5% and 15%; under pure AdamW, the increment is positive but small at 5% and 10% and disappears at 15%.Source: Table 4 · p. 18, §B.5 · p. 18
  • Shrinking only the 29.67% matrix group recovers most of the full TSA gain, while shrinking the complement yields a substantially smaller improvement.The asymmetry persists under pure AdamW, so the partition is a proxy for parameter-role or trajectory heterogeneity rather than a causal Muon-versus-AdamW explanation.Source: §B.5 · pp. 18–19, Figure 9 · p. 19, Table 5 · p. 19

B.6 EXTENDING STRUCTURED TSA TO PARAMETER-ROLE GROUPINGS

The section extends groupwise TSA from a hidden-matrix/rest split to parameter-role partitions, then tests frozen rules on fresh trajectories. The role-based embedding/hidden/unembedding/rest rule adds confirmed gains over scalar TSA in both optimizer regimes, while hidden matrices retain strong shrinkage.

B.6 EXTENDING STRUCTURED TSA TO PARAMETER-ROLE GROUPINGS

  • The hidden-matrix/rest split localized much of the averaging value, but its two-parameter estimator improved only modestly over scalar TSA; the broader role split added out-of-sample gain.The selected grouping is not established as universally optimal; alternative partitions had smaller confirmation signals.Source: §B.6 · pp. 19, 22
  • The confirmation froze selected rules after evaluating 290 predeclared rules on eight development trajectories, then tested them on fresh native and paired pure-AdamW trajectories.The confirmation used an untouched 96-batch holdout; pure-AdamW rules transferred without retuning.Source: §B.6 · p. 20
  • The role-based rule separates embeddings, hidden transformation matrices, the unembedding matrix, and remaining parameters, refining the earlier hidden-matrix/rest split.Most additional gain came from splitting the former complement, especially separating input embeddings from the unembedding matrix and remaining parameters.Source: §B.6 · pp. 19–20
  • 0.000153 ± 0.000051 BPB: the frozen role-based rule improved over scalar α = 0.50 on fresh native trajectories, and the role-based rule also gained in pure AdamW.The native result’s 95% interval was [0.000103, 0.000204]; the transferred rule was not retuned.Source: §B.6 · p. 21, Figure 11 · p. 22

B.7 GROUPWISE TERMINAL ACTIVITY AS A MECHANISM PROBE

This mechanism probe maps a frozen groupwise estimator to heterogeneous terminal learning-rate floors and compares it with fixed uniform schedules. The mapped allocation worsened the raw final iterate but its penalty was essentially removed by groupwise TSA; the experiment does not establish a better final recipe.

B.7 GROUPWISE TERMINAL ACTIVITY AS A MECHANISM PROBE

  • The frozen estimator’s coefficients were mapped to floors of 5.0% for embeddings, 15.0% for hidden matrices, 11.92% for the unembedding matrix, and 8.85% for remaining parameters.Eight new paired native trajectories compared these with uniform 5%, uniform 10%, and a parameter-count-weighted uniform floor of 8.576%.Source: §B.7 · pp. 22–23
  • The mapped allocation made the raw final iterate worse by 0.003404 BPB, while groupwise TSA removed essentially all of that relative penalty.The tested heterogeneous schedule’s difference from the parameter-count-matched uniform schedule was small and unresolved under the frozen groupwise estimator.Source: §B.7 · p. 23, Table 7 · p. 23
  • Within the mapped arm, groupwise TSA improved over the raw final iterate by 0.005114 ± 0.000174 BPB and over scalar TSA at α = 0.55 by 0.000322 ± 0.000061 BPB.The complete mapped-schedule plus groupwise-TSA output improved over the paired uniform-5% raw final iterate by 0.000991 ± 0.000375 BPB.Source: §B.7 · p. 23
  • The parameter-count-matched control equalizes mean floor, not optimizer-update energy, displacement, or noise injection, so the result concerns estimator value under different distributions of terminal activity.It does not establish that this particular mapped cooldown should replace a uniform schedule.Source: §B.7 · p. 23

B.8 FULL PRELIMINARY SINGLE-TRAJECTORY SWEEPS

Preliminary single-trajectory sweeps show an interior scalar-shrinkage optimum and worsening raw-final-iterate BPB as the terminal floor rises. Replicated fits across optimizers and floors show near-quadratic shrinkage responses, with preferred shrinkage increasing with terminal activity.

B.8 FULL PRELIMINARY SINGLE-TRAJECTORY SWEEPS

  • On one fixed trajectory, partial shrinkage produced a broad interior optimum, while full LAWA eventually became worse than the raw final iterate.This was preliminary evidence from one trajectory; uncertainty reflects evaluation batches rather than independent training runs.Source: §B.8 · p. 23, Table 8 · p. 24
  • The raw final iterate worsened monotonically across the tested terminal-floor range, with steeper degradation at larger floors.The sweep varied only the terminal floor and used one training trajectory; the figure’s quadratic curve is descriptive, not implied by the TSA theory.Source: §B.8 · p. 23, Figure 12 · p. 24, Table 9 · p. 25

B.9 REPLICATED QUADRATIC RESPONSE

  • The replicated fit modeled mean paired BPB improvement over the raw final iterate across five confirmation trajectories as a quadratic in shrinkage coefficient α.Because the response is improvement rather than risk, a beneficial interior optimum corresponds to a concave curve; the fitted vertex is bα⋆ = −b/(2c).Source: §B.9 · p. 24
  • The replicated quadratic description had R2 ≥ 0.99950 across all six optimizer-floor conditions, and fitted optima increased with terminal floor under both optimizers.Under Muon+AdamW the fitted optimum moved from 0.419 to 0.547 to 0.701 across 5%, 10%, and 15%; under pure AdamW it moved from 0.330 to 0.499 to 0.668.Source: §B.9 · p. 24, Table 10 · p. 25, Figure 13 · p. 26

C PROOF OF THE TSA RISK THEOREM

The proof proceeds from the TSA risk expression by expanding its deterministic term, identifying an unconstrained minimizer, and projecting that minimizer onto the allowed shrinkage interval.

C PROOF OF THE TSA RISK THEOREM

D ONE-STEP QUADRATIC ANALYSIS

Under a one-step quadratic model, TSA changes the preferred terminal learning rate without changing the jointly optimized minimum risk; this equivalence has a limited scope.

D ONE-STEP QUADRATIC ANALYSIS

  • When neither optimum is clipped by ηmax, the preferred terminal learning rate for TSA is twice that for the raw final iterate.The analysis therefore identifies a schedule preference, not an automatic improvement in the best attainable risk.Source: D ONE-STEP QUADRATIC ANALYSIS · p. 28
  • The minimum risk remains constant across α, but multi-step trajectories, longer checkpoint windows, and nonstationary late-training dynamics can break the rescaling equivalence.The result isolates a change in preferred schedule rather than claiming an automatic improvement in minimum risk.Source: D ONE-STEP QUADRATIC ANALYSIS · p. 28

E LLM DISCLOSURE

  • LLMs helped implement and review experiments, critique drafts, and identify related papers that the authors read before citing.Source: E LLM DISCLOSURE · p. 29
Loading 2609.25482v1…