Reproduce “Aligned Diffusion Schrödinger Bridges” on OpenResearch

About this paper

Vignesh Ram Somnath, Matteo Pariset, Ya-Ping Hsieh, Maria Rodriguez Martinez, Andreas Krause, Charlotte Bunne

A framework for Diffusion Schrödinger Bridges rigorously integrates explicit data alignment, allowing for the inference of pathwise optimal solutions without relying on iterative proportional fitting. This approach improves trajectory reconstruction and endpoint prediction, achieving a median RMSD of 3.80 A in protein docking compared to 19.99 A for baselines, and enhances performance across single-cell differentiation tasks.

Key insights

• Explicitly leveraging data alignment simplifies the Diffusion Schrödinger Bridge problem, leading to a direct training objective that avoids the numerical instabilities of Iterative Proportional Fitting (IPF).
• The optimal stochastic process can be characterized by equating the drift of a known conditioned reference process (scaled Brownian bridge) with a learned drift derived from Doob's h-transform, facilitating direct optimization.
• The method achieves pathwise optimality, enabling the prediction of specific terminal states for individual initial samples, which is crucial for applications like protein docking.

Summary by alphaXiv.