Résumé Technique : Attention Interdomaine

Énoncé du Problème

Les Transformers et les modèles d'espaces d'états (SSM) profonds représentent deux paradigmes de conception opposés pour la modélisation de séquences. Les mécanismes d'attention standard (Vaswani et al., 2017) utilisent un cache clé-valeur (KV) croissant, acheminant les requêtes via un appariement basé sur le contenu. Bien que cela offre une récupération précise, il engendre des coûts computationnels quadratiques ($O(N_q N)$) et une croissance linéaire de la mémoire ($O(N)$) par rapport à la longueur de la séquence. À l'inverse, les SSM profonds (par exemple, S4, Mamba) compressent l'ensemble du contexte dans un état récurrent de taille fixe, se mettant à jour en $O(1)$ par étape. Cependant, cette compression se fait au prix de la perte de l'appariement direct requête-clé ; l'état fixe n'est pas directement adressable par des requêtes spécifiques, conduisant souvent à des performances plus faibles dans les tâches nécessitant une récupération précise ou une attention basée sur le contenu.

Les architectures hybrides ont tenté d'alterner ces approches, et les variantes d'attention linéaire ont cherché à réduire directement les coûts de l'attention. Cet article propose une voie différente : intégrer un SSM directement dans le mécanisme d'attention afin de maintenir l'efficacité d'un état de taille fixe tout en permettant à chaque requête de s'attacher à l'historique compressé.

Méthodologie : Attention Interdomaine

La contribution centrale est l'Attention Interdomaine, un mélangeur de tokens unifiant l'attention par noyau et les modèles d'espaces d'états. La méthode est dérivée d'une vue de régression par noyau de l'attention et d'une projection de base de style HiPPO.

Dérivation Théorique

1. Carte de Caractéristiques du Noyau : Le noyau d'attention $K(q, k)$ est approximé par une carte de caractéristiques finie $\xi(q)^\top \xi(k)$.
2. Projection de Base : Au lieu de stocker toutes les clés et valeurs, la méthode projette les caractéristiques de clé $\xi(k)$ et les valeurs $v$ sur un ensemble partagé de fonctions de base maintenu par une seule récurrence SSM (spécifiquement de style HiPPO). Cela résulte en un état de taille fixe indépendant de la longueur de la séquence.
3. Coefficients Compressés : Les intégrales de ces projections génèrent trois composantes d'état :
   • $U$ : Coefficients de projection de base pour les clés (variables d'induction).
   • $\Gamma$ : Projections de valeurs.
   • $\eta$ : Constantes de normalisation.
4. Lecture Conditionnée par la Requête : À chaque position, la requête $q$ est mappée dans l'espace de caractéristiques $\xi(q)$. La sortie est calculée en s'attachant aux coefficients compressés plutôt qu'aux tokens individuels :
   $$\tilde{o} = \frac{\sum_m (\xi(q)^\top u_m) \gamma_m}{\sum_m (\xi(q)^\top u_m) \eta_m}$$
   Sous forme matricielle : $\tilde{O} = \frac{F_q U^\top \Gamma}{F_q U^\top \eta}$, où $F_q$ est la matrice de caractéristiques de la requête.

Implémentation Évolutif

L'implémentation pratique est une relaxation apprise de la dérivation théorique, conçue pour la stabilité et l'efficacité :

• Carte de Caractéristiques : Utilise une carte de caractéristiques neuronale de style DeltaNet $\xi(x) = \text{SiLU}(\text{SC}(x)) / \|\text{SiLU}(\text{SC}(x))\|^2$, où SC est une convolution profonde causale.
• Cœur Récurrent : Approxime la projection HiPPO en utilisant une récurrence S4D à diagonale complexe. L'entrée du SSM concatène la caractéristique de clé $\xi(k)$ et la valeur $v$.
• Normalisation : Applique une RMSNorm d'entrée avec des biais par tête appris aux moitiés clé et valeur avant l'entrée dans le SSM.
• Lecture Sans Dénominateur : Pour éviter de mélanger des échelles incompatibles causées par le redimensionnement de l'entrée, le terme de normalisation $\eta$ est supprimé, résultant en une lecture d'attention linéaire : $\tilde{O} = F_q U^\top \Gamma$.
• Architecture : Intégré dans un décodeur de style Llama avec normalisation préliminaire, feedforwards SwiGLU et embeddings de position rotatifs (RoPE).

Contributions Clés

1. Construction : Un mélangeur de tokens à état fixe conditionné par la requête, dérivé de la régression par noyau et de la projection de base HiPPO, avec une distinction claire entre la dérivation théorique idéale et l'implémentation évolutive.
2. Décomposition du Mécanisme : Une étude d'ablation à 125M de paramètres identifie la projection conditionnée par la requête comme le facteur dominant conduisant aux gains de performance par rapport aux mélangeurs S4D standards. L'entrée double clé/valeur contribue à un gain secondaire plus faible.
3. Étude de Mise à l'Échelle : Une évaluation complète sur FineWeb-Edu, de 125M à 1,3B de paramètres, sous un budget d'état récurrent apparié.

Résultats Expérimentaux

L'étude compare l'Attention Interdomaine à un mélangeur de tokens uniquement S4D et à une ligne de base d'attention softmax standard.

• Performance vs S4D : L'Attention Interdomaine améliore le mélangeur de tokens S4D à chaque échelle (125M–1,3B), atteignant une perplexité de validation inférieure de 13 à 16 %.
• Performance vs Softmax : À 1,3B de paramètres, l'Attention Interdomaine dépasse la ligne de base softmax de même recette sur la perplexité de validation (7,98 contre 8,63) et sur une suite de huit tâches de bon sens (+3,03 points de pourcentage d'exactitude moyenne). À 125M, elle est à peu près à égalité avec le softmax.
• Extrapolation de Longueur : Tirant parti de son cœur à état fixe, l'Attention Interdomaine présente un comportement « plat en longueur ». Alors que la perplexité de la ligne de base softmax se dégrade significativement au-delà du contexte d'entraînement (4K), l'Interdomaine maintient des performances stables jusqu'à 14K (3,5 fois le contexte d'entraînement).
• Limites de Récupération : L'article note que l'Attention Interdomaine, comme les autres modèles à état fixe, lutte avec la récupération associative de chaînes exactes (par exemple, les tâches aiguille-dans-une-foin) par rapport à l'attention softmax, identifiant cela comme une limitation structurelle de la compression à état fixe plutôt qu'un défaut spécifique du mécanisme Interdomaine.

Signification et Revendications

L'article revendique que l'Attention Interdomaine comble avec succès le fossé entre l'efficacité des SSM à état fixe et les capacités de récupération basées sur le contenu de l'attention. En projetant des caractéristiques sur une base SSM, elle permet aux requêtes de s'attacher à un historique compressé sans le coût quadratique ni la croissance linéaire de la mémoire de l'attention standard.

La signification principale réside dans la démonstration qu'une projection conditionnée par la requête sur un état récurrent fixe peut surpasser à la fois les SSM purs et, à plus grande échelle, l'attention softmax standard, tout en héritant des propriétés désirables d'extrapolation de longueur des SSM. Les auteurs suggèrent que cette approche offre une voie viable pour la modélisation du langage à long contexte où l'efficacité mémoire et des performances constantes à travers des longueurs de contexte variables sont critiques. Les directions de travail futures mentionnées incluent l'intégration de cœurs à état fixe plus puissants (par exemple, Mamba-3) ou la combinaison de la lecture avec des règles de mise à jour des poids rapides.