End-to-End Test-Time Training for Long Context

Near the start of a 128-thousand-token book, Mira hides the brass key in the blue atlas. Much later, the model must predict “atlas,” but the sentence has already fallen outside its sliding attention window. The recent window cannot directly revisit the original words.

The obvious repair is global attention: retain every earlier token so the late prediction can inspect the original sentence directly. That preserves the clue, and it also commits the model to revisiting an ever-growing record. Treating long memory as an attention-design problem makes this trade look unavoidable: either keep direct access to the whole book or give up direct access to the early sentence. Is exact access to every old token really the only form memory can take?

The alternative starts by changing the role assigned to context. During ordinary pre-training, text is training data: prediction errors alter weights, and those weights retain useful regularities without storing a searchable copy of every sentence. The paper applies that same logic during inference. As the model reads the given context, it continues next-token training and compresses what it has read into mutable weights, while the underlying architecture remains a Transformer with sliding-window attention. ^[1][1] Introduction“We formulate long-context language modeling as a problem in continual learning rather than architecture design. Under this formulation, we only use a standard architecture -- a Transformer with sli…”Open source

For Mira’s sentence, the intended causal chain is concrete even before specifying the update rule. The sentence appears while it is still locally available. Predicting its tokens creates learning signals. Those signals alter the model’s test-time weights. After the sentence leaves the window, later predictions use those altered weights, so information extracted from the sentence can still affect “atlas.” The old tokens are no longer directly addressable; their useful structure has changed the model that reads the rest of the book.

This is why the recasting matters. Architecture design asks where to route information between token representations. Continual learning asks what the current context can teach the model before that context disappears. The second question permits persistent memory without global attention: a fixed local window handles recent detail, while mutable weights carry compressed effects of earlier text. Compression is not a lossless archive; it can preserve useful structure while leaving out exact details. ^[1][1] Introduction“We formulate long-context language modeling as a problem in continual learning rather than architecture design. Under this formulation, we only use a standard architecture -- a Transformer with sli…”Open source For Mira’s clue, later predictions can use only what the updates preserved, not search the old sentence directly.

But “the sentence changes the weights” is still only a description. Which prediction is scored, what loss is produced, and how does that one token cause a specific weight change?

Mira’s sentence passes beyond the sliding window and disappears from direct attention. A model with fixed weights can use “blue” to predict “atlas” while both are visible, but that prediction cannot alter the weights that process later text. The missing operation must turn the prediction at position $t$ into a weight change that survives after the sentence leaves the window.

The obvious attempt is to treat the context only as input: predict each token, discard the error, and continue. Test-time training instead treats each observed token $x_t$ as a label for a small training example. The model predicts $x_t$ from the preceding token using the current test-time weights $W_{t-1}$, then cross-entropy measures the mismatch as the token loss $\ell_t(W_{t-1})$. The paper performs this update sequentially at every position. ^[2][2] TTT via Next-Token Prediction“One way to give our baseline architecture some memory is to train it on the context. Similar to standard pre-training, we can predict $\hat{p}_t$ and compare it to $x_t$ at every $t = 1, \dots, T$…”Open source

Trace the atlas sentence through that rule. Before observing “atlas,” the model assigns a distribution over possible next tokens. The observed token determines the loss; the loss determines a gradient; the gradient changes the weights. Because those changed weights process later text, the local error can affect a prediction long after the original words leave the window. This is compression rather than an explicit record: the update does not append the sentence to a searchable cache. It changes many weight coordinates according to one prediction error.

A token therefore changes memory through a concrete local chain: prediction, loss, gradient, weight update. But usefulness does not follow from mutability alone. After thousands of context-dependent steps, a poor initial state can drift into weights that no longer predict well. The unresolved requirement is an initialization designed to remain effective after being changed repeatedly.

When the early “atlas” sentence changes the weights, the model that later predicts “atlas” is no longer the model that began the book. The obvious training objective evaluates every token loss at the unchanged initialization $W_0$. ^[4][4] Equation eq:wrong“\begin{equation} \label{eq:wrong} \mathcal{L}_{\texttt{naive}}(W_0; X) = \frac{1}{T}\sum_{t=1}^{T}\ell_t(W_0). \end{equation}”Open source But test-time training updates the weights after each token, so this frozen objective does not evaluate the states used during deployment. ^[3][3] Equation eq:toy“\begin{equation} \label{eq:toy} W_t = W_{t-1} - \eta\,\nabla\ell_t(W_{t-1}), \end{equation}”Open source ^[5][5] Learning to (Learn at Test Time)“For gradient-based optimization, computing $\nabla \mathcal{L}(W_0)$ for the E2E $\mathcal{L}$ entails computing gradients of gradients, since the update rule in Equation~\ref{eq:toy} itself contai…”Open source

