In-Place Test-Time Training

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

cs.LGApr 7, 2026

In-Place Test-Time Training

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

Organizations: ByteDance Seed · Peking University

Abstract

The static train then deploy" paradigm fundamentally limits Large Language Models (LLMs) from dynamically adapting their weights in response to continuous streams of new information inherent in real-world tasks. Test-Time Training (TTT) offers a compelling alternative by updating a subset of model parameters (fast weights) at inference time, yet its potential in the current LLM ecosystem is hindered by critical barriers including architectural incompatibility, computational inefficiency and misaligned fast weight objectives for language modeling. In this work, we introduce In-Place Test-Time Training (In-Place TTT), a framework that seamlessly endows LLMs with Test-Time Training ability. In-Place TTT treats the final projection matrix of the ubiquitous MLP blocks as its adaptable fast weights, enabling a drop-in" enhancement for LLMs without costly retraining from scratch. Furthermore, we replace TTT's generic reconstruction objective with a tailored, theoretically-grounded objective explicitly aligned with the Next-Token-Prediction task governing autoregressive language modeling. This principled objective, combined with an efficient chunk-wise update mechanism, results in a highly scalable algorithm compatible with context parallelism. Extensive experiments validate our framework's effectiveness: as an in-place enhancement, it enables a 4B-parameter model to achieve superior performance on tasks with contexts up to 128k, and when pretrained from scratch, it consistently outperforms competitive TTT-related approaches. Ablation study results further provide deeper insights on our design choices. Collectively, our results establish In-Place TTT as a promising step towards a paradigm of continual learning in LLMs.

Figures & tables

Figure 1: The overall framework of our In-Place Test-Time Training. The module operates sequentially on input chunks. For each chunk, the current fast weights are first applied to the intermediate activations $\mathbf{Z}$ to produce the output. Then, these weights are updated using the activations $\mathbf{Z}$ and a value $\mathbf{V}$ derived from the token embeddings. This "apply-then-update" cycle allows the model to dynamically adapt to incoming context in a strictly causal manner.

	In-Domain Evaluation						Extrapolation
Model	4k	8k	16k	32k	64k	128k	256k
Mistral-7B [ 31 ]	$93.6$	$91.2$	$87.2$	$75.4$	$49.0$	$13.8$	-
GLM3-6B [ 23 ]	$87.8$	$83.4$	$78.6$	$69.9$	$56.0$	$42.0$	-
Phi3-medium-14B [ 1 ]	$93.3$	$93.2$	$91.1$	$86.8$	$78.6$	$46.1$	-
Llama3-8B [ 41 ]	$92.8$	$90.3$	$85.7$	$79.9$	$76.3$	$69.5$	-
Qwen3-4B (Instruct) [ 57 ]	$95.1$	$93.6$	$91.0$	$87.8$	$77.8$	$66.0$	-

Table 1: Evaluation results on RULER [ 28 ] . Average accuracy (%) is reported, with the best results in bold .

Base Model	Method	4k	8k	16k	32k	64k	64k+YaRN
LLaMA-3.1-8B	Baseline	93.9	92.1	92.5	91.1	81.6	–
LLaMA-3.1-8B	In-Place TTT	94.4	93.0	93.3	91.7	83.7	–
Qwen3-14B	Baseline	96.8	95.0	94.6	90.7	67.9	81.3
Qwen3-14B	In-Place TTT	97.2	95.7	95.2	91.2	70.6	82.5

Table 2: Extension of In-Place TTT to LLaMA-3.1-8B and Qwen3-14B-Base on the RULER benchmark. We report the average accuracy (%) with the best results in bold .

Panel of Figure 2 — Figure 2: Sliding Window Perplexity at varying context lengths on the Pile dataset for 500M (left) and 1.5B (right) parameter models. Our In-Place TTT consistently achieves lower perplexity than all competitive baselines.

		Common Sense Reasoning					Long-Context Evaluation
Model Architecture		HellaSwag	ARC-E	ARC-C	MMLU	PIQA	RULER-4k	RULER-8k	RULER-16k
Baselines	Full Attn.	$55.67$	$64.52$	$33.19$	$36.43$	$72.63$	$45.77$	$38.09$	$6.58$
Baselines	SWA	$54.92$	$64.18$	$32.85$	$36.06$	$72.58$	$14.77$	$9.91$	$5.07$
I.P. TTT	Full Attn.	$55.85$	$64.98$	$32.34$	$37.42$	$73.29$	$49.98$	$43.82$	$19.99$
I.P. TTT	SWA	$55.24$	$64.60$	$33.70$	$36.48$	$72.03$	$28.33$	$26.80$	$7.57$

