Learning to (Learn at Test Time): RNNs with Expressive Hidden States
Authors: Yu Sun
Xinhao Li
Karan Dalal
Jiarui Xu
Arjun Vikram
Genghan Zhang
Yann Dubois
Xinlei Chen
Xiaolong Wang
Sanmi Koyejo
Tatsunori Hashimoto
Carlos Guestrin
Venue: ICML 2025 spotlightposter (2025)
Keywords: test-time training
Primary area: deep_learning->foundation_models

Abstract: 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.

Learning to (Learn at Test Time): RNNs with Expressive Hidden States
Yu Sun " Xinhao Li
Karan Dalal *3 Jiarui Xu²
Arjun Vikram
Genghan Zhang'
Yann Dubois'
Xinlei Chen
Xiaolong Wang Sanmi Koyejo"
Tatsunori Hashimoto
Carlos Guestrin
Abstract
we re-evaluate these findings in Figure
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 practical framework for instantiating sequence modeling layers with linear complexity and expressive hid- den states. The key idea is to make the hidden state machine learning model itself, and the up- date rule step of self-supervised learning. Since the hidden state is updated by training test sequences, our layers are called Test-Time Training (TTT) layers. We consider two instantia- tions: TTT-Linear and TTT-MLP, whose hidden state is a linear model and a two-layer MLP re- spectively. We evaluate our instantiations the scale of 125M to 1.3B parameters, comparing with a strong Transformer and Mamba, modern RNN. Similar to Transformer, TTT-Linear and TTT-MLP can keep reducing perplexity by con- ditioning 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 fu- ture research.
On the left, we observe that Mamba (Gu & Dao, 2023) of the most popular RNNs today scales similarly to a strong Transformer, showing great progress since the LSTMs in 2020. However, on the right, we observe the same issue with Mamha as Kaplan et al. did with LSTMMs. Tokens later in a sequence should be casier to predict on average, since they condition on more information. This indeed the case for Transformer, whose average perplexity at each token index decreases throughout its 32k context. In contrast, the same metric plateaus for Mamba after 16k.
on
This result represents an awkward reality for existing RNNs. On one hand, the main advantage of RNNs (vs. Trans- formers) is their linear (vs. quadratic) complexity. This asymptotic advantage only realized practice for long context, which according to Figure after 8k. On the other hand, once context is long enough, existing RNNs such as Mamba struggle to actually take advantage of the extra information being conditioned on.
The difficulty with long context is inherent to the very na- ture of RNN layers: Unlike self-attention, RNN layers have to compress context into hidden state of fixed size. As a compression heuristic, the update rule needs to discover the underlying structures and relationships among thousands or potentially millions of tokens. This need is inherently challenging. In this paper, we begin with the observation that self-supervised learning can compress a massive train- ing set into the weights of a model such as an LLM, which often exhibits deep understanding about the semantic con- nections among its training data exactly what need from compression heuristic.
1. Introduction
This version of the paper has been abridged fit the page limit of ICML camera ready. Please read our arXiv version instead: https: /tt///aaivvv www.riiwwww.raeeeooooooe abs / 2407. 04620. 620.
TTT layers. Motivated by this observation, we make the hidden state 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, these RNN layers are called Test-Time Training (TTT) layers. We introduce simple instantiations: TTT-Linear and TTT- MLP, where the hidden state is linear model and two- layer MLP, respectively. TTT layers can be integrated into any network architecture and optimized end-to-end, similar to RNNs layers and self-attention.
In 2020, the OpenAl scaling law paper (Kaplan et. al (Ka- plan et al., 2020)) showed that LSTMs (a type of RNN) could not scale similarly to Transformers or effectively use long context. Now, with modern RNNs and best practices,
Core contributors. t: Joint advising. 1 Stanford University 'UC San Diego *UC Berkeley "Meta AL.
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Learning to (Learn at Test Time); RNNs with Expressive Hidden States
Transformer
11.0
Transformer
Perplexity fleg scale!
Mamba
Perplexity (log scale)
10.5
Mamba
TTT-Linear
TTT-Linear
TTT-MLP
10.0
TTT MLP
9.0
8.5
6,0
2e+19
5e+19 FLOPS (log scale)
1e+20
2e+20
128
512
2k
16k
Token Index (log scale)
Figure 1. Comparing Mamba, TTT-Linear and TTT-MIPP have similar perplexity in 8k context (left) and better use of long context (right). Evaluations follow Kaplan et al. (Kaplan et al.. 2020). Left: Scaling trends on the Pile with 8k context, zoomed in between 350M and 1.3B parameters. Right: Similar to Transformer, TTT-Linear and TTT-MLP can keep reducing perplexity by conditioning on more tokens, while Mamba cannot after 16k context. All methods have matched training FLOPs as Mamba 1.4B.
Wall-clock time. We apply two techniques to make TTT layers more efficient on modern GPUs and TPUs. First, similar to the standard practice of taking gradient steps on mini-batches of sequences during regular training for mini-batches of tokens during
Self-attention can also be viewed from the perspective above, except that its hidden state, commonly known as the Key- list that grows linearly with Its update rule simply appends the current KV tuple to this list, and the output rule scans over all tuples up to & to form the attention matrix. The hidden state explicitly stores all historic context without compression, making self-attention more expressive than RNN layers for long context. How- ever, scanning this linearly growing hidden state also takes linearly growing time per token.
Value (KV) cache,
better parallelism, TTT. Second, we develop dual form for operations inside each TTT mini-batch. The dual form is equivalent in output to the naive implementation, but trains more than 5x faster on our TPUs.
Contributions and limitations. The idea of using linear models hidden states has already been well studied in DeltaNet (Schlag al., 2021; Yang et al., 2024). Since our first version was released, RNN layers with matrix (linear) hidden states have also been further advanced in Mamba 2 (Dao & Gu, 2024) and Gated DeltaNet (Yang et al., 2023). Compared to this line of work, our contribution is a practi- cal framework that can instantiate arbitrary neural networks as hidden states. However, such instantiations can still re- quire substantial wall-clock time, even after applying our improvements in efficiency. It remains to be seen whether our framework can produce instantiations that either over- come this limitation or offer benefits outweighing it.
To remain both efficient and expressive in long context, need better compression heuristic. Specifically, need to compress thousands potentially millions of tokens into hidden state that can effectively capture their underlying structures and relationships
2.1. TTT as updating hidden state
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 (Le, 2013) - exactly what need from compression heuristic.
2. Method
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 deep understanding of the seman- tic connections among existing knowledge express new pieces of reasoning (Achiam et al., 2023).
All sequence modeling layers can be viewed from the per- spective of storing historic context into a hidden state, shown in Figure 2. For example, RNN layers such as LSTM (Hochreiter & Schmidhuber, 1997) and Mamba (Gu & Dao, 2023) layers compress context into a state of fixed size across time. This compression has two consequences. On one hand, mapping an input token 2, output token 2; efficient, because both the update rule and output rule take constant time per token. On the other hand, the performance of RNN layers in long context is limited by the expressive power of its hidden state
Our key idea is to use self-supervised learning to compress the historic context ø1, 2y into a hidden state 8t, by making the context an unlabeled dataset and the state a model. Concretely, the hidden state 8; is now equivalent to
2

---

