扩散模型数学基础

一. 基础知识

1. 马尔可夫假设

马尔可夫过程（Markov Process）是一种随机过程，其中系统的未来状态只依赖于当前状态，而与过去的历史状态无关

这个假设使得复杂的前向与反向联合概率分布能够被分解为一系列简单条件概率的乘积，从而让模型可以通过训练一个网络来近似每一步的转移概率，最终实现高质量的数据生成。

2. 高斯分布的KL散度

KL散度（Kullback-Leibler Divergence）是一种衡量两个概率分布之间差异的指标，取值范围是 \([0, +\infty]\)，越小说明两个概率分布越相似。

KL散度的定义

KL散度的定义为两个概率分布 \(P\) 和 \(Q\) 之间的差异度量，用来量化从“真实分布”到“近似分布”所损失的平均信息量。定义如下：

\[D_{\text{KL}}(P \parallel Q) = \int_{-\infty}^{\infty} p(x) \log \frac{p(x)}{q(x)} \, dx \]

• 非负性：\(D_{KL}(P || Q) \ge 0\)，当且仅当 \(P = Q\) 时等于0。这保证了“距离”概念合理。
• 不对称性：\(D_{KL}(P || Q) \ne D_{KL}(Q || P)\)。即KL散度不是真正的“距离”，而是一种方向性的信息损失度量。在扩散模型中，我们总是用真实的后验分布 \(P\) 去衡量模型学习的近似分布 \(Q\)，方向固定。

高斯分布的概率密度函数

我们现在将 \(P\) 和 \(Q\) 设为高斯分布，分别表示为：

• \(P(x) = \mathcal{N}(x; \mu_P, \sigma_P^2)\)
• \(Q(x) = \mathcal{N}(x; \mu_Q, \sigma_Q^2)\)

\[p(x) = \frac{1}{\sqrt{2\pi \sigma_P^2}} \exp\left( -\frac{(x - \mu_P)^2}{2 \sigma_P^2} \right) \]

\[q(x) = \frac{1}{\sqrt{2\pi \sigma_Q^2}} \exp\left( -\frac{(x - \mu_Q)^2}{2 \sigma_Q^2} \right) \]

推导

1. 将高斯分布的密度函数代入KL散度的定义中：

   \[D_{\text{KL}}(P \parallel Q) = \int_{-\infty}^{\infty} p(x) \log \frac{p(x)}{q(x)} \, dx \]

2. 先计算 \(\log \frac{p(x)}{q(x)}\)：

   \[\log \frac{p(x)}{q(x)} = \log \left( \frac{ \frac{1}{\sqrt{2\pi \sigma_P^2}} \exp\left( -\frac{(x - \mu_P)^2}{2 \sigma_P^2} \right) }{ \frac{1}{\sqrt{2\pi \sigma_Q^2}} \exp\left( -\frac{(x - \mu_Q)^2}{2 \sigma_Q^2} \right) } \right) \]

3. 化简后得到：

   \[\log \frac{p(x)}{q(x)} = \log \frac{\sqrt{\sigma_Q^2}}{\sqrt{\sigma_P^2}} + \frac{(x - \mu_Q)^2}{2 \sigma_Q^2} - \frac{(x - \mu_P)^2}{2 \sigma_P^2} \]

4. 代入KL散度的积分公式：

   \[\begin{aligned} D_{\text{KL}}(P \parallel Q) &= \int_{-\infty}^{\infty} \frac{1}{\sqrt{2\pi \sigma_P^2}} \exp\left( -\frac{(x - \mu_P)^2}{2 \sigma_P^2} \right) \left[ \log \frac{\sqrt{\sigma_Q^2}}{\sqrt{\sigma_P^2}} + \frac{(x - \mu_Q)^2}{2 \sigma_Q^2} - \frac{(x - \mu_P)^2}{2 \sigma_P^2} \right] dx \\ &= \int_{-\infty}^{\infty} \frac{1}{\sqrt{2\pi \sigma_P^2}} \exp\left( -\frac{(x - \mu_P)^2}{2 \sigma_P^2} \right) \log \frac{\sqrt{\sigma_Q^2}}{\sqrt{\sigma_P^2}} \, dx \\ &\quad + \int_{-\infty}^{\infty} \frac{1}{\sqrt{2\pi \sigma_P^2}} \exp\left( -\frac{(x - \mu_P)^2}{2 \sigma_P^2} \right) \left[ \frac{(x - \mu_Q)^2}{2 \sigma_Q^2} - \frac{(x - \mu_P)^2}{2 \sigma_P^2} \right] dx \end{aligned} \]

