\documentclass{article}

\usepackage[bottom]{footmisc}
\usepackage{float}
\usepackage{sidecap}
\usepackage{multirow}
\usepackage{parskip}

\usepackage[margin=1.25in]{geometry}
\usepackage[utf8]{inputenc} % allow utf-8 input
\usepackage[T1]{fontenc}    % use 8-bit T1 fonts
\usepackage{url}            % simple URL typesetting
\usepackage{booktabs}       % professional-quality tables
\usepackage{amsfonts}       % blackboard math symbols
\usepackage{nicefrac}       % compact symbols for 1/2, etc.
\usepackage{microtype}      % microtypography
\usepackage{kpfonts}
\usepackage{graphicx}
\usepackage{multicol}
\usepackage{enumitem}

\usepackage[small]{caption}
\captionsetup{
  labelsep=period  % Changes the separator to a period
}

\usepackage[table]{xcolor}
\definecolor{codegreen}{rgb}{0,0.6,0}
\definecolor{codegray}{rgb}{0.3607843137,
0.4823529412,
0.5725490196}
\definecolor{codeblue}{rgb}{0,0.28,0.67}
\definecolor{backcolour}{rgb}{0.9882352941,
0.9725490196,
0.9294117647}

\usepackage{hyperref}
\hypersetup{
    colorlinks,
    linkcolor={codeblue},
    citecolor={codeblue},
    urlcolor={codeblue}
}

\usepackage[scaled]{beramono}

\usepackage{lipsum}
\usepackage{listings}

\lstdefinestyle{mystyle}{
    escapechar=\%,
    backgroundcolor=\color{backcolour},   
    commentstyle=\color{codegray},
    keywordstyle=\color{codeblue},
    basicstyle=\ttfamily\small,
    breakatwhitespace=false,         
    breaklines=true,                 
    captionpos=b,                    
    keepspaces=true,                 
    numbers=left,                    
    numbersep=5pt,                  
    showspaces=false,                
    showstringspaces=false,
    showtabs=false,                  
    tabsize=2
}
\lstset{style=mystyle}

\usepackage{amsthm}
\newtheorem{theorem}{Theorem}
\newtheorem{fact}{Fact}

\usepackage{setspace}

\interfootnotelinepenalty=1000000

\input{notations}

