STree: Speculative Tree Decoding for Hybrid State Space Models
Authors: Yangchao Wu
Zongyue Qin
Alex Wong
Stefano Soatto
Venue: NeurIPS 2025 poster ()
Keywords: State Space Models
hybrid models
speculative decoding
tree-based decoding
Primary area: deep_learning
TL;DR: We enable tree-based decoding on SSMs to facilitate speculative decoding with tree-based verification with SSMs

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.

STree: Speculative Tree Decoding for Hybrid State-Space Models
Yangchao Wu!*
Zongyue Qin'
Alex Wong
Stefano Soatto
UCLA
2 Yale University
Abstract
Speculative decoding is a technique to leverage hardware concurrency in order to enable multiple steps of token generation 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 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 accumulated state transition matrices to facilitate tree- based speculative decoding with minimal overhead relative current SSM implementations. Along with the al algorithm, we describe a hardware-awarr 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 opportunities for further speed up with SSM and hybrid model inference. Code can be find https::/ betts://githut wiggttut.b... www. wwtp://wwc199////// stree
Introduction
Recursive sequence models, such autoregressive (AR) Transformers, produce single token with each forward pass. Speculative decoding is a technique to make this process more efficient by leveraging smaller 'draft model' to generate multiple tokens, and separate 'verifier model' to validate the proposed drafts [13]. The efficiency stems from exploiting the concurrency of parallel computer hardware to verify multiple tokens in single model call, while ensuring that the accepted drafts are identical to those that would have been generated by repeated calls to the original model. AR Transformers leverage sliding window of input tokens (context) their 'state", which is ideal for speculative decoding since their context can be easily edited to remove unverified tokens.
On the other hand, State-Space Models (SSMs) maintain an explicit Markov state that is designed and trained to be a sufficient statistic of the past for the purpose of predicting one-step ahead [17]. They are not naturally suited for speculative decoding, since the past states are discarded and the verifier would need to backtrack each speculated state, hampering the efficiency Recently, hardware-
whole process. efficient methods for speculative decoding in SSMs have been proposed 24 25). but they do not leverage rree-based verification that drives the most efficient methods for
"Correspondence to vuyangchao19976gg ucla.edi
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Prefix token tree:
Unrolled Input:
Packed Input with mask: O repeated tokens sequence
7 tokens
5 repeated tokens 4 sequences
1
2
3
2
Current SSMs
Specinfer, STree
Figure 1: Methods to decode prefix token tree with AR Transformers. A prefix token tree can be unrolled into multiple sequences (used by current SSMs) and computed as a batch or packed into one sequence using mask to indicate the tree structure (first used by SpecInfer [19] and extended to SSMs here). The former leads to inefficiency due to repeatedly computed tokens (in orange).
AR Transformers. This paper proposes the first algorithm do so, with method that is applicable to both SSMs and hybrid architectures interleaving SSMs and Transformer layers.
To perform tree-based verification in AR Transformers 19 14
22), prefix token tree is built with the draft model and verified with the target model. Existing works pack the token tree into one sequence and leverage topology-aware mask [19] with self-attention to compute the output of the tree in one model call (Fig. Right), thus avoiding repeated computation when naively unrolling the tree into individual sequences.
However, modern SSM realizations 66 are designed to scan through all tokens in the past causally to obtain the state, lacking mechanism to specify the tree structure. As consequence, current SSMs can only use the unrolled input (Fig. Left), leading to incfficient repeated token generation. Moreover, SSMs require one state per input sequence in the batch, leading to explosion of the memory footprint when using the unrolled input. As the size of the tree grows, unrolling the tree quickly becomes infeasible. Therefore, we seek ways to leverage the characteristics of the hardware to perform multiple steps of state update following the tree structure through a single pass through the model.
We propose State-Space Speculative Tree (STree) decoding, a method to facilitate tree-based specula- tive decoding in SSMs through accumulating state transition matrices according to the tree structure. The tree is easily computed with minimal overhead relative to current SSM state update implementa- tions. We describe hardware-aware implementation that already in its simplest instantiation shows improvement over the baseline naive application of
speculative decoding to SSMs.
Our contributions in this paper are to (i) propose what, to the best of our knowledge, is the first scalable method to leverage tree decoding in the speculative decoding for both SSMs and hybrid architectures; we also (ii) provide a simplified analysis of the trade-off between acceptance length and model runtime to help determine whether we should scale tree size or even use tree decoding. Finally, we (iii) empirically demonstrate that with baseline drafting model and static tree structure, there are already improvements in generation speed, thus opening the door to further investigation of more advanced speculative decoding methods employed with transformers.
Related Work
State-space models (SSMs) are sequence models that maintain hidden state and use predict
the next datum in the sequence given the previous ones [17]. Such models can be stacked, so
the output of one is the input of another [10, 26), and their parameters can be input-dependent In order to scale model size, transition matrices are often restricted to be diagonal [6] 77),
12 g 27
while to leverage efficient hardware-aware implementations, SSM layers are often interleaved with Transformers in Hybrid architectures [16 8) also useful for initial distillation [24] Our method is

---

