arXiv:2506.01206cs.CLcs.AI2025-06EMNLP被引 3

用Mamba模型做快速草稿生成,提速降存还兼容多模型。

Mamba Drafters for Speculative Decoding

  • 基于Mamba的线性结构设计新草稿生成器,避免传统Transformer的二次复杂度。
  • 在多个模型上测试,速度优于现有外部草稿方法,内存占用更低。
  • 适合需要高效推理且跨模型通用的LLM应用,如在线服务和边缘部署。

推测解码已成为一种有前景的加速大语言模型生成的方法,通过使用快速草稿模型同时保持与目标模型分布一致。然而,现有方法存在权衡:外部草稿模型虽灵活但生成较慢,自推测方法虽适配目标模型但需重新训练。本文提出基于Mamba(一种先进的状态空间模型,SSM)的新草稿模型,融合两者优势。利用SSM的线性结构,该方法避免了传统Transformer中的二次复杂度,实现更快的草稿生成与更低内存消耗,同时具备跨不同目标模型的灵活性。我们进一步提出一种新的运行时树搜索算法,以生成高质量草稿候选。实证评估表明,基于Mamba的草稿模型不仅优于现有外部草稿方法,且在性能上可媲美最先进的自推测方法,同时内存占用更少并保持跨模型适应性。

原文摘要 · Abstract (English)

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.

推理加速Mamba推测解码低内存

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。