Efficient Diffusion(一):从多步采样到 Few-Step 生成

2 minute read

Published:

本系列想重点回顾一下 few-step diffusion 的发展历史。

众所周知,Diffusion Model 能够生成十分精美的图像和视频。与此同时,World Model 的发展也在很大程度上受到了 Diffusion 的影响。虽然 Diffusion 未必就是未来的最终答案,但至少在当前阶段,它依然是视觉生成和世界建模中最重要的技术路线之一。

相比于纯文字,图像天然具备更高的信息密度,尤其擅长承载空间结构、几何关系和视觉细节。

很多事情,一张图片就能讲清楚;但如果换成文字,可能需要成千上万倍的信息量也难以描述准确。比如,你试着用文字描述一个人的长相,然后让别人仅凭这段描述把那个人画出来,就会发现这件事几乎不可能做好。

这种得天独厚的优势,使得视觉生成模型逐渐成为 World Model 的重要 backbone。在某种程度上,WAM,也就是 World Action Model 的兴起,也很大程度上受到了 Video Generation Model 的启发:如果模型能够生成未来的视频帧,那么它或许也能够模拟动作对世界状态的影响。

但 Diffusion 老生常谈的问题就是:

我们知道,当前的 Diffusion Model 还很难被广泛用于机器人实时推理,或者直接作为世界模拟器使用。一个重要原因就是它的推理速度太慢。你敢信,一个 5 秒、720p 的视频,可能需要接近 1GB 的 token、30GB 的模型权重,并进行几十次去噪计算,最终耗时数千秒吗?

而机器人的控制频率通常在 50Hz 左右,也就是每 0.02 秒就要完成一次感知、推理和控制决策。这意味着,如果 Diffusion Model 想真正成为实时的 World Simulator,中间还有着巨大的鸿沟。

目前已有的 Efficient Diffusion 技术路线大致可以分成几类:

  1. 压缩数据量:例如在 latent space 中去噪,而不是直接在 pixel space 中建模。
  2. 减少去噪步数:通过 few-step 或 one-step 采样降低 latency,提高 FPS。
  3. 优化网络结构:例如稀疏注意力、局部注意力、低秩分解等。
  4. 量化与蒸馏:降低模型权重和推理计算成本。
  5. Cache 复用:复用历史特征或 attention 计算结果,加速视频生成或连续帧生成。

本系列希望尽量全面地覆盖主流 Efficient Diffusion 方案,同时提出一些自己的思考:这些策略各自解决了什么问题?它们是否足够本质?未来还有哪些有潜力的方向?

Why Diffusion?

我们首先来聊一个最根本的问题:为什么是 Diffusion?

对啊,Diffusion 明明有这么多弊端,为什么大家还要用它?

答案很直接:因为效果好。

那为什么效果好?这就要从 Diffusion 最初的 insight 说起。

Diffusion 的灵感部分来自物理学中的非平衡热力学。比如,在一杯水中滴入一滴红墨水,随着时间推移,红墨水会逐渐扩散到整杯水中。这个过程是自然界熵增的体现。

那这和图像生成有什么关系?

我们可以做一个类比:如果把一张 $256 \times 256 \times 3$ 的图片的所有可能取值看成一个极高维空间,那么真实图像并不是均匀分布在整个空间中的。真正有意义的图片,其实只占据了所有可能像素组合中的极小一部分。

换句话说,真实图像更像是高维空间中的一个低维流形。

如果把整个高维空间看成一杯水,那么真实图像就像其中那一滴红墨水。Diffusion 要做的事情,就是从高熵的噪声状态出发,逐步逆转扩散过程,最终回到真实图像所在的低维流形附近。

不过,这里有一个关键点:扩散过程并不是一一对应的。

比如你在水杯左边滴一滴红墨水,或者在右边滴一滴红墨水,经过足够长的时间之后,它们可能都会扩散成类似的稳定状态。因此,单纯从最终状态反推初始状态,本质上是一个不适定问题。

但是,如果我们引入一些条件,比如图像类别、文本描述、视频上下文,甚至物理约束,那么逆过程就会变得更有方向。

这时候模型学习的目标就不是某一滴具体的红墨水,而是整个扩散与反扩散的规律。给定任意 condition,模型都可以依靠这个规律,从高熵噪声中逐步回到符合条件的数据流形上。

那么,相比于之前的 GAN 和 VAE,Diffusion Model 的优势在哪里?

我们先回顾一下 VAE。VAE 的训练过程可以概括为:

\[image \rightarrow Encoder \rightarrow z \rightarrow Decoder \rightarrow image\]

也就是说,VAE 先把图像压缩成 latent variable $z$,再从 $z$ 解码回图像。

诶?这好像和现在很多 one-step generation 有点像?

但注意,VAE 中的 $z$ 是由 encoder 从真实图像中提取出来的,它已经包含了大量关于原图的统计信息。而从零开始生成图像时,我们并没有这样一个现成的充分统计量。模型必须从随机噪声出发,自己找到通往数据流形的路径。

至于 GAN,它的优势是生成速度快,但由于采用对抗训练,训练过程往往不稳定,也容易出现 mode collapse。

