Title: Go With the Flow: Fast Diffusion for Gaussian Mixture Models
Venue: NeurIPS 2025
Accepted: yes
Keywords: Schrodinger Bridge problem

Optimal transport
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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors proposes a training‑free method to solve Schrödinger Bridge problems when both endpoints are Gaussian Mixture Models (GMMs). The method expresses a global policy as a weighted mixture of closed‑form gaussian bridges, and chooses weights via a low‑dimensional linear program whose size grows with the number of mixture components. The paper gives convincing numerical results on a set of toy and real world problems including image-to-image translation.\n\nSTRENGTHS AND WEAKNESSES:\n*Strengths*\n- The paper is clear and easy to read.\n- Training-free methods and Gaussian Mixture Models are probably under-rated in a world dominated by expensive generative methods. I like to see this kind of paper. \n- The paper server as a good introduction to the area\n- The derivation of the method is clear and the theory is substantial.\n- The results for image-to-image translation tasks are compelling.\n\n*Weaknesses*\n- The clarity of the paper would be enhanced with an Algorithm block describing the full implementation. \n- It would be valuable to know more about the sensitivity when boundary data are poorly approximated by a GMM. It would be great to see an experiment with error on the y-axis and number of GMM components on the x-axis for some of the examples. Is most of the error induced by fitting the marginals, or even if the marginal are well approximated by the GMM can other sources of error dominate?\n- The paper is sold as a \"fast\" method yet there is no runtime analysis, empirical or theoretical. This type of analysis would shed light on key questions such as how does the runtime scale with dimension and number of samples etc. Similarly an experiment with runtime on the y-axis and components on the x-axis would be useful.\n\nQUESTIONS:\n- The results for image-to-image translation tasks are compelling. Moreover it is surprising that you can get away with only a 10-component diagonal GMM. Do the authors imagine you do full image generation on some constrained datasets in a latent space? At what point does the method break down? It would be interesting to see the GM assumptions pushed to its limits.\n\n*Possible action items (see weaknesses above)*\n- Runtime analysis in terms of number of components, number samples, etc. (what ever is most relevant) for both training and inference.\n- Algorithm block \n- Error vs number of mixture components.\n\n(I am not insisting that all of these must be done, I understand time constraints, but I do believe they would improve the paper.)\n\nLIMITATIONS:\n- The author discuss, but obvious limit is the Gaussian assumption.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors address the Schrödinger Bridge (SB) problem for Gaussian mixture endpoints (SB-GM). They observe that the SB between two individual Gaussians (SB-G) admits a closed-form solution, and hence propose to parametrize the feasible SB-GM solutions as a weighted mixture of SB-G solutions. By doing so, they mitigate the original optimization problem and optimize only over the mixture weights to receive the feasible solution. They the authors define a linear optimization problem for optimization over the mixture weights and rigorously show that the linear optimization problem objective upper-bounds the true optimal cost for the SB-GM. Building on this insight, the authors extend their framework to both the multi-marginal momentum SB and to SBs between infinite Gaussian mixtures.\n\nIn practice, they approximate arbitrary distribution endpoints $p_0$ and $p_1$ by finite Gaussian mixtures, construct SB-G solutions for every pair of mixture components, and recover the approximation of overall SB-GM by solving a linear optimization problem over mixture weights.  An analogous scheme is developed for the multi-marginal momentum SB. The method is validated on 2D toy problem, Image-to-Image translation in the latent space of ALAE and learning dynamics of cell populations.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**:\n\n- The proposed algorithm is both novel and computationally efficient, demonstrating strong performance in low-dimensional settings.\n- The authors’ approximation strategies for diverse SB formulations—including the multi-marginal momentum SB and the infinite-mixture SB—represent valuable conceptual contributions.\n- Empirical results on image-to-image translation, EOT benchmark (Appendix C)  and cell-population dynamics modeling show that GMMflow mostly matches state-of-the-art methods.\n- The proposed algorithm tries to advance the area of SB methods specifically crafted for low dimensional cases.\n\n**Weaknesses**:\n\n- Authors do not provide guarantees for a mixture of SB-G solutions to reach the optimal solution for the SB-GM problem. In general, only the upper bound on the optimal transport cost can be received that way. For the general SB the approximation of its distribution endpoints by the gaussian mixtures is already an approximation and the mixture of SB-G solutions adds a second approximation on top.\n    - Would be nice to know more about the solutions gotten as the result of solving linear optimization problems in Theorems 2, 4, 6 and their relation to the optimal ones. Maybe there is some cases where you can find exactly the optimal solutions. For other related questions see **Questions** section.\n- The number of variables in the linear optimization problem (mixture weights) grows linearly with the number of components in Gaussian Mixtures used for approximation of distributions. In the two distribution case number of variables is: $(M_1 \\times M_2)$, or in the multi-marginal case  $(M_1 \\times M_2 … \\times M_k)$. This seems like a lot of mixture weights, especially when one considers complex distributions that require a lot of Gaussian Mixture components to be approximated accurately. Additionally one has to solve the linear optimization problem, with its own computational complexity, to optimize over these mixture weights.\n    - Would like to see a study on the growth of Gaussian Mixtures components along with  training time and method quality.\n- As far as I understand the inference of the SB-G mixture solution requires the simulation of the SDE. Which is rather unconvenient, especially since LSB/LSBM and Gaussian to Gaussian SB solution transport plans can be inferred without the SDE simulation.\n- The evaluation:\n    - In the Image-to-Image translation experiment only the target distribution matching  metrics (FID, ALAE-BW) are shown. Which is weird since the target matching is guaranteed, up to GMM approximation error, by the construction of the method. On the other hand, the optimality of transport cost is not guaranteed. Notably by the Figure 3 one can see that LightSB transport plan has lower transport cost.\n    - The only experiment that quantitatively evaluates the transport cost is EOT benchmark shown in Appendix C. The performance of method becomes better with growth of $\\epsilon$, while the lower $\\epsilon$ examples do mostly evaluate the transport cost of the method and higher $\\epsilon$ examples do mostly evaluate the ability of method to fit the ending ($p_1$) distribution. This experiment results can be seen as proof that the proposed method is bad as transport cost but good at distribution fitting.\n    - In the EOT benchmark experiment some baselines are missing: DSBM [1], SF2M-sink [2], LightSB [3], LightSBM [4]. All the metrics are already computed, see Table 1 [4].\n\n**Typos/Minor text concerns**:\n\n- Line 128: $p_i, p_f$ should it be $p_0, p_1$?\n- Line 298, 300: Bures-Wasserstein distance\n- Line 332, 337, Table 3: NSBL → NLSB\n\n[1] Shi, Y., De Bortoli, V., Campbell, A., & Doucet, A. (2023). Diffusion schrödinger bridge matching. *Advances in Neural Information Processing Systems*, *36*, 62183-62223.\n\n[2] Tong, A. Y., Malkin, N., Fatras, K., Atanackovic, L., Zhang, Y., Huguet, G., ... & Bengio, Y. (2024, April). Simulation-Free Schrödinger Bridges via Score and Flow Matching. In *International Conference on Artificial Intelligence and Statistics*(pp. 1279-1287). PMLR.\n\n[3] Korotin, A., Gushchin, N., & Burnaev, E. Light Schrödinger Bridge. In *The Twelfth International Conference on Learning Representations*.\n\n[4] Gushchin, N., Kholkin, S., Burnaev, E., & Korotin, A. (2024, July). Light and optimal schrödinger bridge matching. In *Forty-first International Conference on Machine Learning*.\n\n[5] Koshizuka, T., & Sato, I. Neural Lagrangian Schr\\\"{o} dinger Bridge: Diffusion Modeling for Population Dynamics. In *The Eleventh International Conference on Learning Representations*.\n\nQUESTIONS:\n- Is the optimal solution for the SB problem with gaussian mixture endpoints lies in the set of solutions defined in Theorems 1, 3, 5?\n- Is the $J_{\\text{OT}}$ in Theorems 2, 4, 6 transport cost for the corresponding $\\lambda$ defined solutions?\n- Are there any results on tightness of upper bounds in Theorems 2,4,6?\n- What is the computational complexity of the linear optimization algorithms used, or that could be potentially used? Including worst cases. Is there any results on training time and method quality w.r.t. number of gaussian mixtures used for approximation (preferably on EOT bench)?\n- Is there results on transport cost for the ALAE experiment?\n- Why do you have relatively high SWD and very low MMD on cell population dynamics learning experiment?\n- How do you apply NLSB [5] to multi-marginal SB problem in the cell population dynamics learning experiment? As far as I understand NLSB is general SB method. If you consider the general SB method, then I think you also should consider methods [1, 2, 3, 4].\n- What potential algorithms can be considered for solving the stated in Theorem 6 optimization problem?\n\nLIMITATIONS:\nMostly yes, but I would like authors to clarify the complexity of linear optimization methods they use w.r.t. number of components in Gaussian Mixtures. Depending on the answer, there is a possibility one could consider it a limitation.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-01"}

