Parallax: Parameterized Local Linear Attention for Language Modeling

2026-05-27

AI Summary

This work addresses the efficiency and expressivity bottlenecks of conventional attention mechanisms, which hinder scalable deployment in large language models despite the theoretical promise of Local Linear Attention (LLA). We propose Parallax, a scalable and parameterized LLA variant that replaces numerical solvers with learnable query projections and integrates hardware-aware decoding kernels with high arithmetic intensity. Parallax enables, for the first time, effective pretraining of LLA-based large models and reveals a synergistic co-design effect with the Muon optimizer, achieving Pareto improvements. Evaluated on 0.6B and 1.7B parameter models, Parallax consistently reduces perplexity, enhances downstream performance, and matches or exceeds the decoding speed of FlashAttention-2/3, maintaining comprehensive advantages under matched parameter and compute budgets.

Abstract

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.

Problem

• Local Linear Attention
• Large Language Models
• Attention Mechanism
• Numerical Stability
• Scalability

Innovation

• Local Linear Attention
• Parameterized Attention
• Hardware-aware Algorithm
• Architecture-Optimizer Co-design
• KV Covariance Estimation