提出稀疏调度策略,让扩散语言模型并行生成文本,速度提升近6倍且质量不降。
Plan for Speed: Dilated Scheduling for Masked Diffusion Language Models
- 将序列位置分组为非相邻块,分批并行解码以减少信息干扰。
- 在多个基准上实现最高5.8倍的加速,同时保持与逐词生成相当的质量。
- 无需修改模型,可直接用于现有扩散模型,适合追求高效生成的场景。
掩码扩散语言模型(MDLMs)有望实现快速的非自回归文本生成,但现有采样器基于模型置信度选择解码位置,忽略并行解码时的相互影响,导致实际退化为缓慢的自回归行为。本文提出稀疏解码调度器(DUS),一种仅用于推理、无需规划模型的无参数方法:将序列位置划分为非相邻的稀疏组,并在每步去噪中并行解码,以最小化联合熵增的上界。通过显式权衡网络调用次数与生成质量,DUS 恢复了传统并行策略损失的大部分性能。在数学(GSM8K, MATH500)、代码(HumanEval, MBPP)、通用知识(BBH, MMLU-Pro)和指令遵循(IFEval)等任务上,DUS 均优于基于置信度的规划器,并将扩散模型特有的质量-速度权衡转化为由块大小 $B$ 决定的确定性、可预测加速,相比逐词解码最多获得5.8倍的墙钟时间加速,且不需修改底层去噪器。作为即插即用后处理,稀疏间距亦能提升自适应采样器性能。代码已开源:https://github.com/omerlux/DUS。
原文摘要 · Abstract (English)
Masked diffusion language models (MDLMs) promise fast, non-autoregressive text generation, yet existing samplers, which pick tokens to unmask based on model confidence, ignore interactions when unmasking multiple positions in parallel and effectively reduce to slow, autoregressive behavior. We propose the Dilated Unmasking Scheduler (DUS), an inference-only, planner-model-free method that partitions sequence positions into non-adjacent dilated groups and unmasks them in parallel so as to minimize an upper bound on joint entropy gain at each denoising step. By explicitly trading off the number of network calls against generation quality, DUS recovers most of the performance lost under traditional parallel unmasking strategies. Across math (GSM8K, MATH500), code (HumanEval, MBPP), general-knowledge (BBH, MMLU-Pro), and instruction following (IFEval) benchmarks, DUS outperforms confidence-based planners and turns the diffusion-specific quality-speed trade-off into a deterministic, predictable speedup set by the block size $B$, yielding up to $5.8\times$ wall-clock speedup over token-by-token MDLM decoding without modifying the underlying denoiser. Applied as a drop-in post-filter, dilated spacing also improves adaptive samplers. Code is available at https://github.com/omerlux/DUS.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。