PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Rachit Bansal, Aston Zhang, Rishabh Tiwari, Lovish Madaan, Sai Surya Duvvuri, Devvrit Khatri, David Brandfonbrener, David Alvarez-Melis, Prajjwal Bhargava, Mihir Sanjay Kale, Samy Jelassi

arXiv:2512.13898v1cs.LGcs.CL

TL;DR

Long-context models can consume millions of tokens yet fail to retrieve buried evidence because static self-attention causes score dilution, while thinking-token scaling has diminishing returns. The paper proposes query-only test-time training, using targeted query updates on the given context to increase target–distractor separation. Across long-context benchmarks, qTTT consistently outperforms thinking-token strategies under matched compute, with reported gains of 12.6 and 14.1 percentage points for Qwen3-4B.

  • Problem

    Long-context LLMs can fail to retrieve relevant buried evidence, and generating additional thinking tokens cannot reliably overcome static-attention limitations.

  • Method

    Query-only test-time training performs a prefill with cached keys and values, then applies a few gradient updates exclusively to query projections for the given context.

  • Results

    qTTT consistently outperforms in-context and thinking strategies under matched FLOPs across LongBench-v2 and ZeroScrolls, with the largest gains on retrieval and multi-hop reasoning.

  • Takeaways & Limitations

    For long-context retrieval and reasoning, reallocating inference compute from generating more tokens to a small number of query updates is more effective.

  • Takeaways & Limitations

    The evaluation covers a single point on the (k, N_TTT) trade-off, and gains are task-dependent.

Abstract

from arXiv · show

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.

1 Introduction

Long-context LLMs fail to reliably use buried evidence as contexts grow, while additional thinking tokens face diminishing returns under static attention. The paper diagnoses score dilution and introduces query-only test-time training (qTTT), which reallocates inference compute to targeted query updates and improves long-context performance.

  • 1 Introduction: As context length grows, in-context accuracy drops and thinking tokens show diminishing returns, whereas qTTT improves performance at the same FLOP budget.Controlled tasks cover bug localization in large code repositories and anomaly detection in transaction logs.
  • 1 Introduction: The method adapts the model to the given context rather than generating more text from unchanged weights, and can complement existing long-context strategies.The paper positions qTTT as compatible with architectural changes, longer-window training, and retrieval augmentation.
  • 1 Introduction: qTTT consistently improves long-context performance across 15+ ZeroScrolls and LongBench-v2 datasets and Qwen3 models from 1.7B–8B parameters, outperforming thinking-token baselines under FLOP-matched budgets.The evaluation spans real-world retrieval, code comprehension, multi-document QA, and multi-hop reasoning tasks.
  • 1 Introduction: Static self-attention suffers score dilution: avoiding vanishing target probability requires a target–distractor logit gap scaling as Ω(log T).Distractor tokens weaken the target's attention mass as context length increases.
  • 1 Introduction: Current inference-time scaling strategies cannot reliably retrieve buried targets because they generate tokens using the same static attention mechanism that under-allocates mass to evidence.Theoretical analysis shows generated tokens cannot materially increase the final answer's effective margin unless an intermediate token first attends non-trivially to the target.
  • 1 Introduction: qTTT performs one prefill, caches keys and values, and applies a few gradient updates only to query projections while reusing the cache.These targeted updates increase target–distractor separation for the specific context without changing other parameters.

2 Vanilla Compute-Scaling Strategies Fail for Long Contexts

Controlled long-context retrieval experiments show that standard in-context performance deteriorates sharply as context grows, while generating thinking tokens provides only short-context gains and saturates at large lengths. The paper attributes this failure to score dilution in static self-attention: maintaining target attention requires a target–distractor margin that grows as Ω(log T), which decoding-based strategies cannot reliably create.

  • Empirical Findings: As context length increases, in-context accuracy drops sharply, while thinking-token gains show clear diminishing returns and converge near standard performance at large lengths.These trends hold across controlled code-repository and transaction-log retrieval tasks with the relevant evidence held fixed as the surrounding haystack grows.
  • Score Dilution: Score dilution makes the needle’s attention mass vanish when a constant fraction of distractors lies within O(1) logit of it.The bound α_i,j⋆ ≤ 1/(1 + me^-∆) tends to zero when m grows proportionally with T and ∆ remains O(1).
  • Logarithmic Margin Requirement: Avoiding vanishing target probability requires a needle–distractor logit margin that scales as Ω(log T) with context length.The logarithmic requirement is derived for guaranteeing a fixed target attention mass against worst-case distractors.
  • Limits of Thinking Tokens: Thinking tokens cannot reliably recover buried evidence because each generated token uses the same static attention and can carry at most its own needle-attention fraction.Under small margins, that fraction is provably tiny, so intermediate tokens cannot materially increase the final answer’s effective margin unless one first attends substantially to the needle.
  • Implication: Successful inference-time scaling must alter the query–key similarity rather than merely generate more tokens with unchanged parameters.This conclusion motivates gradient-based query adaptation as an alternative to decoding-based compute scaling.

