Reasoning in Action: MCTS-Driven Knowledge Retrieval for Large Language Models

TL;DR

The paper addresses the challenge of simultaneously achieving retrieval relevance and logical reasoning in LLMs by proposing a reasoning-aware knowledge retriever that uses a coarse-to-fine search guided by MCTS-inspired traversal. It introduces three modules—Reasoner, Concept Bridging, and Reasoning-Aware Knowledge Retrieval—to map conversation context to a context-relevant subregion ${K}_{c}$ and retrieve sentences supporting specific inferences, framed as a multi-objective optimization with an epsilon-constraint. Across two multi-turn dialogue datasets, the method yields higher alignment with human logical structure and greater knowledge diversity than semantic-only baselines, leading to more informative and creative responses. The approach demonstrates robust improvements in similarity to LLM inferences, pairwise diversity, human-logic alignment, and response quality, highlighting its practical potential for enhancing dialogue systems with reasoning-aware external knowledge. Future work may extend the framework with deeper logical reasoning modules and broader knowledge sources to further enhance open-domain generation.

Abstract

Large language models (LLMs) typically enhance their performance through either the retrieval of semantically similar information or the improvement of their reasoning capabilities. However, a significant challenge remains in effectively integrating both retrieval and reasoning strategies to optimize LLM performance. In this paper, we introduce a reasoning-aware knowledge retrieval method that enriches LLMs with information aligned to the logical structure of conversations, moving beyond surface-level semantic similarity. We follow a coarse-to-fine approach for knowledge retrieval. First, we identify a contextually relevant sub-region of the knowledge base, ensuring that all sentences within it are relevant to the context topic. Next, we refine our search within this sub-region to extract knowledge that is specifically relevant to the reasoning process. Throughout both phases, we employ the Monte Carlo Tree Search-inspired search method to effectively navigate through knowledge sentences using common keywords. Experiments on two multi-turn dialogue datasets demonstrate that our knowledge retrieval approach not only aligns more closely with the underlying reasoning in human conversations but also significantly enhances the diversity of the retrieved knowledge, resulting in more informative and creative responses.

Figures (3)

• Figure 1: Illustration of the improved response quality achieved by integrating reasoning-aware knowledge into LLM. The retrieved knowledge exhibits higher pairwise diversity and aligns more closely with the logical structure of human conversations, resulting in more informative and multifaceted responses.
• Figure 2: Illustration of our Monte Carlo Tree Search (MCTS)-inspired reasoning-aware knowledge retriever, comprising three interconnected modules: (1) Reasoner, (2) Concept Bridging, and (3) Reasoning-Aware Knowledge Retrieval. The Concept Bridging and Reasoning-Aware Knowledge Retrieval modules utilize the MCTS-inspired knowledge retrieval method, each with its own task-specific policy and critic models.
• Figure 3: Comparison of human logic alignment coverage between our method and ChatGPT-inferred knowledge across varying entailment thresholds.