Learning to (Learn at Test Time):
RNNs with Expressive Hidden States

The process of parametric learning can be viewed as compressing a massive training set into the weights of a model. Specifically, we know that models trained with self-supervision can capture the underlying structures and relationships behind their training data [48] – exactly what we need from a compression heuristic.

LLMs themselves are great examples. Trained with the self-supervised task of next-token prediction, their weights can be viewed as a compressed form of storage for existing knowledge on the internet. By querying LLMs, we can extract knowledge from their weights. More importantly, LLMs often exhibit a deep understanding of the semantic connections among existing knowledge to express new pieces of reasoning [1].

Our key idea is to use self-supervised learning to compress the historic context $x_{1},\dots,x_{t}$ into a hidden state $s_{t}$, by making the context an unlabeled dataset and the state a model. Concretely, the hidden state $s_{t}$ is now equivalent to $W_{t}$, the weights of a model $f$, which can be a linear model, a small neural network, or anything else. The output rule is simply:

Intuitively, the output token is just the prediction on $x_{t}$, made by $f$ with the updated weights $W_{t}$. The update rule is a step of gradient descent on some self-supervised loss $\ell$:

with learning rate $\eta$.²²2 For now, consider $W_{0}=0$. We will discuss more sophisticated techniques for initializing $W$ in Subsection 2.7. From the compression point of view, every heuristic needs to decide which input to remember or forget. Our $W$ remembers inputs that produce large gradients – intuitively, inputs that make $W$ learn a lot.

One choice of $\ell$ is reconstructing $x_{t}$ itself. To make the learning problem nontrivial, we first process $x_{t}$ into a corrupted input $\tilde{x}_{t}$ (details in Subsection 2.3), then optimize:

Similar to denoising autoencoders [75], $f$ needs to discover the correlations between dimensions of $x_{t}$ in order to reconstruct it from partial information $\tilde{x}_{t}$.³³3 In past experiments, we have also tried adding another model $g$ (decoder) after $f$ (encoder), such that the reconstruction is produced by $g\circ f$ instead of only $f$ itself. While this heftier design did slightly improve results, it made overall training less stable and added significant computational cost. Therefore we focus on the encoder-only design. As shown in Figure 5, gradient descent is able to reduce $\ell$, but cannot reduce it to zero. We discuss more sophisticated formulations of the self-supervised task in Subsection 2.3.

As with other RNN layers and self-attention, our algorithm that maps an input sequence $x_{1},\dots,x_{T}$ to output sequence $z_{1},\dots,z_{T}$ can be programmed into the forward pass of a sequence modeling layer, using the hidden state, update rule, and output rule above. Even at test time, our new layer still trains a different sequence of weights $W_{1},\dots,W_{T}$ for every input sequence. Therefore, we call it the Test-Time Training (TTT) layer.

The forward pass of a TTT layer also has a corresponding backward pass. Our forward pass only consists of standard differentiable operators except the gradient operator $\nabla$. However, $\nabla$ just maps one function to another, in this case $\ell$ to $\nabla\ell$, and $\nabla\ell$ is also composed of differentiable operators. Conceptually, calling backward on $\nabla\ell$ means taking gradients of gradients – a well explored technique in meta-learning [51].

TTT layers have the same interface as RNN layers and self-attention, therefore can be replaced in any larger network architecture, which usually contains many of these sequence modeling layers. Training a network with TTT layers also works the same way as training any other language model, such as a Transformer. The same data, recipe, and objective such as next-token prediction can be used to optimize parameters of the rest of the network.

We refer to training the larger network as the outer loop, and training $W$ within each TTT layer as the inner loop. An important difference between the two nested learning problems is that the inner-loop gradient $\nabla\ell$ is taken w.r.t. $W$, the parameters of $f$, while the outer-loop gradient is taken w.r.t the parameters of the rest of the network, which we will denote by $\theta_{\text{rest}}$. Throughout this paper, outer-loop parameters are always denoted by $\theta$ with various subscripts.

