Parallaxe : Attention locale linéaire paramétrée pour la modélisation du langage

Résumé

Les grands modèles de langage (LLM) sont devenus le paradigme central de l'intelligence artificielle, mais la primitive computationnelle fondamentale qu'est l'attention est restée structurellement inchangée. L'Attention Linéaire Locale (Local Linear Attention, LLA) est un mécanisme d'attention issu des statistiques non paramétriques dans le cadre de la régression au moment du test. Contrairement aux recherches antérieures sur les variantes d'attention efficaces, la LLA améliore l'estimation constante locale de l'attention softmax en une estimation linéaire locale, offrant des compromis biais-variance prouvés supérieurs pour la mémoire associative. Cependant, la LLA n'a pas été mise à l'échelle dans le pré-entraînement des LLM en raison de problèmes de stabilité numérique et computationnelle. Nous présentons Parallax, une Attention Linéaire Locale paramétrée, évolutive pour les LLM. Parallax élimine le solveur numérique présent dans la LLA et apprend un projecteur supplémentaire de type requête qui sonde la covariance des clés-valeurs (KV). Nous positionnons Parallax au sein d'une famille de mécanismes d'attention reliés par la largeur de bande, la construction de la sonde et la structure affine. Nous proposons un algorithme tenant compte du matériel qui augmente l'intensité arithmétique par rapport à FlashAttention, faisant basculer l'attention vers un régime plus limité par le calcul. Notre noyau de décodage prototype égale ou surpasse FlashAttention 2/3 pour diverses tailles de lot et longueurs de contexte. Nous pré-entraînons Parallax aux échelles de 0,6B et 1,7B paramètres et observons des améliorations constantes de la perplexité tout au long du pré-entraînement, avec des gains qui se transfèrent aux benchmarks aval. L'avantage persiste sous des contrôles appariés à la fois en nombre de paramètres et en charge de calcul, démontrant une amélioration de Pareto. Nous effectuons des ablations minutieuses du pré-entraînement et identifions un phénomène inédit par lequel Muon déverrouille la capacité de Parallax. À notre connaissance, il s'agit de la première démonstration empirique d'une co-conception forte entre architecture et optimiseur pour les mécanismes d'attention dans la littérature de recherche en architecture.

English

Large Language Models (LLMs) have become the central paradigm in artificial intelligence, yet the core computational primitive of attention has remained structurally unchanged. Local Linear Attention (LLA) is an attention mechanism derived from nonparametric statistics in the test-time regression framework. In contrast to prior research on efficient attention variants, LLA upgrades the local constant estimate in softmax attention to a local linear estimate, yielding provably superior bias-variance tradeoffs for associative memory. However, LLA has not been scaled in LLM pretraining due to computational and numerical stability concerns. We introduce Parallax, a parameterized Local Linear Attention that is scalable for LLMs. Parallax eliminates the numerical solver in LLA and learns an extra query-like projector that probes the KV covariance. We place Parallax within a family of attention mechanisms connected by the bandwidth, the probe construction and the affine structure. We propose a hardware-aware algorithm that increases the arithmetic intensity over FlashAttention, shifting attention into a more compute bound regime. Our prototype decode kernel matches or outperforms FlashAttention 2/3 across diverse batch sizes and context lengths. We pretrain Parallax at 0.6B and 1.7B scales and find consistent perplexity improvements throughout pretraining with gains that transfer to downstream benchmarks. The advantage persists under both parameter-matched and compute-matched controls, demonstrating a Pareto improvement. We perform careful pretraining ablations and identify a novel phenomenon whereby Muon unlocks the capacity of Parallax. To our knowledge, this is the first empirical demonstration of strong architecture-optimizer codesign for attention mechanisms in the architecture research literature.