Learning to (Learn at Test Time); RNNs with Expressive Hidden States
Output tokens
RL-1 Ay-1
#: output st))
Output rule
Hidden state
8g
8; update (1-1 (pdate(#
Update rulc
Imput tokens
Initial state
Update rule
Output rule
Cost
Naive RNN
8g vector O
St (0,581-1 22x
2 0258t Ozwlt
O(1)
Self-attention
Sg 1ist ()
St sp-..append((k, ut)
2 Visoftmax (KF Q:)
O(t)
Naive TTT
Wo = f.params ()
W = Wi-1 - (VV((W(11))))
2 = (()
0(1)
A generic sequence modeling layer expressed as hidden state that transitions according to an update rule. All sequence modeling layers can be viewed as different instantiations of three components in this figure: the initial state, update rulc and output rule. Bottom: Examples of sequence modeling layers and their instantiations of the three components. Self-attention has hidden state growing with context, therefore growing cost per token. Both the naive RNN and TTT layer compress the growing context into a hidden state fixed size, therefore their cost per token stays constant.
Figure 2. Top:
Wr the weights of a model f. which can be a linear model, a small neural network, or anything else. The output rule is simply: 2, = ((x; W.). Intuitively, the output token just the prediction on 24, made by f with the updated weights We The update rule is step of gradient descent some superviid loss E:
2.2. Training a network with TTT layers
The forward pass of a TTT layer also has a correspond- ing backward pass. Our forward pass only consists of stan- dard differentiable operators except the gradient operator V. However, V just maps one function to another, in this case to Ve, and VI is also composed of differentiable operators. Conceptually, calling backward on Ve means taking gradients of gradients a well explored technique in meta-learning (Maclaurin et al., 2015).
W. = W V((W.-1;25)
with learning rate 17." 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.
TTT layers have the same interface as RNN layers and self- attention, therefore can be replaced any larger network architecture, which usually contains many of these sequence modeling layers. Training network with TTT layers also works the same way as training any other language model, such as Transformer. The same data, recipe, and objec- tive such as next-token prediction can be used to optimize parameters of the rest of the network.
One choice of f is reconstructing x, itself. To make the learning problem nontrivial, we first process 2, into cor- rupted input 2, (details in Subsection 2.3), then optimize:
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 V 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 0rest- Throughout this paper, outer-loop parameters are always denoted by with various subscripts.
E(W: 2 = [f W)- - -rel
(2)
Similar to denoising autoencoders (Vincent et al., 2008), f needs to discover the correlations between dimensions of Ze in order to reconstruct it from partial information dt. 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 21..... Zy to output sequence . 2T 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 WI...., Wr for every input sequence. Therefore, we call it the Test-Time Training (TTT) layer.
2.3. Learning a self-supervised task for TTT
Arguably the most important part of TTT is the self- supervised task, because 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's is for 2, (31,; W.) to perform well on language modeling. Instead of handcraft- ing a self-superviser 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.
For now, consider Wo 0. We will discuss more sophisti- cated techniques for initializing W in Subsection 2.7.
3

---

Learning to (Learn at Test Time); RNNs with Expressive Hidden States
Concretely, we learn the self-supervised task part of the outer loop. Starting from the naive reconstruction task in Equation 2, we add some outer-loop parameters to make this task learnable. In Subsection 2.1, we did not specify the corruption that produces , from Zt- One design is to make a low-rank projection ã = 2g where 0x is a learnable matrix. 2 Following the terminology of multi-view reconstruction, 0g is called training view.
W
Wh
Figure 3. High-level computation graph of the first TTT mini- batch, where nodes are variables and edges are computations. The blue nodes are input variables, and yellow are output.
Moreover, perhaps not all the information in ze is worth remembering, so the reconstruction label can be another low-rank projection zy instead of z4. Here Iy is called the label view, where By is also learnable. In summary, our new self--uue loss is:
Peop begp sogi
200
Time Cme
100
11
&(W; 24) || (xx2; - 2
(3)
misbbudd dic log
TTT mi sotchatdhsee & morl)
Since both W and various As appear together in Equation 3, we emphasize again their difference in nature. the inner loop, only Wi is optimized, therefore written as an argument of é; the Os are "hyper-parrmeters" of this loss function. In the outer loop, 0x, fy, Qq are optimized alongside frest, and W is merely a hidden state, not parameter.
Figure 4. Ablations TTT mini-batch size b, where 6 1 is online GD and = T is batch GD. We choose b = 16 for all experiments in this paper. Left: Smaller b improves perplexity since more GD steps are taken. The perplexity of 11.09 at b 16 corresponds to the final result of TTT-Linear Figure 6. Right: Forward time in dual form, with context length T = 2048. Total time (orange) can be decomposed into time for computing the Ws at the end of every mini-batch (blue) and time for 21, ,.,.
Lastly, the training view 2g has fewer dimensions than Zt. so we can no longer use the output rule in Subsection 2.1.
The simplest solution is to create a test view OgIt, and change our output rule to:
where G2 is the descent direction. Note that once we have calculated G, for 1...., T, we can then obtain all the Wis through a cums um by the second half of Equation 5. Our naive update rule, known online gradient descent, uses G. = (i(
(4) This solution has an additional benefit. The training and label views specify the information in I; that is compressed into W, and propagated forward through time. The test view specifies potentially different information that is mapped to the current output token 2 and propagated forward through network layers, therefore adds more flexibility to the self- supervised task.
2 f (fQ=;; W.).
To parallelize Ge for t = I,...,T, we can take all of them w.r.t. Wg. This variant with G. = Ve (Wo; v.2) is known as batch gradient descent, since > V (Wo; xx) is the same as the gradient w.r.t. Wo over . .., De as batch. However, in batch GD, W, effectively only one gradient step away from Wo, in contrast to online GD, where W, is steps away from Wo. Therefore, batch GD has smaller effective search space, which ends up hurting performance for language modeling.
2.4. Parallelization with mini-batch TTT
The naive TTT layer developed so far is already efficient in the number of floating point operations (FLOPs). However, its update rule W, = W.- - (;; cannot be parallelized, because W. depends on W in two places: before the minus sign and inside VL. Since V/ contains the bulk of the computation, we focus on making this second part parallel.
Our proposed solution - mini-batch gradient descent - shown in Figure 3. Denote the TTT batch size by b. We use G2 = V (Wr; where 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 4. We chose b = 16 for all experiments in this paper.
We approach this systems challenge through concepts in the TTT framework. There many variants of gradient descent (GD). Its general update rule can be expressed as:
W = Wi - G = Wo
Σ G.
(5)
2.5. Dual form
s=1
The parallelization introduced above is necessary but not sufficient for efficiency in wall-clock time. Modern acceler- ators specialize in matrix-matris multiplications, known as
The subscript K hints connection self-attention, as will establish Subsection 2.6.
4

---

Learning to (Learn at Test Time); RNNS with Expressive Hidden States
matmmls Is. For example, the NVIDIA A100 GPU contains highly optimized units called TensorCores that can only perform single operation multiplying two matrices each of size 16 x 16. Without enough of these matmuls, the TensorCores are idle, and most of the potential for the A 100 is unrealized.
is O(b dP) for computing W, alone, then an additional O(b2xd for computing 21;-..; 26- Compared to the primal, the dual form sacrifices theoretical complexity for hardware utilization. In practice, d is typically few hundred and b is chosen to be only 16. As consequence, wall-clock time for computing 21,.... 2n is relatively small, as observed in the right panel Figure 4. In our JAX implementation. training with the dual form more than 5x faster than with primal.
Unfortunately, the TTT layer developed so far even with mini-batch still has very few matmuls. Consider the sim- plest case of l, where 0x = By Aq = I, for only the first TTT mini-batch of size b. In addition, consider f as linear model. Copying Equation 2, our loss at time t is: & (Wo; = 1 - 211.². As discussed in Subsection 2.4, we can parallelize the computation of: G, = 2(Wo - 2y) )a?" for 1 = 1,...,b. b. However, we cannot compute all 6 of the G5s through a single matmul. Instead, need b outer products to compute them one by one. To make matters worse, for each de E Rd, Ge is d x d, which incurs much heavier memory footprint and I/O cost than Te for large d.
2.6. Theoretical equivalences
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 x 3 combinations induces a different instantiation of the TTT layer. We now show that among these induced instantiations, the TTT layer with a linear model and batch GD is equivalent to linear attention (Katharopoulos et al., 2020).
Theorem 1. Consider the TTT layer with [(x) = Wx as the inner-loop model, batch gradient descent with 1/2 as the update rule, and Wo = O. Then, given the same imput sequence T I1....,21 DT, the output rule defined in Equation 4 produces the same output sequence 21,.....3 2T as linear attention.
To solve these two problems, we make a simple observa- tion: We do not actually need to materialize Gi,..., Gb as long as we can compute W, at the end of the mini- batch, and the output tokens 21, : 26 (see Figure 3). Now we demonstrate these computations with the simplified then:
TTT-Linear case above. Denote X = [si,..., So W, can be conveniently
W, = Wn- - (W, computed with a matmul. To compute Z- = [21....,. 20].
Proof. By definition of & in Equation 3, Vi (Wo: = 2(0x))) By definition of batch GD: W; (0y2,)
we know that:
Plugging W; into the output rule in Equation we obtain the output token: 2t
24 = far = Woxz - 2n (Woz= - 2.)xfar. (6)
f (Bqzz; WT) E.=1(0 (0xx0)))) which is the definition of linear attention.
Denote 81 = -' (Woz, - xx x2)x²" ae and the matrix A
86). We can derive that:
In Table 1, we first empirically verify the equivalence above with an improved implementation of linear attention. Then, to illustrate the contribution of each of our components (in- cluding 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 in- stantiation called TTT-Linear. The change from batch GD to mini-batch GD contributes the most improvement by a large margin.
A = (WoX- mask (x x),
(7)
where mask is the upper triangular mask with zeros (similar to the attention mask, but with zeros instead of infinities), and the term WoX- - X can be reused from the computa- tion of Wg. Now A is also conveniently computed with ma tmu 1s. Plugging A back into Equation 6, we obtain Z = WoX- - 21/A.
We call this procedure the dual form. in contrast to the form form before this subsection, where the Gs and Ws 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 (Irie et al., 2022; Bishop & Nasrabadi, 2006; Rosenblatt, 1958). In Appendix A, we show that the dual form still works when f is a neural network with nonlinear layers.
While the space of models x optimizers is already large, machine learning is much richer than optimizing the parame- ters W, 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,, and instead directly uses training data Ii...., 2t. Hence we use the notation [(;; (1.;..,...40 R2). We now show that for particular non- parametric learner, the induced TTT layer is equivalent to self-attention.
Time complexity of the primal form within a TTT mini- batch is O(b x d Time complexity of the dual form
Theorem 2. Consider the TTT layer with the Nadaraya-
5

---

Learning to (Learn at Test Time); RNNs with Expressive Hidden States
Sequence modeling layers
Configuration
Ppl.
Diff.
Linear attention (Katharopoulos et al., 2020)
15.91
RNN lavers
FIT layers
Linear atin. improved
15.23
0.68
Mambe
parametric
w/ nonparametric
TTT equivalence
15.23
0
RWKV
learnens
learners
learnable Wo
15.27
+0.04
LSTM
Detail Pigure
Sel-attentio
LN and residual in
14.05
- -1.22
mini-batch TTT
12.35
- 1170
learnable
11.99
-0.36
Figure 5. RNN layers and TTT layers are both subsets of sequence modeling layers. RNN layers have hidden state that is fixed in size across time. TTT layers with parametric learn- ers also RNN layers, since their hidden state is also fixed in size. TTT layers with nonparametric learners can represent self- attention, discussed in Subsection 2.6.
Mamba backbone
11.09
-0.90
Table 1. Ablations on improving from linear attention. All models here have 125M parameters, and are trained according to the recipe in Subsection 3.1. The last row, with perplexity 11.09, is the final result of TTT-Linear in Figure 6. Starting from the equivalence discussed in Subsection 2.6, learnable Wo hurts slightly, but the rows below cannot train stably without it. The biggest improve- ment comes from mini-batch TTT (changing from b = = 2048 to 16). The second comes from instantiating the imner model f with LN and residual connection.
Watson estimator (Bierens, 1988; Cai, 2001), defined as:
(1; E.mn 6(2, v.))
Σ x(x,,,)) x.) Mar (8)
where yo = I2. and
adds negligible amount of parameters comparing to the reconstruction views 0x, Qq. Bv, because both its input and output are low dimensional. Empirically, we observe that learning Wo significantly improves training stability.
1xxxx';0K: (0:;0K (2) x (9xx)) eqx
(9)
is kernel with bandwidth hyper-parameters 0x and lq- Then given the same input sequence , Ur, the out- put rule defined in Equation 4 produces the same outpul 2T as self-attention.
Learnable 77. The learning rate usually the most important hyper-parameter for gradient descent, with learning the inner-loop learning rate 1/ in Equation 5 as part of the outer loop. We make rj function of the input token (therefore different across time) for additional flexibility. Concretely, we design 77x)) = Thase ( I), where the learnable vector Orr is an outer-loop parameter, is the sigmoid function, and the scalar Thase the base learning rate, set to I for TTT-Linear and 0.1 for TTT-MLP. Alternatively, 77(z) can also be interpreted as gate for VŁ.
experiment
sequence ,T
Proof. Plugging V/s and above into Equation 8 gives us the definition of self-attention.
Appendix B contains a detailed explanation of the Nadaraya- Watson estimator and kernel x above. In contrast to Theo- rem Theorem does not produce a different implementa- tion from attention.
Backbone architecture. The cleanest way to integrate any RNN layer into a larger architecture would be to di- rectly replace self-attention in Transformer, known in this context as a backbone. However, existing RNNs such as Mamba (Gu & Dao, 2023) and Griffin (De al., 2024) all use different backbone from Transformers. Most tably, their backbone contains temporal convolutions before the RNN layers, which might help collect local information across time. After experimenting with the Mamba backhone, we find that it also improves perplexity for TTT layers, so we incorporate it into our proposed method. See Figure (in Appendix) for details.
2.7. Implementation details
Instantiations of f. We propose two variants TTT lay- ers TTT-Linear and TTT-MLP, differing only their instantiations of f. For TTT-Linear, / (x)) Wa, where W is square. For TTT-MLP, Jue has two layers similar to the MLPs in Transformers. Specifically, the hidden di- mension is 4x the input dimension, followed by GELU activation (Hendrycks Gimpel, 2016). For better stability during TTT, f always contains a Layer Normalization (LN) and residual connection. That is, [(x) = x + LN(J rea (x)),
where f fees can be fiin or fuLE-
Learnable Wo. The TTT initialization Wo is shared between all sequences, even though subsequent weights WI.... Wr are different for each input sequence. Instead of setting Wn = O, we can learn as part of the outer loop. Since outer-loop parameters are always denoted by fs in- stead of Ws, we assign an alias fm = Wo. In practice, Omm
3. Experiments
We evaluate TTT-Linear and TTT-MLP by comparing with two baselines - Transformer and Mamba, a modern RNN. Our main codebase is based on EasyLM (Geng, 2023), an open-source project for training and serving LLMS JAX.
6

---

Learning to (Learn at Test Time): RNNs with Expressive Hidden States
Sealin rerde on Mo
Sestinemenoe in Pir
Scvinn lig treeds Sower
Scxling swnde on Boses S0k
Tarshere
Thansbreet
Tessberea Mat
Profeee Nacka
MI
24
ETS HDR (N)
ITAAEE (T)
12 24
12
12
Figure 6. Evaluations for context lengths 2k and 8k on the Pile. De- tails in Subsection 3.1. TTT-Linear has comparable performance as Mamha at 2k context, and better performance at 8k.
Figure 7. Evaluations for context lengths 2k and 32k on Books. Details in Subsection 3.2. Our complete results for context lengths 1k, 2k, 4k, 8k, 16k, 32k, including Transformer finetuning, are in Figure 11 (in Appendix).
Datasets. Following the Mamba paper (Gu & Dao, 2023), we perform standard experiments with 2k and 8k context lengths on the Pile (Gao et al., 2020), a popular dataset of documents for training open-source LLMs (Black al., 2022). However, the Pile contains few sequences of length greater than 8k (de Vries, 2023). To evaluate capabilities in long context, we also experiment with context lengths ranging from 1k to 32k in 2x increments, on subset of the Pile called Books3, which has been widely used to train LLMs in long context (Liu et al., 2024).
former backbone performs slightly better than Mamba around 1.3B. A robust phenomenon we observe through- out 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 competi- tive because of the cost in FLOPs.
Effect of backbone. Switching the TTT layers from Mamba backbone into Transformer backbone has two effects. First, TTT layers with Mamba backbone perform better in our evaluations so far. Second, with Mamba backbone, TTT- MLP at best is only comparable to TTT-Linear; but with Transformer backbone, TTT-MLP is clearly better. We hy- pothesize that the temporal convolutions in the Mamba back- bone help more when the sequence modeling layer has a less expressive hidden state. The linear model is less expressive than the MLP, therefore benefits more from the convolutions. We will revisit this hypothesis in the next subsection.
Backbone architecture. As discussed in Subsection 2.7, Transformer and Mamba use different backbones, and TTT- Linear and TTT-MLP always use the Mamba backbone unless noted otherwise. As an ablation study, Figure 6 and Figure 7 contain TTT layers within the Transformer backbone. When a figure contains both the Transformer backhone and Mamba backbone, denote them by (T) and (M), respectively.
Protocols. To ensure fairness to our baselines, we strietly follow the evaluation protocols in the Mamba paper when possible. For each evaluation setting (e.g., dataset, context length, and method), we experiment with four model sizes: 125M, 350M 760M, and 1.3B parameters. For Mamba, the corresponding sizes are 130M, 370M, 790M, and 1.4B, Mamba does not follow the Transformer configurations. All models are trained with the Chinchilla recipe described in the Mamba paper and reproduced in our Appendix C.
3.2. Long context: Books
To evaluate capabilities in long context, we experiment with context lengths ranging from 1k to 32k in 2x increments, using popular subset of the Pile called Books3. The train- ing recipe here is the same as that for Pile. From the subset of results in Figure 7, 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).
3.1. Short context: the Pile
From Figure 6, 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 FLOPSs offsets the advantage.
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.
At 8k context, both TTT-Linear (M) and TTT-MLP (M) perform significantly better than Mamba, in contrast to the observation 2k, Even TTT-MLP (T) with the Trans-
TTT-MLP (T) is only slightly worse than TTT-MLP (M) at 1.3B scale. As discussed, it is hard to derive an em- pirical scaling law due to the lack of a clean linear fit.
7

---

Learning to (Learn at Test Time); RNNS with Expressive Hidden States
Frewart (port linetry zatre diw
Generass (tecodes istrney for hair? dre 512
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. How- ever, Mamba (implemented in PyTorch, Triton, and CUDA) can only run on GPUs, so for fair comparison, we also rewrite our method into GPU kernels. We only write infer- ence kernels for this work because the training kernel would require substantial effort and cannot be used on our TPUs.
Notha
15
Tib
Cortuctary
Figure& 8. Latency an NVIDIA A10 GPU with 80G HBM and PCIe connections.
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.
Figure 8 shows the latency of our inference kernel for for- ward (prefill) and generate (decode). All models are 1.3B (1.4B for Mamba). As expected, time per token grows lin- early for Transformer as the context length increases, but stays roughly constant for the other methods. Note that our Transformer baseline is significantly faster that in the Mamba paper, because we use vL..M (Kwon et al., 2023), a tfateoof-the-ar- serving system, instead of the HuggingFace Transformer (Wolf et al., 2019).
We only ablate the backbones for 2k and 32k due the cost of training LLMs. For future work, we believe that given TTT layers with even more expressive hidden states, the Mamba backbone with convolutions will be unnecessary.
Transformer finetuning, While we have been training Transformers from scratch following the Mamba paper, in practice this approach is rarely used for long context. The standard practice is train a Transformer in short context, then finetune in long context. To reflect this practice, we add another baseline, TF finetune, for context lengths 4k and above. This baseline starts from the model trained (accord- ing to the Chinchilla recipe) on Books 2k, then uses 20% more tokens to finetune at the designated context length, following the Llama Long paper (Xiong et al., 2023). See details of the TF finetune recipe in Appendix C.
4. Related Work
4.1. Learning at Test Time
The idea of learning at test time has a long history in ma- chine learning. One of the earliest versions of this idea is called local learning (Bottou and Vapnik (Bottou & Vapnik, 1992)): For each test input, train on its neighbors before making a prediction. This procedure has been effectively applied to models ranging from SVMs (Zhang et al., 2006) to modern LLMs (Hardt & Sun, 2023). Next, we discuss two relevant lines of work in detail: test-time training and fast weights.
Experiments in Figure 1 (right). Compared to TTT-Linear, TTT-MLP with matched FLOPs performs worse at short context but better at long context. This observation matches our expectation that the MLP as hidden state is more ex- pressive than the linear model: The larger capacity of a more expressive hidden state is well-utilized in long con- text (therefore an advantage), but redundant in short con- text (therefore disadvantage in our setting with matched FL.OPs.. The Transformer in this figure is TF finetune, which is the stronger baseline in 32k context. Details of the experiments in Figure 1 are included in Appendix C. Our complete results for context lengths 1k, 2k, 4k, 8k, 16k, 32k, including TF finetune, are in Figure 11 (in Appendix).
4.1.1. TEST-TIME TRAINING
The core idea of Test-Time Training (TTT) is that each test instance defines its own learning problem, where this test in- stance alone is the target of generalization (Sun et al., 2020). Concretely, for each test instance ær, the conventional prac- tice is predict fxx), using a predictor that is optimized for all training instances on average. TTT first formulates learning problem defined by æ. then trains model fa on (often with f as initialization). and predicts fx)..
3.3. Wall-clock time
Since the test instance comes without its label, the learning problem can only be formulated with a self-super task. Prior work has shown that TTT with reconstruction signif- icantly improves performance especially on outliers (Gan- delsman et al., 2022). Improvements become even more pronounced when testing on video frames that arrive in a stream and TTT is autoregressive (Wang et al., 2023), as ft is trained on past frames æ ; .t. The autoregressive connection makes (Wang et al., 2023) most relevant to our paper. Conceptually, the biggest difference between our pa-
LLM training and inference can be decomposed into for- ward, backward, and generate. Prompt processing during inference, also known as prefill, is the same operation as forward during training, except that the intermediate acti- vations 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 to- kens, also known as decode, is inherently sequential, so we use the primal form.
8

---

Learning to (Learn at Test Time); RNNs with Expressive Hidden States
per and prior work is that our reconstruction task is learned in outer loop, instead of handcrafted with human priors.
improvements in efficiency. For example, TTT-MLP is effective in terms of FLOPs, as shown in Figure 1. But the additional complexity of the MLP structure increases wall-clock time much more relative to FLOPs, as shown in Figure& 8. It remains to be seen whether our framework can produce instantiations that either overcome this limitation or offer benefits outweighing it.
4.1.2. FAST WEIGHTS
The general idea of fast weights is to update the parameters of a "fast" model on only the most relevant data, as opposed to the conventional practice of updating a "slow" model on all data (Tieleman & Hinton, 2009). This idea has existed since the 1980s (Hinton & Plaut, 1987). The most relevant data can be the test instance itself, therefore TTT can be viewed as special case of fast weights. Compared to fast weights, TTT embraces the idea of formulating an explicit learning problem, where the test instance is the target of generalization. Our update rule is also an explicit step of optimization.
4.3. Learning to Learn
For decades, researchers have been arguing that learning to learn, also known as meta-learning or bi-level optimization, should be a critical component of intelligence (Schmidhuber, 1987; Bengio et al., 1990; Thrun & Pratt, 1998; Lake et al., 2017). In prior work such as (Andrychowicz et al., 2016), (Finn et al., 2017) 3017) and (Metz et al., 2018), the inner loop learns from an entire dataset at a time instead of a sequence, so the outer loop needs a collection of datasets or tasks. In short, the outer loop is "one level above" regular training. Since is hard collect millions of datasets, this outer loop is hard to scale.
The idea of fast weight programmers (FWPs) is update the fast weights with a "slow" model (Schmidhuber, 1992). As modern example for language modeling, Clark et al. (Clark et al., 2022) give Transformer a final layer of fast weights, whose initialization is trained as slow weights. Our inner- loop weights W can be viewed as "fast" and outer-loop weights A as "slow". Therefore, networks containing TTT layers can be viewed as a special case of FWPs (Kirsch & Schmidhuber, 2021), similar to how TTT can be viewed as special case of fast weights.
In contrast, for TTT, each sequence itself is a dataset and defines its own generalization problem. The inner loop is "one level below" regular training, so our outer loop is only another solution to the canonical problem of supervised learning, instead of a new problem setting like generaliza- tion across datasets.
Modern RNN layers such as linear attention (Katharopoulos et al., 2020; Schlag et al., 2020) and DeltaNet (Schlag et al., 2021; Yang et al., 2024) are inspired by the idea of FWPs. Given their relevance to our work, we diseuss these modern RNN layers in detail in the next subsection.
5. Future work
The search space for effective instantiations inside this framework is huge, and our paper has only taken a baby step. Fortunately, if our perspective holds, then heuristics from regular training can transfer to test-time training. and search can be efficient. Next we outline some especially promising directions for future work:
4.2. Modern RNN layers
Our baseline, Manba (Gu & Dao, 2023), is only one of the many recent RNN layers that inherit the linear (ma-
trix) hidden states of linear attention (Katharopoulos al.,
2020; Schlag et al., 2020). Some more recent examples are RWKV (Peng et al., 2024), xLSTM (Beck et al., 2024), and Gated Linear Attention (GLA) (Yang et al., 2023). The most relevant work is DeltaNet (Schlag et al., 2021), which equivalent to TTT-Linear with inner-loop mini- size without the Layer Norm and residual connection. (Yang et al., 2024) further improve the performance of DeltaNet and enable parallelized updates across tokens (in our terms, across inner loop mini-batched Since our first version was released, RNN layers with matrix (linear) hidden states have also been further advanced in Mamba 2 (Dao & Gu, 2024) and Gated DeltaNet (Yang et al., 2023).
Systems optimization. Our systems optimization in Sub- section 3.3 has been preliminary at best, and there are many ways to improve it. In addition, pipeline parallelism through time might allow us process long sequences of millions of tokens on multiple devices together.
Longer context and larger models. Constrained by our academic resources, we have not trained with millions or billions in context length, which would also require larger models according to Figure 12. The advantage of TTT layers should become more pronounced in longer context.
More ambitious instantiations of f. When context length becomes longer, f would also need to be larger. For video tasks and embodied agents, whose context length can easily scale up to millions or billions, f could be a convolutional neural network.
Compared this line of work, our contribution is a practical framework that can instantiate arbitrary neural networks as hidden states. However, such instantiations can still require substantial wall-clock time, even after applying our

---

Learning to (Learn at Test Time); RNNs with Expressive Hidden States
Impact statement
Bengio, Y., Bengio, S., and Cloutier, J. Learning synaptic learning rule. Citeseer, 1990.
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
Bierens, H. J. The nadaraya-watson kemel regression func- tion estimator. (Serie Research Memoranda; No. 1988- 58). Faculty of Economics and Business Administration. Vrije Universiteil Amsterdam., 1988.
Acknowledgements
Bishop, C. M. and Nasrabadi, N. M. Pattern recognition and machine learning, volume Springer, 2006.
Part of the compute for this project is generously supported by the Google TPU Research Cloud program and Hyper- bolic Labs. XW is supported, in part, by the Amazon Re- search Award, the Cisco Faculty Award and the Qualcomm Innovation Fellowship. SK acknowledges support by NSF 2046795 and 2205329, NIFA award 2020-67021- 111-670211327999 1332999, the Alfred P. Sloan Foundation, and Google Inc. TH is sup- ported by& a Sony Faculty Innovation Award and a gift from Panasonic. CG acknowledges support by the Air Force OF- fice of Scientific Research (AFOSR), FA9550-20-1-008 -0-1-77 Stanford Human-Centered Artificial Intelligence (HAT) In- stitute, and gifts from Google and IBM.
Black, S., Biderman, S., Hallahan, E., Anthony, Q., Gao, L., Golding, L., He, H., Leahy, C., McDonell, K., Phang, J., et al. Gpt-neox-205 20b:: An open-source autoregressive language model. arXiv preprint arXiv:2204. 06745, 2022.
Bottou, L. and Vapnik, V. Local learning algorithms. Neural computation, 4(6): 1992.
Breiman, L., Meisel, W., and Purcell, E. Variable kemnel estimates of multivariate densities. Technometrics, 19(2): 135-144, 1977.
We would like to thank Rohan Taori, Xuechen Li, Allan Zhou, Ke Chen, and Guandao Yang for many helpful discus- sions, Menghao Guo for help with code release, Xinyang Geng for help with EasyLM, Hao Liu for help with the LWM codebase, David Hall for help with Levanter, Yossi Gandelsman and Yutong Bai for help at early stage of the project, Mert Yuksekgonul for help with figures in the paper, Horace He, Ben Spector, and Azalia Mirhoseini for help with systems, Sharad Vikram and Roy Frostig for an- swering our questions about JAX and Pallas, Albert Gu and Tri Dao for helping us reproduce experiments in the Mamba paper, and Kilian Weinberger and Percy Liang for advice on presentation. Yu Sun is grateful his PhD advisors, Alexei A. Efros and Moritz Hardt, for their many insights from years ago that eventually became part of this paper.
Cai, Weighted nadaraya- watson regression estimation. Statistics & probability letters, 51(3):307-318. 2001.
Chen, T., Xu, B., Zhang, C., and Guestrin, C. Training deep nets with sublinear memory cost, 2016.
Chen, Y.-. A tutorial kernel density estimation and recent advances. Biostatistics & Epidemiology. 1(1):161- 187, 2017.
Clark, K., Guu, K., Chang, M.-W., Pasupat, P., Hinton, G., and Norouzi, M. Meta-learning fast weight language models. arXiv preprini 1rXiv:2212.02475. 202.
Dao, T. and Gu, A. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality. arXiv preprint arXiv:2405.2 1060, 2024.
References
De, S., Smith, S. L., Fernando, A., Botev, A., Cristian- Muraru, G., Gu, A., Haroun, R., Berrada, L., Chen, Y., Srinivasan, S.,, et al. Griffin: Mixing gated linear recur- rences with local attention for efficient language models. arXiv preprint arXiv:2402.
Achiam, J., Adler, S., Agarwal, S., Ahmad, I.., Akkaya, I., Aleman, F. L., Almeida, D., Altenschmidt, J., Altman, S., Anadkat, S., et al. Gpt-4 technical report. arXiv preprint arXiv:23333088774, 2023.
2024.
de Vries, H. URL https: https://www. httppp:/wwwwwwh.ndvi..i... www.dhvriesee...eewwwww.wwwae..m com/ /ptpt//s///ttt context- length// Accessed: 2024-06-24.
the long (context) run, 2023.
Andrychowicz, M., Denil, M., Gomez, S., Hoffman, M. W., Pfau, D., Schaul, T., Shillingford, B., and De Freitas, N. Learning to learn by gradient descent by gradient descent. Advances in neural information processing systems, 29, 2016.
Finn, C., Abbeel, P., and Levine, S. Model-agnostic meta- learning for fast adaptation of deep networks. In Interna- tional conference on machine learning, pp. 1135. PMLR, 2017.
Beck, M.. Pöppel, K., Spanring, M., Auer, A., Prudnikova, O., Kopp, M., Klambauer, G., Brandstetter, J., and Hochreiter, S. xlstm: Extended long short-term mem- ory. arXiv preprint arXiv:2405. 045/7, 2024.
Gandelsman, Y., Sun, Y., Chen, X., and Efros, A. A. Test- time training with masked autoencoders. Advances in Neural Information Processing Systems, 2022.
10

---

Learning to (Learn at Test Time); RNNS with Expressive Hidden States
Gao, L., Biderman, S., Black, S., Golding, L., Hoppe, T., Foster, C., Phang. J., He, H., Thite, A., Nabeshima, N., Presser, S., and Leahy, C. The pile: An 800gb dataset of diverse text for language modeling, 2020.
Le, Q. V. Building high-level features using large scale unsupervised learning. In 2013 IEEE international con- ference acoustics, speech and signal processing, pp.
IEEE, 2013.
Geng, X. EasyLM: A Simple And Scalable Training Frame- work for Large Language Models. ht https: /github.. com/young-gae geng/ EasyLM, mar 2023, https:// github.
Liu, H., Yan, W. Zaharia, M.. and Abbeel, P. World model on million-length video and language with blockwise ringattention. arXiv preprint arXiv:2402.08266, 2024.
Maclaurin, D., Duvenaud, D., and Adams, R. Gradient- based hyperparameter optimization through reversible learning. In International conference on machine learn- ing. pp- 2113-2122. PMLR, 2015.
Gu, A. and Dao, T. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752 2023.
Hardt, M. and Sun, Y. Test-time training on nearest neighbors for large language models. arXiv preprint arXiv:2305.18-465 2023.
Metz, L., Maheswaranathan, N., Cheung, B., and Sohl- Dickstein, J. Meta-learning update rules for unsupervised representation learning. arXiv preprint urXiv: 000222, 2018.
Hendrycks, D. and Gimpel, K. Gaussian error linear units (gelus). arXiv preprint arXiv:1606.088 1606.08415 2016.
Peng, B., Goldstein, D., Anthony, Q. Albalak, A., Al- caide, E., Biderman, S., Cheah, E., Ferdinan, T., Hou, H., Kazienko, P., al. Eagle and finch: Rwkv with matrix- valued states and dynamic recurrence. arXiv preprint arXivv2404.05882. 2024.
Hinton, G. E. and Plaut, D. C. Using fast weights to deblur old memories. In Proceedings of the ninth annual con- ference of the Cognitive Science Society, pp. 177-186, 1987.
Hochreiter, S. and Schmidhuber, J. Long short-term memory. Neural computation, 98)::7335-188) 1997.
Rosenblatt, F.
The perceptron: probabilistic model for information storage and organization in the brain. Psy- chological review, 65(6):386, 1958.
Irie, K., Csordás, R., and Schmidhuber, J. The dual form of neural networks revisited: Connecting test time pre- dictions to training patterns via spotlights of attention. In International Conference on Machine Learning, pp- 9639-9659. PMLR, 2022.
Schlag, I., Munkhdalai, T., and Schmidhuber, J. Learning associative inference using fast weight memory. arXiv preprint arXiv:2011.0783/ 2020.
Kaplan, J., McCandlish, S., Henighan, T., Brown, T. B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., and Amodei, D. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361 2020.
Schlag, I., Irie, K.,, and Schmidhuber, J. Linear transform- ers are secretly fast weight programmers. In Interna- tional Conference on Machine Learning, pp. 9355-9366. PMLR, 2021.
Katharopoulos, A.,
Vyas, A., Pappas, N., and Fleuret, F. Transformers are rnns: Fast autoregressive transformers with linear attention. In International conference ma- chine learning, pp. 5156-5165. PMLR, 2020.
Schmidhuber, J. Evolutionary principles in self-referential learning, or on learning how to learn: the meta--eta-- hook. PhD thesis, Technische Universität München, 1987.
Kirsch, L. and Schmidhuber, J. Meta learning backpropaga- tion and improving it. Advances in Neural Information Processing Systems, 14:14122-1411 14134, 2021.
Schmidhuber, J. Learning to control fast-weight memories: An alternative to dynamic recurrent networks. Neural Computation, 4(1):131-139, 1992.
Kwon, W., Li, Z., Zhuang, S., Sheng. Y., Zheng, L., Yu, C. H., Gonzalez, J., Zhang, H., and Stoica, I. Efficient memory management for large language model serving with pagedattention. In Proceedings of the 29th Sym- posium on Operating Systems Principles, pp. 611-626, 2023.
Shazeer, N. Glu variants improve transformer, 2020.
Shleifer, S., Weston, J., and Ott, M. Normformer: Improved transformer pretraining with extra normalization. arXiv preprint arXiv:21 10,09456, 2021.
Lake, B. M., Ullman, T. D., Tenenbaum, B., and Gersh- man, S. Building machines that learn and think like people. Behavioral and brain sciences, 40:e253, 2017.
Su, J., Lu, Y., Pan, S., Murtadha, A., Wen, B., and Liu, Y. Roformer: Enhanced transformer with rotary position embedding, 2023.

---

Learning to (Learn at Test Time); RNNs with Expressive Hidden States
Sun, Y., Wang, X., Liu, Z.., Miller, J., Efros, A., and Hardt, M. Test-time training with self-supervision for generaliza- tion under distribution shifts. In International Conference on Machine Learning, pp. 9229-92442 9248. PMLR, 2020.
Zhang, B. and Sennrich, R. Root mean square layer normal- ization, 2019.
Zhang, H., Berg, A. C., Maire, M., and Malik, J. Svm-knn: Discriminative nearest neighbor classification for visual category recognition. In 2006 IEEE Computer Society Conference on Computer Vision and Pattern Recognition (CVPR'06), volume 2. pp. 2126-2 IEEE, 2006.
Thrun, S. and Pratt, L. Learning to learn: Introduction and overview. In Learning to learn, pp. 3-17. Springer, 1998.
Tieleman, T. and Hinton, G. Using fast weights improve persistent contrastive divergence. In Proceedings of the 26th anmual international conference on machine learn- ing, pp. 1033-1040, 2009.
Touvron, H., Martin, L., Stone, K., Albert, P., Almahairi,
A., Babaei, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., Bikel, D., Blecher, L., Ferrer, C.,, Chen, M., Cucurull, G., Esiobu, D., Fernandes, J., Fu, J., Fu, W.,
Fuller, B., Gao, C., Goswami, V., Goyal, N., Hartshorn, A., Hosseini, S., Hou, R., Inan, H., Kardas, M., Kerkez, V., Khabsa, M., Kloumann, I., Korenev, A., Koura, P. S., Lachaux, M.-A., Lavril, T., Lee, J., Liskovich, D., Lu, Y., Mao, Y., Martinet, X., Mihaylov, T., Mishra, P., Molybog, I. Nie, Y., Poulton, A., Reizenstein, J., Rungta, R., Saladi, K., Schelten, A., Silva, R., Smith, M., Subramanian, R., Tan, X. E., Tang, B., Taylor, R., Williams, A., Kuan, J. X., Xu, P., Yan, Z., Zarov, I., Zhang, Y., Fan, A., Kambadur,
M.. Narang, S., Rodriguez, A., Stojnic, R., Edunov, S., and Scialom, T. Llama 2: Open foundation and fine-tuned chat models, 2023.
Vincent, P., Larochelle, H., Bengio, Y., and Manzagol, P.-A. Extracting and composing robust features with denoising autoencoders. In ICMI. pp. 1096-1100 1003, 2008.
Wang, R., Sun, Y., Gandelsman, Y., Chen, X., Efros, A. A., and Wang, X. Test-time training on video streams. arXiv preprint ar vrXiv:2307.00014 2023.
Wolf, T., Debut, I.., Sanh, V., Chaumond, J., Delangue, C., Moi, A., Cistac, P., Rault, T., Louf, R., Funtowicz, M., et al. Huggingface's 's transformers: State-of of-e--of-theann natural language processing. arXiv preprint arXiv:1910003771, 2019.
Xiong, W., Liu, J.. Molybog, I., Zhang, H., Bhargava, P., Hou, R., Martin, L., Rungta, R., Sankararaman, K. A., Oguz, B., Khabsa, M., Fang. H., Mehdad, Y., Narang, S., Malik, K., Fan, A., Bhosale, S., Edunov, S., Lewis, M., Wang, S., and Ma, H. Effective long-context scaling of foundation models, 2023.
Yang, S., Wang, B., Shen, Y., Panda, R., and Kim, Y. Gated linear attention transformers with ing, arXiv preprint arXiv:2312.0000 1,660655,, 2023.
train-
Yang, S., Wang, B., Zhang, Y., Shen, Y., and Kim, Y. Par- allelizing linear transformers with the delta rule over se- quence length. arXiv preprint arXiv:2406. 06444,, 2024.
12

---

Learning to (Learn at Test Time); RNNs with Expressive Hidden States
Appendix
A. Dual Form
Here we derive the dual form for general MLPs of arbitrary depth, with nonlinear activations.
Without loss of generality, consider ?) for convenience, and consider only the first mini-batch, where 1 ,....p b. Denote:
f = KT
y/ = Wy Ft 0g
Also denote X = [ 86]. and and X analogously. In general, uppercase letters denote matrices whose columns are vectors denoted by the corresponding lowercase letter.
For network with K layers, denote the initial parameters in layer & by W Our convention is use superscripts for the layer and subscripts for time.
A.1. Forward pass
During the initial forward pass of TTT, we denote the input to layer & by X [ff.... at). with X X. Now we write the forward pass of TTT using these notations.
For h
-,....K' K:
Z = Wet X
gk+l - ok (2*)
where o for = 1,..., K can be any dlement-wise operation (R r> R) with derivative a'.
Given XK+1, we compute the loss:
1 l (W)
2 0: W x) 2 1|8t1|- 1
1 =
4.
where 4 = 2 12 2 is the same defined in Equation 3, except scaled by 1/2 for convenience.
All the operations above (except o) are matmuls and s sums, therefore are hardware efficient. Both the primal form and the dual form share these initial operations.
A.2. Primal form
The primal form first computes G V Vwele for = 1.... b, then updates W W E'm G.. Finally, given X b......āf ..]] X, the primal form repeats the forward pass with the updated Ws.
For -l...,K. K:
a = - =t. for I....,T T
x+1 = (z for = I....,T
where 8 k++ = |ø k+
k+1 1
contains the output tokens.
Note that a standard backward pass only computes the sum of the gradients:
V Wal
Σ Wihh Σ Gh
so the computation of the individual terms in the sum G for t 1,..., cannot be batched together into ma mutuls.. Similarly, the forward pass in primal form uses a different W, for each G,, therefore also cannot be batched in the same way as standard forward pass. These non non-standard passes have poor hardware efficiency.
13

---

Learning to (Learn at Test Time); RNNs with Expressive Hidden States
A.3. Dual form
As discussed in Subsection 2.5, the goal of the dual form is compute X and W b..... W with only matmuls and light- weight operations such sums, 0, and o'. To achieve this goal, avoid explicitly computing the intermediate variables: G and W for t I....,b.
The dual form first computes V Vxx+l= RK+1_ Y, then takes a standard backward pass.
For = ,
Vzel o (7*) o V X
Vxal= (W/)² T 2
Vwyl= = 2al (x*)²
where o' is applied element-wise and eiement-wise multiplication.
Now already compute W W V wgll To compute the output tokens, do another forward pass.
For h = -.....K. K:
T 8x =* = W -vzal mask (xx))²
8k+1 = (zt)
By the end of the forward pass, have computed X K+1
While this forward pass is non-standard, it only contains matmuls, sums, d, and mask, therefore is efficient like the standard forward pass.
A.4. Derivation
To derive the dual form, we show that:
((xA)" V
z²= Vzll l mask fmak
xx)
is the same as what would be computed in the primal form. Specifically, we show that each column zt of 7h in the second forward pass of the dual equals to W = in the forward pass of the primal. We invoke a simple fact.
Fact 1. Define matrices A [a_..... aul, Q Iq.....,qul :ul, and V = [ui....,ool vo).3 Define én = E. a qu and V pi,.... 06), then V- V mask ( QQ)
Now plug A - 8 Q 8 V 1 and V x* 2 into the fact above, have shown the desired equality.
Note that the 0% and o used above can be extended to arbitrary functions that not necessarily element-wise operations, including normalization layers. This extension can be achieved through, for example, vip (vector-Jacobie) product) in standard libraries for automatic differentiation such as JAX and PyTorch. However, the dual form cannot accelerate operations inside a or its vjp.
"Our matrix would usually denoted hy another context. We use to avoid confusion with the layer number
14

---

Learning to (Learn at Test Time); RNNs with Expressive Hidden States
B. Nadaraya- Watton estimator
Derivation for the Nadaraya- Watson estimator. Throughout this section, we use x to denote the input token æ as random variable. Our desired output is the corresponding output token, another random variable z. This is formulated as estimating the conditional expectation of z:
E x] p(z)1) do /
[ p(x)
dz.
Since the true probability distributions p(xx) and p(z, are unknown, replace them with their kernel density estimations. Specifically, the kernel density estimation for p(x)) is:
p(x) = Σ x(),
where each i; a piece of training data general. (Recall that for our paper, vs is specifically training data for the inner loop, i.e. token, which matches our notation in the main text.)
For estimating p(., v), the product kernel:
p(x, 2)
Σ x(2, x4) (2(2,z))
At first sight, it seems absurd to factor the joint probability into two seemingly independent kernels. But in this case, n' can actually be any ni dependent on Ii, since will be integrated out. So the two kernels do not need to be independent. Plugging in those estimations, we obtain the Nadaraya- Hadaraya-Watson Watson estimator:
(I, 2) j(x)
E E|z|x x|
dz
()x) / 11xx z) dz
/ Σ klaa 2z) (2 a dz
El= N(),
1
k(0, xx) I (2) de
Ei= 16(I,, xz)
1
<(x xz)) Zi.
E K(I,
Asymmetric kernels. In modern days, people think of kernels as positive semi-definite, which might not be guaranteed for is unless 0K = 0Q. However, people working on kernels decades ago, around the time when the Nadaraya-Watson estimator was popular, have been very lenient with the choice of kernels, and asymmetric kernels such as our x in Equation 9 have enjoyed long tradition: When kernel estimator uses 0x * lq. is known as a balloon estimator (Chen, 2017). Papers such Breiman al. (Breiman et al., 1977) have even used 0g as function of x', known sample-adaptive smoothing.
15

---

Learning to (Learn at Test Time); RNNs with Expressive Hidden States
Residual block
Transformer backbone
Mamba backbone
O
MLP block
LayerNorm
LayerNorm
LayerNorm
TTT layer
TTT layer
Soquence modeling block
Conv
Conv Q
LayerNorm
V
K
Q
V
K Q
Gate
Figure 9. Left: A residual block, the basic building block for Transformers. The sequence modeling block is instantiated into variants: the Transformer backbone and Mamba backbone. Middle: TTT layer in the Transformer backbone. The LN before O comes from NormFormer (Shleifer al., 2021). Right: TTT layer in the backbone inspired Mamba (Gu & Dao, 2023) and Griffin (De al., 2024). Following these two architectures, a here GELU (Hendrycks & Gimpel, 2016). To accommodate the extra parameters of the gate without changing the embedding dimension, simply combine 0g and 0o into single projection.
C. Experiment details
Architectures. Our Transformer strictly follows the construction in the Mamba paper, where Transformer is called Transformer++. Specifically, the Transformer architecture is based on Llama (Touvron et al., 2023), with rotary positional encodings (RoPE) (Su et al., 2023), SwiGLU MLP blocks (Shazeer, 2020), and RMSNNrm (Zhang & Sennrich, 2019) instead of LayerNorm. Our Mamba baseline uses the public code provided by the authors. We have verified that our baselines can reproduce the numbers reported (Gu Dao, 2023).
Training configurations. Our training configurations are in Table 2, which simply reproduces Table 12 in the Mamba paper. All models are trained with a batch size of 0.5M tokens regardless of context length. All of our optimization hyper-parameters follow the "improved recipe" in Appendix E.2 of the Mamba paper, reproduced below:
AdamW optimizer: 3 = (0.9, 0.95)
Cosine schedule: decay to end learning rate 1e 5
Linear learning rate warmup 10% of the training steps
Weight decay: 0.1
Gradient clipping: 1.0
No Dropout
Mixed Precision
For experiments on the Pile, this the only difference with the recipe in the Mamba paper, which uses two other tokenizers.
For experiments on Books, we find that the original angle of the RoPE encoding (Su al., 2023) 0 10, 000 sub-optimal
16

---

Learning to (Learn at Test Time); RNNs with Expressive Hidden States
Params.
Blocks
Embed. dim.
Heads
Train steps
Peak LR
Tokens
125M
12
768
12
4800
3e-3
2.5B
350M
24
1024
13500
1.5e-3
7B
760M
24
1536
16
29000
1.25e-3
15B
1.3B
24
2048
32
50000
1e-3
26B
Table Training configurations for all experiments. This table reproduces Table 12 in the Mamba paper. The only difference is that the learning rate they use for Mamba and Transformer is 5x the values in their Table 12, and we report the actual values (5x). Note that this table only applics to TTT-Linear, TTT-MPP, and Transformers, Mamba does not follow the multi-hcad residual black structure inherited from Transformers.
for our Transformer baseline in long context. Starting at context length 4k, we try 0 = 500, 000 following the Llama Long paper (Xiong al., 2023), and use the better perplexity for Transformer (both pretrain and finetune).
Transformer finetuning, Finetuning starts from scratch, except the peak learning rate. We try three peak learning rates for finetuning: 1e-5, le-4, and 1e-3, and select for the best perplexity. We observe that le-4 works the best for the 125M models, while le-5 works the best for 350M and larger. This observation is reasonable considering that the end learning rate for the Chinchilla recipe is 1e-5.
cosine schedule with the same optimization syper-parameter training
Learning rate for TTT. As mentioned in Subsection 2.7, the inner-loop base learning rate Yhuse is set to for TTT-Linear and 0.1 for TTT-MLP. Our heuristic for setting Time is similar to how people set the outer-loop learning rate for regular training: We tried Thase E {0.01, 0.1, 1, 10} and used the largest value that does not cause instabilities. For TTT-MLP, we use linear warmup for Thase over 10% of the training steps. similar regular training. The number of training steps in the inner loop is T (assume divisible). For TTT-Linear, we tried linear warmup in the inner loop but did not observe a difference.
Experiments in Figure (right). To ensure fairness to Mamba, all methods in these experiments have matched training FLOPs and are trained with the same recipe (last row of Table 2) as Mamba 1 1.4B, For TTT-Linear and TTT-MLP, matched training FLOPs also imply matched inference FLOPs. Transformer (TF finetune) has 2.8x the inference FLOPs, giving it an advantage baseline. To match training FL.OPs with Mamba, Transformer has 19 blocks instead of 24. For TTT-Linear and TTT-MLP, their training FLOPs are already close to those of Mamba, we only need to change the hidden dimension of the MLP blocks from 5504 to 5808 for TTT-Linear and 5248 for TTT-MPP.
Gradient checkpointing through time. By default, libraries such as JAX and PyTorch save the intermediate activations during forward pass so they can be reused during the backward pass. However, for TTT layer with W as hidden state, this default saves Wi.... Wr. which uses too much memory. With TTT mini-batch and the dual form, we still need save (assume divisible) T/b Ws at the end of the mini-batches. A standard technique to save memory in this scenario gradient checkpointing (Chen et al., 2016), which is usually applied through layers, but apply it through time.
17

---

Learning to (Learn at Test Time); RNNs with Expressive Hidden States
1.8
MWlLyfeeMtet MANN
1.0
1.0
1.6
loss
TTT loss
0.9
K[Wo: xy)
losst
0.9
#(Ws: x:)
KW; x;
0.8
R:
R(W xx)
0.8
1.2
X(Wex)))
((W6
0.7
1.0
0.6
0.6
500
1000
1500
2000
500
1000
1500
2000
500
1000
1500
2000
Token index
Token index
Token Index
1.1
1.0
loss7
TTT loss!
0.9
F(Wo: xt)
loss
E(Wo: X.)
F(W xt
E(W:11:X)
0.8
I(W;
2(W;
0.6
KW;
0.7
1.0
((W))))
0.6
500
1000
1500
2000
500
1000
1500
2000
500
1000
1500
2000
Token index
Token index
Token index
1.1
1.6
1.0
1.6
loss
0.9
Z(Wo: x
loss
1.4
TTT loss
z(Wo:
X(W,-1:x))
[(W,-1:x)
1.4
0.8
x(W):
E
1.2
2(W, xx)
RLW; xi)
W
0.7
1.0
R(WG X))
1.2
500
1000
1500
2000
500
1000
1500
2000
500
1000
1500
2000
Token index t
Token index
Token index
wwwm///weeeett///eoooom Ma
2.0
1.6
1.8
loss
1.8
ITT loss
1.4
X(Wg:
loss
2(Wo: x))
R(W; 1: x))
2(W; 1: Kt)
1.6
Z(W; xt)
2(W5 x.)
kW):)))
=1W Xi)
K(W:
1.0
1.2
500
1000
1500
2000
500
1000
1500
2000
500
1000
1500
2000
Token index
Token Index
Token Index
Figure 10. The self-supervised TTT loss averaged all sequences of the form #i....,47 where 2048, for all 12 TTT layers in a network with 125M parameters train on the Pile. The same network also used for (online GD) in the left panel of Figure 4, For layers the middle, we observe that ||>,]] || rises steadily, causing all three losses to rise with Even for these layers, the between ((Wo; xx) and E(We; xx) still increases with t. For visual clarity, loss values have been averaged over a sliding window of 10 timesteps
18

---

Learning to (Learn at Test Time); RNNs with Expressive Hidden States
1k
2k
18
TF pretrain
17
TF pretrain
17
Mamba
16
Mamba
TTT-Linear
TTT-Linear
TTT-MLP
15
TIT-MLP
Perplexity (log scale)
15
Perplexity (log scale)
14
14
13
13
12
12
11-
11
10a
10a
10
10a0
10
1020
FLOPs (log scale)
FLOPs (log scale)
4k
8k
TF finetune
TF finetune
16
TF pretrain
16
TF pretrain
15
Mamba
Mamba
TTT-Linear
TTT-Linear
Perplexity (log scale)
14
TTT-MLP
Perplexity (log scale)
14
TTT-MLP
13
12
12
11
101
10a
9
10
1020
10
1020
FLOPs (log scale)
FLOPs (log scale)
16k
32k
18
TF finetune
TF finetune
18-
TF pretrain
TF pretrain
16
Mamba
Mamba
16-
TTT-Linear
TTT-Linear
Perplexity (log scale)
Perplexity (log scale)
TTT-MLP
TTT-MLP
14:
12-
12
10a
101
1018
10
102c
1018
10
1020
FLOPs (log scale)
FLOPs (log scale)
Figure/ 11. Complete results Books, presented by context lengths. Figure 7 in Subsection 3.2 presents the subset of results for context lengths 2k and 32k.
19

---

Learning to (Learn at Test Time): RNNs with Expressive Hidden States
125M
350M
TF finetune
12.8
TF pretrain
12.6
perplexity (log scale)
Mamba
perplexity (log scale)
TTT-Linear
12.4
17.5
TTT-MLP
17.0
12.0
TF finetune
11.8
TF pretrain
Mamba
16.0
11.6
TTT-Linear
TTT-MLP
11.4
2k
4k
8k
16k
32k
2k
4k
8k
16k
32k
context length (log scale)
context length (log scale)
760M
1.3B
10.4
perplexity (log Scale)
perplexity (log scale)
10.2
9.0
8.8
9.8
TF finetune
finetune
TF pretrain
TF pretrain Mamba
9.6
Mamba
TTT-Linear
TTT-Linear
8.4
9.4
TTT-MLP
TTT-MLP
2k
4k
8k
16k
32k
2k
4k
8k
16k
32k
context length (log scale)
context length (log scale)
Figure 12. An alternative view of complete results Books, presented model sizes, with context length the x-axis. For all methods trained from scratch, perplexity becomes worse once the context length becomes too large. This trend is not observed with TF finetune, except for one case the 125M scale. The best context length increases for larger models (trained from scratch).
20


Reviews:
[{"review_id": "KHXvEIQ1Kb", "reviewer": "Reviewer_SWT9", "summary": "This paper presents a novel approach to sequence modeling by enhancing the hidden states of RNNs through a general method called Test-Time Training (TTT). The key idea is to frame the online updating of hidden states as a self-supervised learning process, using the update loss $\\ell: W_t = W_{t−1} − \\eta \\nabla \\ell(W_{t−1}; x_t)$, where $W_t$ is the hidden states learnt at each step, reminiscent of fast weights in RNN literature, and $\\nabla$ denotes gradients from the chosen objective function. This framework unifies several previous approaches, including naive RNNs, linear attention, and even self-attention, under a common paradigm.\n\nOne main challenge of TTT is the difficulty of parallelization, as the proposed objective may introduce non-linearities and dependencies on previous steps. To address this, the authors propose a batch update approach, allowing $W_t$  to depend on the previous N steps, or in extreme cases, on $W_0$. Interestingly, the authors demonstrate that under certain conditions, TTT becomes equivalent to vanilla linear attention.\n\nBuilding on this foundation, the authors introduce two variants: TTT-linear and TTT-MLP, both incorporating layer normalization and mini-batch updating. The experimental results on approximately 1B parameter models show promising performance in both short and long-context scenarios, demonstrating the potential of TTT for sequence modeling tasks.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for recognizing our framework’s potential to guide the development of RNNs with better memory. We also thank the reviewer for the concrete suggestions and questions, which we address below.\n\n***Benchmark on downstream tasks beyond perplexity***\n\nThe focus of our paper is on long context. However, all the downstream tasks in Mamba (Gu and Dao, 2024), as the reviewer suggested, have sequence length <1000. In fact, most of them only need ~100 tokens. There also exist downstream tasks with long context, such as book summarization and solving software issues in large repositories. However, these tasks require post-training for instruction following capability, and pre-training frontier models with trillions of tokens. Unfortunately, the cost of this evaluation exceeds our budget.\n\n***Comparison with Mamba 2***\n\nMamba 2 130M with 2k context trained with Chinchilla scaling law on the Pile has perplexity 10.77, which is significantly lower than our number of 11.09 reported in Table 1 of the manuscript. However, 11.09 was based on the older Mamba backbone. During the rebuttal period, we experimented with putting TTT-Linear into the Mamba 2 backbone. This model has perplexity 10.65, which is slightly better than Mamba 2.\n\n***Concurrent work***\n\nWe appreciate these references highlighted by the reviewer. We will add the discussion below to the final version if the paper is accepted (ICML does not allow submitting a revision):\nConcurrent work such as [1] [2] [3] has further improved the update rule for RNNs with matrix hidden states. Given that TTT-Linear is already similar to DeltaNet and linear attention, it is also similar to these more recent variants. As discussed in the opening of the paper, our core contribution is a practical framework to use any neural network as the hidden state of an RNN. This contribution is still salient in comparison with concurrent work.\n\n***Would LayerNorm create obstacles for parallelism? How to implement the dual form for TTT-Linear and TTT-MLP?***\n\nSorry we did not include these details in the main text due to space constraints. We chose to defer them to Appendix A (page 14 and 15), which includes a derivation and pseudocode of the dual form for neural networks of arbitrary depth with arbitrary nonlinear operations. We will expand Subsection 2.5 in the final version to include at least some of these details.\n\nFollowing the explanation in Appendix A, it should be clear that LayerNorm (or any other nonlinear operations with few FLOPs) would not create an obstacle for parallelizing the rest of the hidden state. While LayerNorm runs on CUDA Cores, compute-intensive operations, specifically matrix multiplications, will still run on Tensor Cores  As an intuitive analogy, regular neural networks, such as Transformers, also have many nonlinear operations, but all of their linear operations can still be parallelized.\n\nDoes the reviewer feel that we have adequately addressed their concerns? Please let us know if we can provide more information.", "claims_and_evidence": "No.\n\nI really appreciate the TTT framework proposed in this paper. However, the experiments conducted are somewhat limited, making the validation of TTT less comprehensive and persuasive. \n\nI strongly recommend the authors include additional common benchmarks found in related works like Mamba or GLA, such as hellaswag, lambada_openai, and winogrande, beyond just measuring perplexity on text length. \n\nFurthermore, more detailed comparisons with other relevant models like DeltaNet (which is heavily discussed in the paper) and Mamaba2 would significantly strengthen the evaluation.", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "Yes", "experimental_designs_or_analyses": "Yes", "relation_to_broader_scientific_literature": "I can see potential benefits in using TTT to develop more efficient memory management strategies for RNNs, guided by the principles of TTT.", "essential_references_not_discussed": "Recent concurrent works should be included and discussed:\n1. Unlocking State-Tracking in Linear RNNs Through Negative Eigenvalues\n2. gated delta Networks: Improving Mamba2 with Delta Rule\n3. Longhorn: State Space Models are Amortized Online Learners", "other_strengths_and_weaknesses": "1. The experimental section of this paper is insufficiently comprehensive, making the results less convincing.\n2. The detailed implementations of TTT-MLP and TTT-Linear are not well-explained, which is quite different from the naive cases in Sec 2.5.", "other_comments_or_suggestions": "See below.", "questions_for_authors": "I'm having trouble understanding how TTT-Linear and TTT-MLP work. \nGiven that both methods include Layernorm, which might create obstacles for parallelism and make it difficult to accelerate using tensor cores.\n\nCould the authors explain how the dual forms introduced in Section 2.5 can be implemented for TTT-Linear and TTT-MLP?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2l53Ff55FM", "reviewer": "Reviewer_L8fP", "summary": "This paper introduces Test-Time Training (TTT) layers, a 'clever' way to handle long sequences without the heavy cost of Transformers. The key idea is to treat the hidden state as a learning model that updates itself during inference, allowing it to capture complex patterns over long contexts. The authors propose two versions, TTT-Linear and TTT-MLP, which outperform existing models like Mamba but also runs faster on modern hardware.", "questions": "", "limitations": "", "rating": 5, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for appreciating our paper as a sufficient proof-of-concept. We also thank the reviewer for the insightful questions, which we answer below.\n\n***1 - Comparison with Mamba 2***\n\nMamba 2 130M with 2k context trained with Chinchilla scaling law on the Pile has perplexity 10.77, which is significantly lower than our number of 11.09 reported in Table 1 of the manuscript. However, 11.09 was based on the older Mamba backbone. During the rebuttal period, we experimented with putting TTT-Linear into the Mamba 2 backbone. This model has perplexity 10.65, which is slightly better than Mamba 2.\n\n***2 - Discussion on Irie et al. 2022***\n\nWe thank the reviewer for bringing up this very relevant work. We will add the discussion below to the final version if the paper is accepted (ICML does not allow submitting a revision):\n\nIrie et al. 2022 is similar to DeltaNet, except that the key, query and value vectors are also produced by the hidden state matrix. In the terminology of TTT, all outer-loop parameters are the initializations of inner-loop parameters. This approach naturally situates Irie et al. 2022 in the same setting as MAML (Finn et al. 2016), and their experiments are both on few-shot learning.\n\n***3 - Results with Mamba 2 backbone***\n\nAnswer for your question 1 in this rebuttal switches to Mamba 2 backbone. As the reviewer expected, the improvement in backbone is indeed orthogonal to the improvement we had with TTT layers. Mamba 2 had not been released when we conducted most of the experiments in May, 2024.\n\n***4 - Incorrect claim about related work***\nIndeed, we have missed some of the recent developments in SSMs, making our claim incorrect. We appreciate the reviewer for carefully reading our text and pointing out this error. We will make sure to fix this in the final version if the paper is accepted.\n\n***5 - Improvements on top of linear hidden states***\n\nThe hidden states of linear attention and Mamba 2 are already matrices. Compared to linear attention, TTT-Linear brought significant improvements because of the added nonlinear operations. Compared to Mamba 2, TTT-Linear still brings a small improvement.\n\n***6 - Sharing code for reproducibility***\n\nSorry we forgot to upload a zip file for ICML submission. We had to delete the link to our online code repo in the manuscript for anonymity. Many people have been able to reproduce our results using the code online. We will add a link to the final version if the paper is accepted.\n\n***7 - Tradeoff between sensitivity to noise vs. over-smoothing***\n\nWe thank the reviewer for pointing out this interesting paper. More expressive hidden states will improve both ends of this tradeoff. As discussed in the paper, the key that controls this tradeoff is recency bias: too much recency bias results in sensitivity to noise, but too little results in over-smoothing. RNN layers with expressive hidden states will be able to remember useful information (therefore less over-smoothing) without using aggressive recency bias (therefore less sensitivity to noise), simply because it has a larger capacity for information. Take self-attention as an example. It does not need any recency bias because it simply stores all key-value pairs in a cache. We believe that very expressive inner models for TTT can achieve similar effects.", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "Yes", "experimental_designs_or_analyses": "The experimental design is limited but sufficient for a proof of concept. I think they could have included another task beside language modeling, e.g., time series.", "relation_to_broader_scientific_literature": "This paper builds on prior work in sequence modeling,(especially  Schlag et al., 2021 and Yang et al. 2024), while addressing their limitations in long-context settings. By leveraging self-supervised learning and meta-learning techniques, the proposed Test-Time Training (TTT) framework introduces adaptive hidden states which might be the way to improve both efficiency and expressiveness of SSMs", "essential_references_not_discussed": "I think (Irie, Kazuki, et al. \"A modern self-referential weight matrix that learns to modify itself.\" ICML, 2022.), which also combines ideas of FWP, RNNs, meta-learning through recursive self-improvement,  is extremely relevant here and should be discussed in this paper, with key differences highlighted.", "other_strengths_and_weaknesses": "Paper is very well-written and structured\n\nWeaknesses: \n- The main weakness of this work is the limited experimental setting and comparisons. Including results with relevant baselines, e.g., DeltaNet and Mamba2 would make the claims stronger and help us see how good these models are compared to these baselines.\n\n- The authors don't share any code for reproducing their results. So hard to judge the reproducibility.", "other_comments_or_suggestions": "Nan", "questions_for_authors": "1- Why not Including competitive results for Mamba2 and DeltaNet?\n\n2- Can you comment on the similarities and differences compared to  (Irie, Kazuki, et al. \"A modern self-referential weight matrix that learns to modify itself.\" ICML, 2022.)?\n\n3- Why did use mamba as backbone for TTTLinear and TTT-MLP? why not recent SSM models, such as mamba2? How does the model perform with such backbones? I understand that (mamba-->mamba2) is an orthogonal improvement to the contribution proposed here but would it be better to include the results with the top performant backbones to assess the full potential of TTT?\n\n4- the claim in first paragraph in Section 4.1. '' 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'' is not valid for all SSMs. In fact, recent SSMs, such as (FACTS: A Factored State-Space Framework For World Modelling) the hidden state is already a matrix. \n\n5- When the hidden state  of the backbone is already a Matrix (e.g., FACTS), does using TTT still improve the expressiveness of the model or this gain is reduced? would it affect the computational aspects of TTT?\n\n6- Is it possible to share code to reproduce some parts of your results (e.g., on the Pile) otherwise, it is hard to assess the reproducibility of the approach's results.\n\n7- One of the major issues of SSMs models and Mamba in particular (https://arxiv.org/pdf/2501.00658) is the over-smoothing and sensitivity to noise, especially at the end of the sequence (See table 1 in https://arxiv.org/pdf/2501.00658). I wonder how does test-time training affect this behavior? Can the authors share their thoughts on it. Because if it is the case, this can mitigate one of the current major bottlenecks of SSMs.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "lAdq3lfzko", "reviewer": "Reviewer_7Qqf", "summary": "The paper proposes the Test-Time-Training (TTT) layer, with the goal to overcome the limitations on the expressive power of modern RNNs. The idea consists in linking the hidden state of an RNN to the parameters of a layer, so that input-driven updates on the hidden state translate into updates of model parameters: by drawing a connection between these updates and gradient descent, the authors effectively manage to setup a training procedure for (a subset of) parameters in the model, which unfolds at test-time. As proven by the authors, the proposed TTT formulation encompasses and generalises the Attention mechanism. \nHardware considerations are employed to speedup the application of the TTT layer, resulting in wall-clock times comparable (or even outperforming) Mamba’s. A detailed analysis of the performance of the layer is reported, including numerous ablations, and scaling behaviour of perplexity of the model when varying context length and size.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the insightful questions, which we answer below.\n\n***1 - Flexibility w.r.t. inner model***\n\nWe agree. We will add the discussion below to the final version if the paper is accepted (ICML does not allow submitting a revision):\n\nIn principle, any differentiable model and loss function can be instantiated as the inner-loop model and loss in our framework, but the practical effectiveness of such instantiations can be limited by their requirements in wall-clock time. While TTT-MLP is an effective instantiation in FLOPs, the additional complexity of the MLP structure makes the increase in wall-clock time much larger than the increase in FLOPs, as shown in Figure 6. It remains to be seen whether our framework can produce instantiations that overcome this limitation.\n\n***2 - On using a different loss***\n\nYes, a different inner loss among the most important things to try next. For the past few months, we have been experimenting with using next-token prediction as the inner loss and a Transformer as the inner model. Our initial results so far have been quite promising.\n\n***3 - Figure 6***\n\nFigure 6 is for training (forward and backward). Timing for inference is in Figure 13 in the Appendix, which has a separate plot for prefill (forward) and decode (generation). We appreciate that the reviewer is clear about the difference between training, prefill and decode. \n\nThe reviewer said: *\"I was expecting to see a (slight) increase of wall-clock time of TTT with context length.\"*\n\nA linear-complexity method that is completely parallelizable across the entire sequence should see a decrease in latency as context length increases, while one that is completely sequential should have constant latency. This is because latency = processing time / tokens processed. As tokens processed increases, the processing time for a parallelizable method increases at a much slower rate, while that for a sequential method increases at the same rate. In Figure 6, none of the methods exhibit significant decrease in latency, showing that given the compute resources in Figure 6, they are all far from parallelizable.\n\n***4 - Figure 7***\n\nWe appreciate the reviewer’s attention to details. We use RoPE, but we do not think the increase in l(x_t;  W_0) has to do with the position embedding. As the reviewer already knows, the increase in l(x_t;  W_0) is a consequence of the increase in \\|x_t\\|, whose reason is unclear to us. We have measured \\|x_t\\| in the Transformer baseline and observed a similar increase, so we know that it is not caused by using TTT layers.\n\nWe will also change the running title, thanks for catching that. Regarding reproducibility, we had to delete the link to our online code repo in the manuscript for anonymity. Many people have been able to reproduce our results using the code online. We will add a link to the final version if the paper is accepted.", "claims_and_evidence": "For the most part, the evidence proposed is complete enough to allow for a proper assessment of the robustness of the method, and convincing in supporting its validity.\nThe only claim I would reframe refers to the applicability of TTT to “any neural network as inner model”. While I believe this to be true in principle (as indeed the framework introduced is general enough), in my opinion it fails to address practical considerations which could actually hinder this, and I think it would be fair to better discuss and highlight this limitation in the paper. See also Q1. This is particularly relevant as the generality of the TTT framework is brought up multiple times as one of the main contribution of the paper, setting it aside from previous work (see eg end of Sec4.1 and 4.2).", "methods_and_evaluation_criteria": "The main revolve around comparing perplexity of the proposed model against Mamba and Transformer baselines, after training on the Pile dataset on a fixed FLOPs budget. The sizes of the models and baseline compared are kept similar. The evaluation is in my opinion fair and meaningful. The only improvement I would recommend consists in reporting performances for downstream tasks as well (using, eg, eval_harness framework), to further confirm that perplexity gain indeed do transfer.", "theoretical_claims": "The main theoretical claims aim at drawing correspondences between the proposed method and linear / softmax attention, when specific choices on the TTT loss function and inner model are made. They also show the equivalence of primal and dual formulation of their method. The proofs mostly reduce to algebraic manipulations: I’ve checked their derivations and they appear correct.", "experimental_designs_or_analyses": "To my knowledge, the authors haven’t made their code available, so I can’t check it directly. Regarding the experiment setup, it is quite standard, so I don’t have any particular issue to highlight (see also Methods and Evaluation Criteria, and Q3).", "relation_to_broader_scientific_literature": "The project expands upon two main branches of work: modern RNNs, and Test-Time-Adaptation. Regarding modern RNNs, for its mini-batching procedure the paper reuses a similar chunk-splitting concept as in GLA (Yang et al, 23); however, in GLA this splitting is mostly a reorganisation of the operations to improve hardware efficiency, while in TTT it’s used to introduce an actual sequential update. Most important are the similarities to DeltaNet (Yang et al, 24): indeed, for specific choices of its inner model and loss function, TTT is exactly equivalent to DeltaNet. The authors however expand on this framework (with some caveats, see Q1). Regarding Test-Time-Adaptation, the method directly inherits the core idea of training a net during inference. In either case, in my opinion the discussion in Sec4 covers these connections reasonably well and fairly.", "essential_references_not_discussed": "The paper seems to reference the most relevant literature already.", "other_strengths_and_weaknesses": "The paper is clearly written for the most part, with its core method being described in a concise but effective manner. The idea is interesting and original, although it does build upon other work on modern RNNs and Test-Time Learning. The validity of the method is properly corroborated by the results provided, without over-claiming.", "other_comments_or_suggestions": "Please change the running title to something more meaningful than “Submission and Formatting Instructions for ICML 2025”", "questions_for_authors": "Main Questions \n1. *Flexibility wrt inner model / loss function* -  While it’s true your model formulation is flexible enough to accommodate for generic inner model and loss functions in theory, in practice wall-clock time remains a concern. As your layer aims to replace Attention or the Mamba mixer, it is reasonable to expect its application to require a comparable computational time. I really appreciated in this sense the comparison in Fig6, but the results therein already highlight how much just introducing the simplest nonlinearity (an MLP) impacts wall-clock time, due to the added complexity in both (inner) forward and backward computations. More complicated loss functions, or layer whose gradient is not readily computable, or that involve additional matrix-matrix-multiplications, are all factors which will further exacerbate this. If you agree with this, it should be properly reported as a limitation of the method; if not, please elaborate. \n\n\nMinor  \n\n2. *On using a different loss* - Given the improvement that mini-batching provides to Linear Attention (ie, TTT-linear with $b=T$, as per Thm1), I was looking forward to seeing a similar application to self-attention, (leveraging Thm2). Also in light of my comment in Q1, I understand why you decided to stick with the simpler inner-loss definition, but it would be interesting to try (especially as I agree with your remark that the choice of inner loss would likely impact the performance of the method), and was wondering whether you had already started experimenting on it.\n\n3. *Fig6* - Is Fig6 reporting average time per token at *training* or *inference*? The reason I’m asking this, is because when processing input data at training, I was expecting to see a (slight) increase of wall-clock time of TTT with context length. This is in light of the fact that, while computing gradients updates *within* a mini-batch is parallelisable (as you show and take advantage of), the TTT gradient updates *along* mini-batches is a sequential procedure (as you need to compute the various updated $W_{0}, W_{b}, W_{2b},…$ in order). Of course this doesn’t hold for generation, since no form of sequence length parallelisation is applied anyway. If Fig6 actually reports average inference time at generation, as I suspect, could you correct its title and also include average time for a forward pass during training (just to get an idea of how much more/less time consuming it is to train your model vs Mamba or Transformer). If not, what am I missing?\n\n4. *Fig7* -  Why would $l(W_0,x_t)$ increase over $t$ in Fig7? What causes the rise in $||x_t||$? Which Positional Encoding do you use?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "xTIMaDIGpn", "reviewer": "Reviewer_i78k", "summary": "This paper casts the sequence modeling problem as a meta-learning problem at training time. The resulting model is a model which minimizes a loss, i.e. learns at test time.\nThe authors show that Linear Attention and Attention are special instances in their Learning to Learn at Test Time Framework.\nBuilding on this framework the paper introduces new neural network layers called Test-Time Training (TTT) layers. The paper considers two instantiations of these layers TTT-Linear and TTT-MLP, whose hidden state is a linear model and a two layer MLP, respectively.\nIn experiments with model sizes ranging from 125M to 1.3B parameters these novel layers match or exceed the performance of Transformer and Mamba baselines.\n\n## Update after Rebuttal:\nThe authors addressed all my questions and I still recommend acceptance of this paper.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for recognizing the impact of our framework. We also thank the reviewer for the concrete suggestions and questions, which we address below.\n\n***Explicit update rule for the final instantiations***\n\nSorry we did not include these formulas in the main text due to space constraints. We chose to defer them to Appendix A (page 14 and 15), which includes a derivation and pseudocode of the update rules in dual form. We will expand Subsection 2.5 in the final version to include at least some of these details.\n\n***Learning rate parameter***\n\nThe learning rate parameter \\theta_{lr} is a vector of the same dimension as the embeddings x_t. We have tried making it a scalar and found performance to be slightly worse. We have not tried making it a matrix since that would significantly increase our number of parameters.\n\n***More fine-grained gating***\n\nWe agree with the interpretation that \\eta(x) is a gate (more precisely, an input gate). However, we believe that the trend in Figure 9 is not a consequence of \\eta(x_t) being input dependent, since we have produced Figure 9 with a fixed \\eta and the trend is the same. Intuitively, \\eta(x_t) is used for every x_t regardless of the inner-loop mini-batch size b, so larger b will not result in more fine-grained gating.\n\n***Answering the numbered questions***\n\n1 - Yes, sorry for the confusion. Perhaps a better phrase would be “efficient in computational complexity.”\n\n2 - No, it does not mean what the reviewer has guessed. We are sorry for the confusion. Sometimes researchers would have multiple runs of the same experiment but with different random seeds to overcome instability, since some runs would get unlucky. We meant to say that we did not have multiple runs with different random seeds. We will delete this sentence in the final version to avoid possible confusion.", "claims_and_evidence": "The claims made by this paper are supported by convincing evidence.", "methods_and_evaluation_criteria": "The choice training the models and baselines on the Pile dataset are reasonable. There exist newer, more cleaned open source datasets such as SlimPajama or DCLM Baseline, but these would very likely not change the overall results.\n\nTraining on two different context lengths 2k and 8k underlines the fair benchmark setup. For 2k context length the FLOP difference between Transformers (quadratic compute scaling with sequence length) and RNNs (linear scaling with sequence length) is smaller, while for 8k it is larger.", "theoretical_claims": "I checked the proofs for Theorem 1 and Theorem 2 for correctness.", "experimental_designs_or_analyses": "Their experimental design is sound and their ablations in Table 1 shows careful model design.", "relation_to_broader_scientific_literature": "Even before publication/acceptance at a conference, this paper has inspired several other follow up works. \n\nMoreover, the paper unifies several other linear RNN architectures, such as Mamba, xLSTM (mLSTM), DeltaNet, and many more in a unified framework.", "essential_references_not_discussed": "None.", "other_strengths_and_weaknesses": "Strengths: \n- Extensive empirical validation of the models. \n- Extensive discussion of related work.\n- Their proposed model architectures show competitive performance.\n\nWeaknesses: \n- The update formulas of their new models (including the dimensions) could be stated explicitly.", "other_comments_or_suggestions": "- It would improve understandability if the final instantiation update rules for the Linear and the MLP variant would be stated somewhere in the paper (including the layernorm and the residual connection).\n\n- You choose the learning rate parameter learnable, so is the learning rate parameter a scalar? Did you think of different parametrizations?\n- If one views the learning rate as inter chunk forget gate, can we then explain the improved performance with smaller b in Figure 9 due to a more finegrained (forget-)gating?", "questions_for_authors": "1) L. 156: Why is the naive layer already efficient in FLOPs? FLOPs not mentioned so far. Do the authors mean the non-quadratic scaling in sequence length compared to transformers?\n2) L. 286: What does it mean „all experiments for the TTT layers are performed in one training run“?\nDoes that mean that there is one model trained with 32k context and then evaluated at different context lengths?\nIf so which FLOP budget is then plotted on the x-Axis? The training FLOPs?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: **Overall review**
There is consensus towards accepting this work. Reviewers coincide this is a well-written paper, with good results and solid theory. They all lean towards acceptance. I also believe this work is interesting as it generalizes transformers and RNNs and can spark new research in the topic.
 
**Summary**
The paper introduces a novel sequence modeling approach . The key idea is to treat the hidden state of a recurrent neural network (RNN) becomes the weights of a learnable model that updates itself during inference using a meta-learned self-supervised task. This allows the model to capture complex patterns over long contexts more efficiently than traditional Transformers.

**Reviews**

Reviewer SWT9:

Strengths:
Appreciates the Test-Time Training (TTT) framework.
Sees potential benefits for efficient memory management in RNNs.

Weaknesses:
Limited experimental validation. (Partially addressed: Authors provided new results comparing to Mamba2)
Missing comparisons with relevant models (DeltaNet, Mamba2). (Partially addressed: Authors provided new results comparing to Mamba2)
Unclear implementation details of TTT-MLP and TTT-Linear. (Addressed: Authors mentioned that details were in the appendix and would expand the section in the final version.)


Reviewer L8fP:

Strengths:
Well-written and structured paper.
TTT offers a "clever" way to handle long sequences efficiently.
Potential to improve efficiency and expressiveness of SSMs.

Weaknesses:
Limited experimental setting (only language modeling). (Not addressed)
No code provided for reproducibility. (Addressed: Authors said they would add a link to the online code repo in the final version.)
Missing comparisons with relevant baselines (DeltaNet, Mamba2). (Partially addressed: Authors provided new results comparing to Mamba2)


Reviewer 7Qqf:

Strengths:
Complete evidence to assess the robustness of the method.
Interesting and original idea.
Clear and concise writing.
Properly corroborated results without over-claiming.

Weaknesses:
Applicability of TTT to "any neural network as inner model" needs more practical consideration. (Addressed: Authors said they would add a discussion about this limitation in the final version.)
Downstream task performance not reported. (Not addressed)


Reviewer i78k:

Strengths:
Extensive empirical validation.
Extensive discussion of related work.
Competitive performance.
Unifies several linear RNN architectures in a framework.

Weaknesses:
Explicit update formulas (including dimensions) for new models are missing. (Addressed: Authors mentioned that details were in the appendix and would expand the section in the final version.)
Decision: Accept (spotlight poster)