Parallax: Parametrisierte lokale lineare Aufmerksamkeit für Sprachmodellierung

Zusammenfassung

Große Sprachmodelle (Large Language Models, LLMs) haben sich zum zentralen Paradigma der künstlichen Intelligenz entwickelt, doch die grundlegende Recheneinheit der Aufmerksamkeit (Attention) ist strukturell unverändert geblieben. Lokale Lineare Aufmerksamkeit (Local Linear Attention, LLA) ist ein aus der nichtparametrischen Statistik im Testzeit-Regressionsrahmen abgeleiteter Aufmerksamkeitsmechanismus. Im Gegensatz zu früheren Arbeiten über effiziente Aufmerksamkeitsvarianten verbessert LLA die lokale konstante Schätzung der Softmax-Aufmerksamkeit zu einer lokalen linearen Schätzung, was nachweislich überlegene Bias-Varianz-Kompromisse für assoziatives Gedächtnis liefert. Allerdings wurde LLA aufgrund von Bedenken hinsichtlich Recheneffizienz und numerischer Stabilität nicht im Pretraining von LLMs skaliert. Wir führen Parallax ein, eine parametrisierte Lokale Lineare Aufmerksamkeit, die für LLMs skalierbar ist. Parallax eliminiert den numerischen Löser in LLA und lernt einen zusätzlichen query-ähnlichen Projektor, der die KV-Kovarianz abtastet. Wir ordnen Parallax in eine Familie von Aufmerksamkeitsmechanismen ein, die durch die Bandbreite, die Sondenkonstruktion und die affine Struktur verbunden sind. Wir schlagen einen hardwarebewussten Algorithmus vor, der die arithmetische Intensität gegenüber FlashAttention erhöht und die Aufmerksamkeit in einen rechengebundeneren Bereich verschiebt. Unser Prototyp-Dekodierungskernel erreicht oder übertrifft FlashAttention 2/3 über verschiedene Batch-Größen und Kontextlängen hinweg. Wir trainieren Parallax in den Größenordnungen 0,6B und 1,7B vor und stellen durchgängige Perplexitätsverbesserungen während des gesamten Pretrainings fest, wobei die Gewinne auf nachgelagerte Benchmarks übertragbar sind. Der Vorteil bleibt sowohl bei parameterangepassten als auch bei rechenangepassten Kontrollen bestehen, was eine Pareto-Verbesserung darstellt. Wir führen sorgfältige Pretraining-Ablationen durch und identifizieren ein neuartiges Phänomen, bei dem Muon die Kapazität von Parallax freisetzt. Unseres Wissens ist dies die erste empirische Demonstration eines starken Architektur-Optimierer-Co-Designs für Aufmerksamkeitsmechanismen in der Architekturforschungsliteratur.

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.