Speculative Streaming: Fast LLM Inference without Auxiliary Models

Speculative Streaming performance on very long sequences

To evaluate the performance of Speculative Streaming on long sequences, we trained speculative adapters on the Arxiv-summarization dataset and tested it on the Summarization task from the LongBench dataset. Since the KV cache is shared between the main and speculative streams, there is no additional runtime memory overhead. While compute in attention layers increases due to longer context, the compute in MLP layers remains the same, and decoding is still memory bandwidth bound. We achieved a 2.64X speedup on the Summarization test set using gamma = 3 and treeFactor = 4. We include these experiments in Appendix K of the revised version.

Effect of multi-stream attention mechanism on the model's understanding of context and coherence

The multi-stream attention mechanism improves context understanding and coherence by enabling the primary stream $M_t$ to incorporate information from speculative streams $S_{tj}$, which approximate future residual states. This mechanism affects the model in three key ways:

1. Improved Predictive Transformations: By attending to $S_{tj}$, $M_t$ refines its residual transformations, incorporating predictive signals about token relationships over the next $\gamma$ tokens. This enhances anticipatory planning during training, leading to more coherent token transitions.

2. Dynamic Context Alignment: Speculative streams progressively align with future residual states, helping $M_t$ better encode dependencies between current and future tokens. This reduces mismatches in long-range coherence.

3. Enhanced Representation Fidelity: Multi-stream attention ensures that $M_t$ captures a richer, temporally-aware context, leading to representations that align closely with the actual ground truth dynamics, thus improving generation quality metrics while maintaining coherence.

Sensitivity of the method to Number of streams and hyperparameters

We observe diminishing returns in speculative prediction accuracy with an increasing number of speculative streams. As detailed in the following table, stream prediction accuracy decreases by approximately 5-10% with each additional stream:

It’s important to note that with each added stream, the size of the speculative window grows quadratically, as discussed in Section 3.1.3. Based on this trade-off, we find that using 3 streams is near optimal for most tasks.

Regarding other hyperparameters, extensive experiments indicate that the following setup yields robust performance, achieving substantial speedups across different tasks and model sizes:

• Number of streams: 3
• Number of MSA layers: 1/8 of base model layers
• LoRA adapter rank: 8
• Tree Factor: 4
• $\alpha_0$: 1
• $\alpha_1, \alpha_2, \alpha_3$: 0.1

We sincerely appreciate the insightful feedback and thoughtful inquiries provided by the reviewer. In this response, we aim to comprehensively address the concerns and questions raised, hoping to clarify and enhance the understanding of our research.

The original LLM distribution is altered

Thank you for your thoughtful feedback. We acknowledge the importance of preserving the original LLM distribution in speculative decoding frameworks. As detailed in Section 3.1.4, to address different application scenarios, Speculative Streaming (SS) supports both "lossy" mode a.k.a. share mode (that you are referring to), and "lossless" mode (that you might not have noticed) and preserves output distributions similar to [1, 2, 3] . More specifically:

• In the lossy (shared) mode, parameter sharing occurs between LoRA adapters for next-token prediction and speculative draft generation. This design minimizes parameter overhead on constrained devices and is supported by observed gains in metric performance with shared parameters. Although, as you mentioned, this mode does not guarantee the preservation of the output distribution, it still outperforms the baseline, as demonstrated in Table 1.
• In the lossless mode, speculative adapters operate independently of the base adapters and only speculative adapters in MSA layers are trained. Similar to [1, 2, 3], this ensures the model’s output distribution remains identical to that of the original LLM, preserving generation quality. Details of this approach are provided in Section 3.1.4 and the Metrics subsection of the Experiments section.

For generic instruction-tuning tasks, where the base model is already well-trained, the lossless mode is employed, as highlighted in the Metrics subsection. This approach ensures the integrity of the output distribution and addresses concerns about sensitivity to fine-tuning and evaluation domains, as evidenced by the results in Figure 3, 4.

Lack of Experimental Detail

Thank you for raising this point. As noted in Section 3.1.4, Speculative Streaming (SS) supports both lossy (shared) and lossless modes. In the lossless mode, speculative adapters operate independently of base adapters, ensuring that the model’s output distribution matches the original LLM without requiring fine-tuning of base adapters. We only tune speculative adapter parameters on instruction following dataset - Vicuna as noted in section 3.1.4. The multi-stream mechanism and speculative candidate generation remains consistent across both modes, as detailed in Section 3.1.1. We are happy to include additional experimental details to clarify this setup further if needed.

