In-Place TTT repurposes MLP projections as fast weights, adding test-time adaptation to LLMs

In-Place Test-Time Training

Guhao Feng, Shengjie Luo, Kai Hua, Ge Zhang, Di He, Wenhao Huang, Tianle Cai

ICLR 2026 Oral Presentation

cs.LG, cs.AI, cs.CL, stat.ML

2026-04-08

Treats each MLP's final projection as fast weights updated online at inference with a next-token-aligned objective, lifting Qwen3-4B from 74.8 to 77.0 on RULER at 128k and holding the edge at 256k extrapolation.

What problem this solves

LLMs follow a train-then-deploy script: weights freeze once training ends. Information arriving at inference can only live inside the context window, and quadratic attention cost plus position extrapolation make that window expensive to stretch. Test-Time Training (TTT) points at an alternative: designate a small subset of parameters as fast weights and update them online during inference, compressing context into the weights themselves. Existing TTT methods carry three handicaps. They are built as specialized layers that replace attention, so adopting them means pretraining from scratch. Their canonical per-token update is sequential and starves modern accelerators. And their self-supervised objective, reconstructing the current token, has no explicit tie to next-token prediction, the actual job of a language model. This paper (ByteDance Seed and Peking University, ICLR 2026 Oral) argues you can get TTT into an LLM without swapping any parts.

Method

Three design choices, one per handicap.

• Where the fast weights live: no new layer. The final projection matrix Wdown of each gated MLP block becomes the fast weight, updated in place at inference; Wup, Wgate and all attention layers stay frozen. MLPs are already viewed as key-value memory holding general pretraining knowledge, so the same matrix now doubles as an online absorber of transient context. Architecture untouched, pretrained weights intact, warm start from a checkpoint works.
• Chunk-wise updates: with attention left alone, updates need not happen per token. The sequence is cut into chunks; each chunk is first processed with the current Wdown (apply), then Wdown takes one gradient step from that chunk's activations Z and target V (update). The update delta is associative, so a parallel scan with a prefix sum computes all chunks under context parallelism, mathematically equivalent to a strictly causal sequential pass. Ablations put the optimal chunk size at 512 to 1024, well beyond the small chunks that standalone TTT layers depend on.
• The objective: instead of reconstructing the current token's representation, the target V carries future-token information, computed as Conv1D over token embeddings followed by a learned projection. A theorem in the induction-head setting shows one update with this next-token target raises the correct token's logit in expectation by λ·cnorm²·calign while leaving other logits nearly unchanged; the reconstruction target fails to move the correct logit at all. With an inner-product loss the update is closed-form, W ← W + ηVZᵀ, so one step costs no backprop.

Results

Drop-in setting: Qwen3-4B-Base goes through the same continual training (20B tokens at 32k, then 15B at 128k, with YaRN); the only variable is whether In-Place TTT is on. Average accuracy on RULER:

At 4k the baseline is slightly ahead (96.6 vs 96.1); from 8k on the TTT model leads and the gap widens with length. In the same table Llama3-8B scores 69.5 at 128k and Phi3-medium-14B 46.1, though the baseline Qwen3-4B already reaches 74.8, so that cross-model edge mostly reflects the base model plus continual training; the net TTT contribution at 128k is +2.2. LLaMA-3.1-8B and Qwen3-14B show the same pattern (+2.1 and +2.7 at 64k), and the gains survive on top of YaRN.

From scratch (500M/1.5B, 32k context, against SWA, GLA, DeltaNet and LaCT): sliding-window perplexity stays below every baseline and keeps improving up to the full 32k. At 4B (120B tokens, 8k training context) the long-context gains are largest: RULER-16k goes from 6.58 to 19.99 for the full-attention model, RULER-8k from 9.91 to 26.80 for the sliding-window model. Common-sense reasoning barely moves (HellaSwag 55.67 to 55.85). Prefill throughput and peak-memory overhead are negligible per the efficiency analysis.

Ablations on a 1.7B model: more TTT-enabled MLP layers (larger state) is monotonically better; chunk sizes 512/1024 are optimal; removing either the Conv1D or the projection hurts, with the convolution mattering most at long context and the projection at short context.

Why it matters

TTT has been an attractive idea that required rebuilding the model to use. This paper shows the capability can be added to an existing checkpoint for the price of continual pretraining, with no architectural surgery. For long-context teams it is a route orthogonal to RoPE extension and sparse attention, so it stacks. And since MLPs hold most of a Transformer's parameters, treating them as a scalable memory state hands the method a natural scaling axis that grows with the model.

Limitations

Drop-in is not free: the Qwen3-4B experiments consumed roughly 35B tokens of continual training, cheaper than pretraining from scratch but far from plug-and-play. RULER gains are a few points, and 43.9 at 256k extrapolation is not usable on its own. The theory rests on a toy induction-head setup plus two assumptions (near-orthogonal embeddings, key-query alignment); it explains direction, not final scores. Common-sense gains sit within noise, so the benefit concentrates in long context. The authors leave the loss and optimizer design space unexplored and defer integration with linear-attention and SSM backbones. There is also no direct comparison against other ways of bolting TTT layers onto a pretrained model; the drop-in advantage is argued, not benchmarked.

Terms

• TTT (Test-Time Training): updating a small parameter subset with a self-supervised objective at inference so the model adapts to the current input.
• fast weights: the parameter subset updated quickly at inference, acting as memory that evolves with context; the frozen remainder are slow weights.
• RULER: a long-context benchmark averaging retrieval and reasoning tasks across context lengths.
• YaRN: a RoPE scaling method that extends effective context beyond the pretrained length.
• Context parallelism: splitting a sequence along its length across accelerators.
• Sliding-window perplexity: perplexity on a fixed final block of tokens as preceding context grows; a falling curve means the model actually uses the context.

Source

• arXiv 2604.06169v1
• PDF
• Hugging Face

What people are saying

• Video explainer: sleep-time compute and offline memory consolidation as LLM's next scaling axis bycloud

Related papers

• Naive RL on few-step diffusion trails the adapted losses by up to 0.013 GenEval2
• Rescaling two adjacent ReLU layers fakes any sharpness, challenging the flat-minima hypothesis
• On-policy distillation tolerates 200x rollout mismatch; FP8 plus async rollouts train 1.62x faster
• DeepMind factors 68-bit semiprimes with discrete diffusion; compute doubles per extra bit
• Mechanistic Shaping turns probes into training penalties, curbing sycophancy and reward hacking
• Verified future sketches cut GPT-5.6 Sol relative MASE by 40.7%

All paper explainers