Parallax: Parameterized Local Linear Attention for Language Modeling

Yifei Zuo; Dhruv Pai; Zhichen Zeng; Alec Dewulf; Shuming Hu; Zhaoran Wang

Source-linked AI summary

Parallax: Parameterized Local Linear Attention for Language Modeling

Yifei Zuo, Dhruv Pai, Zhichen Zeng, Alec Dewulf, Shuming Hu, Zhaoran Wang

arXiv:2605.29157v1May 27, 2026cs.LGcs.AIcs.CL

TL;DR

Efficient attention variants can underperform Softmax Attention on in-context retrieval, while Local Linear Attention has not been demonstrated at LLM pretraining scale. Parallax parameterizes Local Linear Attention for scalable training and inference, consistently improving perplexity and downstream accuracy while matching or outperforming FlashAttention 2/3 in decode kernels.

Problem
Local Linear Attention has theoretical advantages but lacks evidence of effectiveness in large-scale LLM pretraining because its solver creates computational, I/O, and numerical challenges.
Method
Parallax replaces Local Linear Attention’s numerical solve with a learned query-like projector and a hardware-aware streaming algorithm for scalable attention.
Results
Parallax consistently improves perplexity and downstream accuracy over Softmax Attention at 0.6B and 1.7B scales, while its decode kernel matches or outperforms FlashAttention 2/3.
Takeaways & Limitations
The results support Parallax as a scalable attention mechanism and identify a strong optimizer–architecture interaction, with substantial advantages under Muon.
Takeaways & Limitations
Validation of Parallax at larger scales, longer contexts, and with components such as MoE remains future work.

Abstract

from arXiv · show

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.

1 Introduction

Parallax addresses the scalability gap between Local Linear Attention’s theoretical advantages and its unresolved large-scale pretraining challenges. It combines a parameterized architecture with hardware-aware computation and demonstrates efficiency and empirical gains in LLM training.

Motivation: Replacing Softmax Attention’s local constant estimator with a local linear estimator yields a strictly richer predictor through the bias-variance and associative-memory perspective.
Architecture: Parallax preserves LLA’s local linear principle while removing the per-token conjugate-gradient solve’s computation, I/O, and numerical-sensitivity barriers to scalable pretraining.
Efficiency: A hardware-aware streaming algorithm enables a custom decode kernel that matches or outperforms FlashAttention 2/3 across diverse batch sizes and context lengths.
Experiments: Parallax improves perplexity and downstream accuracy over Softmax Attention at 0.6B and 1.7B scales, with gains persisting under parameter-matched and compute-matched controls.
Experiments: Parallax’s advantage depends on the optimizer: it shows substantial gains under Muon, while remaining comparable to Softmax Attention under AdamW.

2 Preliminary

The preliminary framework interprets attention as test-time regression over key-value pairs, with attention variants distinguished by hypothesis space, objective, and optimization. Local Linear Attention improves on Softmax Attention’s constant estimator but introduces substantial scaling challenges, motivating attention–optimizer considerations such as Muon.

Test-Time Regression: Attention predicts the value associated with query q_i by treating preceding key vectors as training data and value vectors as labels.The framework formulates attention as a regression solver over the KV pairs D_i.
Attention Design Space: Attention mechanisms differ through their hypothesis spaces, objective functions, and optimization methods, including parametric Linear Attention, ridge-regression MesaNet, and one-step DeltaNet.Linear Attention uses F = {W x + b} with context-independent weighting; MesaNet uses Ω(f) = λ∥W∥2_F, while DeltaNet uses one-step stochastic gradient descent without regularization.
Local Linear Attention: Local Linear Attention replaces Softmax Attention’s local constant function with a query-centered linear estimator, providing a second-order geometric correction and strictly smaller integrated MSE under the stated theory.LLA fits f ∈ F(q_i) = {b + W(x − q_i)} with kernel weighting and can degenerate to related mechanisms by tuning λ and h.
Challenges for LLM Training with LLA: Exact LLA requires solving a linear system for every query with parallel conjugate gradient, causing intensive I/O, a regularization–expressiveness tradeoff, and low-precision incompatibility.CG requires T L d memory access versus 2 L d for Softmax Attention; large λ reduces expressiveness, while small λ risks ill-conditioning and instability.
Muon: Muon updates matrix parameters using momentum and an approximate polar factor, with Newton–Schulz iterations making the optimizer hardware-aligned and feasible at scale.Its polar-factor updates have condition number exactly one, and prior work links this update conditioning to better-conditioned weight matrices.

