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

Yue Wu; Yang Zhang; Weinan Gan; Chongming Gao; Xinsen Guo; Xiaoyuan Li; Fuli Feng

cs.LGOct 27, 2026

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

Authors: Yue Wu, Yang Zhang, Weinan Gan, Chongming Gao, Xinsen Guo, Xiaoyuan Li, Fuli Feng

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

Abstract

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.

Figures & tables

Figure 1. Motivation experiment on LongMemEval-S (Llama-3.1-8B-Instruct). (a) Accuracy vs. the answer-chunk learning rate. (b) Accuracy vs. the irrelevant-chunk learning rate. The default In-Place TTT configuration uses

0.3

for both learning rates. The dashed line marks the zero-shot Baseline (

33.4\%

Figure 2. Overview of Select-TTT. (1) The frozen LM head maps chunk hidden states to logit space, where query-token scores are aggregated by top-

k

mean pooling to assess chunk relevance. (2) Relevance scores are converted into learning-rate multipliers. (3) In-Place TTT computes the NTP-aligned fast-weight updates. (4) The multipliers rescale these updates to selectively write query-relevant information into fast-weight memory.

Backbone	Method	LongMemEval-S
Backbone	Method	KU	MS	TR	S-Pref.	S-Asst.	S-User	Avg.
Llama-3.1-8B	Baseline	50.00	18.80	16.54	3.33	69.64	58.57	33.4
	IP-TTT	47.44	18.80	20.30	3.33	66.07	58.57	33.6
	Select-TTT	52.56	22.56	28.57	13.33	64.29	54.29	37.4
Qwen3-4B	Baseline	28.21	21.80	22.56	6.67	57.14	30.00	27.2
	IP-TTT	26.92	15.79	20.30	3.33	67.86	27.14	25.4

Table 1. Protocol I: zero-shot inference-time TTT on LongMemEval-S with two backbones. Accuracy (%) is reported; the best result within each backbone block is in bold .

Backbone	Method	LoCoMo
Backbone	Method	Multi	Temp.	Open	Single	Adv.	Avg.
Llama-3.1-8B	Baseline	36.60	29.30	17.78	57.91	18.83	39.55
	IP-TTT	34.92	30.12	19.51	58.22	15.47	38.90
	Select-TTT	36.38	30.74	20.74	58.34	20.40	40.42
Qwen3-4B	Baseline	25.66	20.12	18.31	38.88	4.04	25.15
	IP-TTT	25.84	20.60	18.19	39.25	4.26	25.45

Table 2. Protocol I: zero-shot inference-time TTT on LoCoMo with two backbones. Overall score (%) is reported; the best result within each backbone block is in bold .

Figure 3. Protocol II: LongMemEval-S overall accuracy vs. continual-training steps for No Training, CPT, IP-TTT, and Select-TTT (Qwen3-4B backbone).

Figure 4. Ablation of Select-TTT on LongMemEval-S (both backbones): (a) temperature

\tau\in\{0.1,0.3,0.5\}

; (b) pooling ratio

\rho\in\{0.01,0.05,0.10,0.20\}

; (c) pooling strategy (Max, Mean, Top-

k

). The default configuration is

\tau=0.1

\rho=0.10

, top-

k

mean.

Ablation Setting	LongMemEval-S
Ablation Setting	KU	MS	TR	S-Pref.	S-Asst.	S-User	Avg.
w/ soft write, w/ clip	52.56	22.56	28.57	13.33	64.29	54.29	37.4
w/ soft write, w/o clip	46.15	21.05	23.31	3.33	66.07	57.14	34.6
w/o soft write, w/ clip	42.31	15.04	18.05	6.67	66.07	51.43	30.4
w/o soft write, w/o clip	21.79	13.53	17.29	3.33	53.57	37.14	23.0

Table 3. Ablation of the multiplier design (Eq. ( 10 )) on LongMemEval-S with the Llama-3.1-8B backbone. Accuracy (%) is reported; the best result per column is in bold .

Chunk size	Signal	Hit@ $k$			Rec@ $k$
Chunk size	Signal	$k{=}1$	$k{=}3$	$k{=}5$	$k{=}1$	$k{=}3$	$k{=}5$
2048	Query-CE (ours)	0.200	0.440	0.582	0.092	0.201	0.278
	Cosine	0.094	0.212	0.293	0.061	0.136	0.189
	Cross-encoder	0.090	0.245	0.372	0.020	0.059	0.098
4096	Query-CE (ours)	0.256	0.519	0.654	0.128	0.295	0.393
	Cosine	0.133	0.308	0.415	0.085	0.198	0.267

Table 4. Comparing relevance signals across different chunk sizes on LongMemEval-S QA instances. Hit@

k

: at least one of the

k

highest-ranked chunks contains the answer. Rec@

k

: fraction of answer chunks covered by the

k

highest-ranked chunks.

Figure 5. Efficiency analysis of Select-TTT. Prefill throughput (a) and peak GPU memory (b) are compared across the Baseline, IP-TTT, and Select-TTT on Qwen3-4B model.

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

Parameter	Llama-3.1-8B-Instruct	Qwen3-4B
Parameters (approx.)	$8$ B	$4$ B
Hidden size $d_{\text{model}}$	$4{,}096$	$2{,}560$
Num. layers	$32$	$36$
Attention heads	$32$	$32$
Key-value heads	$8$	$8$
FFN intermediate size $d_{\text{ff}}$	$14{,}336$	$9{,}728$

Appendix

Table 5. Architectural configurations of the two backbones.

Component	Value
TTT layers	$\{0,6,12,18,24,30\}$
Fast weight	MLP down-projection $\mathbf{W}_{\mathrm{down}}$
Conv1D kernel size	$5$ (causal, depthwise, zero-init)
Target projection $\mathbf{W}_{\mathrm{target}}$	$d_{\text{model}}\times d_{\text{model}}$ , sparse-diagonal init
Objective	NTP-aligned (negative Frobenius inner product)
Pooling strategy	top- $k$ mean

Appendix

Table 6. Configuration of the TTT module and Select-TTT scoring.

Statistic	LongMemEval-S	LoCoMo
Questions	$500$	$1{,}986$
Context organization	$\approx 40$ sessions	$19$ – $32$ sessions
Context scale	$\approx 115$ K tokens	$369$ – $689$ turns/dialogue
Subtasks / query types	$6$ subtasks	$5$ query types
Purpose	inference evaluation	inference evaluation

Appendix

Table 7. Inference-time evaluation benchmarks.

Statistic	Value
Original corpus	$\approx 44{,}600$ QA instances
Languages	Chinese and English (mixed)
Subset used for CPT	English only
Context-length filter	$8$ K– $65$ K tokens
Annotations	sentence-level citation spans
Purpose	continual pre-training

Appendix

Table 8. LongCite-45k subset used for continual pre-training.

Hyperparameter	Llama-3.1-8B-Instruct	Qwen3-4B
Optimizer	AdamW	AdamW
Peak learning rate	$5\times 10^{-6}$	$5\times 10^{-6}$
Weight decay	$0.1$	$0.1$
Gradient clipping	$1.0$	$1.0$
LR decay ratio	0.90	0.90
Global batch size	8	8

Appendix

Table 9. Continual-training hyperparameters (Protocol II).

Step	Method	LongMemEval-S
Step	Method	KU	MS	TR	S-Pref.	S-Asst.	S-User	Avg.
1000	CPT	43.59	20.30	24.81	3.33	76.79	34.29	32.40
	IP-TTT	35.90	24.06	28.57	3.33	73.21	37.14	33.20
	Select-TTT	44.87	19.55	29.32	0.00	69.64	30.00	32.00
2000	CPT	43.59	26.32	25.56	0.00	71.43	34.29	33.40
	IP-TTT	51.28	27.07	27.07	6.67	71.43	35.71	35.80

Appendix

Table 10. Continual-training results on LongMemEval-S with Qwen3-4B. All values are percentages (%).

Step	Method	LongMemEval-S
Step	Method	KU	MS	TR	S-Pref.	S-Asst.	S-User	Avg.
1000	CPT	29.49	22.56	30.08	0.00	71.43	35.71	31.60
	IP-TTT	44.87	18.80	24.81	6.67	71.43	32.86	31.60
	Select-TTT	43.59	21.05	25.56	3.33	76.79	35.71	33.00
2000	CPT	44.87	20.30	25.56	3.33	73.21	37.14	32.80
	IP-TTT	52.56	24.81	32.33	0.00	69.64	37.14	36.40

Appendix

Table 11. Continual-training results on LongMemEval-S with Llama-3.1-8B-Instruct. All values are percentages (%). Avg. is the question-weighted overall accuracy. Bold Avg. values indicate the best result at each training step.

Method	$\tau$	LongMemEval-S
Method	$\tau$	KU	MS	TR	S-Pref.	S-Asst.	S-User	Avg.
Select-TTT	0.1	52.56	22.56	28.57	13.33	64.29	54.29	37.40
	0.3	46.15	21.05	23.31	3.33	66.07	57.14	34.60
	0.5	46.15	20.30	19.55	10.00	64.29	57.14	33.60

Appendix

Table 12. Temperature ablation on LongMemEval-S with Llama-3.1-8B-Instruct.

Method	$\tau$	LongMemEval-S
Method	$\tau$	KU	MS	TR	S-Pref.	S-Asst.	S-User	Avg.
Select-TTT	0.1	29.49	26.32	22.56	3.33	57.14	32.86	28.80
	0.3	33.33	22.56	21.05	0.00	60.71	25.71	27.20
	0.5	28.21	21.05	21.05	0.00	60.71	27.14	26.20

Appendix

Table 13. Temperature ablation on LongMemEval-S with Qwen3-4B.

Method	$\rho$	LongMemEval-S
Method	$\rho$	KU	MS	TR	S-Pref.	S-Asst.	S-User	Avg.
Select-TTT	0.01	53.85	23.31	24.06	6.67	64.29	54.29	36.20
	0.05	53.85	21.05	25.56	13.33	60.71	54.29	36.00
	0.10	52.56	22.56	28.57	13.33	64.29	54.29	37.40
	0.20	53.85	22.56	25.56	13.33	66.07	50.00	36.40

Appendix

Table 14. Pooling-ratio ablation on LongMemEval-S with Llama-3.1-8B-Instruct.

Method	$\rho$	LongMemEval-S
Method	$\rho$	KU	MS	TR	S-Pref.	S-Asst.	S-User	Avg.
Select-TTT	0.01	37.18	18.80	20.30	3.33	64.29	21.43	26.60
	0.05	32.05	18.05	20.30	3.33	57.14	32.86	26.40
	0.10	29.49	26.32	22.56	3.33	57.14	32.86	28.80
	0.20	35.90	20.30	18.80	3.33	60.71	35.71	28.00

Appendix

Table 15. Pooling-ratio ablation on LongMemEval-S with Qwen3-4B.

Method	Pooling	LongMemEval-S
Method	Pooling	KU	MS	TR	S-Pref.	S-Asst.	S-User	Avg.
Select-TTT	Max	53.85	21.80	18.80	6.67	67.86	60.00	35.60
	Mean	50.00	19.55	20.30	3.33	64.29	45.71	32.20
	Top- $k$ mean	52.56	22.56	28.57	13.33	64.29	54.29	37.40

Appendix

Table 16. Pooling-strategy ablation on LongMemEval-S with Llama-3.1-8B-Instruct.

Method	Pooling	LongMemEval-S
Method	Pooling	KU	MS	TR	S-Pref.	S-Asst.	S-User	Avg.
Select-TTT	Max	38.46	20.30	18.80	3.33	60.71	35.71	28.40
	Mean	33.33	20.30	21.05	0.00	62.50	28.57	27.20
	Top- $k$ mean	29.49	26.32	22.56	3.33	57.14	32.86	28.80

Appendix

Table 17. Pooling-strategy ablation on LongMemEval-S with Qwen3-4B.

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

EASE-TTT: Evidence-Aligned Selective Test-Time Training for Long-Context Question Answering
Jun 5, 2026Xiaopeng Yuan, Zebin Wang, Suwen Wang +3Long-Context Retrieval Long-Context QA
University of Illinois Urbana-Champaign · 2Harvard University · 3Brion, ASML US LP +1
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
In-Place Test-Time Training
Apr 7, 2026Guhao Feng, Shengjie Luo, Kai Hua +4Test-Time Training Continual Learning for LLMs
ByteDance Seed · Peking University

Jun 5, 2026J/K move · Enter open · S save

EASE-TTT: Evidence-Aligned Selective Test-Time Training for Long-Context Question Answering

Xiaopeng Yuan, Zebin Wang, Suwen Wang, Zongxin Yang+2

University of Illinois Urbana-Champaign · 2Harvard University · 3Brion, ASML US LP+1

Long-Context Retrieval Long-Context QA Long-Context Modeling Test-Time Training

Open paper AskPDF arXiv

Long-context question answering (QA) remains challenging for smaller language models even when answer-bearing evidence is already present in the input. Existing within-context retrieval methods localize and expose candidate evidence chunks for the question, but they stop at input-level evidence exposure rather than adapting the query-side attention parameters that control how the model allocates attention over full-context positions. In contrast, lightweight test-time adaptation methods, such as query-only test-time training (qTTT), leave evidence localization unresolved because their generic span-level self-supervised objectives do not identify which context positions support the current answer. In this paper, we propose Evidence-Aligned SElective Test-Time Training (EASE-TTT), a within-context retrieval-augmented test-time training framework that converts selected evidence chunks into a soft attention supervision target over their token positions. Instead of replacing the full context with retrieved chunks, EASE-TTT uses the resulting attention target to guide query-side adaptation, with the adapted model generating the final answer from the original full context. Experiments on six LongBench QA tasks and three small decoder-only language models show that EASE-TTT achieves the strongest macro-average performance among full-context inference, retrieval-only baselines, and qTTT, supporting evidence-aligned test-time adaptation in long-context QA.

Abstract

Test-Time Training

Explore similar work

EASE-TTT: Evidence-Aligned Selective Test-Time Training for Long-Context Question Answering

Test-Time Training with Next-Token Prediction

In-Place Test-Time Training