based on scalable SSMs of the general form ①T.
a++1 = A(u)) + B(u)) U/ C(u)xx D(u)uu
(I)
Speculative decoding denotes set of techniques designed to parallelize sequential inference in large language models (LLMs) by utilizing a smaller 'draft model' produce multiple candidate trajectories and large 'verifier model' to test them for consistency with the original model [13]. Sub-networks can also be used for drafting, a form of "self-speculation"
T8] while speculation can
be applied to latent features [I5 14] in addition to the output. Structural changes to the speculative processes include rese verification 20 4. multi-head decoding (2) and beam sampling [2] 22
The same methods can also be used for 'lossy verification using judge decoding [
Tree-based verification drafted sequence [15 14
19 has been shown effective in improving the acceptance length of the 2 23 21
22 EAGLE [15] reported an increase in acceptance length by using a static tree structure, leading to an overall speed up ratio improvement. Sequoia [ generates the optimal tree for the speculated tokens using a dynamic programming algorithm and achieved further improvement. The underlying mechanism that enables tree-based verification is Transformers' ability to specify which tokens to attend to with an arbitrary attention mask individually for each token. prefix token tree can be packed into sequence and topology aware mask [9 can be used to inform the self attention block the tree structure.
4
Speculative SSMs have been independently championed by 24 25 Since the state SSMs
updated causally to summarize the past history, extending speculative decoding to SSMs require backtracking the state, which non-trivial scale. Both [24 25 proposed hardware-awarc efficient algorithms to backtrack the state while performing the forward pass. However, neither leveraged tree-based verification, since currently available algorithms for SSM updates require unrolling the tree into individual sequences, which leads to inefficiencies. These include repeated token computation and extra states to be maintained.
Method
Formalization. Given a prefix token tree T with tokens (ti..... (N) as vertices and t as the root
node, we can pack the tokens into one sequence S (l1,...,4x) [N] and represent the topology of the tree using a special attention mask L E (0, [0,1].XXN ('tree' mask). Each vertex in T has one unique path to the root node, and we denote these paths 8; {lalt, is in the path from i to 4). Then the tree mask L can be constructed by the indicator function:
Li.j Le, {ty}.
(2)
Given the Pre-SSM input features u. = .u....., un (u E Rdd } of the packed sequence S and an initial state zo Rdd,, our goal is to compute the output sequence y = {yI.... .N. ly; E Id..) J without repeatedly computing any tokens or requiring extra SSM states. We present our approach below.
Tree decoding with State-space models. The input feature is first mapped onto the parameters using a linear projection. Here we abuse the notation to let X, == X (w+):
A WAlg E Xdd
B, WBu E R4 xd,
C, = Wou E Rd. xd2
(3)
Ignoring the last term D, in Eqn. for simplicity of notation, for each yt = C,Tt, we can expand recursively into:
yr = C.(AA Lea heae- ... ...A
Afri xo + Le,1 Afre
Alaa- ... Afra B + .... B.ve)
= CC(A A
baa-1 ... A.AA A so + Σ La, II
Ahw)),uu
(4)
j=s+1
3

---

As done in 9 6 16 we enforce a diagonal structure on A, to reduce the products in Equ. to sum of logarithms:
V/2 (expp
Σ L expl
E Lis log(A)))
j=st1
= C. expl
T log((:))3a A:)
+
Σ s=1
T expl
E |=++1
Tn log( (ogg(,))]CC C,
(5)
Since A; is diagonal, we can represent the diagonal of log A, as vector and assemble new matrix Alog [diag(log (el)....,di(((uu, An)]² E RRxx,, We define:
Atree == (LAtog) (Atrec)t = Σ The x diag(log A.) E RR
We note that (Atree)s is equivalent to Eiz Le, log(A,) in Equ. S.. since A, is diagonal and we can always organize the tokens in the sequence such that Lrs 0 Vi Furthermore, Ej=sti Le, log(A,) is equivalent to (Atree)t (Atree)s. Then, Eqn. (5) becomes:
Yt C( o xo) +
Σ ==1
Li,e exp( (Atree)e) (C,),
where denotes element-wise multiplication. Then the entire sequence of outputs y can be written as: y STTr
u) M.To + M,
where:
(M,); C. x diag(exp((AAree)))
(M.), = L, x diag(exp( (Atree); -(Arre)))))
We note that our method is generalization of the matrix transformation form of SSM in Mamba2 [6]. When L is lower triangular causal attention mask, our method the same as 66 with non-zero initial state. We introduced the quantity Atree which accumulates the state transition matrices A according to the tree structure, enabling tree decoding with SSM. The computation for Atree is simple, and incorporating it into the original SSM imposes little overhead. We further note that our method is not limited a tree structure, but can potentially be used with an arbitrary mask to specify the structure, allowing greater flexibility with SSMs.
Implementation. Based on the methods presented in previous paragraphs, we propose an algorithm in Alg. [ to perform State-space speculative decoding with tree decoding. The algorithm centers on tree scan kernel that computes the output of a packed prefix tree input and also caches the intermediate activation values (i.e., Augs Bags Cag):
Legz Cache <- TREESCAN(L, 1sij
The kernel is hardware-aware as it avoids instantiating any intermediate results as well as the SSM states off from the fast GPU shared memory. During the verification process, we only generate the output Yio for the input sequence, but not the state after the input. This is because, for speculative decoding, the state at the end of an input sequence most likely incorrect unless all tokens in the input sequence are accepted. We use the activation replay method [ to recompute the correct state before the first rejected tokens at the start of the next iteration.
Analysis
The wall time of speculative decoding is jointly determined by the runtime of the target model and the average acceptance length т. Specifically, we have
t
Wall Time x

---

Algorithm Speculative decoding with Tree Scan for SSMs
1:
function SECCUAATVEDECCOONGVTTTIIT
2:
Initialize L* mask to indicate last accepted token
3: Initialize Cache activation cache to facilitate recomputation of state
4: Initialize x=" : the correct state
5: while should covildd__ottiueeeee do
6:
Lie 14:j 6- (RAFT((;-1)
P Draft tree with last accepted token > Recompute state up to the rejected tokens Getting output and cache from target model Accept/Reject Drafted tokens
7:
x' + -AATIVATIONREPLLYLL Cache) V.yg. Cache +- TREESCAN(L, fej,z²)
8:
9:
17,4 +- FiRSTREJECTEDD (4yy 415)
10:
end while
end function
19-4
Shoo L EOOOATS 2.0OW
32
#55: 0.3TH 15.99
Videe foownd nuntore
8
2.0
12
12
300
115st A00 N
(a)
(b)
(C)
Figure 2: Left: The runtime for a call to Mamba2-2.7B model vs. input size with STree and Fuse
Selective Scan (FSS). A linear regression is performed to obtain the slope and intercept. Middle: The runtime for a call to MambalnL lama-8B model vs. input size with STree and FSS. A polynomial regression with degree is used obtain the parameters. Right: The ratio of acceptance length required to get wall-clock time improvement vs. input length N2, with a fixed N₃
The success of speculative decoding relies on the fact that computing input of length N doesn't increase too much. As SSMs are not as good as transformers in performing parallel computation due their recurrent nature, the runtime of the target model calls sees a non non-negligible increase even at small input length. Therefore, there exists a trade-off between the gain in acceptance length and the increase in runtime of target model calls as the size of the tree grows. Given that SSM compute and memory cost is linear in the input length, we assume simple linear model between the model runtime and the input size, where f kN C, where for loading model weights and &:N is proportional to input size N, contributed by loading inputs and computation. When the models are large, the constant dominates the runtime, leading to negligible effect runtime by input length N when N is small.
a constant contributed by the time
Therefore, to gain wall-clock time improvement we require:
N +
1g + C2
TI
T2
N2 + C2
(7)
K Ni +
We note that when applying the above analysis to hybrid models, one key difference is that the transformer blocks in the hybrid models have quadratic space and time complexity in input length, and therefore would require a quadratic model = a.22 + bN + e for modeling the runtime and input size, which leads to:
T2
N2 by + c2
(8)
T1 01 N
N1 +
To determine whether we should use tree scan for speculative decoding, we want to evaluate whether the gain in acceptance length can outweigh the increase runtime. For example, we measured runtime for Mamba2-2.77 and MambalnI.lama-81 model calls using both STree and Fused Selective
Scan (FSS), which an algorithm optimized for short sequence inference (25), with different input
sizes in Fig. 2a and Fig. 2b We performed linear and quadratic regression to obtain approximation
5

---

Table 1: Latency of Mamba2 forward pass using STree vs. autoregressive forward with selective scan. Note that for 7B, 13B, and 23B models, we initialize the model with random weight to measure the runtime since there is no pretrained model of those sizes.
2.3B
7B
13B
23B
Autoregressive STree Vanilla SD
10.95
22.94
40.69 55.90 (1.37x)
72.41
22.36 (2.04x)
33.79 (1.47x)
91.08 (1.26x)
15.05 (1.37x)
26.46 (1.15x)
45.95 (1.13x)
76.74 (1.06x)
of the parameters (k, C, a, b, c). Assuming that we are comparing STree against vanilla state-space speculative decoding that uses FSS to compute forward pass where we draft I sequence with 5 tokens, we can let N1 and plot the ratio of acceptance length T2 Tı required to achieve wall-clock time improvement for both models in Fig. 2c] We see that for STree with Mamba2-2.788 to gain wall-clock time improvement at a tree size N₂ of 15 tokens, we need the drafted tree to achieve at least 1.5x acceptance rate as compared vanilla speculative decoding. On the other hand, STree with LamballnLllma-24
requires at least 1.1 x acceptance rate with the same tree size, which is much easier to achieve than Mamba2-2.7B. This means that the performance gain that can be achieved with tree decoding with Mamba2InLLama- is much more than that with Mamba2-2.7B This could be due to Mamba2InLlamal
being a larger model with a larger constant overhead for model call.
Models of bigger size Theoretically, the effectiveness of speculative decoding depends the runtime difference between the draft model and the target model as well as the acceptance rate of the drafted tokens. Everything else being equal, using a bigger target model leads bigger runtime difference. Therefore, the effectiivneess of our method on bigger models should hold or improve. In Table [ we provide the results of latency of Mamba2 forward pass using STree vs. autoregressive forward with selective scan with same input tree configuration. Note that for 7B, 13B, and 23B models, we initialize the model with random weight to measure the runtime since there no pretrained model of those sizes. We can see that as the model size grows, the relative overhead of using STree decreases from 2.04x to 1.26x, signaling that as model sizes increase, our method is likely going to be even faster, given that the average acceptance length stays the same. Meanwhile, the gap between the Vanilla SD and STree is closing as model size increases (from 48.9% with 2.3B to 18.9% with 23B), which is another signal that our method is scalable to larger models.
Experimental Results
In this section, we aim to demonstrate the efficiency of STree on speculative decoding. In Sec. 5.1 we compare STree against the unrolled input baseline to show that STree is more efficient at decoding tree with SSMs. Then, in Sec. [5.2] we compare against speculative decoding without tree-based verification to show that STree is able to achieve speed improvement with baseline drafting models and tree construction methods. All experiments are run on an Nvidia RTX 3090 GPU.
5.1 Efficiency of STree against unrolled baseline
Forward pass runtime. Since there no efficient algorithm currently available for computing tree decoding with SSMs, we evaluate STree against the baseline method of unrolling the token tree into different sequences and computing the output of these sequences using the currently available Fused Selective Scan (FSS) (25,, and Chunk Scan (6] kernels. We measure the runtime of a forward pass through a Mamba2-2. 2-B.BB model. For the input token tree, we use full binary trees of 4/5/6 layers deep, which contain 15/31/63 tokens in the tree respectively. The results are shown in Fig. Right.
STree performs on par with FSS at small tree size (4-layer), as FSS is optimized for short sequences. However, as the size of the tree increases, the forward pass with STree increases slightly from 27. 6ms to 34.0ms, while FSS increases drastically from 27.3ms to 59.8ms. This due to unrolling prefix tree introduced repeated tokens and extra states as shown in Fig. Left. This slows down the forward pass for both FSS as well as chunk sean. As STree is able to directly decode packed tree sequence, it demonstrates good scalability as the input tree size grows because it avoids repeated computations.

---

Tree depth
Methods
#. of
Tokens Computed
Sieee
States
Foood solect scan
320
Chunki Scan
Packed Unrolled
15
4
Runtire O000
8
32
Packed Unrolled
31
5
16
90
22.6 27.5
Packed
63
6
Unrolled
32
192
C.byerra tall Liruty Lrdu
Bayerhul Iinary looe
Jayer hii Linwry looz
Figure 3: Left: Comparison of number of states required and number of tokens computed to get an output for full binary tree in forward pass. STree is able to decode packed tree sequence, while other methods need to unroll the tree into multiple sequences. Right: Runtime in milliseconds (ms) for forward pass for different full binary trees using different algorithms.
Table 2:
Generation speed (tokens/second) and memory usage (GB) using STree on packed tree input and Fused Selective Scan (FSS) unrolled tree input. M is the number of beams we keep at each step and is the number of steps we heam-searched for. M=4/5 with N=16 for FSS ran Out Of Memory (OOM) 3090 GPU with 24GB of GPU memory.
M=2
M=3
M=4
M-5
STree
PSS
Slieee
PSS
102.69 (1.05x)
97.44
101.68
(1.12x)
90.64
94.61 (1.10x).
85.33
91.80 (1.16x)
78.94
Speed
8 107.28 (1.09x)
97.53
104.99 (1.21x)
86.68
100.67 (1.31))
76.72
98.75 (1.4x))
68.44
16 105.62 (1.25x)
84,26
101.97 (1.49x)
68.14
94.47
OOM
88.41
00M
4 7.67 (0.91x)
8.47
7.68 (0.85
8.97
7.73 (0.83x)
9.36
7.75 (0.78x)
9.85
Memory
7.79 (0.86x)
9.05
7,85 (0.80x)
7.88 (0.72x)
10.89
7.91 (0.64x)
(0.75x)
(0.57xx
Meanwhile, forward pass with chunk scan is consistently slower than both fused selective scan and STree, as it is optimized for parallel training for long sequences and pays a big overhead at a short sequence length 25 24 Hence, it should not be considered for use in speculative decoding.
Generation speed against unrolled baseline. Taking one step further, we measure the generation speed of STree with packed tree input against FSS with unrolled tree input. We used Mamba2-2.7B model as the target model and a Mamba2-1 30M model as the drafting model. We generate the token tree using beam search the tokens generated at each step of beam search, even if the beam is later discarded. This results in an N-layer tree with M tokens each layer, where the number of beams. We verify the target model with greedy search, where at each step, the token with the maximum conditional likelihood from the target model is compared to the corresponding child nodes in the token tree to see if we should accept the draft. The acceptance length for the two methods is the same, as the same drafted tree used for verification. We perform this experiment on the MT Bench 28 benchmarks for generating 100 tokens. The results are presented in Tab. 2
23 where we perform beam search with the drafting model and keep all
the number beam search steps and M is
We can see that STree with packed input is both faster and more memory efficient than FSS with unrolled input across all tree sizes. We also notice that as the size of the tree grows, the advantage with STree becomes bigger (1.05x 1.49x for speed and 0.91 0.57x for memory), which coincides with our previous finding that STree forward pass is more scalable to larger trees. Meanwhile, we note that unrolling the tree also adds overhead the execution. Memory-wise, the increase memory for STree is almost negligible when tree size increases, mainly due to the increase in input size. On the contrary, the repeated tokens and states required by FSS contribute to a large increase GPU memory, making large token tree (N 16, M 4/5) infeasible to compute.
7

---

Table 3: Generation speed (tokens/sec.) and average number of tokens accepted T for Vanilla speculative decoding (SD) and STreewith MambalnLlama 8B model with 50% mix of transformers.
MT-Bench
HvaanEval
GSM-8K
Methods
Speed
Speed
Speed
Temperature=0
Autoregressive
39.98
39.98
40.47
Vanilla SD
67.68 (1.69x)
2,04
77.18 (193x)
2.32
78.38 (1.93x)
2.34
STree
69.84 (1.74x)
2.47
78.35 (1.95x)
2.79
80.03 (1.98x)
2.83
Temperature=l
Autoregressive
39.72
40.04
40.16
Vanilla SD
49.24 (1.23
1.55
52.88 (1.32x)
1.65
51.80 (1.28x)
1.61
STree
54.08 (1.36x)
2.03
60.34 (1.50x)
2.26
58.25 (1.45x)
2.16
5.2 End-to-end generation with hybrid models
Generation with MambaInL models. Having verified the effectiveness of STree in per- forming tree decoding, we then demonstrate the potential of STree in further boosting the speed of generation for SSMs and hybrid models. We choose a Hamba2lnnllama-3 50% mix of transformers and SSMs, as the target model. As there is a lack of smaller models in the same family, we distill 2-layer SSM from the target model using data that is used finetune the target model. We show the results of using a checkpoint from 48000 steps into the distillation. Results from different checkpoints are shown in the ablation study below.
24 hybrid model, with
For the vanilla speculative decoding baseline, we use the 2-layer model as the drafting model to draft 1 sequence of 4 tokens every step (target imput size 1 5) and verify with the target model output using speculative sampling algorithm !? J. For STree, we use the draft model to draft a static tree structure shown in Fig. 4a and use multi-step speculative sampling (MSS sampling) [19] verify with the target model output. Both methods use activation replay to backtrack the state. We evaluate the speed of generating 1024 tokens three different benchmarks: MT Bench (28). HumanEval (3) and GSM8K 5) with two different temperatures: 0 (greedy) and 1. The results are shown in Tab. 3
We can see that across all three benchmarks used, STree is able to outperform vanilla speculative decoding. The advantage becomes more obvious when the temperature of the generation increases and the acceptance length drops result. We note that this performance improvement is achieved with a baseline tree generation strategy (a static tree) and draft model still early in its distillation process. With more advanced tree generation strategy and draft model, the acceptance length with STree will likely improve, which will translate into more generation speed improvement. We include the results with H100 GPUs and beam search tree in the appendix due to space constraint.
Ablation studies
Effect of temperature We study the effect of sampling temperature on the acceptance rate and speed-up. We use the same setup as the end-to-end generation experiment with MambaInLlama model and vary the temperature used with sampling. From Fig. 4b and Fig. 4c we can that average acceptance length drops as temperature goes up for both STree and vanilla speculative decoding, leading decrease in generation speed. We notice that the absolute difference between the acceptance length of STree and vanilla speculative decoding is relatively stable (~~ 0.4), which means that the acceptance length drops for both STree and vanilla speculative decoding, the improvement in generation speed becomes bigger, re-affirming our results in the previous sections.
Effect of sampling algorithm Besides multi-step speculative sampling, we test our method with naive sampling used in [3]. Specifically, we use the top-k tokens from the drafting distribution to build our token tree, and perform sampling with the target model. If at any token, the sampled token from the target model falls within the top-k tokens drafted, accept that token and continue to verify
"Checkpoint https//tuggiiggfiee.c//umiooio Wang//aaaa3,, - -Mamba2-8B 8B-2-Btti

---

Genervtton Sronedt concersoo
Aveeeay acceddanc
2.0
Stree MSS samrpingg Show nais sampling
sampling
0.5
0.
0,4
Temperoture
Temperature
(b)
(c)
Structure of static tree used to generate prefix token tree with the drafting model. We draft 4 steps for each iteration and tokens for each layer, resulting in 13 tokens in every input sequence. Middle: Generation speed of STree and Vanilla Specultaive Decoding (SD) under different temperature. Right: Average acceptance length of STree and Vanilla Speculative Decoding (SD) under different temperature.
Figure 4: Left:
Table 4: Inference speed (tokens/sec.) and average number of tokens accepted 7 for STree with MambalnL.lama models with different static tree configurations (sampling temperature 1).
Static tree configuration
A
B
C
D
E
Max tree width
3
16
4
Tree depth
4
4
Number of tokens
13
29
16
13
16
Speed
54.08
46.90
51.19
54.39
56.79
2.03
2.15
2.01
2.07
2.09
the next token. We refer the readers to (23] for more details. The results are shown together in Fig. and Fig. 4c We can see that low temperatures, naive sampling has relatively high acceptance length and fast run time. As temperature increases, the acceptance length drops significantly, because the sampled token by the target model is less likely to come from the top-k tokens with the drafted model. This leads to a drop in overall generation speed as well. We note that this drop in generation speed is due to the characteristic of the sampling algorithm. but not our framework. With a high acceptance length at lower temperatures, we can still achieve speed improvements.
Effect of static tree structure We study the effect of different static structures the acceptance rate and speed-up. We note that producing a wider and deeper tree also slows the drafting model and contributes the overall runtime. The static tree configurations are shown in Fig. S in the Appendix. The results are presented in Tab. 4 To determine whether to use deeper/widd wider tree, we need to consider only the effect on the acceptance rate, but also the runtime of the target and drafting models. Increasing the width of the tree drastically does improve the acceptance rate (Column B vs. C), but it also slows down the models and therefore overall speed. When we increase the tree depth (Column A vs. E), we can see that the improvement in outweighs the runtime cost, and therefore we see an improvement. Comparing Column and D, where the number of tokens is the same with different connectivity, improvement can be seen in both speed and 7 when better tree structure used. This signifies the importance having good tree structure in using STree. This result also agrees with our analysis on the trade-off between runtime and acceptance rate. It provides the promise that when we use better-trained drafting model and better drafting algorithms, which would give us a better tree and token candidates, the generation speed with STree will improve.
Effect of percentage of transformer blocks target model Hybrid models are mix of trans- former blocks and SSM blocks. As Transformers have better scaling in parallel compute, we hypothesize that more transformer blocks the model will lead a larger improvement in generation speed by STree. We perform an ablation study using Mamba2-Llamas models (4] with 50%. 25%, and 0% transformer blocks, using the same drafting model we distilled. The results are shown in Tab. 5) With autoregressive generation, models with more SSM blocks are faster, which agrees with our expectation. For generation with STree, we are still able to outperform vanilla speculative decod- ing using models with fewer transformer blocks. If we eliminate the effect of different acceptance

---

Table 5: Inference speed (tokens/sec.) and average number of tokens accepted T for Vanilla speculative decoding (SD) and STreewith Kamba2-Llamai models with sampling temperature of 1. The percentage number indicated with the model name the percentage of transformer blocks in the model. Speed-up is obtained by dividing ratio of speed-up against auto-regressive speed in the bracket by
Mamba2-I Jama3 5)%
Mamba2-I (25%)
Mamba2- Bamba2-I.ama3 (0%)
Methods
Speed
Spees-up
Speed
Speod-up
Speed
Speed-up
Autoregressive
39.71
1
41.65
1
43.79
Vanilla SD
46.56
(L.17))
1.47
0.80
47.69
(1.14x)
1.45
0.78
53.72 (1.22x)
1.57
0.77
STree
49.04
(1.23))
1.84
0.66
48.08
(1.15x)
1.77
0.64
55.32 (1.26x)
2,00
Table End-to-end generation speed Tokens/seconds) and acceptance length for Vanilla Specula- tive decoding and STree using drafting model trained for different steps.
12000 step
48000 step
264000 step
Speed (Vanilla SD)
44.52 (1.21x)
53.11 (1.24x)
56.49 (1.42x)
(Vanilla SD)
1.34 42.96 (1.19x)
1.56
1.68
Speed (STree)
58.65 (1.36x)
64.30 (1.62x)
(STree)
1.51
2.02
2.23
lengths in the Speed-up column, acceptance length (0.66 to 0.63), agreeing with our hypothesis.
that models with fewer transformers less speed-up per
Effect of different drafting model checkpoints In Table we provide ablation study drafting model distilled to different steps. We can see that as the distillation goes on, the acceptance length for the draft model gets longer. As the acceptance length increases, the end-to-end generation speed of STree improves from being 1% slower than Vanilla SD at step 12000 to 9.67% faster at step 48000 and further 14.8% speedup at step 264000. This is because the distribution of the drafting model is more aligned with the target model, suggesting that with better drafting model and technique that gives better acceptance length, the end-to-end runtime of STree will become better.
Discussion
Our work proposes STree, an efficient algorithm to unlock the potential of speculative tree decoding for SSMs and hybrid models. We close by discussing the limitations STree. As presented in our analysis, STree has an overhead at a short input length compared to previous methods, and will not universally improve generation speed when applied without careful consideration. Reducing this overhead through further algorithmic innovation could lead better trade-off. Mcanwhile, we only demonstrate improvement on 8B models, which are still considered small LLMs. scaling up our method to bigger models and more advanced speculative decoding methods may bring more benefits, including energy savings for LLM inference. We leave that for the exploration future work.
Acknowledgements, This work was supported by ONR award 02000012221222 and NSF 21 125522 Athena Al Institute.
10

---

References
[1] Gregor Bachmann. Sotiris Anagnostidis, Albert Pumarola, Markos Georgopoulos, Artsiom Sanakoyeu. Yuming Du, Edgar Schönfeld, Ali Thabet, and Jonas Kohler. Judge decoding: Faster speculative sampling requires going beyond model alignment. arXiv preprint arXiv:2501. 19309, 2025.
[2] Tianle Cai, Yuhong Li. Zhengyang Geng, Hongwu Peng, Jason D Lee, Deming Chen, and Tri Dao. Medusa: Simple IIm inference acceleration framework with multiple decoding heads. arXiv preprint arXiv:2401. 10774, 2024.
[3] Mark Chen al. Evaluating Large Language Models Trained on Code, July 2021. arXiv.2107.03374 [es].
[4] Zhuoming Chen, Avner May, Ruslan Svirschevski, Yuhsun Huang, Max Ryabinin, Zhihao Jia, and Beidi Chen. SEQUOIA: Scalable and robust speculative decoding. In Advances Neural Information Processing Systems (NeurIPS), 2024.
[5] Karl Cobbe. Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training Verifiers to Solve Math Word Problems, November 2021. arXiv:21 1:::21100,, 10.111 14168 [cs].
[6] Tri Dao and Albert Gu. Transformers are SSMs: Generalized models and efficient algorithms through structured state space duality. In Proceedings of the s International Conference on Machine Learning. pages 10041-10071, 2024.
[7] Soham De, Samuel L Smith, Anushan Fernando, Aleksandar Botev, George Cristian-Murard. Albert Gu, Ruba Haroun, Leonard Berrada, Yutian Chen, Srivatsan Srinivasan, et al. Griffin: Mixing galed linear recurrences with local attention for efficient language models. arXiv preprint arXiv:2402. 19427, 2024.
[8] Paolo Glorioso, Quentin Authony, Yury Tokpanov. James Whittington, Jonathan Pilault, Adam Ibrahim, and Beren Millidge. Zamba: A compact 7b ssm hybrid model. arXiv preprint arXiv:2405.167) 2024.
[9] Albert Gu and Tri Dao. Mamba: Lincar-time sequence modeling with selective state spaces, arXiv preprint 2rXiv:2312.00052: 2023.
[10] Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. arXiv preprint arXiv:2111.00396, 2021.
[11] Wonseok Jeon, Mukul Gagrani, Raghavy Goel, Junyoung Park, Mingu Lee, and Christopher Lou. Recursive speculative decoding: Accelerating I1m inference via sampling without replacement. arXiv preprint arXiv:2402. 14160, 2024.
[12] Arthur J Krener. Bilinear and nonlinear realizations of input-output maps. SIAM Journal on Control. ((4)):877-34, 1975.
[13] Yaniv Leviathan, Matan Kalman, and Yossi Matias. Fast inference from transformers via speculative decoding. In International Conference on Machine Learning, pages 19274-19286. PMLR, 2023.
[14] Yuhui Li, Fangyun Wei, Chao Zhang, and Hongyang Zhang. Fagle-2: Faster inference of language models with dynamic draft trees. In Proceedings the 2024 Conference on Empirical Methods in Natural Language Processing (EMNLP), pages 7421-7432. Association for Computational Linguistics, 2024.
[15] Yuhui Li, Fangyun Wei, Chao Zhang, and Hongyang Zhang. Eagle: Speculative sampling requires rethinking feature uncertainty, arXiv preprint arXiv:2401.15077 2024.
[16] Opher Lieber, Barak Lenz, Hofit Bata, Gal Cohen, Jhonathan Osin, Itay Dalmedigos, Erez Safahi, Shaked Mcirom, Yonatan Belinkov, Shai Shwaar-Shwarzz, et al. Jamba: A hybrid transformer-mamba language model. arXiv preprint ar0iv:2403. 19887, 2024.
[17] Anders Lindquist and Giorgio Picci. On the stochastic realization problem. SIAM Journal on Control and Optimization, 17())365-389, 1979.
[18] Fangcheng Liu, Yehui Tang. Zhenhua Liu. Yunsheng Ni, Duyu Tang, Kai Han, and Yunhe Wang. Kangaroo: Lossless self-speculative decoding for accelerating lmms via double early exiting. Advances Neural Information Processing Systems (NeurIPS), 2024.
[19] Xupeng Miao, Gabriele Oliaro, Zhihao Zhang, Xinhao Cheng, Zeyu Wang. Rae Ying Yee Wong, Zhuoming Chen, Daiyaan Arfeen, Reyna Abhyankar, and Zhihao Jia. Specinfer: Accelerating generative 1lm serving with speculative inference and token tree verification. arXiv preprint 097788, 1(2)4, 2023.
11

---

[20] Xupeng Miao, Gabriele Oliaro, Zhihao Zhang, Xinhao Cheng, Zeyu Wang, Zhengxin Zhang, Rae Ying Yee Wong, Alan Zhu, Lijie Yang. Xiaoxiang Shi, Chunan Shi, Zhuoming Chen, Daiyaan Arfeen, Reyna Abhyankar, and Zhihao Jia. Specinfer: Accelerating generative large language model serving with tree- based speculative inference and verification. Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems (ASPLOS), 2024.
[21] Zongyue Qin, Zifan He. Neha Prakriya, Jason Cong, and Yizhou Sun. Dynamic- width speculative beam decoding for Ilm inference. In Proceedings of the 39th AAAI Conference on Artificial Intelligence (AAAI), 2025.
[22] Zongyue Qin. Ziniu Hu, Zifan He, Neha Prakriya, Jason Cong, and Yizhou Sun. Optimized multi- token joint decoding with auxiliary model for IIm inference. In International Conference on Learning Representations (ICI.), 2025.
[23] Ruslan Svirschevski, Avner May, Zhuoming Chen, Beidi Chen, Zhihao Jia, and Max Ryabinin. Specexee: Massively parallel speculative decoding for interactive Ilm inference on consumer devices. Advances Neural Information Processing Systems, 37:16342-16368, 2024.
[24] Junxiong Wang, Daniele Paliotta, Avner May, Alexander M Rush, and Tri Dao. The mamba in the Ilama: Distilling and accelcrating hybrid models. arXiv preprint arXiv:2408. 15237, 2024.
[25] Yangehao Wu, Yonatan Dukler, Matthew Trager, Alessandro Achille, Wei Xia, and Stefano Soatto. Snakes and ladders: Accelerating ssm inference with speculative decoding. In NeurlPS Efficient Natural Language and Speech Processing Workshop, pages 292-304. PMLR, 2024.
[26] Luca Zancato, Alessandro Achille, Giovanni Paolini, Alessandro Chiuso, and Stefano Soatto. Stacked residuals of dynamic layers for time series anomaly detection. arXiv preprint arXiv:2202. 12457, 2022.
[27] Luca Zancato, Arjun Seshadri, Yonatan Dukler, Aditya Golatkar, Yantao Shen, Benjamin Bowman, Matthew Trager, Alessandro Achille, and Stefano Soatto. B'MOJO: Hybrid state space realizations of foundation models with cidetic and fading memory. Proc. of NewrIPS, 2024.
[28] Lianmin Zheng, Wei-Lin Chiang, Ying Sheng, Siyuan Zhuang, Zhanghao Wu, Yonghao Zhuang. Zi Lin, Zhuohan Li, Dacheng Li, Eric P. Xing, Hao Zhang, Joseph E. Gonzalez, and Ion Stoica. Judging LLM-as- a-Judge with MT-Bench and Chatbot Arena, December 2023. arXiv:2306. 05665 [cs].
12

---

NeurIPS Paper Checklist
The checklist is designed encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No]. or [NA].
[NA] means either that the question Not Applicable for that particular paper or the relevant information is Not Available.
Please provide short (1-2 sentence) justification right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include (after eventual revisions) with the final version of your paper, and its final version will published with the paper.
The reviewers of your paper will be asked use the checklist as one of the factors their evaluation. While "[Yes] is generally preferable to "[No]" ", it is perfectly acceptable to answer "[No] providedd proper justification is given (e.g.. "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering binary way, we
"[No] or "[NA]" is not grounds for rejection. While the questions are phrased acknowledge that the true answer is often more nuanced, please just use your best judgment and write justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can be found.
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS Paper Checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify the questions and only use the provided macros for your answers.
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Claims in the abtract and paper reflects what is done in the paper.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize other settings.
is fine include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Limitation is discussed in the Last discussion section.
13

