arXiv:2607.02805cs.LGcs.AI2026-07

提出新型混合扩散模型,实现长文本生成的高效并行计算。

Training Hybrid Block Diffusion Language Models with Partial Bidirectionality

论文配图:Training Hybrid Block Diffusion Language Models with Partial Bidirectionality
图 1 · 摘自论文原文
  • 仅在去噪块内进行反向Mamba扫描,支持精确缓存复用。
  • 65K上下文长度下推理速度达全序列DiffuMamba-H的19.7倍。
  • 适合需要高吞吐长文本生成的研究与工业应用。

大语言模型在长上下文生成中面临高吞吐挑战,生成过程常受内存带宽限制:每步解码需从内存读取累积的键值缓存,带宽随上下文长度增长而上升,但每次仅输出一个词元。为此,现有方法分为两类:通过高效注意力或线性时间混合器(如Mamba)减少内存访问,或通过一次性生成词元块提升并行度。然而二者结合存在技术难题。早期混合扩散模型如DiffuMamba采用双向Mamba混合,包含与因果生成相反方向的扫描,该反向扫描需遍历整个序列,状态非前缀独有,无法在分块去噪时精确复用缓存。本文提出BDLM Mamba-attention混合模型,将反向Mamba扫描限制在活跃去噪块内,实现跨块精确缓存。在8700万参数的DCLM实验中,BDLM Mamba-H在C4-en验证集上达到最优困惑度,3.5亿参数版本仍优于BDLM attention。长上下文推理中,其在6.5万词元时吞吐率达全序列DiffuMamba-H的19.7倍,在26.2万词元时为BDLM attention的3.7倍,表明Mamba混合架构是长上下文扩散生成的可行路径。

原文摘要 · Abstract (English)

High-throughput long-context generation is one of the central challenges for large language models. Generation is typically memory-bandwidth-bound rather than compute-bound: each decoding step must stream the accumulated key/value (KV) cache from memory, so bandwidth demand grows with context length while only one token is emitted. Two parallel approaches have therefore emerged: reducing memory access with efficient attention variants and linear-time mixers such as Mamba, or increasing parallel computation by generating blocks of tokens at once. However, technical challenges arise when combining these two ideas. Earlier hybrid diffusion models such as DiffuMamba use bidirectional Mamba mixing, including a reverse-direction scan relative to causal generation. This reverse scan needs to scan the entire sequence, so its states are not prefix-only and cannot be precisely reused as a cache even when diffusion is performed block by block. We propose a BDLM Mamba--attention hybrid that addresses this challenge by restricting the reverse Mamba scan to the active denoising block, which enables exact caching across blocks. In an 87M-parameter DCLM sweep, BDLM Mamba-H achieves the best C4-en validation perplexity compared to BDLM attention and full-sequence baselines. At 350M parameters, it remains competitive with BDLM attention. For long-context inference, BDLM Mamba-H reaches 19.7x the throughput of full-sequence DiffuMamba-H at 65K tokens and 3.7x the throughput of BDLM attention at 262K, showing that Mamba hybrids are a potential long-context diffusion architecture.

扩散模型长文本生成Mamba高效推理

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