Accelerated Diffusion Models via Speculative Sampling

📄 arXiv: 2501.05370v2 📥 PDF

作者: Valentin De Bortoli, Alexandre Galashov, Arthur Gretton, Arnaud Doucet

分类: cs.LG, stat.ML

发布日期: 2025-01-09 (更新: 2025-07-06)

💡 一句话要点

提出基于推测采样的加速扩散模型方法，显著提升生成速度。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 扩散模型 推测采样 加速算法 图像生成 采样方法

📋 核心要点

  1. 扩散模型生成高质量样本但计算成本高昂，限制了其应用。
  2. 利用快速草稿模型进行推测采样，加速扩散模型的生成过程。
  3. 实验表明，该方法在多种扩散模型上显著提高了生成速度，同时保证了样本质量。

📝 摘要（中文）

本文将推测采样技术扩展到扩散模型，该技术通常用于加速大型语言模型的推理，通过快速草稿模型生成候选token，并根据目标模型的分布接受或拒绝它们。在扩散模型的背景下，目标模型是高质量但计算成本高的模型。我们提出了多种草稿策略，包括一种简单有效的方法，该方法不需要训练草稿模型，并且可以直接应用于任何扩散模型。实验表明，在各种扩散模型上，该方法显著提高了生成速度，在生成目标模型的精确样本的同时，将函数评估次数减半。

🔬 方法详解

问题定义：扩散模型虽然能生成高质量的图像，但其迭代式的采样过程计算量巨大，严重限制了其应用。现有的加速方法通常需要复杂的模型蒸馏或者专门设计的加速架构，缺乏通用性和易用性。

核心思路：借鉴大型语言模型中推测采样的思想，利用一个快速但精度较低的“草稿模型”来预测扩散过程中的中间步骤，然后使用计算代价更高的目标模型来验证这些预测，从而减少目标模型的调用次数。

技术框架：该方法的核心在于使用一个快速的草稿模型来生成一系列的候选样本点，然后使用目标扩散模型来评估这些候选样本点的接受概率。如果接受概率高，则直接采用草稿模型生成的样本点；如果接受概率低，则重新从目标模型采样。通过这种方式，可以减少对目标模型的采样次数，从而加速整个生成过程。整体流程包括：1. 使用草稿模型生成候选样本；2. 使用目标模型计算接受概率；3. 根据接受概率决定是否接受候选样本；4. 重复以上步骤直到生成最终样本。

关键创新：将推测采样这一加速技术从离散序列生成（如语言模型）推广到连续向量生成（如扩散模型）。提出了一种无需训练额外草稿模型的推测采样策略，可以直接应用于现有的扩散模型，具有很强的通用性。

关键设计：论文提出了多种草稿策略，其中一种简单而有效的方法是直接使用目标扩散模型在较少的迭代步数下进行采样作为草稿模型。关键在于如何设计接受/拒绝的准则，使其能够保证最终生成的样本仍然服从目标模型的分布。具体的接受概率计算公式依赖于扩散模型的具体形式，需要根据不同的扩散模型进行调整。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，该方法在多种扩散模型上实现了显著的加速效果，可以将函数评估次数减半，同时保证生成样本的质量。该方法不需要训练额外的草稿模型，可以直接应用于现有的扩散模型，具有很强的实用性。

🎯 应用场景

该研究成果可广泛应用于图像生成、音频合成、视频编辑等领域，尤其是在对生成速度有较高要求的场景下，例如实时图像生成、交互式设计等。通过加速扩散模型的生成过程，可以降低计算成本，提高用户体验，促进扩散模型在更多实际应用中的落地。

📄 摘要（原文）

Speculative sampling is a popular technique for accelerating inference in Large Language Models by generating candidate tokens using a fast draft model and accepting or rejecting them based on the target model's distribution. While speculative sampling was previously limited to discrete sequences, we extend it to diffusion models, which generate samples via continuous, vector-valued Markov chains. In this context, the target model is a high-quality but computationally expensive diffusion model. We propose various drafting strategies, including a simple and effective approach that does not require training a draft model and is applicable out of the box to any diffusion model. Our experiments demonstrate significant generation speedup on various diffusion models, halving the number of function evaluations, while generating exact samples from the target model.