技術要約：Polyak型のステップサイズを用いた適応的シャープネス認識最小化（Adaptive Sharpness-Aware Minimization）

問題提起

Sharpness-Aware Minimization (SAM) およびその変種である Unnormalized SAM (USAM) は、損失の平坦さ（sharpness）を最小化することでディープニューラルネットワーク（DNN）の汎化性能を向上させる最適化手法として注目を集めている。しかし、これらの手法には、学習率の選択に対して非常に敏感であるという重大な実用的制限がある。ほとんどの学習アルゴリズムと同様に、SAMの変種は、最適な性能を得るために広範なハイパーパラメータ・チューニングや、定義済みのスケジューラ（例：コサインアニーリング）を必要とするのが一般的である。このような手動チューニングへの依存は、コストが高く、時間も要する。

著者らは、手動での学習率調整を必要としないSAM変種を設計するという課題に取り組んでいる。彼らは、特に、最近のStochastic Gradient Descent (SGD) において非常に効果的であることが示された「確率的Polyakステップサイズ（stochastic Polyak step sizes）」が、シャープネス認識最適化の設定において適用可能であるかどうかを調査している。

手法

本論文は、古典的なPolyakステップサイズから派生した、USAM（およびその拡張としてのSAM）のための新しいクラスの適応的スケジューラを提案している。コアとなる手法は以下の通りである：

1. 適応的ステップサイズの導出:

   • 決定論的な設定: USAMの更新規則から出発し、最適点への二乗距離（$\|x_{t+1} - x^*\|^2$）の上界を最小化することにより、閉形式のステップサイズを導出する。これにより、摂動を加えた点における損失 $f(e_t)$、最適値 $f^*$、摂動を加えた点の勾配、およびシャープネス半径 $\rho_t$ に依存するステップサイズ $\gamma_t$ が得られる。
   • 確率的な設定: この決定論的な導出を、有限和の設定（ミニバッチ）へと拡張する。提案される Stochastic Polyak Scheduler は、フルバッチの量をミニバッチの推定値（$f_{S_t}, \nabla f_{S_t}$）に置き換え、安定性を確保するために下限 $\ell^*_{S_t}$（非負の損失に対しては通常0）とキャップ $\gamma_b$ を含んでいる。
   • 正規化: 主要な理論的展開はUSAM（非正規化）に焦点を当てているが、著者らは同じロジックを正規化されたSAMの更新にも適応させている。ただし、USAMにおける理論的な非負性の保証は正規化されたケースには直接転用できないため、非負性を保証するためのセーフガード（ReLU）を保持している。

2. 理論的枠組み:

   • 分析は、標準的な滑らかさ（$L$-smooth）および凸性（または $\mu$-強凸性）の条件を仮定している。
   • 著者らは、シャープネス半径 $\rho_t$ が十分に小さい場合（具体的には $\rho_t \leq 1/L$）、導出されたステップサイズは自然に非負となり、ReLUによるセーフガードは冗長であることを示している。
   • この枠組みは、標準的な勾配法とシャープネス認識型の手法を統一するものである。$\rho_t = 0$ と設定することで、GDにおける古典的なPolyakステップサイズとその確率的類似物（SPS$_{max}$）が回収される。

主な貢献

1. 新規な適応的スケジューラ: 本論文は、SAM型の更新に特化した初のPolyak型のステップサイズを導入している。これらのスケジューラは完全に適応的であり、手動の学習率チューニングを必要とせず、現在のイテレーションで利用可能な量（摂動を加えた点の損失と勾配）のみに依存する。
2. 決定論的な収束保証:
   • $\mu$-強凸な目的関数に対して、著者らは二乗距離（$\|x_t - x^*\|^2$）における線形収束を証明している。その収束率は、GDにおけるPolyakステップサイズの既知のレートと一致し、収束係数はシャープネス半径の影響を明示的に定量化している。
   • 凸な目的関数に対して、関数値に関する劣線形な $O(1/T)$ の収束率を証明している。
   • 非増加の半径スケジュール（$\rho_t \to 0$）を用いることで、この手法が勾配ノルムをゼロに導くことを示している。
3. 確率的な収束保証:
   • 確率的な設定において、本手法は最適点の近傍に収束する。この近傍の大きさは、ミニバッチの下限に関連する分散の尺度（$\sigma^2$）によって定量化される。
   • 極めて重要な点として、これらの保証は、従来の適応的SAMの分析で一般的に必要とされる追加の仮定（有界な勾配、有界な分散、または成長条件など）を必要とせずに成立する。
   • 補間レジーム（過剰パラメータ化されたモデルで一般的な $\sigma^2 = 0$ の場合）では、この近傍は消失し、手法は最適点へ線形に収束する。
   • また、本分析は、定数ステップサイズのUSAMにおける確率的設定での収束保証も提供しており、それが先行研究よりも厳密に大きなステップサイズを許容することを示している。
4. 実証的検証:
   • 合成実験: リッジ回帰問題において、提案されたスケジューラは予測された線形収束の傾向を示し、チューニングされた定数ステップのベースラインや他の最新の適応的SAM変種（AdaSAM, LightSAMなど）と同等またはそれ以上の性能を発揮した。
   • ディープラーニング・ベンチマーク: ResNetアーキテクチャを用いたCIFAR-10およびCIFAR-100における実験により、提案されたStochastic Polyak Schedulerが、注意深くチューニングされた定数学習率やコサインアニーリングと同等、あるいはそれ以上のテスト精度を達成することが示された。特筆すべきは、提案されたスケジューラはシャープネス半径 $\rho$ が増加しても堅牢な性能を維持するのに対し、コサインアニーリングは大幅な精度の低下を招く点である。

結果と意義

本論文は、提案されたPolyakスケジューラが、SAM型の手法における学習率チューニングの負担を軽減するための原理的なアプローチを提供すると主張している。

• パフォーマンス: 数値実験は、提案された適応的スケジューラが、ハイパーパラメータ探索の必要性を大幅に減らしつつ、チューニングされたベースラインと同等以上の性能を達成することを示している。
• 堅牢性: 本手法は、標準的なSAMの実装において注意深いチューニングを必要とするパラメータであるシャープネス半径 $\rho$ の選択に対して、堅牢であることが示されている。
• 理論的進展: 本研究は、Polyak型のステップサイズに基づく適応的USAMの初の収束理論を提供するものである。これは、SGDにおけるPolyakステップサイズの理論的理解と、シャープネス認識最適化への適用との間の溝を埋めるものであり、既存の適応的SAMの文献よりも強力な（仮定が少ない）保証を提供している。

著者らは、理論は凸設定に焦点を当てているものの、ディープラーニングのベンチマーク（非凸）における実用的な成功は、幅広い適用可能性を示唆していると結論付けている。また、非凸な目的関数（例：PL条件）への分析の拡張や、シャープネス半径 $\rho_t$ をスケジューリングするための原理的な戦略の開発が、重要な将来の方向性であると特定している。