Hexo

凡事预则立,不预则废


  • Home

  • Tags

  • Archives

  • Navigation

  • Search

DL——DDPM

文本介绍DDPM(Denoising Diffusion Probabilistic Models)的理论介绍

  • 参考链接:
    • 原始论文:Denoising Diffusion Probabilistic Models, NeurIPS 2020
    • What are Diffusion Models?:Lilian Weng的博客
    • 54、Probabilistic Diffusion Model概率扩散模型理论与完整PyTorch代码详细解读
      • 视频中有部分小bug,评论区已修复
    • 一文解决你关于扩散模型ddpm的所有疑惑,回答了很多疑惑
      • 一个视频看懂扩散模型DDPM原理推导|AI绘画底层模型,

推导过程

这部分推导参考自由浅入深了解Diffusion Model - ewrfcas的文章 - 知乎和What are Diffusion Models?

Diffusion前向过程

前向过程的整体描述
  • 给定真实图片 \(x_0 \sim q(x_0)\),Diffusion前向过程通过 \(T\) 次对其添加高斯噪声,得到 \(x_T\),第 \(t\) 次加噪后得到的样本为 \(x_t\)
  • 加噪过程是增加一个服从高斯分布的误差,每一步中高斯分布方差的超参数是提前设定的 \(\{\beta_t \in (0,1)\}_{t=1}^T\)
  • 前向过程由于每个时刻 \(t\) 只与 \(t-1\) 时刻有关,所以也可以看做马尔可夫过程
    $$
    \begin{align}
    q(x_t|x_{t-1}) = \mathcal{N}(x_t; \sqrt{1-\beta_t}x_{t-1}, \beta_t\boldsymbol{I}) \\
    q(x_{1:T}|x_0) = \prod_{t=1}^T q(x_t|x_{t-1}) \\
    \end{align}
    $$
  • 实际实现中 \(\{\beta_t \in (0,1)\}_{t=1}^T\) 中的元素是递增的,即 \(\beta_1 < \beta_2 < \cdots < \beta_T\),DDPM实验时设置了 \(T=1000\),且 \(\beta_t\) 是 \([0.0001, 0.02]\) 之间的线性的固定值,即 \(\beta_1 = 0.0001 \ \text{to} \ \beta_T = 0.02\)
  • 加噪过程中,随着加噪时间步骤 \(t\) 的增加, \(x_t\) 越来越接近纯随机噪声,且 \(T \rightarrow \infty\) 时, \(x_T\) 完全是噪声
可以用 \(x_0\) 和方差 \(\beta\) 来表示 \(x_t\)
  • 即 \(\forall t, \ x_t = f(x_0, \{\beta_i\}_{i=1}^t)\) \(x_t\)
  • 引入新变量 \(\alpha_t = 1- \beta_t\),且 \(\bar{\alpha}_t = \prod_{i=1}^t\alpha_i\) (理解: \(\alpha_t\) 和 \(\bar{\alpha}_t\) 都随着扩散时间步逐步变小)
  • 将 \(x_t\) 进行展开可得得到用 \(x_0\) 和方差 \(\beta\) 来表示 \(x_t\):
    $$
    \begin{align}
    \mathbf{x}_t
    &= \sqrt{\alpha_t}\mathbf{x}_{t-1} + \sqrt{1 - \alpha_t}\mathbf{\epsilon}_{t-1} \quad \quad \quad \quad \quad \quad \quad \quad \text{ ;where } \mathbf{\epsilon}_{t-1}, \mathbf{\epsilon}_{t-2}, \dots \sim \mathcal{N}(\mathbf{0}, \boldsymbol{I}) \\
    &= \sqrt{\alpha_t}(\sqrt{\alpha_{t-1}}\mathbf{x}_{t-2} + \sqrt{1 - \alpha_{t-1}}\mathbf{\epsilon}_{t-2}) + \sqrt{1 - \alpha_t}\mathbf{\epsilon}_{t-1} \\
    &= \sqrt{\alpha_t \alpha_{t-1}}\mathbf{x}_{t-2} + (\sqrt{\alpha_t(1 - \alpha_{t-1})}\mathbf{\epsilon}_{t-2} + \sqrt{1 - \alpha_t}\mathbf{\epsilon}_{t-1}) \\
    &= \sqrt{\alpha_t \alpha_{t-1}} \mathbf{x}_{t-2} + \sqrt{1 - \alpha_t \alpha_{t-1}} \bar{\mathbf{\epsilon}}_{t-2} \quad \quad \quad \quad \text{ ;where } \bar{\mathbf{\epsilon}}_{t-2} \text{ merges two Gaussians (*).} \\
    &\quad \vdots \\
    &= \sqrt{\bar{\alpha}_t}\mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t}\mathbf{\epsilon} \\
    q(\mathbf{x}_t \vert \mathbf{x}_0) &= \mathcal{N}(\mathbf{x}_t; \sqrt{\bar{\alpha}_t} \mathbf{x}_0, (1 - \bar{\alpha}_t)\boldsymbol{I})
    \end{align}
    $$
    • 推导中两个高斯分布的合并使用使用到: \(\mathcal{N}(\mu_1,\sigma_1^2\boldsymbol{I}) + \mathcal{N}(\mu_2,\sigma_2^2\boldsymbol{I}) \sim \mathcal{N}\Big((\mu_1+\mu_2),(\sigma_1^2+\sigma_2^2)\boldsymbol{I}\Big)\)
    • 显然,当 \(T \rightarrow \infty\) 时, \(x_T\) 完全是噪声,因为 \(\sqrt{\bar{\alpha}_T} \rightarrow 0,\ \sqrt{1-\bar{\alpha}_T} \rightarrow 1\)
    • Merge后的噪声 \(\epsilon\) 也是高斯噪声,且服从标准正态分布 \(\epsilon \sim \mathcal{N}(0,\boldsymbol{I})\)
  • 同时,也可以用 \(x_t\) 和方差 \(\beta\) 来表示 \(x_0\):
    $$ x_0 = \frac{1}{\sqrt{\bar{\alpha}_t}}(x_t - \sqrt{1-\bar{\alpha}_t}\epsilon) $$
  • 问题:为什么前向过程的需要使用 \(\sqrt{1-\beta_t}x_{t-1}\) 来表示均值?
    • 回答:因为这样可以使得 \(x_T\) 的均值中不含有原始样本 \(x_0\) 的信息,也就是说,这样使得 \(x_T\) 接近噪声
    • 补充问题:使用其他小于0的数 \(k_t\) 替代 \(\sqrt{1-\beta_t}\) 吗?答案是可以的,但是使用 \(\sqrt{1-\beta_t}\) 是为了满足 \((\sqrt{1-\beta_t})^2+\beta_t = 1\) 这样的形式,方便后续的推导,如果不使用这样的形式,推导起来会更复杂些,详情见苏神的博客生成扩散模型漫谈(一):DDPM = 拆楼 + 建楼

