No title
January 01, 1970
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.
1 Introduction↩︎
Many modern deep neural networks are composed of repeated architectural blocks with shared computational structure, such as Residual Networks (ResNets) [1] and transformers [2]. 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 [3] and Residual Alignment [4] 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 [5], as well as analyses showing smooth evolution of information-theoretic quantities across depth [6].
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.
1.0.0.1 Contributions.
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.
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.
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.
2 Background and Setup↩︎
2.1 Block-Structured Optimization↩︎
We consider neural networks composed of \(L\) repeated 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, let \(F:\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 by \(F(\cdot;\theta_\ell)\), where \(\theta_\ell \in \mathbb{R}^p\) are layer-specific parameters. Given the initial embedding \(h_0\), the hidden states \((h_\ell)_{\ell=1}^{L} \subset \mathbb{R}^d\) then evolve according to \[\label{eq:residual95model} h_{\ell+1} = h_\ell + F(h_\ell;\theta_\ell), \qquad \ell = 1,\dots,L-1.\tag{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 \[(\theta, \phi), \qquad \theta = (\theta_1, \dots, \theta_L), \quad \theta_l \in \mathbb{R}^{p},\] where \(\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 each \(h_{\ell}\) be a sequence of vectors). We denote the training objective by \(\mathcal{L}(\theta,\phi)\) and assume it is differentiable.
At optimization step \(t\), the block-wise gradients are given by \[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 by \(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 \(\{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 vector \(g^{(t)}\) possesses meaningful structure along the depth dimension.
Figure 2: Nanochat pretraining with Gradient Smoothing. Validation loss, BPB, and DCLM CORE metric during nanochat pretraining of GPT2 with the default Adam + NorMuon optimizer setup (Table 1). 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 \(\alpha = 0.1\) are replicated against the baseline for \(\boldsymbol{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 depth \(30\).. a — Depth \(24\): Validation Loss., b — Depth \(24\): Validation BPB.,
c — Depth \(24\): CORE Metric., d — Depth \(30\): Validation Loss., e — Depth \(30\): Validation BPB., f — Depth \(30\):
CORE Metric.
2.2 General First-Order Updates↩︎
Rather than committing to a specific optimizer, we consider a general first-order update rule of the form \[(\theta^{(t+1)},\phi^{(t+1)}) = (\theta^{(t)},\phi^{(t)}) - \eta\, (u^{(t)}, u_\phi^{(t)}),\] where \[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)}\)), momentum methods, adaptive optimizers such as Adam or AdamW, as well as Muon and related variants. Importantly, \(u^{(t)}\) may depend on both the current gradient and internal optimizer state, and may involve nonlinear or time-varying transformations of \(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.
3 Depth-wise Gradient Augmentation↩︎
3.1 General 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 iteration \(t\), let \[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)} = (u_1^{(t)}, \dots, u_L^{(t)}) \in (\mathbb{R}^{p})^{L},\] together with the corresponding update \(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 coordinate \(i=1,\ldots,p\), let \[\mathcal{C}^{(i)}:\mathbb{R}^{L}\rightarrow\mathbb{R}^{L},\] be a depth-wise augmentation operator. Then \[\bigl(\tilde{u}_{1}^{(i)},\ldots,\tilde{u}_{L}^{(i)}\bigr) = \mathcal{C}^{(i)} \bigl(u_{1}^{(i)},\ldots,u_{L}^{(i)}\bigr).\] Equivalently, \[\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 \(\ell\) may 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.2 Gradient 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 \[\mathcal{C}^{(i)} = S, \qquad i=1,\ldots,p,\] where \(S \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 \otimes I \in \mathbb{R}^{D_\theta \times D_\theta},\] where \(I\) denotes the identity operator on the within-block parameter space \(\mathbb{R}^p\).
The transformed block-wise update is then given by \[\tilde{u}^{(t)} := P\,u^{(t)} = (S \otimes I)\,u^{(t)}.\]
To describe the update on the full parameter vector \((\theta,\phi)\), we lift \(P\) to the block-diagonal operator \[\bar P := \begin{pmatrix} P & 0 \\ 0 & I_\phi \end{pmatrix} \in \mathbb{R}^{(D_\theta + D_\phi)\times(D_\theta + D_\phi)},\] where \(I_\phi\) denotes the identity operator on the \(\phi\)-coordinates.
Given the full optimizer update \[u_{\mathrm{full}}^{(t)} := \bigl(u^{(t)},\,u_\phi^{(t)}\bigr),\] Gradient Smoothing applies \(\bar P\) to obtain \[\tilde{u}_{\mathrm{full}}^{(t)} = \bar P\,u_{\mathrm{full}}^{(t)} = \bigl(Pu^{(t)},\,u_\phi^{(t)}\bigr).\]
The resulting optimization step becomes \[\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.3 Smoothing Operator↩︎
In this work, we instantiate Gradient Smoothing with a simple local averaging operator, which we term Window Smoothing. Taking \(S\) to be a local window averaging operator across depth, the smoothed update is given by \[\tilde{u}_l = \begin{cases} \big(1-\tfrac{\alpha}{2}\big)u_1 + \tfrac{\alpha}{2}u_2, & l=1,\\[4pt] (1-\alpha)u_l + \tfrac{\alpha}{2}(u_{l-1}+u_{l+1}), & 2\le l\le L-1,\\[4pt] \big(1-\tfrac{\alpha}{2}\big)u_L + \tfrac{\alpha}{2}u_{L-1}, & l=L, \end{cases} \label{eq:local95avg}\tag{2}\] where \(\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 at \(l=1\) and \(l=L\) use a half-weight on the single available neighbor so that \(S\) remains row-stochastic (\(S\mathbf{1}=\mathbf{1}\)). The resulting \(S\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.
block weights \(\{\boldsymbol{\theta}_\ell\}_{\ell=1}^{L}\), non-block weights \(\boldsymbol{\phi}\), loss \(L\), base optimizer \(\mathcal{U}\), learning rate \(\eta\), weight decay \(\lambda\), smoothing strength \(\alpha \in [0,1)\) \(\mathbf{u}_\ell \gets \mathcal{U}_\theta\!\left(\nabla_{\boldsymbol{\theta}_\ell} L\right)\) \(\ell = 1,\ldots,L\) \(\mathbf{u}_\phi \gets \mathcal{U}_\phi\!\left(\nabla_{\boldsymbol{\phi}} L\right)\) \(\boldsymbol{\theta}_\ell \gets \boldsymbol{\theta}_\ell { \color{red} -\eta\lambda\,\boldsymbol{\theta}_\ell} - \eta\,\tilde{\mathbf{u}}_\ell\) optional w.d. decouple \(\boldsymbol{\phi} \gets \boldsymbol{\phi} - \eta\,\mathbf{u}_\phi\)
3.4 Normalization Variants↩︎
We additionally consider two normalization variants that separate the effects of smoothing on update directions and update magnitudes.
Norm-Preserving Smoothing. After applying the smoothing operator, each block update is rescaled to match the norm of the original optimizer update: \[\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.
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: \[\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).
4 Experiments↩︎
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.1 Experimental Setup↩︎
4.1.0.1 Gradient Smoothing Configuration.
Gradient smoothing is applied across block-wise updates at each optimization step (Section 3.3) 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 (attention \(W_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 as Full in Section 4.2), (ii) smoothing applied only to linear layers within each block (the non-normalization layers in the models we use) which we call Proj.
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.
4.1.0.2 RLVR for LLM Reasoning:
The base model we use is DeepSeek-R1-Distill-Qwen-1.5B [7], [8]. We train with Group Relative Policy Optimization (GRPO) [9] on the dataset from Open-RS [10] (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 of \(10^{-6}\) with a cosine learning-rate schedule and a minimum learning-rate ratio of \(0.1\), including a linear warmup over the first \(10\%\) of training steps.
Evaluation sampling is done with temperature \(0.6\) and top_p \(0.95\).
4.1.0.3 LLM Pretraining:
We further evaluate gradient smoothing in a language-model pretraining setting using the nanochat recipe [11]. nanochat
provides 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 depths \(24\) and \(30\) (approximately \(1.38\)B and \(2.40\)B parameters respectively, Table 1) using the default nanochat architecture and Muon-style [12] optimization setup with NorMuon [13]. Both
depths use the same auto-selected global batch size of \(\approx1\)M 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 the nanochat evaluation suite.
4.1.0.4 Vision Transformer Image Classification:
ViT-B/4 models following DeiT training [14] with default image augmentations. We train for 1700 epochs using AdamW with a cosine learning rate schedule, initial learning rate of \(5 \times 10^{-4}\), and weight decay of 0.05. CIFAR-100 evaluated via top-1 accuracy.
| Variant | Layers | \(d_{\text{model}}\) | Params | Batch (tokens) | |
|---|---|---|---|---|---|
| \(d{=}24\) | 24 | 1536 | 1.38 B | 1,048,576 | |
| \(d{=}30\) | 30 | 1920 | 2.40 B | 1,048,576 |
4.1.0.5 Diffusion:
We evaluate gradient smoothing on U-ViT [15], a ViT-based diffusion backbone with long skip connections proposed for image generation on CIFAR-10. The U-ViT architecture has \(12\) transformer blocks with embedding dimension \(512\) and \(8\) attention heads. Optimization uses AdamW with learning rate \(2\times 10^{-4}\), weight decay \(0.03\), batch size \(128\), with a \(2500\)-step warmup for \(500{,}000\) steps. We sample via Euler–Maruyama SDE at \(50\) NFE and report FID against both \(10\)k and \(50\)k generated samples.
4.2 Experimental 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.
| Method | AIME24 | AIME25 | AMC23 | MATH-500 | Avg |
|---|---|---|---|---|---|
| Base model | 26.67 | 16.67 | 70.00 | 83.00 | 49.09 |
| Baseline (AdamW) | 35.00 | 24.17 | 78.12 | 84.50 | 55.45 \(\pm\) 0.99 |
| + Window (\(\alpha=0.1\), Standard, Full) | 35.83 | 28.33 | 78.75 | 84.30 | 56.80 \(\pm\) 0.21 |
| + Window (\(\alpha=0.1\), Standard, Proj) | 40.83 | 25.00 | 80.00 | 84.55 | 57.60 \(\pm\) 1.12 |
| + Window (\(\alpha=0.2\), Standard, Proj) | 36.67 | 30.00 | 77.50 | 85.35 | 57.38\(\pm\) 0.96 |
4.2.0.1 RLVR Fine-Tuning for LLM Reasoning:
Gradient Smoothing improves the reasoning accuracy of DeepSeek-R1-Distill-Qwen-1.5B after RLVR fine-tuning across mathematical benchmarks (Table 2). All Window Smoothing configurations outperform the
AdamW baseline in average pass@1 accuracy, with the best configuration achieving \(57.60\%\) versus \(55.45\%\). Across individual benchmarks, Gradient Smoothing yields
substantial improvements, including gains on AIME24 (40.83% vs.%), AIME25 (30.00% vs.%), AMC23 (80.00% vs.%), and MATH-500 (85.35% vs.%) 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.
4.2.0.2 LLM Pretraining:
Gradient Smoothing also improves the optimization trajectory in language-model pretraining with the nanochat recipe (Figure 2). During training, across both depth-\(24\) and depth-\(30\) models, 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 from nanochat,
indicating that the benefits of smoothing extend beyond standard Adam-family training alone. The stronger improvements observed for the deeper depth-\(30\) model further suggest that coupling updates across depth may become
increasingly beneficial as the model depth increases.
4.2.0.3 Vision Transformer Image Classification:
| Method | CIFAR-100 |
|---|---|
| Baseline (AdamW) | 74.56 |
| + Window (\(\alpha=0.1\), Norm, Proj) | 75.44 |
| + Window (\(\alpha=0.2\), Dir, Proj) | 75.62 |
Gradient Smoothing improves test accuracy of ViT-B on CIFAR-100 (Table 3, Figure 3). Window smoothing with \(\alpha = 0.2\) achieves notable improvements (75.62% vs 74.56% baseline), while \(\alpha = 0.1\) also 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.
4.2.0.4 Diffusion:
We evaluate Gradient Smoothing on U-ViT [15], 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 from \(6.58\) to \(5.82\) and FID@50k from \(4.01\) to \(3.74\) (Table 4). 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.
| Method | FID@10k | FID@50k |
|---|---|---|
| Baseline (AdamW) | 6.58 | 4.01 |
| + Window (\(\alpha=0.2\), | 5.82 | 3.74 |
| Norm, Proj) |
4.3 Effect 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 depth \(x_\ell\), for \(1 \leq \ell \leq L\). For \(\ell\in \{1, \ldots, L-1 \}\) the difference \(d_\ell = x_{\ell+1} - x_\ell\) corresponds to the contribution of layer \(\ell+1\). The similarity of contributions between layers is quantified by \(c_\ell = \frac{ \langle{d_\ell, d_{\ell+1} \rangle}}{ \|d_\ell\| \|d_{\ell+1}\| }\) (Figure 6).
Line shape score. We measure the linearity of token trajectories via the line shape score (LSS) [16], measuring the closeness of a trajectory to a line. Given representations \(x_0, \ldots, x_L\), the normalized trajectory \(\tilde{x}_0 = x_0\) with \[\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 by \(\text{LSS} = \frac{L}{\|\tilde{x}_L - \tilde{x}_0\|_2}\). LSS \(\geq 1\), with equality if and only if the trajectory is perfectly co-linear (Figure 7).
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 \(\alpha \in \{0.1, \ldots, 0.4\}\) relative to the non-smoothed baseline (\(\alpha = 0.0\)). As the smoothing strength increases, layer-wise contributions to the CLS token become more aligned (Figure 5), and the trend persists when measuring mean cosine similarity of the contribution at individual layers (Figure 6). Smoothing also increases the linearity of the CLS token trajectory across layers (Figure 7). Together, these results suggest that Gradient Smoothing may implicitly regularize representation paths, encouraging more coherent evolution of representations through the model depth.
4.4 Effect of Smoothing on Training Stability↩︎
4.4.0.1 Gradient variance.
To assess downstream improvements on the optimization trajectory, we measure the total stochastic gradient variance within the microbatch. We quantify \[V_{\mathrm{mb}}^{(t)} \;=\; \frac{1}{N}\sum_{n=1}^{N}\big\|g^{(t)}_n - \bar g^{(t)}\big\|_2^2,\] where at step \(t\), \(g_n^{(t)}\) is the per-sample parameter gradient and \(\bar g^{(t)}\) is the sample mean over the \(N\) batch elements at that step. Window smoothing with \(\alpha \in \{ 0.05, 0.1 \}\) maintains consistently lower variance than the baseline throughout nanochat pretraining (Figure 4), suggestive of improved training stability. Full measurement details are deferred to Appendix 9.4, and a variance contraction bound is shown across layer depth (Proposition 4).
5 Alignment 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 increments \(r_\ell := h_{\ell+1}-h_\ell,\) and the cosine similarity \(\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.
5.0.0.1 Residual increments and alignment metric.
For any depth-stacked parameters \(\vartheta=(\vartheta_1,\dots,\vartheta_L)\), define for \(\ell=1,\dots,L\), \[\label{eq:r95def95main} r_\ell(\vartheta):=h_{\ell+1}(\vartheta)-h_\ell(\vartheta) =F\!\big(h_\ell(\vartheta);\vartheta_\ell\big)\in\mathbb{R}^d.\tag{3}\] We measure directional alignment of consecutive residual increments by \[\label{eq:avg95cos95def95main} \overline{\cos}(\vartheta) :=\frac{1}{L-1}\sum_{\ell=1}^{L-1}\cos\!\big(r_{\ell+1}(\vartheta),r_\ell(\vartheta)\big).\tag{4}\] This metric is invariant to the magnitudes \(\|r_\ell\|\) and isolates the alignment of representation directions. Empirically, we consistently observe that \(\overline{\cos}\) increases under smoothing, as shown in Figure 5.
5.0.0.2 Baseline and smoothed one-step updates.
Let \(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 \[\label{eq:one95step95updates95main} \theta^b=\theta-\eta\,u(\theta), \qquad \theta^s=\theta-\eta\,S u(\theta),\tag{5}\] where \(S\) is a depth-smoothing operator. We focus on the local-averaging smoother (Window Smoothing) with \(\alpha\in(0,\tfrac12]\): \[\label{eq:S95local95avg95main} (Su)_\ell = \begin{cases} \big(1-\tfrac{\alpha}{2}\big)u_1+\tfrac{\alpha}{2}u_2, & \ell=1,\\[2pt] (1-\alpha)u_\ell+\tfrac{\alpha}{2}(u_{\ell-1}+u_{\ell+1}), & 2\le \ell\le L-1,\\[2pt] \big(1-\tfrac{\alpha}{2}\big)u_L+\tfrac{\alpha}{2}u_{L-1}, & \ell=L. \end{cases}\tag{6}\]
5.0.0.3 Assumptions.
We assume standard Lipschitz regularity of the block map \(F\) and 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 9.2.
5.1 Main 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 variation \(D\theta\) and the update-induced variation \(Du(\theta)\), where \(D\) is the first-difference operator \((Dv)_\ell=v_{\ell+1}-v_\ell\).
Theorem 1 (Alignment Comparison under Smoothing). Assume the regularity conditions of Appendix 9.2 (Assumption 1). Let \(S\) be the local-averaging smoother 6 with \(\alpha\in(0,\tfrac12]\), and let \(D\) be the first-difference operator \((Dv)_\ell=v_{\ell+1}-v_\ell\). Define \[\delta := D\theta, \qquad w := Du(\theta),\] where norms and inner products are taken over the corresponding depth-stacked block vectors. Let \(S^{(1)}\) denote the induced averaging operator on first differences, so that \[D(Su)=S^{(1)}Du\] as shown in Appendix 9.2, Lemma 4.
Then the smoothed and baseline one-step updates satisfy \[\begin{align} \tag{7} \overline{\cos}(\theta^s) &\ge 1-\frac{4L_h^2M^2}{m^2} -\frac{4L_\theta^2}{m^2}\cdot\frac{1}{L-1}\, \big\|\delta-\eta S^{(1)}w\big\|^2,\\ \tag{8} \overline{\cos}(\theta^b) &\ge 1-\frac{4L_h^2M^2}{m^2} -\frac{4L_\theta^2}{m^2}\cdot\frac{1}{L-1}\, \big\|\delta-\eta w\big\|^2. \end{align}\] Moreover, the difference between these two lower bounds is exactly \[\label{eq:cert95gap95exact95main} \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].\tag{9}\] Consequently, if \(\eta>0\) and \(L_\theta>0\), smoothing gives a strictly improved lower bound whenever \[\label{eq:improvement95condition95main} \big\langle \delta,(I-S^{(1)})w\big\rangle < \frac{\eta}{2}\Big(\|w\|^2-\|S^{(1)}w\|^2\Big).\tag{10}\] In particular, since \(S^{(1)}\) is a contraction with \[\label{eq:mu95star95main} \|S^{(1)}\|_{\mathrm{op}} = \mu_\star := 1-\alpha\Big(1-\cos\Big(\frac{\pi}{L}\Big)\Big) <1,\tag{11}\] the update-induced depth variation satisfies \[\label{eq:update95roughness95contraction95main} \|S^{(1)}w\|^2 \le \mu_\star^2\|w\|^2.\tag{12}\]
The proof is given in Appendix 9.2.
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 difference \(w=Du(\theta)\), while the smoothed bound depends on the averaged update difference \(S^{(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\] measures the reduction in depth-to-depth update variation caused by smoothing, while \[\langle \delta,(I-S^{(1)})w\rangle\] measures how the removed component of the update variation aligns with the existing parameter variation \(\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 \(\delta=0\), this obstruction vanishes, and the improvement follows directly from the contraction of update differences by \(S^{(1)}\).
6 Discussion↩︎
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). 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.
7 Related Work↩︎
7.0.0.1 Transformer Layer Regularity.
Recent empirical work shows that trained transformers exhibit strong layer-wise structure in their representations and Jacobians [3]–[5], [17]–[21], with some links of such structure to improved generalization. A closely related line of work on compression valleys [6], [22] 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 [5], [23]–[25]. Our work connects to this literature by showing that gradient smoothing amplifies these empirical alignment tendencies, while leading to improved performance.
7.0.0.2 Adaptive Optimization.
Modern deep learning relies heavily on adaptive optimizers that scale updates based on gradient statistics. Adam [26] and AdamW [27] maintain adaptive per-parameter estimates of first and second moments, enabling effective training across all modern architectures. Second-order methods such as Shampoo [28] and SOAP [29] approximate full-matrix preconditioning for improved convergence, while Muon [12] 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.
7.0.0.3 Pre-conditioned Optimization.
Natural gradient descent [30] and its approximations such as K-FAC [31]–[35] 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 [36]–[50], especially in transformers. In LLMs, recent works identify uniform token structures [6], [22], [51]–[53] and low-dimensional hidden trajectories [54], [55]. Our work proposes an approach to regularize representation trajectories across depth by augmenting the gradient updates.
8 Conclusion↩︎
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.
9 Appendix↩︎
9.1 Specialization of Smoothing to Adam, AdamW, and Muon↩︎
In our experiments, the base optimizer \(\mathcal{U}^{(t)}\) is typically Adam [26], AdamW [27], or Muon [12]. For completeness, we briefly recall the update rule and describe how gradient smoothing composes with it.
9.1.0.1 Adam updates.
Let \(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 \(\phi\) are 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, \[\begin{align} m^{(t)} &= \beta_1 m^{(t-1)} + (1-\beta_1) g^{(t)}, \\ v^{(t)} &= \beta_2 v^{(t-1)} + (1-\beta_2)\big(g^{(t)}\odot g^{(t)}\big), \end{align}\] with bias-corrected versions \(\hat{m}^{(t)}\) and \(\hat{v}^{(t)}\). The update direction is \[u^{(t)} := \frac{\hat{m}^{(t)}}{\sqrt{\hat{v}^{(t)}} + \varepsilon},\] and the block parameters are updated as \[\theta^{(t+1)} = \theta^{(t)} - \eta\, u^{(t)}.\]
Equivalently, Adam applies a diagonal, coordinate-wise preconditioner \[u^{(t)} = A^{(t)} g^{(t)}, \qquad A^{(t)} := \mathrm{diag}\!\left(\frac{1}{\sqrt{\hat{v}^{(t)}}+\varepsilon}\right).\]
9.1.0.2 AdamW.
AdamW applies decoupled weight decay, yielding the block-parameter update \[\theta^{(t+1)} = \theta^{(t)} - \eta\big(u^{(t)} + \lambda\, \theta^{(t)}\big),\] with non-block parameters \(\phi\) updated by the base optimizer in the standard way.
9.1.0.3 Adam/AdamW with Gradient Smoothing.
The smoothed block update is then \[\tilde{u}^{(t)} = P\, u^{(t)}.\] Gradient smoothing yields \[\begin{align} \theta^{(t+1)} &= \theta^{(t)} - \eta\, \tilde{u}^{(t)} \quad \text{(Adam)}, \\ \theta^{(t+1)} &= \theta^{(t)} - \eta\big(\tilde{u}^{(t)} + \lambda\, \theta^{(t)}\big) \quad \text{(AdamW)}. \end{align}\]
Equivalently, the effective update can be written as \[\theta^{(t+1)} = \theta^{(t)} - \eta\, (P A^{(t)}) g^{(t)},\] highlighting that gradient smoothing composes a depth-structured preconditioner with Adam’s coordinate-wise preconditioning, with \(\phi\) being updated as in standard Adam/AdamW.
9.1.0.4 Muon.
Muon [12] treats each repeated-block parameter as a matrix and produces an approximately orthogonal update direction. It maintains a momentum buffer \[M^{(t)} = \mu\, M^{(t-1)} + g^{(t)},\] and forms the update direction by orthogonalizing \(M^{(t)}\) via a Newton–Schulz iteration, \[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 \(\theta^{(t+1)} = \theta^{(t)} - \eta\,u^{(t)}\). In contrast to Adam’s diagonal preconditioning, \(u^{(t)}\) has singular values approximately equal to one, so its per-step magnitude is largely decoupled from that of the raw gradient.
9.1.0.5 Muon with Gradient Smoothing.
The smoothed block update composes with Muon in the same way as with Adam: \[\tilde{u}^{(t)} = P\, u^{(t)}, \qquad \theta^{(t+1)} = \theta^{(t)} - \eta\, \tilde{u}^{(t)}.\] For matrix-valued block parameters, \(P\) is 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, \(P\) is applied directly to the base-optimizer update for that parameter.
9.2 Alignment of Representation Residuals under Gradient Smoothing↩︎
9.2.0.1 Residual increments and alignment metric.
For any depth-stacked parameter vector \(\vartheta=(\vartheta_1,\dots,\vartheta_L)\), define the residual increments \[\label{eq:r95def} 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,\tag{13}\] and the normalized directions \[\label{eq:d95def} d_\ell(\vartheta)\;:=\;\frac{r_\ell(\vartheta)}{\|r_\ell(\vartheta)\|}\in\mathbb{R}^d \quad\text{(defined whenever r_\ell(\vartheta)\neq 0).}\tag{14}\] Our empirical metric is \[\label{eq:cos95metric} \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.\tag{15}\] Working with \(d_\ell\) isolates directional alignment from magnitude effects (e.g.\(\|r_{\ell+1}-r_\ell\|\) may be large even when \(\cos(r_{\ell+1},r_\ell)\) is close to \(1\) due to norm mismatch).
9.2.0.2 Assumptions.
We work on a region containing the trajectories \(\{h_\ell(\theta^b)\}\) and \(\{h_\ell(\theta^s)\}\) and impose the following standard regularity assumptions.
Assumption 1. There exist constants \(L_h,L_\theta\ge 0\) such that for all \(h,\bar h\in\mathbb{R}^d\) and \(\theta,\bar\theta\in\mathbb{R}^p\) in the region of interest, \[\label{eq:F95Lip} \|F(h;\theta)-F(\bar h;\bar\theta)\| \;\le\;L_h\|h-\bar h\|+L_\theta\|\theta-\bar\theta\|.\tag{16}\]
Moreover, there exist constants \(0<m\le M<\infty\) such that, for each \(\vartheta\in\{\theta^b,\theta^s\}\), \[m \le \min_{\ell\in[L]}\|r_\ell(\vartheta)\|, \qquad \max_{\ell\in[L]}\|r_\ell(\vartheta)\|\le M .\]
Lemma 1. For any nonzero \(a,b\in\mathbb{R}^d\), \[\label{eq:cos95identity} 1-\cos(a,b)\;=\;\frac{1}{2}\left\|\frac{a}{\|a\|}-\frac{b}{\|b\|}\right\|^2.\tag{17}\]
Proof. Expand \(\|a/\|a\|-b/\|b\|\|^2=2-2\langle a/\|a\|,b/\|b\|\rangle\) and rearrange. ◻
Lemma 2. For any nonzero \(a,b\in\mathbb{R}^d\), \[\label{eq:norm95lipschitz} \left\|\frac{a}{\|a\|}-\frac{b}{\|b\|}\right\| \;\le\;\frac{2\|a-b\|}{\min(\|a\|,\|b\|)}.\tag{18}\]
Proof. Write \[\frac{a}{\|a\|}-\frac{b}{\|b\|} =\frac{a-b}{\|a\|}+b\Big(\frac{1}{\|a\|}-\frac{1}{\|b\|}\Big),\] so by the triangle inequality, \[\left\|\frac{a}{\|a\|}-\frac{b}{\|b\|}\right\| \le \frac{\|a-b\|}{\|a\|}+\|b\|\cdot\frac{\big|\|b\|-\|a\|\big|}{\|a\|\|b\|} \le \frac{\|a-b\|}{\|a\|}+\frac{\|a-b\|}{\|a\|} =\frac{2\|a-b\|}{\|a\|},\] and symmetrizing in \((a,b)\) yields 18 . ◻
9.2.0.3 A bound on cosine misalignment.
Lemma 3. Under Assumption 1, for any \(\vartheta\in\{\theta^b,\theta^s\}\) and any \(\ell\in[L-1]\), \[\label{eq:turning95bound95pointwise} 1-\cos\!\big(r_{\ell+1}(\vartheta),r_\ell(\vartheta)\big) \;\le\;\frac{4}{m^2}\Big(L_h^2\|r_\ell(\vartheta)\|^2 + L_\theta^2\|\vartheta_{\ell+1}-\vartheta_\ell\|^2\Big).\tag{19}\] Consequently, \[\label{eq:avg95turning95bound} \frac{1}{L-1}\sum_{\ell=1}^{L-1}\cos\!\big(r_{\ell+1}(\vartheta),r_\ell(\vartheta)\big) \;\ge\; 1-\frac{4L_h^2M^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.\tag{20}\]
Proof. Fix \(\vartheta\) and \(\ell\). By Lemma 1 and Lemma 2, \[1-\cos(r_{\ell+1},r_\ell) =\frac{1}{2}\|d_{\ell+1}-d_\ell\|^2 \le \frac{1}{2}\left(\frac{2\|r_{\ell+1}-r_\ell\|}{\min(\|r_{\ell+1}\|,\|r_\ell\|)}\right)^2 \le \frac{2}{m^2}\|r_{\ell+1}-r_\ell\|^2.\] Using 13 and the Lipschitz property 16 , \[\|r_{\ell+1}-r_\ell\| =\big\|F(h_{\ell+1};\vartheta_{\ell+1})-F(h_\ell;\vartheta_\ell)\big\| \le 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\le 2a^2+2b^2\) yields \[\|r_{\ell+1}-r_\ell\|^2 \le 2L_h^2\|r_\ell\|^2+2L_\theta^2\|\vartheta_{\ell+1}-\vartheta_\ell\|^2.\] Combining the last two inequalities gives 19 . Averaging 19 over \(\ell\) and using \(\|r_\ell(\vartheta)\|\le M\) gives 20 . ◻
9.2.0.4 Local averaging smoother and its action on update differences.
We now take \(S\) to be the local averaging (tridiagonal) smoother with strength \(\alpha\in(0,\tfrac12]\): \[\label{eq:S95local95avg} (Su)_\ell = \begin{cases} \big(1-\tfrac{\alpha}{2}\big)u_1+\tfrac{\alpha}{2}u_2, & \ell=1,\\[2pt] (1-\alpha)u_\ell+\tfrac{\alpha}{2}(u_{\ell-1}+u_{\ell+1}), & 2\le \ell\le L-1,\\[2pt] \big(1-\tfrac{\alpha}{2}\big)u_L+\tfrac{\alpha}{2}u_{L-1}, & \ell=L. \end{cases}\tag{21}\] Let \(D\in\mathbb{R}^{(L-1)\times L}\) denote the first-difference operator \((Dv)_\ell=v_{\ell+1}-v_\ell\).
Lemma 4 (Induced local averaging on first differences). Let \(D\) be the first-difference operator \((Du)_\ell=u_{\ell+1}-u_\ell\) for \(\ell=1,\dots,L-1\). Let \(S\) be the local averaging operator defined in 21 . Define \(S^{(1)}\in\mathbb{R}^{(L-1)\times(L-1)}\) by \[(S^{(1)}w)_\ell= \begin{cases} (1-\alpha)w_1+\tfrac{\alpha}{2}w_2, & \ell=1,\\[2pt] (1-\alpha)w_\ell+\tfrac{\alpha}{2}(w_{\ell-1}+w_{\ell+1}), & 2\le \ell\le L-2,\\[2pt] (1-\alpha)w_{L-1}+\tfrac{\alpha}{2}w_{L-2}, & \ell=L-1. \end{cases}\] Then for all block-vectors \(u\), \[D(Su)=S^{(1)}(Du).\]
Proof. Let \(w:=Du\), i.e.\(w_\ell=u_{\ell+1}-u_\ell\). For interior \(\ell\in\{2,\dots,L-2\}\), expanding \((Su)_{\ell+1}-(Su)_\ell\) using 21 gives \[(D(Su))_\ell=(1-\alpha)w_\ell+\tfrac{\alpha}{2}w_{\ell-1}+\tfrac{\alpha}{2}w_{\ell+1}.\] At \(\ell=1\), \[(D(Su))_1=(Su)_2-(Su)_1=(1-\alpha)w_1+\tfrac{\alpha}{2}w_2.\] At \(\ell=L-1\), \[(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 of \(S^{(1)}\) applied to \(w\), hence \(D(Su)=S^{(1)}w=S^{(1)}(Du)\). ◻
Lemma 5 (Eigenvalues of induced averaging on first differences). Let \(n\ge 1\) and define \(\widetilde{S}^{(n)}\in\mathbb{R}^{n\times n}\) by the tridiagonal Toeplitz form \[(\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 \(\alpha\in(0,1)\). Then \(\widetilde{S}^{(n)}\) is symmetric and has eigenpairs \[v^{(k)}_j=\sin\!\Big(\frac{\pi k j}{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, \(\|\widetilde{S}^{(n)}\|_{\mathrm{op}}=\widetilde{\mu}_1^{(n)}<1\).
Proof. Fix \(k\in\{1,\dots,n\}\) and set \(t:=\pi k/(n+1)\). For \(2\le j\le n-1\), \[(\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),\] using \(\sin((j-1)t)+\sin((j+1)t)=2\sin(jt)\cos t\). At \(j=1\) and \(j=n\) the same identity holds since the missing neighbor corresponds to \(\sin(0)=\sin((n+1)t)=0\). Thus \(v^{(k)}\) is an eigenvector with eigenvalue \(1-\alpha+\alpha\cos t=1-\alpha(1-\cos t)\). The operator norm equals the largest absolute eigenvalue \(\widetilde{\mu}_1^{(n)}\) since \(\widetilde{S}^{(n)}\) is symmetric. ◻
Proposition 2. Let \(S^{(1)}\in\mathbb{R}^{(L-1)\times(L-1)}\) be the induced operator on differences from Lemma 4 (i.e.with boundary diagonal \(1-\alpha\)). Then for all \(w\in\mathbb{R}^{L-1}\), \[\label{eq:S195contraction95new} \|S^{(1)}w\| \;\le\;\mu_\star\,\|w\|, \qquad \mu_\star := 1-\alpha\Big(1-\cos\Big(\frac{\pi}{L}\Big)\Big)\;<\;1,\qquad{(1)}\] and hence \[\label{eq:energy95drop95new} \|w\|^2-\|S^{(1)}w\|^2 \;\ge\;\big(1-\mu_\star^2\big)\,\|w\|^2.\qquad{(2)}\]
Proof. By Lemma 5 with \(n=L-1\), the eigenvalues of \(S^{(1)}=\widetilde{S}^{(L-1)}\) are \(\widetilde{\mu}_k^{(L-1)} = 1-\alpha(1-\cos(\pi k/L))\) for \(k=1,\dots,L-1\). Thus \(\|S^{(1)}\|_{\mathrm{op}}=\widetilde{\mu}_1^{(L-1)}=\mu_\star<1\), giving ?? . Then ?? follows from \(\|S^{(1)}w\|^2\le \mu_\star^2\|w\|^2\). ◻
9.2.0.5 General one-step bound for the average cosine alignment.
Define the average cosine across consecutive residual increments: \[\label{eq:avg95cos95def} \overline{\cos}(\vartheta) := \frac{1}{L-1}\sum_{\ell=1}^{L-1} \cos\!\big(r_{\ell+1}(\vartheta),r_\ell(\vartheta)\big).\tag{22}\] We now state the general one-step comparison, allowing the current parameters to have pre-existing depth variation.
Corollary 1. Assume Assumption 1. Let \(S\) be the local-averaging smoother 21 with parameter \(\alpha\in(0,\tfrac12]\), and let \(S^{(1)}\) be the induced \((L-1)\times(L-1)\) averaging operator from Lemma 4. Define \[\delta := D\theta, \qquad w := Du(\theta),\] where \(D\) is 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 \[\begin{align} \tag{23} \overline{\cos}(\theta^s) &\ge 1-\frac{4L_h^2M^2}{m^2} -\frac{4L_\theta^2}{m^2}\cdot\frac{1}{L-1}\, \big\|\delta-\eta S^{(1)}w\big\|^2,\\ \tag{24} \overline{\cos}(\theta^b) &\ge 1-\frac{4L_h^2M^2}{m^2} -\frac{4L_\theta^2}{m^2}\cdot\frac{1}{L-1}\, \big\|\delta-\eta w\big\|^2. \end{align}\] Moreover, the difference between the two lower bounds admits the exact identity \[\label{eq:cert95gap95exact} \Big[\text{RHS of \eqref{eq:avg95cos95general95s}}\Big] - \Big[\text{RHS of \eqref{eq:avg95cos95general95b}}\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].\tag{25}\] In particular, if \(\eta>0\) and \(L_\theta>0\), smoothing yields a strictly improved lower bound whenever \[\label{eq:improvement95condition} \big\langle \delta,(I-S^{(1)})w\big\rangle < \frac{\eta}{2}\Big(\|w\|^2-\|S^{(1)}w\|^2\Big).\tag{26}\] A simple sufficient condition for such is \[\label{eq:improvement95condition95sufficient} \|\delta\| < \frac{\eta}{2}\cdot \frac{1-\mu_\star^2}{\|I-S^{(1)}\|_{\mathrm{op}}}\,\|w\|, \qquad \mu_\star:=\|S^{(1)}\|_{\mathrm{op}}<1.\tag{27}\]
Proof. The bounds 23 –24 follow by applying Lemma 3 with \(\vartheta=\theta^s\) and \(\vartheta=\theta^b\). Indeed, \[D\theta^s = D(\theta-\eta Su) = D\theta-\eta D(Su) = \delta-\eta S^{(1)}w,\] where we used Lemma 4, and similarly \[D\theta^b = D(\theta-\eta u) = \delta-\eta w.\] Substituting these identities into 20 gives 23 and 24 .
For 25 , subtract the right-hand sides of 23 and 24 . The only changing term is the squared depth-difference term, and \[\|\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 \[-\frac{4L_\theta^2}{m^2}\cdot\frac{1}{L-1}\] yields 25 . Condition 26 makes the right-hand side strictly positive.
Finally, 27 follows from Cauchy–Schwarz and the contraction of \(S^{(1)}\). Specifically, \[\big\langle \delta,(I-S^{(1)})w\big\rangle \le \|\delta\|\,\|I-S^{(1)}\|_{\mathrm{op}}\|w\|,\] while Proposition 2 gives \[\|w\|^2-\|S^{(1)}w\|^2 \ge (1-\mu_\star^2)\|w\|^2.\] Combining these two inequalities gives the stated sufficient condition. ◻
Remark 3 (Depth-flat special case). If the current parameters are depth-flat, \(\theta_1=\cdots=\theta_L\), then \(\delta=D\theta=0\). In this case, Corollary 1 reduces to a purely update-induced comparison: \[\label{eq:avg95cos95lower95smoothed95depth95flat} \overline{\cos}(\theta^s) \ge 1-\frac{4L_h^2M^2}{m^2} -\frac{4L_\theta^2\eta^2}{m^2}\cdot \frac{\|S^{(1)}w\|^2}{L-1}.\tag{28}\] The corresponding improvement over the baseline lower bound is \[\begin{align} \label{eq:avg95cos95cert95gap95depth95flat} \Big[\text{RHS of \eqref{eq:avg95cos95lower95smoothed95depth95flat}}\Big] - \Big[\text{same bound with \|S^{(1)}w\|^2 replaced by \|w\|^2}\Big] &= \frac{4L_\theta^2\eta^2}{m^2}\cdot \frac{\|w\|^2-\|S^{(1)}w\|^2}{L-1} \\ &\ge \frac{4L_\theta^2\eta^2}{m^2}\cdot \frac{1-\mu_\star^2}{L-1}\,\|w\|^2. \end{align}\tag{29}\] Thus, when \(\eta>0\), \(L_\theta>0\), and \(w\neq 0\), depth-flat parameters yield a strict improvement in the lower bound. This special case isolates the mechanism: local averaging improves the bound solely by contracting depth-to-depth variation in the optimizer update. The main result, Corollary 1, does not require this depth-flat condition and additionally accounts for interactions with pre-existing depth variation \(D\theta\).
9.3 Variance-Reduction of Gradients with Smoothing↩︎
Here we provide a simple variance-based analysis of Gradients with the local window operator. Let \(U \in \mathbb{R}^{Ld}\) denote the stacked stochastic (e.g.from minibatches) block update, and define the smoothed update by \[\tilde{U} = (S \otimes I_d) U,\] where \(S\) is the depth-smoothing matrix.
Proposition 4 (Variance under a single step of window smoothing). Let \(U \in \mathbb{R}^{Ld}\) be a random vector with \(\mathbb{E}\|U\|_2^2 < \infty\) and covariance \(\Sigma = \mathrm{Cov}(U)\), and let \(\tilde{U} = (S \otimes I_d) U\) with \(S\) the local window operator 21 . For any \(\alpha \in [0,1]\), \[\mathbb{E}\|\tilde{U} - \mathbb{E}\tilde{U}\|_2^2 \;\le\; \mathbb{E}\|U - \mathbb{E}U\|_2^2.\]
Proof. Since \(S\) is symmetric with eigenvalues \[\lambda_k(S) = 1 - \alpha\!\left(1 - \cos\!\frac{k\pi}{L}\right), \qquad k = 0, \dots, L-1,\] we have \(\|S\|_2 = 1\) for \(\alpha \in [0,1]\).
By linearity, \(\mathrm{Cov}(\tilde{U}) = (S \otimes I_d)\,\Sigma\,(S \otimes I_d)^\top\), so the total stochastic update variance is \[\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).\] Let \(A := (S^\top S) \otimes I_d\). Since \(\Sigma\) is a covariance matrix, \(\Sigma \succeq 0\), and clearly \(A \succeq 0\) as well. Then \(A \preceq \|A\|_2\,I\), which implies \[\mathrm{tr}(A\Sigma) = \mathrm{tr}\!\bigl(\Sigma^{1/2} A\,\Sigma^{1/2}\bigr) \le \|A\|_2\,\mathrm{tr}(\Sigma).\] Since \[\|A\|_2 = \|(S^\top S) \otimes I_d\|_2 = \|S^\top S\|_2 = \|S\|_2^2 = 1,\] we obtain \[\mathbb{E}\|\tilde{U} - \mathbb{E}\tilde{U}\|_2^2 \le \mathrm{tr}(\Sigma) = \mathbb{E}\|U - \mathbb{E}U\|_2^2. \qedhere\] ◻
Thus, for \(\alpha \in [0,1]\), window smoothing update variance is at most that of the blockwise stochastic update.
9.4 Empirical Measurement of Gradient Variance↩︎
To complement the variance contraction bound of Proposition 4, we directly measure two notions of stochastic gradient variance during pretraining and compare them across baseline and smoothing runs.
9.4.0.1 Setup.
We train three otherwise-identical nanochat d24 models (1.38B parameters, 24 transformer blocks) for 7,308 optimization steps: a baseline run and two window-smoothing runs with \(\alpha \in \{0.05,\,0.1\}\).
Every \(50\) steps, before applying the optimizer update, we record per-microbatch gradients for each parameter group \((r,\ell)\), where \(r\) indexes one
of \(R=6\) roles within a block (the four attention projections and the two MLP linears) and \(\ell \in \{1,\ldots,L\}\) indexes block depth. Let \(g_{r,\ell,n}\) denote the flattened gradient for role \(r\), layer \(\ell\), and microbatch \(n\), with \(N\) microbatches per logging step.
9.4.0.2 Microbatch variance.
The microbatch variance measures the across-sample noise of the stochastic gradient at each parameter group: \[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 \[V^{\mathrm{mb}}_{\mathrm{total}} \;:=\; \sum_{r=1}^{R}\sum_{\ell=1}^{L}V^{\mathrm{mb}}_{r,\ell},\] which is the empirical estimate of \(\mathbb{E}\|U-\mathbb{E}U\|_2^2\) in the notation of Proposition 4.
9.4.0.3 Depth variance.
The depth variance measures the across-layer dispersion of the gradient within each role, evaluated per microbatch and averaged over microbatches: \[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 is \(V^{\mathrm{depth}}_{\mathrm{total}}:=\sum_{r}V^{\mathrm{depth}}_{r}\), and is the empirical estimate of the quantity from Proposition 2.
9.4.0.4 Observations.
Figures 4 and 8 report \(V^{\mathrm{mb}}_{\mathrm{total}}\) and \(V^{\mathrm{depth}}_{\mathrm{total}}\) over training for the three runs. After the initial warmup transient (steps \(\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 4 acting at every step along the smoothed trajectory.
