Schrödinger Bridge Flow for Unpaired Data Translation

Summary

This paper proposes $\alpha$-IMF, which is an incremental version of IMF method in DSBM paper. Specifically, this work demonstrate the convergence properties of $\alpha$-IMF and implement it through an online learning method. Here, $\alpha$ is implicitly reflected. Moreover, the functional flow interpretation provides a basis for online learning, allowing for online learning as opposed to the iterative approach of DSBM. From an engineering perspective, the authors suggest parameterizing forward and backward control simultaneously by adding a forward/backward indicator (0 or 1) as an input variable to the neural network. The method shows feasible results on toy (Gaussian) data and image-to-image translation tasks.

Strengths

- This work is theoretically fruitful. The introduction of $\alpha$-IMF leaded to a more expansive theoretical framework. This flow-based perspective enables online learning. - Reduced the sensitivity to hyperparameters such as the iteration number, the number of phase in DSBM. - Expands the interpretation of the Schrödinger bridge by demonstrating connections with reinforcement learning (RL) and Sinkhorn flow.

Weaknesses

- 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. Totally, it needs more evaluation on other kinds of dataset, e.g. Male <-> Female / Handbags <-> Shoes. - It would be beneficial to quantify the advantages of bidirectional learning or online learning over DSBM. For instance, comparing the actual wall-clock time or GPU memory usage would be useful. - The current algorithm reflects $\alpha$ implicitly, and so there is no analysis provided on $\alpha$. If there were any intuition or insights into the role of $\alpha$ (step size) or if an algorithm could be developed to explicitly control $\alpha$, it would offer more valuable insights.

Questions

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

Rating

6

Confidence

3

Soundness

3

Presentation

2

Contribution

3

Limitations

The authors adequately addressed the limitations.

Summary

The authors consider developing a new algorithm for a Schrödinger Bridge problem of translating between two probability distributions. Motivated by the fact that current Schrödinger Bridge (SB) approaches either use mini-batch optimal transport techniques or require training diffusion at every iteration, the authors propose discretizing the Schrödinger Bridge Flow (SB Flow) of path measures that converge to SB. They call this discretization $\alpha$-Iterative Markovian Fitting ($\alpha$-IMF) procedure. Using the fact that Markovian path measures characterizing the considered flow can be parametrized by diffusion vector field $v$, authors propose a non-parametric scheme for updating $v_n$ corresponding to $n$ discretization step of SB Flow and prove its correctness. To implement the procedure in practice, authors prove a parametric counterpart to the non-parametric update scheme that allows for optimization over a vector field $v_{\theta}$ parametrized as a neural network called $\alpha$ Diffusion Schrödinger Bridge Matching ($\alpha$-DSBM). Furthermore, the authors computationally validate the proposed algorithm on both vector and image state spaces. The authors compare their $\alpha$-DSBM algorithm to other DSBM versions [3] in the Gaussian setting. Image state space experiments were held on image-to-image translation tasks, i.e., MNIST $\leftrightarrow$ EMNIST, Wild $\leftrightarrow$ Cat on AFHQ 64/256, with quantitative evaluation and comparison with the original DSBM algorithm [3].

Strengths

- The overall theory is novel and interesting. The authors prove that the proposed $\alpha$-IMF procedure with non-parametric updates converges to the solution of SB. - The parametric counterpart to $\alpha$-IMF ($\alpha$-DSBM), which can be used with neural networks, is proposed and empirically verified. - The possibility of online $\alpha$-DSBM finetuning seems beneficial in terms of computational resources compared to regular DSBM iterative training. - A sufficient comparison of the proposed $\alpha$-DSBM with the previously known DSBM algorithm [3] is presented for a wide set of hyperparameter $\epsilon$. - The Appendix presents an extensive additional theoretical review of the proposed concept and its connections to existing research. - The study on the parameterization of diffusion in both directions by one neural network with varied conditional vectors is proposed and tested.

Weaknesses

- The proposed $\alpha$-DSBM algorithm [3] is, in some sense, an advancement of the previously known DSBM algorithm [3]. It is expected to be more efficient because it removes the necessity to learn a distinct model at each iteration. However, the paper lacks a study of image quality boost given the same computational budget. Thus, it is not clear whether the authors achieved the goal of developing a more efficient algorithm. - The convergence of parametric $\alpha$-DSBM is not proved. While it may be a hard theoretical problem to prove, at least a more extensive empirical evaluation of the setup with the known ground truth solution would be beneficial. The presented setup with a scalar covariance matrix is too simple. It would be a more solid argument to evaluate the Gaussian setup with a full covariance matrix [6] or use a mixture benchmark [7] and compare it with other neural network methods. - The FID metric is similar for DSBM and $\alpha$-DSBM for the AFHQ experiment in Table 1 considering measured standard deviation. Thus again highlighting the lack of study with the same computational budget to justify that $\alpha$-DSBM is more efficient. The chosen AFHQ dataset may not be the best choice considering the small test set, which may introduce bias in measuring the FID metric. In turn, the visual comparison is hard since most images produced by DSBM and $\alpha$-DSBM and given in Figure 19 are similar, and it is hard to deduce which algorithm is better. A bigger dataset such as Celeba [5] may help to solve this issue. At least a bigger dataset would give a good estimate of FID metrics. - It seems that some hyperparameters for DSBM AFHQ experiments are missing. 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? - There is no study on *cornerstone* hyperparameter $\alpha$. What is the value of $\alpha$ used in the experiments? Is it equal to the learning rate of the diffusion model parameters? - The absence of code may cause trouble with reproducing the presented results.