Diffusion逆向过程

  • 符号定义 :为了与其他文章常用习惯保持一致,论文在前向过程中使用 \(q(x_t|x_0)\),逆向过程中 \(q(x_0|x_t)\),需要按照区分前后下标大小来区分前向过程和后向过程

  • 逆向过程的目标 :逆向过程是前向过程的相反过程,其目标是从随机噪声 \(x_T \sim \mathcal{N}(0, \boldsymbol{I})\) 中一步步恢复得到原始样本 \(x_0\)

    • 思考 :如果能够得到 \(q(x_{t-1}|x_t)\),则可以一步步恢复出来原始样本 \(x_0\)
    • 注:可以证明,如果 \(q(x_t|x_{t-1})\) 是高斯分布且方差 \(\beta_t\) 足够小,那么 \(q(x_{t-1}|x_t)\) 也还是一个高斯分布【TODO: On the theory of stochastic processes, with particular reference to applications.?】
  • 逆向过程的拟合 :我们使用 \(p_\theta(x_{t-1}|x_t)\) 来拟合逆向分布,由于 \(p_\theta(x_{t-1}|x_t)\) 是高斯分布,所以可以拆解成建模均值 \(\mu_\theta(x_t, t)\) 和方差 \(\Sigma_\theta(x_t, t)\) 即可
    $$
    \begin{align}
    p_\theta(\mathbf{x}_{0:T}) &= q(\mathbf{x}_T) \prod^T_{t=1} p_\theta(\mathbf{x}_{t-1} \vert \mathbf{x}_t) \\
    p_\theta(\mathbf{x}_{t-1} \vert \mathbf{x}_t) &= \mathcal{N}(\mathbf{x}_{t-1}; \mathbf{\mu}_\theta(\mathbf{x}_t, t), \mathbf{\Sigma}_\theta(\mathbf{x}_t, t))
    \end{align}
    $$

  • 难题 :实际上,我们是无法直接得到 \(q(x_{t-1}|x_t)\) 的(理解:不知道 \(x_0\) 的情况下,噪声并不知道要朝着什么方向去生成 \(x_0\) )

  • 如果 \(x_0\) 已知,结合 \(x_t \sim q(x_t|x_0) = \mathcal{N}(\mathbf{x}_t; \sqrt{\bar{\alpha}_t} \mathbf{x}_0, (1 - \bar{\alpha}_t)\boldsymbol{I})\),则可以通过贝叶斯公式求得 \(q(x_{t-1}|x_t,x_0)\):
    $$q(x_{t-1}|x_t,x_0) = \mathcal{N}(x_{t-1};\tilde{\mu}(x_t, x_0), \tilde{\beta}_t\boldsymbol{I})$$

    • \(\tilde{\mu}(x_t,x_0)\) 是均值,是 \(x_t,x_0\) 的函数,而 \(\tilde{\beta}_t\) 是方差,与 \(x_t,x_0\) 无关,仅与时间片有关
  • 目标是均值和方差 :至此,我们的目标是求得 \(\tilde{\mu}(x_t, x_0)\) 和 \(\tilde{\beta}_t\),为了这个目标,我们首先将 \(q(x_{t-1}|x_t,x_0)\) 展开:
    $$
    \begin{aligned}
    q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0)
    &= \frac{q(x_{t-1},x_t,x_0)}{q(x_t, x_0)} \\
    &= \frac{q(x_t|x_{t-1},x_0) \cdot q(x_{t-1},x_0)}{q(x_t, x_0)} \\
    &= q(x_t|x_{t-1},x_0) \frac{q(x_{t-1}|x_0) \cdot q(x_0)}{q(x_t| x_0) \cdot q(x_0)} \\
    &= q(\mathbf{x}_t \vert \mathbf{x}_{t-1}, \mathbf{x}_0) \frac{ q(\mathbf{x}_{t-1} \vert \mathbf{x}_0) }{ q(\mathbf{x}_t \vert \mathbf{x}_0) } \\
    &= q(\mathbf{x}_t \vert \mathbf{x}_{t-1}) \frac{ q(\mathbf{x}_{t-1} \vert \mathbf{x}_0) }{ q(\mathbf{x}_t \vert \mathbf{x}_0) } \quad \quad \quad \quad \text{;马尔可夫性}\\
    &\propto \exp \Big(-\frac{1}{2} \big(\frac{(\mathbf{x}_t - \sqrt{\alpha_t} \mathbf{x}_{t-1})^2}{\beta_t} + \frac{(\mathbf{x}_{t-1} - \sqrt{\bar{\alpha}_{t-1}} \mathbf{x}_0)^2}{1-\bar{\alpha}_{t-1}} - \frac{(\mathbf{x}_t - \sqrt{\bar{\alpha}_t} \mathbf{x}_0)^2}{1-\bar{\alpha}_t} \big) \Big) \\
    &= \exp \Big(-\frac{1}{2} \big(\frac{\mathbf{x}_t^2 - 2\sqrt{\alpha_t} \mathbf{x}_t \color{blue}{\mathbf{x}_{t-1}} \color{black}{+ \alpha_t} \color{red}{\mathbf{x}_{t-1}^2} }{\beta_t} + \frac{ \color{red}{\mathbf{x}_{t-1}^2} \color{black}{- 2 \sqrt{\bar{\alpha}_{t-1}} \mathbf{x}_0} \color{blue}{\mathbf{x}_{t-1}} \color{black}{+ \bar{\alpha}_{t-1} \mathbf{x}_0^2} }{1-\bar{\alpha}_{t-1}} - \frac{(\mathbf{x}_t - \sqrt{\bar{\alpha}_t} \mathbf{x}_0)^2}{1-\bar{\alpha}_t} \big) \Big) \\
    &= \exp\Big( -\frac{1}{2} \big( \color{red}{(\frac{\alpha_t}{\beta_t} + \frac{1}{1 - \bar{\alpha}_{t-1}})} \mathbf{x}_{t-1}^2 - \color{blue}{(\frac{2\sqrt{\alpha_t}}{\beta_t} \mathbf{x}_t + \frac{2\sqrt{\bar{\alpha}_{t-1}}}{1 - \bar{\alpha}_{t-1}} \mathbf{x}_0)} \mathbf{x}_{t-1} \color{black} + C(\mathbf{x}_t, \mathbf{x}_0) \big) \Big)
    \end{aligned}
    $$

    • 其中 \(C(\mathbf{x}_t, \mathbf{x}_0)\) 是与 \(x_{t-1}\) 无关的常数,在指数上,可以变成乘法,所以能忽略
  • 转换为高斯分布 ,上面的结果本质是在用等式表达一个高斯分布来,回顾高斯分布的概率密度函数的指数部分是 \(\exp\Big( -\frac{(x-\mu)^2}{2\sigma^2} \Big) = \exp\Big( -\frac{1}{2}(\frac{1}{\sigma^2}x^2 - \frac{2\mu}{\sigma^2}x + \frac{\mu^2}{\sigma^2})\Big) \),上面的推导结果可以进一步得到 \(\tilde{\mu}(x_t, x_0)\) 和 \(\tilde{\beta}_t\):

    • 方差的推导:
      $$
      \begin{aligned}
      \frac{1}{\sigma^2} &= \frac{1}{\tilde{\beta}_t} = (\frac{\alpha_t}{\beta_t} + \frac{1}{1 - \bar{\alpha}_{t-1}}) \\
      \color{red}{\tilde{\beta}_t} &= 1/(\frac{\alpha_t}{\beta_t} + \frac{1}{1 - \bar{\alpha}_{t-1}})
      = 1/(\frac{\alpha_t - \bar{\alpha}_t + \beta_t}{\beta_t(1 - \bar{\alpha}_{t-1})})
      = \color{red}{\frac{1 - \bar{\alpha}_{t-1}}{1 - \bar{\alpha}_t} \cdot \beta_t} ;
      \end{aligned}
      $$
    • 均值的推导:
      $$
      \begin{aligned}
      \frac{2\mu}{\sigma^2} &= \frac{2\tilde{\mu}(x_t,x_0)}{\tilde{\beta}_t} = \frac{2\sqrt{\alpha_t}} {\beta_t} \mathbf{x}_t + \frac{2\sqrt{\bar{\alpha}_{t-1}}}{1 - \bar{\alpha}_{t-1}} \\
      \color{red}{\tilde{\boldsymbol{\mu}} (\mathbf{x}_t, \mathbf{x}_0)} &= (\frac{\sqrt{\alpha_t}}{\beta_t} \mathbf{x}_t + \frac{\sqrt{\bar{\alpha}_{t-1} }}{1 - \bar{\alpha}_{t-1}} \mathbf{x}_0)/(\frac{\alpha_t}{\beta_t} + \frac{1}{1 - \bar{\alpha}_{t-1}}) \\
      &= (\frac{\sqrt{\alpha_t}}{\beta_t} \mathbf{x}_t + \frac{\sqrt{\bar{\alpha}_{t-1} }}{1 - \bar{\alpha}_{t-1}} \mathbf{x}_0) \color{red}{\frac{1 - \bar{\alpha}_{t-1}}{1 - \bar{\alpha}_t} \cdot \beta_t} \\
      &= \color{red}{\frac{\sqrt{\alpha_t}(1 - \bar{\alpha}_{t-1})}{1 - \bar{\alpha}_t} \mathbf{x}_t + \frac{\sqrt{\bar{\alpha}_{t-1}}\beta_t}{1 - \bar{\alpha}_t} \mathbf{x}_0} \\
      \end{aligned}
      $$
  • 去掉结果中的\(\boldsymbol{x_0}\),结合用 \(x_t\) 和方差 \(\beta\) 来表示 \(x_0\) 的公式 : \(x_0 = \frac{1}{\sqrt{\bar{\alpha}_t}}(x_t - \sqrt{1-\bar{\alpha}_t}\epsilon) \) ,我们可以进一步可以将 \(\tilde{\mu}(x_t,x_0)\) 中的 \(x_0\) 替换为 \(x_t\),得到 \( \tilde{\mu}(x_t,x_0) = \tilde{\mu}(x_t)\):
    $$
    \begin{aligned}
    \color{red}{\tilde{\boldsymbol{\mu}}(x_t)}
    &= \frac{\sqrt{\alpha_t}(1 - \bar{\alpha}_{t-1})}{1 - \bar{\alpha}_t} \mathbf{x}_t + \frac{\sqrt{\bar{\alpha}_{t-1}}\beta_t}{1 - \bar{\alpha}_t} \frac{1}{\sqrt{\bar{\alpha}_t}}(\mathbf{x}_t - \sqrt{1 - \bar{\alpha}_t}\boldsymbol{\epsilon}) \\
    &= \color{red}{\frac{1}{\sqrt{\alpha_t}} \Big( \mathbf{x}_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar{\alpha}_t}} \boldsymbol{\epsilon} \Big)}
    \end{aligned}
    $$

    • 其中 \(\epsilon\) 是使用深度学习预测得到的混合噪声,表示 \(x_0 \rightarrow x_t\) 过程中添加的 \(t\) 个噪声的Merge结果(仍然是一个高斯噪声且服从标准正太分布),可以表示为 \(\epsilon_\theta(x_t, t)\)
  • 最终得到逆向推导结果 :我们已经求得了 \(\tilde{\mu}(x_t, x_0) = \tilde{\mu}(x_t)\) 和 \(\tilde{\beta}_t\),即可根据下面的方式采样:
    $$q(x_{t-1}|x_t,x_0) = \mathcal{N}(x_{t-1};\tilde{\mu}(x_t, x_0), \tilde{\beta}_t\boldsymbol{I})$$

    • 将 \(x_0\) 替换为 \(x_t\) 后得到:
      $$q(x_{t-1}|x_t,x_0) = q(x_{t-1}|x_t) = \mathcal{N}(x_{t-1};\tilde{\mu}(x_t), \tilde{\beta}_t\boldsymbol{I})$$
    • 进一步地带入详细表征形式有(其中 \(z_t \sim N(0,\boldsymbol{I})\)):
      $$
      \begin{align}
      x_{t-1} &= \tilde{\mu}(x_t) + \sqrt{\tilde{\beta}_t} z_t \\
      \color{red}{x_{t-1}} &= \color{red}{\frac{1}{\sqrt{\alpha_t}} \Big( \mathbf{x}_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar{\alpha}_t}} \boldsymbol{\epsilon}_\theta(x_t, t) \Big)} + \color{blue}{\sqrt{\frac{1 - \bar{\alpha}_{t-1}}{1 - \bar{\alpha}_t} \cdot \beta_t}} \cdot z_t
      \end{align}
      $$
  • 对推导结果进行简化 :采样(生成)过程可以总结为如下的流程(下面的流程表示已知 \(x_t\) 和时间步 \(t\) 时采样 \(x_{t-1}\) 的步骤):

    • 计算均值 \(\tilde{\boldsymbol{\mu}}(x_t, t)\):将 \(x_t\) 和时间步 \(t\) 输入神经网络预测 \(x_0 \rightarrow x_t\) 过程中添加的 \(t\) 个噪声的Merge结果 \(\epsilon_\theta(x_t, t)\),接着将 \(\epsilon_\theta(x_t, t)\) 带入 \(\tilde{\boldsymbol{\mu}}(x_t) = \color{red}{\frac{1}{\sqrt{\alpha_t}} \Big( \mathbf{x}_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar{\alpha}_t}} \boldsymbol{\epsilon}_\theta(x_t, t) \Big)}\) 得到均值
    • 计算方差 \(\boldsymbol{\Sigma}_\theta(x_t,t)\):在DDPM中,使用固定值作为方差 \(\boldsymbol{\Sigma}_\theta(\mathbf{x}_t, t) = \sigma^2_t \boldsymbol{I}\),此时所有 \(\mathbf{x}_t\) 中所有维度的变量共享同一个方差 ,同时在实践中,使用 \(\color{blue}{\sigma_t^2 = \tilde{\beta}_t = \frac{1 - \bar{\alpha}_{t-1}}{1 - \bar{\alpha}_t} \cdot \beta_t \approx \beta_t} \),即简单实现下可以取 \(\boldsymbol{\Sigma}_\theta(\mathbf{x}_t, t) = \beta_t \boldsymbol{I}\)
      • 其他有些文章会使用可学习的神经网络表示 \(\Sigma_\theta(x_t, t)\),此时不同维度的变量方差可能不同,但是一般仍然假设不同维度变量独立 ,即始终保持 \(\Sigma_\theta(x_t, t)\) 是对角矩阵,比如[GLIDE: Towards Photorealistic Image Generation and Editing with Text-Guided Diffusion Models]和[Improved Denoising Diffusion Probabilistic Models, ICML 2021]等
    • 采样 \(x_{t-1} \sim \mathcal{N}(\tilde{\boldsymbol{\mu}}(x_t, t), \Sigma_\theta(x_t,t))\)
  • 一些符号说明 :在数学中,常用 \(\sigma\) 表示标准差、 \(\Sigma\) 表示方差或协方差矩阵(多变量情况下,如果变量间独立,则有 \(\Sigma\) 矩阵上每个元素是 \(\sigma_i^2\) ),Diffusion相关的文章也沿用了这种表示。同时,由于特定的设定在Diffusion中, \(\beta\) 常常等于方差,所以常常也会用 \(\beta\) 直接表达方差含义

