The Origin of Edge of Stability

Titre : L'Origine de la Stabilité à la Limite (Edge of Stability)

1. Problématique

Le phénomène de la Stabilité à la Limite (Edge of Stability - EoS) a été observé empiriquement lors de l'entraînement de réseaux de neurones par descente de gradient par lots complets (full-batch). Contrairement à la théorie classique de l'optimisation qui garantit une décroissance monotone de la perte tant que le taux d'apprentissage $\eta$ satisfait $\eta < 2/\lambda_{\max}(\nabla^2 L)$ (où $\lambda_{\max}$ est la "sharpness" ou courbure maximale), les observations montrent que :

• La sharpness augmente pendant l'entraînement jusqu'à atteindre la valeur critique $2/\eta$.
• À ce seuil, la sharpness se sature et la perte d'entraînement commence à osciller à court terme tout en continuant de diminuer à long terme.
• Ce phénomène se produit systématiquement, quel que soit l'initialisation.

Le problème central : Bien que des travaux antérieurs aient établi des mécanismes d'auto-régulation locaux près de la limite (par exemple via des termes cubiques de Taylor), ils n'expliquent pas pourquoi la trajectoire est forcée globalement vers ce seuil $2/\eta$ à partir d'une initialisation arbitraire. Il manquait une explication unifiée et globale de ce phénomène.

2. Méthodologie : Le "Edge Coupling"

L'auteur introduit un nouvel outil mathématique appelé Edge Coupling (Couplage de l'arête), noté $\mathcal{A}_\eta(x, y)$, défini sur les paires d'itérations consécutives $(x, y)$ d'une trajectoire de descente de gradient.

La fonctionnelle est définie comme suit :
$$\mathcal{A}_\eta(x, y) \triangleq L(x) + L(y) - \frac{1}{2\eta} \|x - y\|^2$$

Points clés de la méthodologie :

• Caractérisation variationnelle : En annulant le gradient par rapport à $x$ ($\nabla_x \mathcal{A}_\eta = 0$), on retrouve exactement la mise à jour de la descente de gradient : $y = x - \eta \nabla L(x)$. Ainsi, la trajectoire de GD est l'ensemble des points critiques partiels de cette fonctionnelle.
• Analyse des conditions critiques : En étudiant les conditions où les gradients partiels par rapport à $x$ et $y$ sont nuls simultanément, l'auteur classe les points fixes et les orbites de période deux.
• Réduction au centre : Pour analyser les orbites de période deux (oscillations), le problème est réparamétré en termes d'un centre $m$ et d'une demi-amplitude $a$, réduisant l'analyse à un problème d'éigenvalue non linéaire dépendant uniquement de $a$.
• Théorèmes de concentration : L'auteur utilise le théorème des accroissements finis (Mean Value Theorem) pour localiser les moyennes de Hessienne sur des points réels de la trajectoire, permettant de passer des moyennes intégrales aux valeurs propres réelles du Hessian.

3. Contributions Clés

1. Unification par le Edge Coupling :
   L'article démontre que toutes les dynamiques de l'EoS découlent des conditions critiques de $\mathcal{A}_\eta$. Cela fournit un cadre unifié reliant la récurrence de l'incrément de pas et la formule de changement de perte.

2. Preuve de l'Attraction Globale vers $2/\eta$ :
   Contrairement aux travaux précédents limités au voisinage local, l'auteur prouve que la trajectoire est forcée globalement à visiter la région où la courbure effective est proche de $2/\eta$.

   • Une formule de changement de perte (théorème 2.2) montre que la variation de perte à chaque étape est proportionnelle à $(2/\eta - \tilde{r}_k)$, où $\tilde{r}_k$ est une courbure effective.
   • En sommant cette formule sur toute la trajectoire (téléscopage), on obtient une loi de conservation : la somme pondérée des écarts à $2/\eta$ est bornée par la chute totale de la perte.
   • Puisque la perte est bornée inférieurement, la courbure moyenne pondérée doit converger vers $2/\eta$.

