用Mamba提升扩散语言模型推理速度,长序列下快8.2倍
DiffuMamba: High-Throughput Diffusion LMs with Mamba Backbone
- 用双向Mamba替代Transformer做扩散语言建模,实现线性时间处理
- 1.3B参数模型在长序列上推理吞吐量最高提升8.2倍,性能持平传统方法
- 首次系统分析各类扩散模型效率,证明Mamba混合器是未来关键方向
扩散语言模型(DLMs)作为自回归生成的有力替代,但其依赖Transformer架构导致推理效率受限,存在二次注意力或键值缓存开销。我们提出DiffuMamba,一种基于双向Mamba骨干网络的掩码扩散语言模型,融合扩散目标与线性时间序列建模;并推出DiffuMamba-H,一种交错注意力的混合变体。在参数规模达1.3B的模型中,我们的方法在下游任务表现与基于Transformer的扩散模型相当,同时在长序列上推理吞吐量分别提升最多8.2倍和4.3倍。我们进一步对现代DLM变体进行系统性效率分析,结合渐近复杂度与实测数据。值得注意的是,采用缓存高效块扩散与Mamba混合器的策略是唯一能随序列长度线性扩展且超越所有基线的方法,为未来基于扩散的生成系统指明了重要方向。
原文摘要 · Abstract (English)
Diffusion language models (DLMs) have emerged as a promising alternative to autoregressive (AR) generation, yet their reliance on Transformer backbones limits inference efficiency due to quadratic attention or KV-cache overhead. We introduce DiffuMamba, a masked diffusion language model built on a bidirectional Mamba backbone that combines the diffusion objective with linear-time sequence modeling, and DiffuMamba-H, a hybrid variant with interleaved attention. Across scales up to 1.3B parameters, our models match Transformer-based diffusion in downstream performance while achieving up to 8.2x and 4.3x higher inference throughput, respectively, on long sequences. We further present a systematic analysis of inference efficiency across modern DLM variants combining asymptotic complexity with empirical measurements. Notably, cache-efficient block diffusion with Mamba mixers emerges as the only strategy that scales linearly with sequence length and achieves the strongest performance across all baselines, suggesting a promising direction for future diffusion-based generation systems.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。