Diffusion模型训练目标函数推导

  • 之前已经知道了,我们Diffusion模型的目标是通过训练学习 \(\mu_\theta(x_t,t)\) 和 \(\Sigma_\theta(x_t,t)\),这个模型训练时的目标是在已知数据集中数据出现的概率最大(极大似然法),即优化在分布 \(x_0 \sim q(x_0)\) 下 \(p_\theta(x_0)\) 的交叉熵:
    $$ L(\theta) = \mathbb{E}_{x_0 \sim q(x_0)}[-\log p_\theta(x_0)] = \mathbb{E}_{q(x_0)}[-\log p_\theta(x_0)] $$
  • 其中(下面的推导使用了Jenson不等式):
    $$
    \begin{aligned}
    - \log p_\theta(\mathbf{x}_0)
    &\leq - \log p_\theta(\mathbf{x}_0) + D_\text{KL}(q(\mathbf{x}_{1:T}\vert\mathbf{x}_0) | p_\theta(\mathbf{x}_{1:T}\vert\mathbf{x}_0) ) \\
    &= -\log p_\theta(\mathbf{x}_0) + \mathbb{E}_{\mathbf{x}_{1:T}\sim q(\mathbf{x}_{1:T} \vert \mathbf{x}_0)} \Big[ \log\frac{q(\mathbf{x}_{1:T}\vert\mathbf{x}_0)}{p_\theta(\mathbf{x}_{0:T}) / p_\theta(\mathbf{x}_0)} \Big] \\
    &= -\log p_\theta(\mathbf{x}_0) + \mathbb{E}_{\mathbf{x}_{1:T}\sim q(\mathbf{x}_{1:T} \vert \mathbf{x}_0)} \Big[ \log\frac{q(\mathbf{x}_{1:T}\vert\mathbf{x}_0)}{p_\theta(\mathbf{x}_{0:T})} + \log p_\theta(\mathbf{x}_0) \Big] \\
    &= \mathbb{E}_{\mathbf{x}_{1:T}\sim q(\mathbf{x}_{1:T} \vert \mathbf{x}_0)} \Big[ \log \frac{q(\mathbf{x}_{1:T}\vert\mathbf{x}_0)}{p_\theta(\mathbf{x}_{0:T})} \Big] \\
    \text{Let }L_\text{VLB}
    &= \mathbb{E}_{\mathbf{x}_{0:T} \sim q(\mathbf{x}_{0:T})} \Big[ \log \frac{q(\mathbf{x}_{1:T}\vert\mathbf{x}_0)}{p_\theta(\mathbf{x}_{0:T})} \Big] \geq - \mathbb{E}_{x_0 \sim q(\mathbf{x}_0)} [\log p_\theta(\mathbf{x}_0)]
    \end{aligned}
    $$
  • 交叉熵损失函数(下面的推导使用到了重积分中的Fubini定理和Jenson不等式)
    $$
    \begin{aligned}
    L_\text{CE}
    &= - \mathbb{E}_{x_0 \sim q(\mathbf{x}_0)} \Big[\log p_\theta(\mathbf{x}_0) \Big]\\
    &= - \mathbb{E}_{x_0 \sim q(\mathbf{x}_0)} \Big[\log \Big( \int p_\theta(\mathbf{x}_{0:T}) d\mathbf{x}_{1:T} \Big)\Big] \\
    &= - \mathbb{E}_{x_0 \sim q(\mathbf{x}_0)} \Big[\log \Big( \int q(\mathbf{x}_{1:T} \vert \mathbf{x}_0) \frac{p_\theta(\mathbf{x}_{0:T})}{q(\mathbf{x}_{1:T} \vert \mathbf{x}_{0})} d\mathbf{x}_{1:T} \Big)\Big] \\
    &= - \mathbb{E}_{x_0 \sim q(\mathbf{x}_0)} \Big[\log \Big( \mathbb{E}_{\mathbf{x}_{1:T} \sim q(\mathbf{x}_{1:T} \vert \mathbf{x}_0)} \frac{p_\theta(\mathbf{x}_{0:T})}{q(\mathbf{x}_{1:T} \vert \mathbf{x}_{0})} \Big)\Big] \\
    &\leq - \mathbb{E}_{x_0 \sim q(\mathbf{x}_0)} \Big[\mathbb{E}_{\mathbf{x}_{1:T} \sim q(\mathbf{x}_{1:T} \vert \mathbf{x}_0)}\Big[ \log \frac{p_\theta(\mathbf{x}_{0:T})}{q(\mathbf{x}_{1:T} \vert \mathbf{x}_{0})}\Big]\Big] \\
    &= - \mathbb{E}_{\mathbf{x}_{0:T} \sim q(\mathbf{x}_{0:T})} \Big[\log \frac{p_\theta(\mathbf{x}_{0:T})}{q(\mathbf{x}_{1:T} \vert \mathbf{x}_{0})} \Big]\\
    &= \mathbb{E}_{\mathbf{x}_{0:T} \sim q(\mathbf{x}_{0:T})}\Big[\log \frac{q(\mathbf{x}_{1:T} \vert \mathbf{x}_{0})}{p_\theta(\mathbf{x}_{0:T})} \Big] = L_\text{VLB}
    \end{aligned}
    $$
  • 进一步推导:
    $$
    \begin{aligned}
    L_\text{VLB}
    &= \mathbb{E}_{\mathbf{x}_{0:T} \sim q(\mathbf{x}_{0:T})} \Big[ \log\frac{q(\mathbf{x}_{1:T}\vert\mathbf{x}_0)}{p_\theta(\mathbf{x}_{0:T})} \Big] \\
    &= \mathbb{E}_{\mathbf{x}_{0:T} \sim q(\mathbf{x}_{0:T})} \Big[ \log\frac{\prod_{t=1}^T q(\mathbf{x}_t\vert\mathbf{x}_{t-1})}{ p_\theta(\mathbf{x}_T) \prod_{t=1}^T p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t) } \Big] \\
    &= \mathbb{E}_{\mathbf{x}_{0:T} \sim q(\mathbf{x}_{0:T})} \Big[ -\log p_\theta(\mathbf{x}_T) + \sum_{t=1}^T \log \frac{q(\mathbf{x}_t\vert\mathbf{x}_{t-1})}{p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t)} \Big] \\
    &= \mathbb{E}_{\mathbf{x}_{0:T} \sim q(\mathbf{x}_{0:T})} \Big[ -\log p_\theta(\mathbf{x}_T) + \sum_{t=2}^T \log \frac{q(\mathbf{x}_t\vert\mathbf{x}_{t-1})}{p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t)} + \log\frac{q(\mathbf{x}_1 \vert \mathbf{x}_0)}{p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1)} \Big] \\
    &= \mathbb{E}_{\mathbf{x}_{0:T} \sim q(\mathbf{x}_{0:T})} \Big[ -\log p_\theta(\mathbf{x}_T) + \sum_{t=2}^T \log \Big( \frac{q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0)}{p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t)}\cdot \frac{q(\mathbf{x}_t \vert \mathbf{x}_0)}{q(\mathbf{x}_{t-1}\vert\mathbf{x}_0)} \Big) + \log \frac{q(\mathbf{x}_1 \vert \mathbf{x}_0)}{p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1)} \Big] \\
    &= \mathbb{E}_{\mathbf{x}_{0:T} \sim q(\mathbf{x}_{0:T})} \Big[ -\log p_\theta(\mathbf{x}_T) + \sum_{t=2}^T \log \frac{q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0)}{p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t)} + \sum_{t=2}^T \log \frac{q(\mathbf{x}_t \vert \mathbf{x}_0)}{q(\mathbf{x}_{t-1} \vert \mathbf{x}_0)} + \log\frac{q(\mathbf{x}_1 \vert \mathbf{x}_0)}{p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1)} \Big] \\
    &= \mathbb{E}_{\mathbf{x}_{0:T} \sim q(\mathbf{x}_{0:T})} \Big[ -\log p_\theta(\mathbf{x}_T) + \sum_{t=2}^T \log \frac{q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0)}{p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t)} + \log\frac{q(\mathbf{x}_T \vert \mathbf{x}_0)}{q(\mathbf{x}_1 \vert \mathbf{x}_0)} + \log \frac{q(\mathbf{x}_1 \vert \mathbf{x}_0)}{p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1)} \Big] \\
    &= \mathbb{E}_{\mathbf{x}_{0:T} \sim q(\mathbf{x}_{0:T})} \Big[ \log\frac{q(\mathbf{x}_T \vert \mathbf{x}_0)}{p_\theta(\mathbf{x}_T)} + \sum_{t=2}^T \log \frac{q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0)}{p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t)} - \log p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1) \Big] \\
    &= \mathbb{E}_{\mathbf{x}_{0:T} \sim q(\mathbf{x}_{0:T})} [\underbrace{D_\text{KL}(q(\mathbf{x}_T \vert \mathbf{x}_0) \parallel p_\theta(\mathbf{x}_T))}_{L_T} + \sum_{t=2}^T \underbrace{D_\text{KL}(q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0) \parallel p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t))}_{L_{t-1}} \underbrace{- \log p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1)}_{L_0} ]
    \end{aligned}
    $$
  • 我们可以进一步将上面的表达式简化写为下面的形式:
    $$
    \begin{aligned}
    L_\text{VLB} &= L_T + L_{T-1} + \dots + L_0 \\
    \text{where } L_T &= D_\text{KL}(q(\mathbf{x}_T \vert \mathbf{x}_0) \parallel p_\theta(\mathbf{x}_T)) \\
    L_t &= D_\text{KL}(q(\mathbf{x}_t \vert \mathbf{x}_{t+1}, \mathbf{x}_0) \parallel p_\theta(\mathbf{x}_t \vert\mathbf{x}_{t+1})) \text{ for }1 \leq t \leq T-1 \\
    L_0 &= - \log p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1)
    \end{aligned}
    $$
    • \(L_T\):由于前向过程 \(q(x_T|x_0)\) 没有可学习参数,且 \(x_T\) 是纯粹的高斯噪声,所以优化参数 \(\theta\) 时 \(L_T\) 可以当做常数忽略
    • \(L_0\): \(\mathcal{N}(\mathbf{x}_0; \boldsymbol{\mu}_\theta(\mathbf{x}_1, 1), \boldsymbol{\Sigma}_\theta(\mathbf{x}_1, 1))\) , \(L_0\) 是扩散模型训练过程中变分下界的一部分,定义为 \(-\log p_{\theta}(x_0 | x_1)\) ,与反向过程的最后一项有关。实际上,虽然 \( L_0 \) 确实与模型参数 \( \theta \) 有关,但在论文的推导和训练过程中,作者通过离散解码器的设计和简化训练目标,将 \( L_0 \) 视为一个相对不重要的常数项。这种做法简化了训练过程,并且实验结果表明,这种简化并不会显著影响模型的生成质量。更多讨论见附录
    • \(L_t\): \(L_t\) 可以看做是两个高斯分布的KL散度, \(q(x_{t-1}|x_t,x_0) = \mathcal{N}(x_{t-1};\tilde{\mathbf{\mu}}(x_t,x_0),\tilde{\beta}_t\boldsymbol{I})\) 和 \(p_\theta(\mathbf{x}_{t-1} \vert \mathbf{x}_t) = \mathcal{N}(\mathbf{x}_{t-1}; \boldsymbol{\mu}_\theta(\mathbf{x}_t, t), \boldsymbol{\Sigma}_\theta(\mathbf{x}_t, t))\),其中 \(\tilde{\boldsymbol{\mu}}(x_t,t) = \frac{1}{\sqrt{\alpha_t}} \Big( \mathbf{x}_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar{\alpha}_t}} \boldsymbol{\epsilon}_t \Big)\)
      • 参照 \(\tilde{\boldsymbol{\mu}}(x_t,t) = \frac{1}{\sqrt{\alpha_t}} \Big( \mathbf{x}_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar{\alpha}_t}} \boldsymbol{\epsilon}_t \Big)\) 的定义,可以将 \(\mu_\theta(x_t,t)\) 定义为如下形式(也就是说,只需要学习噪声 \(\epsilon_\theta(x_t, t)\) 即可):
        $$
        \begin{aligned}
        \boldsymbol{\mu}_\theta(\mathbf{x}_t, t) &= \color{red}{\frac{1}{\sqrt{\alpha_t}} \Big( \mathbf{x}_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar{\alpha}_t}} \boldsymbol{\epsilon}_\theta(\mathbf{x}_t, t) \Big)} \\
        \text{Thus }\mathbf{x}_{t-1} &= \mathcal{N}(\mathbf{x}_{t-1}; \frac{1}{\sqrt{\alpha_t}} \Big( \mathbf{x}_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar{\alpha}_t}} \boldsymbol{\epsilon}_\theta(\mathbf{x}_t, t) \Big), \boldsymbol{\Sigma}_\theta(\mathbf{x}_t, t))
        \end{aligned}
        $$
      • 使得KL散度最小,可以直接使得两者的均值和方差尽可能相等,其中均值目标可以进一步优化为(DDPM中方差直接使用固定值)
        $$\begin{aligned}
        L_t
        &= \mathbb{E}_{\mathbf{x}_0, \boldsymbol{\epsilon}} \Big[\frac{1}{2 | \boldsymbol{\Sigma}_\theta(\mathbf{x}_t, t) |^2_2} | \color{blue}{\tilde{\boldsymbol{\mu}}_t(\mathbf{x}_t, \mathbf{x}_0)} - \color{green}{\boldsymbol{\mu}_\theta(\mathbf{x}_t, t)} |^2 \Big] \\
        &= \mathbb{E}_{\mathbf{x}_0, \boldsymbol{\epsilon}} \Big[\frac{1}{2 |\boldsymbol{\Sigma}_\theta |^2_2} | \color{blue}{\frac{1}{\sqrt{\alpha_t}} \Big( \mathbf{x}_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar{\alpha}_t}} \boldsymbol{\epsilon}_t \Big)} - \color{green}{\frac{1}{\sqrt{\alpha_t}} \Big( \mathbf{x}_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar{\alpha}_t}} \boldsymbol{\boldsymbol{\epsilon}}_\theta(\mathbf{x}_t, t) \Big)} |^2 \Big] \\
        &= \mathbb{E}_{\mathbf{x}_0, \boldsymbol{\epsilon}} \Big[\frac{ (1 - \alpha_t)^2 }{2 \alpha_t (1 - \bar{\alpha}_t) | \boldsymbol{\Sigma}_\theta |^2_2} |\boldsymbol{\epsilon}_t - \boldsymbol{\epsilon}_\theta(\mathbf{x}_t, t)|^2 \Big] \\
        &= \mathbb{E}_{\mathbf{x}_0, \boldsymbol{\epsilon}} \Big[\frac{ (1 - \alpha_t)^2 }{2 \alpha_t (1 - \bar{\alpha}_t) | \boldsymbol{\Sigma}_\theta |^2_2} |\boldsymbol{\epsilon}_t - \boldsymbol{\epsilon}_\theta(\sqrt{\bar{\alpha}_t}\mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t}\boldsymbol{\epsilon}_t, t)|^2 \Big]
        \end{aligned}
        $$
  • 最终训练时的目标函数可以表示为:
    $$
    \begin{aligned}
    L_t^\text{simple}
    &= \mathbb{E}_{t \sim [1, T], \mathbf{x}_0, \boldsymbol{\epsilon}_t} \Big[|\boldsymbol{\epsilon}_t - \boldsymbol{\epsilon}_\theta(\mathbf{x}_t, t)|^2 \Big] \\
    &= \mathbb{E}_{t \sim [1, T], \mathbf{x}_0, \boldsymbol{\epsilon}_t} \Big[|\boldsymbol{\epsilon}_t - \boldsymbol{\epsilon}_\theta(\sqrt{\bar{\alpha}_t}\mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t}\boldsymbol{\epsilon}_t, t)|^2 \Big]
    \end{aligned}
    $$
  • 如果考虑 \(L_T\) 和 \(L_0\),则有:
    $$ L_\text{simple} = L_t^\text{simple} + C $$

