Technische Samenvatting: Adaptieve Sharpness-Aware Minimization met een Polyak-type Stapgrootte

Probleemstelling

Sharpness-Aware Minimization (SAM) en varianten daarvan, zoals Unnormalized SAM (USAM), zijn prominente optimizers geworden voor het verbeteren van de generalisatie van Deep Neural Networks (DNN's) door de scherpte (sharpness) van het verlieslandschap te minimaliseren. Echter, een aanzienlijke praktische beperking van deze methoden is hun hoge gevoeligheid voor de keuze van de leersnelheid (learning rate). Net als de meeste trainingsalgoritmen vereisen SAM-varianten doorgaans uitgebreide hyperparameter-afstemming of het gebruik van vooraf gedefinieerde schedulers (bijv. cosine annealing) om optimale prestaties te behalen. Deze afhankelijkheid van handmatige afstemming is kostbaar en tijdrovend.

De auteurs pakken de uitdaging aan om SAM-varianten te ontwerpen die geen handmatige afstemming van de leersnelheid vereisen. Ze onderzoeken specifiek de toepasbaarheid van stochastische Polyak-stapgroottes — recentelijk aangetoond zeer effectief te zijn voor Stochastic Gradient Descent (SGD) — binnen de context van sharpness-aware optimalisatie.

Methodologie

Het artikel stelt een nieuwe klasse van adaptieve schedulers voor USAM (en daarmee USAM) voor, afgeleid van de klassieke Polyak-stapgrootte. De kernmethodologie omvat:

1. Afleiding van Adaptieve Stapgroottes:

   • Deterministische Setting: Vertrekkend vanuit de USAM-update regel, leiden de auteurs een gesloten vorm voor de stapgrootte af door een bovengrens te minimaliseren op de gekwadrateerde afstand tot het optimum ($\|x_{t+1} - x^*\|^2$). Dit levert een stapgrootte $\gamma_t$ op die afhankelijk is van de loss bij het geperturbeerde punt $f(e_t)$, de optimale waarde $f^*$, de gradiënt bij het geperturbeerde punt, en de scherptestraal (sharpness radius) $\rho_t$.
   • Stochastische Setting: De deterministische afleiding wordt uitgebreid naar de finite-sum setting (mini-batches). De voorgestelde Stochastic Polyak Scheduler vervangt de full-batch grootheden door mini-batch schattingen ($f_{S_t}$, $\nabla f_{S_t}$) en bevat een ondergrens $\ell^*_{S_t}$ (vaak 0 voor niet-negatieve losses) en een cap $\gamma_b$ om stabiliteit te waarborgen.
   • Normalisatie: Hoewel de primaire theoretische ontwikkeling zich richt op USAM (unnormalized), passen de auteurs dezelfde logica toe op de genormaliseerde SAM-update, waarbij een veiligheidsmaatregel (ReLU) wordt behouden om niet-negativiteit te garanderen, aangezien de theoretische garantie voor niet-negativiteit voor USAM niet direct overdraagbaar is naar het genormaliseerde geval.

2. Theoretisch Kader:

   • De analyse gaat uit van standaard smoothnes ($L$-smooth) en convexiteit (of $\mu$-sterke convexiteit) condities.
   • De auteurs stellen vast dat wanneer de scherptestraal $\rho_t$ voldoende klein is (specifiek $\rho_t \leq 1/L$), de afgeleide stapgroottes van nature niet-negatief zijn, waardoor de ReLU-veiligheidsmaatregel overbodig is in het deterministische convexe regime.
   • Het framework verenigt standaard gradiëntmethoden en sharpness-aware methoden: het instellen van $\rho_t = 0$ herstelt de klassieke Polyak-stapgrootte voor GD en de stochastische tegenhanger (SPS$_{max}$).

Belangrijkste Bijdragen

1. Nieuwe Adaptieve Schedulers: Het artikel introduceert de eerste Polyak-type stapgroottes die specifiek zijn afgestemd op SAM-stijl updates. Deze schedulers zijn volledig adaptief, vereisen geen handmatige leersnelheid-afstemming en vertrouwen enkel op grootheden die beschikbaar zijn bij de huidige iteratie (loss en gradiënt bij het geperturbeerde punt).
2. Deterministische Convergentiegaranties:
   • Voor $\mu$-sterk convexe objectieven bewijzen de auteurs lineaire convergentie in de gekwadrateerde afstand ($\|x_t - x^*\|^2$). De snelheid komt overeen met de bekende snelheid voor GD met de Polyak-stapgrootte, waarbij een contractiefactor expliciet het effect van de scherptestraal kwantificeert.
   • Voor convexe objectieven bewijzen ze een sublineaire $O(1/T)$ convergentiesnelheid voor de functiewaarde.
   • Ze tonen aan dat met een niet-toenemend straalschema ($\rho_t \to 0$), de methode de gradiëntnorm naar nul drijft.
3. Stochastische Convergentiegaranties:
   • In de stochastische setting convergeert de methode naar een buurt van het optimum. De omvang van deze buurt wordt gekwantificeerd door een variantie-maat ($\sigma^2$) gerelateerd aan de mini-batch ondergrenzen.
   • Cruciaal is dat deze garanties standhouden zonder aanvullende aannames die vaak vereist zijn in eerdere adaptieve SAM-analyses, zoals begrensde gradiënten, begrensde variantie of groeicondities.
   • In geïnterpoleerde regimes (waar $\sigma^2 = 0$, gebruikelijk in over-geparameteriseerde modellen), verdwijnt de buurt en convergeert de methode lineair naar het optimum.
   • De analyse levert ook een convergentiegarantie op voor constante stapgrootte USAM in het stochastische regime, waarbij wordt aangetoond dat dit strikt grotere stapgroottes toestaat dan eerdere werken.
4. Empirische Validatie:
   • Synthetische Experimenten: Op ridge regressie problemen vertonen de voorgestelde schedulers de voorspelde lineaire convergentietrends en presteren ze gelijk aan of beter dan getunede constante-stap baselines en andere recente adaptieve SAM-varianten (bijv. AdaSAM, LightSAM).
   • Deep Learning Benchmarks: Experimenten op CIFAR-10 en CIFAR-100 met ResNet-architecturen laten zien dat de Stochastic Polyak Scheduler een testnauwkeurigheid bereikt die vergelijkbaar is met of beter is dan zorgvuldig afgestemde constante leersnelheden en cosine annealing. Opvallend is dat de voorgestelde scheduler robuust presteert naarmate de scherptestraal $\rho$ toeneemt, terwijl cosine annealing significante nauwkeurigheidsdegradatie vertoont.

Resultaten en Betekenis

Het artikel beweert dat de voorgestelde Polyak-schedulers een geprincieerd aanpak bieden om de last van het afstemmen van de leersnelheid bij SAM-type methoden te verminderen.

• Prestaties: Numerieke experimenten wijzen uit dat de adaptieve schedulers prestaties leveren die vergelijkbaar zijn met of beter zijn dan getunede baselines, terwijl ze de noodzaak voor hyperparameter-zoekopdrachten aanzienlijk verminderen.
• Robuustheid: De methode blijkt robuust tegen de keuze van de scherptestraal $\rho$, een parameter die vaak zorgvuldige afstemming vereist in standaard SAM-implementaties.
• Theoretische Vooruitgang: Het werk biedt de eerste convergentietheorie voor adaptieve USAM gebaseerd op Polyak-type stapgroottes. Het overbrugt de kloof tussen het theoretische begrip van Polyak-stapgroottes in SGD en hun toepassing op sharpness-aware optimalisatie, waarbij garanties worden geboden die sterker zijn (minder aannames) dan de bestaande adaptieve SAM-literatuur.

De auteurs concluderen dat hoewel hun theorie zich richt op convexe settings, het praktische succes op deep learning benchmarks (niet-convex) wijst op brede toepasbaarheid. Ze identificeren het uitbreiden van de analyse naar niet-convexe objectieven (bijv. de PL-conditie) en het ontwikkelen van geprincipeerde strategieën voor het plannen van de scherptestraal $\rho_t$ als belangrijke toekomstige richtingen.