Résumé technique : Deux vitesses d'apprentissage : une décomposition représentation-lecture du grokking et de la double descente

Énoncé du problème

Les signaux standards pour surveiller l'entraînement des réseaux de neurones profonds, tels que la perte d'entraînement et la précision, échouent souvent à capturer la dynamique complexe de la généralisation. Deux phénomènes saillants compliquent cette image :

1. Grokking : Un modèle atteint une performance d'entraînement quasi parfaite tandis que la précision de test reste au niveau du hasard pendant une période prolongée, suivie d'une émergence soudaine de la généralisation.
2. Double descente par époque : La perte d'entraînement diminue de manière monotone, tandis que la perte ou l'erreur de test présente une hausse et une baisse non monotones.

Les explications existantes sont fréquemment spécifiques à la tâche (par exemple, reposant sur des structures algébriques dans l'arithmétique modulaire) ou reposent sur des transitions de régime binaires (par exemple, de l'apprentissage « paresseux » à « riche »). Il manque un cadre unifié, agnostique de la tâche, capable de diagnostiquer ces phénomènes à travers des architectures et des tâches diverses, distinguant la véritable généralisation des artefacts d'apprentissage spurius induits par des recettes d'entraînement non standard.

Méthodologie : Décomposition représentation-lecture

Les auteurs proposent un cadre qui décompose un réseau de neurones $f(x)$ en deux composantes distinctes :

1. Encodeur ($\phi_\theta$) : Une application non linéaire de l'espace d'entrée vers un espace de représentation à $N$ dimensions.
2. Lecteur ($W$) : Une application linéaire de l'espace de représentation vers les logits de sortie.

Cette décomposition permet de suivre séparément l'apprentissage de la représentation (changements dans $\phi_\theta$) et le calibrage du lecteur (changements dans $W$). Pour analyser l'interdépendance entre ces composantes, les auteurs emploient quatre mesures diagnostiques agnostiques de la tâche, dérivées de la géométrie représentative et de la théorie du Noyau Tangent aux Neurones (NTK) :

1. Dimension critique ($N_{crit}$) : Quantifie la robustesse de la séparabilité linéaire dans l'espace de représentation. Elle est définie comme la plus petite dimension de projection attendue requise pour préserver la séparabilité avec une probabilité $\geq 0,5$. Une $N_{crit}$ plus faible indique de meilleurs variétés « démêlées » et une plus grande généralisabilité.
2. Mesures géométriques liées à la tâche (GLUE) : Un ensemble de quatre métriques ($D, R, \rho_c, \rho_a$) dérivées de la théorie de la Géométrie Liée à l'Efficacité de Démêlage (GLUE). Ces métriques caractérisent la dimension effective de la variété, son rayon, ainsi que l'alignement des centres et des axes entre les variétés conditionnées par l'étiquette.
3. Précision de la sonde linéaire : Un classifieur linéaire régularisé entraîné sur des représentations figées. L'écart entre la précision de la sonde et la précision du modèle diagnostique un désalignement du lecteur (surapprentissage ou sous-apprentissage par rapport aux capacités de l'encodeur).
4. Alignement NTK-Étiquette de la dernière couche : Mesure l'alignement entre la matrice de Gram des caractéristiques (NTK de la dernière couche) et le noyau des étiquettes. Cela quantifie dans quelle mesure la géométrie des caractéristiques actuelle soutient le flux de gradient vers des solutions discriminantes des étiquettes, évalué séparément sur les ensembles d'entraînement et de test.

Contributions et résultats clés

1. Deux vitesses d'apprentissage dans le grokking

En appliquant la décomposition à des tâches canoniques de grokking (addition modulaire, composition de permutations, parité sparse) sur des architectures MLP et Transformer, les auteurs réfutent l'explication stricte « paresseux-vers-riche », qui postule que l'apprentissage de la représentation est latent jusqu'au début du grokking.

