Diffusion & Adversarial Schrödinger Bridges via Iterative Proportional Markovian Fitting

Theoretical part. From the theoretical point of view, we show that the ideas behind the 1D proof can be generalized to multivariate Gaussians as well. Firstly, we introduce a new optimality matrix $A$ that, for a given Gaussian joint distribution, forms a cost function for which this distribution is a EOT solution. During IPF steps, this optimality matrix remains the same, while marginals become exponentially closer to desired values (Lemmas F.1, F.2). We believe that IMF step changes the optimality matrix so that it exponentially becomes closer to the matrix of the solution of $\varepsilon$-EOT as we prove in $1D$ case. We verify it in the experiments (Section 4.1, Figure 2), where we run IPMF and track the norm between current optimality matrix and $\varepsilon$-EOT matrix.

For more general distributions, we believe that the logic can be saved too. One can use forward KL divergence between the ground truth marginals and current one, since it drops after IPF step and does not change during IMF. However, one needs only to propose a new optimality parameter instead of optimality matrix which will be kept during IPF step and will be improved after IMF step.

(3) Finally, I find the experiments to be quite poor. ... For instance, FID scores ... it seems that changing the coupling yields worse results regarding the quality of the obtained coupling. I think, the authors should investigate other couplings ...

We provided the answer to this in the previous answer.

(4) In the introduction, the authors mention that the reverse Kullback-Leibler between the iterates of the IMF and the target EOT coupling decreases at each step of the IMF ...

We meant the reverse Kullback-Leibler between the iterates of the IMF $q^k$ (or $T^k$ for the continuous-time version) and the ground-truth Schrödinger Bridge process $q*$ (or $T*$ for the continuous-time version). Both corresponding proofs of convergence for discrete-time [2, Theorem 3.6] and continuous-time [3, Theorem 8] versions are based on this fact. However, the mentioned fact for the joint distribution on $t=0,t=1$ of IMF iterates ($q^k(x_0, x_1)$) and the target EOT coupling ($q^*(x_0, x_1)$) can be derived from the fact for the processes. We provide the derivation below.

The reverse Kullback-Leibler between the joint distribution on $t=0,t=1$ of IMF iterates ($q^k(x_0, x_1)$) and the target EOT coupling ($q^*(x_0, x_1)$) does not increase at each step of the IMF. Indeed, for any $q^{n}(x_0, x_{\text{in}}, x_1)$ in the discrete IMF sequence, we can write down the following:

$$\text{KL}(q^{n}(x_0, x_{\text{in}}, x_1)|| q^*(x_0, x_{\text{in}}, x_1)) =$$ $$\text{KL}(q^{n}(x_0, x_1)|| q^*(x_0, x_1)) + \int \text{KL}({q^{n}(x_{\text{in}}|x_0, x_1)}||q^*(x_{\text{in}}|x_0, x_1))q(x_0,x_1)dx_0 dx_1,$$ $$\text{KL}(q^{n}(x_0, x_1)||q^*(x_0, x_1)) =$$ $$\text{KL}(q^{n}(x_0, x_{\text{in}}, x_1)|| q^*(x_0, x_{\text{in}}, x_1)) - \int \text{KL}({q^{n}(x_{\text{in}}|x_0, x_1)}||q^*(x_{\text{in}}|x_0, x_1))q(x_0,x_1)dx_0 dx_1,$$

hence

$$\text{KL}(q^{n}(x_0, x_1)||q^*(x_0, x_1)) \leq \text{KL}(q^{n}(x_0, x_{\text{in}}, x_1)|| q^*(x_0, x_{\text{in}}, x_1)).$$

Since $\text{KL}(q^{n}(x_0, x_{\text{in}}, x_1)|| q^*(x_0, x_{\text{in}}, x_1)) \rightarrow 0$ it follows, that $\text{KL}(q^{n}(x_0, x_1)||q^{*}(x_0, x_1)) \rightarrow 0$.

The continuous case can be proved exactly by using the disintegration theorem, as in Eq. 15 in our paper.