3 Efficient Test-Time Adaptation via Query-Only Updates

Query-only test-time training makes long-context adaptation practical by caching keys and values once, then updating only query projections on short spans. These targeted updates increase the target–distractor margin while preserving the cache, addressing score dilution with far less computation than full-parameter adaptation.

  • 3 Efficient Test-Time Adaptation via Query-Only Updates: Full-parameter TTT is impractical for long contexts because updates invalidate cached keys and values, requiring fresh full-context passes with prohibitive compute and activation memory.One full-parameter update is estimated to cost about 1.2 × T decoding tokens; for T ≈105 tokens, this is about 120K decoding tokens.
  • 3.1 Query-Only TTT for Long Context: qTTT reuses a single prefilled key–value cache and applies lightweight gradient updates only to query projections on short spans.The adapted model then generates the final answer without repeatedly re-encoding the long context.
  • 3.2 Why Query-Only Test-Time Training is Effective: qTTT moves each query toward the target key and away from the attention-weighted mean, directly counteracting score dilution.The update holds keys and values fixed while reshaping query-to-key similarity.
  • 3.2 Why Query-Only Test-Time Training is Effective: The logit margin strictly increases for sufficiently small steps, with the largest gains when attention is diffuse and long-context dilution is most severe.The theoretical result links the improvement to the squared distance between the target key and the attention-weighted mean.
  • 3.3 FLOP Equivalence: Thinking Tokens vs. Query-Only TTT: Under matched compute, qTTT uses query updates instead of generating thinking tokens that expand the cache without changing attention.For T = 10^5 and an 8K-token budget, the paper gives approximately 16 updates with span length 128 or 8 updates with span length 512.

4 Experimental Results

Across long-context benchmarks, qTTT consistently improves performance over in-context decoding and FLOP-matched thinking, especially on retrieval-heavy tasks.

  • Synthetic and benchmark evaluation: qTTT improves performance consistently across context lengths, whereas thinking tokens show rapidly diminishing returns in the synthetic long-context setup.The evaluation compares standard in-context decoding with a compute-matched thinking baseline and qTTT.
  • Mechanistic analysis: Attention-mass analysis examines target-token scores across layers as input length increases to verify whether qTTT reduces score dilution.The analysis compares target-token attention with and without qTTT on synthetic tasks.
  • ZeroScrolls: Across six ZeroScrolls subsets and Qwen3-{1.7/4/8B} models, qTTT improves over in-context learning and FLOP-matched thinking, especially for retrieval-based multi-hop reasoning and long-form comprehension.The full benchmark suite includes diverse long-context tasks and model abilities.
  • LongBench-v2: 30.8 →43.6 on Long Dialogue History and 40.0 →46.0 on Multi-Document QA show substantial Qwen3-4B gains under compute-matched budgets.For Qwen3-8B on Code Repositories, the reported progression is 30.0 →44.0 →52.0.
  • ZeroScrolls: qTTT consistently outperforms vanilla thinking on multi-hop QA and comprehension, with stronger gains as model size increases.Improvements are smaller and comparable to thinking on summarization-style datasets, where generation quality is the primary bottleneck.
  • Additional comparisons: The study reports additional best-of-N, beam-search, latency, and wall-clock comparisons in appendices.These comparisons supplement the main benchmark evaluation.

5 Prior Work

Prior work established long-context scaling and documented retrieval failures, while inference-time scaling typically spends extra compute on decoding rather than context adaptation.

  • Long-Context LLMs: Long-context models reach million-token windows through scaling and attention-structure strategies, yet remain sensitive to evidence position and distractors.The related work highlights the “lost in the middle” effect and overwhelmed relevant spans.
  • Long-Context LLMs: The paper targets long-context retrieval failure by addressing how attention mass is allocated over very long inputs.This focus connects prior observations about distractor interference to the paper’s method.
  • Inference-Time Compute Scaling: Chain-of-thought, self-consistency, best-of-n, and related methods can help but scale decoding with compute-heavy diminishing returns.Test-time training is presented as an alternative way to spend inference compute.