\title{\Large
Learning to (Learn at Test Time):\\
RNNs with Expressive Hidden States
\vspace{1.5ex}
}
\author{
\normalsize
Yu Sun\thanks{
\,Core contributors.
~$^\dag$\,Joint advising. ~ See author contributions at the end of the paper.\\
Correspondence to: 
\texttt{ys646@stanford.edu},
\texttt{xil202@ucsd.edu},
\texttt{kdalal@berkeley.edu}.\\
Code available in \href{https://github.com/test-time-training/ttt-lm-jax}{JAX} and \href{https://github.com/test-time-training/ttt-lm-pytorch}{PyTorch}.\\
The first version of this paper was submitted to arXiv on July 5, 2024. 
The current version contains updates on related work and limitations. 
All experiments were completed in the first version.
}$~\,^1$,
Xinhao Li$^*$$^2$, 
Karan Dalal$^*$$^3$, \\
\normalsize
Jiarui Xu$^2$, 
Arjun Vikram$^1$, 
Genghan Zhang$^1$, 
Yann Dubois$^1$, \\
\normalsize
Xinlei Chen$^\dag$$^4$,
Xiaolong Wang$^\dag$$^2$, 
Sanmi Koyejo$^\dag$$^1$, 
Tatsunori Hashimoto$^\dag$$^1$, 
Carlos Guestrin$^\dag$$^1$\\
\normalsize
$^1$\,Stanford University\quad 
~$^2$\,UC San Diego\quad
~$^3$\,UC Berkeley\quad
~$^4$\,Meta AI
\vspace{.5ex}
}
\date{}

\begin{document}

\maketitle

\begin{abstract}
\noindent
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 \emph{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.
\end{abstract}
\vspace{5ex}

\begin{figure}
    \centering
    \hspace*{0.15in}
    \includegraphics[width=0.95\textwidth]{figs/simple_teaser.pdf}
    \caption{All sequence modeling layers can be expressed as a hidden state that transitions according to an update rule.
    Our key idea is to make the hidden state itself a model $f$ with weights $W$, and the update rule a gradient step on the self-supervised loss $\ell$.
    Therefore, updating the hidden state on a test sequence is equivalent to training the model $f$ at test time. 
    This process, known as Test-Time Training (TTT), is programmed into our TTT layers.
    }
    \label{fig:ttt-layer}
\end{figure}

\clearpage

\begin{figure}[t]
    \centering
    \includegraphics[width=0.49\textwidth]{figs/teaser_2025.pdf}
    \includegraphics[width=0.49\textwidth]{figs/token_full.pdf}
    \caption{
    Comparing to Mamba, TTT-Linear and TTT-MLP have similar perplexity in 8k context (left) and better use of long context (right).
    Evaluations follow Kaplan et al.~\cite{kaplan2020scaling}.
    \textbf{Left:} 
    Scaling trends on the Pile with 8k context, zoomed in between 350M and 1.3B parameters.
    \textbf{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.
    }
    \label{fig:teaser}
\end{figure}

\section{Introduction}
\label{intro}
In 2020, the OpenAI scaling law paper (Kaplan et. al~\cite{kaplan2020scaling}) 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, we re-evaluate these findings in Figure~\ref{fig:teaser}.

On the left, we observe that Mamba~\cite{gu2023mamba}
-- one 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 Mamba as Kaplan et al. did with LSTMs.
Tokens later in a sequence should be easier to predict on average, since they condition on more information.
This is 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.

This result represents an awkward reality for existing RNNs.
On one hand, the main advantage of RNNs (vs. Transformers) is their linear (vs. quadratic) complexity. This asymptotic advantage is only realized in practice for long context, which according to Figure~\ref{fig:latency-full} is 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 nature of RNN layers:
Unlike self-attention, RNN layers have to compress context into a 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 training set into the weights of a model such as an LLM, which often exhibits deep understanding about the semantic connections among its training data -- exactly what we need from a compression heuristic.

\paragraph{TTT layers.} Motivated by this observation, we 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, these RNN layers are called \emph{Test-Time Training (TTT) layers}.
We introduce two simple instantiations: TTT-Linear and TTT-MLP, where the hidden state is a linear model and a 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.

\paragraph{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 better parallelism, we use mini-batches of tokens during TTT.
Second, we develop a dual form for operations inside each TTT mini-batch.
The dual form is equivalent in output to the naive implementation, but trains more than $5\times$ faster on our TPUs.

\paragraph{Contributions and limitations.}
The idea of using linear models as hidden states has already been well studied in DeltaNet~\cite{schlag2021linear, yang2024parallelizing}.
Since our first version was released, RNN layers with matrix (linear) hidden states have also been further advanced in Mamba 2~\cite{dao2024transformers} and Gated DeltaNet~\cite{yang2023gated}.
Compared to 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 improvements in efficiency.
It remains to be seen whether our framework can produce instantiations that either overcome this limitation or offer benefits outweighing it.

\begin{figure}
\vspace{-1ex}
    \centering
    \begin{minipage}{\textwidth}
    \centering
    \includegraphics[width=\textwidth]{figs/simple.pdf}
    \end{minipage}
    \renewcommand{\arraystretch}{1.8}
    \begin{tabular}{|l|l|l|l|l|}
    \hline
    & \textbf{Initial state}
    & \textbf{Update rule}
    & \textbf{Output rule}
    & \textbf{Cost}\\
    \hline
    \textbf{Naive RNN} & $s_0 = \texttt{vector()}$ 
    & $s_t = \sigma\left(\theta_{ss}s_{t-1} + \theta_{sx}x_t \right)$ 
    & $z_t = \theta_{zs}s_t + \theta_{zx}x_t$ 
    & $O(1)$ \\
    \textbf{Self-attention} & $s_0 = \texttt{list()}$ 
    & $s_t = s_{t-1}.\texttt{append}(k_t, v_t)$ 
    & $z_t = V_t\texttt{softmax}\left(K^T_tq_t\right)$
    & $O(t)$ \\
    \textbf{Naive TTT} & $W_0 = f.\texttt{params()}$
    & $W_t = W_{t-1} - \eta\nabla\ell(W_{t-1}; x_t)$
    & $z_t = f(x_t; W_t)$
    & $O(1)$ \\
    \hline
    \end{tabular}
    \vspace{0.5ex}
    \caption{\textbf{Top}: A generic sequence modeling layer expressed as a 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 rule and output rule.
    \textbf{Bottom}: Examples of sequence modeling layers and their instantiations of the three components. The naive TTT layer was shown in Figure~\ref{fig:ttt-layer}.
    Self-attention has a 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 of fixed size, therefore their cost per token stays constant.
    }
    \label{fig:all-layers}
\end{figure}

\section{Method}
\label{sec:method}

All sequence modeling layers can be viewed from the perspective of storing historic context into a hidden state, as shown in Figure~\ref{fig:all-layers}.\footnote{
\label{footnote:1}
We define a sequence modeling layer as an autoregressive mapping from one sequence to another.
}
For example, RNN layers -- such as LSTM~\cite{hochreiter1997long}, RWKV~\cite{peng2024eagle} and Mamba~\cite{gu2023mamba} layers -- compress context into a state of fixed size across time.
This compression has two consequences. 
On one hand, mapping an input token $x_t$ to output token $z_t$ is 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 $s_t$.
% ($s_{t-1}, x_t \mapsto s_t$)
% ($s_t, x_t \mapsto z_t$)

Self-attention can also be viewed from the perspective above, except that its hidden state, commonly known as the Key-Value (KV) cache, is a list that grows linearly with $t$.
Its update rule simply appends the current KV tuple to this list, and the output rule scans over all tuples up to $t$ 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.
However, scanning this linearly growing hidden state also takes linearly growing time per token.

To remain both efficient and expressive in long context, we need a better compression heuristic.
Specifically, we need to compress thousands or potentially millions of tokens into a hidden state that can effectively capture their underlying structures and relationships.
This might sound like a tall order, but all of us are actually already familiar with such a heuristic.

\subsection{TTT as updating a hidden state}
\label{subsec:hidden}

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~\cite{le2013building} -- exactly what we need from a compression heuristic.

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

Our key idea is to use self-supervised learning to compress the historic context $x_1,\dots,x_t$ into a hidden state $s_t$, by making the context an unlabeled dataset and the state a model.
Concretely, the hidden state $s_t$ is now equivalent to $W_t$, the weights of a model $f$, which can be a linear model, a small neural network, or anything else.
The output rule is simply:
\begin{equation}
\label{eq:output_naive}
z_t = f(x_t; W_t).
\end{equation}
Intuitively, the output token is just the prediction on $x_t$, made by $f$ with the updated weights $W_t$.
The update rule is a step of gradient descent on some self-supervised loss $\ell$: 
\begin{equation}
\label{eq:update_naive}
W_t = W_{t-1} - \eta\,\nabla\ell(W_{t-1}; x_t),
\end{equation}
with learning rate $\eta$.\footnote{\label{footnote:2}
For now, consider $W_0 = 0$. We will discuss more sophisticated techniques for initializing $W$ in Subsection~\ref{subsec:tricks}.}
From the compression point of view, every heuristic needs to decide which input to remember or forget. Our $W$ remembers inputs that produce large gradients -- intuitively, inputs that make $W$ learn a lot.

One choice of $\ell$ is reconstructing $x_t$ itself. 
To make the learning problem nontrivial, we first process $x_t$ into a corrupted input $\tilde{x}_t$ (details in Subsection~\ref{subsec:task}), then optimize:
\begin{equation}
\label{eq:recon}
\ell(W; x_t) = \| f(\tilde{x}_t; W) - x_t \|^2.
\end{equation}
Similar to denoising autoencoders~\cite{denoisingautoencoder}, $f$ needs to discover the correlations between dimensions of $x_t$ in order to reconstruct it from partial information $\tilde{x}_t$.\footnote{\label{footnote:3}
In past experiments, we have also tried adding another model $g$ (decoder) after $f$ (encoder), such that the reconstruction is produced by $g\circ f$ instead of only $f$ itself.
While this heftier design did slightly improve results, it made overall training less stable and added significant computational cost.
Therefore we focus on the encoder-only design.
}
As shown in Figure~\ref{fig:inner}, gradient descent is able to reduce $\ell$, but cannot reduce it to zero.
We discuss more sophisticated formulations of the self-supervised task in Subsection~\ref{subsec:task}.

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

\begin{figure}[t]
\vspace{-2ex}
    \centering
    \includegraphics[width=\textwidth]{figs/inner_loss.pdf}
    \caption{The self-supervised TTT loss $\ell$ averaged over all test sequences of the form $x_1,\dots,x_T$ where $T=2048$, for the first three TTT layers in a network with 125M parameters.
    One step of gradient descent is able to reduce TTT loss from $\ell(W_{t-1}; x_t)$ to $\ell(W_t; x_t)$.
    As $t$ moves further along the test sequence, $\ell(W_t; x_t)$ also improves further from $\ell(W_0; x_t)$.
    For visual clarity, loss values have been averaged over a sliding window of 10 timesteps.
    See Figure~\ref{fig:inner_full} (in Appendix) for complete results on all 12 layers.
    }
    \label{fig:inner}
\end{figure}

\subsection{Training a network with TTT layers}
\label{subsec:layer}

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

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

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


So far, the TTT layer has no outer-loop parameters, in contrast to other RNN layers and self-attention.
In Subsection~\ref{subsec:task}, we add outer-loop parameters to the TTT layer to improve its self-supervised task. 
Then in Subsection~\ref{subsec:mini-batch} and~\ref{subsec:dual}, we discuss two ways to improve the wall-clock time of TTT layers.

\input{code}

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

Concretely, we learn the self-supervised task as part of the outer loop.
Starting from the naive reconstruction task in Equation~\ref{eq:recon}, we add some outer-loop parameters to make this task learnable.
In Subsection~\ref{subsec:hidden}, we did not specify the corruption that produces $\tilde{x}_t$ from $x_t$. One design is to make it a low-rank projection $\tilde{x}_t = \theta_Kx_t$, where $\theta_K$ is a learnable matrix.\footnote{\label{footnote:4}
The subscript $K$ hints at a connection to self-attention, as we will establish in Subsection~\ref{subsec:general}.
}
Following the terminology of multi-view reconstruction, $\theta_Kx_t$ is called a \emph{training view}~\cite{chen2020improved}.

Moreover, perhaps not all the information in $x_t$ is worth remembering, so the reconstruction label can be another low-rank projection $\theta_Vx_t$ instead of $x_t$.
Here $\theta_Vx_t$ is called the \emph{label view}, where $\theta_V$ is also learnable.
In summary, our new self-supervised loss is:
\begin{equation}
\label{eq:multi}
\ell(W; x_t) = \big\| f\left(\theta_K x_t; W\right) - \theta_V x_t \big\|^2.
\end{equation}

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

Lastly, the training view $\theta_Kx_t$ has fewer dimensions than $x_t$, so we can no longer use the output rule in Equation~\ref{eq:output_naive}.
The simplest solution is to create a \emph{test view} $\theta_Qx_t$, and change our output rule to:
\begin{equation}
\label{eq:output}
z_t = f\left(\theta_Qx_t; W_t\right).
\end{equation}
This solution has an additional benefit. 
The training and label views specify the information in $x_t$ that is compressed into $W_t$ and propagated forward through time.
The test view specifies potentially different information that is mapped to the current output token $z_t$ and propagated forward through network layers, therefore adds more flexibility to the self-supervised task.

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

\subsection{Parallelization with mini-batch TTT}
\label{subsec:mini-batch}

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

We approach this systems challenge through concepts in the TTT framework. There are many variants of gradient descent (GD). 
The general update rule of GD can be expressed as:
\begin{equation}
\label{eq:gd}
W_t = W_{t-1} - \eta\,G_t = W_0 - \eta\sum_{s=1}^t G_s,
\end{equation}
where $G_t$ is the descent direction.
Note that once we have calculated $G_t$ for $t=1,\dots,T$, we can then obtain all the $W_t$s through a \texttt{cumsum} by the second half of Equation~\ref{eq:gd}.
Our naive update rule, known as \emph{online gradient descent}, 
uses \mbox{$G_t = \nabla l(W_{t-1}; x_t)$}.

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

\begin{figure}[t]
    \centering
    \includegraphics[width=0.65\textwidth]{figs/inner-batch.pdf}
    \caption{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.
    \textbf{Subsection~\ref{subsec:mini-batch}}: Since $G_1, \dots, G_b$ are not connected, they have no sequential dependency on each other, therefore can be computed in parallel.
    \textbf{Subsection~\ref{subsec:dual}}:
    We do not actually materialize the white nodes -- the intermediate $G$s and $W$s -- to compute the output variables in the dual form.
    }
    \label{fig:inner-batch}
\end{figure}

Our proposed solution -- \emph{mini-batch gradient descent} -- is shown in Figure~\ref{fig:inner-batch}.
Denote the TTT batch size by $b$. We use $G_t = \nabla \ell\left( W_{t'}; x_t \right)$, where $t' = t \,- \,$\texttt{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~\ref{fig:batch-size}.
We chose $b=16$ for all experiments in this paper.

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

% This trade-off is opposite to what people observe for chunk-wise attention~\cite{ma2024megalodon}.

\begin{figure}[t]
    \centering
    \includegraphics[width=\textwidth]{figs/bs.pdf}
    \caption[Caption for LOF]{
    Ablations on TTT mini-batch size $b$, where $b=1$ is online GD and $b=T$ is batch GD. 
    We choose $b=16$ for all experiments in this paper.
    \textbf{Left}: Smaller $b$ improves perplexity since more GD steps are taken.\protect\footnotemark\,
    The perplexity of 11.09 at $b=16$ corresponds to the final result of TTT-Linear in Figure~\ref{fig:pile}.
    \textbf{Right}: Forward time in dual form, with context length $T=2048$.
    Total time (orange) can be decomposed into time for computing the $W$s at the end of every mini-batch (blue) and time for $z_1,\dots,z_T$ (orange $-$ blue).\protect\footnotemark\,
    Time complexity for the $W$s is $O(T\times d^2)$, constant in $b$, but the blue line decreases as larger $b$ allows more parallelization until hardware utilization saturates.
    Time complexity for $z_1,\dots,z_T$ is $O(T\times b\times d)$, so the orange line first decreases with more parallelization, then increases as the extra computation for $z_1,\dots,z_T$ becomes dominant.
}
    \label{fig:batch-size}
\end{figure}
\setcounter{footnote}{\value{footnote}-1}
\footnotetext{\label{footnote:5}
    In theory, $b$ can potentially be too small such that the variance between mini-batches is too high, hurting optimization. 
    However, we have not observed such an effect in practice.
    % maybe because of the autoregressive nature of the updates.
    }
\setcounter{footnote}{\value{footnote}+1}
\footnotetext{\label{footnote:6}
    For Figure~\ref{fig:batch-size}, we use a single TTT layer in TTT-Linear 1.3B, implemented in pure PyTorch. 
    Our fused kernel significantly improves time efficiency, but makes it difficult to cleanly decompose the time for computing $W_b$ vs. $z_1,\dots,z_b$.
    }

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

Unfortunately, the TTT layer developed so far even with mini-batch still has very few \texttt{matmul}s.
Consider the simplest case of $\ell$, where $\theta_K = \theta_V = \theta_Q = I$, for only the first TTT mini-batch of size $b$.
In addition, consider $f$ as a linear model.
Copying Equation~\ref{eq:recon},
our loss at time $t$ is:
$$
\ell\left( W_0; x_t \right) = \|f(x_t; W_0) - x_t\|^2 
= \|W_0 x_t - x_t\|^2.
$$
As discussed in Subsection~\ref{subsec:mini-batch}, we can parallelize the computation of:
$$
G_t = \nabla \ell\left( W_0; x_t \right) = 2(W_0 x_t - x_t)x_t^T,
$$
for $t=1,\dots,b$.
However, we cannot compute all $b$ of the $G_t$s through a single \texttt{matmul}.
Instead, we need $b$ outer products to compute them one by one.
To make matters worse, for each $x_t\in \R^d$, $G_t$ is $d\times d$, which incurs much heavier memory footprint and I/O cost than $x_t$ for large $d$.

To solve these two problems, we make a simple observation: We do not actually need to materialize $G_1, \dots, G_b$ as long as we can compute $W_b$ at the end of the mini-batch, and the output tokens $z_1, \dots, z_b$ (see Figure~\ref{fig:inner-batch}).
Now we demonstrate these computations with the simplified TTT-Linear case above.
Denote $X = [x_1, \dots, x_b]$, then:
$$
W_b = W_0 - \eta\sum_{t=1}^b G_t
= W_0 - 2\eta\sum_{t=1}^b (W_0 x_t - x_t)x_t^T
= W_0 - 2\eta(W_0X - X)X^T.
$$
So $W_b$ can be conveniently computed with a \texttt{matmul}.
To compute $Z = [z_1, \dots, z_b]$, we know that:
\begin{equation}
\label{eq:zmat}
z_t = f(x_t; W_t) = W_tx_t = \left(W_0 - \eta \sum_{s=1}^t G_t\right)x_t
= W_0x_t - 2\eta \sum_{s=1}^t (W_0x_s - x_s)x_s^Tx_t.
\end{equation}
Denote $\delta_t = \sum_{s=1}^t (W_0x_s - x_s)x_s^Tx_t$ and the matrix $\Delta = [\delta_1, \dots, \delta_b]$.
We can derive that:
\begin{equation}
\Delta = \left( W_0X - X \right)\texttt{mask}\left( X^TX \right),
\end{equation}
where \texttt{mask} is the upper triangular mask with zeros (similar to the attention mask, but with zeros instead of infinities), and the term $W_0X - X$ can be reused from the computation of $W_b$.
Now $\Delta$ is also conveniently computed with \texttt{matmul}s. 
Plugging $\Delta$ back into Equation~\ref{eq:zmat}, we obtain $Z = W_0X - 2\eta\Delta$.

We call this procedure the \emph{dual form}, in contrast to the \emph{primal form} before this subsection, where the $G$s and $W$s are explicitly materialized.
As discussed, the two forms are equivalent in output. 
The terminology of primal and dual follows prior work that has explored similar mathematical formulations outside of TTT~\cite{irie2022dual, bishop2006pattern, rosenblatt1958perceptron}.
In Appendix~\ref{app:dual}, we show that the dual form still works when $f$ is a neural network with nonlinear layers, except with more complicated notation.

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

\subsection{Theoretical equivalences}
\label{subsec:general}

\begin{figure}
    \centering
    \includegraphics[width=0.75\textwidth]{figs/local.pdf}
    \caption{Parametric learners need to define two attributes: model and optimizer (left), and each learner uniquely induces a TTT layer (right). 
    Two of the induced TTT layers: TTT-Linear and TTT-MLP, are proposed in this paper.
    The TTT layer with a linear model and batch GD is equivalent to linear attention~\cite{katharopoulos2020transformers}.
    }
    \label{fig:local}
\end{figure}

In Subsection~\ref{subsec:hidden}, we mentioned that $f$ can be a linear model or a neural network.
In Subsection~\ref{subsec:mini-batch}, we also discussed three variants of the update rule: online GD, batch GD, and mini-batch GD.
Each of these $2\times 3$ combinations induces a different instantiation of the TTT layer, as illustrated in Figure~\ref{fig:local}.
We now show that among these induced instantiations, the TTT layer with a linear model and batch GD is equivalent to linear attention~\cite{katharopoulos2020transformers}, a widely known RNN layer.\footnote{\setstretch{1.25}
In a nutshell, linear attention~\cite{katharopoulos2020transformers} is simply self-attention without the softmax.
Recall the definition of self-attention:
$z_t = V_t\texttt{softmax}\left(K^T_tq_t\right)$.
Without \texttt{softmax}, this becomes 
$z_t = V_t\left(K^T_tq_t\right)
=\sum_{s=1}^t v_s k_s^T q_t$, which is the simplest formulation of linear attention.
Similar to other RNN layers, it can be written in a recurrent form, where $\sum_{s=1}^t v_s k_s^T$ is the hidden state.
Since $\sum_{s=1}^t v_s k_s^T$ can be computed in a \texttt{cumsum} for every $t=1,\dots,T$, linear attention also has linear complexity w.r.t. $T$.
}
\begin{theorem}
\label{theorem:la}
Consider the TTT layer with $f(x) = Wx$ as the inner-loop model, batch gradient descent with $\eta=1/2$ as the update rule, and $W_0=0$. Then, given the same input sequence $x_1,\dots,x_T$, the output rule defined in Equation~\ref{eq:output} produces the same output sequence $z_1,\dots,z_T$ as linear attention.
\end{theorem}
\begin{proof}
By definition of $\ell$ in Equation~\ref{eq:multi},
$\nabla \ell\left(W_0; x_t\right) = -2(\theta_V x_t) (\theta_K x_t)^T$.
By definition of batch GD in Equation~\ref{eq:gd}
\vspace{-2ex}:
$$
W_t = W_{t-1} - \eta\nabla \ell\left(W_0; x_t\right)  
= W_0 - \eta\sum_{s=1}^t\nabla \ell\left(W_0; x_s\right)
= \sum_{s=1}^t (\theta_V x_s) (\theta_K x_s)^T.
$$
Plugging $W_t$ into the output rule in Equation~\ref{eq:output}, we obtain the output token:
$$
z_t = f\left(\theta_Qx_t; W_t\right) = \sum_{s=1}^t (\theta_V x_s) (\theta_K x_s)^T (\theta_Qx_t),
\vspace{-1ex}
$$
which is the definition of linear attention.
\end{proof}

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

\begin{figure}
\vspace{-5ex}
    \centering
    \begin{minipage}[c]{0.60\textwidth}
    \includegraphics[width=0.92\textwidth]{figs/global.pdf}
    \end{minipage}\hfill
    \begin{minipage}[c]{0.38\textwidth}
    \vspace{2ex}
    \caption{
    RNN layers and TTT layers are both subsets of \mbox{sequence modeling layers.}
    RNN layers have a hidden state that is fixed in size across time.
    TTT layers with parametric learners are also RNN layers, since their hidden state is also fixed in size.
    TTT layers with nonparametric learners can represent self-attention, as discussed in Subsection~\ref{subsec:general}.
    }
    \label{fig:global}
    \end{minipage}
    \vspace{1ex}
\end{figure}

While the space of models $\times$ optimizers in Figure~\ref{fig:local} is already large, 
machine learning is much richer than optimizing the parameters $W_t$ of a model $f$.
There are also nonparametric learners, such as nearest neighbors, support vector machines (SVMs), and kernel ridge regression.
By definition, nonparametric learners do not have parameters $W_t$, and instead directly uses training data $x_1,\dots,x_t$. 
Hence we use the notation $f(x; x_1, \dots, x_t)$.
We now show that for a particular nonparametric learner, the induced TTT layer is equivalent to self-attention. 
\begin{theorem}
\label{theorem:sa}
Consider the TTT layer with the Nadaraya-Watson estimator~\cite{bierens1988nadaraya, cai2001weighted}, defined as:
\begin{equation}
\label{eq:nwe}
    f(x; x_1, \dots, x_t) = \frac{1}{\sum_{s=1}^t \kappa(x, x_s)} \sum_{s=1}^t \kappa(x, x_s) ~y_s,
\end{equation}
where $y_s = \theta_Vx_s$ is the label view discussed in Subsection~\ref{subsec:task}, and 
\begin{equation}
\label{eq:kernel}
\kappa\left(x, x'; \theta_K, \theta_Q\right) \propto e^{(\theta_Kx)^T \theta_Qx'}
\end{equation}
is a kernel with bandwidth hyper-parameters $\theta_K$ and $\theta_Q$.
Then given the same input sequence $x_1,\dots,x_T$, the output rule defined in Equation~\ref{eq:output} produces the same output sequence $z_1,\dots,z_T$ as self-attention.
\end{theorem}
\begin{proof}
Plugging $y_s$ and $\kappa$ above into Equation~\ref{eq:nwe} gives us the definition of self-attention.
\end{proof}
Appendix~\ref{app:theory} contains a detailed explanation of the Nadaraya-Watson estimator and kernel $\kappa$ above.
In contrast to Theorem~\ref{theorem:la}, Theorem~\ref{theorem:sa} does not produce a different implementation from attention.

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

Similar to its definition in standard machine learning packages~\cite{scikit-learn}, all learners need to implement two methods: \mbox{train and predict.}
Now we redefine the hidden state of the induced TTT layer as the internal storage of the learner, and the update and output rules as the train and predict methods.
Under this new definition of TTT layers, both parametric learners such as that in Theorem~\ref{theorem:la} and nonparametric learners such as that in Theorem~\ref{theorem:sa} can be included.
Figure~\ref{fig:global} summarizes this general definition of TTT layers in the broader scope of all sequence modeling layers.

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

\renewcommand{\arraystretch}{1.5}
\begin{table}[t]
\centering
\begin{minipage}{0.45\textwidth}
\scalebox{0.92}{
\begin{tabular}{|l|c|c|}
\hline
\textbf{Configuration} & \textbf{Ppl.} & \textbf{Diff.} \\
\hline
Linear attention~\cite{katharopoulos2020transformers} & 15.91 & - \\
Linear attn. improved   & 15.23 & $-$0.68 \\
\hline
TTT equivalence             & 15.23 & 0 \\
+ learnable $W_0$           & 15.27 & +0.04 \\
+ LN and residual in $f$    & 14.05 & $-$1.22 \\
+ mini-batch TTT            & 12.35 & $-$1.70 \\
% + learnable $\eta$ & 12.16 \\
% + learnable cidx & 11.99 \\
+ learnable $\eta$      & 11.99 & $-$0.36 \\
+ Mamba backbone          & 11.09 & $-$0.90 \\
\hline
\end{tabular}
}
\end{minipage}
\hspace{2ex}
\begin{minipage}{0.5\textwidth}
\vspace{2ex}
\caption{Ablations on improving from linear attention. 
        All models here have 125M parameters, and are trained according to the recipe in Subsection~\ref{subsec:pile}.
        The last row, with perplexity 11.09, is the final result of TTT-Linear in Figure~\ref{fig:pile}.
        Starting from the equivalence discussed in Subsection~\ref{subsec:general}, learnable $W_0$ hurts slightly, but the rows below cannot train stably without it.
        The biggest improvement comes from mini-batch TTT (changing from $b=T=2048$ to $b=16$).
        The second comes from instantiating the inner model $f$ with LN and residual connection.
        Both of these designs would be difficult to come across without the conceptual framework of TTT.
}
\label{tab:ablations}
\end{minipage}
\end{table}

\subsection{Implementation details}
\label{subsec:tricks}

\paragraph{Instantiations of $f$.} 
We propose two variants of TTT layers -- TTT-Linear and TTT-MLP, differing only in their instantiations of $f$.
For TTT-Linear, $f_{\,\texttt{lin}}(x) = Wx$, where $W$ is square.
For TTT-MLP, $f_{\,\texttt{MLP}}$ has two layers similar to the MLPs in Transformers.
Specifically, the hidden dimension is $4 \times$ the input dimension, followed by a GELU activation~\cite{hendrycks2016gaussian}.
For better stability during TTT, $f$ always contains a Layer Normalization (LN) and residual connection. That is,
$f(x) = x + \texttt{LN}(f_{\,\texttt{res}}(x)),$
where $f_{\,\texttt{res}}$ can be $f_{\,\texttt{lin}}$ or $f_{\,\texttt{MLP}}$.

\paragraph{Learnable $W_0$.} The TTT initialization $W_0$ is shared between all sequences, even though subsequent weights $W_1,\dots,W_T$ are different for each input sequence. 
Instead of setting $W_0 = 0$, we can learn it as part of the outer loop. 
Since outer-loop parameters are always denoted by $\theta$s instead of $W$s, we assign an alias $\theta_\text{init} = W_0$.
In practice,
$\theta_\text{init}$ adds a negligible amount of parameters comparing to the reconstruction views $\theta_K,\theta_Q,\theta_V$, because both its input and output are low dimensional.
Empirically, we observe that learning $W_0$ significantly improves training stability.

\paragraph{Learnable $\eta$.} The learning rate is usually the most important hyper-parameter for gradient descent, so we experiment with learning the inner-loop learning rate $\eta$ in Equation~\ref{eq:gd} as part of the outer loop.
We make $\eta$ a function of the input token (therefore different across time) for additional flexibility.
Concretely, we design $\eta(x) = \eta_\text{base}\,\sigma(\theta_\text{lr} \cdot x)$, where the learnable vector $\theta_\text{lr}$ is an outer-loop parameter, $\sigma$ is the sigmoid function, and the scalar $\eta_\text{base}$ is the base learning rate, set to 1 for TTT-Linear and 0.1 for TTT-MLP.
% Again, $\theta_\text{lr}$ only adds a negligible amount of parameters.
Alternatively, $\eta(x)$ can also be interpreted as a gate for $\nabla\ell$.

\paragraph{Backbone architecture.}
The cleanest way to integrate any RNN layer into a larger architecture would be to directly replace self-attention in a Transformer, known in this context as a backbone. 
However, existing RNNs such as Mamba~\cite{gu2023mamba} and Griffin~\cite{de2024griffin} all use a different backbone from Transformers.
Most notably, their backbone contains temporal convolutions before the RNN layers, which might help collect local information across time.
After experimenting with the Mamba backbone, we find that it also improves perplexity for TTT layers, so we incorporate it into our proposed method.
See Figure~\ref{fig:backbone} (in Appendix) for details.

\section{Experiments}
\label{sec:results}

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~\cite{geng2023easylm}, an open-source project for training and serving LLMs in JAX.
All experiments can be reproduced using the publicly available code and datasets provided at the bottom of the first page.

\paragraph{Datasets.}
Following the Mamba paper~\cite{gu2023mamba}, we perform standard experiments with 2k and 8k context lengths on the Pile~\cite{gao2020pile}, a popular dataset of documents for training open-source LLMs~\cite{black2022gpt}. 
However, the Pile contains few sequences of length greater than 8k~\cite{devries2023contextlength}.
To evaluate capabilities in long context, we also experiment with context lengths ranging from 1k to 32k in $2\times$ increments, on a subset of the Pile called Books3, which has been widely used to train LLMs in long context~\cite{liu2024world, authorsguild2023}.

\paragraph{Backbone architecture.}
As discussed in Subsection~\ref{subsec:tricks}, 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~\ref{fig:pile} and Figure~\ref{fig:books} contain TTT layers within the Transformer backbone.
When a figure contains both the Transformer backbone and Mamba backbone, we denote them by \emph{(T)} and \emph{(M)}, respectively.

\paragraph{Protocols.}
To ensure fairness to our baselines, we strictly follow the evaluation protocols in the Mamba paper when possible:
\begin{itemize}[leftmargin=*]
    \item 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, as Mamba does not follow the Transformer configurations.
    \item All models are trained with the Chinchilla recipe\footnote{\label{footnote:10}
    The Chinchilla paper is another highly influential study of empirical scaling laws~\cite{hoffmann2022training}.
    From large-scale experiments with many hyper-parameters, they observe that the compute-optimal models follow a particular training recipe.
    We only follow the Chinchilla recipe used in the Mamba paper, which may be slightly different from the original recipe in~\cite{hoffmann2022training}.
    } described in the Mamba paper and reproduced in our Appendix~\ref{app:details}.
    Our Transformer baseline, based on the Llama architecture~\cite{touvron2023llama}, also follows the baseline in the Mamba paper.
    As verification, our baselines can reproduce the numbers reported in the Mamba paper in their evaluation settings.\footnote{\label{foot:tokenizer}
    The only difference between our protocol and that in the Mamba paper is the tokenizer.
    The Mamba paper uses two different tokenizers -- GPT-2 and GPT-NeoX -- for various experiments. 
    For consistency, we adhere to a single tokenizer throughout this paper and choose the Llama tokenizer~\cite{touvron2023llama}, which is the modern state-of-the-art.
    }    
    \item We do not experiment with hybrid architectures (e.g. Griffin~\cite{de2024griffin}), because our baselines are not hybrid.
    While hybrid architectures that use both self-attention and TTT layers may improve performance, they would reduce the clarity of our academic evaluation.
\end{itemize}

\subsection{Short context: the Pile}
\label{subsec:pile}

\begin{figure}
    \centering
    \includegraphics[width=\textwidth]{figs/pile.pdf}
    \caption{Evaluations for context lengths 2k and 8k on the Pile. 
    Details in Subsection~\ref{subsec:pile}.
    TTT-Linear has comparable performance as Mamba at 2k context, and better performance at 8k.
    }
    \label{fig:pile}
\end{figure}

From Figure~\ref{fig:pile}, we make a few observations:
\begin{itemize}[leftmargin=*]
    \item At 2k context, TTT-Linear~(M), Mamba, and Transformer have comparable performance, as the lines mostly overlap. 
    TTT-MLP~(M) performs slightly worse under large FLOP budgets.
    Even though TTT-MLP has better perplexity than TTT-Linear at every model size, the extra cost in FLOPs offsets the advantage.
    \item At 8k context, both TTT-Linear~(M) and TTT-MLP~(M) perform significantly better than Mamba, in contrast to the observation at 2k. Even TTT-MLP~(T) with the Transformer backbone performs slightly better than Mamba around 1.3B.
    A robust phenomenon we observe throughout this paper is that as context length grows longer, the advantage of TTT layers over Mamba widens.
    \item At 8k context, Transformer still has good (if not the best) perplexity at every model size, but its line is not competitive because of the cost in FLOPs.
\end{itemize}

\paragraph{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 hypothesize that the temporal convolutions in the Mamba backbone 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.
\mbox{We will revisit this hypothesis in the next subsection.}

\paragraph{Lack of linear fit.}
The Chinchilla paper empirically observed that the compute-optimal models following their recipe fall onto a line in the log-log plot of FLOPs vs. perplexity, as is often the case for scaling law experiments~\cite{hoffmann2022training}.
However, we do not observe a clean linear fit in Figure~\ref{fig:pile} or Figure~\ref{fig:books} (the analogous experiments in Books), not even for Transformers.
This is not surprising given the differences in dataset, context length, tokenizer, and architecture.
Following the Mamba paper, we connect the points instead of fitting them with linear regression due to the large error.\footnote{\label{footnote:12}
Ideally, we would have rerun all the hyper-parameters and derived a potentially new recipe for each method based on our evaluation setting, following the process in the Chinchilla paper. 
If the new compute-optimal models do fall onto a line, we could then predict performance beyond the current FLOPs regime~\cite{kaplan2020scaling, hoffmann2022training}. However, this empirical study would require orders of magnitude more resources than ours.
}

\subsection{Long context: Books}
\label{subsec:books}

\begin{figure}[t!]
    \centering
    \includegraphics[width=\textwidth]{figs/books.pdf}
    \caption{Evaluations for context lengths 2k and 32k on Books. 
    Details in Subsection~\ref{subsec:books}.
    Our complete results for context lengths 1k, 2k, 4k, 8k, 16k, 32k, including Transformer finetuning, are in Figure~\ref{fig:booksscaling} (in Appendix).
    Most observations from the Pile still hold.}
    \label{fig:books}
\end{figure}

To evaluate capabilities in long context, we experiment with context lengths ranging from 1k to 32k in $2\times$ increments, using a popular subset of the Pile called Books3. 
The training recipe here is the same as that for Pile.\footnote{\label{foot:bs}
Following the Mamba paper, we always use 0.5M tokens per training batch regardless of the context length. In other words, for context length $T$ we have $0.5\text{M}\,/\,T$ sequences per batch (assume divisible).
}
From the subset of results in Figure~\ref{fig:books}, we make a few observations:
\begin{itemize}[leftmargin=*]
    \item 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).
    \item 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.
    \item TTT-MLP (T) is only slightly worse than TTT-MLP (M) at 1.3B scale. 
    As discussed, it is hard to derive an empirical scaling law due to the lack of a clean linear fit. However, the strong trend for TTT-MLP (T) suggests that the Transformer backbone might be more suitable for larger models and longer context beyond our evaluations.
\end{itemize}
We only ablate the backbones for 2k and 32k due to the cost of training LLMs. For future work, we believe that given TTT layers with even more expressive hidden states, the Mamba backbone with temporal convolutions will become unnecessary.

\paragraph{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 to train a Transformer in short context, then finetune in long context.
To reflect this practice, we add another baseline, \emph{TF finetune}, for context lengths 4k and above.
This baseline starts from the model trained (according to the Chinchilla recipe) on Books 2k, then uses 20\% more tokens to finetune at the designated context length, following the Llama Long paper~\cite{xiong2023effective}.
See details of the TF finetune recipe in Appendix~\ref{app:details}.

\paragraph{Experiments in Figure~\ref{fig:teaser} (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 expressive than the linear model:
The larger capacity of a more expressive hidden state is well-utilized in long context (therefore an advantage), but redundant in short context (therefore a disadvantage in our setting with matched FLOPs).
The Transformer in this figure is TF finetune, which is the stronger baseline in 32k context.
Details of the experiments in Figure~\ref{fig:teaser} are included in Appendix~\ref{app:details}.

Our complete results for context lengths 1k, 2k, 4k, 8k, 16k, 32k, including TF finetune, are in Figure~\ref{fig:booksscaling} (in Appendix).

\subsection{Wall-clock time}
\label{subsec:time}

\begin{figure}[t!]
\vspace{-2ex}
    \centering
    \includegraphics[width=0.95\textwidth]{figs/latency_full.pdf}
    \vspace{-1ex}
    \caption{
    Latency on an NVIDIA A100 GPU with 80G HBM and PCIe connections.
    }
    \label{fig:latency-full}
\end{figure}

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

Due to resource constraints, our experiments are written in JAX and run on TPUs. 
On a v5e-256 TPU pod, the Transformer baseline takes 0.30s per iteration of training at context 2k, while TTT-Linear takes 0.27s per iteration, already 10\% faster without any systems optimization.
However, Mamba (implemented in PyTorch, Triton, and CUDA) can only run on GPUs, so for fair comparison, we also rewrite our method into GPU kernels.
We only write inference kernels for this work because the training kernel would require substantial effort and cannot be used on our TPUs.

Figure~\ref{fig:latency-full} shows the latency of our inference kernel for forward (prefill) and generate (decode). 
All models are 1.3B (1.4B for Mamba).
As expected, time per token grows linearly for Transformer as the context length increases, but stays roughly constant for the other methods.\footnote{\label{footnote:15}
We observe that the forward latency of the network increases slightly for TTT-Linear, TTT-MLP, and Mamba, even though the latency of each sequence modeling layer alone stays constant.
Consider the operation $\theta X$, where $\theta$ is $d\times d$ and $X$ is $d\times T$.
Its latency (normalized over $T$) is expected to be constant, but in practice grows slightly with $T$.
One possible cause of this phenomenon is the GPU throttling after $T$ gets very large~\cite{horace2024}.
}
Note that our Transformer baseline is significantly faster that in the Mamba paper, because we use vLLM~\cite{kwon2023efficient}, a state-of-the-art serving system, instead of the HuggingFace Transformer~\cite{wolf2019huggingface}.

\section{Related Work}
\label{sec:rel}

\subsection{Learning at Test Time}
\label{subsec:rel-ttt}

The idea of learning at test time has a long history in machine learning.
One of the earliest versions of this idea is called local learning (Bottou and Vapnik~\cite{bottou1992local}):
For each test input, train on its neighbors before making a prediction.
This procedure has been effectively applied to models ranging from SVMs~\cite{zhang2006svm} to modern LLMs~\cite{hardt2023test}.

Another early version of learning at test time is called \emph{transductive learning}~\cite{Gammerman98learningby}.
The principle of transduction, as stated by Vladimir Vapnik~\cite{vapnik2013nature}, is to
"... get the answer that you really need, but not a more general one."
Practical implementations of transductive learning use test data to add constraints to the margin of SVMs~\cite{joachims2002learning, collobert2006large}.
% This is an example of non-parametric learning at test time, similar to our kernel estimator in Subsection~\ref{subsec:general}. 
However, transductive learning usually needs multiple test instances to be empirically effective, unlike many instantiations of test-time training, which only need a test single instance (image, video, or natural language sequence) at a time.

In computer vision, the idea of learning at test time has been applied for decades to applications such as face detection~\cite{jain2011online}, object detection~\cite{mullapudi2018online}, image super-resolution~\cite{shocher2018zero}, and 3D reconstruction~\cite{luo2020consistent}.
More recently, the same idea has also been applied to natural language processing, where it is called dynamic evaluation~\cite{krause2018dynamic, krause2019dynamic}.
The basic approach is to directly finetune a language model on the test sequence, which often comes in the form of a prompt.

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

\subsubsection{Test-Time Training}

The core idea of \emph{Test-Time Training} (TTT) is that each test instance defines its own learning problem, where this test instance alone is the target of generalization~\cite{sun2020test}.
Concretely, for each test instance $x$, the conventional practice is to predict $f(x)$, using a predictor $f$ that is optimized for all training instances on average.
TTT first formulates a learning problem defined by $x$, then trains a model $f_x$ on $x$ (often with $f$ as initialization), and predicts $f_x(x)$.

Since the test instance comes without its label, the learning problem can only be formulated with a self-supervised task. Prior work has shown that TTT with reconstruction significantly improves performance especially on outliers~\cite{ttt-mae}. 
Improvements become even more pronounced when testing on video frames that arrive in a stream and TTT is autoregressive~\cite{wang2023test}, as $f_t$ is trained on past frames $x_1, \dots, x_t$.
The autoregressive connection makes~\cite{wang2023test} most relevant to our paper.

Conceptually, the biggest difference between our paper and prior work is that our reconstruction task is learned in an outer loop, instead of handcrafted with human priors.
Follow-up work has explored applications such as 
robot manipulation~\cite{hansen2020self} and locomotion~\cite{sun2021online},
among others, that often require different designs for the self-supervised task.
In a preliminary manuscript~\cite{sun2023learning}, we explore the idea of learning to (learn at test time) from the perspective of TTT with reconstruction, with experiments in object recognition.

\subsubsection{Fast Weights and Fast Weight Programmers}
\label{subsubsec:fw}

The general idea of \emph{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~\cite{tieleman2009using}.
This idea has existed since the 1980s~\cite{feldman1982dynamic, hinton1987using, von1994correlation}.
The most relevant data often includes the test instance itself, therefore TTT can be viewed as a 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.

The idea of \emph{fast weight programmers} (FWPs) is to update the fast weights at test time with a ``slow'' model that is updated less frequently, if at all~\cite{schmidhuber1992learning}.
Our inner-loop weights $W$ can be viewed as ``fast'' and outer-loop weights $\theta$ as ``slow''.
Therefore, networks containing TTT layers can be viewed as a special case of FWPs~\cite{kirsch2021meta}, similar to how TTT can be viewed as a special case of fast weights.
Notably, one instantiation in Irie et al.~\cite{irie2021going} makes the fast weights an MLP, preceding our TTT-MLP. 

Many other modern RNN layers such as linear attention~\cite{katharopoulos2020transformers, schlag2020learning} and DeltaNet~\cite{schlag2021linear, yang2024parallelizing} are also inspired by the idea of FWPs.
Given their relevance to our work, we discuss these modern RNN layers in detail in the next subsection.
For the rest of this subsection, we briefly outline a few other forms of FWPs.
Clark et al.~\cite{clark2022meta} give a Transformer a final layer of fast weights, whose initialization is trained as slow weights.
Irie et al.~\cite{irie2022modern} design the fast weights to be programmed by themselves, which can be interpreted as recursive self-improvement.
In addition,
\cite{irie2022images} builds an image generator using the images as fast weights,
\cite{irie2022neural} applies continuous-time extensions of FWPs to time-series classification,
while 
\cite{irie2023practical} and \cite{grazzi2024unlocking} demonstrate how the choice of update rules affects the expressiveness of FWPs on formal language recognition tasks.

% The definition of FWPs is very broad.
% In fact, all networks with some gating mechanism, such as Transformers with SwiGLU blocks~\cite{shazeer2020glu}, can also be viewed as a special case of FWPs\footnote{\label{footnote:17}
% Consider a simple gate $z=\sigma(\theta x)\odot (\theta'x)$, where $x$ is the input, $z$ is the output, $\theta$ and $\theta'$ are learnable weight matrices, $\odot$ is element-wise multiplication, and $\sigma$ is the sigmoid function.
% A well known interpretation is to view $W = \texttt{diag}(\theta' x)$ as the fast weights controlled by slow weights $\theta'$, then equivalently, $z = W\sigma(\theta x)$ is simply a two-layer MLP with fast weights~\cite{gu2023mamba}.
% }.
% Recent work has been experimenting with FWPs for language modeling: 
% Our contribution relative to existing work on FWPs, again, is formulating an explicit learning problem for the update, which enables us to borrow tools from learning such as mini-batch and LN.

\subsection{Modern RNN layers}

Our baseline, Mamba~\cite{gu2023mamba}, is only one of the many recent RNN layers that inherit the  linear (matrix) hidden states of linear attention~\cite{katharopoulos2020transformers, schlag2020learning}.
Some more recent examples are RWKV~\cite{peng2023rwkv, peng2024eagle}, xLSTM~\cite{beck2024xlstm}, and Gated Linear Attention (GLA)~\cite{yang2023gated}.
The most relevant work is DeltaNet~\cite{schlag2021linear}, which is equivalent to TTT-Linear with inner-loop mini-batch size 1, without the Layer Norm and residual connection.
In \cite{yang2024parallelizing}, Yang et al. further improve the performance of DeltaNet and enable parallelized updates across tokens (in our terms, across inner loop mini-batches).
Since our first version was released, RNN layers with matrix (linear) hidden states have also been further advanced in Mamba 2~\cite{dao2024transformers} and Gated DeltaNet~\cite{yang2023gated}.

Compared to 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 improvements in efficiency.
For example, TTT-MLP is effective in terms of FLOPs, as shown in Figure~\ref{fig:teaser}. 
But the additional complexity of the MLP structure increases wall-clock time much more relative to FLOPs, as shown in Figure~\ref{fig:latency-full}.
It remains to be seen whether our framework can produce instantiations that either overcome this limitation or offer benefits outweighing it.

\subsection{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~\cite{schmidhuber1987evolutionary, bengio1990learning, thrun1998learning, lake2017building}.
In prior work such as \cite{andrychowicz2016learning}, \cite{finn2017model} and \cite{metz2018meta}, 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 it is hard to collect millions of datasets, this outer loop is hard to scale.

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 generalization across datasets.
As illustrated in Table~\ref{tab:in-out}, our outer loop is ``at the same level'' as regular training.
This makes our outer loop easier to scale.

\renewcommand{\arraystretch}{1.6}
\begin{table}[t]
\vspace{-4ex}
\centering
\scalebox{0.95}{
\begin{tabular}{|l|l|l|l|}
\hline
& \textbf{Inner loop}
& \textbf{Outer loop}
& \textbf{Subsection} \\
\hline
\textbf{Piece of data} & Token $x_t$ 
& \cellcolor{backcolour}Sequence $x_1, \dots, x_T$ 
& \multirow{4}{*}{\ref{subsec:hidden}, \ref{subsec:layer}} \\
\cline{1-3}
\textbf{Training set} & Sequence $x_1, \dots, x_T$ 
& \cellcolor{backcolour}Dataset of sequences, e.g., Books & \\
\cline{1-3}
\textbf{Objective} & Reconstruction (loss $\ell)$ 
& \cellcolor{backcolour}Next-token prediction & \\
\cline{1-3}
\multirow{3}{*}{\textbf{Parameters}} 
& \multirow{3}{*}{$W$ (weights of $f$)} 
& \cellcolor{backcolour}$\theta_{\text{rest}}$ (rest of the network) & \\
\cline{3-4}
& & $\theta_K, \theta_Q, \theta_V$ (reconstruction views) 
& \ref{subsec:task} \\
\cline{3-4}
& & $\theta_{\text{init}}$ and $ \theta_\text{lr}$ 
& \ref{subsec:tricks} \\
\hline
\end{tabular}
}
\caption{
In summary, our paper reformulates supervised learning as learning to learn, with two nested loops.
Highlighted rows of the outer loop are the same as in the regular training.
Parameters of the outer loop become hyper-parameters of the inner loop.
Intuitively, the inner loop, \textit{i.e.} TTT, is ``one level below'' regular training.
}
\label{tab:in-out}
\end{table}

\section{Discussion}
We have reformulated the canonical problem of supervised learning as learning to (learn at test time).
Our formulation produces an alternative conceptual framework for building what is traditionally known as network architectures.
We summarize our current instantiation in Table~\ref{tab:in-out}.

\paragraph{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:
\begin{itemize}[leftmargin=*]
\item \textbf{Outer-loop parameterization.}
There are many other ways to parameterize a family of multi-view reconstruction tasks, or perhaps a more general family of self-supervised tasks. It would be a big coincidence if the first one we have tried turns out to be the best.
\item \textbf{Systems optimization.} Our systems optimization in Subsection~\ref{subsec:time} has been preliminary at best, and there are many ways to improve it.
In addition, pipeline parallelism through time might allow us to process long sequences of millions of tokens on multiple devices together.
\item \textbf{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~\ref{fig:ctx}.
The advantage of TTT layers should become more pronounced in longer context.
\item \textbf{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.
\item \textbf{Multi-level learning to learn.} If $f$ itself is a self-attention layer, then by Theorem~\ref{theorem:sa} it can be interpreted as yet another inner loop nested inside the existing one. 
In this fashion, we can potentially build many levels of nested learning problems.
Versions of this idea have already been explored in \cite{irie2021going} and \cite{irie2022modern}.
\end{itemize}

\paragraph{Why do we study TTT?} First a more basic question: Why study AI? For some of us, AI is a playground to probe about the nature of human intelligence. 
Prior work often tries to model human learning with machine learning, where training is on a shuffled dataset with i.i.d. instances, and inference is on a separate test set. 
However, humans do not naturally learn with i.i.d. instances or have a train-test split.
We believe that human learning has a more promising connection with TTT, our inner loop, whose data is a potentially very long sequence with strong temporal dependencies, and any piece of data can be used for both training and testing. This is why we study TTT.

\clearpage

\section*{Author Contributions}
\label{app:authors}

\textbf{Yu Sun} started this project with Xinhao Li in November 2022, and has been working on it full-time since June 2023.
Yu proposed the conceptual framework of the project, designed mini-batch TTT and the dual form, wrote the paper with help from others, and led the daily operations of the team.

\textbf{Xinhao Li} started this project with Yu Sun in November 2022, and has been working on it full-time since then.
Xinhao and Karan co-led the development of our current codebase.
Before March 2024, Xinhao was the primary contributor to our earlier codebases that shaped this project.
Xinhao made significant contributions to the project direction in discussions.

\textbf{Karan Dalal} joined this project full-time in June 2023.
In collaboration with Xinhao, he co-led the development of our current codebase.
Karan managed the experiments in Section~\ref{sec:results}, helped write the paper, and made significant contributions to the project direction in discussions.

\textbf{Jiarui Xu} joined this project in March 2024. He led our architectural development since he joined, and made significant contributions to the project direction in discussions.

\textbf{Arjun Vikram} joined this project in September 2023. He made significant contributions to our systems optimization, as well as current and earlier codebases.

\textbf{Genghan Zhang} joined this project in January 2024. He provided critical insights and made significant improvements to our systems optimization.

\textbf{Yann Dubois} joined this project in February 2024. He proposed our current instantiation of $f$, and made significant contributions to the project direction in discussions.

\textbf{Xinlei Chen and Xiaolong Wang} have been supporting this project since November 2022, and the direction of test-time training for many years. Without their support in compute and organization, this project could not have survived its early stage.
They gave invaluable advice to our experiments.

\textbf{Sanmi Koyejo, Tatsunori Hashimoto, and Carlos Guestrin} have been supporting this project since May 2023. They gave invaluable advice to our experiments and presentation.
For example, Sanmi suggested us to focus on TTT-Linear,
Tatsu suggested the experiments in Figure~\ref{fig:teaser} (left), and Carlos outlined Section~\ref{sec:method}.

\section*{Acknowledgements}
Part of the compute for this project is generously supported by the Google TPU Research Cloud program and Hyperbolic Labs.
XW is supported, in part, by the Amazon Research Award, the Cisco Faculty Award and the Qualcomm Innovation Fellowship.
SK acknowledges support by NSF 2046795 and 2205329, NIFA award 2020-67021-32799, the Alfred P. Sloan Foundation, and Google Inc.
TH is supported by a Sony Faculty Innovation Award and a gift from Panasonic.
CG acknowledges support by the Air Force Office of Scientific Research (AFOSR), FA9550-20-1-0427, Stanford Human-Centered Artificial Intelligence (HAI) Institute, and gifts from Google and IBM.

We would like to thank Rohan Taori, Xuechen Li, Allan Zhou, Ke Chen, and Guandao Yang for many helpful discussions,
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 an 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 answering 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 to his PhD advisors,
Alexei A. Efros and Moritz Hardt, for their many insights from years ago that eventually became part of this paper.


\clearpage

\bibliographystyle{plain}
\bibliography{refs}

\clearpage
\appendix

\section{Dual Form}
\label{app:dual}

Here we derive the dual form for general MLPs of arbitrary depth, with nonlinear activations.

Without loss of generality, consider $\eta = 1$ for convenience, and consider only the first mini-batch, where $t=1,\dots,b$.
Denote:
\vspace{-1ex}
$$
\hat{x}_t = \theta_Kx_t, ~~~
y_t = \theta_Vx_t, ~~~
\bar{x}_t = \theta_Qx_t.
$$
Also denote
$\hat{X} = [\hat{x}_1, \dots, \hat{x}_b]$,
and $Y$ and $\bar{X}$ analogously.
In general, uppercase letters denote matrices whose columns are vectors denoted by the corresponding lowercase letter.

For a network with $K$ layers, denote the initial parameters in layer $k$ by $W_0^k$.
Our convention is to use superscripts for the layer and subscripts for time.

\subsection{Forward pass}
During the initial forward pass of TTT,
we denote the input to layer $k$ by
$\hat{X}^k = [\hat{x}^k_1, \dots, \hat{x}^k_b]$, with $\hat{X}^1 = \hat{X}$.
Now we write the forward pass of TTT using these notations.

For $k = 1,\dots,K$:
\begin{itemize}
    \item $Z^k = W_0^k \, \hat{X}^k$
    \item $\hat{X}^{k+1} = \sigma_k\left(Z^k\right)$
\end{itemize}
where $\sigma_k$ for $k = 1,\dots,K$ can be any element-wise operation ($\R\mapsto\R$) with derivative $\sigma'$.

Given $\hat{X}^{K+1}$, we compute the loss:
\vspace{-1ex}
$$
l = \frac{1}{2}\ell\left(W_0^1,\dots,W_0^K; X\right) = \frac{1}{2}\big\|\hat{X}^{K+1} - Y\big\|_F^2 = \sum_{t=1}^b l_t,
$$
where $l_t = \frac{1}{2}\|\hat{x}^K_t - y_t\|^2$ is the same as defined in Equation~\ref{eq:multi}, except scaled by $1/2$ for convenience.

All the operations above (except $\sigma$) are \texttt{matmul}s and \texttt{sum}s, therefore are hardware efficient.
Both the primal form and the dual form share these initial operations. 

\subsection{Primal form}
The primal form first computes $G^k_t = \nabla_{W^k_0}l_t$ for $t=1,\dots,b$, then updates $W^k_t = W_0^k - \sum_{s=1}^t G^k_s$. 
Finally, given $\bar{X}^1 = [\bar{x}_1^1,\dots,\bar{x}_b^1] = \bar{X}$, the primal form repeats the forward pass with the updated $W$s.

For $k = 1,\dots,K$:
\begin{itemize}
    \item $\bar{z}_t^k = W_t^k \, \bar{x}_t^k$, for $t=1,\dots,T$
    \item $\bar{x}_t^{k+1} = \sigma_k(\bar{z}_t^k)$, for $t=1,\dots,T$
\end{itemize}
where $\bar{X}^{K+1} = [\bar{x}_1^{k+1},\dots,\bar{x}_b^{k+1}]$ contains the output tokens.

Note that a standard backward pass only computes the sum of the gradients:
$$
\nabla_{W_0^k}l  = \sum_{t=1}^b \nabla_{W^k_0}l_t = \sum_{t=1}^b G^k_t,
$$
so the computation of the individual terms in the sum $G^k_t$ for $t=1,\dots,b$ cannot be batched together into \texttt{matmul}s.
Similarly, the forward pass in primal form uses a different $W_t$ for each $\bar{x}_t$, therefore also cannot be batched in the same way as a standard forward pass.
These non-standard passes have poor hardware efficiency.

\subsection{Dual form}
As discussed in Subsection~\ref{subsec:dual}, the goal of the dual form is to compute $\bar{X}^{K+1}$ and $W_b^1,\dots,W_b^K$ with only \texttt{matmul}s and light-weight operations such as \texttt{sum}s, $\sigma$, and $\sigma'$.
To achieve this goal, we avoid explicitly computing the intermediate variables: $G^k_t$ and $W^k_t$ for $t=1,\dots,b$.

The dual form first computes 
$\nabla_{\hat{X}^{K+1}}l = \hat{X}^{K+1} - Y$, then takes a standard backward pass.

For $k = K,\dots,1$:
\begin{itemize}
    \item $\nabla_{Z^k}l = \sigma'_k\left( Z^k \right)\odot\nabla_{\hat{X}^{k+1}}l$
    \item $\nabla_{\hat{X}^k}l = \left(W_0^k\right)^T \, \nabla_{Z^k}l$
    \item $\nabla_{W_0^k}l = \nabla_{Z^k}l \, \left(\hat{X}^k\right)^T$
\end{itemize}
where $\sigma'$ is applied element-wise, and $\odot$ is element-wise multiplication.

Now we can already compute $W_b^k = W_0^k - \nabla_{W_0^k}l$.
To compute the output tokens, we do another forward pass.

For $k = 1,\dots,K$:
\begin{itemize}
    \item $\bar{Z}^k = W_0^k\bar{X}^k - \nabla_{Z^k}l\cdot \texttt{mask}\left( \left(\hat{X}^k\right)^T\bar{X}^k \right)$
    \item $\bar{X}^{k+1} = \sigma\left(\bar{Z}^k\right)$
\end{itemize}
By the end of the forward pass, we have computed $\bar{X}^{K+1}$.

While this forward pass is non-standard, it only contains \texttt{matmul}s, \texttt{sum}s, $\sigma$, and \texttt{mask}, therefore is efficient like the standard forward pass.

\subsection{Derivation}
To derive the dual form, we show that:
$$\bar{Z}^k = W_0^k\bar{X}^k - \nabla_{Z^k}l\cdot \texttt{mask}\left( \left(\hat{X}^k\right)^T\bar{X}^k \right)$$
is the same as what would be computed in the primal form.
Specifically, we show that each column $\bar{z}_t^k$ of $\bar{Z}^k$ in the second forward pass of the dual equals to $W_t^k \, \bar{x}_t^k$ in the forward pass of the primal.
We invoke a simple fact.
\begin{fact}
Define matrices $A = [a_1,\dots,a_b]$, $Q = [q_1,\dots,q_b]$, and $V = [v_1,\dots,v_b]$.\footnote{
Our matrix $A$ would usually be denoted by $K$ in another context.
We use $A$ to avoid confusion with the layer number $K$.
}
Define $\hat{v}_t = \sum_{s=1}^t a_s^Tq_t v_s$, and $\hat{V} = [\hat{v}_1,\dots,\hat{v}_b]$, then $\hat{V} = V \cdot \texttt{mask}(A^TQ)$.
\end{fact}

Now plug $A = \hat{X}^k$, $Q = \bar{X}^k$, $V = \nabla_{Z^k}l$, and $\hat{V} = W^k\bar{X}^k - \bar{Z}^k$ into the fact above, we have shown the desired equality.

Note that the $\sigma_k$ and $\sigma'_k$ used above can be extended to arbitrary functions that are not necessarily element-wise operations, including normalization layers.
This extension can be achieved through, for example, \texttt{vjp} (vector-Jacobian product) in standard libraries for automatic differentiation such as JAX and PyTorch. 
However, the dual form cannot accelerate operations inside $\sigma$ or its \texttt{vjp}.
\clearpage

\section{Nadaraya-Watson estimator}
\label{app:theory}

\paragraph{Derivation for the Nadaraya-Watson estimator.}
Throughout this section, we use $\mathbf{x}$ to denote the input token $x$ as a random variable.
Our desired output is the corresponding output token, another random variable $\mathbf{z}$.
This is formulated as estimating the conditional expectation of $\mathbf{z}$:
$$
\E[\mathbf{z}|\mathbf{x} = x] = \int p(z|x)~z~dz = \int \frac{p(x, z)}{p(x)}~z~dz.
$$
Since the true probability distributions $p(x)$ and $p(x, z)$ are unknown, we replace them with their kernel density estimations.
Specifically, the kernel density estimation for $p(x)$ is:
$$
\hat{p}(x) = \frac{1}{n}\sum_{i=1}^n \kappa(x, x_i),
$$
where each $x_i$ is a piece of training data in general.
(Recall that for our paper, $x_i$ is specifically training data for the inner loop, \textit{i.e.} a token, which matches our notation in the main text.)

For estimating $p(x,y)$, we use the product kernel:
$$
\hat{p}(x, z) = \frac{1}{n}\sum_{i=1}^n \kappa(x, x_i)~\kappa'(z, z_i).
$$
At first sight, it seems absurd to factor the joint probability into two seemingly independent kernels.
But in this case, $\kappa'$ can actually be any $\kappa_i'$ dependent on $x_i$, since it will be integrated out. So the two kernels do not need to be independent.

Plugging in those estimations, we obtain the Nadaraya-Watson estimator:
\begin{align*}
\hat{\E}[\mathbf{z}|\mathbf{x} = x] &= \int \frac{\hat{p}(x, z)}{\hat{p}(x)}~z~dz \\
&= \frac{1}{\hat{p}(x)} \int \hat{p}(x, z)~z~dz \\
&= \frac{1}{\sum_{i=1}^n \kappa(x, x_i)} \int \sum_{i=1}^n \kappa(x, x_i)~\kappa'(z, z_i)~z~dz \\
&= \frac{1}{\sum_{i=1}^n \kappa(x, x_i)} \sum_{i=1}^n \kappa(x, x_i)~\int \kappa'(z, z_i)~z~dz \\
&= \frac{1}{\sum_{i=1}^n \kappa(x, x_i)} \sum_{i=1}^n \kappa(x, x_i) ~z_i.
\end{align*}

\paragraph{Asymmetric kernels.}
In modern days, people think of kernels as positive semi-definite, which might not be guaranteed for $\kappa$ unless $\theta_K=\theta_Q$.
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 $\kappa$ in Equation~\ref{eq:kernel} have enjoyed a long tradition:
When a kernel estimator uses $\theta_K\neq \theta_Q$, it is known as a balloon estimator~\cite{chen2017tutorial}.
Papers such as Breiman et al.~\cite{breiman1977variable} have even used $\theta_Q$ as a function of $x'$, known as sample-adaptive smoothing.

\clearpage

\begin{figure}
    \centering
    \includegraphics[width=\textwidth]{figs/backbone.pdf}
    \caption{\textbf{Left}: A residual block, the basic building block for Transformers. The sequence modeling block is instantiated into two variants: the Transformer backbone and Mamba backbone.
    \textbf{Middle}: TTT layer in the Transformer backbone. The LN before $O$ comes from NormFormer~\cite{shleifer2021normformer}.
    \textbf{Right}: TTT layer in the backbone inspired by Mamba~\cite{gu2023mamba} and Griffin~\cite{de2024griffin}. 
    Following these two architectures, $\sigma$ here is GELU~\cite{hendrycks2016gaussian}.
    To accommodate the extra parameters of the gate without changing the embedding dimension, we simply combine $\theta_K$ and $\theta_Q$ into a single projection.
    }
    \label{fig:backbone}
    \vspace{2ex}
\end{figure}

\renewcommand{\arraystretch}{1.5}
\begin{table*}
\centering
\begin{tabular}{lccccccc}
\toprule
Params. & Blocks & Embed. dim. & Heads & Train steps & Peak LR & Tokens \\
\midrule
125M & 12 & 768 & 12 & 4800 & 3e-3 & 2.5B \\
350M & 24 & 1024 & 16 & 13500 & 1.5e-3 & 7B \\
760M & 24 & 1536 & 16 & 29000 & 1.25e-3 & 15B \\
1.3B & 24 & 2048 & 32 & 50000 & 1e-3 & 26B \\
\bottomrule
\end{tabular}
\caption{Training configurations for all experiments. This table reproduces Table 12 in the \mbox{Mamba paper.}
The only difference is that the learning rate they use for Mamba and Transformer is $5\times$ the values in their Table 12, and we report the actual values ($5\times$).
Note that this table only applies to TTT-Linear, TTT-MLP, and Transformers, as Mamba does not follow the multi-head residual block structure inherited from Transformers.
}
\label{tab:configs}
\end{table*}

\section{Experiment details}
\label{app:details}
\paragraph{Architectures.}
Our Transformer strictly follows the construction in the Mamba paper, where \emph{Transformer} is called \emph{Transformer++}.
Specifically, the Transformer architecture is based on Llama~\cite{touvron2023llama}, with rotary positional encodings (RoPE)~\cite{su2023roformer}, SwiGLU MLP blocks~\cite{shazeer2020glu}, and RMSNorm~\cite{zhang2019root} 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 in~\cite{gu2023mamba}.

\paragraph{Training configurations.}
Our training configurations are in Table~\ref{tab:configs}, which simply reproduces Table~12 in the Mamba paper.
As discussed in Footnote~\ref{foot:bs}, 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:
\begin{itemize}
\item AdamW optimizer: $\beta = (0.9, 0.95)$
\item Cosine schedule: decay to end learning rate \(1e-5\)
\item Linear learning rate warmup over 10\% of the training steps
\item Weight decay: 0.1
\item Gradient clipping: 1.0
\item No Dropout
\item Mixed Precision
\end{itemize}
As discussed in Footnote~\ref{foot:tokenizer}, all models are trained using the Llama tokenizer~\cite{touvron2023llama}. 
For experiments on the Pile, this is 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~\cite{su2023roformer} $\theta=10,000$ is sub-optimal for our Transformer baseline in long context.
Starting at context length 4k, we try $\theta=500,000$ following the Llama Long paper~\cite{xiong2023effective}, and use the better perplexity for Transformer (both pretrain and finetune).

\paragraph{Transformer finetuning.}
Finetuning starts a new cosine schedule with the same optimization hyper-parameter as training from scratch, except the peak learning rate.
We try three peak learning rates for finetuning: 1e-5, 1e-4, and 1e-3, and select for the best perplexity.
We observe that 1e-4 works the best for the 125M models, while 1e-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.

\paragraph{Learning rate for TTT.}
As mentioned in Subsection~\ref{subsec:tricks}, the inner-loop base learning rate $\eta_\text{base}$ is set to 1 for TTT-Linear and 0.1 for TTT-MLP.
Our heuristic for setting $\eta_\text{base}$ is similar to how people set the outer-loop learning rate for regular training: We tried $\eta_\text{base} \in \{0.01, 0.1, 1, 10\}$ and used the largest value that does not cause instabilities.
For TTT-MLP, we use linear warmup for $\eta_\text{base}$ over 10\% of the training steps, similar to regular training. 
The number of training steps in the inner loop is $T / b$ (assume divisible).
For TTT-Linear, we tried linear warmup in the inner loop but did not observe a difference.

\paragraph{Experiments in Figure~\ref{fig:teaser} (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~\ref{tab:configs}) as Mamba 1.4B.
For TTT-Linear and TTT-MLP, matched training FLOPs also imply matched inference FLOPs.
Transformer (TF finetune) has $2.8\times$ the inference FLOPs, giving it an advantage as our baseline.
To match training FLOPs 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, so we only need to change the hidden dimension of the MLP blocks from 5504 to 5808 for TTT-Linear and 5248 for TTT-MLP.

\paragraph{Gradient checkpointing through time.} 
By default, libraries such as JAX and PyTorch save the intermediate activations during a forward pass so they can be reused during the backward pass.
However, for a TTT layer with $W$ as hidden state, this default saves $W_1,\dots,W_T$, which uses too much memory.
With TTT mini-batch and the dual form, we still need to save (assume divisible) $\kappa = T/b$ $W$s at the end of the mini-batches.
A standard technique to save memory in this scenario is gradient checkpointing~\cite{chen2016training}, which is usually applied through layers, but we apply it through time.

\clearpage
\begin{figure}
    \centering
    \includegraphics[width=\textwidth]{figs/inner_loss_full.pdf}
    \caption{
    The self-supervised TTT loss $\ell$ averaged over all test sequences of the form $x_1,\dots,x_T$ where $T=2048$, for all 12 TTT layers in a network with 125M parameters train on the Pile.
    The same network is also used for $b=1$ (online GD) in the left panel of Figure~\ref{fig:batch-size}.
    For layers in the middle, we observe that $\|x_t\|$ rises steadily, causing all three losses to rise with it.
    Even for these layers, the gap between $\ell(W_0; x_t)$ and $\ell(W_t; x_t)$ still increases with $t$ .
    For visual clarity, loss values have been averaged over a sliding window of 10 timesteps.
}
    \label{fig:inner_full}
\end{figure}

\begin{figure}
    \centering
    \includegraphics[width=\textwidth]{figs/scaling.pdf}
    \caption{Complete results on Books, presented by context lengths. 
    Figure~\ref{fig:books} in Subsection~\ref{subsec:books} presents the subset of results for context lengths 2k and 32k.
}
    \label{fig:booksscaling}
\end{figure}

\begin{figure}
    \centering
    \includegraphics[width=\textwidth]{figs/ctx.pdf}
    \caption{An alternative view of our complete results on Books, presented by model sizes, with context length as 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 at the 125M scale.
    The best context length increases for larger models (trained from scratch).
}
    \label{fig:ctx}
\end{figure}

\end{document}