Résumé Technique : Modèles de Langage à État d'Espace en Boucle avec Sélection Adaptative de l'État de Sortie

Énoncé du Problème

Les progrès récents des grands modèles de langage (LLM) ont été principalement portés par le passage à l'échelle de la taille des modèles, des données d'entraînement et du calcul intensif lors de l'inférence. Cependant, ce paradigme de mise à l'échelle introduit des défis de déploiement importants, notamment une consommation de mémoire accrue et des coûts d'inférence qui croissent rapidement en raison de la maintenance de larges caches clé-valeur dans les mécanismes d'auto-attention des Transformers. Bien que les modèles d'espace d'état (SSM), particulièrement Mamba-2, offrent une alternative efficace avec une mise à l'échelle favorable pour les contextes longs, il reste incertain si le principe de « profondeur de calcul plutôt que paramètres indépendants » — récemment démontré par les Transformers en boucle (Looped Transformers) — s'applique également aux SSM. De plus, la recherche existante sur les architectures en boucle s'est concentrée presque exclusivement sur les architectures de type Transformer, laissant un vide dans la compréhension de la manière dont le partage récurrent de paramètres fonctionne au sein des architectures Mamba et hybrides Mamba-Transformer.

Méthodologie

Les auteurs proposent les architectures Looped Mamba et Looped Hybrid Mamba–Transformer. Le mécanisme central consiste à partager les paramètres d'un bloc Mamba (ou hybride) unique à travers les couches et à l'appliquer de manière répétée via des boucles explicites à profondeur finie.

Architecture

• Notation : Un modèle est dénoté $N \otimes R$, où $N$ est le nombre de couches aux paramètres indépendants et $R$ est le nombre d'étapes récurrentes. La profondeur effective est $NR$.
• Looped Mamba : Applique de manière répétée une structure de base Mamba $M_\theta$ de $N$ couches partagées.
  • $h^{(r)} = M_\theta(h^{(r-1)})$ pour $r = 1, \dots, R$.
  • La sortie finale est dérivée de l'état caché à l'étape $R$.
• Looped Hybrid : Étend la boucle à une structure contenant à la fois des couches Mamba-2 et un sous-ensemble restreint de couches d'auto-attention causale. La boucle s'applique à l'ensemble du bloc hybride, permettant un raffinement itératif des représentations cachées à travers des mises à jour d'espace d'état et des interactions basées sur l'attention.
• Protocole d'Entraînement :
  1. Pré-entraînement : Optimisé en utilisant un objectif de boucle finale (la perte est calculée uniquement à l'étape $R$).
  2. Étape de la Porte de Sortie (Exit-Gate) 1 : Introduit un objectif de calcul adaptatif régularisé par l'entropie (adapté d'Ouro) pour apprendre une distribution par jeton (token) sur les profondeurs de sortie en utilisant les sorties de toutes les étapes récurrentes.
  3. Étape de la Porte de Sortie (Exit-Gate) 2 : Gèle la structure de base et la tête, et entraîne uniquement la porte de sortie pour encourager la continuation uniquement lorsque la prochaine étape améliore la perte de prédiction.

Sélection Adaptative de l'État de Sortie

Le papier adapte la porte de sortie à deux étapes d'Ouro à Mamba. Lors de l'inférence, un seuil $q$ contrôle la sélection de l'étape de sortie $z_i$ pour chaque jeton en fonction de la probabilité de sortie cumulée.

• Note Sémantique : Les auteurs précisent que dans les SSM, la chaîne d'états de toutes les $R$ étapes doit être préservée en raison des contraintes de continuité d'état. Par conséquent, la « sortie » sélectionne l'étape récurrente dont l'état caché est transmis à la tête de modélisation du langage ; elle ne réduit pas intrinsèquement le temps de calcul effectif ou le nombre d'étapes exécutées sans mécanismes supplémentaires de gestion d'état.

Principales Contributions

1. Étude Empirique Contrôlée : Les auteurs présentent une comparaison rigoureuse des modèles Mamba et hybrides en boucle sous des protocoles iso-paramètres (même $N$, $R$ variable) et iso-FLOPs (même profondeur effective $NR$) contre des modèles de référence non-bouclés (Mamba et Transformer).
2. Évaluation des Tâches de Raisonnement : L'étude valide l'approche sur des tâches de raisonnement synthétiques (Mano et $p$-hop induction) et sur le pré-entraînement de modèles de langage avec des benchmarks en aval.
3. Mécanisme de Sortie Adaptatif : L'introduction d'un mécanisme de sélection d'état de sortie adaptatif contrôlé par un seuil pour les modèles de langage à espace d'état en boucle, démontrant sa capacité à améliorer les performances en aval à des profondeurs intermédiaires.