Questions

- Since the AFHQ dataset is quite small, experiments on unpaired male-to-female translation on the Celeba [5] dataset could provide a more meaningful comparison (as it was in DSBM[3]). - How do DSBM and $\alpha$-DSBM compare in terms of computational budget? Could authors provide a quantitative comparison? - How does variation of $\alpha$ affect speed of convergence and deviation from the SB solution? - Can authors provide experimental results of $\alpha$-DSBM on more complex Gaussian distributions or multimodal low dimensional data, i.e., gaussian mixtures? - Is it possible to extend the proposed approach to GSBM [4], where, as far as my understanding goes, authors change the reciprocal projection step compared to DSBM [3]? - Can this method with small $\epsilon$ can be applied to the generation problem? This could drastically reduce the computational budget of such and be applied in modern text-to-image generative models [2] similar to the work [1]. It would be great to see such experimental results. - Is it possible to perform *parametric* updates for Sinkhorn flow in a similar to $\alpha$-DSBM way? Overall, the paper looks very promising. I will consider increasing my score if the authors address some of the listed weaknesses and questions, especially related to the same computational budget comparison, and evaluate their method on one of the listed more complex setups with the known ground truth distributions. [1] Liu, Xingchao, et al. "Instaflow: One step is enough for high-quality diffusion-based text-to-image generation." The Twelfth International Conference on Learning Representations. 2023. [2] Esser, Patrick, et al. "Scaling rectified flow transformers for high-resolution image synthesis." Forty-first International Conference on Machine Learning. 2024. [3] Shi, Yuyang, et al. "Diffusion Schrödinger bridge matching." Advances in Neural Information Processing Systems 36 (2024). [4] Liu, Guan-Horng, et al. "Generalized Schrödinger Bridge Matching." The Twelfth International Conference on Learning Representations. [5] Liu Z. et al. Deep learning face attributes in the wild //Proceedings of the IEEE international conference on computer vision. – 2015. – С. 3730-3738. [6] Hicham Janati, Boris Muzellec, Gabriel Peyré, and Marco Cuturi. Entropic optimal transport between unbalanced gaussian measures has a closed form. Advances in neural information processing systems, 33:10468–10479, 2020. [7] Gushchin N. et al. Building the bridge of schrödinger: A continuous entropic optimal transport benchmark //Advances in Neural Information Processing Systems. – 2023. – Т. 36. – С. 18932-18963.

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors adequately addressed the limitations and potential negative societal impact of their work.

Summary

This work introduces $\alpha$-DSBM, a new way of training DSBM-like models, which does not require a Markovian projection at each step and eliminates the need to train multiple models. The main advantage over previous DSBM-based approaches is that $\alpha$-DSBM only needs to train a single model with a single model, thus exhibiting a more stable training procedure. The authors thoroughly contextualize their work within related work and provide detailed theoretical derivations to motivate their approach. Empirically $\alpha$-DSBM is validated by comparing to existing DSBM-based approaches across different unpaired image translation tasks as well as some toy data examples.

Strengths

- The new $\alpha$-DSBM is introduced based on thorough theoretical motivation and derivations. - The authors manage to present and contextualize their method very clearly within the scope of related work across the different bridge matching approaches (Appendix E) and highlight connections accordingly (e.g. connection to the reflow procedure). - In general, the paper is well-written and, given the complexity of the topic, quite clear to follow. - Empirical validation includes important ablations giving further insights into the hyperparameter choice of e.g. $\varepsilon$.

Weaknesses

