扩散模型（diffusion model）笔记

扩散模型（diffusion model）

扩散过程

对初始数据分布$x_0$~q(x)，不断添加高斯噪声，最终使数据分布$X_T$变成各项独立的高斯分布。

• 前向扩散过程的定义

  $q(x_t|x_{t-1})=N(x_t;\sqrt{1-{\beta }_t}{x}_{t-1},{\beta }_{t}I)$

  $q(x_{1:T}|x_0)={\prod }_{t=1}^{T}q(x_t|x_{t-1})$（马尔科夫链过程）

• 通过重参数化技巧，可以推导出任意时刻的$q(x_t)$，无需做迭代

  $x_t=\sqrt{{\alpha }_t}{x}_{t-1}+\sqrt{1-{\alpha }_t}{z}_{t-1}=...=\sqrt{{\overline{\alpha }}_t}{x}_0+\sqrt{1-{\overline{\alpha }}_t}z$

  其中$\overline{{\alpha }_t}={\prod }_{i=1}^T{\alpha }_i$；参数重整化体现为$\sqrt{{\alpha }_t}(\sqrt{{\alpha }_{t-1}}{x}_{t-2}+\sqrt{1-{\alpha }_{t-1}}{z}_{t-2})+\sqrt{1-{\alpha }_t}{z}_{t-1}$中，$\sqrt{{\alpha }_t-{\alpha }_t{\alpha }_{t-1}}{z}_{t-2}+\sqrt{1-{\alpha }_t}{z}_{t-1}$为两个正态分布叠加，可以重参数化为$\sqrt{1-{\alpha }_t{\alpha }_{t-1}}{\overline{z}}_{t-2}$

• 每个时间步所添加的噪声的标准差${\beta }_t$给定，且随t增大而增大

• 每个时间步所添加的噪声的均值与${\beta }_t$有关，为了使$x_T$稳定收敛到$N(0,1)$

• 由${\mathbf{x}}_{\mathbf{t}}=\sqrt{\overline{{\alpha }_{\mathbf{t}}}}{\mathbf{x}}_0+\sqrt{1-\overline{{\alpha }_{\mathbf{t}}}}\mathbf{z}$可得

  • $\mathbf{q}({\mathbf{x}}_{\mathbf{t}}|{\mathbf{x}}_0)=\mathbf{N}({\mathbf{x}}_{\mathbf{t}};\sqrt{{\overline{\alpha }}_{\mathbf{t}}}{\mathbf{x}}_0,(1-{\overline{\alpha }}_{\mathbf{t}})\mathbf{I})$

  • 随着不断加噪，$x_t$逐渐接近纯高斯噪声

  • ${\mathbf{x}}_0=\frac{1}{\sqrt{{\overline{\alpha }}_{\mathbf{t}}}}({\mathbf{x}}_{\mathbf{t}}-\sqrt{1-{\overline{\alpha }}_{\mathbf{t}}}{\mathbf{z}}_{\mathbf{t}})$

