Simon Dufort-Labbé Mila, Université de Montréal
Mehrab Hamidi Mila, Université de Montréal
Razvan Pascanu Mila, Université de Montréal
Ioannis Mitliagkas Mila, Université de Montréal
Damien Scieur Samsung – SAIL Montreal Mila, Université de Montréal
Aristide Baratin Samsung – SAIL Montreal Mila, Université de Montréal
Abstract Sharpness-aware minimization (SAM) encourages flat minima by perturbing parameters along directions of high loss curvature, but treats all parameter directions uniformly, ignoring the underlying loss geometry. We introduce LLQR+SAM, which combines SAM with a learned preconditioner obtained from the recently proposed LLQR framework–a second-order method that recasts steepest descent as a layerwise linear-quadratic regulator problem. The preconditioner is updated sparsely and maintained as a slow exponential moving average, so it captures a smoothed, low-resolution picture of the loss landscape geometry. The SAM perturbation then operates on top of this learned geometry, probing curvature at a faster timescale. We show that this two-timescale structure is not merely a computational convenience: theoretically, the preconditioner amplifies the SAM escape signal in directions that are flat under the average geometry but locally sharp (“potholes”). Wide, flat basins, by contrast, remain stable. Empirically, LLQR+SAM gives consistent gains over both SAM and LLQR alone across standard vision and sequence modeling benchmarks, supporting the view that slow learned geometry and fast sharpness correction are genuinely complementary. Preprint. 80
5
70
23
60
22
50
475
480
485
490
495
500
40
Test Loss
4
Top-1 Error
arXiv:2605.16134v1 [cs.LG] 15 May 2026
Navigating Potholes with Geometry-Aware Sharpness Minimization
3
2
30 1
20
0
100
200
300
ImageNet Steps (×103) Baseline--SGDM
400
500
With FSAM
0
With LLQR
15
30
45
Elapsed Time (hours)
60
With LLQR+FSAM
Figure 1: Interaction between FSAM and LLQR on ResNet-50/ImageNet. Left: Top-1 error for SGDM and FSAM, a SAM variant, with and without LLQR. Although both SAM-style methods and LLQR can be interpreted as curvature-correcting mechanisms, their combination yields gains over either component alone, suggesting complementary effects. Right: Test loss versus elapsed training time. Despite the usual concern that second-order methods are prohibitively expensive, LLQR remains efficient enough not to become the bottleneck: whenever the slowdown of SAM is acceptable, topping it with LLQR likely preserves affordability.
1
Introduction
Modern deep-learning loss landscapes are widely thought to combine smooth, slowly-varying structure at large scales—wide basins, valleys, and ridges shaped by the architecture and data—with much sharper features at smaller scales: narrow valleys, spurious minima, and localized potholes that act as attractors of poor generalization. Two largely independent lines of optimization research target these two scales. Geometry-aware methods such as natural gradient [Amari, 1998], Newton, K-FAC [Martens and Grosse, 2015] and Shampoo [Gupta et al., 2018] use estimates of curvature to reshape the descent direction, addressing the ill-conditioning of the average loss geometry. Sharpness-aware methods, beginning with SAM [Foret et al., 2021], perturb the parameters before evaluating the gradient, biasing optimization toward flatter minima. These two mechanisms operate on different aspects of the geometry, and at very different time scales. A learned preconditioner accumulated over many gradient steps can only capture the slow average curvature of the landscape; it is by construction insensitive to sharp localized features. A SAM perturbation, evaluated on a single minibatch, probes local sharpness directly—but treats all parameter directions uniformly, ignoring the underlying optimization geometry. Each captures something the other misses. In this work, we study the combination of the two. We use LLQR [Dufort-Labbé et al., 2026], a recently proposed framework that learns a structured inverse preconditioner U by minimizing a linear-quadratic-regulator (LQR) objective derived from the layerwise network dynamics and a chosen divergence on the loss. The preconditioner is updated sparsely (a few times per epoch) and maintained as an exponential moving average; it captures a smoothed, low-resolution picture of the average loss geometry. On top of this slow geometry we apply a SAM perturbation evaluated on each minibatch in the U -induced norm, and the resulting outer update uses the same learned geometry. The combined algorithm, LLQR+SAM, applies the SAM displacement and the gradient transport in a single learned metric. Our central claim is that this two-time-scale combination is synergistic, not merely additive. We make this precise on a model landscape consisting of a quadratic loss with a smooth average geometry, supplemented by a sharp localized pothole that creates a spurious minimum invisible to U . On this landscape the LLQR+SAM dynamics is closed-form and exposes a clean mechanism: SAM prevents the iterate from settling at any minimum, forcing it to hover at the scale of the SAM probe itself, while U shapes that scale according to the average geometry. Around a wide basin the hovering is appropriate to the basin width and the iterate is stable; around a pothole the same hovering scale exceeds the pothole’s basin of attraction and the iterate escapes. Crucially, the escape signal is amplified relative to vanilla SAM by a factor that grows as the surrounding basin becomes flatter— precisely the regime where preconditioning matters most. The same U that enables fast convergence in flat directions also enables pothole escape: U and SAM are not independent ingredients but two complementary uses of one learned geometry. These predictions are borne out by the toy experiments in Figures 2 and 3 and by the empirical study of Section 5, which shows consistent gains over either component alone across CIFAR, TinyImageNet, ImageNet, and IWSLT14. Contributions. (i) We propose LLQR+SAM, combining a slowly-updated, EMA-smoothed second-order preconditioner with a SAM perturbation evaluated in the induced geometry. (ii) On a quadratic landscape with a sharp pothole perturbation we derive the exact LLQR+SAM dynamics and quantify the escape-amplification of the preconditioner relative to vanilla SAM, with the gain growing as the surrounding basin becomes flatter. (iii) We validate empirically that the combination yields consistent gains over either component alone across standard vision and sequence benchmarks.
2
Related Work
Sharpness, flatness, and SAM. The relation between generalization and loss-landscape sharpness has long motivated methods that bias training toward flatter minima [Hochreiter and Schmidhuber, 1997, Keskar et al., 2017, Neyshabur et al., 2015], although this connection is neither universal nor invariant: sharpness can change under reparameterization without changing the represented function [Dinh et al., 2017]. SAM [Foret et al., 2021] is the canonical sharpness-aware optimizer, replacing the training objective by a local min–max problem that probes the worst-case loss near the cur2
rent parameters and descends using the gradient at the perturbed point. Many variants modify either the perturbation neighborhood or the adversarial-point estimate: ASAM [Kwon et al., 2021] uses a scale-adaptive radius to reduce sensitivity to parameter rescaling, while Friendly SAM [Li et al., 2024] reduces minibatch-noise sensitivity by averaging perturbation information across batches. A complementary view comes from gradient-norm penalization: Zhao et al. [2022] argue that penalizing the gradient norm encourages smaller local Lipschitz constants and show that SAM can be recovered as a special case of a loss augmented by a gradient-norm regularizer. Geometry-aware sharpness neighborhoods. Vanilla SAM defines sharpness through Euclidean neighborhoods in parameter space, which may poorly reflect the effective geometry of neural networks [Li et al., 2018]. Fisher SAM addresses this by replacing the Euclidean ball with an ellipsoid induced by the Fisher information matrix, so the adversarial probe is measured in a geometry tied to the model distribution [Kim et al., 2022]. Riemannian SAM further generalizes this idea by defining the sharpness neighborhood through an arbitrary Riemannian metric [Yun and Yang, 2023]. In this view, Fisher SAM is a special case, and sharpness-aware perturbations are interpreted as steepestascent directions under a non-Euclidean local norm. Geometry-aware optimization and robustness. A separate line of work uses geometry to transport gradients rather than to define adversarial neighborhoods. Natural-gradient descent defines steepest descent with respect to the Fisher information metric [Amari, 1998]. K-FAC makes this principle scalable through layerwise Kronecker-factored Fisher approximations [Martens and Grosse, 2015], while Shampoo maintains tensor-structured preconditioners along parameter dimensions [Gupta et al., 2018]. These methods use curvature estimates to precondition the optimization trajectory, but they do not introduce a SAM-style local adversarial probe. Another related direction controls sensitivity to input perturbations rather than parameter perturbations: spectral normalization and Lipschitz regularization reduce the local sensitivity of the network function to data-space changes [Yoshida and Miyato, 2017, Virmaux and Scaman, 2018]. Thus, while geometry-aware optimizers act on the parameter-space training dynamics and Lipschitz methods target input-space robustness, sharpness-aware methods study loss variation under parameter perturbations.
3
Method
Metric and notation. Let L : Rd → R denote the training loss, with g(θ) := ∇L(θ) and H(θ) := ∇2 L(θ). We write P (θ) ≻ 0 for a local optimization metric and U (θ) := P (θ)−1 for its inverse. The associated dual norm on gradients is ∥g∥U := (g ⊤ U g)1/2 . LLQR inverse metric. The LLQR framework [Dufort-Labbé et al., 2026] (See Appendix A) learns a structured inverse preconditioner U from layerwise network dynamics and a divergence on the loss, such as KL/Fisher or Newton/Bregman. In practice, U is represented by diagonal or K-FAC blocks, trading expressivity for cost. It is updated only every nLLQR ∼ 500 optimizer steps and smoothed by an exponential moving average, so applying U during training is a structured matrix-vector product with modest overhead. For the analysis in Section 4, the relevant properties are that Ut is positive definite, uniformly spectrally bounded, and frozen during each perturb-and-recompute step. Euclidean SAM. Standard SAM [Foret et al., 2021] evaluates the update gradient at the worst-case first-order perturbation in a Euclidean ball: θ+ = θ + ρ
g(θ) , ∥g(θ)∥2
θ ← θ − η g(θ+ ).
LLQR+SAM. LLQR+SAM replaces the Euclidean perturbation geometry by the learned metric P = U −1 and uses the same inverse metric to transport the outer gradient: θ+ = θ + ρ
U g(θ) , ∥g(θ)∥U
θ ← θ − η U g(θ+ ).
(1)
The perturbation is therefore the normalized steepest-ascent direction under the metric P , while the descent step applies the same geometry to the gradient evaluated at the perturbed point. 3
Algorithm 1 LLQR+SAM 1: Input: initial parameters θ0 , learning rate η, SAM radius ρ, LLQR metric state, damping and update schedule. 2: for t = 0, 1, 2, . . . do 3: Draw minibatch bt and compute gt = ∇θ Lbt (θt ). 4: Update or query the slow LLQR state to obtain Ut ≈ Pt−1 . 5: Set ϵt = ρ Ut gt /(gt⊤ Ut gt )1/2 . 6: Form the probe parameters θt+ = θt + ϵt . 7: Compute the sharpness-aware gradient g̃t = ∇θ Lbt (θt+ ). 8: Apply the geometry-aware update θt+1 = θt − η Ut g̃t . 9: end for 10: Output: trained parameters θT .
Interpretation. Equation (1) can be viewed locally as a Riemannian SAM step in the LLQR metric, with Ut fixed during the perturb-and-recompute step. The corresponding fixed-metric transfer from Riemannian SAM [Yun and Yang, 2023] gives stationarity in the instantaneous metric; Appendix B.1 gives the formal statement. This also separates LLQR+SAM from Fisher SAM [Kim et al., 2022]: Fisher SAM uses a Fisher-scaled perturbation to probe sharpness, but typically applies the resulting gradient with the base optimizer. LLQR+SAM instead uses the learned inverse metric Ut both to define the perturbation neighborhood and to transport the outer gradient. This coupling is central to the pothole-escape mechanism analyzed in Section 4: the amplification depends on the metric acting on both the probe and the descent step. Algorithm 1 states the core update without momentum, weight decay, or optimizer-specific bookkeeping. In implementation, these terms are composed around the same two operations: use the current LLQR state to define the perturbation geometry, then apply the preconditioned sharpnessaware gradient through the geometry-aware base update. The wall-clock overhead of LLQR+SAM relative to the corresponding non-SAM, nonpreconditioned run is approximately one extra forward/backward pass per step (the SAM probe), plus the cost of refreshing and applying the LLQR preconditioner. As shown in Fig. 1 (right), this additional geometry does not become the bottleneck in practice: when the cost of SAM is acceptable, adding LLQR preserves the same affordability regime while yielding stronger accuracy. For reproducibility, the implementation is available at github.com/SimonDufLab/LLQR.
4
Analysis: pothole navigation in a two-scale landscape
We analyze LLQR +SAM on a minimal two-scale quadratic model: an average geometry tracked by LLQR, plus a sharp localized perturbation not captured by the slowly updated preconditioner. Two-scale model.
Place a local minimum at the origin and consider L(θ) = 12 θ⊤ Hθ,
H = H̄ + Hϵ ,
H ≻ 0,
(2)
where H̄ ≻ 0 encodes the smooth, slowly-varying part of the geometry and Hϵ ⪰ 0 is a localized sharp component. We make the following modeling assumption, motivated by the slow, exponential moving average smoothed nature of the LLQR preconditioner. Assumption 4.1. The LLQR inverse metric captures the average geometry: U = H̄ −1 . For simplicity, assume that H̄ and Hϵ share the same eigenbasis. Define λi = λ̄i + λϵ,i ,
µi := λi λ̄−1 = 1 + λϵ,i λ̄−1 i i .
Thus µi ≈ 1 in average directions and µi ≫ 1 in pothole directions. The non-commuting case is handled by whitening with H̄ −1/2 ; see Appendix B.2. Dynamics. For (2), g(θ) = Hθ, the LLQR+SAM recursion is exact for the quadratic loss (no Taylor approximation; see Appendix B.3): et+1 = (I − ηU H)et − ηρ ∥Het ∥−1 U U HU Het , 4
(3)
Toy SAM sharp-minima escape Level curves and optimizer paths
Loss during optimization
55
50
60
0 65 75 7
50
80
70
30
40
20
45
101
60
15
10
10
20
3 80
10
60
50
5
Loss
SGDM SGDM+SAM SGDM+LLQR SGDM+LLQR+SAM
1.5 × 101
45
15
55
6
75 60 6 7 5 0
3D half-surface and paths
x2
40 0
40
5
20 0
30 15
−6
−3
45
10
5 60
55
0
3
0
75
−6
60
−6
20
SGDM: Sharp trapped SGDM+SAM: Flat centered SGDM+LLQR: Sharp trapped SGDM+LLQR+SAM: 50 Flat centered
5
50
6
35
3
x2
0 75 65 7
x1
−3
55
20
0 0
70 6
−3
45
−6
−3
25
6 3
0
6
0
x1
20
40
60
80
100
Step
Figure 2: Toy sharp-well escape mechanism. The surface has a flat basin at the origin and a sharp annular basin near radius 5. All four optimizers use the same learning rate, and the SAM variants use the same radius, with starts chosen in the sharp basin, but not at minima. The non-SAM variants remain trapped, while the SAM variants leave the sharp well; the LLQR +SAM trajectory reaches the flat region with faster loss decay in this setting. where et = θt is the displacement from the minimum. Restricted to a single eigendirection vi (i.e., et ∝ vi ), −1/2
zi,t+1 = (1 − ηµi )zi,t − ηρ µi λ̄i
sign(zi,t ),
zi,t := ⟨vi , et ⟩.
(4)
Hovering and escape. The scalar map (4) has no nonzero fixed point: the SAM perturbation makes the iterate oscillate around the minimum instead of settling exactly at it. Its hovering envelope is −1/2 |zi,t | ≲ ρ λ̄i , (5) with the exact two-cycle amplitude given in Appendix B.4. This scale comes from a cancellation: −1/2 the SAM kick is of size ηρµi λ̄i , while the contraction scale is ηµi , hence −1/2
ηρµi λ̄i ηµi
−1/2
= ρλ̄i
.
Thus the hovering scale is controlled by the average curvature λ̄i , not by the localized sharpness λϵ,i . Consequently, if a pothole has basin radius rϵ , the iterate is no longer trapped in the pothole whenever ρ λ̄−1/2 > rϵ . (6) ϵ Comparison with vanilla SAM. For vanilla SAM, the same calculation gives a Euclidean hovering envelope of order ρ in every direction; see Appendix B.6. Therefore LLQR +SAM amplifies the −1/2 pothole-direction escape scale by λ̄ϵ . Corollary 4.2 (Pothole-escape amplification; see Cor. B.2). Under Assumption 4.1 and the commuting hypothesis, the LLQR +SAM hovering envelope around a pothole minimum is ρ λ̄−1/2 = λ̄−1/2 × vanilla-SAM hovering envelope . ϵ ϵ The amplification depends only on the average curvature λ̄ϵ in the pothole direction. Thus retuning vanilla SAM cannot reproduce the effect selectively: increasing ρ enlarges the probe in all directions, whereas LLQR +SAM uses the same U ≈ H̄ −1 for direction-adaptive probing and direction-adaptive contraction. Toy illustrations. Figures 2 and 3 instantiate the mechanism on a two-dimensional sharp-well landscape. The non-SAM variants remain trapped, whereas SAM variants escape. With injected gradient noise, LLQR +SAM reaches the flat basin with shorter path length than vanilla SAM, consistent with Corollary 4.2. A stochastic selection model is given in Appendix B.7. 5
Toy SAM sharp-minima escape with gradient noise Level curves and optimizer paths 5 150
15
175
12
150
13
5
16
Loss during optimization 1.25 × 102 102
13
8
5 16
5
3D half-surface and paths
12
0
0
SGDM SGDM+SAM SGDM+LLQR SGDM+LLQR+SAM
10
0
10
5
5
45
4
125
15
15
160
80
x2
Loss
120 100
0
40
−4
60
5 10
12 0 5
5
0
16
0
15
16
−10
25
13
−8
5
10
50
SGDM: Sharp trapped 75 SGDM+SAM: Flat centered 90 SGDM+LLQR: Sharp trapped SGDM+LLQR+SAM: Flat centered
13
5
x2
0 15
x1
−5
5
0 0
75
10
−5
0 12
5
−10
30
10
15
0
5
0
−8
−4
0
4
8
0
x1
80
160
240
320
400
Step
Figure 3: Gradient-noise escape from the sharp minimum. All variants start at the bottom of the sharp well and receive a shared deterministic Gaussian perturbation schedule in the update gradient. The non-SAM variants remain near the sharp well, while SAM variants are ejected. At variance 10−9 , both SAM variants reach the flat basin, but LLQR +SAM has substantially shorter path length than Euclidean SAM, consistent with the stochastic selection picture in Appendix B.7. Table 1: CIFAR-10 standard supervised benchmark. Entries are best top-1 test accuracy (mean ± std, 5 seeds). Blue denotes NGD-induced geometry, Orange denotes Newton-induced geometry. Base geometry Architecture VGG16-BN ResNet-18 WRN-28-10 PyramidNet-110
Sharpness baselines
Combined geometry and sharpness
SGDM
LLQR
LLQR
SAM
F ISHER SAM
LLQR+SAM
LLQR+SAM
95.10 ± 0.14 96.14 ± 0.19 96.94 ± 0.09 97.18 ± 0.17
95.27 ± 0.19 96.37 ± 0.11 97.04 ± 0.08 97.21 ± 0.08
95.41 ± 0.11 96.37 ± 0.11 97.10 ± 0.03 97.26 ± 0.10
95.54 ± 0.14 96.65 ± 0.16 97.43 ± 0.09 97.63 ± 0.09
95.45 ± 0.04 96.74 ± 0.14 97.41 ± 0.04 97.57 ± 0.12
95.92 ± 0.07 96.85 ± 0.09 97.64 ± 0.08 97.86 ± 0.04
95.73 ± 0.08 96.78 ± 0.03 97.52 ± 0.07 97.66 ± 0.08
Stochastic counterpart: selection between wells. The deterministic analysis above shows that pothole minima are not attracting fixed points of LLQR+SAM. Under stochastic gradients this refines into a selection statement: sharp wells are short-lived (entered briefly through gradient noise, exited through the SAM kick), while flat wells support long visits. The same hovering-vs-basin comparison (6) controls the expected exit time, with LLQR+SAM benefiting from the amplification of Corollary 4.2. We give the regenerative model and the exit-time bound in Appendix B.7; the gradient-noise toy of Figure 3 is consistent with the prediction that LLQR+SAM exits the sharp well faster than vanilla SAM and reaches the flat basin with a markedly shorter path length.
5
Experiments
In this section, we evaluate LLQR+SAM on standard supervised vision benchmarks (CIFAR10/100, TinyImageNet, ImageNet), a sequence-modeling task (IWSLT14 De-En), and a scalability study on ViT models up to ViT-Large/16. Throughout, LLQR+SAM reuses the hyperparameters of its LLQR and SAM components without method-specific retuning, isolating the effect of the geometry-sharpness coupling. Full training protocols are deferred to Appendix C. Additional experiments about noise injection and method scalability are presentend in Appendix D. 5.1
Standard Supervised Training Benchmarks
CIFAR Datasets. We first evaluate the method on CIFAR-10 and CIFAR-100 [Krizhevsky and Hinton, 2009] with VGG16-BN [Simonyan and Zisserman, 2015], ResNet-18 [He et al., 2016], WRN-28-10 [Zagoruyko and Komodakis, 2016], and PyramidNet-110 [Han et al., 2017]; the full protocol is deferred to Appendix C.1. Tables 1 and 2 show that LLQR+SAM consistently improves over either component alone, indicating that coupling SAM with the learned LLQR geometry is more effective than applying either mechanism in isolation. Across CIFAR-10 and CIFAR-100, the NGD variant delivers the strongest results, suggesting that the natural-gradient geometry provides the more effective steepest descent metric in those setups. TinyImageNet. We next scale the ResNet family networks to TinyImageNet-200 to test whether the geometry–sharpness coupling remains effective beyond CIFAR. Table 3 shows that the CIFAR 6
Table 2: CIFAR-100 standard supervised benchmark. Entries are best top-1 test accuracy (mean ± std, 5 seeds). Blue denotes NGD-induced geometry, Orange denotes Newton-induced geometry. Base geometry Architecture VGG16-BN ResNet-18 WRN-28-10 PyramidNet-110
Sharpness baselines
Combined geometry and sharpness
SGDM
LLQR
LLQR
SAM
F ISHER SAM
LLQR+SAM
LLQR+SAM
75.65 ± 0.28 79.27 ± 0.18 82.50 ± 0.18 83.96 ± 0.25
76.32 ± 0.35 79.53 ± 0.39 82.80 ± 0.37 84.40 ± 0.46
76.31 ± 0.11 79.83 ± 0.18 82.72 ± 0.24 84.00 ± 0.16
76.70 ± 0.18 81.13 ± 0.21 84.83 ± 0.06 86.20 ± 0.18
76.70 ± 0.10 80.89 ± 0.31 84.71 ± 0.13 85.92 ± 0.10
78.14 ± 0.27 81.82 ± 0.14 85.30 ± 0.14 86.67 ± 0.21
77.67 ± 0.27 81.65 ± 0.10 85.05 ± 0.12 86.53 ± 0.05
Table 3: TinyImageNet-200 ResNet-family benchmark. Entries are best top-1 test accuracy (mean ± std, 3 seeds). Blue denotes NGD-induced geometry, Orange denotes Newton-induced geometry. Base geometry
Sharpness baseline
Combined geometry and sharpness
Architecture
SGDM
LLQR
LLQR
SAM
LLQR+SAM
LLQR+SAM
ResNet-18 WRN-28-10
57.97 ± 0.28 61.06 ± 0.37
57.02 ± 0.49 60.08 ± 0.32
57.62 ± 0.47 60.37 ± 0.33
61.17 ± 0.27 63.91 ± 0.41
62.24 ± 0.49 64.65 ± 0.51
61.98 ± 0.22 64.59 ± 0.25
trend persists when scaling to TinyImageNet. SAM gives a strong gain over SGDM, while LLQR alone is not sufficient or not well-tuned in this setting. Yet, the combined LLQR+SAM variants improve over SAM on both architectures, with the NGD-induced geometry giving the best result in each case. This is a particularly stringent test of complementarity: even when LLQR alone does not improve over SGDM, its geometry still strengthens SAM, suggesting that the learned geometry provides useful information that is not captured by the sharpness perturbation alone. Synergy with SAM variants. The SAM framework has led to a broad family of perturb-andrecompute optimizers, many differing only in how the adversarial direction is filtered or normalized. To check that LLQR is not merely tuned to vanilla SAM, we pair it with F-SAM [Li et al., 2024], a strong SAM variant. Performance is compared with and without the learned LLQR geometry, using vanilla SAM as a reference in Table 4. ImageNet. We also extend the comparison to ResNet-50 on ImageNet [Deng et al., 2009]. As shown in Table 5, adding LLQR to F-SAM again improves performance over the sharpness-aware baseline. The full top-1 error trajectory in Fig. 1 (left) makes the effect even clearer: the paired method not only reaches better generalization, but also converges faster across early and mid training. These results suggest that LLQR and SAM-family pairing synergize throughout the entire training process, even at ImageNet scale. IWSLT14 German-to-English. To assess if the geometry–sharpness coupling transfers beyond vision, we evaluate on the fairseq IWSLT14 German-to-English Transformer benchmark. As shown in Table 6, the gains are more modest than in image classification, but LLQR+SAM still achieves the best BLEU. More strikingly, the convergence curves in Fig. 4 show faster optimization when
34
Token Error (Validation)
Best BLEU (Validation)
70
33 32 31 30 15
30
45
Steps (×103) AdamW
60
65 60 55 50 45 40 35 0
SAM
LLQR
15
30
45
Steps (×103)
60
LLQR+SAM
Figure 4: IWSLT14 German-to-English convergence. Validation BLEU and token error curves for the fairseq Transformer benchmark. Pairing LLQR with SAM accelerates optimization while offering the modest best-performance gains reported in Table 6 7
Table 4: CIFARs Friendly-SAM synergy benchmark. Entries are best top-1 test accuracy (mean ± std, 5 seeds). Blue denotes NGD-induced LLQR geometry. Without LLQR Dataset
Architecture
CIFAR-10 CIFAR-10 CIFAR-10 CIFAR-10 CIFAR-100 CIFAR-100 CIFAR-100 CIFAR-100
With LLQR geometry
SAM
F-SAM
LLQR+SAM
LLQR +F-SAM
VGG16-BN ResNet-18 WRN-28-10 PyramidNet-110
95.54 ± 0.14 96.65 ± 0.16 97.43 ± 0.09 97.63 ± 0.09
95.58 ± 0.09 96.60 ± 0.05 97.49 ± 0.08 97.65 ± 0.07
95.92 ± 0.07 96.85 ± 0.09 97.64 ± 0.08 97.86 ± 0.04
95.81 ± 0.11 96.88 ± 0.11 97.68 ± 0.07 97.89 ± 0.06
VGG16-BN ResNet-18 WRN-28-10 PyramidNet-110
76.70 ± 0.18 81.13 ± 0.21 84.83 ± 0.06 86.20 ± 0.18
76.95 ± 0.32 81.16 ± 0.12 84.86 ± 0.16 86.23 ± 0.08
78.14 ± 0.27 81.82 ± 0.14 85.30 ± 0.14 86.67 ± 0.21
78.22 ± 0.31 82.00 ± 0.28 85.21 ± 0.20 86.76 ± 0.13
Table 5: ImageNet/ResNet-50 benchmark. Entries are best top-1 test accuracy (mean ± std, 3 seeds). Blue denotes NGD-induced geometry, Orange denotes Newton-induced geometry. LLQR rows use the E-KFAC block structure. Base geometry
Sharpness baseline
Combined geometry and sharpness
Architecture
SGDM
LLQR
LLQR
F-SAM
LLQR +F-SAM
LLQR +F-SAM
ResNet-50
77.60 ± 0.31
78.05 ± 0.12
76.93 ± 0.42
77.92 ± 0.07
78.61 ± 0.16
78.19 ± 0.14
LLQR and SAM are paired, extending the same pattern observed in vision to a sequence-modeling benchmark. 5.2
Isolating the role of geometry in probe and transport: FisherSAM Ablation
FisherSAM uses its geometry only to define the perturbation; LLQR+SAM additionally uses it to transport the outer update. To separate these two roles, we introduce LLQR∆ +SAM, which applies the learned preconditioner only in the perturbation step (matching the FisherSAM pattern). On ResNet18/CIFAR-100 (Table 7), LLQR∆ +SAM already improves over FisherSAM–the LLQR geometry is itself a better probe– and full LLQR+SAM improves further, showing that transporting the outer update in the same geometry adds a separate gain. 5.3
Table 7: ResNet-18/CIFAR-100 FisherSAM-style perturbation ablation. Entries are best top-1 accuracy (mean ± std, five seeds). LLQR∆ +SAM uses LLQR preconditioner only in the perturbation step, similar to FisherSAM. Method
Top-1 accuracy
F ISHER SAM LLQR∆ +SAM LLQR+SAM
80.89 ± 0.31 81.23 ± 0.14 81.82 ± 0.14
Scalability
To assess scalability, we evaluate the computational and memory footprint of LLQR and LLQR+SAM on ViT models of increasing size [Dosovitskiy et al., 2021], from ViT-Tiny/16 to ViT-Large/16, on ImageNet-scale data. We compare LLQR and LLQR+SAM against SGD, AdamW, and SAM, measuring wall-clock time across update cadences. Figure 6 reports the steady preconditioner-update time as a function of model size: empirically, the recurring LLQR cost scales as O(P 1.27–1.28 ) on the number of parameters P , closer to linear than to the quadratic scaling typically associated with second-order methods. Pairing LLQR with SAM does not change this scaling. Memory is controlled by the choice of preconditioner block structure (E-KFAC, K-FAC, diag-KFAC, diagonal, etc.); Appendix D.2 details the storage tradeoffs and reports the cadence sweep across update intervals. We also provide cadence results in Table 10 (Appendix D.2) and Fig. 5, to be read as wall-clock overhead measurements, not as accuracy or convergence comparisons. As expected, the cost decreases when LLQR updates are performed less frequently, since fewer preconditioner refreshes are executed per epoch. This steady-update scaling trend is further exposed in Fig. 6. 8
Table 6: IWSLT14 German-to-English Transformer benchmark. Entries are best validation BLEU (mean ± std, 5 seeds). Blue denotes NGD-induced LLQR geometry. Base geometry
Per-epoch seconds (compile excluded)
IWSLT14 De-En
A DAM W
LLQR
SAM
LLQR+SAM
34.24 ± 0.27
34.51 ± 0.12
34.20 ± 0.19
34.57 ± 0.09
vit_b16: LLQR cadence sweep, compile-adjusted per-epoch runtime
1000 800 600 400 200 0
50 100
200
Combined geometry and sharpness
500
1000 update_preconditioner_every LLQR SGD SAM LLQR+SAM AdamW
1500
Per-epoch seconds (compile excluded)
Dataset
Sharpness baseline
vit_l16: LLQR cadence sweep, compile-adjusted per-epoch runtime 1750 1500 1250 1000 750 500 250 0
50 100
200
500
1000 update_preconditioner_every LLQR SGD SAM LLQR+SAM AdamW
1500
Figure 5: LLQR cadence sweep for ViT-B/16 and ViT-L/16. Each panel shows compile-adjusted seconds per epoch versus the LLQR update interval, with first-order baselines as horizontal references. Note that the ImageNet ResNet-50 runs in Fig. 1 use 1500-step updates.
6
Conclusion Steady LQR preconditioner update time (sec, log scale)
Measured LQR Time-Complexity Scaling Across ViT Sizes
We introduced LLQR+SAM, which pairs a slowly-updated LLQR preconditioner with a SAM perturbation evaluated and transported in the induced geometry. On a quadratic two-scale model the dynamics is closed-form: SAM prevents the iterate from localizing at any minimum, and the learned U shapes the hovering scale to match the average geometry– large enough to escape sharp potholes, small enough to stay inside wide basins. The potholedirection probe is amplified relative to vanilla SAM by a factor that grows as the surrounding basin becomes flatter–the regime where preconditioning matters most. The same U that improves conditioning therefore also improves escape: U and SAM are not independent ingredients but two complementary uses of one learned geometry.
Linear reference O(P) Quadratic reference O(P^2) LLQR measured LLQR+SAM measured
103
L/16
102
B/16
101
100
S/16
Ti/16
LLQR: t ~ P^1.27 LLQR+SAM: t ~ P^1.28 101
102 Approx. trainable parameters (M, log scale)
Figure 6: Across ViT scales, preconditionerupdate time grows nearly linearly for both LLQR and LLQR+SAM, far from the quadratic scaling typically associated with second-order methods.
Empirically, LLQR+SAM yields consistent gains over either component alone across CIFAR10/100, TinyImageNet, ImageNet, and IWSLT14, including when paired with stronger SAM variants such as F-SAM, and the per-step overhead remains close to first-order training across ViT scales. Limitations Our analysis relies on an idealized two-timescale quadratic model in which the slow LLQR metric captures average geometry and the SAM perturbation probes localized sharpness. This setting is useful for exposing the mechanism, but it cannot certify that the same decomposition transfers cleanly to the highly nonconvex and stochastic regime of deep networks. The empirical scope is also finite: we cover standard supervised vision benchmarks and one sequence-modeling task, but not large language model pre-training or fine-tuning, nor every architecture family. Finally, LLQR+SAM inherits the extra gradient evaluation of SAM and adds periodic LLQR metric updates; while the resulting overhead is practical in our experiments, the best block structure, update cadence, and memory budget remain scale-dependent design choices.
9
References Shun-ichi Amari. Natural gradient works efficiently in learning. Neural Comput., 10(2): 251–276, 1998. doi: 10.1162/089976698300017746. URL https://doi.org/10.1162/ 089976698300017746. Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE Computer Society Conference on Computer Vision and Pattern Recognition (CVPR 2009), 20-25 June 2009, Miami, Florida, USA, pages 248–255. IEEE Computer Society, 2009. doi: 10.1109/CVPR.2009.5206848. URL https://doi.org/ 10.1109/CVPR.2009.5206848. Terrance Devries and Graham W. Taylor. Improved regularization of convolutional neural networks with cutout. CoRR, abs/1708.04552, 2017. URL http://arxiv.org/abs/1708.04552. Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. Sharp minima can generalize for deep nets. In Proceedings of the 34th International Conference on Machine Learning, ICML 2017, Sydney, NSW, Australia, 6-11 August 2017, Proceedings of Machine Learning Research, pages 1019–1028. PMLR, 2017. URL http://proceedings.mlr.press/v70/dinh17b.html. Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021, 2021. URL https://openreview.net/forum?id=YicbFdNTTy. Simon Dufort-Labbé, Pierre-Luc Bacon, Razvan Pascanu, Simon Lacoste-Julien, and Aristide Baratin. Layerwise lqr for geometry-aware optimization of deep networks, 2026. URL https: //arxiv.org/abs/2605.04230. Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimization for efficiently improving generalization. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL https://openreview.net/forum?id=6Tm1mposlrM. Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor optimization. In Proceedings of the 35th International Conference on Machine Learning, ICML 2018, Stockholmsmässan, Stockholm, Sweden, July 10-15, 2018, Proceedings of Machine Learning Research, pages 1837–1845. PMLR, 2018. URL http://proceedings.mlr.press/v80/ gupta18a.html. Dongyoon Han, Jiwhan Kim, and Junmo Kim. Deep pyramidal residual networks. In 2017 IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2017, Honolulu, HI, USA, July 21-26, 2017, pages 6307–6315. IEEE Computer Society, 2017. doi: 10.1109/CVPR.2017.668. URL https://doi.org/10.1109/CVPR.2017.668. Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In 2016 IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2016, Las Vegas, NV, USA, June 27-30, 2016, pages 770–778. IEEE Computer Society, 2016. doi: 10.1109/CVPR.2016.90. URL https://doi.org/10.1109/CVPR.2016.90. Sepp Hochreiter and Jürgen Schmidhuber. Flat minima. Neural Comput., 9(1):1–42, 1997. doi: 10.1162/neco.1997.9.1.1. URL https://doi.org/10.1162/neco.1997.9.1.1. Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 2426, 2017, Conference Track Proceedings. OpenReview.net, 2017. URL https://openreview. net/forum?id=H1oyRlYgg. Minyoung Kim, Da Li, Shell Xu Hu, and Timothy M. Hospedales. Fisher SAM: information geometry and sharpness aware minimisation. International Conference on Machine Learning, ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA, pages 11148–11161, 2022. URL https://proceedings.mlr.press/v162/kim22f.html. 10
Alex Krizhevsky and Geoffrey Hinton. Learning multiple layers of features from tiny images. Technical report, University of Toronto, 2009. Jungmin Kwon, Jeongseop Kim, Hyunseo Park, and In Kwon Choi. ASAM: adaptive sharpnessaware minimization for scale-invariant learning of deep neural networks. In Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event, Proceedings of Machine Learning Research, pages 5905–5914. PMLR, 2021. URL http:// proceedings.mlr.press/v139/kwon21b.html. Bingcong Li and Georgios B. Giannakis. Enhancing sharpness-aware optimization through variance suppression. Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurIPS 2023, New Orleans, LA, USA, December 10 - 16, 2023, 2023. Hao Li, Zheng Xu, Gavin Taylor, Christoph Studer, and Tom Goldstein. Visualizing the loss landscape of neural nets. Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NeurIPS 2018, December 3-8, 2018, Montréal, Canada, pages 6391–6401, 2018. Tao Li, Pan Zhou, Zhengbao He, Xinwen Cheng, and Xiaolin Huang. Friendly sharpness-aware minimization. In IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR 2024, Seattle, WA, USA, June 16-22, 2024, pages 5631–5640. IEEE, 2024. doi: 10.1109/CVPR52733. 2024.00538. URL https://doi.org/10.1109/CVPR52733.2024.00538. James Martens and Roger B. Grosse. Optimizing neural networks with kronecker-factored approximate curvature. In Proceedings of the 32nd International Conference on Machine Learning, ICML 2015, Lille, France, 6-11 July 2015, JMLR Workshop and Conference Proceedings, pages 2408– 2417. JMLR.org, 2015. URL http://proceedings.mlr.press/v37/martens15.html. Peng Mi, Li Shen, Tianhe Ren, Yiyi Zhou, Xiaoshuai Sun, Rongrong Ji, and Dacheng Tao. Make sharpness-aware minimization stronger: A sparsified perturbation approach. Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NeurIPS 2022, New Orleans, LA, USA, November 28 - December 9, 2022, 2022. Behnam Neyshabur, Ruslan Salakhutdinov, and Nathan Srebro. Path-sgd: Path-normalized optimization in deep neural networks. Advances in Neural Information Processing Systems 28: Annual Conference on Neural Information Processing Systems 2015, December 7-12, 2015, Montreal, Quebec, Canada, pages 2422–2430, 2015. URL https://proceedings.neurips.cc/ paper/2015/hash/eaa32c96f620053cf442ad32258076b9-Abstract.html. Myle Ott, Sergey Edunov, David Grangier, and Michael Auli. Scaling neural machine translation. In Proceedings of the Third Conference on Machine Translation: Research Papers, WMT 2018, Belgium, Brussels, October 31 - November 1, 2018, pages 1–9. Association for Computational Linguistics, 2018. doi: 10.18653/V1/W18-6301. URL https://doi.org/10.18653/ v1/w18-6301. Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015. URL http://arxiv.org/abs/ 1409.1556. Aladin Virmaux and Kevin Scaman. Lipschitz regularity of deep neural networks: analysis and efficient estimation. Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NeurIPS 2018, December 3-8, 2018, Montréal, Canada, pages 3839–3848, 2018. Yuichi Yoshida and Takeru Miyato. Spectral norm regularization for improving the generalizability of deep learning. CoRR, abs/1705.10941, 2017. URL http://arxiv.org/abs/1705.10941. Jihun Yun and Eunho Yang. Riemannian SAM: sharpness-aware minimization on riemannian manifolds. In Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine, editors, Advances in Neural Information Processing Systems 36: Annual Conference on 11
Neural Information Processing Systems 2023, NeurIPS 2023, New Orleans, LA, USA, December 10 - 16, 2023, 2023. URL http://papers.nips.cc/paper_files/paper/2023/hash/ cf701db0e3b4d0b8681ca6915ac3e87e-Abstract-Conference.html. Sergey Zagoruyko and Nikos Komodakis. Wide residual networks. In Proceedings of the British Machine Vision Conference 2016, BMVC 2016, York, UK, September 19-22, 2016. BMVA Press, 2016. URL https://bmva-archive.org.uk/bmvc/2016/papers/paper087/index.html. Yang Zhao, Hao Zhang, and Xiuyuan Hu. Penalizing gradient norm for efficiently improving generalization in deep learning. International Conference on Machine Learning, ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA, pages 26982–26992, 2022. URL https: //proceedings.mlr.press/v162/zhao22i.html.
12
Contents 1
Introduction
2
2
Related Work
2
3
Method
3
4
Analysis: pothole navigation in a two-scale landscape
4
5
Experiments
6
5.1
Standard Supervised Training Benchmarks . . . . . . . . . . . . . . . . . . . . . .
6
5.2
Isolating the role of geometry in probe and transport: FisherSAM Ablation . . . . .
8
5.3
Scalability . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
8
6
Conclusion
9
A Layerwise LQR: Learning a Reusable Geometry
14
B Detailed Analysis for the Two-Scale Pothole Model
16
B.1 Riemannian-SAM transfer under the frozen LLQR metric . . . . . . . . . . . . . .
16
B.2 Whitened form and non-commuting case . . . . . . . . . . . . . . . . . . . . . . .
17
B.3 Coordinate recursion . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
17
B.4 Scalar Hovering Envelope . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
18
B.5 Pothole Escape Criterion . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
18
B.6 Vanilla SAM Comparison . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
18
B.7 Stochastic Selection Between Wells . . . . . . . . . . . . . . . . . . . . . . . . .
19
C Experiments Details
21
C.1 Standard Supervised Training Benchmarks . . . . . . . . . . . . . . . . . . . . . . D Additional Experimental Results
21 22
D.1 Noise injection robustness . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
22
D.2 Scalability . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
22
13
A
Layerwise LQR: Learning a Reusable Geometry
We briefly summarize the Layerwise LQR (LLQR) methodology of Dufort-Labbé et al. [2026], which provides the geometry used throughout this work. The motivation is to obtain the benefits of second-order or natural-gradient descent without explicitly forming, storing, or inverting a global curvature matrix. At an iterate θk , a broad class of geometry-aware descent directions can be written as the solution of the local quadratic problem 1 (7) ∆θ⋆ = arg min ∇L(θk )⊤ ∆θ + ∆θ⊤ H(θk )∆θ , ∆θ 2 where H may be a regularized Hessian, Gauss–Newton matrix, Fisher matrix, or more general divergence-induced metric. Directly solving H∆θ = −∇L(θk ) is infeasible for modern networks, and standard scalable approximations typically impose block structure on H before solving the problem. LLQR instead first rewrites the dense quadratic problem as a layerwise optimal-control problem, and only later imposes structure on the learned inverse update. Consider a depth-N network i = 0, . . . , N − 1.
xi+1 = fi (xi , θi ),
Linearizing the forward pass at (xki , θik ) gives the perturbation dynamics δxi+1 = Ai δxi + Bi δθi ,
δx0 = 0,
(8)
with Ai = ∂fi /∂xi and Bi = ∂fi /∂θi . For divergence-induced quadratic models, the global quadratic form in (7) can then be decomposed into layerwise costs δx⊤ N QN δxN +
N −1 X i=0
δxi δθi
⊤
Qi Mi
Mi⊤ Ri
δxi , δθi
subject to the linearized dynamics (8). The resulting problem is a finite-horizon Linear Quadratic Regulator: the network defines the dynamics, while the chosen descent geometry defines the quadratic costs. Its exact Riccati solution recovers the corresponding second-order update, including Newton, Gauss–Newton, and natural-gradient variants. In large networks, however, the exact Riccati recursion remains too expensive, so LLQR uses it as a reference objective rather than as the deployed optimizer. The scalable relaxation parameterizes the update by a structured inverse preconditioner δθi = −Ui ∇θi L(θk ),
U = diag(U0 , . . . , UN −1 ),
Algorithm 2 LLQR: relaxed periodic preconditioner Require: model f , loss L, data loader D, outer optimizer Oout , step size η Require: recompute period n, inner steps T , structure S, inner optimizer Oin , inner step size α, EMA parameter β 1: Initialize θ 0 and U ← I projected to S 2: for k = 0, 1, 2, . . . do 3: Sample minibatch (xk , y k ) ∼ D and compute gk ← ∇θ L(θk ) 4: if k mod n = 0 then 5: Linearize the network: (Ai , Bi ) ← (∇x , ∇θ )fi (xki , θik ) 6: Form the LQR blocks (Qi , Ri , Mi ) for the chosen divergence-induced metric 7: Set U0 ← U 8: for t = 0 to T − 1 do 9: Ut+1 ← Oin (Ut , ∇Ut J, α), where J is the relaxed objective (9) 10: end for 11: U ← βU + (1 − β)UT 12: end if 13: Precondition the gradient: gek ← U gk 14: Update parameters with the base optimizer: θk+1 ← Oout (θk , gek , η) 15: end for
14
where each Ui may be diagonal, Kronecker-factored, E-KFAC-like, or another structured block. Substituting this restricted update into the LQR objective yields a direct objective over U : N −1 X 1 ⊤ 1 ⊤ 1 ⊤ ⊤ ⊤ min ∇xN ℓ(xN ) δxN + δxN QN δxN + δx Qi δxi + δθi Ri δθi + δθi Mi δxi U 2 2 i 2 (9) i=0 s.t.
δxi+1 = Ai δxi − Bi Ui ∇θi L(θk ),
δx0 = 0.
Thus, structure is imposed on the reusable inverse action U , not on the curvature model before the layerwise objective is derived. This distinction is important for LLQR +SAM: the learned geometry still comes from a layer-coupled objective that encodes the dense quadratic model through the forward dynamics, while the deployed update only requires applying U to gradients. In the LLQR +SAM setting, U should be viewed as a slowly updated transport geometry. Between preconditioner refreshes, it is held fixed and applied cheaply to the current gradient. This makes the method compatible with standard optimizers and with perturb-and-recompute schemes: SAM determines a local sharpness-aware gradient, while LLQR supplies the anisotropic geometry used to probe and/or transport that gradient.
15
B
Detailed Analysis for the Two-Scale Pothole Model
This appendix contains the supporting arguments used in Sections 3 and 4. The first subsection records the fixed-metric Riemannian-SAM transfer for LLQR +SAM. The remaining subsections justify the two-scale pothole calculation: the coordinate recursion, the scalar hovering envelope, the vanilla-SAM comparison, and the stochastic selection picture. B.1
Riemannian-SAM transfer under the frozen LLQR metric
This subsection records the convergence-transfer statement used in Section 3. The statement is conditional on the LLQR metric being frozen during each perturb-and-recompute step; it gives stationarity in the instantaneous metric, not in a fixed limiting metric unless additional metric convergence is assumed. Let Pt = Ut−1 . Assume that Ut is measurable with respect to the optimization history Ft , is fixed during the SAM perturb-and-recompute step, and satisfies the uniform spectral bounds 0 < umin I ⪯ Ut ⪯ umax I < ∞. Assume also that the smoothness, retraction-Lipschitz, and stochastic-gradient assumptions of Riemannian SAM [Yun and Yang, 2023] hold uniformly for the compact family of metrics {Pt }. Proposition B.1 (Riemannian-SAM transfer for LLQR +SAM). With the same decaying stepsize and perturbation radius schedule as in Riemannian SAM, for τ ∼ Unif{0, . . . , T − 1}, E ∥ gradPτ L(θτ )∥2Pτ = E ∇L(θτ )⊤ Uτ ∇L(θτ ) = O(T −1/2 ), up to the same stochastic-gradient and minibatch-variance terms as in the Riemannian-SAM theorem. Consequently, E∥∇L(θτ )∥22 = O(T −1/2 ). Proof. Condition on Ft . Since Ut is Ft -measurable and fixed during the perturb-and-recompute step, Pt = Ut−1 is a fixed Riemannian metric for the current update. For the constant metric ⟨ξ, ζ⟩Pt = ξ ⊤ Pt ζ, the Riemannian gradient is gradPt L(θ) = Pt−1 ∇L(θ) = Ut ∇L(θ), and hence
∥ gradPt L(θ)∥2Pt = ∇L(θ)⊤ Ut ∇L(θ). The LLQR + SAM perturbation Ut gt ϵt = ρt ⊤ (gt Ut gt )1/2 is therefore the normalized Riemannian steepest-ascent direction under the frozen metric Pt . The update θt+1 = θt − αt Ut ∇Lbt (θt + ϵt ) is the corresponding Riemannian steepest-descent step on the probed loss. Conditioned on Ft , this is one fixed-metric Riemannian-SAM step. The spectral bounds on Ut place all metrics Pt = Ut−1 in a compact subset of the positive-definite cone. Hence the norm-equivalence, smoothness, retraction-Lipschitz, and stochastic-gradient constants in the Riemannian-SAM one-step descent inequality can be chosen uniformly over t. Applying that inequality conditionally on Ft , then taking expectation and summing over t = 0, . . . , T − 1, gives T −1 1 X E∥ gradPt L(θt )∥2Pt = O(T −1/2 ), T t=0 up to the same stochastic-gradient and minibatch-variance terms. Unif{0, . . . , T − 1}, E∥ gradPτ L(θτ )∥2Pτ = O(T −1/2 ). 16
Equivalently, for τ
∼
Finally, since Ut ⪰ umin I, ∇L(θt )⊤ Ut ∇L(θt ) ≥ umin ∥∇L(θt )∥22 , so stationarity in the instantaneous LLQR metric implies Euclidean first-order stationarity: ⊤ −1/2 E∥∇L(θτ )∥22 ≤ u−1 ). min E ∇L(θτ ) Uτ ∇L(θτ ) = O(T If Ut → Uf and ∥∇L(θt )∥2 ≤ G, then ∇L(θt )⊤ Ut ∇L(θt ) − ∇L(θt )⊤ Uf ∇L(θt ) ≤ G2 ∥Ut − Uf ∥op . Averaging gives T −1 T −1 1 X 1 X E ∇L(θt )⊤ Uf ∇L(θt ) ≤ O(T −1/2 ) + G2 E∥Ut − Uf ∥op . T t=0 T t=0 Thus convergence in the limiting metric follows whenever this average discrepancy vanishes, and the original rate is preserved when the discrepancy is O(T −1/2 ). B.2
Whitened form and non-commuting case
The commuting hypothesis used in the main text is not essential. Define yt := H̄ 1/2 et , A := H̄ −1/2 H H̄ −1/2 = I + H̄ −1/2 Hϵ H̄ −1/2 . Since A ≻ 0 is symmetric, it admits an orthonormal eigendecomposition. Its eigenvalues are the perceived sharpnesses µi . In the commuting case, these reduce to λi λϵ,i µi = =1+ . λ̄i λ̄i Using U = H̄ −1 , the exact LLQR+SAM recursion becomes, in whitened coordinates, ηρ yt+1 = (I − ηA)yt − A2 yt . (10) ∥Ayt ∥2 Thus the non-commuting case is identical after diagonalizing the normalized curvature matrix A. The shared-eigenbasis presentation in the main text is only a notational simplification. B.3
Coordinate recursion
For the quadratic loss L(θ) = 12 θ⊤ Hθ, we have g(θ) = Hθ. The LLQR+SAM perturbation gives HU Hθ g(θ+ ) = Hθ + ρ , ∥Hθ∥U and therefore, with et = θt , ηρ U HU Het . (11) et+1 = (I − ηU H)et − ∥Het ∥U In P the commuting case, let vi be the common eigenvectors of H̄, Hϵ , and H, and write et = i zi,t vi . Since H̄vi = λ̄i vi , Hvi = λi vi , we get λi U HU Hvi = µ2i vi . U Hvi = H̄ −1 Hvi = vi = µi vi , λ̄i Moreover, X X 2 2 ∥Het ∥2U = (Het )⊤ U (Het ) = λ2j λ̄−1 µ2j λ̄j zj,t . j zj,t = j
j
Substituting these identities into (11) gives the full coordinate recursion ηρµ2i zi,t+1 = (1 − ηµi )zi,t − qP zi,t . (12) 2 λ̄ z 2 µ j j j,t j p If one eigendirection vi dominates, the denominator reduces to µi λ̄i |zi,t |, yielding the scalar recursion ηρµi zi,t+1 = (1 − ηµi )zi,t − p sign(zi,t ). (13) λ̄i This is the one-dimensional model used in the main text. 17
B.4
Scalar Hovering Envelope
Consider the scalar map zt+1 = azt − b sign(zt ),
a ∈ (0, 1),
b > 0.
(14)
It has no nonzero fixed point. Indeed, if z > 0, then b < 0, 1−a which contradicts z > 0; the negative half-line gives the symmetric contradiction. Hence the map cannot localize at a nonzero point and instead oscillates around the origin. z = az − b
z=−
=⇒
Once the signs alternate, the limiting two-cycle has amplitude r⋆ satisfying r⋆ = b − ar⋆ , so r⋆ =
b . 1+a
(15)
For (13),
ηρµi bi = p , λ̄i and therefore the exact asymptotic two-cycle amplitude is ai = 1 − ηµi ,
lim sup |zi,t | = t→∞
ηρµi bi p . = 1 + ai (2 − ηµi ) λ̄i
(16)
The main text uses the simpler hovering envelope |zi,t | ≲
ρ bi =p . 1 − ai λ̄i
(17)
This envelope is the scale relevant for the cancellation argument: the SAM kick bi and the contraction scale 1 − ai = ηµi both grow linearly with the perceived sharpness µi , leaving a radius controlled only by the average curvature λ̄i . B.5
Pothole Escape Criterion
The quadratic model is local, so it does not by itself describe the global motion after leaving a well. It gives a localization criterion. If a pothole has basin radius rϵ in the relevant eigendirection, then the local model predicts loss of localization when the hovering envelope exceeds that radius: ρ p > rϵ . (18) λ̄ϵ The key point is that this criterion depends on the average curvature λ̄ϵ , not on the localized pothole sharpness λϵ . A sharper localized component increases the perceived sharpness µϵ , but the SAM kick and contraction scale increase together, producing the cancellation in (17). In a wide basin, small λ̄i is paired with a large basin radius, so the same envelope can remain inside the basin. In a pothole, the basin radius is small while the average direction remains flat, so the same envelope can exceed the local basin radius. B.6
Vanilla SAM Comparison
For vanilla SAM, U = I. In an eigendirection of H, the scalar recursion is zt+1 = (1 − ηλi )zt − ηρλi sign(zt ).
(19)
The same envelope calculation gives ηρλi = ρ. (20) ηλi Thus vanilla SAM has a Euclidean hovering envelope p of order ρ in every direction, whereas LLQR+SAM has the direction-dependent envelope ρ/ λ̄i . 18
Corollary B.2 (Pothole-escape amplification). Under Assumption 4.1 and the commuting hypothesis, the LLQR+SAM hovering envelope around a pothole minimum is 1 ρ p = p × vanilla-SAM hovering envelope . λ̄ϵ λ̄ϵ The amplification depends only on the average curvature λ̄ϵ in the pothole direction, not on the localized pothole sharpness λϵ . This amplification cannot be reproduced selectively by retuning vanilla SAM. Choosing ρ ρSAM = p λ̄ϵ would match the pothole-direction envelope, but it would enlarge the probe in all directions, including non-pothole directions where this extra scale is not needed. The contraction behavior also differs: vanilla SAM contracts at rate 1 − ηλi , so its stable step size is constrained by the largest curvature of H. LLQR+SAM contracts at rate 1 − ηµi , which removes the conditioning of the average geometry H̄ from this constraint. B.7
Stochastic Selection Between Wells
The deterministic calculation above gives a local non-localization criterion: when the hovering envelope is comparable to the radius of a sharp well, the SAM kick prevents long residence near its bottom. With stochastic gradients, this becomes a selection effect. Noise may inject the iterate into a sharp well, but the SAM-induced hovering ejects it quickly; wide flat wells support longer visits because their basin radius is large relative to the hovering scale. The following regenerative model formalizes this interpretation. It is not intended as a full global model of training dynamics; it isolates the residence-time consequence of the local escape criterion. Proposition B.3 (Regenerative selection between wells). Consider a regenerative idealization with wells indexed by m. At each cycle, a well Jn is sampled from a distribution ν with positive mass on each well; the process starts near the well center, evolves under local gradient-noise dynamics with noise scale σ, and exits when ∥et ∥ ≥ R. Suppose flat wells satisfy (m)
2 −2 E[τR ] ≥ cflat , m R σ
whereas sharp unstable wells satisfy (m) sharp E[τR ] ≤ Cm 1 + log(R/σ) . If at least one flat well is present, then the long-run occupation mass of sharp wells satisfies 2 X σ log(R/σ) µm (σ) = O −−−→ 0. σ→0 R2 m∈Msharp
Proof. By the renewal-reward formula, (m)
νm E[τR ] µm (σ) = P . (ℓ) ℓ νℓ E[τR ] The total sharp-well numerator is O(1+log(R/σ)), while the denominator is at least the contribution of one flat well, namely Ω(R2 σ −2 ). Dividing gives the stated bound. For a pothole direction, the relevant dimensionless escape ratio is ρ p . rϵ λ̄ϵ Corollary B.2 increases this ratio for LLQR +SAM relative to vanilla SAM, predicting shorter visits to sharp wells and faster selection of the wide flat basin. The next calculation supports the path-length behavior observed in the noisy toy experiment. It states that, in a stable scalar noisy mode, increasing the effective metric denominator damps both stationary variance and one-step motion. 19
Proposition B.4 (Metric damping under gradient noise). Consider the scalar noisy recursion B −1 zt+1 = qiB ztB − η(dB εi,t , i )
−1 qiB := 1 − ηλi (dB , i )
Var(εi,t ) = τi2 ,
where B denotes a metric choice and dB i > ηλi /2. Then VarπB (zi ) =
ητi2 , λi (2dB i − ηλi )
B B 2 EπB (zi,t+1 − zi,t ) =
2η 2 τi2 . B dB i (2di − ηλi )
Q Consequently, if dP i ≥ di > ηλi /2, then metric P has no larger stationary variance or one-step Q motion than metric Q, with strict improvement when τi2 > 0 and dP i > di .
Proof. The AR(1) variance equation gives ViB = (qiB )2 ViB + Since 1 − (qiB )2 = we obtain ViB = Moreover,
η 2 τi2 . 2 (dB i )
ηλi (2dB i − ηλi ) , 2 (dB i )
ητi2 . λi (2dB i − ηλi )
B B −1 B −1 zi,t+1 − zi,t = −ηλi (dB zi,t − η(dB εi,t . i ) i )
B At stationarity, zi,t is independent of the fresh noise, so the cross term vanishes and
B B 2 EπB (zi,t+1 − zi,t ) =
2η 2 τi2 η2 2 B 2 λ V + τ = . i i i B 2 (dB dB i ) i (2di − ηλi )
Both denominators are increasing for dB i > ηλi /2, which proves the comparison.
20
C
Experiments Details
C.1
Standard Supervised Training Benchmarks
CIFAR protocol. For CIFAR-10 and CIFAR-100, all methods use the same supervised classification protocol: a 200-epoch cosine schedule with initial learning rate 0.05, Polyak momentum 0.9, batch size 128, random crop and horizontal flip augmentation, normalization, and Cutout regularization [Devries and Taylor, 2017]. Weight decay is selected in {10−4 , 5 × 10−4 , 10−3 } based on baseline best performance. SAM uses the recommended perturbation radius ρ = 0.1 on CIFAR-10 and ρ = 0.2 on CIFAR-100 [Foret et al., 2021, Li and Giannakis, 2023, Mi et al., 2022, Li et al., 2024]; FisherSAM uses its recommended ρ = 0.1 and inverse-Fisher regularization η = 0.1 [Kim et al., 2022]. The LLQR rows keep the outer SGDM recipe fixed and use the recommended inner preconditioner-learning settings of Dufort-Labbé et al. [2026]: inner batch size 128, preconditioner learning rate 10−3 , 50 inner momentum-solver steps, and update period 500 optimizer steps. The reported NGD rows use the NGD-induced divergence with EMA 0.95, while the Newton rows use the Newton-induced divergence with EMA 0.9. LLQR+SAM reuses the same hyperparameters as its LLQR and SAM components, without additional method-specific tuning. Each CIFAR entry is the best test accuracy observed during training, aggregated over five independent seeds. All CIFAR experiments were run on NVIDIA L40S GPUs. TinyImageNet protocol. TinyImageNet-200 uses the same training and reporting protocol and optimizer hyperparameters as the CIFAR benchmarks, with the comparison restricted in the main text to ResNet-18 and WRN-28-10. Entries in Table 3 are best top-1 test accuracy aggregated over three seeds. All TinyImageNet experiments were run on NVIDIA L40S GPUs. SAM-variant synergy protocol. The CIFAR SAM-variant comparison pairs LLQR with FSAM [Li et al., 2024]. F-SAM uses the recommended exponential moving average of gradient accumulation: λ = 0.9 for WRN-28-10 and PyramidNet-110, and λ = 0.6 for the other architectures. The comparison otherwise follows the CIFAR protocol above and reports best top-1 test accuracy over five seeds. All experiments were run on NVIDIA L40S GPUs. ImageNet protocol. The ImageNet comparison uses ResNet-50 on ImageNet [Deng et al., 2009]. Relative to the preceding vision experiments, the only hyperparameter changes are ρ = 0.075 and λ = 0.95, following the recommendation of Li et al. [2024]. The LLQR rows use the E-KFAC block structure. All ImageNet experiments were run on NVIDIA A100 GPUs. IWSLT14 German-to-English protocol. The sequence-modeling benchmark uses the fairseq IWSLT14 German-to-English Transformer training recipe of Ott et al. [2018]. The SAM perturbation radius is tuned over [0.0005, 0.2], with ρ = 0.005 giving the best validation result. SAM pertubation only starts being applied after learning rate warmup phase for stability. LLQR hyperparameters are kept unchanged from the vision experiments except for the preconditioner EMA decay, which is set to 0.925. Table 6 reports best validation BLEU aggregated over five seeds. All IWSLT14 experiments were run on NVIDIA L40S GPUs.
21
Table 8: Robustness to symmetric training-label corruption on CIFAR-100. Entries are best test accuracy on clean test labels (mean ± std, 3 seeds). The corruption rate γ denotes the fraction of randomly flipped training labels. Blue denotes NGD-induced LLQR geometry. Base geometry
Sharpness baseline
Combined geometry and sharpness
Dataset
γ
SGDM
A DAM
LLQR
SAM
LLQR+SAM
CIFAR-100 CIFAR-100 CIFAR-100 CIFAR-100
0.2 0.6 0.7 0.8
63.41 ± 0.46 41.21 ± 1.27 32.52 ± 0.72 22.97 ± 0.86
41.01 ± 0.63 9.76 ± 1.46 4.16 ± 0.52 2.23 ± 0.12
64.32 ± 0.22 41.21 ± 1.13 34.51 ± 0.71 24.54 ± 0.46
67.15 ± 0.43 46.81 ± 1.46 38.96 ± 0.72 31.19 ± 1.74
68.64 ± 0.36 48.54 ± 0.81 43.27 ± 0.82 33.06 ± 1.22
D
Additional Experimental Results
D.1
Noise injection robustness
Since SAM is known to improve robustness under label noise, we ask whether learned geometry strengthens this effect. On CIFAR-100, we corrupt only the training labels with symmetric random flips at rate γ, keep the test set clean, and reuse the standard supervised training recipe. Table 8 reports the best clean-test accuracy across γ ∈ {0.2, 0.6, 0.7, 0.8}. LLQR+SAM improves over SAM at every corruption level, with especially clear gains at moderate-to-high noise, showing that transporting the sharpness-aware update through the learned LLQR geometry strengthens robustness under corrupted supervision. The gains remain positive even at γ = 0.8, where supervision is nearly degenerate, the learned-geometry sharpness correction remains beneficial. D.2
Scalability
Our LLQR implementation attaches a blockwise preconditioner to network layers. The storage footprint is therefore determined by the chosen parameterization of each block rather than by LLQR as a framework. Richer blocks can encode more geometry but require more optimizer state, while simpler parameterizations reduce memory at the cost of expressivity. Table 9 summarizes the resulting fp32 storage across ViT scales for the preconditioner families considered in this scalability study. For comparison, we also include a hypothetical diag-kfac structure, which approximates each Kronecker factor by its diagonal: Diag(a) X Diag(b) for a reshaped kernel X ∈ Rm×n . This reduces the per-kernel preconditioner storage from dense Kronecker-style scaling to O(m + n) scalars, providing a lightweight option for trading geometric expressivity for memory. We also provide cadence results in Table 10 and Fig. 5, to be read as wall-clock overhead measurements, not as accuracy or convergence comparisons. As expected, the cost decreases when LLQR updates are performed less frequently, since fewer preconditioner refreshes are executed per epoch. This steady-update scaling trend is further exposed in Fig. 6.
22
Table 9: Cross-model analytical fp32 storage for e-KFAC and diag-KFAC preconditioners under ImageNet-scale ViT geometry. Model counts include heads and classification layers; #precond reports the number of stored preconditioner scalars for each parameterization.
Model
Storage
S
#model params
#precond params
ViT-Ti/16 ViT-Ti/16 ViT-Ti/16 ViT-Ti/16
e-kfac e-kfac diag-kfac diag-kfac
1 4 1 4
5,427,080 5,427,080 5,427,080 5,427,080
24,235,082 24,235,082 73,426 73,426
92.45 23.11 0.28 0.07
113.15 43.82 20.98 20.77
ViT-B/16 ViT-B/16 ViT-B/16 ViT-B/16
e-kfac e-kfac diag-kfac diag-kfac
1 4 1 4
86,567,656 86,567,656 86,567,656 86,567,656
387,826,882 387,826,882 294,038 294,038
1,479.44 369.86 1.12 0.28
1,809.67 700.09 331.35 330.51
ViT-L/16 ViT-L/16 ViT-L/16 ViT-L/16
e-kfac e-kfac diag-kfac diag-kfac
1 4 1 4
304,326,632 1,367,114,178 304,326,632 1,367,114,178 304,326,632 771,990 304,326,632 771,990
5,215.13 1,303.78 2.94 0.74
6,376.04 2,464.70 1,163.86 1,161.65
precond total (MiB/dev.) (MiB/dev.)
Table 10: Measured LLQR cadence details for ViT-B/16 and ViT-L/16. Times are compile-adjusted seconds per epoch, excluding the first preconditioner-update overhead. Baseline rows report firstorder optimizer runtimes; LLQR rows additionally report ratios relative to SGD and SAM on the same architecture. Arch
Method
Update interval
Updates
Epoch time
vs. SGD
vs. SAM
Steady update
ViT-B/16 ViT-B/16 ViT-B/16
SGD AdamW SAM
– – –
– – –
33.20 38.51 93.84
1.00× 1.16× 2.83×
0.35× 0.41× 1.00×
– – –
ViT-B/16 ViT-B/16 ViT-B/16 ViT-B/16 ViT-B/16 ViT-B/16 ViT-B/16 ViT-B/16 ViT-B/16 ViT-B/16 ViT-B/16 ViT-B/16
LLQR LLQR LLQR LLQR LLQR LLQR LLQR+SAM LLQR+SAM LLQR+SAM LLQR+SAM LLQR+SAM LLQR+SAM
50 100 200 500 1000 1500 50 100 200 500 1000 1500
32 16 8 4 2 2 32 16 8 4 2 2
1007.31 408.09 204.11 118.97 143.95 91.24 813.98 390.59 426.81 155.50 173.25 112.65
30.34× 12.29× 6.15× 3.58× 4.34× 2.75× 24.51× 11.76× 12.85× 4.68× 5.22× 3.39×
10.73× 4.35× 2.18× 1.27× 1.53× 0.97× 8.67× 4.16× 4.55× 1.66× 1.85× 1.20×
59.38 44.25 39.38 38.85 86.55 42.52 46.05 39.91 85.33 40.58 87.41 41.63
ViT-L/16 ViT-L/16 ViT-L/16
SGD AdamW SAM
– – –
– – –
62.59 67.84 146.54
1.00× 1.08× 2.34×
0.43× 0.46× 1.00×
– – –
ViT-L/16 ViT-L/16 ViT-L/16 ViT-L/16 ViT-L/16 ViT-L/16 ViT-L/16 ViT-L/16 ViT-L/16 ViT-L/16 ViT-L/16 ViT-L/16
LLQR LLQR LLQR LLQR LLQR LLQR LLQR+SAM LLQR+SAM LLQR+SAM LLQR+SAM LLQR+SAM LLQR+SAM
50 100 200 500 1000 1500 50 100 200 500 1000 1500
32 16 8 4 2 2 32 16 8 4 2 2
1647.13 868.47 476.77 275.63 299.48 301.14 1809.84 914.70 667.77 470.89 238.29 372.23
26.32× 13.88× 7.62× 4.40× 4.79× 4.81× 28.92× 14.62× 10.67× 7.52× 3.81× 5.95×
11.24× 5.93× 3.25× 1.88× 2.04× 2.06× 12.35× 6.24× 4.56× 3.21× 1.63× 2.54×
97.46 97.34 97.52 95.18 96.46 95.91 95.78 96.36 96.20 96.57 96.11 95.50
23