6 Discussion

The discussion identifies score dilution as a cause of long-context failure and presents qTTT as a compute-efficient alternative to generating more tokens, while noting unresolved budget and task-selection questions.

  • Discussion: The paper identifies score dilution in static quadratic attention and uses query-only updates to increase the target–distractor margin.This mechanism links the theoretical diagnosis to the proposed inference-time adaptation.
  • Discussion: qTTT consistently outperforms in-context decoding and thinking under matched FLOPs, with the largest gains on retrieval and multi-hop reasoning.The paper concludes that adapting queries is more effective than generating more tokens for long-context tasks.
  • Future directions: The evaluation samples one point on the (k, N_TTT) trade-off and leaves broader budget schedules for future work.The authors also identify task-dependent gains and the need for predictors of when to prefer qTTT.
  • Future directions: Future comparisons should extend beyond thinking tokens to self-consistency and best-of-n under the same compute-matched framework.These methods are explicitly listed as future extensions.

A Synthetic Tasks

The synthetic evaluations use concise structured targets to test long-context retrieval: exact code locations for bug reports and the first violating transaction for consistency logs.

  • Code bug localization: The study presents the code bug-localization task as one representative synthetic long-context evaluation.Figure 6 introduces the task format before the example outcome.
  • Code bug localization: The code bug-localization task supplies a natural-language bug description and line-numbered code, requiring the exact offending file and line.The example identifies olmo/model.py:L345, where attention scores lack proper normalization.
  • Transaction-log consistency: The transaction-log task provides an initial account state, invariants, and transfers, requiring one bug type and the first offending transaction.The example labels NEGATIVE_BAL at TX004 when an account balance becomes negative.
  • Task design: Together, the tasks provide structured contexts with concise targets: a specific code line or a {bug_type, TX_id} pair.The same input/output format is used throughout the evaluation.

B Proofs for Section 2

The proofs formalize score dilution in static self-attention: as distractors accumulate, target attention vanishes unless the target–distractor margin grows logarithmically with context length. They also show that query-only updates move queries toward the target key and away from distractors, improving the margin.

  • Score dilution: A bounded target–distractor margin causes target probability to approach zero when the number of comparable distractors grows proportionally with context length.For m ≥ cT and constant Δ, the target attention weight tends to zero as T grows.
  • Score dilution: As distractors increase, target attention can vanish; preventing this requires a target–distractor logit margin of at least Ω(log T).The bound follows from softmax competition when many distractors have logits near the target.
  • Thinking-token limitation: For any thinking token, the target contribution is bounded by its attention weight, so diluted target attention limits what additional generated tokens can recover.The proof decomposes the attention output into the target and renormalized distractor contributions.
  • Query adaptation: A descent step updates each query toward the target key and away from the distractor mean.This directional update is the mechanism used to counteract score dilution.
  • Query adaptation: Under a Lipschitz-gradient condition and a sufficiently small step size, query-only updates monotonically improve the target margin unless target and distractor means coincide.The multi-head extension applies the same argument headwise and preserves it through concatenation and output projection.

C FLOP Derivations for §3.3

The FLOP derivation compares autoregressive thinking with query-only test-time training after a shared prefill. In long contexts, matched compute can be reallocated from generated tokens to lightweight query updates over cached keys and values.

  • Compute matching: For long contexts with T ≫ d and k ≪ T, matched compute equates generated-token cost with N_qTTT query-update cost after canceling the shared prefill.The resulting approximation uses the dominant terms in the long-context regime.
  • Cost model: The FLOP model separates shared prefill cost from added inference cost for autoregressive thinking and query-only TTT.The model uses context length T, thinking-token count T_think, update count N_qTTT, and span size k.
  • Query-only TTT: Query-only TTT recomputes queries for k positions, attends to cached {K, V}, and backpropagates only into the query projection matrices.All other parameters remain fixed during these lightweight passes.
  • Compute matching: With L=32, d=4096, r=4, T=10^5, and 8,000 thinking tokens, 10 updates of span 400 provide an example matched schedule.The reallocation keeps the KV cache length fixed while spending equivalent FLOPs on reshaping queries.