Résultats Expérimentaux

Tâches de Raisonnement Synthétique

• Tâche Mano (Arithmétique Modulaire) : Le modèle Looped Mamba surpasse systématiquement les modèles de référence non-bouclés avec des paramètres équivalents. Dans le cas le plus difficile ($L=24$), un modèle $4 \otimes 6$ a atteint une précision de 99,35 % contre 7,50 % pour la base $4 \otimes 1$. Sous des conditions iso-FLOPs (profondeur effective de 24), le modèle $4 \otimes 6$ a nettement surpassé la base non-bouclée $24 \otimes 1$ (99,35 % contre 37,45 %), suggérant que le partage récurrent de paramètres fournit un biais inductif plus fort pour le calcul compositionnel qu'une augmentation des couches non partagées.
• $p$-hop Induction (Récupération Récursive) : Les modèles Mamba non-bouclés peu profonds ont éprouvé des difficultés (ex: $1 \otimes 1$ à environ 38-41 %), tandis que l'ajout de boucles a considérablement amélioré les performances. Le modèle $1 \otimes 12$ a atteint une précision quasi parfaite (~99,89 %), indiquant que la profondeur récursive est plus critique que le nombre de paramètres distincts pour la récupération récursive.

Modélisation de Langage et Benchmarks Aval

• Iso-Paramètre : La profondeur récurrente a amélioré la perplexité de validation (PPL) pour les structures de base Mamba-2 et LLaMA-2 sous un budget de paramètres fixe.
• Iso-FLOPs : Les modèles non-bouclés plus profonds ($24 \otimes 1$) ont conservé un avantage en termes de perplexité de validation par rapport aux modèles en boucle avec une profondeur effective identique. Cependant, les modèles en boucle sont restés compétitifs sur les benchmarks en aval, surpassant parfois les modèles non-bouclés plus profonds.
• Sortie Adaptative :
  • Échelle 140M : La porte de sortie adaptative a amélioré les performances en aval. Pour le modèle Mamba-2 de 140M, la porte de sortie de l'étape 2 a obtenu un score moyen de 38,87 avec une étape de sortie sélectionnée moyenne de 3,00, surpassant la base fixed-loop4 (37,79).
  • Échelle 370M : La meilleure performance en aval est restée concentrée près du réglage fixed-loop4. Les variantes avec porte de sortie n'ont pas surpassé les scores de pointe des boucles fixes, suggérant que les modèles plus larges peuvent dépendre davantage de la dernière étape récurrente ou nécessitent un meilleur calibrage de la porte.
  • Compromis : La porte de sortie offre une frontière contrôlable entre la profondeur sélectionnée et la performance, permettant au modèle d'utiliser efficacement les états récurrents intermédiaires plutôt que de forcer une profondeur maximale fixe.

Signification et Revendications

Le papier affirme que la profondeur récurrente offre un axe de mise à l'échelle prometteur et efficace en termes de paramètres pour Mamba et les modèles hybrides de langage. Les conclusions clés incluent :

• Profondeur vs Paramètres : De nombreuses tâches de raisonnement bénéficient d'une plus grande profondeur de calcul plutôt que de paramètres indépendants supplémentaires. La réutilisation récurrente des couches Mamba améliore considérablement les capacités de calcul compositionnel et récursif sans augmenter le nombre de paramètres.
• Viabilité Hybride : Le principe de la profondeur en boucle s'applique efficacement aux architectures hybrides Mamba-Transformer, où les couches d'attention sont bouclées conjointement avec les couches Mamba.
• Sélection Adaptative : Bien que les implémentations actuelles des SSM nécessitent l'exécution complète de l'état, la sélection adaptative de l'état de sortie permet de choisir les profondeurs de prédiction optimales, améliorant les performances en aval à des profondeurs intermédiaires (particulièrement à l'échelle 140M).

Les auteurs font preuve de modestie concernant leurs revendications, notant que les expériences sont limitées à des modèles allant jusqu'à 370 millions de paramètres et que les économies réelles de calcul lors de l'inférence nécessiteront des travaux futurs sur les mécanismes de gestion d'état. Les résultats sont présentés comme des preuves empiriques contrôlées plutôt que comme une caractérisation complète des performances à grande échelle.