Résumé Technique : Attention Stochastique

Énoncé du Problème

Les mécanismes d'attention efficaces sont essentiels pour mettre à l'échelle les modèles Transformer, mais les approches existantes font face à un compromis entre le coût computationnel et la couverture du champ récepteur. L'Attention à Fenêtre Glissante (SWA) réduit la complexité à $O(nw)$ par couche en restreignant les jetons à une fenêtre locale de taille $w$. Cependant, cette localité déterministe limite la croissance effective du champ récepteur à une progression linéaire ($\ell w$ après $\ell$ couches), laissant de grandes portions de la séquence inaccessibles lorsque $w \ll n$. Bien que des remèdes tels que des jetons globaux, des motifs épars conçus à la main ou un routage au niveau des blocs existent, ils introduisent souvent une complexité architecturale ou échouent à fournir des garanties de couverture prouvables. À l'inverse, l'attention complète offre une connectivité globale mais engendre des coûts prohibitifs de $O(n^2)$.

L'article s'inspire du connectome de la drosophile (Drosophila melanogaster), un réseau d'environ 130 000 neurones avec une probabilité de connexion de seulement 0,02 %. Malgré cette éparsité, le réseau atteint une longueur de chemin moyenne courte d'environ 4,4 sauts. Cette topologie de « petit monde » résulte de l'interplay entre un regroupement local dense et des connexions à longue portée largement distribuées qui fonctionnent comme des raccourcis stochastiques. Les auteurs postulent qu'un flux d'information global efficace peut émerger d'un calcul local structuré augmenté par des connexions à longue portée distribuées et stochastiques, plutôt que d'exiger une connectivité dense de tous à tous.

Méthodologie : Attention Stochastique (SA)

Les auteurs proposent l'Attention Stochastique (SA), une amélioration sans paramètres pour la SWA qui imite la stratégie de routage du connectome.

Mécanisme Central

La SA transforme une fenêtre locale fixe en une fenêtre globale stochastique grâce à un processus en trois étapes appliqué avant et après le calcul standard de la SWA :

1. Permutation : Une permutation aléatoire $\sigma \sim \text{Uniform}(S_n)$ est tirée indépendamment pour chaque couche. Les représentations de la séquence de jetons ($Q, K, V$) sont réarrangées selon $\sigma$.
2. Attention à Fenêtre : La SWA standard est appliquée dans l'espace permuté. Parce que la permutation est aléatoire, la fenêtre fixe de taille $w$ couvre désormais un sous-ensemble aléatoire de la séquence originale.
3. Annulation de la Permutation : L'ordre original est restauré en appliquant $\sigma^{-1}$ à la sortie.

Dans l'espace des jetons original, la position $i$ porte attention à un voisinage aléatoire $\tilde{N}^\sigma_w(i)$ de taille attendue $w$, uniformément distribué sur toute la séquence indépendamment de la distance originale $|i-j|$. Cela crée une probabilité de connexion par paire d'environ $w/n$ pour toute paire de jetons, analogue aux raccourcis à longue portée distribués dans le connectome.

Propriétés Théoriques

• Croissance du Champ Récepteur : Alors que le champ récepteur de la SWA croît linéairement ($O(n/w)$ couches pour une couverture complète), les permutations échantillonnées indépendamment à travers les couches dans la SA produisent des champs récepteurs à croissance exponentielle. L'article démontre qu'une couverture complète de la séquence est atteinte en $O(\log_w n)$ couches.
• Complexité : Les opérations de permutation et de permutation inverse sont des réarrangements d'indices de $O(n)$. Le calcul de l'attention reste à $O(nw)$. Ainsi, la SA maintient le même budget par couche de $O(nw)$ que la SWA standard.
• Contraintes Causales : Dans les contextes autorégressifs, le masque causal est appliqué dans l'espace permuté, garantissant que les jetons ne portent attention qu'aux prédécesseurs causalement accessibles tout en permettant toujours des connexions à longue portée stochastiques.

Combinaison SA à Portes + SWA

