Parallax: Parameterized Local Linear Attention for Language Modeling

Abstract: 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.

• The paper presents Parallax, a novel parameterized local linear attention that replaces softmax with a learned probe to achieve lower integrated MSE and improved recall.
• It leverages dual streaming and covariance correction branches to ensure hardware efficiency, stability in low-precision formats, and superior throughput in LLM pretraining.
• Empirical benchmarks reveal that Parallax consistently lowers perplexity and enhances associative memory compared to traditional Transformer models.

Parallax: Parameterized Local Linear Attention for Language Modeling

Motivation and Theoretical Foundations

The Transformer architecture dominates sequence modeling tasks, hinging on Softmax Attention as its primary computational mechanism. However, the attention mechanism's quadratic complexity and lack of structural innovation have prompted the search for efficient alternatives that maintain associative memory capacity. Linear Attention variants and State Space Models (SSMs) offer subquadratic complexity but demonstrably underperform in-context recall tasks due to their limited memory capacity. The test-time regression framework interprets attention mechanisms as regression solvers over $(k_j, v_j)$ pairs, with the prediction for each query $q_i$ formulated as a regularized regression. Local Linear Attention (LLA) emerges from this framework by replacing the local constant estimator (Softmax) with a local linear estimator, which is theoretically proven to yield strictly smaller integrated MSE than both global linear and local constant estimators.

LLA, despite its theoretical superiority, has not been scaled to LLM pretraining due to computational bottlenecks (per-token conjugate gradient solver) and numerical instability in low-precision formats.

Parallax Mechanism and Architectural Innovations

Parallax is introduced as a scalable, parameterized variant of LLA. Its key architectural contribution is eliminating the numerical solve by learning a probe vector via an additional query-like projection $\bm W_R x_i$, which interacts with the KV covariance to produce a correction to Softmax Attention output:

• The correction branch directly computes a covariance-probe product, avoiding unstable iterative solvers.
• Parallax maintains the affine structure of LLA but parameterizes the probe, enabling compatibility with reduced precision and hardware-optimized streaming.

This mechanism is positioned within a unified family of attention mechanisms parameterized by bandwidth, probe construction, and affine structure.

Figure 1: A taxonomy of attention mechanisms spanning softmax-weighted, uniform-weighted, and solution-based probes, contextualizing Parallax within both parametric and nonparametric designs.

Parallax is reformulated as an additive correction to Softmax Attention, grounded in projected KV covariance. The output takes the form:

$$o_i^{\textsf{PLX}} = o_i^{\textsf{SA}} - \Sigma_{KV}^{(i)}\,\bm W_R x_i$$

This formulation yields a robust bias-variance tradeoff, amplifies boundary corrections when the query lies near the KV distribution's edge, and provides more expressive associative memory.

Hardware-Aware Algorithm and Streaming Efficiency

To address efficiency, Parallax adopts a streaming algorithm similar to FlashAttention (FA) but increases arithmetic intensity through:

• Dual streaming branches (Softmax and covariance correction) sharing tiled KV streams.
• Hardware alignment with NVIDIA Hopper tensor core matmuls, exploiting the minimal row requirements by joint computation of QK and RK.

Profiling demonstrates that Parallax's custom decode kernel matches (and often outperforms) FlashAttention 2/3 across batch sizes and context lengths, enabled by its doubled arithmetic intensity and resource reuse.

Figure 2: Dependency graph and hardware mapping for Parallax, showing operator relationships and resource alignment for efficient kernel execution.

Empirical Validation: Synthetic and Language Modeling Benchmarks

Parallax is validated on:

• MAD synthetic benchmarks probing recall, compression, and memorization. Parallax achieves highest accuracy in recall-oriented tasks (ICR, FCR, NCR, SC) and retains performance as task difficulty increases.
• LLM pretraining at 0.6B and 1.7B parameter scales on Ultra-FineWeb with Qwen-3 backbone. Parallax consistently improves training and evaluation perplexity over the Transformer baseline. Gains persist under parameter-matched and compute-matched controls, demonstrating that its improvement is not due to parameter or compute scaling, but the mechanism itself.

Strong numerical results include:

• At 0.6B scale, Parallax achieves LAMBADA perplexity of $18.56$ vs. $22.15$ for Transformer, and WikiText perplexity of $22.25$ vs. $23.43$.
• At 1.7B scale, Parallax achieves LAMBADA perplexity of $10.80$ vs. $13.07$ and consistently higher average downstream accuracy.

Optimizer-architecture codesign is empirically demonstrated: Parallax's correction branch is strongly activated by Muon optimizer, while its benefit vanishes under AdamW. The probe norm and correction-to-output ratio (COR) substantially increase with Muon, confirming Parallax's reliance on the optimizer to activate its affine structure.

Figure 3: COR (correction-to-output ratio) quantifies Parallax's covariance correction branch activation layer-wise, showing optimizer-dependent amplification.

Figure 4: Temporal evolution of COR under different optimizers, revealing Muon's capacity to unlock correction strength in deep layers.

Mechanism Analysis and Score Distribution

Parallax produces distinct score distributions:

• The correction branch enables negative attention weights, allowing explicit subtraction of irrelevant token values and unbounded expressive capacity.
• Attention entropy is higher, and attention sink concentration is reduced compared to Softmax, implying broader contextual aggregation and improved routing.
• Score ranges expand with layer depth, evidencing the mechanism's increased expressive power.

Figure 5: Layer-wise Parallax score ranges, illustrating negative and extreme values enabled by covariance correction.

Practical Implications and Limitations

Parallax offers the following practical and theoretical implications:

• Its parameterized correction branch allows post-training adaptation of pretrained Transformers, unlike Linear Attention variants.
• The mechanism's optimizer dependence suggests that future affine attention designs must co-evolve with optimizer selection for full capacity utilization.
• The doubled arithmetic intensity opens avenues for kernel-level optimization, head dimension tuning, and hardware-aligned throughput maximization.
• The family of attention mechanisms described here clarifies that intercept-enabled variants (affine linear, affine MesaNet, affine DeltaNet) may offer superior memory and recall performance if matched with suitable optimizers.

Figure 6: Decode latency benchmarking demonstrates Parallax kernel efficiency, matching or surpassing FlashAttention on Hopper GPUs.

Figure 7: Activation norms for $q_i$0, $q_i$1, $q_i$2, and probe show optimizer and architectural effects during Parallax training.

Conclusion

Parallax generalizes local linear attention into a parameterized, hardware-efficient form for LLMs, achieving improved bias-variance tradeoffs and associative memory. Empirically, it persistently outperforms Softmax Attention under both parameter and compute-matched regimes, contingent on optimizer codesign. The analysis exposes both architectural and optimizer-driven spectral structure, confirming Parallax's correction branch activation via Muon. The approach unlocks new axes for attention mechanism design, kernel-level efficiency, and theoretical investigation into optimizer-architecture interaction, motivating further exploration of affine attention variants and scalable training recipes.