训练和推理

  • 训练和推理伪代码:
  • 训练过程:
    • 从数据集中挑选出一个样本 \(x_0 \sim q(x_0)\)
    • 均匀采样,随机选择一个扩散时间步 \(t \sim \text{Uniform}(\{ 1,\cdots,T \})\) (扩散时间步的总长度是提前设置的)
    • 从标准正态分布采样误差 \(\epsilon \sim \mathcal{N}(0, \boldsymbol{I})\),注意这里的误差跟样本 \(x_0\) 维度相同
    • 使用梯度下降训练Diffusion的误差模型 \(\epsilon_\theta(x_t,t)\): \(L(\theta) = \mathbb{E}_{t,x_0,\epsilon}[|\epsilon - \epsilon_\theta(\sqrt{\bar{\alpha}}x_0+\sqrt{1-\bar{\alpha}}\epsilon, t)|]^2 \)
  • 采样过程:
    • 从标准正态分布采样加噪后的样本 \(x_T\): \(x_T \sim \mathcal{N}(0, \boldsymbol{I})\),注:可以这样采样的原因是实际上加噪后的样本可以看做是已经变成了随机值了
    • 逆向过程包含 \(T\) 步, \(t = T,\cdots,1\):
      • 如果 \(t > 0\),则采样 \(z \sim \mathcal{N}(0, \boldsymbol{I})\),否则 \(z = 0\)
      • 恢复 \(x_{t-1}\): \(x_{t-1} = \frac{1}{\sqrt{\alpha_t}}\Big( x_t - \frac{1-\alpha_t}{\sqrt{1-\bar{\alpha}_t}}\epsilon_\theta(x_t, t) \Big) + \sigma_t z\)
        • 其中 \(\epsilon_\theta(x_t, t)\) 的含义是在预测从 \(x_0 \rightarrow x_t\) 的过程中加入的 \(t\) 个混合噪声, \(\sigma_t\) 是方差,DDPM中直接取近似值 \(\sigma_t = \beta_t\)
        • \(z\) 是为了添加生成图片的随机性加入的,与 \(\epsilon_\theta(x_t, t)\) 没有任何关系,如果在每一步中保持 \(z=0\),对于相同的噪声 \(x_T\),生成的图片 \(x_0\) 将是固定的,且图片质量往往不好

一些问题和思考

为什么学习的是 Merge噪声 \(\epsilon_{\theta}(x_t, t)\)

  • 训练时预估的 \(\epsilon_{\theta}(x_t, t)\) 的含义是什么?学到的是从 \(x_0\) 到 \(x_t\) 加的噪声,还是从 \(x_{t-1}\) 到 \(x_t\) 加的噪声?
    • 回答:是从 \(x_0\) 到 \(x_t\) 加的噪声,本质是多个噪声Merge以后得结果(还是高斯噪声),因为从训练流程的公式就可以看出, \(\epsilon_{\theta}(x_t, t)\) 中的 \(x_t=\sqrt{\bar{\alpha}_{t}}x_0 + \sqrt{1-\bar{\alpha}_{t}}\epsilon\),其中 \(\epsilon\) 就是从 \(x_0\) 到 \(x_t\) 加的噪声,而损失函数的目标就是让 \(\epsilon_{\theta}\rightarrow\epsilon\)
  • 在推理时,既然 \(\epsilon_{\theta}(x_t, t)\) 是从 \(x_0\) 到 \(x_t\) 加的累计噪声,为什么可以使用 \(\epsilon_{\theta}(x_t, t)\) 来完成从 \(x_{t}\) 到 \(x_{t-1}\) 的过程?
    • 回答:因为这里从 \(x_{t}\) 到 \(x_{t-1}\) 的本质也是分两步的,第一步是从 \(x_{t}\) 到 \(x_{0}\) (由 \(x_t=\sqrt{\bar{\alpha}_{t}}x_0 + \sqrt{1-\bar{\alpha}_{t}}\epsilon_{\theta}\) 反推即可得到 \(x_0 = \frac{1}{\sqrt{\bar{\alpha}_{t}}}(x_t-\sqrt{1-\bar{\alpha}_{t}}\epsilon_{\theta})\),注意这里得到的 \(x_0\) 可能质量不太好,不能直接作为最终结果);第二步是已知从 \(x_{t}\) 和 \(x_{0}\) 后,可以得到 \(x_{t-1}\) 的分布 \(q(x_{t-1}|x_t,x_0)\),进一步对分布进行采样,就能得到一个 \(x_{t-1}\) 的实例。两步合并以后就是DDPM的采样伪代码中的公式

为什么不能一步到位得到 \(x_0\)?

  • 补充一步到位的公式 :\(x_0 = \frac{1}{\sqrt{\bar{\alpha}_{t}}}(x_t-\sqrt{1-\bar{\alpha}_{t}}\epsilon_{\theta})\)
  • 回答1(训练损失函数视角):可以,但效果不好。从训练的损失函数看,这样做可以,但是效果不一定好(实践来看效果不好,效果不好的原因是一步到位难度较高?),此时相当于简单的把Diffusion模型作为一个单步去噪器了,甚至不需要采样过程
  • 回答2(一种直观理解):为什么效果不好? :模型直接学习 \(q(x_0|x_t)\) 是困难的,即误差 \(\epsilon_\theta(x_t,t)\) 虽然可以用于直接恢复 \(x_0\),但是不太精确(相当于是“粗略预估”),为了得到较好的 \(x_0\),在预估误差 \(\epsilon_\theta(x_t,t)\) 后,我们需要进行“修正”,所以,逐步采样的过程就是逐步“粗略预估”,逐步“修正”的过程(“预估”+“修正”的视角来自生成扩散模型漫谈(十):统一扩散模型(理论篇))
    • 更进一步的理解:可以将多步采样看做是一个ensemble的过程,可以提升效果
  • 回答3(一种理论视角的讨论,待商榷):不可以。在整个推导过程中,我们有 \(q(x_t|x_0) = \mathcal{N}(x_t;\sqrt{\bar{\alpha}_{t}}x_0, (1-\bar{\alpha}_{t})\boldsymbol{I})\) (这里是通过高斯过程的叠加实现的,整个过程遵循马尔可夫过程),所以才有公式 \(x_t=\sqrt{\bar{\alpha}_{t}}x_0 + \sqrt{1-\bar{\alpha}_{t}}\epsilon\),这不代表我们可以通过 \(x_0 = \frac{1}{\sqrt{\bar{\alpha}_{t}}}(x_t-\sqrt{1-\bar{\alpha}_{t}}\epsilon_{\theta})\) 来得到 \(x_0\),因为此时的 \(x_0\) 是无法一步导出的(即 \(q(x_0|x_t)\) 是未知的,未知的原因是不遵循马尔可夫过程,训练时使用的损失函数是在满足马尔可夫过程假设的情况下推导出来的,推理时也不能违背该假设),必须遵循马尔可夫过程(即 \(q(x_{t-1}|x_t)\) )
    • \(q(x_t|x_0)\) 已知但是 \(q(x_0|x_t)\) 未知的原因是因为采样是不可逆的
    • 既然 \(x_0 = \frac{1}{\sqrt{\bar{\alpha}_{t}}}(x_t-\sqrt{1-\bar{\alpha}_{t}}\epsilon_{\theta})\) 不准确(不遵循马尔可夫过程),为什么按照 \(x_0 = \frac{1}{\sqrt{\bar{\alpha}_{t}}}(x_t-\sqrt{1-\bar{\alpha}_{t}}\epsilon_{\theta})\) 的到的 \(x_0\) 可以作为中间变量来帮助生成 \(x_{t-1}\) 呢?公式推导带入的时候使用了这个式子,本质上也说明推导过程不遵循马尔可夫过程了吧?(一种理解是中间使用的 \(x_0 = \frac{1}{\sqrt{\bar{\alpha}_{t}}}(x_t-\sqrt{1-\bar{\alpha}_{t}}\epsilon_{\theta})\) 本质是 \(x_0\) 的一个不精确近似,使用这个来作为条件引导生成可行【引导需要的准确度不高】,但是直接用来作为最终的 \(x_0\) 则效果不太行)

DDPM可以不加入噪声吗?

  • 实践1 :不可以,图片生成质量会特别差
  • 实践2 :不加入噪声,没有随机性,生成图片的多样性受损
  • 理解1 :可以类比LLM采样不能使用纯贪心策略(argmax),因为这样可能生成的是重复的无意义文本,噪音反而能一定程度跳开重复
  • 理解2 :类似于随机梯度下降(SGD)中的噪声帮助跳出局部最优解,扩散模型的噪声让生成过程有机会“绕开”累积的预测偏差
  • 理解3 :扩散模型的生成是一个多步迭代过程,如果每一步仅依赖预测的均值(不加噪声),前一步的误差会逐步累积,最终偏离真实数据分布。通过加入噪声,模型能够在后续步骤中修正误差,提升生成稳定性
  • 理解4 :扩散模型的推导基于随机微分方程(SDE)或变分推断。在反向过程中,噪声项的引入是数学推导的自然结果。例如,在基于分数的生成模型(Score-Based Models)中,采样过程需要引入朗之万动力学(Langevin Dynamics)的随机性,对应扩散模型中的噪声添加
  • 其他参考:一个视频看懂DDIM凭什么加速采样|扩散模型相关【这个观点无法解释DDIM为什么可以直接生成,DDIM(DDIM生成时不加入噪声)生成1000步效果也不会太差】

为什么模型不直接拟合 \(x_t\rightarrow x_{t-1}\) 的噪音?

  • DDPM现状 :目前DDPM训练时模拟拟合的是从 \(x_0\) 到 \(x_t\) 加的噪声,本质是多个噪声Merge以后得结果(还是高斯噪声)
  • 原因(效率视角) :如果直接一步步采样,会导致没张图片都要采样T次(T一般很大)才能训练,训练效率较低,都2023年了,我不允许你还不懂DDPM! - Defa Zhu的文章 - 知乎 中是这个观点
  • 原因(理论视角) :因为逆向过程实际上是 \(q(x_{t-1}|x_t,x_0)\),而 \(q(x_{t-1}|x_t)\) 是在近似拟合 \(q(x_{t-1}|x_t,x_0)\),噪音应该是与 \(x_0\) 相关的,只是在推导过程中用 \(x_t\) 和 Merge噪声 把 \(x_0\) 替换了
    • 特别注意: \(q(x_{t-1}|x_t)\) 本身可以被预估误差,但是本质上是无法直接求解的(一个采样的逆过程,没有梯度,也无法求解),只有 \(q(x_{t-1}|x_t,x_0)\) 可以求解
  • 如果不考虑效率,直接拟合 \(x_t\rightarrow x_{t-1}\) 的噪音是否可行呢?
    • 直观感觉是可行的,但是预估结构可能无法保障,此时相当于假设了 \(q(x_{t-1}|x_t)\) 是一个高斯分布,是否还能继续在采样过程中加入随机有待实践考证
  • 参考:为什么DDPM反向去噪时不直接用q(x_t|x_{t-1})从x_t反推回x_{t-1}? - gpenai的回答 - 知乎 是一个知乎回答,不一定准确,但可以参考看一下

