Test-Time Training and Adaptive Inference

A TTT layer maintains a weight matrix $W_t$ that is updated at each timestep via a self-supervised gradient step. To preserve causality (no leakage from $x_t$ into the prediction at step $t$), the update at step $t$ is computed from the previous input $x_{t-1}$:

$W_t = W_{t-1} - \eta \nabla_W \ell(W_{t-1}; x_{t-1})$

where $\ell$ is a self-supervised loss (e.g., reconstruction loss, masked-prediction loss) evaluated on $x_{t-1}$. The output at step $t$ is computed using these causally-updated weights:

$y_t = f_{W_t}(x_t)$

so $y_t$ depends on $x_t$ and on $W_t$, which itself depends only on $x_1, \ldots, x_{t-1}$. (Implementations may equivalently compute $W_{t+1}$ from $x_t$ and use it to predict $x_{t+1}$; the constraint is that the update step and the prediction step use disjoint input information.)

This replaces the linear recurrence $h_t = Ah_{t-1} + Bx_t$ of SSMs with a gradient-based update: the "state" is the weight matrix $W_t$, and the "transition function" is one step of gradient descent.

In a standard SSM, the hidden state $h_t$ is a compressed summary of the context, and the compression is linear (matrix multiply). This limits what can be stored. In a TTT layer, the "state" is an entire weight matrix, and the "compression" is gradient descent on a loss function. Gradient descent is a much richer update rule: it can store key-value associations, learn patterns, and adapt to the specific distribution of the current context.

The self-supervised loss acts as the "what to remember" signal. If the loss is reconstruction ($\ell = \|W x - x\|^2$), the weights learn to reconstruct the input distribution seen so far. If it is masked prediction, the weights learn to predict missing tokens.

TTT layers achieve the expressiveness benefits of attention (content-addressable memory, ability to store and retrieve arbitrary information) with the linear-time complexity of SSMs (the state update is $O(d^2)$ per step, independent of context length). The key insight: when the inner model and the outer model share the same width $d$, the $d \times d$ weight matrix carries $O(d^2)$ real-valued parameters compared to $O(d \cdot N)$ for Mamba (with $N = 16$ typically), so the state-parameter budget scales quadratically with $d$ instead of linearly with $d$. Parameter count is not the same as information capacity in bits (that depends on precision and on how many distinguishable patterns can be recovered by the inner loss), but it sets an upper bound on retrievable associations. This budget argument is what drives TTT's advantage on long context. The comparison assumes a common $d$: in practice the outer model width can be hundreds while the TTT inner width is set much smaller (e.g., $d_{\text{inner}} = 64$), and the effective-capacity comparison against Mamba must use the specific architectural dimensions.

The gradient step at each position adds computation: you must compute $\nabla_W \ell$ at every token, which involves a forward and backward pass through the TTT layer's inner model. This is more expensive per step than Mamba's linear recurrence. The learning rate $\eta$ is critical: too large and the weights oscillate, too small and the layer does not adapt. In practice, a linear self-supervised model (TTT-Linear, where $\ell = \|Wx - x\|^2$) keeps the cost manageable while still outperforming fixed recurrences.