So far, the TTT layer has no outer-loop parameters, in contrast to other RNN layers and self-attention. In Subsection 2.3, we add outer-loop parameters to the TTT layer to improve its self-supervised task. Then in Subsection 2.4 and 2.5, we discuss two ways to improve the wall-clock time of TTT layers.

Arguably the most important part of TTT is the self-supervised task, because it determines the kind of features that $W$ will learn from the test sequence. So how should we design this task? The final goal of TTT is for $z_{t}=f(x_{t};W_{t})$ to perform well on language modeling. Instead of handcrafting a self-supervised task from human priors, we take a more end-to-end approach – directly optimizing the self-supervised task for the final goal of next-token prediction.

Concretely, we learn the self-supervised task as part of the outer loop. Starting from the naive reconstruction task in Equation 3, we add some outer-loop parameters to make this task learnable. In Subsection 2.1, we did not specify the corruption that produces $\tilde{x}_{t}$ from $x_{t}$. One design is to make it a low-rank projection $\tilde{x}_{t}=\theta_{K}x_{t}$, where $\theta_{K}$ is a learnable matrix.⁴⁴4 The subscript $K$ hints at a connection to self-attention, as we will establish in Subsection 2.6. Following the terminology of multi-view reconstruction, $\theta_{K}x_{t}$ is called a training view [14].

Moreover, perhaps not all the information in $x_{t}$ is worth remembering, so the reconstruction label can be another low-rank projection $\theta_{V}x_{t}$ instead of $x_{t}$. Here $\theta_{V}x_{t}$ is called the label view, where $\theta_{V}$ is also learnable. In summary, our new self-supervised loss is:

Since both $W$ and various $\theta$s appear together in Equation 4, we emphasize again their difference in nature. In the inner loop, only $W$ is optimized, therefore written as an argument of $\ell$; the $\theta$s are "hyper-parameters" of this loss function. In the outer loop, $\theta_{K},\theta_{V},\theta_{Q}$ are optimized alongside $\theta_{\text{rest}}$, and $W$ is merely a hidden state, not a parameter. Figure 6 illustrates this difference with code, where $\theta_{K}$ and $\theta_{V}$ are implemented as parameters of the TTT layer, analogous to the Key and Value parameters of self-attention.

Lastly, the training view $\theta_{K}x_{t}$ has fewer dimensions than $x_{t}$, so we can no longer use the output rule in Equation 1. The simplest solution is to create a test view $\theta_{Q}x_{t}$, and change our output rule to:

This solution has an additional benefit. The training and label views specify the information in $x_{t}$ that is compressed into $W_{t}$ and propagated forward through time. The test view specifies potentially different information that is mapped to the current output token $z_{t}$ and propagated forward through network layers, therefore adds more flexibility to the self-supervised task.

Altogether, the set of all possible choices for $\theta_{K},\theta_{Q},\theta_{V}$ induces a family of multi-view reconstruction tasks, and the outer loop can be interpreted as selecting a task from this family. Here we have designed all views as linear projections for simplicity. Future work might experiment with more flexible transformations, or bigger and different families of self-supervised tasks.

The naive TTT layer developed so far is already efficient in the number of floating point operations (FLOPs). However, its update rule $W_{t}=W_{t-1}-\eta\nabla l(W_{t-1};x_{t})$ cannot be parallelized, because $W_{t}$ depends on $W_{t-1}$ in two places: before the minus sign and inside $\nabla l$. Since $\nabla l$ contains the bulk of the computation, we focus on making this second part parallel.

We approach this systems challenge through concepts in the TTT framework. There are many variants of gradient descent (GD). The general update rule of GD can be expressed as:

where $G_{t}$ is the descent direction. Note that once we have calculated $G_{t}$ for $t=1,\dots,T$, we can then obtain all the $W_{t}$s through a cumsum by the second half of Equation 6. Our naive update rule, known as online gradient descent, uses $G_{t}=\nabla l(W_{t-1};x_{t})$.

