arXiv:2604.18471cs.LG2026-04中稿 · ICLR被引 1

优化采样顺序,让离散扩散模型生成速度提升14倍

NI Sampling: Accelerating Discrete Diffusion Sampling by Token Order Optimization

论文配图:NI Sampling: Accelerating Discrete Diffusion Sampling by Token Order Optimization
图 1 · 摘自论文原文
  • 用神经指示器动态决定每步采样哪些词元
  • 相比全量采样减少90%迭代步数,准确率几乎不变
  • 适合需要快速生成的场景,如对话系统、内容创作

离散扩散语言模型(dLLMs)作为传统自回归方法的替代方案,具备任意顺序生成词元和并行解码的潜力。然而现有启发式采样策略效率低下:每步仅采样少量词元,改进空间巨大。本文研究词元采样顺序优化问题,发现充分利用每步正确预测可使采样迭代次数减少一个数量级,而不影响准确性。为此提出神经指示采样(NI Sampling)框架,通过神经指示器决定每步应采样的词元,并设计一种新的轨迹保持目标函数进行训练。在LLaDA和Dream模型上的多任务基准测试表明,该方法相比全步采样最高实现14.3×加速,性能损失微乎其微,且在准确率-步数权衡上持续优于置信度阈值采样。代码已公开于https://github.com/imagination-research/NI-Sampling。

原文摘要 · Abstract (English)

Discrete diffusion language models (dLLMs) have recently emerged as a promising alternative to traditional autoregressive approaches, offering the flexibility to generate tokens in arbitrary orders and the potential of parallel decoding. However, existing heuristic sampling strategies remain inefficient: they choose only a small part of tokens to sample at each step, leaving substantial room for improvement. In this work, we study the problem of token sampling order optimization and demonstrate its significant potential for acceleration. Specifically, we find that fully leveraging correct predictions at each step can reduce the number of sampling iterations by an order of magnitude without compromising accuracy. Based on this, we propose Neural Indicator Sampling (NI Sampling), a general sampling order optimization framework that utilize a neural indicator to decide which tokens should be sampled at each step. We further propose a novel trajectory-preserving objective to train the indicator. Experiments on LLaDA and Dream models across multiple benchmarks show that our method achieves up to 14.3$\times$ acceleration over full-step sampling with negligible performance drop, and consistently outperforms confidence threshold sampling in the accuracy-step trade-off. Code is available at https://github.com/imagination-research/NI-Sampling.

扩散模型生成加速采样优化

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