Full attention can use long contexts well, but its cost grows with the context length. Efficient alternatives such as sliding-window attention have bounded cost, but lose direct access to older tokens. TTT-E2E reframes this as a continual-learning problem: instead of retaining the entire context in an attention cache, the model compresses what it reads into selected MLP weights.
The method is end-to-end in two senses. At test time, it updates those MLP weights using the model’s real next-token loss rather than a layer-local proxy. At training time, it meta-learns an initialization that is optimized for the sequence of updates the model will perform at test time. The resulting model combines sliding-window attention for recent context with plastic MLP weights for compressed longer-term context.
The paper’s main empirical claim is that this combination keeps improving with longer language-modeling contexts like full attention, while keeping the bounded per-token cost of SWA. The trade-off is lossy recall and a more expensive meta-training procedure.
This notebook explains how that proposal becomes an actual loop. The key implementation detail is that the model does not update after every token. It predicts a chunk of \(b\) tokens with one version of the fast weights, averages their losses, and then takes one SGD step. Sliding-window attention (SWA), with window size \(k\), supplies the recent context that those temporarily stale weights do not yet contain.
We develop this mechanism from the objective to the actual loop used in the released JAX code. Prerequisites: foundations M1–M7 and NL-1.
Learning goals
After this module, you should be able to:
distinguish the layer-local KV Binding objective from TTT-E2E’s next-token objective;
trace one inner-loop chunk in the correct order: forward pass, loss, gradient, update;
distinguish the TTT mini-batch size \(b\) from the outer training batch size \(B\);
explain why the paper chooses \(k \ge b\), and what fails when \(k < b\);
identify exactly which parameters and states are carried across chunks; and
explain why the outer loop must differentiate through the inner updates.
Notation: two different batch sizes
The word batch refers to two unrelated axes in this method. Keeping them separate removes much of the confusion.
symbol
meaning
released configuration
\(T\)
tokens in one sequence
8K during pre-training; up to 128K in extension runs
\(b\)
inner-loop TTT mini-batch, or chunk; one SGD update per \(b\) tokens
1K
\(k\)
SWA window and per-layer KV-cache length
8K
\(B\)
outer-loop batch of independent sequences used for an AdamW update
configured separately as global_batch_size
\(W_i\)
per-sequence fast MLP weights after chunk \(i\)
reset to \(W_0\) for every sequence
\(\theta\)
all slow, outer-loop parameters, including the initialization \(W_0\)
shared across sequences
Changing \(B\) changes the variance and cost of the meta-gradient. Changing \(b\) changes the model’s computation and predictions. They are not interchangeable.
1. From KV Binding to the real task loss
Earlier TTT-style memories use a layer-local reconstruction objective. At layer \(l\), the basic KV Binding (KVB) loss is
The outer loop learns the key and value projections, so this proxy can be useful. But it is not the language-model loss: each layer trains its memory to reconstruct a projected value, rather than training the complete network to predict the next token.
TTT-E2E replaces all of those local objectives with one cross-entropy loss at the output of the network:
For online updates (\(b=1\)), the update is \(W_t=W_{t-1}-\eta\nabla\ell_t(W_{t-1})\). The scalable implementation uses chunks instead. For chunk \(i\),
Every prediction in chunk \(i\) is scored with \(W_{i-1}\). The update to \(W_i\) happens after those predictions, so the target tokens do not leak into their own predictions. \(W_i\) is used by the next chunk.
A small objective check
The following toy model isolates the difference between the two losses. Hold the current representation fixed and consider two possible successor tokens. The KVB gradient is determined by the current representation and its learned projections. The E2E gradient changes when the observed successor changes.
This is not an argument that KVB is useless. It shows precisely where the supervision enters: indirectly through outer-learned projections in KVB, and directly through the next-token label in E2E.
import mathimport torchimport torch.nn.functional as FV, D, H, T =16, 32, 64, 64# vocabulary, embedding, hidden width, sequence lengthETA =1.0# inner-loop learning ratedef suffix_mlp(W, e):"""A tiny stand-in for the plastic MLPs. W is batched by sequence.""" W1, b1, W2, b2 = W hidden = torch.relu(torch.bmm(e, W1) + b1[:, None])return torch.bmm(hidden, W2) + b2[:, None]def per_sequence(W0, n):"""Create n logical fast-weight copies, one for each independent sequence."""return [w.unsqueeze(0).expand(n, *w.shape) for w in W0]def init_memory(gen):return [ (torch.randn(D, H, generator=gen) / D**0.5).requires_grad_(), torch.zeros(H, requires_grad=True), (torch.randn(H, V, generator=gen) / H**0.5).requires_grad_(), torch.zeros(V, requires_grad=True), ]gen = torch.Generator().manual_seed(1)emb = torch.randn(V, D, generator=gen) *0.5W0 = init_memory(gen)theta_K = torch.randn(D, D, generator=gen) / D**0.5theta_V = torch.randn(D, V, generator=gen) / D**0.5e_prev = emb[torch.tensor([3])][:, None] # representation at the current positiondef e2e_gradient(next_token): W = per_sequence(W0, 1) logits = suffix_mlp(W, e_prev).reshape(1, V) loss = F.cross_entropy(logits, torch.tensor([next_token]))return torch.autograd.grad(loss, W)def kvb_gradient():# There is intentionally no next-token argument: Eq. 7 uses the current# layer representation, projected once as a key and once as a value. W = per_sequence(W0, 1) key, value = e_prev @ theta_K, e_prev @ theta_V loss = (suffix_mlp(W, key) - value).square().sum()return torch.autograd.grad(loss, W)flat =lambda trees: torch.cat([x.reshape(-1) for x in trees])target_a, target_b =5, 11g_kvb = flat(kvb_gradient())g_e2e_a = flat(e2e_gradient(target_a))g_e2e_b = flat(e2e_gradient(target_b))assertnot torch.allclose(g_e2e_a, g_e2e_b)print(f"KVB gradient norm (same current representation): {g_kvb.norm():.4f}")print(f"distance between E2E gradients for targets {target_a} and {target_b}: "f"{(g_e2e_a - g_e2e_b).norm():.4f}")
KVB gradient norm (same current representation): 30.4396
distance between E2E gradients for targets 5 and 11: 6.3737
For the prediction of \(x_t\), a KVB layer at position \(t-1\) constructs its update from the representation of \(x_{t-1}\); changing the candidate label \(x_t\) does not change that update. TTT-E2E first scores \(x_t\), then uses the observed \(x_t\) in cross-entropy to construct an update for future predictions.
The distinction is about the training signal, not merely the optimizer. Both methods can use gradient descent and both can be meta-learned.
2. What is actually plastic?
The released model is a Transformer with SWA in every block. Its inner loop does not update the whole Transformer.
Only the last quarter of the blocks are plastic: 3 of 12 blocks in the 125M model and 8 of 32 in the 3B model.
A plastic block contains two MLPs. The added feed_forward_prime MLP is updated by TTT; the ordinary MLP remains static as a protected store of pre-trained knowledge.
Embeddings, attention projections, normalization parameters, the language-model head, and all earlier blocks remain fixed during the inner loop.
The plastic weights start from the meta-learned \(W_0\) at the beginning of every sequence. They are not carried between independent documents.
Inner optimization is plain SGD with gradient clipping in the reported configurations. The outer optimizer is AdamW.
The model therefore has three kinds of state while reading one sequence: plastic MLP weights \(W_i\), per-layer SWA KV caches, and the inner optimizer state. Plain SGD has no momentum buffer, but the code still carries the optimizer state through the scan.
3. The inner loop in the released implementation
The first three quarters of the Transformer never change inside a sequence. The code therefore evaluates this frozen prefix over the full sequence once, using causal SWA, and then splits its hidden states into chunks of length \(b\). It scans the plastic suffix one chunk at a time. This is a safe optimization because the prefix is causal and has no fast weights.
# Code-aligned pseudocode; one independent sequenceprefix = frozen_prefix(full_sequence) # causal SWA, computed oncechunks = split((sequence, prefix), size=b)W = copy(meta_learned_W0)opt_state = inner_sgd.init(W)suffix_cache = empty_kv_caches(length=k)losses = []for token_chunk, prefix_chunk in chunks:# All b logits use the same pre-update W. The forward pass also# consumes the old cache and produces the cache for the next chunk. logits, suffix_cache = plastic_suffix( prefix_chunk, W, suffix_cache ) loss_i = mean_cross_entropy(logits, token_chunk.targets) grad_W = grad(loss_i, W) # only feed_forward_prime W, opt_state = inner_sgd.step(W, grad_W, opt_state) losses.append(loss_i) # loss was measured before updateouter_loss_for_sequence = mean(losses)
Before each chunk, the implementation combines the carried plastic parameters with a fresh view of the unchanged outer parameters. This is how it enforces that only the selected MLP weights persist as fast state.
At training time, the outer objective includes the complete sequence of inner updates. At evaluation time, the same scan adapts \(W\), but no outer update is taken. The scan also computes an SGD step after the final chunk. That step cannot change an already-recorded loss. A streaming decoder would retain it for a continuation; loss_for_sequence discards the final carry because this evaluation API returns only losses and metrics.
Where this appears in the code
The pseudocode above is a direct simplification of commit a4fc478 in the paper repository:
operation
implementation
compute frozen prefix, reshape to \(b\)-token chunks, scan suffix
Inside each suffix SWA layer, the current query tensor has shape \([b,h,d]\). The code concatenates the previous cache and current keys/values to obtain \([k+b,h,d]\), applies a \([b,k+b]\) causal sliding-window mask, and retains the newest \(k\) keys and values. The cache produced by the forward pass is carried as-is; it is not recomputed after the subsequent MLP update.
4. How \(k\) and \(b\) interact
\(b\) and \(k\) control different mechanisms:
knob
directly controls
increasing it does what?
TTT chunk \(b\)
number of losses averaged per inner SGD step
fewer sequential updates and better hardware utilization, but fast weights remain stale for longer
SWA window \(k\)
direct attention span and KV-cache length at each layer
gives each layer a longer local path, but increases attention/cache cost; it does not make SGD more frequent
While chunk \(i\) is being scored, \(W_{i-1}\) contains information only from completed chunks. Tokens already seen inside the current chunk cannot enter the fast weights until the chunk ends. SWA must cover this delay. At the final prediction in a chunk, as many as \(b-1\) newly observed tokens from that chunk may be needed. A window of \(k\ge b\) keeps all of them on a direct attention path.
The released code does not assert \(k\ge b\); it can construct a mask for \(k<b\). The inequality is a design condition from §2.3 of the paper. If \(k<b\), an early token in the current chunk can leave the attention window before the SGD update that could consolidate it. That creates a recurring coverage gap.
A concrete timeline
Let \(b=4\). Chunk 1 scores targets \(x_1,\ldots,x_4\) with \(W_0\), then produces \(W_1\). Chunk 2 scores \(x_5,\ldots,x_8\) with \(W_1\), then produces \(W_2\). Consider the last prediction in chunk 2: the model is at input \(x_7\) and predicts \(x_8\).
source of context for predicting \(x_8\)
\(x_1\)
\(x_2\)
\(x_3\)
\(x_4\)
\(x_5\)
\(x_6\)
\(x_7\)
eligible to be compressed in \(W_1\)
✓
✓
✓
✓
direct SWA path when \(k=4\)
✓
✓
✓
✓
direct SWA path when \(k=2\)
gap
✓
✓
With \(k=b=4\), \(x_5\) is still directly visible when the chunk’s last prediction is made. With \(k=2<b\), \(x_5\) is too new to be in \(W_1\) and too old to be in the local window. The next update \(W_2\) comes too late to help the prediction of \(x_8\).
The update clock
phase
operation
fast weights used for the logits
score chunk 1
predict \(x_1, x_2, x_3, x_4\) and collect their losses
\(W_0\) for all four predictions
boundary after \(x_4\)
average the four losses and take one SGD step
\(W_0 \rightarrow W_1\)
score chunk 2
predict \(x_5, x_6, x_7, x_8\) and collect their losses
\(W_1\) for all four predictions
boundary after \(x_8\)
average the four losses and take one SGD step
\(W_1 \rightarrow W_2\)
The weights change at the chunk boundary, not between tokens. SWA is what lets the four predictions inside a chunk use one another while they are all sharing the same version of \(W\).
For the paper’s defaults, \(k/b=8192/1024=8\): the attention buffer spans roughly eight chunk lengths. The extra overlap is useful; \(k\) does not have to equal \(b\). Also, \(k\) is a per-layer direct window, not a strict end-to-end receptive-field cutoff. Information can propagate farther through a stack of SWA layers, even without TTT.
Practical consequences
Larger \(b\): more tokens are processed in parallel and each gradient is averaged over more examples, but there are fewer chances to consolidate context. If \(b=T\), an update is still taken after the sequence, but no token in that sequence benefits from it.
Smaller \(b\): the fast weights are fresher, but there are more sequential gradient steps. The learning rate and initialization are coupled to this update schedule; changing \(b\) only at test time is a distribution shift.
Larger \(k\): local context improves independently of TTT. It also increases the cache and attention work. In the paper’s ablation, larger windows helped SWA, Gated DeltaNet, and TTT-E2E.
Outer batch \(B\): each sequence in \(B\) has its own \(W_i\) and caches. Gradients are averaged only for the outer AdamW update. Increasing \(B\) does not change the inner timeline above.
Long decoding: in the procedure described in §2.3.2, the model waits until \(b\) generated tokens have accumulated, then takes one TTT step on those self-generated tokens before continuing. Before the chunk fills, decoding is ordinary SWA decoding with the current \(W_i\).
5. Why the outer loop must include the inner loop
The fast weights need a good starting point. Ordinary pre-training minimizes the loss of a model that never changes:
Because every \(W_i\) depends on \(W_0\) through the earlier updates, this objective trains the initialization for two jobs: making good immediate predictions and producing useful future weights after SGD. This is the meta-learning step. See the backpropagation primer for the differentiation machinery.
The toy task below isolates this distinction. Every sequence follows its own random permutation \(x_t=\pi(x_{t-1})\). Since \(\pi\) is redrawn for each sequence, a static model cannot learn it across sequences. It must infer the mapping from the current context and store it in its fast weights.
def make_batch(n, gen):"""Generate n sequences, each with its own random transition permutation.""" pi = torch.argsort(torch.rand(n, V, generator=gen), dim=1) x = torch.zeros(n, T +1, dtype=torch.long) x[:, 0] = torch.randint(V, (n,), generator=gen)for t inrange(T): x[:, t +1] = pi.gather(1, x[:, t:t +1]).squeeze(1)return xdef inner_chunk_objective(W, inputs, targets):"""Mean chunk loss per sequence, summed across independent sequences.""" n = inputs.shape[0] logits = suffix_mlp(W, inputs) per_token = F.cross_entropy( logits.reshape(-1, V), targets.reshape(-1), reduction="none" ).reshape(n, -1)return per_token.mean(1).sum()def token_losses(emb, W0, x, b):"""Token losses with one inner SGD update per b tokens.""" n = x.shape[0] inputs, targets = emb[x[:, :-1]], x[:, 1:] W, losses = per_sequence(W0, n), []for start inrange(0, T, b): logits = suffix_mlp(W, inputs[:, start:start + b]) chunk_loss = F.cross_entropy( logits.reshape(-1, V), targets[:, start:start + b].reshape(-1), reduction="none", ).reshape(n, -1) losses.append(chunk_loss)if start + b < T:# Average within each sequence. Summing over sequences is safe because# W has a separate leading slice for every sequence. grad_W = torch.func.grad(inner_chunk_objective)( W, inputs[:, start:start + b], targets[:, start:start + b] ) W = [w - ETA * g for w, g inzip(W, grad_W)]# The repository's scan computes a final update too, although loss_for_sequence# discards that carry. We omit the unused update in this finite-loss demo.return torch.cat(losses, dim=1)def train_init(objective, b=4, steps=300, seed=0): gen = torch.Generator().manual_seed(seed) learned_emb = (torch.randn(V, D, generator=gen) *0.5).requires_grad_() learned_W0 = init_memory(gen) optimizer = torch.optim.Adam([learned_emb] + learned_W0, lr=3e-3)for _ inrange(steps): x = make_batch(64, gen)if objective =="e2e": loss = token_losses(learned_emb, learned_W0, x, b).mean()else:# One T-token chunk: no scored token sees adapted weights. loss = token_losses(learned_emb, learned_W0, x, T).mean() optimizer.zero_grad() loss.backward() optimizer.step()return learned_emb.detach(), [w.detach() for w in learned_W0]def evaluate(emb, W0, b, n=256, seed=999): x = make_batch(n, torch.Generator().manual_seed(seed)) W0 = [w.clone().requires_grad_() for w in W0]return token_losses(emb, W0, x, b).detach().mean(0)emb_naive, W_naive = train_init("naive")emb_e2e, W_e2e = train_init("e2e")chance = math.log(V)static = evaluate(emb_e2e, W_e2e, T).mean().item()naive = evaluate(emb_naive, W_naive, 4)e2e = evaluate(emb_e2e, W_e2e, 4)print(f"chance (ln {V}) {chance:.3f}")print(f"no scored token sees an update {static:.3f}")print(f"TTT from the naive init {naive.mean():.3f}")print(f"TTT from the E2E init {e2e.mean():.3f}")print("\nloss by token index first 8 last 8")print(f" TTT, naive init {naive[:8].mean():7.3f}{naive[-8:].mean():7.3f}")print(f" TTT, E2E init {e2e[:8].mean():7.3f}{e2e[-8:].mean():7.3f}")assertabs(static - chance) <0.02assert e2e.mean() <0.6* naive.mean()assert e2e[-8:].mean() <0.05assert e2e[:8].mean() >4* e2e[-8:].mean()
chance (ln 16) 2.773
no scored token sees an update 2.778
TTT from the naive init 0.888
TTT from the E2E init 0.374
loss by token index first 8 last 8
TTT, naive init 2.566 0.104
TTT, E2E init 2.267 0.000
The static loss is at chance because the transition rule is private to each sequence. Both adapted models improve later in the sequence, which means the inner updates are storing the permutation. The E2E initialization is substantially better because it was trained for the trajectory it will actually follow.
This toy model has no attention. Its purpose is to isolate the fast-weight and meta-learning mechanism, not to reproduce the paper’s architecture or reported losses.
Changing \(b\) after meta-training
The initialization above was trained with \(b=4\). Evaluating it with another \(b\) changes both the update frequency and the scale/variance of the inner gradients. The following sweep demonstrates that coupling in the toy problem.
print("test-time b mean loss")for b in (1, 2, 4, 8, 16, 64): value = evaluate(emb_e2e, W_e2e, b).mean().item() note = {1: " <- unstable with this learning rate",4: " <- meta-training value",64: " <- no scored token benefits from the update", }.get(b, "")print(f"{b:>10}{value:10.3f}{note}")sweep = {b: evaluate(emb_e2e, W_e2e, b).mean().item() for b in (1, 4, 16, 64)}assert sweep[4] < sweep[16] < sweep[64]assert sweep[1] >10
test-time b mean loss
1 53944.078 <- unstable with this learning rate
2 0.372
4 0.374 <- meta-training value
8 0.455
16 0.781
64 2.778 <- no scored token benefits from the update
Two effects are visible. Large chunks postpone useful updates, so performance approaches the static baseline. At \(b=1\), this particular initialization and learning rate are unstable because they were tuned for gradients averaged over four tokens. That is an illustration of hyper-parameter coupling, not a claim that single-token TTT must always diverge.
The paper reports the same qualitative pressure at scale: larger \(b\) significantly hurts loss, while values below 1K make hardware utilization and stability poor. It therefore uses \(b=1\)K. The paper separately chooses \(k=8\)K because larger SWA windows helped and smaller ones did not substantially improve runtime.
6. What the method buys—and what it does not
For 3B models trained on 164B tokens, TTT-E2E’s language-model loss continues to improve with longer context in a way that the efficient baselines in the paper do not. With fixed \(k\) and \(b\), prefill is linear in sequence length and decode has constant cost per token apart from the periodic TTT step; the paper reports 2.7× faster inference than full attention at 128K.
The updated weights are a lossy compressor, not an exact archive. On Needle-in-a-Haystack, full attention is much better, especially at long context. SWA offers a direct path to recent tokens; older information can influence the model through updated weights, but the method does not guarantee verbatim recall.
Training is also expensive. The outer loop differentiates through inner gradients and a long recurrent scan. The reported implementation is 3.4× slower than full-attention training at 8K, although it becomes 1.2× faster at 128K. Gradient checkpointing through time controls memory use at the cost of recomputation.
Finally, the architecture alone is not the result. In the 8K pre-training ablation, setting \(b=8\)K means no scored token benefits from TTT; that modified architecture performs almost identically to full attention. The gain comes from the adaptation procedure and the initialization trained for it.
7. Synthesis
TTT with KV Binding
TTT-E2E
inner objective
layer-local key-to-value reconstruction
one next-token cross-entropy at the network output
fast state
weights of a memory module
added MLP weights in the last quarter of Transformer blocks
update schedule
often token-wise or mini-batched
one plain-SGD step per \(b\) tokens
short-range path
depends on the hybrid architecture
SWA with per-layer window/cache \(k\)
long-range path
compressed in memory weights
compressed in plastic MLP weights
outer learning
learns projections and initialization through the memory process
learns all slow parameters and \(W_0\) through the complete inner scan
The clean mental model is local buffer plus periodic consolidation. SWA supplies a direct, finite-span path. The plastic MLPs are a recurrent, lossy state updated every \(b\) tokens. The condition \(k\ge b\) aligns the buffer with the consolidation interval so that newly read information does not fall between them.
There is no explicit decay gate on \(W\). Old KV entries leave each layer’s sliding cache, while old information in \(W\) can be overwritten indirectly by later SGD updates. That is interference, not a programmed forgetting rate.
Exit check
1. When are the \(b\) predictions in a chunk made relative to the SGD update?
All are made with \(W_{i-1}\). Their mean next-token loss produces one gradient, and only then is \(W_i\) created. Therefore, the write from chunk \(i\) first helps chunk \(i+1\).
2. Why does the paper choose \(k\ge b\)?
The fast weights do not yet contain tokens from the current chunk. A direct SWA window at least as long as the chunk keeps those tokens reachable until the boundary update can consolidate them. If \(k<b\), an early current-chunk token can be outside both stores for later predictions in that chunk.
3. Does \(k=8\)K mean the whole network has a hard 8K receptive field?
No. It is the direct window of each attention layer and the length of each layer’s cache. Stacked SWA layers can propagate information farther, and TTT adds a separate path through updated weights.
4. Which weights change at test time?
Only the added feed_forward_prime MLPs in the last quarter of blocks. Attention, embeddings, normalization, the language-model head, and the ordinary MLPs stay fixed in the inner loop.
5. What are the two meanings of “end-to-end”?
At test time, the inner objective is the final next-token loss rather than a layer-local proxy. At training time, the outer loss is differentiated through the actual sequence of inner updates, so \(W_0\) and the rest of the model are optimized for post-update behavior.
6. What resets between documents?
The plastic weights return to the shared meta-learned \(W_0\), and the SWA caches and inner optimizer state are reinitialized. The released training loop does not carry memory across independent sequences.
