DreamFusion

1. 使用神经网路进行数据生成

使用神经网络生成一个高维度数据是机器学习中非常重要的一个工作。我们假设数据集 $\left\{\boldsymbol{x_1}, \boldsymbol{x_2}, \dots, \boldsymbol{x_n}\right\}$ 为一个大小为$n$的数据集，该数据集统一的服从一个概率分布 $p_{data}(\boldsymbol{x})$ 。我们假设对数据集的抽样都是独立同分布的，即：

$$ \left\{\boldsymbol{x_1}, \boldsymbol{x_2}, \dots, \boldsymbol{x_n}\right\} \sim p_{data}(\boldsymbol{x}) $$

那么丛现有数据生成新的数据的核心就是使用神经网络学习这个概率分布。不妨假设学习的概率分布为 $\hat p_\theta(\boldsymbol x)$。我们会希望 $\hat p_\theta(\boldsymbol x)$ 尽可能的接近 $p_{data}(\boldsymbol(x))$ 。为了衡量真是分布和我们学习的分布之间的差距，我们需要定义一个距离函数 $D(\cdot \mid \cdot)$ 我们可以定义优化目标：

$$ \hat \theta = \arg \min_{\theta} D\left(p_{data}(\boldsymbol{x}) \mid \hat p_\theta(\boldsymbol x) \right) $$

关于距离函数，我们可以定义 $D(\cdot \mid \cdot)$ 为 f-divergence 定义为：

$$ D_f(p_{data}(\boldsymbol(x)) \mid \hat p_\theta(\boldsymbol x)) = \int p_\theta(\boldsymbol x) f \left(\frac{p_{data}(\boldsymbol x)}{p_\theta(\boldsymbol x)}\right) d\boldsymbol x $$

不妨取 $f(x) = x\log x$ ，我们可以得到 KL 散度：

$$ \begin{aligned} D_{KL}(p_{data}(\boldsymbol(x)) \mid \hat p_\theta(\boldsymbol x)) &= \int p_{data}(\boldsymbol x) \log \frac{p_{data}(\boldsymbol x)}{p_\theta(\boldsymbol x)} d\boldsymbol x\\ &= \mathbb E_{p_{data}(\boldsymbol x)}\left[ \log \frac{p_{data}(\boldsymbol x)}{p_\theta(\boldsymbol x)} \right] \end{aligned} $$

我们可以用抽样的均值来代替期望，有：

$$ \begin{aligned} D_{KL}(p_{data}(\boldsymbol(x)) \mid \hat p_\theta(\boldsymbol x)) &\approx \frac{1}{n} \sum_{i=1}^n \log \frac{p_{data}(\boldsymbol x_i)}{p_\theta(\boldsymbol x_i)}\\ &= \frac{1}{n} \sum_{i=1}^n \log p_{data}(\boldsymbol x_i) - \frac{1}{n} \sum_{i=1}^n\log p_\theta(\boldsymbol x_i)\\ &\stackrel{i}{=} -\frac{1}{n} \sum_{i=1}^n\log p_\theta(\boldsymbol x_i) + \text{const} \end{aligned} $$

其中，因为 $\frac{1}{n} \sum_{i=1}^n \log p_{data}(\boldsymbol x_i)$ 是一个与优化参数无关的分布，我们可以直接抛弃。那么显然的，我们可以通过下面的方式训练我们的生成式模型的概率分布，即在所有的数据点上计算最大似然：

$$ \hat \theta = \arg \max_{\theta} \frac{1}{n} \sum_{i=1}^n\log p_\theta(\boldsymbol x_i) $$

但是问题没有解决。在优化过程中，我们需要保证训练的生成式神经网络为一个概率分布，那么需要满足概率分布的基本条件：

1. 非负性： $\forall \boldsymbol x , p_\theta(\boldsymbol x) > 0$
2. 归一化性质： $\int p_\theta(\boldsymbol x) d\boldsymbol x = 1$

这两个条件都不容易满足，但是和归一化性质需要计算全概率的积分相比，非负性都显得不那么困难了。因此我们需要思考一些数学技巧来摆脱归一化条件的桎梏。

1.1. 归一化常数近似

第一个方法收到了玻尔兹曼函数的启发。我们使用神经网络学习一个标量的能量函数 $E_\theta(\boldsymbol x)$ ，有概率分布：

$$ \begin{array}{c} p_\theta(\boldsymbol x) = \frac{\exp \left(-E_\theta(\boldsymbol x)\right)}{Z_\theta}\\ \text{where: }Z_\theta = \int \exp \left(-E_\theta(\boldsymbol x)\right) d\boldsymbol x \end{array} $$

显然其中能量函数 $E_\theta(\boldsymbol x)$ 不需要受到概率分布的性质约束。对于前面的最大自然函数，带入玻尔兹曼函数的概率分布，有：

$$ \begin{aligned} \hat \theta &= \arg \max_{\theta} \frac{1}{n} \sum_{i=1}^n\log p_\theta(\boldsymbol x_i)\\ &= \arg \max_{\theta} -\frac{1}{n} \sum_{i=1}^n E_\theta(\boldsymbol x) - \log Z_\theta\\ &= \arg \min_{\theta} \frac{1}{n} \sum_{i=1}^n E_\theta(\boldsymbol x) + \log Z_\theta \end{aligned} $$

不妨定义损失函数为：

$$ \mathcal L(\theta) = \frac{1}{n} \sum_{i=1}^n E_\theta(\boldsymbol x) + \log Z_\theta $$

我们有：

