On Approaches to Building Surrogate ODE Models for Diffusion Bridges

TL;DR

This paper introduces surrogate models for diffusion bridges that significantly improve efficiency and interpretability, using sparse regression and neural-ODE reformulations, demonstrated on Gaussian and MNIST tasks.

Contribution

It proposes two novel algorithms, SINDy-FM and DSBM-NeuralODE, to create simpler, faster, and more interpretable diffusion bridge models, advancing practical deployment.

Findings

SINDy-FM achieves high interpretability with sparse, symbolic models.

Surrogates drastically reduce parameter counts and inference time.

Models maintain competitive performance on benchmark tasks.

Abstract

Diffusion and Schr\"{o}dinger Bridge models have established state-of-the-art performance in generative modeling but are often hampered by significant computational costs and complex training procedures. While continuous-time bridges promise faster sampling, overparameterized neural networks describe their optimal dynamics, and the underlying stochastic differential equations can be difficult to integrate efficiently. This work introduces a novel paradigm that uses surrogate models to create simpler, faster, and more flexible approximations of these dynamics. We propose two specific algorithms: SINDy Flow Matching (SINDy-FM), which leverages sparse regression to identify interpretable, symbolic differential equations from data, and a Neural-ODE reformulation of the Schr\"{o}dinger Bridge (DSBM-NeuralODE) for flexible continuous-time parameterization. Our experiments on Gaussian transport tasks and MNIST latent translation demonstrate that these surrogates achieve competitive performance while offering dramatic improvements in efficiency and interpretability. The symbolic SINDy-FM models, in particular, reduce parameter counts by several orders of magnitude and enable near-instantaneous inference, paving the way for a new class of tractable and high-performing bridge models for practical deployment.