Computation overhead Thank you for your observation. While Speculative Streaming (SS) does slightly increase FLOPs overhead due to the additional stream representations, it remains substantially more efficient than LooKAhead decoding, as shown in Figure 6. Additionally, decoding is often memory bandwidth-bound on mainstream accelerators, meaning the slight increase in compute overhead does not lead to increased decoding latency, as demonstrated in Figures 3 and 4.

tree drafting strategy for Medusa

we use optimal configurations for speculative tree generation for both Medusa and Speculative streaming for fair comparison as noted in Appendix G.1.

Call Reduction Ratio Metic CR ratio denotes ratio between number of model.forward() calls required for auto-regressive decoding to those required for speculative-streaming.

Reporting of extra parameters in Table 1

Thank you for pointing this out. We agree that Medusa-2 also tunes the base LLM parameters, which highlights the parameter efficiency of our approach. We will update both tuned and extra parameters for Medusa-2 in the final revision to make this distinction.

Comparison with Lora Fine-tuning : In line 403, the authors claim that "the generation quality of Speculative Streaming consistently outperforms that of next-token prediction-based fine-tuning, positioning it as a compelling alternative to LoRA-based fine-tuning approaches." However, no comparison with next-token prediction-based fine-tuning methods (such as LoRA tuning) is provided, which makes this an overstatement.

Thank you for the feedback. We would like to clarify that Table 1 provides a direct comparison between Speculative Streaming and next-token prediction-based fine-tuning (LoRA). As shown in the table 1, Speculative Streaming consistently achieves higher generation metrics compared to LoRA-based fine-tuning, supporting our claim.

Thanks for the authors' prompt response. I still have the following concerns:

1. Q1: The original LLM distribution is altered. Thanks for your clarification. I acknowledged the chatbot-like settings that $\alpha_0$ is set to 0. However, as you stated in Equation (1) that "We enable the main stream to attend to speculative streams, allowing it to plan its residual transformations based on anticipated future residual states." In this way, the original attention is altered, leading to changes in the LLM output distributions. How do you ensure that the target LLM's original output distribution remains unchanged given these architectural changes?
2. Q2: Lack of Experimental Detail. In the abstract, Section 3.1.4, and Metrics, the authors claim that "in lossless settings for generic chatbot applications that do not necessitate supervised fine-tuning." However, in your response, you claimed that "We only tune speculative adapter parameters on instruction following dataset - Vicuna as noted in section 3.1.4." This apparent contradiction suggests a potential overclaim in the manuscript. Could you clarify whether fine-tuning is indeed required for both lossy and lossless modes?
3. Q3: tree drafting strategy. The referenced Appendix G.1 appears to be missing from the manuscript. There is only Appendix G and I could not find comparisons between Medusa and SS. Besides, please provide the average draft length after pruning.
4. Q4: the CR (call reduction) ratio metric. You define the CR (call reduction) ratio as"the ratio between number of model.forward() calls required for auto-regressive decoding to those required for speculative-streaming." Is it equal to the mean number of generated tokens per step? Why don't you use this metric for consistency with prior work?
5. Q5: Comparison with Lora Fine-tuning. Thanks for the clarification that "Table 1 provides a direct comparison between Speculative Streaming and next-token prediction-based fine-tuning (LoRA)." Do you mean that the baseline is trained by LoRA? I noticed that the "# Extra Parameters" is set to 0 for the baseline. Do you mean that LoRA parameters are not counted as extra parameters in SS?

We sincerely appreciate the insightful feedback and thoughtful inquiries provided by the reviewer. In this response, we aim to comprehensively address the concerns and questions raised, hoping to clarify and enhance the understanding of our research.

Model Evaluation Thank you for your comment. Speculative Streaming operates in both lossy(shared) and lossless modes, as described in Section 3.1.4. In lossless mode, speculative adapters operate independently of the base adapters and the model output distribution remains identical to the base model. We provide extensive evaluation in Figures 3 and 4, showcasing speedups on Vicuna-7B, Vicuna-13B, LLaMA-7B, and LLaMA-13B, comparing our approach with state-of-the-art speculative decoding methods. We employ “lossy” (shared) mode in application specific settings, where parameters are shared between the LoRA adapters for next-token prediction and speculative draft generation to minimize the trainable parameter overhead on resource-constrained devices. In application specific settings we report metrics recommended for those tasks while in generic chatbot settings, we keep the output distribution same and only report speedup as noted inn Metrics subsection.

