Au-delà de l'auto-attention : l'après Transformers

L'architecture Transformer fait tourner l'IA depuis huit ans maintenant. Tous les grands modèles de langage, la plupart des systèmes de génération d'images et un nombre croissant de modèles audio et vidéo reposent sur le mécanisme d'auto-attention introduit dans l'article « Attention Is All You Need ». Mais l'auto-attention a un problème fondamental : son coût de calcul et de mémoire croît quadratiquement avec la longueur de la séquence. Doubler la longueur de l'entrée, c'est quadrupler le coût.

Pour des séquences courtes, ça ne pose pas de problème. Pour les fenêtres de contexte de 128K tokens vers lesquelles on tend — et les fenêtres d'un million de tokens que les gens veulent — c'est un goulot d'étranglement sérieux. Une vague de recherches explore des alternatives : les résiduels d'attention qui réutilisent le calcul d'une couche à l'autre, les variantes d'attention linéaire qui suppriment le coût quadratique, et les architectures hybrides qui mêlent attention et mécanismes moins coûteux. Le Transformer ne disparaît pas, mais il est en train d'être remodelé.

Pourquoi l'auto-attention coûte cher

Pour comprendre les alternatives, il faut savoir ce que calcule réellement l'auto-attention. Étant donnée une séquence de N tokens, l'auto-attention calcule un score de pertinence entre chaque paire de tokens. Token 1 contre token 2, token 1 contre token 3, ..., token 1 contre token N, puis token 2 contre tous les autres, et ainsi de suite. Ça fait N² paires.

import torch
import torch.nn.functional as F
def self_attention(Q, K, V):
"""
Standard self-attention.
Q, K, V: (batch, seq_len, d_model)
The attention matrix is seq_len × seq_len.
For seq_len = 1024:   ~1M entries   (manageable)
For seq_len = 32768:  ~1B entries   (expensive)
For seq_len = 131072: ~17B entries  (very expensive)
"""
d_k = Q.size(-1)
# This matmul creates the N×N attention matrix
scores = torch.matmul(Q, K.transpose(-2, -1)) / (d_k ** 0.5)
weights = F.softmax(scores, dim=-1)
return torch.matmul(weights, V)

À 4K tokens, la matrice d'attention compte 16 millions d'entrées — rien d'insurmontable pour les GPU modernes. À 128K tokens, elle en compte 16 milliards. À un million de tokens, on dépasse le trillion. Même avec Flash Attention (qui ne réduit pas le calcul, mais améliore considérablement les accès mémoire), la croissance quadratique finit par l'emporter.

C'est pour ça que les premiers Transformers étaient limités à 512 ou 1024 tokens. Chaque génération de matériel et d'optimisations a repoussé le plafond, mais on se heurte à un mur mathématique. Une mise à l'échelle linéaire (O(N)) serait fondamentalement meilleure qu'une mise à l'échelle quadratique (O(N²)), et c'est ce que poursuit la plupart des architectures alternatives.

Résiduels d'attention : réutiliser ce qui a déjà été calculé

L'une des approches les plus pragmatiques pour réduire le coût de l'attention ne remplace pas l'attention : elle rend chaque couche d'attention moins chère en réutilisant le calcul des couches précédentes.

L'observation : dans un Transformer profond (disons 32 couches), les motifs d'attention des couches adjacentes se ressemblent souvent énormément. La couche 15 et la couche 16 ont tendance à porter leur attention sur des positions similaires, avec de légers ajustements. Recalculer la matrice N² complète à chaque couche est redondant : une grande partie du travail a déjà été faite une couche plus tôt.

Les résiduels d'attention exploitent ça en calculant un motif d'attention « résiduel » : la différence entre ce sur quoi la couche veut porter son attention et ce que la couche précédente a calculé. Si la différence est faible (ce qui est généralement le cas dans les couches intermédiaires), le calcul est moins coûteux. Le motif d'attention complet est la somme du motif de la couche précédente et du résiduel de la couche courante.

