终端收缩平均揭示LLM预训练中的调度-估计器交互

# Terminal Shrinkage Averaging Reveals a Schedule-Estimator Interaction in LLM Pretraining
Source: [https://arxiv.org/html/2609.25482](https://arxiv.org/html/2609.25482)
###### Abstract

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*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\.

## 1Introduction

Large language model \(LLM\) pretraining typically follows a parameter trajectory under a decaying learning\-rate schedule and ultimately deploys the raw final iterate\([Brown et al\., 2020](https://arxiv.org/html/2609.25482#bib.bib4);[Hoffmann et al\., 2022](https://arxiv.org/html/2609.25482#bib.bib10);[Touvron et al\., 2023](https://arxiv.org/html/2609.25482#bib.bib27)\)\. This convention answers two design questions with a single choice:*which trajectory should be followed near the end of training*, and*which function of that trajectory should be deployed?*

The end of the schedule is usually a*cooldown*, a terminal reduction of the learning rate\. Its length and specific implementation materially change endpoint loss and can even reorder optimizer comparisons\([Hoffmann et al\., 2022](https://arxiv.org/html/2609.25482#bib.bib10);[Bergsma et al\., 2025](https://arxiv.org/html/2609.25482#bib.bib3);[Wen et al\., 2025](https://arxiv.org/html/2609.25482#bib.bib29)\)\. A small terminal learning rate reduces the variation of the raw final iterate, and it also limits how far the parameters can still move\. A trajectory that stays more active late in training can therefore keep making progress while arriving at a worse raw final iterate\.

![Refer to caption](https://arxiv.org/html/2609.25482v1/figures/fig1aaveraging.png)Figure 1:Averaging weights enables more aggressive terminal optimization\.Raising the terminal learning\-rate floor from5%5\\%to15%15\\%of peak worsens validation bits per byte \(BPB\) for the raw final iterate but improves it for TSA\.Weight averaging separates the trajectory from the model returned from it\([Polyak and Juditsky, 1992](https://arxiv.org/html/2609.25482#bib.bib24);[Izmailov et al\., 2018](https://arxiv.org/html/2609.25482#bib.bib12)\)\. Recent checkpoint\-averaging methods such as Latest Weight Averaging \(LAWA\) have also been used to accelerate pretraining\([Kaddour, 2022](https://arxiv.org/html/2609.25482#bib.bib16)\); in LLM pretraining, averaging can become more effective at higher learning rates\([Sanyal et al\., 2024](https://arxiv.org/html/2609.25482#bib.bib26)\)\. These findings motivate the question we study: if the deployed model averages checkpoints, does a weaker cooldown become preferable, and do the strength of the averaging and the strength of the cooldown need to be chosen together?

To study these questions, we propose terminal shrinkage averaging \(TSA\), an interpolation between the raw final iterate and the average of the lastKKcheckpoints, with a single shrinkage coefficient setting how much of the checkpoint average enters the returned model\.

This interpolation allows us to have control over the strength of averaging\. The raw final iterate and uniform LAWA are its two endpoints\. Under a local quadratic approximation of the loss, we derive its exact surrogate risk, which is quadratic in the shrinkage coefficient\. The risk separates the displacement of the mean output from the reduction in stochastic variation, demonstrating a bias\-variance tradeoff, which gives conditions under which the minimizer is interior, so that neither the raw final iterate nor LAWA is optimal\. We use the same approximation to show that an averaged output prefers a more active terminal trajectory than the raw final iterate does\.

Our experiments are designed to isolate the effect of averaging, the effect of the terminal learning rate, and their interaction\. Depth\-12 NanoChat runs first show that the best shrinkage coefficient lies strictly between the two endpoints\. A paired factorial design then crosses the terminal learning\-rate floor \(5%5\\%,10%10\\%, and15%15\\%of the peak learning rate\) with the output estimator \(the raw final iterate and TSA\)\. Across five repetitions, raising the floor worsens validation BPB when the raw final iterate is deployed and improves it when TSA is deployed, and both effects grow with the size of the increase\. A replication under pure AdamW reproduces the direction of this interaction, so this effect does not depend on NanoChat’s Muon\+AdamW optimizer\. Finally, we apply the higher floor and TSA to the recipe behind the best depth\-22 NanoChat time\-to\-GPT\-2 run, which gives preliminary evidence that the interaction survives at scale, along with a qualifying record run\.

#### Contributions\.

This work makes the following contributions\.

- •We separate two choices in pretraining: the learning\-rate schedule, which determines the parameter trajectory, and the*output estimator*, which maps that trajectory to the deployed model\. We introduce TSA, a shrinkage estimator between the raw final iterate and LAWA, and characterize its bias\-variance tradeoff under a local quadratic approximation\.
- •Through controlled depth\-12 NanoChat experiments, we identify that increasing terminal learning\-rate activity is more favorable when TSA is returned, rather than the raw final iterate\. Paired repetitions isolate this interaction from the effects of averaging or the terminal schedule alone\.
- •We demonstrate that the resulting gains transfer to the depth\-22 NanoChat setting\. At a matched training budget, the full schedule\-estimator recipe improves validation quality over the underlying public recipe and produces a qualifying time\-to\-GPT\-2 benchmark run\.

## 2Related work

#### Iterate and checkpoint averaging\.

Averaging has classical asymptotic guarantees in convex problems\([Polyak and Juditsky, 1992](https://arxiv.org/html/2609.25482#bib.bib24);[Neu and Rosasco, 2018](https://arxiv.org/html/2609.25482#bib.bib23)\), while finite\-sample analyses of tail averaging show how averaging can reduce the variance of the final output\([Jain et al\., 2018](https://arxiv.org/html/2609.25482#bib.bib2)\)\. In deep learning, stochastic weight averaging \(SWA\) returns an average over a late training trajectory and often finds solutions with improved generalization\([Izmailov et al\., 2018](https://arxiv.org/html/2609.25482#bib.bib12)\)\. Latest Weight Averaging \(LAWA\) applies uniform averaging to a finite set of recent checkpoints and frames averaging as a way to accelerate training\([Kaddour, 2022](https://arxiv.org/html/2609.25482#bib.bib16)\)\.

For language\-model pretraining,[Sanyal et al\. \(2024\)](https://arxiv.org/html/2609.25482#bib.bib26)show that checkpoint averaging becomes more effective at higher learning rates, while[Ajroldi et al\. \(2025\)](https://arxiv.org/html/2609.25482#bib.bib1)characterize when averaging improves large\-scale training\. These results ask how the benefit of averaging changes under a given training trajectory\. Our focus is complementary: we ask whether committing to an averaged output changes which terminal trajectory should be trained in the first place\.

Lookahead also interpolates between fast and slow parameter states\([Zhang et al\., 2019](https://arxiv.org/html/2609.25482#bib.bib31)\)\. It performs this interpolation repeatedly during optimization and resets the fast weights after each inner loop\. In contrast, TSA leaves the optimizer updates unchanged and performs a single terminal interpolation between the raw final iterate and a finite late\-checkpoint average to construct the deployed model\.

#### Learning\-rate schedules designed with averaging\.

Several lines of work connecting sustained optimizer activity, cooldowns, and averaging\. Warmup\-stable\-decay schedules retain a long constant\-learning\-rate phase before a terminal decay\([Hu et al\., 2024](https://arxiv.org/html/2609.25482#bib.bib11)\)\. This paper doesn’t touch on averaging\.[Hägele et al\. \(2024\)](https://arxiv.org/html/2609.25482#bib.bib9)study constant\-rate trajectories with cooldowns and show that SWA can be used in place of cooldowns\. By allowing joint tuning of cooldowns and averaging, our method takes their observation that averaging plays a similar role to cooldowns and operationalizes it to improve optimization\. Schedule\-Free optimization asks the question of how to perform pretraining without knowing the length of the training a priori and strongly relies on learning\-rate scheduling and iterate averaging\([Defazio et al\., 2024](https://arxiv.org/html/2609.25482#bib.bib5)\)\. Anytime pretraining falls in a similar regime\([Meterez et al\., 2026b](https://arxiv.org/html/2609.25482#bib.bib21)\)\. Our work is different from these two because it is in the setting where you know the training length in advance\.[Tian et al\. \(2026\)](https://arxiv.org/html/2609.25482#bib.bib7)more formally study the interchangeability of averaging and decay\. Their Warmup\-Stable and Merge framework uses weighted checkpoint merging to emulate decay from a stable\-rate trajectory\. However, they do not study the interpolation and interaction between these two techniques\.[Au and Block \(2026\)](https://arxiv.org/html/2609.25482#bib.bib8)utilize averaging throughout training and develop PACE, which modifies the live AdamW trajectory by pulling parameters toward an exponential moving average\. They do not study the interaction between the decay and the averaging, however\. TSA leaves the live optimizer iterate unmodified and separately intervenes on terminal learning\-rate activity\. These are orthogonal directions\.

Our setting is very controlled\. We retain the underlying horizon\-dependent pretraining recipe, alter only its terminal learning\-rate activity, and separately vary the returned estimator\. This design isolates whether the effect of the terminal schedule depends on which model is ultimately returned\.

#### Pretraining optimizers\.

Adam and AdamW remain standard scalar\-preconditioned optimizers\([Kingma and Ba, 2015](https://arxiv.org/html/2609.25482#bib.bib17);[Loshchilov and Hutter, 2019](https://arxiv.org/html/2609.25482#bib.bib20)\)\. Matrix\-preconditioned methods include Shampoo\([Gupta et al\., 2018](https://arxiv.org/html/2609.25482#bib.bib6)\), SOAP\([Vyas et al\., 2024](https://arxiv.org/html/2609.25482#bib.bib28)\), and Muon, which orthogonalizes momentum updates for matrix\-valued parameters\([Jordan et al\., 2024](https://arxiv.org/html/2609.25482#bib.bib13);[Liu et al\., 2025](https://arxiv.org/html/2609.25482#bib.bib19)\)\. Optimizer rankings can depend on model scale, data ratio, tuning budget, and evaluation horizon\([Wen et al\., 2025](https://arxiv.org/html/2609.25482#bib.bib29);[Wen et al\., 2026](https://arxiv.org/html/2609.25482#bib.bib30)\)\. TSA does not alter the underlying optimizer update rule and is therefore complementary to optimizer design\. Our primary experiments use NanoChat’s native Muon\+AdamW optimizer, and the appendix examines how the averaging gain is distributed across the corresponding parameter groups\.

#### Compute\-optimal pretraining and capability thresholds\.

Scaling\-law studies relate model size, data, and compute\([Kaplan et al\., 2020](https://arxiv.org/html/2609.25482#bib.bib14);[Hoffmann et al\., 2022](https://arxiv.org/html/2609.25482#bib.bib10)\)\. NanoChat provides a compact full\-stack training harness and a time\-to\-GPT\-2 benchmark on one8×8\\timesH100 node\([Karpathy and contributors, 2026](https://arxiv.org/html/2609.25482#bib.bib15)\)\. Qualification uses the 22\-task DCLM CORE metric introduced with DataComp\-LM\([Li and others, 2024](https://arxiv.org/html/2609.25482#bib.bib18)\), with the original GPT\-2 checkpoint as the reference capability\([Radford et al\., 2019](https://arxiv.org/html/2609.25482#bib.bib25)\)\. We begin from the exact PR \#830 recipe available at the time of our experiments\([Zinzi, 2026](https://arxiv.org/html/2609.25482#bib.bib32)\)and use it as the base for our matched\-endpoint scale\-transfer comparisons\.

## 3Method

The model returned from optimization depends on both the saved terminal checkpoints and the rule used to combine them\. We first define that rule and analyze its risk for a fixed trajectory\. We then ask how changing the rule can change the preferred terminal learning\-rate activity\.

### 3\.1Terminal Shrinkage Averaging

Consider optimization of a model withppparameters\. Letθt∈ℝp\\theta\_\{t\}\\in\\mathbb\{R\}^\{p\}denote the iterate at steptt\. At a terminal stepTT, we saveKKcheckpoints separated byssoptimizer steps and define their average as

θ¯T:=1K​∑j=0K−1θT−j​s\.\\bar\{\\theta\}\_\{T\}:=\\frac\{1\}\{K\}\\sum\_\{j=0\}^\{K\-1\}\\theta\_\{T\-js\}\.\(1\)Terminal Shrinkage Averaging returns

θ^T​\(α\):=\(1−α\)​θT\+α​θ¯T,α∈\[0,1\]\.\\widehat\{\\theta\}\_\{T\}\(\\alpha\):=\(1\-\\alpha\)\\theta\_\{T\}\+\\alpha\\bar\{\\theta\}\_\{T\},\\qquad\\alpha\\in\[0,1\]\.\(2\)The endpointα=0\\alpha=0returns the raw final iterate, whileα=1\\alpha=1recovers uniform LAWA\. We use LAWA as the averaging endpoint, although the interpolation can use any other averaging estimator in its place\. Our goal is to characterize when the optimal model lies strictly between these two endpoints\.

BecauseθT\\theta\_\{T\}is itself included inθ¯T\\bar\{\\theta\}\_\{T\}, TSA assigns total weight

1−α\+αK1\-\\alpha\+\\frac\{\\alpha\}\{K\}to the final checkpoint and weightα/K\\alpha/Kto each of the otherK−1K\-1checkpoints\.

#### Local quadratic model\.

To study how the interpolation affects performance, we approximate the loss over the short terminal region visited by the saved checkpoints using a positive\-semidefinite quadratic model,

Lquad​\(θ\):=L⁡\(θ⋆\)\+12​\(θ−θ⋆\)⊤​H​\(θ−θ⋆\),H⪰0,L\_\{\\mathrm\{quad\}\}\(\\theta\):=L\(\\theta^\{\\star\}\)\+\\frac\{1\}\{2\}\(\\theta\-\\theta^\{\\star\}\)^\{\\top\}H\(\\theta\-\\theta^\{\\star\}\),\\qquad H\\succeq 0,\(3\)whereθ⋆\\theta^\{\\star\}is a minimizer of the local model\. SinceHHis positive semidefinite, this optimum may not be unique\. This approximation is intended to describe the short late\-training window used by TSA rather than the full optimization trajectory\. Recent work finds that local linear and quadratic models become particularly accurate late in language\-model pretraining\([Meterez et al\., 2026a](https://arxiv.org/html/2609.25482#bib.bib22)\)\.

Forx,y∈ℝpx,y\\in\\mathbb\{R\}^\{p\}, define the curvature\-weighted inner product and associated seminorm

⟨x,y⟩H:=x⊤​H​y,‖x‖H2:=x⊤​H​x\.\\langle x,y\\rangle\_\{H\}:=x^\{\\top\}Hy,\\qquad\\\|x\\\|\_\{H\}^\{2\}:=x^\{\\top\}Hx\.\(4\)

#### Mean trajectory and stochastic residuals\.

To separate progress along the terminal trajectory from variation across training runs, we decompose each checkpoint into its mean and a stochastic residual\. For notational convenience, write

θj:=θT−j​s,j=0,…,K−1,\\theta\_\{j\}:=\\theta\_\{T\-js\},\\qquad j=0,\\ldots,K\-1,so thatθ0=θT\\theta\_\{0\}=\\theta\_\{T\}\. Define

μj:=𝔼⁡\[θj\],ϵj:=θj−μj,𝔼⁡\[ϵj\]=0\.\\mu\_\{j\}:=\\mathbb\{E\}\[\\theta\_\{j\}\],\\qquad\\epsilon\_\{j\}:=\\theta\_\{j\}\-\\mu\_\{j\},\\qquad\\mathbb\{E\}\[\\epsilon\_\{j\}\]=0\.\(5\)The expectation is over the training stochasticity, including the order of minibatches\.

Define the average of the expected checkpoints and the average stochastic residual as

μ¯:=1K​∑j=0K−1μj,ϵ¯:=1K​∑j=0K−1ϵj\.\\bar\{\\mu\}:=\\frac\{1\}\{K\}\\sum\_\{j=0\}^\{K\-1\}\\mu\_\{j\},\\qquad\\bar\{\\epsilon\}:=\\frac\{1\}\{K\}\\sum\_\{j=0\}^\{K\-1\}\\epsilon\_\{j\}\.\(6\)We also define

m:=μ0−θ⋆,A:=μ0−μ¯\.m:=\\mu\_\{0\}\-\\theta^\{\\star\},\\qquad A:=\\mu\_\{0\}\-\\bar\{\\mu\}\.\(7\)mmis the displacement of the mean of the raw final iterate from the minimizer of the local quadratic model\.AAis the displacement between the mean of the raw final iterate and the expected checkpoint average\.

Substituting these definitions into Eq\.[2](https://arxiv.org/html/2609.25482#S3.E2)gives the exact decomposition

θ^T​\(α\)−θ⋆\\displaystyle\\widehat\{\\theta\}\_\{T\}\(\\alpha\)\-\\theta^\{\\star\}=\(1−α\)​\(μ0\+ϵ0\)\+α⁡\(μ¯\+ϵ¯\)−θ⋆\\displaystyle=\(1\-\\alpha\)\(\\mu\_\{0\}\+\\epsilon\_\{0\}\)\+\\alpha\(\\bar\{\\mu\}\+\\bar\{\\epsilon\}\)\-\\theta^\{\\star\}\(8\)=m−α​A\+ϵ0\+α​q,\\displaystyle=m\-\\alpha A\+\\epsilon\_\{0\}\+\\alpha q,\(9\)where

q:=ϵ¯−ϵ0\.q:=\\bar\{\\epsilon\}\-\\epsilon\_\{0\}\.\(10\)Equation[9](https://arxiv.org/html/2609.25482#S3.E9)separates the effect of TSA into two parts\. The deterministic term−α​A\-\\alpha Adescribes how averaging moves the mean output model relative to the local minimizer\. The stochastic termα​q\\alpha qdescribes how averaging changes the random deviation around that mean\.

We measure performance under the local quadratic surrogate through

ℛ⁡\(α\):=𝔼⁡\[Lquad​\(θ^T​\(α\)\)−Lquad​\(θ⋆\)\]=12​𝔼​‖θ^T​\(α\)−θ⋆‖H2\.\\mathcal\{R\}\(\\alpha\):=\\mathbb\{E\}\\\!\\left\[L\_\{\\mathrm\{quad\}\}\\bigl\(\\widehat\{\\theta\}\_\{T\}\(\\alpha\)\\bigr\)\-L\_\{\\mathrm\{quad\}\}\(\\theta^\{\\star\}\)\\right\]=\\frac\{1\}\{2\}\\mathbb\{E\}\\left\\\|\\widehat\{\\theta\}\_\{T\}\(\\alpha\)\-\\theta^\{\\star\}\\right\\\|\_\{H\}^\{2\}\.\(11\)Althoughθ⋆\\theta^\{\\star\}need not be unique whenH⪰0H\\succeq 0, this quantity is invariant to the choice of local minimizer because distinct minimizers differ only along directions in the null space ofHH\.

###### Theorem 1\(Risk of Terminal Shrinkage Averaging\)\.

Under the local quadratic model, let

m:=μ0−θ⋆,A:=μ0−μ¯,q:=ϵ¯−ϵ0,m:=\\mu\_\{0\}\-\\theta^\{\\star\},\\qquad A:=\\mu\_\{0\}\-\\bar\{\\mu\},\\qquad q:=\\bar\{\\epsilon\}\-\\epsilon\_\{0\},and define

Vq:=𝔼​‖q‖H2,C:=𝔼​⟨ϵ0,q⟩H\.V\_\{q\}:=\\mathbb\{E\}\\\|q\\\|\_\{H\}^\{2\},\\qquad C:=\\mathbb\{E\}\\langle\\epsilon\_\{0\},q\\rangle\_\{H\}\.Then the expected excess risk of TSA satisfies

ℛ⁡\(α\)=ℛ⁡\(0\)\+α⁡\(−⟨m,A⟩H\+C\)\+α22​\(‖A‖H2\+Vq\)\.\\mathcal\{R\}\(\\alpha\)=\\mathcal\{R\}\(0\)\+\\alpha\\bigl\(\-\\langle m,A\\rangle\_\{H\}\+C\\bigr\)\+\\frac\{\\alpha^\{2\}\}\{2\}\\left\(\\\|A\\\|\_\{H\}^\{2\}\+V\_\{q\}\\right\)\.\(12\)Whenever‖A‖H2\+Vq\>0\\\|A\\\|\_\{H\}^\{2\}\+V\_\{q\}\>0, this is a convex quadratic inα\\alphaand

α⋆=Π\[0,1\]​\(⟨m,A⟩H−C‖A‖H2\+Vq\)\.\\alpha^\{\\star\}=\\Pi\_\{\[0,1\]\}\\left\(\\frac\{\\langle m,A\\rangle\_\{H\}\-C\}\{\\\|A\\\|\_\{H\}^\{2\}\+V\_\{q\}\}\\right\)\.\(13\)

The linear term in[Equation12](https://arxiv.org/html/2609.25482#S3.E12)exposes the central tradeoff\. SinceAApoints from the average of the saved checkpoints toward the raw final iterate, moving toward the checkpoint average can incur*lag*when the optimizer is still making useful progress\. Define

Glag:=−⟨m,A⟩H\.G\_\{\\mathrm\{lag\}\}:=\-\\langle m,A\\rangle\_\{H\}\.At the same time, averaging can remove stochastic error from the raw final iterate\. Define

Gnoise:=−C\.G\_\{\\mathrm\{noise\}\}:=\-C\.Then

ℛ⁡\(α\)=ℛ⁡\(0\)\+α⁡\(Glag−Gnoise\)\+α22​\(‖A‖H2\+Vq\)\.\\mathcal\{R\}\(\\alpha\)=\\mathcal\{R\}\(0\)\+\\alpha\\left\(G\_\{\\mathrm\{lag\}\}\-G\_\{\\mathrm\{noise\}\}\\right\)\+\\frac\{\\alpha^\{2\}\}\{2\}\\left\(\\\|A\\\|\_\{H\}^\{2\}\+V\_\{q\}\\right\)\.\(14\)
A small amount of averaging therefore helps whenGnoise\>GlagG\_\{\\mathrm\{noise\}\}\>G\_\{\\mathrm\{lag\}\}: the stochastic error removed by averaging outweighs the cost of moving toward older checkpoints\. Full averaging need not be optimal because the quadratic term makes the marginal benefit of increasingα\\alphadecrease\. In particular, an interior optimum occurs when

0<Gnoise−Glag<‖A‖H2\+Vq\.0<G\_\{\\mathrm\{noise\}\}\-G\_\{\\mathrm\{lag\}\}<\\\|A\\\|\_\{H\}^\{2\}\+V\_\{q\}\.This gives the intuition behind TSA: retain enough weight on the raw final iterate to limit checkpoint lag while averaging strongly enough to reduce terminal variation\.

The proof in[AppendixC](https://arxiv.org/html/2609.25482#A3)formalizes this fixed\-trajectory result, and[SectionB\.9](https://arxiv.org/html/2609.25482#A2.SS9)validates empirically that loss varies quadratically inα\\alpha\. The next question is whether the same tradeoff changes which trajectory should be trained\.

### 3\.2A sufficient condition for a more active terminal schedule

The previous analysis holds the training trajectory fixed and changes only the returned estimator\. We next ask how changing the estimator can change which terminal trajectory is preferable\.

Letρ\\rhoparameterize terminal optimizer activity; in our experiments,ρ\\rhocorresponds to the terminal learning\-rate floor\. For fixedα\\alpha, write the local risk as

ℛα​\(ρ\)=Bα​\(ρ\)\+Sα​\(ρ\),\\mathcal\{R\}\_\{\\alpha\}\(\\rho\)=B\_\{\\alpha\}\(\\rho\)\+S\_\{\\alpha\}\(\\rho\),\(15\)where

Bα​\(ρ\)\\displaystyle B\_\{\\alpha\}\(\\rho\):=12​‖mρ−α​Aρ‖H2,\\displaystyle:=\\frac\{1\}\{2\}\\\|m\_\{\\rho\}\-\\alpha A\_\{\\rho\}\\\|\_\{H\}^\{2\},\(16\)Sα​\(ρ\)\\displaystyle S\_\{\\alpha\}\(\\rho\):=12​𝔼​‖ϵ0,ρ\+α​qρ‖H2\.\\displaystyle:=\\frac\{1\}\{2\}\\mathbb\{E\}\\\|\\epsilon\_\{0,\\rho\}\+\\alpha q\_\{\\rho\}\\\|\_\{H\}^\{2\}\.\(17\)The first term captures mean displacement, while the second captures stochastic variation\.

Let

ρraw:=arg⁡minρ​ℛ0​\(ρ\)\\rho\_\{\\mathrm\{raw\}\}:=\\arg\\min\_\{\\rho\}\\mathcal\{R\}\_\{0\}\(\\rho\)denote the terminal activity preferred when the raw final iterate is returned\.ℛ0′​\(ρraw\)=0\\mathcal\{R\}\_\{0\}^\{\\prime\}\(\\rho\_\{\\mathrm\{raw\}\}\)=0\. Therefore,

ℛα′​\(ρraw\)=dd​ρ​\[Bα​\(ρ\)−B0​\(ρ\)\]ρ=ρraw−dd​ρ​\[S0​\(ρ\)−Sα​\(ρ\)\]ρ=ρraw\.\\mathcal\{R\}\_\{\\alpha\}^\{\\prime\}\(\\rho\_\{\\mathrm\{raw\}\}\)=\\frac\{d\}\{d\\rho\}\\left\[B\_\{\\alpha\}\(\\rho\)\-B\_\{0\}\(\\rho\)\\right\]\_\{\\rho=\\rho\_\{\\mathrm\{raw\}\}\}\-\\frac\{d\}\{d\\rho\}\\left\[S\_\{0\}\(\\rho\)\-S\_\{\\alpha\}\(\\rho\)\\right\]\_\{\\rho=\\rho\_\{\\mathrm\{raw\}\}\}\.\(18\)
Consequently, if the stochastic benefit supplied by averaging grows with terminal activity faster than its additional deterministic cost, thenℛα′​\(ρraw\)<0\\mathcal\{R\}\_\{\\alpha\}^\{\\prime\}\(\\rho\_\{\\mathrm\{raw\}\}\)<0\. Under local convexity, the risk\-minimizing averaged trajectory therefore lies at a larger value ofρ\\rho\.

#### A one\-dimensional corollary\.

We can demonstrate the schedule shift for a simple model using[Equation18](https://arxiv.org/html/2609.25482#S3.E18)\. Consider one noisy terminal update onL⁡\(x\)=h2​x2L\(x\)=\\tfrac\{h\}\{2\}x^\{2\}and construct TSA from the pre\-update and post\-update checkpoints\. Whenever the risk\-minimizing step sizes are interior, the preferred step size under shrinkage strengthα\\alphais

ηα⋆=ηraw⋆1−α/2\.\\eta\_\{\\alpha\}^\{\\star\}=\\frac\{\\eta\_\{\\mathrm\{raw\}\}^\{\\star\}\}\{1\-\\alpha/2\}\.\(19\)Thus everyα\>0\\alpha\>0strictly increases the preferred terminal step size, and full averaging of the two checkpoints doubles it\. The exact risk and proof are provided in[AppendixD](https://arxiv.org/html/2609.25482#A4)\.

This does not assert that increasing terminal activity is always beneficial under averaging, nor does it predict the magnitude of the optimal shift\. Rather, it identifies the condition under which changing the returned estimator changes the preferred terminal schedule\. The experiments in[Section4\.3](https://arxiv.org/html/2609.25482#S4.SS3)test this interaction directly\.

## 4Experiments

We design the experiments as a sequence of controlled intervention tests\. We first ask whether partial averaging improves the output model under a fixed training trajectory\. We then ask whether increasing terminal learning\-rate activity helps without averaging\. Then, we show that our method is robust to choice of optimizer by replicating our results on AdamW\. Finally, we combine the interventions to test the central prediction of our theory: an output estimator that reduces terminal variance can prefer a more active learning\-rate schedule than the raw final iterate\.

We use depth\-12 NanoChat models for controlled ablations and a depth\-22 model for a final time\-to\-GPT\-2 experiment to demonstrate that our method scales\. Unless otherwise stated, all depth\-12 comparisons use the same model initialization, optimizer, training batches, and paired validation data so that differences can be attributed to the intervention being studied\. We report validation bits per byte \(BPB\), where lower is better\. For the depth\-22 speedrun, we additionally report DCLM CORE, for which the qualification threshold is0\.2565250\.256525\. Full details are provided in[AppendixA](https://arxiv.org/html/2609.25482#A1)\.

### 4\.1Partial shrinkage has a replicated interior optimum

We first hold each training trajectory fixed and vary only the returned estimator\. To test whether the single\-trajectory response from our preliminary sweep replicates, we evaluate a denseα\\alphagrid on five independently ordered confirmation trajectories at the nanochat baseline5%5\\%floor\. Every value ofα\\alphais evaluated on the same held\-out batches within a trajectory\. The grid is evaluated only after the confirmation estimator has been frozen and is not used to select that rule\.

Figure 2:Partial shrinkage has a replicated interior optimum\.Points show mean paired BPB improvement over the raw final iterate across five independently ordered depth\-12 trajectories at the5%5\\%terminal floor; bars are 95%tt\-intervals across trajectories\. The best tested point isα=0\.40\\alpha=0\.40, whereas full LAWA \(α=1\\alpha=1\) is worse than the raw final iterate\.The response is broad and clearly interior\. The best tested coefficient,α=0\.40\\alpha=0\.40, improves BPB by0\.000794±0\.0000870\.000794\\pm 0\.000087; the predeclared confirmation settingα=0\.55\\alpha=0\.55improves BPB by0\.000711±0\.0000980\.000711\\pm 0\.000098\. In contrast, uniform LAWA worsens BPB by0\.000734±0\.0001720\.000734\\pm 0\.000172\. Thus, the benefit of partial rather than full averaging persists across training streams and does not require precise tuning ofα\\alpha\. The full response at all three terminal floors is reported in[SectionB\.1](https://arxiv.org/html/2609.25482#A2.SS1)and we compare our method to other common forms of averaging in[AppendixB](https://arxiv.org/html/2609.25482#A2)\.

### 4\.2More terminal activity worsens the raw final iterate

We next vary only the terminal learning\-rate floor and always return the raw final iterate\. Using the same five paired stream seeds, increasing the floor from5%5\\%to10%10\\%worsens BPB for the raw final iterate by0\.000915±0\.0008360\.000915\\pm 0\.000836, while increasing it to15%15\\%worsens BPB for the raw final iterate by0\.003757±0\.0007690\.003757\\pm 0\.000769\. All five paired differences are positive for both interventions\.

Figure 3:More terminal activity worsens the raw final iterate\.Points show the mean change in BPB for the raw final iterate relative to the5%5\\%floor across the five streams; bars are 95%tt\-intervals across streams\. The5%5\\%point is zero by construction\. Increasing terminal activity degrades the raw final iterate at both tested intervention strengths, with a substantially larger effect at the15%15\\%floor\.A denser preliminary single\-trajectory sweep through17\.5%17\.5\\%shows the same monotone, increasingly steep degradation of the raw final iterate\. We report that descriptive sweep separately in[SectionB\.8](https://arxiv.org/html/2609.25482#A2.SS8)\.

These comparisons establish the cost of a more active terminal schedule when the returned model is the raw final iterate\. We next test whether TSA changes that cost\.

### 4\.3The output estimator changes the effect of terminal activity

We now vary the schedule and estimator together\. We hold TSA fixed atα=0\.55\\alpha=0\.55,K=8K=8, ands=32s=32, and compare the5%5\\%baseline floor with both10%10\\%and15%15\\%floors over five paired data\-order repetitions\. For an active floorρ\\rho, define the paired interaction

ℐ⁡\(ρ\):=\[Lraw​\(ρ\)−Lraw​\(5%\)\]−\[LTSA​\(ρ\)−LTSA​\(5%\)\]\.\\mathcal\{I\}\(\\rho\):=\\bigl\[L\_\{\\mathrm\{raw\}\}\(\\rho\)\-L\_\{\\mathrm\{raw\}\}\(5\\%\)\\bigr\]\-\\bigl\[L\_\{\\mathrm\{TSA\}\}\(\\rho\)\-L\_\{\\mathrm\{TSA\}\}\(5\\%\)\\bigr\]\.\(20\)Positive values mean the active schedule is more favorable under TSA than under the raw final iterate\.

Figure 4:The schedule–estimator interaction replicates across intervention strengths and optimizers\.Left:under NanoChat’s native Muon\+AdamW optimizer, points show the BPB effect of an active floor relative to the5%5\\%floor; positive values mean worse BPB\. The15%15\\%floor worsens the raw final iterate but improves TSA\.Right:paired interactions from[Equation20](https://arxiv.org/html/2609.25482#S4.E20)under the native optimizer and pure AdamW, using the same fixedα=0\.55\\alpha=0\.55\. Positive values mean the active floor is more favorable under TSA\. Bars are 95%tt\-intervals over five paired streams\.Table 1:Two replicated schedule–estimator interventions\.Effects are active\-floor BPB minus5%5\\%\-floor BPB, so positive simple effects indicate worse BPB\. The interaction is the effect for the raw final iterate minus that for TSA\. Intervals are 95%tt\-intervals over five paired repetitions\.The smaller10%10\\%intervention already yields a clear interaction, even though the TSA simple effect is not individually resolved\. At15%15\\%, the simple effects themselves point in opposite and statistically resolved directions: the raw final iterate worsens, whereas TSA improves\. The paired interaction is positive for every stream at both floors, showing that TSA makes a weaker cooldown more beneficial\.

### 4\.4The interaction persists under pure AdamW

The native optimizer assigns Muon to the main matrix parameter block\. To test whether that assignment is necessary for the interaction, we repeat the same5%/10%/15%5\\%/10\\%/15\\%paired design under pure AdamW\. We retain the fixed TSA coefficientα=0\.55\\alpha=0\.55and measure the within\-stream interaction\. It is0\.000782±0\.0000890\.000782\\pm 0\.000089BPB for the10%10\\%floor and0\.002536±0\.0002040\.002536\\pm 0\.000204BPB for the15%15\\%floor, and is positive in all five repetitions at both floors \([Figure4](https://arxiv.org/html/2609.25482#S4.F4), right\)\. The interaction is smaller than under the native optimizer, but persists without Muon updates, demonstrating that the phenomenon is not specific to Muon\+AdamW\.

### 4\.5Scale transfer to the NanoChat time\-to\-GPT\-2 benchmark

The depth\-12 results establish the interaction under controlled training conditions\. We next test whether the combined schedule and estimator intervention of TSA transfers to the depth\-22 NanoChat speedrun setting\. We start from the PR \#830 recipe available at the time of our experiments\([Zinzi, 2026](https://arxiv.org/html/2609.25482#bib.bib32)\), whose learning\-rate schedule is calibrated for a data\-to\-parameter ratio of9\.49\.4\. We compare three configurations at the same physical ratio\-9\.09\.0endpoint \(step 10,172\): an early stop of the original ratio\-9\.49\.4schedule, the public recipe recalibrated directly to ratio9\.09\.0, and the ratio\-9\.09\.0schedule with a15%15\\%terminal learning\-rate floor and TSA\. For the treatment, we choseα\\alphabased on the best\-performingα\\alphain the depth\-12 experiment at the15%15\\%floor, which is0\.70\.7as seen in[AppendixB](https://arxiv.org/html/2609.25482#A2)\.

Across three training repetitions, TSA improves validation BPB from0\.720235±0\.0007350\.720235\\pm 0\.000735to0\.719580±0\.0005620\.719580\\pm 0\.000562relative to the calibrated baseline \([Table2](https://arxiv.org/html/2609.25482#S4.T2)\)\. The intervention improves BPB by0\.000654±0\.0001840\.000654\\pm 0\.000184in the paired comparison, with a favorable difference in all three seeds\.111Training\-time differences mostly reflect throughput variation across allocations; we do not interpret them as an algorithmic speed difference\.

Table 2:Matched\-endpoint depth\-22 comparison over three seeds\.All configurations stop at step1017210172\. Time is accumulated training time\. Values are means±\\pm95% Student\-ttintervals\.Relative to the calibrated baseline, the seedwise CORE changes are\+0\.0033\+0\.0033,\+0\.0124\+0\.0124, and\+0\.0108\+0\.0108\. The NanoChat qualification threshold is0\.2565250\.256525\. While the average CORE value exceeds this, only one of the three seeds actually surpassed this threshold\. For context, PR \#830 reports81\.835±0\.13881\.835\\pm 0\.138minutes over six runs\. Because CORE is discrete and noisy, and only one of our three repetitions individually qualifies, we treat this as a preliminary result with respect to an official record\.

## 5Conclusion

Pretraining need not return the raw final iterate\. Once the deployed model is allowed to be a function of the terminal trajectory, the trajectory itself can be optimized differently\. We study one simple realization of this idea: keeping optimization more active near the end of training while using Terminal Shrinkage Averaging \(TSA\) to construct the returned model\. A local quadratic analysis characterizes the resulting tradeoff between optimization progress and stochastic variation and explains why partial, rather than uniform, averaging can be optimal\.

Our controlled depth\-12 experiments isolate the components of the method and establish the central schedule\-estimator interaction: increasing terminal learning\-rate activity can worsen the raw final iterate while improving the averaged output\. 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\. More broadly, these results suggest that the optimization schedule and the rule used to construct the deployed model should be treated as a joint design problem rather than optimized independently\.

#### Limitations\.

TSA explores only a small part of the space of possible output estimators and training trajectories\. We use a single global shrinkage coefficient, although different layers or parameter tensors may benefit from different amounts of averaging; we do preliminary explorations of this in[SectionB\.4](https://arxiv.org/html/2609.25482#A2.SS4)and[SectionB\.5](https://arxiv.org/html/2609.25482#A2.SS5)\. Likewise, we modify terminal optimization through a simple learning\-rate floor, but other changes to cooldown shape, duration, optimizer activity, or checkpoint placement could interact differently with the returned estimator\. An issue with averaging is the memory required, which the observations in[SectionB\.6](https://arxiv.org/html/2609.25482#A2.SS6)help mitigate by showing that averaging’s benefit is heterogeneous between tensors\. Our theoretical analysis uses a local quadratic approximation and is intended to explain the observed late\-training tradeoff rather than provide a complete model of neural\-network optimization\. Finally, our experiments focus on NanoChat language\-model pretraining, and establishing how these interactions transfer across architectures, optimizers, datasets, and substantially larger training regimes remains an important direction for future work\.

## References

- Ajroldiet al\.\(2025\)N\. Ajroldi, A\. Orvieto, and J\. GeipingWhen, where and why to average weights?\.InInternational Conference on Machine Learning,Cited by:[§2](https://arxiv.org/html/2609.25482#S2.SS0.SSS0.Px1.p2.1)\.
- Au and Block \(2026\)K\. C\. Au and A\. BlockTraining for the model you return: improving optimization for iterate\-averaged language models\.External Links:2606\.25086,[Link](https://arxiv.org/abs/2606.25086)Cited by:[§2](https://arxiv.org/html/2609.25482#S2.SS0.SSS0.Px2.p1.1)\.
- Bergsmaet al\.\(2025\)S\. Bergsma, N\. Dey, G\. Gosal, G\. Gray, D\. Soboleva, and J\. HestnessStraight to zero: why linearly decaying the learning rate to zero works best for LLMs\.InInternational Conference on Learning Representations,Cited by:[§1](https://arxiv.org/html/2609.25482#S1.p2.1)\.
- Brownet al\.\(2020\)T\. B\. Brown, B\. Mann, N\. Ryder, M\. Subbiah, J\. D\. Kaplan, P\. Dhariwal, A\. Neelakantan, P\. Shyam, G\. Sastry, A\. Askell, S\. Agarwal, A\. Herbert\-Voss, G\. Krueger, T\. Henighan, R\. Child, A\. Ramesh, D\. M\. Ziegler, J\. Wu, C\. Winter, C\. Hesse, M\. Chen, E\. Sigler, M\. Litwin, S\. Gray, B\. Chess, J\. Clark, C\. Berner, S\. McCandlish, A\. Radford, I\. Sutskever, and D\. AmodeiLanguage models are few\-shot learners\.InAdvances in Neural Information Processing Systems,Vol\.33\.Cited by:[§1](https://arxiv.org/html/2609.25482#S1.p1.1)\.
- Defazioet al\.\(2024\)A\. Defazio, X\. A\. Yang, H\. Mehta, K\. Mishchenko, A\. Khaled, and A\. CutkoskyThe road less scheduled\.arXiv preprint arXiv:2405\.15682\.Cited by:[§2](https://arxiv.org/html/2609.25482#S2.SS0.SSS0.Px2.p1.1)\.
- Guptaet al\.\(2018\)V\. Gupta, T\. Koren, and Y\. SingerShampoo: preconditioned stochastic tensor optimization\.InInternational Conference on Machine Learning,Cited by:[§2](https://arxiv.org/html/2609.25482#S2.SS0.SSS0.Px3.p1.1)\.
- Hägeleet al\.\(2024\)A\. Hägele, E\. Bakouch, A\. Kosson, L\. Ben Allal, L\. von Werra, and M\. JaggiScaling laws and compute\-optimal training beyond fixed training durations\.InAdvances in Neural Information Processing Systems,External Links:[Link](https://arxiv.org/abs/2405.18392)Cited by:[§2](https://arxiv.org/html/2609.25482#S2.SS0.SSS0.Px2.p1.1)\.
- Hoffmannet al\.\(2022\)J\. Hoffmann, S\. Borgeaud, A\. Mensch, E\. Buchatskaya, T\. Cai, E\. Rutherford, D\. de Las Casas, L\. A\. Hendricks, J\. Welbl, A\. Clark, T\. Hennigan, E\. Noland, K\. Millican, G\. van den Driessche, B\. Damoc, A\. Guy, S\. Osindero, K\. Simonyan, E\. Elsen, J\. W\. Rae, O\. Vinyals, and L\. SifreTraining compute\-optimal large language models\.arXiv preprint arXiv:2203\.15556\.Cited by:[§1](https://arxiv.org/html/2609.25482#S1.p1.1),[§1](https://arxiv.org/html/2609.25482#S1.p2.1),[§2](https://arxiv.org/html/2609.25482#S2.SS0.SSS0.Px4.p1.1)\.
- Huet al\.\(2024\)S\. Hu, Y\. Tu, X\. Han, C\. He, G\. Cui, X\. Long, Z\. Zheng, Y\. Fang, Y\. Huang, W\. Zhao, X\. Zhang, Z\. L\. Thai, K\. Zhang, C\. Wang, Y\. Yao, C\. Zhao, J\. Zhou, J\. Cai, Z\. Zhai, N\. Ding, C\. Jia, G\. Zeng, D\. Li, Z\. Liu, and M\. SunMiniCPM: unveiling the potential of small language models with scalable training strategies\.External Links:2404\.06395,[Link](https://arxiv.org/abs/2404.06395)Cited by:[§2](https://arxiv.org/html/2609.25482#S2.SS0.SSS0.Px2.p1.1)\.
- Izmailovet al\.\(2018\)P\. Izmailov, D\. Podoprikhin, T\. Garipov, D\. Vetrov, and A\. G\. WilsonAveraging weights leads to wider optima and better generalization\.InProceedings of the Thirty\-Fourth Conference on Uncertainty in Artificial Intelligence,Cited by:[§1](https://arxiv.org/html/2609.25482#S1.p3.1),[§2](https://arxiv.org/html/2609.25482#S2.SS0.SSS0.Px1.p1.1)\.
- Jainet al\.\(2018\)P\. Jain, S\. M\. Kakade, R\. Kidambi, P\. Netrapalli, and A\. SidfordParallelizing stochastic gradient descent for least squares regression: mini\-batching, averaging, and model misspecification\.Journal of Machine Learning Research18\(223\),pp\. 1–42\.Cited by:[§2](https://arxiv.org/html/2609.25482#S2.SS0.SSS0.Px1.p1.1)\.
- Jordanet al\.\(2024\)K\. Jordan, Y\. Jin, V\. Boza, Y\. Jiacheng, F\. Cesista, L\. Newhouse, and J\. BernsteinMuon: an optimizer for hidden layers in neural networks\.Note:[https://github\.com/KellerJordan/Muon](https://github.com/KellerJordan/Muon)Cited by:[§2](https://arxiv.org/html/2609.25482#S2.SS0.SSS0.Px3.p1.1)\.
- Kaddour \(2022\)J\. KaddourStop wasting my time\! saving days of ImageNet and BERT training with latest weight averaging\.arXiv preprint arXiv:2209\.14981\.Cited by:[§1](https://arxiv.org/html/2609.25482#S1.p3.1),[§2](https://arxiv.org/html/2609.25482#S2.SS0.SSS0.Px1.p1.1)\.
- Kaplanet al\.\(2020\)J\. Kaplan, S\. McCandlish, T\. Henighan, T\. B\. Brown, B\. Chess, R\. Child, S\. Gray, A\. Radford, J\. Wu, and D\. AmodeiScaling laws for neural language models\.arXiv preprint arXiv:2001\.08361\.Cited by:[§2](https://arxiv.org/html/2609.25482#S2.SS0.SSS0.Px4.p1.1)\.
- Karpathy and contributors \(2026\)A\. Karpathy and contributorsNanochat: a minimal experimental harness for training language models\.Note:[https://github\.com/karpathy/nanochat](https://github.com/karpathy/nanochat)Cited by:[§2](https://arxiv.org/html/2609.25482#S2.SS0.SSS0.Px4.p1.1)\.
- Kingma and Ba \(2015\)D\. P\. Kingma and J\. BaAdam: a method for stochastic optimization\.InInternational Conference on Learning Representations,Cited by:[§2](https://arxiv.org/html/2609.25482#S2.SS0.SSS0.Px3.p1.1)\.
- Liet al\.\(2024\)J\. Liet al\.DataComp\-LM: in search of the next generation of language model pretraining datasets\.InAdvances in Neural Information Processing Systems, Datasets and Benchmarks Track,Note:arXiv:2406\.11794Cited by:[§2](https://arxiv.org/html/2609.25482#S2.SS0.SSS0.Px4.p1.1)\.
- Liuet al\.\(2025\)J\. Liu, J\. Su, X\. Yao, Z\. Jiang, G\. Lai, Y\. Du, Y\. Qin, W\. Xu, E\. Lu, J\. Yan,et al\.Muon is scalable for LLM training\.arXiv preprint arXiv:2502\.16982\.Cited by:[§2](https://arxiv.org/html/2609.25482#S2.SS0.SSS0.Px3.p1.1)\.
- Loshchilov and Hutter \(2019\)I\. Loshchilov and F\. HutterDecoupled weight decay regularization\.InInternational Conference on Learning Representations,Cited by:[§2](https://arxiv.org/html/2609.25482#S2.SS0.SSS0.Px3.p1.1)\.
- Meterezet al\.\(2026a\)A\. Meterez, P\. A\. Nair, D\. Morwani, C\. Pehlevan, S\. Kakade, and A\. DamianA defense of the quadratic model\.External Links:2607\.21716,[Link](https://arxiv.org/abs/2607.21716)Cited by:[§3\.1](https://arxiv.org/html/2609.25482#S3.SS1.SSS0.Px1.p1.2)\.
- Meterezet al\.\(2026b\)A\. Meterez, P\. A\. Nair, D\. Morwani, C\. Pehlevan, and S\. KakadeAnytime pretraining: horizon\-free learning\-rate schedules with weight averaging\.arXiv preprint arXiv:2602\.03702\.Cited by:[§2](https://arxiv.org/html/2609.25482#S2.SS0.SSS0.Px2.p1.1)\.
- Neu and Rosasco \(2018\)G\. Neu and L\. RosascoIterate averaging as regularization for stochastic gradient descent\.arXiv preprint arXiv:1802\.08009\.Cited by:[§2](https://arxiv.org/html/2609.25482#S2.SS0.SSS0.Px1.p1.1)\.
- Polyak and Juditsky \(1992\)B\. T\. Polyak and A\. B\. JuditskyAcceleration of stochastic approximation by averaging\.SIAM Journal on Control and Optimization30\(4\),pp\. 838–855\.External Links:[Document](https://dx.doi.org/10.1137/0330046)Cited by:[§1](https://arxiv.org/html/2609.25482#S1.p3.1),[§2](https://arxiv.org/html/2609.25482#S2.SS0.SSS0.Px1.p1.1)\.
- Radfordet al\.\(2019\)A\. Radford, J\. Wu, R\. Child, D\. Luan, D\. Amodei, and I\. SutskeverLanguage models are unsupervised multitask learners\.Technical reportOpenAI\.Cited by:[§2](https://arxiv.org/html/2609.25482#S2.SS0.SSS0.Px4.p1.1)\.
- Sanyalet al\.\(2024\)S\. Sanyal, A\. T\. Neerkaje, J\. Kaddour, A\. Kumar, and S\. SanghaviEarly weight averaging meets high learning rates for LLM pre\-training\.InFirst Conference on Language Modeling,External Links:[Link](https://openreview.net/forum?id=IA8CWtNkUr)Cited by:[§1](https://arxiv.org/html/2609.25482#S1.p3.1),[§2](https://arxiv.org/html/2609.25482#S2.SS0.SSS0.Px1.p2.1)\.
- Tianet al\.\(2026\)C\. Tian, J\. Wang, Q\. Zhao, K\. Chen, J\. Liu, Z\. Liu, J\. Mao, W\. X\. Zhao, Z\. Zhang, and J\. ZhouWSM: decay\-free learning rate schedule via checkpoint merging for LLM pre\-training\.InInternational Conference on Learning Representations,External Links:[Link](https://openreview.net/forum?id=HhThhjKyfw)Cited by:[§2](https://arxiv.org/html/2609.25482#S2.SS0.SSS0.Px2.p1.1)\.
- Touvronet al\.\(2023\)H\. Touvron, L\. Martin, K\. Stone, P\. Albert, A\. Almahairi, Y\. Babaei, N\. Bashlykov, S\. Batra, P\. Bhargava, S\. Bhosale, D\. Bikel, L\. Blecher, C\. C\. Ferrer, M\. Chen, G\. Cucurull, D\. Esiobu, J\. Fernandes, J\. Fu, W\. Fu, B\. Fuller, C\. Gao, V\. Goswami, N\. Goyal, A\. Hartshorn, S\. Hosseini, R\. Hou, H\. Inan, M\. Kardas, V\. Kerkez, M\. Khabsa, I\. Kloumann, A\. Korenev, P\. S\. Koura, M\. Lachaux, T\. Lavril, J\. Lee, D\. Liskovich, Y\. Lu, Y\. Mao, X\. Martinet, T\. Mihaylov, P\. Mishra, I\. Molybog, Y\. Nie, A\. Poulton, J\. Reizenstein, R\. Rungta, K\. Saladi, A\. Schelten, R\. Silva, E\. M\. Smith, R\. Subramanian, X\. E\. Tan, B\. Tang, R\. Taylor, A\. Williams, J\. X\. Kuan, P\. Xu, Z\. Yan, I\. Zarov, Y\. Zhang, A\. Fan, M\. Kambadur, S\. Narang, A\. Rodriguez, R\. Stojnic, S\. Edunov, and T\. ScialomLlama 2: open foundation and fine\-tuned chat models\.arXiv preprint arXiv:2307\.09288\.Cited by:[§1](https://arxiv.org/html/2609.25482#S1.p1.1)\.
- Vyaset al\.\(2024\)N\. Vyas, D\. Morwani, R\. Zhao, M\. Kwun, I\. Shapira, D\. Brandfonbrener, L\. Janson, and S\. KakadeSOAP: improving and stabilizing shampoo using adam\.arXiv preprint arXiv:2409\.11321\.Cited by:[§2](https://arxiv.org/html/2609.25482#S2.SS0.SSS0.Px3.p1.1)\.
- Wenet al\.\(2026\)K\. Wen, X\. Dang, K\. Lyu, T\. Ma, and P\. LiangFantastic pretraining optimizers and where to find them II: hyperball optimization\.arXiv preprint arXiv:2606\.16899\.Cited by:[§2](https://arxiv.org/html/2609.25482#S2.SS0.SSS0.Px3.p1.1)\.
- Wenet al\.\(2025\)K\. Wen, D\. Hall, T\. Ma, and P\. LiangFantastic pretraining optimizers and where to find them\.arXiv preprint arXiv:2509\.02046\.Cited by:[§1](https://arxiv.org/html/2609.25482#S1.p2.1),[§2](https://arxiv.org/html/2609.25482#S2.SS0.SSS0.Px3.p1.1)\.
- Zhanget al\.\(2019\)M\. R\. Zhang, J\. Lucas, J\. Ba, and G\. E\. HintonLookahead optimizer: k steps forward, 1 step back\.InAdvances in Neural Information Processing Systems,Vol\.32\.Cited by:[§2](https://arxiv.org/html/2609.25482#S2.SS0.SSS0.Px1.p3.1)\.
- Zinzi \(2026\)G\. ZinziSpeed up the 8xH100 GPT\-2 run to 81\.8 minutes\.Note:nanochat pull request \#830,[https://github\.com/karpathy/nanochat/pull/830](https://github.com/karpathy/nanochat/pull/830)Cited by:[§2](https://arxiv.org/html/2609.25482#S2.SS0.SSS0.Px4.p1.1),[§4\.5](https://arxiv.org/html/2609.25482#S4.SS5.p1.1)\.

## Appendix AExperimental details

This section specifies the training trajectories, estimator selection, evaluation protocol, and scale\-transfer comparison underlying the main experiments\.

### A\.1Depth\-12 model and common training setup

The depth\-12 experiments use the NanoChat architecture with 12 transformer layers and 286,261,730 trainable parameters\. Runs use sequence length 512, device batch size 4, gradient accumulation 16, and therefore 32,768 tokens per optimizer step\. Every trajectory is trained forT=3000T=3000optimizer steps, with the first 40 steps used for learning\-rate warmup\. The model initialization seed is 1337 in every run; training stochasticity is introduced by an explicit permutation of the 3000 stored optimizer\-step batch blocks\.

The native optimizer follows NanoChat’s heterogeneous parameter assignment: Muon updates the main matrix\-valued parameter block, while AdamW updates the remaining parameter groups\. The corresponding learning\-rate settings are matrix learning rate0\.020\.02, embedding learning rate0\.20\.2, unembedding learning rate0\.0040\.004, scalar learning rate0\.50\.5, and weight decay0\.280\.28\. In the pure\-AdamW control, the matrix block is instead updated by AdamW with learning rate0\.0030\.003,β1=0\.9\\beta\_\{1\}=0\.9,β2=0\.95\\beta\_\{2\}=0\.95,ϵ=10−10\\epsilon=10^\{\-10\}, and matrix weight decay0\.10\.1; the remaining group\-specific settings are inherited from the same pack\. We use the pure\-AdamW arm only for within\-optimizer estimator comparisons, not to rank the two optimizers by absolute BPB\.

A terminal floorρ\\rhomeans that we follow the original cooldown until its learning\-rate multiplier reachesρ\\rhoand then hold the multiplier fixed for the remainder of training\. Unless otherwise stated, TSA usesK=8K=8checkpoints separated bys=32s=32steps\. The selected steps end at step 3000 and span 224 optimizer steps, approximately 7\.5% of training\. Snapshots are stored in bfloat16 and all returned candidates are materialized and evaluated in float32\.

### A\.2Development and confirmation split

The new replicated suite separates rule development from confirmation\. The development stage contains five streams for each optimizer at the10%10\\%floor, with stream seeds

\{6103,7203,8303,9403,10503\}\.\\\{6103,7203,8303,9403,10503\\\}\.On a calibration block, it evaluates a scalar gridα∈\{0,0\.05,…,1\}\\alpha\\in\\\{0,0\.05,\\ldots,1\\\}and a two\-group gridαA,αB∈\{0,0\.125,…,1\}\\alpha\_\{A\},\\alpha\_\{B\}\\in\\\{0,0\.125,\\ldots,1\\\}\. Rules maximize mean paired BPB gain over the raw final iterate across the five development streams, with deterministic ties favoring less extreme coefficients\.

The confirmation stage contains five new stream seeds,

\{11103,12203,13303,14403,15503\},\\\{11103,12203,13303,14403,15503\\\},and crosses two optimizers with terminal floors5%5\\%,10%10\\%, and15%15\\%, for 30 independently trained trajectories\. Within each stream seed, all floor and optimizer arms begin from the same sampled initial parameter vector and use the same explicit data\-order permutation\. A smoke test also verifies matching sampled training and evaluation data fingerprints across the native and pure\-AdamW packs\. Confirmation trajectories are trained and saved before the rule\-freeze job, but are not evaluated until the frozen\-rule file exists\.

The native development stage selects scalarα=0\.55\\alpha=0\.55\. This is also the pre\-existing main\-paper setting, and it is used unchanged in both optimizer regimes in[Section4\.4](https://arxiv.org/html/2609.25482#S4.SS4)\. Pure\-AdamW development separately selectsα=0\.50\\alpha=0\.50; optimizer\-specific selected results are reported only in the structured\-estimator appendix below\. The distinction lets the main text test transfer of the same estimator while retaining the full development study for diagnostics\.

### A\.3Evaluation protocol and uncertainty

Each pack contains 256 stored evaluation batches\. The replicated suite reserves batches 0–31 for the earlier curve block, batches 32–95 for coefficient calibration, batches 96–159 for floor\-development bookkeeping, and batches 160–255 as the untouched 96\-batch confirmation holdout\. All estimators constructed from a given trajectory are evaluated on exactly the same holdout batches\.

For single\-trajectory sweeps, intervals quantify paired variation across held\-out evaluation batches and do not measure training\-run uncertainty\. For the replicated experiments, all intervals in the main text and this appendix are 95% Student\-ttintervals across five independently ordered training streams\. Estimator gains and schedule–estimator interactions are computed within stream before averaging, preserving the paired design\.

The schedule effects for the raw final iterate and TSA are strongly correlated across paired streams \(r=0\.995r=0\.995for the5%→10%5\\%\\\!\\to\\\!10\\%intervention andr=0\.994r=0\.994for5%→15%5\\%\\\!\\to\\\!15\\%under the native optimizer\)\. This common\-mode variation is removed by the within\-stream interaction, explaining why its uncertainty is substantially smaller than that of either simple schedule effect\.

### A\.4Depth\-22 scale\-transfer protocol

The depth\-22 experiments start from NanoChat PR \#830 at commite09bc164162f35da0b5b8315be791e9e974a4c3a\. The recipe uses a depth\-22 model, a 49,152\-token tokenizer, FP8 training, fused Liger cross\-entropy, device batch size 32, and learnable RMSNorm scales\. A training ratio denotes the target data\-to\-parameter ratio; ratio9\.09\.0corresponds to optimizer step 10,172, whereas the original PR \#830 ratio\-9\.49\.4schedule has a horizon of 10,624 steps\.

We compare three configurations at the same physical ratio\-9\.09\.0endpoint\. First, the*early\-stop*control follows the original ratio\-9\.49\.4PR \#830 schedule but stops at step 10,172 and returns the raw final iterate\. Second, the*direct\-9\.09\.0*control recalibrates the schedule itself to ratio9\.09\.0, again with no added terminal floor, and returns the raw final iterate\. Third, the treatment uses the same direct ratio\-9\.09\.0schedule, clamps the terminal learning\-rate multiplier from below at15%15\\%of peak, and returns TSA\. Each configuration is evaluated for experiment seeds4242,4343, and4444\.

For TSA, we useK=8K=8checkpoints at

9381,9494,9607,9720,9833,9946,10059,10172,9381,\\ 9494,\\ 9607,\\ 9720,\\ 9833,\\ 9946,\\ 10059,\\ 10172,corresponding to a spacing of 113 optimizer steps and a total terminal window of 791 steps, approximately7\.8%7\.8\\%of the ratio\-9\.09\.0endpoint horizon\. The 113\-step spacing was originally obtained by transferring the depth\-12 spacing fraction,

323000≈1\.07%,\\frac\{32\}\{3000\}\\approx 1\.07\\%,to the ratio\-9\.49\.4depth\-22 schedule horizon\. When recalibrating the treatment schedule to ratio9\.09\.0, we intentionally retain these same physical checkpoint steps rather than rescaling the averaging window\. Thus, the direct\-9\.09\.0control and treatment differ in terminal learning\-rate activity and returned estimator, while the TSA checkpoint window remains fixed\.

We select the depth\-22 TSA coefficient by selecting the best\-performingα\\alphain the depth\-12 experiment for a15%15\\%floor\. We therefore usedα=0\.70\\alpha=0\.70from[AppendixB](https://arxiv.org/html/2609.25482#A2)\.

For each configuration, we report canonical validation BPB and DCLM CORE over the three experiment seeds\. Training time is NanoChat’s accumulated training\-iteration time through step 10,172 rather than end\-to\-end wall\-clock duration\. TSA introduces no additional optimizer steps\. Because the repetitions span separate allocations, small timing differences mostly include hardware\-session variation; we therefore use training time primarily to establish the matched compute budget rather than to claim an algorithmic speed difference between configurations\.

## Appendix BAdditional depth\-12 results

The main text uses fixed scalar TSA to isolate the schedule–estimator interaction\. Here we examine the coefficient response across conditions, sensitivity to checkpoint selection, and whether alternative estimators or parameter\-specific rules change that picture\.

### B\.1Replicated shrinkage responses across floors and optimizers

After freezing the confirmation estimators, we evaluate the complete scalarα\\alphagrid on the untouched holdouts\. These curves are descriptive and do not alter the predeclared confirmation rule\.

Figure 5:The shrinkage response remains interior and shifts with terminal activity\.Curves show mean paired BPB improvement over the raw final iterate across five confirmation trajectories; shaded regions are 95%tt\-intervals across streams\. Crosses mark the best tested coefficients\. Under the native optimizer the best testedα\\alphamoves from0\.400\.40to0\.550\.55to0\.700\.70as the floor rises from5%5\\%to10%10\\%to15%15\\%; under pure AdamW it moves from0\.300\.30to0\.500\.50to0\.650\.65\.Table 3:Descriptive optima of the post\-freeze scalar grids\.The grids were evaluated on confirmation holdouts after the estimator rules had been frozen and were not used for selection\.The systematic rightward shift is consistent with the paper’s comparative statics: when the terminal trajectory is more active, the returned model can benefit from stronger shrinkage\. It also clarifies why a single moderate value such asα=0\.55\\alpha=0\.55is robust across the tested settings even though it is not the pointwise optimum at every floor\. Although the pointwise optimum shifts with terminal activity, the fixedα=0\.55\\alpha=0\.55rule remains effective across the native\-optimizer conditions: it retains approximately 90%, 100%, and 96% of the best observed scalar\-TSA gain at the 5%, 10%, and 15% floors, respectively\. Thus the interaction does not depend on precise coefficient tuning\.

### B\.2Sensitivity to checkpoint count and spacing

We train five additional depth\-12 trajectories with the10%10\\%floor and stream seeds\{1103,2203,3303,4403,5503\}\\\{1103,2203,3303,4403,5503\\\}\. Every estimator is evaluated post hoc on the same five trajectories\. We fixα=0\.55\\alpha=0\.55and varyK∈\{4,8,16\}K\\in\\\{4,8,16\\\}ands∈\{16,32,64\}s\\in\\\{16,32,64\\\}\.

![Refer to caption](https://arxiv.org/html/2609.25482v1/fig_appendix_ks_heatmap.png)Figure 6:TSA remains beneficial across checkpoint counts and spacings\.Cells show mean BPB improvement over the raw final iterate across five training repetitions\. The outlined cell is the main\-paper settingK=8,s=32K=8,s=32\.TSA improves over the raw final iterate for all nine windows\. The mainK=8,s=32K=8,s=32setting gains0\.002070±0\.0002050\.002070\\pm 0\.000205BPB; the best tested window,K=16,s=32K=16,s=32, gains0\.002363±0\.0002070\.002363\\pm 0\.000207\. Uniform LAWA is more sensitive: atK=8,s=32K=8,s=32it gains only0\.000684±0\.0002610\.000684\\pm 0\.000261, and atK=16,s=64K=16,s=64it is worse than the raw final iterate by0\.005222±0\.0006260\.005222\\pm 0\.000626\. Partial shrinkage therefore broadens the useful range of checkpoint windows\.

These trajectories and their evaluation block are separate from the main confirmation suite, so their nominally identicalK=8,s=32K=8,s=32estimate need not numerically equal[Table3](https://arxiv.org/html/2609.25482#A2.T3)\.

### B\.3Alternative output estimators

The checkpoint\-window results motivate a comparison with other ways to weight recent checkpoints\. We use the following returned estimators\.

Letθ\(1\),…,θ\(K\)\\theta^\{\(1\)\},\\ldots,\\theta^\{\(K\)\}denote selected checkpoints ordered from oldest to newest, withθ\(K\)=θT\\theta^\{\(K\)\}=\\theta\_\{T\}\.

#### Raw final iterate\.

θ^raw=θT\\widehat\{\\theta\}\_\{\\mathrm\{raw\}\}=\\theta\_\{T\}\.

#### Latest Weight Averaging \(LAWA\)\.

θ^LAWA=1K​∑i=1Kθ\(i\)\.\\widehat\{\\theta\}\_\{\\mathrm\{LAWA\}\}=\\frac\{1\}\{K\}\\sum\_\{i=1\}^\{K\}\\theta^\{\(i\)\}\.

#### Terminal Shrinkage Averaging\.

θ^TSA​\(α\)=\(1−α\)​θT\+α​θ^LAWA\.\\widehat\{\\theta\}\_\{\\mathrm\{TSA\}\}\(\\alpha\)=\(1\-\\alpha\)\\theta\_\{T\}\+\\alpha\\widehat\{\\theta\}\_\{\\mathrm\{LAWA\}\}\.Thus the raw final iterate and LAWA are the endpointsα=0\\alpha=0andα=1\\alpha=1\.

#### Finite\-window EWA\.

θ^EWA​\(β\)=∑i=1KβK−i​θ\(i\)∑i=1KβK−i,0<β<1\.\\widehat\{\\theta\}\_\{\\mathrm\{EWA\}\}\(\\beta\)=\\frac\{\\sum\_\{i=1\}^\{K\}\\beta^\{K\-i\}\\theta^\{\(i\)\}\}\{\\sum\_\{i=1\}^\{K\}\\beta^\{K\-i\}\},\\qquad 0<\\beta<1\.Matched\-window controls use the sameK=8,s=32K=8,s=32checkpoints andβ∈\{0\.50,0\.75,0\.90,0\.95\}\\beta\\in\\\{0\.50,0\.75,0\.90,0\.95\\\}\.

#### Checkpoint EMA\.

The checkpoint\-EMA control uses the same normalized exponential rule over a denser, longer sequence withβ=0\.95,K=16,s=16\\beta=0\.95,K=16,s=16\.

#### SWA\-style late average\.

The SWA\-style control is a broader uniform average withK=32,s=16K=32,s=16\. It is mathematically a late checkpoint average but spans a substantially longer terminal window than the main LAWA comparison\.

More generally, shrinkage can be applied toward any fixed weighted checkpoint estimator\. The same local quadratic decomposition follows after redefining the averaged endpoint\. We focus on the uniform finite\-window average to keep the method and controlled comparisons simple\.

Figure 7:Alternative output estimators at depth 12\.Points show mean paired BPB improvement over the raw final iterate, with 95%tt\-intervals across five training repetitions\. Matched EWA denotes the best tested finite\-window EWA \(β=0\.75\\beta=0\.75\) on the sameK=8,s=32K=8,s=32checkpoints as TSA\.The strongest matched EWA gains0\.001910±0\.0002090\.001910\\pm 0\.000209BPB\. TSA gains0\.002070±0\.0002050\.002070\\pm 0\.000205, with a paired TSA\-over\-EWA advantage of0\.000159±0\.0001010\.000159\\pm 0\.000101\. Checkpoint EMA gains0\.001113±0\.0002500\.001113\\pm 0\.000250, while the long SWA\-style average is worse than the raw final iterate by0\.000775±0\.0003750\.000775\\pm 0\.000375\. Recency weighting is useful, but the simple endpoint\-shrinkage family remains competitive\.

### B\.4Exploratory tensorwise adaptive TSA

The preceding comparisons use one weighting rule across all parameters, implicitly treating their terminal lag–variation tradeoffs alike\. To explore whether terminal updates can inform tensor\-specific coefficients, letΔi\(g\)=θi\(g\)−θi−1\(g\)\\Delta\_\{i\}^\{\(g\)\}=\\theta\_\{i\}^\{\(g\)\}\-\\theta\_\{i\-1\}^\{\(g\)\}be the update between adjacent saved checkpoints for tensorgg, and letΔ¯\(g\)\\bar\{\\Delta\}^\{\(g\)\}be the mean update\. Define coherent drift energy

Dg=‖Δ¯\(g\)‖22D\_\{g\}=\\\|\\bar\{\\Delta\}^\{\(g\)\}\\\|\_\{2\}^\{2\}and residual update energyNgN\_\{g\}after subtracting the mean drift\. The empirical noise fraction is

νg=NgNg\+Dg\+ϵ\.\\nu\_\{g\}=\\frac\{N\_\{g\}\}\{N\_\{g\}\+D\_\{g\}\+\\epsilon\}\.As a directional\-persistence statistic, letcgc\_\{g\}be mean cosine similarity between consecutive updates and define

og=clip⁡\(1−max⁡\(cg,0\),0,1\)\.o\_\{g\}=\\operatorname\{clip\}\\\!\\left\(1\-\\max\(c\_\{g\},0\),0,1\\right\)\.The heuristic adaptive rule is

αg=clip⁡\(νg\+og2,0,1\),θ^T\(g\)=\(1−αg\)​θT\(g\)\+αg​θ¯T\(g\)\.\\alpha\_\{g\}=\\operatorname\{clip\}\\\!\\left\(\\frac\{\\nu\_\{g\}\+o\_\{g\}\}\{2\},0,1\\right\),\\qquad\\widehat\{\\theta\}\_\{T\}^\{\(g\)\}=\(1\-\\alpha\_\{g\}\)\\theta\_\{T\}^\{\(g\)\}\+\\alpha\_\{g\}\\bar\{\\theta\}\_\{T\}^\{\(g\)\}\.It is not an estimator of the exact optimum in[Theorem1](https://arxiv.org/html/2609.25482#Thmtheorem1); it is a trajectory\-statistic heuristic inspired by the same lag–variation intuition\.

Across the five appendix trajectories, adaptive TSA over all tensors gains0\.002002±0\.0001700\.002002\\pm 0\.000170BPB\. Applying the adaptive rule only to the Muon\-managed tensors gains0\.001951±0\.0001400\.001951\\pm 0\.000140, recovering approximately 97% of the all\-tensor benefit\. This observation motivated the initial optimizer\-aligned two\-group study below\.

### B\.5Initial optimizer\-aligned structured TSA and parameter localization

The adaptive results suggest that shrinkage gains differ across tensors\. We test a simpler, predeclared split based on NanoChat’s optimizer parameter groups\.

LetPAP\_\{A\}project onto the flattened coordinates assigned to Muon by NanoChat’s native optimizer, and letPB=I−PAP\_\{B\}=I\-P\_\{A\}\. GroupAAcontains 84,935,088 parameters, or 29\.67% of the model\. In pure\-AdamW runs we retain the identical coordinate partition even though all parameters are trained with AdamW\. Structured TSA is

θ^T​\(αA,αB\)=θT\+αA​PA​\(θ¯T−θT\)\+αB​PB​\(θ¯T−θT\)\.\\widehat\{\\theta\}\_\{T\}\(\\alpha\_\{A\},\\alpha\_\{B\}\)=\\theta\_\{T\}\+\\alpha\_\{A\}P\_\{A\}\(\\bar\{\\theta\}\_\{T\}\-\\theta\_\{T\}\)\+\\alpha\_\{B\}P\_\{B\}\(\\bar\{\\theta\}\_\{T\}\-\\theta\_\{T\}\)\.\(21\)Scalar TSA is the diagonal restrictionαA=αB\\alpha\_\{A\}=\\alpha\_\{B\}\.

The development grids select the same off\-diagonal pair,\(αA,αB\)=\(0\.625,0\.250\)\(\\alpha\_\{A\},\\alpha\_\{B\}\)=\(0\.625,0\.250\), under both optimizers \([Figure8](https://arxiv.org/html/2609.25482#A2.F8)\)\. The rule and group\-schema hashes are written before any confirmation evaluation\.

![Refer to caption](https://arxiv.org/html/2609.25482v1/fig_appendix_structured_development_surfaces.png)Figure 8:Development surfaces for the initial optimizer\-aligned two\-group TSA\.Each panel reports mean paired BPB gain over the raw final iterate across five development streams at the10%10\\%floor\. The dashed line is scalar TSA\. The same off\-diagonal pair is selected under the native and pure\-AdamW optimizers\.Table 4:Frozen structured TSA versus frozen scalar TSA\.Development selected\(αA,αB\)=\(0\.625,0\.250\)\(\\alpha\_\{A\},\\alpha\_\{B\}\)=\(0\.625,0\.250\)for both optimizers, scalarα=0\.55\\alpha=0\.55for the native optimizer, and scalarα=0\.50\\alpha=0\.50for pure AdamW\. Positive values in the final column favor structured TSA\.On the native confirmation runs, structured TSA provides only0\.000040±0\.0000500\.000040\\pm 0\.000050BPB beyond the best frozen scalar rule at the predeclared10%10\\%decision gate\. The interval includes zero, and the additional gains at5%5\\%and15%15\\%are similarly small\. We therefore retain scalar TSA as the main method\. Under pure AdamW, the structured increment is positive and small at5%5\\%and10%10\\%, but disappears at15%15\\%\. These findings show measurable heterogeneity without establishing a practically important advantage for the extra parameter\.

The group\-only ablation is more revealing\. Shrinking only the 29\.67% matrix group recovers most of the full TSA gain, whereas shrinking only the much larger complement produces a substantially smaller improvement\. This asymmetry persists when all tensors are trained with AdamW \([Figures9](https://arxiv.org/html/2609.25482#A2.F9)and[5](https://arxiv.org/html/2609.25482#A2.T5)\)\.

Figure 9:The initial two\-group split localizes most averaging gain to the matrix parameter group\.Points show paired BPB improvement over the raw final iterate when only one predeclared parameter group is shrunk\. The same coordinate partition is used under both optimizers\. Bars are 95%tt\-intervals across five confirmation streams\.Table 5:Where the averaging gain is localized\.“Matrix group” is the 29\.67% of parameters assigned to Muon in the native optimizer; the identical parameter partition is retained in pure\-AdamW runs\. All values are paired BPB improvements over the raw final iterate\.Because the same pattern survives pure AdamW, it cannot be attributed solely to Muon updates\. The optimizer\-aligned matrix/complement partition is therefore best viewed as a useful first proxy for parameter\-role or trajectory heterogeneity rather than as a causal Muon\-versus\-AdamW explanation\. It establishes two facts that motivate the extension below: hidden matrices prefer substantially stronger shrinkage than the complement, while the corresponding two\-parameter estimator adds only a small amount beyond scalar TSA\.

### B\.6Extending structured TSA to parameter\-role groupings

The optimizer\-aligned matrix/complement partition in[SectionB\.5](https://arxiv.org/html/2609.25482#A2.SS5)provides a useful first localization: the hidden matrix block prefers substantially stronger shrinkage than the complement, and this asymmetry remains when all parameters are trained with AdamW\. At the same time, the corresponding two\-parameter estimator improves only modestly over scalar TSA\. We therefore retain that split as an explicit baseline and ask whether a broader, but still low\-dimensional family of parameter\-role partitions explains additional out\-of\-sample gain\.

Equation \([21](https://arxiv.org/html/2609.25482#A2.E21)\) is theG=2G=2instance of a general groupwise estimator\. Let\{Pg\}g=1G\\\{P\_\{g\}\\\}\_\{g=1\}^\{G\}be disjoint coordinate projections satisfying∑g=1GPg=I\\sum\_\{g=1\}^\{G\}P\_\{g\}=I\. We define

θ^T​\(𝜶\)=θT\+∑g=1Gαg​Pg​\(θ¯T−θT\),𝜶∈\[0,1\]G\.\\widehat\{\\theta\}\_\{T\}\(\\bm\{\\alpha\}\)=\\theta\_\{T\}\+\\sum\_\{g=1\}^\{G\}\\alpha\_\{g\}P\_\{g\}\(\\bar\{\\theta\}\_\{T\}\-\\theta\_\{T\}\),\\qquad\\bm\{\\alpha\}\\in\[0,1\]^\{G\}\.\(22\)Scalar TSA is the restricted familyα1=⋯=αG\\alpha\_\{1\}=\\cdots=\\alpha\_\{G\}; the matrix/complement rule above is one structured member of this larger family\.

#### Development and confirmation protocol\.

All trajectories in this study use the10%10\\%terminal floor and the main checkpoint windowK=8,s=32K=8,s=32\. On eight native Muon\+AdamW development trajectories, we evaluated 290 predeclared rules comprising a scalar baseline and 12 structured families\. These included the earlier hidden\-matrices\-versus\-rest partition, several functional and geometric subdivisions of the hidden block, trajectory\-statistic rules, and partitions that separately expose embedding, unembedding, and remaining parameters\. Rules were compared on a 32\-batch calibration block\. We then froze the scalar baseline, the winner from each structured family, and the overall structured winner before evaluating any new trajectory\.

Confirmation used eight entirely fresh native trajectories and eight paired pure\-AdamW trajectories, each evaluated on the untouched 96\-batch holdout\. The same parameter partitions and coefficients were transferred to pure AdamW without retuning\. Unless noted otherwise, intervals in this subsection are 95% Student\-ttintervals across the eight trajectories\.

Figure 10:Development search over parameter\-group structures\.Each point is the best candidate within one predeclared structured family, measured relative to the development\-selected scalar TSA rule; bars are 95%tt\-intervals across eight development trajectories\. Because these data were used for selection, the estimates are descriptive rather than confirmatory\. The “Hidden matrices vs\. rest” row is the extension’s direct version of the initial optimizer\-aligned split in[SectionB\.5](https://arxiv.org/html/2609.25482#A2.SS5); the embedding/hidden/unembedding/rest family produced the overall development winner\.The selected rule separates embedding parameters, hidden transformation matrices, the unembedding matrix, and the remaining scalar/other parameters:

\(αemb,αhidden,αunemb,αrest\)=\(0,0\.65,0\.45,0\.25\)\.\(\\alpha\_\{\\mathrm\{emb\}\},\\alpha\_\{\\mathrm\{hidden\}\},\\alpha\_\{\\mathrm\{unemb\}\},\\alpha\_\{\\mathrm\{rest\}\}\)=\(0,0\.65,0\.45,0\.25\)\.\(23\)The initial localization is therefore preserved rather than overturned: hidden matrices still receive strong shrinkage\. The additional gain comes primarily from refining the former complement, especially by separating the input embedding from the unembedding matrix and the remaining parameters\.

Because the final checkpoint is included in theKK\-checkpoint average, its total weight in groupggiswT,g=1−\(K−1\)​αg/Kw\_\{T,g\}=1\-\(K\-1\)\\alpha\_\{g\}/K\. ForK=8K=8, the frozen rule therefore has the interpretation shown in[Table6](https://arxiv.org/html/2609.25482#A2.T6)\.

Table 6:Frozen parameter\-role shrinkage rule\.The final column is the total weight placed on the final checkpoint after accounting for its inclusion in the checkpoint average\.
#### Fresh confirmation\.

On the eight new native trajectories, the frozen role\-based rule improved over the suite\-selected scalar baselineα=0\.50\\alpha=0\.50by0\.000153±0\.0000510\.000153\\pm 0\.000051BPB, with 95% interval\[0\.000103,0\.000204\]\[0\.000103,0\.000204\]\. It also improved directly over the paper’s established scalar settingα=0\.55\\alpha=0\.55by0\.000138±0\.0000450\.000138\\pm 0\.000045BPB and over the earlier hidden\-matrix/rest rule by0\.000108±0\.0000130\.000108\\pm 0\.000013BPB\. Relative to the raw final iterate of the same10%10\\%\-floor trajectories, the selected groupwise estimator gained0\.002316±0\.0000930\.002316\\pm 0\.000093BPB\. Thus scalar TSA supplies most of the total averaging benefit, while parameter\-role conditioning provides a smaller but resolved additional gain\.

The same frozen rule improved over scalarα=0\.50\\alpha=0\.50by0\.000178±0\.0000230\.000178\\pm 0\.000023BPB under pure AdamW and improved over the hidden/rest rule by0\.000130±0\.0000130\.000130\\pm 0\.000013BPB\. This transfer reinforces the conclusion of the initial localization experiment: the signal cannot be explained solely by Muon updates\. It is instead consistent with heterogeneity associated with parameter role, tensor geometry, group\-specific learning\-rate scale, or the resulting terminal dynamics\. Pure AdamW does not distinguish among those mechanisms because the parameter groups retain their different base learning rates\.

Figure 11:Frozen groupwise rules on fresh confirmation trajectories\.Points show the gain of each development\-selected family winner over the frozen scalarα=0\.50\\alpha=0\.50baseline\. Circles use eight fresh native Muon\+AdamW trajectories; squares transfer the identical rules to eight paired pure\-AdamW trajectories without retuning\. Bars are 95%tt\-intervals\. The “Hidden matrices vs\. rest” row carries forward the initial two\-group experiment, while the role\-based embedding/hidden/unembedding/rest rule provides a further confirmed gain in both optimizer regimes\.Several alternative partitions also carried smaller confirmation signals, including attention\-versus\-MLP, square\-versus\-rectangular, and four\-way hidden\-matrix role splits\. Their gains were substantially below that of the selected role\-based split\. Consistently, the measured trajectory scores of QKV, attention\-output, MLP\-up, and MLP\-down matrices were tightly clustered between0\.7080\.708and0\.7150\.715, providing little support for a large functional separation within the hidden matrix block\. The combined evidence therefore supports a progression rather than a replacement: the initial optimizer\-aligned split correctly localized most averaging value to hidden matrices, and the broader search subsequently identified additional heterogeneity within the complement\. These experiments do not establish that the selected grouping is universally optimal\.

### B\.7Groupwise terminal activity as a mechanism probe

The averaging experiments raise a separate question: if parameter groups prefer different returned estimators, should they also receive different amounts of terminal optimizer activity? We tested this as a fixed mechanism probe rather than a new hyperparameter search\. The estimator in[Equation23](https://arxiv.org/html/2609.25482#A2.E23)remained frozen, and its coefficients were mapped monotonically to terminal learning\-rate floors via

ρg=0\.05\+0\.10​αg0\.65\.\\rho\_\{g\}=0\.05\+0\.10\\frac\{\\alpha\_\{g\}\}\{0\.65\}\.\(24\)This gives floors of5\.0%5\.0\\%for embedding parameters,15\.0%15\.0\\%for hidden matrices,11\.92%11\.92\\%for the unembedding matrix, and8\.85%8\.85\\%for the remaining parameters\. The schedule implementation directly routed each optimizer parameter group to one of these four buckets\.

We trained eight new paired native trajectories under four fixed schedules: uniform5%5\\%, uniform10%10\\%, the mapped groupwise schedule, and a uniform floor equal to the parameter\-count\-weighted mean of the mapped floors \(8\.576%8\.576\\%\)\. Within each stream, all arms shared initialization, data order, checkpoint window, and the 96\-batch holdout\. For every trajectory, we evaluated the raw final iterate, scalar TSA atα=0\.55\\alpha=0\.55, and the frozen role\-based estimator\. No schedule or estimator was selected using these trajectories\.

For an estimatorEE, define the mapped\-schedule effect relative to the parameter\-count\-matched uniform schedule asΔE=LE​\(mapped\)−LE​\(matched\)\\Delta\_\{E\}=L\_\{E\}\(\\mathrm\{mapped\}\)\-L\_\{E\}\(\\mathrm\{matched\}\), so positive values mean that the mapped schedule is worse\. The corresponding interaction is

ℐgroup=Δraw−Δgroupwise\.\\mathcal\{I\}\_\{\\mathrm\{group\}\}=\\Delta\_\{\\mathrm\{raw\}\}\-\\Delta\_\{\\mathrm\{groupwise\}\}\.\(25\)
Table 7:Predeclared groupwise terminal\-floor intervention\.Schedule effects are mapped\-floor BPB minus the parameter\-count\-matched uniform\-floor BPB; positive values indicate worse BPB\. The interaction is the effect for the raw final iterate minus that for groupwise TSA\. Intervals are 95%tt\-intervals over eight paired trajectories\.The tested heterogeneous schedule therefore did not establish a better final training recipe: under the frozen groupwise estimator, its difference from the matched uniform schedule was small and unresolved\. Its effect nevertheless depended strongly on the returned estimator\. The mapped allocation made the raw final iterate worse by0\.0034040\.003404BPB, while the groupwise estimator removed essentially all of that relative penalty, yielding the precisely estimated interaction in[Table7](https://arxiv.org/html/2609.25482#A2.T7)\.

Parameter\-specific averaging also became more valuable on the heterogeneous trajectory\. Within the mapped arm, frozen groupwise TSA improved over the raw final iterate by0\.005114±0\.0001740\.005114\\pm 0\.000174BPB and over scalar TSA atα=0\.55\\alpha=0\.55by0\.000322±0\.0000610\.000322\\pm 0\.000061BPB\. The complete mapped\-schedule plus groupwise\-TSA output improved over the paired uniform\-5%5\\%raw final iterate by0\.000991±0\.0003750\.000991\\pm 0\.000375BPB\. As an ancillary check, these eight new streams independently reproduced the ordinary uniform\-5%5\\%to uniform\-10%10\\%schedule–estimator interaction at0\.001474±0\.0000860\.001474\\pm 0\.000086BPB, close to the main confirmation estimate\.

The parameter\-count\-matched control equalizes the mean floor across parameter coordinates, not optimizer\-update energy, displacement, or noise injection; the corresponding optimizer groups have different base learning rates and update geometries\. We therefore interpret this experiment as evidence that the value of parameter\-specific output estimation depends on how terminal activity is distributed across parameter groups, not as evidence that the particular mapped cooldown should replace a uniform schedule\.

### B\.8Full preliminary single\-trajectory sweeps

Before the replicated confirmation experiments, we performed two denser diagnostic sweeps on individual depth\-12 trajectories\. The first holds the training trajectory fixed and varies the scalar TSA coefficientα\\alpha\. The second returns the raw final iterate and varies only the terminal learning\-rate floor\. These experiments provide higher\-resolution views of the response surfaces, but their uncertainty reflects evaluation batches within a single trained trajectory rather than variation across independent training runs\. We therefore treat them as descriptive preliminary evidence; the corresponding replicated main\-text experiments provide the stronger inferential results\.

#### Scalar shrinkage sweep\.

Holding the baseline5%5\\%\-floor trajectory fixed, partial shrinkage produces a broad interior optimum, while full LAWA eventually becomes worse than the raw final iterate\. This preliminary response motivated the replicated estimator sweep in[Figure2](https://arxiv.org/html/2609.25482#S4.F2)\.

Table 8:Preliminary fixed\-trajectory scalar sweep\.Improvement is BPB for the raw final iterate minus estimator BPB\. Intervals quantify paired evaluation\-batch uncertainty for one trajectory\.
#### Terminal\-floor sweep for the raw final iterate\.

We separately vary only the terminal learning\-rate floor while returning the raw final iterate\. The raw final iterate worsens monotonically over the tested range, and the degradation becomes steeper at the larger floors \([Figure12](https://arxiv.org/html/2609.25482#A2.F12)\)\. This shape is consistent with the replicated5%5\\%,10%10\\%, and15%15\\%comparison in[Figure3](https://arxiv.org/html/2609.25482#S4.F3)\.

Figure 12:Preliminary response of the raw final iterate across terminal floors\.Points show BPB differences for the raw final iterate relative to the5%5\\%floor along one training trajectory; bars quantify 95% paired evaluation\-batch uncertainty\. The dashed curve is a least\-squares quadratic guide, constrained to pass through the5%5\\%baseline \(R2=0\.982R^\{2\}=0\.982\)\. The fit is descriptive only: unlike the exact quadratic dependence of the local TSA surrogate on shrinkage coefficientα\\alpha, the theory does not imply that BPB for the raw final iterate must be quadratic in the terminal floor\.Table 9:Single\-trajectory floor sweep without averaging\.Differences are relative to the5%5\\%floor\. Intervals quantify paired evaluation\-batch uncertainty for one trajectory\.

### B\.9Replicated quadratic response

Under the local quadratic surrogate, the TSA risk is exactly quadratic in the shrinkage coefficientα\\alphafor a fixed trajectory\. We therefore ask whether the observed validation response has the corresponding shape\. For each optimizer and terminal floor, we fit

g⁡\(α\)=b​α\+c​α2,g⁡\(0\)=0,g\(\\alpha\)=b\\alpha\+c\\alpha^\{2\},\\qquad g\(0\)=0,\(26\)to the mean paired BPB improvement over the raw final iterate across the five confirmation trajectories\. Becauseggdenotes improvement rather than risk, a beneficial interior optimum corresponds to a concave response\. The fitted vertex isα^⋆=−b/\(2c\)\\widehat\{\\alpha\}^\{\\star\}=\-b/\(2c\)\.

Table 10:Quadratic fits to the replicated shrinkage responses\.For each optimizer and terminal floor, the best tested coefficient on the post\-freeze scalar grid is compared with the vertex of the fitted quadratic\. Fits use the mean paired BPB improvement over the raw final iterate across five confirmation trajectories\.Figure 13:The replicated shrinkage responses are nearly quadratic\.Each panel shows mean paired BPB improvement over the raw final iterate across five confirmation trajectories; bars are 95%tt\-intervals across trajectories\. Solid curves are fits ofg⁡\(α\)=b​α\+c​α2g\(\\alpha\)=b\\alpha\+c\\alpha^\{2\}constrained byg⁡\(0\)=0g\(0\)=0\. Stars mark the best tested scalar\-grid coefficients, and dashed vertical lines mark the fitted vertices\.The quadratic description is accurate across all six optimizer\-floor conditions, withR2≥0\.99950R^\{2\}\\geq 0\.99950\. The fitted vertices closely track the independently observed grid optima\. Under Muon\+AdamW, the fitted optimum moves from0\.4190\.419to0\.5470\.547to0\.7010\.701as the terminal floor rises from 5% to 10% to 15%; under pure AdamW it moves from0\.3300\.330to0\.4990\.499to0\.6680\.668\. Thus the near\-quadratic response is not specific to one trajectory, floor, or optimizer regime, and the preferred amount of shrinkage systematically increases with terminal activity\. This provides an empirical consistency check for the local quadratic description and its comparative\-static interpretation; it does not imply that global neural\-network training dynamics are quadratic\.

## Appendix CProof of the TSA risk theorem

###### Proof of[Theorem1](https://arxiv.org/html/2609.25482#Thmtheorem1)\.

By[Equation9](https://arxiv.org/html/2609.25482#S3.E9),

θ^T​\(α\)−θ⋆=m−α​A\+ϵ0\+α​q\.\\widehat\{\\theta\}\_\{T\}\(\\alpha\)\-\\theta^\{\\star\}=m\-\\alpha A\+\\epsilon\_\{0\}\+\\alpha q\.The random componentϵ0\+α​q\\epsilon\_\{0\}\+\\alpha qhas mean zero\. Therefore,

2​ℛ​\(α\)\\displaystyle 2\\mathcal\{R\}\(\\alpha\)=𝔼​‖m−α​A\+ϵ0\+α​q‖H2\\displaystyle=\\mathbb\{E\}\\left\\\|m\-\\alpha A\+\\epsilon\_\{0\}\+\\alpha q\\right\\\|\_\{H\}^\{2\}\(27\)=‖m−α​A‖H2\+𝔼​‖ϵ0\+α​q‖H2\\displaystyle=\\\|m\-\\alpha A\\\|\_\{H\}^\{2\}\+\\mathbb\{E\}\\\|\\epsilon\_\{0\}\+\\alpha q\\\|\_\{H\}^\{2\}\(28\)=‖m−α​A‖H2\+𝔼​‖ϵ0‖H2\+2​α​C\+α2​Vq\.\\displaystyle=\\\|m\-\\alpha A\\\|\_\{H\}^\{2\}\+\\mathbb\{E\}\\\|\\epsilon\_\{0\}\\\|\_\{H\}^\{2\}\+2\\alpha C\+\\alpha^\{2\}V\_\{q\}\.\(29\)Expanding the deterministic term,

‖m−α​A‖H2=‖m‖H2−2​α​⟨m,A⟩H\+α2​‖A‖H2\.\\\|m\-\\alpha A\\\|\_\{H\}^\{2\}=\\\|m\\\|\_\{H\}^\{2\}\-2\\alpha\\langle m,A\\rangle\_\{H\}\+\\alpha^\{2\}\\\|A\\\|\_\{H\}^\{2\}\.Since

ℛ⁡\(0\)=12​\(‖m‖H2\+𝔼​‖ϵ0‖H2\),\\mathcal\{R\}\(0\)=\\frac\{1\}\{2\}\\left\(\\\|m\\\|\_\{H\}^\{2\}\+\\mathbb\{E\}\\\|\\epsilon\_\{0\}\\\|\_\{H\}^\{2\}\\right\),dividing by two yields

ℛ⁡\(α\)=ℛ⁡\(0\)\+α⁡\(−⟨m,A⟩H\+C\)\+α22​\(‖A‖H2\+Vq\)\.\\mathcal\{R\}\(\\alpha\)=\\mathcal\{R\}\(0\)\+\\alpha\\bigl\(\-\\langle m,A\\rangle\_\{H\}\+C\\bigr\)\+\\frac\{\\alpha^\{2\}\}\{2\}\\bigl\(\\\|A\\\|\_\{H\}^\{2\}\+V\_\{q\}\\bigr\)\.Differentiating gives

ℛ′​\(α\)=−⟨m,A⟩H\+C\+α⁡\(‖A‖H2\+Vq\)\.\\mathcal\{R\}^\{\\prime\}\(\\alpha\)=\-\\langle m,A\\rangle\_\{H\}\+C\+\\alpha\\bigl\(\\\|A\\\|\_\{H\}^\{2\}\+V\_\{q\}\\bigr\)\.When‖A‖H2\+Vq\>0\\\|A\\\|\_\{H\}^\{2\}\+V\_\{q\}\>0, the unconstrained minimizer is

αunc=⟨m,A⟩H−C‖A‖H2\+Vq\.\\alpha\_\{\\mathrm\{unc\}\}=\\frac\{\\langle m,A\\rangle\_\{H\}\-C\}\{\\\|A\\\|\_\{H\}^\{2\}\+V\_\{q\}\}\.Projection onto\[0,1\]\[0,1\]gives[Equation13](https://arxiv.org/html/2609.25482#S3.E13)\. The optimum is strictly interior exactly when

0<⟨m,A⟩H−C<‖A‖H2\+Vq\.0<\\langle m,A\\rangle\_\{H\}\-C<\\\|A\\\|\_\{H\}^\{2\}\+V\_\{q\}\.∎

## Appendix DOne\-step quadratic analysis

Consider the one\-dimensional quadratic loss

L⁡\(x\)=L⁡\(0\)\+h2​x2,h\>0\.L\(x\)=L\(0\)\+\\frac\{h\}\{2\}x^\{2\},\\qquad h\>0\.\(30\)Letx−≠0x^\{\-\}\\neq 0be a deterministic parameter value immediately before one terminal stochastic\-gradient update, and let

x\+=\(1−η​h\)​x−−η​ξ,𝔼⁡\[ξ\]=0,𝔼⁡\[ξ2\]=σ2\.x^\{\+\}=\(1\-\\eta h\)x^\{\-\}\-\\eta\\xi,\\qquad\\mathbb\{E\}\[\\xi\]=0,\\qquad\\mathbb\{E\}\[\\xi^\{2\}\]=\\sigma^\{2\}\.\(31\)Using the two checkpointsx−x^\{\-\}andx\+x^\{\+\}, TSA returns

x^α\\displaystyle\\widehat\{x\}\_\{\\alpha\}=\(1−α\)​x\+\+α2​\(x−\+x\+\)\\displaystyle=\(1\-\\alpha\)x^\{\+\}\+\\frac\{\\alpha\}\{2\}\\left\(x^\{\-\}\+x^\{\+\}\\right\)\(32\)=\(1−α2\)x\+\+α2x−,α∈\[0,1\]\.\\displaystyle=\\left\(1\-\\frac\{\\alpha\}\{2\}\\right\)x^\{\+\}\+\\frac\{\\alpha\}\{2\}x^\{\-\},\\qquad\\alpha\\in\[0,1\]\.Defineaα:=1−α/2a\_\{\\alpha\}:=1\-\\alpha/2\. Substitution gives

x^α=\(1−aα​η​h\)​x−−aα​η​ξ\.\\widehat\{x\}\_\{\\alpha\}=\\left\(1\-a\_\{\\alpha\}\\eta h\\right\)x^\{\-\}\-a\_\{\\alpha\}\\eta\\xi\.\(33\)Consequently, the expected excess quadratic risk is exactly

Rα​\(η\):=𝔼⁡\[L⁡\(x^α\)−L⁡\(0\)\]=h2​\[\(1−aα​η​h\)2​\(x−\)2\+aα2​η2​σ2\]\.R\_\{\\alpha\}\(\\eta\):=\\mathbb\{E\}\\\!\\left\[L\(\\widehat\{x\}\_\{\\alpha\}\)\-L\(0\)\\right\]=\\frac\{h\}\{2\}\\left\[\\left\(1\-a\_\{\\alpha\}\\eta h\\right\)^\{2\}\(x^\{\-\}\)^\{2\}\+a\_\{\\alpha\}^\{2\}\\eta^\{2\}\\sigma^\{2\}\\right\]\.\(34\)
#### Corollary\.

Supposeη\\etais selected from\[0,ηmax\]\[0,\\eta\_\{\\max\}\]\. Then

ηα⋆=Π\[0,ηmax\]​\(h​\(x−\)2aα​\(h2​\(x−\)2\+σ2\)\)\.\\eta\_\{\\alpha\}^\{\\star\}=\\Pi\_\{\[0,\\eta\_\{\\max\}\]\}\\left\(\\frac\{h\(x^\{\-\}\)^\{2\}\}\{a\_\{\\alpha\}\\left\(h^\{2\}\(x^\{\-\}\)^\{2\}\+\\sigma^\{2\}\\right\)\}\\right\)\.\(35\)If the optima for the raw final iterate and TSA are not clipped byηmax\\eta\_\{\\max\}, then

ηα⋆=η0⋆1−α/2\.\\eta\_\{\\alpha\}^\{\\star\}=\\frac\{\\eta\_\{0\}^\{\\star\}\}\{1\-\\alpha/2\}\.\(36\)In particular,ηα⋆\>η0⋆\\eta\_\{\\alpha\}^\{\\star\}\>\\eta\_\{0\}^\{\\star\}for everyα∈\(0,1\]\\alpha\\in\(0,1\], andη1⋆=2​η0⋆\\eta\_\{1\}^\{\\star\}=2\\eta\_\{0\}^\{\\star\}\.

#### Proof\.

Differentiating \([34](https://arxiv.org/html/2609.25482#A4.E34)\) yields

Rα′​\(η\)=h​aα​\[−h​\(x−\)2\+aα​η​\(h2​\(x−\)2\+σ2\)\],R\_\{\\alpha\}^\{\\prime\}\(\\eta\)=ha\_\{\\alpha\}\\left\[\-h\(x^\{\-\}\)^\{2\}\+a\_\{\\alpha\}\\eta\\left\(h^\{2\}\(x^\{\-\}\)^\{2\}\+\\sigma^\{2\}\\right\)\\right\],\(37\)while

Rα′′​\(η\)=h​aα2​\(h2​\(x−\)2\+σ2\)\>0\.R\_\{\\alpha\}^\{\\prime\\prime\}\(\\eta\)=ha\_\{\\alpha\}^\{2\}\\left\(h^\{2\}\(x^\{\-\}\)^\{2\}\+\\sigma^\{2\}\\right\)\>0\.\(38\)The unconstrained minimizer is therefore

ηα,unc⋆=h​\(x−\)2aα​\(h2​\(x−\)2\+σ2\)\.\\eta\_\{\\alpha,\\mathrm\{unc\}\}^\{\\star\}=\\frac\{h\(x^\{\-\}\)^\{2\}\}\{a\_\{\\alpha\}\\left\(h^\{2\}\(x^\{\-\}\)^\{2\}\+\\sigma^\{2\}\\right\)\}\.\(39\)Projection onto\[0,ηmax\]\[0,\\eta\_\{\\max\}\]gives \([35](https://arxiv.org/html/2609.25482#A4.E35)\)\. Sincea0=1a\_\{0\}=1andaα=1−α/2a\_\{\\alpha\}=1\-\\alpha/2, the interior ratio follows\.□\\square

#### Remark\.

In this minimal one\-update model,aα​ηα⋆a\_\{\\alpha\}\\eta\_\{\\alpha\}^\{\\star\}is constant inα\\alpha, so the jointly optimized minimum risk is also constant\. The result therefore isolates a change in the preferred schedule rather than claiming an automatic improvement in the best attainable risk\. Multi\-step trajectories, longer checkpoint windows, and nonstationary late\-training dynamics can break this exact rescaling equivalence\.

## Appendix ELLM disclosure

We used LLMs in three ways\. First, essentially all experiments were set up and implemented with LLM coding, followed by a separate LLM review of code correctness\. Second, we asked LLMs to critique manuscript drafts and incorporated suggestions when we judged them appropriate\. Third, we used LLMs to find related papers; we read each paper before citing it to confirm its relation to our work\.