附录:为什么\(L_0\)可以舍弃?

  • 结论:不论是否丢弃 \(L_0\),最终求出来的损失函数形式是一样的,都是 MSE 的形式;
    • 一方面:通过其他推导方式(生成扩散模型漫谈(一):DDPM = 拆楼 + 建楼或 DDPM 推导视频),可以做到不丢弃任何项得到相同的最终结果
    • 另一方面:最小化 \(\mathcal{N}(\mathbf{x}_0; \boldsymbol{\mu}_\theta(\mathbf{x}_1, 1), \boldsymbol{\Sigma}_\theta(\mathbf{x}_1, 1))\) 等价于最小化 \(\vert x_0 - \boldsymbol{\mu}_\theta(\mathbf{x}_1, 1)\vert^2\),进一步推导后可得 \(t=1\) 时刻的最终形式和以上化简后的最优目标函数其实是相同的
  • 在论文中,损失函数的变分下界(VLB)被分解为多个项,其中 \( L_0 \) 是最后一项,表示在给定 \(\mathbf{x}_1\) 的情况下,模型对原始数据 \( \mathbf{x}_0 \) 的重建误差。具体来说,\( L_0 \) 的形式如下:
    $$
    L_0 = -\log p_{\theta}(\mathbf{x}_0 | \mathbf{x}_1)
    $$
    • 进一步可表达为:\(\mathcal{N}(\mathbf{x}_0; \boldsymbol{\mu}_\theta(\mathbf{x}_1, 1), \boldsymbol{\Sigma}_\theta(\mathbf{x}_1, 1))\)
  • 在论文的推导中,有一些误区
    • 问题一 :作者提到 \( L_0 \) 可以被视为常数,这实际上是一种简化,不是严谨的表达,因为 \( L_0 \) 显然与模型参数 \( \theta \) 有关,直接说是化简不合适
    • 问题二(问题一的扩展):虽然 \( L_0 \) 确实与模型参数 \( \theta \) 有关,但在论文的推导和训练过程中,作者认为自己通过离散解码器的设计和简化训练目标,将 \( L_0 \) 视为一个相对不重要的常数项,这种做法简化了训练过程,并且实验结果表明,这种简化并不会显著影响模型的生成质量
      • 实际上作者使用的损失函数已经是最优的了,且训练会采样到 \(t=1\) 的样本,所以不存在丢弃或者简化的操作
    • 问题三 :作者声称自己丢弃了 \(p_\theta(x_0|x_1)\) 的损失学习,如果真的是丢弃,是否会导致采样到最后一步的时候模型不知道应该如何从 \(x_1\) 生成 \(x_0\)?实验证明是不会的(从理论上来看也不会,因为作者其实没丢弃 \(p_\theta(x_0|x_1)\) 的损失学习),具体原因是:
      • 从目标函数看:简化后的训练的目标函数为:
        $$
        L_\text{simple} = \mathbb{E}_{t \sim [1, T], \mathbf{x}_0, \boldsymbol{\epsilon}_t} \Big[|\boldsymbol{\epsilon}_t - \boldsymbol{\epsilon}_\theta(\sqrt{\bar{\alpha}_t}\mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t}\boldsymbol{\epsilon}_t, t)|^2 \Big]
        $$
        • 当训练时遇到 \(t=1\) 时(训练采样时 \(t\in \{1,2,\cdots,T\}\),所以可以遇到这种样本),这里这个损失函数的目标是建模误差就是 \(\epsilon_\theta(x_1, 1)\),即已知 \(x_1\) 时,自然可以恢复 \(x_0\)
      • 考虑到最小化 \(\mathcal{N}(\mathbf{x}_0; \boldsymbol{\mu}_\theta(\mathbf{x}_1, 1), \boldsymbol{\Sigma}_\theta(\mathbf{x}_1, 1))\) 等价于让 \(\vert x_0 - \boldsymbol{\mu}_\theta(\mathbf{x}_1, 1)\vert^2\) 最小化,进一步推导后得到 \(t=1\) 时刻的最终形式和以上化简后的最优目标函数其实是相同的(这在苏神的博客 生成扩散模型漫谈(一):DDPM = 拆楼 + 建楼 中给出了证明)
  • 苏神的博客中,从另一个视角推导(不一定很严谨),得到了简化的目标函数,详情参考:生成扩散模型漫谈(一):DDPM = 拆楼 + 建楼,这说明 \(L_0\) 的简化是合理的
  • DDPM 推导视频 也提供了另一种推导方式,在不丢弃任何一项的情况下,最终目标函数是一样的

附录:其他推导过程参考

DL——GPipe

  • 参考链接:
    • 原始论文:GPipe: Easy Scaling with Micro-Batch Pipeline Parallelism, 2019, Google

Paper Summary

  • 已有结论:模型越大,模型效果上限越高
  • 模型容量超过单个加速器(GPU 或 TPU)的内存限制时,需要开发特殊算法和工程设计(这些解决方案通常是针对特定架构设计的,难以迁移)
  • 论文提出了 GPipe
    • 是一个 流水线并行(Pipeline Parallelism)库
    • 能够扩展任何可表示为层序列的网络
    • 高效且任务无关的模型并行
  • GPipe 将不同的层子序列分配到不同的加速器上流水执行
    • 可以灵活高效地将多种网络扩展至超大规模
  • GPipe 采用了一种新颖的批次分割流水线算法 ,在模型跨多个加速器分区时实现了近乎线性的加速比
  • 论文通过训练两种不同任务的大规模神经网络来展示 GPipe 的优势,这些任务具有截然不同的网络架构:
    • (i) 图像分类(Image Classification) :
      • 训练了一个包含 5.57 亿参数的 AmoebaNet 模型,在 ImageNet-2012 上达到了 84.4% 的 top-1 准确率;
    • (ii) 多语言神经机器翻译(Multilingual Neural Machine Translation) :
      • 在涵盖 100 多种语言的语料库上训练了一个包含 60 亿参数、128 层的 Transformer 模型,其质量优于所有双语模型
      • 实验表明,该模型在 100 种语言对上的性能优于单独训练的 3.5 亿参数双语 Transformer Big 模型 (2017)
  • 评价:
    • Google 出品,必属精品,GPipe 已成为各种框架的加速选项

Introduction and Discussion

  • 近年来深度学习的巨大进步部分归功于神经网络有效容量扩展方法的发展
  • 这一趋势在图像分类任务中最为明显,例如 ImageNet 的准确率随着模型容量的增加而提升(图 0(a))
    • 类似现象也出现在自然语言处理领域(图 0(b)),简单的浅层句子表示模型 (2017; 2018) 被更深、更大的模型 (2018; 2019) 超越
  • 更大的模型为多个领域带来了显著的性能提升,但扩展神经网络也带来了实际的挑战
    • 硬件限制(包括加速器(GPU 或 TPU)的内存和通信带宽)迫使用户将大模型分割为多个分区,并将不同分区分配到不同的加速器上
    • 高效的模型并行算法设计和实现极为困难,通常需要用户在扩展容量、灵活性(或对特定任务和架构的适应性)以及训练效率之间做出艰难抉择
    • 大多数高效的模型并行算法都是针对特定架构和任务的
    • 随着深度学习应用的不断增加,对可靠且灵活的基础设施的需求日益增长
  • 为解决这些挑战,论文提出了 GPipe,一个灵活的库,支持高效训练大规模神经网络
    • GPipe 通过将模型分配到不同加速器上并支持在每个加速器上重新计算中间结果 (2000; 2016),能够突破单个加速器的内存限制 ,扩展任意深度神经网络架构
    • GPipe 中,每个模型可以表示为层序列,连续的层组可以划分为单元(cell),每个单元分配到单独的加速器上
  • 基于这种分区设置,论文提出了一种新颖的批次分割流水线并行算法
    • 首先,论文将训练样本的小批次(mini-batch)分割为更小的微批次(micro-batch) ,然后在单元上流水执行每组微批次
    • 训练采用同步小批次梯度下降,梯度在所有微批次中累积,并在小批次结束时统一更新
    • GPipe 的梯度更新与分区数量无关,保证了训练的一致性 ,使研究人员能够通过部署更多加速器轻松训练更大的模型
    • GPipe 还可以与数据并行结合,进一步扩展训练规模

GPipe 库(The GPipe Library)

  • 论文现在描述 GPipe 的接口和主要设计特性
  • 该开源库是在 Lingvo(2019)框架下实现的
  • GPipe的核心设计特性具有通用性,可以为其他框架(2017; 2018; 2019)实现

Interface

  • 任何深度神经网络都可以定义为一个由 \(L\) 个层组成的序列
    • 每个层 \(L_i\) 由一个前向计算函数 \(f_i\) 和对应的参数集 \(w_i\) 组成
  • GPipe 允许用户指定一个可选的计算成本估计函数 \(c_i\)
  • 在给定分区数量 \(K\) 的情况下,可以将 \(L\) 个层的序列划分为 \(K\) 个复合层或单元(cell)
    • 设复合单元 \(p_k\) 由层 \(i\) 到层 \(j\) 之间的连续层组成
    • 与复合单元 \(p_k\) 对应的参数集等价于 \(w_i, w_{i+1}, \ldots, w_j\) 的并集,其前向函数为 \(F_k = f_j \circ \ldots \circ f_{i+1} \circ f_i\)
    • 对应的反向传播函数 \(B_k\) 可以通过自动符号微分从 \(F_k\) 计算得到
    • 成本估计器 \(C_k\) 设置为 \(\Sigma^{j}_{l=i}c_l\)
  • GPipe 的接口非常简单直观,要求用户指定:
    • (i) 模型分区数量 \(K\)
    • (ii) 微批次(micro-batch)数量 \(M\)
    • (iii) 定义模型的 \(L\) 个层的序列和定义
  • 具体示例请参见补充材料

Algorithm

  • 当用户通过模型参数 \(w_i\) 、前向计算函数 \(f_i\) 和成本估计函数 \(c_i\) 定义网络中层的序列后,GPipe 将网络划分为 \(K\) 个单元,并将第 \(k\) 个单元放置在第 \(k\) 个加速器上
  • 在分区边界处自动插入通信原语,以允许相邻分区之间的数据传输
  • 分区算法最小化所有单元的估计成本方差,以通过同步所有分区的计算时间来最大化流水线的效率
  • 在前向传播期间:
    • GPipe 首先将每个大小为 \(N\) 的小批次(mini-batch)划分为 \(M\) 个相等的微批次
    • 这些微批次通过 \(K\) 个加速器进行流水线处理
  • 在反向传播期间:
    • 每个微批次的梯度基于前向传播使用的相同模型参数计算
    • 在每个小批次结束时,所有 \(M\) 个微批次的梯度被累积并应用于更新所有加速器上的模型参数
  • 这一系列操作如图1(c)所示
  • 如果网络中使用了 Batch Normalization(BatchNorm),则训练期间输入的充分统计量(Sufficient Statistic)是在每个微批次上计算的 ,并在需要时在副本上计算(2017)
    • 论文还跟踪整个小批次上充分统计量的移动平均值,以用于评估
      • 注:充分统计量的介绍见附录
    • 理解:使用 GPipe 时,对 BatchNorm 不太友好,因为太小的微批次会导致训练时使用的均值和方差波动太大(注:BatchNorm 训练时始终使用的是当前微批次的统计值)

