Navigating Potholes with Geometry-Aware Sharpness Minimization

2026-05-15

AI Summary

This work addresses a key limitation of traditional Sharpness-Aware Minimization (SAM), which neglects the geometric structure of the loss landscape and struggles to distinguish sharp local minima (“pits”) within otherwise flat regions from genuinely wide minima. To overcome this, the authors propose LLQR+SAM, the first method to integrate learned loss geometry with SAM. It employs a layer-wise, sparse preconditioner derived from the LLQR framework, updated via exponential moving averages on a slow timescale to capture low-resolution geometric information, while SAM operates on a fast timescale to probe local curvature using this geometric context. Theoretical analysis shows that the approach amplifies escape signals along pit directions while preserving stability in broad flat regions. Experiments across multiple vision and sequence modeling benchmarks demonstrate that LLQR+SAM significantly outperforms either SAM or LLQR alone, confirming the complementary benefits of slow geometric learning and rapid sharpness-aware correction.

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.

Problem

• sharpness-aware minimization
• loss landscape geometry
• flat minima
• parameter curvature
• optimization

Innovation

• sharpness-aware minimization
• preconditioning
• loss landscape geometry
• two-timescale optimization
• LLQR