Grokking as a Phase Transition between Competing Basins: a Singular Learning Theory Approach

En appliquant la Théorie de l'Apprentissage Singulier (SLT), cette étude interprète le phénomène de « grokking » comme une transition de phase entre des bassins de solutions compétitifs dans les réseaux quadratiques, démontrant que le coefficient d'apprentissage local sert d'indicateur fiable pour suivre la dynamique de généralisation et prédire ces transitions.

L'essentiel

Imaginez que vous apprenez un nouveau jeu de société très complexe, comme les échecs ou le Sudoku, mais avec des règles mathématiques très précises (l'addition modulo un nombre premier).

Vous avez un élève très doué, un modèle d'intelligence artificielle. Au début, il apprend par cœur toutes les parties qu'on lui montre. Il est excellent en classe : il répond parfaitement à toutes les questions de l'examinateur. Mais dès qu'on lui pose une question un peu différente (un test), il échoue lamentablement. Il a mémorisé les réponses, mais il n'a pas compris la logique.

C'est ce qu'on appelle le phénomène de "Grokking" (un mot d'argot américain qui signifie "comprendre soudainement et profondément").

Ce papier de recherche explique pourquoi et comment cet élève passe brutalement de la mémorisation à la compréhension, en utilisant une théorie mathématique appelée Théorie de l'Apprentissage Singulier (SLT).

Voici l'explication simple, avec des analogies :

1. Le Paysage des Solutions : Deux Vallées

Imaginez que l'apprentissage de l'IA se déroule dans un immense paysage de montagnes et de vallées. Le but du jeu est de trouver le point le plus bas (le "minimum") pour avoir le meilleur score.

• La Vallée de la Mémorisation (La Vallée Étroite) : Au début, l'IA tombe dans une petite vallée très profonde et très étroite. Elle y trouve une solution parfaite pour les questions d'entraînement. C'est comme si elle avait trouvé un trou de souris : elle rentre parfaitement, mais elle est coincée. Elle ne peut pas bouger sans sortir du trou. C'est une solution "rigide".
• La Vallée de la Généralisation (La Vallée Large) : Plus loin, il y a une immense plaine, large et plate. Si l'IA s'y installe, elle peut bouger un peu sans que son score ne chute. C'est une solution "flexible". Elle a compris la règle générale, pas juste les exemples.

Le problème, c'est que l'IA commence souvent par tomber dans la petite vallée (mémorisation) parce qu'elle y arrive vite. Mais elle reste coincée là pendant très longtemps, même si on continue à l'entraîner.

2. Le "Coût de l'Apprentissage Local" (LLC) : La Mesure de la Flexibilité

Les auteurs utilisent un outil mathématique appelé le Coefficient d'Apprentissage Local (LLC).

• Imaginez que le LLC soit une mesure de l'espace disponible dans votre vallée.
• Une vallée étroite (mémorisation) a un LLC élevé (peu d'espace, peu de liberté).
• Une vallée large (généralisation) a un LLC faible (beaucoup d'espace, beaucoup de liberté).

La théorie dit que plus le temps passe (plus on a de données), plus l'IA "préfère" naturellement les vallées larges (faible LLC), car elles sont statistiquement plus probables et plus robustes.

3. Le "Saut" (Grokking) : Un Changement de Phase

Le Grokking, c'est le moment où l'IA décide soudainement de quitter la petite vallée étroite pour traverser la montagne et s'installer dans la grande plaine.

• Avant le Grokking : L'IA est coincée dans la vallée étroite. Elle a un score parfait en classe, mais un score nul à l'examen.
• Pendant le Grokking : L'IA explore le paysage. Elle trouve un chemin vers la grande plaine. C'est comme si elle avait soudainement "compris" la règle du jeu.
• Après le Grokking : Elle s'installe dans la grande plaine. Son score en classe reste bon, mais maintenant, son score à l'examen explose aussi ! Elle a généralisé.

4. Ce que les auteurs ont découvert

Les chercheurs ont étudié ce phénomène sur des réseaux de neurones très simples (des réseaux "quadratiques") qui font des calculs de modulo.

1. Ils ont fait les maths exactes : Ils ont calculé précisément la taille de ces "vallées" (le LLC) pour savoir exactement quand l'IA devrait changer de vallée.
2. Ils ont observé la transition : En regardant comment le LLC évolue pendant l'entraînement, ils ont vu que le moment où le LLC commence à baisser correspond exactement au moment où l'IA commence à bien réussir les tests.
3. Le rôle de l'apprentissage (Learning Rate) : Ils ont découvert que si on change la "vitesse" à laquelle l'IA apprend (le taux d'apprentissage), on change la difficulté du saut.
   • Analogie : Si l'IA apprend trop lentement, elle reste coincée dans la petite vallée trop longtemps. Si elle apprend avec la bonne vitesse (ni trop lent, ni trop vite), elle saute plus facilement vers la grande plaine.

En résumé

Ce papier nous dit que le Grokking n'est pas de la magie. C'est une transition physique (comme l'eau qui gèle ou bout).

L'IA commence par une solution "coincée" (mémorisation) qui est statistiquement probable au début. Mais à mesure qu'elle accumule de l'expérience, la nature mathématique du problème la pousse vers une solution "libre" (généralisation) qui est plus stable.

Les auteurs ont prouvé qu'en mesurant la "flexibilité" de la solution (le LLC), on peut prédire exactement quand l'IA va faire ce saut de compréhension, même avant qu'elle ne réussisse ses tests ! C'est comme avoir un thermomètre qui vous dit quand l'eau va bouillir, avant même de voir les bulles.

Résumé technique

1. Problématique et Contexte

Le Grokking est un phénomène observé lors de l'entraînement de modèles d'apprentissage automatique, particulièrement sur des tâches algorithmiques comme l'arithmétique modulaire. Il se caractérise par une dynamique en deux temps :

1. Le modèle atteint rapidement une perte empirique (erreur d'entraînement) proche de zéro, indiquant qu'il a mémorisé les données.
2. Cependant, la performance de généralisation (sur les données de test) reste faible pendant une longue période.
3. Soudainement, après un entraînement prolongé, la précision sur le test s'améliore brusquement, marquant l'entrée dans une phase de généralisation.

L'article pose la question fondamentale : Qu'est-ce qui détermine quel bassin de solution (bassin de mémorisation vs bassin de généralisation) est statistiquement préféré par l'optimiseur ?

Les hypothèses existantes suggèrent que les minima "plats" du paysage de perte généralisent mieux, mais les fondements théoriques de cette intuition restent incomplets, en particulier pour les réseaux de neurones qui sont des modèles singuliers (non identifiables localement, avec des matrices d'information de Fisher de rang déficient).

2. Méthodologie : La Théorie de l'Apprentissage Singulier (SLT)

Les auteurs adoptent le cadre de la Théorie de l'Apprentissage Singulier (Singular Learning Theory - SLT), développée par Sumio Watanabe, pour analyser ce phénomène.

• Le Coefficient d'Apprentissage Local (LLC) : Au cœur de l'approche se trouve le LLC, noté $\lambda$. C'est une mesure de la dégénérescence locale de la surface de perte autour d'un optimum.
  • Pour un modèle régulier, $\lambda = d/2$ (où $d$ est le nombre de paramètres).
  • Pour un modèle singulier, $\lambda < d/2$. Un $\lambda$ plus faible indique un bassin de solution plus "plat" et plus dégénéré (plus grand volume de paramètres équivalents).
• Transition de Phase Bayésienne : La SLT établit que la masse postérieure se concentre asymptotiquement dans les bassins ayant le plus faible LLC. La différence d'énergie libre entre deux bassins est dominée par le terme $(\lambda_a - \lambda_b) \log n$. Ainsi, lorsque la taille de l'échantillon $n$ augmente, le modèle effectue une transition de phase du bassin de mémorisation (LLC élevé) vers le bassin de généralisation (LLC faible).
• Cadre Expérimental : Pour rendre les calculs analytiques possibles, les auteurs étudient des réseaux quadratiques (quadratic networks) entraînés sur des tâches d'arithmétique modulaire. Ce choix permet de dériver des expressions fermées pour le LLC, comblant ainsi le fossé entre la théorie abstraite et l'analyse empirique.

3. Contributions Clés

L'article apporte deux contributions majeures :

1. Définitions fermées du LLC pour les réseaux quadratiques :
   Les auteurs dérivent des expressions analytiques exactes du LLC pour deux régimes de paramétrisation :

   • Régime sur-paramétrisé ($K \ge d(d+1)/2$) : Le LLC est donné par $\lambda = p \cdot \frac{d(d+1)}{4}$, où $p$ est le nombre de sorties et $d$ la dimension d'entrée.
   • Régime sous-paramétrisé ($K < d(d+1)/2$) : Le LLC est donné par $\lambda = K \cdot \frac{d + p - 1}{2}$.
     Ces résultats montrent comment la géométrie du bassin de solution change en fonction de la largeur du réseau et de la structure de la tâche.

2. Interprétation du Grokking comme transition de phase :
   En combinant ces calculs théoriques avec des mesures empiriques, l'article interprète le grokking non pas comme un simple changement de dynamique d'optimisation, mais comme une transition de phase bayésienne. Le modèle commence dans un bassin de mémorisation (solution "paresseuse" ou NTK) avec un LLC élevé, puis, sous l'effet de l'optimisation stochastique (SGD), il bascule vers un bassin de généralisation (solution structurée) avec un LLC plus faible, qui est statistiquement préféré à long terme.

4. Résultats Expérimentaux

Les auteurs valident leur théorie à travers plusieurs expériences :

• Validation des lois d'échelle : Les LLC estimés empiriquement suivent les prédictions théoriques en fonction de la dimension du réseau ($K$) et de la taille du groupe modulaire ($p$).
• Suivi de la généralisation : Les trajectoires du LLC, calculées uniquement à partir des données d'entraînement, suivent étroitement l'évolution de la perte de validation. Une baisse du LLC précède ou coïncide avec l'amélioration soudaine de la généralisation.
• Effet du taux d'apprentissage (Learning Rate) :
  • Ils introduisent une mesure de sévérité du grokking (GSM).
  • Ils observent une corrélation négative entre le taux d'apprentissage et la sévérité du grokking : des taux d'apprentissage plus élevés réduisent le délai avant la généralisation.
  • Explication SLT : Un taux d'apprentissage plus élevé favorise l'exploration de bassins plus dégénérés (LLC plus faible) dès le début, évitant ainsi de rester piégé dans les vallées aiguës du bassin de mémorisation.
• Robustesse : Ces résultats sont confirmés sur différentes tailles de réseaux, différents hyperparamètres (poids de régularisation, taille de lot) et même sur des architectures plus complexes comme les Transformers.

5. Signification et Impact

Ce travail est significatif pour plusieurs raisons :

• Unification Théorique : Il fournit un cadre mathématique rigoureux (SLT) pour expliquer un phénomène empirique mystérieux (le grokking), reliant la géométrie du paysage de perte, la complexité effective du modèle et la généralisation.
• Outil de Diagnostic : Le LLC s'avère être un indicateur fiable et prédictif de l'émergence de la généralisation, calculable uniquement sur les données d'entraînement. Cela offre un outil potentiel pour surveiller l'entraînement des modèles sans avoir besoin de données de validation.
• Au-delà de la "Platitude" : L'article renforce l'idée que la "platitude" des minima est une propriété géométrique profonde liée à la dégénérescence des paramètres (mesurée par le LLC) plutôt qu'à la simple courbure locale (Hessienne), qui n'est pas invariante par reparamétrisation.
• Limites et Perspectives : L'analyse repose sur un cadre asymptotique bayésien. Bien que les résultats empiriques soutiennent le lien avec l'optimisation par SGD, une connexion théorique complète entre la concentration postérieure et la dynamique stochastique reste un défi ouvert. De plus, l'extension à des architectures plus complexes (comme les grands Transformers) nécessite des travaux futurs.

En résumé, l'article démontre que le grokking est une transition de phase naturelle où l'optimiseur, guidé par la géométrie singulière du paysage de perte, finit par sélectionner la solution de généralisation grâce à sa complexité statistique inférieure (LLC plus faible).