To parallelize $G_{t}$ for $t=1,\dots,T$, we can take all of them w.r.t. $W_{0}$. This variant with $G_{t}=\nabla\ell\left(W_{0};x_{t}\right)$ is known as batch gradient descent, since $\sum_{s=1}^{t}\nabla\ell\left(W_{0};x_{t}\right)$ is the same as the gradient w.r.t. $W_{0}$ over $x_{1},\dots,x_{t}$ as a batch. However, in batch GD, $W_{t}$ is effectively only one gradient step away from $W_{0}$, in contrast to online GD, where $W_{t}$ is $t$ steps away from $W_{0}$. Therefore, batch GD has a smaller effective search space, which ends up hurting performance for language modeling.

Our proposed solution – mini-batch gradient descent – is shown in Figure 7. Denote the TTT batch size by $b$. We use $G_{t}=\nabla\ell\left(W_{t^{\prime}};x_{t}\right)$, where $t^{\prime}=t\,-\,$mod$(t,b)$ is the last timestep of the previous mini-batch (or 0 for the first mini-batch), so we can parallelize $b$ gradient computations at a time. Empirically, $b$ controls a trade-off between speed and quality, as shown in Figure LABEL:fig:batch-size. We chose $b=16$ for all experiments in this paper.

In summary, there are two potential channels to propagate information from $W_{s}$ to $W_{t}$ where $s<t$: cumsum and the gradient operator. The cumsum is always active, but the gradient channel is only active when $W_{s}$ is from a previous mini-batch. Different variants of gradient descent only affect the gradient channel, i.e., the descent direction $G_{t}$, specifically w.r.t. which $W$ the gradient is taken. However, the descent step $W_{t}=W_{t-1}-\eta\,G_{t}$ always starts from $W_{t-1}$, due to the autoregressive nature of the update rule, which is orthogonal to the choice of $G_{t}$.

The parallelization introduced above is necessary but not sufficient for efficiency in wall-clock time. Modern accelerators specialize in matrix-matrix multiplications, known as matmuls. For example, the NVIDIA A100 GPU contains highly optimized units called TensorCores that can only perform a single operation – multiplying two matrices each of size $16\times 16$. Without enough of these matmuls, the TensorCores are idle, and most of the potential for the A100 is unrealized.

Unfortunately, the TTT layer developed so far even with mini-batch still has very few matmuls. Consider the simplest case of $\ell$, where $\theta_{K}=\theta_{V}=\theta_{Q}=I$, for only the first TTT mini-batch of size $b$. In addition, consider $f$ as a linear model. Copying Equation 3, our loss at time $t$ is:

As discussed in Subsection 2.4, we can parallelize the computation of:

for $t=1,\dots,b$. However, we cannot compute all $b$ of the $G_{t}$s through a single matmul. Instead, we need $b$ outer products to compute them one by one. To make matters worse, for each $x_{t}\in\mathbb{R}^{d}$, $G_{t}$ is $d\times d$, which incurs much heavier memory footprint and I/O cost than $x_{t}$ for large $d$.

To solve these two problems, we make a simple observation: We do not actually need to materialize $G_{1},\dots,G_{b}$ as long as we can compute $W_{b}$ at the end of the mini-batch, and the output tokens $z_{1},\dots,z_{b}$ (see Figure 7). Now we demonstrate these computations with the simplified TTT-Linear case above. Denote $X=[x_{1},\dots,x_{b}]$, then:

So $W_{b}$ can be conveniently computed with a matmul. To compute $Z=[z_{1},\dots,z_{b}]$, we know that:

Denote $\delta_{t}=\sum_{s=1}^{t}(W_{0}x_{s}-x_{s})x_{s}^{T}x_{s}$ and the matrix $\Delta=[\delta_{1},\dots,\delta_{b}]$. We can derive that:

where mask is the lower triangular mask with zeros (similar to the attention mask, but with zeros instead of infinities), and the term $W_{0}X-X$ can be reused from the computation of $W_{b}$. Now $\Delta$ is also conveniently computed with matmuls. Plugging $\Delta$ back into Equation 7, we obtain $Z=W_{0}X-2\eta\Delta$.

We call this procedure the dual form, in contrast to the primal form before this subsection, where the $G$s and $W$s are explicitly materialized. As discussed, the two forms are equivalent in output. The terminology of primal and dual follows prior work that has explored similar mathematical formulations outside of TTT [36, 8, 59]. In Appendix A, we show that the dual form still works when $f$ is a neural network with nonlinear layers, except with more complicated notation.

Time complexity of the primal form within a TTT mini-batch is $O(b\times d^{2})$. Time complexity of the dual form is $O(b\times d^{2})$ for computing $W_{b}$ alone, then an additional $O(b^{2}\times d)$ for computing $z_{1},\dots,z_{b}$. Compared to the primal, the dual form sacrifices theoretical complexity for hardware utilization. In practice, $d$ is typically a few hundred and $b$ is chosen to be only 16. As a consequence, wall-clock time for computing $z_{1},\dots,z_{b}$ is relatively small, as observed in the right panel of Figure LABEL:fig:batch-size. In our JAX implementation, training with the dual form is more than $5\times$ faster than with primal.

In Subsection 2.1, we mentioned that $f$ can be a linear model or a neural network. In Subsection 2.4, we also discussed three variants of the update rule: online GD, batch GD, and mini-batch GD. Each of these $2\times 3$ combinations induces a different instantiation of the TTT layer, as illustrated in Figure 9. We now show that among these induced instantiations, the TTT layer with a linear model and batch GD is equivalent to linear attention [41], a widely known RNN layer.⁷⁷7 In a nutshell, linear attention [41] is simply self-attention without the softmax. Recall the definition of self-attention: $z_{t}=V_{t}\texttt{softmax}\left(K^{T}_{t}q_{t}\right)$. Without softmax, this becomes $z_{t}=V_{t}\left(K^{T}_{t}q_{t}\right)=\sum_{s=1}^{t}v_{s}k_{s}^{T}q_{t}$, which is the simplest formulation of linear attention. Similar to other RNN layers, it can be written in a recurrent form, where $\sum_{s=1}^{t}v_{s}k_{s}^{T}$ is the hidden state. Since $\sum_{s=1}^{t}v_{s}k_{s}^{T}$ can be computed in a cumsum for every $t=1,\dots,T$, linear attention also has linear complexity w.r.t. $T$.

In Table 1, we first empirically verify the equivalence above with an improved implementation of linear attention.⁸⁸8 The original formulation of linear attention in [41] contains a normalizer and a feature expansion on $x_{t}$, which can still be included in an equivalent TTT layer. However, prior work has found that these two additions can hurt performance [58], which we have verified in our own experiment (first vs. second row of Table 1). Therefore, we only construct a TTT layer equivalent to the simplest formulation of linear attention without the two additions. Then, to illustrate the contribution of each of our components (including some that will be introduced in the next subsection), we add them row by row to the TTT layer that is equivalent to linear attention, and ultimately obtain our proposed instantiation called TTT-Linear. The change from batch GD to mini-batch GD contributes the most improvement by a large margin.

While the space of models $\times$ optimizers in Figure 9 is already large, machine learning is much richer than optimizing the parameters $W_{t}$ of a model $f$. There are also nonparametric learners, such as nearest neighbors, support vector machines (SVMs), and kernel ridge regression. By definition, nonparametric learners do not have parameters $W_{t}$, and instead directly uses training data $x_{1},\dots,x_{t}$. Hence we use the notation $f(x;x_{1},\dots,x_{t})$. We now show that for a particular nonparametric learner, the induced TTT layer is equivalent to self-attention.

Appendix B contains a detailed explanation of the Nadaraya-Watson estimator and kernel $\kappa$ above. In contrast to Theorem 1, Theorem 2 does not produce a different implementation from attention.

For the TTT layer above, the hidden state is $x_{1},\dots,x_{t}$ or a similar list of processed training data, the update rule adds $x_{t}$ to the list, and the output rule scans the list with $\kappa$. In previous subsections, our hidden state has been defined as $W_{t}$, the update rule a gradient step, and the output rule a call to $f$. To unify these two constructions, we define a new abstraction called a learner, which uniquely induces a TTT layer.

