pith. sign in

arxiv: 2605.29157 · v1 · pith:RKOPRCTFnew · submitted 2026-05-27 · 💻 cs.LG · cs.AI· cs.CL

Parallax: Parameterized Local Linear Attention for Language Modeling

Yifei Zuo , Dhruv Pai , Zhichen Zeng , Alec Dewulf , Shuming Hu , Zhaoran Wang This is my paper

Pith reviewed 2026-06-29 13:18 UTC · model grok-4.3

classification 💻 cs.LG cs.AIcs.CL
keywords local linear attentionparameterized attentionlanguage modelingLLM pretrainingefficient attentionbias-variance tradeoffMuon optimizerPareto improvement
0
0 comments X

The pith

Parallax replaces softmax attention with a learned local linear estimator that lowers perplexity at 0.6B and 1.7B scales under matched controls.

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The paper establishes that standard softmax attention rests on a local constant estimator whose bias-variance properties can be improved by moving to a local linear estimator derived from nonparametric statistics. Parallax renders this upgrade trainable at LLM scale by introducing a learned query-like projector that probes KV covariance, removing the need for a numerical solver, and embedding the mechanism in a family parameterized by bandwidth and affine structure. Pretraining runs at 0.6B and 1.7B parameters produce lower perplexity throughout training than baselines, with the gains carrying to downstream benchmarks. The advantage survives both parameter-matched and compute-matched comparisons, indicating a Pareto improvement. The work additionally reports that the Muon optimizer is required to unlock the full capacity of the new attention form.

Core claim

Local Linear Attention upgrades the local constant estimate inside softmax attention to a local linear estimate, which carries provably superior bias-variance tradeoffs for associative memory. Parallax makes this form scalable by learning an extra projector that probes the KV covariance, eliminating the numerical solver, and supplying a hardware-aware algorithm that raises arithmetic intensity. When models using Parallax are pretrained at 0.6B and 1.7B scales, they exhibit consistent perplexity reductions relative to standard attention; the reductions transfer to downstream tasks and remain visible under both parameter-matched and compute-matched controls. The same experiments identify that

What carries the argument

The learned query-like projector that probes KV covariance to realize local linear estimation without an explicit numerical solver.

If this is right

  • Perplexity improves consistently throughout pretraining at both 0.6B and 1.7B scales.
  • Gains on downstream benchmarks follow from the perplexity reductions.
  • The improvement constitutes a Pareto advance under both parameter-matched and compute-matched controls.
  • The Muon optimizer is necessary to realize the full capacity of Parallax.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

  • The projector-based parameterization may be portable to other nonparametric estimators inside attention variants.
  • The reported architecture-optimizer interaction suggests that systematic codesign experiments could be run on additional attention mechanisms.
  • Hardware-aware kernels that shift attention toward a compute-bound regime may be worth exploring for other attention replacements on modern accelerators.

Load-bearing premise

The learned projector and bandwidth parameterization preserve the theoretical bias-variance advantage of local linear estimation while remaining numerically stable and trainable at the reported scales without post-hoc data exclusions or hidden hyperparameter tuning that would invalidate the matched-control comparisons.

What would settle it

A 1.7B-scale pretraining run in which Parallax produces equal or higher perplexity curves than the softmax baseline under identical data, optimizer, and compute settings would falsify the claimed improvement.

read the original 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.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit. Tearing a paper down is the easy half of reading it; the pith above is the substance, this is the friction.

Referee Report

3 major / 2 minor

Summary. The manuscript introduces Parallax, a parameterized Local Linear Attention mechanism derived from nonparametric local linear estimation. It replaces LLA's numerical solver with a learned query-like projector probing KV covariance, places the method in a bandwidth/probe/affine family, and presents a hardware-aware kernel with higher arithmetic intensity than FlashAttention. Pretraining results at 0.6B and 1.7B scales report consistent perplexity gains that transfer to downstream tasks and persist under both parameter-matched and compute-matched controls, alongside a novel Muon optimizer interaction that unlocks Parallax capacity.

Significance. If the perplexity gains are shown to arise from the local-linear bias-variance improvement rather than added projector capacity or Muon-specific interactions, the work would provide the first scaled empirical validation of nonparametric local-linear attention in LLMs, together with a practical kernel that shifts attention toward compute-bound regimes.