Performance Optimization

  • 为了减少激活内存需求,GPipe 支持重计算(re-materialization)(2016)
    • 在前向计算期间,每个加速器仅存储分区边界处的输出激活
    • 在反向传播期间,第 \(k\) 个加速器重新计算复合前向函数 \(F_k\)
  • 峰值激活内存需求减少到
    $$O(N + \frac{L}{K} \times \frac{N}{M})$$
    • 其中 \(\frac{N}{M}\) 是微批次大小
    • \(\frac{L}{K}\) 是每个分区的层数
  • 如果不使用重计算和分区,内存需求将为
    $$ O(N \times L)$$
    • 因为计算梯度 \(b_i\) 需要上层梯度 \(b_{i+1}\) 和缓存的激活 \(f_i(x)\)
  • 如图2(c)所示,分区会在每个加速器上引入一些空闲时间,论文称之为气泡(bubble)开销
    • 这种气泡时间在微步骤数量 \(M\) 上平均为 \(O(\frac{K-1}{M+K-1})\)
    • 在论文的实验中,当 \(M \geq 4 \times K\) 时,气泡开销可以忽略不计
    • 这部分也是因为反向传播期间的重计算可以更早调度,而无需等待来自上层的梯度
  • GPipe 还引入了低通信开销,因为论文只需在加速器之间传递分区边界处的激活张量
    • 即使在没有高速互连的加速器上,也可以实现高效的扩展性能
  • 图2(c)假设分区是均匀平衡的
    • 但不同层的内存需求和计算浮点操作往往非常不平衡
    • 在这种情况下,不完美的分区算法可能导致负载不平衡
    • 更好的分区算法可能会在论文的启发式方法基础上进一步提升性能

Performance Analyses

  • 通过两种模型架构评估 GPipe 的性能(研究它们的可扩展性、效率和通信成本):
    • AmoebaNet(2018)卷积模型
    • Transformer(2017)Sequence2Sequence模型
  • 论文预计重计算和流水线并行都会有益于内存利用,从而使训练巨型模型成为可能
  • 在表1中,报告了在合理大的输入大小下 GPipe 可以支持的最大模型大小
  • 对于 AmoebaNet 模型 ,在每个加速器内存为 8GB 的 Cloud TPUv2 上运行实验
    • 使用固定的输入图像大小 \(224 \times 224\) 和小批次大小 128
    • 在没有 GPipe 的情况下,受设备内存限制,单个加速器最多可以训练 8200万 参数的 AmoebaNet,
    • 利用反向传播中的重计算和批次分割:
      • GPipe 将中间激活内存需求从 6.26GB 减少到 3.46GB,从而在单个加速器上支持 3.18亿 参数的模型
    • 通过模型并行
      • 能够在 8 个加速器上将 AmoebaNet 扩展到 18亿 参数 ,比不使用 GPipe 时多25倍
      • 在这种情况下,最大模型大小没有完全线性扩展,因为 AmoebaNet 中不同层的模型参数分布不平衡
  • 对于 Transformer 模型 ,每个加速器核心内存为 16GB 的 Cloud TPUv3 训练
    • 模型配置如下:
      • 使用固定的词汇大小 32k
      • 序列长度 1024 和批次大小 32
      • 每个Transformer层的模型维度为2048
      • 前馈隐藏维度为 8192
      • 注意力头数量为 32
    • 注:论文通过改变层数量来扩展模型
    • 重计算允许在单个加速器上训练比原来大 2.7 倍的模型
    • 通过 128 个分区,GPipe 可以将 Transformer 扩展到 839亿 参数,比单个加速器上可能的规模增加 298 倍
    • 与 AmoebaNet 不同,Transformer 的最大模型大小随加速器数量线性扩展,因为每个层具有相同数量的参数和输入大小
  • 表2给出了 GPipe 的训练效率效率(主要评估 AmoebaNet-D 和 Transformer-48 的标准化训练吞吐量)
    • 使用不同数量的分区和不同数量的微批次
    • 每个分区分配给一个独立的加速器
  • 一些结论:
    • 当微批次数量 \(M\) 至少是分区数量 \(K\) 的 4 倍时,气泡开销几乎可以忽略不计
      • 换算一下可以知道,就以 4 倍为例
      • 气泡量是(注:气泡总数仅与 \(K\) 有关,与 \(M\) 无关)
        $$ 6 + 6 = 12 $$
      • 总的小区域数量是
        $$ (16+3) \times 4 = 76$$
      • 气泡率约为:
        $$12 / 76 \approx 15.8%$$
      • 且微批次数量 \(M\) 越大,气泡开销越小,因为气泡总数是固定的
    • 对于 Transformer 模型,当分区数量增加 4 倍时,速度提升 3.5 倍
    • 由于计算在 Transformer 层之间均匀分布,训练吞吐量几乎随设备数量线性扩展
    • 由于计算在 AmoebaNet模型上分布不平衡,训练吞吐量仅实现了次线性加速
    • 当 微批次数量 \(M\) 相对较小时,气泡开销不再可以忽略
    • 当 微批次数量 \(M\) 为 1 时,实际上没有流水线并行
      • 论文观察到训练吞吐量相对恒定,无论使用多少加速器,这表明任何时候只有一个设备在主动计算
  • 为了测量 GPipe 的通信开销影响,论文在没有 NVLink 的多个 NVIDIA P100 GPU 的单个主机上运行实验
    • 跨 GPU 的数据传输必须通过 PCI-E 进行相对较慢的设备到主机和主机到设备传输
    • 微批次数量固定为 32
  • 如表3所示,当分区数量从 2 增加到 8 时:
    • AmoebaNet-D 的速度提升 2.7 倍
    • 对于 24 层 Transformer,速度提升为 3.3 倍
  • 与配备高速互连的 TPU 上观察到的情况类似,存在类似的线性加速
    • 由于 GPipe 仅在分区边界传输激活张量,设备之间的通信带宽不再是模型并行的瓶颈

Image Classification

  • 作为概念验证,本节首先使用 GPipe 扩展 AmoebaNet
    • 增加了 AmoebaNet 的通道数,并将输入图像尺寸扩展到 \(480 \times 480\)
    • 在 ImageNet 2012 数据集上训练了这个包含 5.57 亿参数的 AmoebaNet-B 模型,使用的超参数与 (2018) 中描述的相同
    • 该网络被划分为 4 个分区
    • 这一单一模型在单次裁剪下实现了 84.4% 的 top-1 准确率和 97% 的 top-5 验证准确率
  • 论文进一步通过迁移学习 (2022, 2023) 证明了巨型卷积网络在其他图像数据集上的有效性
    • 使用预训练的 ImageNet 模型在多个目标数据集上进行微调,这些数据集涵盖了一般分类到细粒度分类任务
    • 将最后一个 softmax 分类层的输出单元数更改为目标数据集的类别数,并随机初始化新的 softmax 层
    • 其余所有层均从 ImageNet 预训练中初始化
    • 训练期间,输入网络的图像被调整为 \(480 \times 480\) ,并随机水平翻转,同时使用 cutout (2017) 进行数据增强
    • 训练超参数与 ImageNet 训练时相同(详细的训练设置见补充材料)
  • 在表4 中,论文报告了每个目标数据集 5 次微调运行的平均单次裁剪测试准确率
    • 巨型模型在所有目标数据集上均取得了具有竞争力的结果
      • CIFAR-10 的错误率降至 1%
      • CIFAR-100 的错误率降至 8.7%
    • 这些结果验证了 Kornblith 等人 (2018) 的发现,即更好的 ImageNet 模型具有更好的迁移能力

Massively Multilingual Machine Translation

  • 本节还通过扩展用于 NLP 的模型来展示 GPipe 的灵活性,论文在一个大规模多语言 NMT 任务上继续 GPipe 的实验
    • 因为并行语料库的丰富性,神经机器翻译(NMT)已成为评估 NLP 架构的基准任务 (2017, 2018, 2019)
  • 论文使用包含 102 种语言与英语的平行文档语料库,总计 250 亿训练样本,每种语言的样本量从 \(10^4\) 到 \(10^9\) 不等 (2019)
    • 该数据集通过涵盖从低资源到高资源的多样化语言,为可扩展性实验提供了真实的测试环境
    • 论文首次在机器翻译中证明,足够大的 NMT 模型可以同时学习超过 100 种语言对的映射,并且在所有语言上均优于双语模型性能
    • 这进一步凸显了高效且灵活的模型并行工具的重要性
  • 论文的比较基于在该语料库上所有语言对训练的单一 Transformer (2017) 的性能
    • 通过两个维度扩展架构以强调 GPipe 的灵活性:
      • (i) 通过增加模型的层数扩展深度;
      • (ii) 通过增加前馈层的隐藏维度以及多头注意力层中的注意力头数(和注意力通道数)扩展宽度
      • 类似于 Shazeer 等人 (2018) 的方法(数据集、基线、训练配置和优化超参数的详细描述见补充材料)
  • 论文从标准的 4 亿参数 Transformer Big 模型 \(T(6,8192,16)\) (如 Chen 等人 (2018) 所述)开始,词表为 64k
    • 注:论文用 \(T(L,\ H,\ A)\) 表示 \(T(层数,FFN隐藏层维度,Attention头数)\)
  • 在图3 中,论文将其性能与下列模型进行了比较:
    • 13 亿参数的深层模型 \(T(24,\ 8192,\ 16)\)
    • 13 亿参数的宽模型 \(T(12,\ 16384,\ 32)\)
    • 30 亿参数模型 \(T(32,\ 16384,\ 32)\)
    • 60 亿参数模型 \(T(64,\ 16384,\ 32)\)
  • 所有模型均使用基于温度的多语言 BERT (2018) 采样方法在所有语言对上同时训练
    • \(T(12,\ 16384,\ 32)\) 、 \(T(24,\ 8192,\ 32)\) 、 \(T(32,\ 16384,\ 32)\) 和 \(T(64,\ 16384,\ 32)\) 分别被划分为 2、4、8 和 16 个加速器
  • 从图3 中可以看出
    • 将模型容量从 4 亿参数增加到 13 亿参数显著提升了所有语言的性能
    • 将模型从 13 亿参数扩展到 60 亿参数进一步提升了性能,尤其是高资源语言
      • 注:从 13 亿到 30 亿和 60 亿参数时出现了收益递减现象
  • 以下是论文基于这些大规模实验的一些实证发现:
    • 深度-宽度权衡(Depth-Width Trade-off) :
      • 论文研究了多语言设置中深度与宽度的权衡,并比较了 13 亿参数的宽模型 \(T(12,\ 16384,\ 32)\) 和 13 亿参数的深模型 \(T(24,\ 8192,\ 16)\) 的性能
      • 虽然这两种模型在高资源语言(图3 左侧)上的质量非常接近,但深层模型在低资源语言上的表现显著优于宽模型 ,这表明增加模型深度可能更有利于泛化
        • 注:图3是 100+ 种语言按照数据量从左到右逆序排序的结果(high-resource language 代表数据量较大的语言),每个点代表当前语言上的表现?
      • 与 4 亿参数模型相比,13 亿参数深层模型在低资源语言(图3 右侧)上的质量提升几乎与高资源语言相当,这表明增加深度可能会增强对低资源任务的迁移效果
    • 深层模型的训练挑战(Trainability Challenges with Deep Models) :
      • 深度增加了神经网络的表示能力,但也使优化问题复杂化
      • 在大规模实验中,论文遇到了由尖锐激活(sharp activations)且是正峰度的(positive kurtosis)和数据集噪声组合引起的严重训练问题
      • 在训练几千步后,模型预测会变得极其尖锐且对噪声敏感,这通常会导致非有限或大梯度,最终破坏学习进程
      • 为了解决这些问题,论文采用了两种方法:
        • (i) 遵循 Zhang 等人 (2019) 的方法,按层数比例缩小所有 Transformer 前馈层的初始化;
        • (ii) 当 logit 预测(softmax 预激活)的幅度超过特定值时,对其进行裁剪
        • 这两种方法的结合缓解了因模型深度扩展带来的训练不稳定性
    • 大批量训练(Large Batches) :
      • 数据并行使用简单,是扩展神经网络训练的主导方法 (2016, 2017)
      • 论文通过显著增加标准 Transformer Big 训练的批量大小来测试大批量训练的极限
      • 从每批 26 万词开始,论文将有效批量增加到 400 万词,并观察高资源语言对(德语-英语)的验证损失和 BLEU 分数(其他语言对也观察到类似趋势)
        • 此处使用的优化参数与之前实验相同
        • 据论文所知,400 万词每批是文献中迄今为止用于训练 NMT 模型的最大批量 (2018)
      • 表5 显示,随着批量增加,两项指标均显著改善
      • 作者相信进一步增加批量可能会带来更多改进

