The Origin of Edge of Stability

Titolo: The Origin of Edge of Stability (L'Origine dell'Edge of Stability)

Autore: Elon Litman (Stanford University)

---

1. Il Problema

Il paper affronta il fenomeno noto come Edge of Stability (EoS) nell'ottimizzazione delle reti neurali tramite discesa del gradiente (Gradient Descent - GD) su batch completi (full-batch).

• Contesto: La teoria classica dell'ottimizzazione garantisce che il gradiente discenda monotonicamente solo se il tasso di apprendimento $\eta$ soddisfa la condizione $\eta < 2/\lambda_{\max}$, dove $\lambda_{\max}$ è il più grande autovalore dell'Hessiana (definito come "sharpness" o nitidezza).
• Osservazione Empirica: Cohen et al. (2021) hanno scoperto che, in pratica, durante l'addestramento la sharpness aumenta fino a raggiungere esattamente il valore soglia $2/\eta$. A questo punto, il sistema non diverge, ma entra in una fase di "stabilità al limite": la sharpness si satura attorno a $2/\eta$, la perdita di addestramento inizia ad oscillare su scale temporali brevi, pur continuando a diminuire su scale temporali lunghe.
• Il Gap Teorico: Lavori precedenti (es. Damian et al., Agarwala et al.) hanno dimostrato che il sistema si "autoregola" vicino a questo limite, ma non hanno spiegato perché la traiettoria sia costretta a muoversi verso $2/\eta$ partendo da inizializzazioni arbitrarie. Mancava una spiegazione unificata e globale del meccanismo di "forzatura" verso questa soglia.

2. Metodologia: L'Edge Coupling

L'autore introduce un nuovo strumento analitico chiamato Edge Coupling (Accoppiamento al Bordo), che permette di analizzare la dinamica del GD non passo-passo, ma attraverso le proprietà funzionali di coppie di iterati consecutivi.

• Definizione del Funzionale: Viene definito un funzionale simmetrico su coppie di punti $(x, y)$:
  $$\mathcal{A}_\eta(x, y) \triangleq L(x) + L(y) - \frac{1}{2\eta} \|x - y\|^2$$
  Dove $L$ è la funzione di perdita.
• Condizioni di Criticità:
  • Ponendo il gradiente rispetto a $x$ a zero ($\nabla_x \mathcal{A}_\eta = 0$), si ottiene esattamente l'aggiornamento del gradiente: $y = x - \eta \nabla L(x)$.
  • Ponendo il gradiente rispetto a $y$ a zero, si ottiene la condizione per un'orbita periodica di periodo due.
  • I punti critici completi di $\mathcal{A}_\eta$ corrispondono quindi ai punti fissi e alle orbite periodiche di periodo due del GD.
• Approccio Analitico: L'autore utilizza questo funzionale per derivare:
  1. Una ricorrenza per l'incremento del passo (propagatore).
  2. Una formula esatta per la variazione della perdita in un singolo passo.
  3. Un'analisi di biforcazione delle orbite periodiche vicino ai punti critici.

3. Contributi Chiave e Risultati Teorici

A. Forzatura Globale verso $2/\eta$ (Teorema 4.1)

Il contributo principale è la dimostrazione che la traiettoria è globalmente costretta a visitare la soglia $2/\eta$.

• Derivando una formula di variazione della perdita basata sull'edge coupling, l'autore ottiene una relazione di "telescopio" (telescoping sum):
  $$\sum_{k=0}^{K-1} \|d_k\|^2 \left( \frac{2}{\eta} - \bar{r}_k \right) = 2(L(w_0) - L(w_K))$$
  Dove $d_k$ è il passo e $\bar{r}_k$ è una media pesata della curvatura (Hessiana) lungo il passo.
• Poiché la perdita totale $L(w_0) - L(w_K)$ è limitata (la perdita è limitata inferiormente), la somma pesata delle deviazioni della curvatura da $2/\eta$ deve essere finita.
• Conseguenza: Se il numero di passi tende all'infinito, la curvatura efficace $\bar{r}_k$ deve concentrarsi attorno a $2/\eta$. Il sistema non può rimanere stabilmente al di sopra o al di sotto di questa soglia senza violare i limiti della perdita.

B. Localizzazione Esatta (Teorema 4.3)

Mentre le formule precedenti coinvolgono medie dell'Hessiana, l'autore utilizza il Teorema del Valore Medio per dimostrare che queste medie corrispondono esattamente al valore dell'Hessiana in un punto interno al segmento del passo.

• Questo elimina qualsiasi "gap" residuo: la forzatura non è solo asintotica o approssimata, ma agisce sull'autovettore reale dell'Hessiana in punti specifici della traiettoria.

C. Analisi delle Orbite Periodiche e Biforcazione (Sezione 2)

L'autore riduce il problema delle orbite periodiche di periodo due a un problema di autovalore non lineare in termini di "semi-ampiezza" ($a$).

• Viene mostrato che la biforcazione (nascita di oscillazioni periodiche) avviene quando $2/\eta$ entra nello spettro dell'Hessiana.
• Per le reti lineari a due strati, viene dimostrato che la branca di period-doubling emerge continuamente dal lato del tasso di apprendimento più grande ($\eta > \eta_c$), confermando la natura supercritica della biforcazione.

D. Meccanismi di Stabilità (Sezione 5)

Il paper spiega perché il sistema non diverge una volta raggiunto il limite:

1. Crescita sopra la soglia: Se la curvatura supera $2/\eta$, la norma del passo cresce geometricamente. Su una traiettoria limitata, questo è insostenibile, costringendo il sistema a rientrare sotto la soglia.
2. Cancellazione Oscillatoria: All'interno della finestra di stabilità (vicino a $2/\eta$), il moltiplicatore del passo è vicino a $-1$, causando oscillazioni. L'autore dimostra che piccole asimmetrie in queste oscillazioni si cancellano a vicenda (non c'è deriva secolare), mantenendo il sistema confinato vicino al limite.

4. Significato e Impatto

• Unificazione Teorica: Il paper risolve il limite della letteratura precedente (che era puramente locale) fornendo una spiegazione globale del perché il GD converge verso $2/\eta$ da qualsiasi inizializzazione.
• Nuovo Strumento Matematico: L'introduzione dell'Edge Coupling offre un nuovo quadro teorico per analizzare la dinamica discreta del GD, collegandolo alla meccanica hamiltoniana e alle funzioni generatrici discrete.
• Precisione: La dimostrazione che la forzatura avviene esattamente sugli autovalori dell'Hessiana (senza approssimazioni di media) è un risultato tecnico rigoroso che chiarisce la natura esatta del fenomeno.
• Implicazioni per l'Addestramento: Conferma che l'uso di tassi di apprendimento elevati (che portano il sistema all'EoS) non è un comportamento caotico, ma una fase strutturata e stabile di ottimizzazione, dove le oscillazioni sono parte integrante del meccanismo di discesa della perdita.

In sintesi, Litman dimostra che l'Edge of Stability non è un accidente numerico, ma una proprietà intrinseca e inevitabile della dinamica del gradiente su funzioni lisce, guidata da una legge di conservazione della perdita che "spinge" la curvatura verso il valore critico $2/\eta$.