3. Localisation Exacte (Sans "Gap") :
   L'article résout le problème de la différence entre les moyennes de Hessienne et les valeurs propres réelles. En utilisant le théorème des accroissements finis, il est démontré que pour chaque étape, la courbure effective est exactement égale à la valeur propre directionnelle du Hessian en un point intérieur du segment de pas. Cela élimine toute erreur d'approximation ("gap") dans le forçage de la courbure.

4. Analyse de la Bifurcation de Période Deux :
   L'auteur caractérise l'apparition des orbites de période deux (oscillations) comme une bifurcation de type fourche (pitchfork).

   • Près d'un point fixe, le problème se réduit à une fonction de la demi-amplitude.
   • La direction de la bifurcation et le côté du seuil critique ($\eta > \eta_c$ ou $\eta < \eta_c$) où elle apparaît sont déterminés par le terme quartique de l'expansion de Taylor de la perte réduite.
   • Pour les réseaux linéaires à deux couches, il est prouvé que la bifurcation est width-invariant (indépendante de la largeur du réseau) et apparaît continûment pour $\eta > \eta_c$.

4. Résultats Principaux

• Théorème 4.1 (Concentration de la courbure) : La courbure effective $\tilde{r}_k$ converge vers $2/\eta$ dans un sens pondéré par la norme des pas. Si la perte est bornée et que la somme des carrés des pas diverge, la moyenne pondérée de la courbure tend exactement vers $2/\eta$.
• Théorème 4.3 (Localisation) : Il existe des points réels sur le segment de pas où la valeur propre du Hessian atteint exactement la valeur de la courbure effective. Ainsi, le forçage vers $2/\eta$ s'applique aux valeurs propres réelles du Hessian, pas seulement à des moyennes.
• Mécanismes de Stabilité (Section 5) :
  • Si la courbure dépasse $2/\eta$, la norme du pas croît géométriquement (instabilité), forçant la trajectoire à sortir de cette région.
  • Si la courbure est proche de $2/\eta$, les itérations alternent de signe (multiplieur proche de -1), créant des oscillations qui s'annulent mutuellement (cancellation oscillatoire) plutôt que de dériver, maintenant ainsi la stabilité.
• Validation Expérimentale : Les prédictions théoriques sont validées sur des réseaux de neurones (MLP) et des réseaux linéaires à deux couches, montrant une concordance parfaite entre l'amplitude des oscillations observée et la prédiction théorique $\sqrt{\eta - \eta_c}$.

5. Signification et Impact

• Explication Fondamentale : Cet article fournit la première explication unifiée et rigoureuse de l'Edge of Stability, passant d'une observation empirique à une loi de conservation mathématique dérivée d'une fonctionnelle variationnelle.
• Au-delà de la théorie locale : Il comble le fossé entre les analyses locales (stabilité près du minimum) et le comportement global de l'optimisation, expliquant pourquoi le système "trouve" et "reste" à la limite de stabilité.
• Implications pour l'entraînement : La compréhension de ces mécanismes (oscillations, forçage de la courbure) pourrait guider le développement de nouveaux algorithmes d'optimisation ou de stratégies de réglage du taux d'apprentissage pour exploiter ou contrôler l'EoS.
• Généralité : La méthode du "Edge Coupling" et les théorèmes de concentration s'appliquent non seulement à la descente de gradient par lots complets, mais sont étendus dans les annexes à la SGD par mini-lots, aux paires de trajectoires (stabilité algorithmique) et au temps continu.

En résumé, ce travail établit que l'Edge of Stability n'est pas un artefact accidentel, mais une conséquence inévitable de la structure géométrique de la descente de gradient, régie par une loi de conservation de la courbure qui force le système à s'auto-réguler à la valeur critique $2/\eta$.