diffusion model 学习笔记

条件引导的 diffusion

对于无条件的DDPM 而言

$$p(x_t|x_0)\sim \mathcal{N}(\sqrt{\overline{{\alpha }_t}}{x}_0,1-\overline{{\alpha }_t}\cdot \mathrm{I})$$

可以得到

$$\begin{array}{rl}\log p(x_t|x_0)& =-\frac{1}{2}\frac{(x_t-\sqrt{\overline{{\alpha }_t}}{x}_0{)}^2}{1-\overline{{\alpha }_t}}\end{array}$$

计算其 score func, 可以得到

$$\begin{array}{rl}{\nabla }_x\log p(x_t|x_0)& =-\frac{(x_t-\sqrt{\overline{{\alpha }_t}}{x}_0)}{1-\overline{{\alpha }_t}}\\ & =-\frac{ϵ}{\sqrt{1-\overline{{\alpha }_t}}}\\ & \approx -\frac{{ϵ}_{\theta }(x_t,t)}{\sqrt{1-\overline{{\alpha }_t}}}\end{array}$$

也就是,我们训练的网络conditionalUnet 输出的是噪声的估计, 这个估计的噪声${ϵ}_{\theta }(x_t,t)$可以用来计算当前数据点在整个概率空间中的score-func: ${\nabla }_x\log p(x_t|x_0)$.

---

对于 conditional Diffusion. 其 score func 可以写成

$$\begin{array}{rl}{\nabla }_{x_t}\log p(x_t|y)& ={\nabla }_{x_t}\log \left(\frac{p(x_t)\cdot p(y|x_t)}{p(y)}\right)\\ & ={\nabla }_{x_t}\log p(x_t)+{\nabla }_{x_t}\log p(y|x_t)-\underset{=0和x_t无关}{\underbrace{{\nabla }_{x_t}\log p(y)}}\\ & =\underset{\mathrm{unconditional}\mathrm{score}}{\underbrace{{\nabla }_{x_t}\log p(x_t)}}+\underset{\mathrm{adversial}\mathrm{gradient}}{\underbrace{{\nabla }_{x_t}\log p(y|x_t)}}\end{array}$$

现在,我们需要估计 ${\nabla }_{x_t}\log p(x_t|y)$ 中的 ${ϵ}^{\prime }$, 然后使用 DDPM/DDIM 采样即可. 不妨设

$$\begin{array}{rl}{\nabla }_{x_t}\log p(x_t|y)& =-\frac{{ϵ}^{\prime }}{\sqrt{1-\overline{{\alpha }_t}}}\end{array}$$

则我们有

$$\begin{array}{rl}{\nabla }_{x_t}\log p(x_t|y)& ={\nabla }_{x_t}\log p(x_t)+\underset{令=g}{\underbrace{{\nabla }_{x_t}\log p(y|x_t)}}\\ -\frac{{ϵ}^{\prime }}{\sqrt{1-\overline{{\alpha }_t}}}& =-\frac{ϵ}{\sqrt{1-\overline{{\alpha }_t}}}+g\\ 可得& :{ϵ}^{\prime }=ϵ-\sqrt{1-{\bar{\alpha }}_t}\cdot g\end{array}$$

即我们在无条件模型 DDPM 估计的噪声中添加一个微小的扰动($\sqrt{1-{\bar{\alpha }}_t}\cdot g$),就可以作为条件模型的噪声估计.

---

• $g={\nabla }_{x_t}\log p(y|x_t)$ 的含义:

  条件概率的梯度: ${\nabla }_{x_t}\log p(y|x_t)$ 表示的是在已知$x_t$的情况下，微小变化$x_t$如何影响条件$y$的对数概率。这是一个 向量场，指向增加条件$y$出现概率的方向。

