cs.CLApr 10, 2026
SaveRecaLLM: Addressing the Lost-in-Thought Phenomenon with Explicit In-Context Retrieval
Organizations: University of Massachusetts Amherst
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.
Figures & tables
Figure 1: Illustration of lost-in-thought and how RecaLLM mitigates it with explicit, faithful recall spans . A reasoning model may recover the correct key yet still hallucinate the value.
Figure 2: Lost in thought : retrieval accuracy before and after reasoning. Injected accuracy measures faithful copying after providing the correct key and prefix.
| Category | Datasets | Count | Context Length |
|---|---|---|---|
| Multi-Hop QA | HotpotQA, MuSiQue, 2WikiMQA | 3,000 | 8–10K |
| Single-Hop QA | NQ, TriviaQA | 2,000 | 8–10K |
| Retrieval | KV Retrieval, Multi-NIAH | 2,000 | 8–10K |
| Reasoning Retrieval | Math Retrieval | 2,000 | 8–10K |
| Short-Context Math | DAPO Math, MCQA Math | 1,000 | 1K |
| In-Context Learning | Banking77, MASSIVE | 3,000 | 8–10K |
Table 1: RL training dataset composition. Detailed descriptions in Appendix C.4 .
Figure 3: Per-category answer scores across context lengths (4K–128K) on the in-domain evaluation. RecaLLM models maintain strong performance as context length increases, while baselines degrade sharply beyond 32K tokens.
| RULER | HELMET | ||||||||||||||
| Model | 4K | 8K | 16K | 32K | 64K | 128K | Avg | Recall | RAG | ICL | Cite | Re-rank | LongQA | Summ | Avg |
| ProLong-8B-Instruct-512K | 92.3 | 86.0 | 73.7 | 59.5 | 60.5 | 46.8 | 69.8 | 74.3 | 68.7 | 18.6 | 12.4 | 12.5 | 36.4 | 36.5 | 37.0 |
| QwenLong-L1-32B | 91.7 | 88.5 | 87.6 | 86.8 | 80.6 | 70.2 | 84.2 | — | 52.9 | — | 14.3 | — | — | 47.0 | — |
| LoongRL-7B | 95.1 | 94.3 | 93.4 | 91.4 | 86.2 | 76.8 | 89.5 | — | 61.6 | — | 11.4 | — | — | 44.5 | — |
| LoongRL-14B | 97.6 | 96.1 | 95.4 | 95.1 | 87.1 | 79.9 | 91.9 | — | 63.7 | — | 18.7 | — | — | 49.8 | — |
| Llama-3.1-8B-Instruct | 95.7 | 93.5 | 89.8 | 85.6 | 78.6 | 42.2 | 80.9 | 77.0 | 63.0 | 3.0 | 10.0 | 21.3 | 36.3 | 39.2 | 35.7 |
Table 2: Performance on RULER ( Hsieh et al., 2024 ) (averaged across tasks) and HELMET ( Yen et al., 2025a ) (averaged across context lengths). Models with — indicate results not reported in their papers.
Figure 4: Accuracy and recall token usage rate of RecaLLM models across context lengths on the retrieval and reasoning-retrieval tasks from Section 2 .
| Model | 1-Hop | N-Hop | Ret. | R Ret. | Math | ICL | LQA | Agg. | Rerank | Ent. Cite | Avg |
|---|---|---|---|---|---|---|---|---|---|---|---|
| RecaLLM-Qwen2.5-7B | 63.4 | 62.2 | 96.7 | 92.8 | 61.0 | 77.9 | 72.2 | 86.5 | 38.8 | 61.8 | 71.3 |
| – No Recall Reward | 53.8 | 54.8 | 94.8 | 88.6 | 62.0 | 75.8 | 76.0 | 70.1 | 28.0 | 51.6 | 65.5 |
| – No Logit Masking | 69.5 | 70.2 | 89.5 | 72.1 | 58.2 | 78.4 | 72.2 | 90.6 | 35.3 | 58.1 | 69.4 |
Table 3: Performance of ablated variants of RecaLLM-Qwen on validation sets, averaged over all context lengths.
Appendix figures & tables8 assets
Supplementary material from the paper’s appendix.
Appendix
| Reas.-Retr. | Injected | Retrieval | Rel. Gap (%) | Cov. | Reas. Toks | |||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|
| Model | S | L | S | L | S | L | S | L | S | L | S | L |
| Llama-3.1-8B-Instruct | 27.1 | 2.2 | 40.7 | 7.9 | 80.6 | 25.5 | 66.3 | 91.4 | 95.8 | 60.7 | 1264 | 866 |
| R1-Distill-Llama-8B | 32.0 | 2.0 | 32.8 | 5.7 | 61.1 | 15.8 | 47.7 | 87.2 | 41.0 | 25.5 | 1973 | 823 |
| ProLong-8B-512K | 7.7 | 0.0 | 22.2 | 9.2 | 79.5 | 30.3 | 90.3 | 100.0 | 65.9 | 37.1 | 91 | 184 |
| Qwen2.5-7B-Instruct | 33.0 | 11.9 | 37.7 | 16.5 | 67.1 | 33.8 | 50.9 | 64.7 | 97.0 | 80.3 | 828 | 526 |
| Qwen3-8B (No Think) | 47.6 | 8.8 | 53.2 | 16.2 | 82.2 | 43.2 | 42.2 | 79.7 | 97.0 | 99.0 | 236 | 230 |
Table 4: Lost in thought and injection analysis. S hort = 4K–32K average; L ong = 64K–128K average. Reas.-Retr. and Retrieval report accuracy on the Reasoning-Retrieval and Retrieval datasets respectively; Rel. Gap is the relative difference between them. Injected reports accuracy after injecting the correct key and its exact lexical prefix mid-generation, showing that even with the exact prefix provided, retrieval recovers only partially. Cov. is the fraction of correct-reasoning examples where the model attempts retrieval; Reas. Toks is the mean extra reasoning tokens compared to the Retrieval baseline.
| Category | Answer Reward | Recall Reward | # Gold | Free Spans | |
| Multi-Hop QA | subEM | Standard F1 | 2–4 | 0.4 | 4 |
| Single-Hop QA | subEM | Top-1 F1 | 1 | 0.4 | 2 |
| Retrieval | subEM (KV) / Net Recall (NIAH) | Standard F1 | 1–6 | 0.9 | 2 (KV) / 6 (NIAH) |
| Reasoning Retr. | subEM | Standard F1 | 1 | 0.9 | 2 |
| Short-Ctx Math | Exact match | Always 1.0 | 0 | — | 2 |
| In-Context Learn. | Exact match | Top-2 F1 | 4–12 | 0.95 | 2 |
Table 5: RL reward configuration per category. Free spans indicates the task-dependent number of initial recall spans exempt from the density penalty.
| Datasets | Description | Additional Augmentations |
|---|---|---|
| HotpotQA, MuSiQue, 2WikiMQA | Multi-hop QA requiring 2–4 supporting passages | Document format, negative source, passage type (paragraphs vs. fixed-window chunks) |
| NQ, TriviaQA | Single-hop factoid QA with one gold passage | Document format, negative source |
| KV Retrieval | Structured key-value lookup | KV format (CSV / JSON / lines) |
| Multi-NIAH | Multi-needle-in-a-haystack retrieval among distractors | Distractor count, target keys, values per key, value type |
| Math Retrieval | Solve a math problem to determine a retrieval key, then look up its value | KV format (CSV / JSON / lines), math problem type |
| DAPO Math, MCQA Math | Short, difficult math problems with no retrieval required | None |
Table 6: Training dataset descriptions and per-dataset augmentations. All long-context datasets share three base augmentations: instruction template variation, question position variation, and gold position randomization for any dataset with known gold information. Additional per-dataset augmentations are listed in the rightmost column.
| 1-Hop | N-Hop | Ret. | R Ret. | Math | ICL | LQA | Agg. | Rerank | Ent. Cite | |||||||||
| Model | S | L | S | L | S | L | S | L | S | S | L | S | S | L | S | L | S | L |
| ProLong-8B-512K | 21.7 | 5.8 | 16.6 | 4.7 | 61.5 | 20.7 | 1.8 | 0.0 | 19.0 | 1.6 | 0.1 | 37.3 | 72.4 | 52.5 | 6.3 | 1.0 | 2.0 | 1.4 |
| Llama-3.1-8B-Instruct | 59.9 | 10.3 | 50.2 | 10.8 | 81.5 | 29.5 | 13.3 | 1.2 | 37.8 | 24.2 | 4.4 | 73.0 | 60.1 | 45.8 | 24.9 | 5.6 | 4.1 | 0.3 |
| RecaLLM-Llama-8B | 62.4 | 57.8 | 63.8 | 51.0 | 98.3 | 86.8 | 90.9 | 51.7 | 46.5 | 74.4 | 63.2 | 72.5 | 92.7 | 81.4 | 48.5 | 34.0 | 83.6 | 71.5 |
| Qwen2.5-7B-Instruct | 53.6 | 33.8 | 42.0 | 28.5 | 73.6 | 47.1 | 23.0 | 7.5 | 55.8 | 31.3 | 18.7 | 71.0 | 45.8 | 45.5 | 22.5 | 10.9 | 5.2 | 1.4 |
| RecaLLM-Qwen2.5-7B | 68.8 | 56.3 | 69.9 | 51.9 | 98.3 | 94.5 | 97.6 | 86.3 | 61.0 | 76.7 | 79.4 | 72.3 | 87.7 | 85.0 | 44.6 | 31.0 | 72.9 | 47.2 |
Table 7: In-domain evaluation results (%) across training task categories at short (4K–32K) and long (64K–128K) context lengths. Categories with only one column are short-context only. Bold indicates best in each model family.
Figure 10: Training dynamics over 150 GRPO steps for both RecaLLM models. (a) Overall score rises rapidly, with Llama converging faster than Qwen. (b) Recall span usage starts very high, then decreases and stabilizes as the models learn selective retrieval. (c) Gold document overlap increases alongside score, indicating that models learn to recall relevant evidence rather than arbitrary context. (d) Per-category breakdown for Qwen, showing that all categories improve, albeit at different rates.
Figure 11: Per-category answer scores (solid) and recall usage rates (dotted) for RecaLLM-Qwen2.5-7B and two ablations across context lengths. No Mask follows the same training procedure but disables constrained decoding throughout, allowing recall spans to generate freely. No Recall Reward uses full constrained decoding but sets unconditionally during RL, removing gold document supervision.
| Model | 4K | 8K | 16K | 32K | 64K | 128K | Avg |
|---|---|---|---|---|---|---|---|
| RecaLLM-Qwen2.5-7B | 98.8 | 96.7 | 95.8 | 93.1 | 89.7 | 82.9 | 92.8 |
| – No Recall Reward | 96.7 | 93.3 | 89.9 | 88.0 | 81.4 | 77.9 | 87.9 |
| – No Logit Masking | 98.1 | 95.0 | 95.4 | 93.0 | 88.2 | 82.3 | 92.0 |
Table 8: Ablation results on RULER across context lengths.
| Model | Recall | RAG | ICL | Cite | Re-rank |
|---|---|---|---|---|---|
| RecaLLM-Qwen2.5-7B | 96.2 | 69.4 | 69.7 | 13.7 | 46.2 |
| – No Recall Reward | 94.9 | 66.1 | 69.0 | 14.9 | 25.3 |
| – No Logit Masking | 87.7 | 78.3 | 76.7 | 15.4 | 44.9 |
Table 9: Ablation results on HELMET, averaged over context lengths from 8K to 128K. LongQA and Summarization are excluded due to the cost of LLM-as-a-judge evaluation over long outputs.