RecaLLM: Addressing the Lost-in-Thought Phenomenon with Explicit In-Context Retrieval

AI summary

Authors: Kyle Whitecross, Negin Rahimi (University of Massachusetts Amherst) arXiv: 2604.09494v1 [cs.CL], 10 Apr 2026 — CC BY 4.0 Code, data, and models: https://github.com/kswhitecross/RecaLLM

Overview

• Research area: Long-context large language model reasoning, in-context retrieval, and reinforcement-learning-based post-training.
• Technical level: Advanced (assumes familiarity with chain-of-thought reasoning, reinforcement learning from verifiable rewards, GRPO, constrained decoding, and long-context benchmarks such as RULER and HELMET).
• Scope in one sentence: The paper diagnoses a failure mode in reasoning LLMs called "lost-in-thought," then proposes RecaLLM, a family of models post-trained to interleave reasoning with explicit, constrained-decoding-guaranteed retrieval spans, achieving strong performance on RULER and HELMET while training on contexts of at most 10K tokens.

What This Paper Is About

Long-context reasoning models get worse at pulling relevant evidence out of their input context once they have produced a long chain of thought. The authors run preliminary experiments on open-source LLMs showing that in-context retrieval accuracy drops substantially after even a short reasoning span, a bottleneck they name "lost-in-thought." RecaLLM addresses this by breaking generation into alternating segments of reasoning and explicit "recall spans" that copy evidence verbatim from either the input context or the model's own prior generation, so that a long-context retrieval problem becomes a local reasoning problem.

Key Contributions

1. Identification of lost-in-thought: The authors characterize a failure mode in reasoning LLMs where the reasoning steps that improve performance also make subsequent in-context retrieval harder, and show via an injection experiment that the degradation stems largely from an inability to faithfully copy context information after reasoning rather than from failing to identify what to retrieve.
2. RecaLLM: A model family that interleaves reasoning with explicit in-context retrieval through recall spans delimited by two new special tokens, <|start_recall|> and <|end_recall|>, allowing retrieval needs that arise dynamically mid-reasoning rather than only being planned upfront.
3. Context-aware constrained decoding: A negligible-overhead decoding mechanism that restricts generation inside recall spans to valid continuations of the searchable context (the input prompt plus all previously generated tokens), guaranteeing by construction that every recall span is a contiguous substring of that context and thereby making retrieval directly verifiable.
4. Efficient long-context generalization: RecaLLM delivers strong long-context gains, especially on retrieval-intensive tasks, and generalizes from training contexts of at most 10K tokens to evaluations up to 128K tokens, outperforming larger models trained on longer contexts.

Main Findings