• 如何得到 g

  假设我现在已经有了一个回归模型, 即 $y=f(x_t,...)$.

  输入一个数据 $x_t$, 返回其对应一个 logit 值.

  ## 输入: x_t, 模型 f
  x_t = torch.tensor(x_t, requires_grad=True)
  y = f(x_y)            ## pytorch model 预测其结果
  log_p = torch.log(y)  ## 计算 log 概率
  log_p.backward()      ## 计算梯度
  grad_x_t = x_t.grad   ## 获取梯度
  

---

Langevin Dynamics 采样

首先,假设我们已经有了一个训练好的score-func($s_{\theta }(x)$). 已经接近于真实的 $s(x)$, 即 $s_{\theta }(x)={\nabla }_x\log p_{\theta }(x)$, 其中 $p_{\theta }(x)\approx p(x)$. 现在, 我们需要利用 $s_{\theta }(x)$ 对 x 进行采样,使得 $x\sim p_{\theta }(x)$

郎之万公式: 描述了粒子做随机布朗运动 (粒子位置随时间变化的关系), 是一种 SDE, 描述了 由梯度力(即. $U(x(t)$)驱动并受到随机噪声(即. $Z_t$)影响的系统的时间演化.

$$dX(t)=-{\nabla }_{x}U(x(t))\cdot dt+\sigma \sqrt{dt}\cdot Z_t$$

$X(t)$:

---

Diffuser 源码解读

DDPMSchedule

1. $\alpha ,\bar{\alpha },\beta$ 的关系和计算

$$\begin{array}{rl}{x}_t& =\sqrt{{\alpha }_t}\cdot x_{t-1}+\sqrt{1-{\alpha }_t}\cdot {ϵ}_t\\ & =\sqrt{\prod _{i=1}^t{\alpha }_i}\cdot x_0+\sqrt{1-\prod _{i=1}^t{\alpha }_i}\cdot {ϵ}_0\\ & =\sqrt{\overline{{\alpha }_t}}\cdot x_0+\sqrt{1-\overline{{\alpha }_t}}\cdot {ϵ}_0\\ {\beta }_t& =1-{\alpha }_t\end{array}$$

• 在 DDPMSchedule 中， 先计算出 ${\beta }_t$ 的值，然后利用 ${\alpha }_t=1-{\beta }_t$ 计算剩下的系数。

• self.beta 的计算有各种不同的方法，

  • trained_betas： 需要传入自定义的beta
  • linear
  • squaredcos_cap_v2
  • …

• alpha 和 alpha_cumprod 的计算

  self.alphas = 1.0 - self.betas
  self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
  # 累积乘积（cumulative product）
  

• self.betas: 是扩散过程中“噪声量”的度量。控制着每一步扩散过程中加入的噪声量的大小.

• self.alphas: 每一步中保留的原始信号的比例。

上图是 DDPM 中的 $\alpha$, $\beta$, $\bar{\alpha }$ 的关系。

2. 计算前向时刻的采样: $x_t\to x_{t-1}$

$$\begin{array}{rl}q(x_{t-1}|x_t,x_0)& =\frac{q(x_t|x_{t-1},x_0)q(x_{t-1}|x_0)}{q(x_t|x_0)}\\ & =\frac{\mathcal{N}(x_t;\sqrt{{\alpha }_t}{x}_{t-1},(1-{\alpha }_t)\text{I})\mathcal{N}(x_{t-1};\sqrt{{\bar{\alpha }}_{t-1}}{x}_0,(1-{\bar{\alpha }}_{t-1})\text{I})}{\mathcal{N}(x_t;\sqrt{{\bar{\alpha }}_t}{x}_0,(1-{\bar{\alpha }}_t)\text{I})}\\ & \propto \mathcal{N}(x_{t-1};\underset{{\mu }_q(x_t,x_0)}{\underbrace{\frac{\sqrt{{\alpha }_t}(1-{\bar{\alpha }}_{t-1})x_t+\sqrt{{\bar{\alpha }}_{t-1}}(1-{\alpha }_t)x_0}{1-{\bar{\alpha }}_t}}},\underset{{\Sigma }_q(t)}{\underbrace{\frac{(1-{\alpha }_t)(1-{\bar{\alpha }}_{t-1})}{1-{\bar{\alpha }}_t}\text{I}}})\end{array}$$

通过上面的公式，我们可以知道，当知道 $x_0$ 和 $x_t$, 既可以得到 $x_{t-1}$. 但是， 我们没有办法获取真实的 $x_0$, 所以只能估计 $x_t$ 对应的 ${\hat{x}}_0$ 是什么样子的。

$$\begin{array}{rl}{\mu }_q(x_t,x_0)& =\frac{\sqrt{{\alpha }_t}(1-{\bar{\alpha }}_{t-1})x_t+\sqrt{{\bar{\alpha }}_{t-1}}(1-{\alpha }_t)x_0}{1-{\bar{\alpha }}_t}\\ {\Sigma }_q(t)& =\frac{(1-{\alpha }_t)(1-{\bar{\alpha }}_{t-1})}{1-{\bar{\alpha }}_t}\text{I}\end{array}$$

• 第一种： 先利用预测的噪声 ${ϵ}_t$ 估计 ${\hat{x}}_0$. 然后利用 $x_t$ 和 ${\hat{x}}_0$ 估计 $x_{t-1}$

$$\begin{array}{rl}{\hat{x}}_0^{t-1}& =\frac{x_t-\sqrt{1-{\bar{\alpha }}_t}\cdot {ϵ}_{\theta }(x_t,t)}{\sqrt{{\bar{\alpha }}_t}}\end{array}$$

${\hat{x}}_0^{t-1}$ 表示第$t-1$时刻估计的 $x_0$. 基本上,在采样一半时,基本上预测的 $x_0$ 就和真实图片差不多了.

接下来计算 $x_{t-1}$ 的均值和方差. 方差非常重要,不能忽略
$$\begin{array}{rl}\mathrm{mean}(x_{t-1})& =\frac{\sqrt{{\alpha }_t}(1-{\bar{\alpha }}_{t-1})}{1-{\bar{\alpha }}_t}\cdot x_t+\frac{\sqrt{{\bar{\alpha }}_{t-1}}(1-{\alpha }_t)}{1-{\bar{\alpha }}_t}\cdot {\hat{x}}_0^{t-1}\\ \mathrm{var}(x_{t-1})& =\frac{(1-{\alpha }_t)(1-{\bar{\alpha }}_{t-1})}{1-{\bar{\alpha }}_t}\cdot \text{I}\\ x_{t-1}& =\mathrm{mean}(x_{t-1})+\mathrm{var}(x_{t-1})\end{array}$$

for i, t in enumerate(ddpm_scheduler.timesteps):alpha_t = alphas[t] alpha_t_bar = alphas_cumprod[t]alpha_t_bar_prev = alphas_cumprod[t - 1] if t - 1 >= 0 else torch.tensor(1.0)# -------------## 1. x_0 和 x_t的系数pred_ori_sample_coeff = torch.sqrt(alpha_t_bar_prev) * (1-alpha_t) / (1-alpha_t_bar)current_sample_coeff = torch.sqrt(alpha_t) * (1-alpha_t_bar_prev) / (1-alpha_t_bar)## 2. 预测 x_0noise_pred = ddpm_noise_model(x_t, t)['sample']est_x_0 = (x_t - torch.sqrt(1-alpha_t_bar) * noise_pred) / torch.sqrt(alpha_t_bar)## 3. clip x_0est_x_0 = est_x_0.clamp(-1, 1)## 4. 预测 x_{t-1}x_t_prev = pred_ori_sample_coeff * est_x_0 + current_sample_coeff * x_t## 5. 添加噪声if t > 0:std_var_noise = torch.randn_like(x_t).to(device)x_t_var_coeff = (1-alpha_t) * (1-alpha_t_bar_prev) / (1-alpha_t_bar)x_t_var_coeff = torch.sqrt(x_t_var_coeff)x_t_prev = x_t_prev + x_t_var_coeff * std_var_noisex_t = x_t_prev