PaperlayerView exact revision on arXiv

Source-linked AI summary

Don't Get Lost in the Trees: Streamlining LLM Reasoning by Overcoming Tree Search Exploration Pitfalls

Ante Wang, Linfeng Song, Ye Tian, Dian Yu, Haitao Mi, Xiangyu Duan, Zhaopeng Tu, Jinsong Su, Dong Yu

arXiv:2502.11183v2cs.CL

TL;DR

Verifier-guided tree search can improve LLM reasoning, but its computational cost is a practical challenge. FETCH combines semantic state merging with verifier-variance reduction, and experiments report improved accuracy and efficiency across several datasets and algorithms.

  • Problem

    Tree-search reasoning faces computational costs, redundant semantically equivalent states, and high-variance verifier scores that can leave promising paths unfinished.

    Source: §1 · p. 2, §2.2 · p. 4

  • Method

    FETCH merges similar states and reduces verifier-score variance using adjusted λ-returns during training and score ensembling during inference.

    Source: §1 · p. 2

  • Results

    Across four search algorithms and three datasets, FETCH substantially improves accuracy and efficiency; variance reduction improves accuracy by approximately 1∼2 points across tested algorithms.

    Source: §1 · p. 2, §4.2 · p. 8

  • Takeaways & Limitations

    State merging and variance reduction improve the practical feasibility of advanced LLM reasoning without prohibitive computational costs.

    Source: §6 · p. 11

    On MATH, the method slightly increases computational costs, likely because search algorithms struggle to converge on confident solutions for difficult problems.

    Source: §4.2 · p. 8

Abstract

from arXiv · show

Recent advancements in tree search algorithms guided by verifiers have significantly enhanced the reasoning capabilities of large language models (LLMs), but at the cost of increased computational resources. In this work, we identify two key challenges contributing to this inefficiency: $\textit{over-exploration}$ due to redundant states with semantically equivalent content, and $\textit{under-exploration}$ caused by high variance in verifier scoring leading to frequent trajectory switching. To address these issues, we propose FETCH, an e$\textbf{f}$fici$\textbf{e}$nt $\textbf{t}$ree sear$\textbf{ch}$ framework, which is a flexible, plug-and-play system compatible with various tree search algorithms. Our framework mitigates over-exploration by merging semantically similar states using agglomerative clustering of text embeddings obtained from a fine-tuned SimCSE model. To tackle under-exploration, we enhance verifiers by incorporating temporal difference learning with adjusted $\lambda$-returns during training to reduce variance, and employing a verifier ensemble to aggregate scores during inference. Experiments on GSM8K, GSM-Plus, and MATH datasets demonstrate that our methods significantly improve reasoning accuracy and computational efficiency across four different tree search algorithms, paving the way for more practical applications of LLM-based reasoning. The code is available at https://github.com/Soistesimmer/Fetch.

1 Introduction

Verifier-guided tree search can strengthen LLM reasoning, but its computational costs hinder practical use. FETCH targets over-exploration from redundant states and under-exploration from verifier-score variance.

1 Introduction

  • Redundant semantically equivalent states cause over-exploration, while verifier scoring variance can trigger trajectory switching and leave promising paths unfinished.The variance is linked to reasoning style and phrasing differences, particularly on challenging math problems.Source: §1 · p. 2
  • FETCH is a plug-and-play framework that merges similar states through agglomerative clustering using embeddings from a SimCSE model fine-tuned for semantic similarity.The embedding model is fine-tuned using signals from computationally intensive methods such as LLM prompting or consistency checking over sampled subsequent steps.Source: §1 · p. 2
  • For scoring-based verifiers, adjusted λ-returns balance short- and long-term rewards during training, and inference-time ensembling aggregates scores from multiple verifiers.The paper presents these techniques as ways to reduce variance and average out individual verifier biases and variance.Source: §1 · p. 2
  • Experiments across four search algorithms and GSM8K, GSM-Plus, and MATH show substantial improvements in accuracy and efficiency.State merging prevents redundant-state over-exploration, while TD(λ) training and verifier ensembling contribute to more accurate, stable scores.Source: §1 · p. 2

2 Pilot Study

The pilot study uses BFS to examine search costs and two sources of inefficiency: redundant semantic states and unstable verifier scores. It reports that larger expansions increase cost, while score variance is especially pronounced on challenging problems.

