扩散模型笔记(1)
论文基本信息
论文名称:Deep Unsupervised Learning using Nonequilibrium Thermodynamics
出版期刊:ICML 2015
学校机构:Stanford University,University of California, Berkeley
论文翻译
以下翻译结合了AI工具和我对论文的理解,在翻译的过程中会省略我认为不重要的部分,并不是全篇一字一句的翻译。并且可能翻译或理解有误。(此过程同时借助了ChatGPT和Qwen两个AI工具)
标题
利用非平衡热力学的深度无监督学习
摘要
机器学习中的一个核心问题在于:如何利用高度灵活的概率分布族对复杂数据集进行建模,同时保证学习、采样、推断和评估在解析上或计算上仍具有可行性。本文提出了一种能够同时实现灵活性与可处理性的方法。其核心思想受非平衡统计物理学启发,即通过迭代式的前向扩散过程,系统且缓慢地破坏数据分布中的结构。随后,我们学习一个反向扩散过程以恢复数据中的结构,从而得到一个兼具高度灵活性与可处理性的数据生成模型。该方法使我们能够快速地在具有数千层或时间步的深度生成模型中进行学习、采样与概率评估,并可计算所学模型下的条件概率与后验概率。此外,我们还开源发布了该算法的参考实现Github。
引言
历史上,概率模型一直面临两个相互冲突目标之间的权衡:可处理性与灵活性。具有可处理性的模型能够进行解析式评估并易于拟合数据(例如高斯分布或拉普拉斯分布)。然而,这些模型难以恰当地描述丰富数据集中的复杂结构。另一方面,灵活的模型则可以被建模以拟合任意数据中的结构。例如,我们可以通过任意(非负)函数 $\phi(x)$ 来定义模型,从而得到灵活的分布 $p(x) = \frac{\phi(x)}{Z}$,其中 Z 为归一化常数。然而,计算该归一化常数通常是不可行的。对这类灵活模型进行概率评估、训练或采样通常需要代价高昂的蒙特卡洛过程。
目前已存在多种解析近似方法,它们虽可缓解但无法彻底消除这一权衡——例如平均场理论及其扩展、变分贝叶斯、对比散度、最小概率流、最小 KL 收缩、恰当评分规则、分数匹配、伪似然、循环信念传播等等,不一而足。此外,非参数方法亦可取得显著成效。
非参数方法可被视为在可处理模型与灵活模型之间实现平滑过渡的一种途径。例如,非参数高斯混合模型在表示少量数据时可能仅使用单个高斯分布,而在表示无限量数据时则可扩展为由无限多个高斯分布构成的混合模型。
参数化 GMM:需预先设定
K=5个高斯分量,无论数据多复杂都只能用这 5 个分量拟合;非参数 GMM:无需指定
K,算法自动学习需要多少个分量——简单数据可能只用 2 个分量,复杂数据可能用 20 个甚至更多。
扩散概率模型
我们提出了一种定义概率模型的新方法,该方法具备以下特性:
- 模型结构具有极高的灵活性;
- 支持精确采样;
- 易于与其他分布相乘(例如用于计算后验概率);
- 模型的对数似然以及单个状态的概率均可高效计算。
我们的方法利用马尔可夫链逐步将一个分布转换为另一个分布,这一思想源于非平衡统计物理学和序列蒙特卡洛方法。我们构建了一个生成式马尔可夫链,通过扩散过程将一个简单的已知分布(例如高斯分布)转换为目标(数据)分布。与使用马尔可夫链对已预先定义的模型进行近似评估不同,我们显式地将概率模型定义为马尔可夫链的终点。由于扩散链中的每一步都具有可解析计算的概率,因此整个链的概率亦可进行解析式评估。
在该框架下的学习过程涉及对扩散过程中的微小扰动进行估计。相比使用单一的、无法解析归一化的势函数来显式描述完整分布,估计微小扰动在计算上更为可行。此外,由于对任意光滑的目标分布都存在对应的扩散过程,该方法能够捕捉任意形式的数据分布。
Potential function:势函数,源自物理学"势能"(potential energy)概念,在概率模型中表示未归一化的能量/权重函数。
我们通过在多个数据集上训练高对数似然模型,验证了这些扩散概率模型的有效性。所用数据集包括:二维瑞士卷(two-dimensional swiss roll)、二值序列、手写数字(MNIST)以及若干自然图像数据集(CIFAR-10、树皮纹理和枯叶模型)。
相关的其他工作
wake-sleep algorithm 首次提出了让推断模型与生成概率模型相互对抗训练的思想。此后近二十年间,这一思路在很大程度上未被深入探索,尽管存在一些例外工作。近年来,围绕该思想的研究呈现爆发式增长。之后提出的变分学习与推断算法,使得灵活的生成模型与潜在变量的后验分布能够被直接联合训练、相互优化。
有些算法翻译为中文很不好理解,所以不做翻译。
这些论文中使用的变分界与我们训练目标中采用的变分界以及 Sminchisescu 等人(2006)早期工作中的形式相似。然而,我们的动机与模型形式均存在显著差异。相较于上述方法,本文工作保留了以下区别与优势:
- 我们借鉴物理学、准静态过程(quasi-static processes)以及退火重要性采样(annealed importance sampling)的思想构建框架,而非基于变分贝叶斯方法;
- 我们展示了如何便捷地将所学分布与另一概率分布相乘(例如与条件分布相乘以计算后验分布);
- 我们解决了变分推断方法中推断模型训练尤为困难的问题——该困难源于推断模型与生成模型在目标函数中的不对称性。为此,我们将前向(推断)过程限制为简单的函数形式,使得反向(生成)过程具有相同的函数形式;
- 我们训练了包含数千层(或时间步)的模型,而非仅限于少量层次;
- 我们为每一层(或时间步)中的熵产生(entropy production)提供了上界与下界估计。
存在多种与概率模型训练相关的技术(下文将简要总结),这些方法致力于构建高度灵活的生成模型形式、训练随机轨迹,或学习贝叶斯网络的逆过程。
| 技术/作者 | 时间 | 简述 |
|---|---|---|
| Reweighted wake-sleep | 2015 | 对原始wake-sleep algorithm提出了扩展并改进了学习规则。 |
| Generative stochastic networks | 2013/2014 | 通过训练马尔可夫核,使其平稳分布与数据分布相匹配。 |
| Neural autoregressive distribution estimators | 2011/2013 | 将联合分布分解为一系列关于各维度的可处理条件分布序列。 |
| Adversarial networks | 2014 | 通过让生成模型与一个试图区分生成样本与真实数据的分类器进行对抗训练。 |
| Schmidhuber(作者) | 1992 | 提出的类似目标学习了一个双向映射,将数据映射到具有边缘独立单元的表示空间。 |
| Rippel 与 Adams以及 Dinh(作者) | 2013/2014 | 学习了到潜在表示空间的双射确定性映射,该潜在空间具有简单的阶乘密度函数。 |
| Stuhlmüller(作者) | 2013 | 贝叶斯网络学习了随机逆映射。 |
| Mixtures of conditional Gaussian scale mixtures, MCGSMs | 2012 | 利用高斯尺度混合来描述数据集,其参数依赖于一系列因果邻域。 |
此外,还有大量工作致力于学习从简单潜在分布到数据分布的灵活生成映射——早期范例包括 MacKay 将神经网络引入作为生成模型,以及 Bishop 等学习从潜在空间到数据空间的随机流形映射。本文将在实验中与对抗网络及 MCGSMs 进行对比。
来自物理学的相关思想包括 Jarzynski 等式,该等式在机器学习中被称为退火重要性采样(Annealed Importance Sampling, AIS)。AIS 利用马尔可夫链将一个分布缓慢转换为另一个分布,以计算归一化常数的比值。Burda 等人进一步证明,AIS 亦可基于反向轨迹而非前向轨迹执行。朗之万动力学(Langevin dynamics)作为福克-普朗克方程(Fokker-Planck equation)的随机实现形式,展示了如何定义一个高斯扩散过程,使其以任意目标分布作为平稳分布。福克-普朗克方程也被用来进行随机优化。最后,柯尔莫哥洛夫前向与后向方程(Kolmogorov forward and backward equations)表明:对于许多前向扩散过程,其对应的反向扩散过程可用相同的函数形式加以描述。
算法
我们的目标是定义一个前向(或推断)扩散过程,将任意复杂的数据分布转换为简单且易于处理的分布;随后学习该扩散过程的有限时间反向过程,以此定义我们的生成模型分布(见图1)。我们首先描述前向推断扩散过程,继而展示如何训练反向生成扩散过程并利用其进行概率评估。此外,我们还为反向过程推导了熵的上下界,并说明所学分布如何与任意第二个分布相乘(例如在图像修复或去噪任务中计算后验分布时所需的操作)。

图1:在 2-d swiss roll data 上训练的建模框架。第一行展示了前向轨迹 $q(x^{(0:T)})$ 在各时间切片上的分布。数据分布(左)经历高斯扩散过程,逐步转化为单位协方差高斯分布(右)。第二行展示了训练所得反向轨迹 $p(x^{(0:T)})$ 对应的时间切片。单位协方差高斯分布(右)经过具有学习得到的均值与协方差函数的高斯扩散过程,逐步重构回原始数据分布(左)。第三行展示了同一反向扩散过程的漂移项(drift term)$f_\mu(x^{(t)}, t) - x^{(t)}$。
为什么论文这里又称前向过程为 Inference Process?
在这篇论文中,“inference” 指的是:给定一个真实数据样本 x(0) ,推断(或构造)它在整个扩散轨迹中对应的噪声版本序列 x(1),x(2),…,x(T) 。这里的 “inference” 不是指“从训练好的模型生成新样本”,而是指“为训练目的,从已知数据反推其在扩散链中的演化路径”。
前向轨迹
我们将数据分布记为 $q(x^{(0)})$。通过反复应用针对 $\pi(y)$ 的马尔可夫扩散核 $T_\pi(y \mid y'; \beta)$,数据分布被逐步转换为一个性质良好(解析可处理)的分布 $\pi(y)$,其中 $\beta$ 为扩散速率,
$$ \pi (y) = \int dy' \, T_\pi (y|y'; \beta) \, \pi (y') \tag{1} $$$$ q \left(x^{(t)} \mid x^{(t-1)} \right) = T_\pi \left(x^{(t)} \mid x^{(t-1)}; \beta_t \right). \tag{2} $$公式 (1):目标分布的平稳性条件。
该公式定义了 目标分布 $\pi (y)$ 作为扩散过程的平稳分布(equilibrium distribution):
- $T_{\pi}(y∣y′;\beta)$ 是 马尔可夫扩散核(Markov diffusion kernel),表示从状态 $y′$ 到 $y$ 的转移概率,$\beta$ 是扩散速率(diffusion rate)。
- $\pi(y)$ 满足:当系统达到平衡时,应用扩散核 $T_{\pi}$ 后分布保持不变。
- 在扩散模型中,$\pi(y)$ 通常选择为标准高斯分布(如 $\pi(y) = N(y;0,I)$),因其解析性质良好(易于计算)。
公式 (2):前向扩散过程的定义。
该公式定义了 前向扩散过程(inference process)的条件转移概率:
- $q(x^{(t)}∣x^{(t−1)})$ 是在时间步 $t−1$ 给定 $x^{(t−1)}$ 时,状态 $x^{(t)}$ 的条件分布。
- $T_{\pi}(x^{(t)}∣x^{(t−1)};β_t)$ 是时间相关的扩散核,$\beta_t$ 是第 $t$ 步的扩散速率(通常随时间递增)。
- 核心思想:前向过程使用与目标分布 $\pi (y)$ 相同的扩散核 $T_{\pi}$,但 $\pi$ 随时间变化($\beta_t$ 而非固定 $\beta$)。
前向轨迹对应于从数据分布开始并执行 $T$ 步扩散的过程,因此有
$$ q \left( x^{(0 \cdots T)} \right) = q \left( x^{(0)} \right) \prod_{t=1}^T q \left( x^{(t)} \mid x^{(t-1)} \right) \tag{3} $$马尔科夫性质,每一步只依赖与前一步,与其他步独立。
在下文展示的实验中,$q(x^{(t)}∣x^{(t−1)})$ 对应于以下两种扩散过程之一:
- 高斯扩散:将数据逐步扩散至单位协方差高斯分布(identity-covariance Gaussian distribution);
- 二项扩散:将数据逐步扩散至独立二项分布(independent binomial distribution)。
反向轨迹
生成分布将被训练以描述相同的轨迹,但方向相反,
$$ p \left( x^{(T)} \right) = \pi \left( x^{(T)} \right) \tag{4} $$$$ p \left( x^{(0 \cdots T)} \right) = p \left( x^{(T)} \right) \prod_{t=1}^T p \left( x^{(t-1)} \mid x^{(t)} \right). \tag{5} $$反向轨迹也是马尔科夫链,根据马尔科夫性质得到的。
关键意义:该框架通过学习反向扩散轨迹,使模型能够从简单分布 $\pi(x^{(T)})$ 逐步还原数据分布 $q(x^{(0)})$。由于每一步反向转移概率 $p(x^{(t−1)}∣x^{(t)})$ 均可参数化为神经网络,因此整个生成过程既灵活又可解析计算。
对于高斯扩散与二项扩散,在连续扩散情形下(即步长 $\beta$ 趋近于零的极限),扩散过程的反向过程具有与前向过程完全相同的函数形式。由于 $q(x^{(t)}∣x^{(t−1)})$ 服从高斯(或二项)分布,且当 $\beta_t$ 足够小时,反向条件分布 $q(x^{(t−1)}∣x^{(t)})$ 同样服从高斯(或二项)分布。轨迹长度 $T$ 越长,每一步的扩散速率 $\beta$ 即可取得越小,从而使离散扩散过程更逼近连续极限,反向过程的函数形式亦更精确地匹配前向过程。
在学习过程中,仅需估计高斯扩散核的均值与协方差,或二项扩散核的比特翻转概率。如表 App.1 所示,$f_{\mu}(x^{(t)},t)$ 与 $f_{\Sigma}(x^{(t)},t)$ 是定义高斯分布反向马尔可夫转移的均值与协方差的函数,$f_b(x^{(t)},t)$ 是为二项分布提供比特翻转概率的函数。运行该算法的计算成本等于这些函数的计算开销乘以时间步数。本文所有实验结果均采用多层感知机(MLP)来参数化这些函数。然而,大量回归或函数拟合技术均可适用,包括非参数方法。
模型概率
生成模型赋予原始数据的概率为
$$ p \left( x^{(0)} \right) = \int dx^{(1 \cdots T)} \, p \left( x^{(0 \cdots T)} \right). \tag{6} $$直观上,该积分难以处理——但借鉴退火重要性采样(annealed importance sampling)和 Jarzynski 等式的思想,我们转而评估前向与反向轨迹的相对概率,并对前向轨迹进行平均,
$$ \begin{align} p \left( x^{(0)} \right) &= \int dx^{(1 \cdots T)} \, p \left( x^{(0 \cdots T)} \right) \frac{q \left( x^{(1 \cdots T)} \mid x^{(0)} \right)}{q \left( x^{(1 \cdots T)} \mid x^{(0)} \right)} \tag{7} \\ &= \int dx^{(1 \cdots T)} \, q \left( x^{(1 \cdots T)} \mid x^{(0)} \right) \frac{p \left( x^{(0 \cdots T)} \right)}{q \left( x^{(1 \cdots T)} \mid x^{(0)} \right)} \tag{8} \\ &= \int dx^{(1 \cdots T)} \, q \left( x^{(1 \cdots T)} \mid x^{(0)} \right) p \left( x^{(T)} \right) \prod_{t=1}^{T} \frac{p \left( x^{(t-1)} \mid x^{(t)} \right)}{q \left( x^{(t)} \mid x^{(t-1)} \right)}. \tag{9} \end{align} $$公式(7):概率表达式的重写,分子分母同乘相同因子,等式不变;
公式(8):重要性采样形式,将积分改写为对前向轨迹的期望值;
公式(9):轨迹相对概率的显示表达,由将(5)带入(8)得到。
该积分可通过从前向轨迹 $q(x^{(1⋯T)}∣x^{(0)})$ 中采样并取平均值来快速评估。当扩散步长 $\beta$ 趋近于无穷小时,前向与反向轨迹分布可趋于一致。若二者完全相同,则仅需从前向轨迹 $q(x^{(1⋯T)}∣x^{(0)})$ 中采样单一样本,即可通过代入精确计算上述积分。此情形对应于统计物理学中的准静态过程(quasi-static process)。
核心概念解析
概念 说明 快速评估 通过蒙特卡洛采样估计积分 $E_q[⋅]$,避免高维解析积分 轨迹分布一致 当 $\beta \rightarrow 0$ 且 $T \rightarrow \infty$(连续扩散极限)时,$p(x^{(0⋯T)})=q(x^{(0⋯T)})$,此时相对概率比值恒为 1 单样本精确估计 若 $p=q$,则 $\frac{p(x^{(0⋯T)})}{q(x^{(1⋯T)}∣x^{(0)})}=const$,任意单样本即可给出精确期望值(方差为零) 准静态过程 热力学中系统在无限缓慢变化下始终保持平衡态的过程;在扩散模型中,对应于每步扰动极小、前向/反向过程可逆的理想极限 实践意义:虽然实际应用中 $\beta$ 无法取无穷小(计算成本限制),但该理论表明——轨迹越长、每步扰动越小,模型对数似然的估计方差越低。这为扩散模型设计提供了理论指导:增加时间步数 $T$ 可提升概率估计精度。
训练
训练等价于最大化模型的对数似然,
$$ L = \int dx^{(0)} q\left(x^{(0)}\right) \log p\left(x^{(0)}\right) \tag{10} $$该目标函数可以重写为:
$$ L = \int dx^{(0)} q\left(x^{(0)}\right) \log \left[ \int dx^{(1 \cdots T)} q\left(x^{(1 \cdots T)} \mid x^{(0)}\right) p\left(x^{(T)}\right) \prod_{t=1}^T \frac{p\left(x^{(t-1)} \mid x^{(t)}\right)}{q\left(x^{(t)} \mid x^{(t-1)}\right)} \right], \tag{11} $$
该表达式由Jensen不等式提供了一个下界:
公式解析(q不含参数,p含参数)
公式(10):最大似然目标。本质是 期望对数似然(负交叉熵)。
公式(11):将公式(9)带入(10),通过引入重要性采样,将高维积分转化为轨迹期望。
公式(12):Jensen 不等式下界,避免计算轨迹比值,下界可表示为单轨迹的对数似然,通过蒙特卡洛采样估计。
该是扩散模型的训练目标函数。它将复杂的最大似然估计转化为一个可优化的目标,且:
- 无需计算归一化常数
- 仅需对轨迹进行采样和局部计算
- 与变分推断中的 ELBO 本质相同,但通过扩散过程结构实现了更精确的下界
与单样本精确估计的关联
当模型完美学习($p=q$)且处于准静态极限时:
- 轨迹比值 $\frac{p(x^{(0:T)})}{q(x^{(1:T)}∣x^{(0)})}$ 变为常数
- 此时 $logE[X]=E[logX]$,不等式取等,下界达到紧致
- 这解释了为何在理想情况下"单样本即可精确估计"(但实际中需多采样)
如附录 B 所述,对于我们的扩散轨迹,该下界简化为,
$$ L \geq K \tag{13} $$$$ \begin{align*} K = &-\sum_{t=2}^T \int dx^{(0)} dx^{(t)} q\left(x^{(0)}, x^{(t)}\right) D_{KL}\left(q\left(x^{(t-1)} \mid x^{(t)}, x^{(0)}\right) \parallel p\left(x^{(t-1)} \mid x^{(t)}\right)\right) \\ &+ H_q\left(X^{(T)} \mid X^{(0)}\right) - H_q\left(X^{(1)} \mid X^{(0)}\right) - H_p\left(X^{(T)}\right). \tag{14} \end{align*} $$公式(14)求和项:
- 本质是每步的 KL 散度期望
- 衡量反向生成过程 p 与前向推断过程 q 的差异
- 训练目标:最小化该差异(即让 p 尽可能接近 q 的逆过程)
熵项:
$H_q\left(X^{(T)} \mid X^{(0)}\right)$:最终状态 $X^{(T)}$ 在给定初始状态 $X^{(0)}$ 下的条件熵
$H_q\left(X^{(1)} \mid X^{(0)}\right)$:第一步后的条件熵
$H_p\left(X^{(T)}\right)$:模型在最终状态 $X^{(T)}$ 的熵(通常为 0,因为 $X^{(T)} \sim \pi$ 是固定分布)
其中熵和 KL 散度均可解析计算。该下界的推导与变分贝叶斯方法中对数似然界(log likelihood bound)的推导过程平行。
训练过程包括寻找能够最大化对数似然下界的反向马尔可夫转移概率,
$$ \hat{p}\left(x^{(t-1)} \mid x^{(t)}\right) = \underset{p\left(x^{(t-1)} \mid x^{(t)}\right)}{\operatorname{argmax}} K. \tag{15} $$由此,概率分布的估计问题被简化为对函数的回归任务:这些函数用于设定一系列高斯分布的均值与协方差(或设定一系列伯努利试验的状态翻转概率)。
选择扩散比率$\beta_t$
前向轨迹中 $\beta_t$ 的选择对训练后模型的性能至关重要。在退火重要性采样(AIS)中,中间分布的合理调度可显著提升对数配分函数估计的准确性。在热力学中,系统在不同平衡分布之间转移时所采用的路径调度,决定了自由能的损耗程度。
自由能损耗:非平衡过程中因不可逆性导致的能量耗散。
在高斯扩散的情形下,我们通过在 $K$ 上执行梯度上升来学习前向扩散调度 $\beta_{2 \cdots T}$。第一步的方差 $\beta_1$ 被固定为一个较小的常数,以防止过拟合。为了显式表达从 $q(x^{(1 \cdots T)} \mid x^{(0)})$ 中采样的样本对 $\beta_{1 \cdots T}$ 的依赖关系,我们采用"冻结噪声"(frozen noise)技术——如(Kingma 与 Welling,2013)所述,将噪声视为额外的辅助变量,并在计算 $K$ 对参数的偏导数时保持其恒定。
$\beta$ 调度的选择有两种:固定或可学习。这段描述的是可学习的清况:
- 即使在可学习调度中,$\beta_1$ 仍被固定为小常数(防止过拟合)
- 可学习的是 $\beta_2$ 到 $\beta_T$,即从第 2 步开始的扩散速率
“冻结噪声”技术:这是使采样过程对 $\beta_t$ 可微的关键技巧,本质与 VAE 中的"重参数化技巧"相同。
刚开始以为和调度的选择有关,然后发现这个技术是为了使随机过程可微,“冻结"是指计算梯度时将已采样的噪声视为常数,是扩散模型中的必要技术,其实就是"重参数化技巧"的另一个表达形式。
对于二项扩散,离散状态空间使得基于冻结噪声的梯度上升方法无法实现(因离散变量不可微)。我们转而选择前向扩散调度 $\beta_{1 \cdots T}$,使得每一步扩散擦除原始信号的固定比例 $1/T$,由此得到的扩散速率为 $\beta_t = (T - t + 1)^{-1}$。
在扩散模型里(无论 Gaussian 还是 binomial),都会用到:
- $\beta_t$:第 t 步的噪声强度
- $1-\beta_t$:第 t 步信号保留比例
- 累积信号保留量:$\bar{\alpha}_t = \prod_{s=1}^{t}(1-\beta_s)$ (表示第 t 步之后,还剩多少“原始信号”)
在二项扩散情况下,每步抹掉 $\frac{1}{T}$ 原始信1号:
则 $\bar{\alpha}_t = 1 - \frac{t}{T}$
由 $\bar{\alpha}_t = \bar{\alpha}_{t-1}(1-\beta_t)$
可解得 $\beta_t = \frac{1}{T-t+1}$
分布相乘与后验计算
如信号去噪或缺失值推断等任务需要计算后验分布,这要求将模型分布 $p(x^{(0)})$ 与第二个分布(或有界正函数)$r(x^{(0)})$ 相乘,从而生成一个新的分布 $\tilde{p}(x^{(0)}) \propto p(x^{(0)}) \, r(x^{(0)})$。
对于许多技术而言,分布相乘在计算上代价高昂且实现困难,这些技术包括变分自编码器(VAE)、生成式随机网络(GSN)、神经自回归分布估计器(NADE)以及大多数图模型。然而,在扩散模型框架下,分布相乘却十分直接:第二个分布 $r(x^{(0)})$ 既可以被视为对扩散过程中每一步的微小扰动,也常常能够精确地融入每一步扩散操作中。后续章节将详细阐述如何在扩散概率模型的框架下实现分布相乘。
修改后的边际分布
首先,为了计算 $\tilde{p}(x^{(0)})$,我们将每个中间分布乘以一个对应的函数 $r(x^{(t)})$。我们使用波浪号(tilde)表示该分布或马尔可夫转移属于一条经过此类修改的轨迹。$\tilde{p}(x^{(0\cdots T)})$ 是修改后的反向轨迹,它从分布 $\tilde{p}(x^{(T)}) = \frac{1}{\tilde{Z}_T} p(x^{(T)}) r(x^{(T)})$ 开始,并通过一系列修改后的中间分布进行传递:
$$ \tilde{p}\left(x^{(t)}\right) = \frac{1}{\tilde{Z}_t} p\left(x^{(t)}\right) r\left(x^{(t)}\right), \tag{16} $$其中 $\tilde{Z}_t$ 是第 $t$ 个中间分布的归一化常数。
修改后的扩散步骤
逆扩散过程的马尔可夫核 $p(x^{(t)} \mid x^{(t+1)})$ 满足平衡条件:
$$ p(x^{(t)}) = \int dx^{(t+1)} \, p(x^{(t)} \mid x^{(t+1)}) \, p(x^{(t+1)}). \tag{17} $$我们希望扰动的马尔可夫核 $\tilde{p}(x^{(t)} \mid x^{(t+1)})$ 满足扰动分布的平衡条件,
$$ \tilde{p}(x^{(t)}) = \int dx^{(t+1)} \, \tilde{p}(x^{(t)} \mid x^{(t+1)}) \, \tilde{p}(x^{(t+1)}), \tag{18} $$$$ \frac{p(x^{(t)}) \, r(x^{(t)})}{\tilde{Z}_t} = \int dx^{(t+1)} \, \frac{\tilde{p}(x^{(t)} \mid x^{(t+1)}) \, p(x^{t+1}) \, r(x^{(t+1)})}{\tilde{Z}_{t+1}}, \tag{19} $$$$ p(x^{(t)}) = \int dx^{(t+1)} \, \tilde{p}(x^{(t)} \mid x^{(t+1)}) \, \frac{\tilde{Z}_t \, r(x^{(t+1)})}{\tilde{Z}_{t+1} \, r(x^{(t)})} \, p(x^{(t+1)}). \tag{20} $$若满足
$$ \tilde{p}\left(x^{(t)} \mid x^{(t+1)}\right) = p\left(x^{(t)} \mid x^{(t+1)}\right) \frac{\tilde{Z}_{t+1} r\left(x^{(t)}\right)}{\tilde{Z}_t r\left(x^{(t+1)}\right)}. \tag{21} $$则方程 (20) 成立。
方程 (21) 可能不对应于一个归一化的概率分布,因此我们选择 $\tilde{p}(x^{(t)} \mid x^{(t+1)})$ 为相应的归一化分布:
$$ \tilde{p}\left(x^{(t)} \mid x^{(t+1)}\right) = \frac{1}{\tilde{Z}_t\left(x^{(t+1)}\right)} p\left(x^{(t)} \mid x^{(t+1)}\right) r\left(x^{(t)}\right), \tag{22} $$其中 $\tilde{Z}_t(x^{(t+1)})$ 是归一化常数。
对于高斯分布,由于其方差较小,每一步扩散通常相对于 $r(x^{(t)})$ 具有非常尖锐的峰值。这意味着 $\frac{r(x^{(t)})}{r(x^{(t+1)})}$ 可视为对 $p(x^{(t)} \mid x^{(t+1)})$ 的微小扰动。对高斯分布的微小扰动仅影响其均值,而不会改变归一化常数,因此在此情况下方程 (21) 与 (22) 是等价的(见附录 C)。
应用 $r(x^{(t)})$
若 $r(x^{(t)})$ 足够平滑,则可将其视为对反向扩散核 $p(x^{(t)} \mid x^{(t+1)})$ 的微小扰动。在此情况下,$\tilde{p}(x^{(t)} \mid x^{(t+1)})$ 将具有与 $p(x^{(t)} \mid x^{(t+1)})$ 相同的函数形式,但对于高斯核而言其均值将被扰动,而对于二项核而言其翻转率将被扰动。扰动后的扩散核见表 App.1,并在附录 C 中针对高斯核进行了推导。
若 $r(x^{(t)})$ 可以与高斯(或二项)分布以闭式形式相乘,则它可以直接与反向扩散核 $p(x^{(t)} \mid x^{(t+1)})$ 以闭式形式相乘。这种情况适用于 $r(x^{(t)})$ 由某些坐标子集的 delta 函数构成的情形,例如图 5 中的图像修复(inpainting)示例。
选择 $r(x^{(t)})$
通常,$r(x^{(t)})$ 应当在整个轨迹过程中缓慢变化。在本文的实验中,我们选择其为常数,
$$ r\left(x^{(t)}\right) = r\left(x^{(0)}\right). \tag{23} $$另一种便捷的选择是 $r(x^{(t)}) = r(x^{(0)})^{\frac{T-t}{T}}$。在此第二种选择下,$r(x^{(t)})$ 对反向轨迹的起始分布没有贡献。这保证了从 $\tilde{p}(x^{(T)})$ 中抽取反向轨迹的初始样本仍然简单直接。
反向过程的熵
由于前向过程已知,我们可以推导反向轨迹中每一步条件熵的上下界,从而得到对数似然的上下界,公式 24:
$$ \begin{align*} H_q\left(X^{(t)} \mid X^{(t-1)}\right) + H_q\left(X^{(t-1)} \mid X^{(0)}\right) - H_q\left(X^{(t)} \mid X^{(0)}\right) &\leq H_q\left(X^{(t-1)} \mid X^{(t)}\right) \\ &\leq H_q\left(X^{(t)} \mid X^{(t-1)}\right), \tag{24} \end{align*} $$其中上下界仅依赖于 $q(x^{(1 \cdots T)} \mid x^{(0)})$,并且可以解析计算。推导过程见附录 A。
实验
使用的是 Theano 库,现在已淘汰,学术上更多的用 pytorch,看代码时能看懂思路就行。
总结
我们提出了一种新颖的概率分布建模算法,该算法支持精确采样与概率评估,并在多种玩具数据集与真实数据集(包括具有挑战性的自然图像数据集)上验证了其有效性。在所有实验中,我们均采用相似的基础算法,表明该方法能够准确建模广泛多样的分布形态。
现有大多数密度估计技术为保持计算可行性与效率,往往不得不牺牲建模能力,且采样或概率评估的计算成本通常极高。本算法的核心在于估计一个马尔可夫扩散链的逆过程——该扩散链将数据映射至噪声分布。随着扩散步数的增加,每一步的逆分布逐渐变得简单且易于估计。最终得到的算法兼具以下优势:
- 能够拟合任意数据分布(高建模灵活性)
- 训练过程可处理(tractable)
- 支持精确采样与概率评估
- 便于对条件分布与后验分布进行直接操作
这一设计在建模能力与计算效率之间实现了前所未有的平衡。
附录
A 条件熵边界推导
反向轨迹中一步的条件熵 $H_q(X^{(t−1)}∣X^{(t)})$ 为
$$ H_q\left(X^{(t-1)}, X^{(t)}\right) = H_q\left(X^{(t)}, X^{(t-1)}\right) \tag{25} $$$$ H_q\left(X^{(t-1)} \mid X^{(t)}\right) + H_q\left(X^{(t)}\right) = H_q\left(X^{(t)} \mid X^{(t-1)}\right) + H_q\left(X^{(t-1)}\right) \tag{26} $$$$ H_q\left(X^{(t-1)} \mid X^{(t)}\right) = H_q\left(X^{(t)} \mid X^{(t-1)}\right) + H_q\left(X^{(t-1)}\right) - H_q\left(X^{(t)}\right) \tag{27} $$公式 25 :联合熵的对称性
公式 26 :条件熵之间的关系
通过观察 $\pi(y)$ 是最大熵分布,可以构建熵变化的上界。这一结论对二项分布无条件成立,对高斯分布的情形则要求训练数据的方差为 1。对于高斯分布的情况,因此需要将训练数据缩放到单位范数以保证下述等式成立。但不需要进行白化处理。该上界推导如下,
$$ H_q\left(X^{(t)}\right) \geq H_q\left(X^{(t-1)}\right) \tag{28} $$$$ H_q\left(X^{(t-1)}\right) - H_q\left(X^{(t)}\right) \leq 0 \tag{29} $$$$ H_q\left(X^{(t-1)} \mid X^{(t)}\right) \leq H_q\left(X^{(t)} \mid X^{(t-1)}\right). \tag{30} $$公式 28 :随着扩散步数增加,数据分布的熵单调非减
公式 30 :将 27 代入 29 得到,表明了反向过程的条件熵不超过前向过程的条件熵,即给出了反向过程条件熵的紧致上界
通过观察马尔可夫链中的额外步骤不会增加关于链中初始状态的可用信息,因此不会降低初始状态的条件熵,可以建立熵差的下界,
$$ H_q\left(X^{(0)} \mid X^{(t)}\right) \geq H_q\left(X^{(0)} \mid X^{(t-1)}\right) \tag{31} $$$$ \begin{align*} H_q\left(X^{(t-1)}\right) - H_q\left(X^{(t)}\right) \geq \,& H_q\left(X^{(0)} \mid X^{(t-1)}\right) + H_q\left(X^{(t-1)}\right) \\ &- H_q\left(X^{(0)} \mid X^{(t)}\right) - H_q\left(X^{(t)}\right) \tag{32} \end{align*} $$$$ H_q\left(X^{(t-1)}\right) - H_q\left(X^{(t)}\right) \geq H_q\left(X^{(0)}, X^{(t-1)}\right) - H_q\left(X^{(0)}, X^{(t)}\right) \tag{33} $$$$ H_q\left(X^{(t-1)}\right) - H_q\left(X^{(t)}\right) \geq H_q\left(X^{(t-1)} \mid X^{(0)}\right) - H_q\left(X^{(t)} \mid X^{(0)}\right) \tag{34} $$$$ H_q\left(X^{(t-1)} \mid X^{(t)}\right) \geq H_q\left(X^{(t)} \mid X^{(t-1)}\right) + H_q\left(X^{(t-1)} \mid X^{(0)}\right) - H_q\left(X^{(t)} \mid X^{(0)}\right). \tag{35} $$公式 31 :随着扩散步数增加,关于初始状态 $X^{(0)}$ 的条件熵单调非减
公式 32-35 :熵差下界的推导,得到单步条件熵下界
公式 35 :将 27 代入 34 可得
结合上述不等式,我们可以得到单步条件熵的紧致边界
$$ \begin{align*} H_q\left(X^{(t)} \mid X^{(t-1)}\right) &\geq H_q\left(X^{(t-1)} \mid X^{(t)}\right) \\ &\geq H_q\left(X^{(t)} \mid X^{(t-1)}\right) + H_q\left(X^{(t-1)} \mid X^{(0)}\right) - H_q\left(X^{(t)} \mid X^{(0)}\right). \tag{36} \end{align*} $$
- 仅依赖前向轨迹:上下界均只依赖于已知的前向扩散过程 $q(x^{(1⋯T)}∣x^{(0)})$
- 解析可计算:无需蒙特卡洛采样,可直接通过概率分布的解析形式计算
- 紧致性:在准静态极限下($\beta \rightarrow0$),上下界收敛,实现精确计算
B. 对数似然下界
对数似然的下界为
$$ L \geq K \tag{37} $$$$ K = \int dx^{(0 \cdots T)} q\left(x^{(0 \cdots T)}\right) \log \left[ p\left(x^{(T)}\right) \prod_{t=1}^T \frac{p\left(x^{(t-1)} \mid x^{(t)}\right)}{q\left(x^{(t)} \mid x^{(t-1)}\right)} \right] \tag{38} $$这里对应正文公式 (11)(12)
B.1. $p(X^{(T)})$ 的熵
我们可以分离出 $p(X^{(T)})$ 的贡献,并将其重写为熵,
$$ K = \int dx^{(0 \cdots T)} \, q\left(x^{(0 \cdots T)}\right) \sum_{t=1}^T \log \left[ \frac{p\left(x^{(t-1)} \mid x^{(t)}\right)}{q\left(x^{(t)} \mid x^{(t-1)}\right)} \right] + \int dx^{(T)} \, q\left(x^{(T)}\right) \log p\left(x^{(T)}\right) \tag{40} $$$$ = \int dx^{(0\cdots T)} q\left(x^{(0\cdots T)}\right) \sum_{t=1}^T \log \left[ \frac{p\left(x^{(t-1)} \mid x^{(t)}\right)}{q\left(x^{(t)} \mid x^{(t-1)}\right)} \right] + \int dx^{(T)} q\left(x^{(T)}\right) \log \pi\left(x^{(T)}\right) \tag{41} $$根据设计(选择或规范化了核与参数,使得每一步的二阶矩保持与 $\pi$ 相同),相对于 $\pi(x^{(t)})$ 的交叉熵在我们的扩散核下是常数,并且等于 $p(x^{(T)})$ 的熵。
因此,
$$ K = \sum_{t=1}^T \int dx^{(0\cdots T)} q\left(x^{(0\cdots T)}\right) \log \left[ \frac{p\left(x^{(t-1)} \mid x^{(t)}\right)}{q\left(x^{(t)} \mid x^{(t-1)}\right)} \right] - H_q\left(X^{(T)}\right). \tag{43} $$$$ -\int q\left(x^{(t)}\right) \log \pi\left(x^{(t)}\right) \, dx^{(t)} = -\int \pi\left(x^{(t)}\right) \log \pi\left(x^{(t)}\right) \, dx^{(t)} = H(\pi) $$
最大熵性质:在给定方差的约束下,高斯分布具有最大熵
交叉熵恒定:若 $q(x^{(t)})$ 与 $\pi (x^{(t)})$ 具有相同的二阶矩(方差),则:
$$ \int q\left(x^{(T)}\right) \log \pi\left(x^{(T)}\right) \, dx^{(T)} = \int \pi\left(x^{(T)}\right) \log \pi\left(x^{(T)}\right) \, dx^{(T)} = -H(\pi) = -H_p\left(X^{(T)}\right) $$即:任何具有单位方差的分布与标准高斯 $\pi$ 的交叉熵都等于 $\pi$ 的熵
- 特别地:当 $t=T$ 时,$q(x^{(T)})=\pi (x^{(T)})$,此时交叉熵退化为 $\pi$ 的熵:
$$ \pi(x) = (2\pi)^{-d/2} \exp\left(-\frac{1}{2}x^\top x\right) $$只要 $q(x^{(t)})$ 与 $\pi (x^{(t)})$ 具有相同的二阶矩,交叉熵即恒等于 $\pi$ 的熵,与 $q$ 的具体形式无关。这是扩散模型设计中利用高斯分布最大熵性质的关键技巧。
核心推导
设目标分布为标准高斯 $\pi(x)=N(x;0,I)$,其概率密度为:
$$ \log \pi(x) = -\frac{d}{2} \log(2\pi) - \frac{1}{2}x^\top x $$对数形式为:
$$ H(q, \pi) = -\int q(x) \log \pi(x) \, dx $$关键观察:$log\pi(x)$ 仅包含常数项和 $x^⊤x$(即二阶矩)。
交叉熵定义为:
$$ \begin{aligned} H(q, \pi) &= -\int q(x) \left[-\frac{d}{2} \log(2\pi) - \frac{1}{2}x^\top x\right] dx \\ &= \frac{d}{2} \log(2\pi) \underbrace{\int q(x)dx}_{=1} + \frac{1}{2} \int q(x)x^\top x \, dx \\ &= \frac{d}{2} \log(2\pi) + \frac{1}{2} \mathbb{E}_q[x^\top x] \end{aligned} $$所以:
$$ H(\pi) = \frac{d}{2} \log(2\pi) + \frac{1}{2} \mathbb{E}_{\pi}[x^T x] $$同理,$\pi$ 的熵为:
因此,当 $E_q[x^⊤x]=E_{\pi}[x^⊤x]$ 时:交叉熵退化为熵
B.2. 在 t=0 时移除边缘效应
为了避免边缘效应,我们将反向轨迹的最后一步设置为与对应的前向扩散步骤相同,
$$ p\left(x^{(0)} \mid x^{(1)}\right) = q\left(x^{(1)} \mid x^{(0)}\right) \frac{\pi\left(x^{(0)}\right)}{\pi\left(x^{(1)}\right)} = T_{\pi}\left(x^{(0)} \mid x^{(1)}; \beta_1\right). \tag{44} $$然后我们利用这种等价性来移除求和中第一步的贡献,
$$ \begin{align*} K = &\sum_{t=2}^{T} \int dx^{(0\cdots T)} q\left(x^{(0\cdots T)}\right) \log \left[\frac{p\left(x^{(t-1)} \mid x^{(t)}\right)}{q\left(x^{(t)} \mid x^{(t-1)}\right)}\right] \\ &+ \int dx^{(0)} dx^{(1)} q\left(x^{(0)}, x^{(1)}\right) \log \left[\frac{q\left(x^{(1)} \mid x^{(0)}\right) \pi\left(x^{(0)}\right)}{q\left(x^{(1)} \mid x^{(0)}\right) \pi\left(x^{(1)}\right)}\right] - H_p\left(X^{(T)}\right) \tag{45} \end{align*} $$$$ K = \sum_{t=2}^{T} \int dx^{(0\cdots T)} q\left(x^{(0\cdots T)}\right) \log \left[\frac{p\left(x^{(t-1)} \mid x^{(t)}\right)}{q\left(x^{(t)} \mid x^{(t-1)}\right)}\right] - H_p\left(X^{(T)}\right), \tag{46} $$其中我们再次利用了这样一个事实:根据设计,$−∫dx^{(t)}q(x^{(t)})log\pi(x^{(t)})=H_p(X^{(T)})$ 对所有 $t$ 都是一个常数。
B.3. 重写后验分布 q(x(t−1)∣x(0))
由于前向轨迹是一个马尔可夫过程,
$$ K = \sum_{t=2}^{T} \int dx^{(0 \cdots T)} \, q\left(x^{(0 \cdots T)}\right) \log \left[ \frac{p\left(x^{(t-1)} \mid x^{(t)}\right)}{q\left(x^{(t)} \mid x^{(t-1)}, x^{(0)}\right)} \right] - H_p\left(X^{(T)}\right). \tag{47} $$利用贝叶斯规则将条件概率重写为后验分布形式,
$$ K = \sum_{t=2}^{T} \int dx^{(0 \cdots T)} \, q\left(x^{(0 \cdots T)}\right) \log \left[ \frac{p\left(x^{(t-1)} \mid x^{(t)}\right)}{q\left(x^{(t)} \mid x^{(t-1)}, x^{(0)}\right)} \cdot \frac{q\left(x^{(t-1)} \mid x^{(0)}\right)}{q\left(x^{(t)} \mid x^{(0)}\right)} \right] - H_p\left(X^{(T)}\right). \tag{48} $$B.4. 用KL散度和熵重写
然后我们认识到一些项是条件熵,
$$ \begin{align*} K = &\sum_{t=2}^{T} \int dx^{(0 \cdots T)} \, q\left(x^{(0 \cdots T)}\right) \log \left[ \frac{p\left(x^{(t-1)} \mid x^{(t)}\right)}{q\left(x^{(t)} \mid x^{(t-1)}, x^{(0)}\right)} \right] \\ &+ \sum_{t=2}^{T} \left[ H_q\left(X^{(t)} \mid X^{(0)}\right) - H_q\left(X^{(t-1)} \mid X^{(0)}\right) \right] - H_p\left(X^{(T)}\right) \tag{49} \end{align*} $$$$ \begin{align*} K = &\sum_{t=2}^{T} \int dx^{(0 \cdots T)} \, q\left(x^{(0 \cdots T)}\right) \log \left[ \frac{p\left(x^{(t-1)} \mid x^{(t)}\right)}{q\left(x^{(t)} \mid x^{(t-1)}, x^{(0)}\right)} \right] \\ &+ H_q\left(X^{(T)} \mid X^{(0)}\right) - H_q\left(X^{(1)} \mid X^{(0)}\right) - H_p\left(X^{(T)}\right) \tag{50} \end{align*} $$通过**识别条件熵(公式 49)并利用望远镜级数(公式 50)**的特性,将复杂的求和项简化为仅依赖于初始和最终状态的表达式。这是扩散模型能够高效计算对数似然下界的关键数学技巧。
最后,我们将概率分布的对数比转换为KL散度,
$$ \begin{align*} K = &-\sum_{t=2}^{T} \int dx^{(0)} dx^{(t)} q\left(x^{(0)}, x^{(t)}\right) D_{KL}\left(q\left(x^{(t-1)} \mid x^{(t)}, x^{(0)}\right) \parallel p\left(x^{(t-1)} \mid x^{(t)}\right)\right) \\ &+ H_q\left(X^{(T)} \mid X^{(0)}\right) - H_q\left(X^{(1)} \mid X^{(0)}\right) - H_p\left(X^{(T)}\right). \tag{51} \end{align*} $$注意:熵项可以解析计算,且给定 $x^{(0)}$ 和 $x^{(t)}$ 时,KL 散度也可以解析计算。
该公式将复杂的概率计算问题转化为可解析计算的KL散度和熵项。由于扩散过程的前向轨迹已知,所有这些项都可以精确计算,无需蒙特卡洛采样。这是扩散模型区别于其他生成模型(如GAN、VAE)的核心优势。
C. 扰动高斯转移
我们希望计算 $\tilde{p}(x^{(t−1)}∣x^{(t)})$。为简化符号表示,令 $\mu=f_{\mu}(x^{(t)},t)$,$\Sigma =f_{\Sigma}(x^{(t)},t)$,且 $y=x(t−1)$。使用此符号表示,
$$ \tilde{p}(y \mid x^{(t)}) \propto p(y \mid x^{(t)}) r(y) = \mathcal{N}(y; \mu, \Sigma) r(y). \tag{53} $$我们可以将此重写为能量函数的形式,其中 $E_r(y)=−logr(y)$,
$$ \tilde{p}(y \mid x^{(t)}) \propto \exp\left[-E(y)\right] \tag{54} $$$$ E(y) = \frac{1}{2}(y - \mu)^T \Sigma^{-1} (y - \mu) + E_r(y). \tag{55} $$公式 54 :扰动后的分布可表示为能量函数 $E(y)$ 的指数形式;能量越低的区域,概率密度越高(玻尔兹曼分布原理)
公式 55 :第一项表示原始高斯分布的能量项(二次型);第二项表示扰动函数 $r(y)$ 对应的能量项
如果 $E_r(y)$ 相对于 $\frac{1}{2}(y−μ)^T\Sigma^{−1}(y−μ)$ 是平滑的,那么我们可以使用其在 $\mu$ 处的泰勒展开来近似它。一个充分条件是 $E_r(y)$ 的 Hessian 矩阵的特征值处处比 $\Sigma^{-1}$ 的特征值的模小得多。然后我们有
$$ E_r(y) \approx E_r(\mu) + (y-\mu)g \tag{56} $$其中 $g=\frac{\partial E_r(y′)}{\partial y′} ∣_{y′=μ}$。将此代入完整能量函数,
$$ \begin{align*} E(y) &\approx \frac{1}{2}(y - \mu)^T \Sigma^{-1} (y - \mu) + (y - \mu)^T \mathbf{g} + \text{constant} \tag{57} \\ &= \frac{1}{2} y^T \Sigma^{-1} y - \frac{1}{2} y^T \Sigma^{-1} \mu - \frac{1}{2} \mu^T \Sigma^{-1} y + \frac{1}{2} y^T \Sigma^{-1} \Sigma \mathbf{g} + \frac{1}{2} \bf{g}^T \Sigma \Sigma^{-1} y + \text{constant} \tag{58} \\ &= \frac{1}{2} (y - \mu + \Sigma \mathbf{g})^T \Sigma^{-1} (y - \mu + \Sigma \mathbf{g}) + \text{constant}. \tag{59} \end{align*} $$这对应于一个高斯分布,
$$ \tilde{p}(y \mid x^{(t)}) \approx \mathcal{N}(y; \mu - \Sigma g, \Sigma). \tag{60} $$将结果代回原始形式,得到:
$$ \tilde{p}\left(x^{(t-1)} \mid x^{(t)}\right) \approx \mathcal{N}\left(x^{(t-1)}; \mathbf{f}_{\mu}\left(x^{(t)}, t\right) + \mathbf{f}_{\Sigma}\left(x^{(t)}, t\right) \left.\frac{\partial \log r\left(x^{(t-1)'}\right)}{\partial x^{(t-1)'}}\right|_{x^{(t-1)'} = \mathbf{f}_{\mu}\left(x^{(t)}, t\right)}, \mathbf{f}_{\Sigma}\left(x^{(t)}, t\right)\right). \quad (61) $$实验细节
用的卷积神经网络,数据集不再进行介绍。