Gradient Smoothing: Coupling Layer-wise Updates for Improved Optimization
This paper introduces Depth-wise Gradient Augmentation, a paradigm that transforms optimizer updates along depth to leverage inter-layer structure. A simple instantiation, Gradient Smoothing with a window operator, consistently improves optimization and generalization across diverse tasks (LM pretraining, RL post-training, diffusion, ViT) without architectural changes. It promotes structured representation evolution, interpreted as depth-wise preconditioning.
-->
[Submitted on 29 Jun 2026]
Title:Gradient Smoothing: Coupling Layer-wise Updates for Improved Optimization
View a PDF of the paper titled Gradient Smoothing: Coupling Layer-wise Updates for Improved Optimization, by Haoming Meng and 2 other authors
View PDF HTML (experimental)
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 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.
Comments: Published in the Proceedings of the 43rd International Conference on Machine Learning (ICML 2026)
Subjects:
Machine Learning (cs.LG); Artificial Intelligence (cs.AI)
Cite as: arXiv:2606.30813 [cs.LG]
(or arXiv:2606.30813v1 [cs.LG] for this version)
https://doi.org/10.48550/arXiv.2606.30813
arXiv-issued DOI via DataCite (pending registration)
Submission history
From: Haoming Meng [view email] [v1] Mon, 29 Jun 2026 18:37:34 UTC (3,949 KB)
Full-text links:
Access Paper:
View a PDF of the paper titled Gradient Smoothing: Coupling Layer-wise Updates for Improved Optimization, by Haoming Meng and 2 other authors
View PDF
HTML (experimental)
TeX Source
view license
Current browse context:
cs.LG
new | recent | 2026-06
Change to browse by:
cs cs.AI
References & Citations
NASA ADS
Google Scholar
Semantic Scholar
Loading...
Data provided by:
Bibliographic Tools
Bibliographic and Citation Tools
Bibliographic Explorer Toggle
Bibliographic Explorer (What is the Explorer?)
Connected Papers Toggle
Connected Papers (What is Connected Papers?)
Litmaps Toggle
Litmaps (What is Litmaps?)
scite.ai Toggle
scite Smart Citations (What are Smart Citations?)
Code, Data, Media
Code, Data and Media Associated with this Article
alphaXiv Toggle
alphaXiv (What is alphaXiv?)
Links to Code Toggle
CatalyzeX Code Finder for Papers (What is CatalyzeX?)
DagsHub Toggle
DagsHub (What is DagsHub?)
GotitPub Toggle
Gotit.pub (What is GotitPub?)
Huggingface Toggle
Hugging Face (What is Huggingface?)
ScienceCast Toggle
ScienceCast (What is ScienceCast?)
Demos
Demos
Replicate Toggle
Replicate (What is Replicate?)
Spaces Toggle
Hugging Face Spaces (What is Spaces?)
Spaces Toggle
TXYZ.AI (What is TXYZ.AI?)
Related Papers
Recommenders and Search Tools
Link to Influence Flower
Influence Flower (What are Influence Flowers?)
Core recommender toggle
CORE Recommender (What is CORE?)
IArxiv recommender toggle
IArxiv Recommender (What is IArxiv?)
Author
Venue
Institution
Topic
About arXivLabs
arXivLabs: experimental projects with community collaborators
arXivLabs is a framework that allows collaborators to develop and share new arXiv features directly on our website.
Both individuals and organizations that work with arXivLabs have embraced and accepted our values of openness, community, excellence, and user data privacy. arXiv is committed to these values and only works with partners that adhere to them.
Have an idea for a project that will add value for arXiv's community? Learn more about arXivLabs.
Which authors of this paper are endorsers? | Disable MathJax (What is MathJax?)