Similar to its definition in standard machine learning packages [54], all learners need to implement two methods: train and predict. Now we redefine the hidden state of the induced TTT layer as the internal storage of the learner, and the update and output rules as the train and predict methods. Under this new definition of TTT layers, both parametric learners such as that in Theorem 1 and nonparametric learners such as that in Theorem 2 can be included. Figure 10 summarizes this general definition of TTT layers in the broader scope of all sequence modeling layers.

This general definition has an additional benefit for parametric learners: There can be more objects other than $W$ in the internal storage of parametric learners, such as the optimizer state, which will also be included in the hidden state of the induced TTT layer. This extension allows TTT layers to use more sophisticated optimizers such as Adam [42] in future work.

From Figure 11, we make a few observations:

• •

  At 2k context, TTT-Linear (M), Mamba, and Transformer have comparable performance, as the lines mostly overlap. TTT-MLP (M) performs slightly worse under large FLOP budgets. Even though TTT-MLP has better perplexity than TTT-Linear at every model size, the extra cost in FLOPs offsets the advantage.

• •

  At 8k context, both TTT-Linear (M) and TTT-MLP (M) perform significantly better than Mamba, in contrast to the observation at 2k. Even TTT-MLP (T) with the Transformer backbone performs slightly better than Mamba around 1.3B. A robust phenomenon we observe throughout this paper is that as context length grows longer, the advantage of TTT layers over Mamba widens.

• •

  At 8k context, Transformer still has good (if not the best) perplexity at every model size, but its line is not competitive because of the cost in FLOPs.

To evaluate capabilities in long context, we experiment with context lengths ranging from 1k to 32k in $2\times$ increments, using a popular subset of the Pile called Books3. The training recipe here is the same as for the Pile, and all experiments for the TTT layers are performed in one training run.¹²¹²12 Following the Mamba paper, we always use 0.5M tokens per training batch regardless of context length. That means for context length $T$ we have $0.5\text{M}\,/\,T$ sequences per batch (assume divisible). From the subset of results in Figure 12, we make a few observations:

• •

  At 2k context on Books, all the observations from Pile 2k still hold, except that Mamba now performs slightly better than TTT-Linear (whereas their lines roughly overlapped for Pile 2k).

• •

  At 32k context, both TTT-Linear (M) and TTT-MLP (M) perform better than Mamba, similar to the observation from Pile 8k. Even TTT-MLP (T) with the Transformer backbone performs slightly better than Mamba at 32k context.

• •

  TTT-MLP (T) is only slightly worse than TTT-MLP (M) at 1.3B scale. As discussed, it is hard to derive an empirical scaling law due to the lack of a clean linear fit. However, the strong trend for TTT-MLP (T) suggests that the Transformer backbone might be more suitable for larger models and longer context beyond our evaluations.

We only ablate the backbones for 2k and 32k due to the cost of training LLMs. For future work, we believe that given TTT layers with even more expressive hidden states, the Mamba backbone with temporal convolutions will become unnecessary.

LLM training and inference can be decomposed into forward, backward, and generate. Prompt processing during inference, also known as prefill, is the same operation as forward during training, except that the intermediate activations do not need to be stored for backward. Since both forward (during training and inference) and backward can be parallelized, we use the dual form. Generating new tokens, also known as decode, is inherently sequential, so we use the primal form.

Due to resource constraints, our experiments are written in JAX and run on TPUs. On a v5e-256 TPU pod, the Transformer baseline takes 0.30s per iteration of training at context 2k, while TTT-Linear takes 0.27s per iteration, already 10% faster without any systems optimization. However, Mamba (implemented in PyTorch, Triton, and CUDA) can only run on GPUs, so for fair comparison, we also rewrite our method with preliminary systems optimization to run on GPUs.

Specifically, we write a GPU kernel for forward in ThunderKittens [66]. Historically, RNNs have been inefficient during forward and backward due to poor use of parallelism and matmuls. Our goal with the forward kernel is to demonstrate the effectiveness of mini-batch TTT and the dual form for these problems. A kernel for backward should have the same properties in efficiency as forward, but requires more complex logic for manual differentiation, therefore is left for future work.

