Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Go With the Flow: Fast Diffusion for Gaussian Mixture Models

George Rapakoulias, Ali Reza Pedram, Fengjiao Liu, Lingjiong Zhu, Panagiotis Tsiotras

TL;DR

This work introduces GMMflow, a training‑free diffusion method for Schrödinger bridges between Gaussian mixtures. It leverages a componentwise decomposition of Gaussian bridge subproblems and a linear program to select couplings, yielding a tractable, scalable policy where the dimension scales linearly with the number of mixture components. The framework extends to multi‑marginal and linear‑time‑varying prior dynamics, and to continuous Gaussian mixtures and heavy‑tailed priors via mixing measures, with theoretical guarantees that the GMM policy achieves an upper bound on the true SB cost. Empirically, GMMflow delivers faster training and competitive or superior marginal distribution fitting compared with lightweight neural SB solvers, demonstrated on 2D benchmarks, latent space image translation, and multi‑marginal cellular diffusion data. The approach is released as open source and holds promise for rapid SB solutions in smaller problems and distributions well‑captured by Gaussian mixtures.

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.

Go With the Flow: Fast Diffusion for Gaussian Mixture Models

TL;DR

This work introduces GMMflow, a training‑free diffusion method for Schrödinger bridges between Gaussian mixtures. It leverages a componentwise decomposition of Gaussian bridge subproblems and a linear program to select couplings, yielding a tractable, scalable policy where the dimension scales linearly with the number of mixture components. The framework extends to multi‑marginal and linear‑time‑varying prior dynamics, and to continuous Gaussian mixtures and heavy‑tailed priors via mixing measures, with theoretical guarantees that the GMM policy achieves an upper bound on the true SB cost. Empirically, GMMflow delivers faster training and competitive or superior marginal distribution fitting compared with lightweight neural SB solvers, demonstrated on 2D benchmarks, latent space image translation, and multi‑marginal cellular diffusion data. The approach is released as open source and holds promise for rapid SB solutions in smaller problems and distributions well‑captured by Gaussian mixtures.

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.

Paper Structure

This paper contains 38 sections, 10 theorems, 100 equations, 10 figures, 7 tables, 3 algorithms.

Table of Contents

  1. Introduction and Background
  2. Preliminaries
  3. Diffusion Schrödinger Bridge Matching and Flow Matching
  4. Schrödinger Bridges with Gaussian Marginals
  5. Momentum Schrödinger Bridges with multiple Gaussian marginals
  6. Fast Diffusion for Mixture Models
  7. Gaussian Mixture Schrödinger Bridge
  8. Multi-Marginal Problems
  9. Continuous Gaussian Mixtures
  10. Related Work
  11. Experiments
  12. 2D Problems and Benchmarks.
  13. Image-to-Image Translation.
  14. Multi-Marginal Problems.
  15. Conclusion and Limitations
  16. ...and 23 more sections

Key Result

Proposition 1

bunne2023schrodinger The optimal solution to Problem (GSB) is given by $u_t(x) = K_t(x-\mu_t) + v_t$ with $\mu_t = (1-t)\mu_0 + t \mu_1$, $v_t = \mu_1 - \mu_0$, and $K_t = S_t ^{\hbox{\tiny\sf T}} \Sigma_t^{-1}$, where with $C_{\epsilon} = \frac{1}{2} ( \Sigma_0^{\frac{1}{2}} D_{\epsilon} \Sigma_0^{-\frac{1}{2}} - \epsilon I)$ and $D_{\epsilon} =( 4 \Sigma_0^{\frac{1}{2}} \Sigma_{1} \Sigma_0^{\fr

Figures (10)

  • Figure 1: Gaussian to 8-Gaussians with no prior dynamics and for various levels of noise.
  • Figure 2: GT to DCSL Schrödinger Bridge with zero prior dynamics and $\epsilon=0.1$.
  • Figure 3: Adult to Child (left) and Man to Woman (right) image translation task.
  • Figure 4: GMM momentum multimagrinal SB: Real vs Generated data.
  • Figure 5: Comparison between true Wasserstein distance and upper bound (\ref{['J_star']}) for 1D Student-t distributions.
  • ...and 5 more figures

Theorems & Definitions (12)

  • Proposition 1
  • Proposition 2
  • Theorem 1
  • Theorem 2
  • Theorem 3
  • Theorem 4
  • Theorem 5
  • Theorem 6
  • Theorem 7
  • proof
  • ...and 2 more