Sparse Attention et Efficient Transformers : les tendances architecturales

L’attention standard est quadratique. Chaque token fait attention à tous les autres tokens, ce qui rend le coût computationnel croître avec le carré de la longueur de la séquence. À 128K tokens, c’est 16 milliards de calculs d’attention par couche. La quête pour briser cette barrière quadratique a produit certaines des innovations architecturales les plus importantes des dernières années.

Le problème quadratique

En attention multi-tête standard, chaque token calcule un score de similitude avec chaque autre token dans la séquence. Pour une séquence de longueur N, cela produit une matrice d’attention N×N. Les coûts de mémoire et de calcul s’échellent comme O(N²), ce qui devient prohibitif pour les longues séquences.

À 2K tokens, l’attention quadratique est gérable. À 8K, c’est cher. À 128K, c’est le coût dominant de toute la passe avant. Et pourtant les applications exigent de plus en plus de long contextes : analyse de documents, compréhension de code au niveau du dépôt, raisonnement multi-documents. L’architecture doit s’adapter.

Sliding Window Attention

Le motif d’attention clairsemée le plus simple est la fenêtre glissante. Chaque token ne fait attention qu’à ses W voisins les plus proches, où W est la taille de la fenêtre. Cela réduit la complexité de O(N²) à O(N×W) — linéaire dans la longueur de la séquence pour une fenêtre fixe.

Mistral a été pionnier dans cette approche dans les modèles de production. Mistral 7B utilise une fenêtre glissante de 4096 tokens dans ses couches d’attention. L’information se propage à travers la séquence complète à travers plusieurs couches — après L couches, le champ réceptif efficace est de L×W tokens. Avec 32 couches et une fenêtre 4096, c’est théoriquement 128K tokens de flux d’information.

Le compromis est que les dépendances longue portée deviennent indirectes. Le token A à la position 0 ne peut influencer le token B à la position 100K que par une chaîne de calculs d’attention intermédiaires à travers les couches. L’attention directe entre les tokens distants est impossible. Pour de nombreuses tâches c’est bien ; pour les tâches nécessitant un rappel précis longue portée, cela peut dégrader la qualité.

Grouped Query Attention (GQA)

GQA n’est pas l’attention clairsemée au sens traditionnel, mais c’est l’amélioration d’efficacité la plus impactante dans les transformers modernes. Au lieu que chaque tête d’attention ait ses propres projections clé et valeur, GQA partage les têtes clé-valeur à travers des groupes de têtes de requête.

L’attention multi-tête standard avec 32 têtes a 32 requête, 32 clé, et 32 projections de valeur. GQA avec 8 groupes KV a 32 requête mais seulement 8 clé et 8 projections de valeur. Les économies de mémoire pour le cache KV sont proportionnelles — environ 4x réduction dans cet exemple.

Cela compte énormément pour l’inférence. Le cache KV stocke les états clé et valeur pour tous les tokens précédents et est souvent le goulot d’étranglement mémoire. Le réduire par 4x signifie servir 4x plus de requêtes concurrentes sur le même matériel, ou gérer 4x les plus longues séquences. Llama 3, Qwen 2.5, et la plupart des modèles récents utilisent GQA par défaut.

Multi-Query Attention (MQA)

MQA prend GQA à l’extrême : toutes les têtes de requête partagent une seule tête clé et une seule tête valeur. Le cache KV se rétrécit par le nombre complet de têtes — 32x pour un modèle 32-tête. C’est maximalement efficace mais peut réduire la qualité du modèle, particulièrement pour les tâches nécessitant des motifs d’attention divers.

En pratique, GQA avec 4-8 groupes a émergé comme le point sucré, offrant la plupart des avantages d’efficacité de MQA avec une perte de qualité minimale. Le pur MQA est utilisé dans certains modèles optimisés pour la vitesse où la latence compte plus que la qualité.

Linear Attention et State Space Models

L’approche plus radicale remplace complètement l’attention quadratique. L’attention linéaire reformule le calcul d’attention pour éviter la matrice N×N, réalisant O(N) complexité. Les variantes comme RetNet et RWKV utilisent différentes stratégies de linéarisation, échangeant l’expressivité complète de l’attention softmax pour l’efficacité computationnelle.

Les State Space Models (SSMs) comme Mamba prennent un chemin différent. Au lieu de l’attention sur la séquence complète, ils maintiennent un état caché de taille fixe qui est mis à jour à mesure que chaque nouveau token arrive. C’est inhéremment O(N) et nécessite une mémoire constante indépendamment de la longueur de la séquence.

Mamba et ses successeurs (Mamba-2, Jamba) ont montré que les SSMs peuvent correspondre à la qualité du transformer sur de nombreux benchmarks tout en étant nettement plus rapides pour les longues séquences. Cependant, ils peinent sur les tâches nécessitant la récupération précise d’informations de positions spécifiques dans le contexte — le problème « aiguille dans une botte de foin » que l’attention traite bien.

Architectures hybrides

Le consensus émergent est que les architectures hybrides combinant l’attention et les couches linéaires offrent le meilleur compromis. Jamba (AI21) alterne entre les couches de transformer avec l’attention complète et les couches Mamba avec la complexité linéaire. Les couches d’attention gèrent la récupération précise ; les couches Mamba gèrent la modélisation longue portée efficace.

Cette approche hybride s’échelle mieux que les transformers purs pour les longues séquences tout en maintenant les capacités de récupération que les SSMs purs manquent. Le ratio de couches d’attention aux couches linéaires est un paramètre modifiable — plus de couches d’attention pour les tâches récupération-lourdes, plus de couches SSM pour l’efficacité.

Flash Attention : la révolution de l’implémentation

Parfois, la meilleure amélioration architecturale n’est pas architecturale du tout. Flash Attention ne change pas ce que le modèle calcule — cela change comment c’est calculé. En carrelant le calcul d’attention et en gardant les données dans la rapide SRAM plutôt que la lente HBM, Flash Attention réalise des accélérations de 2-4x avec zéro changement de qualité.

Flash Attention a rendu les longs contextes pratiques avant que l’attention clairsemée ne le fasse. Un contexte 32K qui était prohibitif en mémoire avec l’attention naïve fonctionne confortablement avec Flash Attention sur le même matériel. Combiné avec GQA et les fenêtres glissantes, cela permet les longueurs de contexte 128K+ que les modèles modernes supportent.

La leçon : avant de redessiner l’architecture, optimise l’implémentation. L’écart entre l’efficacité théorique et pratique dans le calcul d’attention était énorme, et la fermeture de cet écart à travers une meilleure ingénierie a livré plus d’impact réel que de nombreuses innovations architecturales.

Pour les approfondissements sur l’architecture du transformer et la recherche d’efficacité, visite Laeka Research.