Gradient Smoothing: Coupling Layer-wise Updates for Improved Optimization

# Coupling Layer-wise Updates for Improved Optimization
Source: [https://arxiv.org/html/2606.30813](https://arxiv.org/html/2606.30813)
## Gradient Smoothing: Coupling Layer\-wise Updates for Improved Optimization

###### Abstract

Deep neural networks with repeated architectural blocks, such as transformers, often exhibit structured relationships across layers that emerge during training\. Motivated by this observation, we introduce*Depth\-wise Gradient Augmentation*, a general optimization paradigm in which the update applied to each layer is obtained by transforming the collection of block\-wise optimizer updates along the depth dimension\. Within this framework, we study*Gradient Smoothing*, a family of depth\-wise smoothing methods, and instantiate it with a simple local*Window Smoothing*operator\. The resulting method operates directly on block\-wise updates produced by arbitrary base optimizers \(e\.g\., SGD, Adam, Muon\), incurs minimal computational overhead, and is compatible with existing optimization pipelines\. We evaluate Gradient Smoothing across a diverse set of architectures and training regimes, including language model pretraining, RL post\-training of LLMs for reasoning, diffusion modeling, and image classification with Vision Transformers\. Across these settings, Gradient Smoothing consistently improves optimization and generalization performance without modifying model architectures or training objectives\. We further show that it promotes more structured representation evolution across depth, consistent with its interpretation as a structured depth\-wise preconditioning method\. Together, these results establish Depth\-wise Gradient Augmentation as a promising framework for exploiting cross\-depth structure in optimization and demonstrate Gradient Smoothing as a simple and broadly applicable instantiation\.

Machine Learning, ICML

## 1Introduction

Many modern deep neural networks are composed of repeated architectural blocks with shared computational structure, such as Residual Networks \(ResNets\)\(Heet al\.,[2015](https://arxiv.org/html/2606.30813#bib.bib31)\)and transformers\(Vaswaniet al\.,[2023](https://arxiv.org/html/2606.30813#bib.bib33)\)\. Recent work has shown that networks of this form often exhibit strong structural regularities across depth that emerge during training\. In particular, studies on*Transformer Block Coupling*\(Aubryet al\.,[2025](https://arxiv.org/html/2606.30813#bib.bib29)\)and*Residual Alignment*\(Li and Papyan,[2024](https://arxiv.org/html/2606.30813#bib.bib30)\)demonstrate that singular vectors of block Jacobians and residual representations become aligned across layers, suggesting a form of implicit coordination throughout depth\. Related evidence appears in studies of layer similarity and pruning in large language models\(Gromovet al\.,[2025](https://arxiv.org/html/2606.30813#bib.bib15)\), as well as analyses showing smooth evolution of information\-theoretic quantities across depth\(Skeanet al\.,[2024](https://arxiv.org/html/2606.30813#bib.bib186)\)\.

These observations suggest that repeated blocks in deep networks do not evolve independently, but instead exhibit coordinated representational and dynamical structure across depth\. However, standard optimization methods typically construct and apply block\-wise optimizer updates independently \(aside from their dependencies through the forward and backward passes\), without explicitly leveraging relationships between layers\. This naturally raises the question:*Can optimizer updates themselves be augmented to exploit the structure that emerges across depth?*

In this work, we introduce a general paradigm of*Depth\-wise Gradient Augmentation*\. Rather than applying the optimizer update produced for each block in isolation, the update applied to each layer is obtained by applying a depth\-wise augmentation operator to the collection of block\-wise optimizer updates\. This perspective encompasses a broad class of update transformations acting along the depth dimension\.

As a first instantiation of this framework, we study*Gradient Smoothing*, motivated by the empirical observation that neighboring layers often learn related representations and exhibit increasingly aligned optimization dynamics throughout training\. We focus on a particularly simple realization,*Window Smoothing*, which augments each block update by averaging optimizer updates from a local window of neighboring layers\. Window Smoothing is simple to implement, incurs minimal computational overhead, and is compatible with arbitrary base optimizers\. Moreover, Gradient Smoothing admits a natural interpretation as a structured preconditioning method acting on block\-structured parameter spaces, while providing a foundation for exploring richer forms of depth\-wise update augmentation\.

#### Contributions\.

1. 1\.We introduce*Depth\-wise Gradient Augmentation*, a general optimization paradigm in which the collection of block\-wise optimizer updates is transformed by a depth\-wise augmentation operator before being applied\. Within this framework, we study*Gradient Smoothing*, a family of structured depth\-wise smoothing methods, and instantiate it with a simple, optimizer\-agnostic*Window Smoothing*operator, providing a foundation for exploring more general forms of depth\-wise update augmentation\.
2. 2\.We demonstrate that Gradient Smoothing consistently improves optimization and generalization across a diverse set of architectures and training regimes, including language model pretraining, RL fine\-tuning of LLMs for reasoning, supervised learning with Vision Transformers, and diffusion models\.
3. 3\.We provide both empirical and theoretical evidence that Gradient Smoothing promotes more structured representation evolution across depth\. Empirically, we observe increased layer\-wise trajectory alignment and representation residual similarity, while theoretically we characterize how smoothing influences representational residual structure\.

## 2Background and Setup

### 2\.1Block\-Structured Optimization

We consider neural networks composed ofLLrepeated blocks connected by skip connections, where each block shares the same \(or similar\) functional form but has its own parameters\. In addition, the model may include a \(typically small\) collection of parameters not associated with the repeated structure, such as input embeddings or output heads\.

Formally, letF:ℝd×ℝp→ℝdF:\\mathbb\{R\}^\{d\}\\times\\mathbb\{R\}^\{p\}\\to\\mathbb\{R\}^\{d\}denote the residual mapping associated with a single block\. For blockℓ\\ell, the residual function is given byF​\(⋅;θℓ\)F\(\\cdot;\\theta\_\{\\ell\}\), whereθℓ∈ℝp\\theta\_\{\\ell\}\\in\\mathbb\{R\}^\{p\}are layer\-specific parameters\. Given the initial embeddingh0h\_\{0\}, the hidden states\(hℓ\)ℓ=1L⊂ℝd\(h\_\{\\ell\}\)\_\{\\ell=1\}^\{L\}\\subset\\mathbb\{R\}^\{d\}then evolve according to

hℓ\+1=hℓ\+F​\(hℓ;θℓ\),ℓ=1,…,L−1\.h\_\{\\ell\+1\}=h\_\{\\ell\}\+F\(h\_\{\\ell\};\\theta\_\{\\ell\}\),\\qquad\\ell=1,\\dots,L\-1\.\(1\)Although we present the formulation for standard residual connections, the framework naturally extends to architectures with more general skip connection patterns\. We write the full parameter vector as

\(θ,ϕ\),θ=\(θ1,…,θL\),θl∈ℝp,\(\\theta,\\phi\),\\qquad\\theta=\(\\theta\_\{1\},\\dots,\\theta\_\{L\}\),\\quad\\theta\_\{l\}\\in\\mathbb\{R\}^\{p\},whereϕ∈ℝDϕ\\phi\\in\\mathbb\{R\}^\{D\_\{\\phi\}\}collects all non\-block parameters\. This setting encompasses a broad class of modern architectures, including residual networks with repeated block structure, Vision Transformers, and Transformer\-based language models \(by letting eachhℓh\_\{\\ell\}be a sequence of vectors\)\. We denote the training objective byℒ​\(θ,ϕ\)\\mathcal\{L\}\(\\theta,\\phi\)and assume it is differentiable\.

At optimization steptt, the block\-wise gradients are given by

gl\(t\):=∇θlℒ​\(θ\(t\),ϕ\(t\)\),g\(t\):=\(g1\(t\),…,gL\(t\)\),g\_\{l\}^\{\(t\)\}:=\\nabla\_\{\\theta\_\{l\}\}\\mathcal\{L\}\(\\theta^\{\(t\)\},\\phi^\{\(t\)\}\),\\qquad g^\{\(t\)\}:=\(g\_\{1\}^\{\(t\)\},\\dots,g\_\{L\}^\{\(t\)\}\),and we denote bygϕ\(t\):=∇ϕℒ​\(θ\(t\),ϕ\(t\)\)g\_\{\\phi\}^\{\(t\)\}:=\\nabla\_\{\\phi\}\\mathcal\{L\}\(\\theta^\{\(t\)\},\\phi^\{\(t\)\}\)the gradient with respect toϕ\\phi\.

Standard first\-order methods treat the block\-wise gradients\{gl\(t\)\}l=1L\\\{g\_\{l\}^\{\(t\)\}\\\}\_\{l=1\}^\{L\}independently\. However, repeated\-block architectures induce strong structural relationships across depth: adjacent blocks often learn similar transformations and exhibit aligned representations, particularly after an initial transient phase of training\. This suggests that the stacked gradient vectorg\(t\)g^\{\(t\)\}possesses meaningful structure along the depth dimension\.

fLf\_\{L\}f3f\_\{3\}f2f\_\{2\}f1f\_\{1\}⋮\\vdots𝒰​\(∇θLL\)\\mathcal\{U\}\\\!\\left\(\\nabla\_\{\\theta\_\{L\}\}L\\right\)𝒰​\(∇θ3L\)\\mathcal\{U\}\\\!\\left\(\\nabla\_\{\\theta\_\{3\}\}L\\right\)𝒰​\(∇θ2L\)\\mathcal\{U\}\\\!\\left\(\\nabla\_\{\\theta\_\{2\}\}L\\right\)𝒰​\(∇θ1L\)\\mathcal\{U\}\\\!\\left\(\\nabla\_\{\\theta\_\{1\}\}L\\right\)⋮\\vdotsFilterSS𝒰~​\(∇θLL\)\\tilde\{\\mathcal\{U\}\}\\\!\\left\(\\nabla\_\{\\theta\_\{L\}\}L\\right\)𝒰~​\(∇θ3L\)\\tilde\{\\mathcal\{U\}\}\\\!\\left\(\\nabla\_\{\\theta\_\{3\}\}L\\right\)𝒰~​\(∇θ2L\)\\tilde\{\\mathcal\{U\}\}\\\!\\left\(\\nabla\_\{\\theta\_\{2\}\}L\\right\)𝒰~​\(∇θ1L\)\\tilde\{\\mathcal\{U\}\}\\\!\\left\(\\nabla\_\{\\theta\_\{1\}\}L\\right\)⋮\\vdots

Figure 1:Gradient Smoothing\.Representation of the gradient augmentation scheme applied across depth to the updates in a deep network after backpropagation\. For a deep network withLLidentical architectural blocks \(but with different parameters\), the gradient updates in each parameter blockθl\\theta^\{l\}are reweighted across depth to stabilize information propagation\.![Refer to caption](https://arxiv.org/html/2606.30813v1/plots/nanochat/baseline_vs_smoothing_metrics_d24_val_loss.png)\(a\)Depth2424: Validation Loss\.
![Refer to caption](https://arxiv.org/html/2606.30813v1/plots/nanochat/baseline_vs_smoothing_metrics_d24_val_bpb.png)\(b\)Depth2424: Validation BPB\.
![Refer to caption](https://arxiv.org/html/2606.30813v1/plots/nanochat/baseline_vs_smoothing_metrics_d24_core_metric.png)\(c\)Depth2424: CORE Metric\.
![Refer to caption](https://arxiv.org/html/2606.30813v1/plots/nanochat/baseline_vs_smoothing_metrics_d30_val_loss.png)\(d\)Depth3030: Validation Loss\.
![Refer to caption](https://arxiv.org/html/2606.30813v1/plots/nanochat/baseline_vs_smoothing_metrics_d30_val_bpb.png)\(e\)Depth3030: Validation BPB\.
![Refer to caption](https://arxiv.org/html/2606.30813v1/plots/nanochat/baseline_vs_smoothing_metrics_d30_core_metric.png)\(f\)Depth3030: CORE Metric\.

Figure 2:Nanochat pretraining with Gradient Smoothing\.Validation loss, BPB, and DCLM CORE metric duringnanochatpretraining of GPT2 with the default Adam \+ NorMuon optimizer setup \(Table[1](https://arxiv.org/html/2606.30813#S4.T1)\)\. We compare the baseline against a run with identical hyperparameters using depth\-wise update smoothing applied to the transformer blocks\. Models of depth 24 \(1\.38B parameters, top row\) and depth 30 \(2\.40B parameters bottom row\) withα=0\.1\\alpha=0\.1are replicated against the baseline forseed∈\{42,67,2026\}\\textbf\{seed\}\\in\\\{42,67,2026\\\}\. At both depths, smoothing accelerates convergence of validation loss and improves the CORE metric during training, with a larger gap visible at depth3030\.
### 2\.2General First\-Order Updates

Rather than committing to a specific optimizer, we consider a general first\-order update rule of the form

\(θ\(t\+1\),ϕ\(t\+1\)\)=\(θ\(t\),ϕ\(t\)\)−η​\(u\(t\),uϕ\(t\)\),\(\\theta^\{\(t\+1\)\},\\phi^\{\(t\+1\)\}\)=\(\\theta^\{\(t\)\},\\phi^\{\(t\)\}\)\-\\eta\\,\(u^\{\(t\)\},u\_\{\\phi\}^\{\(t\)\}\),where

u\(t\):=𝒰\(t\)​\(g\(t\)\),uϕ\(t\):=𝒰ϕ\(t\)​\(gϕ\(t\)\)u^\{\(t\)\}:=\\mathcal\{U\}^\{\(t\)\}\(g^\{\(t\)\}\),\\qquad u\_\{\\phi\}^\{\(t\)\}:=\\mathcal\{U\}\_\{\\phi\}^\{\(t\)\}\(g\_\{\\phi\}^\{\(t\)\}\)denote the updates produced by the base optimizer for the block and non\-block parameters, respectively\.

This formulation includes gradient descent \(u\(t\)=g\(t\)u^\{\(t\)\}=g^\{\(t\)\}\), momentum methods, adaptive optimizers such as Adam or AdamW, as well as Muon and related variants\. Importantly,u\(t\)u^\{\(t\)\}may depend on both the current gradient and internal optimizer state, and may involve nonlinear or time\-varying transformations ofg\(t\)g^\{\(t\)\}\.

This optimizer\-agnostic viewpoint allows us to study update augmentation directly at the level of optimizer updates, independent of the specific mechanisms used to construct them\. As a result, the framework can be incorporated as a lightweight, modular extension to a broad class of existing optimizers\.

## 3Depth\-wise Gradient Augmentation

### 3\.1General Framework

Networks built from repeated architectural blocks develop strong cross\-depth structure in their representations during training\. Motivated by this observation, we introduce*Depth\-wise Gradient Augmentation*, a general optimization paradigm in which the update applied to each layer is obtained by transforming the collection of block\-wise optimizer updates along the depth axis\. This perspective provides a unified framework for incorporating cross\-depth information into optimizer updates while remaining agnostic to the particular transformation used\.

At iterationtt, let

uℓ\(t\)∈ℝp,ℓ=1,…,L,u\_\{\\ell\}^\{\(t\)\}\\in\\mathbb\{R\}^\{p\},\\qquad\\ell=1,\\ldots,L,denote the block\-wise update produced by a base optimizer \(e\.g\., SGD, Adam, Muon\) for the parameters of layerℓ\\ell, and write

u\(t\)=\(u1\(t\),…,uL\(t\)\)∈\(ℝp\)L,u^\{\(t\)\}=\(u\_\{1\}^\{\(t\)\},\\dots,u\_\{L\}^\{\(t\)\}\)\\in\(\\mathbb\{R\}^\{p\}\)^\{L\},together with the corresponding updateuϕ\(t\)u\_\{\\phi\}^\{\(t\)\}for the non\-block parameters\.

In its most general form, Depth\-wise Gradient Augmentation applies an operator along the depth dimension independently for each within\-block coordinate\. For each coordinatei=1,…,pi=1,\\ldots,p, let

𝒞\(i\):ℝL→ℝL,\\mathcal\{C\}^\{\(i\)\}:\\mathbb\{R\}^\{L\}\\rightarrow\\mathbb\{R\}^\{L\},be a depth\-wise augmentation operator\. Then

\(u~1\(i\),…,u~L\(i\)\)=𝒞\(i\)​\(u1\(i\),…,uL\(i\)\)\.\\bigl\(\\tilde\{u\}\_\{1\}^\{\(i\)\},\\ldots,\\tilde\{u\}\_\{L\}^\{\(i\)\}\\bigr\)=\\mathcal\{C\}^\{\(i\)\}\\bigl\(u\_\{1\}^\{\(i\)\},\\ldots,u\_\{L\}^\{\(i\)\}\\bigr\)\.Equivalently,

u~ℓ\(i\)=𝒞ℓ\(i\)​\(u1\(i\),…,uL\(i\)\),ℓ=1,…,L\.\\tilde\{u\}\_\{\\ell\}^\{\(i\)\}=\\mathcal\{C\}^\{\(i\)\}\_\{\\ell\}\\bigl\(u\_\{1\}^\{\(i\)\},\\ldots,u\_\{L\}^\{\(i\)\}\\bigr\),\\qquad\\ell=1,\\ldots,L\.Thus, each coordinate of the update for layerℓ\\ellmay depend only on the corresponding coordinate across layers, while remaining independent of all other within\-block coordinates\.

Among the many possible choices of depth\-wise augmentation operators, we focus on operators that smooth optimizer updates across depth\.

### 3\.2Gradient Smoothing

A natural class of Depth\-wise Gradient Augmentation methods is obtained by choosing the coordinate\-wise depth augmentation operators to be a common linear smoothing operator\. We refer to this family of methods as*Gradient Smoothing*\. Specifically, we take

𝒞\(i\)=S,i=1,…,p,\\mathcal\{C\}^\{\(i\)\}=S,\\qquad i=1,\\ldots,p,whereS∈ℝL×LS\\in\\mathbb\{R\}^\{L\\times L\}is a linear operator acting on the layer index\. Equivalently, Gradient Smoothing is represented by the block\-structured operator

P:=S⊗I∈ℝDθ×Dθ,P:=S\\otimes I\\in\\mathbb\{R\}^\{D\_\{\\theta\}\\times D\_\{\\theta\}\},whereIIdenotes the identity operator on the within\-block parameter spaceℝp\\mathbb\{R\}^\{p\}\.

The transformed block\-wise update is then given by

u~\(t\):=P​u\(t\)=\(S⊗I\)​u\(t\)\.\\tilde\{u\}^\{\(t\)\}:=P\\,u^\{\(t\)\}=\(S\\otimes I\)\\,u^\{\(t\)\}\.
To describe the update on the full parameter vector\(θ,ϕ\)\(\\theta,\\phi\), we liftPPto the block\-diagonal operator

P¯:=\(P00Iϕ\)∈ℝ\(Dθ\+Dϕ\)×\(Dθ\+Dϕ\),\\bar\{P\}:=\\begin\{pmatrix\}P&0\\\\ 0&I\_\{\\phi\}\\end\{pmatrix\}\\in\\mathbb\{R\}^\{\(D\_\{\\theta\}\+D\_\{\\phi\}\)\\times\(D\_\{\\theta\}\+D\_\{\\phi\}\)\},whereIϕI\_\{\\phi\}denotes the identity operator on theϕ\\phi\-coordinates\.

Given the full optimizer update

ufull\(t\):=\(u\(t\),uϕ\(t\)\),u\_\{\\mathrm\{full\}\}^\{\(t\)\}:=\\bigl\(u^\{\(t\)\},\\,u\_\{\\phi\}^\{\(t\)\}\\bigr\),Gradient Smoothing appliesP¯\\bar\{P\}to obtain

u~full\(t\)=P¯​ufull\(t\)=\(P​u\(t\),uϕ\(t\)\)\.\\tilde\{u\}\_\{\\mathrm\{full\}\}^\{\(t\)\}=\\bar\{P\}\\,u\_\{\\mathrm\{full\}\}^\{\(t\)\}=\\bigl\(Pu^\{\(t\)\},\\,u\_\{\\phi\}^\{\(t\)\}\\bigr\)\.
The resulting optimization step becomes

θ\(t\+1\)=θ\(t\)−η​u~\(t\),ϕ\(t\+1\)=ϕ\(t\)−η​uϕ\(t\)\.\\theta^\{\(t\+1\)\}=\\theta^\{\(t\)\}\-\\eta\\,\\tilde\{u\}^\{\(t\)\},\\qquad\\phi^\{\(t\+1\)\}=\\phi^\{\(t\)\}\-\\eta\\,u\_\{\\phi\}^\{\(t\)\}\.
Gradient Smoothing thus corresponds to a structured depth\-wise coupling of optimizer updates, or equivalently, a structured preconditioning operator acting on the block\-wise updates\. Importantly, this construction is independent of the specific choice of base optimizer and operates purely at the level of update vectors\.

### 3\.3Smoothing Operator

In this work, we instantiate Gradient Smoothing with a simple local averaging operator, which we term*Window Smoothing*\. TakingSSto be a local window averaging operator across depth, the smoothed update is given by

u~l=\{\(1−α2\)​u1\+α2​u2,l=1,\(1−α\)​ul\+α2​\(ul−1\+ul\+1\),2≤l≤L−1,\(1−α2\)​uL\+α2​uL−1,l=L,\\tilde\{u\}\_\{l\}=\\begin\{cases\}\\big\(1\-\\tfrac\{\\alpha\}\{2\}\\big\)u\_\{1\}\+\\tfrac\{\\alpha\}\{2\}u\_\{2\},&l=1,\\\\\[4\.0pt\] \(1\-\\alpha\)u\_\{l\}\+\\tfrac\{\\alpha\}\{2\}\(u\_\{l\-1\}\+u\_\{l\+1\}\),&2\\leq l\\leq L\-1,\\\\\[4\.0pt\] \\big\(1\-\\tfrac\{\\alpha\}\{2\}\\big\)u\_\{L\}\+\\tfrac\{\\alpha\}\{2\}u\_\{L\-1\},&l=L,\\end\{cases\}\(2\)whereα∈\[0,1\)\\alpha\\in\[0,1\)controls the smoothing strength\. Each block’s update is replaced by a convex combination of its own update and those of its immediate neighbors in depth; the boundary cases atl=1l=1andl=Ll=Luse a half\-weight on the single available neighbor so thatSSremains row\-stochastic \(S​𝟏=𝟏S\\mathbf\{1\}=\\mathbf\{1\}\)\. The resultingS∈ℝL×LS\\in\\mathbb\{R\}^\{L\\times L\}is a symmetric tridiagonal matrix and introduces no additional trainable parameters\.

Many other depth\-wise weighting schemes are possible within this framework\. For instance, exponential smoothing, higher\-order tridiagonal kernels, and local window averaging with larger window sizes all fit the same general paradigm\. Our work focuses on local window averaging as a simple yet effective instantiation of Gradient Smoothing, applicable to many optimizers and across broad tasks\.

Algorithm 1Gradient Smoothing \(window, strengthα\\alpha\)Input:block weights

\{𝜽ℓ\}ℓ=1L\\\{\\boldsymbol\{\\theta\}\_\{\\ell\}\\\}\_\{\\ell=1\}^\{L\}, non\-block weights

ϕ\\boldsymbol\{\\phi\}, loss

LL, base optimizer

𝒰\\mathcal\{U\}, learning rate

η\\eta, weight decay

λ\\lambda, smoothing strength

α∈\[0,1\)\\alpha\\in\[0,1\)
for

t=1,2,…t=1,2,\\ldotsdo

𝐮ℓ←𝒰θ​\(∇𝜽ℓL\)\\mathbf\{u\}\_\{\\ell\}\\leftarrow\\mathcal\{U\}\_\{\\theta\}\\\!\\left\(\\nabla\_\{\\boldsymbol\{\\theta\}\_\{\\ell\}\}L\\right\)ℓ=1,…,L\\ell=1,\\ldots,L

𝐮ϕ←𝒰ϕ​\(∇ϕL\)\\mathbf\{u\}\_\{\\phi\}\\leftarrow\\mathcal\{U\}\_\{\\phi\}\\\!\\left\(\\nabla\_\{\\boldsymbol\{\\phi\}\}L\\right\)

𝐮ℓ←𝐮ℓ‖𝐮ℓ‖2\\mathbf\{u\}\_\{\\ell\}\\leftarrow\\dfrac\{\\mathbf\{u\}\_\{\\ell\}\}\{\\\|\\mathbf\{u\}\_\{\\ell\}\\\|\_\{2\}\}\\,optional normalization

𝐮~1←\(1−α2\)​𝐮1\+α2​𝐮2\\tilde\{\\mathbf\{u\}\}\_\{1\}\\leftarrow\\left\(1\-\\tfrac\{\\alpha\}\{2\}\\right\)\\mathbf\{u\}\_\{1\}\+\\tfrac\{\\alpha\}\{2\}\\,\\mathbf\{u\}\_\{2\}

𝐮~ℓ←\(1−α\)​𝐮ℓ\+α2​\(𝐮ℓ−1\+𝐮ℓ\+1\)\\tilde\{\\mathbf\{u\}\}\_\{\\ell\}\\leftarrow\(1\-\\alpha\)\\,\\mathbf\{u\}\_\{\\ell\}\+\\tfrac\{\\alpha\}\{2\}\\\!\\left\(\\mathbf\{u\}\_\{\\ell\-1\}\+\\mathbf\{u\}\_\{\\ell\+1\}\\right\)ℓ=2,…,L−1\\ell=2,\\ldots,L\{\-\}1

𝐮~L←\(1−α2\)​𝐮L\+α2​𝐮L−1\\tilde\{\\mathbf\{u\}\}\_\{L\}\\leftarrow\\left\(1\-\\tfrac\{\\alpha\}\{2\}\\right\)\\mathbf\{u\}\_\{L\}\+\\tfrac\{\\alpha\}\{2\}\\,\\mathbf\{u\}\_\{L\-1\}

𝐮~ℓ←‖𝐮ℓ‖2‖𝐮~ℓ‖2​𝐮~ℓ\\tilde\{\\mathbf\{u\}\}\_\{\\ell\}\\leftarrow\\dfrac\{\\\|\\mathbf\{u\}\_\{\\ell\}\\\|\_\{2\}\}\{\\\|\\tilde\{\\mathbf\{u\}\}\_\{\\ell\}\\\|\_\{2\}\}\\,\\tilde\{\\mathbf\{u\}\}\_\{\\ell\}optional norm pres\.

𝜽ℓ←𝜽ℓ−η​λ​𝜽ℓ−η​𝐮~ℓ\\boldsymbol\{\\theta\}\_\{\\ell\}\\leftarrow\\boldsymbol\{\\theta\}\_\{\\ell\}\{\\color\[rgb\]\{1,0,0\}\\definecolor\[named\]\{pgfstrokecolor\}\{rgb\}\{1,0,0\}\-\\eta\\lambda\\,\\boldsymbol\{\\theta\}\_\{\\ell\}\}\-\\eta\\,\\tilde\{\\mathbf\{u\}\}\_\{\\ell\}optional w\.d\. decouple

ϕ←ϕ−η​𝐮ϕ\\boldsymbol\{\\phi\}\\leftarrow\\boldsymbol\{\\phi\}\-\\eta\\,\\mathbf\{u\}\_\{\\phi\}

endfor

### 3\.4Normalization Variants

We additionally consider two normalization variants that separate the effects of smoothing on update directions and update magnitudes\.

1. i\)Norm\-Preserving Smoothing\.After applying the smoothing operator, each block update is rescaled to match the norm of the original optimizer update: u~ℓ=‖uℓ‖2‖\(P​u\)ℓ‖2​\(P​u\)ℓ\.\\tilde\{u\}\_\{\\ell\}=\\frac\{\\\|u\_\{\\ell\}\\\|\_\{2\}\}\{\\\|\(Pu\)\_\{\\ell\}\\\|\_\{2\}\}\(Pu\)\_\{\\ell\}\.This preserves the magnitude of each block update while allowing smoothing to modify only its direction\.
2. ii\)Directional Smoothing\.Before smoothing, each block update is normalized to unit norm, so that smoothing depends only on update directions\. The resulting smoothed direction is then rescaled to the original update norm: u~ℓ=‖uℓ‖2​\(P​\(u1/‖u1‖2,…,uL/‖uL‖2\)\)ℓ‖\(P​\(u1/‖u1‖2,…,uL/‖uL‖2\)\)ℓ‖2\.\\tilde\{u\}\_\{\\ell\}=\\frac\{\\\|u\_\{\\ell\}\\\|\_\{2\}\\,\\bigl\(P\(u\_\{1\}/\\\|u\_\{1\}\\\|\_\{2\},\\ldots,u\_\{L\}/\\\|u\_\{L\}\\\|\_\{2\}\)\\bigr\)\_\{\\ell\}\}\{\\left\\\|\\bigl\(P\(u\_\{1\}/\\\|u\_\{1\}\\\|\_\{2\},\\ldots,u\_\{L\}/\\\|u\_\{L\}\\\|\_\{2\}\)\\bigr\)\_\{\\ell\}\\right\\\|\_\{2\}\}\.This variant performs smoothing entirely in the space of update directions while preserving the original update magnitude\.

Throughout the paper, we distinguish three normalization variants:*Standard*\(no normalization\),*Norm*\(Norm\-Preserving Smoothing\), and*Dir*\(Directional Smoothing\)\.

## 4Experiments

We evaluate Gradient Smoothing across four distinct training regimes that stress different aspects of optimization: \(i\) RL fine\-tuning of LLMs on reasoning tasks, \(ii\) language model pretraining, \(iii\) supervised image classification with Vision Transformers, and \(iv\) diffusion model training\. To demonstrate compatibility with existing training setups, in all settings we compare against standard tuned baselines while keeping the underlying training*hyperparameters unchanged when applying smoothing*\. All base\-optimizer hyperparameters are inherited unchanged from the tuned baselines, suggesting that further gains may be possible with dedicated tuning for smoothing\. Because the method operates solely on parameter updates, it requires no architectural modifications and introduces negligible computational overhead, making it easy to integrate across a wide range of models and training regimes\.

### 4\.1Experimental Setup

#### Gradient Smoothing Configuration\.

Gradient smoothing is applied across block\-wise updates at each optimization step \(Section[3\.3](https://arxiv.org/html/2606.30813#S3.SS3)\) withα\\alpha\-window averaging across depth\. In addition to experimenting withα\\alpha\-window settings, we test standard, norm\-preserving, and directional variants for update smoothing\.

We also consider varied parameter groups within each block that are affected by any update smoothing operations\. In Transformers/ViTs, most capacity and gradient signal flows through linear projections \(attentionWQ,WK,WV,WOW\_\{Q\},W\_\{K\},W\_\{V\},W\_\{O\}and MLP layers\), while normalization and embedding parameters are lower\-dimensional and often behave differently under optimization\. Specifically we distinguish between \(i\) smoothing applied to*all*block parameters \(which we refer to asFullin Section[4\.2](https://arxiv.org/html/2606.30813#S4.SS2)\), \(ii\) smoothing applied*only*to linear layers within each block \(the non\-normalization layers in the models we use\) which we callProj\.

Gradient Smoothing follows the weight decay convention of the underlying base optimizer\. For optimizers with decoupled weight decay \(e\.g\., AdamW\), smoothing is applied only to the update produced by the loss gradients, with the weight decay term applied separately\.

![Refer to caption](https://arxiv.org/html/2606.30813v1/plots/vitb_cifar100_test_acc1.png)Figure 3:Test accuracy improvement with smoothing\.Test accuracy of ViT\-B trained on CIFAR\-100 with DeiT training recipe with data augmentations for 1700 epochs\. We compare baseline training \(α=0\\alpha=0, red\) against window smoothing withα=0\.1\\alpha=0\.1\(blue\) andα=0\.2\\alpha=0\.2\(purple\)\. Solid lines show running average trendlines \(25 epochs\) for clarity\. Both smoothing configurations outperform baseline, withα=0\.2\\alpha=0\.2achieving a final test accuracy of 75\.62% compared to the baseline 74\.56% \(Table[3](https://arxiv.org/html/2606.30813#S4.T3)\), and suggests that moderate gradient smoothing provides a consistent performance benefit throughout training\.
#### RLVR for LLM Reasoning:

The base model we use isDeepSeek\-R1\-Distill\-Qwen\-1\.5B\(Guoet al\.,[2025](https://arxiv.org/html/2606.30813#bib.bib4); Yanget al\.,[2024](https://arxiv.org/html/2606.30813#bib.bib63)\)\. We train with Group Relative Policy Optimization \(GRPO\)\(Shaoet al\.,[2024](https://arxiv.org/html/2606.30813#bib.bib21)\)on the dataset from Open\-RS\(Dang and Ngo,[2026](https://arxiv.org/html/2606.30813#bib.bib1)\)\(without the cosine response length reward\), and then evaluate on downstream mathematical reasoning tasks: AIME24, AIME25, AMC23, MATH\-500\. We train the model for up to 200 optimization steps with AdamW using a base learning rate of10−610^\{\-6\}with a cosine learning\-rate schedule and a minimum learning\-rate ratio of0\.10\.1, including a linear warmup over the first10%10\\%of training steps\. Evaluation sampling is done with temperature0\.60\.6and top\_p0\.950\.95\.

#### LLM Pretraining:

We further evaluate gradient smoothing in a language\-model pretraining setting using thenanochatrecipe\(Karpathy,[2025](https://arxiv.org/html/2606.30813#bib.bib17)\)\.nanochatprovides a compact end\-to\-end GPT\-style pretraining pipeline in which the transformer depth is the main scale parameter, while model width, number of heads, learning\-rate scaling, training horizon, batch size, and weight decay are automatically determined from this depth\. We train decoder\-only transformer models with depths2424and3030\(approximately1\.381\.38B and2\.402\.40B parameters respectively, Table[1](https://arxiv.org/html/2606.30813#S4.T1)\) using the defaultnanochatarchitecture and Muon\-style\(Jordanet al\.,[2024](https://arxiv.org/html/2606.30813#bib.bib27)\)optimization setup with NorMuon\(Liet al\.,[2025](https://arxiv.org/html/2606.30813#bib.bib18)\)\. Both depths use the same auto\-selected global batch size of≈1\\approx 1M tokens, and step count is set by the default target parameter\-data ratio\. For each depth, we compare the baseline run with an otherwise identical run where depth\-wise update smoothing is applied to the repeated transformer blocks\. We keep all other training settings fixed and evaluate using validation loss and the DCLM CORE metric from thenanochatevaluation suite\.

#### Vision Transformer Image Classification:

ViT\-B/4 models following DeiT training\(Touvronet al\.,[2021](https://arxiv.org/html/2606.30813#bib.bib53)\)with default image augmentations\. We train for 1700 epochs using AdamW with a cosine learning rate schedule, initial learning rate of5×10−45\\times 10^\{\-4\}, and weight decay of 0\.05\. CIFAR\-100 evaluated via top\-1 accuracy\.

Table 1:nanochatpretraining configurations\. Width \(dmodeld\_\{\\text\{model\}\}\), number of heads, and the global batch size are auto\-derived from the transformer depth by the scaling configuration\.
#### Diffusion:

We evaluate gradient smoothing on U\-ViT\(Baoet al\.,[2023](https://arxiv.org/html/2606.30813#bib.bib121)\), a ViT\-based diffusion backbone with long skip connections proposed for image generation on CIFAR\-10\. The U\-ViT architecture has1212transformer blocks with embedding dimension512512and88attention heads\. Optimization uses AdamW with learning rate2×10−42\\times 10^\{\-4\}, weight decay0\.030\.03, batch size128128, with a25002500\-step warmup for500,000500\{,\}000steps\. We sample via Euler–Maruyama SDE at5050NFE and report FID against both1010k and5050k generated samples\.

### 4\.2Experimental Results

Gradient Smoothing consistently improves performance across the diverse set of training regimes and benchmarks we evaluate\. Notably, even simple local Window Smoothing is effective across domains, suggesting that the benefits do not depend on carefully engineered or highly specialized smoothing operators, though more sophisticated operators may yield further gains\. Importantly, in all experiments, smoothing is applied directly on top of already tuned baseline training setups without modifying the underlying training hyperparameters, yet consistently improves performance without any additional hyperparameter tuning specific to Gradient Smoothing\.

Table 2:LLM mathematical reasoning performance \(pass@1accuracy\) after RLVR fine\-tuning\. Gradient Smoothing improves accuracy on downstream mathematical benchmarks compared to fine\-tuning just with AdamW\. Results are averaged across 4 runs\.#### RLVR Fine\-Tuning for LLM Reasoning:

Gradient Smoothing improves the reasoning accuracy ofDeepSeek\-R1\-Distill\-Qwen\-1\.5Bafter RLVR fine\-tuning across mathematical benchmarks \(Table[2](https://arxiv.org/html/2606.30813#S4.T2)\)\. All Window Smoothing configurations outperform the AdamW baseline in averagepass@1accuracy, with the best configuration achieving57\.60%57\.60\\%versus55\.45%55\.45\\%\. Across individual benchmarks, Gradient Smoothing yields substantial improvements, including gains on AIME24 \(40\.83% vs\. 35\.00%\), AIME25 \(30\.00% vs\. 24\.17%\), AMC23 \(80\.00% vs\. 78\.12%\), and MATH\-500 \(85\.35% vs\. 84\.50%\) under different smoothing configurations\. Notably, these improvements are achieved without modifying the RL objective or reward structure, suggesting that regularizing the optimization dynamics through depth\-wise update smoothing can improve RL\-based reasoning fine\-tuning\.

#### LLM Pretraining:

Gradient Smoothing also improves the optimization trajectory in language\-model pretraining with thenanochatrecipe \(Figure[2](https://arxiv.org/html/2606.30813#S2.F2)\)\. During training, across both depth\-2424and depth\-3030models, smoothing consistently accelerates convergence of the validation loss/BPB while simultaneously improving the CORE metric\. The improvements emerge at early steps and persist across the training trajectory, suggesting that smoothing provides a useful inductive bias for the model\. Notably, these gains are obtained in an optimized pretraining setup using the NorMuon/AdamW optimization pipeline fromnanochat, indicating that the benefits of smoothing extend beyond standard Adam\-family training alone\. The stronger improvements observed for the deeper depth\-3030model further suggest that coupling updates across depth may become increasingly beneficial as the model depth increases\.

#### Vision Transformer Image Classification:

Table 3:Image classification accuracy \(%\) for ViT\-B on CIFAR\-100\.Gradient Smoothing improves test accuracy of ViT\-B on CIFAR\-100 \(Table[3](https://arxiv.org/html/2606.30813#S4.T3), Figure[3](https://arxiv.org/html/2606.30813#S4.F3)\)\. Window smoothing withα=0\.2\\alpha=0\.2achieves notable improvements \(75\.62% vs 74\.56% baseline\), whileα=0\.1\\alpha=0\.1also improves over baseline \(75\.44%\)\. These results suggest that the smoothing regime provides consistent regularization for vision transformer training, with gains emerging early and persisting throughout optimization\.

#### Diffusion:

We evaluate Gradient Smoothing on U\-ViT\(Baoet al\.,[2023](https://arxiv.org/html/2606.30813#bib.bib121)\), a ViT\-based diffusion backbone with long U\-Net\-style skip connections which provides a strong tuned baseline on CIFAR\-10\. With 50 NFE sampling, Gradient Smoothing reduces FID@10k from6\.586\.58to5\.825\.82and FID@50k from4\.014\.01to3\.743\.74\(Table[4](https://arxiv.org/html/2606.30813#S4.T4)\)\. All runs share identical architecture, hyperparameters, and sampling budget as the tuned baseline\. These results demonstrate that the benefits of Gradient Smoothing extend to generative diffusion modeling, and beyond architectures with standard residual connections to transformer backbones with long U\-Net\-style skip connections\.

Table 4:FID scores on CIFAR\-10 for U\-ViT diffusion training with and without Gradient Smoothing\.![Refer to caption](https://arxiv.org/html/2606.30813v1/plots/nanochat_microbatch_var.png)Figure 4:Microbatch gradient variance with window smoothing\.Total microbatch varianceVmbV\_\{\\mathrm\{mb\}\}\(Section[4\.4](https://arxiv.org/html/2606.30813#S4.SS4)\) during nanochat depth 24 pretraining, comparing the baseline \(red\) against window smoothing withα=0\.05\\alpha=0\.05\(green\) andα=0\.1\\alpha=0\.1\(blue\)\. After the initial warmup, both smoothing runs maintain consistently lower microbatch gradient variance than baseline, with the gap widening later in training\.

### 4\.3Effect of Gradient Smoothing on Representations

Smoothing parameter updates suggests that at convergence, models may tend toward more coordinated transformations across depth\. To quantify this effect on learned representations, we measure two geometric properties of token trajectories: \(i\) cosine similarity between successive residual updates, and \(ii\) line shape score \(LSS\), which captures how close the trajectory lies to a one\-dimensional affine space\.

Cosine similarity of residual\.To quantify the change in representation induced by a layer, we consider the differences in token representation across depthxℓx\_\{\\ell\}, for1≤ℓ≤L1\\leq\\ell\\leq L\. Forℓ∈\{1,…,L−1\}\\ell\\in\\\{1,\\ldots,L\-1\\\}the differencedℓ=xℓ\+1−xℓd\_\{\\ell\}=x\_\{\\ell\+1\}\-x\_\{\\ell\}corresponds to the contribution of layerℓ\+1\\ell\+1\. The similarity of contributions between layers is quantified bycℓ=⟨dℓ,dℓ\+1⟩‖dℓ‖​‖dℓ\+1‖c\_\{\\ell\}=\\frac\{\\langle\{d\_\{\\ell\},d\_\{\\ell\+1\}\\rangle\}\}\{\\\|d\_\{\\ell\}\\\|\\\|d\_\{\\ell\+1\}\\\|\}\(Figure[6](https://arxiv.org/html/2606.30813#A1.F6)\)\.

Line shape score\.We measure the linearity of token trajectories via the line shape score \(LSS\)\(Gai and Zhang,[2021](https://arxiv.org/html/2606.30813#bib.bib20)\), measuring the closeness of a trajectory to a line\. Given representationsx0,…,xLx\_\{0\},\\ldots,x\_\{L\}, the normalized trajectoryx~0=x0\\tilde\{x\}\_\{0\}=x\_\{0\}with

x~ℓ=x~ℓ−1\+xℓ−xℓ−1‖xℓ−xℓ−1‖2\\tilde\{x\}\_\{\\ell\}=\\tilde\{x\}\_\{\\ell\-1\}\+\\frac\{x\_\{\\ell\}\-x\_\{\\ell\-1\}\}\{\\\|x\_\{\\ell\}\-x\_\{\\ell\-1\}\\\|\_\{2\}\}The line shape score is then given byLSS=L‖x~L−x~0‖2\\text\{LSS\}=\\frac\{L\}\{\\\|\\tilde\{x\}\_\{L\}\-\\tilde\{x\}\_\{0\}\\\|\_\{2\}\}\. LSS≥1\\geq 1, with equality if and only if the trajectory is perfectly co\-linear \(Figure[7](https://arxiv.org/html/2606.30813#A1.F7)\)\.

We measure the effect of smoothing on representation structure via the Line Shape Score \(LSS\) and cosine similarity of CLS token trajectories for ViT\-B trained with window smoothing forα∈\{0\.1,…,0\.4\}\\alpha\\in\\\{0\.1,\\ldots,0\.4\\\}relative to the non\-smoothed baseline \(α=0\.0\\alpha=0\.0\)\. As the smoothing strength increases, layer\-wise contributions to the CLS token become more aligned \(Figure[5](https://arxiv.org/html/2606.30813#S5.F5)\), and the trend persists when measuring mean cosine similarity of the contribution at individual layers \(Figure[6](https://arxiv.org/html/2606.30813#A1.F6)\)\. Smoothing also increases the linearity of the CLS token trajectory across layers \(Figure[7](https://arxiv.org/html/2606.30813#A1.F7)\)\. Together, these results suggest that Gradient Smoothing may implicitly regularize representation paths, encouraging more coherent evolution of representations through the model depth\.

### 4\.4Effect of Smoothing on Training Stability

#### Gradient variance\.

To assess downstream improvements on the optimization trajectory, we measure the total stochastic gradient variance within the microbatch\. We quantify

Vmb\(t\)=1N​∑n=1N‖gn\(t\)−g¯\(t\)‖22,V\_\{\\mathrm\{mb\}\}^\{\(t\)\}\\;=\\;\\frac\{1\}\{N\}\\sum\_\{n=1\}^\{N\}\\big\\\|g^\{\(t\)\}\_\{n\}\-\\bar\{g\}^\{\(t\)\}\\big\\\|\_\{2\}^\{2\},where at steptt,gn\(t\)g\_\{n\}^\{\(t\)\}is the per\-sample parameter gradient andg¯\(t\)\\bar\{g\}^\{\(t\)\}is the sample mean over theNNbatch elements at that step\. Window smoothing withα∈\{0\.05,0\.1\}\\alpha\\in\\\{0\.05,0\.1\\\}maintains consistently lower variance than the baseline throughout nanochat pretraining \(Figure[4](https://arxiv.org/html/2606.30813#S4.F4)\), suggestive of improved training stability\. Full measurement details are deferred to Appendix[A\.4](https://arxiv.org/html/2606.30813#A1.SS4), and a variance contraction bound is shown across layer depth \(Proposition[A\.10](https://arxiv.org/html/2606.30813#A1.Thmtheorem10)\)\.

## 5Alignment of Representation Residuals under Gradient Smoothing

This section formalizes a consistent empirical observation:*Gradient Smoothing increases the alignment of representation differences across depth*\. We study this through the residual incrementsrℓ:=hℓ\+1−hℓ,r\_\{\\ell\}:=h\_\{\\ell\+1\}\-h\_\{\\ell\},and the cosine similaritycos⁡\(rℓ\+1,rℓ\)\\cos\(r\_\{\\ell\+1\},r\_\{\\ell\}\), which measures how much consecutive residual increments point in the same direction\. Our analysis gives lower bounds on the*average cosine*after applying either the baseline optimizer update or its smoothed version\. The key mechanism is that Window Smoothing contracts depth\-to\-depth variation in the update directions\.

#### Residual increments and alignment metric\.

For any depth\-stacked parametersϑ=\(ϑ1,…,ϑL\)\\vartheta=\(\\vartheta\_\{1\},\\dots,\\vartheta\_\{L\}\), define forℓ=1,…,L\\ell=1,\\dots,L,

rℓ​\(ϑ\):=hℓ\+1​\(ϑ\)−hℓ​\(ϑ\)=F​\(hℓ​\(ϑ\);ϑℓ\)∈ℝd\.r\_\{\\ell\}\(\\vartheta\):=h\_\{\\ell\+1\}\(\\vartheta\)\-h\_\{\\ell\}\(\\vartheta\)=F\\\!\\big\(h\_\{\\ell\}\(\\vartheta\);\\vartheta\_\{\\ell\}\\big\)\\in\\mathbb\{R\}^\{d\}\.\(3\)We measure directional alignment of consecutive residual increments by

cos¯​\(ϑ\):=1L−1​∑ℓ=1L−1cos⁡\(rℓ\+1​\(ϑ\),rℓ​\(ϑ\)\)\.\\overline\{\\cos\}\(\\vartheta\):=\\frac\{1\}\{L\-1\}\\sum\_\{\\ell=1\}^\{L\-1\}\\cos\\\!\\big\(r\_\{\\ell\+1\}\(\\vartheta\),r\_\{\\ell\}\(\\vartheta\)\\big\)\.\(4\)This metric is invariant to the magnitudes‖rℓ‖\\\|r\_\{\\ell\}\\\|and isolates the alignment of representation*directions*\. Empirically, we consistently observe thatcos¯\\overline\{\\cos\}increases under smoothing, as shown in Figure[5](https://arxiv.org/html/2606.30813#S5.F5)\.

#### Baseline and smoothed one\-step updates\.

Letu​\(θ\)=\(u1​\(θ\),…,uL​\(θ\)\)u\(\\theta\)=\(u\_\{1\}\(\\theta\),\\dots,u\_\{L\}\(\\theta\)\)denote the blockwise update direction produced by a base optimizer\. From current parametersθ\\theta, define the baseline and smoothed one\-step updates

θb=θ−η​u​\(θ\),θs=θ−η​S​u​\(θ\),\\theta^\{b\}=\\theta\-\\eta\\,u\(\\theta\),\\qquad\\theta^\{s\}=\\theta\-\\eta\\,Su\(\\theta\),\(5\)whereSSis a depth\-smoothing operator\. We focus on the local\-averaging smoother \(Window Smoothing\) withα∈\(0,12\]\\alpha\\in\(0,\\tfrac\{1\}\{2\}\]:

\(S​u\)ℓ=\{\(1−α2\)​u1\+α2​u2,ℓ=1,\(1−α\)​uℓ\+α2​\(uℓ−1\+uℓ\+1\),2≤ℓ≤L−1,\(1−α2\)​uL\+α2​uL−1,ℓ=L\.\(Su\)\_\{\\ell\}=\\begin\{cases\}\\big\(1\-\\tfrac\{\\alpha\}\{2\}\\big\)u\_\{1\}\+\\tfrac\{\\alpha\}\{2\}u\_\{2\},&\\ell=1,\\\\\[2\.0pt\] \(1\-\\alpha\)u\_\{\\ell\}\+\\tfrac\{\\alpha\}\{2\}\(u\_\{\\ell\-1\}\+u\_\{\\ell\+1\}\),&2\\leq\\ell\\leq L\-1,\\\\\[2\.0pt\] \\big\(1\-\\tfrac\{\\alpha\}\{2\}\\big\)u\_\{L\}\+\\tfrac\{\\alpha\}\{2\}u\_\{L\-1\},&\\ell=L\.\\end\{cases\}\(6\)

#### Assumptions\.

We assume standard Lipschitz regularity of the block mapFFand a mild non\-degeneracy condition requiring the residual increments to have norms bounded away from zero along the relevant layer\-wise trajectories\. The formal assumptions are given in Appendix[A\.2](https://arxiv.org/html/2606.30813#A1.SS2)\.

### 5\.1Main Result

We now state the main alignment comparison between the baseline and smoothing\. Unlike a simplified depth\-flat analysis where the parameters are assumed to be equal across depth, this result allows the current parameters to already have nontrivial depth\-to\-depth variation\. The comparison therefore separates two effects: the pre\-existing depth variationD​θD\\thetaand the update\-induced variationD​u​\(θ\)Du\(\\theta\), whereDDis the first\-difference operator\(D​v\)ℓ=vℓ\+1−vℓ\(Dv\)\_\{\\ell\}=v\_\{\\ell\+1\}\-v\_\{\\ell\}\.

###### Theorem 5\.1\(Alignment Comparison under Smoothing\)\.

Assume the regularity conditions of Appendix[A\.2](https://arxiv.org/html/2606.30813#A1.SS2)\(Assumption[A\.1](https://arxiv.org/html/2606.30813#A1.Thmtheorem1)\)\. LetSSbe the local\-averaging smoother \([6](https://arxiv.org/html/2606.30813#S5.E6)\) withα∈\(0,12\]\\alpha\\in\(0,\\tfrac\{1\}\{2\}\], and letDDbe the first\-difference operator\(D​v\)ℓ=vℓ\+1−vℓ\(Dv\)\_\{\\ell\}=v\_\{\\ell\+1\}\-v\_\{\\ell\}\. Define

δ:=D​θ,w:=D​u​\(θ\),\\delta:=D\\theta,\\qquad w:=Du\(\\theta\),where norms and inner products are taken over the corresponding depth\-stacked block vectors\. LetS\(1\)S^\{\(1\)\}denote the induced averaging operator on first differences, so that

D​\(S​u\)=S\(1\)​D​uD\(Su\)=S^\{\(1\)\}Duas shown in Appendix[A\.2](https://arxiv.org/html/2606.30813#A1.SS2), Lemma[A\.5](https://arxiv.org/html/2606.30813#A1.Thmtheorem5)\.

Then the smoothed and baseline one\-step updates satisfy

cos¯​\(θs\)\\displaystyle\\overline\{\\cos\}\(\\theta^\{s\}\)≥1−4​Lh2​M2m2−4​Lθ2m2⋅1L−1​‖δ−η​S\(1\)​w‖2,\\displaystyle\\geq 1\-\\frac\{4L\_\{h\}^\{2\}M^\{2\}\}\{m^\{2\}\}\-\\frac\{4L\_\{\\theta\}^\{2\}\}\{m^\{2\}\}\\cdot\\frac\{1\}\{L\-1\}\\,\\big\\\|\\delta\-\\eta S^\{\(1\)\}w\\big\\\|^\{2\},\(7\)cos¯​\(θb\)\\displaystyle\\overline\{\\cos\}\(\\theta^\{b\}\)≥1−4​Lh2​M2m2−4​Lθ2m2⋅1L−1​‖δ−η​w‖2\.\\displaystyle\\geq 1\-\\frac\{4L\_\{h\}^\{2\}M^\{2\}\}\{m^\{2\}\}\-\\frac\{4L\_\{\\theta\}^\{2\}\}\{m^\{2\}\}\\cdot\\frac\{1\}\{L\-1\}\\,\\big\\\|\\delta\-\\eta w\\big\\\|^\{2\}\.\(8\)Moreover, the difference between these two lower bounds is exactly

4​Lθ2m2⋅1L−1​\[η2​\(‖w‖2−‖S\(1\)​w‖2\)−2​η​⟨δ,\(I−S\(1\)\)​w⟩\]\.\\frac\{4L\_\{\\theta\}^\{2\}\}\{m^\{2\}\}\\cdot\\frac\{1\}\{L\-1\}\\left\[\\eta^\{2\}\\big\(\\\|w\\\|^\{2\}\-\\\|S^\{\(1\)\}w\\\|^\{2\}\\big\)\-2\\eta\\big\\langle\\delta,\(I\-S^\{\(1\)\}\)w\\big\\rangle\\right\]\.\(9\)Consequently, ifη\>0\\eta\>0andLθ\>0L\_\{\\theta\}\>0, smoothing gives a strictly improved lower bound whenever

⟨δ,\(I−S\(1\)\)​w⟩<η2​\(‖w‖2−‖S\(1\)​w‖2\)\.\\big\\langle\\delta,\(I\-S^\{\(1\)\}\)w\\big\\rangle<\\frac\{\\eta\}\{2\}\\Big\(\\\|w\\\|^\{2\}\-\\\|S^\{\(1\)\}w\\\|^\{2\}\\Big\)\.\(10\)In particular, sinceS\(1\)S^\{\(1\)\}is a contraction with

‖S\(1\)‖op=μ⋆:=1−α​\(1−cos⁡\(πL\)\)<1,\\\|S^\{\(1\)\}\\\|\_\{\\mathrm\{op\}\}=\\mu\_\{\\star\}:=1\-\\alpha\\Big\(1\-\\cos\\Big\(\\frac\{\\pi\}\{L\}\\Big\)\\Big\)<1,\(11\)the update\-induced depth variation satisfies

‖S\(1\)​w‖2≤μ⋆2​‖w‖2\.\\\|S^\{\(1\)\}w\\\|^\{2\}\\leq\\mu\_\{\\star\}^\{2\}\\\|w\\\|^\{2\}\.\(12\)

The proof is given in Appendix[A\.2](https://arxiv.org/html/2606.30813#A1.SS2)\.

The theorem gives a comparison between the baseline and smoothed one\-step updates\. Both bounds have the same form, but the baseline depends on the raw update differencew=D​u​\(θ\)w=Du\(\\theta\), while the smoothed bound depends on the averaged update differenceS\(1\)​wS^\{\(1\)\}w\. Thus, smoothing affects the bound only through the depth\-wise variation of the update\.

The exact gap separates two effects\. The term

‖w‖2−‖S\(1\)​w‖2\\\|w\\\|^\{2\}\-\\\|S^\{\(1\)\}w\\\|^\{2\}measures the reduction in depth\-to\-depth update variation caused by smoothing, while

⟨δ,\(I−S\(1\)\)​w⟩\\langle\\delta,\(I\-S^\{\(1\)\}\)w\\ranglemeasures how the removed component of the update variation aligns with the existing parameter variationδ=D​θ\\delta=D\\theta\. Hence smoothing improves the lower bound whenever the removed component is not too positively aligned with the current depth\-wise parameter variation\. In the depth\-flat caseδ=0\\delta=0, this obstruction vanishes, and the improvement follows directly from the contraction of update differences byS\(1\)S^\{\(1\)\}\.

![Refer to caption](https://arxiv.org/html/2606.30813v1/plots/cos_sim2_edit.png)Figure 5:Layer contributions similarity with increased smoothing\.Average cosine similarity of layer differencesdℓ=xℓ\+1−xℓd\_\{\\ell\}=x\_\{\\ell\+1\}\-x\_\{\\ell\}for ViT\-B CLS token trained on CIFAR\-100 \(1700 epochs\)\. Means are taken across all layer pairs\(di,dj\)\(d\_\{i\},d\_\{j\}\)for1≤i,j≤L−11\\leq i,j\\leq L\-1,i≠ji\\neq j\. Each point shows the median across 100 images from the validation set \(rhombus\) and training set \(circle\); error bars indicate interquartile range \(Q1–Q3\)\. As the window smoothing parameterα∈\{0\.1,…,0\.4\}\\alpha\\in\\\{0\.1,\\ldots,0\.4\\\}increases from the baseline \(α=0\\alpha=0\), mean cosine similarity increases monotonically, suggesting that gradient smoothing encourages more consistent directional updates across layers\.

## 6Discussion

Implicit Biases\.A primary motivation arises from the observed similarity of deep layer representations in transformers and LLMs\. If representations and layer contributions exhibit alignment properties at convergence, Gradient Smoothing may be a helpful tool for understanding the implicit structure induced by preconditioning and a mechanism for injecting helpful inductive biases earlier in training\. There is an optimal balance when introducing similar regularizers or inductive biases: while we find that moderate layer\-wise gradient smoothing is often beneficial, excessive smoothing may reduce important layer\-specific information in the updates\.

General Operators\.Gradient Smoothing represents only one family of the broader paradigm of Depth\-wise Gradient Augmentation\. While we focus on local window averaging, many other coupling operators are possible, including exponential smoothing as well as adaptive or learned schemes based on gradient or representation structure\. The strength of coupling may also vary across depth or throughout training, and could be used to steer other characteristics such as gradient variance \(Figure[4](https://arxiv.org/html/2606.30813#S4.F4)\)\. More generally, while this work studies smoothing\-based coupling, alternative forms of update interaction, including selective coupling, may offer additional benefits\. Exploring this broader design space is a promising potential direction for future work\.

## 7Related Work

#### Transformer Layer Regularity\.

Recent empirical work shows that trained transformers exhibit strong layer\-wise structure in their representations and Jacobians\(Aubryet al\.,[2025](https://arxiv.org/html/2606.30813#bib.bib29); Gromovet al\.,[2025](https://arxiv.org/html/2606.30813#bib.bib15); Li and Papyan,[2024](https://arxiv.org/html/2606.30813#bib.bib30); Patrawalaet al\.,[2026](https://arxiv.org/html/2606.30813#bib.bib13); Kaplet al\.,[2025](https://arxiv.org/html/2606.30813#bib.bib14); Ladet al\.,[2025](https://arxiv.org/html/2606.30813#bib.bib11); Jianget al\.,[2025b](https://arxiv.org/html/2606.30813#bib.bib19); Wolfram and Schein,[2025](https://arxiv.org/html/2606.30813#bib.bib12)\), with some links of such structure to improved generalization\. A closely related line of work on*compression valleys*\(Skeanet al\.,[2024](https://arxiv.org/html/2606.30813#bib.bib186),[2025](https://arxiv.org/html/2606.30813#bib.bib188)\)shows that mid\-layer parameters exhibit lower\-rank structure than boundary layers, with this low\-rank structure often shared across consecutive mid\-layers\. Such observations also motivate the exploration of layer pruning methods that highlight minimal performance loss\(Menet al\.,[2024](https://arxiv.org/html/2606.30813#bib.bib10); Gromovet al\.,[2025](https://arxiv.org/html/2606.30813#bib.bib15); Jianget al\.,[2025a](https://arxiv.org/html/2606.30813#bib.bib2); Krauseet al\.,[2025](https://arxiv.org/html/2606.30813#bib.bib9)\)\. Our work connects to this literature by showing that gradient smoothing amplifies these empirical alignment tendencies, while leading to improved performance\.

#### Adaptive Optimization\.

Modern deep learning relies heavily on adaptive optimizers that scale updates based on gradient statistics\. Adam\(Kingma and Ba,[2017](https://arxiv.org/html/2606.30813#bib.bib22)\)and AdamW\(Loshchilov and Hutter,[2019](https://arxiv.org/html/2606.30813#bib.bib23)\)maintain adaptive per\-parameter estimates of first and second moments, enabling effective training across all modern architectures\. Second\-order methods such as Shampoo\(Guptaet al\.,[2018](https://arxiv.org/html/2606.30813#bib.bib25)\)and SOAP\(Vyaset al\.,[2025](https://arxiv.org/html/2606.30813#bib.bib26)\)approximate full\-matrix preconditioning for improved convergence, while Muon\(Jordanet al\.,[2024](https://arxiv.org/html/2606.30813#bib.bib27)\)applies momentum in the orthogonalized gradient space\. These methods focus on adaptively scaling updates per parameter group, whereas gradient smoothing operates on a complementary axis:*depth\-wise augmentations*of updates across optimization steps\. Our approach is agnostic to the choice of base optimizer and can be combined with existing adaptive methods to stabilize training by amplifying gradient information across depth\.

#### Pre\-conditioned Optimization\.

Natural gradient descent\(Amari,[1998](https://arxiv.org/html/2606.30813#bib.bib50)\)and its approximations such as K\-FAC\(Martens and Grosse,[2020](https://arxiv.org/html/2606.30813#bib.bib8); Georgeet al\.,[2021](https://arxiv.org/html/2606.30813#bib.bib7); Eschenhagenet al\.,[2024](https://arxiv.org/html/2606.30813#bib.bib6); Linet al\.,[2024](https://arxiv.org/html/2606.30813#bib.bib5); Nagwekar,[2025](https://arxiv.org/html/2606.30813#bib.bib3)\)precondition updates using curvature information to accelerate convergence\. Our method adopts a pre\-conditioned perspective but focuses on depth structure rather than per\-layer curvature, applying smoothing across transformer depth with minimal computational overhead\.

Structure in Hidden Representations\.Many works study the emergent structure of deep network representations\(Wanget al\.,[2024a](https://arxiv.org/html/2606.30813#bib.bib177); Parkeret al\.,[2023](https://arxiv.org/html/2606.30813#bib.bib84); Zangrandoet al\.,[2025](https://arxiv.org/html/2606.30813#bib.bib79); Garrod and Keating,[2024](https://arxiv.org/html/2606.30813#bib.bib179); Wanget al\.,[2024b](https://arxiv.org/html/2606.30813#bib.bib180); Hoyt and Owen,[2021](https://arxiv.org/html/2606.30813#bib.bib181); Arouset al\.,[2024](https://arxiv.org/html/2606.30813#bib.bib80); Zarkaet al\.,[2021](https://arxiv.org/html/2606.30813#bib.bib81); Ben\-Shaul and Dekel,[2022](https://arxiv.org/html/2606.30813#bib.bib184); Papyan,[2020](https://arxiv.org/html/2606.30813#bib.bib191); Súkeníket al\.,[2023](https://arxiv.org/html/2606.30813#bib.bib192); Papyanet al\.,[2020](https://arxiv.org/html/2606.30813#bib.bib115); Zhouet al\.,[2025](https://arxiv.org/html/2606.30813#bib.bib128); Wanget al\.,[2026](https://arxiv.org/html/2606.30813#bib.bib127); Fisheret al\.,[2024](https://arxiv.org/html/2606.30813#bib.bib16)\), especially in transformers\. In LLMs, recent works identify uniform token structures\(Wu and Papyan,[2024](https://arxiv.org/html/2606.30813#bib.bib69); Shaiet al\.,[2025](https://arxiv.org/html/2606.30813#bib.bib189); Piotrowskiet al\.,[2025](https://arxiv.org/html/2606.30813#bib.bib187); Skeanet al\.,[2025](https://arxiv.org/html/2606.30813#bib.bib188),[2024](https://arxiv.org/html/2606.30813#bib.bib186)\)and low\-dimensional hidden trajectories\(Songet al\.,[2025](https://arxiv.org/html/2606.30813#bib.bib185); Sarfatiet al\.,[2024](https://arxiv.org/html/2606.30813#bib.bib68)\)\. Our work proposes an approach to regularize representation trajectories across depth by augmenting the gradient updates\.

## 8Conclusion

We introduced*Depth\-wise Gradient Augmentation*, a general optimization paradigm in which the collection of block\-wise optimizer updates is transformed along the depth dimension\. Within this framework, we studied*Gradient Smoothing*, a family of depth\-wise smoothing methods, instantiated through local Window Smoothing\. The resulting method is compatible with arbitrary base optimizers and consistently improves training across a diverse range of settings, including LLM pretraining, RL post\-training for reasoning, image classification, and diffusion modeling\. Beyond these performance gains, we showed that Gradient Smoothing promotes more structured representation evolution across depth, leading to greater alignment and linearity of representation trajectories\. We further provided a theoretical characterization of this phenomenon by interpreting Gradient Smoothing as a structured depth\-wise preconditioning method that contracts variation in block\-wise updates, yielding improved bounds on the alignment of representation differences\.

More broadly, our results suggest that explicitly exploiting cross\-depth structure during optimization can improve both optimization and generalization in modern deep networks\. We hope that the proposed Depth\-wise Gradient Augmentation framework motivates future work on richer depth\-wise update transformations, adaptive smoothing operators, learned augmentation schemes, and other forms of structured optimization that leverage the organization of repeated\-block architectures\.

## Impact Statement

This paper presents work whose goal is to advance the field of Machine Learning\. There are some potential societal consequences of improved optimization methods, including downstream impacts on systems used in science and decision\-making\. At the same time, the techniques introduced in this work are general\-purpose and do not target any specific application domain\. We do not foresee immediate negative societal impacts arising directly from this work\.

## References

- S\. Amari \(1998\)Natural gradient works efficiently in learning\.Neural Computation10\(2\),pp\. 251–276\.External Links:[Document](https://dx.doi.org/10.1162/089976698300017746)Cited by:[§7](https://arxiv.org/html/2606.30813#S7.SS0.SSS0.Px3.p1.1)\.
- G\. B\. Arous, R\. Gheissari, J\. Huang, and A\. Jagannath \(2024\)High\-dimensional SGD aligns with emerging outlier eigenspaces\.InThe Twelfth International Conference on Learning Representations,External Links:[Link](https://openreview.net/forum?id=MHjigVnI04)Cited by:[§7](https://arxiv.org/html/2606.30813#S7.SS0.SSS0.Px3.p2.1)\.
- M\. Aubry, H\. Meng, A\. Sugolov, and V\. Papyan \(2025\)Transformer block coupling and its correlation with generalization in llms\.External Links:2407\.07810,[Link](https://arxiv.org/abs/2407.07810)Cited by:[§1](https://arxiv.org/html/2606.30813#S1.p1.1),[§7](https://arxiv.org/html/2606.30813#S7.SS0.SSS0.Px1.p1.1)\.
- F\. Bao, S\. Nie, K\. Xue, Y\. Cao, C\. Li, H\. Su, and J\. Zhu \(2023\)All are worth words: a vit backbone for diffusion models\.External Links:2209\.12152,[Link](https://arxiv.org/abs/2209.12152)Cited by:[§4\.1](https://arxiv.org/html/2606.30813#S4.SS1.SSS0.Px5.p1.11),[§4\.2](https://arxiv.org/html/2606.30813#S4.SS2.SSS0.Px4.p1.4)\.
- I\. Ben\-Shaul and S\. Dekel \(2022\)Nearest class\-center simplification through intermediate layers\.InProceedings of Topological, Algebraic, and Geometric Learning Workshops 2022,A\. Cloninger, T\. Doster, T\. Emerson, M\. Kaul, I\. Ktena, H\. Kvinge, N\. Miolane, B\. Rieck, S\. Tymochko, and G\. Wolf \(Eds\.\),Proceedings of Machine Learning Research, Vol\.196,pp\. 37–47\.External Links:[Link](https://proceedings.mlr.press/v196/ben-shaul22a.html)Cited by:[§7](https://arxiv.org/html/2606.30813#S7.SS0.SSS0.Px3.p2.1)\.
- Q\. Dang and C\. Ngo \(2026\)Reinforcement learning for reasoning in small llms: what works and what doesn’t\.External Links:2503\.16219,[Link](https://arxiv.org/abs/2503.16219)Cited by:[§4\.1](https://arxiv.org/html/2606.30813#S4.SS1.SSS0.Px2.p1.5)\.
- R\. Eschenhagen, A\. Immer, R\. E\. Turner, F\. Schneider, and P\. Hennig \(2024\)Kronecker\-factored approximate curvature for modern neural network architectures\.External Links:2311\.00636,[Link](https://arxiv.org/abs/2311.00636)Cited by:[§7](https://arxiv.org/html/2606.30813#S7.SS0.SSS0.Px3.p1.1)\.
- Q\. Fisher, H\. Meng, and V\. Papyan \(2024\)Pushing boundaries: mixup’s influence on neural collapse\.External Links:2402\.06171,[Link](https://arxiv.org/abs/2402.06171)Cited by:[§7](https://arxiv.org/html/2606.30813#S7.SS0.SSS0.Px3.p2.1)\.
- K\. Gai and S\. Zhang \(2021\)A mathematical principle of deep learning: learn the geodesic curve in the wasserstein space\.External Links:2102\.09235,[Link](https://arxiv.org/abs/2102.09235)Cited by:[§4\.3](https://arxiv.org/html/2606.30813#S4.SS3.p3.2)\.
- C\. Garrod and J\. P\. Keating \(2024\)Unifying low dimensional observations in deep learning through the deep linear unconstrained feature model\.External Links:2404\.06106Cited by:[§7](https://arxiv.org/html/2606.30813#S7.SS0.SSS0.Px3.p2.1)\.
- T\. George, C\. Laurent, X\. Bouthillier, N\. Ballas, and P\. Vincent \(2021\)Fast approximate natural gradient descent in a kronecker\-factored eigenbasis\.External Links:1806\.03884,[Link](https://arxiv.org/abs/1806.03884)Cited by:[§7](https://arxiv.org/html/2606.30813#S7.SS0.SSS0.Px3.p1.1)\.
- A\. Gromov, K\. Tirumala, H\. Shapourian, P\. Glorioso, and D\. A\. Roberts \(2025\)The unreasonable ineffectiveness of the deeper layers\.External Links:2403\.17887,[Link](https://arxiv.org/abs/2403.17887)Cited by:[§1](https://arxiv.org/html/2606.30813#S1.p1.1),[§7](https://arxiv.org/html/2606.30813#S7.SS0.SSS0.Px1.p1.1)\.
- D\. Guo, D\. Yang, H\. Zhang, J\. Song, P\. Wang, Q\. Zhu, R\. Xu, R\. Zhang, S\. Ma, X\. Bi, X\. Zhang, X\. Yu, Y\. Wu, Z\. F\. Wu, Z\. Gou, Z\. Shao, Z\. Li, Z\. Gao, A\. Liu, B\. Xue, B\. Wang, B\. Wu, B\. Feng, C\. Lu, C\. Zhao, C\. Deng, C\. Ruan, D\. Dai, D\. Chen, D\. Ji, E\. Li, F\. Lin, F\. Dai, F\. Luo, G\. Hao, G\. Chen, G\. Li, H\. Zhang, H\. Xu, H\. Ding, H\. Gao, H\. Qu, H\. Li, J\. Guo, J\. Li, J\. Chen, J\. Yuan, J\. Tu, J\. Qiu, J\. Li, J\. L\. Cai, J\. Ni, J\. Liang, J\. Chen, K\. Dong, K\. Hu, K\. You, K\. Gao, K\. Guan, K\. Huang, K\. Yu, L\. Wang, L\. Zhang, L\. Zhao, L\. Wang, L\. Zhang, L\. Xu, L\. Xia, M\. Zhang, M\. Zhang, M\. Tang, M\. Zhou, M\. Li, M\. Wang, M\. Li, N\. Tian, P\. Huang, P\. Zhang, Q\. Wang, Q\. Chen, Q\. Du, R\. Ge, R\. Zhang, R\. Pan, R\. Wang, R\. J\. Chen, R\. L\. Jin, R\. Chen, S\. Lu, S\. Zhou, S\. Chen, S\. Ye, S\. Wang, S\. Yu, S\. Zhou, S\. Pan, S\. S\. Li, S\. Zhou, S\. Wu, T\. Yun, T\. Pei, T\. Sun, T\. Wang, W\. Zeng, W\. Liu, W\. Liang, W\. Gao, W\. Yu, W\. Zhang, W\. L\. Xiao, W\. An, X\. Liu, X\. Wang, X\. Chen, X\. Nie, X\. Cheng, X\. Liu, X\. Xie, X\. Liu, X\. Yang, X\. Li, X\. Su, X\. Lin, X\. Q\. Li, X\. Jin, X\. Shen, X\. Chen, X\. Sun, X\. Wang, X\. Song, X\. Zhou, X\. Wang, X\. Shan, Y\. K\. Li, Y\. Q\. Wang, Y\. X\. Wei, Y\. Zhang, Y\. Xu, Y\. Li, Y\. Zhao, Y\. Sun, Y\. Wang, Y\. Yu, Y\. Zhang, Y\. Shi, Y\. Xiong, Y\. He, Y\. Piao, Y\. Wang, Y\. Tan, Y\. Ma, Y\. Liu, Y\. Guo, Y\. Ou, Y\. Wang, Y\. Gong, Y\. Zou, Y\. He, Y\. Xiong, Y\. Luo, Y\. You, Y\. Liu, Y\. Zhou, Y\. X\. Zhu, Y\. Huang, Y\. Li, Y\. Zheng, Y\. Zhu, Y\. Ma, Y\. Tang, Y\. Zha, Y\. Yan, Z\. Z\. Ren, Z\. Ren, Z\. Sha, Z\. Fu, Z\. Xu, Z\. Xie, Z\. Zhang, Z\. Hao, Z\. Ma, Z\. Yan, Z\. Wu, Z\. Gu, Z\. Zhu, Z\. Liu, Z\. Li, Z\. Xie, Z\. Song, Z\. Pan, Z\. Huang, Z\. Xu, Z\. Zhang, and Z\. Zhang \(2025\)DeepSeek\-r1 incentivizes reasoning in llms through reinforcement learning\.Nature645\(8081\),pp\. 633–638\.External Links:ISSN 1476\-4687,[Link](http://dx.doi.org/10.1038/s41586-025-09422-z),[Document](https://dx.doi.org/10.1038/s41586-025-09422-z)Cited by:[§4\.1](https://arxiv.org/html/2606.30813#S4.SS1.SSS0.Px2.p1.5)\.
- V\. Gupta, T\. Koren, and Y\. Singer \(2018\)Shampoo: preconditioned stochastic tensor optimization\.External Links:1802\.09568,[Link](https://arxiv.org/abs/1802.09568)Cited by:[§7](https://arxiv.org/html/2606.30813#S7.SS0.SSS0.Px2.p1.1)\.
- K\. He, X\. Zhang, S\. Ren, and J\. Sun \(2015\)Deep residual learning for image recognition\.External Links:1512\.03385,[Link](https://arxiv.org/abs/1512.03385)Cited by:[§1](https://arxiv.org/html/2606.30813#S1.p1.1)\.
- C\. R\. Hoyt and A\. B\. Owen \(2021\)Probing neural networks with t\-sne, class\-specific projections and a guided tour\.External Links:2107\.12547Cited by:[§7](https://arxiv.org/html/2606.30813#S7.SS0.SSS0.Px3.p2.1)\.
- J\. Jiang, J\. Zhou, and Z\. Zhu \(2025a\)On layer\-wise representation similarity: application for multi\-exit models with a single classifier\.InNeurIPS 2024 Workshop on Symmetry and Geometry in Neural Representations,External Links:[Link](https://openreview.net/forum?id=YanMgtZhfY)Cited by:[§7](https://arxiv.org/html/2606.30813#S7.SS0.SSS0.Px1.p1.1)\.
- J\. Jiang, J\. Zhou, and Z\. Zhu \(2025b\)Tracing representation progression: analyzing and enhancing layer\-wise similarity\.External Links:2406\.14479,[Link](https://arxiv.org/abs/2406.14479)Cited by:[§7](https://arxiv.org/html/2606.30813#S7.SS0.SSS0.Px1.p1.1)\.
- K\. Jordan, Y\. Jin, V\. Boza, J\. You, F\. Cesista, L\. Newhouse, and J\. Bernstein \(2024\)Muon: an optimizer for hidden layers in neural networks\.External Links:[Link](https://kellerjordan.github.io/posts/muon/)Cited by:[§A\.1](https://arxiv.org/html/2606.30813#A1.SS1.SSS0.Px4.p1.4),[§A\.1](https://arxiv.org/html/2606.30813#A1.SS1.p1.1),[§4\.1](https://arxiv.org/html/2606.30813#S4.SS1.SSS0.Px3.p1.5),[§7](https://arxiv.org/html/2606.30813#S7.SS0.SSS0.Px2.p1.1)\.
- F\. Kapl, E\. Angelis, T\. Höppe, K\. Maile, J\. von Oswald, N\. Scherrer, and S\. Bauer \(2025\)Do depth\-grown models overcome the curse of depth? an in\-depth analysis\.External Links:2512\.08819,[Link](https://arxiv.org/abs/2512.08819)Cited by:[§7](https://arxiv.org/html/2606.30813#S7.SS0.SSS0.Px1.p1.1)\.
- A\. Karpathy \(2025\)Nanochat: the best chatgpt that $100 can buy\.GitHub\.External Links:[Link](https://github.com/karpathy/nanochat)Cited by:[§4\.1](https://arxiv.org/html/2606.30813#S4.SS1.SSS0.Px3.p1.5)\.
- D\. P\. Kingma and J\. Ba \(2017\)Adam: a method for stochastic optimization\.External Links:1412\.6980,[Link](https://arxiv.org/abs/1412.6980)Cited by:[§A\.1](https://arxiv.org/html/2606.30813#A1.SS1.p1.1),[§7](https://arxiv.org/html/2606.30813#S7.SS0.SSS0.Px2.p1.1)\.
- F\. Krause, T\. Phan, M\. Gui, S\. A\. Baumann, V\. T\. Hu, and B\. Ommer \(2025\)TREAD: token routing for efficient architecture\-agnostic diffusion training\.External Links:2501\.04765,[Link](https://arxiv.org/abs/2501.04765)Cited by:[§7](https://arxiv.org/html/2606.30813#S7.SS0.SSS0.Px1.p1.1)\.
- V\. Lad, J\. H\. Lee, W\. Gurnee, and M\. Tegmark \(2025\)The remarkable robustness of llms: stages of inference?\.External Links:2406\.19384,[Link](https://arxiv.org/abs/2406.19384)Cited by:[§7](https://arxiv.org/html/2606.30813#S7.SS0.SSS0.Px1.p1.1)\.
- J\. Li and V\. Papyan \(2024\)Residual alignment: uncovering the mechanisms of residual networks\.External Links:2401\.09018,[Link](https://arxiv.org/abs/2401.09018)Cited by:[§1](https://arxiv.org/html/2606.30813#S1.p1.1),[§7](https://arxiv.org/html/2606.30813#S7.SS0.SSS0.Px1.p1.1)\.
- Z\. Li, L\. Liu, C\. Liang, W\. Chen, and T\. Zhao \(2025\)NorMuon: making muon more efficient and scalable\.External Links:2510\.05491,[Link](https://arxiv.org/abs/2510.05491)Cited by:[§4\.1](https://arxiv.org/html/2606.30813#S4.SS1.SSS0.Px3.p1.5)\.
- W\. Lin, F\. Dangel, R\. Eschenhagen, K\. Neklyudov, A\. Kristiadi, R\. E\. Turner, and A\. Makhzani \(2024\)Structured inverse\-free natural gradient: memory\-efficient & numerically\-stable kfac\.External Links:2312\.05705,[Link](https://arxiv.org/abs/2312.05705)Cited by:[§7](https://arxiv.org/html/2606.30813#S7.SS0.SSS0.Px3.p1.1)\.
- I\. Loshchilov and F\. Hutter \(2019\)Decoupled weight decay regularization\.External Links:1711\.05101,[Link](https://arxiv.org/abs/1711.05101)Cited by:[§A\.1](https://arxiv.org/html/2606.30813#A1.SS1.p1.1),[§7](https://arxiv.org/html/2606.30813#S7.SS0.SSS0.Px2.p1.1)\.
- J\. Martens and R\. Grosse \(2020\)Optimizing neural networks with kronecker\-factored approximate curvature\.External Links:1503\.05671,[Link](https://arxiv.org/abs/1503.05671)Cited by:[§7](https://arxiv.org/html/2606.30813#S7.SS0.SSS0.Px3.p1.1)\.
- X\. Men, M\. Xu, Q\. Zhang, B\. Wang, H\. Lin, Y\. Lu, X\. Han, and W\. Chen \(2024\)ShortGPT: layers in large language models are more redundant than you expect\.External Links:2403\.03853,[Link](https://arxiv.org/abs/2403.03853)Cited by:[§7](https://arxiv.org/html/2606.30813#S7.SS0.SSS0.Px1.p1.1)\.
- A\. Nagwekar \(2025\)Towards guided descent: optimization algorithms for training neural networks at scale\.External Links:2512\.18373,[Link](https://arxiv.org/abs/2512.18373)Cited by:[§7](https://arxiv.org/html/2606.30813#S7.SS0.SSS0.Px3.p1.1)\.
- V\. Papyan, X\. Y\. Han, and D\. L\. Donoho \(2020\)Prevalence of neural collapse during the terminal phase of deep learning training\.Proceedings of the National Academy of Sciences117\(40\),pp\. 24652–24663\.External Links:[Document](https://dx.doi.org/10.1073/pnas.2015509117),[Link](https://www.pnas.org/doi/abs/10.1073/pnas.2015509117),https://www\.pnas\.org/doi/pdf/10\.1073/pnas\.2015509117Cited by:[§7](https://arxiv.org/html/2606.30813#S7.SS0.SSS0.Px3.p2.1)\.
- V\. Papyan \(2020\)Traces of class/cross\-class structure pervade deep learning spectra\.External Links:2008\.11865Cited by:[§7](https://arxiv.org/html/2606.30813#S7.SS0.SSS0.Px3.p2.1)\.
- L\. Parker, E\. Onal, A\. Stengel, and J\. Intrater \(2023\)Neural collapse in the intermediate hidden layers of classification neural networks\.External Links:2308\.02760Cited by:[§7](https://arxiv.org/html/2606.30813#S7.SS0.SSS0.Px3.p2.1)\.
- A\. Patrawala, J\. Feng, E\. Jones, and J\. Steinhardt \(2026\)LLM layers immediately correct each other\.InThe Thirty\-ninth Annual Conference on Neural Information Processing Systems,External Links:[Link](https://openreview.net/forum?id=7DY7kB8wyZ)Cited by:[§7](https://arxiv.org/html/2606.30813#S7.SS0.SSS0.Px1.p1.1)\.
- M\. Piotrowski, P\. M\. Riechers, D\. Filan, and A\. S\. Shai \(2025\)Constrained belief updates explain geometric structures in transformer representations\.External Links:2502\.01954,[Link](https://arxiv.org/abs/2502.01954)Cited by:[§7](https://arxiv.org/html/2606.30813#S7.SS0.SSS0.Px3.p2.1)\.
- R\. Sarfati, T\. J\. B\. Liu, N\. Boullé, and C\. J\. Earls \(2024\)Lines of thought in large language models\.External Links:2410\.01545,[Link](https://arxiv.org/abs/2410.01545)Cited by:[§7](https://arxiv.org/html/2606.30813#S7.SS0.SSS0.Px3.p2.1)\.
- A\. S\. Shai, S\. E\. Marzen, L\. Teixeira, A\. G\. Oldenziel, and P\. M\. Riechers \(2025\)Transformers represent belief state geometry in their residual stream\.External Links:2405\.15943,[Link](https://arxiv.org/abs/2405.15943)Cited by:[§7](https://arxiv.org/html/2606.30813#S7.SS0.SSS0.Px3.p2.1)\.
- Z\. Shao, P\. Wang, Q\. Zhu, R\. Xu, J\. Song, X\. Bi, H\. Zhang, M\. Zhang, Y\. K\. Li, Y\. Wu, and D\. Guo \(2024\)DeepSeekMath: pushing the limits of mathematical reasoning in open language models\.External Links:2402\.03300,[Link](https://arxiv.org/abs/2402.03300)Cited by:[§4\.1](https://arxiv.org/html/2606.30813#S4.SS1.SSS0.Px2.p1.5)\.
- O\. Skean, M\. R\. Arefin, Y\. LeCun, and R\. Shwartz\-Ziv \(2024\)Does representation matter? exploring intermediate layers in large language models\.External Links:2412\.09563,[Link](https://arxiv.org/abs/2412.09563)Cited by:[§1](https://arxiv.org/html/2606.30813#S1.p1.1),[§7](https://arxiv.org/html/2606.30813#S7.SS0.SSS0.Px1.p1.1),[§7](https://arxiv.org/html/2606.30813#S7.SS0.SSS0.Px3.p2.1)\.
- O\. Skean, M\. R\. Arefin, D\. Zhao, N\. Patel, J\. Naghiyev, Y\. LeCun, and R\. Shwartz\-Ziv \(2025\)Layer by layer: uncovering hidden representations in language models\.External Links:2502\.02013,[Link](https://arxiv.org/abs/2502.02013)Cited by:[§7](https://arxiv.org/html/2606.30813#S7.SS0.SSS0.Px1.p1.1),[§7](https://arxiv.org/html/2606.30813#S7.SS0.SSS0.Px3.p2.1)\.
- Z\. Song, Z\. Li, Q\. Cao, M\. Luo, and H\. X\. Zhu \(2025\)Bridging the dimensional chasm: uncover layer\-wise dimensional reduction in transformers through token correlation\.External Links:2503\.22547,[Link](https://arxiv.org/abs/2503.22547)Cited by:[§7](https://arxiv.org/html/2606.30813#S7.SS0.SSS0.Px3.p2.1)\.
- P\. Súkeník, M\. Mondelli, and C\. Lampert \(2023\)Deep neural collapse is provably optimal for the deep unconstrained features model\.External Links:2305\.13165Cited by:[§7](https://arxiv.org/html/2606.30813#S7.SS0.SSS0.Px3.p2.1)\.
- H\. Touvron, M\. Cord, M\. Douze, F\. Massa, A\. Sablayrolles, and H\. Jégou \(2021\)Training data\-efficient image transformers & distillation through attention\.External Links:2012\.12877,[Link](https://arxiv.org/abs/2012.12877)Cited by:[§4\.1](https://arxiv.org/html/2606.30813#S4.SS1.SSS0.Px4.p1.1)\.
- A\. Vaswani, N\. Shazeer, N\. Parmar, J\. Uszkoreit, L\. Jones, A\. N\. Gomez, L\. Kaiser, and I\. Polosukhin \(2023\)Attention is all you need\.External Links:1706\.03762,[Link](https://arxiv.org/abs/1706.03762)Cited by:[§1](https://arxiv.org/html/2606.30813#S1.p1.1)\.
- N\. Vyas, D\. Morwani, R\. Zhao, M\. Kwun, I\. Shapira, D\. Brandfonbrener, L\. Janson, and S\. Kakade \(2025\)SOAP: improving and stabilizing shampoo using adam\.External Links:2409\.11321,[Link](https://arxiv.org/abs/2409.11321)Cited by:[§7](https://arxiv.org/html/2606.30813#S7.SS0.SSS0.Px2.p1.1)\.
- J\. Wang, L\. Fan, D\. Zhang, W\. Jing, D\. Di, Y\. Song, S\. Liu, and C\. Cong \(2026\)Visual prompt\-agnostic evolution\.External Links:2601\.20232,[Link](https://arxiv.org/abs/2601.20232)Cited by:[§7](https://arxiv.org/html/2606.30813#S7.SS0.SSS0.Px3.p2.1)\.
- P\. Wang, X\. Li, C\. Yaras, Z\. Zhu, L\. Balzano, W\. Hu, and Q\. Qu \(2024a\)Understanding deep representation learning via layerwise feature compression and discrimination\.External Links:2311\.02960Cited by:[§7](https://arxiv.org/html/2606.30813#S7.SS0.SSS0.Px3.p2.1)\.
- S\. Wang, K\. Gai, and S\. Zhang \(2024b\)Progressive feedforward collapse of resnet training\.External Links:2405\.00985Cited by:[§7](https://arxiv.org/html/2606.30813#S7.SS0.SSS0.Px3.p2.1)\.
- C\. Wolfram and A\. Schein \(2025\)Layers at similar depths generate similar activations across llm architectures\.External Links:2504\.08775,[Link](https://arxiv.org/abs/2504.08775)Cited by:[§7](https://arxiv.org/html/2606.30813#S7.SS0.SSS0.Px1.p1.1)\.
- R\. Wu and V\. Papyan \(2024\)Linguistic collapse: neural collapse in \(large\) language models\.External Links:2405\.17767,[Link](https://arxiv.org/abs/2405.17767)Cited by:[§7](https://arxiv.org/html/2606.30813#S7.SS0.SSS0.Px3.p2.1)\.
- A\. Yang, B\. Yang, B\. Hui, B\. Zheng, B\. Yu, C\. Zhou, C\. Li, C\. Li, D\. Liu, F\. Huang, G\. Dong, H\. Wei, H\. Lin, J\. Tang, J\. Wang, J\. Yang, J\. Tu, J\. Zhang, J\. Ma, J\. Yang, J\. Xu, J\. Zhou, J\. Bai, J\. He, J\. Lin, K\. Dang, K\. Lu, K\. Chen, K\. Yang, M\. Li, M\. Xue, N\. Ni, P\. Zhang, P\. Wang, R\. Peng, R\. Men, R\. Gao, R\. Lin, S\. Wang, S\. Bai, S\. Tan, T\. Zhu, T\. Li, T\. Liu, W\. Ge, X\. Deng, X\. Zhou, X\. Ren, X\. Zhang, X\. Wei, X\. Ren, X\. Liu, Y\. Fan, Y\. Yao, Y\. Zhang, Y\. Wan, Y\. Chu, Y\. Liu, Z\. Cui, Z\. Zhang, Z\. Guo, and Z\. Fan \(2024\)Qwen2 technical report\.External Links:2407\.10671,[Link](https://arxiv.org/abs/2407.10671)Cited by:[§4\.1](https://arxiv.org/html/2606.30813#S4.SS1.SSS0.Px2.p1.5)\.
- E\. Zangrando, P\. Deidda, S\. Brugiapaglia, N\. Guglielmi, and F\. Tudisco \(2025\)Provable emergence of deep neural collapse and low\-rank bias inL2L^\{2\}\-regularized nonlinear networks\.External Links:2402\.03991,[Link](https://arxiv.org/abs/2402.03991)Cited by:[§7](https://arxiv.org/html/2606.30813#S7.SS0.SSS0.Px3.p2.1)\.
- J\. Zarka, F\. Guth, and S\. Mallat \(2021\)Separation and concentration in deep networks\.External Links:2012\.10424,[Link](https://arxiv.org/abs/2012.10424)Cited by:[§7](https://arxiv.org/html/2606.30813#S7.SS0.SSS0.Px3.p2.1)\.
- N\. Zhou, J\. Chen, and D\. Huang \(2025\)Sharing task\-relevant information in visual prompt tuning by cross\-layer dynamic connection\.IEEE Transactions on Image Processing34,pp\. 4527–4540\.External Links:ISSN 1941\-0042,[Link](http://dx.doi.org/10.1109/tip.2025.3587587),[Document](https://dx.doi.org/10.1109/tip.2025.3587587)Cited by:[§7](https://arxiv.org/html/2606.30813#S7.SS0.SSS0.Px3.p2.1)\.

## Appendix AAppendix

### A\.1Specialization of Smoothing to Adam, AdamW, and Muon

In our experiments, the base optimizer𝒰\(t\)\\mathcal\{U\}^\{\(t\)\}is typically Adam\(Kingma and Ba,[2017](https://arxiv.org/html/2606.30813#bib.bib22)\), AdamW\(Loshchilov and Hutter,[2019](https://arxiv.org/html/2606.30813#bib.bib23)\), or Muon\(Jordanet al\.,[2024](https://arxiv.org/html/2606.30813#bib.bib27)\)\. For completeness, we briefly recall the update rule and describe how gradient smoothing composes with it\.

#### Adam updates\.

Letg\(t\)=∇θℒ​\(θ\(t\),ϕ\(t\)\)g^\{\(t\)\}=\\nabla\_\{\\theta\}\\mathcal\{L\}\(\\theta^\{\(t\)\},\\phi^\{\(t\)\}\)denote the gradient with respect to the repeated\-block parameters\. Note that the remaining parametersϕ\\phiare updated in the same manner as defined below, but are just not involved in the smoothing process\. Adam maintains exponential moving averages of first and second moments,

m\(t\)\\displaystyle m^\{\(t\)\}=β1​m\(t−1\)\+\(1−β1\)​g\(t\),\\displaystyle=\\beta\_\{1\}m^\{\(t\-1\)\}\+\(1\-\\beta\_\{1\}\)g^\{\(t\)\},\(13\)v\(t\)\\displaystyle v^\{\(t\)\}=β2​v\(t−1\)\+\(1−β2\)​\(g\(t\)⊙g\(t\)\),\\displaystyle=\\beta\_\{2\}v^\{\(t\-1\)\}\+\(1\-\\beta\_\{2\}\)\\big\(g^\{\(t\)\}\\odot g^\{\(t\)\}\\big\),\(14\)with bias\-corrected versionsm^\(t\)\\hat\{m\}^\{\(t\)\}andv^\(t\)\\hat\{v\}^\{\(t\)\}\. The update direction is

u\(t\):=m^\(t\)v^\(t\)\+ε,u^\{\(t\)\}:=\\frac\{\\hat\{m\}^\{\(t\)\}\}\{\\sqrt\{\\hat\{v\}^\{\(t\)\}\}\+\\varepsilon\},\(15\)and the block parameters are updated as

θ\(t\+1\)=θ\(t\)−η​u\(t\)\.\\theta^\{\(t\+1\)\}=\\theta^\{\(t\)\}\-\\eta\\,u^\{\(t\)\}\.\(16\)
Equivalently, Adam applies a diagonal, coordinate\-wise preconditioner

u\(t\)=A\(t\)​g\(t\),A\(t\):=diag​\(1v^\(t\)\+ε\)\.u^\{\(t\)\}=A^\{\(t\)\}g^\{\(t\)\},\\qquad A^\{\(t\)\}:=\\mathrm\{diag\}\\\!\\left\(\\frac\{1\}\{\\sqrt\{\\hat\{v\}^\{\(t\)\}\}\+\\varepsilon\}\\right\)\.

#### AdamW\.

AdamW applies decoupled weight decay, yielding the block\-parameter update

θ\(t\+1\)=θ\(t\)−η​\(u\(t\)\+λ​θ\(t\)\),\\theta^\{\(t\+1\)\}=\\theta^\{\(t\)\}\-\\eta\\big\(u^\{\(t\)\}\+\\lambda\\,\\theta^\{\(t\)\}\\big\),\(17\)with non\-block parametersϕ\\phiupdated by the base optimizer in the standard way\.

#### Adam/AdamW with Gradient Smoothing\.

The smoothed block update is then

u~\(t\)=P​u\(t\)\.\\tilde\{u\}^\{\(t\)\}=P\\,u^\{\(t\)\}\.Gradient smoothing yields

θ\(t\+1\)\\displaystyle\\theta^\{\(t\+1\)\}=θ\(t\)−η​u~\(t\)\(Adam\),\\displaystyle=\\theta^\{\(t\)\}\-\\eta\\,\\tilde\{u\}^\{\(t\)\}\\quad\\text\{\(Adam\)\},\(18\)θ\(t\+1\)\\displaystyle\\theta^\{\(t\+1\)\}=θ\(t\)−η​\(u~\(t\)\+λ​θ\(t\)\)\(AdamW\)\.\\displaystyle=\\theta^\{\(t\)\}\-\\eta\\big\(\\tilde\{u\}^\{\(t\)\}\+\\lambda\\,\\theta^\{\(t\)\}\\big\)\\quad\\text\{\(AdamW\)\}\.\(19\)
Equivalently, the effective update can be written as

θ\(t\+1\)=θ\(t\)−η​\(P​A\(t\)\)​g\(t\),\\theta^\{\(t\+1\)\}=\\theta^\{\(t\)\}\-\\eta\\,\(PA^\{\(t\)\}\)g^\{\(t\)\},highlighting that gradient smoothing composes a depth\-structured preconditioner with Adam’s coordinate\-wise preconditioning, withϕ\\phibeing updated as in standard Adam/AdamW\.

#### Muon\.

Muon\(Jordanet al\.,[2024](https://arxiv.org/html/2606.30813#bib.bib27)\)treats each repeated\-block parameter as a matrix and produces an approximately orthogonal update direction\. It maintains a momentum buffer

M\(t\)=μ​M\(t−1\)\+g\(t\),M^\{\(t\)\}=\\mu\\,M^\{\(t\-1\)\}\+g^\{\(t\)\},and forms the update direction by orthogonalizingM\(t\)M^\{\(t\)\}via a Newton–Schulz iteration,

u\(t\):=NS​\(M\(t\)\)≈M\(t\)​\(\(M\(t\)\)⊤​M\(t\)\)−1/2\.u^\{\(t\)\}:=\\mathrm\{NS\}\\\!\\big\(M^\{\(t\)\}\\big\)\\;\\approx\\;M^\{\(t\)\}\\big\(\(M^\{\(t\)\}\)^\{\\top\}M^\{\(t\)\}\\big\)^\{\-1/2\}\.The block parameters are updated asθ\(t\+1\)=θ\(t\)−η​u\(t\)\\theta^\{\(t\+1\)\}=\\theta^\{\(t\)\}\-\\eta\\,u^\{\(t\)\}\. In contrast to Adam’s diagonal preconditioning,u\(t\)u^\{\(t\)\}has singular values approximately equal to one, so its per\-step magnitude is largely decoupled from that of the raw gradient\.

#### Muon with Gradient Smoothing\.

The smoothed block update composes with Muon in the same way as with Adam:

u~\(t\)=P​u\(t\),θ\(t\+1\)=θ\(t\)−η​u~\(t\)\.\\tilde\{u\}^\{\(t\)\}=P\\,u^\{\(t\)\},\\qquad\\theta^\{\(t\+1\)\}=\\theta^\{\(t\)\}\-\\eta\\,\\tilde\{u\}^\{\(t\)\}\.For matrix\-valued block parameters,PPis applied to the orthogonalized update produced by the Newton–Schulz iteration; for non\-matrix block parameters \(e\.g\., normalization scales\), to which Newton–Schulz does not apply,PPis applied directly to the base\-optimizer update for that parameter\.

### A\.2Alignment of Representation Residuals under Gradient Smoothing

#### Residual increments and alignment metric\.

For any depth\-stacked parameter vectorϑ=\(ϑ1,…,ϑL\)\\vartheta=\(\\vartheta\_\{1\},\\dots,\\vartheta\_\{L\}\), define the residual increments

rℓ​\(ϑ\):=hℓ\+1​\(ϑ\)−hℓ​\(ϑ\)=F​\(hℓ​\(ϑ\);ϑℓ\)∈ℝd,ℓ=1,…,L,r\_\{\\ell\}\(\\vartheta\)\\;:=\\;h\_\{\\ell\+1\}\(\\vartheta\)\-h\_\{\\ell\}\(\\vartheta\)\\;=\\;F\\\!\\big\(h\_\{\\ell\}\(\\vartheta\);\\vartheta\_\{\\ell\}\\big\)\\in\\mathbb\{R\}^\{d\},\\qquad\\ell=1,\\dots,L,\(20\)and the normalized directions

dℓ​\(ϑ\):=rℓ​\(ϑ\)‖rℓ​\(ϑ\)‖∈ℝd\(defined wheneverrℓ​\(ϑ\)≠0\)\.d\_\{\\ell\}\(\\vartheta\)\\;:=\\;\\frac\{r\_\{\\ell\}\(\\vartheta\)\}\{\\\|r\_\{\\ell\}\(\\vartheta\)\\\|\}\\in\\mathbb\{R\}^\{d\}\\quad\\text\{\(defined whenever $r\_\{\\ell\}\(\\vartheta\)\\neq 0$\)\.\}\(21\)Our empirical metric is

cos⁡\(rℓ\+1​\(ϑ\),rℓ​\(ϑ\)\)=⟨dℓ\+1​\(ϑ\),dℓ​\(ϑ\)⟩,ℓ=1,…,L−1\.\\cos\\\!\\big\(r\_\{\\ell\+1\}\(\\vartheta\),r\_\{\\ell\}\(\\vartheta\)\\big\)\\;=\\;\\big\\langle d\_\{\\ell\+1\}\(\\vartheta\),d\_\{\\ell\}\(\\vartheta\)\\big\\rangle,\\qquad\\ell=1,\\dots,L\-1\.\(22\)Working withdℓd\_\{\\ell\}isolates*directional*alignment from magnitude effects \(e\.g\.‖rℓ\+1−rℓ‖\\\|r\_\{\\ell\+1\}\-r\_\{\\ell\}\\\|may be large even whencos⁡\(rℓ\+1,rℓ\)\\cos\(r\_\{\\ell\+1\},r\_\{\\ell\}\)is close to11due to norm mismatch\)\.

#### Assumptions\.

We work on a region containing the trajectories\{hℓ​\(θb\)\}\\\{h\_\{\\ell\}\(\\theta^\{b\}\)\\\}and\{hℓ​\(θs\)\}\\\{h\_\{\\ell\}\(\\theta^\{s\}\)\\\}and impose the following standard regularity assumptions\.

###### Assumption A\.1\.

There exist constantsLh,Lθ≥0L\_\{h\},L\_\{\\theta\}\\geq 0such that for allh,h¯∈ℝdh,\\bar\{h\}\\in\\mathbb\{R\}^\{d\}andθ,θ¯∈ℝp\\theta,\\bar\{\\theta\}\\in\\mathbb\{R\}^\{p\}in the region of interest,

‖F​\(h;θ\)−F​\(h¯;θ¯\)‖≤Lh​‖h−h¯‖\+Lθ​‖θ−θ¯‖\.\\\|F\(h;\\theta\)\-F\(\\bar\{h\};\\bar\{\\theta\}\)\\\|\\;\\leq\\;L\_\{h\}\\\|h\-\\bar\{h\}\\\|\+L\_\{\\theta\}\\\|\\theta\-\\bar\{\\theta\}\\\|\.\(23\)
Moreover, there exist constants0<m≤M<∞0<m\\leq M<\\inftysuch that, for eachϑ∈\{θb,θs\}\\vartheta\\in\\\{\\theta^\{b\},\\theta^\{s\}\\\},

m≤minℓ∈\[L\]⁡‖rℓ​\(ϑ\)‖,maxℓ∈\[L\]⁡‖rℓ​\(ϑ\)‖≤M\.m\\leq\\min\_\{\\ell\\in\[L\]\}\\\|r\_\{\\ell\}\(\\vartheta\)\\\|,\\qquad\\max\_\{\\ell\\in\[L\]\}\\\|r\_\{\\ell\}\(\\vartheta\)\\\|\\leq M\.

###### Lemma A\.2\.

For any nonzeroa,b∈ℝda,b\\in\\mathbb\{R\}^\{d\},

1−cos⁡\(a,b\)=12​‖a‖a‖−b‖b‖‖2\.1\-\\cos\(a,b\)\\;=\\;\\frac\{1\}\{2\}\\left\\\|\\frac\{a\}\{\\\|a\\\|\}\-\\frac\{b\}\{\\\|b\\\|\}\\right\\\|^\{2\}\.\(24\)

###### Proof\.

Expand‖a/‖a‖−b/‖b‖‖2=2−2​⟨a/‖a‖,b/‖b‖⟩\\\|a/\\\|a\\\|\-b/\\\|b\\\|\\\|^\{2\}=2\-2\\langle a/\\\|a\\\|,b/\\\|b\\\|\\rangleand rearrange\. ∎

###### Lemma A\.3\.

For any nonzeroa,b∈ℝda,b\\in\\mathbb\{R\}^\{d\},

‖a‖a‖−b‖b‖‖≤2​‖a−b‖min⁡\(‖a‖,‖b‖\)\.\\left\\\|\\frac\{a\}\{\\\|a\\\|\}\-\\frac\{b\}\{\\\|b\\\|\}\\right\\\|\\;\\leq\\;\\frac\{2\\\|a\-b\\\|\}\{\\min\(\\\|a\\\|,\\\|b\\\|\)\}\.\(25\)

###### Proof\.

Write

a‖a‖−b‖b‖=a−b‖a‖\+b​\(1‖a‖−1‖b‖\),\\frac\{a\}\{\\\|a\\\|\}\-\\frac\{b\}\{\\\|b\\\|\}=\\frac\{a\-b\}\{\\\|a\\\|\}\+b\\Big\(\\frac\{1\}\{\\\|a\\\|\}\-\\frac\{1\}\{\\\|b\\\|\}\\Big\),so by the triangle inequality,

‖a‖a‖−b‖b‖‖≤‖a−b‖‖a‖\+‖b‖⋅\|‖b‖−‖a‖\|‖a‖​‖b‖≤‖a−b‖‖a‖\+‖a−b‖‖a‖=2​‖a−b‖‖a‖,\\left\\\|\\frac\{a\}\{\\\|a\\\|\}\-\\frac\{b\}\{\\\|b\\\|\}\\right\\\|\\leq\\frac\{\\\|a\-b\\\|\}\{\\\|a\\\|\}\+\\\|b\\\|\\cdot\\frac\{\\big\|\\\|b\\\|\-\\\|a\\\|\\big\|\}\{\\\|a\\\|\\\|b\\\|\}\\leq\\frac\{\\\|a\-b\\\|\}\{\\\|a\\\|\}\+\\frac\{\\\|a\-b\\\|\}\{\\\|a\\\|\}=\\frac\{2\\\|a\-b\\\|\}\{\\\|a\\\|\},and symmetrizing in\(a,b\)\(a,b\)yields \([25](https://arxiv.org/html/2606.30813#A1.E25)\)\. ∎

#### A bound on cosine misalignment\.

###### Lemma A\.4\.

Under Assumption[A\.1](https://arxiv.org/html/2606.30813#A1.Thmtheorem1), for anyϑ∈\{θb,θs\}\\vartheta\\in\\\{\\theta^\{b\},\\theta^\{s\}\\\}and anyℓ∈\[L−1\]\\ell\\in\[L\-1\],

1−cos⁡\(rℓ\+1​\(ϑ\),rℓ​\(ϑ\)\)≤4m2​\(Lh2​‖rℓ​\(ϑ\)‖2\+Lθ2​‖ϑℓ\+1−ϑℓ‖2\)\.1\-\\cos\\\!\\big\(r\_\{\\ell\+1\}\(\\vartheta\),r\_\{\\ell\}\(\\vartheta\)\\big\)\\;\\leq\\;\\frac\{4\}\{m^\{2\}\}\\Big\(L\_\{h\}^\{2\}\\\|r\_\{\\ell\}\(\\vartheta\)\\\|^\{2\}\+L\_\{\\theta\}^\{2\}\\\|\\vartheta\_\{\\ell\+1\}\-\\vartheta\_\{\\ell\}\\\|^\{2\}\\Big\)\.\(26\)Consequently,

1L−1​∑ℓ=1L−1cos⁡\(rℓ\+1​\(ϑ\),rℓ​\(ϑ\)\)≥1−4​Lh2​M2m2−4​Lθ2m2⋅1L−1​∑ℓ=1L−1‖ϑℓ\+1−ϑℓ‖2\.\\frac\{1\}\{L\-1\}\\sum\_\{\\ell=1\}^\{L\-1\}\\cos\\\!\\big\(r\_\{\\ell\+1\}\(\\vartheta\),r\_\{\\ell\}\(\\vartheta\)\\big\)\\;\\geq\\;1\-\\frac\{4L\_\{h\}^\{2\}M^\{2\}\}\{m^\{2\}\}\-\\frac\{4L\_\{\\theta\}^\{2\}\}\{m^\{2\}\}\\cdot\\frac\{1\}\{L\-1\}\\sum\_\{\\ell=1\}^\{L\-1\}\\\|\\vartheta\_\{\\ell\+1\}\-\\vartheta\_\{\\ell\}\\\|^\{2\}\.\(27\)

###### Proof\.

Fixϑ\\varthetaandℓ\\ell\. By Lemma[A\.2](https://arxiv.org/html/2606.30813#A1.Thmtheorem2)and Lemma[A\.3](https://arxiv.org/html/2606.30813#A1.Thmtheorem3),

1−cos⁡\(rℓ\+1,rℓ\)=12​‖dℓ\+1−dℓ‖2≤12​\(2​‖rℓ\+1−rℓ‖min⁡\(‖rℓ\+1‖,‖rℓ‖\)\)2≤2m2​‖rℓ\+1−rℓ‖2\.1\-\\cos\(r\_\{\\ell\+1\},r\_\{\\ell\}\)=\\frac\{1\}\{2\}\\\|d\_\{\\ell\+1\}\-d\_\{\\ell\}\\\|^\{2\}\\leq\\frac\{1\}\{2\}\\left\(\\frac\{2\\\|r\_\{\\ell\+1\}\-r\_\{\\ell\}\\\|\}\{\\min\(\\\|r\_\{\\ell\+1\}\\\|,\\\|r\_\{\\ell\}\\\|\)\}\\right\)^\{2\}\\leq\\frac\{2\}\{m^\{2\}\}\\\|r\_\{\\ell\+1\}\-r\_\{\\ell\}\\\|^\{2\}\.Using \([20](https://arxiv.org/html/2606.30813#A1.E20)\) and the Lipschitz property \([23](https://arxiv.org/html/2606.30813#A1.E23)\),

‖rℓ\+1−rℓ‖=‖F​\(hℓ\+1;ϑℓ\+1\)−F​\(hℓ;ϑℓ\)‖≤Lh​‖hℓ\+1−hℓ‖\+Lθ​‖ϑℓ\+1−ϑℓ‖=Lh​‖rℓ‖\+Lθ​‖ϑℓ\+1−ϑℓ‖\.\\\|r\_\{\\ell\+1\}\-r\_\{\\ell\}\\\|=\\big\\\|F\(h\_\{\\ell\+1\};\\vartheta\_\{\\ell\+1\}\)\-F\(h\_\{\\ell\};\\vartheta\_\{\\ell\}\)\\big\\\|\\leq L\_\{h\}\\\|h\_\{\\ell\+1\}\-h\_\{\\ell\}\\\|\+L\_\{\\theta\}\\\|\\vartheta\_\{\\ell\+1\}\-\\vartheta\_\{\\ell\}\\\|=L\_\{h\}\\\|r\_\{\\ell\}\\\|\+L\_\{\\theta\}\\\|\\vartheta\_\{\\ell\+1\}\-\\vartheta\_\{\\ell\}\\\|\.Squaring and applying\(a\+b\)2≤2​a2\+2​b2\(a\+b\)^\{2\}\\leq 2a^\{2\}\+2b^\{2\}yields

‖rℓ\+1−rℓ‖2≤2​Lh2​‖rℓ‖2\+2​Lθ2​‖ϑℓ\+1−ϑℓ‖2\.\\\|r\_\{\\ell\+1\}\-r\_\{\\ell\}\\\|^\{2\}\\leq 2L\_\{h\}^\{2\}\\\|r\_\{\\ell\}\\\|^\{2\}\+2L\_\{\\theta\}^\{2\}\\\|\\vartheta\_\{\\ell\+1\}\-\\vartheta\_\{\\ell\}\\\|^\{2\}\.Combining the last two inequalities gives \([26](https://arxiv.org/html/2606.30813#A1.E26)\)\. Averaging \([26](https://arxiv.org/html/2606.30813#A1.E26)\) overℓ\\elland using‖rℓ​\(ϑ\)‖≤M\\\|r\_\{\\ell\}\(\\vartheta\)\\\|\\leq Mgives \([27](https://arxiv.org/html/2606.30813#A1.E27)\)\. ∎

#### Local averaging smoother and its action on update differences\.

We now takeSSto be the*local averaging*\(tridiagonal\) smoother with strengthα∈\(0,12\]\\alpha\\in\(0,\\tfrac\{1\}\{2\}\]:

\(S​u\)ℓ=\{\(1−α2\)​u1\+α2​u2,ℓ=1,\(1−α\)​uℓ\+α2​\(uℓ−1\+uℓ\+1\),2≤ℓ≤L−1,\(1−α2\)​uL\+α2​uL−1,ℓ=L\.\(Su\)\_\{\\ell\}=\\begin\{cases\}\\big\(1\-\\tfrac\{\\alpha\}\{2\}\\big\)u\_\{1\}\+\\tfrac\{\\alpha\}\{2\}u\_\{2\},&\\ell=1,\\\\\[2\.0pt\] \(1\-\\alpha\)u\_\{\\ell\}\+\\tfrac\{\\alpha\}\{2\}\(u\_\{\\ell\-1\}\+u\_\{\\ell\+1\}\),&2\\leq\\ell\\leq L\-1,\\\\\[2\.0pt\] \\big\(1\-\\tfrac\{\\alpha\}\{2\}\\big\)u\_\{L\}\+\\tfrac\{\\alpha\}\{2\}u\_\{L\-1\},&\\ell=L\.\\end\{cases\}\(28\)LetD∈ℝ\(L−1\)×LD\\in\\mathbb\{R\}^\{\(L\-1\)\\times L\}denote the first\-difference operator\(D​v\)ℓ=vℓ\+1−vℓ\(Dv\)\_\{\\ell\}=v\_\{\\ell\+1\}\-v\_\{\\ell\}\.

###### Lemma A\.5\(Induced local averaging on first differences\)\.

LetDDbe the first\-difference operator\(D​u\)ℓ=uℓ\+1−uℓ\(Du\)\_\{\\ell\}=u\_\{\\ell\+1\}\-u\_\{\\ell\}forℓ=1,…,L−1\\ell=1,\\dots,L\-1\. LetSSbe the local averaging operator defined in \([28](https://arxiv.org/html/2606.30813#A1.E28)\)\. DefineS\(1\)∈ℝ\(L−1\)×\(L−1\)S^\{\(1\)\}\\in\\mathbb\{R\}^\{\(L\-1\)\\times\(L\-1\)\}by

\(S\(1\)​w\)ℓ=\{\(1−α\)​w1\+α2​w2,ℓ=1,\(1−α\)​wℓ\+α2​\(wℓ−1\+wℓ\+1\),2≤ℓ≤L−2,\(1−α\)​wL−1\+α2​wL−2,ℓ=L−1\.\(S^\{\(1\)\}w\)\_\{\\ell\}=\\begin\{cases\}\(1\-\\alpha\)w\_\{1\}\+\\tfrac\{\\alpha\}\{2\}w\_\{2\},&\\ell=1,\\\\\[2\.0pt\] \(1\-\\alpha\)w\_\{\\ell\}\+\\tfrac\{\\alpha\}\{2\}\(w\_\{\\ell\-1\}\+w\_\{\\ell\+1\}\),&2\\leq\\ell\\leq L\-2,\\\\\[2\.0pt\] \(1\-\\alpha\)w\_\{L\-1\}\+\\tfrac\{\\alpha\}\{2\}w\_\{L\-2\},&\\ell=L\-1\.\\end\{cases\}Then for all block\-vectorsuu,

D​\(S​u\)=S\(1\)​\(D​u\)\.D\(Su\)=S^\{\(1\)\}\(Du\)\.

###### Proof\.

Letw:=D​uw:=Du, i\.e\.wℓ=uℓ\+1−uℓw\_\{\\ell\}=u\_\{\\ell\+1\}\-u\_\{\\ell\}\. For interiorℓ∈\{2,…,L−2\}\\ell\\in\\\{2,\\dots,L\-2\\\}, expanding\(S​u\)ℓ\+1−\(S​u\)ℓ\(Su\)\_\{\\ell\+1\}\-\(Su\)\_\{\\ell\}using \([28](https://arxiv.org/html/2606.30813#A1.E28)\) gives

\(D​\(S​u\)\)ℓ=\(1−α\)​wℓ\+α2​wℓ−1\+α2​wℓ\+1\.\(D\(Su\)\)\_\{\\ell\}=\(1\-\\alpha\)w\_\{\\ell\}\+\\tfrac\{\\alpha\}\{2\}w\_\{\\ell\-1\}\+\\tfrac\{\\alpha\}\{2\}w\_\{\\ell\+1\}\.Atℓ=1\\ell=1,

\(D​\(S​u\)\)1=\(S​u\)2−\(S​u\)1=\(1−α\)​w1\+α2​w2\.\(D\(Su\)\)\_\{1\}=\(Su\)\_\{2\}\-\(Su\)\_\{1\}=\(1\-\\alpha\)w\_\{1\}\+\\tfrac\{\\alpha\}\{2\}w\_\{2\}\.Atℓ=L−1\\ell=L\-1,

\(D​\(S​u\)\)L−1=\(S​u\)L−\(S​u\)L−1=\(1−α\)​wL−1\+α2​wL−2\.\(D\(Su\)\)\_\{L\-1\}=\(Su\)\_\{L\}\-\(Su\)\_\{L\-1\}=\(1\-\\alpha\)w\_\{L\-1\}\+\\tfrac\{\\alpha\}\{2\}w\_\{L\-2\}\.These are exactly the defining rules ofS\(1\)S^\{\(1\)\}applied toww, henceD​\(S​u\)=S\(1\)​w=S\(1\)​\(D​u\)D\(Su\)=S^\{\(1\)\}w=S^\{\(1\)\}\(Du\)\. ∎

###### Lemma A\.6\(Eigenvalues of induced averaging on first differences\)\.

Letn≥1n\\geq 1and defineS~\(n\)∈ℝn×n\\widetilde\{S\}^\{\(n\)\}\\in\\mathbb\{R\}^\{n\\times n\}by the tridiagonal Toeplitz form

\(S~\(n\)​w\)i=\(1−α\)​wi\+α2​\(𝟏\{i\>1\}​wi−1\+𝟏\{i<n\}​wi\+1\),i=1,…,n,\(\\widetilde\{S\}^\{\(n\)\}w\)\_\{i\}=\(1\-\\alpha\)w\_\{i\}\+\\frac\{\\alpha\}\{2\}\\big\(\\mathbf\{1\}\_\{\\\{i\>1\\\}\}w\_\{i\-1\}\+\\mathbf\{1\}\_\{\\\{i<n\\\}\}w\_\{i\+1\}\\big\),\\qquad i=1,\\dots,n,withα∈\(0,1\)\\alpha\\in\(0,1\)\. ThenS~\(n\)\\widetilde\{S\}^\{\(n\)\}is symmetric and has eigenpairs

vj\(k\)=sin⁡\(π​k​jn\+1\),μ~k\(n\)=1−α​\(1−cos⁡\(π​kn\+1\)\),k=1,…,n\.v^\{\(k\)\}\_\{j\}=\\sin\\\!\\Big\(\\frac\{\\pi kj\}\{n\+1\}\\Big\),\\qquad\\widetilde\{\\mu\}\_\{k\}^\{\(n\)\}=1\-\\alpha\\Big\(1\-\\cos\\Big\(\\frac\{\\pi k\}\{n\+1\}\\Big\)\\Big\),\\qquad k=1,\\dots,n\.In particular,‖S~\(n\)‖op=μ~1\(n\)<1\\\|\\widetilde\{S\}^\{\(n\)\}\\\|\_\{\\mathrm\{op\}\}=\\widetilde\{\\mu\}\_\{1\}^\{\(n\)\}<1\.

###### Proof\.

Fixk∈\{1,…,n\}k\\in\\\{1,\\dots,n\\\}and sett:=π​k/\(n\+1\)t:=\\pi k/\(n\+1\)\. For2≤j≤n−12\\leq j\\leq n\-1,

\(S~\(n\)​v\(k\)\)j=\(1−α\)​sin⁡\(j​t\)\+α2​\(sin⁡\(\(j−1\)​t\)\+sin⁡\(\(j\+1\)​t\)\)=\(1−α\+α​cos⁡t\)​sin⁡\(j​t\),\(\\widetilde\{S\}^\{\(n\)\}v^\{\(k\)\}\)\_\{j\}=\(1\-\\alpha\)\\sin\(jt\)\+\\frac\{\\alpha\}\{2\}\\big\(\\sin\(\(j\-1\)t\)\+\\sin\(\(j\+1\)t\)\\big\)=\\big\(1\-\\alpha\+\\alpha\\cos t\\big\)\\sin\(jt\),usingsin⁡\(\(j−1\)​t\)\+sin⁡\(\(j\+1\)​t\)=2​sin⁡\(j​t\)​cos⁡t\\sin\(\(j\-1\)t\)\+\\sin\(\(j\+1\)t\)=2\\sin\(jt\)\\cos t\. Atj=1j=1andj=nj=nthe same identity holds since the missing neighbor corresponds tosin⁡\(0\)=sin⁡\(\(n\+1\)​t\)=0\\sin\(0\)=\\sin\(\(n\+1\)t\)=0\. Thusv\(k\)v^\{\(k\)\}is an eigenvector with eigenvalue1−α\+α​cos⁡t=1−α​\(1−cos⁡t\)1\-\\alpha\+\\alpha\\cos t=1\-\\alpha\(1\-\\cos t\)\. The operator norm equals the largest absolute eigenvalueμ~1\(n\)\\widetilde\{\\mu\}\_\{1\}^\{\(n\)\}sinceS~\(n\)\\widetilde\{S\}^\{\(n\)\}is symmetric\. ∎

###### Proposition A\.7\.

LetS\(1\)∈ℝ\(L−1\)×\(L−1\)S^\{\(1\)\}\\in\\mathbb\{R\}^\{\(L\-1\)\\times\(L\-1\)\}be the induced operator on differences from Lemma[A\.5](https://arxiv.org/html/2606.30813#A1.Thmtheorem5)\(i\.e\. with boundary diagonal1−α1\-\\alpha\)\. Then for allw∈ℝL−1w\\in\\mathbb\{R\}^\{L\-1\},

‖S\(1\)​w‖≤μ⋆​‖w‖,μ⋆:=1−α​\(1−cos⁡\(πL\)\)<1,\\\|S^\{\(1\)\}w\\\|\\;\\leq\\;\\mu\_\{\\star\}\\,\\\|w\\\|,\\qquad\\mu\_\{\\star\}:=1\-\\alpha\\Big\(1\-\\cos\\Big\(\\frac\{\\pi\}\{L\}\\Big\)\\Big\)\\;<\\;1,\(29\)and hence

‖w‖2−‖S\(1\)​w‖2≥\(1−μ⋆2\)​‖w‖2\.\\\|w\\\|^\{2\}\-\\\|S^\{\(1\)\}w\\\|^\{2\}\\;\\geq\\;\\big\(1\-\\mu\_\{\\star\}^\{2\}\\big\)\\,\\\|w\\\|^\{2\}\.\(30\)

###### Proof\.

By Lemma[A\.6](https://arxiv.org/html/2606.30813#A1.Thmtheorem6)withn=L−1n=L\-1, the eigenvalues ofS\(1\)=S~\(L−1\)S^\{\(1\)\}=\\widetilde\{S\}^\{\(L\-1\)\}areμ~k\(L−1\)=1−α​\(1−cos⁡\(π​k/L\)\)\\widetilde\{\\mu\}\_\{k\}^\{\(L\-1\)\}=1\-\\alpha\(1\-\\cos\(\\pi k/L\)\)fork=1,…,L−1k=1,\\dots,L\-1\. Thus‖S\(1\)‖op=μ~1\(L−1\)=μ⋆<1\\\|S^\{\(1\)\}\\\|\_\{\\mathrm\{op\}\}=\\widetilde\{\\mu\}\_\{1\}^\{\(L\-1\)\}=\\mu\_\{\\star\}<1, giving \([29](https://arxiv.org/html/2606.30813#A1.E29)\)\. Then \([30](https://arxiv.org/html/2606.30813#A1.E30)\) follows from‖S\(1\)​w‖2≤μ⋆2​‖w‖2\\\|S^\{\(1\)\}w\\\|^\{2\}\\leq\\mu\_\{\\star\}^\{2\}\\\|w\\\|^\{2\}\. ∎

#### General one\-step bound for the average cosine alignment\.

Define the average cosine across consecutive residual increments:

cos¯\(ϑ\):=1L−1∑ℓ=1L−1cos\(rℓ\+1\(ϑ\),rℓ\(ϑ\)\)\.\\overline\{\\cos\}\(\\vartheta\):=\\frac\{1\}\{L\-1\}\\sum\_\{\\ell=1\}^\{L\-1\}\\cos\\\!\\big\(r\_\{\\ell\+1\}\(\\vartheta\),r\_\{\\ell\}\(\\vartheta\)\\big\)\.\(31\)We now state the general one\-step comparison, allowing the current parameters to have pre\-existing depth variation\.

###### Corollary A\.8\.

Assume Assumption[A\.1](https://arxiv.org/html/2606.30813#A1.Thmtheorem1)\. LetSSbe the local\-averaging smoother \([28](https://arxiv.org/html/2606.30813#A1.E28)\) with parameterα∈\(0,12\]\\alpha\\in\(0,\\tfrac\{1\}\{2\}\], and letS\(1\)S^\{\(1\)\}be the induced\(L−1\)×\(L−1\)\(L\-1\)\\times\(L\-1\)averaging operator from Lemma[A\.5](https://arxiv.org/html/2606.30813#A1.Thmtheorem5)\. Define

δ:=D​θ,w:=D​u​\(θ\),\\delta:=D\\theta,\\qquad w:=Du\(\\theta\),whereDDis the first\-difference operator and all norms and inner products are taken over depth\-stacked block vectors\. Then the smoothed and baseline one\-step updates satisfy

cos¯​\(θs\)\\displaystyle\\overline\{\\cos\}\(\\theta^\{s\}\)≥1−4​Lh2​M2m2−4​Lθ2m2⋅1L−1​‖δ−η​S\(1\)​w‖2,\\displaystyle\\geq 1\-\\frac\{4L\_\{h\}^\{2\}M^\{2\}\}\{m^\{2\}\}\-\\frac\{4L\_\{\\theta\}^\{2\}\}\{m^\{2\}\}\\cdot\\frac\{1\}\{L\-1\}\\,\\big\\\|\\delta\-\\eta S^\{\(1\)\}w\\big\\\|^\{2\},\(32\)cos¯​\(θb\)\\displaystyle\\overline\{\\cos\}\(\\theta^\{b\}\)≥1−4​Lh2​M2m2−4​Lθ2m2⋅1L−1​‖δ−η​w‖2\.\\displaystyle\\geq 1\-\\frac\{4L\_\{h\}^\{2\}M^\{2\}\}\{m^\{2\}\}\-\\frac\{4L\_\{\\theta\}^\{2\}\}\{m^\{2\}\}\\cdot\\frac\{1\}\{L\-1\}\\,\\big\\\|\\delta\-\\eta w\\big\\\|^\{2\}\.\(33\)Moreover, the difference between the two lower bounds admits the exact identity

\[RHS of \([32](https://arxiv.org/html/2606.30813#A1.E32)\)\]−\[RHS of \([33](https://arxiv.org/html/2606.30813#A1.E33)\)\]=4​Lθ2m2⋅1L−1​\[η2​\(‖w‖2−‖S\(1\)​w‖2\)−2​η​⟨δ,\(I−S\(1\)\)​w⟩\]\.\\Big\[\\text\{RHS of \\eqref\{eq:avg\_cos\_general\_s\}\}\\Big\]\-\\Big\[\\text\{RHS of \\eqref\{eq:avg\_cos\_general\_b\}\}\\Big\]=\\frac\{4L\_\{\\theta\}^\{2\}\}\{m^\{2\}\}\\cdot\\frac\{1\}\{L\-1\}\\left\[\\eta^\{2\}\\big\(\\\|w\\\|^\{2\}\-\\\|S^\{\(1\)\}w\\\|^\{2\}\\big\)\-2\\eta\\big\\langle\\delta,\(I\-S^\{\(1\)\}\)w\\big\\rangle\\right\]\.\(34\)In particular, ifη\>0\\eta\>0andLθ\>0L\_\{\\theta\}\>0, smoothing yields a strictly improved lower bound whenever

⟨δ,\(I−S\(1\)\)​w⟩<η2​\(‖w‖2−‖S\(1\)​w‖2\)\.\\big\\langle\\delta,\(I\-S^\{\(1\)\}\)w\\big\\rangle<\\frac\{\\eta\}\{2\}\\Big\(\\\|w\\\|^\{2\}\-\\\|S^\{\(1\)\}w\\\|^\{2\}\\Big\)\.\(35\)A simple sufficient condition for such is

‖δ‖<η2⋅1−μ⋆2‖I−S\(1\)‖op​‖w‖,μ⋆:=‖S\(1\)‖op<1\.\\\|\\delta\\\|<\\frac\{\\eta\}\{2\}\\cdot\\frac\{1\-\\mu\_\{\\star\}^\{2\}\}\{\\\|I\-S^\{\(1\)\}\\\|\_\{\\mathrm\{op\}\}\}\\,\\\|w\\\|,\\qquad\\mu\_\{\\star\}:=\\\|S^\{\(1\)\}\\\|\_\{\\mathrm\{op\}\}<1\.\(36\)

###### Proof\.

The bounds \([32](https://arxiv.org/html/2606.30813#A1.E32)\)–\([33](https://arxiv.org/html/2606.30813#A1.E33)\) follow by applying Lemma[A\.4](https://arxiv.org/html/2606.30813#A1.Thmtheorem4)withϑ=θs\\vartheta=\\theta^\{s\}andϑ=θb\\vartheta=\\theta^\{b\}\. Indeed,

D​θs=D​\(θ−η​S​u\)=D​θ−η​D​\(S​u\)=δ−η​S\(1\)​w,D\\theta^\{s\}=D\(\\theta\-\\eta Su\)=D\\theta\-\\eta D\(Su\)=\\delta\-\\eta S^\{\(1\)\}w,where we used Lemma[A\.5](https://arxiv.org/html/2606.30813#A1.Thmtheorem5), and similarly

D​θb=D​\(θ−η​u\)=δ−η​w\.D\\theta^\{b\}=D\(\\theta\-\\eta u\)=\\delta\-\\eta w\.Substituting these identities into \([27](https://arxiv.org/html/2606.30813#A1.E27)\) gives \([32](https://arxiv.org/html/2606.30813#A1.E32)\) and \([33](https://arxiv.org/html/2606.30813#A1.E33)\)\.

For \([34](https://arxiv.org/html/2606.30813#A1.E34)\), subtract the right\-hand sides of \([32](https://arxiv.org/html/2606.30813#A1.E32)\) and \([33](https://arxiv.org/html/2606.30813#A1.E33)\)\. The only changing term is the squared depth\-difference term, and

‖δ−η​S\(1\)​w‖2−‖δ−η​w‖2=η2​\(‖S\(1\)​w‖2−‖w‖2\)\+2​η​⟨δ,\(I−S\(1\)\)​w⟩\.\\\|\\delta\-\\eta S^\{\(1\)\}w\\\|^\{2\}\-\\\|\\delta\-\\eta w\\\|^\{2\}=\\eta^\{2\}\\big\(\\\|S^\{\(1\)\}w\\\|^\{2\}\-\\\|w\\\|^\{2\}\\big\)\+2\\eta\\big\\langle\\delta,\(I\-S^\{\(1\)\}\)w\\big\\rangle\.Multiplying by

−4​Lθ2m2⋅1L−1\-\\frac\{4L\_\{\\theta\}^\{2\}\}\{m^\{2\}\}\\cdot\\frac\{1\}\{L\-1\}yields \([34](https://arxiv.org/html/2606.30813#A1.E34)\)\. Condition \([35](https://arxiv.org/html/2606.30813#A1.E35)\) makes the right\-hand side strictly positive\.

Finally, \([36](https://arxiv.org/html/2606.30813#A1.E36)\) follows from Cauchy–Schwarz and the contraction ofS\(1\)S^\{\(1\)\}\. Specifically,

⟨δ,\(I−S\(1\)\)​w⟩≤‖δ‖​‖I−S\(1\)‖op​‖w‖,\\big\\langle\\delta,\(I\-S^\{\(1\)\}\)w\\big\\rangle\\leq\\\|\\delta\\\|\\,\\\|I\-S^\{\(1\)\}\\\|\_\{\\mathrm\{op\}\}\\\|w\\\|,while Proposition[A\.7](https://arxiv.org/html/2606.30813#A1.Thmtheorem7)gives

‖w‖2−‖S\(1\)​w‖2≥\(1−μ⋆2\)​‖w‖2\.\\\|w\\\|^\{2\}\-\\\|S^\{\(1\)\}w\\\|^\{2\}\\geq\(1\-\\mu\_\{\\star\}^\{2\}\)\\\|w\\\|^\{2\}\.Combining these two inequalities gives the stated sufficient condition\. ∎

![Refer to caption](https://arxiv.org/html/2606.30813v1/plots/cossim_v_layer_3.png)Figure 6:Gradient smoothing increases layer contribution alignment\.CLS token trajectories of ViT\-B trained on CIFAR\-100 for 1700 epochs, comparing baseline \(α=0\\alpha=0, red\) against heavy smoothing \(α=0\.4\\alpha=0\.4, navy\)\. For each layer differencedℓ=xℓ\+1−xℓd\_\{\\ell\}=x\_\{\\ell\+1\}\-x\_\{\\ell\}, we compute the mean cosine similarity with all other layer differencesdjd\_\{j\}\(j≠ℓj\\neq\\ell\)\. Across both validation and training sets, the smoothed model exhibits consistently higher mean cosine similarity at most layers, indicating that gradient smoothing encourages more coherent directional updates throughout the network\.![Refer to caption](https://arxiv.org/html/2606.30813v1/plots/lss_alpha_3.png)Figure 7:Greater linearity with increased smoothing\.Line Shape Score \(LSS\) of CLS token trajectories in ViT\-B trained on CIFAR\-100 \(1700 epochs\)\. Each point shows the median across 100 validation \(rhombus\) and training \(circle\) images, while error bars indicate interquartile range \(Q1–Q3\)\. As the window smoothing parameterα∈\{0\.1,…,0\.4\}\\alpha\\in\\\{0\.1,\\ldots,0\.4\\\}increases from the baseline \(α\\alpha= 0\), median LSS decreases monotonically, indicating that gradient smoothing produces more linear CLS token trajectories\.

### A\.3Variance\-Reduction of Gradients with Smoothing

Here we provide a simple variance\-based analysis of Gradients with the local window operator\. LetU∈ℝL​dU\\in\\mathbb\{R\}^\{Ld\}denote the stacked stochastic \(e\.g\. from minibatches\) block update, and define the smoothed update by

U~=\(S⊗Id\)​U,\\tilde\{U\}=\(S\\otimes I\_\{d\}\)U,whereSSis the depth\-smoothing matrix\.

###### Proposition A\.10\(Variance under a single step of window smoothing\)\.

LetU∈ℝL​dU\\in\\mathbb\{R\}^\{Ld\}be a random vector with𝔼​‖U‖22<∞\\mathbb\{E\}\\\|U\\\|\_\{2\}^\{2\}<\\inftyand covarianceΣ=Cov​\(U\)\\Sigma=\\mathrm\{Cov\}\(U\), and letU~=\(S⊗Id\)​U\\tilde\{U\}=\(S\\otimes I\_\{d\}\)UwithSSthe local window operator \([28](https://arxiv.org/html/2606.30813#A1.E28)\)\. For anyα∈\[0,1\]\\alpha\\in\[0,1\],

𝔼​‖U~−𝔼​U~‖22≤𝔼​‖U−𝔼​U‖22\.\\mathbb\{E\}\\\|\\tilde\{U\}\-\\mathbb\{E\}\\tilde\{U\}\\\|\_\{2\}^\{2\}\\;\\leq\\;\\mathbb\{E\}\\\|U\-\\mathbb\{E\}U\\\|\_\{2\}^\{2\}\.

###### Proof\.

SinceSSis symmetric with eigenvalues

λk​\(S\)=1−α​\(1−cos⁡k​πL\),k=0,…,L−1,\\lambda\_\{k\}\(S\)=1\-\\alpha\\\!\\left\(1\-\\cos\\\!\\frac\{k\\pi\}\{L\}\\right\),\\qquad k=0,\\dots,L\-1,we have‖S‖2=1\\\|S\\\|\_\{2\}=1forα∈\[0,1\]\\alpha\\in\[0,1\]\.

By linearity,Cov​\(U~\)=\(S⊗Id\)​Σ​\(S⊗Id\)⊤\\mathrm\{Cov\}\(\\tilde\{U\}\)=\(S\\otimes I\_\{d\}\)\\,\\Sigma\\,\(S\\otimes I\_\{d\}\)^\{\\top\}, so the total stochastic update variance is

𝔼​‖U~−𝔼​U~‖22=tr​\(\(S⊗Id\)​Σ​\(S⊗Id\)⊤\)=tr​\(\(\(S⊤​S\)⊗Id\)​Σ\)\.\\mathbb\{E\}\\\|\\tilde\{U\}\-\\mathbb\{E\}\\tilde\{U\}\\\|\_\{2\}^\{2\}=\\mathrm\{tr\}\\\!\\bigl\(\(S\\otimes I\_\{d\}\)\\,\\Sigma\\,\(S\\otimes I\_\{d\}\)^\{\\top\}\\bigr\)=\\mathrm\{tr\}\\\!\\bigl\(\(\(S^\{\\top\}S\)\\otimes I\_\{d\}\)\\,\\Sigma\\bigr\)\.LetA:=\(S⊤​S\)⊗IdA:=\(S^\{\\top\}S\)\\otimes I\_\{d\}\. SinceΣ\\Sigmais a covariance matrix,Σ⪰0\\Sigma\\succeq 0, and clearlyA⪰0A\\succeq 0as well\. ThenA⪯‖A‖2​IA\\preceq\\\|A\\\|\_\{2\}\\,I, which implies

tr​\(A​Σ\)=tr​\(Σ1/2​A​Σ1/2\)≤‖A‖2​tr​\(Σ\)\.\\mathrm\{tr\}\(A\\Sigma\)=\\mathrm\{tr\}\\\!\\bigl\(\\Sigma^\{1/2\}A\\,\\Sigma^\{1/2\}\\bigr\)\\leq\\\|A\\\|\_\{2\}\\,\\mathrm\{tr\}\(\\Sigma\)\.Since

‖A‖2=‖\(S⊤​S\)⊗Id‖2=‖S⊤​S‖2=‖S‖22=1,\\\|A\\\|\_\{2\}=\\\|\(S^\{\\top\}S\)\\otimes I\_\{d\}\\\|\_\{2\}=\\\|S^\{\\top\}S\\\|\_\{2\}=\\\|S\\\|\_\{2\}^\{2\}=1,we obtain

𝔼​‖U~−𝔼​U~‖22≤tr​\(Σ\)=𝔼​‖U−𝔼​U‖22\.∎\\mathbb\{E\}\\\|\\tilde\{U\}\-\\mathbb\{E\}\\tilde\{U\}\\\|\_\{2\}^\{2\}\\leq\\mathrm\{tr\}\(\\Sigma\)=\\mathbb\{E\}\\\|U\-\\mathbb\{E\}U\\\|\_\{2\}^\{2\}\.\\qed

Thus, forα∈\[0,1\]\\alpha\\in\[0,1\], window smoothing update variance is at most that of the blockwise stochastic update\.

### A\.4Empirical Measurement of Gradient Variance

To complement the variance contraction bound of Proposition[A\.10](https://arxiv.org/html/2606.30813#A1.Thmtheorem10), we directly measure two notions of stochastic gradient variance during pretraining and compare them across baseline and smoothing runs\.

#### Setup\.

We train three otherwise\-identicalnanochatd24 models \(1\.38B parameters, 24 transformer blocks\) for 7,308 optimization steps: a baseline run and two window\-smoothing runs withα∈\{0\.05,0\.1\}\\alpha\\in\\\{0\.05,\\,0\.1\\\}\. Every5050steps, before applying the optimizer update, we record per\-microbatch gradients for each parameter group\(r,ℓ\)\(r,\\ell\), whererrindexes one ofR=6R=6roles within a block \(the four attention projections and the two MLP linears\) andℓ∈\{1,…,L\}\\ell\\in\\\{1,\\ldots,L\\\}indexes block depth\. Letgr,ℓ,ng\_\{r,\\ell,n\}denote the flattened gradient for rolerr, layerℓ\\ell, and microbatchnn, withNNmicrobatches per logging step\.

#### Microbatch variance\.

The microbatch variance measures the across\-sample noise of the stochastic gradient at each parameter group:

Vr,ℓmb:=1N​∑n=1N‖gr,ℓ,n−g¯r,ℓ‖22,g¯r,ℓ:=1N​∑n=1Ngr,ℓ,n\.V^\{\\mathrm\{mb\}\}\_\{r,\\ell\}\\;:=\\;\\frac\{1\}\{N\}\\sum\_\{n=1\}^\{N\}\\bigl\\\|g\_\{r,\\ell,n\}\-\\bar\{g\}\_\{r,\\ell\}\\bigr\\\|\_\{2\}^\{2\},\\qquad\\bar\{g\}\_\{r,\\ell\}\\;:=\\;\\frac\{1\}\{N\}\\sum\_\{n=1\}^\{N\}g\_\{r,\\ell,n\}\.Summing over roles and layers gives the total

Vtotalmb:=∑r=1R∑ℓ=1LVr,ℓmb,V^\{\\mathrm\{mb\}\}\_\{\\mathrm\{total\}\}\\;:=\\;\\sum\_\{r=1\}^\{R\}\\sum\_\{\\ell=1\}^\{L\}V^\{\\mathrm\{mb\}\}\_\{r,\\ell\},which is the empirical estimate of𝔼​‖U−𝔼​U‖22\\mathbb\{E\}\\\|U\-\\mathbb\{E\}U\\\|\_\{2\}^\{2\}in the notation of Proposition[A\.10](https://arxiv.org/html/2606.30813#A1.Thmtheorem10)\.

#### Depth variance\.

The depth variance measures the across\-layer dispersion of the gradient within each role, evaluated per microbatch and averaged over microbatches:

Vrdepth:=1N​∑n=1N1L​∑ℓ=1L‖gr,ℓ,n−g¯r,⋅,n‖22,g¯r,⋅,n:=1L​∑ℓ=1Lgr,ℓ,n\.V^\{\\mathrm\{depth\}\}\_\{r\}\\;:=\\;\\frac\{1\}\{N\}\\sum\_\{n=1\}^\{N\}\\frac\{1\}\{L\}\\sum\_\{\\ell=1\}^\{L\}\\bigl\\\|g\_\{r,\\ell,n\}\-\\bar\{g\}\_\{r,\\cdot,n\}\\bigr\\\|\_\{2\}^\{2\},\\qquad\\bar\{g\}\_\{r,\\cdot,n\}\\;:=\\;\\frac\{1\}\{L\}\\sum\_\{\\ell=1\}^\{L\}g\_\{r,\\ell,n\}\.The total isVtotaldepth:=∑rVrdepthV^\{\\mathrm\{depth\}\}\_\{\\mathrm\{total\}\}:=\\sum\_\{r\}V^\{\\mathrm\{depth\}\}\_\{r\}, and is the empirical estimate of the quantity from Proposition[A\.7](https://arxiv.org/html/2606.30813#A1.Thmtheorem7)\.

#### Observations\.

Figures[4](https://arxiv.org/html/2606.30813#S4.F4)and[8](https://arxiv.org/html/2606.30813#A1.F8)reportVtotalmbV^\{\\mathrm\{mb\}\}\_\{\\mathrm\{total\}\}andVtotaldepthV^\{\\mathrm\{depth\}\}\_\{\\mathrm\{total\}\}over training for the three runs\. After the initial warmup transient \(steps≲1000\\lesssim 1000\), both smoothing runs consistently exhibit lower gradient variance than baseline, in both metrics\. The reduction is stable across the bulk of training and widens further in late training, consistent with Proposition[A\.10](https://arxiv.org/html/2606.30813#A1.Thmtheorem10)acting at every step along the smoothed trajectory\.

![Refer to caption](https://arxiv.org/html/2606.30813v1/plots/nanochat_depth_var.png)Figure 8:Lower depth gradient variance with smoothing\.Total depth varianceVtotaldepthV^\{\\mathrm\{depth\}\}\_\{\\mathrm\{total\}\}duringnanochatd24 pretraining \(7000 optimization steps, logged every 50 steps\), comparing the baseline \(blue\) against window smoothing withα=0\.05\\alpha=0\.05\(red\) andα=0\.1\\alpha=0\.1\(green\)\. Both smoothing runs show consistently lower depth\-wise gradient dispersion than baseline once training leaves the warmup phase, withα=0\.05\\alpha=0\.05andα=0\.1\\alpha=0\.1producing comparable reductions\.