Suppose this objective finds a $W_0$ that predicts the book well before adaptation. The first token update moves it to $W_1$, the second starts from $W_1$ and produces $W_2$, and the process continues. ^[3][3] Equation eq:toy“\begin{equation} \label{eq:toy} W_t = W_{t-1} - \eta\,\nabla\ell_t(W_{t-1}), \end{equation}”Open source None of the frozen-state losses measures whether those updated states are useful. ^[4][4] Equation eq:wrong“\begin{equation} \label{eq:wrong} \mathcal{L}_{\texttt{naive}}(W_0; X) = \frac{1}{T}\sum_{t=1}^{T}\ell_t(W_0). \end{equation}”Open source A low loss at $W_0$ therefore says little about the state that must later recover the atlas clue.

This changes what training asks of the initialization. The inner loop applies token-driven updates from $W_0$ through the later states, while the outer loop optimizes $W_0$ using the losses along that trajectory. ^[5][5] Learning to (Learn at Test Time)“For gradient-based optimization, computing $\nabla \mathcal{L}(W_0)$ for the E2E $\mathcal{L}$ entails computing gradients of gradients, since the update rule in Equation~\ref{eq:toy} itself contai…”Open source Changing $W_0$ changes the first update; that changes the next state and update; the chain eventually changes the loss on the later “atlas” target. ^[3][3] Equation eq:toy“\begin{equation} \label{eq:toy} W_t = W_{t-1} - \eta\,\nabla\ell_t(W_{t-1}), \end{equation}”Open source Because the objective measures losses at those updated states, it can favor an initialization whose adaptation leads toward lower-loss states rather than one that works only before adaptation. ^[6][6] Equation eq:toy_train“\begin{equation} \label{eq:toy_train} \mathcal{L}(W_0; X) = \frac{1}{T}\sum_{t=1}^{T}\ell_t(W_{t-1}) = \frac{1}{T}\sum_{t=1}^{T} \texttt{CE}\left(f(x_{t-1}; W_{t-1}), x_t\right). \end{equation}”Open source

To optimize $W_0$, differentiation must follow that entire chain. Each state update already contains a gradient of a token loss, so differentiating the outer loss through the update produces gradients of gradients. Standard optimizers can then use the resulting outer gradient to meta-learn the initialization. ^[5][5] Learning to (Learn at Test Time)“For gradient-based optimization, computing $\nabla \mathcal{L}(W_0)$ for the E2E $\mathcal{L}$ entails computing gradients of gradients, since the update rule in Equation~\ref{eq:toy} itself contai…”Open source

Training and deployment now agree about when losses are measured: after the relevant online updates. But agreement does not remove the cost of processing those updates in order, and a step based on one token can be unstable. The remaining problem is how to replace one-token steps without losing nearby context. ^[5][5] Learning to (Learn at Test Time)“For gradient-based optimization, computing $\nabla \mathcal{L}(W_0)$ for the E2E $\mathcal{L}$ entails computing gradients of gradients, since the update rule in Equation~\ref{eq:toy} itself contai…”Open source

Mira’s phrase can fit inside one update batch, yet “brass,” “key,” and “blue” cannot change the weights used to predict later words in that same batch. Increasing TTT batch size therefore buys parallel work by making local predictions share a stale state; the write arrives only at the batch boundary. The obvious move—just group more token losses—turns the model into a bigram within each batch: every prediction uses the pre-batch weights, so later predictions miss more preceding context. ^[7][7] Mini-Batch TTT and Sliding Window“It is important to set $k\geq b$ so our model can remember the context within each mini-batch before TTT has a chance to update its weights.”Open source

This update replaces $b$ sequential steps with one step built from $b$ token contributions. The paper uses the larger batch to improve parallelism and stability, while $b=1$ recovers token-by-token updating. ^[7][7] Mini-Batch TTT and Sliding Window“It is important to set $k\geq b$ so our model can remember the context within each mini-batch before TTT has a chance to update its weights.”Open source But batching has not solved Mira’s local context problem: until the boundary, the mutable weights contain nothing written by the current batch.