Diffusion 的优势在于,它把一个极其困难的生成问题拆成了很多个相对简单的去噪子问题。每一步只需要学习局部的去噪方向,而不是一步到位地从随机噪声映射到高质量图像。因此,Diffusion 的训练通常更加稳定,生成质量也更容易随着模型规模和数据规模提升。

综上,Diffusion 的核心 insight 可以概括为:

Image 和 Video 等视觉数据位于高维空间中的低维流形上,而 Diffusion 学习的是一个从高维随机噪声逐步回到低维数据流形的过程。

当然,Diffusion 可能不会是最终答案。如果未来出现了更强的、能够直接建模高维空间到低维流形映射的方法,也许它会取代 Diffusion。

但至少现在,Diffusion 依然是生成建模中最有效、最稳定的路线之一。

Why Multistep?

那么,为什么 Diffusion 需要多步采样?

因为真实数据位于高维空间中的低维流形上,而从噪声分布到数据流形的路径通常不是一条直线。

每一次去噪,模型实际上都在预测当前状态下应该前进的方向。这个方向可以理解为 probability flow 轨迹上的局部切线方向。

如果采样步数足够多,那么整个生成过程就像是在沿着一条高维曲线不断做局部线性近似。步数越多,近似越精细,最终越容易到达真实数据流形附近。

但如果我们试图一步生成,就相当于从初始噪声点直接沿着某个局部切线方向冲向目标区域。

问题是,切线只在局部有效。一旦步长太大,它就不再能代表整条曲线的真实方向。于是模型很可能直接偏离真实轨迹,最终生成质量下降。

所以,多步采样的本质是:

用多次局部近似,逼近从噪声分布到数据分布的复杂非线性路径。

这也是为什么 few-step diffusion 如此重要。它想解决的问题并不是简单地“少走几步”,而是要回答一个更本质的问题:

如果真实轨迹是一条复杂曲线,我们有没有办法用极少的步数,甚至一步,就逼近这条曲线的终点?

传统加速采样方法是怎么做的?

在进入 few-step distillation 之前,我们需要先回顾传统的采样加速方法。因为 few-step 方法本质上也是在回答同一个问题:如何更快地从噪声到达数据流形?

DDPM

首先从 DDPM 说起。

DDPM 将前向加噪过程建模为一个马尔可夫链。每一步都向数据中加入少量高斯噪声:

\[q(x_t \mid x_{t-1}) = \mathcal{N} \left( x_t; \sqrt{1-\beta_t}x_{t-1}, \beta_t I \right)\]

经过多步加噪之后,原始数据 $x_0$ 会逐渐变成接近标准高斯分布的噪声 $x_T$。

反向生成过程则学习:

\[p_\theta(x_{t-1} \mid x_t)\]

也就是从 $x_t$ 逐步恢复到 $x_{t-1}$。

DDPM 的一个核心假设是:当每一步加噪足够小时,反向过程也可以用一个条件高斯分布近似。这个假设让训练变得稳定,但代价是采样时通常需要几百甚至上千步。

如果直接大幅跳步,反向过程中的条件高斯近似就会变差,生成质量也会明显下降。

所以 DDPM 的特点很明确:

稳定,但慢。

DDIM

DDIM 的核心思想是:重新构造一个非马尔可夫的前向过程,使得它在每个时间点的边缘分布与 DDPM 保持一致。

也就是说,DDIM 不要求采样过程严格按照 DDPM 的马尔可夫链一步一步走,而是允许在保持边缘分布一致的前提下进行跳步采样。

如果把 DDPM 比作“朝一个方向前进一点,校正一下,再继续前进”,那么 DDIM 更像是“沿着相同的大方向,一次走得更远”。

这让 DDIM 可以用更少的步数完成采样。

但是,本质上 DDIM 的方向仍然来自局部预测。如果步长太大,轨迹依然可能偏离真实的数据流形。

所以 DDIM 的优点是简单有效,缺点是加速上限仍然有限。

ODE 视角

ODE 视角进一步统一了 Diffusion 的采样过程。

对于一个前向加噪 SDE,通常存在一个与其在所有时间点上具有相同边缘概率密度的确定性 ODE,这被称为 probability flow ODE:

\[d x = \left[ f(x,t) - \frac{1}{2} g(t)^2 \nabla_x \log p_t(x) \right] dt\]

其中,$\nabla_x \log p_t(x)$ 是 score function,表示当前噪声水平下数据分布对样本的梯度方向。

在常见的噪声预测参数化中,如果有:

\[x_t = \alpha_t x_0 + \sigma_t \epsilon\]

那么噪声预测网络 $\epsilon_\theta(x_t, t)$ 与 score function 的关系可以近似写作:

\[\nabla_{x_t} \log p_t(x_t) \approx -\frac{\epsilon_\theta(x_t,t)}{\sigma_t}\]

这样一来,采样过程就可以从随机微分方程转化为确定性的常微分方程求解问题。

也就是说,我们不再只是把 Diffusion 看成一个离散的去噪链条,而是可以把它看成一条连续的 probability flow。然后,我们就可以借助数值分析中的 ODE Solver 来加速采样。

