Résumé technique : Ce qui importe dans la linéarisation des modèles de langage ?

Énoncé du problème

La linéarisation est devenue une stratégie primaire pour le développement de modèles de langage (LM) efficaces en remplaçant le mécanisme d'auto-attention $O(L^2)$ des Transformers par des mélangeurs de jetons subquadratiques (par exemple, l'attention linéaire, les récurrences à portes). Bien que cette approche vise à conserver les performances des Transformers à grande échelle tout en offrant une inférence en temps linéaire et une surcharge mémoire constante, des lacunes importantes subsistent dans la compréhension du domaine.

La littérature actuelle souffre de conditions d'évaluation hétérogènes, variant en taille de modèle, budgets d'entraînement et objectifs de distillation. Par conséquent, il n'est pas clair quels biais inductifs spécifiques (tels que le fenêtrage/gating, la décroissance ou les mises à jour par la règle delta) sont les mieux adaptés pour hériter des capacités d'un Transformer. De plus, le comportement de mise à l'échelle de ces modèles linéarisés concernant le nombre de paramètres et le budget de jetons, ainsi que leur viabilité pour des tâches en aval comme le réglage d'instructions (instruction tuning) et la récupération de contexte long, restent des domaines de recherche ouverts.

Méthodologie

Pour combler ces lacunes, les auteurs proposent un cadre empirique unifié pour comparer sept architectures subquadratiques représentatives : xLSTM, GLA, Gated DeltaNet, DeltaNet, RetNet, Kimi (KDA) et Linear Attention.

Configuration expérimentale

L'étude isole le mélangeur de jetons comme la variable principale en employant un pipeline de distillation en trois étapes qui aligne les modèles étudiants sur un Transformer pré-entraîné enseignant (famille SmolLM2) tout en maintenant l'enseignant, la distribution de données et la recette d'entraînement constants.

1. Alignement du mélangeur de jetons (Sortie d'Attention) : Le mélangeur de jetons de l'étudiant est aligné avec la sortie d'attention de l'enseignant en utilisant une perte d'erreur quadratique moyenne (MSE). Cela fournit un signal de supervision dense à travers l'ensemble du chemin de mélange.
2. Alignement de l'état caché : Les représentations cachées intermédiaires de l'enseignant et de l'étudiant sont alignées via une perte MSE.
3. Distillation de connaissances de bout en bout : L'étudiant est entraîné de bout en bout en utilisant la distillation de connaissances standard (Prédiction du jeton suivant + Divergence KL).

Dimensions d'évaluation

L'étude évalue les modèles selon trois axes expérimentaux distincts :

1. Lois d'échelle (Scaling Laws) : Les modèles sont entraînés à trois échelles de paramètres (140M, 360M et 1,7B) sur un budget fixe de 3 milliards de jetons pour observer les hiérarchies de performance.
2. Maturité asymptotique : Les modèles sont fixés à environ 1,7 milliard de paramètres et entraînés sur un budget étendu de 10 milliards de jetons pour déterminer si les écarts de performance sont dus à un sous-entraînement ou à des limites architecturales fondamentales.
3. Adaptation en aval : Les étudiants distillés subissent un réglage d'instructions (SFT post-entraînement) et une distillation sensible aux instructions pour évaluer leur capacité à suivre des instructions et à effectuer de la récupération de contexte long (tâches de type "Needle-In-A-Haystack").

Contributions clés et résultats

1. Stabilité des hiérarchies de performance et des exposants d'échelle

L'étude constate que les hiérarchies de performance entre les architectures restent remarquablement stables à travers les échelles (140M à 1,7B).

• Architectures à portes (Gated Architectures) : Les modèles dotés de mécanismes de portes (xLSTM, Gated DeltaNet) surpassent systématiquement les bases sans portes (Linear Attention, RetNet).
• Le point de croisement : Une découverte critique est un "croisement" à l'échelle de 1,7B. Alors que xLSTM domine aux petites échelles, Gated DeltaNet émerge comme le meilleur performeur à 1,7B (94,2 % de récupération de la performance de l'enseignant). Cela suggère que les règles de mise à jour correctrices d'erreurs (règle Delta) possèdent des exposants d'échelle supérieurs par rapport aux alternatives purement additives à mesure que la capacité du modèle augmente.
• Déficiences intrinsèques : L'écart de performance entre l'Attention Linéaire et les modèles à portes s'élargit avec l'échelle (passant de ~14 % à 140M à ~20 % à 1,7B), indiquant que les déficiences de l'attention linéaire additive sont intrinsèques à l'architecture et ne peuvent être surmontées simplement en augmentant le nombre de paramètres.

