Search the Reasoning Space
Derive MCTS over reasoning states from first principles, build it in DSPy, audit whether the tree actually branches, then run the budgeted seven-strategy comparison — with an oracle-leakage control — on two models and separate what it establishes from what it cannot.
One Path Is a Decision
In Chapter 15, we built an agent. It acquired repository evidence safely. It inspected the files it was permitted to inspect. It produced a diagnosis.
It was wrong.
Not completely wrong—the agent scored 0.60. It found the relevant code. It identified that service.py references item.cost. It located the model definition. And then it concluded that Item has no cost attribute, missing the final connection: service.py should use item.price.
The tools worked. The evidence was present. The reasoning stopped one step short.
Why?
The agent committed to a single reasoning path early. Once it began down the path of “what attributes does Item have?” it followed that path to its conclusion without reconsidering alternatives. A human debugging this same code might pause: Wait—maybe the service is using the wrong attribute name. Let me check what the model actually defines.
That pause is a search over reasoning states. And it is exactly what this chapter is about.
Your sequence up to this point has been:
15 Agents Are Programs Too
↓
give the program tools
16 Search, Memory and Long Context
↓
decide what evidence to inspect
17 Search the Reasoning Space ← this chapter
↓
decide which reasoning path to pursue
18 Don't Let the Optimizer Cheat
↓
constrain what search is allowed to know
Chapter 15 established action. Chapter 16 established evidence selection. Now we tackle the next problem: what if the failure is no longer access to evidence, but committing too early to one reasoning path?
A note on what this chapter is. Like the earlier experimental chapters, it ends in a program that was run and measured. Unlike them, the headline result is a null one, and the most useful sentence in the chapter is about a measurement the experiment could not make:
Evidence status for Chapter 17
Element Status Reasoning-search agent built in DSPy done — generation, an LM value function, and reflection (everything LATS has except acting against an environment) Whether the tree actually searches audited — an early version does not branch; four defects are found and fixed; 27 unit tests Budgeted seven-strategy comparison (plus oracle control) run — seeds 17–19, three repetitions, on a local 8B model ( qwen3) and a stronger hosted model (muse-spark-1.3-contributor); see 17.8Whether tree search allocates a budget better than Best-of-N not measured — the fair tree arms never spent their budget (17.8.2); this is the experiment’s honest gap Where you might expect a leaderboard you will find a null result: on this one fixture, none of the strategies produced a substantive diagnosis improvement over a single direct call. The only numerical deviation was a qwen3 Random Tree repetition that earned an extra quarter-point by naming
service.pyexplicitly, a phrasing effect in the rule-based evaluator rather than evidence of a search-policy advantage. Nothing here claims that tree search beats Best-of-N, or that reflection earns its cost — the run could not test either. The one firm claim is structural: a search is only as trustworthy as the evidence its value function is allowed to see.
17.1 From Chain to Tree
Let’s trace the evolution of reasoning strategies we’ve built so far.
Predict: One Continuation
The simplest possible use of a language model produces one answer and stops:
answer = dspy.Predict(GenerateAnswer)(context=evidence, trace=[])
There is no search here. One sample is taken as final. Chapter 15’s agent was a more elaborate version of the same shape — a longer single path, but still one path, followed to its conclusion without reconsidering the alternative framing.
Best-of-N: Several Independent Continuations
A natural next baseline is independent sampling:
results = []
for _ in range(num_candidates):
result = cot_program(context=evidence)
results.append(result)
best = max(results, key=score_diagnosis)
Now we generate multiple complete paths independently. Each path gets the same starting point. Each path runs to completion without looking at the others. Then we pick the best.
This is breadth—but only at the leaves. The paths never interact. A promising early step cannot receive more attention than a dead end.
Tree Search: Continuations of Continuations
What if we could:
- Generate several candidate next steps
- Evaluate which partial path looks promising
- Spend more compute extending promising paths
- Backtrack from weak paths
That is a tree search over reasoning states.
graph TD
A[Start: repository evidence] --> B[Step 1: item.cost missing]
A --> C[Step 1: service uses wrong attribute]
A --> D[Step 1: model definition changed]
B --> E[Step 2: check Item fields]
B --> F[Step 2: check import]
C --> G[Step 2: compare to model.py]
C --> H[Step 2: trace service.py usage]
D --> I[Step 2: git history]
D --> J[Step 2: schema migration]
G --> K[Step 3: price vs cost → fix identified]
E --> L[Step 3: cost genuinely missing]
style K fill:#2ecc71,color:white
style L fill:#e74c3c,color:white
The key difference: partial paths get evaluated before they reach completion, and computational resources are reallocated toward the most promising areas of the tree.
17.2 The Reasoning State
Before we build a search, we need to define what we’re searching over.
State
A state is the current point in the reasoning process. In our implementation:
class TraceStep(dspy.Signature):
"""Produce the next reasoning step given what is known."""
context = dspy.InputField(desc="The repository evidence")
trace = dspy.InputField(desc="Reasoning steps so far")
reasoning = dspy.OutputField(desc="Next reasoning step")
The state contains:
context: The frozen evidence from Chapter 15trace: The sequence of reasoning steps taken so far
For a tree node, the state is represented as:
class MCTSReasoningNode:
def __init__(
self,
state: dict,
trace: list[str] | None = None,
parent: "MCTSReasoningNode" | None = None,
depth: int = 0,
node_id: str | None = None,
):
self.state = state # the accumulated state
self.trace = trace or [] # path of reasoning steps to reach this state
self.parent = parent # parent node in tree
self.children: list[MCTSReasoningNode] = []
self.visits = 0 # times this node was selected
self.reward = 0.0 # accumulated return through this node
self.score = 0.0 # evaluation of just this node's state
self.depth = depth
self.id = node_id or str(uuid.uuid4())
This is deliberately simple. The trace is just a list of strings. The state is a dictionary. No fancy data structures needed.
Trace
The trace is the path taken to reach the current state. It’s the reasoning history:
trace = [
"service.py calls item.cost on line 42",
"model.py defines Item with a price field on line 18",
"Checking if cost is aliased... no alias found",
]
Each element of the trace is one reasoning step generated by the language model. The trace defines where we are in the search space.
Child
A child node is created by:
- Taking the current state and trace
- Asking the LM to produce the next reasoning step
- Creating a new node with the extended trace
def _predict_next(self, node: MCTSReasoningNode) -> dspy.Prediction:
"""Generate the next reasoning step given node's state."""
return self.reasoning(context=node.state["context"], trace=node.trace)
This is the generation half of the search. The LM proposes. The tree structure decides which proposals get extended.
17.3 Selection, Expansion, Evaluation, Backup
MCTS is built from four operations. Let’s derive each one from first principles.
Selection: Where Should We Spend the Next Compute?
At any point, the tree has many frontier nodes—nodes that could be expanded. We have a finite compute budget (LM calls are expensive). We need a policy for choosing which node to expand next.
The simplest policy: always expand the best-scoring current node. This is greedy and purely exploitative. If the current best path is actually a dead end, we waste all our budget on it.
The opposite policy: expand nodes uniformly at random. This explores broadly but wastes budget on paths that are clearly not working.
What we want: balance exploration and exploitation. Spend more compute on promising paths, but keep enough exploration to discover better paths.
This is the exploration-versus-exploitation problem handled by the UCT selection rule. For a child node, the usual form uses the parent’s visit count, not a global tree counter:
def _select(self, node: MCTSReasoningNode) -> MCTSReasoningNode:
"""Select a leaf node for expansion using UCT."""
while node.children and not self._is_terminal(node):
if len(node.children) < self.branching_factor:
# This node hasn't been fully expanded yet
return node
# Move to best child according to UCT
node = max(node.children, key=lambda c: self.uct_value(c))
return node
Two cases:
- The node has fewer children than
branching_factor→ expand it (add another child) - The node is fully expanded → move down to the best child by UCT
This gives us a fixed branching cap: a node may acquire children until it reaches branching_factor, after which selection descends through its existing children. That is not progressive widening in the usual sense. Progressive widening makes the number of permitted children itself grow as a function of visits; this implementation simply caps breadth at a configured value.
Expansion: Generating New States
Once we’ve selected a leaf node, we generate a child:
def _expand(self, node: MCTSReasoningNode) -> MCTSReasoningNode:
"""Create a new child node by generating the next reasoning step."""
prediction = self._predict_next(node)
new_trace = node.trace + [prediction.reasoning]
new_state = {**node.state, "last_step": prediction.reasoning}
child = MCTSReasoningNode(
state=new_state,
trace=new_trace,
parent=node,
depth=node.depth + 1,
)
node.children.append(child)
return child
The child’s trace is the parent’s trace plus one new reasoning step. The child’s state is the parent’s state updated with that step. The child’s depth is one more than the parent’s.
This is where the reasoning generator is called. In a minimal implementation, one uncached expansion consumes one generation call. That is not yet the total inference cost of a simulation: evaluation may also invoke a model-backed scorer, and final answer synthesis may require another model call. The experiment must count all inference work, not only calls tracked by the reasoning generator.
Evaluation: How Good Is This Partial Path?
After expansion, we need to assign a score to the new node. In full MCTS, this typically involves running a simulation (rollout) from the new node to a terminal state. In our simpler version, we evaluate the partial path directly:
def _evaluate(self, node: MCTSReasoningNode) -> float:
"""Evaluate a reasoning state and return a numeric score."""
prediction = self.score_fn(
context=node.state["context"],
trace=node.trace,
)
return float(prediction.score)
For a clean teaching implementation the search-time evaluator is itself a DSPy signature — ScoreReasoning, shown in full in Section 17.5. The point to carry forward is that if that scorer is LM-backed, every node evaluation is a model call, and the budget accounting has to include it (Section 17.4).
Backup: Propagating Evidence
After evaluating a node, we propagate the score backward through the tree:
def _backpropagate(self, node: MCTSReasoningNode, reward: float) -> None:
"""Propagate reward from this node up to the root."""
current: MCTSReasoningNode | None = node
while current is not None:
current.visits += 1
current.reward += reward
current = current.parent
Every ancestor of the evaluated node gets:
- One more visit (they were “part” of this simulation)
- The reward added to their cumulative reward
This means the average reward (reward / visits) of a node reflects the average quality of all paths that pass through it. If a node leads to good outcomes (children and grandchildren that score well), its average reward goes up. If it leads to dead ends, its average reward goes down.
This is how information from deep in the tree propagates back to guide selection at shallow levels.
17.4 The Search Loop
Putting it all together:
def run(self, context: dict, **kwargs) -> MCTSAgentResult:
"""Run MCTS over reasoning paths."""
self._reset_state()
root = self._create_root(context)
root.id = "root"
self._all_nodes.append(root)
for simulation in range(self.num_simulations):
# BUDGET CHECK: stop if we've exceeded LM call budget
if self.lm_call_count >= self.max_lm_calls:
self._log_info(f"LM call budget reached: {self.lm_call_count}")
break
# 1. SELECT
leaf = self._select(root)
# 2. CHECK TERMINAL
if self._is_terminal(leaf):
self._backpropagate(leaf, leaf.score)
continue
# 3. EXPAND
child = self._expand(leaf)
self._all_nodes.append(child)
# 4. EVALUATE
score = self._evaluate(child)
child.score = score
# 5. BACKUP
self._backpropagate(child, score)
# 6. EARLY STOP CHECK
if self.early_stop_threshold and score >= self.early_stop_threshold:
self._log_info(f"Early stop at simulation {simulation + 1}: score {score:.3f}")
break
# After search: find best paths
return self._finalize_result(root)
Each iteration through the loop is one search simulation. In this base version it contains one expansion call plus one evaluation — and Section 17.5 adds a third LM call here, a reflection when the evaluation is low. If ScoreReasoning is LM-backed, that is already two to three model calls per simulation before final answer synthesis. In the source implementation audited for this chapter, calls_used counts generator calls and does not account for model work inside the scoring or reflection paths. For the experiments later in this chapter, the budget therefore has to be defined in terms of total model invocations, tokens, or wall-clock cost rather than trusting one internal counter.
Let’s visualize the full flow:
sequenceDiagram
participant S as Search Controller
participant LM as Language Model
participant SC as Scorer
participant T as Tree
S->>T: Select leaf node
T-->>S: Return selected node
S->>LM: Generate next step (context, trace)
LM-->>S: New reasoning step
S->>T: Add child node
S->>SC: Score new node (context, extended trace)
SC-->>S: Score 0.0 - 1.0
S->>T: Backpropagate score
T-->>S: Update ancestors
Note over S: Continue to next simulation<br/>or stop if budget exceeded
This is the whole algorithm. Four simple operations, composed in a loop. The complexity is in the details—how to select, how to evaluate, how to balance exploration—but the core structure is simple.
17.5 Build It With DSPy
Let’s build the complete agent. We’ll use DSPy for the LM-interaction signatures and ordinary Python for the tree mechanics.
The Signatures
import dspy
class TraceStep(dspy.Signature):
"""Produce the next reasoning step for diagnosing a repository issue."""
context = dspy.InputField(desc="The repository evidence")
trace = dspy.InputField(desc="Reasoning steps taken so far")
reflections = dspy.InputField(desc="Lessons from earlier failed paths; may be empty")
reasoning = dspy.OutputField(desc="Next reasoning step to investigate")
class ScoreReasoning(dspy.Signature):
"""Evaluate how promising a reasoning path is."""
context = dspy.InputField(desc="The repository evidence")
trace = dspy.InputField(desc="Reasoning steps taken so far")
score = dspy.OutputField(desc="Float score 0-1 for this reasoning path")
class Reflect(dspy.Signature):
"""Given a reasoning path that scored poorly, state one concrete lesson for
the next attempt. Do not restate the answer; name the wrong assumption."""
context = dspy.InputField(desc="The repository evidence")
failed_trace = dspy.InputField(desc="A low-scoring reasoning path")
failed_score = dspy.InputField(desc="Its score, 0-1")
lesson = dspy.OutputField(desc="One sentence: what this path got wrong")
class GenerateAnswer(dspy.Signature):
"""Generate the final diagnosis based on the best reasoning path."""
context = dspy.InputField(desc="The repository evidence")
trace = dspy.InputField(desc="Best reasoning path found")
diagnosis = dspy.OutputField(desc="Final diagnosis")
Four LM roles: generate the next step (TraceStep), value a partial path (ScoreReasoning), reflect on a dead end (Reflect), and synthesise the final answer (GenerateAnswer). Generation, value, and reflection all run during the search; synthesis runs once at the end. The one LATS role missing from this list is acting — none of these signatures touches an environment — which Section 17.6 comes back to.
The Agent Class
class MCTSDiagnosisAgent:
def __init__(
self,
lm: dspy.LM | None = None,
branching_factor: int = 3,
max_depth: int = 4,
num_simulations: int = 12,
ucb_weight: float = 1.41,
early_stop_threshold: float = 0.85,
max_lm_calls: int = 20,
top_k_leaves: int = 3,
reflect: bool = True,
reflect_threshold: float = 0.4,
):
# Set up DSPy with the documented public configuration API
dspy.configure(lm=lm)
self.reasoning = dspy.ChainOfThought(TraceStep)
self.score_fn = dspy.Predict(ScoreReasoning)
self.reflect_fn = dspy.Predict(Reflect)
self.answer_fn = dspy.ChainOfThought(GenerateAnswer)
# Search parameters
self.branching_factor = branching_factor
self.max_depth = max_depth
self.num_simulations = num_simulations
self.ucb_weight = ucb_weight
self.early_stop_threshold = early_stop_threshold
self.max_lm_calls = max_lm_calls
self.top_k_leaves = top_k_leaves
self.reflect = reflect
self.reflect_threshold = reflect_threshold
# State tracking
self.lm_call_count = 0
self._all_nodes: list[MCTSReasoningNode] = []
self._node_lookup: dict[str, MCTSReasoningNode] = {}
self._reflections: list[str] = []
self._total_visits = 0
The constructor sets up:
- The DSPy predictors for generation, scoring, reflection, and answer synthesis
- The search hyperparameters, including whether reflection is on and the score below which a path is reflected on
- The state tracking (LM call count, nodes, visits, accumulated reflections)
The agent uses DSPy’s ChainOfThought for generation and answer synthesis, and Predict for scoring and reflection. Where we want multi-step reasoning, use CoT; where we want a single judgement, use Predict.
The Search Methods
def _reset_state(self) -> None:
"""Reset MCTS internal state."""
self.lm_call_count = 0
self._all_nodes = []
self._node_lookup = {}
self._reflections = []
self._total_visits = 0
self._lm_cache: dict[tuple, dspy.Prediction] = {}
def _predict_next(self, node: MCTSReasoningNode) -> dspy.Prediction:
"""Generate the next reasoning step given node's state."""
cache_key = (json.dumps(node.state, sort_keys=True), tuple(node.trace))
if cache_key in self._lm_cache:
return self._lm_cache[cache_key]
self.lm_call_count += 1
result = self.reasoning(
context=node.state["context"],
trace=node.trace,
reflections=self._reflections,
)
self._lm_cache[cache_key] = result
return result
Note the cache on (state, trace). It will make sibling expansions return identical children, and — once reflection is added below — it will also serve a continuation generated before the search had learned anything. Section 17.7 removes it for exactly these reasons.
The full run loop:
def run(self, context: dict, **kwargs) -> MCTSAgentResult:
"""Run MCTS search over reasoning paths."""
self._reset_state()
root = self._create_root(context)
root.id = "root"
self._all_nodes.append(root)
self._node_lookup[root.id] = root
for simulation in range(self.num_simulations):
# Budget check
if self.lm_call_count >= self.max_lm_calls:
self._log_info(f"LM call budget reached: {self.lm_call_count}")
break
# 1. Selection
leaf = self._select(root)
# 2. Terminal check
if self._is_terminal(leaf):
self._backpropagate(leaf, leaf.score)
continue
# 3. Expansion
child = self._expand(leaf)
self._all_nodes.append(child)
self._node_lookup[child.id] = child
# 4. Evaluation
score = self._evaluate(child)
child.score = score
# 5. Backup
self._backpropagate(child, score)
# 6. Early stop
if self.early_stop_threshold and score >= self.early_stop_threshold:
self._log_info(f"Early stop at simulation {simulation + 1}")
break
# Generate final answer from best path
return self._finalize_result(root)
Selecting the Best Path
After the search, we need to extract the best reasoning path and generate a final answer:
def _finalize_result(self, root: MCTSReasoningNode) -> MCTSAgentResult:
"""Extract best paths and generate final answer."""
# Find top paths by score
scored_paths = []
for node in self._all_nodes:
if node.trace and node.score > 0:
scored_paths.append({
"node_id": node.id,
"trace": node.trace,
"score": node.score,
"visits": node.visits,
"depth": node.depth,
"children_count": len(node.children),
})
scored_paths.sort(key=lambda x: x["score"], reverse=True)
top_paths = scored_paths[: self.top_k_leaves]
# Generate the final answer from the best path
best_path = top_paths[0] if top_paths else None
if best_path:
answer = self.answer_fn(
context=root.state["context"],
trace=best_path["trace"],
)
else:
answer = None
return MCTSAgentResult(
best_diagnosis=answer.diagnosis if answer else "",
best_trace=best_path["trace"] if best_path else [],
top_k_paths=top_paths,
total_nodes=len(self._all_nodes),
total_simulations=self.num_simulations,
generator_calls=self.lm_call_count,
# ... total scorer calls, tokens, latency, tree statistics
)
Note the last comment. generator_calls counts calls made by self.reasoning. It does not count the scorer calls in _evaluate, the reflection calls added below, or the final answer_fn call. Any budget expressed in terms of generator_calls alone undercounts the real inference cost — a point Section 17.4 already flagged and Section 17.8 has to design around.
Reflection: Learning From a Dead End
Best-of-N and the plain tree so far share a weakness: a path that fails teaches the search nothing. The next expansion from a sibling starts with exactly the context the failed path started with, and can walk straight into the same mistake. LATS closes that loop with reflection — when a trajectory scores badly, the model writes down why, and that note joins the context every later expansion sees.
The teaching agent adds it with the Reflect signature from above and one method:
def _reflect(self, node: MCTSReasoningNode) -> None:
"""Record one lesson from a poorly scored path."""
self.lm_call_count += 1
lesson = self.reflect_fn(
context=node.state["context"],
failed_trace=node.trace,
failed_score=node.score,
).lesson
self._reflections.append(lesson.strip())
_predict_next already passes self._reflections to the generator (shown above). The only loop change is to call _reflect after a low score, before backup:
score = self._evaluate(child)
child.score = score
if self.reflect and score < self.reflect_threshold:
self._reflect(child)
self._backpropagate(child, score)
Two things to keep honest. Reflection is more model calls — one per bad path — so it belongs in the budget accounting exactly like scoring does; a reflecting MCTS run and a non-reflecting one are different-cost conditions, and Section 17.8 treats them as separate arms. And a reflection is itself an LM output that can be wrong. The search adds it to the context the generator reads; it never adds it to the value function or treats it as a constraint. It changes what gets tried, not what counts as good — the same boundary Chapter 14 drew around GEPA’s feedback channel.
17.6 What This Agent Is — and Isn’t
Let’s be precise about what we’ve built.
What This Agent Is
This is a tree search over reasoning states. It uses:
- A language model to generate candidate next reasoning steps (
TraceStep) - Tree-structured states to hold multiple reasoning paths at once
- UCT to balance exploration and exploitation
- An LM-backed value function consumed during the search (
ScoreReasoning→_evaluate→ backup → UCT) - Reflection: a failed path produces a lesson that later expansions read (
Reflect) - Backpropagation to move evidence from deep in the tree up to shallow selection
What This Agent Is Not
This is not full LATS (Language Agent Tree Search), and the gap is now down to one thing. LATS combines MCTS with LM reasoning, an LM value function, self-reflection, and acting against an environment. The agent here has the first three:
| LATS component | This agent |
|---|---|
| LM next-step generation | ✅ TraceStep |
| Tree states + UCT selection | ✅ |
| LM value function guiding search | ✅ ScoreReasoning, consumed by UCT |
| Backpropagation, budgets, early stop | ✅ |
| Self-reflection during search | ✅ Reflect |
| Acting against an environment | ❌ no tools, no observations mid-search |
Every node here is a thought. In LATS a node can be an action — call a tool, read a file, run a test — and its child is conditioned on the real observation that came back. That is Chapter 15’s tool agent running inside this tree, and it is a genuinely bigger build: the search branches on what the environment returns, not only on what the model imagines. This chapter stops one step short of that on purpose. Adding tools to reasoning search is the natural next project, not a paragraph.
The Critical Distinction: Search-Time Value vs. Post-Hoc Evaluation
The point of this subsection is about where a value function sits in the dataflow, not what it is called.
The teaching agent has a search-time value function. ScoreReasoning runs inside _evaluate(), its result is backed up through the tree, and UCT consumes those backed-up rewards to decide where to search next:
ScoreReasoning → _evaluate() → node.score → _backpropagate() → UCT selection
The source implementation that motivated this chapter also carries a separately named ValueEstimator. It looks like the value function from the LATS paper, but it runs once, after the search, on the selected trace:
ValueEstimator → (runs only after search) → scores the chosen answer
Nothing in the search reads its output. Following Chapter 5’s rule — trace what actually consumes a value — ValueEstimator is post-hoc evaluation with a misleading name. A component is part of the search policy only if the selection rule depends on it.
17.7 Fix the Search Before Measuring It
An implementation this size is where the interesting bugs live, and several of them would quietly invalidate any comparison built on top of the agent. Each is the kind of defect that produces plausible output and a plausible trace while doing something other than what the code appears to do.
Issue 1: Cache Collapses Branching
The _predict_next method caches on (state, trace):
cache_key = (json.dumps(node.state, sort_keys=True), tuple(node.trace))
if cache_key in self._lm_cache:
return self._lm_cache[cache_key]
When the same parent is expanded again, the (state, trace) key is identical, so every sibling call returns the first cached continuation. Instead of three distinct children:
parent
├── child A (thought 1)
├── child B (thought 2)
└── child C (thought 3)
We get:
parent
├── child A (thought 1)
├── child A (thought 1) ← duplicate
└── child A (thought 1) ← duplicate
The tree diagram in Section 17.1 still renders; the search underneath it has one child repeated. Drawing a tree does not mean you are exploring one — which is the whole reason 17.7 exists.
The fix: do not memoize sibling sampling by (state, trace). A cache is right when the request is “give me the same continuation again”; expansion is asking for another candidate continuation from the same parent.
Removing this cache makes diversity possible; it does not guarantee it. With deterministic or low-variance generation, two fresh calls can still return the same continuation. A real search experiment therefore needs an explicit sibling-diversity policy — for example controlled sampling or multiple generations — and it must report unique_traces and duplicate rate rather than assuming that three calls created three branches.
Issue 2: Reward Double-Counting
An earlier version of _evaluate() added the score to the node itself:
def _evaluate(self, node: MCTSReasoningNode) -> float:
score = float(self.score_fn(context=node.state["context"], trace=node.trace).score)
node.reward += score # BUG: _backpropagate adds it again
return score
run() then calls _backpropagate(child, score), which walks from the child to the root adding reward at every node. The evaluated node therefore receives its score twice — once from _evaluate, once from _backpropagate — while its ancestors receive it once. Its average reward is inflated relative to every other node, and UCT over-selects the subtree beneath it.
The fix: _evaluate returns the score and mutates nothing. _backpropagate is the only place reward and visits change:
def _evaluate(self, node: MCTSReasoningNode) -> float:
return float(self.score_fn(context=node.state["context"], trace=node.trace).score)
def _backpropagate(self, node: MCTSReasoningNode, reward: float) -> None:
current = node
while current is not None:
current.visits += 1
current.reward += reward
current = current.parent
Now the evaluated node and each of its ancestors get the score exactly once.
Issue 3: UCB Weight of 0 Doesn’t Mean Zero
Current code:
ucb_weight or self.ucb_weight
This is buggy because 0 or self.ucb_weight returns self.ucb_weight. You cannot set ucb_weight=0 to mean “pure exploitation.”
The fix:
weight = self.ucb_weight if ucb_weight is None else ucb_weight
Now:
ucb_weight=None→ use the defaultucb_weight=0→ pure exploitation (no exploration bonus)ucb_weight=1.41→ standard UCB balanceucb_weight=3.0→ aggressive exploration
Issue 4: Per-Run State Is Not Fully Reset
The source implementation audited for this chapter resets only part of its run-scoped state before a new search. The teaching class above uses different field names, but the invariant is the same: counters, nodes, lookup tables, caches, and best-so-far state must belong to one run, not to the lifetime of the Python object.
If the same agent instance is reused without clearing all of that mutable state, a later condition can inherit call-budget consumption, tree nodes, cached outputs, or selection state from an earlier condition. That contaminates both behavior and measurement.
The fix: every field that a run mutates is reset at the start of run(), not only in __init__:
def _reset_state(self) -> None:
self.lm_call_count = 0
self._all_nodes = []
self._node_lookup = {}
self._total_visits = 0
self._lm_cache = {}
self._best_so_far = None
self._t0 = time.perf_counter()
An agent instance is then safe to reuse across conditions in the comparison below — which matters, because the comparison drives it through several of them in one process.
17.8 Run the Budgeted Search Experiment
The design comes first — conditions, budget accounting, metrics, the leakage control — and then the two runs it produced. Both runs use the same scaffold; the second changes only the model.
There are really two questions:
Question A. Does additional inference-time compute help this diagnosis task?
and, conditional on spending that extra compute:
Question B. Does tree search allocate the budget better than simpler multi-candidate strategies?
Do not collapse those questions. A one-call direct baseline is useful, but it is not an equal-compute competitor to a twelve-expansion tree. As it turns out (17.8.4), the run answers A and cannot answer B.
The Conditions
| Condition | Strategy | Search/generation budget |
|---|---|---|
| Direct | One answer, no explicit search | 1 generation |
| ChainOfThought | One reasoned answer | 1 generation |
| Best-of-N | Independent complete candidates | 12 generations |
| Greedy Tree | Tree expansion, no exploration bonus | 12 expansions |
| MCTS | Tree expansion + UCT | 12 expansions |
| MCTS − reflection | Same, reflect=False | 12 expansions |
| Random Tree | Tree expansion with a genuinely random child-selection policy | 12 expansions |
The fair strategy comparison is primarily Best-of-N vs. Greedy Tree vs. MCTS vs. Random Tree, because those conditions receive the same expansion budget. Direct and ChainOfThought tell us what the additional test-time compute bought relative to a one-call baseline. The MCTS − reflection arm isolates one thing: whether the reflection calls earn their extra cost, or whether tree search alone is doing the work.
Expansion count alone is still not enough. The MCTS conditions also spend scorer calls and reflection calls; every condition ends with a synthesis call. Record generation, scoring, reflection, and synthesis calls separately, plus total tokens and latency, and compare quality against total inference cost — not against expansion count.
The Metrics
For each condition, record:
@dataclass
class ExperimentResult:
condition: str
final_diagnosis_score: float
semantic_failures: list[str]
generation_calls: int
scorer_calls: int
reflection_calls: int
synthesis_calls: int
tokens_used: int
latency_seconds: float
nodes_created: int
unique_nodes: int
unique_traces: int
max_depth_reached: int
duplicate_rate: float
best_path_score: float
chosen_path_score: float
trajectories: list[dict]
The four call counts are separate on purpose. “MCTS used 12 expansions” and “MCTS made 31 model calls” are both true, and only the second is comparable to Best-of-N’s 12.
final_diagnosis_score is an evaluation-only quantity. It may compare the final diagnosis with the known answer after a condition has finished, but it must not be used to choose among Best-of-N candidates or to guide a fair tree search. Those decisions must use the same admissible value policy, with no reference answer, and all of those scorer calls belong in the cost accounting. The oracle arm in Section 17.9 is the deliberate exception, which is why its result is invalid by construction.
The unique_traces metric is equally important. It tells us whether the tree search actually explored diverse paths or just duplicated the same path multiple times.
The Data
Reuse Chapter 15’s broken repository as the test case. The evidence packet is frozen — the relevant code excerpts, and nothing else — and the correct diagnosis is known: service.py uses item.cost where Item defines price. That is what lets a score mean something. It is also n=1: one defect, one repository, exactly the scope limit Chapter 15 already flagged for the agent it measured.
The Scaffold
def run_experiment(condition: str, evidence: dict, lm: dspy.LM) -> ExperimentResult:
if condition == "direct":
# Direct answer
result = dspy.Predict(GenerateAnswer)(
context=evidence, trace=[]
)
return evaluate_direct(result)
elif condition == "cot":
# Single ChainOfThought path
result = dspy.ChainOfThought(GenerateAnswer)(
context=evidence, trace=[]
)
return evaluate_cot(result)
elif condition == "best_of_n":
# N independent CoT paths
paths = []
for _ in range(12):
pred = dspy.ChainOfThought(GenerateAnswer)(
context=evidence, trace=[]
)
paths.append(pred)
# Selection must use the same admissible value policy as the
# tree conditions. These scorer calls count toward total cost.
scored_paths = [
(admissible_score_path(evidence, pred), pred)
for pred in paths
]
best = max(scored_paths, key=lambda item: item[0])[1]
return evaluate_path(best)
elif condition == "greedy_tree":
agent = MCTSDiagnosisAgent(
lm=lm,
ucb_weight=0.0, # pure exploitation
branching_factor=3,
num_simulations=12,
)
return evaluate_agent(agent)
elif condition == "mcts":
agent = MCTSDiagnosisAgent(
lm=lm,
ucb_weight=1.41,
branching_factor=3,
num_simulations=12,
)
return evaluate_agent(agent)
elif condition == "random_tree":
agent = RandomTreeDiagnosisAgent(
lm=lm,
branching_factor=3,
num_simulations=12,
)
return evaluate_agent(agent)
evaluate_agent runs the agent, scores its diagnosis against the known answer, and — critically — records scorer calls, synthesis calls, tokens, and latency alongside the score, so that “12 expansions” cannot hide a tree condition that actually spent three times the model calls of Best-of-N.
17.8.1 The Qwen3 baseline run
The first run used the same local model as the rest of the book: ollama_chat/qwen3:latest, an 8B model, think=False, generation temperature 0.7, scoring temperature 0.0. Three repetitions, seeds 17–19. Artifacts: .artifacts/dspy-book/ch17_reasoning_search/.
Mean over three repetitions:
| Condition | Final score | Model calls | Tokens | Latency | Tree depth |
|---|---|---|---|---|---|
| Direct | 0.75 | 1 | 378 | 4s | — |
| ChainOfThought | 0.75 | 1 | 505 | 5s | — |
| Best-of-N | 0.75 | 24 | 11,686 | 56s | — |
| Greedy Tree | 0.75 | 3 | 1,457 | 7s | 1 |
| MCTS | 0.75 | 3 | 1,419 | 7s | 1 |
| MCTS − reflection | 0.75 | 3 | 1,435 | 7s | 1 |
| Random Tree | 0.83 | 3 | 1,428 | 7s | 1 |
| Oracle MCTS (invalid — 17.9) | 0.75 | 25 | 6,562 | 28s | 2–3 |
Every fair strategy landed on 0.75. That 0.75 is not a wrong diagnosis: every fair arm correctly identified that service.py reads item.cost where Item defines price. The missing quarter-point is the rule-based evaluator’s filename literalism — it wanted the string service.py named explicitly, and the qwen3 diagnoses described the location without quoting it. Random Tree’s 0.83 is one repetition of three (seed 18) whose phrasing happened to name the file; the other two scored 0.75. That is sampling noise, not a policy effect.
Best-of-N spent 24 model calls and about 11.7k tokens to reach the same 0.75 the tree arms reached in 3 calls and about 1.4k tokens — and the same 0.75 a single Direct call reached in one call. On this fixture, with this model, extra inference-time compute bought nothing.
17.8.2 Why the fair trees stopped immediately
This is the finding that determines how much the rest of the comparison can say.
Every fair tree run — Greedy, MCTS, MCTS − reflection, Random — recorded the same shape: nodes_created = 2, max_depth_reached = 1, unique_traces = 1, one generation call, one scorer call, one synthesis call. One expansion. One score. One synthesis. The twelve-simulation budget was never entered.
The reason is the early-stop rule from the search loop: stop when the admissible search-time score clears early_stop_threshold (0.85). The scorer rated the very first expansion at 0.95–1.0 — the reasoning was already good — so the loop exited after simulation 1. Branching factor 3 was never exercised. UCT never chose between siblings. Reflection had no low score to react to and never fired.
So the fair tree arms in this run are not measurements of tree search. They are measurements of “expand once, score once, stop.” The Greedy / MCTS / Random labels are indistinguishable here because the code paths that separate them never ran.
There are two ways to make the comparison actually test budget allocation, and neither was run:
- disable early stopping and impose a fixed call budget, or
- use a fixture hard enough that the first expansion does not clear the threshold.
Both change the experiment. Neither is a knob you turn quietly; each is a declared design change that needs its own run and its own writeup.
17.8.3 Stronger-model replication with Muse Spark
The protocol was frozen — same fixture and fingerprint, same four signatures, same scorer-admissibility rules, same 0.85 threshold, same budgets, same seeds — and one thing changed: the model. muse-spark-1.3-contributor, a stronger hosted model (opencode-zen Responses API, reasoning_effort low), in both the generation and the search-time scoring roles. Three repetitions, seeds 17–19. Artifacts: .artifacts/dspy-book/ch17_reasoning_search/muse-go-full/.
Mean over three repetitions, both models side by side:
| Condition | qwen3 score | qwen3 calls | qwen3 tokens | Muse score | Muse calls | Muse tokens |
|---|---|---|---|---|---|---|
| Direct | 0.75 | 1 | 378 | 1.00 | 1 | 616 |
| ChainOfThought | 0.75 | 1 | 505 | 1.00 | 1 | 667 |
| Best-of-N | 0.75 | 24 | 11,686 | 1.00 | 24 | 16,352 |
| Greedy Tree | 0.75 | 3 | 1,457 | 1.00 | 3 | 1,884 |
| MCTS | 0.75 | 3 | 1,419 | 1.00 | 3 | 1,800 |
| MCTS − reflection | 0.75 | 3 | 1,435 | 1.00 | 3 | 1,835 |
| Random Tree | 0.83 | 3 | 1,428 | 1.00 | 3 | 1,930 |
| Oracle MCTS (invalid — 17.9) | 0.75 | 25 | 6,562 | 1.00 | 25 | 9,598 |
The stronger model raised the ceiling. Every fair strategy scored a clean 1.00 — the Muse diagnoses named service.py explicitly and cleared the evaluator rubric that had docked qwen3. And every fair strategy scored 1.00 equally: Direct in one call and about three seconds, Best-of-N in 24 calls, about 16k tokens, and 66–83 seconds, the tree arms in between. The gap between “search” and “one direct call” is exactly zero.
The early-stop behaviour from 17.8.2 replicated exactly: two nodes, depth 1, one unique trace, one generation call per fair tree run. The stronger model cleared the 0.85 threshold on the first expansion just as the weaker one did, so once again the fair trees never searched.
17.8.4 What the comparison still cannot tell us
The run answered Question A and left Question B open.
Question A — does extra inference-time compute help this diagnosis task? Measured, twice: no. For both models, every strategy scored the same as a single direct call. Best-of-N’s 24 calls and the trees’ 3 calls bought nothing over Direct’s 1.
Question B — does tree search allocate a fixed budget better than Best-of-N? Not measured. The fair tree arms never spent their budgets (17.8.2), so there is no budget-allocation behaviour to compare. Any statement of the form “MCTS ≈ Best-of-N,” “MCTS beats Best-of-N,” or “MCTS loses to Best-of-N” is unsupported by this run — the conditions that would separate those outcomes did not execute.
What the run does establish, precisely:
- The repaired machinery is capable of genuine tree exploration when the stopping rule permits it: the oracle arm reached depth 2–3 across all twelve simulations and produced nine to twelve unique traces per run. Because that arm is oracle-guided, those trajectories establish harness behavior only, not fair-search performance.
- On this one fixture (
n = 1), for two models a generation apart in capability, the correct diagnosis was reachable with no search at all. - Budget honesty holds: the cost columns are real, separately counted per call type, and Best-of-N’s 8× call count over the tree arms is not hidden behind “12 expansions.”
- The 0.75 → 1.00 shift between models is an evaluator-rubric effect (filename literalism), not evidence that one model reasons better about this fault; both identified it.
What it must not be read as: evidence about reflection (it never fired in a fair arm), evidence about duplicate rates in general (0.00 here, but one oracle repetition hit 0.25 at temperature 0.7), or evidence that tree search is unnecessary in general. The honest scope is one sentence: for this fixture, with these two models, tree search was unnecessary because the task did not require search.
17.9 Why Oracle-Guided Search Is Invalid Evidence
The comparison in 17.8 carries one more condition, and it is the conceptually important one. It is also the one that had to be labelled invalid before it ran, not after.
The setup runs MCTS twice:
# A: Fair search
fair_agent = MCTSDiagnosisAgent(
lm=lm,
ucb_weight=1.41,
branching_factor=3,
num_simulations=12,
)
# ScoreReasoning sees only admissible evidence
# B: Oracle-guided search
class OracleScoreReasoning(dspy.Signature):
"""Score reasoning path using known answer."""
context = dspy.InputField()
trace = dspy.InputField()
known_answer = dspy.InputField()
score = dspy.OutputField(desc="Score comparing trace to known answer")
oracle_agent = MCTSDiagnosisAgent(
lm=lm,
ucb_weight=1.41,
branching_factor=3,
num_simulations=12,
)
oracle_agent.score_fn = dspy.Predict(OracleScoreReasoning)
# The oracle scorer sees the reference answer
The oracle scorer compares each reasoning path against the known correct answer and scores it by how close the path is getting. In both runs it did the one thing the fair arms never got to do: it searched. With early stopping disabled by design, all three repetitions ran the full twelve simulations, built thirteen nodes, reached depth 2–3, and produced nine to twelve unique reasoning traces — on both models. Its final diagnosis score was 1.00 on Muse and 0.75 on qwen3, matching the fair arms exactly.
That result is worthless as evidence about the deployable agent, and the reasons do not depend on the number. The condition is invalid regardless of what it scores:
- The search-time value function has information that would not exist at inference time.
- The search is not discovering the answer; it is being steered toward a target it was handed.
- Any score it produces measures the strength of the leak, not the quality of the reasoning.
That is the point of including the condition at all. It is not there to produce a number. It is there so that the comparison has a known-invalid arm — a control whose result must be discarded on principle — sitting next to the fair arms, as a reminder that a value function’s access is part of the experiment’s validity.
A search procedure allowed to consume information unavailable at inference time cannot provide valid evidence about the deployable program, however good its outputs look.
That is the exact argument Chapter 18 then makes at the scale of the whole optimization loop, with 288 constructed attacks instead of one designed control. This section is its small, sharp instance: the firewall is not only for optimizers and retrievers — it is for anything with a value function, reasoning search included.
There is a second reason the oracle arm earns its place. It is the only condition in either run that actually spent a search budget, so it is the only direct evidence that the machinery repaired in 17.7 can explore a tree at all — thirteen nodes, up to twelve distinct traces, depth 3 — rather than collapsing to one child the way the buggy version did. The harness works. It just had nothing to do in the fair arms.
17.10 When a Stronger Model Makes Search Irrelevant
The two runs in 17.8 line up into a single observation, and it is worth stating carefully because it is easy to overclaim.
On this fixture:
- the weaker model (qwen3, 8B) reached the correct diagnosis on the first try, under every strategy, and additional search did not raise the score;
- the stronger model (
muse-spark-1.3-contributor) did the same, at a higher ceiling.
Neither model needed search here because the task did not require it. The evidence packet was frozen and small, the fault was a single attribute-name mismatch, and one direct generation was enough to name it. Search allocates compute across competing reasoning paths; when the first path reaches the answer, there is nothing to allocate.
This experiment shows one important case in which search complexity is unnecessary: the evidence is already in context, the fault is local, and a single generation is enough to identify it. Tasks with that shape may not justify explicit tree search at all. In this experiment, the early-stopped fair tree arms used three model calls versus one for Direct, while the deliberately invalid oracle control used roughly twenty-five. The audit in 17.7 matters more in this light, not less: if you are going to pay for search infrastructure, you need to know the search is real — and you need a task on which it earns its cost.
What this does not show — the scope guards, all four:
- Not “MCTS equals Best-of-N.” The fair arms never spent their budgets; the comparison did not run.
- Not “MCTS is more efficient than Best-of-N.” Same reason. The trees used fewer calls only because they stopped early, not because they allocated a budget more cleverly.
- Not “reflection adds no value.” Reflection never fired in a fair arm; the run contains no information about it either way.
- Not “tree search is unnecessary.” It was unnecessary for this fixture, with these two models. A harder fault, a larger evidence space, or a task that requires chaining several inferences could change every row of the table.
The general lesson is about experiment design rather than about MCTS: before paying for search, confirm the task is one where a single reasoning path is not already enough. When Direct scores the same as MCTS, the interesting question is not “which search strategy wins” — it is “why are we searching.”
What Usually Goes Wrong
The implementation and the two runs give a practical diagnosis table:
| Symptom | Likely Cause | Fix |
|---|---|---|
| Tree has many duplicate traces | Cache on (state, trace) prevents diversity | Remove cache; let LM generate fresh samples |
UCT with c=0 still explores | ucb_weight or default treats 0 as falsy | Use None check: weight if ucb_weight is not None else default |
| Node scores are inflated | Double-counting in _evaluate and _backpropagate | Only _backpropagate modifies reward |
| Search finds great path that final answer ignores | Value function and answer generator misaligned | Ensure GenerateAnswer receives the actual best trace |
| Search quality varies wildly across runs | Generation randomness dominates the search policy | Pin and record the sampling configuration; repeat the comparison; treat temperature or n as declared experimental variables rather than silently tuning them |
| LM calls exceed budget | Scoring and reflection also consume LM calls | Budget all LM interactions — generation, scoring, reflection, synthesis — not just expansions |
| Reflection makes no difference | Lessons are generic (“check the code carefully”) | Read the recorded lessons; tighten the Reflect signature to demand a named wrong assumption; compare against the reflect=False arm |
| Best path found but not selected | Selection based on visits not score | Use score for final selection, visits for search policy |
Conclusion
MCTS does not make the model reason correctly. It allocates inference-time computation toward the reasoning paths that its value function prefers.
If the value function tracks the task, this is useful. If the value function has a blind spot, search concentrates compute around the blind spot. If the value function can see the reference answer, search is not reasoning better—the experiment is leaking.
This is the same structure the optimizer chapters had, applied to a new object:
optimizer chapters this chapter
program candidate reasoning candidate
↓ ↓
metric search-time value
↓ ↓
selection selection
↓ ↓
better-scoring candidate better-scoring candidate
The objects differ; the danger does not: search does not know truth; search knows reward.
That places this chapter in the book’s last arc of “things that search”:
search program states Chapters 11–14
search environment/actions Chapter 15
search evidence Chapter 16
search reasoning states Chapter 17 ← this chapter
protect all search Chapter 18
promote Chapter 19
compose everything Chapter 20
The evidence status at the top of the chapter is the summary: the agent is built, its search behaviour is audited across 27 unit tests, and the seven-strategy comparison was run — seeds 17–19, three repetitions, two models. The comparison returned a null result. On the frozen cost-vs-price fixture, every strategy reached the same substantive diagnosis as a single direct call, for both an 8B local model and a stronger hosted one. The small qwen3 score variation came from whether the generated wording explicitly named service.py, not from a different diagnosis. And the run could not test its second question at all: the fair tree arms never spent their search budgets, because the value function was satisfied by the first expansion.
That is not a disappointing outcome; it is a precise one. It says the chapter’s weight sits in the audit — 17.7 and 17.8.2. Source code that looks like search still has to be read as a causal graph: the cache can collapse branching, a separately named post-hoc ValueEstimator does not guide search merely by existing, ucb_weight=0 can silently fail to disable exploration, and stale per-run state can contaminate the next condition. And an experiment that looks like a strategy comparison still has to be checked for whether the strategies ran — here they did not. The oracle condition in 17.9 shows the other half: the comparison’s validity depends on what the actual search-time value function is allowed to see. A search policy is only as legitimate as the evidence available to the function that guides it. The next chapter closes that boundary across optimization, tools, retrieval, memory, and reasoning search.
A production reasoning-search system with a written-down call budget, a trace cache, and a scheduled evaluation policy is excerpted in the Chapter 21 appendix.
Further Reading
- Tree of Thoughts (Yao et al., 2023): Introduces deliberate reasoning as search over thought trees. (arXiv:2305.10601)
- Reasoning with Language Model is Planning with World Model (Hao et al., 2023): RAP treats reasoning as planning and uses MCTS. (arXiv:2305.14992)
- Language Agent Tree Search (Zhou et al., 2023): LATS combines MCTS with LM reasoning, acting, and value functions. (arXiv:2310.04406)
- Monte Carlo Tree Search: A Review (Browne et al., 2012): The definitive survey of MCTS. (IEEE)