Technical Summary: EASE-TTT

Problem Statement

Long-context question answering (QA) remains a significant challenge for smaller language models (LMs), even when the answer-bearing evidence is explicitly present in the input. The core bottleneck is not the model's capacity to fit the context, but its ability to reliably access and prioritize the specific evidence required to answer a question amidst distractors.

Existing approaches fail to fully resolve this issue:

• Within-Context Retrieval Methods: These localize candidate evidence chunks and modify the input (e.g., by replacing or shortening the context). However, they operate at the input level and do not adapt the model's internal attention parameters or context-access behavior. Furthermore, hard chunk selection risks discarding useful surrounding information or separating distributed evidence.
• Test-Time Training (TTT) / Query-Only TTT (qTTT): These methods adapt model parameters at inference time, specifically updating query-side attention to change how the model allocates attention. However, they typically rely on generic self-supervised objectives (e.g., next-token prediction on random spans). These objectives do not explicitly identify which context positions support the current answer, leading to a mismatch where the model adapts but not necessarily toward the relevant evidence.

There is a gap between evidence localization (finding the right chunks) and model adaptation (learning to attend to them). Existing methods either select evidence without adapting the model, or adapt the model without explicit evidence guidance.

Methodology: EASE-TTT

The authors propose Evidence-Aligned SElective Test-Time Training (EASE-TTT), a framework that bridges this gap by using retrieved evidence to guide query-side test-time adaptation. The method operates in three main stages:

1. Within-Context Evidence Selection

Given a long context $c$ and a question $q$, the system segments the context into candidate spans. It then scores these spans based on their question-conditioned utility.

• Span Detection: Boundaries are detected using spikes in token-level negative log-likelihood (NLL).
• Utility Scoring: For a candidate span $s$, the utility $r(s)$ is defined as the reduction in next-token prediction loss on the question when $s$ is prepended:
  $$r(s) = \text{LNTP}([BOS, q]) - \text{LNTP}([s, BOS, q])$$
  A positive score indicates the span likely contains relevant evidence. The top-$K$ spans are selected as the evidence set $E$.

2. Soft Attention Supervision Target

Instead of replacing the full context with the selected chunks $E$, EASE-TTT constructs a soft attention target distribution $\pi$ over all context token positions.

• Positions covered by the selected evidence chunks ($\Omega(E)$) are assigned a higher probability mass ($\alpha/|\Omega(E)|$).
• Remaining positions are assigned a lower mass ($(1-\alpha)/(T - |\Omega(E)|)$).
• This soft target ($\alpha \in (0,1)$) guides the model to focus on evidence while preserving a non-zero probability for the rest of the context, avoiding the risks of hard masking.

3. Query-Side Test-Time Adaptation

The framework performs lightweight adaptation at inference time:

• Parameter Updates: Only the query-side projection parameters (via LoRA adapters) in the self-attention layers are updated; the base model and key-value representations remain frozen.
• Objective: The adaptation minimizes the Kullback-Leibler (KL) divergence between the model's current attention distribution $a$ and the soft target $\pi$:
  $$L_{attn} = D_{KL}(\pi \parallel a)$$
• Final Generation: After adaptation, the model generates the final answer using the original full context, ensuring that the retrieved evidence serves as a supervision signal for attention alignment rather than a truncated input.

Key Contributions

1. Identification of a Bottleneck: The paper identifies that for smaller LMs, the failure in long-context QA is often an "evidence-use failure" where relevant information is present but inaccessible due to distractors, rather than a lack of context capacity.
2. EASE-TTT Framework: The authors propose a novel framework that converts question-relevant chunks into soft attention supervision for query-side adaptation. Unlike retrieval-only methods, it preserves the full context for generation while using evidence to guide the model's attention behavior.
3. Empirical Validation: Extensive evaluation across six LongBench QA tasks and three small decoder-only models (Qwen3-0.6B, Qwen3-1.7B, Llama-3.2-1B) demonstrates that EASE-TTT outperforms full-context inference, retrieval-only baselines, and standard qTTT.

Experimental Results

• Performance: EASE-TTT achieved the strongest macro-average performance across all tested models.
  • On Qwen3-0.6B, it improved the average score by 4.1 points over full-context inference and 1.2 points over qTTT.
  • On Qwen3-1.7B, it improved by 5.6 points over full-context and 1.9 points over qTTT.
• Comparison with Baselines:
  • vs. Retrieval-Only: Retrieval-only methods (RAG, ICR) showed inconsistent gains and sometimes underperformed full-context inference, likely due to the loss of surrounding context.
  • vs. qTTT: EASE-TTT consistently outperformed qTTT, indicating that evidence-localized supervision is more effective than generic span-based self-supervision for long-context QA.
• Ablation Studies:
  • Loss Objective: Using explicit attention alignment (Attn. KL) significantly outperformed using selected chunks for standard next-token prediction (Chunk NTP), confirming that the value lies in guiding attention, not just exposing data.
  • Attention Layer: Supervising intermediate attention layers yielded better results than early or final layers, suggesting a trade-off between information gathering and consolidation.
• Efficiency: EASE-TTT incurs a moderate latency increase (approx. 2.4s per example) compared to qTTT due to evidence selection and attention map computation, trading this for improved accuracy.

Significance and Claims

The paper claims that EASE-TTT supports evidence-aligned test-time adaptation as a promising direction for smaller long-context models. Its significance lies in:

• Bridging the Gap: It successfully integrates within-context retrieval (localization) with test-time training (adaptation), addressing the limitation where existing methods handle only one side of the problem.
• Preserving Context: By using evidence as a supervision signal rather than an input replacement, it avoids the information loss associated with hard chunk selection while still correcting the model's attention biases.
• Modest Scope: The authors explicitly note that their experiments focus on long-context QA tasks where evidence is expected in the input. They do not claim the method generalizes to other task types (e.g., mathematical reasoning) or larger models without further evaluation, acknowledging that larger models may already possess stronger inherent context utilization abilities.