$$ \nabla_\theta \mathcal L = \frac{1}{n} \sum_{i=1}^n \nabla_\theta E_\theta(\boldsymbol x) + \nabla_\theta\log Z_\theta $$

其中 $\nabla_\theta E_\theta(\boldsymbol x)$ 的计算是简单的，问题是后一项。

$$ \begin{aligned} \nabla_\theta\log Z_\theta &= \nabla_\theta \log \int \exp \left(-E_\theta(\boldsymbol x)\right) d\boldsymbol x\\ &= -\frac{\int \exp \left(-E_\theta(\boldsymbol x)\right) \nabla E_\theta (\boldsymbol x)d \boldsymbol x}{\int \exp \left(-E_\theta(\boldsymbol x)\right) d\boldsymbol x}\\ &= -\int \frac{\exp \left(-E_\theta(\boldsymbol x)\right)}{\int \exp \left(-E_\theta(\boldsymbol x)\right) d\boldsymbol x} \nabla E_\theta (\boldsymbol x)d \boldsymbol x\\ &= -\int p_\theta(\boldsymbol x) \nabla E_\theta (\boldsymbol x)d \boldsymbol x\\ &= E_{p_\theta (\boldsymbol x)}\left[\nabla_\theta E_\theta (\boldsymbol x)\right] \end{aligned} $$

在训练中，我们解决了归一化的问题，但是在推理中并没能避免。此时一般会用蒙特卡洛方法计算。

1.2. Autoregress Module (自回归模型)

通过概率的链式法则，我们可以将一个高维度的概率分布分解为多个低维度的概率分布的乘积。这种方法被称为自回归模型。例如：

$$ \begin{aligned} p_\theta\left(x_1, x_2, \cdots, x_n\right) &= p_\theta(x_1) p_\theta(x_2\mid x_1) p_\theta(x_3\mid x_1, x_2) \cdots p_\theta(x_n\mid x_1, x_2, \cdots, x_{n-1})\\ &= \prod_{i=1}^n p_\theta(x_i\mid x_{i-1}, x_{i-2}, \cdots, x_1) \end{aligned} $$

那么归一化性质应有：

$$ \begin{aligned} \int p_\theta\left(x_1, x_2, \cdots, x_n\right) d\boldsymbol x &= \int \prod_{i=1}^n p_\theta(x_i\mid x_{i-1}, x_{i-2}, \cdots, x_1) d\boldsymbol x\\ &= \int p_\theta(x_1) d\boldsymbol x \int p_\theta(x_2\mid x_1) d\boldsymbol x \cdots \int p_\theta(x_n\mid x_1, x_2, \cdots, x_{n-1}) d\boldsymbol x\\ \end{aligned} $$

此时，我们不需要计算一个很高维度的概率分布的积分，只需要对一些简单的低维概率分布归一化就好了。事实上自回归模型的训练更加巧妙，对于 $n$ 维数据生成而言，我们不需要同时训练 $n$ 个神经网络，而是训练一个神经网络，结合一个编码位置或者无效数据掩码来实现。

1.3. Normalizing Flows (归一化流模型)

假设存在一个平凡的概率分布 $\pi(\boldsymbol z)$，例如我们可以取 $\pi(\boldsymbol z) \sim \mathcal G(\boldsymbol z)$。我们希望存在一个可逆函数 $\boldsymbol x = f_\theta(\boldsymbol z)$，实现潜变量 $\boldsymbol z$ 与数据空间 $\boldsymbol x$ 之间的转换。那么根据概率的链式法则，我们有：

$$ p_\theta(\boldsymbol x) = \pi(\boldsymbol z) = \pi\left(f_\theta^{-1}(\boldsymbol x)\right) | \det \boldsymbol J_{f_\theta^{-1}}(\boldsymbol x)| $$

其中 $\boldsymbol J_{f_\theta^{-1}}(\boldsymbol x)$ 为 $f_\theta^{-1}(\boldsymbol x)$ 的雅可比矩阵。

因为 $\pi(\boldsymbol z)$ 是一个归一化的概率函数，那么 $p_\theta(\boldsymbol x)$ 也是一个归一化的概率函数。但是 Normalizing FLow Module 对函数 $f_\theta(\boldsymbol z)$ 的要求很苛刻，要求它同时满足可逆性和雅可比矩阵的计算是简单的。因此 $f_\theta(\boldsymbol z)$ 的选取往往是很简单的。一般会多层嵌套来使用：

$$ p(x) = p_z(z_0) \prod_{i=1}^{n} \left| \det \frac{\partial f_i^{-1}(z_{i-1})}{\partial z_i} \right| $$

本质上讲，Flow Module 学习了一组可逆且非线性的映射，讲一个复杂的分布映射到了一个简单的数据空间。具体的例子可以参考 Normalizing Flows。

1.4. Variational Autoencoder（变分自编码）

对于原始数据 $\boldsymbol{x}$，和自编码器类似，我们希望可以估计出两个分布 $p_\theta(\boldsymbol{z}\mid \boldsymbol{x})$ 和 $p_\theta(\boldsymbol{x} \mid \boldsymbol{z})$，将原始数据分布变换到一个隐藏层 $\boldsymbol{z}$ 的分布上。

我们希望隐藏层 $\boldsymbol{z}$ 的分布是简单的，例如 $p(\boldsymbol{z}) \sim \mathcal G$。事实上这个思路与自编码器十分类似，主要区别在于我们引入了对概率分布的采样来描述数据的生成过程。那么，VAE 的推断过程我们可以总结为：