• Apprentissage continu de la représentation : $N_{crit}$ diminue de manière graduelle et continue tout au long de l'entraînement, y compris la phase pré-grokking. L'apprentissage de la représentation est lent mais pas absent ; il progresse continûment plutôt que comme une transition de phase soudaine.
• Lecteur biaisé vers l'entraînement : Avant le début du grokking, la direction d'apprentissage du lecteur est systématiquement biaisée vers les données d'entraînement. L'écart d'alignement NTK (entraînement moins test) est grand et positif, indiquant que le lecteur surapprend l'ensemble d'entraînement tandis que les représentations ne sont pas encore suffisamment robustes pour soutenir la généralisation.
• Mécanisme : Le grokking ne résulte pas d'un changement soudain de régime d'apprentissage, mais de l'égalisation progressive de la direction du lecteur à mesure que la qualité de la représentation ($N_{crit}$) finit par rattraper la vitesse de calibrage du lecteur.

2. Différences architecturales dans les mécanismes géométriques

Le cadre révèle que différentes architectures atteignent la généralisation par des voies géométriques distinctes.

• MLP vs Transformers : Dans les tâches d'addition modulaire et de composition de permutations, les MLP et les Transformers présentent des tendances opposées dans l'alignement des centres ($\rho_c$) et l'alignement des axes ($\rho_a$). Bien que les deux architectures réduisent la dimension et le rayon de la variété, elles manipulent la géométrie inter-classe différemment pour atteindre le même comportement de surface de grokking.

3. Diagnostic de l'apprentissage spurius

Le cadre distingue avec succès la véritable généralisation des artefacts causés par des recettes d'entraînement non standard.

• Artefact de grokking sur MNIST : Un exemple de grokking précédemment rapporté sur MNIST (induit par des normes de poids initiales élevées et de petits ensembles d'entraînement) est diagnostiqué comme spurius. Le cadre révèle une dégradation de la représentation (augmentation de $N_{crit}$ et diminution de la précision de la sonde linéaire) et un surapprentissage du lecteur pendant la phase « pré-grokking ». La généralisation retardée apparente est un artefact du lecteur corrigeant éventuellement une représentation dégradée, plutôt qu'une véritable amélioration de l'apprentissage des caractéristiques.
• Double descente par époque : Dans les scénarios de double descente induits par du bruit d'étiquette, le cadre identifie un processus en deux étapes : une phase initiale de détérioration de la représentation, suivie d'une phase où les mises à jour du lecteur sont dominées par des signaux spurius (étiquettes bruyantes). Le retrait des ingrédients non standard (par exemple, le bruit d'étiquette) élimine à la fois le motif de double descente et les signatures diagnostiques.

4. Cadre d'analyse descendante

Les auteurs positionnent cette décomposition comme un complément « descendant » (top-down) à l'interprétabilité mécaniste « ascendante » (bottom-up) (par exemple, découverte de circuits). En caractérisant la structure globale des représentations (géométrie de la variété, séparabilité, alignement), le cadre fournit des observables cohérents à travers les splits d'entraînement/test et des tâches diverses, offrant un signal fiable pour surveiller la dynamique d'apprentissage là où les courbes de perte standard échouent.

Importance

L'article établit la décomposition représentation-lecture comme un cadre unifié et agnostique de la tâche pour comprendre la dynamique d'apprentissage. Son importance principale réside dans :

• Affinement des comptes rendus théoriques : Il va au-delà de la dichotomie binaire paresseux-vers-riche, démontrant que l'apprentissage de la représentation est continu et que le grokking est piloté par les vitesses relatives de l'amélioration de la représentation par rapport au calibrage du lecteur.
• Utilité diagnostique : Il fournit des signatures spécifiques pour identifier les dynamiques d'apprentissage spurius (par exemple, dégradation de la représentation ou désalignement du lecteur) qui pourraient être confondues avec de véritables phénomènes d'apprentissage dans des configurations d'entraînement non standard.
• Généralisabilité : Le cadre s'applique de manière cohérente aux tâches algorithmiques, aux ensembles de données réels (MNIST) et aux architectures (MLP, Transformer, ResNet), offrant un outil principiel pour la recherche en interprétabilité et la surveillance de modèles critiques pour la sécurité.

Les auteurs concluent que, bien que le cadre soit actuellement appliqué à la dernière couche linéaire, il offre une base robuste pour de futures investigations sur les couches intermédiaires, les structures spécifiques à la tâche et les scénarios d'apprentissage continu.