Test-Time Training with Next-Token Prediction

2026-06-19

AI Summary

This work addresses the challenge of enhancing long-context reasoning capabilities in pretrained large language models without modifying their architecture. It proposes a plug-and-play test-time training (TTT) method that, for the first time, directly leverages the self-supervised signal from next-token prediction to guide online updates of fast weights. Specifically, the approach generates target hidden states through causal computation and employs pointwise linear projection for efficient adaptation. Requiring neither additional training nor model restructuring, the method substantially improves long-context performance across multiple mainstream large language models on the RULER Full-13 and LongBench-v2 benchmarks, while preserving their original commonsense and factual knowledge capabilities.

Abstract

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, 8k, 16k, and 32k context lengths), 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: Llama-3.1-8B (+3.9), Mistral-7B-v0.3 (+3.0), and the Qwen3 series (Qwen3-4B +4.1, Qwen3-0.6B +2.9). On the real-world LongBench-v2 long-document QA benchmark, TTT-NTP improves over the base model on both Llama-3.1-8B (+5.6) and Mistral-7B-v0.3 (+3.7), while preserving commonsense and knowledge performance.

Problem

• test-time training
• next-token prediction
• fast-weight adaptation
• pretrained LLMs
• long-context

Innovation

• test-time training
• next-token prediction
• fast-weight adaptation
• pretrained LLMs
• long-context modeling