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

2025/10/01 by Yutong Yin, Zuo, Yifei, Yin, Yutong +8
Computer Science · Engineering · #Artificial Intelligence (cs.AI) #FOS: Computer and information sciences #Face and Expression Recognition #Fault Detection and Control Systems #Machine Learning (cs.LG) #Neural Networks and Applications

paper · pdf · doi:10.48550/arxiv.2510.01450

openalex publication_date 2025/10/01 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28

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 Θ(n² d) and Θ(n d²) 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.

Citations

• MesaNet: Sequence Modeling by Locally Optimal Test-Time Training
• Understanding Transformer from the Perspective of Associative Memory
• Understanding the Skill Gap in Recurrent Language Models: The Role of the Gather-and-Aggregate Mechanism
• It's All Connected: A Journey Through Test-Time Memorization, Attentional Bias, Retention, and Online Optimization
• RWKV-7 "Goose" with Expressive Dynamic State Evolution
• Forgetting Transformer: Softmax Attention with a Forget Gate
• MoBA: Mixture of Block Attention for Long-Context LLMs
• Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention
• DeltaProduct: Improving State-Tracking in Linear RNNs via Householder Products
• Test-time regression: a unifying framework for designing sequence models with associative memory
• Gated Delta Networks: Improving Mamba2 with Delta Rule
• SeerAttention: Learning Intrinsic Sparse Attention in Your LLMs
• Longhorn: State Space Models are Amortized Online Learners
• Parallelizing Linear Transformers with the Delta Rule over Sequence Length
• Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality
• Chameleon: Mixed-Modal Early-Fusion Foundation Models
• Megalodon: Efficient LLM Pretraining and Inference with Unlimited Context Length
• Eagle and Finch: RWKV with Matrix-Valued States and Dynamic Recurrence
• Repeat After Me: Transformers are Better than State Space Models at Copying
• Gated Linear Attention Transformers with Hardware-Efficient Training
• Zoology: Measuring and Improving Recall in Efficient Language Models
• Mamba: Linear-Time Sequence Modeling with Selective State Spaces
• Small-scale proxies for large-scale Transformer training instabilities
• Uncovering mesa-optimization algorithms in Transformers
• FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning
• Retentive Network: A Successor to Transformer for Large Language Models
• One Step of Gradient Descent is Provably the Optimal In-Context Learner with One Layer of Linear Self-Attention
• Trained Transformers Learn Linear Models In-Context
• Hyena Hierarchy: Towards Larger Convolutional Language Models
• Scaling Vision Transformers to 22 Billion Parameters
• Why Can GPT Learn In-Context? Language Models Implicitly Perform Gradient Descent as Meta-Optimizers
• Transformers learn in-context by gradient descent
• General-Purpose In-Context Learning by Meta-Learning Transformers
• What learning algorithm is in-context learning? Investigations with linear models
• What Can Transformers Learn In-Context? A Case Study of Simple Function Classes
• The Parallelism Tradeoff: Limitations of Log-Precision Transformers
• How to Train Your HiPPO: State Space Models with Generalized Orthogonal Basis Projections
• FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
• Efficiently Modeling Long Sequences with Structured State Spaces
• Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation
• Linear Transformers Are Secretly Fast Weight Programmers
• Hopfield Networks is All You Need
• Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention
• Online normalizer calculation for softmax
• Attention Is All You Need
• Learning to Forget: Continual Prediction with LSTM

Related

• Gated Linear Attention Transformers with Hardware-Efficient Training
• Transformer Based Linear Attention with Optimized GPU Kernel Implementation
• Degrees of Freedom for Linear Attention: Distilling Softmax Attention with Optimal Feature Efficiency
• Log-Linear Attention
• Rethinking Attention with Performers