MCTS-KBQA: Monte Carlo Tree Search with Information Gain Rewards for Knowledge Base Question Answering

When an LLM uses Monte Carlo tree search (MCTS) to find answers step by step over a knowledge base, intermediate steps can be given useful scores without training a reward model and without simulating to the end. The method gives the query history up to that step to an open instruction model and computes the perplexity of that history with and without the original question attached. The more the perplexity drops once the question is attached, the more relevant the findings are to the question, and the higher the reward for that step. Earlier options either simulate every step to the end and score the final state, which is expensive, or train a separate reward model, which needs a lot of data. The new method needs only two forward passes per new node. On 4 KBQA (knowledge base question answering) benchmarks it consistently beats methods that follow a single path without search, and compared with classic tree search that simulates to the end it usually offers a better trade-off between accuracy and cost.

People working on agentic retrieval or multi-step tool calling should read this, because it offers a way to score intermediate steps without extra training: no step-level labels and no reward model training, at the cost of two extra forward passes per new node. The core idea in one sentence: the score only asks whether what was found makes the original question easy to guess. A small SFT model proposes the next query, and a separate open instruction model with no fine-tuning does the scoring. On the complex-question benchmark CWQ (ComplexWebQuestions), F1 goes from 37.47 with a single path to 70.15, above the 63.44 of classic tree search that simulates to the end, with only half as many generation calls as classic tree search.

What problem is it solving?

The knowledge bases here are structured ones such as Freebase or Wikidata, and the final output is an executable query written in the knowledge-base query format SPARQL. The paper follows Interactive-KBQA: the LLM acts as an agent, and at each step it writes a short thought, calls a tool, and gets back what the knowledge base returns. There are four tools: find an entity (SearchNodes), look up which relations an entity has (SearchGraphPatterns), run a SPARQL query (ExecuteSPARQL), and stop (Done). GrailQA adds a fifth tool, SearchTypes, for looking up types.

Take the question from Figure 1 of the paper: "What is the TV program with the largest tvrage_id that Shawnee Smith played in?" After the first step finds the actress as an entity, there are several possible next steps:

An agent that follows a single path cannot recover after one wrong choice. Tree search can keep several branches alive, but it must answer one question: before the final SPARQL is written, which branch looks more promising? That is the intermediate-step scoring problem.

The big picture: classic MCTS and Fast MCTS

Each node in the tree stores one step's thought, tool call and returned result, and the path from the root to a node is that branch's full history. Each search round does four things:

1. Selection: walk down from the root and pick a node with UCT, the selection algorithm commonly used in tree search.
2. Expansion: sample 5 candidate next steps from the agent at this node, call the tool for each, and get 5 child nodes.
3. Scoring: give each new child a reward.
4. Backpropagation: add the reward back to every ancestor along the path.

The two kinds of MCTS differ only in step 3. Classic MCTS runs one simulation for each new child that has not finished: the agent keeps generating along one path until Done or the maximum depth, and only that end point is scored. In the paper's experiments the end point is also scored with the IG (information gain) reward described below, without the gold answer, so the two differ only in whether they simulate to the end. One simulation needs several generation calls in a row. Fast MCTS skips the simulation and scores the child directly with the IG reward, which needs only two forward passes.

Each question runs at most 12 rounds of selection; search stops early once 5 terminal branches that execute and return non-empty results have been collected. The final answer is a vote over the execution results of those branches.

The IG reward: does what was found make the question easy to guess?

The paper calls this score the IG reward. The intuition: if a branch has found things closely related to the question, then after reading that query history, the original question becomes "easy to guess".

Concretely, a frozen open instruction model (by default the original Llama-3.1-8B-Instruct, with no fine-tuning) computes perplexity twice.

• First pass: the query history followed by a fixed prompt, "Based on this interaction history, the original question is:", and the perplexity of that text.
• Second pass: the same history and prompt, followed by the original question, and the perplexity of the whole text.

If the history contains things like "TV appearances" and a list of series, the question "...TV program ... played in" reads naturally after it, and the perplexity of the whole text drops. The reward is
$$r=\alpha \sigma \left(\beta \log \frac{\mathrm{PPL}(H^{\prime })}{\mathrm{PPL}(H^{\prime },Q)}\right)$$
where $H^{\prime }$ is the cleaned query history, $Q$ is the original question, $\sigma$ is the sigmoid, and $\alpha =2$ and $\beta =10$ set the range of the score. A perplexity ratio above 1 means the text reads more naturally once the question is attached, so the reward is above 1.

Strictly speaking, the second pass measures the average perplexity of the whole text, so the length and fluency of the history itself also affect it. The ratio is therefore not an exact "how much more likely the question is with the history than without it", Despite the name "information gain", it is an approximate score inspired by mutual information, not an estimate of it. The paper treats it only as an approximate score for ranking sibling branches under the same parent. Sibling branches face the same question, so how guessable the question is on its own is the same for all of them, and what is compared is mainly what each one found.

A small example (illustrative): two sibling branches both get perplexity 20 in the first pass. Branch A found the TV appearances; with the question attached, perplexity falls to 16, the ratio is 1.25, and the reward is about 1.81. Branch B found the films; with the question attached, perplexity rises to 21, the ratio is about 0.95, and the reward is about 0.76. UCT will return to branch A more often.

At inference, scoring uses only the question and the query history, and needs no gold answer. Gold answers are used only offline: to compute F1 at evaluation time and for the after-the-fact ranking analysis below, never during search.

Why clean the history first?

If the raw history were scored directly, the agent's thoughts would often restate the question, and the entity-search parameters would copy names straight from the question. Attaching the question would then tend to lower perplexity no matter what was found, and the score would reflect literal repetition.

So before scoring, the paper cleans the history and keeps only these parts:

