通过并行时间采样加速离散扩散模型

离散扩散模型广泛用于学习和生成离散分布。由于生成过程本质上是连续的，因此采样的加速非常重要。在这项工作中，我们在连续时间马尔可夫链（CTMC）框架中并行化了主流的 $τ$-leaping 算法，用于吸收离散扩散。通过利用$τ$跳跃算法的连续时间随机积分形式和Picard迭代方法，我们实现了时间并行采样加速，并为我们的算法提供了指数阶乘收敛的证明。相对于 NFE，我们将吸收设置下 $τ$ 的总体时间复杂度从 ${\mathcal{O}}(d \log S)$ 提高到 ${\mathcal{O}}(\log (d\log S)\cdot \log d)$。根据经验，我们的方法在合成和真实数据设置中显示出一致的加速。新的采样器在合成分发方面实现了最多 $7$--$9\times$ 运行时加速，并在保持相同质量的情况下减少了 $50\%$ NFE，并且在单个 GPU 上的图像/文本任务中实现了 $1.45$--$1.86\times$ 运行时加速。我们的研究扩展了离散扩散模型在高效并行推理方面的潜力，对分子结构和语言生成等应用具有更广泛的影响。