Parallax: Parameterized Local Linear Attention for Language Modeling

Hugging Face Daily Papers Papers
attention-mechanism language-modeling local-linear-attention hardware-aware-optimization llm-pretraining efficient-attention architecture-optimizer-codesign

Summary

Introduces Parallax, a parameterized local linear attention mechanism with hardware-aware optimization that improves LLM pretraining efficiency and performance, achieving Pareto improvements at 0.6B and 1.7B scales.

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.
Original Article
View Cached Full Text

Cached at: 05/29/26, 02:59 AM

Paper page - Parallax: Parameterized Local Linear Attention for Language Modeling

Source: https://huggingface.co/papers/2605.29157

Abstract

Local Linear Attention is enhanced through parameterization and hardware-aware optimization to improve LLM training efficiency and performance while maintaining computational stability.

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 anattention mechanismderived from nonparametric statistics in the test-time regression framework. In contrast to prior research on efficient attention variants, LLA upgrades the local constant estimate insoftmax attentionto a local linear estimate, yielding provably superior bias-variance tradeoffs forassociative memory. However, LLA has not been scaled in LLMpretrainingdue to computational and numerical stability concerns. We introduce Parallax, a parameterizedLocal Linear Attentionthat is scalable for LLMs. Parallax eliminates the numerical solver in LLA and learns an extraquery-like projectorthat probes theKV covariance. We place Parallax within a family ofattention mechanisms connected by thebandwidth, the probe construction and theaffine structure. We propose a hardware-aware algorithm that increases thearithmetic intensityoverFlashAttention, shifting attention into a more compute bound regime. Our prototypedecode kernelmatches or outperformsFlashAttention2/3 across diverse batch sizes and context lengths. We pretrain Parallax at 0.6B and 1.7B scales and find consistentperplexityimprovements throughoutpretrainingwith gains that transfer to downstream benchmarks. The advantage persists under both parameter-matched and compute-matched controls, demonstrating a Pareto improvement. We perform carefulpretrainingablations and identify a novel phenomenon wherebyMuonunlocks the capacity of Parallax. To our knowledge, this is the first empirical demonstration of strongarchitecture-optimizer codesignforattention mechanisms in the architecture research literature.

View arXiv pageView PDFGitHubAdd to collection

Get this paper in your agent:

hf papers read 2605\.29157

Don’t have the latest CLI?curl \-LsSf https://hf\.co/cli/install\.sh \| bash

Models citing this paper0

No model linking this paper

Cite arxiv.org/abs/2605.29157 in a model README.md to link it from this page.

Datasets citing this paper0

No dataset linking this paper

Cite arxiv.org/abs/2605.29157 in a dataset README.md to link it from this page.

Spaces citing this paper0

No Space linking this paper

Cite arxiv.org/abs/2605.29157 in a Space README.md to link it from this page.

Collections including this paper0

No Collection including this paper

Add this paper to acollectionto link it from this page.

Similar Articles

Lighthouse Attention (11 minute read)

TLDR AI

Lighthouse Attention is a selection-based hierarchical attention mechanism that accelerates long-context pretraining by running forward+backward passes ~17× faster at 512K context and delivering 1.4–1.7× end-to-end speedup at 98K context, validated with Llama-3 530M on 50B tokens.

Data-Efficient Adaptation of LLMs via Attention Head Reweighting

arXiv cs.LG

Introduces Attention Head Reweighting (AHR), a data-efficient method for adapting LLMs to text classification tasks by learning a single scalar per attention head, drastically reducing trainable parameters while outperforming LoRA in limited data settings.

Dynamic Linear Attention

Hugging Face Daily Papers

DLA introduces adaptive state merging and capacity-bounded memory modeling for multi-state linear attention, improving long-context LLM performance.