• All thought text is removed.
• For the entity-search step, only the tool name is kept, and the entity name copied from the question is removed.
• Structured parameters of later steps are kept, such as the full SPARQL query. The paper's reason is that they build on what the knowledge base returned.
• Only the current step's returned result is kept, because the score judges the new evidence this step brings.

The "History for Fast MCTS" box on the right of Figure 1 shows a cleaned history: a SearchNodes call with only its tool name, one SearchGraphPatterns query kept in full, and the list of relations it returned.

Why also decay by depth?

The deeper a branch, the more likely it is long and unproductive. During backpropagation, the paper multiplies each node's reward by a factor that shrinks with depth: the expected depth is set to 7, and beyond 7 the factor drops by 0.1 per extra level, down to 0. For example, a node at depth 9 adds back only 0.8 times its reward. Deep branches that produce nothing are then selected less and less, and rewards never turn negative.

Experiments: how much comes from search, and how much from IG

The agent is Llama-3.1-8B-Instruct after SFT on a small number of annotated trajectories: 300 shared by WebQSP and CWQ, 150 for GrailQA (generated by a closed LLM and checked one by one by hand), and 450 for the Wikidata benchmark KQA Pro, which has complex questions. The CWQ, GrailQA and KQA Pro test sets are random samples of 500, 500 and 900 questions, and all methods use the same questions.

The table shows three things. First, tree search itself helps: with the same 8B agent, classic MCTS beats the 5-run vote on all four benchmarks. Second, the IG score is the key: keeping the Fast MCTS search procedure but replacing the reward with a constant drops CWQ from 70.15 to 57.99, about the same as a simple vote. Keeping the search procedure but removing informative scoring is not enough to hold up the results. Third, on KQA Pro Fast MCTS and classic MCTS are about even.

The paper also tried a level-by-level beam search with the same IG score, keeping only the top few branches at each level. It is below Fast MCTS on all four benchmarks, with 60.82 on CWQ. A branch cut by beam search at one level never comes back, while UCT keeps a chance to revisit branches whose score is low for now.

The closest prior work is KBQA-o1, which needs to train a reward model from final outcomes. On the full WebQSP test set with the same 100 annotations, Fast MCTS scores 66.11 versus 59.8 for KBQA-o1 with Llama-3.1-8B (KBQA-o1 numbers are taken from its own paper), and is also higher with Qwen2.5-7B and Gemma-2-9B.

How much cost does it save?

Fast MCTS makes half as many generation calls as classic MCTS and takes about 16% less time. Time does not fall in proportion to the calls; the time here includes the scoring forward passes and tool execution. It is about 45% slower than the 5-run vote, and the paper's reason is that every expansion needs forward passes for scoring and the search explores more states. That time buys 14 points of F1.

Does the IG score really pick the right branch?

To check, the paper first runs classic MCTS, executes the SPARQL at every terminal node, and computes F1 against the gold answer. It then backpropagates these terminal F1 values up the tree to get a Monte Carlo value estimate for each node, called the "true value" below. This is only an after-the-fact analysis and is not part of inference. It then checks whether, among sibling branches under the same parent, the IG ranking keeps the branch with the highest true value.

On CWQ, the branch IG ranks first is the best branch 75.73% of the time, versus 62.14% when a model is prompted to score with chain-of-thought. Keeping only the IG top 3 retains the best branch 100% of the time on both datasets. But the full order IG produces is mixed. Kendall τ, which measures agreement of the whole ranking, is 0.37 for IG versus 0.12 for chain-of-thought scoring on CWQ, but on KQA Pro IG gets only 0.11, below the 0.19 of chain-of-thought scoring. So IG is good at "not cutting the good branch", not at estimating exactly how much a branch is worth.

This also explains why Fast MCTS does not beat classic MCTS on KQA Pro. On CWQ the main difficulty is picking the right relation among many Freebase relations with similar names, and a wrong branch shows itself as soon as its results come back, so its IG score falls quickly. On KQA Pro a child may be a qualifier, a comparison, a count or a yes/no check, so children are more varied in kind and chains run deeper. What one step finds says less about whether a correct query can be built in the end.

Limitations as I see them

First, it is slower than the simplest approach. Time per question is 1.45 times that of the 5-run vote on CWQ and 1.24 times on KQA Pro. For latency-sensitive online question answering, this cost needs to be counted.

Second, scoring quality depends on local progress predicting the final result. On KQA Pro, with many kinds of operations and long chains, Fast MCTS gets 86.44 versus 86.67 for classic MCTS, with no advantage.

Third, search often finds a good branch but does not pick it. On CWQ the best F1 among the top 5 candidates is 76.29, while the final voted result is 70.15, 6 points lower. The voting step leaves room to improve. Also, apart from WebQSP all test sets are samples, each knowledge base uses its own trained agent, and the paper does not claim transfer across knowledge bases.

How a search or QA team could try it

1. Pick a multi-step tool-calling scenario, such as agentic retrieval or structured query generation, where the current agent runs one line to the end. Use a vote over several samples as the baseline to see whether more tries help at all; if not, tree search is unlikely to help either.
2. Implement the IG score: take an existing open instruction model and compute perplexity twice on the cleaned history. Adapt the cleaning rules to your own tools, with the principle of removing any text that may restate the question.
3. Validate the ranking offline first: on questions with gold answers, check whether the IG top 3 keeps the best branch. If it does not, stop here.
4. Then add UCT search, control the number of candidates per expansion and the total budget, and record both accuracy and time per question to compare cost-effectiveness against the vote baseline.

Verdict

The most useful takeaway: a reward for intermediate steps does not have to be trained, nor computed by simulating to the end. Asking an existing model "after reading this, is the original question easy to guess?" is enough to help tree search keep the promising branches.