(5) In Equation (9) and Equation (10) the forward and backward are done successively. Do you think it is possible to learn the backward and forward processes jointly? ...

There is an online version of the DSBM algorithm [1], which jointly learns the process given by equations (9) and (10) at the same time. Since this online version shows promising results, it would be interesting to generalize IPMF additionally to support both successive and online (joint training) training. We think it is a good avenue for further research.

(6) It would be good to have some kind of diagram to illustrate the differences between IMF, IPF and IPMF as well as some explicit descriptions on when they are equal ...

Thank you for your suggestion. To address this, we have added Figure 1 to the paper, which visually illustrates the differences between Iterative Proportional Fitting (IPF), Iterative Markovian Fitting (IMF), and Iterative Proportional Markovian Fitting (IPMF). The diagram highlights the projection steps for each method and their corresponding properties, such as having the correct marginals, being reciprocal (R), or Markovian (M).

We also added additional clarification at the end of Section 3.1, where we defined the IPMF procedure, for which couplings IPMF becomes IPF and for which couplings IPMF becomes IMF.

Dear Reviewer KCc1, thank you for your comments. Here are the answers to your questions and comments.

(1) I think the paper has limited novelty. ... the IPMF rewriting is not needed to show that the IMF procedure reverts to the IPF one when the coupling is initialized as the IPF one which is one of the main points of the paper.

We agree that the connection between Iterative Markov Fitting (IMF) and Iterative Proportional Fitting (IPF) is a well-established result. However, we would like to clarify that proving this fact is not the primary focus of our work. Instead, we started our discovery by noting that the bidirectional heuristic used in both DSBM and ASBM solvers performs better in practice, while the one-directional training often breaks down due to accumulated errors (Appendix I [1]). We argue that the observed stability arises from the IPF projections "hidden" in this bidirectional procedure. Based on this, we propose a separate IPMF procedure, which generalizes both IPF and IMF procedures. We also make a conjecture that IPMF converges to the Schrödinger Bridge solution for any starting coupling, unlike specific couplings like IPF or IMF. We support our conjecture by providing a rigorous proof of convergence for the one-dimensional Gaussian case (Section 3.2) and demonstrating through extensive experiments that convergence is indeed achieved regardless of the choice of initial coupling (Section 4).

(2) If it should be read as a theory paper then the result is quite weak ... If the paper is to be seen as a methodology paper then I think that there should be a clear motivation to choose couplings that are different from the ones of the IMF.

Practical part. From the practical point of view, considering the generalized IPMF procedure opens a door toward using arbitrary starting processes rather than ones that can be used in IPF or IMF procedures. While, in theory, all such starting processes are expected to converge to the same solution, the dynamics of this convergence will differ. Additionally, approximation errors introduced by using neural networks, which differ for different starting processes, will influence the results obtained in practice. We can use this difference introduced by the various couplings to trade-off between focusing on similarity or quality of modeling of target marginal. For instance, in the original version, we considered the independent $p_0 \rightarrow p_0$ starting process, which shows better similarity than other couplings (Figure 9b) but worse quality of the target samples (Figure 6, 9a).

To further show the practical impact and to answer the related question about what other couplings could be used, we found a new one and conducted experiments with it. Specifically, we found a way to further improve this independent $p_0 \rightarrow p_0$ coupling by improving the quality of target samples. We provide additional experiments on the Celeba dataset with new coupling obtained by using the training-free translation method SDEdit [4] with the pretrained diffusion model (DDPM [5]) for the target distribution. We chose it due to the recent widespread of different pretreated diffusion models used for a wide set of tasks. The coupling provided by SDEdit cannot be used in IMF (since it matches only one of the marginals $p_0$ or $p_1$) and IPF (since it is generated from the Markovian, but not the reciprocal process used in SDEdit). This coupling allows us to obtain a similar quality of generated images as other methods when measured with CMMD (Figure 9a) and slightly worse DSBM-IMF when measured with FID (Figure 6). At the same time, this SDEdit coupling shows a better similarity between the input and generated images than all other couplings except independent $p_0 \rightarrow p_0$. We describe the obtained results and the coupling details in Section 4.4 and in Appendix C.2.