Design Features and Trade-Offs

  • 已有多种方法用于实现高效的大规模模型并行,但每种方法都有其特定的权衡,使其适用于特定硬件约束下的特定架构扩展
  • 本节将重点讨论多种模型并行方法的设计选择与权衡,并比较它们与 GPipe 在灵活性、可扩展性以及不同硬件约束和架构变体下的效率
  • 模型并行的核心思想是将网络划分为不同的计算单元,然后将这些单元分配到不同的设备上(2014; 2014; 2017; 2012)
    • 从概念上讲,这种方法支持将多种模型扩展至巨大规模
    • 但这些方法通常面临硬件利用率低和设备间通信瓶颈的问题
  • 单程序多数据(Single Program Multiple Data, SPMD)和流水线并行被提出以应对这些挑战
    • Mesh-Tensorflow(2018)遵循 SPMD 范式
      • 将数据并行中使用的单指令多数据(Single Instruction Multiple Data, SIMD)方法扩展到其他张量维度
    • SPMD 允许将每个计算拆分到多个设备上,从而使用户能够将单个矩阵乘法(以及单个层的模型参数)的大小随加速器数量线性扩展
      • 但这也引入了设备间的高通信开销,因为需要大量类似 AllReduce 的操作来合并每个并行化矩阵乘法的输出
        • 限制了该方法在配备高速互联的加速器场景下的适用性
      • SPMD 限制了可高效扩展的操作类型,使其仅适用于特定网络架构和机器学习任务
        • 例如,在该范式下沿卷积层的通道维度拆分效率较低,因为通道实际上是全连接的,而沿空间维度拆分则需要复杂的技术来处理边缘区域(halo regions)
      • 尽管 SPMD 允许通过使每个操作更小来扩展模型深度,但它需要将每个层拆分到更多加速器上,这进一步增加了设备间的通信开销
  • 其他方法尝试利用基于流水线并行的技术来扩展神经网络(1993; 2017),最近应用于神经网络训练的流水线并行迭代是 PipeDream(2018)
    • PipeDream 的目标是减少参数服务器(2014)的通信开销
    • PipeDream 通过流水线化前向传播的执行,并将其与反向传播交错,以最大化硬件利用率
    • 这种设计因异步反向更新引入的权重陈旧性(weight staleness)而受到影响
    • 为了避免权重陈旧性导致的优化问题,PipeDream 需要在每个加速器上维护多个版本化的模型参数副本以准确计算梯度更新 ,从而限制了用户扩展至更大模型的能力
  • GPipe 引入了一种新型流水线并行技术
    • 在应用整个小批量(mini-batch)的同步梯度更新之前,对微批次(micro-batches)的执行进行流水线化
    • 论文的新型批拆分流水线并行算法与重计算(re-materialization)相结合,支持扩展到大量微批次
      • 这最小化了“气泡”(bubble)开销,同时无需异步梯度更新
    • GPipe 使用户能够将模型大小随加速器数量线性扩展
      • 与 SPMD 不同,流水线并行在扩展模型时引入的额外通信开销极少
    • 设备间通信仅在每个微批次的划分边界发生,且引入的通信开销可以忽略不计,这使得 GPipe 在缺乏高速设备互联的场景下仍然适用
  • 待提升点:
    • GPipe 目前假设单个层可以适配单个加速器的内存限制
      • 注:突破这一限制的一种可行方法是:将单次矩阵乘法拆分为多个更小的矩阵运算,并依次分散到多个网络层中执行(注:即张量并行)
    • 微批次拆分需要复杂的策略来支持跨批次计算的层(例如,BatchNorm 在训练时使用微批次的统计量,但在评估时累积小批次的统计量)

Conclusion

  • 论文介绍了 GPipe,一个用于训练巨型神经网络的可扩展模型并行库
  • 论文提出并实现了一种新型批拆分流水线并行算法,该算法使用同步梯度更新,实现了高硬件利用率和训练稳定性的模型并行
  • 论文利用 GPipe 训练了大规模卷积和基于 Transformer 的模型,并在图像分类和多语言机器翻译任务上展示了强大的实证结果
  • 论文重点强调了 GPipe 库的三个关键特性:
    • 1)高效性(Efficiency) :通过新型批拆分流水线算法,GPipe 实现了几乎随设备数量线性增长的速度提升
    • 2)灵活性(Flexibility) :GPipe 支持任何可以表示为层序列的深度网络
    • 3)可靠性(Reliability) :GPipe 使用同步梯度下降,并保证无论划分数量多少,训练结果一致

附录1:GPipe GPipe Example Usage

  • GPipe 库的用户首先需要将他们的神经网络表示为 \(L\) 层的顺序列表
    • 任何计算图都可以被划分为一系列子图
  • 示例基础层包括卷积(convolution)、池化(pooling)、批量归一化(batch normalization)、dropout、transformer、softmax 以及其他层
  • 可以按顺序或并行连接的层可以被组合成一个新的复合层
  • 用户可以以任意方式组合任意数量的层,只要正确定义了复合前向函数
  • 图1展示了 GPipe 库的一个示例用例
    • 这是一个用于验证训练一致性的单元测试,它验证了在这个示例网络中,所有梯度的范数在数值误差范围内是相同的,无论分区数量如何
      1
      2
      3
      4
      5
      6
      7
      8
      9
      10
      11
      12
      13
      14
      15
      16
      17
      18
      19
      20
      21
      22
      23
      24
      25
      26
      27
      28
      29
      30
      31
      32
      33
      34
      35
      36
      37
      38
      39
      40
      41
      42
      43
      44
      45
      46
      import tensorflow as tf
      from lingvo.core import base_layer
      from lingvo.core import layers
      from lingvo.core import py_utils
      from TensorPipe import TensorPipeLayer
      def BuildDummyTensorPipeLayer( num_layers=16, num_splits=4, num_mirco_batches=8):
      assert num_layers % num_splits == 0
      layers = []
      # Construct a dummy layer with 16 3x3 conv layers
      for i in range(num_layers):
      layers.append(layers.Conv2DLayer.Params().Set( name=’layer_{}’.format(i)), filter_shape=(3, 3, 1, 1), filter_stride=(1, 1))
      # Evenly distribute layers to partitions.
      partitions = []
      layers_per_split = num_layers // num_splits
      for split in range(num_splits):
      sub = layers[split * layers_per_split: (split + 1) * layers_per_split]
      partitions.append(sub)
      # Build pipeline parallelism model using TensorPipe
      p = TensorPipeLayer.Params().Set(name=’TensorPipe’, num_mirco_batches=num_mirco_batches , partitions=partitions)
      layer = p.cls(p)
      return layer

      class DummyTensorPipeTest(tf.test.TestCase):
      def _verify_consistent_training(self, num_splits):
      g = tf.Graph()
      with g.as_default():
      py_utils.GetOrCreateGlobalStep() tf.set_random_seed(88888888)
      inputs = tf.random_uniform([16, 8, 8, 1])
      net = BuildDummyTensorPipeLayer(num_splits=num_splits) logits = net.FPropDefaultTheta(inputs)
      loss = tf.reduce_mean(logits)
      grads = tf.gradients(loss, tf.trainable_variables()) grad_norm = tf.sqrt(
      py_utils.SumSquared(grads))
      with self.session(graph=g) as sess:
      sess.run(tf.global_variables_initializer())
      grad_norm_val , = sess.run([grad_norm])
      # Verify grad norm is the same regardless of # the number of splits
      self.assertNear(grad_norm_val , 0.269997 , err=1.0e-6)

      def testDummyPipelineCnnOneSplit(self):
      self._verify_timestep_counts(num_splits=1)

      def testDummyPipelineCnnTwoSplits(self):
      self._verify_timestep_counts(nu2m_splits=2)

      def testDummyPipelineCnnFourSplits(self):
      self._verify_timestep_counts(num_splits=4)

附录2:图像分类训练细节 (2 Image Classification Training Details)

Training Hyperparameters

  • 论文在 ImageNet ILSVRC-2012 数据集上训练了一个具有 \(557\) 百万模型参数和输入图像尺寸为 \(480\times 480\) 的 AmoebaNet-B 模型
  • 论文遵循 (2018) 中描述的超参数和输入预处理方法来训练 AmoebaNet-B
    • 衰减率为 \(0.9\)
    • \(\epsilon=0.1\) 的 RMSProp 优化器
    • \(L^{2}\) 正则化系数 \(\lambda=4\times 10^{-5}\)
    • 标签平滑系数 \(0.1\)
    • 权重为 \(0.4\) 的辅助头
    • 对中间层应用了与 NasNet (2018) 中相同的 drop-path 调度,并对最终层应用了概率为 \(0.5\) 的 dropout
    • 使用了学习率调度,该调度每 \(3\) 个周期以 \(0.97\) 的速率衰减,初始学习率为 \(0.00125\) 乘以批量大小
  • 在附录表1中,论文报告了每个迁移学习数据集使用的超参数
    • 从集合 \(\{0.0125,0.00375,0.075,0.115,0.15\}\) 中选择学习率
    • 从集合 \(\{0,4e-8,4e-7,4e-6,4e-5\}\) 中选择 \(L^{2}\) 权重衰减
    • 选择基于训练数据的保留子集(20%)
    • 论文应用选定的超参数进行最终训练,并重复五次
  • 表1中:
    • 论文为每个数据集在训练数据集的保留子集上选择了学习率和 \(L^{2}\) 权重正则化参数
    • 对于其他超参数,论文使用了与 ImageNet 训练中相同的参数

Consistent Training

  • GPipe 在微批次(micro-batches)上执行同步训练
  • 在本节中,验证了以下假设:
    • 使用 GPipe 的端到端收敛精度在统计误差范围内是相同的,无论分区数量如何
  • 实验过程:
    • 多次训练 AmoebaNet-D (6, 256) 模型 35 个 epoch,并在 ImageNet 上测量最终的验证准确率
    • 选择 AmoebaNet-D (6, 256) 是因为它是 DAWNBench 竞赛 (2018) 中训练成本最优的图像模型
    • 采用了 DAWNBench 中报告的超参数和训练流程
    • 作为基线,使用官方开源实现训练了 AmoebaNet-D (6, 256) 5 次,并计算了最终准确率的均值和标准差
    • 使用相同的超参数和训练流程,分别用 1、2、4 和 8 个分区通过 GPipe 训练了相同的网络
  • 结果发现,最终的准确率落在均值 1.6 倍标准差范围内,符合预期

附录3:Machine Translation Training Details

  • 论文研究了大规模多语言神经机器翻译(Neural Machine Translation, NMT),使用的语料库通过从网页爬取并提取平行句子构建
  • 图 2 展示了论文研究的 102 种语言的数据分布情况

Baselines

  • 在双语实验中,论文使用了 Transformer 架构(2017)的变体
  • 对于大多数双语实验,论文采用了一个包含 3.75 亿参数的大型 Transformer Big 模型,并使用了共享的源-目标句子片段模型(SentencePiece Model, SPM)词汇表,包含 32k 个词符
  • 论文根据每种语言对的数据集大小调整了不同的 dropout 值
    • 理解:这里说明是训练了很多个双语 Transformer 模型
    • 问题:图3中的模型是如何和基线比较的呢?想说明什么呢?
      • 理解:基线是谁不重要,重点是实验模型不同配置之间(不同宽度,深度之间)互相比较能拿到结论
  • 对于大多数中低资源语言,论文还尝试了 Transformer Base 模型
    • 注:数据量大对应高资源,数据量小对应低资源
  • 所有模型均采用 Adafactor 优化器(2018)和动量因子化,学习率调度为 \( (3.0, 40\text{K}) \),并设置了每参数范数裁剪阈值为 1.0
  • 对于 Transformer Base 模型,学习率调度为 \( (2.0, 8\text{K}) \)
  • BLEU 分数基于验证集上表现最佳的检查点计算,输出和参考文本均为真实大小写形式

