Table of Contents
Fetching ...

Accelerating Large Language Model Reasoning via Speculative Search

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

TL;DR

SpecSearch tackles the latency bottleneck of tree-search-based reasoning in large language models by introducing a bi-level speculative framework that couples a small draft model with a large model at both thought and token levels. It incorporates a quality-preserving rejection mechanism and a threshold-estimation scheme with exponential moving averages to maintain undegraded reasoning quality, supported by theoretical guarantees. Empirically, SpecSearch achieves up to $2.12\times$ speedup while preserving reasoning quality across Qwen and Llama models on math and benchmark tasks, and it remains compatible with multiple search strategies and thought evaluators. This work enhances the practicality of slow-thinking LLMs by delivering faster, scalable, and flexible reasoning acceleration that generalizes across domains.

Abstract

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.

Accelerating Large Language Model Reasoning via Speculative Search

TL;DR

SpecSearch tackles the latency bottleneck of tree-search-based reasoning in large language models by introducing a bi-level speculative framework that couples a small draft model with a large model at both thought and token levels. It incorporates a quality-preserving rejection mechanism and a threshold-estimation scheme with exponential moving averages to maintain undegraded reasoning quality, supported by theoretical guarantees. Empirically, SpecSearch achieves up to speedup while preserving reasoning quality across Qwen and Llama models on math and benchmark tasks, and it remains compatible with multiple search strategies and thought evaluators. This work enhances the practicality of slow-thinking LLMs by delivering faster, scalable, and flexible reasoning acceleration that generalizes across domains.

Abstract

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 speedup with comparable reasoning quality.
Paper Structure (46 sections, 5 theorems, 61 equations, 10 figures, 11 tables, 2 algorithms)

This paper contains 46 sections, 5 theorems, 61 equations, 10 figures, 11 tables, 2 algorithms.

Key Result

Theorem 4.3

(Quality-Preserving Condition on the Threshold) The generator $G_s\left(\left\{\beta^{(k)}\right\}_{k=1}^{K}\right)$ preserves undegraded quality if the following condition holds: $\beta^{(k)} \ge \mu_p^{(k)}, \;\forall k = 1, 2, \dots, K$.

Figures (10)

  • Figure 1: (a) The inference latency increases by several orders of magnitude with the introduction of tree-search-based reasoning methods. (b) Thought generation acts as an efficiency bottleneck of tree-search-based reasoning methods.
  • Figure 2: (a) Small models can generate thoughts with high reward scores. (b) Simple large model engagement strategies at the thought level struggle to preserve comparable reasoning quality.
  • Figure 3: Illustration of our proposed SpecSearch. SpecSearch proposes a bi-level speculative thought generator with a quality-preserving rejection mechanism, which significantly accelerates LLM reasoning while preserving comparable quality.
  • Figure 4: To verify that our method preserves comparable reward scores for reasoning thoughts, we visualize the average reward scores at each reasoning step during the tree search process.
  • Figure 5: The bound decent rapidly with the reasoning steps. However, the bound remains as high as $0.90$ even at the $10$-th step.
  • ...and 5 more figures

Theorems & Definitions (10)

  • Definition 4.1
  • Definition 4.2
  • Theorem 4.3
  • Theorem 4.5
  • Theorem 4.6
  • Lemma 1.1
  • proof
  • proof
  • Proposition 1.2
  • proof