Rethinking Expressivity and Efficiency in Test-Time Training

Zeyun Zhong; Joya Chen; Manuel Martin; Frederik Diederichs; Juergen Gall; Juergen Beyerer

cs.LGAug 21, 2026

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 ·

Abstract

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 E $^2$ -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 E $^2$ -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\times$ the training context length. Meanwhile, E $^2$ -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 $(\bm{k}_{t},\bm{v}_{t})$ defines a self-supervised loss whose gradient step (the “training” stage) updates $\mathbf{W}_{t-1}$ to $\mathbf{W}_{t}$ ; the query $\bm{q}_{t}$ then reads from $\mathbf{W}_{t}$ to produce $\bm{o}_{t}$ .

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

\mathbf{G}_{t}

. State recurrence : the granularity at which the states

(\mathbf{W}_{t},\mathbf{M}_{t})

evolve. Output computed from : the state applied to the query

\bm{q}_{t}

. 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 $\mathbf{Q},\mathbf{K},\mathbf{V}$ and are fused by a data-dependent gate. The TTT update is modulated by input-dependent $\alpha,\eta,\beta$ .

Model	Attn.	Wiki.	LMB.	LMB.	PIQA	Hella.	Wino.	ARC-e	ARC-c	Avg.
	window	ppl $\downarrow$	ppl $\downarrow$	acc $\uparrow$	acc $\uparrow$	acc_n $\uparrow$	acc $\uparrow$	acc $\uparrow$	acc_n $\uparrow$	$\uparrow$
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 (

\downarrow

) and zero-shot accuracy (

\uparrow

) 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 $\uparrow$	acc $\uparrow$	acc $\uparrow$
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.
Model (1.3B)	NQA	QQA	MFQ	HQA	2WM	Mus	GvR	QMS	MNs	TRC	TQA	SSM	LCC	RBP	Avg.
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
Method	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
Model	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 $\downarrow$	LMB ppl $\downarrow$	Avg $\uparrow$
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

\beta_{t},\gamma_{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 ${}_{\text{SwiGLU}}$	95.6	95.4	93.6	14.1

Appendix

Table 9: Component ablation at the 1.3B scale on S-NIAH-1 (accuracy

\uparrow

at three context lengths) and LongBench (average score

\uparrow

). The full model is trained for 15B tokens at 2K context; ablations reuse identical training data and hyperparameters.

Inner-loop update rule	Wiki ppl $\downarrow$	LMB ppl $\downarrow$	Avg $\uparrow$
$\eta_{t}$ only (no momentum, no decay)	26.6	26.4	48.7
$\eta_{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

\eta_{t}

Configuration	Wiki ppl $\downarrow$	LMB ppl $\downarrow$	Avg $\uparrow$
Base learning rate $\eta_{\text{base}}$ ( $\alpha_{\text{base}}{=}10^{-1}$ )
$\eta_{\text{base}}=10^{-1}$	28.8	30.3	46.5
$\eta_{\text{base}}=10^{-2}$ (default)	25.5	25.0	49.1
$\eta_{\text{base}}=10^{-3}$	26.1	23.0	49.1
Base decay $\alpha_{\text{base}}$ ( $\eta_{\text{base}}{=}10^{-2}$ )
$\alpha_{\text{base}}=10^{-1}$ (default)	25.5	25.0	49.1

Appendix

Table 11: Sensitivity of E 2 -TTT

{}_{\text{MLP}}

(340M) to the base hyperparameters of the input-dependent dynamics.

Shuffled head	Wiki ppl $\downarrow$	$\Delta$ ppl	S-NIAH-1 16K $\uparrow$
— (baseline)	19.94	—	93.6
$\eta$	20.39	$+0.45$	6.0
$\beta$	20.05	$+0.11$	86.2
$\alpha$	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

{}_{\text{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) $\downarrow$		Peak mem. (GB) $\downarrow$
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 $\downarrow$		Commonsense	Retrieval
tokens	ctx.	Wiki.	LMB.	Avg. $\uparrow$	Avg. $\uparrow$
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

Test-time training adapts model parameters or memory during deployment using each input or stream to improve performance under changing tasks and data.

See all 104 papers

Select-TTT: Query-Conditioned Selective Fast Weight Updates for In-Place Test-Time Training
Oct 27, 2026Yue Wu, Yang Zhang, Weinan Gan +4Test-Time Training Inference-Time Adaptation of Language Models
University of Science and Technology of China China · National University of Singapore Singapore · Huawei Noah’s Ark Lab China
Test-Time Training with Next-Token Prediction
Jun 19, 2026Xuan Ouyang, Zefan Cai, Junjie HuLong-Context Language Modeling Next-Token Prediction
University of Wisconsin–Madison
Universal Test-Time Training
Oct 4, 2026Zefan Cai, Qinzhe Hu, Ziqiao Ma +2Memory-Augmented Language Models Recurrent Transformers
University of Wisconsin–Madison · University of Michigan–Ann Arbor · Adobe

Oct 27, 2026J/K move · Enter open · S save

Select-TTT: Query-Conditioned Selective Fast Weight Updates for In-Place Test-Time Training

Yue Wu, Yang Zhang, Weinan Gan, Chongming Gao+3

University of Science and Technology of China China · National University of Singapore Singapore · Huawei Noah’s Ark Lab China

Test-Time Training Inference-Time Adaptation of Language Models Fast Weights Long-Context QA

Open paper AskPDF arXiv

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.

Abstract

Test-Time Training

Explore similar work

Select-TTT: Query-Conditioned Selective Fast Weight Updates for In-Place Test-Time Training

Test-Time Training with Next-Token Prediction

Universal Test-Time Training