A pullback-corrected scalar auxiliary variable optimizer with momentum and adaptive mobility
Authors: Jiahao Zhang, Shiheng Zhang, Guang Lin
Organizations: School of Mechanical Engineering, Purdue University, West Lafayette, IN 47906, USA · Department of Mathematics, University of Washington, Seattle, WA 98195, USA · Department of Mathematics, Purdue University, West Lafayette, IN 47906, USA
Objectives in scientific machine learning are often prescribed as a sum of several terms, such as the residual, boundary, initial, and data losses of a physics-informed neural network. In the pullback-corrected scalar auxiliary variable (PB--SAV) method, one scalar tracks the shifted objective while the component gradients build a positive semidefinite curvature correction of rank at most the number of components. We carry that correction into an optimizer with momentum and an adaptive mobility, applying it to the gradient and the stored momentum in a single implicit solve. A mobility that is nonincreasing in the Loewner order yields an exact modified energy law, covering Euclidean and AMSGrad-type choices; the corresponding identity for momentum appended after the solve carries a cross term of indefinite sign. For a fixed mobility we give a necessary and sufficient condition for local stability at a stationary point, depending on the Hessian minus twice the correction, and show that it also gives local geometric convergence for every scalar relaxation sequence. The implicit solve reduces to a dense system whose order is the number of components. In the forward Burgers comparison, four components reduce the mean tail objective by 64.7% and the final solution error by 50.2% relative to one component at the same learning rate and momentum settings.
Figures & tables
Figure 1: Convergence on the quadratic problem during the first 100 updates. All methods start from the same state.
Figure 2: Fixed-state geometry for the quadratic problem. Left: relative error between the PB–SAV correction and the exact Hessian. Right: relative error between the PB–SAV increment and the damped Newton increment.
Method
Hessian error
Newton-step error
Updates to 10−10
Final gap
Gradient descent (optimal step)
–
–
–
3.102×10−4
Relaxed SAV (1 comp.)
9.900×10−1
1.627×10−1
–
1.873×10−10
Relaxed PB–SAV (2 comp.)
9.899×10−1
3.022×10−12
63
7.879×10−11
Relaxed PB–SAV (100 comp.)
1.414×10−12
3.022×10−12
63
9.623×10−11
Exact damped Newton
–
0
62
1.550×10−48
Table 1: Quadratic comparison after 300 updates. The Hessian and step errors are evaluated at the common initial state. The final gap is F(ϕ300)−F(ϕ∗) , with F(ϕ∗)=0 . A dash indicates that the quantity is not applicable or that the target was not reached within 300 updates. The columns report different quantities and are not ranked, so no entry is bold.
Figure 3: Training objective for one, eight, and 30 PB–SAV components, together with gradient descent and damped Gauss–Newton. The pointwise PB–SAV and Gauss–Newton trajectories are indistinguishable at plotting resolution.
Figure 4: Fixed-state geometry for nonlinear regression. Left: relative error between the PB–SAV correction and the Gauss–Newton matrix. Right: the largest relative step error over η=0.1,1,10,100 . Both errors collapse when each observation is represented by a separate component.
Method
Updates
Final objective
Best noise-free test MSE
Final ∥∇F∥
PB–SAV (30 pointwise comp.)
600
2.525×10−3
2.804×10−3
1.032×10−4
Damped Gauss–Newton
600
2.525×10−3
2.804×10−3
1.032×10−4
Gradient descent (rate 0.22)
600
3.301×10−2
6.040×10−2
2.468×10−2
Adam
600
2.985×10−3
3.113×10−3
6.695×10−4
Table 2: Nonlinear-regression results after 600 accepted updates. No entry is bold: PB–SAV and damped Gauss–Newton agree in every column to the displayed precision.
Figure 5: Five-seed Darcy DeepONet comparison. (a) Mean common training data loss and (b) mean held-out relative L2 error; shaded regions show one sample standard deviation. The legend entry “SGD + momentum” in (a) and (b) is the full-batch heavy-ball gradient descent of Table 3 . The lower row shows a fixed test sample selected by the maximum variation of the reference field: (c) permeability, (d) reference solution, (e) flux magnitude, (f) AdamW absolute error, and (g) PB–SAV absolute error.
Method
Components
η
Train data loss
Test relative L2
Wall time (s)
Heavy-ball gradient descent
–
10−2
0.01771±0.00041
0.15510±0.00095
37.4
AdamW
–
10−3
0.01125±0.00024
0.11264±0.01001
36.5
PB–SAV (data + weight decay)
2
3×10−3
0.00399±0.00011
0.08512±0.00225
74.5
Table 3: Darcy DeepONet results after 5,000 updates over five paired seeds. Losses and errors are reported as mean ± sample standard deviation. Wall time is the mean recorded time for the full training run. Best mean training and test values are bold; wall time is reported for reference and is not ranked.
Figure 6: Five-seed comparison for the forward Burgers PINN. (a) Mean common task objective, recorded at every update, and (b) mean relative L2 solution error. Shaded regions show one sample standard deviation.
Method
η
Tail task objective
Tail objective CV
Heavy-ball GD
0.01
(5.095±0.660)×10−4
130.2%±22.0%
AdamW
0.001
(3.181±0.450)×10−4
161.5%±22.0%
PB–SAV (1 component)
0.001
(1.741±0.430)×10−4
58.9%±31.0%
PB–SAV (4 components)
0.001
(6.154±0.840)×10−5
5.4%±0.3%
Table 4: Forward Burgers results after 10,000 updates over five paired seeds. For each seed, the tail mean and coefficient of variation (CV) are computed over the final 2,000 updates; the table reports their mean ± sample standard deviation across seeds. Here CV=100stail/Fˉtail . Best mean values are bold.
Figure 7: Five-seed comparison for inverse Burgers identification. (a) Mean common task objective and (b) mean relative viscosity error; shaded regions show one sample standard deviation. The lower row uses seed 42, the median four-component PB–SAV seed according to final viscosity error: (c) reference solution, (d) four-component PB–SAV reconstruction, and (e) absolute error with the 100 observation locations.
Method
η
Tail task objective
Heavy-ball GD
0.01
(0.878±1.900)×10−1
AdamW
0.001
(3.515±0.330)×10−4
PB–SAV (1 component)
0.003
(9.507±2.000)×10−5
PB–SAV (4 components)
0.01
(4.816±2.000)×10−5
Table 5: Inverse Burgers results after 10,000 updates over five paired seeds. Values are mean ± sample standard deviation. For each seed, the tail task objective is averaged over the final 2,000 updates of Finv . Best mean values are bold.
Department of Mathematics, Purdue University, West Lafayette, IN 47907, USA · Department of Mathematics & School of Mechanical Engineering, Purdue University, West Lafayette, IN 47907, USA
Department of Pure and Applied Mathematics, Graduate School of Information Science and Technology, The University of Osaka · D3 Center, The University of Osaka