SEED: Accelerating Reasoning Tree Construction via Scheduled Speculative Decoding

Zhenglin Wang; Jialong Wu; Yilong Lai; Congzhi Zhang; Deyu Zhou

Source-linked AI summary

SEED: Accelerating Reasoning Tree Construction via Scheduled Speculative Decoding

Zhenglin Wang, Jialong Wu, Yilong Lai, Congzhi Zhang, Deyu Zhou

arXiv:2406.18200v2Jun 26, 2024cs.CL

TL;DR

Complex reasoning and planning remain difficult for LLMs, while tree-search methods incur high latency from systematic exploration and evaluation. SEED uses scheduled speculative execution to coordinate draft models for thought generation and state evaluation, achieving speed improvements across three reasoning datasets while managing GPU resources. The framework provides a path toward batched, training-free speculative inference.

Problem
Tree-search reasoning improves exploration for complex tasks but introduces substantial inference latency through systematic thought exploration and evaluation.
Method
SEED applies scheduled speculative execution to accelerate thought generation and state evaluation while coordinating multiple draft models with one target model.
Results
SEED achieves speed improvements across GSM8K, Creative Writing, and Blocksworld, including 1.1−1.5× speedups and up to 20 additional tokens per second.
Takeaways & Limitations
SEED provides a viable path for batched inference in training-free speculative decoding while preserving the original distribution.
Takeaways & Limitations
The study optimizes tree-construction inference speed but not search speed, using BFS and leaving integration with more advanced search algorithms for future work.

Abstract

from arXiv · show

Large Language Models (LLMs) demonstrate remarkable emergent abilities across various tasks, yet fall short of complex reasoning and planning tasks. The tree-search-based reasoning methods address this by surpassing the capabilities of chain-of-thought prompting, encouraging exploration of intermediate steps. However, such methods introduce significant inference latency due to the systematic exploration and evaluation of multiple thought paths. This paper introduces SeeD, a novel and efficient inference framework to optimize runtime speed and GPU memory management concurrently. By employing a scheduled speculative execution, SeeD efficiently handles multiple iterations for the thought generation and the state evaluation, leveraging a rounds-scheduled strategy to manage draft model dispatching. Extensive experimental evaluations on three reasoning datasets demonstrate superior speedup performance of SeeD, providing a viable path for batched inference in training-free speculative decoding.

1 Introduction

Tree-search reasoning improves exploration for complex tasks but incurs substantial inference latency, especially during tree construction. SEED addresses this bottleneck with scheduled speculative execution that balances speed and GPU memory while accelerating thought generation and state evaluation.

Motivation: Tree-search methods reduce the limitations of simple prompting by exploring intermediate thoughts, but their systematic lookahead and backtracking create serious inference latency.The paper focuses on accelerating tree construction rather than search algorithms or prompting methods.
Proposed framework: SEED targets both runtime speed and GPU memory management for reasoning-tree construction across repeated same-prompt executions and evaluations with different prompts.Its scheduled speculative decoding manages parallel draft models for these two execution scenarios.
Proposed framework: Rounds-Scheduled execution uses an FCFS queue to coordinate multiple parallel draft models around one shared target model that cannot verify them simultaneously.The scheduling design draws inspiration from operating-system process scheduling.
Results: SEED achieves 1.1−1.5× speedups and generates up to 20 additional tokens per second across GSM8K, Creative Writing, and Blocksworld.These results support batched inference with training-free speculative decoding while preserving the original distribution.

2 Related Work

Related work improves LLM reasoning through tree search and accelerates search with diversity rewards or pruning. However, prior approaches leave the latency of iterative tree construction largely unaddressed, motivating scheduled speculative decoding.

Tree-search reasoning: Tree-search methods such as RAP, ToT, and RoT construct reasoning trees that explore multiple paths and evaluate states with search algorithms including BFS and MCTS.These methods augment reasoning by considering alternative reasoning paths before selecting subsequent actions.
Tree-search reasoning: Diversity rewards and pruning techniques improve search efficiency, but they do not address inference latency caused by iterative tree construction.SEED instead targets the construction process itself.
Parallel decoding: Speculative decoding, early exiting, and non-autoregressive decoding are parallel-decoding approaches developed to reduce LLM inference latency.This paper focuses on training-free speculative decoding because it can integrate with modular methods such as CoT and TSB without modifying model architecture or requiring additional training.

3 Preliminaries

Speculative decoding drafts tokens with a small model and verifies them in parallel with a larger target model, while TSB construction alternates thought generation and state evaluation. SEED applies this paradigm to accelerate the independent branches of both components.