• 扩散过程中的后验条件概率$q(x_{t-1}|x_t,x_0)$可以用公式表达，即给定$x_t$、$x_0$，可计算出$x_{t-1}$

  假设${\beta }_t$足够小时，$\mathbf{q}({\mathbf{x}}_{\mathbf{t}-1}|{\mathbf{x}}_{\mathbf{t}},{\mathbf{x}}_0)=\mathbf{N}({\mathbf{x}}_{\mathbf{t}-1};\tilde{\mu }({\mathbf{x}}_{\mathbf{t}},{\mathbf{x}}_0),\widetilde{{\beta }_{\mathbf{t}}}\mathbf{I})$

  由高斯分布的概率密度函数$f(x)=\frac{1}{\sqrt{2\pi }\sigma }{e}^{-\frac{(x-\mu {)}^2}{2{\sigma }^2}}$和贝叶斯可得

  $\begin{array}{rl}q(x_{t-1}|x_t,x_0)& =q(x_t|x_{t-1},x_0)\frac{q(x_{t-1}|x_0)}{q(x_t|x_0)}\\ & \propto exp(-\frac{1}{2}((\frac{{\alpha }_t}{{\beta }_t}+\frac{1}{1-{\overline{\alpha }}_{t-1}}{x}_{t-1}^2)-(\frac{2\sqrt{{\alpha }_t}}{{\beta }_t}{x}_t+\frac{2\sqrt{{\overline{\alpha }}_t}}{1-{\overline{\alpha }}_t}{x}_0)x_{t-1}+C\end{array}$

  由二次函数的均值和方差计算可得

  $\tilde{\mu }(x_t,x_0)=\frac{\sqrt{{\alpha }_t}(1-{\overline{\alpha }}_{t-1})}{1-{\overline{\alpha }}_t}{x}_t+\frac{\sqrt{{\overline{\alpha }}_{t-1}{\beta }_t}}{1-{\overline{\alpha }}_t}{x}_0$

  $\widetilde{{\beta }_t}=\frac{1-{\overline{\alpha }}_{t-1}}{1-{\overline{\alpha }}_t}{\beta }_t$（DDPM作者使用$\widetilde{{\beta }_{\mathbf{t}}}={\beta }_{\mathbf{t}}$，认为两者结果近似）

  将$x_0$的公式代入得（$z_t\sim N(0,I)$使用了重参数化）

  ${\tilde{\mu }}_{\mathbf{t}}=\frac{1}{\sqrt{{\alpha }_{\mathbf{t}}}}({\mathbf{x}}_{\mathbf{t}}-\frac{{\beta }_{\mathbf{t}}}{\sqrt{1-{\overline{\alpha }}_{\mathbf{t}}}}{\mathbf{z}}_{\mathbf{t}})$

  即在$x_0$条件下，后验条件概率分布可通过$x_t$和$z_t$计算得到

逆扩散过程

从高斯噪声$x_T$中逐步还原出原始数据$x_0$。马尔科夫链过程。

• ${\mathbf{p}}_{\theta }({\mathbf{x}}_{\mathbf{t}-1}|{\mathbf{x}}_{\mathbf{t}})=\mathbf{N}({\mathbf{x}}_{\mathbf{t}-1};{\mu }_{\theta }({\mathbf{x}}_{\mathbf{t}},\mathbf{t}),{\mathbf{\Sigma }}_{\theta }({\mathbf{x}}_{\mathbf{t}},\mathbf{t})$

  $p_{\theta }(x_{0:T})=p(x_T){\prod }_{t-1}^T{p}_{\theta }(x_{t-1}|x_t)$

目标函数

对负对数似然$L=E_{q(x_0)}[-log{p}_{\theta }(x_0)]$使用变分下限（VLB），并进一步推导化简得到最终loss

• ${\mathbf{L}}_{\mathbf{t}}^{\mathbf{simple}}={\mathbf{E}}_{\mathbf{t},{\mathbf{x}}_0,ϵ}[||ϵ-{ϵ}_{\theta }(\sqrt{{\overline{\alpha }}_{\mathbf{t}}}{\mathbf{x}}_0+\sqrt{1-{\overline{\alpha }}_{\mathbf{t}}}ϵ,\mathbf{t})|{|}^2]$

• 在推导的过程中，loss转换为$q(x_{t-1}|x_t,x_0)=N(x_{t-1};\tilde{\mu }(x_t,x_0),\widetilde{{\beta }_t}I)$与$p_{\theta }(x_{t-1}|x_t)=N(x_{t-1};{\mu }_{\theta }(x_t,t),{\Sigma }_{\theta }(x_t,t)$两个高斯分布之间的KL散度，将$\mu$与$x_t$的公式代入将loss转化为$ϵ$、$x_0$、$t$的公式

• DDPM作者采用了预测随机变量（噪声）法，并不直接预测后验分布的期望值或原始数据

• DDPM作者将方差${\Sigma }_{\theta }(x_t,t)$用给定的${\beta }_t$或$\widetilde{{\beta }_t}$代替，训练参数只存在均值中，为了使训练更加稳定

训练过程

1. 给出原始数据$x_0\sim q(x_0)$

2. 设定$t\sim Uniform(1,...,T)$

3. 从标准高斯分布采样一个噪声$ϵ\sim N(0,I)$

4. 采用梯度下降法优化目标函数$||ϵ-{ϵ}_{\theta }(\sqrt{{\overline{\alpha }}_t}{x}_0+\sqrt{1-{\overline{\alpha }}_t}ϵ,t)||$

推断过程

1. 每个时间步通过$x_t$和$t$计算$p_{\theta }(x_{t-1}|x_t)$

   均值${\mu }_{\theta }(x_t,t)=\frac{1}{\sqrt{{\alpha }_t}}(x_t-\frac{{\beta }_t}{\sqrt{1-{\overline{\alpha }}_t}}{z}_{\theta }(x_t,t))$，方差${\Sigma }_{\theta }(x_t,t)=\widetilde{{\beta }_t}={\beta }_t$

   $p_{\theta }(x_{t-1}|x_t)=N(x_{t-1};{\mu }_{\theta }(x_t,t),{\Sigma }_{\theta }(x_t,t))$

2. 通过重参数从$p_{\theta }(x_{t-1}|x_t)$中采样得到$x_{t-1}$

3. 通过不断迭代最终得到$x_0$

代码实现

• 定义时间步数、${\beta }_t$、$\sqrt{\overline{{\alpha }_t}}$等公式计算中需要用到的常量

  ${\mathbf{x}}_{\mathbf{t}}=\sqrt{\overline{{\alpha }_{\mathbf{t}}}}{\mathbf{x}}_0+\sqrt{1-\overline{{\alpha }_{\mathbf{t}}}}\mathbf{z}$

  ${\mu }_{\theta }=\frac{1}{\sqrt{{\alpha }_{\mathbf{t}}}}({\mathbf{x}}_{\mathbf{t}}-\frac{{\beta }_{\mathbf{t}}}{\sqrt{1-{\overline{\alpha }}_{\mathbf{t}}}}{\mathbf{z}}_{\theta }({\mathbf{x}}_{\mathbf{t}},\mathbf{t}))$

  ${\mathbf{\Sigma }}_{\theta }({\mathbf{x}}_{\mathbf{t}},\mathbf{t})=\widetilde{{\beta }_{\mathbf{t}}}={\beta }_{\mathbf{t}}$

  DDPM论文中作者将时间步数$T$设置为1000，${\beta }_t$为0.0001到0.02之间的线性插值

  num_timesteps = 1000
  schedule_low = 1e-4
  schedule_high = 0.02
  betas = torch.tensor(np.linspace(schedule_low, schedule_high, num_timesteps), dtype=torch.float32)
  
  alphas = 1 - betas
  alphas_cumprod = np.cumprod(alphas)
  sqrt_alphas_cumprod = np.sqrt(alphas_cumprod)
  sqrt_one_minus_alphas_cumprod = np.sqrt(1 - alphas_cumprod)
  reciprocal_sqrt_alphas = np.sqrt(1 / alphas)
  betas_over_sqrt_one_minus_alphas_cumprod = (betas / sqrt_one_minus_alphas_cumprod)
  sqrt_betas = np.sqrt(betas)
  

• 前向扩散过程

  ${\mathbf{x}}_{\mathbf{t}}=\sqrt{\overline{{\alpha }_{\mathbf{t}}}}{\mathbf{x}}_0+\sqrt{1-\overline{{\alpha }_{\mathbf{t}}}}\mathbf{z}$

  $||ϵ-{ϵ}_{\theta }(\sqrt{{\overline{\alpha }}_{\mathbf{t}}}{\mathbf{x}}_0+\sqrt{1-{\overline{\alpha }}_{\mathbf{t}}}ϵ,\mathbf{t})||$

  def forward_diffusion_process(model, x0, num_timesteps, sqrt_alphas_cumprod, sqrt_one_minus_alphas_cumprod):
      batch_size = x0.shape[0]
      t = torch.randint(0, num_timesteps, size=(batch_size,))
      noise = torch.randn_like(x0)
      xt = sqrt_alphas_cumprod[t] * x0 + sqrt_one_minus_alphas_cumprod[t] * noise
      estimated_noise = model(xt, t)
      loss = (noise - estimated_noise).square().mean()
      return loss
  

• 逆向扩散过程

  ${\mathbf{p}}_{\theta }({\mathbf{x}}_{\mathbf{t}-1}|{\mathbf{x}}_{\mathbf{t}})=\mathbf{N}({\mathbf{x}}_{\mathbf{t}-1};{\mu }_{\theta }({\mathbf{x}}_{\mathbf{t}},\mathbf{t}),{\mathbf{\Sigma }}_{\theta }({\mathbf{x}}_{\mathbf{t}},\mathbf{t}))$

  ${\mu }_{\theta }({\mathbf{x}}_{\mathbf{t}},\mathbf{t})=\frac{1}{\sqrt{{\alpha }_{\mathbf{t}}}}({\mathbf{x}}_{\mathbf{t}}-\frac{{\beta }_{\mathbf{t}}}{\sqrt{1-{\overline{\alpha }}_{\mathbf{t}}}}{\mathbf{z}}_{\theta }({\mathbf{x}}_{\mathbf{t}},\mathbf{t}))$

  ${\mathbf{\Sigma }}_{\theta }({\mathbf{x}}_{\mathbf{t}},\mathbf{t})=\widetilde{{\beta }_{\mathbf{t}}}={\beta }_{\mathbf{t}}$

  def reverse_diffusion_process(model, shape, num_timesteps, reciprocal_sqrt_alphas, betas_over_sqrt_one_minus_alphas_cumprod, sqrt_betas):
      current_x = torch.randn(shape)
      x_seq = [current_x]
      for t in reversed(range(num_timesteps)):
          current_x = sample(model, current_x, t, shape[0], reciprocal_sqrt_alphas, betas_over_sqrt_one_minus_alphas_cumprod, sqrt_betas)
          x_seq.append(current_x)
      return x_seq
  
  def sample(model, xt, t, batch_size, reciprocal_sqrt_alphas, betas_over_sqrt_one_minus_alphas_cumprod, sqrt_betas):
      ts = torch.full([batch_size, 1], t)
      estimated_noise = model(xt, ts)
      mean = reciprocal_sqrt_alphas[ts] * (xt - betas_over_sqrt_one_minus_alphas_cumprod[ts] * estimated_noise)
      if t > 0:
          z = torch.randn_like(xt)
      else:
          z = 0
      sample = mean + sqrt_betas[t] * z
      return sample