简化训练损失的推导

神经网络（通常是U-Net）在扩散模型中的核心任务是学习逆向扩散过程。这涉及通过逐步去除噪声来估计原始数据分布。虽然扩散模型的完整概率推导涉及通过证据下界（ELBO）最大化数据似然，但一种常见且有效的训练目标大大简化了这一点。

从ELBO到简单的MSE

扩散模型的完整ELBO包含与逆向马尔可夫链中转换相关的多个项。直接优化它可能很复杂。然而，研究人员（特别是Ho等人在原始DDPM论文中）发现，训练目标可以大大简化，而不会牺牲性能。

与其直接建模复杂的逆向转换概率$p_\theta(x_{t-1}|x_t)$，我们重新参数化网络$\epsilon_\theta$，以预测在时间步$t$从原始图像$x_0$获得$x_t$时所添加的噪声$\epsilon$。回顾允许直接从$x_0$采样$x_t$的前向过程方程：

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

这里，$\epsilon$是一个标准高斯噪声样本$(\epsilon \sim \mathcal{N}(0, I))$，$\bar{\alpha}_t$则根据噪声时间表得出。

其要点在于，如果我们的网络能够在给定带噪声图像$x_t$和时间步$t$的情况下准确预测此噪声$\epsilon$，它实际上就学会了如何逆转扩散步骤。

简化的损失函数

这带来一个简单得多的训练目标：我们最小化用于创建$x_t$的实际噪声$\epsilon$与我们的网络$\epsilon_\theta(x_t, t)$预测的噪声之间的均方误差（MSE）。

损失函数$L_{simple}$被表述为对随机时间步$t$、初始数据样本$x_0$和噪声样本$\epsilon$的期望：

$$L_{simple} = \mathbb{E}_{t \sim [1, T], x_0 \sim q(x_0), \epsilon \sim \mathcal{N}(0, I)} \left[ || \epsilon - \epsilon_\theta(\sqrt{\bar{\alpha}_t} x_0 + \sqrt{1 - \bar{\alpha}_t} \epsilon, t) ||^2 \right]$$

我们来分解一下：

1. 采样 $t$： 从$\{1, ..., T\}$中均匀选择一个随机时间步$t$。
2. 采样 $x_0$： 从训练数据集中获取一个数据样本（例如，一张图像）。
3. 采样 $\epsilon$： 从标准正态分布中抽取一个噪声样本。
4. 计算 $x_t$： 使用公式$x_t = \sqrt{\bar{\alpha}_t} x_0 + \sqrt{1 - \bar{\alpha}_t} \epsilon$计算带噪声的版本$x_t$。
5. 预测噪声： 将$x_t$和时间步$t$输入网络$\epsilon_\theta$，以获得预测噪声$\epsilon_\theta(x_t, t)$。
6. 计算MSE： 计算实际噪声$\epsilon$与预测噪声$\epsilon_\theta(x_t, t)$之间的平方差。
7. 平均： 在实际中，这个期望是通过在训练期间对小批量数据上的MSE求平均来近似的。

为何如此有效

最小化此MSE损失会促使网络$\epsilon_\theta$成为对于任何给定的噪声水平$t$都有效的噪声预测器。直观地看，如果模型能够精确识别$x_t$中的噪声成分，它就隐式理解了逆转噪声过程所需的底层数据$x_0$的结构。

虽然我们跳过了从ELBO详细的数学推导，但可以证明，这个简化的MSE损失对应于变分下界中项的特定加权。这为此较简单的目标为何能有效训练高质量扩散模型提供了理论依据。

实际优势显著：训练归结为一个标准的回归问题，其中网络学习将带噪声的输入$x_t$和时间步$t$映射到所添加的噪声$\epsilon$。这比直接处理完整ELBO的复杂分布要简单得多，更易于实现和优化。这个简化损失是我们将在下一节中描述的训练算法的依据。