Stabilized Gradient Residual Loss

• SGR Loss is a loss construction that interpolates between the MSE gradient and a residual-based gradient while keeping the loss value identical to MSE.
• It addresses the loss conditioning bottleneck by modulating the optimization landscape between the normal equations and the original PDE operator, thereby enhancing convergence rates.
• The method is effective in both ODIL and PINNs settings, with a tunable parameter β that balances convergence speed and stability across varied PDE examples.

Stabilized Gradient Residual (SGR) loss is a loss construction for optimization-based partial differential equation (PDE) solvers in which the scalar loss value remains identical to the mean squared error (MSE), while the gradient is replaced by an interpolation between the MSE gradient and a residual-based gradient. It was introduced for both Optimizing a Discrete Loss (ODIL) and Physics-Informed Neural Networks (PINNs) to address the “loss conditioning bottleneck”: when a PDE residual is optimized with MSE, the resulting optimization landscape is governed by the normal equations and inherits the squared condition number of the original operator, which can make convergence much slower than that of classical iterative solvers (Cao et al., 24 Jul 2025).

1. Problem formulation and motivation

After discretizing a PDE

$\mathcal{N}[q](x)=f(x),$

a linearization around an iterate yields an algebraic residual system

$$f(q)=Aq-b=0,\qquad A\in\mathbb{R}^{n\times n},\; q\in\mathbb{R}^n,\; b\in\mathbb{R}^n.$$

In ODIL, the optimization variables are the discrete solution values $q$ themselves. In PINNs, the unknown is represented by a neural network $q_\theta(x)$, and optimization proceeds over network parameters $\theta$ while the loss is built from PDE residuals together with boundary-condition, initial-condition, or data terms. In both settings, the solver is ultimately driven by minimization of a scalar loss derived from the residual (Cao et al., 24 Jul 2025).

For the linear system, the standard MSE loss is

$L_{\text{MSE}}(q)=\|Aq-b\|_2^2.$

Its gradient and Hessian are

$$\nabla L_{\text{MSE}}(q)=2A^T(Aq-b),\qquad \nabla^2 L_{\text{MSE}}(q)=2A^TA.$$

The conditioning of the optimization problem is therefore governed by $A^TA$, not by $A$. Since

$\kappa(A^TA)=\kappa(A)^2,$

optimization over the MSE loss reproduces the conditioning degradation associated with the normal equations. This is the central bottleneck identified by the SGR formulation: the loss chosen for optimization can be substantially worse conditioned than the original discrete PDE operator.

For symmetric positive definite (SPD) systems, the paper contrasts MSE with the quadratic programming loss

$$f(q)=Aq-b=0,\qquad A\in\mathbb{R}^{n\times n},\; q\in\mathbb{R}^n,\; b\in\mathbb{R}^n.$$0

whose gradient is $$f(q)=Aq-b=0,\qquad A\in\mathbb{R}^{n\times n},\; q\in\mathbb{R}^n,\; b\in\mathbb{R}^n.$$1 and whose Hessian is $$f(q)=Aq-b=0,\qquad A\in\mathbb{R}^{n\times n},\; q\in\mathbb{R}^n,\; b\in\mathbb{R}^n.$$2. In that SPD case, optimization over $$f(q)=Aq-b=0,\qquad A\in\mathbb{R}^{n\times n},\; q\in\mathbb{R}^n,\; b\in\mathbb{R}^n.$$3 preserves $$f(q)=Aq-b=0,\qquad A\in\mathbb{R}^{n\times n},\; q\in\mathbb{R}^n,\; b\in\mathbb{R}^n.$$4 rather than $$f(q)=Aq-b=0,\qquad A\in\mathbb{R}^{n\times n},\; q\in\mathbb{R}^n,\; b\in\mathbb{R}^n.$$5. The difficulty is that this construction does not extend safely to general nonsymmetric systems, which motivates SGR.

2. Formal definition of GR and SGR

The immediate precursor to SGR in the paper is the Gradient Residual (GR) construction. For a residual $$f(q)=Aq-b=0,\qquad A\in\mathbb{R}^{n\times n},\; q\in\mathbb{R}^n,\; b\in\mathbb{R}^n.$$6, GR is defined as

$$f(q)=Aq-b=0,\qquad A\in\mathbb{R}^{n\times n},\; q\in\mathbb{R}^n,\; b\in\mathbb{R}^n.$$7