major comments (3)
  1. [Experimental results at 0.6B/1.7B scales] Experimental results (0.6B and 1.7B scales): the parameter-matched control must explicitly report total parameter counts for Parallax versus baseline and confirm that the extra learned projector parameters are exactly offset; without this, the Pareto claim cannot be attributed to the local-linear structure.
  2. [Ablations on Muon interaction] Ablations section: the reported Muon-Parallax interaction requires an ablation that holds the optimizer fixed (e.g., AdamW for both Parallax and baseline) to isolate whether the perplexity gains are driven by the attention mechanism itself rather than the optimizer-architecture coupling.
  3. [Theoretical connection to LLA] Theoretical framing: the claim that Parallax inherits the provably superior bias-variance tradeoffs of LLA must be supported by explicit reduction equations showing the settings of bandwidth, probe, and affine parameters under which Parallax recovers standard LLA or softmax attention.
minor comments (2)
  1. [Kernel implementation] The hardware-aware kernel description would benefit from a short pseudocode listing or arithmetic-intensity calculation to substantiate the claim of outperforming FlashAttention 2/3.
  2. [Downstream evaluation] Downstream benchmark tables should report the exact number of evaluation runs and standard deviations to allow assessment of transfer reliability.

Simulated Author's Rebuttal

3 responses · 0 unresolved

We thank the referee for the constructive and detailed comments, which help clarify the presentation of our results and theoretical connections. We address each major comment below and will incorporate revisions to strengthen the manuscript accordingly.

read point-by-point responses

  1. Referee: [Experimental results at 0.6B/1.7B scales] Experimental results (0.6B and 1.7B scales): the parameter-matched control must explicitly report total parameter counts for Parallax versus baseline and confirm that the extra learned projector parameters are exactly offset; without this, the Pareto claim cannot be attributed to the local-linear structure.

    Authors: We agree that explicit reporting of total parameter counts is required to fully substantiate the parameter-matched controls. In the revised manuscript, we will add a dedicated table (or expanded section in the experimental results) listing the precise total parameter counts for Parallax and the baseline at both 0.6B and 1.7B scales. This will confirm that the additional parameters introduced by the learned query-like projector are exactly offset by corresponding reductions in other components, ensuring the comparison isolates the effect of the local-linear structure. revision: yes

  2. Referee: [Ablations on Muon interaction] Ablations section: the reported Muon-Parallax interaction requires an ablation that holds the optimizer fixed (e.g., AdamW for both Parallax and baseline) to isolate whether the perplexity gains are driven by the attention mechanism itself rather than the optimizer-architecture coupling.

    Authors: We concur that an optimizer-fixed ablation is necessary to isolate the contribution of the Parallax mechanism. Although our existing ablations highlight the Muon-Parallax interaction, the revised manuscript will include a new set of pretraining runs in which both Parallax and the baseline attention are trained exclusively with AdamW. This will allow direct comparison of perplexity gains attributable to the attention architecture independent of optimizer coupling, and we will report these results alongside the existing Muon findings. revision: yes

  3. Referee: [Theoretical connection to LLA] Theoretical framing: the claim that Parallax inherits the provably superior bias-variance tradeoffs of LLA must be supported by explicit reduction equations showing the settings of bandwidth, probe, and affine parameters under which Parallax recovers standard LLA or softmax attention.

    Authors: We appreciate this suggestion for making the theoretical inheritance explicit. The manuscript already situates Parallax in the bandwidth/probe/affine family, but the revised version will add a new subsection (or appendix) containing the explicit reduction equations. These will specify the exact settings (e.g., bandwidth = 1 with identity probe and constant affine for LLA recovery; further degenerations to recover softmax attention) under which Parallax reduces to standard LLA or softmax, thereby rigorously supporting the claimed bias-variance advantages. revision: yes

Circularity Check

0 steps flagged

No circularity in derivation or empirical claims

full rationale

