Gradient Smoothing: Coupling Layer-wise Updates for Improved Optimization

2026-06-29

AI Summary

This work addresses the limited coordination among layer-wise updates during deep neural network training, which hinders optimization efficiency and generalization. To overcome this, the authors propose a Deep Gradient Enhancement framework that explicitly incorporates structural information along the depth dimension into the optimization process. By applying a windowed smoothing operator to block-level parameter updates generated by any base optimizer—such as SGD, Adam, or Muon—the method enables inter-layer coordination without altering the model architecture or objective function. Empirical results demonstrate consistent improvements in both training efficiency and generalization across diverse tasks, including language modeling, reinforcement learning fine-tuning, diffusion models, and Vision Transformer–based image classification.

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.

Problem

• depth-wise optimization
• gradient smoothing
• structured layer relationships
• optimization in deep networks
• cross-layer dependencies

Innovation

• Depth-wise Gradient Augmentation
• Gradient Smoothing
• Window Smoothing
• Structured Optimization
• Cross-layer Preconditioning