# Gradient Smoothing: Coupling Layer-wise Updates for Improved Optimization

arXiv: https://arxiv.org/abs/2606.30813 · HTML: https://mlsanity.com/paper/2606.30813/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 37 of its 50 references that are in the map.
- First posted on arXiv 29 Jun 2026.

## Builds on (most cited first, 30 of 37)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://mlsanity.com/paper/2402.03300/) (2024)
- [Training data-efficient image transformers & distillation through attention](https://mlsanity.com/paper/2012.12877/) (2020)
- [Optimizing Neural Networks with Kronecker-factored Approximate Curvature](https://mlsanity.com/paper/1503.05671/) (2015)
- [Prevalence of Neural Collapse during the terminal phase of deep learning training](https://mlsanity.com/paper/2008.08186/) (2020)
- [All are Worth Words: A ViT Backbone for Diffusion Models](https://mlsanity.com/paper/2209.12152/) (2022)
- [ShortGPT: Layers in Large Language Models are More Redundant Than You Expect](https://mlsanity.com/paper/2403.03853/) (2024)
- [Layer by Layer: Uncovering Hidden Representations in Language Models](https://mlsanity.com/paper/2502.02013/) (2025)
- [The Unreasonable Ineffectiveness of the Deeper Layers](https://mlsanity.com/paper/2403.17887/) (2024)
- [SOAP: Improving and Stabilizing Shampoo using Adam](https://mlsanity.com/paper/2409.11321/) (2024)
- [Fast Approximate Natural Gradient Descent in a Kronecker-factored Eigenbasis](https://mlsanity.com/paper/1806.03884/) (2018)
- [The Remarkable Robustness of LLMs: Stages of Inference?](https://mlsanity.com/paper/2406.19384/) (2024)
- [Traces of Class/Cross-Class Structure Pervade Deep Learning Spectra](https://mlsanity.com/paper/2008.11865/) (2020)
- [NorMuon: Making Muon more efficient and scalable](https://mlsanity.com/paper/2510.05491/) (2025)
- [Transformers represent belief state geometry in their residual stream](https://mlsanity.com/paper/2405.15943/) (2024)
- [Reinforcement Learning for Reasoning in Small LLMs: What Works and What Doesn't](https://mlsanity.com/paper/2503.16219/) (2025)
- [Kronecker-Factored Approximate Curvature for Modern Neural Network Architectures](https://mlsanity.com/paper/2311.00636/) (2023)
- [Deep Neural Collapse Is Provably Optimal for the Deep Unconstrained Features Model](https://mlsanity.com/paper/2305.13165/) (2023)
- [Linguistic Collapse: Neural Collapse in (Large) Language Models](https://mlsanity.com/paper/2405.17767/) (2024)
- [Tracing Representation Progression: Analyzing and Enhancing Layer-Wise Similarity](https://mlsanity.com/paper/2406.14479/) (2024)
- [Does Representation Matter? Exploring Intermediate Layers in Large Language Models](https://mlsanity.com/paper/2412.09563/) (2024)
- [Understanding Deep Representation Learning via Layerwise Feature Compression and Discrimination](https://mlsanity.com/paper/2311.02960/) (2023)
- [Separation and Concentration in Deep Networks](https://mlsanity.com/paper/2012.10424/) (2020)
- [Neural Collapse in the Intermediate Hidden Layers of Classification Neural Networks](https://mlsanity.com/paper/2308.02760/) (2023)
- [TREAD: Token Routing for Efficient Architecture-agnostic Diffusion Training](https://mlsanity.com/paper/2501.04765/) (2025)
- [Nearest Class-Center Simplification through Intermediate Layers](https://mlsanity.com/paper/2201.08924/) (2022)
- [Constrained belief updates explain geometric structures in transformer representations](https://mlsanity.com/paper/2502.01954/) (2025)
- [Residual Alignment: Uncovering the Mechanisms of Residual Networks](https://mlsanity.com/paper/2401.09018/) (2024)
- [Progressive Feedforward Collapse of ResNet Training](https://mlsanity.com/paper/2405.00985/) (2024)
- [Layers at Similar Depths Generate Similar Activations Across LLM Architectures](https://mlsanity.com/paper/2504.08775/) (2025)

## Abstract (arXiv)

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.