The paper's central claims consist of empirical perplexity gains from pretraining Parallax at 0.6B/1.7B scales under parameter- and compute-matched controls, plus ablations identifying a Muon interaction. These outcomes are measured results from training runs rather than quantities derived by construction from fitted parameters or self-referential equations. The background derivation of LLA from nonparametric local linear estimation is presented as external motivation and does not reduce the reported gains to inputs inside the same chain. No self-citations, uniqueness theorems, or ansatzes are invoked in a load-bearing way that collapses the result to its own definitions.

Axiom & Free-Parameter Ledger

0 free parameters · 0 axioms · 0 invented entities

Abstract-only review yields no explicit free parameters, axioms, or invented entities; the central claim rests on the unstated assumption that the learned projector faithfully approximates the local-linear estimator.

pith-pipeline@v0.9.1-grok · 5835 in / 1139 out tokens · 37588 ms · 2026-06-29T13:18:50.690640+00:00 · methodology

discussion (0)

Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.

Reference graph

Works this paper leans on

6 extracted references · 3 canonical work pages · 1 internal anchor

  1. [1]

    Simran Arora, Sabri Eyuboglu, Aman Timalsina, Isys Johnson, Michael Poli, James Zou, Atri Rudra, and Christopher Ré

    URLhttps://openreview.net/forum?id=3ciBX6oWXP. Simran Arora, Sabri Eyuboglu, Aman Timalsina, Isys Johnson, Michael Poli, James Zou, Atri Rudra, and Christopher Ré. Zoology: Measuring and improving recall in efficient language models, 2023. URLhttps://arxiv.org/abs/2312.04927. Ali Behrouz, Peilin Zhong, and Vahab Mirrokni. Titans: Learning to memorize at t...

    work page arXiv 2023

  2. [2]

    Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge

    URLhttps://arxiv.org/abs/1803.05457. Damai Dai, Yutao Sun, Li Dong, Yaru Hao, Shuming Ma, Zhifang Sui, and Furu Wei. Why can gpt learn in-context? language models implicitly perform gradient descent as meta-optimizers, 2023. URLhttps://arxiv.org/abs/2212.10559. Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. In 1...

    work page internal anchor Pith review Pith/arXiv arXiv 2023

  3. [3]

    Szegedy, C., Vanhoucke, V ., Ioffe, S., Shlens, J., and Wojna, Z

    URLhttps://arxiv.org/abs/2410.06511. Jingyuan Liu, Jianlin Su, Xingcheng Yao, Zhejun Jiang, Guokun Lai, Yulun Du, Yidao Qin, Weixin Xu, Enzhe Lu, Junjie Yan, et al. Muon is scalable for llm training, 2025. URLhttps://arxiv. org/abs/2502.16982. Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In7th International Conference on Learni...

    work page doi:10.1016/j.neucom.2023.127063 2025

  4. [4]

    The WGMMA then emitsS1 and S2 in Algorithm 1 jointly in the same accumulator

    WGMMA sharing.Each compute thread array (CTA) loadsQr andR r into a single shared memory tile, withQr as the first row andRr as the second. The WGMMA then emitsS1 and S2 in Algorithm 1 jointly in the same accumulator. After producingP1, we buildP2 =P 1⊙S2 in registers and stack it withP1 in shared memory. The PV WGMMA then emitsO1 andO 2 jointly. The cova...

  5. [5]

    We launch a persistent grid of(B,H,S )CTAs, where theS number of CTAs share a(B,H )partition the⌈L/Bc⌉tile loop of Algorithm 1

    Persistent split over the KV loop.Decoding presents onlyBH query rows, often well below the 132 SMs of an H200 in practical configurations. We launch a persistent grid of(B,H,S )CTAs, where theS number of CTAs share a(B,H )partition the⌈L/Bc⌉tile loop of Algorithm 1. The split countS is set so that the launch fits one wave on the device and is rounded to ...

  6. [6]

    In-kernel reduction.Each CTA writes its unnormalized partials(m, d1, d2, O1, O2)to a small fp32 HBM workspace and atomically increments a per(B,H )counter. The CTA that observes the final increment is elected the merger: it reads theS partials, runs the log-sum-exp rescaling in fp32, evaluates(1 +d 2/d1)O1/d1−O2/d1, and writes the output row in the same k...

    2024