Sliding-window attention supplies the missing path. Its attention window $k$ exposes each prediction to recent tokens directly, without waiting for a gradient update. If $k<b$, the earliest words in a batch can leave the window before its final predictions; they are then absent from both the stale weights and local attention. Setting $k\ge b$ prevents that gap: while the batch is being processed, attention can still reach its earlier tokens. Once the boundary is crossed, the averaged gradient compresses that batch into the updated weights. ^[7][7] Mini-Batch TTT and Sliding Window“It is important to set $k\geq b$ so our model can remember the context within each mini-batch before TTT has a chance to update its weights.”Open source

The two paths now divide the work. Sliding-window attention preserves nearby details such as the relation among “brass key,” “blue,” and “atlas”; updated weights carry compressed information across batch and window boundaries. Mini-batching stabilizes the write, and the window covers context while that write is pending. ^[7][7] Mini-Batch TTT and Sliding Window“It is important to set $k\geq b$ so our model can remember the context within each mini-batch before TTT has a chance to update its weights.”Open source

This hybrid has local memory and persistent memory, but its gradient still needs a destination. Sending the update through more layers increases backpropagation cost, while updating attention and normalization caused outer-loop instability in the reported setup. ^[7][7] Mini-Batch TTT and Sliding Window“It is important to set $k\geq b$ so our model can remember the context within each mini-batch before TTT has a chance to update its weights.”Open source The unresolved choice is which parameters should become the persistent state.

Mira’s atlas clue must survive after its sentence leaves the attention window. Updating only a small part of the Transformer limits storage, but updating more layers requires more computation to backpropagate each test-time gradient update. The obvious choice—make the whole network mutable—therefore incurs the largest inner-loop cost. ^[9][9] Mini-Batch TTT and Sliding Window“In general, less information is lost during compression when we have a larger amount of storage. In our case, the information is the context, and the storage is the updated MLP layers. However, upd…”Open source

The first cut is by parameter type. A Transformer block contains attention, normalization, and a feed-forward multilayer perceptron, or MLP. TTT-E2E freezes embeddings, normalization parameters, and attention parameters because updating them in the inner loop made the outer loop unstable. Only the MLP weights become persistent context state. This keeps the sliding-window mechanism fixed while gradient updates write information such as the relation between “brass key” and “blue atlas” into ordinary model parameters. ^[10][10] Mini-Batch TTT and Sliding Window“We freeze the embedding layers, normalization layers, and attention layers during TTT, since updating them in the inner loop causes instability in the outer loop. Therefore, the MLP layers are the…”Open source

Updating every MLP is still the naive extreme. More mutable MLPs provide more storage, so compression can discard less context, but reaching early blocks requires backpropagating through more of the network. The paper instead updates regular MLPs in only the last quarter of Transformer blocks: fewer blocks shorten the backward path, while each selected MLP offers a larger effective state than the smaller multi-head MLPs considered in the alternative design. ^[9][9] Mini-Batch TTT and Sliding Window“In general, less information is lost during compression when we have a larger amount of storage. In our case, the information is the context, and the storage is the updated MLP layers. However, upd…”Open source ^[11][11] Alternative Derivation“And if we remove difference 4 by reverting to regular MLPs (instead of multi-head MLPs with LoRA), then we have a larger effective state at the cost of more compute (and memory).”Open source

Why the last quarter rather than a smaller slice? In a 24-block ablation, updating only one or three blocks failed to scale with context length like full attention; updating six or twelve blocks did scale, and twelve performed roughly like six. The final design therefore uses six mutable blocks in that model—the last quarter—rather than paying to update twelve. ^[12][12] Ablations on Hyper-Parameters“We experiment with updating the last 1/2, 1/4, and 1/8 of the layers. For our 760M model with a total of 24 layers, these ratios translate to the last 12, 6, and 3 layers. We also experiment with…”Open source