The left panel of Figure 15 shows the latency for batch size 16 of our forward kernel. All models are 1.3B (1.4B for Mamba). Time per token grows linearly for Transformer as context length increases, but stays roughly constant for the other methods.¹⁴¹⁴14 We observe that forward latency of the network increases slightly for TTT-Linear, TTT-MLP, and Mamba, even though latency of each sequence modeling layer alone stays constant. Consider the operation $\theta X$, where $\theta$ is $d\times d$ and $X$ is $d\times T$. Its latency (normalized over $T$) is expected to be constant, but in practice grows slightly with $T$. One possible cause of this phenomenon is the GPU throttling after $T$ gets very large [30]. Note that our Transformer baseline is significantly faster that in the Mamba paper, because we use vLLM [46], a state-of-the-art serving system, instead of the HuggingFace Transformer [77].

In addition, we write another GPU kernel for generate in Triton [72], and benchmark its speed in the right panel of Figure 15 for batch size 512. Another popular metric for wall-clock time is throughput, which takes into account the potential benefit of being able to use a larger batch size.¹⁵¹⁵15 To calculate throughput for each method, we increase its batch size in $2\times$ increments until GPU runs out of memory, measure the tokens per second for every batch size, and select the highest. For completeness, we report the throughput for forward and generate in Figure 20 (in Appendix). All the observations and ordering between methods above still hold for throughput.

Mamba is one of the many Structured State-Space Models [27, 21, 57, 18]. The hidden state in these models is a vector, similar to in LSTMs. For TTT-Linear or TTT-MLP, the hidden state is a matrix or two matrices, therefore larger. In Figure 14, we find that TTT layers can take advantage of their larger hidden states to compress more information in long context, where TTT-MLP outperforms TTT-Linear, which in turn outperforms Mamba.

Similar to TTT-Linear, RWKV [55, 56], xLSTM [5], and Gated Linear Attention (GLA) [79] also have matrix hidden states, which are inherited from linear attention [41]. Modern RNNs such as GLA use chunk-wise parallelism to improve hardware efficiency, so tokens inside a chunk can be processed with matmuls instead of a cumsum. However, chunk-wise parallelism does not change the expressiveness of the model, since all temporal dependencies are still equivalent to a cumsum.

In contrast, mini-batch TTT allows more complex temporal dependencies across mini-batches. Each hidden state $W_{t}$ depends on previous $W$s within its mini-batch still through a cumsum, but depends on $W$s in previous mini-batches also through the gradient operator. As illustrated Figure LABEL:fig:batch-size, mini-batch TTT enables a trade-off between expressiveness and hardware efficiency, since a smaller batch size $b$ leads to better perplexity at the cost of higher latency. This trade-off is a unique and important feature of TTT. As shown in Table 1, the intermediate batch size $b=16$ significantly outperforms $b=T$ which is fully cumsum.

The idea of learning at test time has a long history in machine learning. One of the earliest versions of this idea is called local learning (Bottou and Vapnik [10]): For each test input, train on its neighbors before making a prediction. This procedure has been effectively applied to models ranging from SVMs [81] to modern LLMs [29].

Another early version of learning at test time is called transductive learning [22]. The principle of transduction, as stated by Vladimir Vapnik [74], is to "… get the answer that you really need, but not a more general one." Practical implementations of transductive learning use test data to add constraints to the margin of SVMs [39, 17]. However, transductive learning usually needs multiple test instances to be empirically effective, unlike many instantiations of test-time training, which only need a test single instance (image, video, or natural language sequence) at a time.

In computer vision, the idea of learning at test time has been applied for decades to applications such as face detection [38], object detection [53], image super-resolution [65], and 3D reconstruction [50]. More recently, the same idea has also been applied to natural language processing, where it is called dynamic evaluation [44, 45]. The basic approach is to directly finetune a language model on the test sequence, which often comes in the form of a prompt.

Next, we discuss two relevant lines of work in detail: test-time training and fast weights.