5. 现在进行分步计算：

   对于第一个积分：

   \[\begin{aligned} \text{第一个积分} &= \int_{-\infty}^{\infty} \frac{1}{\sqrt{2\pi \sigma_P^2}} \exp\left( -\frac{(x - \mu_P)^2}{2 \sigma_P^2} \right) \log \frac{\sqrt{\sigma_Q^2}}{\sqrt{\sigma_P^2}} \, dx \\ &= \log \frac{\sqrt{\sigma_Q^2}}{\sqrt{\sigma_P^2}} \int_{-\infty}^{\infty} \frac{1}{\sqrt{2\pi \sigma_P^2}} \exp\left( -\frac{(x - \mu_P)^2}{2 \sigma_P^2} \right) dx \\ &= \log \frac{\sqrt{\sigma_Q^2}}{\sqrt{\sigma_P^2}} \\ &= \log \frac{\sigma_Q}{\sigma_P} \end{aligned} \]

   对于第二个积分：

   1. 首先将中括号内平方项展开：

      \[(x - \mu_Q)^2 = x^2 - 2x\mu_Q + \mu_Q^2 \]
      \[(x - \mu_P)^2 = x^2 - 2x\mu_P + \mu_P^2 \]

   2. 代入积分中整理得：

      \[\int_{-\infty}^{\infty} \frac{1}{\sqrt{2\pi \sigma_P^2}} \exp\left( -\frac{(x - \mu_P)^2}{2 \sigma_P^2} \right) \left[ \left( \frac{1}{2 \sigma_Q^2} - \frac{1}{2 \sigma_P^2} \right) x^2 + \left( \frac{\mu_P}{\sigma_P^2} - \frac{\mu_Q}{\sigma_Q^2} \right) x + \left( \frac{\mu_Q^2}{2 \sigma_Q^2} - \frac{\mu_P^2}{2 \sigma_P^2} \right) \right] dx \]

   3. 积分中 \(x^2\) 的项是高斯分布的二阶矩；积分中 \(x\) 的项是高斯分布的一阶矩（期望）：

      \[\text{二阶矩: } \int_{-\infty}^{\infty} \frac{1}{\sqrt{2\pi \sigma_P^2}} \exp\left( -\frac{(x - \mu_P)^2}{2 \sigma_P^2} \right) x^2 \, dx = \sigma_P^2 + \mu_P^2 \]
      \[\text{一阶矩: } \int_{-\infty}^{\infty} \frac{1}{\sqrt{2\pi \sigma_P^2}} \exp\left( -\frac{(x - \mu_P)^2}{2 \sigma_P^2} \right) x \, dx = \mu_P \]

      所以得到：

      \[\begin{aligned} \text{第二个积分} &= \frac{1}{2}\left( \frac{1}{\sigma_Q^2} - \frac{1}{\sigma_P^2} \right) (\sigma_P^2 + \mu_P^2) \\ &\quad + \left( \frac{\mu_P}{\sigma_P^2} - \frac{\mu_Q}{\sigma_Q^2} \right) \mu_P \\ &\quad + \frac{1}{2}\left( \frac{\mu_Q^2}{\sigma_Q^2} - \frac{\mu_P^2}{\sigma_P^2} \right) \end{aligned} \]

   4. 整理后可得：

      \[\begin{aligned} \text{第二个积分} &= \frac{1}{2} \left[ \frac{\sigma_P^2 + \mu_P^2}{\sigma_Q^2} - 1 - \frac{\mu_P^2}{\sigma_P^2} \right] + \frac{1}{2} \left( \frac{2\mu_P^2}{\sigma_P^2} - \frac{2\mu_P \mu_Q}{\sigma_Q^2} \right) + \frac{1}{2} \left( \frac{\mu_Q^2}{\sigma_Q^2} - \frac{\mu_P^2}{\sigma_P^2} \right) \\ &= \frac{\sigma_P^2 + \mu_P^2 - 2\mu_P \mu_Q + \mu_Q^2}{2\sigma_Q^2} - \frac{1}{2} \\ &= \frac{\sigma_P^2 + (\mu_P - \mu_Q)^2}{2\sigma_Q^2} - \frac{1}{2} \end{aligned} \]

