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

Thank the reviewer for the thoughtful questions. Let’s first address the concerns:

Significance of the efficiency gains: We have further improved the speed of our prefill kernel, as shown in our updated Figure 13 (in Appendix). We hope the differences are now significant :) We have also implemented the training kernels for TTT-Linear and TTT-MLP, as shown in Figure 7 in the revision.

Reconstruction loss has a closed form solution in the linear case: For non-autoregressive problems, yes, but the closed form solution is too compute intensive to be practical, since it needs to solve a linear system of size d-by-d. For autoregressive problems, the hidden state (i.e. the set of parameters of the inner model) needs to be different at each timestep. Unfortunately, the closed form solution only produces a single set of parameters.

Next-token prediction as self-supervised loss: Yes, we have tried next-token prediction at small scale, the result is very slightly worse, changing from 11.09 (as reported in Table 1) to 11.13.

Was convolution through time responsible for gains in the Mamba backbone: Yes. The two main changes in the backbone are: 1) convolutions, 2) SwiGLU around the sequence modeling layer. If we only add SwiGLU, perplexity actually gets worse, changing from 11.99 (in Table 1) to 12.47. Since the backbone is inherited from Mamba and orthogonal to our contributions, we have not investigated further why convolutions are important to the backbone.

Conclusion section is weak: We have added a richer conclusion section in the revision.

We believe that the reviewer’s first three questions have already been answered as we addressed the concerns. For the reviewer’s last question about the kernels, yes, they are necessary for fair comparison. Every widely used sequence modeling layer has been written into GPU kernels, including Mamba and self-attention. Transformers without their custom kernels for self-attention would likely be an order of magnitude slower, because most of the time would be spent on memory I/O. The kernels allow the programmers to control how low-level memory is used on GPUs.

HodN

Thank the reviewer for the detailed review. Addressing the reviewer’s concerns:

Memory and runtime analysis during training: We have recently finished implementing the training kernels for TTT-Linear and TTT-MLP. Please refer to Figure 7 in the revision for a plot of training latency. A 1B TTT-Linear model is 2x faster than Mamba, and 1B TTT-MLP model is only 20% slower. In terms of memory usage (GPU DRAM), our maximum batch size per device during training is the same as that of Mamba. This is reasonable since the outer loop parameters take up most of the memory.

Connection to meta-learning: We have added a comprehensive discussion in the revision.

Second-order gradients are memory intensive: We train TTT layers in dual form, as explained in Subsection 2.5. So we do not have second-order gradients. One of our main technical contributions is deriving the dual form for TTT layers with any neural network as inner model.

Plots comparing the performance and runtime for various values of TTT mini-batch size b: We have already provided this comparison in Figure 9 in the Appendix. Sorry that we could not find enough space to include it in the main text.

Overall: We are glad that the reviewer found our paper “good” in contributions (3 out of 4), and “fair” in soundness and presentation (2 out of 4). Have we been able to adequately address the reviewer concerns? If so, would the reviewer please kindly consider adjusting the score?

Thank the reviewer for participating in discussion. Addressing the additional concerns:

“However, I believe further analysis is needed regarding the hyperparameter $b$ and its impact on the practicality of the method.”

A method is practical as long as any one set of its hyper-parameters produces useful results across a variety of experiments. Knowing the results for other non-default sets of hyper-parameters can help with understanding the method, but is not necessary to determine practicality. Throughout all of our experiments (except those in Figure 9), we have held the inner-loop mini-batch size $b=16$ constant. So the practicality of our method should be determined according to the results for $b=16$.

“Figure 9 currently compares only TTT's performance and runtime with respect to $b$, but it does not include comparisons with the existing baseline methods in the paper (Transformers or Mamba) ...... For instance, the y-axis could include dotted lines representing existing baseline methods like Mamba, showing their respective performance (e.g., perplexity), runtime (ms), and memory usage.”

The inner-loop mini-batch size $b$ is particular to the TTT approach. Transformers and Mamba do not have this hyper-parameter, so we cannot add a non-constant line for these methods to Figure 9. As discussed above, we keep $b=16$ for all experiments. Results in Section 3, such as Figure 3 and 4, compare TTT ($b=16$) with Transformers and Mamba in perplexity. Figure 6 compares them in runtime.

“Additionally, Figure 9 is missing a plot for training memory usage.”

Here is the memory usage in GB of the various methods at 125M scale, on an A100 with 80GB memory, with context length 2048, for various outer-loop mini-batch sizes. The memory footprint of TTT is comparable to Mamba and the most optimized Transformer (with FlashAttention). Sorry we cannot obtain memory usage for every inner-loop mini-batch size $b$, because our kernel is specifically written for $b=16$. As discussed above, as long as memory usage is reasonable for $b=16$, our method should be practical in this aspect.

Have we been able to adequately address your concerns? Thank you again for your participation.

Thank the reviewer for the detailed review. Addressing the reviewer’s concerns:

