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 \emph{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 \emph{Gradient Smoothing}, a family of depth-wise smoothing methods, and instantiate it with a simple local \emph{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 LLM 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.

• The paper introduces Depth-wise Gradient Augmentation via Gradient Smoothing to couple block updates and regularize optimization dynamics.
• Empirical results demonstrate improved convergence, higher test accuracies, and reduced gradient variance across diverse deep network architectures.
• The method is optimizer-agnostic, lightweight, and theoretically supported by bounds on layer update alignment and noise contraction.

Depth-wise Gradient Augmentation via Gradient Smoothing

Introduction

"Gradient Smoothing: Coupling Layer-wise Updates for Improved Optimization" (2606.30813) introduces Depth-wise Gradient Augmentation, an optimization paradigm designed for block-structured deep neural networks (e.g., Transformers, ResNets, Vision Transformers). It targets the block-wise optimizer updates—normally computed independently per layer—and transforms them through depth-wise operators before application. The framework specifically instantiates "Gradient Smoothing," whereby the optimizer updates for each layer are coupled via smoothing along the network depth, facilitating structured evolution of representations and optimization dynamics.

Depth-wise Gradient Augmentation Framework

Block-structured networks exhibit emergent cross-depth coordination in layer representations during training. Traditional optimizers treat block-wise gradients independently, ignoring this structure. In Depth-wise Gradient Augmentation, the collection of block-wise updates $(u_1, ..., u_L)$ generated by the base optimizer (e.g., Adam, Muon) is transformed by a depth-wise operator $S$, which acts independently on each within-block coordinate. The general setting utilizes $S \in \mathbb{R}^{L \times L}$ (typically symmetric/tridiagonal), yielding transformed updates $\tilde u = S u$.

Window Smoothing is the primary operator studied. For strength $\alpha$, the smoothed update for layer $l$ is:

$$\tilde u_l = \begin{cases} (1-\frac{\alpha}{2})u_1 + \frac{\alpha}{2}u_2, & l=1 \ (1-\alpha)u_l + \frac{\alpha}{2}(u_{l-1} + u_{l+1}), & 2 \leq l \leq L-1 \ (1-\frac{\alpha}{2})u_L + \frac{\alpha}{2}u_{L-1}, & l=L \end{cases}$$

This construction maintains row-stochasticity and requires no additional trainable parameters.

Gradient Smoothing: Mechanism and Variants

Gradient Smoothing acts as a structured preconditioning operator along the depth axis, regularizing block updates and leveraging the observed cross-layer dynamical coordination. The implementation is optimizer-agnostic, directly operating on the update vectors.

Normalization variants are considered:

• Standard: No normalization.
• Norm-Preserving: Each block update is rescaled to its original norm after smoothing.
• Directional: Updates are first normalized, smoothed, and then rescaled to original magnitudes.

Empirical Evaluation

Gradient Smoothing was evaluated across:

• RL fine-tuning of LLMs for reasoning
• LLM pretraining (nanochat recipe)
• Supervised image classification with Vision Transformers (ViT)
• Diffusion model training (U-ViT backbone)

In all settings, smoothing (Window operator) was applied directly atop tuned base optimizer pipelines (Adam/W/NormMuon), without hyperparameter changes or architectural modifications.

LLM Mathematical Reasoning

Window smoothing with $\alpha=0.1$ or $0.2$ improved pass@1 accuracy across mathematical benchmarks (AIME24, AIME25, AMC23, MATH-500) over the AdamW baseline. Gains were consistent across all datasets and optimizer configurations.

LLM Pretraining

Gradient Smoothing led to accelerated convergence of validation loss and improved CORE metric during GPT2 nanochat pretraining. The effect was amplified in deeper models (depth 30), indicating benefits scale with network depth.

Figure 1: Depth 24 validation loss demonstrating accelerated convergence with smoothing.

Vision Transformer Classification

Window smoothing with $\alpha = 0.2$ improved top-1 CIFAR-100 accuracy from 74.56% (baseline) to 75.62%. Gains persisted throughout 1700 epochs and across normalization variants.

Figure 2: Test accuracy improvement with smoothing for ViT-B on CIFAR-100.

Diffusion Modeling

Gradient Smoothing reduced FID@10k from 6.58 (baseline) to 5.82 and FID@50k from 4.01 to 3.74 for U-ViT diffusion models, showing applicability to generative architectures with skip connections.

Analysis of Gradient Smoothing Effects

Stability: Gradient Variance Reduction

Window smoothing consistently reduced microbatch gradient variance during nanochat pretraining, indicating increased training stability and noise contraction.

Figure 3: Microbatch gradient variance reduced by window smoothing through training.

Depth variance across layers also decreased, supporting the hypothesis that smoothing contracts stochastic update noise along the depth axis.

Figure 4: Lower depth gradient variance with smoothing in nanochat pretraining.

Representation Alignment and Linearity

Smoothing increases cosine similarity between layer contributions and induces greater linearity in representation trajectories, as quantified by line shape score (LSS).

Figure 5: Layer contributions similarity increases with greater smoothing.

Figure 6: Gradient smoothing yields higher alignment of layer contributions across depth.

Figure 7: Line shape score decreases (greater linearity) with increased smoothing strength.

This regularization effect facilitates coordinated representation evolution through the network and may support generalization.

Theoretical Results

A main theorem quantifies improved alignment of consecutive block residual increments under Window Smoothing, showing that smoothing contracts depth-to-depth variation in update directions. It provides explicit bounds on the improvement, isolating the contraction effect even when parameters exhibit pre-existing depth variation, and contrasts update-induced roughness contraction mechanisms with standard optimization.

Practical and Theoretical Implications

Gradient Smoothing is widely compatible, agnostic to base optimizer, and incurs negligible computational overhead. It regularizes optimization dynamics—with empirical gains in convergence speed, final accuracy, and stability—by exploiting depth-wise structure. Theoretical analysis underscores its effect as a structured preconditioning operator. Excessive smoothing, however, may hinder layer-specific adaptation, indicating an optimal balance must be maintained.

Further developments may include adaptive or learned smoothing operators, richer coupling schemes (e.g., exponential smoothing, selective coupling), and dynamic strength modulation across depth and training epochs. The framework motivates novel structured optimization methods for repeated-block architectures, especially as depth and parameter count scale.

Conclusion

Depth-wise Gradient Augmentation via Gradient Smoothing offers a principled, practical approach to regularizing block-wise updates in deep networks. Its consistent empirical benefits across diverse domains, theoretical guarantees of contraction and alignment, and compatibility with existing pipelines suggest wide applicability for improving optimization and generalization in modern architectures. This paradigm opens opportunities for further structured optimizers that leverage network depth and architectural repetition to improve learning outcomes.