Technical Summary: Adaptive Sharpness-Aware Minimization with a Polyak-type Step size

Problem Statement

Sharpness-Aware Minimization (SAM) and its variants, such as Unnormalized SAM (USAM), have become prominent optimizers for improving the generalization of Deep Neural Networks (DNNs) by minimizing the sharpness of the loss landscape. However, a significant practical limitation of these methods is their high sensitivity to the choice of the learning rate. Like most training algorithms, SAM variants typically require extensive hyperparameter tuning or the use of predefined schedulers (e.g., cosine annealing) to achieve optimal performance. This reliance on manual tuning is costly and time-consuming.

The authors address the challenge of designing SAM variants that do not require manual learning rate tuning. They specifically investigate the applicability of stochastic Polyak step sizes—recently shown to be highly effective for Stochastic Gradient Descent (SGD)—to the sharpness-aware optimization setting.

Methodology

The paper proposes a new class of adaptive schedulers for USAM (and by extension, SAM) derived from the classical Polyak step size. The core methodology involves:

1. Derivation of Adaptive Step Sizes:

   • Deterministic Setting: Starting from the USAM update rule, the authors derive a closed-form step size by minimizing an upper bound on the squared distance to the optimum ($\|x_{t+1} - x^*\|^2$). This yields a step size $\gamma_t$ that depends on the loss at the perturbed point $f(e_t)$, the optimal value $f^*$, the gradient at the perturbed point, and the sharpness radius $\rho_t$.
   • Stochastic Setting: The deterministic derivation is extended to the finite-sum setting (mini-batches). The proposed Stochastic Polyak Scheduler replaces the full-batch quantities with mini-batch estimates ($f_{S_t}$, $\nabla f_{S_t}$) and includes a lower bound $\ell^*_{S_t}$ (often 0 for non-negative losses) and a cap $\gamma_b$ to ensure stability.
   • Normalization: While the primary theoretical development focuses on USAM (unnormalized), the authors adapt the same logic to the normalized SAM update, retaining a safeguard (ReLU) to ensure non-negativity, as the theoretical non-negativity guarantee for USAM does not directly transfer to the normalized case.

2. Theoretical Framework:

   • The analysis assumes standard smoothness ($L$-smooth) and convexity (or $\mu$-strong convexity) conditions.
   • The authors establish that when the sharpness radius $\rho_t$ is sufficiently small (specifically $\rho_t \leq 1/L$), the derived step sizes are naturally non-negative, rendering the ReLU safeguard redundant in the deterministic convex regime.
   • The framework unifies standard gradient methods and sharpness-aware methods: setting $\rho_t = 0$ recovers the classical Polyak step size for GD and its stochastic analogue (SPS$_{max}$).

Key Contributions

1. Novel Adaptive Schedulers: The paper introduces the first Polyak-type step sizes tailored specifically for SAM-style updates. These schedulers are fully adaptive, requiring no manual learning rate tuning, and rely only on quantities available at the current iteration (loss and gradient at the perturbed point).
2. Deterministic Convergence Guarantees:
   • For $\mu$-strongly convex objectives, the authors prove linear convergence in squared distance ($\|x_t - x^*\|^2$). The rate matches the known rate for GD with the Polyak step size, with a contraction factor explicitly quantifying the effect of the sharpness radius.
   • For convex objectives, they prove a sublinear $O(1/T)$ convergence rate for the function value.
   • They show that with a non-increasing radius schedule ($\rho_t \to 0$), the method drives the gradient norm to zero.
3. Stochastic Convergence Guarantees:
   • In the stochastic setting, the method converges to a neighborhood of the optimum. The size of this neighborhood is quantified by a variance measure ($\sigma^2$) related to the mini-batch lower bounds.
   • Crucially, these guarantees hold without additional assumptions commonly required in prior adaptive SAM analyses, such as bounded gradients, bounded variance, or growth conditions.
   • In interpolated regimes (where $\sigma^2 = 0$, common in over-parameterized models), the neighborhood vanishes, and the method converges linearly to the optimum.
   • The analysis also yields a convergence guarantee for constant step size USAM in the stochastic regime, showing it permits strictly larger step sizes than previous works.
4. Empirical Validation:
   • Synthetic Experiments: On ridge regression problems, the proposed schedulers exhibit the predicted linear convergence trends and outperform or match tuned constant-step baselines and other recent adaptive SAM variants (e.g., AdaSAM, LightSAM).
   • Deep Learning Benchmarks: Experiments on CIFAR-10 and CIFAR-100 using ResNet architectures demonstrate that the Stochastic Polyak Scheduler achieves test accuracy comparable to or better than carefully tuned constant learning rates and cosine annealing. Notably, the proposed scheduler maintains robust performance as the sharpness radius $\rho$ increases, whereas cosine annealing suffers significant accuracy degradation.

Results and Significance

The paper claims that the proposed Polyak schedulers offer a principled approach to reducing the learning rate tuning burden in SAM-type methods.

• Performance: Numerical experiments indicate that the adaptive schedulers achieve performance comparable to or better than tuned baselines while substantially reducing the need for hyperparameter search.
• Robustness: The method is shown to be robust to the choice of the sharpness radius $\rho$, a parameter that often requires careful tuning in standard SAM implementations.
• Theoretical Advancement: The work provides the first convergence theory for adaptive USAM based on Polyak-type step sizes. It bridges the gap between the theoretical understanding of Polyak step sizes in SGD and their application to sharpness-aware optimization, offering guarantees that are stronger (fewer assumptions) than existing adaptive SAM literature.

The authors conclude that while their theory focuses on convex settings, the practical success on deep learning benchmarks (non-convex) suggests broad applicability. They identify extending the analysis to non-convex objectives (e.g., PL condition) and developing principled strategies for scheduling the sharpness radius $\rho_t$ as important future directions.