The Origin of Edge of Stability

Titel: The Origin of Edge of Stability

Auteur: Elon Litman
Institution: Stanford University

---

1. Het Probleem

In de klassieke optimalisatietheorie voor neurale netwerken wordt aangenomen dat de verliesfunctie monotoon daalt zolang de leersnelheid $\eta$ voldoet aan $\eta < 2/\lambda_{\max}$, waarbij $\lambda_{\max}$ de grootste eigenwaarde van de Hessian-matrix is (ook wel "sharpness" of scherpte genoemd).

Echter, empirische observaties (Cohen et al., 2021) tonen aan dat bij het gebruik van vaste leersnelheden in de praktijk het tegenovergestelde gebeurt:

1. De scherpte ($\lambda_{\max}$) stijgt tijdens het trainen totdat deze de drempelwaarde $2/\eta$ bereikt.
2. Bij deze waarde "saturatie" de scherpte en begint de trainingsfout te oscilleren op korte tijdschalen, terwijl deze op langere tijdschalen toch blijft dalen.
3. Dit fenomeen wordt de Edge of Stability (EoS) genoemd.

De bestaande uitdaging: Hoewel eerdere werken (zoals Damian et al., 2023; Agarwala et al., 2023) lokale mechanismen hebben aangetoond die zorgen voor zelfregulatie naast de rand, missen ze een verenigde verklaring voor waarom de trajectorie vanuit een willekeurige initialisatie gedwongen wordt naar deze specifieke drempel $2/\eta$ toe te bewegen. Er ontbreekt een globaal mechanisme dat de "aantrekkingskracht" van deze grens verklaart.

---

2. Methodologie: De Edge Coupling

De auteur introduceert een nieuwe theoretische benadering die niet de dynamiek stap voor stap analyseert, maar kijkt naar een scalair functionaal op opeenvolgende paren iteraties.

De Edge Coupling ($A_\eta$):
De kern van de methode is het definiëren van een functioneel genaamd de "edge coupling":
$$A_\eta(x, y) \triangleq L(x) + L(y) - \frac{1}{2\eta} \|x - y\|^2_2$$
Waarbij $L$ de verliesfunctie is en $x, y$ opeenvolgende iteraties zijn.

• Karakterisering: Als men de gradiënt van $A_\eta$ naar $x$ nul stelt, krijgt men exact de update-regel van gradient descent: $y = x - \eta \nabla L(x)$, mits de coefficient $1/\eta$ correct is gekozen.
• Interpretatie: Dit functioneel kan worden gezien als een discrete genererende functie uit de Hamiltoniaanse mechanica. De kritieke punten van dit functioneel corresponderen met de dynamica van gradient descent.

Analyse van Kritieke Punten:
Door beide partiële gradiënten ($\nabla_x A_\eta$ en $\nabla_y A_\eta$) gelijk te stellen aan nul, worden twee soorten oplossingen geïdentificeerd:

1. Vaste punten: Waar $x = y$ (conventionele convergentie).
2. Periode-twee banen: Waar $x \neq y$ (oscillaties).

De auteur reparametriseert het probleem door een "centrum" $m$ en een "half-amplitude" $a$ te definiëren ($x = m-a, y = m+a$). Dit leidt tot een gereduceerd functioneel $\Phi_\eta(a)$ dat alleen afhankelijk is van de amplitude van de oscillatie.

---

3. Belangrijkste Bijdragen

A. Unificatie van Propagatie en Afdaling

Het paper toont aan dat de kritieke voorwaarden van de edge coupling leiden tot twee fundamentele formules:

1. Stap-propagator: Een recurrente relatie voor de stapgrootte $d_k$:
   $$d_{k+1} = (I - \eta \bar{H}_k) d_k$$
   Hierbij is $\bar{H}_k$ een gemiddelde Hessian over het stapsegment. Stabiliteit vereist dat de eigenwaarden van $\bar{H}_k$ binnen $[0, 2/\eta]$ liggen.
2. Verliesveranderingsformule: Een exacte formule voor de verliesverandering per stap:
   $$L(w_{k+1}) - L(w_k) = -\frac{\|d_k\|^2}{2\eta} (2 - \eta \tilde{r}_k)$$
   Hierbij is $\tilde{r}_k$ een ander type Hessian-gemiddelde (driehoekig gewogen).

B. Het Globale Aantrekkingsmechanisme (The Forcing Theorem)