2.1 Definitions and Setup

  • BFS represents reasoning as a tree of partial or complete states, selecting an unexplored node by verifier feedback and expanding it with N sampled reasoning steps.The study uses BFS as a representative tree-search example.Source: §2.1 · p. 3

2.2 Results and Analyses

  • 3× to 9× more tokens than greedy decoding accompanied only 0.5% accuracy improvement when comparing N = 5 with N = 10.BFS performance improved as expansion size N increased, but computational costs rose dramatically.Source: §2.2 · p. 4
  • Search costs were similar across problem-difficulty categories, suggesting resource allocation was poorly balanced between basic and complex problems.The analysis identifies over-exploration in basic problems and under-exploration in complex cases.Source: §2.2 · p. 4
  • Verifier scores varied substantially across correct trajectories, especially on challenging problems, which could trap search in a local branch and leave valuable paths unexplored.The pilot measured score standard deviations for sampled correct trajectories; the authors attribute further randomness to limited Monte Carlo rollouts.Source: §2.2 · p. 4

3 FETCH

FETCH combines redundant-state merging and verifier score-variance reduction in a plug-and-play framework designed to work with existing tree search algorithms. Its methods address over-exploration through semantic clustering and under-exploration through verifier training and inference changes.

3 FETCH

  • FETCH pairs redundant-state merging with score-variance reduction to address over- and under-exploration, respectively, and is designed to complement existing tree search algorithms.The two methods are described as complementary and orthogonal to widely used tree search algorithms.Source: §3 · p. 4

3.1.1 Semantic-Based State Merging

  • After each expansion, FETCH embeds newly created reasoning states and uses agglomerative clustering to group them into hyper-nodes, with a distance threshold that determines cluster count.A hyper-node represents multiple original nodes, and its verifier score can be calculated using max, average, or min aggregation; expansion selects original nodes in verifier-score order.Source: §3.1.1 · pp. 4–5

3.1.2 Embedding Model Fine-Tuning

  • Because pretrained embeddings can cluster math reasoning poorly, including by overlooking important numbers, FETCH post-trains embeddings using prompting-based or consistency-based data collection.The prompting approach uses an instruction-following LLM to judge state equivalence; the consistency approach labels pairs by comparing how swapping states affects subsequent policy rollout predictions.Source: §3.1.2 · p. 5

3.2 Handling Under-Exploration via Score Variance Reduction

  • For verifier training, FETCH uses TD(λ) returns to balance bias and variance, then trains the verifier with mean-square error.The trace-decay parameter λ controls the bias–variance trade-off; verifier outputs may be constrained to [0, 1] for compatibility with some search algorithms.Source: §3.2 · p. 6
  • At inference, FETCH averages predictions from multiple verifiers to reduce score variance and can run them in parallel on separate devices to avoid serial scoring delays.The ensemble can also avoid the need for training when open-source verifiers are available.Source: §3.2 · p. 6

4 Experiment

Experiments evaluate FETCH across three math datasets and four tree-search algorithms. Combining state merging with score variance reduction balances lower token use on simpler tasks with sufficient exploration on harder ones.

4.1 Setup

  • The evaluation uses GSM8K, GSM-Plus, and MATH, with BFS, Beam Search, Tree Search, and MCTS; accuracy and generated tokens measure performance and cost.GSM-Plus contributes one randomly sampled variant per GSM8K test case; MATH evaluation uses a 500-case subset.Source: §4.1 · pp. 6–7

4.2 Main Results

  • On BFS for GSM8K, state merging reduces #Token (k) from 2.47 to 0.88 while maintaining performance.The authors describe the reduction as nearly 3× and say saved computation can prioritize promising paths.Source: §4.2 · pp. 7–8
  • Beam Search has the smallest token savings because its fixed node-expansion count limits optimization, but state merging still improves accuracy.Source: §4.2 · p. 8
  • Score variance reduction improves accuracy by approximately 1–2 points across tested algorithms, with stable token use on GSM8K and GSM-Plus.On MATH, computational costs increase slightly, likely because the method allocates more computation to unfinished trajectories on difficult problems.Source: §4.2 · p. 8
  • Combining state merging and score variance reduction balances lower token use on simple tasks with sufficient exploration on complex ones.The reported results on GSM8K and MATH support this balance.Source: §4.2 · p. 8

