Mamba Drafters for Speculative Decoding (2025)

Abstract: Speculative decoding has emerged as a promising approach to accelerating large language model (LLM) generation using a fast drafter while maintaining alignment with the target model’s distribution. However, existing approaches face a trade-off: external drafters offer flexibility but can suffer from slower drafting, while self-speculation methods use drafters tailored to the target model but require re-training. In this paper, we introduce novel drafters based on Mamba, a state-of-the-art state space model (SSM), as a solution that combines the best aspects of both approaches. By leveraging the linear structure of SSMs, our approach avoids the quadratic complexity inherent in traditional Transformer-based methods, enabling faster drafting and lower memory usage while maintaining the flexibility to work across different target models. We further enhance efficiency with a novel test-time tree search algorithm for generating high-quality draft candidates. Our empirical evaluation demonstrates that Mamba-based drafters not only outperform existing external drafting methods but are also comparable to state-of-the-art self-speculation approaches while using less memory and maintaining their cross-model adaptability.

## introduction recent breakthroughs in large language models (llms) have been largely driven by transformer architectures @xcite , which have enabled exceptional performance across a wide range of tasks @xcite @xcite . however, their capabilities often come with significant computational overhead, primarily due to the autoregressive nature of sequential token generation, while the quadratic complexity of the attention mechanism further exacerbates scalability challenges. speculative decoding (sd) @xcite @xcite has emerged as a promising approach to addressing the inefficiencies of autoregressive models by generating multiple candidate tokens with an efficient drafter and verifying them in parallel with the target model, ensuring identical output with greater efficiency. this approach enables simultaneous decoding of multiple tokens within a single forward pass. existing sd methods are mainly categorized into two types: (i) using an external drafter @xcite applicable to multiple target models, and (ii) adopting a selfspeculation approach, where the target model itself is leveraged for drafting, showing faster drafting speed compared to external drafters @xcite @xcite . for instance, self-speculation involves training a small transformer head on top of the target model's large transformer block to generate multiple candidate tokens, which are then used as a draft for the target model @xcite . while showing a faster drafting speed compared to using an external drafter, training a separate drafter for each target model is computationally expensive. more importantly, to handle distribution shift-i.e., to process novel inputs unseen during training-the drafter should be trained on a large corpus, which is particularly challenging because using the last hidden state as input, still requires forwarding all lower layers of the large target model during training, leading to significant computational overhead. this raises a key question: how can we develop an external drafter to have cross-model adaptability while also avoiding the limitation of the transformer's quadratic computation for fast drafting? this naturally leads us to explore non-quadratic sequential models such as state-space models (ssms) @xcite , as external drafters. ssms leverage a linear recurrence structure with a fixed state size, ensuring per-token computation and memory complexity remain constant during inference, making them more effective drafters than transformers. specifically, we use mamba @xcite , a state-of-the-art ssm, as a drafter in sd and make the following key observations: to further leverage mamba's efficiency, we propose simple yet effective tree decoding strategies with mamba drafters by formulating decoding as a multi-armed bandit (mab) @xcite problem. specifically, we introduce a test-time tree search algorithm that dynamically optimizes the draft tree structure based on the input query. with mamba's low latency, we observe that it benefits from trees of varying widths and lengths (see recent advancements, such as mamba @xcite , leverage hardware-aware algorithms and selection mechanisms to enhance ssms further. these developments have enabled ssms to demonstrate effectiveness in complex tasks across diverse domains, including language, audio, and video @xcite . building on this foundation, our research aims to explore mamba's potential as a drafter and utilize its efficiency to achieve faster and more effective speculative decoding. speculative decoding. speculative decoding follows a draft-and-verify framework @xcite , where a smaller drafter generates candidate tokens that are verified by the target model. this method accelerates generation by increasing parallelism while ensuring alignment with the target model's distribution. later advancements extended this approach to sampling settings @xcite and incorporated various optimization techniques to improve efficiency. speculative decoding approaches can be broadly categorized into two types. one approach utilizes external drafter models @xcite , which provide high flexibility, allowing a single drafter to be directly used for different target models. on the other hand, self-speculation @xcite takes a different approach by utilizing a very small model that uses the target model's internal hidden states for drafting. while being faster, they require expensive re-training of the drafter for every target model, showing limited flexibility. in this work, we demonstrate that mamba can get the best of both worlds, being an external drafter with very fast drafting speed. to further improve the acceptance probability of the drafts, recent works move from sequential drafting to tree-structured drafting @xcite @xcite , which allows verifying multiple draft candidates in parallel. additionally, researchers have explored non-transformer drafters @xcite . @xcite introduced hardware optimizations for applying ssms to speculative decoding, we explored the effectiveness on transformer-based target models (which is a de facto architecture) and developed an efficient inference scheme (i.e., tree-drafting) for mamba. furthermore, we present a more thorough comparison and analysis across different drafter models. in this section, we demonstrate that mamba can serve as powerful external drafters for speculative decoding (sd). we examine this in terms of both efficiency and effectiveness. for efficiency, we compare latency and peak memory usage during drafting, which includes encoding (prefill), initial forwarding of the given input sequence, and per-token generation during decoding (in figure 1 ). for effectiveness, we report throughput, the average number of tokens per unit time, and acceptance length, the average number of tokens accepted per forward pass of the target model (in figures 2 and 3 ). preliminaries: speculative decoding. given a target model m p , an efficient drafter m q and an input sequence x prefix , sd starts by drafter m q generates candidate tokens with a length of γ tokens, denoted as x1 , x2 , ..., xγ , where each are sampled from drafter distribution, xi ∼ q i (x | x1 , ..., xi-1 , x prefix ). after drafting, candidate tokens along with x prefix are passed to the target model m q in parallel and corresponding target distribution p i (x | x1 , ..., xi-1 , x prefix ) and tail distribution p γ+1 is obtained. finally, each candidate token xi is verified by criterion sequentially from @xmath0. here, the criterion is determined based on p i and q i . once a token is rejected at the position of i, a new token is sampled from the adjusted distribution p ′ ## tree-structured drafting with mamba in this section, we introduce an effective drafting strategy for mamba drafters by using treestructured decoding, i.e., hierarchically expanding multiple candidate nodes at each step instead of sequentially generating tokens. specifically, we suggest an efficient way to implement tree search for mamba decoding (in section 4.1) and introduce a test-time tree searching algorithm to adaptively optimize draft tree structure (in section 4.2). ## experiments in this section, we present a comprehensive evaluation of our proposed mamba drafter framework. baselines. we use a wide range of transformerbased drafters as baselines, applying tree drafting for a fair comparison with our mamba drafter. additionally, we consider eagle @xcite , a recent single-layer transformer drafter that leverages tree drafting and is directly trained to align with the target model, as a baseline for selfspeculation methods. for eagle, we follow its own adaptive tree search using confidence-based strategy, i.e., eagle-2 @xcite . ## conclusion in this work, we present mamba-based drafters as an effective solution to the challenges of existing speculative decoding methods. leveraging the linear structure of state-space models, mamba significantly improves drafting speed and memory efficiency. to further enhance drafting quality, we introduce a novel tree-based search algorithm. our experimental results show that mamba-based drafters not only outperform existing external drafting techniques but also match the performance of advanced self-speculation approaches, particularly in long-context scenarios.