Search-in-Context: Efficient Multi-Hop QA over Long Contexts via M onte C arlo Tree Search with Dynamic KV Retrieval (2025)

Abstract: Recent advancements in large language models (LLMs) have demonstrated remarkable capabilities in complex reasoning tasks, such as math problem-solving and code generation. However, multi-hop question answering (MHQA) over long contexts, which demands both robust knowledge-intensive reasoning and efficient processing of lengthy documents, remains a significant challenge. Existing approaches often struggle to balance these requirements, either neglecting explicit reasoning or incurring expensive computational costs due to full-attention mechanisms over long contexts. To address this, we propose **Search-in-Context (SIC)**, a novel framework that integrates Monte Carlo Tree Search (MCTS) with dynamic key-value (KV) retrieval to enable iterative, context-aware reasoning. SIC dynamically retrieves critical KV pairs (e.g., 4K tokens) at each step, prioritizing relevant evidence while mitigating the “lost in the middle” problem. Furthermore, the paper introduces a Process-Reward Model (PRM) trained on auto-labeled data to guide the MCTS process with stepwise rewards, promoting high-quality reasoning trajectories without manual annotation. Experiments on three long-context MHQA benchmarks (HotpotQA, 2WikiMultihopQA, MuSiQue) and a counterfactual multi-hop dataset demonstrate SIC’s superiority, achieving state-of-the-art performance while significantly reducing computational overhead.

## introduction recent advancements in large language models (llms) @xcite have significantly improved their capability to tackle complex, reasoning-intensive tasks across diverse domains, including mathematical problem-solving @xcite , repository-level code generation and correction @xcite * corresponding author @xcite , and scientific reasoning @xcite . such achievements highlight their growing capacity to handle sophisticated tasks that were previously thought to require human-level expertise. despite these advances, applying llms to multihop question answering (mhqa) over long contexts remains a significant challenge @xcite , as it requires models to simultaneously satisfy two key capabilities: strong knowledgeintensive reasoning capability: the model must effectively integrate information across multiple reasoning hops, effectively synthesizing relevant information from intermediate subquestions to enable knowledge-driven inference @xcite . robust long-context processing capability: the model must efficiently handle extensive contexts (often exceeding 10k tokens) while filtering out irrelevant or distracting information @xcite , ensuring the accurate identification and extraction of key information across length context necessary for answering the question. current approaches struggle to meet these dual requirements simultaneously. many existing methods rely on long-context llms to directly answer mhqa tasks @xcite , underestimating the task's complexity and overlooking the importance of explicit test-time reasoning. to bridge this gap, some works adopt chainof-thought (cot) prompting techniques @xcite @xcite . unlike reasoning tasks in mathematics or science, where contexts are typically limited to a few hundred tokens, mhqa over long contexts requires models to generate reasoning chains based on inputs exceeding 10k tokens. in such scenarios, generating a reasoning chain necessitates computing full attention over an increasingly large key-value (kv) cache at each decoding step, leading to quadratic computational complexity growth @xcite . furthermore, models often struggle with vast amounts of distractive information, a phenomenon commonly referred to as the "lost in the middle" problem @xcite . in this paper, we propose a novel framework that integrates monte carlo tree search (mcts) with dynamic kv retrieval to empower llms with iterative, context-aware exploration over long contexts. inspired by test-time scaling algorithms @xcite , our method formulates the retrieval and reasoning as a search tree process, where each node represents a potential reasoning step guided by contextually retrieved evidence. at each iteration, the model selectively utilizes portions of critical kvs (e.g., 4k token budgets) based on a specialized kv retriever, rather than relying on computationally expensive fullattention over the entire kv cache. this process effectively retrieves and prioritizes the most relevant kvs that contribute to uncovering critical information for subsequent reasoning hops, mitigating the "lost in the middle" problem in long context and improving reasoning efficiency. additionally, we incorporate a process-reward model (prm) into the mcts process to guide the model's reasoning. this prm provides step-by-step rewards to encourage the model to follow high-quality reasoning paths. importantly, the prm can be trained using automatically labeled data without requiring manual annotation, ensuring scalability and reducing human intervention. the main contributions of this paper can be summarized as follows: long-context language modeling. scaling llm to process long texts poses significant challenges due to the quadratic computational complexity of attention mechanisms @xcite . to mitigate the computational and memory constraints, recent research has explored various kv compression techniques @xcite . these methods selectively retain subsets of kvs based on predefined reduction strategies, often compressing them to a fixed budget @xcite @xcite . for instance, h2o @xcite employs a policy that discards kvs during generation according to a scoring function derived from cumulative attention. infllm @xcite partitions kvs into fixed-size chunks and retains the top-k most salient chunks based on attention score patterns. ## preliminary in this section, we provide a formal and descriptive definition of our task. problem formulation. multi-hop question answering (mhqa) over long context is a complex reasoning task requiring iterative reasoning across multiple, often disparate sources of information over documents to deduce an answer. we formulate this task as follows: given the input which contains the query q and contexts @xmath0, c 2 , ..., c n } where c i represents a single and independent document, the task is aimed at predicting an answer @xmath1: ## conclusion in this paper, we propose sic, a novel framework integrating monte carlo tree search (mcts) with dynamic key-value (kv) retrieval to address the dual challenges of efficiency and reasoning in large language models (llms) for multi-hop question answering (mhqa) over long contexts. by modeling the reasoning process as a search tree and incorporating dynamic kv retrieval, sic iteratively focuses on critical contextual segments (e.g., 4k tokens per step), mitigating the "lost in the middle" problem while reducing computational complexity. our comprehensive experiments across two models and four datasets validate the superiority of sic in long-context multi-hop qa tasks.