draft model may contain as little as 1% of the parameters of the target model

Thank you for the insightful comment. As highlighted in Section 1, in application-specific settings, a dedicated draft model must be trained to align with the fine-tuned target model for optimal acceptance rates. As the number of applications increases, the cumulative parameter overhead can become significant. The memory resources required to host these application-specific draft models could instead be allocated to training a larger and more capable target model.

While parameter efficiency is an important advantage of our approach, it is not the sole benefit. As demonstrated in Table 2, using a draft model that is 10-100x smaller in traditional speculative decoding results in suboptimal speedups. Appendix G provides a detailed analysis of the model size versus speedup trade-off in two-model speculative decoding, showing that excessively small models compromise acceptance rates, while larger draft models increase generation time, limiting overall speedup. In contrast, Speculative Streaming leverages the full LLM, eliminating the trade-offs between draft generation and draft acceptance. Because speculative tokens are generated within the same forward pass, there is no additional latency, and the decoding process remains memory bandwidth-bound on most mainstream accelerators. This ensures both high acceptance rates and substantial speedups without compromising efficiency or performance.

Formatting Issues Thanks for pointing out the compile issues. We have fixed all formatting issues in the revised submission.

speedup in comparison to 2-model SD? Yes your understanding is correct. Our method achieves 3X speedups relative to 2-Model SD on Vicuna-7B, Vicuna-13B, Llama-7B, LLama-13B as shown in Figure3 and 4.

What prompts are used during the evaluations?

Thank you for your question. For LLM evaluations, we use the prompts from the FastChat repository. As mentioned in the Metrics subsection, for instruction following tasks, our method trains only the speculative adapters, ensuring that the model’s output distribution remains identical to that of the baseline model. In supervised fine-tuning setups, such as for SQL response generation, we compute the exact match between the generated and target responses using the Huggingface library.

Are the target model weights frozen in all experiments?

Yes, the target model weights are frozen in all experiments. In Application specific settings, we use lossy(shared) mode and share LoRa parameters between main and speculative streams and train them jointly. In “lossless” mode, we use separate speculative adapters with no parameter sharing, ensuring the model’s output distribution remains identical to the base model, as noted in Section 3.1.4 and the Metrics subsection.

Evidence of generation quality

In Table 1, we provide a direct comparison of generation metrics between next-token prediction-based fine-tuning (LoRA) and Speculative Streaming. Speculative Streaming consistently achieves superior generation metrics.

Experiments with higher batch sizes

Thank you for your feedback. Speculative Streaming is primarily designed for resource-constrained edge device use cases where batch size is typically 1. However, as noted in the ablation section, we also explore the use of larger batch sizes (2 and 4) for server-based models, with minimal impact on generation speedups. We will consider including a more thorough comparison of throughput at higher batch sizes in final revision to demonstrate the practical performance of Speculative Streaming across different deployment scenarios.

The method of this paper is very similar to BiTA, https://arxiv.org/html/2401.12522v2.

Thanks for bringing this work to our attention. While there are similarities, such as using masked tokens for speculative generation, we highlight the following key differences and novelties in our approach:

Attention Mechanism with Speculative Residual States

Our approach features a novel attention mechanism where the main stream also attends to speculative residual states, in contrast to the BiTA paper, where speculative streams attend to the main stream without reciprocity. This design improves generation metrics by enabling next-token planning (Section 4.1.2).

Tree Pruning for Computational Efficiency

We propose a Tree Pruning mechanism where we prune tree draft nodes that are less likely to be accepted during final verification to reduce both compute and memory overhead, further accelerating decoding (Section 3.1.3, Appendix A.1 (B.1 in new revision)). In contrast, BiTA has no such mechanism.

Trainable Parameters

Unlike BiTA, our work focuses on NAR speculative draft generation using minimal trainable parameters (See Table 1) to transform residuals in each layer in form of speculative LoRA adapters, tailored for resource-constrained settings.

Shared vs. Lossless Modes for LoRA Adapters

Our work introduces shared and lossless modes for LoRA adapter parameters between the main and speculative streams, offering flexibility in balancing the number of trainable parameters and preserving the output distribution's integrity. We employ the shared mode for application-specific tasks and the lossless mode for general-purpose chat scenarios. In contrast, BiTa does not make this distinction.

Training Dataset Bias

