Domain-Specialized Tree of Thought through Plug-and-Play Predictors

Xuanqi Gao; Haoyu Wang; Jun Sun; Shiqing Ma; Chao Shen

cs.AIMar 14, 2026

Domain-Specialized Tree of Thought through Plug-and-Play Predictors

Authors: Xuanqi Gao, Haoyu Wang, Jun Sun, Shiqing Ma, Chao Shen

Organizations: Xi’an Jiaotong University, China · Singapore Management University, Singapore · University of Massachusetts at Amherst, United States

Abstract

While Large Language Models (LLMs) have advanced complex reasoning, prominent methods like the Tree of Thoughts (ToT) framework face a critical trade-off between exploration depth and computational efficiency. Existing ToT implementations often rely on heavyweight LLM-based self-evaluation or rigid heuristics for branch pruning, making them prohibitively expensive and inflexible for broad application. To address this, we introduce DST, an adaptable, plug-and-play predictor that serves as a lightweight, supervised heuristic to guide the ToT search process. Our predictor enables dynamic, context-aware pruning, allowing the search to proceed with near-greedy efficiency on simpler reasoning steps while adaptively expanding the search beam only when encountering uncertainty or task complexity. We evaluate our approach on a diverse suite of benchmarks spanning mathematical reasoning, general reasoning, and complex logical reasoning. Experimental results demonstrate that our method achieves accuracy competitive with or superior to strong baselines, including standard ToT, while reducing computational overhead by 26-75%. Our work effectively resolves the accuracy-efficiency trade-off in tree-based reasoning, transforming ToT from a resource-intensive technique into a scalable and practical paradigm for complex problem-solving in LLMs.

Figures & tables

Method	Dynamic Branching	Evaluation Signal	Token Cost (vs. CoT)	Domain Adaptation
ToT ( Yao et al., 2023 )	✗	LLM Self-Eval	$\sim$ 10x	✗
ProbTree ( Cao et al., 2023 )	✗	Probabilistic	$\sim$ 5x	✗
DPTS ( Ding et al., 2025 )	✓	Confidence	$\sim$ 4x	✗
AGoT ( Pandey et al., 2025 )	✓	Task-specific	$\sim$ 4x	✗
iToT ( Boyle et al., 2024 )	✓	Tool-cost aware	$\sim$ 3x	✗
BPP-Search ( Wang et al., 2025 )	✓	RL preference	$\sim$ 3x-5x	✗

Table 1: Comparison of Tree-of-Thought reasoning methods.

Dataset	Seed Problems	Generated Trees	Training Samples
GSM8K	200	6,000	6,000
SVAMP	100	6,200	6,200
Minerva	28	1,700	1,700
MATH-500	50	3,000	3,000
GPQA	45	1,800	1,800
BoardgameQA	20	2,400	2,400

Table 2: Training data statistics for predictor training across domains.

Figure 2: Accuracy vs. Efficiency Trade-off. Each point represents the performance of a method on a specific task and model, plotted as accuracy gain (percentage points) versus efficiency gain (percentage cost reduction) relative to CoT.

Table 3: Comparison of Tree-of-Thought reasoning methods. Best performance for each metric (highest accuracy, lowest cost) is shown in bold . Results are relative improvements over CoT.

Panel of Table 4 — Table 4: Cross-model transfer performance on GSM8K. Rows indicate the model used for predictor training; columns indicate the evaluation model. Diagonal values (bold) show native performance. Values show accuracy (%) / cost (tokens).

Panel of Table 5 — Table 5: Cross-domain transfer performance using Qwen3-8B. Predictor trained on one dataset (rows) and evaluated on another (columns). Diagonal values (bold) show native training. Values show DST accuracy (%) / cost (tokens).

Method	GSM8K		GPQA
Method	Accuracy (%)	Avg. Tokens	Accuracy (%)	Avg. Tokens
DST	92.4	992	49.7	8230
DST w/o $c_{s}$	90.1	1150	47.0	8500
DST w/o $\mathbf{v}_{s}$	85.7	1300	42.3	9800

Table 6: Impact of State Feature Components on GSM8K and GPQA.

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 3: Accuracy vs. Average Tokens as a function of Beam Width (

b

) on BBEH-BoardgameQA. The red dot marks our default setting of

b=3

, which offers a strong balance between performance and cost.

Figure 4: Top: Accuracy vs. Average Tokens as a function of Pruning Threshold (

\tau

) on GSM8K (left) and BBEH-BoardgameQA (right). Bottom: The corresponding Shortcut Rate for each dataset.

Figure 5: Accuracy vs. Average Tokens as a function of Discount Factor (

\gamma

) on GSM8K (left) and GPQA (right). The red marker indicates the optimal performance point, achieved at

\gamma=0.99

Training Data	10%	25%	50%	100%
Accuracy (%)	88.2	90.1	91.5	92.4
Tokens (avg)	1,180	1,050	1,010	992

Appendix

Table 7: Impact of training data scale on predictor performance (GSM8K, Qwen3-8B). Even with 10% of training data, the predictor achieves competitive performance.

Explore similar work

Beyond Fixed Budgets: Characterizing the Inelasticity and Limitations of Tree-of-Thought Reasoning Strategies
May 19, 2026Atkia Mahila, Avinash Maurya, M. Mustafa Rafique +1LLM Evaluation Inference-Time Search
Rochester Institute of Technology Rochester, USA · Argonne National Laboratory Lemont, USA
Breaking the Reward Barrier: Accelerating Tree-of-Thought Reasoning via Speculative Exploration
May 11, 2026Shuzhang Zhong, Haochen Huang, Shengxuan Qiu +3Inference-Time Search Tree of Thoughts
Institute for Artificial Intelligence, Peking University · School of Integrated Circuits, Peking University · 3ByteDance Seed
SAT: Balancing Reasoning Accuracy and Efficiency with Stepwise Adaptive Thinking
Apr 9, 2026Weiyang Huang, Xuefeng Bai, Kehai Chen +4Efficient Language Model Reasoning Overthinking in Language Models
Harbin Institute of Technology (Shenzhen), Shenzhen, China · Huawei Technologies

May 19, 2026J/K move · Enter open · S save

Beyond Fixed Budgets: Characterizing the Inelasticity and Limitations of Tree-of-Thought Reasoning Strategies

Atkia Mahila, Avinash Maurya, M. Mustafa Rafique, Bogdan Nicolae

Rochester Institute of Technology Rochester, USA · Argonne National Laboratory Lemont, USA

LLM Evaluation Inference-Time Search Tree of Thoughts Monte Carlo Tree Search LLM Mathematical Reasoning LLM Reasoning

Open paper AskPDF arXiv

Tree of Thought (ToT) search has become a promising direction for improving the reasoning capabilities of large language models, but deploying these methods in practice raises a question that has received little systematic attention: how do different search strategies behave under varying compute budgets, model sizes, and problem difficulties? In this work, we evaluate two representative ToT methods; DPTS, a Monte Carlo tree search based approach, and SSDP, a semantic deduplication based approach, across two mathematical reasoning benchmarks (Math500 and GSM8K), two model scales (Llama-3B and Llama-8B), and four token budgets (3k--10k). Our analysis reveals that the two methods exhibit limitations that pull in opposite directions. DPTS suffers from a cold-start bottleneck at low budgets: it requires sufficient exploration before its value estimates become reliable, making it a poor fit for resource-constrained settings despite strong scaling behavior at higher budgets. SSDP, on the other hand, reaches candidate solutions efficiently but is prone to frontier depletion; its aggressive node merging permanently discards unexplored paths, leaving it unable to improve regardless of how much budget remains. Together, these findings suggest that neither a fixed exploration strategy nor a fixed pruning strategy is sufficient across compute continuum. We argue that effective search for scientific reasoning agents requires strategies that can adapt their behavior based on search progress and available resources.

Abstract

Explore similar work

Beyond Fixed Budgets: Characterizing the Inelasticity and Limitations of Tree-of-Thought Reasoning Strategies

Breaking the Reward Barrier: Accelerating Tree-of-Thought Reasoning via Speculative Exploration

SAT: Balancing Reasoning Accuracy and Efficiency with Stepwise Adaptive Thinking