Parallax: Parameterized Local Linear Attention for Language Modeling

• URL: http://arxiv.org/abs/2605.29157v1
• Date: Wed, 27 May 2026 22:50:44 GMT
• Title: Parallax: Parameterized Local Linear Attention for Language Modeling
• Abstract summary: We introduce Parallax, a parameterized Local Linear Attention that is scalable for Large Language Models (LLMs)<n>We place Parallax within a family of attention mechanisms connected by the bandwidth, the probe construction and the affine structure.<n>Our prototype decode kernel matches or outperforms FlashAttention 2/3 across diverse batch sizes and context lengths.
• Score: 19.36455749398556
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• PALoRA: Projection-Adaptive LoRA for Preserving Reasoning in Large Language Models [11.396379339626122]
  We show that information essential for reasoning is not localized only in dominant singular directions, but is instead distributed across the singular spectrum.<n>We introduce PALoRA, a two-stage framework for knowledge injection with reduced interference.
  arXiv  Detail & Related papers  (2026-05-23T12:34:08Z)
• Bottleneck Transformer-Based Approach for Improved Automatic STOI Score Prediction [16.426476430697587]
  We present a novel approach to predict the Short-Time Objective Intelligibility (STOI) metric using a bottleneck transformer architecture.<n>Our model has shown higher correlation and lower mean squared error for both seen and unseen scenarios.
  arXiv  Detail & Related papers  (2026-02-17T10:46:54Z)
• RooflineBench: A Benchmarking Framework for On-Device LLMs via Roofline Analysis [53.90240071275054]
  The transition toward localized intelligence through Small Language Models (SLMs) has intensified the need for rigorous performance characterization on resource-constrained edge hardware.<n>We propose a systematic framework that unifies architectural primitives and hardware constraints through the lens of operational intensity (OI)<n>By defining an inference-potential region, we introduce the Relative Inference Potential as a novel metric to compare efficiency differences between Large Language Models (LLMs) on the same hardware substrate.
  arXiv  Detail & Related papers  (2026-02-12T03:02:22Z)
• LINA: Linear Autoregressive Image Generative Models with Continuous Tokens [56.80443965097921]
  Autoregressive models with continuous tokens form a promising paradigm for visual generation, especially for text-to-image (T2I) synthesis.<n>We study how to design compute-efficient linear attention within this framework.<n>We present LINA, a simple and compute-efficient T2I model built entirely on linear attention, capable of generating high-fidelity 1024x1024 images from user instructions.
  arXiv  Detail & Related papers  (2026-01-30T06:44:33Z)
• The Path Not Taken: RLVR Provably Learns Off the Principals [85.41043469428365]
  We show that sparsity is a surface artifact of a model-conditioned optimization bias.<n>We mechanistically explain these dynamics with a Three-Gate Theory.<n>We provide a parameter-level characterization of RLVR's learning dynamics.
  arXiv  Detail & Related papers  (2025-11-11T18:49:45Z)
• ShishuLM: Lightweight Language Model with Hybrid Decoder-MLP Architecture and Paired Weight Sharing [0.5565728870245015]
  We introduce an efficient language model architecture, referred to as ShishuLM, which reduces both the parameter count and Key-Value (KV) cache requirements.<n>Our results show that ShishuLM provides up to 25% reduction in memory requirements and up to 40% improvement in latency during both training and inference, compared to parent models.
  arXiv  Detail & Related papers  (2025-10-13T04:04:54Z)
• Local Linear Attention: An Optimal Interpolation of Linear and Softmax Attention For Test-Time Regression [35.16407520369906]
  Local Linear Attention is a novel attention mechanism derived from nonparametric statistics through the lens of test-time regression.<n>We introduce FlashLLA, a hardware-efficient, blockwise algorithm that enables scalable and parallel computation on modern accelerators.<n>Experiment results demonstrate that LLA effectively adapts to non-stationarity, outperforming strong baselines in test-time training and in-context learning.
  arXiv  Detail & Related papers  (2025-10-01T20:42:21Z)
• LESA: Learnable LLM Layer Scaling-Up [57.0510934286449]
  Training Large Language Models (LLMs) from scratch requires immense computational resources, making it prohibitively expensive.<n>Model scaling-up offers a promising solution by leveraging the parameters of smaller models to create larger ones.<n>We propose textbfLESA, a novel learnable method for depth scaling-up.
  arXiv  Detail & Related papers  (2025-02-19T14:58:48Z)
• Localized Gaussians as Self-Attention Weights for Point Clouds Correspondence [92.07601770031236]
  We investigate semantically meaningful patterns in the attention heads of an encoder-only Transformer architecture.<n>We find that fixing the attention weights not only accelerates the training process but also enhances the stability of the optimization.
  arXiv  Detail & Related papers  (2024-09-20T07:41:47Z)
• Building on Efficient Foundations: Effectively Training LLMs with Structured Feedforward Layers [16.253898272659242]
  State-of-the-art results in large language models (LLMs) often rely on scale, which becomes computationally expensive. Our study focuses on transformer-based LLMs, specifically targeting the computationally intensive feedforward networks (FFNs) We show that wide and structured networks can utilize training FLOPs more efficiently, with fewer parameters and lower loss than dense models at their optimal trade-off.
  arXiv  Detail & Related papers  (2024-06-24T08:43:21Z)
• Bilevel Fast Scene Adaptation for Low-Light Image Enhancement [50.639332885989255]
  Enhancing images in low-light scenes is a challenging but widely concerned task in the computer vision. Main obstacle lies in the modeling conundrum from distribution discrepancy across different scenes. We introduce the bilevel paradigm to model the above latent correspondence. A bilevel learning framework is constructed to endow the scene-irrelevant generality of the encoder towards diverse scenes.
  arXiv  Detail & Related papers  (2023-06-02T08:16:21Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.