StableGrad: Backward Scale Control without Batch Normalization
arXiv:2605.19856v1 [cs.LG] 19 May 2026
Jose I. Mestre Universitat Politècnica de València [email protected] Cristian Pérez-Corral Universitat Politècnica de València [email protected]
Alberto Fernández-Hernández Universitat Politècnica de València [email protected]
Manuel F. Dolz Universitat Jaume I [email protected]
Enrique S. Quintana-Ortí Universitat Politècnica de València [email protected]
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.
1
Introduction
Depth is a central source of expressive power in Deep Neural Networks (DNNs), but it also makes training increasingly sensitive to the propagation of magnitudes across layers. Activations must remain well-scaled in the forward pass, while gradients must remain useful in the backward pass. Although these two requirements are related, they are not equivalent: preserving scale in one direction does not necessarily preserve it in the other. Classical initialization schemes address this problem by controlling weight variance through choices such as fan-in, fan-out, or combinations between them (Glorot and Bengio, 2010; He et al., 2015b). These choices implicitly define a trade-off between forward and backward stability: fan-in scaling primarily preserves activation magnitudes in the forward pass, whereas fan-out scaling Preprint.
primarily preserves gradient magnitudes in the backward pass. Moreover, initialization only sets the scale of the network at the beginning of training; it does not guarantee that activation or gradient magnitudes will remain well controlled as the weights evolve under optimization. Normalization layers, most notably Batch Normalization, reduce this burden by repeatedly re-scaling intermediate representations (Ioffe and Szegedy, 2015). However, such mechanisms are not always available. In Physics-Informed Neural Networks (PINNs), the network represents a continuous physical field whose derivatives are used to define the training objective (Raissi et al., 2019). Batch-dependent normalization can therefore interfere with the local interpretation of the predicted field and its derivatives. In this work, we propose StableGrad, an optimizer-level mechanism for controlling backward scale without introducing architectural normalization. Rather than relying on initialization to simultaneously balance forward and backward propagation, StableGrad acts after the backward pass and before the optimizer step, directly rescaling the weight gradients of each layer. Specifically, each layer-wise weight gradient is normalized using its own empirical standard deviation and rescaled using the standard deviation of the prediction gradient as a reference scale. This transfers the scale of the output adjoint to the parameter updates, mitigating depth-induced gradient-scale imbalances across layers. StableGrad is complementary to initialization rather than a replacement for it. In our experiments, we use an activation-aware fan-in initialization to obtain a well-scaled forward pass at the start of training, while StableGrad dynamically controls the backward gradient scale throughout optimization. Because the normalization is applied only to gradients before the optimizer update, it does not modify the forward computation, introduce batch-dependent predictions, or alter the physical residual. Our contributions are: • We introduce StableGrad, an optimizer-level gradient rescaling mechanism that stabilizes layer-wise weight-gradient magnitudes using the prediction-gradient scale as a reference. • We analyze the local training dynamics induced by StableGrad through its effective kernel, showing how layer-wise gradient rescaling changes the functional update. • We evaluate StableGrad in two settings where architectural normalization is undesirable or deliberately removed: BatchNorm-free CNNs, where it stabilizes training without BatchNorm, and deep PINNs, where it enables deeper networks and improves solution accuracy.
2
Background and Motivation
2.1
Forward and Backward Signal Propagation
A deep neural network applies a sequence of transformations of the form hℓ+1 = ϕ(zℓ ),
zℓ = Wℓ hℓ ,
(1)
where hℓ denotes the representation at layer ℓ, Wℓ the corresponding weight matrix, and ϕ a nonlinear activation. As depth increases, the scale of hℓ becomes increasingly sensitive to the statistics of the weights and activations. If this scale is not controlled, representations may progressively vanish, explode, or enter saturated regimes where optimization becomes difficult. The backward pass has an analogous propagation problem. Gradients are recursively transformed as ∂L ∂L = Wℓ⊤ ⊙ ϕ′ (zℓ ) , (2) ∂hℓ ∂hℓ+1 where ⊙ denotes the Hadamard, or element-wise, product. Thus, the same weights and nonlinearities that determine forward activation scales also affect backward gradient scales. However, the conditions for stable forward propagation and stable backward propagation are not identical. A choice of weight scale that preserves activation magnitudes does not necessarily preserve gradient magnitudes, especially in deep networks. Depth therefore turns scale propagation into a bidirectional constraint: activations must remain well-scaled in the forward pass, while the backpropagated adjoint signals must retain useful magnitudes before they induce parameter gradients. This issue is closely related to the classical vanishing and exploding gradient problem (Bengio et al., 1994), as well as later analyses of signal propagation, critical initialization, and dynamical isometry in deep networks (Saxe et al., 2013; Schoenholz et al., 2017; Pennington et al., 2017). 2
2.2
Initialization as a Trade-off
Weight initialization is the first mechanism used to control scale propagation. Classical initialization schemes, such as Xavier and Kaiming, set weight variance according to layer dimensions and activation statistics, aiming to preserve signal variance across depth (Glorot and Bengio, 2010; He et al., 2015b). This principle also persists in more recent initializers: even when they replace random sampling with structured constructions, or introduce refined activation-dependent scaling, the final magnitude is still governed by fan-in, fan-out, or combinations of both (Chang et al., 2020; FernándezHernández et al., 2026). Thus, initialization remains tied to a fan-mode choice. Using fan-in primarily favors forward activation stability, whereas fan-out favors backward gradient stability. Averages or other interpolations between fan-in and fan-out reduce the asymmetry, but they do not remove the trade-off; they merely choose a different compromise between the two propagation directions. Activation-dependent gains further adjust the weight scale to account for the expected effect of nonlinearities. However, these gains do not remove the underlying tension: initialization must still decide how much scale preservation to allocate to the forward pass and how much to the backward pass. This trade-off becomes more restrictive as depth increases, because small deviations from the desired scale can compound across many layers. This motivates a different view of initialization. Rather than treating it as a single mechanism that must balance forward and backward stability, we split the responsibility across the training procedure: initialization is used to preserve forward activations, while gradient normalization before the optimizer step is used to control the scale of weight gradients. 2.3
Why Forward Normalization is Problematic in PINNs
PINNs are a representative case where scale propagation becomes especially delicate. A PINN represents a continuous field, for example uθ (x, t), and is trained not only from data, but also by penalizing the residual of a differential equation (Raissi et al., 2019). Given collocation points Nf {(xi , ti )}i=1 , a typical physics loss is Nf
LPDE (θ) =
1 X 2 |N [uθ ](xi , ti )| , Nf i=1
(3)
where N is a differential operator involving derivatives of the network output with respect to its inputs. This changes the gradients received by the optimizer. For the physics loss, Nf
2 X ∇θ LPDE = N [uθ ](xi , ti )∇θ N [uθ ](xi , ti ). Nf i=1
(4)
Thus, if N contains terms such as ux , uxx , or ∆u, the gradient involves quantities such as ∇θ ∂x uθ , ∇θ ∂xx uθ , ∇θ ∆uθ . These are parameter sensitivities of input derivatives, not merely sensitivities of the output itself. Consequently, even in constant-width PINNs where hidden layers satisfy fan-in = fan-out, the physics-informed objective can induce gradient-scale imbalances across depth. The difficulty is therefore not only architectural; it is also introduced by the differential structure of the loss. Appendix A.1 shows that, unlike in standard supervised networks where forward and backward variance propagation can be approximated by scalar recursions, PINN residuals involving input derivatives create coupled derivative-adjoint channels. This makes it unlikely that a single initialization rule can reliably preserve backward scale across all relevant channels. Recent work has similarly shown that PINN training is often dominated by optimization pathologies rather than only by approximation capacity. Gradient-flow and NTK analyses identify imbalanced convergence rates across loss components and stiff training dynamics, while loss-landscape analyses connect these difficulties to ill-conditioning induced by differential operators (Wang et al., 2020, 2022). These observations are consistent with our motivation: differential residuals can create backward-scale imbalances that are not resolved by initialization alone. Batch-dependent normalization is also problematic in this context. The output at a point should represent a local physical quantity, and its derivatives should be consistent with that local interpretation. BatchNorm can interfere with this structure because the prediction at one point may depend on other points in the same batch. This is particularly problematic when collocation points, boundary points, 3
and initial-condition points are sampled from different distributions but are jointly used to define the physical objective. This creates a depth-scaling bottleneck: the standard architectural normalization tools that make deep networks trainable are largely unavailable, while the physics-informed objective can still produce unstable gradient magnitudes across depth. We therefore seek a strategy that preserves the physical forward model unchanged and moves gradient-scale control to the optimization procedure. 2.4
Architectural Mechanisms for Scale Stabilization
Architectural mechanisms provide a widely used practical answer to the scale propagation problem. Normalization layers such as Batch Normalization, Layer Normalization, Group Normalization, and related methods re-scale intermediate representations during training (Ioffe and Szegedy, 2015; Ba et al., 2016; Wu and He, 2018). These layers can be interpreted as scale-resetting mechanisms inserted inside the network. By repeatedly normalizing hidden representations, they reduce the burden placed on initialization and make both forward activations and backward gradients easier to control. Residual connections provide a complementary mechanism by creating shorter paths for signal and gradient propagation across depth (He et al., 2015a). Rather than re-scaling representations, they improve trainability by allowing information and gradients to bypass long chains of transformations. In practice, modern deep architectures often combine residual pathways with normalization layers, making optimization substantially less sensitive to the precise initialization scale. This is especially important in modern convolutional architectures, where Batch Normalization and residual connections are central components of stable deep training. With these architectural stabilizers, scale errors are repeatedly corrected or bypassed. Conversely, when such mechanisms are removed, the network again becomes much more exposed to the accumulation of scale errors across depth. However, these architectural stabilizers do not solve the setting considered here. Batch-dependent normalization can alter the local physical interpretation of PINN outputs and derivatives, while residual connections do not directly control the scale of the weight gradients produced by differential operators. We therefore seek a mechanism that leaves the forward model unchanged, does not modify the physical residual, and acts only on the gradients passed to the optimizer.
3
StableGrad
The previous discussion suggests a simple principle: forward-scale control and backward-scale control need not be imposed by the same mechanism. Initialization can be used to set a well-scaled forward pass at the beginning of training, while the scale of the gradients can be corrected directly before the optimizer update. StableGrad follows this principle. Let the trainable parameters be partitioned into layerwise blocks, θ = (θ1 , . . . , θL ), and let g ℓ = ∇θℓ L denote the gradient of block ℓ. After the backward pass, StableGrad computes the empirical standard deviation σℓ = std(g ℓ ) of each block gradient. It also computes a reference scale from the adjoint signal at the network output, σout = std(∂L/∂uθ ). The gradient passed to the optimizer is then σout ℓ geℓ = g , ℓ = 1, . . . , L, (5) σℓ + ε where ε > 0 is a small numerical constant. Thus, StableGrad transfers the scale of the signal that initiates the backward pass to all parameter blocks. The method acts after the loss and its derivatives have been computed, and before the optimizer update. The rescaled gradients geℓ are then passed to the optimizer in place of the raw gradients g ℓ . The choice of σout is motivated by the role of the output adjoint as the source of the backward signal. StableGrad does not aim to preserve the global norm of the gradient. Instead, it enforces a layerwise notion of backward-scale consistency: all blocks are updated from gradients expressed at the same statistical scale as the output adjoint. This is precisely the scale that is propagated backwards through the network. Appendix A.2 discusses this choice in more detail and describes alternative reference scales that can be used to separate layerwise balancing from changes in global gradient scale. 4
This makes StableGrad a natural companion to fan-in initialization. A fan-in scheme provides an activation-preserving starting point, setting forward magnitudes to propagate stably at initialization, while StableGrad dynamically controls the scale of layer-wise weight gradients throughout optimization. This separation avoids using initialization as a compromise between forward and backward propagation: forward activations are initialized in the regime targeted by fan-in, and gradient-scale imbalances are corrected by StableGrad after each backward pass, without introducing normalization layers into the forward model. The distinction is particularly useful in PINNs. The physical residual is evaluated from the network output and its input derivatives, such as ux , uxx , or ut . StableGrad leaves this forward map untouched: the network uθ , its derivatives, and the physics residual are computed exactly as in the original model. Only the gradient delivered to the optimizer is rescaled. In this sense, StableGrad provides backward stabilization while preserving the physical interpretation of the forward model. Figure 1 illustrates the different roles of initialization, BatchNorm, LayerNorm, and StableGrad in a controlled MLP. With vanilla initialization, fan-in preserves forward activation scales but leaves layer-wise weight gradients highly imbalanced, whereas fan-out improves backward signal scaling at the cost of increasing activation scales. BatchNorm and LayerNorm both reduce forwardscale sensitivity by repeatedly normalizing intermediate representations, making the fan-in and fan-out cases more similar. However, the resulting weight-gradient scales remain uneven across layers. StableGrad, instead, keeps the fan-in forward computation unchanged and acts only after backpropagation, equalizing the layer-wise weight-gradient scales passed to the optimizer.
1e+0 1e+3
1e+0 2e+3
1e+0 5e+3
1e+0 1e+4
6e-2
1e-1
3e-1
5e-1
1e+0
L1 (512)
L2 (256)
L3 (128)
L4 (64)
Weights
1e+0 6e+2
L5 (32)
2e+0 1e+4
4e+0 2e+4
8e+0 4e+4
2e+1 7e+4
3e+1 2e+5
1e+0
1e+0
1e+0
1e+0
1e+0
L1 (512)
L2 (256)
L3 (128)
1e+0 2e+3
1e+0 5e+3
1e+0 1e+4
1e-1
3e-1
5e-1
1e+0
L1 (512)
L2 (256)
L3 (128)
L4 (64)
L5 (32)
1e+0 6e+2
1e+0 1e+3
1e+0 2e+3
1e+0 5e+3
2e+0 1e+4
1e-1
3e-1
5e-1
1e+0
1e+0
L1 (512)
L2 (256)
L3 (128)
L4 (64)
Weights
1e+0 1e+3
6e-2
Weights
1e+0 6e+2
L5 (32)
1e+0 1e+3
1e+0 2e+3
1e+0 5e+3
1e+0 9e+3
7e-2
1e-1
3e-1
5e-1
1e+0
L1 (512)
L2 (256)
L3 (128)
L4 (64)
Weights
1e+0 7e+2
L5 (32)
1e+0 7e+2
1e+0 1e+3
1e+0 2e+3
1e+0 5e+3
2e+0 9e+3
1e-1
3e-1
5e-1
1e+0
1e+0
L1 (512)
L2 (256)
L3 (128)
L4 (64)
Weights
(f) MLP mode: withLayerNorm LayerNorm and fan-out | fan: fan-out Activations
(e) MLP with LayerNorm and fan-in mode: LayerNorm | fan: fan-in Activations
L5 (32)
(d) MLPmode: with BatchNorm and fan-out BatchNorm | fan: fan-out Activations
Activations
(c) MLP with BatchNorm and fan-in mode: BatchNorm | fan: fan-in
L4 (64)
Weights
(b) MLP with fan-out mode: Vanilla | fan: fan-out Activations
Activations
(a) MLP with fan-in mode: Vanilla | fan: fan-in
L5 (32)
1e+0 1e+0
1e+0 1e+0
1e+0 1e+0
1e+0 1e+0
1e+0 1e+0
6e-2
1e-1
3e-1
5e-1
1e+0
Activations Weights
Activations
(g) MLP withmode: fan-in and StableGrad Ours | fan: fan-in
Activation Gradients Weight Gradients
L1 (512)
L2 (256)
L3 (128)
L4 (64)
L5 (32)
Figure 1: Forward and backward scale diagnostics across depth. StableGrad preserves the forward computation and equalizes the scale of layerwise weight gradients before the optimizer update.
5
4
Effective Training Dynamics
We now analyze the local effect of StableGrad on the training dynamics. The full derivations are given in Appendices A.3 and A.4. Consider a weighted least-squares objective L(θ) =
1 ∥r(θ)∥2 , 2
where r(θ) ∈ RN is the vector of weighted residuals. This notation covers supervised residuals as well as the data, boundary, initial-condition, and PDE residuals used in PINNs. Let J = ∂r/∂θ be the residual Jacobian. Under the local linearization r(θ + ∆θ) ≃ r(θ) + J∆θ, a standard gradient step gives the residual dynamics r+ ≃ r − ηJJ ⊤ r. The matrix K = JJ ⊤ is the empirical neural tangent kernel governing this local evolution. StableGrad modifies the gradient before the optimizer step. If we define the block-diagonal matrix σout P = diag(α1 I1 , . . . , αL IL ), αℓ = , σℓ + ε then the idealized StableGrad step is ∆θ = −ηP J ⊤ r. The corresponding residual dynamics are governed by the effective kernel KSG = JP J ⊤ . (6) Writing the Jacobian by parameter blocks as J = [J1 , . . . , JL ], this kernel decomposes as KSG =
L X
αℓ Jℓ Jℓ⊤ .
ℓ=1
StableGrad therefore reweights the contribution of each layer to the functional training dynamics. To measure how strongly a kernel acts on the current residual, we use the residual Rayleigh quotient ρA (r) =
r⊤ Ar . ∥r∥2
We denote by ρ = ρK (r) the standard value and by ρSG = ρKSG (r) the StableGrad value. Since g ℓ = Jℓ⊤ r, the change induced by StableGrad is explicit: PL (αℓ − 1)∥g ℓ ∥2 . ρSG − ρ = ℓ=1 ∥r∥2 Hence, StableGrad increases the kernel action on the current residual whenever the amplified blocks carry enough gradient energy to dominate the blocks that are downscaled. The following result connects this quantity with the decrease of the linearized loss. Theorem 1 (Local decrease under StableGrad). Let K = JJ ⊤ and KSG = JP J ⊤ . Consider the linearized residual dynamics generated by the standard step and by the StableGrad step. If η ηλmax (KSG ) < 2 and ρSG 1 − λmax (KSG ) > ρ, (7) 2 then the StableGrad linearized step produces a larger decrease of 12 ∥r∥2 than the standard linearized gradient step. The proof is given in Appendix A.4. The theorem separates the two quantities that matter locally. The term ρSG measures the useful action of the effective kernel on the current residual, while ηλmax (KSG ) controls the stability of the local step. Thus, StableGrad improves the local linearized decrease when it increases the action of the training kernel on the residual and keeps the step in a stable regime. This result also gives a direct interpretation of the method. StableGrad is beneficial when the backward scale imbalance suppresses layers that still contain useful directions for reducing the current residual. 6
By rescaling those blocks, the method increases their contribution to KSG . The theorem then states when this increased contribution translates into a larger local decrease of the loss. Appendix B.1 provides a controlled Burgers PINN diagnostic that directly measures the quantities appearing in Theorem 1. In that experiment, StableGrad reduces the validation loss much faster than AdamW, reaching by epoch 1000 a lower loss than AdamW attains at the end of training for 5000 epochs. The diagnostic quantities explain this behavior: StableGrad flattens the layer-wise gradient scales, keeps the stability factor ηλmax (KSG ) far below the instability threshold, and satisfies the theorem margin during the phases where the loss decreases most effectively. The few checkpoints where the margin becomes negative coincide with the transient plateau in the loss curve, where the theorem no longer predicts an improved local decrease. Thus, the empirical dynamics match the theoretical picture: StableGrad improves the effective backward dynamics while leaving the forward physical model unchanged. A natural question is whether StableGrad mainly acts as an implicit learning-rate schedule, since rescaling gradients can change the effective step size seen by the optimizer. Appendix B.2 addresses this directly with a control experiment in which AdamW is equipped with a piecewise learning-rate multiplier chosen from the observed spectral ratio λmax (KSG )/λmax (K). This boosted AdamW baseline improves over standard AdamW, but it does not reproduce StableGrad: the residual loss remains higher and the layer-wise update distribution stays much more concentrated. Thus, the effect of StableGrad is not reducible to a global learning-rate increase; it changes the effective geometry of the update.
5
Evaluation
We evaluate StableGrad along two complementary axes. First, we use deep CNNs as a diagnostic setting to test whether optimizer-level gradient normalization can preserve trainability when BatchNorm is removed from architectures that normally rely on it. This setting allows us to observe whether training remains numerically stable without architectural normalization. Second, we evaluate deep PINNs, where BatchNorm is not a suitable stabilizer, and test whether forward-stable initialization combined with StableGrad improves depth scaling and solution accuracy.
5.1
CNNs
We first evaluate StableGrad on image classification models where BatchNorm is a standard component of the architecture. We consider EfficientNetV2-S (Tan and Le, 2021) on CIFAR-100 (Krizhevsky et al., 2009) and ResNet-50 (He et al., 2015a) on ImageNet-1k (Deng et al., 2009). For each model, we compare three variants: the default architecture with BatchNorm (Vanilla), the same architecture with BatchNorm removed (Without BN), and the BatchNorm-free architecture trained with StableGrad. All variants are trained under the same protocol except for the presence of BatchNorm and the use of gradient normalization; Appendix C.1 further analyzes the BatchNorm-free failure mode by tracking how activation scales evolve until training collapse. To check that the effect is not obtained by a trivial scale-removal rule, Appendix C.2 also compares against sign-based gradient preprocessing, which homogenizes gradients more aggressively but fails to train in the same EfficientNet setting. Further implementation and training details are provided in Appendix C.3. Figure 2 shows the diagnostic role of the CNN experiments. In both EfficientNetV2-S and ResNet-50, simply removing BatchNorm makes training collapse almost immediately, confirming that these architectures rely strongly on architectural normalization for stable optimization. StableGrad prevents this collapse without reintroducing any forward normalization. On EfficientNetV2-S, the BatchNormfree model trained with StableGrad converges faster and reaches a higher validation accuracy than the BatchNorm baseline, exceeding 77% compared with roughly 75% for the default model. On ResNet50, StableGrad remains close to the BatchNorm baseline, reaching about 67% validation accuracy while the default BatchNorm model reaches about 71%. The relevant point is not a one-to-one speed comparison, since BatchNorm and StableGrad induce different effective training dynamics, but rather that optimizer-level gradient normalization makes otherwise non-trainable BatchNorm-free CNNs train stably and reach competitive accuracy while leaving the forward architecture unnormalized. 7
(a) EfficientNetV2-S ononCIFAR-100 EfficientNetV2_s CIFAR-100
(b) ResNet-50 ononImageNet-1k ResNet50 ImageNet-1k 77% 75%
71% 67%
70% 60%
60%
Validation Acc. (%)
Validation Acc. (%)
80%
40%
Vanilla StableGrad Without BN
20%
1%
0% 0
20
40
60
80
50% 40% 30% 20%
Vanilla StableGrad Without BN
10%
0%
0%
100
0
Epoch
20
40
60
80
Epoch
Figure 2: Training dynamics of deep CNNs with and without BatchNorm. The default architectures with BatchNorm and the BatchNorm-free variants trained with StableGrad remain trainable, while removing BatchNorm without gradient control leads to training collapse after a few iterations due to forward-pass overflow. Differences between the BatchNorm and StableGrad curves reflect different effective learning-rate dynamics and should not be interpreted as a direct training-speed comparison.
5.2
PINNs
StableGrad is evaluated in the target setting of deep PINNs. The experiments consider three PDE benchmarks commonly used in physics-informed learning: Burgers’ equation with viscosity ν = 10−4 , Poisson’s equation, and a high-frequency k = 10π Helmholtz problem (Raissi et al., 2019; Hao et al., 2024). For each benchmark, fully connected PINNs with depths 6 and 12 are trained. Burgers and Poisson use tanh activations, while Helmholtz uses SiLU activations together with Fourier feature inputs. All PDE constraints are enforced softly through the optimization objective, by penalizing the PDE residual together with the corresponding boundary-condition and initial-condition losses. The comparison is between AdamW (Loshchilov and Hutter, 2019) and AdamW equipped with StableGrad. The AdamW baseline is trained for 50,000 optimization steps. In the StableGrad setting, StableGrad is applied during the first 25,000 steps. Once the layer-wise gradient scales have stabilized, training continues with standard AdamW for an additional 25,000 fine-tuning steps. Further benchmark, architecture, and hyperparameter details are provided in Appendix C.3. Table 1 reports validation metrics for each benchmark: relative L2 error, PDE residual loss, boundarycondition loss, and, for Burgers’ equation, initial-condition loss. All reported PINN values are means over three independent runs with different random seeds. Across all benchmark–depth configurations, AdamW+StableGrad outperforms the corresponding Vanilla AdamW baseline in nearly all reported metrics. This holds for both shallow and deeper PINNs, and for both solution accuracy and physical consistency terms. In some cases the improvement is large, while in others the gains are more modest; nevertheless, the direction is consistent across all reported comparisons. These results indicate that StableGrad improves the matched-depth optimization of PINNs rather than merely shifting error between different components of the objective. The method reduces the validation L2 error while also lowering the PDE, boundary-condition, and, when applicable, initial-condition losses. Thus, the improvements are not obtained by fitting the solution field at the expense of the physical constraints, but by producing solutions that are both more accurate and more consistent with the imposed equations. The depth comparison further shows that simply increasing PINN capacity is not sufficient under standard optimization. Vanilla AdamW does not consistently benefit from moving from depth 6 to depth 12, and in some benchmarks the deeper baseline remains similar to, or worse than, its shallower counterpart. StableGrad mitigates this depth-scaling difficulty: at each evaluated depth, it produces better validation metrics than Vanilla, and the deeper StableGrad models remain trainable and competitive across all three PDE benchmarks. The Poisson case illustrates an important nuance. Since this benchmark is comparatively simple, the depth-12 StableGrad model obtains lower training losses but slightly worse validation metrics than its depth-6 counterpart, suggesting that the additional capacity may lead to overfitting rather than improved generalization. Thus, StableGrad improves the 8
Table 1: Evaluation comparison between vanilla and StableGrad with increasing PINN depth in validation points. Experiment Depth Mode Validation L2 PDE loss IC loss BC loss 6
Vanilla StableGrad
1.3e-1 8.5e-2
1.0e-3 6.3e-4
1.2e-5 7.6e-6
4.6e-8 1.8e-8
12
Vanilla StableGrad
1.6e-1 1.1e-2
6.8e-3 9.0e-4
8.2e-5 8.7e-6
4.7e-8 3.1e-8
6
Vanilla StableGrad
6.5e-5 2.6e-5
8.6e-6 1.8e-6
– –
7.9e-9 1.7e-9
12
Vanilla StableGrad
5.9e-5 3.8e-5
4.0e-6 3.1e-6
– –
7.7e-9 3.5e-9
6
Vanilla StableGrad
4.6e-3 3.1e-3
6.0e-6 6.0e-6
– –
1.9e-8 1.7e-8
12
Vanilla StableGrad
6.5e-3 2.3e-3
1.0e-5 6.6e-6
– –
2.5e-8 1.7e-8
Burgers
Poisson
Helmholtz
optimization of deeper PINNs, but additional depth is not automatically beneficial when the target problem does not require the extra capacity. These results support the main claim of the paper: when architectural normalization is unavailable or undesirable, controlling layer-wise gradient scale before the optimizer step provides a practical mechanism for improving the trainability and accuracy of deep PINNs.
6
Conclusion
Depth turns scale propagation into a bidirectional problem: activations must remain well scaled in the forward pass, while adjoint signals and weight gradients must remain useful in the backward pass. Classical initialization schemes address this tension only at the beginning of training and necessarily make a trade-off between forward- and backward-preserving fan modes. Architectural normalization layers, especially BatchNorm, alleviate this burden by repeatedly resetting intermediate scales, but they also introduce forward-pass dependencies that are undesirable in settings such as PINNs, where the network output and its input derivatives define a physical field and its residual. We introduced StableGrad, an optimizer-level mechanism focused on stabilizing backward scale. It keeps layer-wise gradient magnitudes controlled during training by normalizing weight gradients after backpropagation and before the optimizer update. Since this operation is applied only at the update stage, StableGrad leaves the forward model, its derivatives, and the physical residual unchanged, providing backward scale control without inserting normalization layers into the architecture. It therefore pairs naturally with activation-aware fan-in initializations that preserve forward scale. The experiments support this separation. In BatchNorm-free CNNs, StableGrad prevents the optimization collapse that occurs when BatchNorm is removed from architectures that normally rely on it. In PINNs, where BatchNorm is not an appropriate stabilizer, StableGrad improves matched-depth accuracy across Burgers, Poisson, and Helmholtz benchmarks and mitigates the optimization difficulties observed in deeper models. These results suggest that optimizer-level control of gradient scale can be a practical alternative when architectural normalization is unavailable or undesirable. StableGrad also has important limitations. The method normalizes the weight gradients after the backward pass, but it does not prevent the raw adjoint signals from becoming unstable while they are being propagated through the network. Similarly, although the initialization is designed to start from a forward-stable regime, StableGrad does not by itself guarantee that activation magnitudes will remain stable throughout training. Thus, the method should be viewed as a backward-scale correction mechanism rather than a complete solution to all forward and backward stability problems. Future work should study how to guarantee forward-pass stability during training while preserving the physical interpretation required by PINNs. Another promising direction is to combine layer-wise 9
gradient normalization with automatic learning-rate adaptation, for example with techniques related to Defazio et al. (2024), in order to control the effective step size and reduce dependence on a learning-rate scheduler. Finally, the evaluation should be extended beyond the PINN and CNN settings considered here, including SIREN-style models, implicit neural representations, and other architectures where forward normalization is difficult to use but depth remains essential.
Acknowledgments and Disclosure of Funding TBD
References Ba, J., Kiros, J. R., and Hinton, G. E. (2016). Layer normalization. ArXiv, abs/1607.06450. Bengio, Y., Simard, P., and Frasconi, P. (1994). Learning long-term dependencies with gradient descent is difficult. IEEE Transactions on Neural Networks, 5(2):157–166. Chang, O., Flokas, L., and Lipson, H. (2020). Principled weight initialization for hypernetworks. In International Conference on Learning Representations. Defazio, A., Yang, X. A., Khaled, A., Mishchenko, K., Mehta, H., and Cutkosky, A. (2024). The road less scheduled. In The Thirty-eighth Annual Conference on Neural Information Processing Systems. Deng, J., Dong, W., Socher, R., Li, L.-J., Li, K., and Fei-Fei, L. (2009). Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pages 248–255. IEEE. Fernández-Hernández, A., Mestre, J. I., Dolz, M. F., Duato, J., and Quintana-Orti, E. S. (2026). Sinusoidal initialization, time for a new start. In The Thirty-ninth Annual Conference on Neural Information Processing Systems. Glorot, X. and Bengio, Y. (2010). Understanding the difficulty of training deep feedforward neural networks. volume 9 of Proceedings of Machine Learning Research, pages 249–256. PMLR. Hao, Z., Yao, J., Su, C., Su, H., Wang, Z., Lu, F., Xia, Z., Zhang, Y., Liu, S., Lu, L., and Zhu, J. (2024). Pinnacle: A comprehensive benchmark of physics-informed neural networks for solving pdes. In Globerson, A., Mackey, L., Belgrave, D., Fan, A., Paquet, U., Tomczak, J., and Zhang, C., editors, Advances in Neural Information Processing Systems, volume 37, pages 76721–76774. Curran Associates, Inc. He, K., Zhang, X., Ren, S., and Sun, J. (2015a). Deep residual learning for image recognition. 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pages 770–778. He, K., Zhang, X., Ren, S., and Sun, J. (2015b). Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. In IEEE International Conference on Computer Vision (ICCV). Ioffe, S. and Szegedy, C. (2015). Batch normalization: Accelerating deep network training by reducing internal covariate shift. In Bach, F. and Blei, D., editors, Proceedings of the 32nd International Conference on Machine Learning, volume 37 of Proceedings of Machine Learning Research, pages 448–456, Lille, France. PMLR. Krizhevsky, A., Hinton, G., et al. (2009). Learning multiple layers of features from tiny images. Loshchilov, I. and Hutter, F. (2019). Decoupled weight decay regularization. In International Conference on Learning Representations. Pennington, J., Schoenholz, S. S., and Ganguli, S. (2017). Resurrecting the sigmoid in deep learning through dynamical isometry: theory and practice. In Proceedings of the 31st International Conference on Neural Information Processing Systems, NIPS’17, page 4788–4798, Red Hook, NY, USA. Curran Associates Inc. 10
Raissi, M., Perdikaris, P., and Karniadakis, G. (2019). Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational Physics, 378:686–707. Saxe, A. M., McClelland, J. L., and Ganguli, S. (2013). Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. cite arxiv:1312.6120Comment: Submission to ICLR2014. Revised based on reviewer feedback. Schoenholz, S. S., Gilmer, J., Ganguli, S., and Sohl-Dickstein, J. (2017). Deep information propagation. In International Conference on Learning Representations. Tan, M. and Le, Q. V. (2021). abs/2104.00298.
Efficientnetv2: Smaller models and faster training.
CoRR,
Wang, S., Teng, Y., and Perdikaris, P. (2020). Understanding and mitigating gradient pathologies in physics-informed neural networks. CoRR, abs/2001.04536. Wang, S., Yu, X., and Perdikaris, P. (2022). When and why pinns fail to train: A neural tangent kernel perspective. Journal of Computational Physics, 449:110768. Wu, Y. and He, K. (2018). Group normalization. In Computer Vision – ECCV 2018: 15th European Conference, Munich, Germany, September 8-14, 2018, Proceedings, Part XIII, page 3–19, Berlin, Heidelberg. Springer-Verlag.
A
Additional Theory
This section collects the theoretical details that support the main analysis of StableGrad. Appendix A.1 explains why standard scalar variance-propagation arguments become insufficient for PINNs with differential residuals, while Appendix A.2 discusses alternative global reference scales for the StableGrad normalization. Appendices A.3 and A.4 then derive the StableGrad effective kernel, the associated Rayleigh-quotient identity, and the proof of the local decrease theorem used in the main text. A.1
Signal Propagation with Differential Residuals
This appendix summarizes why the usual variance-preservation argument for standard multilayer perceptrons becomes more involved in PINNs. Consider the feedforward recursion zℓ = Wℓ hℓ and hℓ+1 = ϕ(zℓ ), with independently initialized 2 zero-mean weights of variance σW . Under the standard independence and equal-scale assumptions, the forward variance is approximately 2 Var(hℓ+1 ) ≈ nℓ σW Var(hℓ ).
If δℓ = ∂L/∂hℓ denotes the backpropagated adjoint, then the backward recursion gives 2 Var(δℓ ) ≈ nℓ+1 σW E[ϕ′ (zℓ )2 ] Var(δℓ+1 ).
Thus, in a standard MLP, the forward and backward scales can be approximated by scalar variance recursions. This is precisely the setting in which fan-in, fan-out, and fan-in/fan-out initialization rules are naturally derived. PINNs add a different difficulty. The loss may depend on input derivatives of the network output. To see the effect, consider the first derivative channel pℓ = ∂x hℓ . Differentiating the layer recursion gives pℓ+1 = ϕ′ (zℓ ) ⊙ Wℓ pℓ . Now suppose that the loss depends on both hℓ+1 and pℓ+1 . Define the adjoints aℓ+1 = ∂L/∂hℓ+1 and bℓ+1 = ∂L/∂pℓ+1 . Backpropagation gives bℓ = Wℓ⊤ (ϕ′ (zℓ ) ⊙ bℓ+1 ) , while the activation-channel adjoint satisfies aℓ = Wℓ⊤ (ϕ′ (zℓ ) ⊙ aℓ+1 + ϕ′′ (zℓ ) ⊙ (Wℓ pℓ ) ⊙ bℓ+1 ) . 11
The second term has no analogue in a standard supervised MLP. It couples the derivative-channel adjoint bℓ+1 into the activation-channel adjoint aℓ , and depends on ϕ′′ , on the derivative state pℓ , and on the weights. Higher-order differential residuals introduce higher derivative channels and higher derivatives of the activation. Consequently, the PINN backward pass is not governed by a single scalar variance recursion. A single initialization variance cannot, in general, simultaneously preserve the scale of all derivative-dependent backward channels. This motivates a dynamic correction of gradient scale during training. A.2
Reference Scale and Useful Variants
StableGrad uses the output-adjoint scale, setting αℓ = σout /(σℓ + ε). This choice follows the backward-scale principle: the scale of the signal that initiates the backward pass is transferred to every parameter block. A more general family is obtained by writing αℓ (c) = c/(σℓ + ε), where c > 0 is a global reference scale. Different choices of c preserve different quantities and can be useful for ablations. A norm-preserving reference scale is obtained by imposing ∥e g ∥ = ∥g∥. Since geℓ = cg ℓ /(σℓ + ε), this gives !1/2 PL ℓ 2 ℓ=1 ∥g ∥ cnorm = PL . ℓ 2 2 ℓ=1 ∥g ∥ /(σℓ + ε) An inner-product preserving reference scale is obtained by imposing g ⊤ ge = g ⊤ g. This gives PL ∥g ℓ ∥2 . cip = PL ℓ=1 ℓ 2 ℓ=1 ∥g ∥ /(σℓ + ε) These alternatives separate the effect of layerwise balancing from changes in global scale. The StableGrad choice c = σout follows a different criterion: it anchors every block gradient to the natural scale of the backward signal at the output. A.3
Effective Kernel Induced by StableGrad
We derive the effective kernel used in Section 4 and prove the identity for the residual Rayleigh quotient. Proposition 2 (StableGrad effective kernel). Let L(θ) = 12 ∥r(θ)∥2 , let J = ∂r/∂θ, and partition the parameters as θ = (θ1 , . . . , θL ). Write the Jacobian by blocks as J = [J1 , . . . , JL ]. If StableGrad rescales each block gradient by αℓ , then the linearized residual dynamics are governed by KSG = JP J ⊤ =
L X
αℓ Jℓ Jℓ⊤ ,
ℓ=1
where P = diag(α1 I1 , . . . , αL IL ). Proof. Since L(θ) = 21 ∥r(θ)∥2 , the gradient is g = J ⊤ r. The block gradient is therefore g ℓ = Jℓ⊤ r. StableGrad rescales the gradient as ge = P g, with P = diag(α1 I1 , . . . , αL IL ). Using the local approximation r(θ + ∆θ) ≃ r(θ) + J∆θ, the idealized StableGrad step ∆θ = −ηP J ⊤ r gives r+ ≃ r − ηJP J ⊤ r. Hence the effective kernel is KSG = JP J ⊤ . Finally, because PL J = [J1 , . . . , JL ] and P is block diagonal, JP J ⊤ = ℓ=1 αℓ Jℓ Jℓ⊤ . We now prove the identity used in the main text. Proposition 3 (Residual Rayleigh quotient under StableGrad). Let K = JJ ⊤ and KSG = JP J ⊤ . Define r⊤ Kr r⊤ KSG r ρ= , ρSG = . 2 ∥r∥ ∥r∥2 12
Then
PL ρSG − ρ =
ℓ 2 ℓ=1 (αℓ − 1)∥g ∥ , ∥r∥2
where g ℓ = Jℓ⊤ r. Proof. Using the block decomposition of KSG , we have ⊤
r KSG r =
L X
αℓ r
⊤
Jℓ Jℓ⊤ r =
ℓ=1 ℓ
= Jℓ⊤ r, this is
PL
L X ℓ=1
⊤
ℓ 2
Since g ℓ=1 αℓ ∥g ∥ . Similarly, r Kr = sions and dividing by ∥r∥2 gives the claim. A.4
αℓ ∥Jℓ⊤ r∥2 .
PL
ℓ 2 ℓ=1 ∥g ∥ . Subtracting both expres-
Proof of the Local Decrease Theorem
We prove Theorem 1 stated in Section 4. The result is local: it applies to the linearized residual dynamics around the current parameters. Theorem 1 (Local decrease under StableGrad). Let K = JJ ⊤ and KSG = JP J ⊤ . Consider the + + standard linearized step rstd = r − ηKr and the StableGrad linearized step rSG = r − ηKSG r. Let ρ=
r⊤ Kr , ∥r∥2
ρSG =
r⊤ KSG r . ∥r∥2
If ηλmax (KSG ) < 2 and
η ρSG 1 − λmax (KSG ) > ρ, 2 then the StableGrad linearized step produces a larger decrease of 12 ∥r∥2 than the standard linearized gradient step.
Proof. For any positive semidefinite matrix A, the linearized step r+ = r − ηAr changes the quadratic loss by 1 η2 1 ∥r∥2 − ∥r+ ∥2 = ηr⊤ Ar − r⊤ A2 r. 2 2 2 Applying this identity with A = KSG , the StableGrad decrease is ∆SG = ηr⊤ KSG r −
η2 ⊤ 2 r KSG r. 2
2 r ≤ λmax (KSG )r⊤ KSG r. Therefore, Since KSG is positive semidefinite, r⊤ KSG η ∆SG ≥ ηr⊤ KSG r 1 − λmax (KSG ) . 2
Using r⊤ KSG r = ∥r∥2 ρSG , this becomes η ∆SG ≥ η∥r∥2 ρSG 1 − λmax (KSG ) . 2 For the standard step, the corresponding decrease is ∆std = ηr⊤ Kr −
η2 ⊤ 2 r K r. 2
The second term is nonnegative, so ∆std ≤ ηr⊤ Kr = η∥r∥2 ρ. Hence, the condition η ρSG 1 − λmax (KSG ) > ρ 2 implies ∆SG > ∆std . The additional condition ηλmax (KSG ) < 2 ensures that the stability factor is positive. 13
B
Connecting Theory and Practice
This section connects the local effective-dynamics analysis with controlled empirical diagnostics. Appendix B.1 reports a diagnostic Burgers run in which the quantities appearing in Theorem 1 are measured directly, including the stability factor, theorem margin, residual linearization error, and layer-wise gradient-scale imbalance. Appendix B.2 then tests whether the observed improvement can be explained purely by a global learning-rate increase, using a spectral scheduler control matched to the StableGrad effective-kernel scale. B.1
Controlled Effective-Dynamics Diagnostics
The main experiments in Section 5 evaluate the final performance of the method. Here we use a controlled diagnostic run to test whether the local quantities appearing in Theorem 1 behave in practice as predicted by the analysis in Section 4. The diagnostic problem is a three-dimensional viscous Burgers PINN on the periodic domain [−1, 1]3 , with t ∈ [0, 1] and viscosity ν = 0.05. The reference solution is an exact periodic solution constructed through the Cole–Hopf transformation. The network is a tanh MLP mapping (t, x, y, z) to (u, v, w), using a rescaled time coordinate together with periodic Fourier features in the spatial variables. It has four hidden layers of width 96 and 30,723 trainable parameters. We compare AdamW against AdamW equipped with StableGrad, using the same learning rate η = 10−3 , no weight decay, and a fixed diagnostic batch evaluated every 500 epochs. All diagnostics are computed from an explicit weighted residual vector r(θ) such that 1 ∥r(θ)∥2 . 2 For this Burgers PINN, r(θ) concatenates the weighted PDE, initial-condition, and periodic boundary residuals. In the StableGrad run, each parameter block gradient g ℓ is rescaled with σref , σref = std(r(θ)). αℓ = std(g ℓ ) + ε L(θ) =
Figure 3 shows the train and validation losses for the two optimizers. The practical effect is immediate: StableGrad decreases the loss much faster than AdamW and reaches, by epoch 1000, a validation loss already lower than the final validation loss obtained by AdamW at epoch 5000. The StableGrad trajectory is not strictly monotone, however. After the rapid initial descent, the loss enters a short transient regime around epochs 3000 and 3500, where the validation loss slightly worsens before improving again. This behavior is exactly the regime in which the sufficient condition in Theorem 1 ceases to hold, as shown below. Theorem 1 states that the StableGrad linearized step gives a larger local decrease of 12 ∥r∥2 than the standard linearized gradient step when ηλmax (KSG ) < 2 and
η MSG := ρSG 1 − λmax (KSG ) − ρ > 0. 2 We therefore report the stability factor sSG = ηλmax (KSG ) and the theorem margin MSG . Positive values of MSG indicate that the sufficient condition of Theorem 1 is satisfied. To assess whether the linearized quantities are meaningful for the actual optimizer step, we also measure the relative residual linearization error. Given the realized parameter update ∆θ, we compute Elin =
∥r(θ + ∆θ) − r(θ) − J∆θ∥ . ∥r(θ + ∆θ) − r(θ)∥ + ε
Small values of Elin indicate that the first-order residual model used in the theory is accurate at the scale of the actual update. 14
3D Burgers PINN: AdamW vs StableGrad AdamW
Train residual loss
10
−1
10
−1
10
−2
10
−2
10
−3
10
−3
10
−4
10
−4
Lval
10
0
1000
2000
3000
4000
5000
Validation residual loss
0
0
10
Ltrain
StableGrad
final AdamW / StableGrad = 3.2x 0
1000
2000
Epoch
3000
4000
5000
Epoch
Figure 3: Train and validation losses for AdamW and AdamW+StableGrad on the controlled Burgers diagnostic run. The vertical axis is logarithmic. StableGrad reduces the loss much faster than AdamW and reaches, early in training, validation losses lower than the final loss attained by AdamW at the end of the run. The small non-monotone segment around epochs 3000 and 3500 coincides with the checkpoints where the margin in Theorem 1 becomes negative.
Table 2: Effective-dynamics diagnostics for the StableGrad run on the fixed diagnostic batch. Here sSG = ηλmax (KSG ), MSG is the margin in Theorem 1, and Rstd is the ratio between the largest and smallest layer-wise gradient standard deviations before and after StableGrad rescaling. Epoch 500 1000 2000 2500 3000 3500 4000 5000
Val. loss −4
8.28×10 2.59×10−4 1.42×10−4 1.31×10−4 1.37×10−4 1.57×10−4 1.43×10−4 8.47×10−5
ρ 0.51 2.15 9.49 13.03 15.82 13.87 10.76 14.19
ρSG 3.40 5.89 11.88 13.76 14.99 13.43 11.63 14.50
sSG 0.134 0.097 0.059 0.050 0.044 0.042 0.042 0.034
MSG 2.66 3.45 2.05 0.378 -1.153 -0.722 0.622 0.0646
Rstd raw → scaled
Elin −3
2.63×10 1.32×10−3 1.46×10−3 1.66×10−3 1.77×10−3 2.80×10−3 4.18×10−3 2.88×10−3
42.7 → 1.00 136.0 → 1.00 282.2 → 1.00 316.2 → 1.00 334.3 → 1.00 300.1 → 1.00 227.5 → 1.00 295.0 → 1.00
Table 2 shows that the stability condition is comfortably satisfied throughout the run once the initial transient has passed. The quantity sSG = ηλmax (KSG ) remains between 3.4×10−2 and 1.34×10−1 , far below the threshold value 2. Thus, the StableGrad effective kernel operates well inside the stable regime covered by Theorem 1. The theorem margin is positive for most of the reported checkpoints, including the early and final phases of training. In those regimes, the sufficient condition of Theorem 1 holds: the StableGrad effective kernel acts more strongly on the current residual while remaining stable. The two negative checkpoints, at epochs 3000 and 3500, are particularly informative rather than problematic. They occur exactly when the rapid initial descent has already saturated and the validation loss temporarily stops improving, as seen in Figure 3. This is precisely what one should expect from a local sufficient condition: when the margin becomes negative, the theorem no longer predicts an improved local decrease, and empirically the loss indeed ceases to improve. The condition is then recovered at later checkpoints, and the loss decreases again. The residual linearization error remains small throughout training. After the first checkpoint, Elin stays in the 10−3 range and never exceeds 4.18 × 10−3 . This indicates that the first-order residual model used to define KSG , ρSG , and λmax (KSG ) is accurate at the scale of the actual optimizer update. The diagnostics are therefore not merely formal kernel quantities; they provide a faithful local description of the optimization dynamics observed in the trained PINN. 15
Table 3: Piecewise learning-rate multipliers used for the AdamW scheduler control. Each multiplier is chosen from the observed ratio λmax (KSG )/λmax (K) on the corresponding diagnostic interval. Epoch interval 1–500 501–1000 1001–1500 1501–2000 2001–2500 LR multiplier 5.114545 4.355357 3.303698 2.603073 2.214707 Epoch interval LR multiplier
2501–3000 1.795869
3001–3500 1.487962
3501–4000 1.336904
4001–4500 1.139297
4501–5000 0.950362
Finally, the raw layer-wise gradient scales are highly imbalanced before StableGrad rescaling, with Rstd ranging from 42.7 to 334.3 across the reported checkpoints. StableGrad reduces this ratio to one by construction. This confirms that the method is carrying out the intended correction: it leaves the forward residual unchanged, but equalizes the scale of the layer-wise weight gradients before they are passed to the optimizer. Overall, the controlled diagnostics show a clear agreement between theory and practice. StableGrad operates in a stable regime, satisfies the sufficient condition in Theorem 1 during the phases where the loss decreases most effectively, and temporarily violates it exactly when the empirical trajectory stops improving. The practical effect is substantial: at epoch 5000, AdamW reaches training and validation losses of 2.82 × 10−4 and 2.72 × 10−4 , whereas AdamW+StableGrad reaches 9.14 × 10−5 and 8.47 × 10−5 . More importantly, StableGrad reaches a validation loss below the final AdamW value already around epoch 1000, showing that the improved effective dynamics predicted by the theory translate into a much faster reduction of the PINN training objective in practice. B.2
Learning-Rate Schedule Control
One possible interpretation of StableGrad is that it improves training mainly by inducing a larger effective step size. We therefore include a control experiment designed to test this explanation directly. Besides AdamW and AdamW+StableGrad, we train AdamW with a piecewise learning-rate multiplier chosen to mimic the spectral scale change induced by StableGrad. For each epoch interval, the multiplier is set from the observed ratio m≈
λmax (KSG ) , λmax (K)
K = JJ ⊤ ,
KSG = JPSG J ⊤ .
Thus, if the advantage of StableGrad were only due to a larger effective learning rate, this boosted AdamW control should reproduce its behavior. We do not use the ratio at epoch 1, which is dominated by the initialization transient; the first interval instead uses the ratio measured at epoch 500. Figure 4 compares the resulting train and validation losses. The scheduler control is a strong baseline: it improves clearly over standard AdamW, showing that part of the acceleration can indeed be attributed to increasing the global step scale. However, it does not reproduce StableGrad. StableGrad maintains lower residual losses throughout training: the validation loss of the scheduler control is 6.59×, 7.76×, 4.38×, 2.54×, and 2.08× higher than StableGrad at epochs 500, 1000, 2000, 3000, and 5000, respectively. At the final checkpoint, StableGrad reaches a validation residual loss of 8.473 × 10−5 , compared with 1.762 × 10−4 for the scheduler control. The difference is also visible in the update geometry. The scheduler changes the global step size, but it does not rebalance how the update is distributed across layers. At the final checkpoint, the valid relative-update ratio is 398.7 for AdamW with the scheduler and 77.5 for StableGrad, while the maximum update-energy concentration is 0.782 for the scheduler and 0.350 for StableGrad. Thus, the scheduler update remains much more concentrated in a small number of parameter blocks. StableGrad, in contrast, changes not only the amount of progress made per step, but also the layer-wise geometry of the update. This control shows that StableGrad is not simply a learning-rate scheduler in disguise. Boosting the learning rate using spectral information explains part of the improvement over AdamW, but it does not recover the residual loss or the update distribution obtained by StableGrad. We note that the scheduler control obtains a lower field relative error in this particular run, so this experiment should not be read as a claim that StableGrad dominates every metric. Its purpose is narrower: it shows that the residual-loss gains and layer-wise update dynamics induced by StableGrad cannot be reduced to global learning-rate scaling alone. 16
3D Burgers PINN: StableGrad vs AdamW LR-scheduler control Train residual loss 10
1
10
0
StableGrad
AdamW + lambda scheduler
10
0
10
−1
10
−2
−2
−3
10
−3
10
10
−4
10
−4
−1
10
Lval
Ltrain
10
0
1000
2000
3000
4000
5000
Validation residual loss
final scheduler / StableGrad = 2.08x
0
Epoch
1000
2000
3000
4000
5000
Epoch
Figure 4: Train and validation residual losses for AdamW, AdamW with the spectral learning-rate scheduler, and AdamW+StableGrad. The scheduler improves over AdamW, but does not reproduce the lower residual losses reached by StableGrad.
C
Additional Experimental Details
This section provides additional empirical evidence and reproducibility details for the experiments in the main paper. Appendix C.1 analyzes the early activation-scale instability of BatchNormfree EfficientNetV2-S, and Appendix C.2 compares StableGrad with a more aggressive sign-based gradient preprocessing baseline. Appendix C.3 specifies the evaluation protocols, training setups, benchmark definitions, validation procedures, and hardware requirements used for the reported CNN and PINN experiments. C.1
Evolution of Standard Deviation without BatchNorm
Figure 5 provides a magnified view of the initial training region shown in Figure 2, where the model trained without BatchNorm stops almost immediately after the beginning of training. To better understand this failure, the figure shows the evolution of the activation standard deviation across the convolutional layers of EfficientNetV2-S after removing all BatchNorm layers. The activation standard deviations are plotted in logarithmic scale, and each curve is vertically shifted for readability. Consequently, sharp vertical movements in the plot correspond to large multiplicative changes in activation scale. This makes the sudden spikes and drops especially relevant, since they indicate abrupt changes in the numerical range of the activations. Without BatchNorm, there is no normalization mechanism to recalibrate the scale of the intermediate activations. As training progresses, the scale of the weights can grow, which increases the magnitude of the activations produced by each layer. These activations are then passed to the following layers, so the effect can accumulate throughout the network. As a result, deeper layers tend to exhibit stronger fluctuations, since they are affected by the scale changes introduced by all preceding layers. The figure shows that this instability appears very early in training. After fewer than 500 batches, the activation scale becomes unstable enough to produce a numerical overflow. This overflow propagates as a NaN value, after which the gradients and network weights also become NaNs. At that point, training fails and the curve corresponding to the model without BatchNorm in Figure 2 stops. C.2
Comparison with Sign-Based Gradient Preprocessing
StableGrad modifies the gradients before the optimizer step, but it preserves their internal magnitude structure up to a layer-wise rescaling. A natural alternative is a more aggressive pre-optimizer transformation: replacing each gradient block g ℓ by its element-wise sign, g ℓ ← sign(g ℓ ). 17
Activation std evolution by layer
L0 L10
80
L20 L30
L50 40
L60
Layer idx
60
L40
L70 20
L80 L90 0
20
40
60
80 Batch
100
120
140
0
Figure 5: Magnified view of the activation-scale instability observed at the beginning of training for EfficientNetV2-S without BatchNorm. The activation standard deviation is shown in logarithmic scale across convolutional layers. Each curve corresponds to one Conv2D layer, vertically offset for readability and colored by layer index. This transformation can be viewed as an extreme form of scale homogenization, since it removes gradient-magnitude variation altogether. However, it also discards the relative magnitude information inside each layer, which may be important for the optimizer. We tested this sign-based preprocessing in the same BatchNorm-free EfficientNet setting used in Section 5.1. Under the same training protocol, the sign-gradient variant failed almost immediately: training collapsed by epoch 2 and validation accuracy remained near 1%, i.e. chance level on CIFAR100. In contrast, StableGrad trains the same BatchNorm-free architecture stably without requiring architectural normalization or additional max-norm constraints. This comparison shows that the benefit of StableGrad is not merely due to making all gradient scales similar. A naive sign transformation also removes scale variation, but it destroys too much gradient information to train the model. StableGrad instead equalizes layer-wise gradient scale while preserving the within-layer gradient structure passed to the optimizer. C.3
Evaluation Details
CNN classification experiments. For the CNN classification baselines, we evaluate ResNet-50 on ImageNet-1k and EfficientNetV2-S on CIFAR-100. In both cases, images are processed at an input resolution of 224 × 224, using random resized crops, random horizontal flips, and normalization. Training uses a batch size of 128. ResNet-50 is trained on ImageNet-1k for 70 epochs using SGD with momentum 0.9, learning rate 0.1, and weight decay 2 × 10−5 . EfficientNetV2-S is trained on CIFAR100 for 100 epochs using AdamW with learning rate 10−3 and weight decay 10−3 . Both models are trained with a cosine annealing learning-rate scheduler. For each experiment, the checkpoint with the best validation top-1 accuracy is used for evaluation, and we report top-1 accuracy on the corresponding validation set. Burgers equation. For the Burgers benchmark, we consider the one-dimensional viscous Burgers equation ut + uux = νuxx , on (x, t) ∈ [−1, 1] × [0, 1], with viscosity ν = 10−4 , initial condition u(x, 0) = − sin(πx), and homogeneous Dirichlet boundary conditions u(−1, t) = u(1, t) = 0. The PINN is a fully connected network with width 64, tanh activations, and the depth specified in each experiment. The model is trained with AdamW using initial learning rate 10−3 , zero weight decay, 18
and a cosine annealing learning-rate scheduler. Baseline runs use AdamW for 50,000 optimization steps. StableGrad runs use AdamW with StableGrad for the first 25,000 steps, followed by 25,000 additional fine-tuning steps with standard AdamW. The loss is a weighted sum of the PDE residual, initial-condition, and boundary-condition losses, with weights 1, 10, and 10, respectively. During training, each stochastic batch contains 100,000 PDE collocation points, 2,048 initial-condition points, and 2,048 boundary-condition points. Evaluation is performed on independent validation samples and using the relative L2 error with respect to a high-resolution numerical reference solution. The reference solution is generated with a method-of-lines solver on a uniform grid of 4096 spatial points and 401 time snapshots. Poisson equation. For the Poisson benchmark, we use fully connected PINNs with width 64, tanh activations, and the depth specified in each experiment. The model is trained with AdamW using initial learning rate 10−3 , zero weight decay, and a cosine annealing learning-rate scheduler. Baseline runs use AdamW for 50,000 optimization steps, while StableGrad runs use AdamW with StableGrad for the first 25,000 steps and standard AdamW for the remaining 25,000 fine-tuning steps. The loss consists of an interior PDE residual term and a softly enforced Dirichlet boundary-condition term, weighted by λPDE = 1 and λBC = 100, respectively. During training, each stochastic batch contains 16,384 interior residual points and 4,096 boundary points. For validation, the solution is evaluated on a 256 × 256 grid, and the PDE and boundary losses are computed using 65,536 interior residual points and 16,384 boundary points. Helmholtz equation. For the Helmholtz benchmark, we consider the three-dimensional problem with wave number k = 10π and exact solution u(x, y, z) = sin(mπx) sin(mπy) sin(mπz), with mode m = 10. The Helmholtz PINN uses Fourier feature inputs and SiLU activations. The input coordinates (x, y, z) are augmented with sinusoidal features sin(πf xi ) and cos(πf xi ), for each coordinate xi ∈ {x, y, z} and frequencies f = 1, . . . , 12, while also retaining the original coordinates. The resulting features are passed to a fully connected network with width 64 and the depth specified in each experiment. The model is trained with AdamW using learning rate 10−4 , zero weight decay, and a warm-up period of 1,000 steps. Baseline runs use AdamW for 50,000 optimization steps, while StableGrad runs use AdamW with StableGrad for the first 25,000 steps and standard AdamW for the remaining 25,000 fine-tuning steps. The PDE residual is normalized by k 2 , which stabilizes optimization at high wave numbers. The loss consists of the normalized PDE residual and a softly enforced Dirichlet boundary-condition term, weighted by λPDE = 1 and λBC = 100, respectively. During training, each stochastic batch contains 32,768 interior residual points and 8,192 boundary points. For validation, the PDE and boundary losses are computed using 65,536 interior residual points and 16,384 boundary points, and the relative L2 error is computed on a uniform grid of size 1283 , evaluated in chunks of 65,536 points. Unless otherwise stated, all PINN constraints are imposed softly through penalty terms in the training objective. The numerical values reported in the tables are computed on validation data, whereas the optimization curves shown in the plots use the corresponding training losses. Code availability. The code used to run all experiments reported in this work, together with a reusable implementation of StableGrad for other training pipelines, is available at github.com/anonymized/stablegrad. Hardware requirements. The experiments require a CUDA-compatible NVIDIA GPU supported by recent PyTorch releases. To provide conservative runtime estimates, each individual PINN training run can be reproduced on a single NVIDIA H100, or an equivalent accelerator, in under one hour. For the CNN experiments, each EfficientNetV2-S/CIFAR-100 training run can be reproduced in under four hours, while each full ResNet-50/ImageNet-1k training run can be reproduced in under 24 hours. These times are conservative upper bounds rather than the minimum required runtime. Full-length CNN training is only necessary to reproduce the final reported accuracies; shorter runs are sufficient to verify that the implementation trains correctly without BatchNorm.
19