3 Parallax Mechanism

Parallax replaces LLA’s per-query solve with a learned probe and removes boundary amplification for numerical stability. Its affine formulation connects related attention mechanisms, while a shared-stream covariance branch roughly doubles arithmetic intensity for hardware-efficient decoding.

Parallax formulation: Parallax replaces the per-query solve for ρ⋆_i with a learned projection ρ_i = W_Rx_i from the layer input.The learned projection matrix satisfies W_R ∈ R^dqk×d.
Parallax formulation: Setting η_i = 0 removes boundary amplification because the parameterized probe is not the exact LLA solution and can make the scaling factor diverge or change sign.The instability arises as t̄_i → 1 or when t̄_i > 1.
Attention family: In the wide-bandwidth limit, Parallax, Affine Linear Attention, and Affine MesaNet share an affine regression template and differ in whether the probe is zero, learned, or solved.The template is empirical OLS regression of v on k with intercept v̄_i, evaluated at the query.
Attention family: A poorly aligned or norm-suppressed learned probe makes the covariance correction inert, causing Parallax to collapse toward its Softmax Attention baseline.Probe alignment and norm depend heavily on optimizer choice.
Hardware-aware implementation: Parallax roughly doubles arithmetic intensity by adding covariance computation while reusing the same KV stream, shifting decoding toward a compute-bound regime.Both branches share the online maximum, rescaling factor, and KV tiles, so each iteration requires no extra I/O.

4 Experiment

Experiments show that Parallax improves recall-oriented synthetic tasks and language-modeling performance, with gains persisting across model scales and matched-parameter or matched-compute controls. The results also reveal a strong Muon–Parallax interaction, reflected in learned covariance corrections, projection ranks, and attention behavior.

MAD-Benchmark: Parallax consistently improves recall-oriented MAD tasks, remains competitive on compression and memorization, and achieves the highest overall accuracy.All models use two-layer sequence-mixer/MLP architectures and Muon optimization.
MAD-Benchmark: Under harder MAD conditions with vocabulary size up to 512 and context length up to 2048, Parallax retains accuracy while other baselines degrade dramatically, especially on selective copying.The challenge tasks scale KV pairs and sequence length for ICR, NCR, and SC.
LLM pretraining: At 0.6B and 1.7B scales, Parallax with Muon achieves the best perplexity on both evaluation tasks and the highest average downstream accuracy.Models are pretrained on Ultra-FineWeb with context length 4096; RoPE on ρ remains beneficial at both scales under Muon.
LLM pretraining: Matched-parameter and matched-attention-compute controls show that Parallax’s advantage is not explained by extra parameters or additional attention compute.The matched-parameter Transformer closes only a small fraction of the gap, while matched-compute Parallax significantly outperforms both Transformer variants.
Optimizer–architecture interaction: Muon produces stronger and deeper covariance corrections than AdamW, with COR exceeding 8 in deepest layers under Muon versus below 4 under AdamW.Muon also yields higher probe alignment and richer KV associations, whereas AdamW suppresses or eliminates Parallax’s advantage.
Mechanistic analysis: Under Muon, Parallax develops higher stable ranks in projection circuits, score ranges near ±40 in deepest layers, reduced attention sinks, and higher base-softmax entropy.Under AdamW, the gate suppresses the covariance correction and reaches performance comparable to Transformer.

5 Limitations and Future Directions · Appendix

The paper identifies future work spanning larger-scale validation, efficiency optimization, post-training adaptation, theoretical analysis, and extensions to other attention mechanisms. These directions include testing Parallax under broader model and hardware settings, characterizing optimizer dependence, and developing affine or nonparametric counterparts for related mechanisms.

