Rethinking Expressivity and Efficiency in Test-Time Training
Authors: Zeyun Zhong, Joya Chen, Manuel Martin, Frederik Diederichs, Juergen Gall, Juergen Beyerer
Organizations: Karlsruhe Institute of Technology (KIT) · Fraunhofer IOSB · National University of Singapore · Lamarr Institute for Machine Learning and Artificial Intelligence · University of Bonn
Test-Time Training (TTT) enables long-context processing via continuous weight updates during inference, but current methods struggle to balance the expressivity of per-token update dynamics with the hardware efficiency of chunk-wise approximations. We propose E2-TTT (Expressive and Efficient TTT) to bridge this gap. Under the standard approximation of taking gradients at the chunk-start weights, we derive a closed-form state transition that exactly reproduces the chunk-end fast-weight and momentum states of the per-token recurrence. This enables fully parallelized chunk-level training while preserving the temporal structure of the update rule that prior chunk-wise methods discard. We validate E2-TTT by training models up to 1.3B parameters from scratch. It performs on par with previous TTT and hybrid attention baselines in language modeling while outperforming them on in-context retrieval. Its advantage is most pronounced in length extrapolation: on the standard ``Needle in a Haystack'' passkey test, it retains over 90% accuracy at 8× the training context length. Meanwhile, E2-TTT can match the training throughput of efficient chunk-wise methods, demonstrating that it effectively reconciles expressivity with efficiency. The code is available at https://github.com/zeyun-zhong/E2-TTT.
Figures & tables
Figure 1: Token-wise TTT. At each step, the pair (kt,vt) defines a self-supervised loss whose gradient step (the “training” stage) updates Wt−1 to Wt ; the query qt then reads from Wt to produce ot .
Method
Gradient Evaluated at
State Recurrence
Output Computed from
Efficient Training
Token-wise TTT [ 51 ]
Previous token
Token-wise
Current token
✗
Mini-Batch TTT [ 51 , 8 ]
Batch start
Token-wise
Current token
✗
Chunk-TTT (LaCT) [ 66 ]
Chunk start
Chunk-wise
Chunk start
✓
E 2 -TTT (Ours)
Chunk start
Token-wise
Chunk start
✓
Table 1: Comparison of TTT update schedules. Gradient evaluated at : the fast-weight state used to compute the per-token gradients Gt . State recurrence : the granularity at which the states (Wt,Mt) evolve. Output computed from : the state applied to the query qt . E 2 -TTT matches Chunk-TTT on both efficiency-relevant columns while retaining a token-wise state recurrence.
Figure 2: E 2 -TTT token mixing block. SWA and E 2 -TTT run in parallel over shared Q,K,V and are fused by a data-dependent gate. The TTT update is modulated by input-dependent α,η,β .
Model
Attn.
Wiki.
LMB.
LMB.
PIQA
Hella.
Wino.
ARC-e
ARC-c
Avg.
window
ppl ↓
ppl ↓
acc ↑
acc ↑
acc_n ↑
acc ↑
acc ↑
acc_n ↑
↑
340M params
Transformer++*
2048
26.5
34.9
33.9
67.6
41.0
53.7
60.2
29.0
47.6
DeltaNet*
–
27.6
35.0
32.8
67.1
40.8
52.6
58.5
28.8
46.8
HQLT*
512
27.0
28.0
35.9
66.3
41.3
53.2
60.1
29.0
47.6
LaCT
512
27.0
28.2
37.1
66.8
41.6
51.4
59.0
30.0
47.7
Table 2: General language modeling experiments. We report perplexity ( ↓ ) and zero-shot accuracy ( ↑ ) for models with 340M and 1.3B parameters, all trained on 15B tokens. E 2 -TTT attains the lowest perplexity at both scales; on average zero-shot accuracy all methods here lie within about two points of one another. Results marked with * are cited from HQLT.
Model
Attn.
SWDE
SQuAD
FDA
Avg.
window
acc ↑
acc ↑
acc ↑
340M params
Transformer++*
2048
44.9
36.9
52.3
44.7
Transformer++*
1024
30.4
25.5
31.2
29.0
DeltaNet*
–
18.5
25.2
8.6
17.4
HQLT*
512
22.9
35.7
17.3
25.3
Table 3: Accuracy on real-world retrieval tasks with input truncated to 2K tokens. ∗ : from HQLT.
Figure 3: Length extrapolation of 1.3B models. We visualize the per-token loss metric [ 59 ] on six long-context benchmarks, where E 2 -TTT maintains stability while the LaCT baseline diverges.
Figure 4: Zero-shot accuracy on S-NIAH-1 and S-NIAH-2 for 1.3B models. All methods use sliding window attention with a 512-token window. E 2 -TTT maintains robust retrieval well beyond the training context, while LaCT collapses past 4K and HQLT degrades steadily across both tasks.
Model (1.3B)
Single-Doc QA
Multi-Doc QA
Summarization
Few-shot
Code
Avg.
NQA
QQA
MFQ
HQA
2WM
Mus
GvR
QMS
MNs
TRC
TQA
SSM
LCC
RBP
Mamba2
2.0
4.6
11.4
3.5
8.9
2.8
4.7
16.6
7.7
9.5
16.6
13.0
22.3
20.0
10.3
HQLT
0.8
1.0
9.3
3.3
7.9
2.0
6.7
14.5
12.4
25.8
19.9
23.1
19.8
22.8
12.1
LaCT
1.1
4.5
8.2
3.5
4.7
2.4
9.4
5.1
11.5
10.0
5.6
7.6
19.4
14.5
7.7
\text{E{}^{2}-TTT}_{\text{SwiGLU}}
2.8
4.9
14.2
5.0
9.0
3.8
10.3
16.7
12.1
34.0
16.2
20.1
24.4
23.2
14.1
\text{E{}^{2}-TTT}_{\text{MLP}}
1.7
4.4
12.4
4.4
10.4
3.1
8.8
16.6
13.6
26.5
20.3
19.3
19.5
24.6
13.3
Table 4: Evaluation on Longbench. Accuracy on 14 tasks from LongBench: Narrative QA, QasperQA, MultiField QA, HotpotQA, 2WikiMulti QA, Musique, GovReport, QMSum, MultiNews, TRec, Trivia QA, SamSum, LCC, and RepoBench-P by order.
Method
S-NIAH-1
S-NIAH-2
2K
4K
8K
16K
2K
4K
8K
SWA only
26.8
–
6.8
2.8
–
–
–
LaCT
100.0
81.0
37.4
3.0
91.4
80.6
5.6
LaCT-Matched
99.8
99.2
51.6
0.0
98.6
65.8
4.8
Chunk-averaged
25.2
10.8
6.8
3.0
30.4
18.2
7.2
\text{E{}^{2}-TTT}_{\text{SwiGLU}}
95.6
95.0
95.4
93.6
96.8
87.6
40.6
Table 5: Controlled comparison of update rules (1.3B). The last three rows share our framework and differ only in how the inner loop is aggregated. Training context is 2K tokens. Columns to the right of 2K are extrapolation.
Model
VideoMMMU
LongVideoBench
Perc.
Comp.
Adap.
Avg.
Perc.
Rel.
Avg.
Qwen3VL
–
–
–
41.9
63.1
51.5
56.9
Qwen3VL (SFT)
60.3
37.3
33.0
43.5
63.4
54.0
58.3
Qwen3VL+E 2 -TTT
60.7
37.3
33.7
43.9
63.7
54.9
59.0
Table 6: Video understanding benchmarks. Integrating E 2 -TTT into Qwen3VL-2B-Instruct and training only the TTT parameters matches full fine-tuning of the backbone on the same training data.
Appendix figures & tables10 assets
Supplementary material from the paper’s appendix.
Appendix
Model
340M
1.3B
Number of layers
24
Feedforward block multiplier
4
Total hidden size
1024
2048
Number of heads
8
16
Sequence length
2048
2240
Appendix
Table 7: Model hyperparameters.
Configuration
Wiki ppl ↓
LMB ppl ↓
Avg ↑
Averaged-factor baseline ( C=512 )
26.9
25.9
48.6
Closed-form kernel, C=1024
25.6
26.2
48.4
Closed-form kernel, C=512
25.5
25.0
49.1
Appendix
Table 8: Effect of the closed-form kernel and chunk granularity. Trained on 15B tokens; 340M parameters. The averaged-factor baseline collapses per-token βt,γt into chunk-level scalars; our kernel preserves the full per-token dynamics in closed form.
Configuration
S-NIAH-1 2K
S-NIAH-1 8K
S-NIAH-1 16K
LongBench
SWA only (no TTT branch)
26.8
6.8
2.8
11.2
TTT only (no SWA branch)
0.0
0.0
0.0
1.5
Full, w/o post-chunk norm
100.0
99.8
98.8
13.4
Full, w/o fusion gate
99.6
88.8
86.6
13.3
Full E 2 -TTT SwiGLU
95.6
95.4
93.6
14.1
Appendix
Table 9: Component ablation at the 1.3B scale on S-NIAH-1 (accuracy ↑ at three context lengths) and LongBench (average score ↑ ). The full model is trained for 15B tokens at 2K context; ablations reuse identical training data and hyperparameters.
Inner-loop update rule
Wiki ppl ↓
LMB ppl ↓
Avg ↑
ηt only (no momentum, no decay)
26.6
26.4
48.7
ηt + momentum + weight decay
25.5
25.0
49.1
Appendix
Table 10: Effect of the inner-loop optimizer terms recovered by our kernel (340M, C=512 , 15B tokens). The baseline retains only the input-dependent learning rate ηt .
Configuration
Wiki ppl ↓
LMB ppl ↓
Avg ↑
Base learning rate ηbase ( αbase=10−1 )
ηbase=10−1
28.8
30.3
46.5
ηbase=10−2 (default)
25.5
25.0
49.1
ηbase=10−3
26.1
23.0
49.1
Base decay αbase ( ηbase=10−2 )
αbase=10−1 (default)
25.5
25.0
49.1
Appendix
Table 11: Sensitivity of E 2 -TTT MLP (340M) to the base hyperparameters of the input-dependent dynamics.
Shuffled head
Wiki ppl ↓
Δ ppl
S-NIAH-1 16K ↑
— (baseline)
19.94
—
93.6
η
20.39
+0.45
6.0
β
20.05
+0.11
86.2
α
19.99
+0.05
91.8
All three
20.54
+0.60
4.0
Appendix
Table 12: Shuffling each coefficient head along the token axis on the trained 1.3B E 2 -TTT SwiGLU checkpoint. The permutation leaves the marginal distribution of each scalar unchanged and removes only its alignment with the input.
Figure 5: Training throughput on a single H100 GPU at the 340M and 1.3B scales.
Decode (ms/token) ↓
Peak mem. (GB) ↓
Model
2K
16K
2K
16K
Batch size 1
Transformer++
14.8
15.0
3.25
6.10
Titans
56.1
56.2
3.41
3.51
LaCT
36.8
36.5
3.33
3.42
\text{E{}^{2}-TTT}_{\text{MLP}}
33.8
33.6
3.43
3.52
Appendix
Table 13: Inference cost at 1.3B on a single H100 GPU, at batch sizes 1 and 8 . Decode latency is the mean per-token time over 512 steps after warm-up (median of three runs). All recurrent models are flat in both latency and memory between 2K and 16K, while attention is not.
Method
S-NIAH-1
S-NIAH-2
512
1K
2K
4K
8K
16K
512
1K
2K
4K
8K
16K
Mamba2
100.0
100.0
98.6
63.8
31.6
13.0
100.0
99.8
99.4
45.2
27.4
5.0
HQLT
100.0
90.0
72.6
50.2
37.4
25.2
100.0
80.0
57.0
37.2
15.4
4.4
LaCT
100.0
99.4
100.0
81.0
37.4
3.0
100.0
96.8
91.4
80.6
5.6
0.0
E2E-TTT
100.0
60.0
27.6
12.6
6.4
3.0
100.0
68.4
29.8
17.4
10.2
3.4
\text{E{}^{2}-TTT}_{\text{SwiGLU}}
100.0
96.8
95.6
95.0
95.4
93.6
100.0
99.0
96.8
87.6
40.6
6.0
Appendix
Table 14: Zero-shot performance on the RULER S-NIAH benchmark suite for 1.3B models. Training context is 2K tokens; columns to the right of 2K are extrapolation. \text{E{}^{2}-TTT}_{\text{SwiGLU}} retains a substantial margin over all baselines at every extrapolation length. Best result per column in bold.
Training
Train
Perplexity ↓
Commonsense
Retrieval
tokens
ctx.
Wiki.
LMB.
Avg. ↑
Avg. ↑
15B
2K
19.9
15.8
53.6
43.6
100B
4K
16.0
10.0
60.1
58.1
Appendix
Table 15: A 1.3B \text{E{}^{2}-TTT}_{\text{SwiGLU}} checkpoint trained for 100B tokens at a 4K training context, shown next to the 15B-token checkpoint used throughout the main experiments. The Commonsense and Retrieval columns average the task sets of Tables 2 and 3 ; retrieval inputs are truncated to 2K and therefore fit inside the training context of both rows.
Test-time training adapts model parameters or memory during deployment using each input or stream to improve performance under changing tasks and data.
Test-time training (TTT) enables a deployed language model to internalize incoming context on-the-fly into a small set of fast-weight parameters, rather than merely treating it as transient context for attention. In-Place TTT recently made this idea practical for large language models by repurposing the down-projection of standard MLP blocks as fast weights. However, we find that its benefit largely disappears on realistic web long-context question answering. We attribute this to In-Place TTT's uniform write schedule: every chunk is written with equal strength, so irrelevant content dilutes query-relevant information in fast-weight memory. Building on this insight, we propose Select-TTT, a model-agnostic, query-conditioned rule for In-Place TTT that scales each chunk's write strength by its relevance to the user query. This relevance signal requires no supervision and no new parameters: it is read directly from the model's frozen output head, as the probability mass that a chunk's hidden states assign to the query tokens. A temperature-scaled softmax turns these scores into learning-rate multipliers, leaving the chunk-wise, context-parallel computation of In-Place TTT intact. Select-TTT consistently improves long-context question answering on two benchmarks and two backbones. These findings highlight the value of query-conditioned fast weight writes for practical test-time training on long-context question answering, which is also potentially promising for long-horizon agent tasks involving multiple user queries over evolving contexts.
Yue Wu, Yang Zhang, Weinan Gan +4
University of Science and Technology of China China · National University of Singapore Singapore · Huawei Noah’s Ark Lab China
Next-token prediction is the self-supervised signal that trains language models, and every observed prompt token provides the same signal at test time. We study whether this signal can define the inner-loop objective for test-time training (TTT) in pretrained long-context language models. Many TTT architectures require models to be trained with test-time adaptation in mind, limiting their direct applicability to released LLM checkpoints. While recent in-place TTT methods make fast-weight adaptation possible for pretrained LLMs without redesigning the backbone, they leave a central question unresolved: what should each test-time write store? Existing recipes train the fast weight to match a learned local value proxy but they are not directly tied to the self-supervised next-token prediction signal. We introduce Test-Time Training with Next-Token Prediction (TTT-NTP), a drop-in fast-weight adaptation method for pretrained LLMs that instead supervises updates using the model's own next contextual hidden state. This makes each local write follow the same causal computation that supports next-token prediction: the value target is a pointwise linear projection of a single next-position contextual state. On RULER Full-13, averaged over 4k to 32k contexts, TTT-NTP is the only method that consistently improves the released backbone across four models spanning three families and a 0.6-8B size range, by 3.9 points on Llama-3.1-8B, 3.0 on Mistral-7B-v0.3, 4.1 on Qwen3-4B, and 2.9 on Qwen3-0.6B. On the real-world LongBench-v2 long-document QA benchmark, TTT-NTP improves over the base model by 5.6 points on Llama-3.1-8B and 3.7 on Mistral-7B-v0.3, while preserving commonsense and knowledge performance. Our code is publicly available at https://github.com/yancyou/TTT-NTP.
Recent Test-Time Training (TTT) architectures compress context into fast weights that are updated online and queried as memory. Existing TTT designs keep this memory private to each layer: it recurs only over time, and depth merely indexes L separate memories. We argue that memory ownership need not be tied to depth, and introduce Universal Test-Time Training (uTTT), in which all layers read and write one shared memory while retaining layer-specific backbone parameters. The shared memory thus recurs over two dimensions, time and depth, with chunks and layers as their units: a write by a deep layer in one chunk can be read by a shallow layer in the next. We instantiate this idea as uTTT-MoE and uTTT-Dense. uTTT-MoE routes each token head to a few experts in a pool shared by all layers; uTTT-Dense applies the whole shared memory at every layer without routing. In language modeling, uTTT-MoE reaches 15.5 and 27.9 RULER accuracy at 124M and 760M, 2.6 and 2.1 points above its layer-private counterpart at equal state and active compute, the highest among tested bounded-state models, with per-token loss matching or beating full attention. In novel view synthesis, sharing at fixed per-layer compute gains 0.92 dB in view-23 object PSNR in routed models and 0.76 dB in dense models.
Zefan Cai, Qinzhe Hu, Ziqiao Ma +2
University of Wisconsin–Madison · University of Michigan–Ann Arbor · Adobe