C'est comparable à la compression vidéo : plutôt que de stocker chaque image indépendamment, on stocke une image clé puis une série de différences (résiduels) par rapport à cette image clé. Ces différences sont généralement bien plus petites que l'image complète, la compression est donc nettement meilleure.

En pratique, les résiduels d'attention réduisent le coût de calcul de l'attention de 30 à 50 % dans les couches intermédiaires des modèles profonds, avec un impact minime sur la qualité. Les premières et dernières couches ont toujours besoin d'un calcul d'attention complet (leurs motifs sont plus distincts), mais les couches intermédiaires — qui sont majoritaires — gagnent beaucoup en vitesse.

Attention linéaire : en finir avec le coût quadratique

Les variantes d'attention linéaire tentent de reformuler l'attention pour qu'elle croisse en O(N) au lieu de O(N²). Le principe général : au lieu de calculer explicitement la matrice d'attention N×N, on cherche un moyen d'obtenir la même sortie (ou à peu près) avec des opérations linéaires.

L'astuce mathématique repose sur la décomposition en noyau de la softmax. L'attention standard calcule softmax(QK^T)V. Si l'on remplace la softmax par une autre fonction noyau décomposable en φ(Q) · φ(K)^T, on peut réordonner le calcul : au lieu de (φ(Q) · φ(K)^T) · V (qui produit un intermédiaire N×N), on calcule φ(Q) · (φ(K)^T · V) (qui produit un intermédiaire d×d, où d est la dimension du modèle). Comme d << N pour les longues séquences, c'est nettement moins coûteux.

def linear_attention(Q, K, V, feature_map=None):
"""
Linear attention via kernel feature maps.
Cost: O(N * d^2) instead of O(N^2 * d)
"""
if feature_map is None:
# ELU+1 is a common choice (from Katharopoulos et al.)
feature_map = lambda x: F.elu(x) + 1
Q = feature_map(Q)  # (batch, seq_len, d)
K = feature_map(K)  # (batch, seq_len, d)
# Key insight: compute K^T @ V first (d × d matrix)
# instead of Q @ K^T first (N × N matrix)
KV = torch.einsum('bnd,bnm->bdm', K, V)  # (batch, d, d)
# Then multiply by Q
output = torch.einsum('bnd,bdm->bnm', Q, KV)  # (batch, N, d)
# Normalize
Z = torch.einsum('bnd,bd->bn', Q, K.sum(dim=1))  # normalization
output = output / Z.unsqueeze(-1)
return output

Le piège : remplacer la softmax par une autre fonction noyau modifie la distribution de l'attention, et les modèles entraînés avec une attention softmax ne se transposent pas forcément bien à l'attention linéaire. L'écart de qualité s'est nettement réduit — les variantes récentes atteignent 95 à 98 % de la qualité de l'attention softmax — mais il persiste, surtout sur les tâches qui exigent une recherche précise à longue distance.

Modèles à espace d'états : un autre paradigme

Les modèles à espace d'états (SSM) comme Mamba adoptent une approche fondamentalement différente. Au lieu de calculer des relations entre paires de tokens, ils traitent la séquence par récurrence, en maintenant un état caché de taille fixe mis à jour à chaque token. C'est intrinsèquement en O(N) : traiter deux fois plus de tokens prend deux fois plus de temps, pas quatre fois.