5 Limitations and Future Directions: The section frames these unresolved questions as directions opened by the work, rather than presenting additional completed results.This framing introduces the subsequent directions on scaling, efficiency, adaptation, theory, and related mechanisms.
5 Limitations and Future Directions: Larger-scale studies should validate Parallax’s perplexity gains and optimizer–architecture interaction across scale, context length, MoE, and other architectural modifications.The doubled arithmetic intensity also motivates empirical tuning of head dimension, head count, and the attention-to-FFN ratio for specific hardware targets.
5 Limitations and Future Directions: Future kernel work should evaluate Parallax with contextual sparsity patterns, including sliding-window, dilated, and block-sparse attention.Because Parallax inherits Softmax Attention’s streaming structure, these patterns extend directly; compatibility with MLA also remains to be optimized and evaluated.
5 Limitations and Future Directions: Initializing W_R = 0 makes Parallax identical to Softmax Attention at training start, enabling pretrained Transformer checkpoints to be converted by adding W_R and fine-tuning.The passage contrasts this with Linear Attention, which lacks a parameter setting that exactly recovers Softmax Attention and typically requires retraining.
5 Limitations and Future Directions: The precise cause of Parallax’s optimizer dependence remains unresolved despite empirical spectral analysis and diagnosis of the performance gap between Muon and AdamW.It is also unknown whether the phenomenon occurs in other affine mechanisms.
5 Limitations and Future Directions: Reintroducing the intercept into Linear Attention, DeltaNet, and MesaNet would produce affine variants whose performance relative to intercept-free originals remains to be tested.The recurrence of the observed optimizer–architecture interaction in these variants is another open question.
5 Limitations and Future Directions: DeltaNet should be positioned between Linear Attention and MesaNet in the attention-mechanism family, and its nonparametric counterpart remains a natural extension.The current family does not yet include DeltaNet.

A Theorem · B Additional Derivation of Parallax · B.1 Reformulation of LLA

The appendix states regularity assumptions for the theorem, explains the bias distinction between Nadaraya–Watson and local linear estimators, and derives Parallax’s LLA reformulation from the exact forward computation.

A Theorem: The theorem assumes a domain with C2 boundary and uniformly bounded principal curvatures.This is the stated domain-regularity condition.
A Theorem: The assumptions require a positive C1 density, C2 regression functions, a radial bounded compactly supported kernel, and bandwidth conditions h →0 and nh^d →∞.The supplied passages also specify H = h^2B and a condition-number requirement, but its value is truncated.
A Theorem: The boundary-gradient condition requires an inward normal derivative bounded below by m and tangential gradient bounded by M on a positive-measure boundary subset.The constants m and M must satisfy a compatibility condition referenced in the cited definition and lemma.
A Theorem: The global linear estimator has an Ω(1) lower bound when f is not affine, while local linear estimation achieves O(∥H∥) uniform bias and avoids NW’s boundary-induced O(∥H∥^1/2) bias.The NW rates follow from pointwise bias–variance analysis with integrated boundary effects; the LL rates use uniform bias control through the boundary.
B.1 Reformulation of LLA: The derivation obtains equation (7) directly from the exact LLA forward expression in equation (3).It begins by introducing the quantities used in the exact forward computation.
B.1 Reformulation of LLA: The reformulation defines softmax weights as p_ij = w_ij/ω_i and uses μ_i together with z̄_i = k̄_i − q_i = E_p_i[z_ij].These identities connect the weighted quantities in the exact forward expression to the reformulated form.
B.1 Reformulation of LLA: Dividing by 1 − t̄_i and substituting 1/(1 − t̄_i) = 1 + η_i recovers the stated reformulation.This is the final algebraic step reported in the derivation.

B.2 Proof of Proposition 3.1 · C Parallax Decode Kernel · C.1 Kernel Optimization Details

The proof establishes the required positive-definite structure and interprets the resulting expression as a weighted Mahalanobis distance. The Parallax decode kernel reduces overhead through shared WGMMA computation, persistent KV-loop splitting, and in-kernel partial reduction.