Adding error bars: We now have error bars for our largest experiment: 1.3B models with 32k context length on Books. We did 4 additional training runs (5 total) for each method, randomizing over model initialization and data ordering. The standard deviation is 0.01 across all methods. Specifically, the perplexities are TTT-Linear 8.82 ± 0.01 TTT-MLP 8.81 ± 0.01 Mamba 8.93 ± 0.01 Sorry that we could not obtain error bars for all experiments due to the large cost of training LLMs. We also had to exclude Transformer in the runs above due to the prohibitive cost of training at 32k context.

Evaluation on downstream tasks: The focus of our paper is on long context. However, all the tasks in 4.2.2 of 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 (usually > 3B) with trillions of tokens. Unfortunately, the cost of this evaluation exceeds our academic budget by at least 500x.

Comparison with Clark et al. 2022 and Yoshida et al. 2021: The focus of our paper is on long context, specifically models with linear complexity. The two papers above also explore the idea of training at test time, but still have the quadratic complexity of Transformers, therefore are out-of-scope for our evaluations. Meta-Learning Fast Weight Language Models (Clark et al. 2022) adds a single fast weights layer on top of a Transformer. Optimizing Hidden States in Language Models (Yoshida et al. 2021) optimizes the entire KV cache of a Transformer. We have added citation for Yoshida et al. 2021 in our related work section.

Dear authors and reviewer:

As the first author of the parallel DeltaNet paper (Yang et al., 2024), I would like to offer some constructive feedback on the review. I find TTT to be a valuable contribution with an interesting perspective and appreciate the authors' inclusion of TTT-linear and DeltaNet comparisons in the updated draft. Notably, TTT-linear with LayerNorm becomes a nonlinear RNN and is no longer equivalent to DeltaNet. As shown in Table 1, this nonlinear enhancement (through LayerNorm and residual connections in f) leads to substantial performance improvements. The authors' proposed hybrid approach – combining intra-chunk linear with inter-chunk nonlinear training strategies – offers an interesting middle ground for studying nonlinear RNNs while maintaining hardware efficiency.

To strengthen the paper further, I suggest:

• Expanding the discussion of related work from meta-learning and fast-weight programming literature, particularly in relation to the NeurIPS 2019 paper "Metalearned Neural Memory" (Sections 3.3-3.4)

• Including evaluations beyond language modeling perplexity to demonstrate broader applicability

• Adding direct experimental comparisons with DeltaNet in future work

These additions would help readers better understand TTT's contributions within the broader context of efficient RNN architectures.

Best,

Songlin

Thank the reviewer for the critical feedback.

There are three levels of ideas in our paper:

1. Learning at test time and meta-learning at training time.
2. A practical framework to do the above with any neural network as inner model.
3. TTT-Linear and TTT-MLP as instantiations.

Our claim of novelty is at level 2. All the reviewers, including Reviewer o5Ft, have acknowledged this claim. However, the reviewer’s main concern is lack of novelty at levels 1 and 3. We are confused by this concern because it diverges from our main claim.

Specifically, Reviewer o5Ft has implied that TTT-Linear is similar to DeltaNet (Yang et al. 2024). We completely agree. But TTT with linear model - the most degenerate instantiation - is merely a pedagogical data point. How should people deal with longer context that outgrows the capacity of a linear hidden state? This question is what our paper really tries to answer. For DeltaNet, this question is out of scope.

At the very least, demanding experimental comparisons with DeltaNet seems unfair because it is concurrent work. From the perspective of traditional conference cycles, DeltaNet has not yet been published (accepted to NeurIPS this December). From the perspective of arxiv, our paper was released less than a month after DeltaNet. In fact, our first version was released in 2023.

To be honest, we find zero-sum debates of academic priority harmful to the community, especially when none of the RNNs today, including ours, has truly succeeded in long context. DeltaNet is a very valuable paper. It shows how to extract the maximum potential out of a linear hidden state. We have revised our related work section to reflect that. But our paper is also valuable. It shows a practical path beyond linear hidden states.

For the minor questions, sorry that we are also confused. What’s the problem that perplexity “seems very low here on the Pile.”? And perplexity is indeed lower for 8k than 2k in most cases.

Schlag et al. (2021) is completely sequential across the time dimension. The chunk-wise parallelization in Yang et al. (2024) is necessary for the Delta rule to be practical in large models. To further illustrate our point, we compare the training throughput of a 1.3B model with 8k context using

• Schlag et al. (2021): 15042 tokens / sec
• TTT-Linear: 48545 tokens / sec

Schlag et al (2021) is 3x slower than TTT-Linear.

In terms of perplexity, we ran the DeltaNet in Schlag et al. (2021) at 2k context on the Pile, in exactly the same setting as Table 1 in our submission. This model produced NaN even after we have tried a few sets of hyper-parameters. Note that the largest experiment in Schlag et al. (2021) is on WikiText-103 with only 103M words in total. It is common for methods that work in toy settings to be unstable in larger settings. Yang et al. (2021) has significantly modified Schlag et al. (2021) to make it stable in their experiments.