StableGrad: Backward Scale Control without Batch Normalization

• URL: http://arxiv.org/abs/2605.19856v1
• Date: Tue, 19 May 2026 13:49:30 GMT
• Title: StableGrad: Backward Scale Control without Batch Normalization
• Abstract summary: Training very deep neural networks requires controlling the propagation of magnitudes across depth.<n>Modern architectures often mitigate this problem through Batch Normalization, residual connections, or other normalization layers.<n>We propose StableGrad, an alternative mechanism that corrects layer-wise weightgradient imbalances without modifying the forward model.
• Score: 1.1145952934885128
• License: http://creativecommons.org/licenses/by-nc-sa/4.0/
• Abstract: Training very deep neural networks requires controlling the propagation of magnitudes across depth. Without such control, activations and gradients may vanish, explode, or enter unstable regimes that make optimization fail. Modern architectures often mitigate this problem through Batch Normalization, residual connections, or other normalization layers, which repeatedly re-scale or bypass intermediate representations. However, these mechanisms are not always appropriate. In Physics-Informed Neural Networks (PINNs), the network represents a continuous physical field and its input derivatives define the training objective, making batch-dependent normalization problematic because it can introduce non-local dependencies into the predicted field and its derivatives. We propose StableGrad, an optimizer-level scale-control mechanism that corrects layer-wise weight-gradient imbalances without modifying the forward model. Because the normalization is applied only after backpropagation and before the optimizer update, the network output, its derivatives, and the physical residual remain unchanged. We analyze the effective training dynamics induced by this rescaling and evaluate StableGrad on deep PINNs as the target application, with BatchNorm-free convolutional networks serving as a diagnostic stress test. On PINN benchmarks, StableGrad improves matched-depth solution accuracy and makes deeper models more reliable under standard optimization. On ResNet and EfficientNet architectures, where removing Batch Normalization normally leads to training collapse, StableGrad stabilizes optimization without introducing any other architectural change. These results show that optimizer-level control of weight-gradient scale can provide a practical alternative when forward normalization is unavailable or undesirable.

Related papers

• When Does Scale-Invariant Optimization Become Unstable? An Exact Schedule Law with Weight Decay [6.625270757109195]
  We show that a single quantity captures all schedule and decay forcing, while norm growth induces an opposing geometric self-quenching effect.<n>This yields a sharp boundary that cleanly separates contraction- and expansion-dominated effective learning rate regimes.
  arXiv  Detail & Related papers  (2026-09-08T17:46:17Z)
• Reducing ANN-SNN Conversion Error via Residual Membrane Potential Alignment [52.910770415431244]
  Spiking Neural Networks (SNNs) serve as core architectures for neuromorphic computing.<n>ANN-SNN conversion circumvents issues by reusing well-trained ANN weights for low-latency, energy-efficient inference.<n>Existing conversion schemes suffer from severe accuracy drops at small timesteps, large inference delays and cumulative quantization errors.
  arXiv  Detail & Related papers  (2026-08-14T04:46:03Z)
• Intrinsically Stable Spiking Neural Networks: Overcoming the Performance Barrier in the Absence of Batch Normalization [6.399076427640221]
  This work proposes the Intrinsically Stable SNN architecture, which removes activation-normalization layers by enforcing signal homeostasis.<n>By folding the standardization operations into static weights offline, IS-SNN removes the runtime statistics tracking and multiplications introduced by activation normalization.<n>Experiments show that IS-SNN achieves performance competitive with or superior to computationally expensive dynamic BN techniques.
  arXiv  Detail & Related papers  (2026-06-30T14:07:21Z)
