Diffusion Schrödinger Bridge Matching

Solving transport problems, i.e. finding a map transporting one given distribution to another, has numerous applications in machine learning. Novel mass transport methods motivated by generative modeling have recently been proposed, e.g. Denoising Diffusion Models (DDMs) and Flow Matching Models (FMMs) implement such a transport through a Stochastic Differential Equation (SDE) or an Ordinary Differential Equation (ODE). However, while it is desirable in many applications to approximate the deterministic dynamic Optimal Transport (OT) map which admits attractive properties, DDMs and FMMs are not guaranteed to provide transports close to the OT map. In contrast, Schr\"odinger bridges (SBs) compute stochastic dynamic mappings which recover entropy-regularized versions of OT. Unfortunately, existing numerical methods approximating SBs either scale poorly with dimension or accumulate errors across iterations. In this work, we introduce Iterative Markovian Fitting (IMF), a new methodology for solving SB problems, and Diffusion Schr\"odinger Bridge Matching (DSBM), a novel numerical algorithm for computing IMF iterates. DSBM significantly improves over previous SB numerics and recovers as special/limiting cases various recent transport methods. We demonstrate the performance of DSBM on a variety of problems.

Paper

Similar papers

Peer review

Reviewer Lk9k6/10 · confidence 4/52023-07-01

Summary

This paper formulates a new framework of a probabilistic model called Diffusion Schrödinger Bridge Matching (DSBM) and an iterative algorithm which is called iterative Markovian fitting (IMF). Firstly, the paper has contribution of pointing out a relationship between famous generative models (score matching and flow matching) in the form of DSBM. IMF is also novel and might be more suitable to be trained in neural networks than well-established algorithms such as IPF. To achieve these, the authors utilize a tractable bridge structure (Brownian bridge), which enables preserving key properties, i.e., densities of marginal distribution at the start and end of the process. The proposed algorithm often shows better generation results in certain tasks, such as image transfer.

Strengths

1. The paper is well-written despite dealing with a complex subject. As far as I understand, most of the arguments seem to be theoretically sound. 2. SB algorithms have suffered from numerical errors and computational complexity. Since the IMF approach is akin to standard matching algorithms for neural networks, the provided empirical results show that it scales well. 3. The paper provides a novel perspective on the diffusion bridge problem with a novel algorithm. As a result, DSBM-IMF shows promising results in certain tasks.

Weaknesses

1. Significance. In Tables 2, 3 and Figure 10, It is difficult for me to observe that DSBM-IMF excels DSMB-IPF. Therefore, I would say significance of the proposed algorithm is not clearly shown in the experiments. The paper needs more empirical evidence of the actual benefits of choosing this algorithm, such as performance boost and scalability, in more challenging, or large-scale settings. 2. Limited applications. I believe the actual implementation of the algorithm heavily relies on the well-known Brownian bridge. Although I think this is actually a valid contraint for certain applications such as image generation, at the same time it means that the algorithm is not always applicable to general optimal transport problems. Therefore, I recommend putting more effort on analyzing more classes of $\mathbb{Q}$ in the experiment (if there are any), or putting clear justification on certain classes of $\mathbb{Q}$ used in the experiment.

Questions

1. Does $f_t$ is predefined in all tasks? 2. Is $\sigma_t$ constant and **not** scheduled? The authors of [1] mentioned time-symmetric scheduling the diffusion term, as suggested by prior SB models. 3. Could you explain how the IMF scheme is fundamentally better than IPF algorithms? [1] Guan-Horng Liu et al. (2023), I2SB: Image-to-Image Schrödinger Bridge.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Although I think overall paper is good, but I also think the proposed experimental results are not quite sufficient; the paper need more empirical evidence to show the excellence of the model. I hope to see response from the authors for this point.

Reviewer 7Dnn7/10 · confidence 3/52023-07-03

Summary

Flow matching and alpha blending have achieved tremendous attention in the matching problems. Although the straight trajectories based on $X_t = (1-\alpha)X_0+\alpha X_1$ yield **fast inference**, it doesn't necessarily mean they are efficient in score estimations in training. To make the training theoretically more efficient, rectified flow leverages the optimal transport ideas to minimize the transport cost. However, the stochastic alternative is still lacking (stochastic interpolant may not be optimal in terms of optimal transport). To tackle this issue, the authors proposed the stochastic version of flow matching/ alpha blending with entropic optimized transport. A novel scheme IMF that resembles the IPF iteration is proposed.

Strengths

1. Despite the advances of stochastic interpolant, the transport from marginals may not be optimal in general. The **optimal-transport**-based stochastic version of flow matching and alpha blending is **still missing** in the community. 2. In my understanding, the proposal of the IMF algorithm is the key/ most important contribution of this paper. 3. solid experiments and comparisons are studied, which verify the promising potential of DSBM. 4. Connections to flow matching, rectified flow, and stochastic interpolant are studied, which makes the authors easier to understand the benefit of this paper.

Weaknesses