---

Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: The methods section contains our algorithms.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: Experiment results sections contains these details. Code will also be released after paper acceptance.
Guidelines:
The answer NA means that the paper does not include experiments.
14

---

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. w
For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification:
Code will be open sourced after paper acceptance.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (nttps://nips.co cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
/ 7/nips. /ttppp.cc/pbllc/guudeeee http://oeeubiiissioooooooooo
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
15

---

Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes]
Justification: The settings detailed in the experiment results section.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification: [TODO]
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results accompanied by enror bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK to report 1-sigma error bars, but one should state The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Detailed in experiment results section.
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
16

---

The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics httpge//eerrii..ccpbbb////eiiiie
Answer: [Yes]
Justification:
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g.. if there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: It is briefly discussed in the last discussion section.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, scraped datasets)?
Answer: [NA]
Justification: No dataset or new models being proposed.
Guidelines:
The answer NA means that the paper poses such risks.
17

---

Released models that have high risk for misuse or dual use should be released with necessary safeguards to allow for controlled use of the model. for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification:
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re- re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: Code will be open sourced after paper acceptance
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
18

---

Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods this research? Note that if the L.M used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, or originality the research, declaration is not required.
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs any important, original, non-standard components.
Please refer to our LLM policy http::///eerrpp.cccco//feeeeeeee for what should or should not be described.
19