L'innovation des SSM modernes consiste à faire dépendre les paramètres de la récurrence de l'entrée (espaces d'états sélectifs). Le modèle acquiert ainsi une forme d'attention basée sur le contenu : il peut « choisir » quelles informations retenir et lesquelles oublier, sans le coût quadratique. Les modèles de type Mamba égalent la qualité des Transformers sur de nombreux benchmarks, tout en étant nettement plus rapides sur les longues séquences.

Le compromis : les SSM traitent les tokens séquentiellement, ce qui les rend plus difficiles à paralléliser à l'entraînement que les Transformers (qui traitent tous les tokens simultanément). L'efficacité d'entraînement compte : un modèle deux fois plus rapide en inférence mais trois fois plus lent à entraîner n'est pas forcément gagnant, puisque l'essentiel du calcul total part dans l'entraînement.

Architectures hybrides : la voie pragmatique

La tendance actuelle dans les modèles en production, ce sont les architectures hybrides qui combinent différents mécanismes d'attention. Le raisonnement est simple : différentes parties d'un modèle profitent de types de calcul différents.

• Attention complète pour le raisonnement global. Certaines couches doivent porter leur attention sur toute la séquence, pour trouver un contexte pertinent situé à des milliers de tokens. Ces couches utilisent une auto-attention standard (éventuellement optimisée avec Flash).
• Attention locale pour le contexte proche. Beaucoup de couches se concentrent surtout sur les tokens voisins (attention à fenêtre glissante). Une fenêtre fixe de 256 à 1024 tokens réduit le coût à O(N·W), où W est la taille de la fenêtre.
• Attention linéaire pour le contexte large. Certaines couches doivent agréger des informations sur toute la séquence sans avoir besoin de poids d'attention précis. L'attention linéaire le permet pour un coût en O(N).
• Couches SSM pour le traitement séquentiel. Les couches de type Mamba traitent efficacement les dépendances séquentielles sans aucun calcul d'attention.

Des modèles comme Jamba (AI21) et diverses architectures de recherche alternent ces mécanismes selon le rôle de chaque couche. Les premières couches utilisent une attention locale (traitement de la syntaxe et des motifs locaux). Les couches intermédiaires utilisent l'attention linéaire ou les SSM (construction de représentations plus larges). Quelques couches stratégiques utilisent l'attention complète (raisonnement global et recherche). On obtient ainsi une mise à l'échelle quasi linéaire globale, tout en préservant la qualité du modèle, qui nécessite une part d'attention complète.

Ce que les développeurs doivent surveiller

Si vous développez des applications sur des modèles de langage, les changements architecturaux en coulisses ont des effets concrets sur votre travail.

• Les fenêtres de contexte vont continuer de grandir. À mesure que le coût de l'attention baisse, les fenêtres de contexte s'élargissent. Cela change l'architecture des applications : au lieu de construire des pipelines RAG complexes pour faire tenir le contexte pertinent dans une fenêtre de 4K, on pourrait tout simplement tout injecter dans un prompt d'un million de tokens. La simplicité est séduisante, mais la latence et le coût diffèrent selon les architectures.
• Les profils de latence évoluent. Les Transformers ont une latence relativement stable jusqu'à un certain point, puis elle augmente de façon quadratique. Les modèles à attention linéaire et SSM présentent des hausses de latence plus progressives et linéaires. Pour les applications où le temps de réponse compte, comprendre le comportement de mise à l'échelle de votre modèle est essentiel.
• Les différences de qualité dépendent de la tâche. Les modèles à attention linéaire peuvent être légèrement moins bons sur les tâches exigeant une recherche précise à des positions données dans de longs contextes (« quel était le troisième élément de la liste en page 47 ? »). Ils s'en sortent aussi bien sur les tâches de compréhension générale. Connaissez votre cas d'usage.
• L'optimisation de l'inférence compte davantage. À mesure que les modèles gagnent en complexité architecturale (en mêlant différents types d'attention), les moteurs d'inférence doivent gérer efficacement des calculs hétérogènes. vLLM, TensorRT-LLM et des frameworks similaires s'adaptent, mais les architectures sur mesure peuvent ne pas être prises en charge immédiatement.

Le Transformer n'est pas remplacé : il évolue. L'auto-attention reste le mécanisme le plus expressif dont nous disposons pour modéliser les relations entre tokens. Mais il n'a pas besoin d'être utilisé partout, à chaque couche, avec un coût N² complet. Les modèles des prochaines années utiliseront l'attention avec précision : pleine puissance là où elle compte le plus, alternatives moins coûteuses partout ailleurs. Le résultat : des modèles plus rapides, capables de gérer des contextes plus longs et moins chers à faire tourner, tout en égalant ou dépassant la qualité actuelle. Ça vaut le coup de les suivre de près.