Technisches Resümee: Adaptive Sharpness-Aware Minimization mit einer Polyak-typischen Schrittweite

Problemstellung

Sharpness-Aware Minimization (SAM) und deren Varianten, wie etwa Unnormalized SAM (USAM), sind zu bedeutenden Optimierern für die Verbesserung der Generalisierung von Deep Neural Networks (DNNs) geworden, indem sie die Schärfe (Sharpness) der Verlustlandschaft minimieren. Eine wesentliche praktische Einschränkung dieser Methoden ist jedoch ihre hohe Sensitivität gegenüber der Wahl der Lernrate. Wie die meisten Trainingsalgorithmen erfordern SAM-Varianten typischerweise eine umfangreiche Hyperparameter-Abstimmung oder die Verwendung vordefinierter Scheduler (z. B. Cosine Annealing), um eine optimale Leistung zu erzielen. Diese Abhängigkeit von manueller Abstimmung ist kostspielig und zeitaufwendig.

Die Autoren adressieren die Herausforderung, SAM-Varianten zu entwerfen, die keine manuelle Lernraten-Abstimmung erfordern. Sie untersuchen spezifisch die Anwendbarkeit stochastischer Polyak-Schrittweiten – die kürzlich gezeigt wurde, dass sie für den Stochastic Gradient Descent (SGD) hocheffektiv sind – im Kontext der schärfe-bewussten Optimierung.

Methodik

Das Paper schlägt eine neue Klasse von adaptiven Schedulern für USAM (und im weiteren Sinne für SAM) vor, die aus der klassischen Polyak-Schrittweite abgeleitet sind. Die Kernmethodik umfasst:

1. Ableitung adaptiver Schrittweiten:

   • Deterministische Einstellung: Ausgehend von der USAM-Update-Regel leiten die Autoren eine geschlossene Form der Schrittweite her, indem sie eine obere Schranke für den quadratischen Abstand zum Optimum ($\|x_{t+1} - x^*\|^2$) minimieren. Dies ergibt eine Schrittweite $\gamma_t$, die vom Verlust am perturbierten Punkt $f(e_t)$, dem optimalen Wert $f^*$, dem Gradienten am perturbierten Punkt und dem Schärferadius $\rho_t$ abhängt.
   • Stochastische Einstellung: Die deterministische Ableitung wird auf das Finite-Sum-Setting (Mini-Batches) erweitert. Der vorgeschlagene Stochastic Polyak Scheduler ersetzt die Full-Batch-Größen durch Mini-Batch-Schätzungen ($f_{S_t}$, $\nabla f_{S_t}$) und enthält eine untere Schranke $\ell^*_{S_t}$ (oft 0 für nicht-negative Verluste) sowie eine Begrenzung $\gamma_b$, um Stabilität zu gewährleisten.
   • Normalisierung: Während die primäre theoretische Entwicklung auf USAM (unnormalisiert) fokussiert, passen die Autoren dieselbe Logik an das normalisierte SAM-Update an und behalten eine Sicherheitsmaßnahme (ReLU) bei, da die theoretische Garantie der Nicht-Negativität für USAM nicht direkt auf den normalisierten Fall übertragbar ist.

2. Theoretischer Rahmen:

   • Die Analyse setzt Standardbedingungen für Glattheit ($L$-smooth) und Konvexität (oder $\mu$-starke Konvexität) voraus.
   • Die Autoren etablieren, dass die abgeleiteten Schrittweiten eine ausreichend kleine Schärferadius $\rho_t$ (speziell $\rho_t \leq 1/L$) voraussetzen, wodurch die ReLU-Sicherheitsmaßnahme im deterministischen konvexen Regime redundant wird.
   • Das Framework vereinheitlicht Standard-Gradientenverfahren und schärfe-bewusste Methoden: Setzt man $\rho_t = 0$, erhält man die klassische Polyak-Schrittweite für GD und deren stochastisches Analogon (SPS$_{max}$).

Zentrale Beiträge