---

A Static Tree configurations
15 takens
tokers
16 tokens
D: 13 tokers
F: tobens
Figure 5: Static tree structure that we used in our ablation study for effect of differnt tree structure.
Fig. 5 show the configurations of the tree we tried in our ablation study. We attempted different breadth (B C), different depth (A vs. E). same number of tokens but different connectivity (A vs. D). The results are show in the
B Experiments on H100 GPU
We further extend our method to H100 GPUs to demonstrate the our method is still applicable with better hardware. We apply our kernels as is on H100 GPUs and compare it to vanilla speculative decoding.
Table 7: Generation speed (tokens/sec.) and average number of tokens accepted for Vanilla speculative decoding (SD) and STree with MambaInI Jama-88 8B model with 50% mix of transformers on H100 GPU.
MT-Bench
HumanEval
GSM-8K
Methods
Speed
Speed
Speed
Temperature=0)
Autoregressive
78.09
78.44
1
79.33
Vanilla SD STree
113.60 (1.45x)
2.03
129.76
(1,65x)
2,33
131.62 (1.66x)
2.33
113.13 (1.45x)
2.45
127.15 (1.62x)
2.77
131.13
(1.6x)
2.81
Temperature=l
Autoregressive
76,66
77.67
1
78.35
Vanilla SD
80.76 (1.05x)
1.56
85.63 (1.10x)
1.64
87.04
(1.11))
1.63
STree
85.18 (1.11x)
2.04
(1.19))
2.22
(1.19x)
From the above tables, STree is smaller as compared to autoregressive decoding. This is because the memory bandwidth of H100 GPU bigger than that of RTX3090, leading smaller bottleneck memory transfer. With greedy generation, STree achieves a slightly slower speed as compared to Vanilla SD. When temperature=1, STree still holds an advantage. This is due to the relative increase acceptance length is smaller when we are doing greedy generation and the relative overhead between the STree and Vanilla Speculative decoding increases due to the faster GPU. We believe that as the model size increases and the memory bandwidth bottleneck due to the model weight transfer becomes more significant, we will still able to show improvement. We also note that we present general algorithm that is not GPU specific and is not optimized for H100 GPU.
that the extent of improvement from applying both Vanilla SD and
C End-to-end generation result using beam search tree
We note that STree can work with any tree generation strategy and optimizing the tree structure orthogonal to our work. Here, we further show that the end-to-end generation results using a beam search tree for drafting, which is basic dynamic tree structure. The expansion of the tree depends on the joint probability and will vary from sample sample. We use greedy decoding with M 3 and N
20