Table 3: Evaluation results of 4B models on common sense reasoning and long-context evaluation benchmarks. Best performance is in bold . “SWA” is Sliding-Window Attention, “Full Attn.” is Full Attention, and “I.P. TTT” is our In-Place TTT.

Figure 3: Ablation studies on the key design choices of the In-Place TTT framework, evaluated on the RULER benchmark with a 1.7B parameter model. The plots illustrate the impact of: (a) State size, showing that performance improves as the state size scales; (b) Chunk size, demonstrating a performance trade-off where intermediate sizes (e.g., 512, 1024) are optimal; and (c) The LM-Aligned Value objective, confirming that both the convolution (w Conv) and the projection (w Proj) are crucial.

Panel of Figure 4 — Figure 4: Efficiency analysis of In-Place TTT. Both prefill throughput (a, b) and peak memory (c, d) metrics are presented for 4B models with Sliding-Window Attention (SWA) and Full Attention at various context lengths. Our In-Place TTT introduces negligible overhead in practical scenarios.

Hyperparameter	500M Model	1.5B Model
Optimizer	AdamW	AdamW
Learning Rate	5e-4	3e-4
Batch Size	2M tokens	4M tokens
Weight Decay	0.1	0.1
Gradient Clipping	1.0	1.0
Warmup Steps	1024	1024

Table 4: Training hyperparameters for 500M and 1.5B models.

Hyperparameter	value
Optimizer	AdamW
Learning Rate	3e-4
Batch Size	8M tokens
Weight Decay	0.1
Gradient Clipping	1.0
Warm-up Tokens 1 1 1 We use both batch size warm-up and learning rate warm-up here.	1.6B

Table 5: Training hyperparameters for 1.7B models and 4B models pretraining

Hyperparameter	Stage 1 (32k Context)	Stage 2 (128k Context)
Base Model	Qwen3-4B-Base	Qwen3-4B-Base
Optimizer	AdamW	AdamW
Learning Rate	5e-6	5e-6
Weight Decay	0.1	0.1
Sequence Length	32,768	131,072
Tokens Trained	$\sim$ 20B	$\sim$ 15B

Table 6: Hyperparameters for two-stage continual pre-training.

Hyperparameter	LLaMA-3.1-8B	Qwen3-14B-Base
Optimizer	AdamW	AdamW
Learning Rate	5e-6	5e-6
Weight Decay	0.1	0.1
Sequence Length	32,768	32,768
Tokens Trained	$\sim$ 20B	$\sim$ 20B
RoPE Extension	None	None

Table 7: Hyperparameters for continual pre-training of LLaMA-3.1-8B and Qwen3-14B-Base.

Parameter	500M	1.5B
Parameters (Approx.)	500M	1.5B
Hidden Size ( $d_{\text{model}}$ )	1024	2048
Num Layers	24	24
Num Attention Heads	8	16
FFN Hidden Size ( $d_{\text{ff}}$ )	3072	6144
Window Size	2048	4096

Table 8: Model architectural configurations for 500M and 1.5B Model.

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 99 papers

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
No Time Like the Present: Agentic Test-Time Training for LLM Agents
Jul 3, 2026Yanbo Wang, Jinhua Hao, Yuze Shi +2Continual Learning for LLM Agents Test-Time Training
Kuaishou Technology
Query-Conditioned Test-Time Self-Training for Large Language Models
May 13, 2026Chaehee Song, Minseok Seo, Yeeun Seong +2Test-Time Adaptation Test-Time Optimization
School of Electrical Engineering, KAIST · 2Graduate School of Green Growth and Sustainability, KAIST

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

Test-Time Training with Next-Token Prediction

Xuan Ouyang, Zefan Cai, Junjie Hu

University of Wisconsin–Madison

Long-Context Language Modeling Next-Token Prediction Test-Time Training Sequence Modeling Inference-Time Adaptation of Language Models

Open paper AskPDF arXiv

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.

Abstract

Test-Time Training

Explore similar work

Test-Time Training with Next-Token Prediction

No Time Like the Present: Agentic Test-Time Training for LLM Agents

Query-Conditioned Test-Time Self-Training for Large Language Models