• Retrieval degrades after reasoning across all tested models: In a synthetic key-value benchmark spanning 4K to 128K tokens, five open-source 7–8B models — Llama-3.1-8B-Instruct, R1-Distill-Llama-8B, ProLong-8B-512K, Qwen2.5-7B-Instruct, and Qwen3-8B — all show substantially worse retrieval accuracy after reasoning than on direct retrieval, even though math accuracy remains high and stable.
• The bottleneck is faithful copying, not identification: In an injection experiment, models were truncated mid-generation, given the correct key and its exact lexical prefix, and still frequently hallucinated the value. Llama-3.1-8B-Instruct generated the correct value 40.7% of the time under injection, against 80.6% on direct Retrieval at short contexts; at long contexts the recovery was smaller still.
• Best-in-class RULER average for the 7–8B class: RecaLLM-Qwen2.5-7B achieves an average of 92.8 on RULER, the best among 7–8B models, surpassing larger long-context baselines including LoongRL-14B (91.9) and QwenLong-L1-32B (84.2).
• Gains grow with context length on RULER: For Qwen, the improvement over the base model grows from +6.8 at 4K to +16.1 at 128K, and RecaLLM-Qwen reduces the 4K-to-128K degradation from 25.2 points to 15.9. For Llama, the largest gain also appears at 128K (+24.2).
• Strong HELMET improvements: RecaLLM improves its base models by 16.1–17.7 points on average on HELMET, attaining the strongest overall results in the 7–8B class. RecaLLM-Llama rises from 3.0 to 64.1 on ICL and from 21.3 to 53.2 on Re-rank, showing that explicit recall improves reasoning over retrieved evidence, not just evidence retrieval.
• Large in-domain gains on reasoning-then-retrieval tasks: On reasoning-retrieval, RecaLLM-Qwen improves from 23.0% to 97.6% at short contexts and from 7.5% to 86.3% at long contexts. On entity citation, RecaLLM-Llama improves from 4.1% to 83.6% at short contexts and from 0.3% to 71.5% at long contexts. Short-context math improves by 5.2 points for RecaLLM-Qwen, and aggregation benefits substantially despite not requiring explicit recall.
• Smaller, less consistent gains on long-form tasks: Gains on HELMET's LongQA and Summ are smaller; RecaLLM-Qwen improves slightly on both, whereas RecaLLM-Llama maintains LongQA performance but declines on Summ. Both categories use LLM-as-a-judge evaluation and require composing evidence into long-form output.
• Training far below evaluation length: RecaLLM is trained on contexts of at most 10K tokens (8–10K in the RL mixture) yet improves consistently up to 128K. For comparison, ProLong is trained on sequences up to 512K tokens, exceeding its 128K evaluation length, and LoongRL trains at 16K.
• Accuracy tracks recall usage: RecaLLM-Qwen maintains high recall usage over nearly the entire context range and sustains strong performance; RecaLLM-Llama shows a clear drop in recall usage at 96K and 128K with much larger accuracy drops, especially on reasoning-retrieval. The authors interpret this as policy drift away from invoking recall, with the model reverting to implicit retrieval.
• Both training components matter, and they are complementary: On validation sets averaged over all context lengths, RecaLLM-Qwen2.5-7B scores 71.3 on average. Removing the recall reward drops this to 65.5, with especially large losses on aggregation, reranking, entity citation, and QA. Removing logit masking yields a smaller decline to 69.4, but causes much larger losses on retrieval and reasoning-retrieval, the categories requiring exact lexical matching; it improves performance on categories where exact lexical retrieval is less critical.

Methodology in Plain English

The authors start by building a controlled synthetic benchmark to isolate the problem. The prompt contains a large structured key-value dictionary, and the model faces one of two tasks: Retrieval, where a query key is given directly, or Reasoning-Retrieval, where a math problem must be solved first because its answer is the query key. Context lengths from 4K to 128K are controlled by varying the number of distractor key-value pairs, while prompt templates, query placement, math problem type, and dictionary format (CSV, JSON, list) are varied for robustness.

To fix the problem, RecaLLM adds two special tokens that mark the start and end of a recall span. While inside a span, the model is forbidden from emitting any token that would not continue an exact occurrence of the recalled prefix somewhere in the searchable context — the input prompt plus all tokens generated so far. The model therefore decides what and when to recall, while the decoding constraint guarantees the recall is exact. In practice this means retrieval becomes verifiable: a recalled span either appears verbatim in the context or it does not.

Training proceeds in two stages. First, a supervised fine-tuning cold start teaches the new token embeddings and then briefly fine-tunes the full model on traces annotated with recall tokens. These traces come from six teacher models across four reasoning and retrieval tasks, retained only when they produce the correct answer, then rewritten by GPT-5.2 so that references to context become verbatim recall spans. Spans are aligned to source text via fuzzy string matching, and failed alignments are discarded, yielding 1,795 annotated examples. Second, the model is trained with GRPO on a deliberately shallow mixture of 20,000 examples across 10 task categories, from multi-hop QA and retrieval through short-context math, aggregation, reranking, and entity citation. Data augmentation varies instruction phrasing, question placement, passage format, and distractor type (random, BM25, and dense-retrieval negatives) so the model does not learn recall tied to narrow surface patterns.

The reward combines three signals: format (weight 0.2), an additive term combining answer quality and retrieval quality (weight 0.4), and a smoothed geometric mean of answer and retrieval quality (weight 0.4, with epsilon = 0.01). Retrieval reward is computed as an interval-overlap F1 between recalled spans and gold passages, normalized by a hit threshold. Because constrained decoding guarantees verbatim reproduction, span and passage can both be located as character-level position intervals, which makes this reward precise; without the constraint, a single non-verbatim character would break the contiguous match. A density penalty downweights reward when recall spans become too frequent, and a correctness penalty detects malformed spans. Training uses 4 A100 GPUs for a single epoch with 16 rollouts per example.

