Distributionally Robust Schrödinger Bridge

Schrödinger bridge（SB）学习给定初始分布与目标分布之间的随机传输。当初始分布在测试时发生偏移，学习到的动力学可能无法恢复目标分布。我们提出 Distributionally Robust Schrödinger Bridge（DRSB），学习一个能够应对初始分布不确定性的单一控制器。DRSB 的目标函数由控制能量以及所得终端分布与目标分布之间的 KL 惩罚项组成。DRSB 旨在寻找一个单一控制器，使初始分布在名义分布周围的模糊集合内变化时，该目标函数的最坏情况取值最小。我们推导了该目标函数的精确变分形式，并将其终端代价固定的子问题与随机最优控制和分布鲁棒优化联系起来。该形式促成了一种交替算法：更新对抗性初始分布、估计终端对数密度比并训练控制器。我们利用随机控制最优性条件，开发了 Wasserstein 和 Sinkhorn 变体，以近似对抗性更新所需的梯度。在二维传输任务和图像到图像翻译上的实验表明，与标准 SB 相比，该方法对输入扰动具有更强的鲁棒性，但需要牺牲一定的名义性能。在高斯混合传输任务中，Sinkhorn DRSB 在测试的两个未见过的噪声水平下，其平均切片 Wasserstein 距离也低于固定噪声水平的数据增强方法。