6. 将两个积分结果相加：

   \[D_{\text{KL}}(P \parallel Q) = \log \frac{\sigma_Q}{\sigma_P} + \frac{\sigma_P^2 + (\mu_P - \mu_Q)^2}{2\sigma_Q^2} - \frac{1}{2} \]

3. 高斯分布的一阶矩与二阶矩

由方差 \(\mathrm{Var}(X) = E[X^2] - (E[X])^2\)：

• \(E[X^2]\)：二阶矩
• \(E[X]\)：均值

所以二阶矩：

\[E(x^2) = \mathrm{Var}(X) + (E[X])^2 = \mu^2 + \sigma^2 \]

4. 重参数化技巧

对于扩散模型中的 \(x_t\)，是先将向上一步的 \(x_{t-1}\) 添加高斯噪声，再从中采样得到。但是直接采样是不可导的过程（采样是随机选择，没有可微分的数学运算），这会导致反向传播时梯度无法传递，模型无法训练。所以我们需要使用重参数化技巧使其可导。即将 “从带参数的高斯分布中采样” 的过程，拆解为 “从标准高斯分布采样 + 线性变换” 的可导形式。

对于高斯分布 \(P = \mathcal{N}(x; \mu_P, \sigma_P^2)\)，如果想要得到 \(P\)，就可以先从标准高斯分布中采样得到 \(z \sim \mathcal{N}(0, 1)\)。

再由

• \(a \cdot z \sim \mathcal{N}(0, a^2)\)，\(a\) 是常数
• \(c + z \sim \mathcal{N}(c, 1)\)，\(c\) 是常数

得：

\[P = \sigma_P \cdot z + \mu_P \]

4. 基于马尔科夫假设的条件概率

\[P(X_0, X_1, \ldots, X_n) = P(X_n \mid X_{n-1}, \ldots, X_1, X_0) \, P(X_{n-1}, \ldots, X_1, X_0) \]

由马尔科夫假设，有：

\[P(X_n \mid X_{n-1}, \ldots, X_1, X_0) = P(X_n \mid X_{n-1}) \]

所以：

\[P(X_0, X_1, \ldots, X_n) = P(X_n \mid X_{n-1}) \, P(X_{n-1}, \ldots, X_1, X_0) \]

依次展开：

\[\begin{aligned} P(X_0, X_1, \ldots, X_n) &= P(X_n \mid X_{n-1}) \, P(X_{n-1} \mid X_{n-2}) \cdots P(X_1 \mid X_0) \, P(X_0) \\ &= P(X_0) \prod_{t=1}^{n} P(X_t \mid X_{t-1}) \end{aligned} \]

对于条件概率：

\[\begin{aligned} P(X_1, X_2, \ldots, X_n \mid X_0) &= \frac{P(X_n \mid X_{n-1}) \, P(X_{n-1} \mid X_{n-2}) \cdots P(X_1 \mid X_0) \, P(X_0)}{P(X_0)} \\ &= \prod_{t=0}^{n-1} P(X_{t+1} \mid X_t) \end{aligned} \]