- 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 $\alpha$-DSBM through empirical evidence. Specifically, a comparison to bridge/flow matching with mini-batch OT similar to [1, 2], a highly optimized CycleGAN as in [3], and adversarial-based OT methods like [4] could be good candidates for further comparisons. The authors also mention these as competing methods. [1] Alexander Tong and Kilian Fatras and Nikolay Malkin and Guillaume Huguet and Yanlei Zhang and Jarrid Rector-Brooks and Guy Wolf and Yoshua Bengio. "# Improving and generalizing flow-based generative models with minibatch optimal transport". IN TMLR 2024. [2] Luca Eyring and Dominik Klein and Théo Uscidda and Giovanni Palla and Niki Kilbertus and Zeynep Akata and Fabian Theis. "Unbalancedness in Neural Monge Maps Improves Unpaired Domain Translation". In ICLR 2024. [3] Dmitrii Torbunov and Yi Huang and Huan-Hsin Tseng and Haiwang Yu and Jin Huang and Shinjae Yoo and Meifeng Lin and Brett Viren and Yihui Ren. "UVCGAN v2: An Improved Cycle-Consistent GAN for Unpaired Image-to-Image Translation". In Arxiv 2023. [4] Beomsu Kim and Gihyun Kwon and Kwanyoung Kim and Jong Chul Ye. "Unpaired Image-to-Image Translation via Neural Schrödinger Bridge". In ICLR 2024.

Questions

- How does $\alpha$-DSBM perform compared to bridge matching with mini-batch OT sampling empirically? While it is clear that these methods introduce significant errors because of the mini-batch approximation, it is unclear how this affects results empirically compared to $\alpha$-DSBM. I think this is the most important additional competing work (apart from DSBM), as these methods share their overall goal. - Does it make sense to combine $\alpha$-DSBM with mini-batch OT sampling for the initial pretraining? How would this impact $\alpha$-DSBM?

Rating

7

Confidence

4

Soundness

3

Presentation

4

Contribution

3

Limitations

- The main limitation of $\alpha$-DSBM is that new data needs to be generated with the current model during its fine-tuning, making the procedure not simulation-free and, thus, more expensive. This is sufficiently addressed in the paper.

Summary

The paper proposes a novel algorithm for mass transport problems, aiming to compute maps that transport one distribution to another. This paper introduces the Schrödinger Bridge Flow algorithm, a dynamic entropy-regularized version of OT, eliminating the need to train multiple DDM-type models. The algorithm discretizes a flow of path measures, ensuring the Schrödinger Bridge as the only stationary point. The paper demonstrates the algorithm's effectiveness on various unpaired data translation tasks, showcasing its potential to solve Entropic Optimal Transport problems with improved computational efficiency and practical implementation compared to existing methodologies.

Strengths

1.This paper is eloquently written, and its ideas are easy to follow. 2.While I did not strictly follow all the proofs provided by the authors, I found them to be generally detailed, clear, and easy to understand. 3.The authors provide both quantitative and qualitative results on several datasets, along with a relatively detailed discussion and visualization. However, it is worth noting that the figures in the appendix appear to have low resolution, and the positioning of the legends is somewhat odd. 4.I understand that the convergence of the parametric methodology may require several strong assumptions. Therefore, it is reasonable that the authors have not provided further proof of convergence.

Weaknesses

1.Assumptions in Theorem 3.1: Theorem 3.1 only mentions "under mild assumptions" without specifying them in the main text. It is necessary to state these assumptions clearly in the main body of the paper and provide references to justify their rationality. While a more detailed explanation can be provided in the appendix, the main text should at least include a summary of the assumptions to ensure the reader understands the conditions under which the theorem holds. 2.Experimental Results and Comparisons: 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. The results in Table 1 do not show significant progress or state-of-the-art (SOTA) performance on certain datasets. Additionally, it is confusing that the authors explicitly mentioned Rectified Flow, Flow Matching, and OT-Flow Matching in the literature review, but did not include comparisons with these methods in the experiments. For instance, Rectified Flow has conducted experiments on mass transport, and other related methods, such as Denoising Diffusion Bridge Models from ICLR 2024, have tackled similar tasks. A direct comparison of the proposed method's performance with these existing methods would provide a clearer picture of its advantages and limitations. 3.Discussion of Related Work: The paper could benefit from a more comprehensive discussion of related work, including diffusion models, flow matching, and the application of Wasserstein gradient flow-based models in mass transport tasks. This would help in better contextualizing the authors' innovations and contributions. By situating their work more clearly within the existing literature, the authors could highlight the unique aspects and potential advantages of their proposed method, enhancing the readers' understanding of its significance.

Questions

1.See Weaknesses 2.I am curious about the complexity and cost of the training method proposed by the authors. Could the authors provide more detailed quantitative data, such as training time comparisons with related methods?

Rating

6

Confidence

5

Soundness

3

Presentation

3

Contribution

2

Limitations

Although the authors proposed a novel model to address mass transport problems and provided some theoretical proofs, the method is simulation-based. I believe that the training cost is higher compared to simulation-free methods, and the model performance does not achieve state-of-the-art (SOTA) results. Additionally, there is a lack of direct comparison with existing SOTA methods.