Here “detach” means that the quantity is used in the forward pass but treated as constant in backpropagation, so its derivative is suppressed. Under this autograd semantics,

$$f(q)=Aq-b=0,\qquad A\in\mathbb{R}^{n\times n},\; q\in\mathbb{R}^n,\; b\in\mathbb{R}^n.$$8

The resulting update is therefore

$$f(q)=Aq-b=0,\qquad A\in\mathbb{R}^{n\times n},\; q\in\mathbb{R}^n,\; b\in\mathbb{R}^n.$$9

which is a residual-based iteration that avoids the appearance of $q$0.

SGR introduces a convex combination of the GR and MSE gradient fields. Its defining expression is

$q$1

The scalar value of this expression is identical to that of MSE:

$q$2

Its gradient, however, is

$q$3

Up to a constant factor, the effective gradient field can be written as

$q$4

This yields the interpolation structure that defines the method. At $q$5, SGR reduces to pure MSE. At $q$6, it reduces to the residual-based GR update. For intermediate $q$7, it combines the robustness associated with the $q$8 term and the favorable conditioning associated with direct residual iteration.

A common misconception is that SGR changes the objective value being minimized. In the reported construction, it does not: the numerical loss value is the same as MSE. What changes is the backward pass, and therefore the optimization dynamics.

3. Conditioning, effective operator, and stabilization

The SGR gradient admits a preconditioned-residual interpretation

$q$9

The associated effective Hessian is

$q_\theta(x)$0

This formula makes the interpolation precise. The endpoint cases are

$q_\theta(x)$1

Thus SGR modulates the optimization landscape between the normal equations and the original operator (Cao et al., 24 Jul 2025).

For SPD systems, GR is equivalent to the QP formulation in the sense that both yield gradient $q_\theta(x)$2 and preserve the original conditioning $q_\theta(x)$3. MSE, by contrast, produces Hessian $q_\theta(x)$4 and condition number $q_\theta(x)$5. This distinction explains why, in the SPD setting, GR- or QP-based optimization can match the convergence characteristics of classical solvers on $q_\theta(x)$6, whereas MSE-based optimization behaves like a method applied to $q_\theta(x)$7.

For nonsymmetric systems, direct residual iteration is not automatically stable. The paper emphasizes that $q_\theta(x)$8 is not, in general, the gradient field of a scalar potential when $q_\theta(x)$9 is nonsymmetric. GR therefore relies on autograd semantics rather than on a classical scalar energy. SGR stabilizes this situation by retaining an $\theta$0 component. Spectrally, $\theta$1 contributes symmetric positive-definite damping, while the $\theta$2 term restores conditioning closer to that of the original operator.

The paper also gives an explicit convergence condition for the GR iteration in the positive-definite case, where all eigenvalues of $\theta$3 have positive real parts:

$\theta$4

This is described as a CFL-like constraint. It guarantees convergence for sufficiently small step size but does not preclude oscillations. In this sense, SGR with $\theta$5 can be interpreted as a generalized explicit scheme with built-in stabilization from the $\theta$6 component.

4. Deployment in ODIL and PINNs

In ODIL, the unknown vector $\theta$7 is optimized directly. The paper examines a linear SPD Poisson problem, a nonlinear Allen–Cahn equation, and a Navier–Stokes lid-driven cavity problem. In the SPD Poisson case, ODIL with QP or GR converges much faster than ODIL with MSE. In the nonlinear nonsymmetric examples, SGR rather than pure GR is the crucial construction (Cao et al., 24 Jul 2025).

For the Allen–Cahn equation,

$\theta$8

the discretization uses central differences in space and forward Euler in time on a $\theta$9 grid. ODIL treats the full space-time field as the optimization variable. The reported behavior is systematic: $L_{\text{MSE}}(q)=\|Aq-b\|_2^2.$0 is very slow; increasing $L_{\text{MSE}}(q)=\|Aq-b\|_2^2.$1 accelerates convergence; with Adam and $L_{\text{MSE}}(q)=\|Aq-b\|_2^2.$2 around $L_{\text{MSE}}(q)=\|Aq-b\|_2^2.$3–$L_{\text{MSE}}(q)=\|Aq-b\|_2^2.$4, the solver reaches machine-precision-level residuals in orders of magnitude fewer iterations than MSE. By contrast, LBFGS diverges for large $L_{\text{MSE}}(q)=\|Aq-b\|_2^2.$5 because the effective curvature is nonsymmetric.