B.2 Proof of Proposition 3.1: λ > 0 ensures Ai is positive definite, enabling the Sherman–Morrison formula in the proof.The passage explicitly links positive definiteness to λ > 0.
B.2 Proof of Proposition 3.1: The quadratic form ui is nonnegative and equals zero exactly when ¯zi = 0.This establishes the equality condition used in the proof.
B.2 Proof of Proposition 3.1: The ratio ui/(1 + ui) lies in [0, 1), providing the bound used by the proof.The displayed expression gives ui/(1 + ui) ∈ [0, 1).
B.2 Proof of Proposition 3.1: The final expression is ωi times the squared Mahalanobis distance from qi to the conditional key mean ¯ki under A−1_i.This supplies the geometric interpretation stated in the main text.
C Parallax Decode Kernel: WGMMA sharing jointly computes S1, S2, O1, and O2, adding one register-accumulator row without additional HBM traffic.Qr and Rr share one memory tile, while the covariance branch is fused into the existing WGMMA sequence.
C.1 Kernel Optimization Details: Persistent KV-loop splitting launches (B, H, S) CTAs, with S chosen to fit one device wave and rounded to a power of two for vectorized reduction.The split distributes the tile loop across CTAs sharing each (B, H) partition.
C.1 Kernel Optimization Details: In-kernel reduction stores fp32 partials, elects the final incrementing CTA as merger, and performs rescaling and output computation within the same launch.When S = 1, a compile-time branch skips the workspace round trip and writes directly from registers.

C.2 Additional Profiling Results · D Parallax Backward

Additional profiling distinguishes kernel latency from end-to-end latency and finds more consistent speedups from CUDA-graph measurements. Parallax backward derives closed-form gradients and implements them with FA-like tiled streaming in two passes.

C.2 Additional Profiling Results: CUDA-graph measurements isolate kernel latency, whereas Triton do_bench includes kernel-launch overhead, which affects FA3 more.The Figure 2b heatmaps use CUDA-graph measurements rather than end-to-end timings.
C.2 Additional Profiling Results: CUDA-graph measurements show a more consistent speedup pattern across different shapes.The raw H200 profiling results are provided in Figure 6.
D Parallax Backward: Closed-form gradients dQ, dR, dK, and dV are derived by differentiating the Parallax forward and applying the standard softmax derivative.Row scalars compress dependence on the output and value mean, while per-token terms resolve individual contributions.
D Parallax Backward: Parallax forward admits a reweighted softmax form in which the query shapes softmax weights and the probe modulates per-token coefficients.This form exposes the two channels through which qi and ρi enter the output.
D Parallax Backward: The backward kernel uses two passes because its reduction direction differs from the forward structure, while retaining row-tile and column-tile streaming.The forward cache adds only d + 1 values over the FA cache.
D Parallax Backward: The row-tile pass accumulates dQr and dRr over column blocks, while the reverse column-tile pass accumulates dKc and dVc over row blocks.This layout loads column tiles once during the column pass and processes row blocks in reverse order.

E Synthetic Experiment Setup

The synthetic experiments follow Poli et al. (2024) without data modification, changing only the sequence mixer block. Models use two mixer–SwiGLU MLP blocks and train for 60 Muon epochs while sweeping peak learning rates.

Experimental setup: The experiments follow Poli et al. (2024) without data modification, swapping only the sequence mixer block.
Model configuration: Each model stacks two mixer and SwiGLU MLP blocks with hidden size d = 128 in bf16 precision.
Optimization: Training runs for 60 epochs with Muon and the WSD schedule, using 0% warmup and linear decay over the last 20%.
Optimization: The peak learning-rate sweep uses {5 × 10−3, 1 × 10−3, 5 × 10−4}, with the best checkpoint reported per task.

F Pretraining Experiment Setup … G Additional Experiment Results

The pretraining runs use shared Qwen-3 decoder backbones with carefully matched Parallax and Transformer variants, standardized optimization, and H100-based distributed training. Experiments run for 20,000 steps, with scale-specific model, batch, and precision configurations.