2. La résolution d'état comme goulot d'étranglement fondamental

L'extension de l'entraînement à 10 milliards de jetons révèle que les écarts de performance relatifs s'établissent tôt et persistent même lorsque les modèles approchent des limites asymptotiques.

• Maturité : La performance plafonne pour les modèles à portes après 6 à 7 milliards de jetons, tandis que les modèles sans portes saturent plus tôt à des niveaux inférieurs. Cela confirme que les écarts ne sont pas dus à un sous-entraînement mais représentent des limites fondamentales dans la manière dont les différents mélangeurs compressent l'information.
• Capacités de récupération : Dans les tâches de type "Needle-In-A-Haystack" (NIAH), la plupart des architectures souffrent d'un effondrement de la résolution de l'état à mesure que la longueur de la séquence augmente. L'Attention Linéaire et RetNet échouent presque totalement au-delà de 512 jetons.
• Supériorité de la règle Delta à portes : Gated DeltaNet est la seule architecture capable de maintenir une précision de récupération significative dans des contextes étendus (par exemple, 81,8 % de précision à 2k jetons), démontant que la règle delta à portes fournit un mécanisme plus efficace pour les mises à jour ciblées de l'état et la prévention de la saturation de l'état.

3. Adaptation aux instructions et distillation

L'étude évalue si le réglage d'instructions peut compenser les faiblesses architecturales.

• Amplification des différences : Le réglage d'instructions post-entraînement améliore la performance de toutes les architectures mais amplifie les différences existantes. Les modèles plus forts (Gated DeltaNet, xLSTM) bénéficient considérablement, tandis que les modèles plus faibles montrent des gains limités.
• Limites de la distillation : La distillation sensible aux instructions (entraînement sur un enseignant ajusté aux instructions) ne procure pas d'avantage constant par rapport au réglage post-entraînement standard et, dans certains cas, dégrade la performance de récupération en contexte long.
• Conclusion sur l'adaptation : L'adaptation aux instructions ne modifie pas les contraintes architecturales fondamentales imposées lors de la linéarisation. La capacité à conserver et à écraser sélectivement l'information reste le facteur dominant régissant la performance en aval.

Signification et affirmations

L'article soutient que pour une linéarisation réussie, les biais inductifs architecturaux sont la contrainte primaire qui ne peut être surmontée simplement en augmentant le calcul d'entraînement ou le budget de jetons de distillation.

• Hiérarchie de conception : L'étude établit une hiérarchie claire pour la conception subquadratique, identifiant que les règles de mise à jour correctrices d'erreurs (règle Delta) combinées au fenêtrage/gating (spécifiquement le gating par tête ou scalaire) sont plus décisives pour le succès que l'accumulation additive ou le gating par canal.
• Résolution d'état : La recherche pose la résolution d'état comme un goulot d'étranglement plus fondamental que le budget de distillation. Les modèles qui ne peuvent pas gérer la précision de leur état interne (en évitant la saturation) échouent à passer à l'échelle efficacement ou à gérer des contextes longs, quelle que soit la durée de l'entraînement.
• Reproductibilité : Pour faciliter la recherche ultérieure, les auteurs publient tous les modèles de base, ajustés aux instructions et distillés pour les instructions, garantissant la reproductibilité de ces conclusions.

Les auteurs concluent que bien que les modèles subquadratiques puissent être entraînés de manière stable, le choix du mélangeur de jetons dicte le plafond de performance, les formulations à règle delta à portes offrant actuellement la voie la plus robuste vers des modèles de langage efficaces capables de gérer des contextes longs.