For the lid-driven cavity Navier–Stokes problem,

$L_{\text{MSE}}(q)=\|Aq-b\|_2^2.$6

with moving lid boundary $L_{\text{MSE}}(q)=\|Aq-b\|_2^2.$7 on the top and no-slip elsewhere, the reported discretization is a $L_{\text{MSE}}(q)=\|Aq-b\|_2^2.$8 grid with central differences. At $L_{\text{MSE}}(q)=\|Aq-b\|_2^2.$9 and $$\nabla L_{\text{MSE}}(q)=2A^T(Aq-b),\qquad \nabla^2 L_{\text{MSE}}(q)=2A^TA.$$0, ODIL with MSE is described as extremely slow and may not reach acceptable error, whereas ODIL with SGR and $$\nabla L_{\text{MSE}}(q)=2A^T(Aq-b),\qquad \nabla^2 L_{\text{MSE}}(q)=2A^TA.$$1–$$\nabla L_{\text{MSE}}(q)=2A^T(Aq-b),\qquad \nabla^2 L_{\text{MSE}}(q)=2A^TA.$$2 converges much faster.

The PINNs setting is treated in two forms. In PINNs-ND, derivatives are approximated numerically on a grid, exposing an explicit residual operator much like ODIL. In PINNs-AD, derivatives are produced by automatic differentiation and there is no explicit matrix $$\nabla L_{\text{MSE}}(q)=2A^T(Aq-b),\qquad \nabla^2 L_{\text{MSE}}(q)=2A^TA.$$3, but the same operator-level issue appears implicitly. The residual vector for PINNs-AD is assembled as

$$\nabla L_{\text{MSE}}(q)=2A^T(Aq-b),\qquad \nabla^2 L_{\text{MSE}}(q)=2A^TA.$$4

where $$\nabla L_{\text{MSE}}(q)=2A^T(Aq-b),\qquad \nabla^2 L_{\text{MSE}}(q)=2A^TA.$$5 is the PDE residual, $$\nabla L_{\text{MSE}}(q)=2A^T(Aq-b),\qquad \nabla^2 L_{\text{MSE}}(q)=2A^TA.$$6 the boundary-condition residual, and $$\nabla L_{\text{MSE}}(q)=2A^T(Aq-b),\qquad \nabla^2 L_{\text{MSE}}(q)=2A^TA.$$7 the initial-condition residual; example scaling coefficients are $$\nabla L_{\text{MSE}}(q)=2A^T(Aq-b),\qquad \nabla^2 L_{\text{MSE}}(q)=2A^TA.$$8, $$\nabla L_{\text{MSE}}(q)=2A^T(Aq-b),\qquad \nabla^2 L_{\text{MSE}}(q)=2A^TA.$$9, and $A^TA$0.

In Poisson examples, PINNs-ND reproduces the same GR/QP-versus-MSE distinction observed in ODIL, whereas in PINNs-AD the difference is smaller for that relatively mild problem. In Allen–Cahn, both PINNs-ND and PINNs-AD show consistent improvement of SGR over MSE, with convergence speed increasing as $A^TA$1 increases, at the cost of stronger oscillations. In Navier–Stokes at $A^TA$2, PINNs-ND with SGR reaches relative error $A^TA$3, close to the best attainable on that grid, reported as $A^TA$4, while MSE remains stuck at large errors. PINNs-AD attains larger final errors overall, but SGR still improves substantially over MSE.

The implementation detail emphasized in all cases is that SGR is realized with standard backpropagation; the essential requirement is correct use of .detach() in constructing the loss expression.

5. Optimization behavior, empirical performance, and tuning

The main empirical claim of the SGR paper is that, within ODIL, SGR achieves orders-of-magnitude faster convergence than MSE, and that it also consistently outperforms MSE in PINNs despite the additional nonlinearity of neural-network parameterizations (Cao et al., 24 Jul 2025).

The practical parameter is $A^TA$5. The paper describes $A^TA$6 as the slowest regime because it is exactly MSE. The interval $A^TA$7 is reported as the best practical range. Higher $A^TA$8 moves the effective operator closer to $A^TA$9 and therefore improves conditioning, but it also increases oscillatory behavior. This trade-off is optimizer-dependent. Adam is reported to handle high $A$0 robustly, especially when combined with learning-rate decay. LBFGS benefits from second-order curvature information under MSE, but for high $A$1 its symmetric quasi-Hessian is mismatched to the nonsymmetric effective curvature and divergence occurs in nonsymmetric problems such as Allen–Cahn and Navier–Stokes.

