[论文解读] Likelihood Training of Schr\\"odinger Bridge using Forward-Backward SDEs Theory

Tianrong Chen, Guan-Horng Liu+1 more|arXiv (Cornell University)|Oct 21, 2021

Generative Adversarial Networks and Image Synthesis被引用 11

一句话总结

本文提出了 SB-FBSDE，一种基于前向-后向随机微分方程（FBSDE）理论的新型似然训练框架，用于薛定谔桥（Schrödinger Bridge, SB）模型。通过推导出 SB 的精确对数似然表达式，该方法将 SGM 似然目标泛化为更一般的形式，从而支持原则性、现代化的训练技术——在 MNIST、CelebA 和 CIFAR10 上实现了最先进（SOTA）的 FID 分数，且在评估时无需使用朗之万校正（Langevin correction）。

ABSTRACT

Schr\\"odinger Bridge (SB) is an entropy-regularized optimal transport problem that has received increasing attention in deep generative modeling for its mathematical flexibility compared to the Scored-based Generative Model (SGM). However, it remains unclear whether the optimization principle of SB relates to the modern training of deep generative models, which often rely on constructing log-likelihood objectives.This raises questions on the suitability of SB models as a principled alternative for generative applications. In this work, we present a novel computational framework for likelihood training of SB models grounded on Forward-Backward Stochastic Differential Equations Theory - a mathematical methodology appeared in stochastic optimal control that transforms the optimality condition of SB into a set of SDEs. Crucially, these SDEs can be used to construct the likelihood objectives for SB that, surprisingly, generalizes the ones for SGM as special cases. This leads to a new optimization principle that inherits the same SB optimality yet without losing applications of modern generative training techniques, and we show that the resulting training algorithm achieves comparable results on generating realistic images on MNIST, CelebA, and CIFAR10. Our code is available at https://github.com/ghliu/SB-FBSDE.

研究动机与目标

• 建立薛定谔桥（SB）优化与现代深度生成建模（特别是基于得分的生成模型 SGM）之间的原则性联系，后者依赖于对数似然目标。
• 通过将 SB 训练建立在似然框架之上，解决 SB 模型是否可作为 SGM 的理论可靠替代方案的模糊性。
• 开发一种计算框架，利用 FBSDE 理论实现 SB 模型的基于似然的训练，克服以往基于启发式或 IPF 的优化方法的局限性。
• 将 SGM 似然目标泛化至完全非线性扩散过程，从而在统一的似然原则下统一 SB 与 SGM。
• 证明通过此似然框架训练的 SB 模型在标准基准测试中可实现具有竞争力或更优的图像生成性能。

提出的方法

• 该方法将 SB 的最优性条件表示为前向-后向 SDE（FBSDE）系统，利用随机最优控制理论将基于 PDE 的最优性条件转化为可处理的 SDE 形式。
• 通过求解 FBSDE 系统，构建 SB 的精确对数似然表达式，当扩散为线性时，该表达式可退化为 SGM 似然目标的特例。
• 框架采用基于散度的损失函数（如 KL 散度）进行前向与后向过程的训练，取代以往基于均值匹配的回归方法。
• 支持使用标准深度学习技术联合优化前向与后向 SDE，包括引入受 SGM 启发的朗之万校正器进行精细化处理。
• 训练过程通过交替优化前向与后向策略，利用推导出的似然目标，实现稳定且可扩展的训练。
• 该方法基于 PyTorch 实现，使用 U-Net 和 NCSN++ 等标准架构，在 MNIST、CelebA 和 CIFAR10 上进行了评估。

实验结果

研究问题

• RQ1 薛定谔桥能否通过泛化 SGM 似然框架的似然目标进行训练？
• RQ2 基于 FBSDE 的 SB 最优性表述如何实现与现代深度生成建模技术的原则性连接？
• RQ3 通过 FBSDE 实现的 SB 似然训练是否相比以往基于均值匹配回归的 SB 方法，能带来更高的样本质量和多样性？
• RQ4 推导出的似然目标是否可用于整合现代训练改进技术（如朗之万校正器）到 SB 模型中？
• RQ5 基于似然训练的 SB 模型在标准图像生成基准（如 MNIST、CelebA 和 CIFAR10）上的定量性能如何？

主要发现

• SB-FBSDE 框架为薛定谔桥推导出精确的对数似然表达式，其作为 SGM 似然目标的特例，建立了两类模型之间的根本理论联系。
• 在 CIFAR10 上，该方法使用 10,000 个生成样本实现了 11.85 的 Fréchet Inception Distance（FID），显著优于以往的最优传输模型。
• 即使在仅进行 5,000 次基于散度的训练（在完成 50,000 步 SGM 回归之后），FID 也从 33.68 降低至 13.35，表明性能快速提升。
• 定性结果表明，与以往 SB 方法相比，该方法在图像多样性与视觉质量方面表现更优，尤其在 CelebA 数据集的面部属性与背景变化方面。
• 该方法在 MNIST、CelebA 和 CIFAR10 上实现了最先进性能，且在 FID 评估中无需依赖朗之万校正器，表明其样本质量具有内在优越性。
• 该框架支持现代训练技术（如朗之万校正器）的应用，可在不同架构和数据集规模下持续提升样本质量。