Parallax: Parameterized Local Linear Attention for Language Modeling
Yifei Zuo, Dhruv Pai, Zhichen Zeng, Alec Dewulf, Shuming Hu, Zhaoran Wang
cs.LG, cs.AI, cs.CL
2026-05-28
Parallax swaps LLA's per-query solver for a learned covariance branch, beating softmax attention on perplexity at 0.6B/1.7B with a decode kernel matching FlashAttention 2/3.
The attention core of the Transformer has not changed since 2017. The main alternatives, linear attention variants (DeltaNet and friends) and SSMs like Mamba, buy subquadratic complexity with a constant-size recurrent state, and pay for it with consistent losses on in-context retrieval. They got faster and lost the memory.
The test-time regression framework reads every attention mechanism as a regression problem solved over the context's KV pairs. Softmax attention corresponds to the Nadaraya-Watson local-constant estimator, which carries an irreducible bias when queries land near the boundary of the attention mass. Local Linear Attention (LLA) upgraded the estimator to local-linear and proved strictly smaller integrated MSE, but never survived contact with real pretraining: it solves a linear system per query with conjugate gradient, the forward pass touches an order of magnitude more memory, the regularization strength λ either degenerates the model or breaks conditioning, and CG does not tolerate BF16.
Parallax first rewrites LLA as softmax output minus a covariance correction: the correction is the softmax-weighted KV covariance ΣKV times a probe vector ρ. Where LLA solves for ρ per query, Parallax deletes the solver and learns a projection matrix WR, so ρ = WR·x is one matmul from the layer input; the model carries an extra R matrix alongside Q, K, V (in practice ρ gets its own RMSNorm, and RoPE on it helps further). It also drops the boundary amplification factor η, which loses its geometric guarantee once parameterized and can blow up training. The final layer is softmax attention minus ΣKV·ρ. With WR = 0 the layer is exactly softmax attention, which matters later for converting pretrained checkpoints.
The paper maps a whole family of mechanisms onto three axes: bandwidth, how the probe is obtained (zero, learned, solved), and whether the intercept is kept. As bandwidth goes to infinity Parallax degenerates to linear attention with an intercept; as regularization grows it collapses back to softmax. A side observation: what linear attention calls the query is, in this view, the probe.
On the hardware side the two branches share one KV stream, one running max, and one rescale factor, so streaming costs no extra HBM traffic and arithmetic intensity roughly doubles over FlashAttention. For decode there is a neat trick: Hopper's WGMMA instructions operate on at least 64 rows, but a decode step supplies one query row, leaving most of the accumulator idle in standard attention. Parallax packs the QK and RK products into the same tensor-core instruction, so the correction branch's FLOPs are nearly free.
On the six synthetic MAD tasks, all trained with Muon (ICR is in-context recall, SC is selective copying):
| Model | Avg | ICR | SC |
| Parallax | 0.716 | 0.951 | 0.988 |
| Softmax Attention | 0.672 | 0.803 | 0.950 |
| GDN | 0.665 | 0.920 | 0.939 |
| Mesa | 0.659 | 0.998 | 0.568 |
| Mamba | 0.631 | 0.756 | 0.950 |
With vocabulary pushed to 512 and context to 2048, baselines collapse while Parallax holds. The decode kernel, profiled on H200 in BF16 across batch 1–2048 and context 128–32768, matches or beats the better of FA2/FA3 in every configuration: up to 1.54× compute-matched, up to 1.14× I/O-matched.
Pretraining (Qwen-3 architecture, Ultra-FineWeb, 4096 context):
| Model (0.6B, 78.6B tokens) | LMB ppl ↓ | Wiki ppl ↓ | Avg downstream ↑ |
| Transformer | 22.15 | 23.43 | 54.54 |
| Parameter-matched Transformer† | 22.35 | 23.36 | 54.90 |
| Compute-matched Parallax† | 20.29 | 22.49 | 55.79 |
| Parallax | 18.56 | 22.25 | 55.99 |
At 1.7B (157.2B tokens): LAMBADA perplexity 13.07 → 10.80, downstream average 61.43 → 62.45. KDA (52.73) and GDN (53.67) both trail the plain Transformer at 0.6B, which makes Parallax the only listed challenger that beats it. Parameter-matching (extra query heads) recovers only a sliver of the gap; compute-matching (head dim halved) still wins clearly. The gain is the mechanism, not the parameters or FLOPs.
Under AdamW the story flips. At 0.6B, AdamW+WSD is a wash (52.61 vs 52.68); under AdamW+Cosine, Parallax lands at 51.45, below the Transformer's 52.34. The mechanism analysis explains why: the correction-to-output ratio exceeds 8 in deep layers under Muon but stays below 4 under AdamW. A learned sigmoid gate on the probe settles near 0.26 under AdamW, meaning the model actively shuts the correction branch off; under Muon the gate opens. Weight spectra locate the bottleneck in WR, whose stable rank collapses hardest under AdamW. Parallax's composite weights can also go negative and reach about ±40 in deep layers, attention sinks shrink, and the base softmax runs at higher entropy, with fine-grained token discrimination delegated to the correction branch.
This is the first time the LLA line of work has survived LLM pretraining, and it wins under both parameter- and compute-matched controls, which makes it a genuine Pareto improvement. The design is portable: keep the softmax machinery, add one branch. Because WR = 0 recovers softmax exactly, an existing pretrained checkpoint could in principle gain a WR and be fine-tuned into Parallax, far cheaper than the linear-attention family, where no parameter setting recovers softmax. The paper proposes this but does not test it.
The broader lesson is about optimizers. The same architecture shows zero gain or worse under AdamW and delivers only under Muon. Teams benchmarking architectures on an AdamW default may be filtering out mechanisms of this type. The authors claim this is the first empirical demonstration of strong architecture-optimizer codesign for attention; the evidence here supports the claim.
The magnitudes deserve honesty: downstream averages move by about one point. The large gains are in perplexity, LAMBADA especially, not in a downstream blowout.
Author-stated: validated only to 1.7B parameters, 157B tokens, and 4096 context; no MoE, no long-context runs, no sparse variants; the Muon dependence is characterized empirically with no theory; checkpoint conversion is untested.
A close read adds more. The selling point is retrieval and associative memory, but training and evaluation never leave 4096 context, while associative memory is mostly a long-context story. The WikiText perplexity gap (22.25 vs 23.43) is much smaller than the LAMBADA gap (18.56 vs 22.15), so the gains are unevenly distributed. The gating experiment shows AdamW actively suppresses the branch, which is both a finding and a fragility: any change to the optimizer recipe requires revalidation, and the authors concede Muon+WSD may not be optimal. The kernel comparison is a homegrown CuTeDSL prototype against FA2/FA3 on a single H200 in BF16; end-to-end training wall-clock is not reported.