Résumé technique : Comment dompter le « grokking » : La géométrie des représentations comme signal de contrôle

Énoncé du problème

Le « grokking » est un phénomène observé dans les réseaux de neurones surparamétrés où les modèles mémorisent d'abord les données d'entraînement et ne font preuve d'une forte généralisation qu'après une optimisation prolongée. Bien que des études récentes aient lié le grokking à la décroissance des poids (weight decay), à la compression des caractéristiques, à la formation de circuits et à la dynamique d'apprentissage des représentations, les facteurs spécifiques régissant le moment et l'émergence de cette généralisation tardive restent incomplètement compris. Un thème commun dans la littérature existante est que les représentations apprises se compressent progressivement en structures de plus faible dimension. Cependant, les travaux antérieurs traitent largement la géométrie des représentations comme une propriété émergente plutôt que comme un mécanisme de contrôle direct. Cet article examine si la manipulation explicite de la géométrie des représentations cachées peut servir de signal de contrôle pour influencer l'apparition et la vitesse du grokking.

Méthodologie

Les auteurs introduisent la Régularisation de la Dimensionnalité Géométrique (GeomDR), un régularisateur spectral au niveau de la représentation conçu pour contrôler directement la dimensionnalité effective des représentations cachées pendant l'entraînement.

Caractérisation géométrique

Pour une matrice de représentation de couche cachée $Z \in \mathbb{R}^{N \times d}$, la méthode centre d'abord les données pour obtenir $\tilde{Z}$. Elle calcule ensuite la matrice de covariance empirique $C = \frac{1}{N-1}\tilde{Z}^\top\tilde{Z}$. La dimensionnalité effective ($D_{eff}$) est calculée à l'aide du rapport de participation des valeurs propres ($\mu_i$) de $C$ :
$$D_{eff} = \frac{(\sum_i \mu_i)^2}{\sum_i \mu_i^2}$$
De plus, une distance de voisinage locale ($\rho$) est définie comme la distance euclidienne moyenne vers les $k$ plus proches voisins des représentations normalisées.

L'objectif GeomDR