Thus, we provide an example of how one can design new couplings to trade between the similarity and the quality of target samples or even improve one of the properties without harming the other. We believe, that it is possible to find even better couplings, for instance, by usage other diffusion models like Stable Diffusion.

Dear Reviewer 5NKn, thank you for your comments. Here are the answers to your questions and comments.

(1) Although this paper shows an interesting theoretical results, it's only for 1d Gaussian cases. I feel like the contribution is relatively weak. If authors can say something more, i.e, for multi-variate Gaussian, this paper will become stronger.

We modified the theoretical analysis section and added new conjecture about exponential convergence in higher dimensions, which we back up with new experiments (Section 4.1). It completely inherits logic of $1D$ proof. Firstly, we introduce a new optimality matrix $A$ that, for a given gaussian joint distribution, forms a cost function for which this distribution is a EOT solution. During IPF steps, this optimality matrix remains the same, while marginals become exponentially closer to desired values (Lemmas F.1, F.2). We believe that IMF step changes the optimality matrix so that it exponentially becomes closer to the matrix of the solution of $\varepsilon$-EOT as we prove in $1D$ case. We verify it in the experiments in Figure 2, where we run IPMF and track the norm between current optimality matrix and $\varepsilon$-EOT matrix.

(2) What's the technical difficulties for generalizing the proof to multi-variate Gaussians or other situations?

We completely generalized analysis of the IPF step for multivariate Gaussians (Lemmas F.1, F.2). However, the most difficult part for multivariate Gaussians is the IMF step analysis. We need to show that optimality matrix becomes closer to ground truth solution in some norm. This step is rather technical, i.e., due to very complex matrix transaction equations, it is hard to prove even in a 1D case (but we proved a 1D case in the original version of the paper).

For a more general case, we believe that the logic can be saved. One can use forward KL divergence between the ground truth marginals and current one, since it drops after IPF step and does not change during IMF. However, one needs to propose a new optimality parameter instead of optimality matrix which will be kept during IPF step and will be improved after IMF step. Usual for IMF reverse KL divergence between EOT/Schrödinger Bridge solution and current process is not appropriate, since its change during IPF step is unclear.

(3) Can you provide more intuition on why IMPF can usually work for any specific starting process?

As outlined in Section 3.1, the main idea is that the bidirectional IMF procedure combines four types of projections: two IMF-specific projections (reciprocal and Markov) and two IPF-related projections ($p_0$ and $p_1$). Specifically, the IPF projections adjust the marginals $q_0(x_0)$ and $q_1(x_1)$ towards $p_0(x_0)$ and $p_1(x_1)$ but do not influence "optimality" of the process. In turn, the Markov and reciprocal projections adjust the process to make it more optimal (to be Markovian and reciprocal simultaneously) but do not change marginals. Thus, regardless of choosing the initial coupling (it just acts as a starting point) the iterative projections lead to convergence to the Schrödinger Bridge.

Concluding remarks. We would be grateful if you could let us know if the explanations we gave have been satisfactory in addressing your concerns and questions about our work. If so, we kindly ask that you consider increasing your rating. We are also open to discussing any other questions you may have.

Dear Reviewer tu1b, thank you for your comments. Here are the answers to your questions and comments.

(1) There are some inaccuracies in the text (Appenix, markovian, Matkovian, etc).

Thank you for pointing out this. We fixed these and several similar typos in the new revision.

(2) Perhaps due to the lack of familiarity of this reviewer with the SOTA on SB: the main contribution of the paper is the identification that IPF & IMF can be understood as the same procedure. However, what are the practical implications of this? Is this a purely conceptual statement, or does this fact help practical applications of SB? Perhaps the authors can be more specific about the practical advantages of their proposal.

From the practical point of view, considering the generalized IPMF procedure opens a door toward using arbitrary starting processes rather than ones that can be used in IPF or IMF procedures. While, in theory, all such starting processes are expected to converge to the same solution, the dynamics of this convergence will differ. Additionally, approximation errors introduced by using neural networks, which differ for different starting processes, will influence the results obtained in practice. We can use this difference introduced by the various couplings to trade-off between focusing on similarity or quality of modeling of target marginal. For instance, in the original version, we considered the independent $p_0 \rightarrow p_0$ starting process, which shows better similarity than other couplings (Figure 9b) but worse quality of the target samples (Figure 6, 9a).

