Adaptive Sharpness-Aware Minimization with a Polyak-type Step size: A Theory-Grounded Scheduler

Sharpness-Aware Minimization (SAM) has established itself as a powerful and widely adopted optimizer for training machine learning models. By explicitly minimizing the sharpness of the loss landscape, SAM often improves generalization while delivering strong empirical performance. However, SAM and its variants, like most training algorithms, are sensitive to the choice of learning rate, which is typically selected through extensive hyperparameter tuning or predefined schedulers. In this work, motivated by recent advances on the effectiveness of stochastic Polyak step sizes for Stochastic Gradient Descent (SGD), we derive Polyak schedulers tailored to SAM-style updates, yielding novel adaptive algorithms in both deterministic and stochastic settings. In the smooth setting, we prove linear convergence for strongly convex objectives and an $\mathcal{O}(1/T)$ convergence rate for convex objectives in the deterministic case. In the stochastic setting, we establish analogous convergence guarantees up to a neighborhood of the optimum. Numerical experiments demonstrate that the proposed Polyak schedulers achieve performance comparable to or better than carefully tuned SAM baselines, while substantially reducing the need for learning-rate tuning.

Training a modern AI model is a bit like searching a vast, hilly landscape for its lowest valley. Where the search settles matters: a model that lands in a broad, gentle valley tends to perform well on data it has never seen, while one stuck in a narrow ravine often does not. A popular training method called Sharpness-Aware Minimization (SAM) deliberately steers toward these broad valleys. But it hinges on one delicate setting: the step size, or how far to move at each stage, and finding a good value usually demands expensive trial-and-error. We adapted a classic optimization idea, the Polyak step size, to SAM. The result is a simple formula that sets the step size automatically at every step, using only information the method already has on hand, which removes most of the guesswork. We also prove mathematically that the method reliably reaches a good solution, a guarantee that earlier automatic approaches for SAM could not offer without extra assumptions. On standard image-recognition benchmarks, our tuning-free method matches or beats carefully hand-tuned SAM. This makes it cheaper and easier to train models that hold up well on new, real-world data.