3.1 Speculative Decoding: Speculative decoding uses a small draft model to generate tokens sequentially, then has a larger target model validate them.It is a two-stage Draft-then-Verify decoding paradigm.
3.1 Speculative Decoding: The target model verifies draft tokens in parallel, accepts them probabilistically, and resamples after the first rejection.This procedure is equivalent to direct sampling from the target LLM, preserving the generated-text distribution in SEED.
3.1 Speculative Decoding: Multiple candidates per draft position and tree attention can increase expected acceptance length and share generated-token caches during verification.SEED can leverage tree attention for further speedup.
3.3 TSB Task Formulation: TSB reasoning constructs a tree whose nodes represent partial solution states, with a Thought Generator producing paths and a State Evaluator assessing them.Using BFS, the root generates multiple paths from the initial prompt before the evaluator selects a path for continuation.

4 Method

SEED accelerates reasoning-tree construction by combining parallel drafting with sequential target-model verification under a rounds-scheduled execution strategy. Its FCFS queue and draft-label coordination manage multiple thought-generation and state-evaluation iterations without verification conflicts.

Speculative Scheduled Execution: SEED implements speculative scheduled execution in both the Thought Generator and State Evaluator, allowing distinct reasoning-tree branches to process in parallel.The framework targets tree construction rather than search or prompting, and uses separate branch executions to avoid interference.
Parallel Drafting Phase: Multiple small draft models generate token sequences concurrently, while one larger target model verifies those sequences sequentially.During thought generation, draft inputs are identical; during state evaluation, they may differ.
Sequential Verification Phase: During verification, accepted draft tokens are retained, whereas rejected sequences trigger target-model resampling before the draft process continues.The target model remains active across the schedule, reducing idle time while preserving speculative execution correctness.
Rounds-Scheduled Strategy: The Rounds-Scheduled strategy uses an FCFS queue to dispatch draft sequences and return each draft model for another round after verification.This scheduling design controls execution flow when one target model cannot verify multiple draft models simultaneously.
Algorithm: Algorithm 1 coordinates parallel drafting D(n) with cyclic verification V through verify and resampling phases, using a draft-label map to synchronize re-drafting.Execution terminates when every draft reaches the maximum new length, producing n response sequences.

5 Experiments

Experiments evaluate SEED for reasoning-tree construction on creative writing and mathematical reasoning tasks using multiple decoding baselines and model suites. The evaluation uses controlled decoding, search, hardware, and model configurations to assess speedup relative to autoregressive generation.

Datasets: Experiments use three reasoning and planning datasets, including Creative Writing with tree depth 2 and GSM8K with tree depth 4.Creative Writing tests coherent multi-paragraph generation, while GSM8K tests multi-step grade-school mathematical reasoning.
Baselines: The baselines compare vanilla autoregressive generation, speculative sampling, and multi-candidate speculative decoding for reasoning-tree construction.SEED is applied within speculative-decoding approaches to evaluate effectiveness across different acceptance rates.
Results: Table 1 reports speedups relative to vanilla autoregressive generation for SEED and baselines using LLaMA-68M as draft model and LLaMA2-7B as target model.The table evaluates the methods across the three reasoning datasets and marks the best results among methods.
Experimental Setup: The evaluation uses LLaMA Chat and QWen model suites, with BFS search, temperatures of 0.2 and 1.0, and experiments on A100-80G or four RTX 3090 GPUs.Hardware differences between the two platforms are discussed separately.

6 Results and Analysis

SEED consistently improves speedup over baselines across acceptance rates, datasets, and reasoning settings, while its gains differ between Thought Generation and State Evaluation and eventually saturate as reasoning paths increase.

6.1 Main Results: Tree attention generally yields higher speedups because it verifies multiple candidate tokens at each position in parallel.When kconfig elements exceed 1, SEED uses tree attention rather than traditional single sampling.
6.1 Main Results: Up to 1.5× speedup and nearly 30 additional tokens per second are achieved on GSM8K with MCSD inside SEED using tree attention, compared with AR.Across datasets and reasoning depths, SEED consistently outperforms baselines under varied temperatures and kconfig settings.
6.1 Main Results: At the same acceptance rate, SEED outperforms the baseline because scheduled draft-model execution adds speedup without increasing acceptance rates.Both methods show higher speedup as acceptance rate increases, consistent with speculative decoding.
6.2 Analysis: Thought Generation consistently exceeds State Evaluation in acceptance rate and speedup, likely because thought proposals align more closely between target and draft models than decision-making does.Thought Generation uses repeated iterations with the same prompt, whereas State Evaluation evaluates iterations with different prompts.
6.2 Analysis: As the number of reasoning paths increases, target-model idle time decreases and GPU utilization and speedup rise until fixed verification capacity causes saturation at n=6.Excess draft models then wait, reducing parallelism and creating bottlenecks that lower utilization and acceleration.