The learning-rate strategy is correspondingly asymmetric. With SGR and Adam, the paper recommends a relatively large initial learning rate, for example $A$2 or $A$3, followed by multiplicative decay every few hundred or thousand iterations with a floor. This combination is reported to suppress oscillations while preserving fast initial progress. Increasing the learning rate under MSE yields little benefit because the normal-equation conditioning remains dominant.

The paper also compares SGR with explicit time marching (ETM). For Poisson on a $A$4 grid and Navier–Stokes on a $A$5 grid, ETM uses maximal stable time steps such as $A$6 for Poisson and $A$7 for Navier–Stokes, whereas SGR with large $A$8 and learning-rate decay reaches the same residual level in $A$9–$\kappa(A^TA)=\kappa(A)^2,$0 orders of magnitude fewer iterations. The cost per iteration remains close to that of ETM: for Poisson, ETM uses $\kappa(A^TA)=\kappa(A)^2,$1 GB and $\kappa(A^TA)=\kappa(A)^2,$2 s per $\kappa(A^TA)=\kappa(A)^2,$3 iterations versus SGR at $\kappa(A^TA)=\kappa(A)^2,$4 GB and $\kappa(A^TA)=\kappa(A)^2,$5 s; for Navier–Stokes, ETM uses $\kappa(A^TA)=\kappa(A)^2,$6 GB and $\kappa(A^TA)=\kappa(A)^2,$7 s versus SGR at $\kappa(A^TA)=\kappa(A)^2,$8 GB and $\kappa(A^TA)=\kappa(A)^2,$9 s. The reported interpretation is that SGR attains an implicit-like efficiency gain at almost explicit-like per-iteration cost.

6. Interpretation, limitations, and relation to other gradient-dependent losses

SGR is best understood as a prescribed gradient field or iterative scheme rather than as a classical scalar potential for nonsymmetric systems. That distinction is essential. The forward-pass value equals MSE, but the backward pass is deliberately modified through detach operations so that optimization follows a different vector field (Cao et al., 24 Jul 2025).

This construction has several consequences. First, SGR does not require special optimizer modifications. Second, in autograd-based frameworks the $$f(q)=Aq-b=0,\qquad A\in\mathbb{R}^{n\times n},\; q\in\mathbb{R}^n,\; b\in\mathbb{R}^n.$$00 contribution already arises through backpropagation, so the method does not add computational cost beyond MSE; the detach operations themselves are described as insignificant. Third, the method is particularly attractive in settings where residuals and their gradients are readily available but implicit solvers are expensive or cumbersome to implement.

The limitations are equally explicit. Very large $$f(q)=Aq-b=0,\qquad A\in\mathbb{R}^{n\times n},\; q\in\mathbb{R}^n,\; b\in\mathbb{R}^n.$$01 destabilizes second-order methods such as LBFGS on nonsymmetric problems. In PINNs-AD, the nonlinearity of deep networks introduces additional instabilities, so very high $$f(q)=Aq-b=0,\qquad A\in\mathbb{R}^{n\times n},\; q\in\mathbb{R}^n,\; b\in\mathbb{R}^n.$$02 can produce strong oscillations unless it is tuned carefully. More broadly, SGR improves conditioning relative to MSE, but it does not eliminate the need for learning-rate control.

A distinct but related gradient-dependent loss family is provided by Lai loss, introduced in 2024 as a task loss multiplied by a gradient-dependent factor derived from a geometric construction (Lai, 2024). Lai loss penalizes the gradients with the loss itself and uses bounded angle-based factors together with a hyperparameter $$f(q)=Aq-b=0,\qquad A\in\mathbb{R}^{n\times n},\; q\in\mathbb{R}^n,\; b\in\mathbb{R}^n.$$03 to bias local input gradients toward a preferred range; the paper presents this as a means to control smoothness, sensitivity, and the trade-off between accuracy and robustness. SGR addresses a different problem: not input-gradient smoothness in regression, but the conditioning of residual-driven PDE optimization. The comparison nevertheless indicates a broader design pattern in which the optimization behavior of a model is altered not by adding a conventional regularization term, but by re-engineering how gradient information enters the loss.