Schrodinger Bridge Flow for Unpaired Data Translation

Thank you for your comments, we appreciate your acknowledgements of the paper’s merits.

The experiment on practical data is insufficient and the performance improvement is incremental. The paper only presents real-world I2I on Cat <-> Wild. Moreover, the performance of bidirectional online learning improves only incrementally. [...] For instance, comparing the actual wall-clock time or GPU memory usage would be useful.

Bidirectional online learning approximately halves the number of network parameters, thereby reducing GPU memory usage and wall-clock training time. In the case of AFHQ-64, pretraining bidirectional and 2-network models takes 4 and 9 hours, while finetuning takes 14 and 19 hours, respectively. We use the same number of gradient updates for both models. The computational cost difference comes from the use of a bidirectional versus a 2-network model. However, note that bidirectional models are incompatible with the original iterative DSBM and are suitable only for online finetuning. A fair comparison of DSBM and α-DSBM in terms of the convergence speed, is then possible with a 2-networks model. It is best illustrated by the Gaussian example where there is a single evaluation metric.

The current algorithm reflects 𝛼 implicitly, and so there is no analysis provided on 𝛼. If there were any intuition or insights into the role of 𝛼(step size) or if an algorithm could be developed to explicitly control 𝛼, it would offer more valuable insights.

If we perform this algorithm with full gradient descent algorithm, is it possible to explicitly control 𝛼? Can we develop exact algorithm that reflect 𝛼 explicitly? If so, can we obtain some experimental results that compares 𝛼?

We agree with the reviewer that parameter $\alpha$ merits further discussion. In our current manuscript, we state in (l.218) “In Algorithm 1, we specify $\alpha \in (0,1]$ as a stepsize parameter. In practice, we use Adam (Kingma and Ba, 2015) for optimization, thus the choice of $\alpha$ is implicit and adaptive throughout the training.” So, $\alpha$ is linked to the learning rate: in a Stochastic Gradient Update (SGD), $\alpha$ would correspond to the learning rate. To address this point comprehensively (and also answers reviewer Ti9G), we train a model with SGD and a sweep on the values of $\alpha$; see additional one-page rebuttal. We found out that setting explicitly the value of $\alpha$ in SGD can yield similar results as letting $\alpha$ adaptive using Adam. However, for too large values of $\alpha$ the algorithm diverges.

Thank you for your comments and thoughtful questions, we are glad you enjoyed the paper.

The paper lacks a study of image quality boost given the same computational budget.

It would be a more solid argument to evaluate the Gaussian setup with a full covariance matrix [6].

A bigger dataset such as Celeba may help to solve this issue.

We appreciate the reviewer's suggestions regarding our experimental setup. Concerning the image quality improvement given the same computational budget, we'd like to emphasize the inherent challenge in designing appropriate metrics for evaluating OT-based algorithms, as we are considering two distinct metrics: 1) Image quality (FID score): This roughly quantifies the alignment between the marginals $\pi_0$ and $\pi_1$ and the output distributions of the sampling process, 2) Alignment between samples from $\pi_0$ and $\pi_1$: This quantifies the minimization of the OT cost. Our objective is to minimize both the FID score and the alignment cost, necessitating the tracking of both these metrics. Nevertheless, we have conducted experiments comparing DSBM and $\alpha$-DSBM with equivalent computational costs, see the one-page supplementary.

Additionally, we have now evaluated our methodology in a complex Gaussian setting using full covariance matrices, similar to the approach in [1]; see one-page additional rebuttal. The findings from this full covariance matrix setting agree with our earlier conclusions from the scalar setting.

Finally, we have also evaluated our method on CelebA and refer to the one-page supplementary.

[1] Janati et al., “Entropic optimal transport between unbalanced gaussian measures has a closed form”, 2020.

[2] Chen et al., “Gradient flow in parameter space is equivalent to linear interpolation in output space”, 2024

The number of pre-training iterations, grad updates per IMF iteration, and IMF iterations are not specified. What is the training time and computational budget?

We have added the missing hyperparameters for the training of DSBM in the case of the AFHQ experiment to the revised manuscript. The number of pretraining iterations is fixed to 100 000. The number of gradient updates per IMF iteration is 500 and the number of IMF iterations is 40.

