加速了 DDIM 的采样过程

Method

假设我们训好了 $\epsilon_{\theta}$, DDIM sample 为

\[x_{t-1}=\sqrt{\frac{\alpha_{t-1}}{\alpha_t}}x_t + (\sqrt{1-\alpha_{t-1}}-\sqrt{\frac{\alpha_{t-1}(1-\alpha_t)}{\alpha_t}})\epsilon_{\theta}(x_t, t)\]

我们可以把 $x^{T-k}$ 直接写成关于 $x_T, \ldots, x_{T-k+1}$ 的函数

\[x^{T-k} = \sqrt{\frac{\alpha_{T-k}}{\alpha_T}}x_T + \sum_{t=T-k}^{T-1} c(t+1, k) \epsilon_{\theta}(x_{t+1}, t+1)\]

其中 $c$ 是系数，不重要. 这样, 我们可以把这个方程写成一个方程组

\[\begin{bmatrix} x_{T-1} \\ x_{T-2} \\ \vdots \\ x_{0} \\ \end{bmatrix} = \begin{bmatrix} h(x_T) \\ h(x_{T-1:T}) \\ \vdots \\ h(x_{1:T}) \\ \end{bmatrix}\]

其中 $h$ 是一个函数. 这样的好处在于, 我们发现这样的一次 $h$ 的 forward 过程是可以并行的！这样会在 GPU 上更快的采样.

$x_{0:T-1}$ 是在 given $x_T$ 的条件下方程的不动点 (就是所谓 Deep Equilibrium), 我们可以使用 black-box solver (e.g. Anderson) 来在远小于 $T$ 步的时间中求出不动点. (这里, Anderson 是不需要 $h$ 的梯度的.)

  • 需要注意的是, 不考虑并行加速的话, 这个方法的 NFE 是明显更大的. 所以聪明的地方就在于利用并行 forward 来加速采样.

Application

给定图片 $x$, 如何得到对应的 $\epsilon$?

也可以用类似的方法加速. 论文给了一种我没看懂的梯度下降的方法.