TiankaiHang / blog

For self learning
3 stars 0 forks source link

Learning-based Methods for Diffusion Distillation - 基于学习的方法做加速采样 #1

Open TiankaiHang opened 3 months ago

TiankaiHang commented 3 months ago

CONSISTENCY TRAJECTORY MODELS: LEARNING PROBABILITY FLOW ODE TRAJECTORY OF DIFFUSION

ICLR 2024

Motivation

主要在 Consistency Models 的基础上进行改进。Consistency Models 的一个缺点是直接预测 $\mathbf{x} _0$, 没法很好的做trade-off, 即我后期提升采样步数,结果也不见得变得多好。具体可见下面的这个示意图。

于是这篇文章的意图就比较明显,就是说我不止要对真实数据点 $\mathbf{x} _0$ 进行预测, 还要对中间的轨迹也有所监督。

image

前置知识

Diffusion的采样过程可以想象是在一条轨迹上走,从噪声点 $\mathbf{x} _T$ 走到 $\mathbf{x} _0$,我们训练的模型实际上是提供了某个点的方向。Yang Song的论文里面将deterministic的方向过程定义为了一个 PF ODE,

$$ \frac{\mathrm{d} \mathbf{x} _t}{\mathrm{d} t} = -t \nabla \log p_t (\mathbf{x} _t) $$

EDM 框架中用Denoiser 表示

$$ \frac{\mathrm{d} \mathbf{x} _t}{\mathrm{d} t} = \frac{\mathbf{x} t - D {\phi} (\mathbf{x} _t , t)}{t} $$

生成过程实际上就是 $T \rightarrow 0$ 的积分。

具体方法

CTM predicts both infinitesimally small step jump and long step jump of the PF ODE trajectory.

首先定义一个从 $t \rightarrow s (s \leq t)$ 的解

$$ G(\mathbf{x} _t ,t,s) := \mathbf{x} t + \int {t}^s \frac{\mathbf{x} _u - \mathbb{E}[\mathbf{x}|\mathbf{x} _u ]}{u} \mathrm{d} u $$

对连续形式实际训练比较困难,受Euler solver启发表示为

$$ G(\mathbf{x} _t ,t,s) = \frac{s}{t} \mathbf{x} _t + (1 - \frac{s}{t}) g(\mathbf{x} _t ,t,s) $$

其中 $g(\mathbf{x} _t ,t,s) = \mathbf{x} t + \frac{t}{t-s} \int {t}^s \frac{\mathbf{x} _u - \mathbb{E}[\mathbf{x}|\mathbf{x} _u ]}{u} \mathrm{d} u$.

我们预测的target 是

$$ G _{\theta} (\mathbf{x} _t ,t,s) = \frac{s}{t} \mathbf{x} t + (1 - \frac{s}{t}) g {\theta} (\mathbf{x} _t ,t,s) $$

算一下可以发现 $g(\mathbf{x} _t ,t,s)$ 在 s 无限趋向于 t 的时候,预测的就是 $\mathbf{x} _0$. 此外,用泰勒展开,可以发现 $g(\mathbf{x} _t ,t,s)$ 就是 denoiser function 加上残差项 $\mathcal{O}(|t-s|)$.

具体技术

image

image

两者之间计算损失。

总的损失

image

image

image

实验