On Approaches to Building Surrogate ODE Models for Diffusion Bridges
This addresses the problem of slow and complex training in generative modeling for researchers and practitioners, offering a more tractable approach, though it appears incremental by building on existing bridge models.
The paper tackles the high computational costs and complexity of diffusion and Schrödinger Bridge models by introducing surrogate models to create simpler, faster approximations, achieving competitive performance with dramatic efficiency gains, such as reducing parameter counts by several orders of magnitude and enabling near-instantaneous inference.
Diffusion and Schrö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ö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.