1. Neuartige adaptive Scheduler: Das Paper führt die ersten Polyak-typischen Schrittweiten ein, die speziell auf SAM-artige Updates zugeschnitten sind. Diese Scheduler sind voll adaptiv, erfordern keine manuelle Lernraten-Abstimmung und verlassen sich nur auf Größen, die in der aktuellen Iteration verfügbar sind (Verlust und Gradient am perturbierten Punkt).
2. Deterministische Konvergenzgarantien:
   • Für $\mu$-stark konvexe Zielfunktionen beweisen die Autoren eine lineare Konvergenz im quadratischen Abstand ($\|x_t - x^*\|^2$). Die Rate entspricht der bekannten Rate für GD mit der Polyak-Schrittweite, wobei ein Kontraktionsfaktor den Effekt des Schärferadius explizit quantifiziert.
   • Für konvexe Zielfunktionen beweisen sie eine sublineare $O(1/T)$ Konvergenzrate für den Funktionswert.
   • Sie zeigen, dass die Methode mit einem nicht-monoton fallenden Radius-Schema ($\rho_t \to 0$) die Gradientennorm gegen Null treibt.
3. Stochastische Konvergenzgarantien:
   • Im stochastischen Setting konvergiert die Methode in eine Umgebung des Optimums. Die Größe dieser Umgebung wird durch ein Varianzmaß ($\sigma^2$) quantifiziert, das mit den Mini-Batch-unteren Schranken zusammenhängt.
   • Entscheidend ist, dass diese Garantien ohne zusätzliche Annahmen gelten, die in früheren adaptiven SAM-Analysen erforderlich waren, wie etwa beschränkte Gradienten, beschränkte Varianz oder Wachstumsbedingungen.
   • In interpolierten Regimen (wo $\sigma^2 = 0$, häufig in überparametrisierten Modellen), verschwindet die Umgebung, und die Methode konvergiert linear zum Optimum.
   • Die Analyse liefert auch eine Konvergenzgarantie für konstante Schrittweiten-USAM im stochastischen Regime und zeigt, dass dies strikt größere Schrittweiten als bisherige Arbeiten erlaubt.
4. Empirische Validierung:
   • Synthetische Experimente: Bei Ridge-Regression-Problemen zeigen die vorgeschlagenen Scheduler die vorhergesagten linearen Konvergenstrends und übertreffen oder erreichen die abgestimmten Constant-Step-Baselines sowie andere jüngste adaptive SAM-Varianten (z. B. AdaSAM, LightSAM).
   • Deep Learning Benchmarks: Experimente auf CIFAR-10 und CIFAR-100 unter Verwendung von ResNet-Architekturen zeigen, dass der Stochastic Polyak Scheduler eine Testgenauigkeit erreicht, die vergleichbar mit oder besser als sorgfältig abgestimmte konstante Lernraten und Cosine Annealing ist. Bemerkenswert ist, dass der vorgeschlagene Scheduler eine robuste Performance beibehält, wenn der Schärferadius $\rho$ steigt, während Cosine Annealing eine signifikante Genauigkeitsdegradation erleidet.

Ergebnisse und Bedeutung

Das Paper behauptet, dass die vorgeschlagenen Polyak-Scheduler einen prinzipienbasierten Ansatz bieten, um den Aufwand der Lernraten-Abstimmung bei SAM-artigen Methoden zu reduzieren.

• Performance: Numerische Experimente deuten darauf hin, dass die adaptiven Scheduler eine vergleichbare oder bessere Performance als abgestimmte Baselines erzielen und gleichzeitig den Bedarf an Hyperparameter-Suche erheblich reduzieren.
• Robustheit: Die Methode erweist sich als robust gegenüber der Wahl des Schärferadius $\rho$, einem Parameter, der in Standard-SAM-Implementierungen oft eine sorgfältige Abstimmung erfordert.
• Theoretischer Fortschritt: Die Arbeit liefert die erste Konvergenztheorie für adaptive USAM basierend auf Polyak-typischen Schrittweiten. Sie schließt die Lücke zwischen dem theoretischen Verständnis von Polyak-Schrittweiten in SGD und ihrer Anwendung auf die schärfe-bewusste Optimierung und bietet Garantien, die stärker (weniger Annahmen erfordern) sind als die bestehende adaptive SAM-Literatur.

Die Autoren kommen zu dem Schluss, dass, obwohl sich ihre Theorie auf konvexe Settings konzentriert, der praktische Erfolg auf Deep-Learning-Benchmarks (nicht-konvex) auf eine breite Anwendbarkeit hindeutet. Sie identifizieren die Erweiterung der Analyse auf nicht-konvexe Zielfunktionen (z. B. PL-Bedingung) und die Entwicklung prinzipienbasierter Strategien für die Planung des Schärferadius $\rho_t$ als wichtige zukünftige Forschungsrichtungen.