Evaluation uses validation splits of the training datasets across 4K–128K contexts, plus RULER and HELMET, comparing against the base Llama-3.1-8B-Instruct and Qwen2.5-7B-Instruct models, ProLong, LoongRL, and QwenLong. All RULER and HELMET tasks are evaluated using chat templates rather than answer-only prompts, to avoid disadvantaging reasoning models.

Why This Matters

The paper reframes long-context ability as a retrieval problem that can be improved at post-training time, rather than purely a matter of extending the context window through pretraining or continued training. It shows that a 7B model trained on 8–10K-token contexts can beat a 14B model trained at 16K and a 32B model trained up to 60K on RULER, which points to a substantially cheaper path for building long-context systems. Because recall spans make retrieval a verifiable event rather than a latent behavior, the work also opens a route to more interpretable and rewardable grounding.

Real-world applications:

• Long-document analysis: Legal, financial, or scientific workflows where the answer lies in one passage of a long filing or report, and where hallucinating a figure is worse than failing to answer.
• Retrieval-augmented generation pipelines: Agentic search systems that pull documents into context would benefit from stronger in-context use of the evidence they already gathered; the authors describe RecaLLM as complementary to such systems.
• Tool and database lookup agents: Tasks where the model must reason to derive an identifier and then fetch the exact associated value, mirroring the paper's reasoning-retrieval benchmark.
• Enterprise in-context learning with arbitrary labels: HELMET's ICL tasks use randomly generated numeric labels rather than semantic class names, and RecaLLM-Llama improved from 3.0 to 64.1 there, suggesting value wherever exact example-label pairs must be retrieved rather than inferred.

Industry relevance centers on cost: if long-context gains can be obtained from 8–10K-token training data, organizations avoid the expense of curating and training on 512K-token sequences. A single-epoch GRPO run on 4 A100 GPUs is a comparatively accessible post-training budget.

Future Directions

• More flexible faithful retrieval: The conclusion calls for retrieval mechanisms that preserve faithfulness without the rigidity of strict constrained decoding, which the ablation shows can restrict useful flexibility on tasks where exact lexical matching matters less.
• Self-recall over the model's own prior generation: Extending recall to operate over the model's earlier output for long-horizon reasoning consistency. The introduction notes this is an increasingly important but underexplored setting, since retrieval needs can emerge after several intermediate steps and cannot be planned upfront.
• Scaling to larger, more capable models: The authors expect larger models to learn better when and how to invoke retrieval, and to close the gap seen on long-form LongQA and Summ tasks where gains were smaller.
• New reward design for verifiable retrieval: Making retrieval an explicit, verifiable step opens opportunities for reward design, interpretability, and trustworthy grounding in long-context applications. A related open question is the observed policy drift away from recall at 96K and 128K for RecaLLM-Llama.

Target Audience

This paper is most useful to researchers and engineers working on long-context language models, reinforcement learning post-training, and retrieval-augmented generation. It will also interest practitioners who need models to ground answers in long inputs without paying for long-context training data, and readers studying constrained decoding, reward design for reasoning models, or chain-of-thought faithfulness. A background in LLM training and evaluation is needed to follow the reward formulation and the ablation analysis; the core diagnosis and benchmark design are accessible to a broader technical audience.

Authors’ abstract

We propose RecaLLM, a set of reasoning language models post-trained to make effective use of long-context information. In-context retrieval, which identifies relevant evidence from context, and reasoning are deeply intertwined: retrieval supports reasoning, while reasoning often determines what must be retrieved. However, their interaction remains largely underexplored. In preliminary experiments on several open-source LLMs, we observe that in-context retrieval performance substantially degrades even after a short reasoning span, revealing a key bottleneck for test-time scaling that we refer to as lost-in-thought: reasoning steps that improve performance also make subsequent in-context retrieval more challenging. To address this limitation, RecaLLM interleaves reasoning with explicit in-context retrieval, alternating between reasoning and retrieving context information needed to solve intermediate subproblems. We introduce a negligible-overhead constrained decoding mechanism that enables verbatim copying of evidence spans, improving the grounding of subsequent generation. Trained on diverse lexical and semantic retrieval tasks, RecaLLM achieves strong performance on two long-context benchmarks, RULER and HELMET, significantly outperforming baselines. Notably, we observe consistent gains at context windows of up to 128K tokens using training samples of at most 10K tokens, far shorter than those used by existing long-context approaches, highlighting a promising path toward improving long-context performance without expensive long-context training data.