通过层间梯度平滑提升深度模型优化效果,无需改架构或目标。
Gradient Smoothing: Coupling Layer-wise Updates for Improved Optimization

- 在层间梯度上做平滑处理,融合多层更新信息。
- 在语言模型、扩散模型等任务中稳定提升性能与泛化能力。
- 兼容主流优化器,适合追求训练稳定的深度学习研究者。
具有重复结构的深层神经网络(如Transformer)在训练过程中常表现出层间结构关系。受此启发,我们提出一种通用优化范式——深度梯度增强(Depth-wise Gradient Augmentation),其中每层的更新由沿深度方向的块级优化器更新进行变换得到。在此框架下,我们研究了梯度平滑(Gradient Smoothing)这一类深度平滑方法,并以简单的局部窗口平滑(Window Smoothing)算子为例进行实现。该方法直接作用于任意基础优化器(如SGD、Adam、Muon)产生的块级更新,计算开销极小,且与现有优化流程完全兼容。我们在多种架构和训练场景中评估该方法,包括语言模型预训练、大模型强化学习后训练、扩散建模及视觉Transformer图像分类。实验表明,梯度平滑在所有设置中均能一致提升优化与泛化性能,且不改变模型结构或训练目标。进一步分析显示,该方法促进了跨层表征演化的结构性,符合其作为结构化深度预条件化方法的定位。这些结果确立了深度梯度增强为利用优化中跨层结构的有力框架,并验证了梯度平滑作为简单、普适的实例有效性。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。