arXiv:2509.01025cs.LG2025-09被引 49

让离散扩散模型支持可变长度生成,还能并行快速推理。

Any-Order Flexible Length Masked Diffusion

  • 通过插入掩码标记并逐步解码,实现任意顺序生成
  • 在合成迷宫任务中成功率比基线高约60%
  • 可轻松将预训练模型转为可变长生成,适合代码和数学任务

掩码扩散模型(MDMs)作为离散域上的自回归模型替代方案崭露头角。它们以任意顺序、并行方式生成序列,实现快速推理并在非因果任务上表现优异。然而,其关键局限在于不支持插入新标记,仅限固定长度生成。为此,我们提出柔性掩码扩散模型(FlexMDMs),一种可同时建模可变长度序列且在理论上保持任意顺序推理灵活性的离散扩散范式。基于随机插值框架的扩展,FlexMDMs 通过插入掩码标记并逐步解码生成序列。实证表明,FlexMDMs 在困惑度上与 MDMs 相当,同时对长度分布的建模精度显著提升。在合成迷宫规划任务中,其成功率比 MDM 基线高出约 60%。此外,我们证明预训练的 MDM 可高效转化为 FlexMDM:仅用 16 张 H100 显卡,三天内即可将 LLaDA-8B 微调为 FlexMDM,数学任务(GSM8K)性能从 58% 提升至 67%,代码补全性能从 52% 提升至 65%。

原文摘要 · Abstract (English)

Masked diffusion models (MDMs) have recently emerged as a promising alternative to autoregressive models over discrete domains. MDMs generate sequences in an any-order, parallel fashion, enabling fast inference and strong performance on non-causal tasks. However, a crucial limitation is that they do not support token insertions and are thus limited to fixed-length generations. To this end, we introduce Flexible Masked Diffusion Models (FlexMDMs), a discrete diffusion paradigm that simultaneously can model sequences of flexible length while provably retaining MDMs' flexibility of any-order inference. Grounded in an extension of the stochastic interpolant framework, FlexMDMs generate sequences by inserting mask tokens and unmasking them. Empirically, we show that FlexMDMs match MDMs in perplexity while modeling length statistics with much higher fidelity. On a synthetic maze planning task, they achieve $\approx 60 \%$ higher success rate than MDM baselines. Finally, we show pretrained MDMs can easily be retrofitted into FlexMDMs: on 16 H100s, it takes only three days to fine-tune LLaDA-8B into a FlexMDM, achieving superior performance on math (GSM8K, $58\% \to 67\%$) and code infilling performance ($52\% \to 65\%$).

扩散模型可变长度序列生成代码补全

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