Résumé technique : Attention linéaire d'ordre supérieur (HLA)

Énoncé du problème

Le mécanisme d'attention par produit scalaire pondéré standard dans les architectures Transformer engendre une complexité computationnelle et mémoire de $O(n^2)$ par rapport à la longueur de séquence $n$. Ce coût quadratique constitue un goulot d'étranglement majeur pour la mise à l'échelle des modèles de langage autoregressifs vers des contextes longs. Bien que des alternatives existantes telles que l'attention linéaire et les modèles d'espace d'état (SSM) offrent une complexité par token de $O(n)$ ou $O(1)$, elles sont généralement limitées à des approximations du premier ordre ou à des méthodes basées sur des noyaux. Ces restrictions limitent souvent l'expressivité du modèle, l'empêchant de capturer des interactions complexes d'ordre supérieur entre les tokens que les mécanismes d'attention standards peuvent modéliser.

Méthodologie

L'article introduit l'Attention linéaire d'ordre supérieur (HLA), un mécanisme causal et en flux continu qui réalise des interactions d'ordre supérieur grâce à des statistiques suffisantes de préfixe compactes. L'idée centrale est que les opérateurs d'attention d'ordre supérieur peuvent être factorisés en moments d'ordre inférieur (par exemple, des sommes de produits externes de clés), permettant un flux causal exact sans matérialiser de matrices d'attention $n \times n$.

1. HLA du second ordre

La variante du second ordre généralise l'attention linéaire en maintenant un état de taille constante par tête qui capture les moments du second ordre.

