发表机构
Ecole polytechnique; Institute of Foundation Models; MBZUAI; EPITA(巴黎综合理工学院; 基础模型研究院; 穆罕默德·本·扎耶德人工智能大学; 法国计算机科学与技术高等学院)
机构由 AI 辅助整理,请以论文原文为准。AI 中文总结
研究广义薛定谔桥问题,基于迭代马尔可夫拟合范式引入扭曲薛定谔桥匹配方法TSBM,它能处理连续和离散时间势,扩展了IMF方案,有新的桥匹配损失,还引入方差减少技术,实验证明其对高维轨迹推断有益。
AI 中文摘要
在过去几年中,基于扩散的薛定谔桥模型被提出以近似两个规定边界分布之间的最优传输动力学,并成功应用于生成建模。这些方法旨在估计一个路径测度,其初始和终端边缘与两个边界分布匹配,同时最小化相对于参考马尔可夫过程的库尔贝克-莱布勒散度。在这项工作中,我们考虑广义薛定谔桥问题,其中参考过程是扭曲布朗运动,即由时间依赖可微势诱导的布朗运动的费曼-卡茨变换。基于迭代马尔可夫拟合(IMF)范式,特别是其特殊情况扩散薛定谔桥匹配(DSBM),我们引入了扭曲薛定谔桥匹配(TSBM),一种基于扩散的方法,旨在处理连续和离散时间势。与以前的方法不同,TSBM为广义薛定谔桥问题提供了IMF方案的严格扩展。这一推导导致了一种新的桥匹配损失,它明确依赖于势的梯度,并在势消失时恢复DSBM目标,从而提高了性能。我们还引入了基于轨迹的方差减少技术,该技术大大稳定了优化,并且可能在当前设置之外有用。最后,我们通过实验证明了TSBM在包括人群导航和单细胞数据在内的越来越高维设置中的轨迹推断的好处。代码可在这个https URL上获取。
英文摘要
Over the past few years, diffusion-based Schrödinger bridge models have been proposed to approximate optimal transport dynamics between two prescribed boundary distributions, with successful applications to generative modeling. More precisely, these methods aim to estimate a path measure whose initial and terminal marginals match the two boundary distributions, while minimizing the Kullback-Leibler divergence with respect to a reference Markov process. In this work, we consider the generalized Schrödinger bridge problem, in which the reference process is a twisted Brownian motion, that is, a Feynman-Kac transform of a Brownian motion induced by a time-dependent differentiable potential. Building on the Iterative Markovian Fitting (IMF) paradigm, and in particular on its special case Diffusion Schrödinger Bridge Matching (DSBM), which corresponds to the zero potential case, we introduce Twisted Schrödinger Bridge Matching (TSBM), a diffusion-based method designed to handle both continuous- and discrete-time potentials. Unlike previous approaches, TSBM provides a rigorous extension of the IMF scheme to the generalized Schrödinger bridge problem. This derivation leads to a new bridge-matching loss that depends explicitly on the gradient of the potential and recovers the DSBM objective when the potential vanishes, yielding improved performance. We further introduce trajectory-based variance-reduction techniques that substantially stabilize optimization and may be useful beyond the present setting. Finally, we empirically demonstrate the benefits of TSBM for trajectory inference across increasingly high-dimensional settings, including crowd navigation and single-cell data. Code available at https://github.com/maxencenoble/twisted-sb-matching.
CommentsPreprint version