Mutable state creates a second risk: the atlas update could overwrite knowledge acquired during pretraining. Each updated block therefore receives a separate static MLP as safe storage for pretrained knowledge, while the mutable MLP carries context-specific changes. The network’s MLP widths are reduced so this duplication does not increase the total parameter count. ^[13][13] Mini-Batch TTT and Sliding Window“In the blocks updated during TTT, we add a static, second MLP layer as a ``safe'' storage for pre-trained knowledge. For fair comparison with the baselines, we reduce the hidden dimension of the M…”Open source

The resulting division is concrete: frozen attention retains the sliding-window mechanism, mutable MLPs in the last quarter store compressed context, and parallel static MLPs protect pretrained knowledge. ^[10][10] Mini-Batch TTT and Sliding Window“We freeze the embedding layers, normalization layers, and attention layers during TTT, since updating them in the inner loop causes instability in the outer loop. Therefore, the MLP layers are the…”Open source ^[9][9] Mini-Batch TTT and Sliding Window“In general, less information is lost during compression when we have a larger amount of storage. In our case, the information is the context, and the storage is the updated MLP layers. However, upd…”Open source ^[13][13] Mini-Batch TTT and Sliding Window“In the blocks updated during TTT, we add a static, second MLP layer as a ``safe'' storage for pre-trained knowledge. For fair comparison with the baselines, we reduce the hidden dimension of the M…”Open source That capacity may retain the atlas relation. But this design alone does not establish whether an unrelated UUID will survive compression.

Mira’s sentence creates two different memory tests. Predicting “atlas” can exploit patterns learned from the surrounding book, while reproducing an unrelated UUID requires preserving an arbitrary string exactly. The obvious hope is that low next-token loss implies both abilities. It does not: a fixed-size compressed state can retain regularities that improve prediction while discarding details that looked useless when they arrived. ^[14][14] Needle in a Haystack“From Table~\ref{tab:niah}, we observe that Transformer with full attention dramatically outperforms the other methods, including ours, especially in long context. This observation, combined with fi…”Open source

For language modeling, that compression works unusually well. On the paper’s token-index evaluation at $32\mathrm{K}$ and $128\mathrm{K}$, TTT-E2E stays below full attention in loss throughout each context; most of its aggregate advantage comes from earlier tokens. ^[15][15] Scaling with Context Length“We make the following observations from both panels of Figure~\ref{fig:token}: \vspace{-0.5ex} \begin{itemize}[itemsep=2pt, topsep=0pt, parsep=0pt, partopsep=0pt] \item TTT-E2E is the only method t…”Open source Across context lengths, the three-billion-parameter model scales like full attention while retaining constant inference latency, reaching a reported $2.7\times$ speedup over full attention at $128\mathrm{K}$ on an H100. ^[1][1] Introduction“We formulate long-context language modeling as a problem in continual learning rather than architecture design. Under this formulation, we only use a standard architecture -- a Transformer with sli…”Open source The mutable weights therefore capture enough book-level structure for earlier text to keep helping later predictions without an expanding full-attention cache.

Now ask for the UUID instead of “atlas.” Needle in a Haystack is a recall evaluation in which the target string is deliberately irrelevant to its surrounding passage. Full attention can revisit the keys and values of every previous token, so irrelevance does not force deletion. TTT-E2E instead compresses the passage into weights; an arbitrary UUID supplies little reusable next-token structure and may be omitted as noise. The paper reports that full attention dramatically outperforms TTT-E2E and the other compressed-state methods on these tasks, especially at long context. ^[14][14] Needle in a Haystack“From Table~\ref{tab:niah}, we observe that Transformer with full attention dramatically outperforms the other methods, including ours, especially in long context. This observation, combined with fi…”Open source This is not a contradiction with the loss result. Average language-model loss rewards useful regularities; exact retrieval penalizes losing even one otherwise useless detail.

The hook resolves as a trade. TTT-E2E compresses context into mutable weights through next-token prediction, and its reported inference latency stays constant with context length. ^[1][1] Introduction“We formulate long-context language modeling as a problem in continual learning rather than architecture design. Under this formulation, we only use a standard architecture -- a Transformer with sli…”Open source That compression can leave out irrelevant details such as a recall target, while full attention keeps keys and values for all previous tokens and dominates the long-context recall tests. ^[14][14] Needle in a Haystack“From Table~\ref{tab:niah}, we observe that Transformer with full attention dramatically outperforms the other methods, including ours, especially in long context. This observation, combined with fi…”Open source Efficient learned memory is not nearly lossless recall.