Parallax: Parameterized Local Linear Attention for Language Modeling

Softmax attention is a local constant estimate over the KV cache; LLA upgrades it to a local linear estimate. Parallax makes that linear estimate scalable, beats FlashAttention 2/3 on decode, and reveals an architecture-optimizer codesign with Muon that gives Pareto perplexity gains at 0.6B and 1.7B pretraining.

Softmax attention            ►  local CONSTANT estimate of Y from (K,V)
                                 (kernel-smoothed value lookup, bias-heavy at sharp signals)

Local Linear Attention (LLA) ►  local LINEAR estimate
                                 (slope + intercept ► provably better bias-variance trade)
                                 (needs numerical solver ► fragile, doesn't scale)

Parallax (this paper)        ►  parameterized LLA
                                 ┌─────────────────────────┐
                                 │ extra query-like projector│ probes KV covariance
                                 │ NO numerical solver        │ stable, scalable
                                 │ FlashAttention-style kernel│ higher arithmetic intensity
                                 └─────────────────────────┘
                                  ▲                         ▲
                                  Architecture          Optimizer (Muon)
                                  ─────────────  CODESIGN  ─────────────
                                  Adam-W: no gain        Muon: unlocks Parallax capacity

TL;DR

Softmax attention has remained structurally unchanged since the original transformer. Parallax shows that the attention readout is mathematically a local constant estimate (Nadaraya-Watson, kernel-smoothed value lookup). Replacing it with a local linear estimate gives provably superior bias-variance trade-offs for associative memory; the textbook version (LLA) was known but had a numerical-stability barrier blocking pretraining-scale use. Parallax removes the solver entirely, adds a learned query-like projector that probes the KV covariance, and ships a hardware-aware kernel that pushes arithmetic intensity over FlashAttention so the kernel is more compute-bound, not less (the opposite of most efficient-attention proposals). On a 0.6B-parameter pretrain and on 1.7B, perplexity improves throughout training; gains transfer to downstream benchmarks under both parameter-matched and compute-matched controls. The paper also names a surprising phenomenon: Muon unlocks Parallax, while Adam-W does not. This is the first concrete demonstration of architecture-optimizer codesign in the attention literature.

What changes vs prior work

Most efficient-attention variants (linear attention, sliding window, sparse) move toward memory-bound to cut compute. Parallax goes the other way: by adding the parameterized linear estimator, the kernel does more work per byte loaded, sitting in a more favorable regime on Blackwell/H200 tensor cores. This is the same direction as kernel-fusion work on FlashAttention-3 and the SemiAnalysis "miscompile" thread of bug-hunting kernels at speed of light: the binding constraint on next-gen attention is no longer memory bandwidth, it is arithmetic intensity.

The local-constant-vs-linear framing also gives a clean theoretical story for why softmax attention has long-tail failures around sharp signal transitions (the constant estimate over-smooths). Local linear estimators have textbook bias-variance superiority for exactly the regimes attention struggles with.

Connections to prior wiki

Massive activations / ME-layer (2026-05-13): identified specific "memory" layers that store key associations. Parallax operates on the same readout primitive at every layer. Whether the ME-layer phenomenon collapses or sharpens under local-linear readout is a clean open question.
Hope nested-learning architecture (2026-04-28): also reframed attention as a learned estimator. Parallax is more conservative (single-pass linear estimate, full backprop compatible) and ships measured pretraining wins.
AccelOpt GPU kernel optimization (2026-04-20) + doubleAI Blackwell SOL kernels (2026-05-28): both showed AI-written kernels reaching speed-of-light. Parallax requires hand-written hardware-aware kernels to realize its arithmetic-intensity gain; whether AI-kernel systems can author a Parallax kernel at SoL is the natural follow-up.

Research angle

The Muon-unlock observation deserves replication. If architecture-optimizer codesign is a real lever, it implies a non-trivial slice of efficient-attention research has been bottlenecked by an arbitrary optimizer choice rather than by the architecture itself. Worth checking at 7B+ scale whether the Pareto gain persists or saturates.

The bandwidth-probe-affine taxonomy the authors propose (a family of attention mechanisms connected by the bandwidth parameter, the probe construction, and the affine structure) provides a clean lattice to traverse. Expect a year of papers walking through this lattice.