4.3 Ablation Studies and Analyses

  • For state merging, larger d yields fewer clusters but raises the risk of grouping different states; f = max(*) performs best on Accuracy and #Token (k).With f = max(*), changing d from 0.1 to 0.2 has limited effect on performance and efficiency; agglomerative clustering outperforms k-means.Source: §4.3 · pp. 8–9, Figure 4 · p. 8, Table 2 · p. 9
  • Prompting- and consistency-based embedding post-training outperform pretrained SimCSE, while edit distance is sensitive to dataset variation.Increasing the embedding model from RoBERTa-Base to Large yields measurable performance gains.Source: §4.3 · p. 9, Table 2 · p. 9
  • TD(λ) and verifier ensembling improve search and step-level alignment; ensembling performs better than TD(λ) but requires multiple verifiers.Combining them yields further gains, and in BFS inference-time scaling, the combined methods achieve the best overall performance as expansion budgets increase.Source: §4.3 · p. 9, Table 3 · p. 9, Figure 5 · p. 10

5 Related Work

Prior work uses inference-time scaling and verifier-guided tree search to allocate computation across reasoning steps. The section identifies unresolved exploration problems and high-variance verifier training as gaps addressed by FETCH.

5 Related Work

  • Tree search methods such as Beam Search, MCTS, and A* allocate computation to critical reasoning steps, unlike full-path sampling in Self-Consistency (Wang et al., 2022).The cited tree-search work includes Beam Search (Yao et al., 2024; Zhu et al., 2024; Yu et al., 2024), MCTS (Tian et al., 2024; Zhang et al., 2024), and A* (Wang et al., 2024b;a).Source: §5 · p. 10
  • Most existing methods leave over- and under-exploration unresolved; Tian et al. (2024) prune repeated states but rely on edit distance or computationally costly LLM similarity checks.The section frames this as a need for a lightweight, robust mechanism for repeated-state handling.Source: §5 · p. 10
  • Value networks and process reward models use LLMs with classification or regression heads to produce scalar scores for guiding search.The passage cites value networks (Yu et al., 2024; Tian et al., 2024) and PRMs (Lightman et al. 2023; Wang et al. 2023).Source: §5 · p. 10
  • Prior MC-based verifier-label collection overlooks estimation variance, which can destabilize verifiers and lead to under-exploration; this study introduces TD(λ) and verifier ensembling.Human annotation of intermediate-step errors was described as cost-prohibitive for scaling training datasets (Lightman et al. (2023)); subsequent MC approaches include Yu et al. (2024), Wang et al. (2023), Tian et al. (2024), and Wang et al. (2024a).Source: §5 · p. 10

6 Conclusion

FETCH targets over-exploration and under-exploration in LLM tree search with state merging and variance reduction. The paper reports improved accuracy and efficiency, and identifies further optimization and extension to more complex tasks and domains as future directions.

6 Conclusion

  • FETCH addresses tree-search over- and under-exploration with state merging and variance reduction.The conclusion associates these enhancements with mitigating redundant states and stabilizing verifier score estimation.Source: §6 · p. 11
  • The paper reports that these methods improve the practical feasibility of advanced reasoning strategies and significantly enhance accuracy and efficiency.The conclusion frames the techniques as refining LLM deductive reasoning without prohibitive computational costs.Source: §6 · p. 11
  • Future work may further optimize the strategies and extend them to more complex tasks and domains.Source: §6 · p. 11

A Parameter Settings

The appendix specifies policy, verifier, embedding, and search-related parameter settings for GSM8K and MATH. It also describes the semantic-equivalence prompt and consistency-based embedding training choices.