To further show the practical impact, we found a new one and conducted experiments with it. Specifically, we found a way to further improve this independent $p_0 \rightarrow p_0$ coupling by improving the quality of target samples. We provide additional experiments on the Celeba dataset with new coupling obtained by using the training-free translation method SDEdit [1] with the pretrained diffusion model (DDPM [2]) for the target distribution. We chose it due to the recent widespread of different pretreated diffusion models used for a wide set of tasks. The coupling provided by SDEdit cannot be used in IMF (since it matches only one of the marginals $p_0$ or $p_1$) and IPF (since it is generated from the Markovian, but not the reciprocal process used in SDEdit). This coupling allows us to obtain a similar quality of generated images as other methods when measured with CMMD (Figure 9a) and slightly worse DSBM-IMF when measured with FID (Figure 6). At the same time, this SDEdit coupling shows a better similarity between the input and generated images than all other couplings except independent $p_0 \rightarrow p_0$. We describe the obtained results and the coupling details in Section 4.4 and in Appendix C.2.

Thus, we provide an example of how one can design new couplings to trade between the similarity and the quality of target samples or even improve one of the properties without harming the other. We believe, that it is possible to find even better couplings.

Concluding remarks. We would be grateful if you could let us know if the explanations we gave have been satisfactory in addressing your concerns and questions about our work. If so, we kindly ask that you consider increasing your rating. We are also open to discussing any other questions you may have.

References:

[1] Meng C. et al. SDEdit: Guided Image Synthesis and Editing with Stochastic Differential Equations //International Conference on Learning Representations.

[2] Ho J., Jain A., Abbeel P. Denoising diffusion probabilistic models //Advances in neural information processing systems. – 2020. – Т. 33. – С. 6840-6851.

Dear Reviewer BjZj, thank you for your comments. Here are the answers to your questions and comments.

(1) Section 4, presenting the experiments, seems to be a bit disconnected from the previous sections. ...

Thank you for this comment. This impression likely arises because, in the first experimental setup (“High Dimensional Gaussians”), the discussion emphasizes differences in convergence speed. We agree that we should better highlight that the main observation is the algorithm’s convergence to the ground-truth solution across all couplings. We have adjusted the text in this section (Section 4) to stress this point more clearly.

We also would like to note that we follow the initial aim stated at the beginning of Section 4 and consistently conclude the discussions in other experiments with observations of convergence, as seen in the following lines:

• “In all the cases, we observe convergence in the target distribution.” (line 418)

• “… yield quite similar results.” (line 430)

• “We observe that both DSBM and ASBM algorithms starting from both IMF and Inverted 7 starting process fit the target distribution of colored MNIST digits of class ‘2’ and preserve the color of the input image during translation.” (lines 458–460)

• “We see from the qualitative results in Figure~6 that all trained models: 1) converge to the target distribution, 2) keep alignment between the features of the input images and generated images (e.g., the color of hair, background, etc.).” (lines 484–485)

If you are referring to the ASBM and DSBM algorithms (i.e., discrete-time and continuous-time approaches to solving Schrödinger Bridge) by “different procedures,” we would like to clarify that we included both to demonstrate that the IPMF generalization is applicable to both frameworks.

To resolve this weakness further, we additionally added in the beginning of the experimental section (Section 4) that our goal is to achieve the same or similar results for all used starting coupling and for both discrete-time (ASBM) and continuous-time (DSBM) solvers.

(2) [Minor] Typos

Thank you for pointing out this. We fixed these and several similar typos in the new revision.

Dear reviewers,

We thank you for your review and appreciate your time reviewing our paper.

The end of the discussion period is close. We would be grateful if we could hear your feedback regarding our answers to the reviews. We are happy to address any remaining points during the remaining discussion period.

Thanks in advance,

Paper authors