Dit is de belangrijkste theoretische doorbraak. De auteur bewijst dat de som van de verliesveranderingen over de hele trajectorie een "telescopische som" vormt:
$$\sum_{k=0}^{K-1} \|d_k\|^2 \left( \frac{2}{\eta} - \tilde{r}_k \right) = 2(L(w_0) - L(w_K))$$
Omdat het totale verlies $L$ begrensd is, kan de som van de afwijkingen $(\frac{2}{\eta} - \tilde{r}_k)$ niet willekeurig groot worden. Dit dwingt de effectieve kromming $\tilde{r}_k$ om zich te concentreren rond de waarde $2/\eta$.

• Conclusie: De dynamica wordt globaal gedwongen om de drempel $2/\eta$ te bezoeken. Dit lost het probleem op van eerdere werken die alleen lokale stabiliteit konden verklaren.

C. Exacte Lokalisatie (Mean Value Theorem)

Een veelgestelde vraag is of deze forcing geldt voor de ware Hessian-eigenwaarde op een specifiek punt, of alleen voor gemiddelden.

• De auteur gebruikt de Mean Value Theorem om aan te tonen dat zowel de uniforme als de driehoekig gewogen gemiddelde Hessian exact overeenkomen met de richtingskromming van de ware Hessian op een specifiek punt binnen het stapsegment.
• Hierdoor is de forcing naar $2/\eta$ exact en zonder enige "gap" of benaderingsfout.

D. Bifurcatie-analyse en Periode-twee Banen

De paper analyseert wanneer periode-twee banen (oscillaties) ontstaan.

• Door het probleem te reduceren tot een functie van de amplitude alleen, wordt de bifurcatie een niet-lineair eigenwaardeprobleem.
• Voor tweelaagse lineaire netwerken wordt bewezen dat de periodeverdubbeling continu optreedt aan de kant van de grote leersnelheid ($\eta > \eta_c$) en dat dit onafhankelijk is van de breedte van het netwerk (width-invariant).

---

4. Resultaten en Validatie

• Theoretische Validatie: De theorie voorspelt dat de effectieve kromming zich concentreert rond $2/\eta$ zodra de trainingsstappen groot genoeg zijn ($E_K \to \infty$).
• Experimentele Validatie:
  • Simulaties op een 3-laags MLP (CIFAR-10) tonen aan dat de effectieve kromming inderdaad convergeert naar $2/\eta$.
  • De voorspelde verliesverandering (gebaseerd op de edge coupling) komt nauwkeurig overeen met de daadwerkelijke verliesverandering per stap.
  • Bij een tweelaags lineair netwerk wordt de continuïteit van de periodeverdubbeling bij de kritieke leersnelheid $\eta_c$ bevestigd, met een amplitude die schaalt als $\sqrt{\eta - \eta_c}$.
• Stabiliteitsmechanismen: Het paper verklaart waarom het systeem niet divergeert:
  1. Als de kromming boven $2/\eta$ komt, groeit de stapgrootte exponentieel, wat het systeem dwingt terug te keren naar gebieden met lagere kromming.
  2. Binnen het stabiliteitsvenster zorgen oscillerende annuleringen (cancellation) voor secular drift, waardoor het systeem stabiel blijft rond de rand.

---

5. Betekenis en Impact

Dit paper biedt een verenigde en rigoureuze verklaring voor het Edge of Stability-fenomeen, een van de meest opvallende observaties in de moderne diepe leertheorie.

• Van Lokaal naar Globaal: Het overbrugt de kloof tussen lokale stabiliteitsanalyses en het globale gedrag van gradient descent. Het bewijst dat $2/\eta$ geen toevallige saturatie is, maar een fundamentele attractor die voortvloeit uit de structuur van de update-regels zelf.
• Nieuw Kader: De introductie van de "edge coupling" biedt een krachtig nieuw wiskundig kader (gebaseerd op variatierekening en discrete mechanica) om de dynamica van neurale netwerken te analyseren.
• Praktische Implicaties: Het inzicht dat de scherpte actief wordt gedwongen naar $2/\eta$ helpt bij het begrijpen van waarom grote leersnelheden vaak leiden tot betere generalisatie (door het vinden van "flattere" minima) en waarom training stabiel blijft ondanks het overschrijden van de klassieke stabiliteitsgrens.
• Uitbreidbaarheid: De theorie wordt uitgebreid naar mini-batch SGD, trajectoormparen en continue tijd, wat aangeeft dat het mechanisme robuust is.

Kortom, Litman demonstreert dat de Edge of Stability geen anomalie is, maar een inherente eigenschap van gradient descent die wordt gestuurd door een behoudswet (conservation law) in de verliesverandering, wat leidt tot een exacte forcing van de kromming naar de kritieke drempel $2/\eta$.