Résumé Technique : MESANET : MODÉLISATION DE SÉQUENCES PAR ENTRAÎNEMENT OPTIMAL AU TEMPS DE TEST (TEST-TIME TRAINING)

Énoncé du Problème

La modélisation de séquences est actuellement dominée par les architectures de transformeurs causaux utilisant l'attention auto-softmax. Bien qu'efficaces, ces modèles souffrent d'une croissance linéaire des exigences en mémoire et en calcul par rapport à la longueur de la séquence lors de l'inférence. Des efforts récents ont cherché à remédier à cela en linéarisant l'opération softmax, ce qui a donné naissance à des modèles de réseaux de neurones récurrents (RNN) avec des coûts de mémoire et de calcul constants (par exemple, DeltaNet, Mamba, xLSTM). Ces modèles sont souvent unifiés par la perspective que leur dynamique récurrente dérive d'un objectif de régression en contexte, optimisé approximativement via des règles d'apprentissage en ligne (par exemple, des étapes de descente de gradient). Cependant, ces approches existantes reposent généralement sur des mises à jour incrémentielles (méthodes du premier ordre) plutôt que sur la résolution des poids optimaux à chaque pas de temps, ce qui peut limiter leur capacité à s'adapter rapidement à de nouvelles preuves tout en conservant les informations passées. De plus, les itérations précédentes de couches d'entraînement optimal au temps de test (telles que l'originale couche Mesa) étaient limitées par une instabilité numérique lors de l'utilisation de portes d'oubli (forget gates) et par un manque de parallélisabilité, empênant un entraînement scalable.

Méthodologie

Les auteurs introduisent MesaNet, une architecture de modèle de langage centrée sur une nouvelle couche Mesa, numériquement stable et parallélisable par blocs (chunkwise). Cette couche est conçue pour effectuer un entraînement optimal local au temps de test à chaque point temporel.

La Couche Mesa