There is no study on cornerstone hyperparameter 𝛼. What is the value of 𝛼 used in the experiments? Is it equal to the learning rate of the diffusion model parameters?

We agree that we should have provided more discussion about the choice of parameter $\alpha$. Currently, we say (l.218) “In Algorithm 1, we specify $\alpha \in (0,1]$ as a stepsize parameter. In practice, we use Adam for optimization, thus the choice of $\alpha$ is implicit and adaptive throughout the training.” So indeed, $\alpha$ is linked to the learning rate and if one were to use SGD and not Adam, $\alpha$ would be exactly the learning rate. To complement this answer and also fully answer reviewer KqVg, we train a model with SGD and a sweep on the values of $\alpha$. The results are reported in the one-page rebuttal. We found that setting explicitly the value of $\alpha$ in SGD can yield similar results as letting $\alpha$ adaptive using Adam. However, for too large values of $\alpha$ the algorithm diverges.

The absence of code may cause trouble with reproducing the presented results.

We are working on releasing the notebooks to reproduce our experiments on Gaussian data. Unfortunately, due to confidentiality agreements and intellectual property concerns, we are unable to open source the full source code. We have included a code snippet of our main training loop implementing our online methodology ($\alpha$-DSBM) in the appendix of the revised paper.

Is it possible to extend the proposed approach to GSBM

We thank the reviewer for this remark. Indeed our online approach can be used to improve the computational efficiency of every method derived from DSBM. For example Generalized SBM [1] would also benefit from the techniques used in $\alpha$-DSBM. We have added a comment on this in the revised manuscript.

Can this method with small epsilon can be applied to the generation problem? [...] [2] similar to the work [1].

You are correct: $\alpha$-DSBM can be used with $\varepsilon$ small (or even $\varepsilon = 0$, even though in that case, our theoretical framework does not guarantee the convergence to the solution of the OT problem). Hence, our technique can potentially enhance the performance of the 2-Reflow of [1] in a generative model context. However, we emphasize that in [2], there is no finetuning stage and only a flow-matching pretraining (i.e. this text2img model only corresponds to the pretraining of our model). As the pretraining and finetuning of such models is long and expensive, we could not conduct such experiments during the rebuttal period, however we plan to include such text2img experiments in the final version of the paper.

[1] Liu et al., InstaFlow: One Step is Enough for High-Quality Diffusion-Based Text-to-Image Generation, 2024.

[2] Esser et al., Scaling Rectified Flow Transformers for High-Resolution Image Synthesis, 2024.

Is it possible to perform parametric updates for Sinkhorn flow in a similar to $\alpha$-DSBM way?

The applicability of a methodology similar to $\alpha$-DSBM to Sinkhorn-flow [1] is not straightforward. First, $\gamma$-Sinkhorn is an inherently static algorithm. In order to apply a methodology closer to $\alpha$-DSBM, we would need to develop a dynamic version of $\gamma$-Sinkhorn. This would correspond to extending DSB [2] to an online setting. It is unclear whether the objective of $\gamma$-Sinkhorn, see Lemma 1 in [1], can be modified to yield a tractable loss for an hypothetical $\gamma$-DSB version.

[1] Karimi et al., Sinkhorn Flow: A Continuous-Time Framework for Understanding and Generalizing the Sinkhorn Algorithm, 2023.

[2] De Bortoli et al., Diffusion Schrödinger Bridge with Applications to Score-Based Generative Modeling, 2021.

Thanks a lot for your comments which have improved the overall quality of the paper. Here are the additional metrics requested below (the hardware setup is identical to AFHQ):

• Base model training time: 16 hours
• Finetuning alpha-DSBM/Finetuning DSBM: 7 hours

The finetuning of DBSM corresponds to two IMF iterations (one forward and one backward).

Regarding the L2, LPIPS evals (lower is better):

• DSBM: 0.159 (L2) / 0.451 (LPIPS)
• alpha-DSBM: 0.05 (L2) / 0.376 (LPIPS)

We are also reporting the Inception Score (higher is better) evals for the base models, DSBM and alpha-DSBM:

• Base training: 2.29
• DSBM: 2.88
• alpha-DSBM: 3.13