Multilingual Baselines

  • 接下来,论文描述训练多语言模型的方法
  • 由于训练数据集的严重不平衡(图 2),论文首先设计了一种采样策略,以同时在 200 多个语言对上训练单一模型
    • 直接从数据分布中采样会导致高资源语言表现良好,但低资源语言表现较差;
    • 平等采样所有语言对会显著提升低资源语言的翻译性能,但高资源语言的表现会明显低于其双语基线
  • 为了平衡高资源和低资源语言对,论文采用了训练多语言 BERT(2018)时使用的基于温度的采样策略
    • 对于语言对 \( l \),设 \( D_l \) 为可用平行语料库的大小
    • 若从数据集的并集中采样,样本来自语言 \( l \) 的概率为 \( p_l = \frac{1}{S_l D_l} \)
    • 论文将采样分布的概率设置为与 \( p_l^{\frac{1}{T} } \) 成正比,其中 \( T \) 为采样温度
      • 当 \( T = 1 \) 时,对应真实数据分布;
      • 当 \( T = 100 \) 时,几乎对每种语言采样数量相等
    • 论文的多语言模型采用了 \( T = 5 \)
  • 在所有多语言实验中,论文使用与双语模型相同的超参数,训练了一个同时在所有语言上训练的单一 Transformer 模型
  • 论文采用了共享的 SPM 词汇表,包含 64k 个词符,生成时使用了与训练相同的采样分布(\( T = 5 \))
  • 论文设置了字符覆盖率为 0.999995,以确保词汇表涵盖所有 103 种语言的大部分字母

Effects of Large Batch Size

  • 由于数据并行(Data Parallelism)的简单性,它成为扩展神经网络训练的主要方法(2016, 2017)
  • 论文通过显著增加标准 Transformer Big 训练的批量大小来测试大批量训练的极限
  • 从每批 260k 词符开始,论文将有效批量增加到 400 万,并观察高资源语言对(德英)的验证损失和 BLEU 分数(其他语言对也呈现相似趋势)
  • 优化参数与此前的实验相同
  • 据论文所知,400 万词符每批是文献中迄今为止用于训练 NMT 模型的最大批量(2018)
  • 表 2 显示,随着批量增加,两项指标均显著提升
  • 作者认为进一步增加批量可能带来更多改进

附录4:Discussion

  • 拥有一个灵活的大规模深度学习实验框架,为理解大规模模型的底层机制和原理提供了令人兴奋的机会
  • 本节中,论文将实验结果与深度学习领域的最新研究联系起来,并与从业者分享一些额外的实证发现
  • 表达力与泛化(Expressivity and Generalization) :
    • 深度学习理论的最新发现(2018, 2018)假设,随着网络表达能力的增长,泛化性能也会提升
      • 论文在此通过实验进行实证验证
      • 论文通过增加深度作为提高网络表达能力(2017)的手段,同时控制批量大小
        • 论文能够观察到网络在以往从未实验过的规模上的泛化行为
      • 从 6 层 Transformer Big(编码器+解码器共 12 层)开始,论文逐步将深度增加到 64 层(共 128 层)
        • 论文发现,64 层模型的表现几乎呈现出与 6 层模型相同的上升趋势,而中间深度模型的表现介于两者之间
      • 尽管结果支持理论,但论文也观察到收益递减现象,这引发了可训练性的担忧
      • 论文可能尚未掌握进一步降低训练误差的工具或技术,理解可训练性挑战对进一步发展是必要的
  • 深度-宽度权衡(Depth-Width Trade-off) :
    • 另一个吸引深度学习理论研究者关注的领域是模型宽度和深度对泛化的影响(2018, 2019)
    • 接下来,论文在多语言设置中研究深度与宽度的权衡,并比较 13 亿参数的宽模型 \( T(12, 16384, 32) \) 和 13 亿参数的深模型 \( T(24, 8192, 16) \) 的性能
    • 虽然这两种模型在高资源语言上的表现非常相似 ,但深模型在低资源语言上的表现显著优于宽模型 ,这表明增加模型深度可能更有利于泛化
    • 将 13 亿参数深模型与 4 亿参数模型进行比较时,低资源语言(图3 右侧)的性能提升几乎与高资源语言相当,这表明增加深度还可能扩大对低资源任务的迁移效果
  • 深度加速收敛(Faster Convergence with Depth) :
    • 论文报告了一个与模型深度相关的有趣现象
    • 在保持有效批量大小、优化器超参数和模型宽度不变的情况下 ,增加模型深度会优化加速 ,如附录图3 所示
      • 注:层数越深,损失下降越快
    • 此前报道类似现象的工作推测,深度通过过参数化实现了预条件(2018)
  • 尽管论文主要关注与扩展神经网络相关的系统挑战,但为了更好地解决泛化和可训练性问题,对深度神经网络的理论理解需求日益增长
    • 作者相信,像 GPipe 这样灵活的扩展工具对于弥合深度学习理论与实践的差距至关重要
    • 希望论文的实证发现和讨论能够激励更多研究沿着这一方向展开

附录:Sufficient Statistic(充分统计量)

  • 在统计学中,充分统计量 的本质是对原始样本数据的“压缩”
    • 即在保留所有关于未知参数的信息的前提下,将复杂的样本数据简化为更简洁的统计量
  • 这种简化不会丢失用于推断未知参数的任何关键信息,因此是统计推断中简化数据、提高效率的重要工具

定义

  • 假设
    • 总体的概率分布依赖于未知参数 \(\theta\)(\(\theta\)可以是单参数或多参数,比如分布的均值)
    • \(X_1, X_2, \dots, X_n\)是来自该总体的样本
  • 若
    • 一个统计量\(T = T(X_1, X_2, \dots, X_n)\) 满足:给定\(T\)的取值后,样本\(X_1, \dots, X_n\)的条件分布不再依赖于\(\theta\)
  • 则
    • 称\(T\)是参数\(\theta\)的充分统计量
  • 简单来说:一旦知道了充分统计量 \(T\) 的值,原始样本中就不再包含任何关于 \(\theta\) 的额外信息
    • 例如,伯努利分布中,\(T = \sum x_i\) 和样本均值 \(\bar{X} = T/n\)(因\(n\) 固定,二者一一对应)都是\(p\)的充分统计量

充分统计量的总结

  • 充分统计量的核心思想是:“压缩数据但不丢失信息”**
  • 统计推断的核心是利用样本信息推断未知参数 \(\theta\)
    • 原始样本往往包含大量冗余信息(例如,100个数据点中可能有重复或无关细节)
  • 充分统计量的作用是:
    • 简化数据 :将高维样本(如 \(n\) 个数据)压缩为低维统计量(如1个或2个值);
    • 保留信息 :压缩后的统计量包含推断 \(\theta\) 所需的全部信息,原始样本的其他细节对推断 \(\theta\) 无意义
  • 注:充分统计量不唯一性 :一个参数可能有多个充分统计量
    • 例如,伯努利分布中,\(T = \sum x_i\) 和样本均值 \(\bar{X} = T/n\)(因 \(n\) 固定,二者一一对应)都是 \(p\) 的充分统计量

附录:其他流水线策略介绍

各种流水线整体介绍

  • 三种流水线策略(pipeline strategies),Gpipe, PipeDream 1F1B, Interleaved 1F1B
  • PipeDream 1F1B 论文:
    • PipeDream: Fast and Efficient Pipeline Parallel DNN Training, 2018, Microsoft :第一篇,主要介绍概念等
    • PipeDream: Generalized Pipeline Parallelism for DNN Training, SOSP 2019, Microsoft:第二篇,包含更多细节
    • PipeDream 1F1B 也叫做经典 1F1B(One-Forward-One-Backward),也称为 1F1B 或 PipeDream 1F1B
  • Interleaved 1F1B(Megatron-2)论文:Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM, SC 2021, NVIDIA
  • 下图来自 Interleaved 1F1B(Megatron-2)论文:Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM, SC 2021, NVIDIA
  • 上图中内容解读:
    • 图中数字表示不同 micro-batch 的过程(共 8 个 micro-batches 和 4 张 GPU),颜色区分前向后向过程,Interleaved 1F1B 中的灰色和深色则表示多分出来的虚拟 Stage
    • 核心要诀:
      • 对每个 micro-batch,前向过程是必须 GPU 正序进行的
      • 对每个 micro-batch,后向过程是必须 GPU 逆序进行的
      • 对任意 micro-batch,需要一个完整的 Stage 完成前向过程,才能开始后向过程
      • 需要所有 micro-batch 都走完后向过程,才能执行一次梯度更新(所以叫做 1F1B)
        • 吐槽:实际上 Gpipe 也是一次正向一次反向的,也可以叫做 1F1B?
      • Interleaved 1F1B 中,同一个 Stage 内部的 Stage 分片是按照模型切分,所以需要以大的 Stage 为单位遵循上面的过程
        • 比如:对每个 micro-batch,需要所有的虚拟 Stage 分片完成前向过程,才能开始后向过程
      • Interleaved 1F1B 中,对同一个 Stage 内部的分片来说,理论上同一个 Stage 内部的 Stage 分片是按照模型切分,所以也是有依赖顺序的
      • 对每个 micro-batch,1F1B 保证激活值在反向结束后立即释放,降低缓存

GPipe -> 1F1B -> Interleaved 1F1B

  • TLDR:Interleaved 1F1B 是在经典 1F1B(PipeDream 1F1B)流水线并行调度基础上,把“模型纵向切成更多 Stage ,再循环交错地放到 GPU 上”,用更细粒度、更紧凑的时序进一步压缩“流水线气泡(bubble)”,从而提升 GPU 利用率与训练吞吐,同时保持 1F1B 的低显存优势
    • 注意:切割的是模型
  • GPipe, 1F1B, Interleaved 1F1B 三种流水线并行调度对比
    方案 层划分方式 执行顺序 主要优点 主要缺点
    GPipe 连续层,每 GPU 一个 Stage 先跑完所有 micro-batch 的 Forward,再依次跑 Backward 易实现、等价性好 激活值全保留,显存爆炸;气泡大
    1F1B 同上 每个 Stage 交替执行 1 Forward / 1 Backward 激活值及时释放,省显存;气泡已减半 气泡仍随 Stage 数线性增加
    Interleaved 1F1B 把原 Stage 均分 v 份,循环交错地放到 GPU 上 仍然是 1F1B 顺序,但 Stage 粒度变细 气泡进一步缩小;通信-计算可重叠 代码/通信更复杂;micro-batch 数需为 v·pp 的整数倍
    • 气泡比例(p 为 Pipeline Stage 数量,即物理 GPU 数, m 为 micro-batch 数量)
      • GPipe:(p−1) / m
      • 1F1B:(p−1) / m
      • Interleaved 1F1B: (p−1) / (v·m)
  • Interleaved 1F1B 的特点:
    • 时间轴上,Forward/Backward 任务被拆得更碎,空隙(bubble)被其他 micro-batch 的计算填满
    • 仍然遵守“1 Forward -> 1 Backward”的稳态节奏,保证激活值在反向结束后立即释放
    • 由于 Stage 变细,单个 micro-batch 在每个 GPU 上的停留时间变短,更多 micro-batch 可以同时在流水线上“滑动” ,
    • 当前 micro-batch 的 Backward 计算 与 下一个 micro-batch 的 P2P 接收 无依赖,可并行;
    • 当前 Send-Forward 与 下一次循环的 Forward 计算 也无依赖,可并行
    • 在带宽充足时,通信几乎被计算完全掩盖

Interleaved 1F1B 的优缺点

  • Interleaved 1F1B显存:与 1F1B 相同,稳态下每 GPU 最多保存 v 份激活 ,仍远低于 GPipe
  • 吞吐:Megatron-LM 实验显示,在 1T+ 参数模型、1024 GPU 上,Interleaved 1F1B 比 1F1B 提升 8~15% 吞吐;

其他相关细节

  • Megatron-LM 中 Interleaved 1F1B 叫做 interleaved schedule

最新技术:Zero Bubble Pipeline Parallelism

  • Zero Bubble Pipeline Parallelism(ZB),即 零气泡流水线并行
  • 参考链接:
    • 原始论文:Zero Bubble Pipeline Parallelism, arXiv 2023, Sea AI Lab
    • 开源地址:github.com/sail-sg/zero-bubble-pipeline-parallelism
  • 基本思路是让 GPU 不要停止,让气泡消失
    • 首先,如下图所示,将原始 1F1B 的后向过程拆开为 B 和 W 来表示(B 表示 对输入 \(x\) 求梯度,W 表示对权重 \(W\) 求梯度)
1…227228229…352
San Ye

San Ye

Stay Hungry. Stay Foolish.

704 posts
53 tags
© 2026 San Ye
Powered by Hexo
|
Theme — NexT.Gemini v5.1.4