Go With the Flow: Fast Diffusion for Gaussian Mixture Models

TL;DR

This paper introduces an efficient analytic method for computing Schrödinger Bridges between Gaussian Mixture Models, avoiding expensive training and enabling applications like image translation and cellular dynamics modeling.

Contribution

It proposes a low-dimensional linear program approach for optimal policies in Schrödinger Bridges between GMMs, generalizing to linear time-varying systems and multi-marginal problems.

Findings

Efficient computation of SBs for GMMs using linear programming.

Application to image translation and cellular dynamics.

Method scales linearly with the number of mixture components.

Abstract

Schrodinger Bridges (SBs) are diffusion processes that steer, in finite time, a given initial distribution to another final one while minimizing a suitable cost functional. Although various methods for computing SBs have recently been proposed in the literature, most of these approaches require computationally expensive training schemes, even for solving low-dimensional problems. In this work, we propose an analytic parametrization of a set of feasible policies for steering the distribution of a dynamical system from one Gaussian Mixture Model (GMM) to another. Instead of relying on standard non-convex optimization techniques, the optimal policy within the set can be approximated as the solution of a low-dimensional linear program whose dimension scales linearly with the number of components in each mixture. The proposed method generalizes naturally to more general classes of dynamical systems, such as controllable linear time-varying systems, enabling efficient solutions to multi-marginal momentum SBs between GMMs, a challenging distribution interpolation problem. We showcase the potential of this approach in low-to-moderate dimensional problems such as image-to-image translation in the latent space of an autoencoder, learning of cellular dynamics using multi-marginal momentum SBs, and various other examples. The implementation is publicly available at https://github.com/georgeRapa/GMMflow.