Hence, alpha-DSBM improves over DSBM for the same amount of compute. We see that the main reason for the underperformance of DSBM is that for that amount of compute there is little deviation of the alignement score, i.e. the model is still close to the base model. We will incorporate these numbers in the final version of the paper. We hope that these numbers resolve the concerns of the reviewer.

Thank you for your positive comments and feedback.

While the experimental section of the paper thoroughly analyzes and compares DSBM and its different flavors, including the proposed method, a comparison with other competing methods could further support-DSBM through empirical evidence.

How does 𝛼-DSBM perform compared to bridge matching with mini-batch OT sampling empirically?

We thank the reviewer for their suggestion to evaluate our approach against competing methods. While we agree that comparisons with [2,3,4] would be valuable, we have chosen, in this rebuttal, to focus on [1] which was also mentioned by reviewer ovZP and is most similar to our approach. However, we will mention [2,3,4] as possible alternatives with similar goals in our extended related work section. Our comparison with OT-bridge matching in the Gaussian setting is reported in the attached one-page rebuttal. We will include a comparison to [2,4] in the final version.

[1] Tong et al., Improving and generalizing flow-based generative models with minibatch optimal transport, TMLR, 2024.

[2] Eyring et al., Unbalancedness in Neural Monge Maps Improves Unpaired Domain Translation, ICLR, 2024.

[3] Torbunov et al., UVCGAN v2: An Improved Cycle-Consistent GAN for Unpaired Image-to-Image Translation, arXiv:2303.16280, 2023

[4] Kim et al., Unpaired Image-to-Image Translation via Neural Schrödinger Bridge, ICLR, 2024.

Does it make sense to combine 𝛼-DSBM with mini-batch OT sampling for the initial pretraining? How would this impact 𝛼-DSBM?

This is indeed doable and is a great suggestion. This would require changing our pretraining strategy to include the mini-batch OT sampling during the training. Our guess is that the convergence of $\alpha$-DSBM (and even DSBM) would be faster. To support this result, we refer to the results of [1,2] which show that the convergence of the Sinkhorn algorithm depends on the Kullback-Leibler divergence between the original guess (here the output of the (E)OT-FM methodology) and the target Schr\”odinger Bridge.

[1] Bernton, Schrödinger Bridge Samplers, arXiv:1912.13170, 2019.

[2] L\’eger, A gradient descent perspective on Sinkhorn, Applied Mathematics & Optimization, 2021.

Thanks for the answers and added experiments. It would be very interesting to see comparisons to [1,2,4] also in the image translation settings. I think this setting would provide a more meaningful empirical comparison. Also thanks for pointing to [1,2] for the convergence of DSBM with mini-batch OT.

Thank you for your comments, we appreciate your acknowledgements of the paper’s merits.

It is necessary to state these assumptions in the main body of the paper and provide references to justify their rationality.

Due to space limitations, we have detailed most of the technical assumptions for our theorems in the supplementary material rather than in the main paper. However, if our submission is accepted, we will be granted an additional page, which will allow us to include the assumptions as detailed below.

Theorem 3.1

The assumptions for the result to hold are stated in Lemma D.2. We recall them here. Let $\pi_0$ and $\pi_1$ be the distributions at time 0 and 1. We require that $\pi_0$, $\pi_1$ have distributions w.r.t. Lebesgue measure, with finite entropy and finite 2nd moments. This finite 2nd moment assumption is quite weak. While the Lebesgue density hypothesis is commonly considered in the diffusion model literature, see e.g. [1,2].

Proposition 3.2

This proposition is detailed in Proposition D.1.The sole assumption in this case is that $\mathbb{P}_{v^n}$ is defined for all $n$. This means that the SDE with drift $v^n$ admits a weak solution. The existence of weak solutions of SDE is a well-known topic in the literature of SDE; see e.g. [3]. In practice this condition is satisfied since $v^n$ is replaced by a NN approximation, which satisfies the conditions for the SDE to admit a solution.

Proposition 4.1

The proof is similar to the one of Proposition 3.2 and therefore holds under the assumption that $v^{n, \rightarrow}$ and $v^{n, \leftarrow}$ define solutions of SDEs.

[1] Lee et al., Convergence of score-based generative modeling for general data distributions, 2023

[2] Conforti et al., Score diffusion models without early stopping: finite Fisher information is all you need, 2023

[3] Stroock et al., Multidimensional Diffusion Processes, 2006.

