GMMflow

GMMflow solves Schrödinger bridge and optimal transport problems whose boundary distributions are Gaussian mixture models, without training a neural network. The drift of the controlled SDE carrying one mixture to the other is given a closed-form parameterisation, and its optimal parameters are recovered by solving a linear program, which makes the solver light enough to run where a learned bridge would be impractical.

The control-theoretic formulation extends to the multi-marginal momentum Schrödinger bridge and to problems with general linear prior dynamics. The repository is the official PyTorch implementation, with notebooks for the Gaussian mixture examples, the multi-marginal case and an image-translation experiment.

Reference

G. Rapakoulias, A. R. Pedram, F. Liu, L. Zhu and P. Tsiotras, “Go With the Flow: Fast Diffusion for Gaussian Mixture Models,” Thirty-Ninth Annual Conference on Neural Information Processing Systems (NeurIPS), 2025.

Visit the GitHub repository for the source code.