Skip to content

Gradient Smoothing: Coupling Layer-wise Updates for Improved Optimization

Jun 2026 · arXiv.org · Vol abs/2606.30813 · 0 citations · 50 references
Computer Science

TL;DR

Depth-wise Gradient Augmentation is established as a promising framework for exploiting cross-depth structure in optimization and Gradient Smoothing is demonstrated as a simple and broadly applicable instantiation.

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.

View source

Similar papers

Preprint Jul 2026

Building a Neural Network from Scratch: Implementation, Evaluation, and Optimization

This paper presents a self-contained neural network framework implemented entirely from scratch -- without relying on automatic differentiation or pre-built deep learning modules, including multi-layer architectures, diverse activation functions, regularization techniques, and state-of-the-art optimizers.

Yuanzhe Jia · 0 citations
Conference Jul 2026

Adaptive Gating Dynamics for Layer-Wise Sparsity Learning and Neural Network Pruning

Deep neural networks often contain substantial parameter redundancy, resulting in unnecessary computational cost and energy consumption. This work presents a dynamically adaptive gating mechanism for learning layer-wise sparsity through differentiable masking. In the proposed framework, each network parameter is associated with a self-learning gate that controls its contribution during training. The gating function follows a progressive soft-to-hard transition in which the slope is gradually annealed, enabling the model to move smoothly from continuous parameter weighting to near-binary pruning decisions while simultaneously learning an adaptive threshold parameter. The framework is evaluated across multilayer perceptrons (MLPs), deep neural networks (DNNs), Tabular Transformer models, and benchmarked against $\mathrm{L}_{0}$ regularization and Variational Dropout on convolutional neural networks. Experimental results demonstrate substantial model compression without degrading predictive performance, and in some cases improving it. On MNIST, the gated MLP retains only 14% of weights (86% pruning) while maintaining 98% accuracy. The gated DNN and Tab Transformer similarly outperform their dense counterparts while retaining only 46% and 46.22% of parameters, respectively. Benchmarking further shows that the proposed adaptive gating achieves competitive or superior accuracy-sparsity trade-offs compared with $\mathbf{L}_{0}$ regularization and Variational Dropout. These results demonstrate that the proposed dynamically adaptive gating framework provides an efficient and interpretable pathway for sparsity learning, enabling high-performing lightweight neural network deployment.

Raunak Dev, Mydhily Sankar, Devaprabha Biju S et al. · 0 citations
Book Open access Aug 2026

Towards Network-Efficient Cross-Regional Inference via Learned Activation Compression

Feather is presented, a system that reduces this overhead by compressing intermediate activations before transmission and reconstructing them before downstream layers resume execution, consistently outperforming existing compression schemes, including linear autoencoder, PCA, and SVD.

Regan McDonald, M. Rego, Ertza Warraich et al. · 0 citations
Jul 2026

ChACo: Channel-Wise Adaptive Competitive Layer-Wise Learning.

Local layer-wise learning offers modular optimization, layer-level transparency, and training without end-to-end error transport. However, its scalability remains limited by three coupled difficulties: local objectives can be weak or poorly aligned with the final task, shallow layers are often forced into premature fine-class discrimination, and fully local optimization is sensitive to activation and variance drift across depth. In this work, we present ChACo, a channel-wise adaptive competitive framework for fully local learning in convolutional networks. Each layer is trained with a local discriminative objective, while the internal block-to-class association is allowed to vary with depth. Early layers can map fewer competitive blocks to fine-label logits through a learned association policy, whereas later layers can recover direct fine-class competition as a special case. This design reduces the shallow-layer channel burden in many-class settings while preserving a direct local path to the final label space. The framework is supported by an analysis of local optimization dynamics, showing why even-power goodness functions and the ordering of rectification, convolution, and normalization are important for stable activation conditioning. We further incorporate lightweight stabilization components, including block-wise normalization and modular training schedules. Experiments on standard and many-class image-classification benchmarks show that ChACo improves over prior local-learning baselines, transfers across WAN and ResNet convolutional architectures, and narrows the gap to matched backpropagation (BP) references, while end-to-end BP remains stronger in the most depth-dependent settings. The source code of ChACo is available at https://github.com/andreaspapac/ChACo.

A. Papachristodoulou, C. Kyrkou, S. Timotheou et al. · 0 citations
Preprint Aug 2026

Scaling Muon for Diffusion Transformers

Periodic Row-wise Muon is introduced, which performs a full NS5 spectral update once every \(K\) steps and applies a low compute and communication cost row-wise constrained update based on the current momentum at the remaining steps to preserve Muon's generative quality advantage while translating it into end-to-end training efficiency for large DiTs.

Chenghao Li, Xiao Han, Xinxin Huang et al. · 0 citations