A Single-Layer Model Can Do Language Modeling

Résumé technique : Un modèle à couche unique peut réaliser la modélisation du langage

Énoncé du problème
Les modèles de langage modernes augmentent généralement leur capacité en augmentant la profondeur, en empilant plusieurs couches où chaque couche maintient son propre état (par exemple, les caches KV par couche dans les Transformers ou les états matriciels par couche dans Mamba, GDN, RWKV et xLSTM). Cela répartit le calcul et la mémoire sur la pile. En revanche, les réseaux de neurones biologiques reposent fortement sur la récurrence — la réutilisation des mêmes circuits au fil du temps — plutôt que sur l'empilement. L'article examine la viabilité de cette forme biologique pour la modélisation du langage : spécifiquement, jusqu'où une seule couche récurrente peu profonde peut aller lorsqu'elle revisite un seul vecteur d'état à chaque étape, plutôt que de répartir l'état sur une pile profonde.

Méthodologie : Réseaux de prédiction ancrés (GPN)
Les auteurs proposent les Réseaux de prédiction ancrés (GPN), une architecture définie par un seul vecteur d'état revisité à chaque pas de temps. La boucle principale comprend trois opérations :

1. Ancrage : L'état prédit précédent ($s^p_{t-1}$) est fusionné avec le token observé actuel ($x_t$) via une connexion résiduelle à porte pour produire un état ancré ($s^g_t$).
2. Prédiction : L'état ancré est traité pour prédire l'état suivant ($s^p_t$). Cette prédiction est décodée en une distribution de tokens pour l'étape suivante.
3. Mémoire (GPN+M) : Les auteurs étendent le GPN de base en GPN+M en ajoutant une mémoire matricielle centralisée ($M$). Cette mémoire est mise à jour via une règle delta à porte et lue en parallèle avec un réseau feed-forward (FFN). Tant le FFN (paramètres hors ligne) que la lecture de la mémoire (paramètres en ligne) agissent comme des prédicteurs additifs pour la mise à jour de l'état.

Contrairement aux modèles récurrents standards qui maintiennent un état par couche, le GPN utilise une seule couche large et peu profonde avec une seule mémoire matricielle partagée. La profondeur est atteinte par le temps (la réutilisation de la boucle) plutôt que par l'espace (l'empilement de couches). Le vecteur d'état joue un double rôle : il agit comme une prévision pour l'étape suivante et comme un encodage de l'historique observé jusqu'alors, l'« ancrage » complétant l'encodage pour le cycle suivant.

Contributions clés

1. Architecture : L'introduction des GPN et GPN+M, qui utilisent un seul vecteur d'état, un seul FFN et une seule mémoire matricielle centralisée revisités à chaque étape, en contraste avec les états par couche empilés des Transformers et des RNN linéaires.
2. Vue conceptuelle : Un cadre où la prédiction de l'état pour l'étape suivante est considérée comme faisant partie de l'encodage de cette étape, un seul vecteur jouant les deux rôles.
3. Preuves empiriques : Des benchmarks de performance à 130 millions de paramètres montrant qu'un GPN+M à 1 couche atteint une perplexité compétitive par rapport à des bases profondes, avec une variante à 2 couches réduisant encore l'écart.
4. Observabilité : Parce que l'ensemble du contexte de travail est un seul vecteur, les auteurs peuvent inspecter directement sa géométrie sans avoir à fouiller dans les flux résiduels. Ils identifient des structures émergentes non explicitement conçues dans l'architecture.

Résultats
Évalués sur le jeu de données FineWeb-Edu (130 millions de paramètres, fenêtre de contexte de 2048) :

• GPN+M à 1 couche : Atteint une perplexité (PPL) de 18,06. Cela se situe à 13 % d'un Transformer++ à 12 couches (PPL 16,05) et à 18 % d'un Gated DeltaNet (GDN) à 10 couches (PPL 15,34).
• GPN+M à 2 couches : Atteint une PPL de 16,95, réduisant l'écart à 6 % du Transformer++ et 11 % du GDN.
• Tâche LAMBADA : Les modèles peu profonds montrent une précision inférieure sur les tâches dépendantes du contexte par rapport aux bases profondes, mettant en évidence un coût lié à la profondeur de calcul par token.
• Sans mémoire (GPN) : Une version sans mémoire matricielle performe nettement moins bien (PPL 23,51), confirmant la nécessité du composant de mémoire associative.

Structures émergentes (Analyse)
L'état à vecteur unique permet une analyse géométrique directe, révélant trois structures qui sont apparues spontanément sans régularisation explicite :

1. Géométrie de rang élevé : Le vecteur d'état est large ($d=2496$) et utilise toute sa largeur. Il présente des dimensions quasi-décorrélées et aucune dimension « morte », avec un rang stable d'environ 40.
2. Direction persistante du token par défaut : La moyenne du vecteur d'état ($\bar{s}$) est un grand décalage fixe qui se décode en une distribution dominée par des tokens fonctionnels fréquents (par exemple, « the », « and », « , »). Cette direction persiste même lorsqu'un biais explicite est ajouté au décodeur, suggérant qu'elle remplit des rôles structurels au-delà du simple décodage (par exemple, l'ancrage de RMSNorm).
3. Séparation mémoire rapide/lente : Les 15 têtes de mémoire architecturalement identiques se divisent spontanément en deux pools pendant l'entraînement : un pool « rapide » (6 têtes) avec des demi-vies de 1 à 3 tokens, et un pool « lent » (9 têtes) avec des demi-vies allant de 5 à 322 tokens. Cela crée un système de mémoire à deux échelles de temps sans conception architecturale explicite.

Signification et limites
L'article affirme qu'une seule couche récurrente peu profonde peut effectuer efficacement la modélisation du langage, atteignant environ 13 % des bases profondes pour des comptes de paramètres équivalents. La signification réside dans la démonstration que la « profondeur » peut être remplacée par une récurrence revisitée dans le temps et qu'un seul vecteur peut soutenir des structures de mémoire complexes et émergentes (séparations rapide/lente) et des propriétés géométriques (ancres persistantes) qui sont généralement cachées dans les flux résiduels profonds.

Cependant, les auteurs sont modestes quant à l'état actuel du travail :

• Pas prêt pour la production : Le modèle s'exécute séquentiellement en Python sans parallélisation le long de la dimension de la séquence, rendant les coûts d'entraînement significativement plus élevés que ceux des Transformers.
• Échelle : Les résultats sont basés sur une seule graine à 130 millions de paramètres, sans balayages d'extension d'échelle.
• Nature descriptive : L'analyse des structures émergentes (comme la direction du token par défaut) est observationnelle ; l'article ne revendique pas de preuve causale de leur utilité par ablation.
• Coûts de contexte : Le modèle peu profond paie un coût mesurable sur les tâches nécessitant un contexte profond (comme LAMBADA), indiquant que bien que l'approche soit viable, les piles profondes conservent toujours des avantages pour des prédictions spécifiques riches en contexte.

Ce travail sert de preuve de concept pour la boucle de « prédiction ancrée » et la viabilité des modèles d'état à vecteur unique, offrant un nouveau prisme pour observer comment les modèles de langage organisent l'information.