GeomDR impose une pénalité sur la variance contenue dans les directions au-delà d'un sous-espace cible de dimension $d^*$. Étant donné les valeurs propres triées par ordre décroissant ($\mu_1 \ge \mu_2 \ge \dots \ge \mu_d$), la perte par couche est :
$$L_{GeomDR}^{(\ell)} = \sum_{j=d^*+1}^{d} \mu_j^{(\ell)}$$
L'objectif d'entraînement total combine la perte de la tâche ($L_{task}$) et le régularisateur géométrique :
$$L = L_{task} + \lambda(t)L_{GeomDR}$$
Crucialement, le régularisateur n'est pas actif dès le début. Il est introduit après une phase initiale d'entraînement non contraint (étape d'activation $t_s$) et sa force $\lambda(t)$ est augmentée progressivement à l'aide d'une rampe cosinus. Cette conception imite la trajectoire naturelle du grokking où la mémorisation précède la généralisation.

Protocole expérimental

Le cadre a été évalué sur trois tâches algorithmiques connues pour présenter du grokking :

1. Addition modulaire : Prédire $(a + b) \mod 97$.
2. Division modulaire : Prédire $(a \cdot b^{-1}) \mod 97$.
3. Composition de permutations : Prédire la composition de deux permutations dans $S_5$.

Les expériences ont utilisé deux architectures : un Perceptron Multicouche (MLP) et un Transformer. L'étude a fait varier systématiquement les calendriers d'intervention (force $\lambda_{max}$ et durée de la rampe $T_{ramp}$), les dimensionnalités cibles ($d^*$) et les temps de début d'intervention ($t_s$).

Contributions clés

1. Observation de l'effondrement précurseur : Les auteurs démontrent que l'effondrement de la dimensionnalité précède systématiquement la transition de la mémorisation vers la généralisation (grokking) dans les configurations de référence. Cela suggère que la compression géométrique n'est pas seulement un corrélat mais un moteur mécaniste de la transition.
2. Introduction de GeomDR : Ils proposent GeomDR, un régularisateur spectral qui supprime directement le spectre de covariance en dehors d'un sous-espace cible, contrôlant ainsi la dimensionnalité effective des représentations cachées.
3. Contrôle systématique des dynamiques : L'article fournit une étude empirique complète montrant que les interventions géométriques peuvent modifier substantiellement les dynamiques de généralisation tardive. Selon le calendrier d'intervention et la dimensionnalité cible, GeomDR peut accélérer le grokking, le retarder ou le supprimer entièrement.

Résultats

Accélération du grokking

Dans tous les contextes évalués, GeomDR a systématiquement modifié les dynamiques de grokking :

• Performance des MLP : Sur l'addition modulaire, GeomDR a accéléré le grokking jusqu'à 52 fois (réduisant les étapes de 362K à ~7K). Des améliorations similaires ont été observées pour la division modulaire (39x) et la composition de permutations (~23x). Dans tous les cas de MLP, GeomDR a atteint un taux de succès de 100 % là où la référence était instable ou échouait.
• Performance des Transformers : Bien que la dynamique de grokking de référence pour les Transformers soit déjà plus rapide que celle des MLP, GeomDR a tout de même fourni une accélération constante (ex: ~2x sur l'addition modulaire), maintenant un taux de succès de 100 %.

Influence des hyperparamètres

• Dimensionnalité cible ($d^*$) : Les dimensions cibles modérées à larges ($16 \le d^* \le 64$) ont produit l'accélération la plus forte. Une compression agressive ($d^* \le 8$) a souvent ralenti ou déstabilisé l'apprentissage.
• Temps d'intervention ($t_s$) : Une « fenêtre temporelle critique » a été identifiée. Les interventions appliquées trop tôt (avant que des représentations utiles ne se forment) ou trop tard (après que la mémorisation soit rigide) étaient moins efficaces. Les temps d'activation intermédiaires (ex: $t_s = 2000$ pour les MLP) ont produit le grokking le plus rapide et le plus stable.
• Calendrier (Schedule) : Les forces de régularisation et les durées de rampe intermédiaires ont généralement obtenu les meilleurs résultats ; des interventions trop fortes pouvaient déstabiliser l'entraînement.

Corrélations géométriques

L'étude a confirmé que les interventions accélérant le grokking accéléraient également l'apparition de l'effondrement de la dimensionnalité. Dans les modèles GeomDR, l'écart temporel entre l'effondrement de la dimensionnalité et la généralisation réussie a été significativement réduit par rapport aux références, suggérant un couplage plus étroit entre la réorganisation géométrique et l'émergence de la généralisation.

Scalabilité et efficacité

Les auteurs ont démontré que GeomDR reste efficace même lorsque le spectre de covariance est estimé à partir d'un sous-ensemble aléatoire de représentations (ex: 5 % – 25 % des données), offrant une accélération du temps d'exécution de 2,4x à 2,7x sans affecter matériellement l'accélération du grokking.

Signification et affirmations

L'article affirme que la géométrie des représentations sert de signal de contrôle efficace pour le grokking. En modifiant directement le spectre de covariance des activations cachées, les chercheurs peuvent influencer systématiquement quand et comment rapidement un réseau passe de la mémorisation à la généralisation.

Les auteurs postulent que ce travail fournit la preuve que les interventions géométriques offrent une approche pratique pour étudier et influencer la généralisation tardive. Plutôt que d'observer passivement les changements de représentation, la méthode permet de manipuler directement les propriétés géométriques pour tester des hypothèses sur les dynamiques d'apprentissage. Les résultats suggèrent que la relation entre la géométrie des représentations et le grokking est robuste à travers différentes tâches et architectures, bien que l'ampleur de l'amélioration varie (étant plus prononcée pour les MLP que pour les Transformers, probablement en raison des biais inductifs plus forts de ces derniers).

L'article conclut que, bien qu'il n'établisse pas la dimensionnalité de la représentation comme l'unique mécanisme causal, il démontre que les interventions géométriques sont suffisantes pour altérer le moment du grokking, offrant une nouvelle méthodologie expérimentale pour la recherche sur l'optimisation et la généralisation.