Title: STree: Speculative Tree Decoding for Hybrid State Space Models
Venue: NeurIPS 2025
Accepted: yes
Keywords: State Space Models

hybrid models

speculative decoding

tree-based decoding
Abstract: Speculative decoding is a technique to leverage hardware concurrency in order to enable multiple steps of token generation in a single forward pass, thus improving the efficiency of large-scale autoregressive (AR) Transformer models. State-space models (SSMs) are already more efficient than AR Transformers, since their state summarizes all past data with no need to cache or re-process tokens in the sliding window context. However, their state can also comprise thousands of tokens; so, speculative decoding has recently been extended to SSMs. Existing approaches, however, do not leverage the tree-based verification methods, since current SSMs lack the means to compute a token tree efficiently. We propose the first scalable algorithm to perform tree-based speculative decoding in state-space models (SSMs) and hybrid architectures of SSMs and Transformer layers. We exploit the structure of accumulated state transition matrices to facilitate tree-based speculative decoding with minimal overhead relative to current SSM implementations. Along with the algorithm, we describe a hardware-aware implementation that improves naive application of AR Transformer tree-based speculative decoding methods to SSMs. Furthermore, we outperform vanilla speculative decoding with SSMs even with a baseline drafting model and tree structure on three different benchmarks, opening up opportunities for further speed up with SSM and hybrid model inference. Code can be find at: https://github.com/wyc1997/stree.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper extends the tree-based speculative decoding from transformer-based auto-regressive models to the mamba models (SSMs) or hybird architectures.\nAuthors build their method, STree, based on accumulated state transition matrices.\nThey design a hardware-aware TreeScan kernel for the naive application of AR Transformer tree-based speculative decoding methods to SSMs.\nThe proposed method can not only accelerate the inference of SSMs, but also reduce the memory usage during inference.\nThis work outperforms vanilla speculative decoding with SSMs on three different benchmarks.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n\n1. This paper extends the tree-based speculative decoding to mamba models\n\n2. The special kernel TreeScan kernel is a hardware-aware design, which minimizes the memory reduction.\n\n3. Compared to Fused Selective Scan (FSS), STree achieves faster inference and memory reduction.\n\n\nWeakness:\n\n1. Limited model scales, this paper only reported models: Mamba2-2.7B and MambaInLlama-8B.\n\n2. Static tree configurations: the manually setting is required for the tree configuration.\n\n3. The acceleration for the model inference compared to the Vanilla SD is not obvious.\n\nQUESTIONS:\n1. Is there any design to make the configuration of the tree in dynamic ways?\n\n2. As compared to the Vanilla SD method, the work does not achieve obvious acceleration improvements, is there any latency and memory results for other scales like 370M, 790M, 1.4B, or other new mamba models.\n\n3. In table 2, can you provide more results with T=2, 3, 4...\n\n4. In table 2, can you provide more results with more benchmarks? like MMLU, HellaSwag, SWAG, BBH, BoolQ\n\n5. Latency comparison on other powerful GPUs? like A100, H100\n\n6. Ablation study for different draft models?\n\nLIMITATIONS:\n1. The novelty of the work is limited.\n\n2. The tree configuration is static, which limits the model performance\n\n3. This work only tests results with RTX 3090, while the results with powerful GPUs like A100 and H100 remains unknown.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 1, "DATE": "2025-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces State-Space Speculative Tree Decoding (STree), to perform tree-based speculative decoding in state-space models (SSMs) and hybrid SSM/Transformer architectures. By accumulating state-transition matrices according to a draft-generated prefix token tree, STree packs the entire tree into a single SSM forward pass (rather than unrolling it into many sequences), dramatically reducing redundant computation and memory footprint. the contribution is,\n\n1. a concise derivation of the tree-accumulation mechanism,\n\n2. a hardware-aware TreeScan kernel to implement it efficiently,\n\n3. an analysis quantifying the trade-off between tree size and acceptance length, and\n\n4. extensive experiments on pure SSMs and hybrid MambaInLlama models showing up to ~1.5× speed-ups and major memory savings over unrolled baselines.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\n\nS1. They extend tree-based speculative decoding—previously exclusive to Transformers—to SSMs and hybrid models, leveraging the diagonal structure of SSM transition matrices, they provide a step-by-step derivation of the tree accumulation and how it generalizes existing matrix‐SSM formulations.\n\nS2. The evaluation is intensive. Forward‐pass benchmarks (Fig. 3): STree scales gracefully from 15→63 tokens (27.6 ms→34.0 ms) vs. a doubling for unrolled FSS (27.3 ms→59.8 ms). End-to-end generation (Table 2): On MT-Bench, HumanEval, GSM8K, STree delivers up to ~1.98× speed-up vs. autoregressive and ~1.36×–1.50× vs. vanilla speculative decoding, especially at higher temperatures.\n\nWeaknesses\n\nW1. Results report mean speeds and acceptance lengths but lack error bars or multiple runs.\n\nW2. The experiments were only run on an NVIDIA RTX 3090 GPU, and it is unclear how they perform on other hardware—such as the H100 using fast GEMM optimizations. It should report the performance on H100 as well.\n\nQUESTIONS:\nN/A\n\nLIMITATIONS:\nSee weaknesses", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper aims to accelerate state-space models (SSMs) using speculative decoding, a technique commonly employed in transformer-based large language models (LLMs). In speculative decoding, tree verification is often used to further enhance acceleration speed. However, due to fundamental differences between transformer-based LLMs and SSMs, the authors observe that token tree verification cannot be directly applied to SSMs accelerated with speculative decoding. Therefore, the authors propose STree, facilitating speculative tree decoding in SSMs.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\n\n1.\tThe motivation of this work is clear and meaningful. As the authors stated, it is necessary to facilitate speculative tree decoding in SSMs.\n\n2.\tThe proposed method is natural and logically sound. The experimental results provide preliminary validation of its effectiveness.\n\nWeaknesses and questions\n\n1.\tSpeculative tree decoding has been explored for transformer-based LLMs, and speculative decoding has also been applied to SSMs. This work addresses a minor gap by proposing speculative tree decoding for SSMs; however, the challenges involved in this extension are limited, which reduces the novelty of the contribution.\n\n2.\tIn line 26, the statement “AR Transformers leverage a sliding window of input tokens (context) as their ‘state’, which is ideal for speculative decoding since their context can be easily edited to remove unverified tokens” is unclear. Does this mean that AR Transformers retain all past tokens and use the attention mechanism to dynamically focus on them? The authors are encouraged to provide a clearer explanation.\n\n3.\tCode availability: I understand that the code will be released upon acceptance. However, NeurIPS allows and encourages authors to submit associated code files during the review process. This would help reviewers better assess the reproducibility of the work.\n\nQUESTIONS:\nSee above weaknesses and questions.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 2, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-03"}