Contra contrairement aux RNN linéaires standards qui mettent à jour leur état via une seule étape de descente de gradient (par exemple, DeltaNet), la couche Mesa définit la carte linéaire $\Phi_t$ comme la solution d'un problème d'optimisation au temps de test. Plus précisément, elle minimise une erreur quadratique cumulée et régularisée sur toutes les données vues jusqu'à présent ($e_{1:t}$) :
$$\hat{\Phi}^{mesa}_t = \arg \min_{\Phi} \left( \frac{1}{2} \sum_{t'=1}^t \zeta_{tt'} \|v_{t'} - \Phi k_{t'}\|^2 + \frac{1}{2} \text{Tr}(\Phi \Lambda \Phi^\top) \right)$$
où $\zeta_{tt'}$ représente les facteurs d'oubli cumulés, et $\Lambda$ est un régularisateur symétrique défini positif.

La sortie est calculée par $\Delta^{mesa}_t = \hat{\Phi}^{mesa}_t q_t$. Crucialement, au lieu d'utiliser les moindres carrés récursifs (qui souffrent d'instabilité numérique avec les portes d'oubli et ne sont pas parallélisables), les auteurs résolvent pour la réponse optimale à la requête $q^*_t = (H_t + \Lambda)^{-1} q_t$ en utilisant la méthode du Gradient Conjugué (CG).

• Parallélisation : Les auteurs tirent parti de l'équivalence entre la méthode CG sur plusieurs pas de temps et l'attention linéaire à porte (Gated Linear Attention - GLA). En maintenant des variables d'état $G_t$ et $H_t$ et en résolvant le système linéaire en parallèle sur des blocs de la séquence, la couche atteint une complexité en $O(T)$ pendant l'entraînement et en $O(1)$ pendant l'inférence.
• Calcul Dynamique : Le nombre d'itérations de CG ($k$) nécessaires pour atteindre un critère d'arrêt varie selon la séquence et le jeton (token). Cela permet au modèle d'allouer dynamiquement le calcul au temps de test, en échangeant le coût d'inférence contre la performance.

Architecture (MesaNet)

MesaNet adopte une architecture de transformeur de type décodeur uniquement standard (similaire à Llama 2) mais remplace le composant de mélange de séquence (typiquement la Multi-Head Attention) par la couche Mesa. L'architecture comprend :

• Blocs Résiduels : Empilement de $N$ blocs contenant une couche de mélange de séquence (Mesa) et une couche de mélange de canaux (MLP SwiGLU).
• Prétraitement : Les clés (keys) et les requêtes (queries) sont normalisées et convoluées avec une petite fenêtre avant la couche Mesa.
• Variante Hybride : Un modèle "Hawk-Mesa" est également exploré, alternant entre des couches LRU (Hawk) et des couches Mesa pour combiner les forces de la modélisation de séquences courtes et longues.

Principales Contributions

1. Une Couche Mesa Parallélisable et Stable : L'article présente une version de la couche Mesa qui surmonte le goulot d'étranglement séquentiel et l'instabilité numérique de la formulation originale. Elle utilise un solveur CG avec parallélisation par blocs, permettant un entraînement efficace sur le matériel moderne (GPU/TPU) tout en maintenant les avantages théoriques de la régression optimale au temps de test.
2. Allocation Dynamique du Calcul au Temps de Test : La couche invoque explicitement un solveur externe où le nombre d'itérations dépend de la difficulté du système linéaire pour une séquence donnée. Cela fournit un mécanisme fondé sur des principes pour varier le calcul d'inférence en fonction de la complexité de l'entrée.
3. Évaluation Empirique Étendue : Les auteurs entraînent des MesaNets à des échelles de 140M, 440M et 1B de paramètres sur le jeu de données SlimPajama. Ils fournissent une comparaison rigoureuse, de type "pommes avec pommes", contre Mamba2, xLSTM, DeltaNet, Gated DeltaNet et les Transformeurs, en contrôlant l'architecture, le tokenizer et les hyperparamètres.

Résultats

• Performance de Modélisation de Langage : Les MesaNets atteignent une perplexité (PPL) de validation plus faible que les autres bases RNN (Mamba2, xLSTM, DeltaNet) à toutes les échelles. À l'échelle de 1B de paramètres, MesaNet égale ou dépasse légèrement la performance d'un transformeur de base avec la même architecture.
• Analyse de la Position des Jetons : Une analyse détaillée de la log-vraisemblance négative (NLL) par rapport à la position du jeton révèle un motif distinct :
  • Jetons Précoces : La plupart des RNN, y compris MesaNet, surpassent les Transformeurs sur les premiers jetons d'une séquence.
  • Contextes Longs : Alors que les Transformeurs excellent sur les jetons ultérieurs, MesaNet et Hawk-Mesa étendent leur avantage de manière significative (au-delà de 512 jetons) par rapport aux autres RNN, bien qu'ils finissent par être dépassés par les Transformeurs à des profondeurs très longues.
• Benchmarks Aval (Downstream) :
  • Raisonnement Global : MesaNet surpasse tous les autres RNN sur les benchmarks nécessitant des dépendances à longue portée (ex: Lambada, HellaSwag, RACE), bien qu'il reste derrière le transformeur de référence.
  • Rappel In-Context : MesaNet surpasse les autres RNN et même un modèle de Sliding Window Attention (SWA) avec une fenêtre de 1024 jetons, mais reste derrière le transformeur à attention complète.
  • Apprentissage à Quelques-Tir (Few-Shot Learning) : MesaNet démontre de solides performances sur les tâches de mélange de mots (word scrambling) et de traduction, dépassant souvent les transformeurs sur les tâches de mélange.
• Efficacité : Malgré la nécessité de FLOPs supplémentaires pour le solveur CG pendant l'inférence, MesaNet maintient un débit de jetons compétitif sur les GPU H100 et les TPUv5, particulièrement lorsqu'on utilise un nombre modéré d'étapes de CG (ex: 10-30).

Signification et Revendications

L'article positionne MesaNet comme une étape importante vers un entraînement optimal au temps de test au sein des réseaux de neurones. Les auteurs affirment qu'en résolvant des problèmes d'optimisation séquentiels (spécifiquement, trouver les meilleurs poids de mémoire associative linéaire) au moment de l'inférence, le modèle peut s'adapter plus efficacement aux nouvelles preuves qu'un modèle utilisant de simples mises à jour de gradient.

Le travail met en évidence un compromis : les gains de performance en perplexité et en raisonnement aval proviennent au prix d'un calcul d'inférence accru (FLOPs). Les auteurs suggèrent que cela s'aligne sur les tendances récentes consistant à augmenter le calcul au temps de test pour améliorer les performances. Ils soutiennent que la couche Mesa offre un mécanisme flexible pour allouer dynamiquement ce calcul, capable de résoudre des problèmes séquentiels complexes au sein même du réseau.

L'article conclut que, bien que les MesaNets fassent actuellement face à un surcoût de calcul, leur capacité à surpasser les autres RNN sur la compréhension de contextes longs et les tâches de raisonnement global suggère que l'entraînement optimal local au temps de test est une direction prometteuse pour la modélisation de séquences efficace, capable potentiellement de combler l'écart entre l'efficacité des RNN et la performance des Transformeurs. Les auteurs reconnaissent que le coût de calcul de l'implémentation actuelle est un inconvénient et suggèrent que des travaux futurs pourraient impliquer des stratégies de "warm-start" pour l'optimisation ou des stratégies d'inférence hybrides pour l'atténuer.