高阶求解器在做什么?

简单来说,高阶求解器想做的事情是:让 probability flow 不再只沿着一阶切线方向前进,而是利用更高阶的信息去近似真实曲线。

从数值分析的角度看,它们本质上是在拟合更高阶的泰勒展开项,从而降低局部截断误差。这样在相同步数下,采样轨迹可以更接近真实的 probability flow。

常见的方法包括:

  • Euler Solver:最基础的一阶方法,只使用当前点的梯度信息。
  • Heun Solver:二阶方法,先预测一步,再用预测点的梯度进行校正。
  • DPM-Solver:将 Diffusion ODE 拆解为线性部分和非线性部分,线性部分用积分因子法精确求解,非线性部分用高阶数值方法近似。
  • PNDM:借鉴传统数值方法中的多步法思想,利用历史步的预测信息构造更高阶的更新方向。

高阶求解器的优势是:在相同步数下,它通常比一阶方法更准确。

但它也有明显的局限。

首先,高阶泰勒近似成立的前提是步长不能太大。如果 $\Delta t$ 太大,局部近似就会失效。

其次,当网络预测本身存在误差时,高阶项可能会放大误差,导致采样轨迹震荡,甚至偏离数据流形。

尤其是在采样刚开始和快结束的时候,probability flow ODE 的轨迹往往具有较高曲率。此时不同时间点的梯度变化较大,如果直接使用多步历史信息,就很容易不稳定。

因此,很多方法会在采样初期使用 warm-up,比如先用低阶方法稳定进入轨迹,然后再切换到高阶方法。

所以,高阶求解器解决的是:

如何在有限步数内,更准确地沿着曲线走。

但它仍然没有彻底跳出“沿着曲线一步步走”的框架。

换句话说,它是在让每一步走得更准,而不是从根本上解决“为什么必须走这么多步”的问题。

时间步调度:把步数集中在更重要的位置

在连续时间的 probability flow ODE 中,控制轨迹流动的核心变量之一是噪声标准差 $\sigma$。

Karras 等人提出的时间步调度方法,本质上是对噪声水平进行非均匀分配。其常见形式为:

\[\sigma_i = \left( \sigma_{\max}^{1/\rho} + \frac{i}{N-1} \left( \sigma_{\min}^{1/\rho} - \sigma_{\max}^{1/\rho} \right) \right)^\rho, \quad i = 0,1,\dots,N-1\]

其中,$N$ 是采样步数,$\sigma_{\max}$ 和 $\sigma_{\min}$ 分别表示最大与最小噪声水平,$\rho$ 控制时间步分布的形状。

简单来说,时间步调度的目标是:

在轨迹曲率高的地方分配更多步数,在轨迹近似线性的地方使用更大步长。

这与 Diffusion 的一个经验现象有关:在高噪声区域,样本还远离数据流形,整体结构比较模糊,轨迹曲率相对较小,因此可以使用较大的步长。

而在低噪声区域,样本逐渐靠近数据流形,细节开始形成,轨迹曲率通常更高,因此需要更密集的采样。

这个方法非常有效,但也有明显局限。

最主要的困难在于,调度参数 $\rho$ 如何选择并不简单。它与模型结构、训练时的加噪分布、数据类型以及采样器本身都有关系。一个在图像任务上有效的时间步调度,不一定能直接迁移到视频、3D 或机器人世界模型中。

我个人认为,时间步调度仍然是一个很 promising 的方向。它不改变模型本身,只改变采样路径的离散方式,因此代价低、兼容性好。

但从本质上看,无论怎么调度,低噪声阶段通常仍然需要比较细的采样。

也就是说,它依然是在贴着曲线走,只是把有限的步数分配得更聪明。

因此,它的加速上限可能不会特别高。相比于蒸馏方法中常见的 4-6 步,甚至 one-step generation,传统 solver 与时间步调度仍然存在明显差距。

小结

这一节我们主要回顾了 Diffusion 为什么有效,以及它为什么天然需要多步采样。

Diffusion 的核心优势在于,它把复杂的生成问题拆解成了多个局部去噪问题,因此训练稳定、生成质量高。但也正因为它依赖逐步去噪,从噪声分布到数据流形的路径往往需要许多离散步来近似,这就带来了严重的推理延迟问题。

传统的加速方法,比如 DDIM、高阶 ODE Solver 和时间步调度,本质上都在试图更高效地沿着 probability flow 前进:

  • DDIM 允许跳步采样;
  • ODE Solver 用更高阶的数值方法逼近采样轨迹;
  • 时间步调度把有限的步数集中到更重要的噪声区间。

但这些方法仍然没有完全摆脱一个共同限制:

它们依然假设生成过程需要沿着轨迹一步步走,只是走得更聪明、更快、更准。

而 few-step diffusion distillation 想解决的正是这个问题:

能不能让模型直接学会“大步走”,甚至“一步到位”?

下一节,我们将正式回顾 2022-2024 年间 few-step diffusion distillation 的代表性方法,看看它们是如何一步步突破采样步数限制的。