• 💡 Central premise of this paper is to reimagine how LLMs handle long contexts

Goal: Formulate long-context language modeling as a continual learning problem rather than architecture design (e.g., building new types of Transformer caches or recurrent layers).

Problem statement

Self-attention requires linearly growing costs to process tokens because of nearly lossless recall.

• 🔑 Question: Can we train the model via next-token prediction at test time s.t. it can absorb information with only a constant cost per token?

Outline of paper

2. Method: Mathematically + structurally defining TTT-E2E
   1. 2.1 and 2.2 (TTT & Meta-learning): Explains the inner loop (test-time training via next-token prediction) and the outer loop (optimizing base weights during training using gradients of gradients)

Architecture setup

Model is a standard transformer, with the following tweaks:

1. Sliding window attention: instead of normal attention ($N \times N$ matrix relating every token to every other token), this model only ever attends to a recent window of words (e.g., the last 8000 tokens)
2. Frozen weights: during test-time, the embedding matrices (words → numbers), normalization layers, and attention matrices are frozen.
3. Fluid weights: the only weights that change during test-time are the MLP matrices in the last 1/4 of the transformer blocks
4. Dual-MLP trick: to make sure the model doesn't "forget" its pre-trained knowledge while learning your prompt, put two MLPs in the final blocks:
   1. MLP1 is a static, frozen matrix containing long-term pre-trained knowledge
   2. MLP2 is a dynamic matrix that gets updated live

What happens at inference (test time)?

Imagine you give the model a massive 128,000-word prompt (e.g., a whole book). Here is what happens step-by-step in the TTT-E2E setup:

1. Mini-batching the prompt: The model groups the prompt into "mini-batches" of tokens (e.g., chunks of 1000 tokens)
   1. ^(model does this instead of updating its weights after every single token)
2. Forward pass ("practicing"): The model reads the first chunk of 1000 tokens. As it reads, it uses its initial base weights from pre-training ($W_0$) to "play a game": model tries to predict token 2 from token 1, token 3 from token 2, ...
3. Inner-loop loss ("checking the answers"): The model checks its predictions against the actual (known!) next words in your prompt and calculates errors (cross-entropy loss)
4. Weight update (compressing context): The model performs standard GD live on the prompt! It calculates the gradients and updates the dynamic MLP matrices in the last 1/4 of the network.
   1. Mathematically: $W_1 = W_0 - \text{learning\_rate} \times \text{gradient\_or\_errors}$
   2. Concretely speaking: The numbers in the fluid MLP matrices are altered! The context of those first 1000 words is now baked into the new weights ($W_1$)
5. Repeating ^ through the document: The model moves to the next chunk of 1000 words, using the new weights $W_1$ to predict the next tokens.
   1. ^ It calculates the error, takes another gradient step, and creates $W_2$ → do this sequentially until model reaches the end of your 128,000-word prompt.
6. Generating the answer: Now that it has processed your entire prompt,

---

Notes from discussion

Full-attention takes $O(T)$ time Amortized $O(1)$ time -- at each token, use next-token loss as the way to get "context" into the weights of the model

Instead of doing weight updates for every token, batch it into $B$ tokens (mini-batch) and update the weights for every $B$ tokens. (Latency is kind of scuffed! )

Idea of using your training to match your use case is really cool!

2.1-2.3 (Deriving real method) 2.4 (Alternate derivation)

Highlighted sentence:

Our primary derivation starts from TTT via next-token prediction, which is E2E at test time, and focused on making it E2E at training time via meta-learning in Subsection 2.2. Our alternative derivation, on the other hand, starts from TTT-KVB, which is E2E at training time, and focused on making it E2E at test time via next-token prediction.

• TTT-naive: single next-token prediction once at training time

  • At testing time: go through the sequence, do a bunch of weight updates per prompt token, and then output the final output
  • ^ there's a disconnect between how we are training the model (single next-token prediction) and how it is being used at inference (chained continual param update next-token prediction)!
  • 🔑 Contribution of this paper: Make the loss at training time match that of test-time (inference)!
    • Core concept: Make your training lost match your test time usage

• ⚠️ Skeptical of this paper ? Can't really parallelize training (so can't scale to 100s of billions of params)

  • Is this a general issue with this type of approach to long-context models? DeltaNet also can't be parallelized
  • Mamba2 has "state-space duality" (wdk what this is lol) but it can be parallelized

Long-context training for video models?

Drifting model = contrastive learning within a feature space