---

Table 8: End-to-end generation speed and acceptance length (r) for different method of drafting.
Autoregressive
Vanilla SD
STree static tree
STree beam search tree
Speed
39.98
67.68 (1.69x)
69.84 (1.75x)
71.60 (1.79x)
2.04
2.47
2.60
4 (M, N explained in table 1). The results are shown in Table 8 We can see that with dynamic tree generation strategy, we are able to get better acceptance length and thus faster end-to-end generation speed. We believe that more advance tree construction technique would bring even more benefits.
21


Reviews:
[{"review_id": "QEp1mbp0xQ", "reviewer": "Reviewer_wh4v", "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.", "questions": "1. Could adaptive or learned tree structures (e.g., from a policy model) significantly boost performance?\n\n2. 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.\n2. Efficiency relies on static tree structures; adapting to diverse prompt types is not explored.", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "Strengths:\n1. Stree is the first to enable tree-based speculative decoding for SSMs, implying the potential of using SD in hybrid models\n2. The paper presents clear derivations showing how diagonal SSM transition matrices can be aggregated for tree decoding.\n\n\nWeaknesses:\n1. Experiments only conducted on 8B models, which may not generalize to much larger LLMs\n2. The model Relies on handcrafted or beam-based static trees; does not explore adaptive or learned trees.", "was_revised": true, "final_justification": "I choose to maintain my current score.", "rebuttal": "We thank the reviewer for recognizing the novelty in tree-base speculative decoding for SSMs. We address the concerns raised by the reviewer below:\n\n**W1: Larger LLMs**: It is challenging to find open-source, pretrained hybrid models suitable for speculative decoding that are bigger than 30B parameters. We explored several options:\n  - Nemotron-H-47B: A hybrid model released by Nvidia in April. The smaller Nemotron-H-8B model can potentially act as the draft model. However, the model does not support generation with KV cache/SSM states, making it impossible for us to use.\n  - Jamba-mini-50B: A hybrid model from AI21. We cannot find any suitable draft model that is small in size and uses the same tokenizer as the target model. \n\nIf the reviewer has any specific open-source models that are suitable we would be happy to run experiments with them. \n\nAbsent end-to-end results on larger models, we analyze the runtime (in ms) of a forward pass through Mamba2 models using different methods below. \n |    | 2.7B | 7B | 13B | 23B |\n |:---|:-------|:--------|:-------|:--------|\n | Autoregressive |  10.95 |  22.94 | 40.69 | 72.41 |\n | STree |   22.36 (2.04x) | 33.79 (1.47x) | 55.90 (1.37x) | 91.08 (1.26x) |\n | Vanilla SD |  15.05 (1.37x) | 26.46 (1.15x) | 45.95 (1.13x) | 76.74 (1.06x) |\n\nTheoretically, the effectiveness of speculative decoding is determined by the runtime overhead of forward pass using different methods as well as the acceptance length. Both STree and Vanilla SD are paying an overhead (having a longer runtime than autoregressive forward pass) to compute multiple tokens. From the table below, we can see that the overhead of an STree forward pass is decreasing from 2.04x to 1.26x as we scale the model from 2.7B to 23B. The ratio between STree and Vanilla SD is also shrinking (from 1.48 to 1.19) as model size increases. This means that, with a larger model, it is much easier to achieve an end-to-end performance gain.\n\nGiven that we already achieved end-to-end performance gain in 8B models as shown in Table 2 in the paper, there is a reason to believe that our performance with bigger models would be even better, assuming that the acceptance length is about the same.   \n\n**W2: Dynamic tree**: Our method does not rely on static trees and is applicable to any tree generation strategy. We note that optimizing the tree structure is orthogonal to our work. Our work focuses on enabling efficient computation of token tree with SSMs. With a baseline static tree, we are already able to obtain runtime improvement, as shown in Table 2 in the paper. We also note that the beam search tree we used in Table 1 is not static, as the expansion of the tree depends on the joint probability and will vary from sample to sample. Here, we additionally show the end-to-end result for drafting with a beam search tree. We use greedy decoding with M = 3 and N = 4 (M, N explained in table 1). The results are as follow: \n\n |    | Autoregressive | Vanilla SD | STree w. static tree | STree w. beam search tree|\n |:---|:-------|:--------|:-------|:--------|\n | Speed |  39.98 | 67.68 (1.69x) | 69.84 (1.75x) | 71.60 (1.79x) |\n | $\\tau$ |   1 | 2.04 | 2.47 | 2.60 |\n\nWe can see that with a beam-search tree, we are able to get a better acceptance length and thus a faster end-to-end generation speed. We believe that more advance tree construction technique would bring even more benefits."}, {"review_id": "q5dZxeo8WI", "reviewer": "Reviewer_7Ar2", "summary": "This 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.", "questions": "See above weaknesses and questions.", "limitations": "yes", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths\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.", "was_revised": true, "final_justification": "I have read the reply and will keep my current score.", "rebuttal": "We thank the reviewer for recognizing the motivation and methods of our work. We hereby address the concerns and questions raised by the reviewer:\n\n**W1: reduced novelty**: To the best of our knowledge, we are the first to apply speculative tree decoding to SSMs and hybrid models, which makes our work novel. While tree based speculative decoding exists for transformer based models, it heavily relies on the transformer's property that computing a tree of tokens doesn't incur much overhead with the use of attention masks. Without this property, it is much slower to compute a tree of tokens in SSMs due to their recurrent state update. Our work addresses this critical challenge by improving the efficiency of computing a tree of tokens with SSMs, which is not done in previous works.\n\n**W2: AR transformer clarification**: Yes, the reviewer is right that AR transformers retain the key and value of past tokens in memory (a structure called KV cache) to avoid recomputing them again for future generation. With speculative decoding, certain tokens computed might not be verified, which then needs to be removed from the KV cache to avoid contaminating future predictions. This can be done easily with KV cache, since it is partitioned on token boundaries and removing unverified tokens is equivalent to zeroing out the corresponding KV cache entries.\n\n**W3: Code availability**: We thank the reviewer for their interest in our work. Here we provide code segments where we apply the tree mask in the kernel to compute SSM output. We will publish the complete code base upon paper acceptance.\n\n```\n    ...\n    # We don't need to iterate K since tree scan sequence length is probably short. \n    for k in range(0, K_MAX, BLOCK_SIZE_K):\n        cb = tl.load(cb_ptrs, mask=(offs_m[:, None] < chunk_size) & (offs_k[None, :] < chunk_size - k), other=0.0).to(tl.float32)\n        dA_cs_k = tl.load(dA_cumsum_ptrs, mask=offs_k < chunk_size - k, other=0.0).to(tl.float32)\n        # If there's seq_idx, we already set cb[i, j] = 0 for seq_idx[i] != seq_idx[j].\n        # So we don't need masking wrt seq_idx here.\n        cb *= tl.exp((dA_cs_m[:, None] - dA_cs_k[None, :]))\n        dt_k = tl.load(dt_ptrs, mask=offs_k < chunk_size - k, other=0.0).to(tl.float32)\n        cb *= dt_k\n\n        # apply tree masking here\n        tree_mask = tl.load(tree_mask_ptrs, mask=(offs_m[:, None] < chunk_size) & (offs_k[None, :] < chunk_size - k), other=0.0).to(tl.int1)\n        cb = tl.where(tree_mask, cb, 0)\n\n        cb = cb.to(x_ptr.dtype.element_ty)\n        x = tl.load(x_ptrs, mask=(offs_k[:, None] < chunk_size_limit - k) & (offs_n[None, :] < hdim), other=0.0)\n        acc += tl.dot(cb, x)\n        cb_ptrs += BLOCK_SIZE_K * stride_cb_csize_k\n        x_ptrs += BLOCK_SIZE_K * stride_x_seqlen\n        dt_ptrs += BLOCK_SIZE_K * stride_dt_csize\n        dA_cumsum_ptrs += BLOCK_SIZE_K * stride_dA_cs_csize\n        tree_mask_ptrs += BLOCK_SIZE_K * stride_tm_csize2\n        ...\n```"}, {"review_id": "5bRuJauAIG", "reviewer": "Reviewer_bAPx", "summary": "This 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.", "questions": "N/A", "limitations": "See weaknesses", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "Strengths\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.", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for recognizing the presentation of method and evaluation. We address the concerns raised by the reviewer below:\n\n**W1: Multiple Runs**: Thank you for you suggestion. Results in table 1 and table 2 (temperature=0) is performed with greedy decoding so re-running the experiment for multiple runs will yield the same results modulo hardware fluctuations. For Table 2 (temperature=1), we additionally provide the average and standard deviation for 4 runs here. We will include that in the updated version of the paper. \n\n | | MT_Bench |  | HumanEval|  | GSM8K | |\n |:---|:---|:---|:---|:---|:---|:---|\n |                 | Speed | $\\tau$ | Speed|   $\\tau$ | Speed | $\\tau$ |\n | Autoregressive  |39.72         |   1 | 39.83        |   1  |   40.16    |  1   |\n | Vanilla SD      | 51.88 $\\pm$ 1.92| 1.56 $\\pm$ 0.01| 54.66 $\\pm$ 1.24| 1.65 $\\pm$ 0.00|55.02 $\\pm$ 2.15 | 1.61 $\\pm$ 0.0|\n | STree           | 57.15 $\\pm$ 2.16| 2.02 $\\pm$ 0.01| 62.68 $\\pm$ 1.68| 2.26 $\\pm$ 0.02 | 61.69 $\\pm$ 2.30 | 2.16 $\\pm$ 0.0|\n\n**W2: Better GPU**: While we do not have access to H100/A100 GPUs, we provide results obtained using A6000 GPU, which is the best we have. The results are presented below. We note that inference for SSM and hybrid models are still memory bound and computation speed is not the bottleneck. Therefore, faster computation units from better GPUs will not invalidate our improvements. \n\t\n | | MT_Bench |  | HumanEval|  | GSM8K | |\n |:---|:---|:---|:---|:---|:---|:---|\n |                 | Speed | $\\tau$ | Speed|   $\\tau$ | Speed | $\\tau$ |\n |                 |       | T=0 |       |      |       |      |\n | Autoregressive  |37.78  |   1 | 37.96 |   1  |       |  1   |\n | Vanilla SD      | 64.24 (1.70x) | 2.06| 73.90 (1.95x) | 2.34 | 74.10 \t(1.94x) | 2.33 |\n | STree           | 65.71 (1.74x) | 2.42 | 74.81 (1.97x) | 2.77 | 76.47 (2.01x) | 2.81 |\n |                 |       | T=1 |       |      |       |      |\n | Autoregressive  |37.45  |   1 | 37.49 |   1  |       |  1   |\n | Vanilla SD      | 47.09 (1.26x) | 1.57 | 49.71 (1.33x) | 1.66 | 49.30 | 1.63 (1.30x)|\n | STree           | 51.80 (1.38x) | 2.04 | 56.57 (1.51x)\t| 2.23 | 55.23 | 2.16 (1.46x)|"}, {"review_id": "wiHaCa8z3i", "reviewer": "Reviewer_7vpj", "summary": "This 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.", "questions": "1. 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?", "limitations": "1. 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.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 1, "strengths_and_weaknesses": "Strengths:\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.", "was_revised": true, "final_justification": "Although I think the proposed method designed for SSMs is not novel, as there is similar method designed for LLMs. \n\nThe authors' rebuttal soloved my questions with several experiments.\n\nI think the paper could be borderline accepted.", "rebuttal": "We thanks the reviewer for their reviews. We address the concerns and questions raised below:\n\n**W1**: See Q2.\n\n**W2**: See Q1.\n\n**W3**: See Q2.\n\n**Q1. Dynamic tree**: We note that optimizing the tree structure is orthogonal to our work. Our work focuses on enabling efficient computation of token tree with SSMs. With a baseline static tree, we are already able to obtain runtime improvement, as shown in Table 2 in paper. We also note that the beam search tree we used in Table 1 is not static, as the expansion of the tree depends on the joint probability and will vary from sample to sample. Here, we additionally show the end-to-end result for drafting with a beam search tree. We use greedy decoding with the M = 3 and N = 4 (M, N explained in table 1). The results are as follow: \n\n |    | Autoregressive | Vanilla SD | STree w. static tree | STree w. beam search tree|\n |:---|:-------|:--------|:-------|:--------|\n | Speed |  39.98 | 67.68 (1.69x) | 69.84 (1.75x) | 71.60 (1.79x) |\n | $\\tau$ |   1 | 2.04 | 2.47 | 2.60 |\n \nWe can see that with a beam search tree, we are able to get a better acceptance length and thus a faster end-to-end generation speed. We believe that more advance tree construction technique would bring even more benefits.\n\n**Q2: improvement not obvious**: At temperature=0, our STree is always better than Vanilla SD, though by a smaller margin. However, at temperature=1, our method achieves an average 12\\% gain, which is an obvious difference. Averaging over all 6 experiment, there is still a 7\\% gain, which is non-trivial. Meanwhile, we show the runtime for forward passes for bigger Mamba2 models. STree and Vanilla SD both pay an overhead over autoregressive runs to compute multiple tokens. As model size increase, the ratio between STree overhead and Vanilla SD overhead decreases from 1.48 to 1.19. This means that when we have a bigger model，it is much easier to achieve an end-to-end performance gain. Given that we already achieved end-to-end performance gain in 8B models as show in Table 2 in the paper, there is reason to believe that performance with bigger models will further improve with similar acceptance rate. \n\n |    | 2.7B | 7B | 13B | 23B |\n |:---|:-------|:--------|:-------|:--------|\n | Autoregressive |  10.95 |  22.94 | 40.69 | 72.41 |\n | STree |   22.36 (2.04x) | 33.79 (1.47x) | 55.90 (1.37x) | 91.08 (1.26x) |\n | Vanilla SD |  15.05 (1.37x) | 26.46 (1.15x) | 45.95 (1.13x) | 76.74 (1.06x) |\n\n**Q3: Generation Temperature=2,3,4**: In practice, no previous works in speculative decoding has used generation with temperature beyond 1. Using too high a temperature can lead to high perplexity in generated output and therefore hinders downstream tasks. We have the ablation study for temperature between 0 and 1 in figure 4b and 4c in the paper. Here, we additionally provide experiment with even higher temperature=1.1 and 1.2. We can see that the acceptance length and hence generation speed continues to decrease with such high temperature. \n\n| MT_Bench | | | \n |:---|:---:|:---:|\n |                 | Vanilla SD | STree |\n | Speed (T=1.1)  |49.90  |   55.25| \n | $\\tau$ (T=1.1)  | 1.59 | 2.07| \n | Speed (T=1.2)  | 48.32 | 52.04| \n | $\\tau$ (T=1.2)  | 1.51 | 1.94| \n\n**Q4: More benchmarks**: The three benchmarks we provide in Table 2 (MT\\_Bench, GSM8K, Human\\_Eval) is extensively used by previous works to benchmark speculative decoding performance with text completion as the task. We note that MMLU, HellaSwag, SWAG, BBH, BoolQ are all datasets on common sense reasoning and Natural Language inference. Some of them feature multiple choice questions, requiring only a single word as answer, which can lead to premature termination in generation. This could lead to abnormal token distribution if we continue to generate beyond the eos token, affecting the acceptance length of both the draft and target model. This is probably why these datasets are often not used in speculative decoding works. Nonetheless, we additionally provide BoolQ's reults here. We note that as we employ a lossless verification strategy with temperature=0, the results we generated will be identical to the result of the target model and we only focus on the runtime. We use the question in BoolQ as a prompt and let it generate for 1024 tokens to measure the runtime. \n\n |  BoolQ | |  | \n |:---|:---:|:---:|\n |                 | Speed | $\\tau$ |\n | Autoregressive  |39.98  |   1 | \n | Vanilla SD      | 70.52 | 1.91| \n | STree           | 71.35 | 2.29| \n\n**Q5: Larger GPUs**: While we do not have access to H100/A100 GPUs, we provide results obtained using A6000 GPU, which is the best we have. The results are presented below. We note that inference for SSM and hybrid models are likely still memory bound and computation speed is not the bottleneck. Therefore, faster computation units from better GPUs should not invalidate our improvements. \n\n | | MT_Bench |  | HumanEval|  | GSM8K | |\n |:---|:---|:---|:---|:---|:---|:---|\n |                 | Speed | $\\tau$ | Speed|   $\\tau$ | Speed | $\\tau$ |\n |                 |       | T=0 |       |      |       |      |\n | Autoregressive  |37.78  |   1 | 37.96 |   1  |       |  1   |\n | Vanilla SD      | 64.24 (1.70x) | 2.06| 73.90 (1.95x) | 2.34 | 74.10 \t(1.94x) | 2.33 |\n | STree           | 65.71 (1.74x) | 2.42 | 74.81 (1.97x) | 2.77 | 76.47 (2.01x) | 2.81 |\n |                 |       | T=1 |       |      |       |      |\n | Autoregressive  |37.45  |   1 | 37.49 |   1  |       |  1   |\n | Vanilla SD      | 47.09 (1.26x) | 1.57 | 49.71 (1.33x) | 1.66 | 49.30 | 1.63 (1.30x)|\n | STree           | 51.80 (1.38x) | 2.04 | 56.57 (1.51x)\t| 2.23 | 55.23 | 2.16 (1.46x)|\n\n**Q5: Draft model Ablation**: For MambaInLlama-8B model, there are no currently viable open-source drafting models that uses the same tokenizer. Therefore we distilled our own drafting model. If the reviewer would like to suggest any suitable model for drafting, we are happy do experiments. Here, we additionally provide an ablation study on our drafting model distilled to different steps. \n\n |  | 12000 step | 48000 step | 264000 step |\n |:---|:---|:---|:---|\n | Speed (Vanilla SD) |   44.52 (1.21x)| 53.11 (1.24x) | 56.49 (1.42x) |\n | $\\tau$ (Vanilla SD)|   1.34           | 1.56          |  1.68       |\n | Speed (STree) |   42.96 (1.19x)| 58.65 (1.36x) | 64.30 (1.62x) |\n | $\\tau$ (STree)|   1.51           | 2.02         |  2.23        |\n\nWe can see that as the distillation goes on, the quality of the drafting model gets better, which is reflected by the increase in acceptance length. The result also suggest that with better drafting model and technique, the end-to-end runtime of STree will become better.\n\n**L1: Novelty**: To the best of our knowledge, we are the first to apply speculative tree decoding to SSMs and hybrid models, which makes our work novel. While tree based speculative decoding exists for transformer based models, it heavily relies on the transformer's property that computing a tree of tokens doesn't incur much overhead with the use of attention masks. Without this property, it is much slower to compute a tree of tokens in SSMs due to their recurrent state update. Our work addresses this critical challenge by improving the efficiency of computing a tree of tokens with SSMs, which is not done in previous works.\n\n**L2**: See Q1\n\n**L3**: See Q4"}]

Author rebuttal:

Meta-review: This paper proposes STree, the first algorithm to perform tree-based speculative decoding in state-space models (SSMs) and hybrid architectures of SSMs and Transformer layers. The proposed method addresses the challenge of using SSMs to compute a token tree efficiently. The empirical evaluation demonstrates that STree outperforms vanilla autoregressive decoding and speculative decoding.

The reviewers generally believed the proposed method is technically sound, and the hardware-aware implementation will be practically useful for the community. Some reviewers asked about the scalability to larger models, and the authors provided additional experiments to address it. While a few reviewers still believed that the paper somewhat lacks novelty, they also regarded the contribution of the paper as satisfying the standard of NeurIPS. Reviewer wh4v's concerns, including the performance of STree on larger model sizes and the reliance on static trees, have been well-addressed by the authors during the rebuttal.

On balance, the AC recommends accepting the submission, as it could serve as a meaningful baseline for speculative decoding in SSMs.
Decision: Accept (poster)