Source-linked AI summary
Learning to (Learn at Test Time): RNNs with Expressive Hidden States
Yu Sun, Xinhao Li, Karan Dalal, Jiarui Xu, Arjun Vikram, Genghan Zhang, Yann Dubois, Xinlei Chen, Xiaolong Wang, Sanmi Koyejo, Tatsunori Hashimoto, Carlos Guestrin
TL;DR
Existing RNNs retain linear complexity but struggle to exploit information in sufficiently long contexts. TTT makes hidden states learnable models updated through self-supervised optimization, and its layers continue reducing perplexity across longer contexts than Mamba, while TTT-MLP remains costly in wall-clock time.
Problem
Existing RNNs struggle to exploit additional information in long contexts despite their linear-complexity advantage.
Method
TTT provides a practical framework for arbitrary neural networks as hidden states, updated through an explicit self-supervised optimization problem during sequence processing.
Results
Transformer's average perplexity decreases throughout its 32k context, whereas Mamba's metric plateaus after 16k.
Takeaways & Limitations
TTT connects sequence modeling with inner-loop learning on temporally dependent test sequences, supporting continued conditioning on long contexts.
Takeaways & Limitations
TTT-MLP can be effective in FLOPs but incurs substantially greater wall-clock time because of its additional MLP complexity.
Abstract
from arXiv · showhide
Self-attention performs well in long context but has quadratic complexity. Existing RNN layers have linear complexity, but their performance in long context is limited by the expressive power of their hidden states. We present a practical framework for instantiating sequence modeling layers with linear complexity and expressive hidden states. The key idea is to make the hidden state a machine learning model itself, and the update rule a step of self-supervised learning. Since the hidden state is updated by training even on test sequences, our layers are called Test-Time Training (TTT) layers. We consider two instantiations: TTT-Linear and TTT-MLP, whose hidden state is a linear model and a two-layer MLP respectively. We evaluate our instantiations at the scale of 125M to 1.3B parameters, comparing with a strong Transformer and Mamba, a modern RNN. Similar to Transformer, TTT-Linear and TTT-MLP can keep reducing perplexity by conditioning on more tokens, while Mamba cannot after 16k context. TTT-MLP still faces challenges in memory I/O, but shows larger potential in long context, pointing to a promising direction for future research.
1 Introduction
Existing RNNs retain linear complexity by compressing context into fixed-size hidden states, but this limits long-context performance. TTT layers instead make the hidden state a trainable model updated through self-supervised learning, while practical efficiency improvements target accelerator execution.
- Motivation: Mamba scales similarly to a strong Transformer, but its average perplexity plateaus after 16k context while Transformer perplexity decreases across 32k.This exposes a long-context limitation despite Mamba’s progress over earlier LSTMs.
- Motivation: Fixed-size hidden states make RNNs efficient per token but limit their ability to exploit additional information in long contexts.The compression must represent structures and relationships across thousands or millions of tokens.
- TTT layers: TTT layers make the hidden state a machine-learning model and update it by self-supervised training on test sequences.The paper introduces TTT-Linear and TTT-MLP, using a linear model and a two-layer MLP as hidden states.
- TTT layers: Sequence-modeling layers differ through their initial state, update rule, and output rule; TTT and naive RNNs use fixed-size states, unlike self-attention’s growing state.Fixed-size states keep per-token cost constant, whereas self-attention’s context-dependent state increases cost per token.
- Efficiency: The dual form preserves the naive implementation’s output while training more than 5× faster on TPUs.The paper also uses mini-batches of tokens to improve parallelism on modern accelerators.
- Limitations: The framework can instantiate arbitrary neural networks as hidden states, but these instantiations may still require substantial wall-clock time.The paper identifies this efficiency limitation as unresolved even after its improvements.
2 Method
TTT layers compress context into a learned model state that is updated by self-supervised learning during inference, aiming to retain linear efficiency with greater long-context expressiveness. The framework supports different learners and optimizers, with efficient mini-batch and dual-form implementations.
- Motivation: RNN layers process tokens efficiently with fixed-size hidden states, but their long-context performance is limited by the expressive power of that compressed state.Self-attention stores historic context explicitly in a growing KV cache, making it more expressive but increasingly costly to scan.
- Learning a self-supervised task for TTT: The self-supervised learner reconstructs corrupted inputs, requiring it to discover correlations among input dimensions; gradient descent reduces but does not eliminate the loss.The self-supervised task is optimized end-to-end for next-token prediction rather than handcrafted from human priors.
- TTT as updating a hidden state: TTT makes the hidden state a machine learning model and updates its weights through self-supervised learning, even separately for each test sequence.The output uses the updated weights, so the layer trains a sequence of weights W1,...,WT during its forward pass.
- Learning a self-supervised task for TTT: TTT uses nested optimization: the inner loop updates the learner weights W, while the outer loop optimizes the remaining network parameters and task-related hyper-parameters.This lets TTT layers be integrated and trained in larger architectures using the same general setup as other language models.
- Parallelization with mini-batch TTT: Mini-batch gradient descent balances quality and parallelism, while the dual form avoids materializing intermediate gradients and weights to improve accelerator utilization.The dual form is output-equivalent to the primal form and is reported as more than 5× faster in the JAX implementation.
- Dual form: A learner defined by a model and optimizer induces a TTT layer; linear-model batch GD recovers linear attention, while a Nadaraya-Watson learner recovers self-attention.The framework therefore unifies parametric and nonparametric constructions under the learner abstraction.
3 Experiments
The experiments compare TTT-Linear and TTT-MLP with Transformer and Mamba across Pile and Books contexts, model sizes, backbones, and inference latency. TTT layers increasingly outperform Mamba as context grows, while their efficiency depends on implementation and their scaling-law behavior remains unclear.
- Experimental setup: The evaluation compares TTT-Linear and TTT-MLP with Transformer and Mamba across 125M–1.3B-scale models and Pile and Books contexts.Experiments use standard 2k and 8k Pile contexts and Books contexts from 1k to 32k, with specified backbone and fairness protocols.
- Pile results: At 2k Pile context, TTT-Linear, Mamba, and Transformer perform comparably, while TTT-MLP performs slightly worse under large FLOP budgets.Although TTT-MLP has better perplexity than TTT-Linear at every model size, its additional FLOP cost offsets that advantage.
- Pile results: At 8k Pile context, both TTT-Linear and TTT-MLP significantly outperform Mamba, and the advantage of TTT layers widens as context length grows.TTT-MLP with the Transformer backbone also slightly outperforms Mamba around 1.3B, while Transformer remains competitive in perplexity but costly in FLOPs.
- Books results: On Books at 2k context, Mamba slightly outperforms TTT-Linear; at 32k, both TTT-Linear and TTT-MLP outperform Mamba.TTT-MLP with the Transformer backbone also slightly outperforms Mamba at 32k context.
- Books results: TTT-MLP with matched FLOPs performs worse at short context but better at long context than TTT-Linear, consistent with greater hidden-state expressiveness being more useful in long context.At 1.3B, TTT-MLP with the Transformer backbone is only slightly worse than its Mamba-backbone counterpart, suggesting potential beyond the evaluated scale and context range.
- Efficiency and limitations: Transformer time per token grows linearly with context length, whereas the other methods remain roughly constant; TTT implementations nevertheless face efficiency and scaling limitations.The dual form trains more than 5× faster than the naive implementation on TPUs, but TTT-MLP can still require substantial wall-clock time and the experiments lack a clean linear FLOP–perplexity fit.
4 Related Work
The paper situates TTT within test-time learning and fast-weight research, emphasizing its explicit self-supervised optimization and broader hidden-state framework. It also distinguishes TTT from prior approaches and acknowledges substantial wall-clock costs for richer instantiations.
- Test-Time Training: Prior test-time learning includes local learning, transductive learning, computer-vision adaptation, and NLP dynamic evaluation.Dynamic evaluation directly fine-tunes a language model on the test sequence, often represented by a prompt.
- Test-Time Training: Test-time training formulates each test instance as its own learning problem and trains a model on that instance before prediction.The test instance is the target of generalization, often using the conventional predictor as initialization.
- Test-Time Training: TTT uses a learned outer-loop reconstruction task rather than a handcrafted task based on human priors.This design supports different self-supervised-task choices across applications.
- Fast Weights and Fast Weight Programmers: Fast weights update a model on relevant data, while TTT additionally formulates an explicit learning problem and uses an explicit optimization step.TTT can therefore be viewed as a special case of fast weights.
- Modern RNN layers: TTT extends modern matrix-state RNN work with a practical framework for instantiating arbitrary neural networks as hidden states.The framework includes TTT-MLP, but richer hidden states can require substantial wall-clock time even after efficiency improvements.
- Learning to Learn: TTT treats each sequence as a dataset, placing its inner loop below regular training and making the outer loop easier to scale.The outer loop remains at the same level as ordinary supervised training, while its parameters become inner-loop hyperparameters.
5 Discussion
The discussion presents TTT as a reformulation of supervised learning with nested loops and identifies broader scalability, systems, context-length, and model-design questions for future work.
- Discussion: TTT reformulates supervised learning as learning to learn with two nested loops, providing an alternative framework for sequence-modeling architectures.The paper summarizes its current instantiation in Table 2.
- Future work: The search space for effective TTT instantiations is large, and the paper describes its current work as an initial step.The authors expect heuristics from regular training may transfer to test-time training.
- Outer-loop parameterization: The first outer-loop parameterization may not be optimal, motivating broader families of multi-view reconstruction or self-supervised tasks.The paper frames finding a better parameterization as an open direction rather than a settled design choice.
- Systems optimization: Systems optimization remains preliminary, and pipeline parallelism could distribute million-token sequences across multiple devices.This direction targets more efficient processing of very long sequences.
- Longer context and larger models: The experiments do not train with millions or billions of context tokens, and longer contexts would require larger models according to Figure 16.The paper expects TTT-layer advantages to become more pronounced at longer context lengths.
- More ambitious instantiations: Longer contexts may require larger hidden models, including convolutional networks for video and embodied-agent workloads.The paper identifies these workloads as contexts that can reach millions or billions of tokens.
- Multi-level learning to learn: Nested learning could continue by making the hidden model f itself a self-attention layer.Theorem 2 is cited as the basis for interpreting this as another inner loop.
- Why study TTT?: The paper connects TTT to human learning because its inner loop processes temporally dependent sequences without a conventional train-test split.Each piece of sequence data can serve both training and testing roles.
A Dual Form
The dual-form appendix derives a hardware-efficient implementation of TTT-MLP’s forward and update computations. It replaces per-token, state-dependent operations with batched matrix operations while preserving the primal form’s outputs.
- A Dual Form: The appendix derives a dual form for general-depth MLPs with nonlinear activations.The derivation introduces layer-indexed parameters and token-batch matrix notation.
- A.1 Forward pass: The initial TTT forward pass processes each layer over a token mini-batch using element-wise operations and differentiable matrix computations.The input to each layer is represented as a matrix whose columns are token vectors, and the loss is computed from the final-layer output.
- A.1 Forward pass: The initial operations are hardware efficient because, apart from activations, they consist of matrix multiplications and sums.Both primal and dual forms share these initial computations.
- A.2 Primal form: The primal form updates parameters for each token and then repeats the forward pass with the updated weights.This produces output tokens but prevents standard batching because different tokens use different parameter states.
- A.2 Primal form: The primal form has poor hardware efficiency because individual per-token gradient terms cannot be batched into matrix multiplications.Its forward pass likewise uses a different weight state for each test view.
- A.3 Dual form: The dual form computes the final output and updated weights using batched operations while avoiding explicit intermediate gradient variables.It first forms the output gradient, applies a standard backward pass, and then performs another forward pass.
- A.3 Dual form: The dual forward pass computes the same output as the primal form using matrix multiplications, sums, activations, and a causal mask.The appendix reports that the resulting pass is efficient like a standard forward pass.
A.4 Derivation
The derivation establishes the dual form by expressing each token’s cumulative gradient contribution as a masked matrix product, then substituting the TTT variables into that identity.
- A.4 Derivation: The derivation aims to prove that the dual-form computation matches the primal form’s per-token forward-pass quantities.It compares each dual-form column with the corresponding primal-form quantity.
- A.4 Derivation: Fact 1 expresses cumulative vector contributions as a matrix product with a causal mask.For matrices A, Q, and V, the masked product produces each column’s sum over preceding indexed terms.
- A.4 Derivation: Substituting the TTT layer’s activations, queries, and gradient terms into Fact 1 yields the desired equality between dual and primal computations.The substitution identifies A with the layer input, Q with the test-view input, V with the output gradient, and the result with the weight-state correction.
- A.4 Derivation: The derivation can extend beyond element-wise activations to operations such as normalization using vector-Jacobian products.The dual form does not accelerate computations inside the operation or its vector-Jacobian product.
- A.4 Derivation: The notation uses A rather than K to avoid confusion with the layer-number index K.Superscripts identify layers, while subscripts identify time.
B Nadaraya-Watson estimator
This section derives the Nadaraya-Watson estimator by replacing unknown probability distributions with kernel density estimates and combining them through product kernels. It also discusses asymmetric kernels used in the estimator.
- The derivation formulates the desired output token as the conditional expectation of z given input token x.
- Unknown distributions p(x) and p(x,z) are replaced with kernel density estimates.
- The joint distribution estimate uses a product kernel, whose factors need not be independent because one factor is integrated out.
- Plugging the density estimates into the derivation yields the Nadaraya-Watson estimator.
- The estimator can use asymmetric kernels; when it uses θK and θQ, it is described as a balloon estimator.
C Experiment details
The experiments reproduce the Mamba training setup while specifying architecture, optimization, tokenizer, context-length, and memory-management choices. Evaluations also match methods by training FLOPs in the main comparison.
- Architectures: The Transformer follows Mamba’s construction, using Llama-based components, while the Mamba baseline uses the authors’ public code and reproduces reported numbers.
- Training configurations: All models use a 0.5M-token batch size regardless of context length, with optimization hyperparameters following Mamba’s improved recipe.
- Training configurations: The training-configuration table reproduces Mamba’s setup, with reported Mamba and Transformer learning rates set to 5× its table values.
- Data and context: For Books, the experiments test RoPE θ = 500,000 from context length 4k because θ = 10,000 is sub-optimal for the Transformer baseline.
- Optimization: Transformer finetuning tests peak learning rates of 1e-5, 1e-4, and 1e-3, selecting the value with the best perplexity.
- Optimization: TTT-Linear uses ηbase = 1 and TTT-MLP uses ηbase = 0.1, selected as the largest tested values that avoid instabilities.
- Evaluation fairness: Figure 2 experiments match training FLOPs across methods, while the Transformer baseline has 2.8× the inference FLOPs.
- Memory: Gradient checkpointing reduces memory use by avoiding storage of every TTT hidden-state intermediate across the sequence.