comments: SUMMARY:
The authors provide a specific formulation of optimal transport that is useful in low-dimensional regimes and avoids the computationally complexity suffered by more established models in high-dimensional regimes. In particular, they propose a policy which is theoretically feasible and an associated probability flow for the setting of transportation between mixed Gaussian conditionals, by solving Gaussian bridge sub-problems. They support their theoretical findings with experiments comparing their methodology to others' on toy data, image matching, and single-cell data with their methodology outperforming contemporary methods in every instance.

STRENGTHS AND WEAKNESSES:
The authors concisely summarize contemporary methods for optimal transport between probability distributions and describe the gap their paper intends to close. They provide theoretical evidence supporting their proposed transport plan and make a clear delineation between existing methods and what their method offers computationally.  The empirical results are strong; their proposed transport plan outperforms contemporary methods across a diverse set of experiments across different metrics. A potential improvement, as mentioned by the authors, is to not constrain the proposed methodology to the setting of Gaussian Mixture Models. It could also be worth empirically exploring around what dimension the proposed model loses its efficacy through ablation studies.

QUESTIONS:
N/A

LIMITATIONS:
Yes, the authors have addressed limitations. The other limitations are included in the weaknesses mentioned in the previous paragraph.
TITLE: Official Review
CLARITY: 4
IMPACT: 4
ORIGINALITY: 3
DATE: 2025-07-03

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper develops a computational method to compute suboptimal solutions of the schrödinger bridge problem between gaussian mixture models, and showcases the potential of the problem in low-dimensional settings. The key idea is to exploit the decoupling of particles dynamics with the assignment problem, building on top of recent work on the topic.\n\nSTRENGTHS AND WEAKNESSES:\nThe paper is very focused on a specific problem with clearly stated limitations. The theory is correct (beside a few minor clarifications needed) and gradually introduces more complexity, rendering the paper very pleasant to read. I identified the following weaknesses, that the authors may consider to address to strengthen their contribution:\n\n1. Although reading the entire paper makes the scope of the work very clear, I believe it could be stated better in the introduction: low-dimensional problem in which the marginals can be well approximated by GMM, and in which there is a need for a good solution, fast, but not necessarily optimal. To build on this, the authors could also exemplify a few of these applications.\n2. In the contributions, it is stated that the paper proposes a solution to the SB problem with GMM, but technically, it only provides a suboptimal solution. Moreover, the suboptimality gap is not quantified and may be very large. Carefully addressing this gap or clarifying this potential limitation will strengthen the paper for other people to build on top.\n3. The authors cite https://epubs.siam.org/doi/abs/10.1137/23M1560902, but perhaps they could also contextualize their result with respect to the recipe provided in https://epubs.siam.org/doi/abs/10.1137/23M1560902. Namely, the mixing discussed in Go with the Flow is an approximation of the mixing derived in https://epubs.siam.org/doi/abs/10.1137/23M1560902, hence the suboptimality. Note that the work cited is for deterministic case, another suboptimality comes from the stochastic dynamics in this type of mixing. When comparing with the literature, in the deterministic case, they could consider comparing to https://epubs.siam.org/doi/abs/10.1137/23M1560902, which provides an exact solution. They could compare it also in the stochastic case, to have an idea of a different cost approximation.\n4. The suboptimality of the derived bound is unclear. Can the authors comment on it?\n5. The efficiency benefits are somewhat minor (the algorithms considered are all very fast), so it is hard to tell if the benefits come from algorithmic speedups or from implementation details without a complexity analysis.\n6. The setting for the RNA dataset is not very clear (it is almost never clear in the literature either….). Can you please clarify the objective, how do you evaluate, etc? Also, for the experiments, while I appreciate the numerous experiments I believe that it would be better to defer more to the appendix and discuss more thoroughly less experiments in the main text.\n\nSome other comments:\n\n1. The authors present 6 theorems. While I appreciate the increasing complexity, reversing the order of the theorem would allow to discuss some of the results as corollary and gain some space for more thorough discussions in the main text. It would also allow to avoid essentially duplicated proofs. I would also make a Lemma instead of a Theorem for the first theorem of each “pair” of theorems.\n2. Can you please double check Proposition 1? The expression for the velocity field seems wrong (independent of t). Or could the authors perhaps clarify?\n3. Before equation 11, the authors refer to Problem 6 and 10, but then mention replacing the first constraint only in 6 (typo? perhaps you want to add 10 as well). I would also not use the word sequel as it made me think of a follow up paper.\n4. Perhaps stop the sentence “Problem (24) is challenging to solve in general.” instead of the “,”.\n5. By the way, I find this decoupling approach very interesting. It also relates to a recent work on a similar, yet slightly different setting https://arxiv.org/abs/2406.12616. There the decoupling is somewhat dual (for estimation).\n\nQUESTIONS:\n1. Can the authors discuss further their suboptimality gap? Related: How does the approximation of GMM further exacerbate this gap?\n2. Can the authors provide comparisons (at least in the deterministic case) with the approach in https://epubs.siam.org/doi/abs/10.1137/23M1560902?\n3. Can the authors clarify the setting for the RNA dataset?\n\nLIMITATIONS:\nThe authors adequately discuss the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose a lightweight method for approximating solutions to SB problems. Their approach has them approximate the boundary distributions by Gaussian Mixtures (GM) and then approximately solving the SB problems for the GMs. Their theoretical results motivate the algorithm for doing the latter, showing that any solution to the GM SB problems has a policy and flow that can be described as a combination of the policy and flow for the SB solution between each of the individual Gaussians. They provide experiments comparing their results to existing SB algorithms, demonstrating improvements.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\n- The algorithm appears to be highly efficient and most SB algoriths are highly inefficient\n- The approach of using GMs as the endpoints is simple and yet appears to be quite novel to the literature. The proposal of this could motivate future work\n- The paper is largely well-written with the preliminaries section especially serving as a high quality introduction to the field\n- The theoretical results regarding SB for GMs is interesting in its own right and could motivate future work\n- The extension of the work to momentum SB and continuous GMs demonstrates a deep exploration of the ideas considered in the work.\n\nWeaknesses\n- The algorithm they ultimately use is not very clearly presented or discussed much. It would be good to include a clear presentation of this following the theoretical results. For example, it would be good to know how J_OT should be optimised\n- The Comparison seems somewhat limited, focussing primarily on comparison with other light methods. A more thorough comparison with standard SB algorithms would help.\n- It is not clear to me how close J_GMM is to the proxy J_OT. There are some experiments in table 4 that provide some results in this direction, but a more thorough discussion and set of experiments about how much the proposed method actually solves the GM problem would greatly improve to the validity of the paper.\n\nQUESTIONS:\n- Does the final algorithm optimise J_OT to obtain \\lambda_{ij}?\n- If so, is it clear how close the objective is to J_GMM or how well the approximation is actually solving the GM SB problem (not just the OT problem). Could you please point to evidence for this?\n- Is \\rho_{ij} and u_{ij} obtainable in closed-form? If not, how difficult is it to approximate?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-07"}