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

TL;DR

RecaLLM is a reasoning language model that interleaves explicit in-context retrieval with reasoning steps to improve long-context understanding and address the lost-in-thought phenomenon, achieving strong results on benchmarks.

Contribution

It introduces a novel interleaving approach of reasoning and retrieval, along with a constrained decoding mechanism, to enhance long-context reasoning without extensive long-context training data.

Findings

RecaLLM significantly outperforms baselines on RULER and HELMET benchmarks.

It maintains strong performance with context windows up to 128K tokens using relatively short training samples.

Explicit in-context retrieval mitigates the lost-in-thought problem in long-context reasoning.

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.