Pour reproduire le régime de « petit monde » (fort regroupement local + chemins courts), les auteurs proposent une architecture à double voie combinant SA et SWA :
$$Y = g_{sa} \odot Y_{sa} + g_{swa} \odot Y_{swa}$$
où $Y_{sa}$ et $Y_{swa}$ sont les sorties des voies stochastique et à fenêtre glissante, respectivement. Les poids de porte ($g_{sa}, g_{swa}$) sont calculés via des projections sigmoïdes indépendantes, permettant au modèle d'apprendre à équilibrer la cohérence locale (provenant de la SWA) et la couverture globale (provenant de la SA). Cette combinaison ajoute des paramètres négligeables et maintient une complexité de $O(nw)$.

Contributions Clés

1. Attention Stochastique (SA) : Une amélioration intégrable et sans paramètres pour la SWA qui applique des permutations aléatoires pour atteindre une croissance exponentielle du champ récepteur ($O(\log_w n)$) dans le même budget computationnel de $O(nw)$.
2. SA à Portes + SWA : Une architecture à double voie qui reproduit théoriquement et empiriquement le régime de petit monde observé dans les connectomes biologiques, équilibrant le regroupement local et les raccourcis stochastiques.
3. Validation Empirique :
   • Pré-entraînement : Sur des modèles de 360M de paramètres entraînés à partir de zéro, la combinaison SA à portes + SWA atteint la meilleure précision moyenne en zéro-shot, surpassant la SWA pure et égalant l'attention complète en perplexité tout la surpassant dans les tâches en aval.
   • Inférence sans Entraînement : Appliquée aux modèles pré-entraînés Qwen3-8B et Qwen3-30B-A3B, la SA surpasse constamment la SWA et égale ou dépasse le Mélange d'Attention par Blocs (MoBA) à des budgets de calcul comparables. Elle récupère la qualité de l'attention complète significativement plus vite à mesure que la taille de la fenêtre augmente.

Résultats

• Pré-entraînement : Le modèle SA + SWA a atteint une précision moyenne en zéro-shot de 35,9 (contre 35,1 pour la SWA et 34,9 pour l'Attention Complète). Il a notablement amélioré les performances sur LAMBADA (nécessitant un contexte de discours large) tout en maintenant une faible perplexité sur WikiText. La SA pure a souffert d'une perplexité plus élevée, confirmant la nécessité d'une cohérence locale pour la modélisation du langage.
• Inférence sur Qwen3 :
  • À une taille de fenêtre effective de 128, la SA sur Qwen3-8B a atteint 70,9 % de précision moyenne (à moins d'un point de la référence à attention complète), tandis que la SWA traînait à 62,2 %.
  • Sur Qwen3-30B-A3B, la SA a atteint 73,2 % de précision à une taille de fenêtre de 64, surpassant significativement la SWA (47,0 %) et le MoBA (66,3 %).
  • La SA a démontré une robustesse sur des tâches riches en connaissances (MMLU, BoolQ) même à des tailles de fenêtre très petites (par exemple, $w=32$), où les performances de la SWA s'effondraient.
• Efficacité : La SA maintient une mise à l'échelle linéaire $O(nw)$. Sur un GPU A100, la SA a montré des accélérations de 1,5× à une longueur de séquence de 2 048, augmentant jusqu'à 28× à 32 768 par rapport à l'attention complète.

Signification et Revendications

L'article revendique que le routage stochastique inspiré du connectome est un primitif pratique pour améliorer l'expressivité de l'attention efficace. Les auteurs soutiennent que le flux d'information global ne nécessite pas de connectivité dense de tous à tous mais peut émerger de l'interplay entre un calcul local structuré et des raccourcis à longue portée stochastiques et épars accumulés sur la profondeur.

La signification de ce travail réside dans :

1. Simplicité : La SA ne nécessite aucun paramètre apprenable dans le mécanisme d'attention lui-même et seulement une surcharge d'indices de $O(n)$, ce qui en fait une mise à niveau « intégrable » pour les noyaux SWA existants (par exemple, dans Mistral, Gemma 2, gpt-oss).
2. Garantie Théorique : Elle fournit une profondeur de couverture prouvable de $O(\log_w n)$, une amélioration significative par rapport à la couverture linéaire de la SWA.
3. Complémentarité : Elle est présentée comme complémentaire aux approches linéaires et éparse existantes, offrant un nouveau principe de conception qui comble le fossé entre l'efficacité locale et l'expressivité globale sans la complexité des motifs épars conçus à la main ou des jetons globaux.