1. Clearer clarification between IPF and IMF would be required to facilitate the understanding. IPF is based on forward K; IMF is based on backward KL, this is good in math, but more detailed information is lacking. for example, if we solely check proposition 2, we can derive almost the same loss based on forward KL (which leads to the IPF algorithm) except that the underlying measure is not based on $\Pi_{0, t}$. So **what is the real benefit/ advantages of simulating $\Pi_{0, t}$ to train the loss in proposition 2 would clarify our concerns** 2. Algorithm 1 is not clear enough. The authors may consider providing us more details in the learning of $\upsilon_{\phi^{\star}}$ in Alg 1. For example, in the training of alpha-blending (line 60-64 in https://github.com/tchambon/IADB/blob/main/iadb.py), given some samples of $x_1$ and $x_0$, very simple code would be used to train the models. However, when it comes to DSBM, how is it implemented? pseudocode would work if it clarifies our concerns. 3. My concern is that this algorithm may be quite expensive. Minor 1. missing one related work: Discussions on the accumulations of IPF errors based on general optimal-transport cost functions studied in [1] would be appreciated. [1] Provably Convergent Schrödinger Bridge with Applications to Probabilistic Time Series Imputation. ICML'23. 2. writing of Equation between line 637 and line 638 is weird, because given X_t=x_t and x_T, what is the randomness there and why do we need expectation w.r.t. $\Pi_{T|t}$? 3. derivation of eq after line 113 is not clear to me. why $E_{\Pi_{0, T}}[||X_0-X_T||]^2$? is there any reference for that? 4. first part of section 4 in the main paper, the need the alternating schemes is still not well motivated 5. what is IMF-b?

Questions

1. DDMs can be seen as the first iteration of DSM, which implies that we can use DDMs model weights to initialize the model weights of DSM. Does it mean that empirically we can also use weights from stochastic interpolant / the stochastic version of flow matching to train the first step of DSBM? 2. In line 993 in the appendix, the authors mentioned that "The advantage of this DSM loss is that it does **not rely on any divergence computation**". I believe this is quite an important part of scalability, but why it is only emphasized in the joint training instead of stating it in the main paper?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

4 excellent

Limitations

NA

Reviewer 24jx7/10 · confidence 2/52023-07-05

Summary

The submission suggests a numerically effective approach for solving the Schroedinger Bridge (SB) problem and illustrates its potential for generative modeling. The approach generalizes in some sense some recent flow matching methods. Furthermore, it provides a more efficient (no full trajectory caching) algorithm with less error accumulations that better remembers the prior reference measure for solving SB problems compared to previous works. Empirical results indicate for instance that it can yield better generative performance for MNIST-EMNIST transfer compared to previous work, or better reconstructions for downscaling geophysical fluid dynamics.

Strengths

Originality: The proposed methodology and algorithm is novel as far as I am aware. The work provides additional contributions compared to cited concurrent work [Peluchetti (2023)] that shares some similar ideas. Related literature and approaches are discussed in detail. Quality: Claims are adequately supported by proofs, and numerical experiments seem to support this. Clarity: The paper is well written. The presentation is quite rigorous. Significance: In my opinion, this work is of interest to the community as the method (i) generalizes recent flow matching approaches with non-degenerate noise and (ii) improves experimentally previous Diffusion Schrödinger Bridge approaches and (iii) provides some theoretical guarantees for solving the SP problem (under idealized assumptions).

Weaknesses

It was not clear to me how the computational cost of the suggested method compares to related works per iteration or diffusion step (e.g. in Figures 5 and 16).

Questions

In Figure 16, it appears that FID can be improved by using more than 100 steps. Does DBSM obtain better FID values compared to Bridge matching for more than 100 steps? It was not clear to me what to take away from some parts of the appendix, such as as E Discrete-Time Markovian Projection or the Brownian/OU bridge representations in B.2. Are the results in B.2 helpful for saying anything about how to specify the SDE in the forward process $X$ for the suggested approach?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Yes

Reviewer J5AJ5/10 · confidence 4/52023-07-06

Summary

The paper introduces Iterative Markovian Fitting (IMF) as a new method to compute Schrödinger Bridges (SBs), which are dynamic versions of entropy-regularized optimal transport. IMF alternates between projecting on the space of Markov processes and the reciprocal class, and it preserves the initial and terminal distributions. The paper also proposes Diffusion Schrödinger Bridge Matching (DSBM) as an algorithm to approximate SB solutions derived from IMF. DSBM overcomes issues of previous techniques and solves a simple regression problem at each iteration. The performance of DSBM is demonstrated in various transport tasks. The paper provides theoretical results, notations, and definitions related to the topics discussed.

Strengths

1, The paper introduces Iterative Markovian Fitting (IMF) as a new procedure for computing Schrödinger Bridges (SBs). This approach provides a fresh perspective and contributes to the field of optimal transport and generative modeling. 2, The paper establishes various theoretical results for IMF, demonstrating the validity and effectiveness of the proposed method. These results enhance the understanding of SBs and their applications in machine learning.

Weaknesses

1, While the paper mentions that IMF and DSBM address scalability issues, it lacks a detailed analysis of their scalability in terms of computational resources and dataset sizes. Understanding the scalability limitations of the proposed methods is crucial for their applicability to real-world, large-scale problems. 2, The paper assumes locally Lipschitz continuity of drift and limited settings for the proposed methods. These assumptions may limit the applicability of IMF and DSBM to certain scenarios. A discussion of the limitations imposed by these assumptions would provide a more comprehensive view of the proposed methods.

Questions

1, How does the Iterative Markovian Fitting (IMF) method differ from the traditional Iterative Proportional Fitting (IPF) approach in computing Schrödinger Bridges (SBs)? 2, Can you provide more insights into the theoretical results established for IMF in the paper? 3, How does Diffusion Schrödinger Bridge Matching (DSBM) overcome the time-discretization and "forgetting" issues of previous diffusion-based techniques?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

As stated in "Weakness"

Reviewer Axbu6/10 · confidence 3/52023-07-08

Summary

This paper proposes IMF, an iterative method for solving Schrodinger Bridge problems where the solution at each iteration preserves the correct marginal distributions at times 0 and T. This is in contrast to the existing method IPF which only satisfies this in the limit (and is very difficult to reach in practice).

Strengths

- An interesting marginal-preserving algorithm for solving Schrodinger bridge problems. The popular method IPF can have convergence problems in practice, while this seems to fix some of that.

Weaknesses

- Presentation. While I eventually understood the resulting IMF algorithm, I felt the presentation can be much more simplified. The algorithm is actually quite simple and closely related to priors works on SB and OT. While I understand the need to prove convergence and showcase novelty, the heavily front-loaded mathematical notation made reading on the first pass very difficult. I imagine the current paper may be off-putting for many readers who are not extremely familiar with related works. - Novelty. The proposed algorithm seems very similar to Rectified Flow but just with an additional noise parameter. - Limitation. The proposed algorithm only works for the standard formulation of Schrodinger bridge with quadratic costs. However, in this setting, it isn't clear to me why the extra stochasticity is useful, as opposed to solving deterministic maps (optimal transport problems).

Questions

- "Most notably, we adopt the SDE approach which is crucial for the validity of Proposition 5" - Can the authors expand on what breaks when there is no stochasticity? From what I understand, RF should also preserve marginal distributions and converge to the optimal transport solution, which is discussed more in [1]. The solutions are also unique similar to any Flow Matching-based approach. - All of the experiments in the main paper are currently only for transferring between two data distributions, but being able to find optimal transport maps is useful for generative modeling as well. I think it could really make the paper stronger by having experiments in generative modeling (Gaussian -> data) and showing some empirical metrics from e.g. [2] such as sampling cost and consistency, since these are important problems in generative modeling at the moment. [Edit: I found the experiments in I.5. I see that sampling efficiency is on par with RF but Table 6 seems pretty good! Although it does look like if we extrapolate to NFE > 100, the other methods are descending quite a bit faster in FID...] - Can the authors expand on the differences between the two proposed methods, DSBM-IPF and DSBM-IMF? If I understood correctly, the main difference between DSBM-IPF and DSBM-IMF is only in the first iteration, where the pairs (x0, xT) are sampled from (either an OU process Q or the independent joint distribution). And the real difference is that Q has a slight bias at time T because it may not have reached its stationary distribution. Practically, this difference seems very negligible, yet there seems to be differences in empirical results (e.g. Table 3 and Figure 5). Theoretically, since Q does not match the correct marginal at T, this is also problematic for the proof of convergence to SB right? [1] "Rectified Flow: A Marginal Preserving Approach to Optimal Transport" https://arxiv.org/abs/2209.14577 [2] "Multisample Flow Matching: Straightening Flows with Minibatch Couplings" https://arxiv.org/abs/2304.14772

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

2 fair

Limitations

The authors do mention limitations such as marginal improvement, computational issues, and difficulties when sigma --> 0. I think these are sufficient.

Area Chair DjHQ2023-08-11

Dear reviewers and authors, Thank you very much for your work on this submission and its evaluation. Now that the authors have responded to the reviews, I **strongly encourage** the reviewers to acknowledge the review, to look at other reviews and rebuttals for this submission, and to adjust their scores if needed. Thanks to those that have already done so. Authors have the possibility to reply if further questions are needed, until the 16th. Thank you very much to all, Area Chair

Reviewer 7Dnn2023-08-12

reply

I appreciate the authors' detailed reply and thanks for the suggested reference. I am satisfied with the response.

Authorsrebuttal2023-08-18

Thank the reviewer for your response and your acknowledgement of the rebuttal.

Reviewer 24jx2023-08-15

Response to authors

I thank the authors for their detailed response that have fully addressed my concerns or questions. I have also read the other reviews and intend to keep my accept score.

Authorsrebuttal2023-08-18

Thank the reviewer for your response and your acknowledgement of the rebuttal.

Authorsrebuttal2023-08-18

Thank the reviewer for your response and your acknowledgement of the rebuttal.

Program Chairsdecision2023-09-21

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC