ICML 2024 | 通过随机微分方程统一贝叶斯流网络和扩散模型

公众号新闻

2024-05-14 05:05

©作者 | 薛凯文
单位 | GSAI-ML

论文题目：

Unifying Bayesian Flow Networks and Diffusion Models through Stochastic Differential Equations

论文作者：

薛凯文*、周聿浩*、聂燊、闵旭、张晓露、周军、李崇轩

论文链接：

https://arxiv.org/abs/2404.15766

代码链接：

https://github.com/ML-GSAI/BFN-Solver

背景

如今深度生成式模型被广泛运用在计算机视觉和自然语言处理等领域。这些模型面临的主要挑战在于如何有效地表示复杂的概率分布，这些分布通常高度依赖于具体的数据或模态，同时能需要够进行有效的训练和推理。

自回归模型（ARM）擅长对序列和离散数据（如文本）进行建模，但在推理速度上存在明显瓶颈，另一方面，扩散模型（DM）采用从粗糙到精细的渐进式方法，实现了生成质量与效率的更好平衡。

目前扩散模型在图像生成领域取得了很大的进展，但它在处理离散变量时面临挑战，分数匹配算法在这里并不直接适用。贝叶斯流网络（BFN）是一种新的生成式模型，与扩散模型（DM）不同，BFN 不直接处理样本，而是通过贝叶斯推理迭代地细化不同噪声级别的分布参数。这种模型因其可微性质显示出在连续与离散数据建模上的巨大潜力，并且能够实现快速采样。

概述

为了更深入地理解和改进 BFN，本文利用随机微分方程（SDE）将其与扩散模型建立了联系。我们发现了与 BFN 中的加噪过程相对应的线性 SDE，证明了 BFN 的回归损失等价于降噪分数匹配（DSM）。

基于此发现，我们验证 BFN 中的采样器近似等同于反向 SDE 的一阶 Solver。此外，通过借鉴扩散模型中现有的快速采样方案，我们提出了一种新的采样方法——BFN-Solvers，通过在图像和文本数据集上的测试，BFN-Solvers 在进行有限次数（例如 10 次）的函数评估下，能显著提升样本质量，并超越原始 BFN 采样器。特别地，我们的最优采样器实现了 5~20 倍的速度提升如图 3、4 所示。

贝叶斯流网络

BFN通过迭代不同噪声水平下的分布参数，这种策略使 BFN 在连续和离散数据上都是可微分的，同时使得少步数生成成为可能。给定一个从真实数据分布采样得到的样本点，BFN 定义了一个贝叶斯更新过程，通过不同噪声水平的带噪数据和贝叶斯公式更新样本的先验分布参数，总共步，产生了参数序列。每一步的噪声水平由准确率决定：

其中可以被设置为一个简单的先验。随着的增加，趋近于样本点的狄拉克分布，当趋近于无穷时，分布收敛到样本点的狄拉克分布。

与扩散模型类似，贝叶斯更新过程同样具有单步采样的性质，即我们可以得到分布的解析形式（具体推导可以参考 BFN 论文 sec3.4）：

其中准确率时间表 (accuracy scheduel) 。

上述的贝叶斯更新过程定义了给定样本作为条件时，从固定先验到样本的狄拉克分布的过程，如果我们可以不依赖 , 从中采样，我们能够通过从固定先验出发得到样本点的狄拉克分布，并从中采样得到生成数据。

困难在于估计的计算开销是巨大的，因为需要用到整个数据集，我们可以用神经网络学习这些条件概率，这样我们得到了由神经网络定义的参数更新过程。特别地，为了方便之后损失函数的化简，我们定义成以下期望的形式：

到这里，我们完成了 BFN 模型的定义，接下来我们考虑如何训练 BFN 和 BFN 的具体参数化形式。

BFN 的优化目标是负对数似然的变分下界：

其中被称为重建损失，我们可以对做进一步化简：

作者从信息论角度理解。数据所有者根据噪声时间表向接受者传输有损信息，接受者根据当前时刻的先验分布接受此信息，并通过贝叶斯更新得到后验分布。表示传输所需的 nat (natural unit of information) 的期望数量。

接着我们需要考虑和的具体参数化形式。这里直接列出了关于离散数据的参数化形式。令：

其中代表神经网络，输出维度为。我们可以计算得到：

假设加噪分布为正态分布给训练带来了立即的好处，我们可以进一步简化损失函数中的 KL 项。在这里我们直接给出了连续时间的化简后的损失函数，具体推导见 BFN 论文 Sec 3.9：

其中：

准确率时间表需要满足单调递增的性质，这里经验上被设置为作为一个超参数。

通过对化简后的损失函数做蒙特卡洛估计，我们得到了可以直接用于训练的损失函数。具体训练和采样流程如算法 8、9 所示。

我们基本完成了 BFN 的介绍。接下来介绍我们近期对 BFN 的一些研究。简而言之，我们通过对应于 BFN 加噪策略的线性随机微分方程建立 BFN 和扩散模型的联系，将 BFN 训练损失函数对应于去噪得分匹配损失（DSM）。基于这个认识，我们发现了 BFN 的原始采样算法对应于反向 SDE 的离散化的某种近似，并为 BFN 开发了加速采样算法。

通过SDE统一BFN和DM

我们发现连续时间 BFN 对离散数据的噪声添加过程唯一求解了一个线性 SDE，总结为定理 5.1。

song 等人（2021）指出线性 SDE 对应一个由未知得分函数定义的反向 SDE：

值得注意的是定理 5.1 描述的是隐变量而不是的动态，如图 2 所示，这暗示了 BFN 原始采样算法不是通过直接离散化 SDE 进行采样。

我们证明连续时间 BFN 在离散数据上的训练目标是 DSM 的重新参数化形式，总结为定理 5.2。

定理 5.1 和定理 5.2 将 BFN 与现有的离散状态扩散模型区分开来。具体来说，应用于离散数据的 BFN 求解线性 SDE，并使用 DSM 进行训练，与连续状态扩散模型无缝对齐。因此，在不改变离散数据的情况下，BFN 可以直接利用连续状态扩散模型的现有的经验进行改进。在这个工作中我们尝试改进 BFN 的采样效率。