comments: SUMMARY:
This paper introduces STree, the first scalable algorithm that enables tree-based speculative decoding for SSMs and hybrid SSM-Transformer architectures. Existing speculative SSMs do not leverage tree-based verification due to inefficiencies in handling tree structures. The authors propose a hardware-aware implementation that constructs and utilizes a tree mask like prior works do, but specifically designed to allow efficient verification by accumulating state transition matrices through a novel formulation.

STRENGTHS AND WEAKNESSES:
Strengths:
1. Stree is the first to enable tree-based speculative decoding for SSMs, implying the potential of using SD in hybrid models
2. The paper presents clear derivations showing how diagonal SSM transition matrices can be aggregated for tree decoding.


Weaknesses:
1. Experiments only conducted on 8B models, which may not generalize to much larger LLMs
2. The model Relies on handcrafted or beam-based static trees; does not explore adaptive or learned trees.

QUESTIONS:
1. Could adaptive or learned tree structures (e.g., from a policy model) significantly boost performance?

2. How well would STree perform on larger models like 30B or 70B?

LIMITATIONS:
1. STree has higher overhead for short sequences; benefits are only realized at sufficient tree depth or acceptance rate.
2. Efficiency relies on static tree structures; adapting to diverse prompt types is not explored.
TITLE: Official Review
CLARITY: 3
IMPACT: 3
ORIGINALITY: 2
DATE: 2025-07-03