• AdamFLIP: Adaptive Momentum Feedback Linearization Optimization for Hard Constrained PINN Training [9.082213877961541]
  Physics-informed neural networks (PINNs) provide a flexible framework for solving forward and inverse problems governed by partial differential equations (PDEs)<n>Standard PINN training typically relies on soft penalty formulations that combine PDE residuals, data mismatch, and initial/boundary conditions using manually chosen weights.<n>In this work, we reformulate PINN training as an equality-constrained optimization problem and propose a novel Adaptive Momentum Feedback Linearization Optimization for Hard Constrained PINN (AdamFLIP)
  arXiv  Detail & Related papers  (2026-05-08T19:13:15Z)
• AutoBalance: An Automatic Balancing Framework for Training Physics-Informed Neural Networks [10.223108587188808]
  PINNs provide a powerful and general framework for solving Partial Differential Equations.<n> PINNs balance multiple loss terms, such as PDE residuals and boundary conditions, which often have conflicting objectives and vastly different curvatures.<n>Existing methods address this issue by manipulating gradients before optimization.<n>We introduce AutoBalance, a novel "post-combine" training paradigm.
  arXiv  Detail & Related papers  (2025-10-08T06:13:03Z)
• Precise Bayesian Neural Networks [0.0]
  We develop a lightweight, implementation-ready variational unit that fits modern normalized architectures and improves calibration without sacrificing accuracy.<n>In short, by aligning the variational posterior with the network's intrinsic geometry, BNNs can be simultaneously principled, practical, and precise.
  arXiv  Detail & Related papers  (2025-06-24T15:42:00Z)
• Achieving Constraints in Neural Networks: A Stochastic Augmented Lagrangian Approach [49.1574468325115]
  Regularizing Deep Neural Networks (DNNs) is essential for improving generalizability and preventing overfitting. We propose a novel approach to DNN regularization by framing the training process as a constrained optimization problem. We employ the Augmented Lagrangian (SAL) method to achieve a more flexible and efficient regularization mechanism.
  arXiv  Detail & Related papers  (2023-10-25T13:55:35Z)
• Globally Optimal Training of Neural Networks with Threshold Activation Functions [63.03759813952481]
  We study weight decay regularized training problems of deep neural networks with threshold activations. We derive a simplified convex optimization formulation when the dataset can be shattered at a certain layer of the network.
  arXiv  Detail & Related papers  (2023-03-06T18:59:13Z)
• Distribution Mismatch Correction for Improved Robustness in Deep Neural Networks [86.42889611784855]
  normalization methods increase the vulnerability with respect to noise and input corruptions. We propose an unsupervised non-parametric distribution correction method that adapts the activation distribution of each layer. In our experiments, we empirically show that the proposed method effectively reduces the impact of intense image corruptions.
  arXiv  Detail & Related papers  (2021-10-05T11:36:25Z)
• GradInit: Learning to Initialize Neural Networks for Stable and Efficient Training [59.160154997555956]
  We present GradInit, an automated and architecture method for initializing neural networks. It is based on a simple agnostic; the variance of each network layer is adjusted so that a single step of SGD or Adam results in the smallest possible loss value. It also enables training the original Post-LN Transformer for machine translation without learning rate warmup.
  arXiv  Detail & Related papers  (2021-02-16T11:45:35Z)
• Improve Generalization and Robustness of Neural Networks via Weight Scale Shifting Invariant Regularizations [52.493315075385325]
  We show that a family of regularizers, including weight decay, is ineffective at penalizing the intrinsic norms of weights for networks with homogeneous activation functions. We propose an improved regularizer that is invariant to weight scale shifting and thus effectively constrains the intrinsic norm of a neural network.
  arXiv  Detail & Related papers  (2020-08-07T02:55:28Z)
• Revisiting Initialization of Neural Networks [72.24615341588846]
  We propose a rigorous estimation of the global curvature of weights across layers by approximating and controlling the norm of their Hessian matrix. Our experiments on Word2Vec and the MNIST/CIFAR image classification tasks confirm that tracking the Hessian norm is a useful diagnostic tool.
  arXiv  Detail & Related papers  (2020-04-20T18:12:56Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.