F Pretraining Experiment Setup: Each training run uses one node with 8×H100 GPUs.
F.1 Backbone Architecture: All language-modeling runs share the Qwen-3 decoder backbone, tied embeddings, RMSNorm on q and k, and RoPE with base θ = 106.
F.1 Backbone Architecture: Parallax adds WR, sharing WQ’s head dimension and head grouping, and applies RMSNorm to ρ.
F.1 Backbone Architecture: Transformer† matches Parallax’s parameter count under GQA by increasing query heads while keeping KV heads fixed; a 3712-wide FFN alternative performs similarly.
F.2 Optimizer and Scheduler: Both optimizers clip gradient norms at 1.0, while Muon uses five Newton–Schulz iterations with the standard quintic coefficient and spectral scaling.
F.2 Optimizer and Scheduler: Training lasts 20,000 optimizer steps, and 1.7B runs double the global batch size to reach approximately 157.2 B total tokens.
F.3 Precision and Parallelism: All runs use fully sharded data parallelism without tensor, context, or pipeline parallelism, with torchao dynamic fp8 applied to linear layers except the bf16 LM head.

G.1 Training Dynamics · G.2 Advantage Shrinkage During Decay

Muon increasingly activates Parallax’s correction branch during training, unlike AdamW, while Parallax’s advantage over Transformer shrinks during WSD decay as weight norms decline. Weight decay annealing mitigates norm shrinkage and improves final loss, but late-training gains converge and only partially resolve the issue.

G.1 Training Dynamics: Muon and AdamW follow similar early trajectories, but Muon’s activation and projection norms continue growing while AdamW saturates.The largest optimizer separation occurs in ∥v∥ and ∥ρ∥.
G.1 Training Dynamics: Under Muon, the correction branch opens progressively and reaches its highest COR values in deepest layers, whereas AdamW largely suppresses it.COR is averaged over the sequence dimension.
G.2 Advantage Shrinkage During Decay: During WSD’s final linear decay, Parallax’s advantage over Transformer shrinks as weight norms decline throughout the decay phase.This norm decline may partially explain the shrinking advantage.
G.2 Advantage Shrinkage During Decay: Weight decay annealing replaces constant λ during decay; γ = 0 recovers WSD, while γ = 1 and γ = 2 provide linear and quadratic annealing.Larger γ suppresses weight decay more aggressively toward training’s end.
G.2 Advantage Shrinkage During Decay: At 0.6B scale, WDA with γ ∈{0.5, 1, 2} mitigates norm shrinkage, raises final WR norms with larger γ, and yields monotonic final-loss improvements.The comparison uses otherwise identical hyperparameters to standard Muon with WSD.
G.2 Advantage Shrinkage During Decay: WDA variants widen their loss gap during the first half of decay, then lose advantage in the second half as final losses converge to a narrow range.The turnaround occurs at similar step counts across all three variants.
G.2 Advantage Shrinkage During Decay: Because WDA holds back denominator shrinkage while η_t approaches zero, relative step size collapses faster than under WSD, reducing late-training progress.This provides a proposed mechanism for the observed late convergence.
G.2 Advantage Shrinkage During Decay: WDA’s training gain confirms weight norm shrinkage contributes mechanistically to advantage erosion, but the method only partially mitigates the issue.The result suggests standard Muon with WSD is not optimal for Parallax in its current form.

H Parallax Score Visualizations

This section visualizes Parallax and Transformer attention score maps to complement aggregate score statistics. The maps show both the top-left and bottom-right corners of 1024-token pretraining sequences.

Visualization setup: Each visualization block contains 64 × 64 tokens from a pretraining-data sequence of length 1024.The visualization uses the same block size for the displayed map regions.
Visualization setup: Figure 9a shows the top-left corner, while Figure 9b shows the bottom-right corner of the attention map.The Parallax AdamW visualization uses the WSD scheduler.
Score-map comparison: Figure 9 compares the attention score maps of the Transformer baseline and Parallax.The figure presents both the top-left and bottom-right corners.