Diffusion Schrödinger Bridge with Applications to Score-Based Generative Modeling
Valentin De Bortoli, James Thornton, Jeremy Heng, Arnaud Doucet
NeurIPS 2021 (spotlight)
stat.ML, cs.LG, math.PR
2021-06-02
用神经网络做 IPF 来解薛定谔桥：第一轮就是 Song 的 SGM，后面几轮把短时程的两端边缘补上。MNIST 12 步、CelebA 用 50 步 10 轮能出图，还没追上当时 SOTA。
Score-based 生成模型把数据逐步加噪，直到接近一个好采样的先验（通常是高斯），再学反向去噪。Song et al. (2021) 把加噪写成 SDE，反向漂移由各时刻扰动分布的 score（对数密度梯度）给出，用神经网络做 score matching。
问题卡在时间预算上。正向过程必须跑够久，终点才接近高斯；步长还得够小，Euler 离散才准。两件事叠在一起，采样步数被拉长。Song and Ermon (2019, 2020) 的早期设定里，N 要到 100。
薛定谔桥换了问题本身：给定参考扩散、数据分布 pdata 和先验 pprior，在路径空间上找与参考过程 KL 最近、且 t=0、t=T 边缘分别等于 pdata、pprior 的扩散。生成不必等正向混到高斯，有限时间两端就能对上。静态版本是熵正则最优传输，噪声趋于零时收回二次 Wasserstein 计划。
求解工具是 Iterative Proportional Fitting（IPF）。从参考加噪过程出发，奇数轮强制终点等于 pprior，偶数轮强制起点等于 pdata，来回投影。有限离散状态上，这就是 Cuturi (2013) 的 Sinkhorn。
DSB 把每一轮的正向、反向转移近似成各向同性高斯，均值由两个神经网络 F 和 B 给出。对所有历史 score 做回归会把显存和计算量堆得很高，所以改成 mean-matching：用当前正向漂移把样本推到下一步，让反向网络拟合这一步实际从哪来。
T 可以短。短了以后第一轮会失败，因为 pN 不是高斯；多轮 IPF 把这座桥修直。
Theorem 1 给出 Song 方法的第一个定量全变差界。score 一致逼近误差不超过 M 时，生成分布与 pdata 的 TV 距离拆成离散化加拟合误差，以及正向混合误差。OU（α>0）的混合项按 exp(-√α T) 衰减，纯布朗运动按 T^{-1}+T^{-1/2}。混合时间只由正向布朗或 OU 决定，与数据维数无关，高维 Langevin 那种混不匀在这条时间非齐次链上不出现。α 和 T 变大时离散化常数一起变差，步长必须更小。
Proposition 4 和 5 把 IPF 从紧支撑推到非紧连续空间：相邻迭代的 KL 与 TV 单调下降，n 乘两端边缘 KL 趋于 0，比 Léger (2020) 的 C/n 边际界更强，联合分布在 TV 下收敛。再加密度正则，极限就是薛定谔桥。高斯参考过程还能收回几何速率。DSB 是这条连续时间动态 IPF 的时间离散。
高斯可控实验：pdata = N(a,I)、pprior = N(-a,I)，a=0.1 乘全 1 向量，参考过程为布朗运动，静态 SB 有闭式协方差。batch 128、N=20、γ=1/40。d=5 时约 3 万参数网络能对上均值、方差、协方差；d=50 时同样结构对不上后两项，隐层加宽到约 24 万参数才跟上。
二维玩具把 T 压到 0.2（N=20，γk=0.01）。第一轮就是原版 SGM，终点离高斯很远，反演回不到数据。Figure 1 里 5 轮后两端接上；20 轮后合成明显好于第一轮。
图像用缩小版 U-Net（基于 Nichol and Dhariwal, 2021）：
| 设置 | N（时程 T） | DSB 轮次 | 结果 |
| 2D 玩具 | 20（0.2） | 1 vs 5/20 | 短 T 下第一轮失败，多轮后接到数据 |
| MNIST | 12 | 8 | 相对第一轮有可见改善；早期 SGM 常用 N=100 |
| CelebA 32×32，d=3072 | 50（0.63） | 10 | 能出可辨人脸；论文称这是高维 SB 首个可跑通的算法 |
| EMNIST→MNIST | 30（1.5） | 10 | 先验换成另一份数据，走出插值桥 |
FID 随 DSB 轮次下降，对多组 N 都成立。正文只有 Figure 6 的曲线，没有把具体 FID 和 DDPM、Score SDE 列成对照表。论文写明实现尚未追上当时 SOTA。
少步数采样不再依赖正向已经混成高斯。短 T 时多做几轮 DSB，用训练轮次换采样步数。MNIST 12 步、CelebA 50 步，对照早期 SGM 的 100 步，方向是对的。
先验可以不是高斯。pprior 换成另一份数据，得到两条数据分布之间的熵正则传输，连续状态空间里就是 Sinkhorn 的扩散版。论文点名可往多边际 SB、Wasserstein 重心、熵正则 Gromov–Wasserstein、域适配上用。解仍是扩散，可按 Song et al. 改写成等边际神经 ODE 算精确似然。
这是把 SGM 接到最优传输上的框架。图像质量停在能出图、加轮次有帮助，不是新的 SOTA 采样器。加速采样还要和噪声日程、更好离散化、蒸馏一起看。
T 不能任意接近 0。正向终点离 pprior 太远时，先验支撑附近几乎没有训练信号，score 仍然难学。实现也还没追上 SOTA。
训练是串行的一列网络。后轮可以微调前轮，总计算仍明显高于单次 score matching。F 和 B 若联合训练，平衡点是某座桥，不一定是薛定谔桥。
Theorem 1 要求 score 误差有一致上界 M，实践里最多能期望在数据分布上高概率成立。IPF 的 o(1/n) 针对精确迭代；DSB 每轮还有网络拟合误差和欧拉误差，三层误差没有合成一条可计算的生成误差界。
图像实验缺 FID 数字表，CelebA 只用 32×32。同期 Vargas et al. (2021) 用高斯过程估漂移，这篇改神经网络，两边没有直接对比。