
Parallax : attention locale linéaire paramétrée avec softmax et correction de covariance apprise
Une équipe de chercheurs de Northwestern University, Tilde Research et l'Université de Washington a présenté Parallax, une nouvelle architecture d'attention pour les grands modèles de langage qui cherche à dépasser les limites du mécanisme softmax utilisé sans changement majeur depuis 2017. Contrairement à la plupart des travaux récents qui tentent de remplacer entièrement l'attention softmax, Parallax adopte une stratégie différente : conserver softmax et lui adjoindre une branche de correction apprise. Concrètement, le mécanisme ajoute un terme correctif basé sur la covariance clé-valeur, calculé via une matrice de projection apprise notée WR. Lorsque cette matrice est nulle, Parallax se réduit exactement à de l'attention softmax classique, ce qui signifie qu'un checkpoint existant peut être converti par simple ajout de WR et fine-tuning.
L'enjeu principal est d'ordre théorique et pratique à la fois. Théoriquement, Parallax s'appuie sur le cadre de régression locale linéaire (LLA), qui interprète l'attention comme un estimateur statistique : les clés sont des points d'entraînement, les valeurs sont des labels, et la requête est le point de test. L'attention softmax correspond à un estimateur local constant (Nadaraya-Watson) ; LLA l'étend à une estimation linéaire locale, prouvée plus précise en termes d'erreur quadratique moyenne intégrée. Sur le plan matériel, Parallax exploite la structure de FlashAttention en ajoutant une seconde branche de scoring qui réutilise exactement le même flux clé-valeur, sans I/O supplémentaire. Le résultat est un doublement approximatif de l'intensité arithmétique, c'est-à-dire du ratio opérations flottantes sur trafic mémoire, poussant le calcul vers un régime davantage limité par le compute que par la mémoire. C'est précisément dans ce régime que les optimisations de kernels GPU sont les plus efficaces.
La présentation inclut un prototype de kernel de décodage en CuTeDSL sur les GPU NVIDIA Hopper, où les instructions matmul opèrent sur des tuiles d'au minimum 64 lignes alors qu'un pas de décodage n'en fournit qu'une seule : Parallax contourne cette contrainte en fusionnant les produits QK et RK dans les mêmes instructions que l'attention standard. Ce travail s'inscrit dans un contexte de recherche très actif autour de l'efficacité des Transformers, où la plupart des alternatives comme Mamba ou Linear Attention sacrifient la qualité pour gagner en vitesse. Parallax parie sur une voie complémentaire : ajouter délibérément du calcul, mais rendre ce calcul moins coûteux à exécuter sur le matériel moderne. Les chercheurs indiquent également une co-conception avec l'optimiseur Muon, ce qui suggère une intégration pensée pour le pré-entraînement à grande échelle plutôt que pour des ajustements marginaux en inférence.
Dans nos dossiers
Vu une erreur factuelle dans cet article ? Signalez-la. Toutes les corrections valides sont publiées sur /corrections.




