Test-Time Training for Long-Context LLMs
Free while signed in. Answers cite the passages they came from.
This paper shows that long-context LLMs can access millions of tokens but often fail to meaningfully use that information. The authors propose query-only test-time training (qTTT), which adapts models during inference through targeted gradient updates rather than generating more thinking tokens. - **Score dilution problem:** The authors identify that static self-attention suffers from score dilution, where target token probabilities vanish as context grows. They prove that target-distractor logit margins must scale logarithmically with context length to maintain performance. - **Query-only TTT:** qTTT performs a single prefill to cache keys and values, then applies lightweight gradient updates exclusively on query projection matrices. This keeps other parameters fixed and reuses the key-value cache, making it computationally efficient. - **Massive improvements:** qTTT achieves 12.6 and 14.1 percentage point improvements for Qwen3-4B on LongBench-v2 and ZeroScrolls benchmarks. Gains exceed 20% on code comprehension, multi-document QA, and multi-hop reasoning tasks. - **Better compute allocation:** Under matched inference-time compute budgets, qTTT consistently outperforms thinking token strategies. The practical takeaway is that a small amount of context-specific training beats generating thousands of thinking tokens for long-context tasks.
Score dilution problem: The authors identify that static self-attention suffers from score dilution, where target token probabilities vanish as context grows. They prove that target-distractor logit margins must scale logarithmically with context length to maintain performance.
Query-only TTT: qTTT performs a single prefill to cache keys and values, then applies lightweight gradient updates exclusively on query projection matrices. This keeps other parameters fixed and reuses the key-value cache, making it computationally efficient.
Massive improvements: qTTT achieves 12.6 and 14.1 percentage point improvements for Qwen3-4B on LongBench-v2 and ZeroScrolls benchmarks. Gains exceed 20% on code comprehension, multi-document QA, and multi-hop reasoning tasks.
Better compute allocation: Under matched inference-time compute budgets, qTTT consistently outperforms thinking token strategies. The practical takeaway is that a small amount of context-specific training beats generating thousands of thinking tokens for long-context tasks.
Abstract
Progress on training and architecture strategies has enabled LLMs with millions of tokens in context length. However, empirical evidence suggests that such long-context LLMs can consume far more text than they can reliably use. On the other hand, it has been shown that inference-time compute can be used to scale performance of LLMs, often by generating thinking tokens, on challenging tasks involving multi-step reasoning. Through controlled experiments on sandbox long-context tasks, we find that such inference-time strategies show rapidly diminishing returns and fail at long context. We attribute these failures to score dilution, a phenomenon inherent to static self-attention. Further, we show that current inference-time strategies cannot retrieve relevant long-context signals under certain conditions. We propose a simple method that, through targeted gradient updates on the given context, provably overcomes limitations of static self-attention. We find that this shift in how inference-time compute is spent leads to consistently large performance improvements across models and long-context benchmarks. Our method leads to large 12.6 and 14.1 percentage point improvements for Qwen3-4B on average across subsets of LongBench-v2 and ZeroScrolls benchmarks. The takeaway is practical: for long context, a small amount of context-specific training is a better use of inference compute than current inference-time scaling strategies like producing more thinking tokens.
Get next week’s papers.
The same picks and the same summaries, in your inbox. Free, and 176 issues deep.
Subscribe on Substack