Résumé Technique : Attention Linéaire Variationnelle (VLA)

1. Énoncé du Problème

La modélisation de séquences à contexte long fait face à un compromis fondamental entre l'efficacité computationnelle et la stabilité de la mémoire. L'attention standard des Transformers réalise une récupération associative optimale mais souffre d'une complexité temporelle de $O(T^2)$ et d'une utilisation mémoire de $O(Td)$, la rendant inapplicable pour des séquences dépassant $10^5$ jetons. L'attention linéaire résout l'étranglement computationnel en réduisant la complexité à $O(T)$ grâce à des cartes de caractéristiques de noyaux et des mises à jour d'état récurrentes. Cependant, elle introduit un mode de défaillance critique : l'état de mémoire interne $S_t$ croît de manière illimitée en norme de Frobenius ($O(T)$) en raison d'une accumulation additive inconditionnelle. Cette croissance illimitée provoque une interférence progressive entre les associations stockées, dégradant la précision de récupération à mesure que la longueur du contexte augmente.

Bien que DeltaNet tente d'atténuer ce problème via un mécanisme de porte scalaire qui décroît uniformément l'état, il manque d'un contrôle sélectif directionnel. Une porte scalaire ne peut pas oublier préférentiellement les directions récemment écrasées tout en préservant les associations plus anciennes et stables, conduisant à un effondrement rapide de la précision sous de fortes charges mémoire. Les auteurs soutiennent que le problème central n'est pas le coût computationnel mais la stabilité de la mémoire, spécifiquement la géométrie de la mise à jour de la mémoire.

2. Méthodologie : Attention Linéaire Variationnelle

L'article propose l'Attention Linéaire Variationnelle (VLA), qui reformule la mise à jour de la mémoire de l'attention linéaire comme un problème de moindres carrés régularisé en ligne.

2.1 Formulation Principale

Au lieu d'une accumulation inconditionnelle ($S_t = S_{t-1} + v_t k_t^\top$), la VLA cherche une matrice de mémoire $S_t$ qui minimise l'erreur de prédiction pénalisée :
$$S^*_t = \arg \min_S \sum_{s=1}^t \|v_s - S\hat{k}_s\|^2 + \text{tr}(S M_t S^\top)$$
où $\hat{k}_s$ est la clé normalisée et $M_t$ est une matrice de pénalité variant dans le temps encodant la géométrie des clés précédemment observées.

2.2 Géométrie de Pénalité Adaptative

La VLA maintient une matrice de pénalité inverse $A_t = M_t^{-1}$ exactement en utilisant la formule de rang-1 de Sherman-Morrison à un coût de $O(d^2)$ par étape.