A Parameter Settings

  • For GSM8K, the Llama-3-8B policy trains for 2 epochs at 5e-6, while its value network trains for 1 epoch at 2e-6 on 16 rollouts per question.The policy uses AdamW, a linear scheduler, and cross-entropy; the value network uses mean square error and temperature 1 rollouts from the initial state.Source: A Parameter Settings · p. 13
  • For MATH, the setup adopts the Mistral-7B policy and PRM training data released by Wang et al. (2023), then trains the PRM with the same value-training script.The trained PRM shows competitive performance with the official results (Wang et al., 2023).Source: A Parameter Settings · p. 13
  • State-merging embeddings use RoBERTa-Large pretrained with SimCSE and are post-trained on around 100k step pairs per dataset.The post-training uses binary cross-entropy with a learning rate of 1e-6; clustering uses d = 0.15 and max(∗) by default.Source: A Parameter Settings · p. 13
  • The prompting-based embedding method asks whether two reasoning steps from separate rationales are semantically equivalent and requests a Yes-or-No answer.Source: A Parameter Settings · p. 13
  • The consistency-based embedding method uses α = 0.02 and β = 0.08, with K = 1 for GSM8K and K = 2 for MATH.It is the default because it does not rely on another model.Source: A Parameter Settings · p. 13
  • GSM8K uses TD(λ) with λ = 0.8 and verifier ensembling with Nnv = 2; MATH uses ensembling with Nnv = 2 only.The MATH setup uses provided process rewards collected with heavy MC sampling and enough rollouts per intermediate state to reduce variance.Source: A Parameter Settings · p. 13

B Implementation of Search Algorithms

The appendix reports search-algorithm budgets and describes figures and a table used to examine verifier scores and similarity methods. The provided captions specify what is measured, not comparative outcomes.

B Implementation of Search Algorithms

  • Search settings use temperature τ = 0.8 except for greedy decoding, with algorithm-specific expansion budgets, beam size, expected accuracy, and rollout counts.BFS and Tree Search use N = 10 for GSM8K/GSM-Plus and N = 5 for MATH; Beam Search uses budget and beam size 5, while MCTS uses root N = 8, other-node N = 4, and 2 rollouts.Source: B Implementation of Search Algorithms · p. 13
  • Table 4 reports Pearson correlations between similarity-method scores and human annotations for 200 sampled pairwise examples from GSM8K or MATH.Source: Table 4 · p. 14
  • Figure 6 compares verifier-estimated-value Brier scores across reasoning steps and final accuracy, with Brier scores calculated via MSE.Source: Figure 6 · p. 14
  • Figure 7 compares score-estimation standard deviations across verifiers for 20 sampled GSM8K reasoning paths at different difficulty levels.Source: Figure 7 · p. 14

C Human Evaluation on Embedding Models

The human evaluation compares embedding-model similarity scores with annotations of whether reasoning steps are the same. The authors report that post-training is effective and that a larger RoBERTa model may improve performance further.

C Human Evaluation on Embedding Models

  • 200 step pairs from GSM8K and MATH were manually labeled by NLP-background annotators as same or different steps, with each annotation double-checked.Embedding-model scores were compared with these labels using Pearson Correlation Coefficient (PCC).Source: C Human Evaluation on Embedding Models · p. 14
  • The human-annotation evaluation validates the effectiveness of post-training the embedding models.The passage does not provide the PCC values.Source: C Human Evaluation on Embedding Models · p. 14
  • Comparing RoBERTa-Base and RoBERTa-Large suggests that scaling model parameters may further improve performance.The authors leave this direction for future work.Source: C Human Evaluation on Embedding Models · p. 14
  • Figure 8 covers ablations of the methods on token consumption across GSM8K and MATH problems of different difficulties using BFS.The caption states the figure's scope but supplies no comparative outcomes.Source: Figure 8 · p. 15

D Analysis on Variance of Verifier Scores

TD(λ) training and verifier ensembles are associated with lower verifier-score variance than the baseline across difficulty subsets; combining them makes further progress in most cases.

D Analysis on Variance of Verifier Scores

  • TD(λ) and verifier ensembles both yield lower verifier-score variance than the baseline across subsets of varying difficulty.The comparison is reported for scores from verifiers trained using different methods.Source: D Analysis on Variance of Verifier… · p. 15
  • Combining TD(λ) and a verifier ensemble makes further progress in reducing score variance in most cases.The passage does not claim improvement in every case.Source: D Analysis on Variance of Verifier… · p. 15
  • The authors conclude that TD(λ), verifier ensembling, and their combination can effectively reduce verifier-score variance.Figure 7 presents variance across difficulty subsets.Source: D Analysis on Variance of Verifier… · p. 15

Related papers

Loading 2502.11183v2…