Gradient Smoothing: Coupling Layer-wise Updates for Improved Optimization

Deep neural networks with repeated blocks, such as transformers and ResNets, often exhibit structured relationships across layers that emerge during training. Motivated by this observation, we introduce a general paradigm of Depth-wise Gradient Augmentation, in which the update applied to a layer may depend on the base optimizer updates computed for other layers. We study an instantiation of this idea, termed Gradient Smoothing, which couples optimizer updates across depth and admits a natural interpretation as a structured preconditioning method. Our framework operates directly on block-wise update vectors produced by arbitrary base optimizers (e.g., SGD, Adam, Muon), applying structured depth-wise smoothing operators such as local weighted averaging with minimal computational overhead. We evaluate Gradient Smoothing across a diverse set of architectures and training regimes, including language model pretraining, RL post-training of LLMs on reasoning tasks, diffusion modeling, and image classification with Vision Transformers. Across these settings, Gradient Smoothing consistently improves convergence and generalization performance without modifying model architectures or training objectives. We further show that smoothing promotes more structured representation evolution across depth, suggesting a connection between structured update coupling and the internal organization of learned representations. These results position Gradient Smoothing as a simple and broadly applicable approach for improving training in modern deep networks.

Modern AI systems such as large language models, vision transformers, and diffusion models are built from many repeated architectural layers. Although these layers often learn related behaviors during training, today's optimization algorithms typically update each layer independently. We introduce a new training framework called Depth-wise Gradient Augmentation, which allows layers to share information through their optimization updates. We study a simple instance of this idea called Gradient Smoothing, where updates from neighboring layers are combined before being applied to the model. Gradient Smoothing can be added on top of existing training methods with very little computational cost and does not require changing the model architecture or training objective. We tested the method on a wide range of machine learning tasks, including language model pretraining, reasoning-focused reinforcement learning, image classification, and image generation. Across these settings, Gradient Smoothing consistently improved training efficiency and model performance. We also found that it encourages representations to evolve more smoothly and coherently across layers. These results suggest that leveraging the structure shared across layers can be a simple and broadly useful way to train modern deep learning systems more effectively.