D Experimental Details

Experiments evaluate Qwen3 models across synthetic tasks and long-context benchmarks using standardized prompts, official metrics, and compute-matched thinking and qTTT settings. Results indicate that qTTT preserves attention mass and accuracy as context grows and performs strongly across domains and model sizes.

  • Setup: Qwen3-1.7B, 4B, and 8B are evaluated with native tokenizers, maximum supported context windows, standardized section-delimited prompts, and official validation or development splits.Metrics follow benchmark-specific scripts, including EM/F1, accuracy, and ROUGE variants.
  • qTTT implementation: qTTT updates only W_Q with AdamW over a learning-rate sweep, while a single prefill caches {K, V} for 32 span updates of length 128.The reported learning rate is selected on held-out validation data; gradients are clipped and computation uses bf16.
  • Setup: Compute matching uses T_think=8192 for Thinking and (k, N_TTT)=(128, 32) for qTTT, with T_think ≈ 2N_TTTk.Random span sampling and decoding seeds are fixed across methods, with one averaged run per configuration.
  • Synthetic-task results: qTTT sustains higher attention mass and accuracy as context grows, whereas Thinking-only performance and attention mass decay sharply; disabling RoPE accelerates the collapse.These comparisons are reported for Bank Transactions and OLMo Code Bugs using Qwen3-4B.
  • LongBench-v2: TTT achieves the best accuracy across LongBench-v2 domains, with largest gains on long-dialogue and document-QA tasks and benefits increasing with model size.Figure 8 summarizes the FLOP-matched comparison across six domains for Qwen3-1.7B/4B/8B.
  • Larger-model results: The same qTTT trends hold for Qwen3-32B across subsets of LongBench-v2 and ZeroScrolls.The appendix presents these results as evidence that efficacy extends across model sizes.

G Additional Test-Time Scaling Baselines

Additional baselines compare qTTT with Best-of-N/self-consistency and beam search under strict compute parity. qTTT is competitive with or better than both, while the alternatives show task-dependent or modest gains.

  • Compute matching: The baseline protocol fixes an 8,192-token extra budget, allocating approximately 8192/N tokens per self-consistency sample and 8192/k tokens per beam.All methods use the same prompt and 128-token output length.
  • Overall comparison: Under strict compute parity, qTTT is competitive with or better than Best-of-N and Beam Search across LongBench-v2 and ZeroScrolls for Qwen3-4B.The comparison uses the same added reasoning budget rather than unconstrained extra computation.
  • Best-of-N / self-consistency: Self-consistency helps when single-run accuracy is already high but often degrades when per-sample accuracy falls below 50%.The conclusion gives Single Document QA and QUALITY as examples of the high-accuracy regime.
  • Beam search: Beam search provides only modest gains under equal budgets because beams are correlated and ranking is imperfect, and it frequently trails qTTT.Beam width k receives a fixed share of the total added thinking-token budget.

H Latency and Compute-Matched Measurements

Latency measurements compare qTTT, thinking, and best-of-N under matched compute budgets across Qwen3 model sizes and context lengths. The strategies have similar wall-clock costs, while KV-cache prefilling dominates decoding time at longer sequences.

  • Setup: The latency analysis matches qTTT's 32 steps with span length k = 128 against thinking-token and best-of-N computation using a fixed FLOP budget.Nthink and NBoN are defined as compute-matched quantities, while tqTTT measures the time for 32 qTTT steps with span length 128.
  • Setup: The reported metrics separate wall-clock costs for vanilla in-context processing, thinking-token generation, best-of-N trajectories, and qTTT updates.The setup defines tICL, tthink, tBoN, and tqTTT as the corresponding single-example timings.
  • Results: Similar wall-clock time was measured for qTTT, thinking, and best-of-N under the reported fixed-FLOP comparisons across Qwen3 model sizes.Measurements covered Qwen3-1.7B, 4B, and 8B; the accompanying tables report these comparisons for each model size.
  • Results: KV-cache prefilling, approximately represented by tICL, dominates most decoding time, especially at longer sequence lengths.The setup defines tICL as the vanilla in-context pass time and reports it as roughly corresponding to prefill time.
  • Results: Strict FLOP-matched comparisons are also reported for Qwen3-4B on LongBench-v2 and ZeroScrolls.The tables specify accuracies for sampling and beam-search variants under matched token budgets.
Loading 2512.13898v1…