• Mécanisme : $M_t$ accumule les produits extérieurs des directions de pénalité apprises. Par conséquent, $A_t$ se contracte dans les directions qui ont récemment reçu une « masse de pénalité » (c'est-à-dire qui ont été écrites).
• Effet : Les sous-espaces récemment écrits reçoivent des mises à jour plus petites, tandis que les sous-espaces intacts restent disponibles. Cela agit comme un analogue à valeur matricielle de la porte scalaire de DeltaNet, permettant des mises à jour de mémoire sélectives directionnellement.

2.3 Règle de Mise à Jour Récursive

La mise à jour récursive exacte dérivée de cette formulation est :

1. Mise à jour de la Pénalité Inverse : $A_t = A_{t-1} - \frac{z_t z_t^\top}{1 + u_t^\top z_t}$ (Sherman-Morrison).
2. Calcul du Vecteur de Porte : $\hat{\alpha}_t = \frac{A_t \hat{k}_t}{\|A_t \hat{k}_t\|}$.
3. Mise à jour de l'État : $S_t = S_{t-1} + (v_t - S_{t-1}\hat{k}_t) \hat{\alpha}_t^\top$.

Crucialement, tant la clé $\hat{k}_t$ que le vecteur de porte $\hat{\alpha}_t$ sont normalisés à une longueur unitaire. La sortie est calculée comme un rapport de récurrences similaire à l'attention linéaire standard, utilisant la carte de caractéristiques ELU+1.

2.4 Efficacité de l'Implémentation

Bien que la mise à jour de Sherman-Morrison introduise un facteur constant de surcharge (environ $5\times$ plus d'opérations que l'attention linéaire standard), les auteurs implémentent une formulation de balayage Blelloch parallèle et un noyau Triton fusionné. Cela fusionne toutes les étapes $T$ en un seul lancement de noyau GPU, éliminant la surcharge de dispatch par jeton et réalisant une accélération de $14\times$ par rapport aux implémentations Python séquentielles.

3. Contributions Clés

3.1 Garanties Théoriques

L'article fournit des preuves rigoureuses pour deux propriétés de stabilité :

1. Croissance Bornée de l'État (Proposition 1) : Sous des entrées bornées, la norme de Frobenius $\|S_t\|_F$ est auto-limitante. Contrairement à l'attention linéaire standard où $\|S_T\|_F = O(T)$, l'état de la VLA converge vers un plateau fini à mesure que les résidus de prédiction s'annulent.
2. Norme Spectrale Unitaires du Jacobien (Proposition 2) : Le jacobien de récurrence $J_t = \partial S_t / \partial S_{t-1} = I - \hat{\alpha}_t \hat{k}_t^\top$ a une norme spectrale exactement égale à 1 pour toutes les longueurs de séquence et dimensions de tête, à condition que les vecteurs soient normalisés unitairement. Cela garantit un flux de gradient stable (pas d'explosion ni de vanissement) à une profondeur arbitraire. Sans normalisation, la norme spectrale croît comme $d_h/\lambda_0$, conduisant à une amplification du gradient (par exemple, $10^{32}$ après 25 étapes pour $d_h=96$).

3.2 Capacité de Mémoire Associative

La VLA maintient la capacité théorique de $d_h$ associations par tête (où $d_h$ est la dimension de la tête). Cependant, contrairement à l'accumulation additive, la VLA achemine les nouvelles écritures vers des directions orthogonales aux sous-espaces récemment utilisés. Cela lui permet de préserver les anciennes associations sans dilution uniforme, maintenant un rappel exact jusqu'à la limite de capacité.

3.3 Efficacité

Le noyau fusionné Triton proposé atteint une mise à l'échelle $O(T)$ et franchit la latence de l'attention softmax à environ 43 000 jetons, la rendant viable pour l'inférence à contexte long.

4. Résultats Empiriques

Les expériences ont été menées sur une colonne vertébrale de Transformer à deux couches partagée en utilisant la tâche de Rappel Associatif Multi-Requêtes (MQAR) et une tâche de Copie.

• Stabilité de la Norme d'État : À $T=1\,000$, l'attention linéaire standard atteint $\|S_T\|_F \approx 1\,600$, tandis que la VLA reste en dessous de 15 (une réduction de $109\times$). DeltaNet est également borné mais manque de la sélectivité directionnelle de la VLA.
• Rappel Associatif (Capacité) : Dans le régime où le nombre de paires $n_{pairs} < d_h$ (spécifiquement $n=24, d_h=32$) :
  • VLA : Maintient une précision de correspondance exacte de 1,000.
  • DeltaNet : S'effondre à 0,010.
  • Attention Linéaire Standard : Chute à 0,074.
  • Softmax : Chute à 0,26.
• Rétention à Contexte Long : Avec $n=8$ fixe et $T$ augmentant jusqu'à 512, la VLA maintient une précision de 1,000. DeltaNet s'effondre à $\approx 0,01$ sur toutes les longueurs en raison de sa porte scalaire qui décroît les anciennes associations avant l'arrivée des requêtes.
• Régime de Surcharge : Lorsque $n > d_h$, la VLA se dégrade plus gracieusement que l'attention linéaire additive car la matrice de pénalité adaptative écrase préférentiellement les directions récemment écrites, préservant partiellement les anciennes associations.

5. Importance et Revendications

L'article postule que la fiabilité à contexte long est un problème de géométrie de mémoire, et non simplement computationnel.

• Contrôle Géométrique : La contribution principale est de démontrer que contrôler quelles directions l'état est autorisé à mettre à jour est plus critique que de simplement réduire le coût computationnel de la mise à jour. En remplaçant l'accumulation inconditionnelle par une mise à jour d'erreur résiduelle régie par une inverse de pénalité adaptative, la VLA réalise une modélisation de séquence à mémoire constante et numériquement stable.
• Lien avec la Théorie Classique : La VLA relie l'attention linéaire aux filtres adaptatifs Moindres Carrés Récursifs (RLS), fournissant une dérivation fondée et une théorie de stabilité absentes dans les approches antérieures de « poids rapides ».
• Portée Pratique : Les auteurs affirment que la VLA fournit un résumé de taille fixe et numériquement stable des $d_h$ associations les plus récemment renforcées, adapté à l'inférence en flux sans infrastructure de cache KV. Ils notent explicitement que la VLA ne surmonte pas la limite fondamentale de capacité de $d_h$ associations par tête mais offre une stabilité et une récupération supérieures dans cette borne.

L'ouvrage conclut que bien que la VLA entraîne une surcharge de facteur constant, la combinaison de la stabilité théorique (jacobien unitaire, état borné) et des performances empiriques (rappel parfait jusqu'à la capacité) l'établit comme une solution robuste pour les Transformers à contexte long où l'attention linéaire standard échoue en raison d'interférences.