We want to note that, as mentioned in the datasets subsection, we studied two settings: 1) Application-specific scenarios as you are referring to, where we perform lightweight supervised fine-tuning with LoRA adapters that are shared between main and speculative streams on a specific task (e.g., SQL), introducing minimal trainable parameter overhead (Table 1 and 2) and general instruction-following setting where adapters are not shared between main and speculative streams and only speculative adapters are trained on instruction set similar to Eagle. As reported in Figure 3, without any task-specific fine-tuning of base model parameters, we achieve 3X speedup on MT Bench, which includes tasks across categories like writing, coding, and reasoning while maintaining integrity of base model’s output.

Thank your thoughtful feedback and critique, and participating in the discussion. Below, we address each of the concerns raised and aiming to clarify the contributions of our work.

1. Acknowledgment of Concurrent Work:

We appreciate references to BiTA, Glide with CAPE, and Eagle-2.

First, please note that BiTA and “Glide with CAPE” were appeared publicly within 20 days of our work, and Eagle-2 was appeared 4 months later and we were unaware of them during the development of our work. That said, regardless of the timeline, we acknowledge the relevance of all these interesting works and clarify our technical differences below. We will also include a detailed discussion in the revised paper to acknowledge their contributions and to clearly delineate the differences and novelties of our approach.

2. Comparison to BiTA: While we acknowledge similarities with BiTA, our approach has some key distinctions:

   2.1 Training-Inference Consistency: BiTA’s Section 3.2 highlights a mismatch in attention, where mask tokens attend to all previous tokens during training but not during inference. As mentioned in section 3.2 mask tokens do NOT attend to draft tokens in speculative window. To illustrate this inconsistency, we quote two relevant sentences from Section 3.2:

   • “Note that they do not “attend” to the current draft token candidates, generating future predictions in the context of only both query tokens and previously generated tokens. “
   • “As stated, the draft token candidates are predicted based on the context of the output sequence until the last forward pass, rather than incorporating the currently accepted draft tokens.”

   On the other hand, as demonstrated in Section 3.1.2, our method employs γ streams per token during the verification step. Each of these γ streams attends to both the previous streams and the preceding context, ensuring consistency between training and inference.

   2.2 Flexibility between shared and lossless modes: Unlike BiTa, which enforces a strict separation between mask tokens and base model parameters to keep generation quality same, our approach in shared mode improves both generation quality and speedup in application specific settings. We find that, generation quality of responses can be improved significantly by having main stream attend to speculative streams that resemble residual states of future tokens leading to better next token planing (see section 4.1.2) and subsequently better generation quality as demonstrated in Table 1. On the other hand, we also offer a lossless mode in generic chatbot settings that is guaranteed to keep output distribution same as base model.

   2.3 Memory and Compute Efficiency: BiTA applies Masked attention across all layers, whereas we restrict the use of multi-stream attention to Ns=4 layers (Sections 3.1.2 and 3.1.3). This design reduces forward-pass latency and FLOP overhead, particularly as model transitions into compute-bound regimes with larger tree sizes.

We believe these design differences contribute to our higher speedup compared to BiTA (3.0x vs. 2.38x on LLaMA-7B as measured on MT-Bench).

3. Tree Pruning for Computational Efficiency:

   • While Glide with CAPE and Eagle-2 use confidence-based pruning directly on speculative draft logits, our approach introduces a separate filter for speculative token pruning. Speculative tokens generated in a non-autoregressive (NAR) manner are pruned for dependency enforcement using the initial layers of the target model via early exiting in next forward pass. As illustrated in Figure 11, we rely on early-exited logits for pruning rather than those produced by the NAR pass. By introducing dependency among consecutive nodes in speculative tree, we improve pruning efficacy leading to better speedups (See Figure 9).

4. Trainable Parameters:

   • We respectfully clarify that our method introduces fewer trainable parameters compared to BiTA. Specifically, while BiTA adds soft prompt tokens across all layers, we limit speculative adapter training to $N_s = 4$ layers while still surpassing BiTa in terms of speedup on MT Bench (3.0x vs. 2.38x on LLaMA-7B).

5. Dataset Scope: We appreciate your suggestion to evaluate on SPEC-Bench. We are actively extending our experiments to include broader datasets and will incorporate these results in the final version to further substantiate our claims.

6. lossless relative to the original LLM? Yes, that’s correct. It is lossless relative to the original LLM and guarantees no change to the output distribution of the original model.

We hope this response clarifies and addresses your concerns, but we would be happy to answer any follow up questions you might have. We sincerely value your feedback and are committed to improving the clarity and quality of our work.