TTT as Linear Attention

We analytically show that most TTT formulations — even those with multi-layer MLPs and complex designs — when used for key-value binding, induce a linear attention-like operator. Below we progressively generalize from the most basic form to a very complex one.

Rewrite Basic TTT Update-Apply Rule as Linear Attention

Consider a TTT model whose inner-loop function has a linear, bias-free final layer: $f(x) = \phi(x; \Theta) W$, where $\phi(x; \Theta) \in \mathbb{R}^{D_h}$ is the hidden representation and $W \in \mathbb{R}^{D_h \times D_{out}}$ is the final layer weight.

After one gradient descent step on an objective $\mathcal{L}$ with learning rate $\eta$, updating all trainable parameters:

$(W_{t+1}, \Theta_{t+1}) = (W_t, \Theta_t) - \eta \nabla_{(W_t, \Theta_t)} \mathcal{L}(f_t(k))$

The output for any query $q$ can be written as:

$o = \phi_{t+1}(q) \left( W_t + \phi_t(k)^\top g_t(k) \right), \quad g_t(k) \triangleq -\eta \frac{\partial \mathcal{L}}{\partial f_t(k)}$

This is exactly the linear attention form $o = \hat{q}(S_0 + \hat{k}^\top \hat{v})$, where:

$\hat{q} = \phi_{t+1}(q), \quad \hat{k} = \phi_t(k), \quad \hat{v} = g_t(k), \quad S_0 = W_t$
Rewrite TTT with Sequence Inputs as Linear Attention

Given a sequence of query-key pairs $\{(q_0, k_0), (q_1, k_1), \ldots, (q_t, k_t)\}$, suppose the TTT model performs one gradient descent step per input in sequence. By repeated application of Theorem 1, the parameters after processing token $t$ are:

$(W_{t+1}, \Theta_{t+1}) = (W_0, \Theta_0) - \eta \sum_{i=0}^{t} \nabla_{(W_i, \Theta_i)} \mathcal{L}(f_i(k_i))$

Evaluating the TTT model on query $q_t$ yields:

$o_t = \phi_{t+1}(q_t) \left( W_0 + \sum_{i=0}^{t} \phi_i(k_i)^\top g_i(k_i) \right)$

This corresponds to the extended linear attention form on sequential inputs:

$o_t = \hat{q}_t \left( S_0 + \sum_{i=0}^{t} \hat{k}_i^\top \hat{v}_i \right)$
More Complex: TTT with Momentum Update as Linear Attention

Given the momentum-augmented gradient accumulator:

$(\Delta W_t, \Delta \Theta_t) = \nabla_{(W,\Theta)} \mathcal{L}(f_t(k_t)) + \alpha_t (\Delta W_{t-1}, \Delta \Theta_{t-1})$

where $\alpha_t$ is the (possibly token-dependent) momentum factor. The parameters are updated as:

$(W_{t+1}, \Theta_{t+1}) = (W_t, \Theta_t) - \eta (\Delta W_t, \Delta \Theta_t)$

Define the cumulative momentum coefficient:

$\beta_i^j \triangleq \begin{cases} \prod_{s=i+1}^{j} \alpha_s & \text{if } i \lt j \\ 1 & \text{if } i = j \end{cases}$

Unrolling and evaluating on query $q_t$ yields:

$o_t = \phi_{t+1}(q_t) \left( W_0 + \sum_{i=0}^{t} \phi_i(k_i)^\top m_i(k_i) \right)$

This induces a linear-attention form identical to Theorem 2, with the effective value being a momentum-weighted sum:

$\hat{v}_i = m_i(k_i) \triangleq g_i(k_i) \cdot \sum_{j=i}^{t} \beta_i^j$

This explains the Paradoxes: