PaperlayerView exact revision on arXiv

Source-linked AI summary

Accelerating Large Language Model Reasoning via Speculative Search

Zhihai Wang, Jie Wang, Jilai Pan, Xilin Xia, Huiling Zhen, Mingxuan Yuan, Jianye Hao, Feng Wu

arXiv:2505.02865v2cs.CLcs.AI

TL;DR

Tree-search reasoning improves LLM problem solving but incurs substantial latency from generating many thoughts. SpecSearch uses small–large model collaboration and quality-preserving rejection to accelerate thought generation, achieving up to 2.12× speedup with comparable reasoning quality across Qwen and Llama models.

  • Problem

    Tree-search-based reasoning methods face substantial inference latency because they must generate numerous reasoning thoughts, limiting practical LLM deployment.

    Source: §1 · p. 1

  • Method

    SpecSearch combines small and large models at thought and token levels and rejects thoughts whose quality is below the large model’s outputs.

    Source: §1 · pp. 1–2, §4.2 · p. 4

  • Results

    Up to 2.12× speedup was achieved with comparable reasoning quality across Qwen and Llama models on MATH and GSM8K.

    Source: §1 · p. 2, §6 · p. 9

  • Takeaways & Limitations

    SpecSearch can substantially accelerate existing tree-search reasoning methods while preserving high reasoning quality and integrating with multiple search algorithms and thought evaluators.

    Source: §1 · p. 2, §6 · p. 9

    The theoretical quality guarantee assumes descending quality and bounded variance during beam search.

    Source: §4.4 · p. 6

Abstract

from arXiv · show

Tree-search-based reasoning methods have significantly enhanced the reasoning capability of large language models (LLMs) by facilitating the exploration of multiple intermediate reasoning steps, i.e., thoughts. However, these methods suffer from substantial inference latency, as they have to generate numerous reasoning thoughts, severely limiting LLM applicability. To address this challenge, we propose a novel Speculative Search (SpecSearch) framework that significantly accelerates LLM reasoning by optimizing thought generation. Specifically, SpecSearch utilizes a small model to strategically collaborate with a large model at both thought and token levels, efficiently generating high-quality reasoning thoughts. The major pillar of SpecSearch is a novel quality-preserving rejection mechanism, which effectively filters out thoughts whose quality falls below that of the large model's outputs. Moreover, we show that SpecSearch preserves comparable reasoning quality to the large model. Experiments on both the Qwen and Llama models demonstrate that SpecSearch significantly outperforms state-of-the-art approaches, achieving up to 2.12$\times$ speedup with comparable reasoning quality.

1. Introduction

Tree-search-based reasoning improves LLM problem solving by exploring structured chains of intermediate thoughts, but generating many thoughts creates severe latency. SpecSearch addresses this bottleneck with bi-level collaboration, quality-preserving rejection, and comparable reasoning quality at substantially higher speed.

1. Introduction

  • Tree-search methods improve complex, multi-step reasoning by systematically exploring diverse paths built from intermediate thoughts.They combine thought generation and evaluation with search algorithms such as beam search and MCTS.Source: §1 · p. 1
  • Inference latency can increase by several orders of magnitude because tree-search methods must generate vast numbers of reasoning thoughts.This latency limits practical deployment, especially in real-time applications, while efficient quality-preserving acceleration remains underexplored.Source: §1 · p. 1, Figure 1 · p. 1
  • SpecSearch uses a small model with a large model at both coarse-grained thought and fine-grained token levels to accelerate thought generation.Its bi-level speculative thought generator targets the primary efficiency bottleneck in tree-search reasoning.Source: §1 · pp. 1–2
  • SpecSearch filters thoughts below the large model’s quality using non-parametric estimates based on historical large-model reasoning thoughts.The framework also provides a theoretical guarantee of reasoning quality comparable to the large model.Source: §1 · p. 2
  • Up to 2.12× speedup is achieved on MATH and GSM8K with Qwen and Llama while maintaining comparable reasoning quality and compatibility with multiple search algorithms and evaluators.The experiments report substantial acceleration over state-of-the-art approaches without compromising reasoning quality.Source: §1 · p. 2

2. Related Work

Prior acceleration methods adapt speculative execution to LLM decoding and tree-search reasoning, but mainly operate at the token level. SpecSearch instead treats both thoughts and tokens as speculative units to better match the structure of search-based reasoning.

Speculative Decoding

  • Speculative decoding drafts tokens with a small model and verifies them in parallel with a large model to accelerate autoregressive inference.This draft-then-verify paradigm preserves alignment with the target model’s distribution.Source: §2 · p. 2
  • Token-level speculative decoding can be poorly aligned with tree-search reasoning because search involves intricate, non-sequential thought generation.The mismatch can produce suboptimal acceleration for reasoning methods organized around thoughts.Source: §2 · p. 2

Tree-Search-Based Reasoning Acceleration

  • Existing tree-search acceleration methods primarily speculate at the token level, whereas SpecSearch uses a bi-level generator spanning thought and token levels.This design exploits thoughts as fundamental reasoning units and targets significant acceleration with comparable quality.Source: §2 · p. 2

3. Background

Tree-search reasoning represents intermediate steps as nodes and candidate solutions as paths, using generation, evaluation, and search to construct a final reasoning path. Speculative sampling accelerates token decoding through draft-and-verify execution while preserving the target distribution.

Speculative Sampling in LLM Decoding

  • Speculative Sampling drafts γ tokens autoregressively with a small model and verifies them in parallel with the target model.Rejected tokens are resampled, preserving alignment with the target model’s distribution while improving inference speed.Source: §3 · pp. 2–3

Tree-Search-Based Reasoning Methods

  • Tree-search reasoning models thoughts as nodes and solution candidates as paths, combining a Thought Generator, Thought Evaluator, and search algorithm.The generator expands child thoughts, the evaluator scores them, and iterative search produces a final reasoning path.Source: §3 · p. 3

4. Speculative Search for LLM Reasoning

SpecSearch accelerates tree-search reasoning through bi-level collaboration between small and large models, while targeting comparable quality through quality-preserving rejection and theoretical guarantees.

4.1. Overview of Speculative Search Framework

  • Thought generation dominates tree-search computation, consuming over 91% of total reasoning runtime and motivating targeted acceleration.Small models can outperform the larger model’s average reward on over 40% of analyzed thoughts, while simple thought-level refinement still struggles to preserve quality.Source: §4.1 · p. 3, Motivation 3: Simple large model… · p. 3, Figure 2 · p. 3

4.2. Bi-Level Speculative Thought Generator

  • SpecSearch combines thought-level drafting, evaluation, rejection, and token-level correction into a bi-level speculative generator.A small model drafts multiple thoughts, an evaluator rejects low-quality candidates, and lossless speculative decoding corrects rejected thoughts.Source: Motivation 3: Simple large model… · p. 3, §4.2 · p. 4
  • Thought-level speculation exploits reasoning-tree structure and typically longer-than-fifty-token thoughts to increase small-model token generation and efficiency.It also leverages small-model reasoning ability to potentially maintain or improve quality relative to the large model.Source: §4.2 · pp. 3–4
  • Large-model verification of thoughts is difficult because reasoning paths can be semantically ambiguous, exponentially numerous, and valid despite differing from the large model’s path.The framework instead uses an existing thought evaluator, such as a process reward model, to assess thought quality.Source: §4.2 · p. 4

4.3. Quality-Preserving Rejection Mechanism

  • Rejection balances acceleration and quality: accepting all small-model thoughts maximizes speed but can reduce quality, whereas rejecting all reverts to the slower large model.The mechanism therefore seeks an intermediate rejection criterion.Source: §4.3 · p. 5
  • The rejection mechanism estimates a step-wise large-model quality threshold from historical corrected thoughts using non-parametric estimation and an exponential moving average.Passed small-model thoughts can also be incorporated as approximate upper-bound evidence when large-model samples are limited.Source: §4.3 · pp. 5–6

4.4. Theoretical Guarantee of Undegraded Quality

  • SpecSearch provides a quality-preserving condition: when its threshold meets or exceeds the large model’s quality, generated reasoning remains undegraded.The framework’s theorem establishes this guarantee under the stated threshold condition.Source: §4.4 · p. 6
  • The overall quality-preservation probability lower bound increases with draft sample size N and approaches 1 as N →∞.This result is derived for beam search with up to K reasoning steps and threshold weight θ ≥γ.Source: §4.4 · pp. 6–7

5. Experiments

Experiments show that SpecSearch accelerates tree-search reasoning across models, search algorithms, and thought evaluators while largely preserving reasoning quality.

Experiment 1. Main Evaluation

  • Up to 1.72× speedup over SpS and 3.35× over AR is achieved on MATH-100 with Qwen.SpecSearch consistently outperforms the baselines in the main evaluation.Source: §5 · p. 7, Table 1 · p. 8

Search Algorithms

  • SpecSearch maintains comparable accuracy while reducing inference latency across beam search and MCTS.The comparison includes AR and SpS on both search algorithms.Source: §5 · p. 7, Table 2 · p. 8

Different Thought Evaluators

  • Up to 2.12× speedup with nearly unchanged accuracy is obtained across the Math-Shepherd and MATH-psa process reward models.This evaluation uses beam search on GSM8K-100.Source: §5 · p. 7

Contribution of Each Component

  • Both evaluation and rejection modules are essential for preserving reasoning quality.The ablation compares large-model log-probability evaluation with fixed-engagement, fixed-threshold, and random-rejection variants.Source: §5 · pp. 7–8, Table 3 · p. 8

Sensitivity Analysis

  • SpecSearch achieves 2.18×–2.87× speedups across draft-model sizes and similar average performance across a wide range of EMA weights.These sensitivity results support robustness across small-model settings and EMA-weight choices.Source: §5 · p. 8

Experiment 4. Visualization Analysis

  • Average reward scores remain comparable to those of the large model at every reasoning step.The visualization uses the MATH-100 dataset during tree search.Source: Figure 4 · p. 8, §5 · p. 9

6. Conclusion

SpecSearch concludes that small-large model collaboration can accelerate reasoning while preserving comparable quality, supported by theoretical analysis of its rejection mechanism.

6. Conclusion

  • SpecSearch combines small-model speculative thoughts with large-model collaboration at thought and token levels, achieving up to 2.12× speedup with comparable reasoning quality.The conclusion presents this framework and outcome as the paper’s central result.Source: §6 · p. 9

A. Theoretical Analysis

  • The theoretical analysis assumes independent normal quality distributions for small- and large-model thoughts at each reasoning step.The analysis is simplified by restricting the nonparametric estimator to the sample mean and assuming an additional large-model thought when M = 0.Source: A. Theoretical Analysis · p. 12

A.1. Proof for Theorem 4.3

  • The proof of Theorem 4.3 models threshold passing, retained thoughts, and large-model refinement to establish a lossless thought generator.Passing thoughts follow a binomial count, while their qualities are treated as a truncated normal distribution.Source: §A.1 · pp. 12–14

A.2. Further Discussion on Theorem 4.3

  • The lossless condition is characterized as a necessary and sufficient threshold condition for every reasoning step.The proposition connects the threshold inequality to the definition of a lossless thought generator.Source: §A.2 · pp. 14–15

A.3. Proof for Theorem 4.5

  • The proof of Theorem 4.5 derives a recursive variance expression for the estimated threshold using variance decomposition, Jensen’s inequality, and Cantelli’s inequality.The retained draft-thought count U^(k) and the sampling probability p_s enter the variance analysis.Source: §A.3 · pp. 16–17

A.5.1. NUMERICAL ANALYSIS FOR PROBABILITY BOUND

The analysis evaluates SpecSearch’s probability bound and threshold estimators, then provides background on speculative sampling and tree-search reasoning. Under stated idealized settings, the quality-preservation bound remains 0.90 after 10 reasoning steps.

A.5.1. NUMERICAL ANALYSIS FOR PROBABILITY BOUND

  • 0.90 probability lower bound remains at the 10th reasoning step under the stated parameter settings.The analysis notes that these conditions are highly idealized but provide theoretical support for high-quality generated thoughts.Source: §A.5.1 · p. 18, Figure 5 · p. 19

A.5.2. THRESHOLD ESTIMATOR WITH MAXIMUM ESTIMATION

  • Maximum estimation increases the probability of producing quality-preserved thoughts, but requires additional computation.Source: §A.5.2 · p. 19

A.5.3. THRESHOLD ESTIMATOR WITH NO ADDITIONAL Gp SAMPLES

  • With no additional large-model samples, the analysis derives the same conclusion as Theorem 4.3 and updates the probability bound for quality preservation.Source: §A.5.3 · pp. 19–20

B. More Background

  • Speculative sampling drafts tokens with a small model and verifies them in parallel with a large model, while beam search and MCTS explore reasoning trees differently.Source: B. More Background · p. 20

C. Implementation Details of the Baselines

The baseline implementation uses OpenR code for autoregressive reasoning and combines speculative sampling with Beam Search and MCTS. Speculative sampling is implemented with vLLM for efficient parallel decoding.

C.1. AR

C.2. SpS

  • Speculative sampling baselines use vLLM and a small model to accelerate Beam Search and MCTS generation.Source: §C.2 · p. 21

D.1. The Datasets Used in the Main Evaluation

The main evaluation uses 100 randomly selected problems each from MATH and GSM8K, while ablations use 50 MATH problems. Qwen and Llama model pairs provide the large and small thought generators, with two PRMs serving as evaluators.

D.1. The Datasets Used in the Main Evaluation

  • 100 randomly selected problems from MATH and 100 from GSM8K form the main evaluation sets because tree-search inference remains costly.Source: §D.1 · p. 21

D.2. The Datasets Used in the Ablation Study

  • The ablation study uses 50 MATH problems to support repeated variant and hyperparameter experiments.Source: §D.2 · p. 21

E. Illustration of Using Models

F.1. Comparison with Existing Speculative Decoding Techniques

SpecSearch extends speculative decoding from tokens to reasoning thoughts, verifies contextual thought quality, and rejects lower-quality thoughts. Compared with Treebon, it uses flexible thought units, adaptive rejection, and a formal reasoning-quality guarantee.

