Let's (not) just put things in Context: Test-Time Training for Long-Context LLMs (2512.13898v1)

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.

• The paper introduces query-only test-time training (qTTT) to adapt LLMs during inference and overcome score dilution in long-context settings.
• It develops a theoretical framework proving that an Ω(log T) logit margin is essential to counteract the vanishing attention mass in massive distractor contexts.
• Empirical results on LongBench-v2 and ZeroScrolls benchmarks demonstrate qTTT’s significant accuracy gains over standard methods, especially in multi-hop reasoning tasks.

Query-Only Test-Time Training for Long-Context LLMs

Introduction and Motivation

The exponential growth in context window sizes in modern LLMs, enabled by pretraining and architectural advances, has not translated to equivalent improvements in long-context reasoning and retrieval. Empirical findings consistently show that while LLMs can accept context windows of tens of thousands to millions of tokens, their effective usage of context is limited: models often fail to retrieve relevant “needle” information embedded deep within massive “haystack” distractor contexts, with sharp performance degradation as sequence length increases. Conventional inference-time compute scaling methods, such as generating chain-of-thought or "thinking" tokens, best-of-$n$ decoding, and related inference-time strategies, have demonstrated benefits for short- and medium-length contexts but exhibit rapidly diminishing returns in the long-context regime.

This paper dissects such scaling failures by formalizing the static self-attention bottleneck inherent to vanilla transformers, identifying “score dilution”—whereby attention to the correct evidence token vanishes as the number of distractors grows—as a fundamental limitation. The authors mathematically prove an $\Omega(\log T)$ margin requirement on logit separation to avoid vanishing attention mass, which is not achievable through additional autoregressive decoding with frozen parameters.

Empirical Analysis of Long-Context Failure Modes

The authors design two controlled synthetic tasks to empirically diagnose long-context failure modes: (1) bug localization in large code repositories, and (2) anomaly detection in transaction logs. Across Qwen3 models ($1.7$B–$8$B), results show that increasing context length leads to a monotonic decrease in accuracy for both standard in-context learning and thinking-token strategies, confirming that inference-time compute spent solely on decoding yields diminishing benefits as context expands.

Figure 1: Illustration of the bug tracing task in a code repository, isolating the effect of context length on model retrieval ability.

This synthetic analysis is extended theoretically to quantify the “score dilution” effect, demonstrating that even with an unambiguous best evidence token, the softmax normalization in static attention layers allocates vanishing mass unless the target-distractor margin increases at least logarithmically with context length. Additional decode-time tokens do not affect this underlying attention allocation and thus cannot reliably recover access to the buried signal.

Query-Only Test-Time Training: Algorithmic Approach

Motivated by this analysis, the paper advances a new inference-time adaptation approach: query-only test-time training (qTTT). Instead of further decoding, the available compute is used to adapt only the query projection matrices on the input’s specific context, with keys and values held fixed via KV cache reuse. This preserves computational efficiency while directly targeting the logit margin bottleneck associated with long-context static attention.

The procedure requires a single full-context prefill to cache K/V tensors for all positions, after which lightweight gradient updates are performed only on query projections for random small token spans. The evidence presented shows that each such update strictly increases the logit margin, especially when the attention distribution is maximally diffuse—precisely the challenging regime for long-context tasks.

Figure 2: Visual schematic of how qTTT gradient updates adjust query projections, shifting query vectors toward key vectors of true evidence tokens and counteracting score dilution.

A theoretical result formalizes that each gradient step with respect to the next-token loss on a span moves queries closer to the key of the relevant “needle,” and away from the centroid of attention mass, thus provably increasing the required target-distractor gap to re-concentrate attention.

FLOP Accounting and Comparison to Other Inference-Time Strategies

The paper provides an exact FLOP-matching analysis between thinking-token generation and qTTT adaptation. For long contexts, generating $T_{\text{think}}$ thinking tokens is computationally equivalent to approximately $T_{\text{think}} / (2k)$ qTTT steps on spans of length $k \ll T$. Crucially, qTTT leaves the KV cache unchanged and avoids ballooning the memory footprint, in contrast to thinking-token strategies that continue to expand sequence length with each additional token.

Experimental Results on LongBench-v2 and ZeroScrolls

Evaluations are performed across all major subsets of the LongBench-v2 and ZeroScrolls benchmarks, with Qwen3 models ranging from 1.7B to 8B parameters. Results highlight consistent, large performance improvements of qTTT over both standard in-context and FLOP-matched thinking-token baselines for long-context retrieval and multi-hop reasoning.

Figure 3: Accuracy improvements of qTTT vs. in-context and thinking-token methods on LongBench-v2 across code, dialogue, and reasoning tasks for Qwen3-8B, with strongest gains in tasks requiring distributed evidence retrieval.

Figure 4: Model performance on ZeroScrolls benchmarks; qTTT exhibits robust gains on retrieval and reasoning tasks, with effects more pronounced as model size increases.

In retrieval-dependent tasks (e.g., long-dialogue history, multi-document QA, code repositories), average improvement margins for Qwen3-4B are as high as 12.6%–14.1% absolute. qTTT effectiveness consistently scales with model size.

On summarization-style datasets with less severe retrieval bottlenecks, improvements are mild, underlining that the margin provided by qTTT specifically targets the dilution regimes identified in theory.

Score Dilution Evidence and Attention Mass Diagnostics

Ablation analyses confirm that the gain from qTTT is directly attributable to sustained attention mass on the target tokens: for increasing context length, both accuracy and target attention mass drop precipitously for vanilla and thinking strategies, but remain stable under qTTT adaptation. This empirically validates the theoretical explanation.

Practical and Theoretical Implications

qTTT is an inference-side procedure requiring no changes to training data, architecture, or pretraining recipes. It is compatible with and complementary to advances in long-context modeling, including alternatives to static attention kernels or improved positional encoding, and can be combined with retrieval-augmented methods or other test-time adaptation paradigms.

Theoretically, these findings elevate the importance of adapting similarity computation (i.e., query-key alignment) as opposed to merely allocating compute for additional decoding. The static self-attention mechanism in its standard form is provably insufficient for long-context retrieval beyond a critical context size, and inference-time adaptation—especially targeted, parameter-efficient updates—offers an effective mitigation strategy.

Future Directions

Further research avenues include investigation of $(k, N_{\text{TTT}})$ adaptation schedules to optimize the compute-accuracy frontier, extension to other test-time adaptation frameworks (e.g., self-consistency, beam search, or reinforcement-based test-time training), and automatic prediction of when query adaptation is most beneficial for a given task or instance.

Conclusion

This work establishes that inference-time compute for LLMs in long-context settings is most effectively spent via parameter-efficient adaptation—specifically, qTTT updates to queries—rather than further static decoding. Both the theoretical characterization and strong empirical evidence suggest that overcoming the logit margin requirement induced by score dilution is the key to scaling LLM effectiveness as context sizes grow. The approach is broadly compatible and promises to inform both practical deployment and architectural evolution of future long-context models.