• Maintenance de l'état : Au lieu d'une simple somme cumulative, le HLA maintient des résumés de préfixe incluant le moment clé $S^K_t = \sum_{i \le t} k_i k_i^\top$, l'accumulateur requête-valeur $C^{QV}_t = \sum_{i \le t} q_i v_i^\top$, et la masse requête $m^Q_t = \sum_{i \le t} q_i$.
• Causalité stricte : Pour imposer une causalité autoregressive stricte (empêchant les tokens futurs d'influencer les sorties actuelles) sans construire la matrice d'attention complète, les auteurs dérivent des résumés étendus. Plus précisément, ils introduisent des résumés croisés $G_t$ et $h_t$ qui corrigent l'inclusion de termes futurs dans le produit naïf $S^K_t C^{QV}_t$.
• Formulation de la sortie : La sortie non normalisée au temps $t$ est calculée comme $o_t = q_t^\top (S^K_t C^{QV}_t - G_t)$. Cette formulation permet un calcul de $O(d^2 + d d_v)$ par token, indépendant de la longueur de séquence.
• Normalisation : Une variante normalisée par ratio est également définie, divisant le numérateur par un dénominateur masqué dérivé de résumés étendus similaires.

2. Attention linéaire d'ordre supérieur asymétrique (AHLA)

L'article propose une variante asymétrique complémentaire, l'AHLA, qui implémente l'opérateur $Q(K^\top Q)(K^\top V)$. Cette structure pondère les valeurs via un passage unique acheminé par un index de clé intermédiaire, se distinguant de la structure symétrique $AA^\top V$ du HLA standard. Elle utilise un ensemble distinct de résumés en flux continu ($P^{KV}_t, m^K_t, E_t, n_t$) et atteint la même complexité linéaire par token.

3. HLA du troisième ordre

Le cadre est étendu aux interactions du troisième ordre, définies via le produit matriciel $AA^\top A$ où $A = QK^\top$.

• Algèbre en flux continu : La sortie masquée du troisième ordre nécessite le maintien de trois moments de préfixe principaux ($S^K_t, S^Q_t, P^{KV}_t$) et de six résumés croisés supplémentaires ($G^{(1)}_t, \dots, G^{(3)}_t$ et $h^{(1)}_t, \dots, h^{(3)}_t$) pour imposer une causalité stricte.
• Complexité : Le coût de mise à jour en ligne reste linéaire par rapport aux dimensions des caractéristiques, bien que la taille de l'état augmente pour accommoder les moments d'ordre supérieur.

4. Entraînement parallèle via des scans associatifs

Pour remédier à l'inefficacité de l'entraînement de modèles purement récurrents sur les GPU, les auteurs développent un schéma d'entraînement parallèle par blocs.

• Opérateurs associatifs : Ils définissent des opérateurs associatifs (spécifiquement des produits semi-directs) pour la concaténation de segments de tokens. Ces opérateurs prennent en compte les termes croisés entre les segments (par exemple, $S^K_B C^{QV}_A$) qui surgissent lors de la combinaison de blocs.
• Équivalence de scan : En appliquant un scan exclusif de Blelloch sur ces résumés de segments, la méthode produit des activations forward exactement identiques à celles d'une récurrence série. Cela permet un entraînement parallèle sans recourir à une rétropropagation approximative dans le temps.
• Scans du troisième ordre : Pour le cas du troisième ordre, l'opérateur de scan est augmenté de cartes linéaires au niveau des segments ($M^{KQP}, M^{KQm}$) pour gérer la composition des états corrigés du troisième ordre de manière exacte.

Contributions clés

1. Flux continu masqué exact au second ordre : L'article fournit une algèbre complète de résumés étendus qui produit un HLA du second ordre strictement causal avec un coût constant par token. Il inclut des preuves formelles établissant des identités de flux continu masqué et des mises à jour en ligne, distinguant les variantes non normalisées et normalisées par ratio.
2. Scans associatifs correspondant aux activations série : Les auteurs définissent des opérateurs associatifs pour les contextes masqués et non masqués (avec et sans décroissance exponentielle). Ils prouvent que les scans exclusifs standards sur ces opérateurs produisent des activations forward identiques aux récurrences série, permettant un entraînement parallèle efficace.
3. Extension du troisième ordre : Le travail présente l'état masqué complet du troisième ordre, les mises à jour en ligne, un noyau de flux continu strictement causal et un algorithme parallèle par blocs exact. Cela inclut la dérivation des cartes de segments requises pour composer les états du troisième ordre de manière exacte.
4. Variante asymétrique (AHLA) : Un nouveau mécanisme d'attention d'ordre supérieur asymétrique est introduit, offrant un biais inductif différent tout en conservant les mêmes propriétés de flux continu et de complexité.

Résultats et affirmations

L'article se concentre principalement sur la structure algorithmique et les propriétés théoriques plutôt que sur des benchmarks empiriques sur des jeux de données spécifiques. Les affirmations sont fondées sur une dérivation mathématique :

• Complexité : Le HLA atteint une complexité temporelle et mémoire de $O(d^2 + d d_v)$ par token pour l'inférence, où $d$ est la dimension requête/clé et $d_v$ la dimension valeur. Cela est indépendant de la longueur de séquence $n$.
• Exactitude : Le schéma d'entraînement parallèle par blocs est prouvé pour reproduire exactement les activations d'une récurrence série, évitant les erreurs d'approximation courantes dans d'autres méthodes d'entraînement d'attention linéaire.
• Causalité : Les formulations masquées proposées imposent strictement la causalité autoregressive, garantissant que la sortie au temps $t$ dépend uniquement des tokens $j \le t$.
• Évolutivité : En évitant la matérialisation de matrices $n \times n$ et en utilisant des scans associatifs, le HLA est positionné comme un bloc de construction évolutif pour les modèles à contexte long.

Importance

L'article positionne le HLA comme un "bloc de construction fondé sur des principes et évolutif" qui comble le fossé entre le mélange dépendant des données de type attention et l'efficacité des architectures récurrentes modernes.

• Expressivité vs Efficacité : Contrairement à l'attention linéaire du premier ordre qui repose sur des cartes de caractéristiques fixes, le HLA apprend une métrique dépendante des données (via $S^K_t$) et capture des interactions d'ordre supérieur (polynômes de degré 2 ou 3) tout en conservant le paradigme de mise à jour d'état de $O(1)$ par token.
• Implémentation : Le travail souligne que le HLA peut servir de remplacement direct pour la sous-couche d'attention standard dans les blocs Transformer, nécessitant uniquement un échange de noyau tout en préservant les encodages positionnels et la logique de masquage.
• Fondement théorique : En fournissant des identités de flux continu exactes et des opérateurs de scan associatifs pour les termes d'ordre supérieur, l'article établit un cadre algébrique rigoureux pour étendre l'attention linéaire au-delà du premier ordre sans sacrifier les contraintes causales ni l'efficacité de l'entraînement.