F.1. Comparison with Existing Speculative Decoding Techniques

  • SpecSearch treats both thoughts and tokens as speculative units, aligning speculation with reasoning structure while retaining token-level compatibility.Source: §F.1 · p. 22
  • Contextual verification accepts correct but non-aligned reasoning outputs, increasing acceptance rates and speedup relative to strict token alignment.Source: §F.1 · p. 22
  • SpecSearch rejects entire thoughts only when their contextual quality is below the large model’s output quality.Source: §F.1 · p. 22

F.2. Comparison with Treebon (Qiu et al., 2024)

  • Unlike Treebon’s fixed-length token sequences, SpecSearch uses dynamically sized, semantically coherent reasoning thoughts alongside token-level speculation.Source: §F.2 · p. 22
  • SpecSearch adaptively rejects thoughts with lower contextual quality than the large model, whereas Treebon rejects a fixed proportion using a preset threshold.Source: §F.2 · p. 22
  • SpecSearch formally assures reasoning quality comparable to the large model, unlike Treebon’s lack of theoretical guarantees.Source: §F.2 · p. 22, §F.1 · p. 22

G. Implementation Details of Our SpecSearch

SpecSearch combines process-level evaluation with parallel small-model generation, quality-based rejection, and fallback large-model generation to accelerate tree-search reasoning. Its implementation uses adaptive thresholds and standard beam-search settings.

G.1. Discussion on Advantages of Our Evaluation Method

  • The process reward model captures complete thought semantics, structures thought-quality distributions, and scores multiple valid reasoning paths.Source: §G.1 · p. 22

G.2.1. SMALL MODEL PARALLEL THOUGHT GENERATION

  • Small models generate 2*N thoughts in parallel, exploiting their low memory footprint while retaining high-quality candidates for Tree of Thoughts.Source: §G.2.1 · p. 22

G.2.3. ALGORITHM IMPLEMENTATION

  • SpecSearch evaluates large-model seeds, initializes thresholds from quality estimates, iteratively updates the beam, and stops when all beams reach <STOP>.Source: §G.2.1 · p. 23

G.2.2. ACCEPTANCE-REJECTION MECHANISM

  • Each small-model thought is retained when its process reward exceeds a dynamic threshold, with only the top N retained when necessary.Source: §G.2.2 · p. 23

G.2.4. SPECULATIVE MODEL SERIAL THOUGHT GENERATION

  • If fewer than N small-model thoughts survive rejection, the large model serially generates additional thoughts until N candidates are available.Source: §G.2.4 · p. 23

G.3. Hyperparameters

  • The experiments use θ = 0.9, beam-search tree width 6, depth 50, and beam size 2, while the reported full-GSM8K setup uses quantized Qwen models with MATH-psa.Source: §G.3 · p. 23, Table 4 · p. 24

H.1. More Main Evaluation

Across broader datasets and settings, SpecSearch maintains comparable accuracy while accelerating reasoning, including diverse tasks, search algorithms, evaluators, and draft-model sizes. Supporting analyses explain its dynamic thresholds, model collaboration, and remaining PRM failure mode.

H.1. More Main Evaluation

  • Speedups range from 2.04× to 2.84× across GSM8K, AIME, Olympiad, and code-generation evaluations while maintaining comparable reasoning accuracy.The full GSM8K result reaches 2.84× speedup with a 0.83% accuracy reduction, while AIME and Olympiad retain identical accuracy to SpS at 1.21× and 1.37× speedups.Source: §H.1 · p. 24

H.2. More Motivating Results

  • Average reward decreases across later reasoning steps, supporting dynamic thresholds that account for increasing step difficulty.Source: §H.2 · p. 25, Figure 7 · p. 26

H.3. Case Study

  • Different reasoning steps vary in difficulty, enabling small models to handle simpler steps while large models handle harder ones; their similar token lengths support collaboration.Source: §H.3 · p. 25, §H.2 · p. 25, Table 8 · p. 26

H.4. More Broad Compatibility Results

  • SpecSearch accelerates both beam search and MCTS with minimal accuracy loss and remains consistent across different process reward models and datasets.Source: §H.4 · p. 26, Table 9 · p. 27

H.5. Sensitivity Analysis to Draft Model’s Size

  • Using multiple small draft models, SpecSearch achieves speedups from 2.18× to 2.87× while maintaining stable accuracy.Source: §H.5 · p. 26, Table 10 · p. 27

H.6. More Ablation Study Results

  • As θ increases from 0.8 to 0.95, accuracy remains largely unchanged while latency decreases, indicating robust efficiency gains near θ = 1.Source: §H.6 · p. 26, Table 11 · p. 27

Related papers

Loading 2505.02865v2…