Research map / Probabilistic and scientific modeling / Molecules and proteins
STree: Speculative Tree Decoding for Hybrid State Space Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes STree, a tree-based speculative decoding algorithm for SSMs and hybrid SSM-Transformer models that packs a token tree into one sequence using an accumulated state-transition matrix, and shows speedups over vanilla speculative decoding on MT-Bench, HumanEval and GSM8K.
Problem
Tree-based speculative decoding is efficient for Transformers but SSMs cannot process a packed token tree, so current SSM speculative decoding unrolls trees into repeated sequences with extra states and memory cost.
Approach
Accumulate diagonal log state-transition matrices along the tree's ancestor paths (A_tree = L · log A) so that a packed tree sequence with a tree mask yields exact per-path SSM outputs in one scan, with a hardware-aware tree scan kernel and activation replay for state backtracking.
Key results
- On MT-Bench at temperature 0, STree reaches 69.84 tokens/s (1.74x over autoregressive) versus 67.68 for vanilla SD, with acceptance length 2.47 vs 2.04.
- At temperature 1 on HumanEval, STree achieves 60.34 tokens/s (1.50x) versus 52.88 (1.32x) for vanilla SD with MambaInLlama-8B.
- For Mamba2-2.7B with a 15-token packed tree, STree forward pass grows only from 27.6 ms to 34.0 ms while unrolled FSS grows from 27.3 ms to 59.8 ms.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace vanilla speculative decoding with STree static-tree verification | MambaInLlama-8B · MT-Bench, HumanEval, and GSM-8K; temperatures 0 and 1 | tokens/sec.: 67.68 (1.69$\times$) -> 69.84 (1.74$\times$); 77.18 (1.93$\times$) -> 78.35 (1.95$\times$); 78.38 (1.93$\times$) -> 80.0 | Helped |
| Use packed tree input instead of unrolled tree input | Mamba2-2.7B speculative decoding · MT-Bench generation; tested tree sizes | tokens/second: 1.05$\times$ to 1.49$\times$ | Helped |
| Increase static tree width substantially (ablation) | STree with MambaInLlama-8B · Temperature 1, tree B versus tree C | tokens/sec.: 51.19 -> 46.90 | Hurt |
| Increase tree depth (ablation) | STree with MambaInLlama-8B · Temperature 1, tree A versus tree E | tokens/sec.: 54.08 -> 56.79 | Helped |
| Use a different connectivity with the same number of tokens (ablation) | STree with MambaInLlama-8B · Temperature 1, tree A versus tree D | tokens/sec.: 54.08 -> 54.39 | Helped |
| Reduce Transformer-block percentage from 50% to 25% (ablation) | STree with Mamba2-Llama3 · Temperature 1 | speed-up per accepted token: 0.66 to 0.64 | Hurt |
| Reduce Transformer-block percentage from 50% to 0% (ablation) | STree with Mamba2-Llama3 · Temperature 1 | speed-up per accepted token: 0.66 to 0.63 | Hurt |
| Train the drafting model for more steps (ablation) | STree with MambaInLlama-8B · End-to-end generation | tokens/seconds: 42.96 (1.19x) -> 58.65 (1.36x) -> 64.30 (1.62x) | Helped |
| Raise sampling temperature (ablation) | STree and Vanilla SD with MambaInLlama · MT-Bench, HumanEval, and GSM-8K | generation speed | Hurt |
| Use naive sampling instead of multi-step speculative sampling (ablation) | STree · MambaInLlama-8B; varying sampling temperature | generation speed | Mixed |
| Use a beam-search tree instead of a static tree (ablation) | STree with MambaInLlama-8B · End-to-end generation | tokens/sec.: 69.84 (1.75x) -> 71.60 (1.79x) | Helped |
| Use STree rather than Vanilla SD on H100 at temperature 0 | MambaInLlama-8B · MT-Bench, HumanEval, and GSM-8K | tokens/sec.: 113.60 -> 113.13; 129.76 -> 127.15; 131.62 -> 131.13 | Mixed |
Surprising findings
- STree's relative overhead over selective scan shrinks as model size grows (2.04x at 2.3B to 1.26x at 23B), though these larger models use random weights for timing.
- On H100 GPUs, STree is slightly slower than vanilla SD at greedy decoding (113.13 vs 113.60 tokens/s on MT-Bench).
Limitations
- STree adds overhead at short input lengths and does not universally improve speed without careful tree and hardware choice.
- Only demonstrated on 8B-scale and smaller models; larger models are not evaluated end-to-end.
- The draft model is an early-stage 2-layer distilled SSM and a static tree, so gains are based on baseline drafting choices.
Experimental setup
Datasets: MT-Bench HumanEval GSM8K
Benchmarks: MT-Bench HumanEval GSM8K
Metrics: tokens per second average acceptance length forward pass latency (ms) GPU memory (GB)
Baselines: Vanilla speculative decoding Autoregressive decoding Fused Selective Scan with unrolled tree Chunk Scan
Builds on: SpecInfer Mamba2 speculative sampling Fused Selective Scan Sequoia
Closest papers by meaning
- RheoSampling: Resolving the One-Hot Dilemma in Stochastic Dynamic-Tree Speculative Decoding (NeurIPS 2026 · similarity 0.91)
- Speculative Speculative Decoding (ICLR 2026 · similarity 0.91)
- SpecBranch: Speculative Decoding via Hybrid Drafting and Rollback-Aware Branch Parallelism (ICLR 2026 · similarity 0.90)
- BASTION: Budget-Aware Speculative Decoding with Tree-structured Block Diffusion Drafting (NeurIPS 2026 · similarity 0.89)
- TALON: Confidence-Aware Speculative Decoding with Adaptive Token Trees (NeurIPS 2026 · similarity 0.89)
- A Theoretical Perspective for Speculative Decoding Algorithm (NeurIPS 2024 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.