1. 首先，从 $\boldsymbol{z}$ 中采样一个先验分布 $p_\theta(\boldsymbol{z})$
2. 通过神经网络预测 $x$ 的分布 $p_\theta(\boldsymbol{x}) = \int p_\theta(\boldsymbol{x}\mid \boldsymbol{z})p(\boldsymbol{z})d\boldsymbol{z}$

训练时可以直接使用最大似然：

$$ \theta = \arg \max_{\theta} \frac{1}{n} \sum_{i=1}^n \log p_\theta(\boldsymbol{x}_i) $$

那么现在唯一的问题是，$p_\theta(\boldsymbol{x}) = \int p_\theta(\boldsymbol{x}\mid \boldsymbol{z})p(\boldsymbol{z})d\boldsymbol{z}$该怎么算。下面我们需要对这个问题进行一定的化简，即变分下界或者证据下界。

变分下界是变分贝叶斯中的一个重要问题，它源于贝叶斯推断中一个十分重要的问题，我们如何计算或近似给定数据的模型证据（也称为边缘似然）$p(x)$，即在所有可能的潜在变量 $z$ 上积分模型概率 $p(x,z)$。显然的，对于复杂模型，这个积分的形式是极其复杂的。变分推断则引入了一个近似于后验分布 $q_\phi(\boldsymbol{z} \mid \boldsymbol{x}) \approx p(\boldsymbol{z} \mid \boldsymbol{x})$，并且定义了一个新的函数 ELBO，证明了它是模型证据 $p(\boldsymbol{x})$ 的下界。

对于原始 $p_\theta(\boldsymbol{x})$，我们有：

$$ \begin{aligned} p_\theta(\boldsymbol{x}) &= \int p_\theta(\boldsymbol{x}, \boldsymbol{z}) d\boldsymbol{z} &= \int \frac{p_\phi(\boldsymbol{z} \mid \boldsymbol{x})}{p_\phi(\boldsymbol{z} \mid \boldsymbol{x})}p_\theta(\boldsymbol{x}, \boldsymbol{z}) d\boldsymbol{z} \end{aligned} $$

此处，我们使用贝叶斯定理分解 $p_\theta(\boldsymbol{x}, \boldsymbol{z})$，有：

$$ \int p_\phi(\boldsymbol{z} \mid \boldsymbol{x})\frac{p_\theta(\boldsymbol{x} \mid \boldsymbol{z}) p(\boldsymbol{z})}{p_\phi(\boldsymbol{z} \mid \boldsymbol{x})} d\boldsymbol{z} $$

此处我们引如 Jensen 下界：

$$ \log \int p_\phi(\boldsymbol{z} \mid \boldsymbol{x})\frac{p_\theta(\boldsymbol{x} \mid \boldsymbol{z}) p(\boldsymbol{z})}{p_\phi(\boldsymbol{z} \mid \boldsymbol{x})} d\boldsymbol{z} \geq \int p_\phi(\boldsymbol{z} \mid \boldsymbol{x}) \log\frac{p_\theta(\boldsymbol{x} \mid \boldsymbol{z}) p(\boldsymbol{z})}{p_\phi(\boldsymbol{z} \mid \boldsymbol{x})} d\boldsymbol{z} $$

分解积分项，可以得到：

$$ \begin{aligned} \log p_\theta(\boldsymbol{x}) &\geq \int p_\phi(\boldsymbol{z} \mid \boldsymbol{x}) \log\frac{p_\theta(\boldsymbol{x} \mid \boldsymbol{z}) p(\boldsymbol{z})}{p_\phi(\boldsymbol{z} \mid \boldsymbol{x})} d\boldsymbol{z}\\ &= \int p_\phi(\boldsymbol{z} \mid \boldsymbol{x})\log p_\theta(\boldsymbol{x} \mid \boldsymbol{z}) d\boldsymbol{z}+ \int p_\phi(\boldsymbol{z} \mid \boldsymbol{x}) \log\frac{p(\boldsymbol{z})}{p_\phi(\boldsymbol{z} \mid \boldsymbol{x})} d\boldsymbol{z}\\ &= \mathbb E_{p_\phi(\boldsymbol{z} \mid \boldsymbol{x})}\left[\log p_\theta(\boldsymbol{x} \mid \boldsymbol{z})\right] - \mathbb D_{KL}\big(p_\phi(\boldsymbol{z} \mid \boldsymbol{x}) \mid\mid p(\boldsymbol{z})\big) \end{aligned} $$

这就是ELBO的数学定义：

$$ ELBO(\theta, \phi) = \mathbb E_{p_\phi(\boldsymbol{z} \mid \boldsymbol{x})}\left[\log p_\theta(\boldsymbol{x} \mid \boldsymbol{z})\right] - \mathbb D_{KL}\big(p_\phi(\boldsymbol{z} \mid \boldsymbol{x}) \mid\mid p(\boldsymbol{z})\big) $$

1.5. Score Function Based Data Generating

前面已经介绍了很多解决概率化模型的方法，但是往往对模型有很多甚至有些过分苛刻的要求。下面将介绍 Score Function Based Method （基于得分函数的模型）。事实上这也是常见的 Diffusion 方法所使用的方法。对于一个概率分布 $p_\theta(\boldsymbol{x})$，可以定义得分函数：

$$ s_\theta(\boldsymbol{x}) = \nabla_{\boldsymbol{x}} \log p(\boldsymbol{x}) $$

而 Score Function Based Method 则认为是通过学习的分函数来生成数据的方法。求前面提到的玻尔兹曼函数的得分函数，有：

$$ \begin{aligned} s_\theta(\boldsymbol{x}) &= \nabla_{\boldsymbol{x}} \log p(\boldsymbol{x})\\ &= -\nabla_{\boldsymbol{x}} f_\theta(\boldsymbol{x}) - \nabla_x\log Z_\theta\\ &= -\nabla_x f_\theta(\boldsymbol{x}) \end{aligned} $$

可以看到求得的的分函数并不依赖于归一化常数。这是一个非常重要的性质，显然的，使用 Score Funcction Based Method 时，我们没有对概率分布 $p_\theta(\boldsymbol x)$ 的性质有任何假设，这极大的扩展了模型的适用范围。与基于似然函数的模型类似，我们可以使用费歇尔信息（Fisher）信息来衡量数据集与模型的分数函数之间的相似水平：

$$ \mathbb E \left[\| \nabla_{\boldsymbol{x}} \log p(\boldsymbol{x}) - s_\theta(\boldsymbol{x}) \|^2\right] $$

使用费歇尔信息来衡量两个分布之间的 $\mathcal{l}^2$ 距离是一个很直觉的方法，但是它并不易于操作，因为我们需要计算一个未知的分布的得分函数 $\nabla_{\boldsymbol{x}} \log p(\boldsymbol{x})$。幸运的是，我们有一种名为 Score Match 的方法在不需要求 $\nabla_{\boldsymbol{x}} \log p(\boldsymbol{x})$ 时最小化费歇尔信息。费希尔散度本身并不要求 $s_\theta(\boldsymbol{x})$ 是任何归一化分布的实际得分函数，它只是比较真实数据得分与基于得分的模型之间的 $\mathcal l ^2$ 距离。事实上我们对 $s_\theta(\boldsymbol{x})$ 的唯一要求只是它的维度和真实分布相同而已。

1.6. Langevin Dynamics

在学习了得分函数 $s_\theta(\boldsymbol{x}) \approx \nabla_{\boldsymbol{x}} \log p(\boldsymbol{x})$ 后，我们还需要将其转换回原本的概率分布，这实际上是要解决一个随机微分方程。这同样并不容易，但是我们可以直接跳过求原本的分布这一步，直接从 Score Function 采样，这个方法被称为郎之万动力学 (Langevin Dynamics)。郎之万动力学提供了一种马尔可夫的方法；具体而言，从任意先验分布 $x_0 \sim \pi(x)$ 初始化，然后进行如下迭代：

$$ \begin{array}{ll} x_{i+1} = x_i + \epsilon \nabla_{\boldsymbol{x}} \log p(\boldsymbol{x}) + \sqrt{2\epsilon} \boldsymbol{z}_i & i = 0, 1, \cdots K \end{array} $$

其中 $\boldsymbol{z_i} \sim \mathcal N(0, I)$，$\epsilon \to 0$, $K\to \infty$。实际上这个行为类似与数值微分方程的求解，后面的高斯分布项可以有效的避免陷入局部最优。事实上只要这个过程足够，误差几乎可以忽略。

2. Diffusion

推荐大家看Lil大佬的博客，我后面的内容知识对这篇博客的拙略模仿。

2.1. Forward diffusion process

好了，现在我们已经掌握了 diffusion 的所有原料，可以开始进入 diffusion 的世界了。假设我们有一个从真实数据分布采样的数据 $\boldsymbol{x}_0 \sim q(\boldsymbol{x})$。我们定义一步的 diffusion 向前过程为向数据集添加一个很小的高斯噪声，有 $\beta \in (0, 1)$

$$ p(x_t \mid x_{t-1}) = \mathcal N (x_t \mid \sqrt{1 - \beta} x_{t-1}, \beta_t I) $$

我们需要酱这个过程重复 $T$ 次，从而产生一个数据序列：$X_1, X_2, \cdots, X_T$。只要步骤数 $T$ 足够多，我们可以获得一个从原始数据缓慢增加噪声，直到一个服从各向异性的高斯分布。

对于上面的加噪序列中的任意一时刻 $t$，令 $\alpha_t = 1 - \beta_t$，我我们可以直接卸出其加噪之后的形式:

$$ \begin{aligned} x_t &= \sqrt{\alpha_t} x_{t-1} + \sqrt{1 - \alpha_t}\epsilon_{t-1}\\ &= \sqrt{\alpha_t \alpha_{t-1}} x_{t-2} + \sqrt{1 - \alpha_t\alpha_{t-1}}\epsilon_{t-2}\\ &= \cdots\\ &= \sqrt{\hat \alpha_t} x_0 + \sqrt{1 - \hat \alpha_t}\epsilon \end{aligned} $$

其中 $\epsilon_i \sim \mathcal N(0, I)$, $\hat \alpha_t = \prod_{i=1}^t \alpha_i$。我们可以写出，对于序列中任意一帧 $x_t$，有分布：

$$ p(x_t\mid x_0) = \mathcal N(x_t\mid \sqrt{\hat \alpha_t} x_0, \sqrt{1 - \hat \alpha_t} I) $$

2.2. Reverse diffusion process

如果说向前过程是向无噪声的数据增加噪声的过程，逆向过程就是一个去除噪声的过程。显然的，如果我们掌握了 $p(x_{t-1} \mid x_t)$，那么可以从中抽样出向前过程中所添加的噪声，那么简单的减去这个噪声，就实现了数据集的降噪。事实上如果加噪的参数 $\beta$ 足够小，应有 $p(x_{t-1} \mid x_t)$ 服从高斯分布。事实上 diffusion 通过增加高斯噪声的方式将一个很难的任务（直接估计高维分布）变为多个简单的任务（逐渐降低一个数据的噪声），从而实现的高维数据估计。那么逆向过程可以写成：

$$ \begin{array}{c} p_\theta(\mathbf{x}_{0:T}) = p(\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}; \boldsymbol{\mu}_\theta(\mathbf{x}_t, t), \boldsymbol{\Sigma}_\theta(\mathbf{x}_t, t)) \end{array} $$

在知道 $x_0$ 的情况下，通过贝叶斯法则，有：

$$ p(x_{t-1}\mid x_t, x_0) = p(x_t\mid x_{t-1}, x_0) \frac{p(x_{t-1} \mid x_0)}{p(x_{t}\mid x_0)} $$

这一步将向后过程转换为向前的过程，向前过程我们都是已知的，有：

$$ \begin{array}{rcl} p(x_t\mid x_{t-1}, x_0) &\sim& \mathcal N(x_t \mid \sqrt{\alpha_t} x_{t-1}, \beta_t I)\\ p(x_{t-1}\mid x_0)&\sim& \mathcal N(x_{t-1} \mid \sqrt{\hat \alpha_{t-1}} x_0, (1 - \hat \alpha_{t-1}) I)\\ p(x_{t}\mid x_0)&\sim& \mathcal N(x_t \mid \sqrt{\hat \alpha_{t}} x_0, (1 - \hat \alpha_t) I)\\ \end{array} $$

那么我们可以直接将逆向的形式展开：

$$ \begin{aligned} &p(x_t\mid x_{t-1}, x_0)\\ &\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{\alpha_t}{\beta_t} + \frac{1}{1 - \bar{\alpha}_{t-1}}) \mathbf{x}_{t-1}^2 - (\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} + C(\mathbf{x}_t, \mathbf{x}_0) \big) \Big)\\ \end{aligned} $$

其中 $C(\mathbf{x}_t, \mathbf{x}_0)$ 是不包含 $x_{t-1}$ 的函数。对于一般的高斯函数 $\mathcal G(x\mid \mu, \sigma) \propto \exp\left(-\frac 1 2 \left(\frac{1}{\sigma^2}x^2 + \frac{2\mu}{\sigma^2} + \frac{\mu^2}{\sigma^2}\right)\right)$。不妨先假设：

$$ \begin{array}{c} q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0) = \mathcal{N}(\mathbf{x}_{t-1}; \tilde{\boldsymbol{\mu}}(\mathbf{x}_t, \mathbf{x}_0), \tilde{\beta}_t \mathbf{I}) \end{array} $$

与上面的形式对应，我们可以总结出：

$$ \begin{aligned} \frac{1}{\sigma^2} &= \frac{1}{\tilde{\beta}_t} = \frac{\alpha_t}{\beta_t} + \frac{1}{1 - \bar{\alpha}_{t-1}}\\ \tilde{\beta}_t &= \frac{1 - \bar{\alpha}_{t-1}}{1 - \bar{\alpha}_t} \beta_t\\ \frac{\mu}{\sigma^2} &= \frac{\mu(x_t, x_0)}{\tilde{\beta_t}}=\frac{\sqrt{\alpha_t}}{\beta_t} \mathbf{x}_t + \frac{\sqrt{\bar{\alpha}_{t-1}}}{1 - \bar{\alpha}_{t-1}} \mathbf{x}_0\\ \mu(x_t, x_0) &= \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} $$

根据 $x_t = \sqrt{\hat \alpha_t} x_0 + \sqrt{1 - \hat \alpha_t}\epsilon$，有：

$$ x_0 = \frac{1}{\sqrt{\bar\alpha_t}}\left(x_t - \sqrt{1 - \bar\alpha_t} \epsilon_t\right) $$

带入上面求出的参数，有：

$$ \begin{aligned} \tilde{\boldsymbol{\mu}}_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}_t) \\ &= \frac{1}{\sqrt{\alpha_t}} \Big( \mathbf{x}_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar{\alpha}_t}} \boldsymbol{\epsilon}_t \Big) \end{aligned} $$

其中 $\epsilon_t$ 为模型需要预测的用于降噪的噪声。现在问题似乎解决了，DDPM（denoising diffusion probabilistic models）的推理可以总结成下面几个步骤：

1. 通过训练好的神经网络，结合 $x_t, t$（作为参数）预测高斯噪声 $\epsilon_t$，通过 $\tilde{\mu_t} = \frac{1}{\sqrt{\alpha_t}} \Big( \mathbf{x}_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar{\alpha}_t}} \boldsymbol{\epsilon}_t \Big)$ 求出均值。

2. 根据 $\tilde{\beta}_t = \frac{1 - \bar{\alpha}_{t-1}}{1 - \bar{\alpha}_t} \beta_t$ 得到方差。在 GLIDE 中，则通过神经网络预测出该方差。

3. 根据 $q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0) = \mathcal{N}(\mathbf{x}_{t-1}; \tilde{\boldsymbol{\mu}}(\mathbf{x}_t, \mathbf{x}_0), \tilde{\beta}_t \mathbf{I})$ 实现去噪。

2.3. Training

在了解了 diffusion 的推断过程后，下一个问题是如何训练我们的模型让它预测出正确的 $\mu_\theta(x_t, t)$ 和 $\Sigma_\theta(x_t, t)$。可以回忆一下在 VAE 介绍过的ELBO，此处亦有应用。由于 KL 散度非负，总有：

$$ \begin{aligned} -\log p_\theta(x) &\leq -\log p_\theta(x) + D_{KL}\left(q_\phi(\boldsymbol{x}_{1:T} \mid \boldsymbol{x}_0) \mid p_\theta(\boldsymbol{x}_{1:T} \mid \boldsymbol{x}_0)\right)\\ &= -\log p_\theta(x) + \mathbb E_{q_\phi(\boldsymbol{x}_{1:T} \mid \boldsymbol{x}_0)}\left[\log \frac{q_\phi(\boldsymbol{x}_t \mid \boldsymbol{x}_0)p_\theta(\boldsymbol{x}_0)}{p_\theta(\boldsymbol{x}_{0:T}) }\right]\\ &= -\log p_\theta(x) + \mathbb E_{q_\phi(\boldsymbol{x}_{1:T} \mid \boldsymbol{x}_0)}\left[\log \frac{q_\phi(\boldsymbol{x}_t \mid \boldsymbol{x}_0)}{p_\theta(\boldsymbol{x}_{0:T}) } + \log p_\theta(\boldsymbol{x}_0)\right]\\ &= \mathbb E_{q_\phi(\boldsymbol{x}_{1:T} \mid \boldsymbol{x}_0)}\left[\log \frac{q_\phi(\boldsymbol{x}_t \mid \boldsymbol{x}_0)}{p_\theta(\boldsymbol{x}_{0:T})}\right] \end{aligned} $$

令：

$$ \boldsymbol{L}_{VIB} = \mathbb E_{q_\phi(\boldsymbol{x}_{1:T} \mid \boldsymbol{x}_0)}\left[\log \frac{q_\phi(\boldsymbol{x}_t \mid \boldsymbol{x}_0)}{p_\theta(\boldsymbol{x}_{0:T})}\right] \geq -\mathbb E_{q}\left[\log p_\theta(x)\right] $$

为了保证每一项都是解析的，需要做如下变形：

$$ \begin{aligned} L_\text{VLB} &= \mathbb{E}_{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}_q \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}_q \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}_q \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}_q \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}_q \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}_q \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}_q \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}_q [\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} $$

可以看到，除去$- \log p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1)$项外，其余项均均为两个高斯分布的 KL 散度，可以解析的计算。并且由于向前过程 $q$ 没有可以学习的参数，$x_T$ 为纯高斯分布，所以 $L_T$ 在优化过程总可以忽略。我们几乎只需要关注 $L_t$ 项。由于 $L_t$ 项几乎可以认为是两个高斯函数 $q(\mathbf{x}_t \vert \mathbf{x}_{t+1}, \mathbf{x}_0)$ 和 $p_\theta(\mathbf{x}_t \vert\mathbf{x}_{t+1})$ 的KL散度，可以直接解析的写出：

$$ L_t = \mathbb{E}_{q} \Big[\frac{1}{2 \| \boldsymbol{\Sigma}_\theta(\mathbf{x}_t, t) \|^2_2} \| \tilde{\boldsymbol{\mu}}_t(\mathbf{x}_t, \mathbf{x}_0) - \boldsymbol{\mu}_\theta(\mathbf{x}_t, t) \|^2 \Big] + C\\ $$

其中 $C$ 为与待优化变量 $\theta$ 无关的常数，可以忽略。将之前求出的均值的形式：

$$ Math $$

带入 $L_t$ 中，有：

$$ \begin{aligned} L_t &= \mathbb{E}_{\mathbf{x}_0, \boldsymbol{\epsilon}} \Big[\frac{1}{2 \| \boldsymbol{\Sigma}_\theta(\mathbf{x}_t, t) \|^2_2} \| \tilde{\boldsymbol{\mu}}_t(\mathbf{x}_t, \mathbf{x}_0) - \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} \| \frac{1}{\sqrt{\alpha_t}} \Big( \mathbf{x}_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar{\alpha}_t}} \boldsymbol{\epsilon}_t \Big) - \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} $$

从上式可以看出，diffusion训练的核心就是取学习高斯噪声 $\bar{z}_t$ 与 $z_\theta$ 之间的最大似然。经验上讲，Denoising Diffusion Probabilistic Models 发现简化掉加权项后效果更好：

$$ \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_{simple} = L_t^{simple} + C $$

好了，现在训练的问题也解决了。diffusion 的训练问题可以总结为下面几个步骤：

1. 获得输入 $x_0$，从 $1, 2, \cdots, T$ 中抽样出一个 $t$
2. 从标准高斯分布中抽样出一个噪声 $\epsilon_t \sim \mathcal N(0, I)$
3. 最小化 $\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]$

现在还有最后一个问题，如何用自然语言控制图片的生成。事实上这也很简单，哪著名的 stable diffusion 为例，它使用了 CLIP 自然语言端的编码器，将自然语言编码为一个embedding。将这个embedding 作为 Unet 的输入监督噪声的预测。这样就实现建立了 CLIP 的 embedding 与生成的图片之间的关系。

3. DreamFusion

现在我们已经介绍了使用 diffusion 实现的 2D AIGC。但是想将之直接应用于 3D AIGC 是不现实的。局限显而易见，相比人类已经积累了数亿计的二维图像数据（事实上 clip 在训练时就用了 4 亿规模的数据），但是三维数据却小的可怜。因此，3D AIGC 主要分为直接使用三维你数据训练和从二维 diffusion 迁移两条路。关于这两条技术路线的具体讨论可以参考 胡渊明 的博客 Taichi NeRF (下): 关于 3D AIGC 的务实探讨。我们要介绍的 DreamFusion 可以说是第二条路线的奠基之作。

3.1. 基本实现

一个直观的想法是直接使用训练完成的 2D diffusion 模型生成照片，之后使用 nerf/3DGS 来三维重建。但是事实上这几乎是不可能的。因为不论是 nerf 还是 3DGS，对照片的几何一致性和位姿都有很高的要求，同时还需要提供相机的位姿和焦距。但是现有的 diffusion 模型只能通过很粗糙的自然语言来约束生成图片的几何，更遑论按照一定的相机内参来生成图片了。

但是事实上 diffusion module 中确实存储了巨量的真实的三维几何在二维的投影，我们目前面临的问题是如何把其中的信息 “取出来”。在前面的讨论中我们知道，diffusion module 生成图片的过程是通过一个 unet（事实上也有用transformer的）来预测噪声。通过不断预测高斯噪声来恢复图片原有的结构。那么只要预测的噪声足够准确，从噪声中恢复结构的图片就会足够准确。那么 diffusion 的 lossfunciton 可以写成一个 UNet 预测的噪声 $\rm{UNet}(\alpha x_{render} - \sigma_t \epsilon_t )$ 与 添加的噪声 $\epsilon$ 之间差异。其中 $\|\cdot\|$ 表示衡量两个概率分布的差异：

$$ \mathcal L_{diff} = \mathbb E\left[ \| \epsilon_{pred} - \epsilon \|\right] $$

事实上，对于一个训练好的 diffusion module，与一个随机初始的 Nerf 网络相比，我们可以认为它几乎是准确的。那么一个显然的思路是模仿 diffusion 的训练过程，通过给 Nerf 渲染的照片添加噪声，再用训练好的 diffusion module 降噪。我们可以写出此时的 Loss：

$$ \mathcal L_{diff} = \mathbb E\left[ \omega(t) \| \rm{UNet}(\alpha x_{render} - \sigma_t \epsilon_t ) - \epsilon \|\right] $$

其中 $\omega(t)$ 是一个与 timestep 有关的函数，是一个收到超参数控制的用于调整生成过程的函数，$x_{render}$ 为 Nerf 渲染的图片。

这是 dreamfusion 的 pipline。首先通过 Nerf 出的照片 $rendering$ 添加噪声 $\epsilon \sim \mathcal N(0, I)$。添加噪声后的照片可以写作：$z_t \sim U(0, 1)$。之后通过 Unet 预测出降噪后的图片 $\check x (z_t \mid y; t)$；之后将求出该步骤中去除的噪声 $\epsilon_\phi(z_t\mid y; t)$。有了预测的噪声后，我们将预测的噪声与设计添加的噪声做差，就可以作为训练 Nerf 的 Loss 了。我们可以看到，dreamfusion 过程实际上是用 Nerf 作为几何的监督，通过 diffusion 来产生新的 “数据分布”，从而解决传统的 diffusion 缺乏几何一致性的问题。实际上是用 diffusion 生成的图片作为监督，来驱动优化 Nerf 模型，即：

$$ \boldsymbol{\theta} = \arg\min_\theta \mathbb E\left[ \omega(t) \| \rm{UNet}(\alpha x_{render} - \sigma_t \epsilon_t ) - \epsilon_t \|_2^2 \right] $$

3.2. SDS Loss：从 diffusion module 中提取知识

事情似乎解决了。但是自己观察 Loss 后不难发现，这个 Loss 几乎是无法计算的。我们希望通过 $\mathcal L_{diff}$ 传递梯度，需要计算 $\frac{\partial \mathcal L_{diff}}{\partial x_{render}}$，这个过程需要计算整个 UNet 的梯度。而一般 UNet 的参数量及其巨大，只是单纯的进行 diffusion 推理可能还能吃得消，但是我们需要 nerf 每一轮训练的每个 timestep 都完整的计算整个 UNet 的梯度，这在实现上几乎是不可行的。因此我们需要找到一个办法，跳过 UNet 的梯度传播，直接计算 $\mathcal L_{diff}$。

我们不妨写出 $\mathcal L_{diff}$ 对 $\boldsymbol{\theta}$ 的梯度：

$$ \nabla_{\theta} \mathcal{L}_{Diff }(\phi, x=g(\theta))=\mathbb{E}_{t, \epsilon}[w(t) \underbrace{\left(\hat{\epsilon}_{\phi}\left(z_{t} ; y, t\right)-\epsilon\right)}_{\text{Noise Residual} } \underbrace{\frac{\partial \hat{\epsilon}_{\phi}\left(z_{t} ; y, t\right)}{z_{t}}}_{\text{U-Net Jacobian} } \underbrace{\frac{\partial x}{\partial \theta}}_{\text{Nerf Jacobian} }] $$

其中从 $\alpha x_{render} - \sigma_t \epsilon_t$ 到 $x$ 的梯度传播为 $\frac{\partial \alpha x_{render} - \sigma_t \epsilon_t}{\partial x} = \alpha I$，直接归纳入 $\omega(t)$ 。此外，直接省略 UNet 的梯度，有：

$$ \nabla_{\theta} \mathcal{L}_{SDS}(\phi, x=g(\theta)) \triangleq \mathbb{E}_{t, \epsilon}\left[w(t)\left(\hat{\epsilon}_{\phi}\left(z_{t} ; y, t\right)-\epsilon\right) \frac{\partial x}{\partial \theta}\right] $$

这就是 SDS Loss Function。事实上在实际操作时，我们通过 KL 三度来衡量两个分布之间的误差。

$$ \nabla_{\theta} \mathcal{L}_{SDS}(\phi, x = g(\theta))=\nabla_{\theta} \mathbb{E}_{t}\left[\sigma_{t} / \alpha_{t} w(t) \text{KL}\left(q\left(z_{t} | g(\theta) ; y, t\right) \| p_{\phi}\left(z_{t} ; y, t\right)\right)\right] $$