The experimental results on some datasets are not particularly impressive, even though the authors claim that their method approximates OT maps more accurately than existing methods

We want to emphasize that we do not claim our approach yields better results than DSBM (or any other method that recovers the EOT). There is no expectation of better results from our method compared to DSBM, as both converge to the same solution. Currently, we state “$\alpha$-DSBM is easier to implement than existing SB methodologies while exhibiting similar performance” (l.313).

The main advantages of our methods are computational: 1) significantly simpler implementation, requiring only one fine-tuning stage and eliminating the need to alternate between two optimisation problems 2) utilization of only one NN (with the bidirectional implementation) 3) faster convergence, as shown in Fig. 2 (and in the Gaussian study in the one-page rebuttal). We hope this explanation clarifies the key computational contributions of our paper. We will revise our manuscript to further underline this point if it remains unclear.

It is confusing that the authors explicitly mentioned Rectified Flow (RF), Flow Matching (FM), and OT-FM in the literature review, but did not include comparisons with these methods.

RF corresponds to DSBM for $\varepsilon = 0$. Our $\alpha$-DSBM algorithm thus provides an alternative online implementation of RF for $\varepsilon = 0$ and results for this online implementation are provided in Fig 3 (first column of right figure). For unpaired data translation, RF is not competitive and adding noise improves performance. However, for generative modeling, RF outperforms DSBM and $\alpha$-DSBM using $\epsilon>0$. We conjecture that when one marginal distribution is Gaussian, adding further noise to the interpolant hurts the generative capabilities. In addition, the fixed points of RF exhibit straight paths but are not necessarily solutions to the OT problem, see [2] for a counterexample.

FM [3] can be seen as the first iteration of DSBM with $\varepsilon = 0$. It consistently produces inferior results compared to RF, this is why we did not include it. OT-Flow Matching [4] is closer to our setting and, if the batch size has infinite size, then we recover the OT solution. In the one-page rebuttal we show that $\alpha$-DSBM outperforms this method in the Gaussian case.

Finally, we also highlight that DDBM [5] do not solve the (E)OT problem but requires paired data to train the model and therefore is not comparable to our approach.

[1] Liu et al., Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow, 2023.

[2] Liu, Rectified Flow: A Marginal Preserving Approach to Optimal Transport, 2022.

[3] Lipman et al., Flow Matching for Generative Modeling, 2023.

[4] Tong et al., Improving and generalizing flow-based generative models with minibatch optimal transport, 2024.

[5] Zhou et al., Denoising Diffusion Bridge Models, 2024

The paper could benefit from a more comprehensive discussion of related work

In the revised manuscript, we have expanded the related work section in the appendix to offer a more comprehensive overview. Notably, we discuss a taxonomy graph of diffusion methods and their connection to OT, see the one-page rebuttal.

Could the authors provide more detailed quantitative data, such as training time comparisons with related methods?

Bidirectional online learning approximately halves the number of network parameters, thereby reducing GPU memory usage and wall-clock training time. For AFHQ-64, pretraining bidirectional and 2-network models takes 4 and 9 hours, while fine-tuning takes 14 and 19 hours, respectively. We use the same number of gradient updates for both models. The computational cost difference comes from the use of a bidirectional versus a 2-network model. Note that bidirectional models are incompatible with the original iterative DSBM and are suitable only for online finetuning. We have included these cost and complexity considerations in the appendix.

Thank you for addressing many of the initial concerns in your rebuttal. However, I believe that if the main contribution is claimed to be a simpler implementation with only one neural network and faster convergence, it is crucial to provide more quantitative data to support these claims.

I recommend:

Including additional experiments that compare your model with other single neural network setups, highlighting computational efficiency and performance (In the rebuttal, the author provides some additional comparative experiments). Providing details on the reduction in neural network parameters to quantitatively demonstrate the simplicity of your approach. Clarifying the convergence rates in comparison with baseline methods, particularly considering the fine-tuning duration mentioned in your rebuttal experiments (I noticed that in the rebuttal experiments, the authors' fine-tuning time was much longer than the training time. Does this contradict the authors' statement?). If these points are addressed in the revised version and the theoretical assumptions are further clarified, I would be inclined to increase my score to a 5-6. Thank you for your efforts to improve the manuscript.