# Test-time regression: a unifying framework for designing sequence models with associative memory

arXiv: https://arxiv.org/abs/2501.12352 · HTML: https://mlsanity.com/paper/2501.12352/

## Key facts

- Cited by 74 papers in the map (data through 14 Oct 2026).
- Builds on 93 of its 136 references that are in the map.
- First posted on arXiv 21 Jan 2025.

## Builds on (most cited first, 30 of 93)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Pointer Sentinel Mixture Models](https://mlsanity.com/paper/1609.07843/) (2016)
- [Training Compute-Optimal Large Language Models](https://mlsanity.com/paper/2203.15556/) (2022)
- [Sequence to Sequence Learning with Neural Networks](https://mlsanity.com/paper/1409.3215/) (2014)
- [Convex optimization](https://mlsanity.com/paper/2106.01946/) (2021)
- [Efficiently Modeling Long Sequences with Structured State Spaces](https://mlsanity.com/paper/2111.00396/) (2021)
- [Matching Networks for One Shot Learning](https://mlsanity.com/paper/1606.04080/) (2016)
- [Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention](https://mlsanity.com/paper/2006.16236/) (2020)
- [Is Space-Time Attention All You Need for Video Understanding?](https://mlsanity.com/paper/2102.05095/) (2021)
- [Rethinking Attention with Performers](https://mlsanity.com/paper/2009.14794/) (2020)
- [GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints](https://mlsanity.com/paper/2305.13245/) (2023)
- [Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality](https://mlsanity.com/paper/2405.21060/) (2024)
- [DeepAR: Probabilistic Forecasting with Autoregressive Recurrent Networks](https://mlsanity.com/paper/1704.04110/) (2017)
- [In-context Learning and Induction Heads](https://mlsanity.com/paper/2209.11895/) (2022)
- [Efficient Transformers: A Survey](https://mlsanity.com/paper/2009.06732/) (2020)
- [Simplified State Space Layers for Sequence Modeling](https://mlsanity.com/paper/2208.04933/) (2022)
- [Chronos: Learning the Language of Time Series](https://mlsanity.com/paper/2403.07815/) (2024)
- [Transformers learn in-context by gradient descent](https://mlsanity.com/paper/2212.07677/) (2022)
- [Fast Transformer Decoding: One Write-Head is All You Need](https://mlsanity.com/paper/1911.02150/) (2019)
- [What Can Transformers Learn In-Context? A Case Study of Simple Function Classes](https://mlsanity.com/paper/2208.01066/) (2022)
- [Scaling Vision Transformers to 22 Billion Parameters](https://mlsanity.com/paper/2302.05442/) (2023)
- [Retentive Network: A Successor to Transformer for Large Language Models](https://mlsanity.com/paper/2307.08621/) (2023)
- [HiPPO: Recurrent Memory with Optimal Polynomial Projections](https://mlsanity.com/paper/2008.07669/) (2020)
- [YaRN: Efficient Context Window Extension of Large Language Models](https://mlsanity.com/paper/2309.00071/) (2023)

## Built on by (most cited first, 30 of 74)

- [Test-Time Training Done Right](https://mlsanity.com/paper/2505.23884/) (2025)
- [TTT3R: 3D Reconstruction as Test-Time Training](https://mlsanity.com/paper/2509.26645/) (2025)
- [One-Minute Video Generation with Test-Time Training](https://mlsanity.com/paper/2504.05298/) (2025)
- [Mamba-3: Improved Sequence Modeling using State Space Principles](https://mlsanity.com/paper/2603.15569/) (2026)
- [DeltaProduct: Improving State-Tracking in Linear RNNs via Householder Products](https://mlsanity.com/paper/2502.10297/) (2025)
- [Nested Learning: The Illusion of Deep Learning Architectures](https://mlsanity.com/paper/2512.24695/) (2025)
- [ATLAS: Learning to Optimally Memorize the Context at Test Time](https://mlsanity.com/paper/2505.23735/) (2025)
- [It's All Connected: A Journey Through Test-Time Memorization, Attentional Bias, Retention, and Online Optimization](https://mlsanity.com/paper/2504.13173/) (2025)
- [MesaNet: Sequence Modeling by Locally Optimal Test-Time Training](https://mlsanity.com/paper/2506.05233/) (2025)
- [Understanding Transformer from the Perspective of Associative Memory](https://mlsanity.com/paper/2505.19488/) (2025)
- [Scal3R: Scalable Test-Time Training for Large-Scale 3D Reconstruction](https://mlsanity.com/paper/2604.08542/) (2026)
- [ZipMap: Linear-Time Stateful 3D Reconstruction via Test-Time Training](https://mlsanity.com/paper/2603.04385/) (2026)
- [In-Place Test-Time Training](https://mlsanity.com/paper/2604.06169/) (2026)
- [Lattice: Learning to Efficiently Compress the Memory](https://mlsanity.com/paper/2504.05646/) (2025)
- [Comba: Improving Bilinear RNNs with Closed-loop Control](https://mlsanity.com/paper/2506.02475/) (2025)
- [Memory Caching: RNNs with Growing Memory](https://mlsanity.com/paper/2602.24281/) (2026)
- [Scaling Linear Attention with Sparse State Expansion](https://mlsanity.com/paper/2507.16577/) (2025)
- [Fast-weight Product Key Memory](https://mlsanity.com/paper/2601.00671/) (2026)
- [Test-Time Training with KV Binding Is Secretly Linear Attention](https://mlsanity.com/paper/2602.21204/) (2026)
- [LoLA: Low-Rank Linear Attention With Sparse Caching](https://mlsanity.com/paper/2505.23666/) (2025)
- [Gated KalmaNet: A Fading Memory Layer Through Test-Time Ridge Regression](https://mlsanity.com/paper/2511.21016/) (2025)
- [Spatial-TTT: Streaming Visual-based Spatial Intelligence with Test-Time Training](https://mlsanity.com/paper/2603.12255/) (2026)
- [Mechanistic evaluation of Transformers and state space models](https://mlsanity.com/paper/2505.15105/) (2025)
- [RAD: Redundancy-Aware Distillation for Hybrid Models via Self-Speculative Decoding](https://mlsanity.com/paper/2505.22135/) (2025)
- [Fast weight programming and linear transformers: from machine learning to neurobiology](https://mlsanity.com/paper/2508.08435/) (2025)
- [TNT: Improving Chunkwise Training for Test-Time Memorization](https://mlsanity.com/paper/2511.07343/) (2025)
- [Kalman Linear Attention: Parallel Bayesian Filtering For Efficient Language Modelling and State Tracking](https://mlsanity.com/paper/2602.10743/) (2026)
- [SplatPainter: Interactive Authoring of 3D Gaussians from 2D Edits via Test-Time Training](https://mlsanity.com/paper/2512.05354/) (2025)
- [Preconditioned DeltaNet: Curvature-aware Sequence Modeling for Linear Recurrences](https://mlsanity.com/paper/2604.21100/) (2026)
- [On State Reduction in Linear Attention](https://mlsanity.com/paper/2602.04852/) (2026)

## Abstract (arXiv)

Sequence models lie at the heart of modern deep learning. However, rapid advancements have produced a diversity of seemingly unrelated architectures, such as Transformers and recurrent alternatives. In this paper, we introduce a unifying framework to understand and derive these sequence models, inspired by the empirical importance of associative recall, the capability to retrieve contextually relevant tokens. We formalize associative recall as a two-step process, memorization and retrieval, casting memorization as a regression problem. Layers that combine these two steps perform associative recall via ``test-time regression'' over its input tokens. Prominent layers, including linear attention, state-space models, fast-weight programmers, online learners, and softmax attention, arise as special cases defined by three design choices: the regression weights, the regressor function class, and the test-time optimization algorithm. Our approach clarifies how linear attention fails to capture inter-token correlations and offers a mathematical justification for the empirical effectiveness of query-key normalization in softmax attention. Further, it illuminates unexplored regions within the design space, which we use to derive novel higher-order generalizations of softmax attention. Beyond unification, our work bridges sequence modeling with classic regression methods, a field with extensive literature, paving the way for developing more powerful and theoretically principled architectures.