可以说 $\nabla_{\theta} \mathcal{L}_{SDS}(\phi, x = g(\theta))$ 是从 diffusion 估计出的分数函数 (Score Function)。

4. Reference

1. 由浅入深了解 Diffusion Model

2. A Tutorial of Normalizing Flow Module

3. Generative Modeling by Estimating Gradients of the Data Distribution —- Yang Song

4. What are Diffusion Models?

5. Sliced score matching: A scalable approach to density and score estimation

6. From Autoencoder to Beta-VAE

7. DreamFusion: Text-To-3D Using 3D Diffusion

5. Append

5.1. Boltzmann Distribution

玻尔兹曼分布是统计力学中描述粒子在不同能级上分布概率的重要概念。以下是详细的推导步骤和依据。

系统的熵定义为：

$$ S = -k_B \sum P_i \ln P_i $$

其中：

• $P_i$ 是系统处于第 $i$ 个能量状态的概率，
• $k_B$ 是玻尔兹曼常数。

推导过程中需要满足以下约束条件：

1. 概率归一化条件：

$$ \sum P_i = 1 $$

2. 固定平均能量：

$$ \sum P_i E_i = \langle E \rangle $$

其中：

• $E_i$ 是第 $i$ 个能量状态的能量，
• $\langle E \rangle$ 是系统的平均能量。

为了在约束条件下最大化熵，引入拉格朗日乘子 $\alpha$ 和 $\beta$，构造拉格朗日函数：

$$ \mathcal{L} = -k_B \sum P_i \ln P_i + \alpha \left( \sum P_i - 1 \right) + \beta \left( \sum P_i E_i - \langle E \rangle \right) $$

对每个 $P_j$ 求偏导并设为零：

$$ \frac{\partial \mathcal{L}}{\partial P_j} = -k_B (\ln P_j + 1) + \alpha + \beta E_j = 0 $$

解得：

$$ \ln P_j = \frac{\alpha}{k_B} - 1 - \frac{\beta E_j}{k_B} $$

因此：

$$ P_j = A e^{-\gamma E_j} $$

其中：

• $A = e^{(\alpha / k_B - 1)}$，
• $\gamma = \beta / k_B$。

通过归一化条件 $\sum P_j = 1$，得到：

$$ A = \frac{1}{Z} $$

其中，配分函数 $Z = \sum e^{-\gamma E_j}$。

通过热力学关系确定 $\gamma = 1 / (k_B T)$，其中 $T$ 是温度。

最终的玻尔兹曼分布为：

$$ P_j = \frac{e^{-E_j / (k_B T)}}{Z} $$

其中，配分函数 $Z = \sum e^{-E_j / (k_B T)}$。

5.2. DDPM 和 Score Module 之间的关系

我们不妨先复习一下，前面有讲到的 DDPM 的更新步骤为：

$$ x_{t-1} = \frac{1}{\sqrt{\alpha_t}} \left( x_t - \frac{1-\alpha_t}{\sqrt{1-\bar{\alpha}t}}\, \epsilon\theta(x_t,t) \right) + \sqrt{\sigma_t^2}\, z $$

此外，Langevin Dynamics 的迭代步骤为：

$$ x_{k+1} = x_k + \frac{\eta}{2} \nabla_{x} \log p(x_k) + \sqrt{\eta}\, z_k,\quad z_k\sim \mathcal{N}(0,\mathbf{I}) $$

对于一个高斯分布 $q(x_t|x_0) = \mathcal{N}(x_t; \sqrt{\bar{\alpha}t}\,x_0,\, (1-\bar{\alpha}t)\mathbf{I})$，其 score function（即 $\nabla{x_t} \log q(x_t|x_0)$）可以直接计算：

$$ \nabla{x_t} \log q(x_t|x_0) = -\frac{1}{1-\bar{\alpha}_t} \left(x_t - \sqrt{\bar{\alpha}_t}\,x_0\right) $$

利用前向过程的表达式 $x_t = \sqrt{\bar{\alpha}_t}\,x_0 + \sqrt{1-\bar{\alpha}_t}\,\epsilon$，我们有：

$$ x_t - \sqrt{\bar{\alpha}_t}\,x_0 = \sqrt{1-\bar{\alpha}t}\,\epsilon $$

因此，score function 可写为：

$$ \nabla{x_t} \log q(x_t|x_0) = -\frac{\sqrt{1-\bar{\alpha}_t}}{1-\bar{\alpha}_t}\,\epsilon = -\frac{1}{\sqrt{1-\bar{\alpha}_t}}\,\epsilon $$

在 DDPM 的训练过程中，我们用 $\epsilon_\theta(x_t, t)$ 来预测前向过程中实际添加的噪声 $\epsilon$。如果网络训练得足够好，就有：

$$ \epsilon_\theta(x_t, t) \approx \epsilon $$

那么，根据上面的推导，我们可以得到：

$$ \nabla_{x_t} \log q(x_t|x_0) \approx -\frac{1}{\sqrt{1-\bar{\alpha}t}}\,\epsilon\theta(x_t,t) $$

这就说明，网络预测的噪声与数据在当前噪声水平下的 score function 之间只相差一个与时间（或噪声水平）有关的比例因子 $\frac{1}{\sqrt{1-\bar{\alpha}_t}}$。

因此我们知道，DDPM实质上是学习了一个隐式的 Score Function。DDPM的反向过程（去噪过程）在数学上可以看作是对目标数据分布（或其在不同噪声水平下的变体）进行梯度上升采样的离散化实现，而Langevin Dynamics正是利用score function（即 $\nabla_x \log p(x)$）进行采样的方法。