DiffuMamba: High-Throughput Diffusion LMs with Mamba Backbone

📄 arXiv: 2511.15927 📥 PDF

作者: Vaibhav Singh, Oleksiy Ostapenko, Pierre-André Noël, Eugene Belilovsky, Torsten Scholak

分类: cs.LG, cs.AI

发布日期: 2026-07-20


💡 一句话要点

提出DiffuMamba以提升扩散语言模型的推理效率

🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture)

关键词: 扩散语言模型 推理效率 Mamba骨干 长序列生成 自然语言处理 深度学习 模型优化

📋 核心要点

  1. 现有的扩散语言模型依赖于Transformer骨干,导致推理效率低下,尤其是在长序列生成时。
  2. 本文提出DiffuMamba,利用双向Mamba骨干实现掩蔽扩散语言模型,结合线性时间序列建模以提高推理效率。
  3. 实验结果表明,DiffuMamba在长序列推理上相较于传统模型提升了8.2倍和4.3倍的推理吞吐量,表现出色。

📝 摘要(中文)

扩散语言模型(DLMs)作为自回归生成的有力替代方案,因其依赖于Transformer骨干网络而在推理效率上受到限制。本文提出了DiffuMamba,一种基于双向Mamba骨干的掩蔽扩散语言模型,结合了扩散目标与线性时间序列建模。同时,提出了DiffuMamba-H,一种具有交错注意力的混合变体。在参数规模达到13亿的情况下,模型在下游任务中的表现与基于Transformer的扩散模型相当,同时在长序列推理上实现了最高8.2倍和4.3倍的推理吞吐量提升。我们还对现代DLM变体的推理效率进行了系统分析,发现Mamba混合器的缓存高效块扩散是唯一能够线性扩展序列长度并在所有基线中表现最佳的策略,指明了未来扩散生成系统的有希望方向。

🔬 方法详解

问题定义:现有的扩散语言模型(DLMs)通常依赖于Transformer架构,这导致在推理时面临二次注意力或KV缓存开销的问题,尤其在处理长序列时效率低下。

核心思路:本文提出DiffuMamba,通过引入双向Mamba骨干网络,结合掩蔽扩散目标与线性时间序列建模,旨在提升推理效率并保持生成质量。

技术框架:DiffuMamba的整体架构包括一个双向Mamba骨干网络,采用掩蔽策略进行扩散生成,同时引入DiffuMamba-H变体,增加交错注意力机制以进一步优化性能。

关键创新:DiffuMamba的主要创新在于其线性时间复杂度的序列建模能力,显著提升了推理效率,并且在长序列生成中表现优异,区别于传统的基于Transformer的扩散模型。

关键设计:模型设计中,采用了高效的块扩散策略,结合Mamba混合器,确保在推理过程中能够有效利用缓存,优化了参数设置与损失函数,提升了整体性能。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示,DiffuMamba在长序列推理上的吞吐量提升显著,最高可达8.2倍,而DiffuMamba-H变体也实现了4.3倍的提升。这些结果表明,DiffuMamba在保持生成质量的同时,极大地提高了推理效率,超越了现有的基于Transformer的扩散模型。

🎯 应用场景

DiffuMamba的研究成果在自然语言生成、对话系统和文本摘要等领域具有广泛的应用潜力。其高效的推理能力使得在实时生成任务中能够更好地满足用户需求,未来可能推动更复杂的生成系统的发展。

📄 摘要(原文)

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.