7 Conclusion and Discussion

SEED accelerates reasoning-tree construction while managing runtime speed and GPU memory through scheduled speculative execution. Experiments across three reasoning settings show improved inference speed, and the framework is presented as a direction for extending test-time scaling.

7 Conclusion and Discussion: SEED optimizes runtime speed and GPU memory usage during reasoning-tree construction by coordinating multiple draft models with one target model through scheduled speculative execution.The approach applies scheduling principles similar to operating-system process management.
7 Conclusion and Discussion: Across three reasoning datasets, SEED improves inference speed by generating up to 20 additional tokens per second.
7 Conclusion and Discussion: SEED offers a potential extension path for advancing test-time scaling of LLM reasoning and improving efficiency in real-world applications.The paper frames this as a promising direction rather than an established outcome.

Limitations

SEED's parallel drafting increases KV-Cache requirements, while the study leaves search-speed optimization and advanced scheduling for future work.

SEED adds n −1 drafting models and an equivalent number of KV-Cache entries, though the authors do not optimize KV-Cache management because the draft models are small.
The study accelerates tree construction but not search, using BFS while leaving integration with A∗ and MCTS for future research.
SEED currently uses FCFS scheduling, with more advanced real-time scheduling algorithms identified as a future direction.
The framework offers a potential implementation of batched speculative decoding that could integrate with other KV-Cache-based batch methods.

A.1 Selection of Baselines

The baseline selection distinguishes SEED from methods requiring training or external knowledge and excludes approaches that are not lossless or are otherwise unsuitable for direct comparison.

The appendix notes that several other speculative decoding strategies were not explored as baselines, and Table 3 compares speedups relative to vanilla AR.
The baseline taxonomy distinguishes training-free methods, extra-knowledge-free methods, and lossless methods as separate comparison criteria.
SEED excludes self-drafting methods that require training because they are not plug-and-play and differ from its independent draft-and-target setup.
Methods using external knowledge modules are considered unfair comparisons because their performance depends on those modules.
SS and Medusa are excluded because they are not lossless, whereas SEED preserves the output distribution of autoregressive decoding.

A.2 Scalability and Extensibility

SEED is evaluated with LLaMA and QWen draft-target pairs, and the QWen results are reported as consistent with the main findings.

The LLaMA evaluation includes LLaMA-160M-Chat as draft model and LLaMA-2-Chat-13B as target model in Table 3.
The QWen1.5 evaluation uses QWen1.5-0.5B-Chat as the draft model and QWen1.5-7B-Chat as the target model.
The QWen results align with the findings in Section 6.1, supporting the framework's extensibility across LLM suites.

A.3 Task Performance

SEED's task evaluations report effectively lossless accuracy and applicability beyond reasoning tasks, with speed depending on hardware and task setting.

The GSM8K performance difference between autoregressive QWen1.5-7B and SEED was within ±1.5%, supporting an effectively lossless outcome.
SEED is presented as applicable to both tree-search-based reasoning and non-reasoning tasks.
Tree-search-based tasks typically generate 2-6 reasoning paths concurrently, a scale described as manageable for consumer-grade GPUs.
1.31x speedup over autoregressive decoding was observed on the WMT dataset in early exploratory experiments.
On 4×3090 hardware, SEED's speedup is lower than on 1×A100, likely because communication between multiple GPUs increases.

A.5 Batch Inference

SeeD targets batch-inference inefficiencies caused by variable acceptance lengths and uses scheduled speculative execution with tree attention and breadth-first reasoning-tree construction. Its setup combines parallel drafting, state evaluation, and task-specific tree configurations across reasoning benchmarks.

Batch inference: SeeD preserves the original KV-Cache length by avoiding padding for varying acceptance rates, with one verified draft sequence mapped directly to each batch sequence.The number of draft models is set equal to the batch size, enabling parallel drafting.
Tree attention: Tree attention concurrently processes multiple tokens in a single sequence while masking tokens outside each token’s permitted attention paths.The root represents previous tokens, and check marks indicate active attention connections.
Reasoning-tree construction: The BFS reasoning-tree procedure initializes the prompt as the root state, iteratively generates and evaluates states up to depth T and breadth b, then selects the highest-valued final state.The thought generator produces candidate thoughts, while the state evaluator assigns scalar or heuristic values.
Experimental setup: SeeD uses temperature-controlled sampling for thought generation, with tree depth and task-specific branching chosen for Creative Writing, GSM8K, and related reasoning tasks.The described configurations include depth 2 with three plans for Creative Writing and depth 4 with three sub-questions for GSM8K.