Local Linear Attention: An Optimal Interpolation of Linear and Softmax Attention For Test-Time Regression

• URL: http://arxiv.org/abs/2510.01450v1
• Date: Wed, 01 Oct 2025 20:42:21 GMT
• Title: Local Linear Attention: An Optimal Interpolation of Linear and Softmax Attention For Test-Time Regression
• Authors: Yifei Zuo, Yutong Yin, Zhichen Zeng, Ang Li, Banghua Zhu, Zhaoran Wang,
• Abstract summary: 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.
• Score: 35.16407520369906
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Transformer architectures have achieved remarkable success in various domains. While efficient alternatives to Softmax Attention have been widely studied, the search for more expressive mechanisms grounded in theoretical insight-even at greater computational cost-has been relatively underexplored. In this work, we bridge this gap by proposing Local Linear Attention (LLA), a novel attention mechanism derived from nonparametric statistics through the lens of test-time regression. First, we show that LLA offers theoretical advantages over Linear and Softmax Attention for associative memory via a bias-variance trade-off analysis. Next, we address its computational challenges and propose two memory-efficient primitives to tackle the $\Theta(n^2 d)$ and $\Theta(n d^2)$ complexity. We then introduce FlashLLA, a hardware-efficient, blockwise algorithm that enables scalable and parallel computation on modern accelerators. In addition, we implement and profile a customized inference kernel that significantly reduces memory overheads. Finally, we empirically validate the advantages and limitations of LLA on test-time regression, in-context regression, associative recall and state tracking tasks. Experiment results demonstrate that LLA effectively adapts to non-stationarity, outperforming strong baselines in test-time training and in-context learning, and exhibiting promising evidence for its scalability and applicability in large-scale models. Code is available at https://github.com/Yifei-Zuo/Flash-LLA.

Related papers

• $\ abla$-Reasoner: LLM Reasoning via Test-Time Gradient Descent in Latent Space [71.23672814629448]
  $nabla$-Reasoner is an iterative generation framework that integrates differentiable optimization over token logits into the decoding loop.<n>$nabla$-Reasoner achieves over 20% accuracy improvement on a challenging mathematical reasoning benchmark.
  arXiv  Detail & Related papers  (2026-03-05T08:42:54Z)
• PAC-Bayes Bounds for Multivariate Linear Regression and Linear Autoencoders [10.40360267905115]
  Linear Autoencoders (LAEs) have shown strong performance in state-of-the-art recommender systems.<n>We investigate the generalizability -- a theoretical measure of model performance in statistical learning.
  arXiv  Detail & Related papers  (2025-12-15T01:12:11Z)
• Efficient Machine Unlearning via Influence Approximation [75.31015485113993]
  Influence-based unlearning has emerged as a prominent approach to estimate the impact of individual training samples on model parameters without retraining.<n>This paper establishes a theoretical link between memorizing (incremental learning) and forgetting (unlearning)<n>We introduce the Influence Approximation Unlearning algorithm for efficient machine unlearning from the incremental perspective.
  arXiv  Detail & Related papers  (2025-07-31T05:34:27Z)
• Generalized Linear Bandits: Almost Optimal Regret with One-Pass Update [60.414548453838506]
  We study the generalized linear bandit (GLB) problem, a contextual multi-armed bandit framework that extends the classical linear model by incorporating a non-linear link function.<n>GLBs are widely applicable to real-world scenarios, but their non-linear nature introduces significant challenges in achieving both computational and statistical efficiency.<n>We propose a jointly efficient algorithm that attains a nearly optimal regret bound with $mathcalO(1)$ time and space complexities per round.
  arXiv  Detail & Related papers  (2025-07-16T02:24:21Z)
• Kinetics: Rethinking Test-Time Scaling Laws [18.325591438335007]
  Kinetics Scaling Law suggests that test-time compute is more effective when used on models above a threshold than smaller ones.<n>Motivated by this, we propose a new scaling paradigm centered on sparse attention, which lowers per-token cost and enables longer generations and more parallel samples.
  arXiv  Detail & Related papers  (2025-06-05T17:59:24Z)
• MesaNet: Sequence Modeling by Locally Optimal Test-Time Training [67.45211108321203]
  We introduce a numerically stable, chunkwise parallelizable version of the recently proposed Mesa layer.<n>We show that optimal test-time training enables reaching lower language modeling perplexity and higher downstream benchmark performance than previous RNNs.
  arXiv  Detail & Related papers  (2025-06-05T16:50:23Z)
• Accelerating RL for LLM Reasoning with Optimal Advantage Regression [52.0792918455501]
  We propose a novel two-stage policy optimization framework that directly approximates the optimal advantage function.<n>$A$*-PO achieves competitive performance across a wide range of mathematical reasoning benchmarks.<n>It reduces training time by up to 2$times$ and peak memory usage by over 30% compared to PPO, GRPO, and REBEL.
  arXiv  Detail & Related papers  (2025-05-27T03:58:50Z)
• In-Context Linear Regression Demystified: Training Dynamics and Mechanistic Interpretability of Multi-Head Softmax Attention [52.159541540613915]
  We study how multi-head softmax attention models are trained to perform in-context learning on linear data.<n>Our results reveal that in-context learning ability emerges from the trained transformer as an aggregated effect of its architecture and the underlying data distribution.
  arXiv  Detail & Related papers  (2025-03-17T02:00:49Z)
• FAStEN: An Efficient Adaptive Method for Feature Selection and Estimation in High-Dimensional Functional Regressions [7.674715791336311]
  We propose a new, flexible and ultra-efficient approach to perform feature selection in a sparse function-on-function regression problem. We show how to extend it to the scalar-on-function framework. We present an application to brain fMRI data from the AOMIC PIOP1 study.
  arXiv  Detail & Related papers  (2023-03-26T19:41:17Z)

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.