Unpaired Image-to-Image Translation via Neural Schrödinger Bridge

Summary

The paper presents Unpaired Neural Schrödinger Bridge (UNSB), an innovative approach for unpaired image-to-image translation tasks, combining Schrödinger Bridge model with adversarial training. It expands the applicability of diffusion models by overcoming the Gaussian prior assumption limitation, which typically restricts them in unpaired image-to-image translations.The experiment results signal the potential of the proposed approach in dealing with high dimensional datasets, commonly referred to as the "curse of dimensionality." Experiments have shown the superiority of the proposed method where comparisons are conducted with their method and multiple baseline I2I methods on multiple I2I datasets. An ablation study is conducted with clear analyses.

Strengths

1 The paper is aligned with recent advancements in stochastic generative modeling, image-to-image translations, and optimal transport, pushing the boundary of what's possible in these domains. 2 The experimentation shows promising results in I2I tasks in terms of multiple metrics. 3 The paper clearly explains the underlying mathematical conventions for Schrödinger Bridge and UNSB models and their connection with stochastic differential equations. 4 The authors have taken additional steps to ensure reproducibility of their work by providing code and detailed descriptions of hyperparameters in the paper.

Weaknesses

There are no significant weaknesses. It could be valuable to include the reverse translations (such as Zebra to Horse, Winter to Summer) in the supplementary material, providing more comprehensive experimental results.

Questions

See the weakness section.

Rating

6: marginally above the acceptance threshold

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

Summary

The paper proposes a method leveraging Schrödinger Bridge (SB) to bypass the limitation of applying Diffusion Models (which assume a Gaussian prior) to Image-to-Image (I2I) translation. Although previous approaches for utilizing SB to I2I translation have been explored, the proposed method, UNSB, is designed to be directly applied to I2I and scalable to high resolution images. The method is implemented as a composition of generators learned via adversarial learning that overcome the curse of dimensionality with advanced discriminators, showing success in several unpaired I2I translation tasks.

Strengths

1. The formulation of SB as a composition of generators learned via adversarial learning seems novel. Moreover, the ability to apply patch-level discriminators is crucial for data with diverse textures as the ones considered in the paper. 2. The translation quality and faithfulness seem to be good. 3. The paper is well written and easy to follow.

Weaknesses

1. I think that some relevant baselines were not included in the experimental section. Could the author compare the proposed method to EGSDE [1] (Diffusion-based model)? [1] Zhao et al. EGSDE: Unpaired Image-to-Image Translation via Energy-Guided Stochastic Differential Equations. In NeurIPS, 2022. 2. The authors attribute the failure of SB for unpaired image-to-image translation to the curse of dimensionality. However, methods such as EGSDE and SDEdit do achieve successful translations in translating high resolution images. Could the authors elaborate on the difference between the cases? 3. Following my previous point, the datasets considered in the paper are all of “scene-level" image to image translation (in which mostly color and texture need to be altered during translation). In most of other I2I translation works, many "object-level" datasets are evaluated e.g. CelebA-HQ Male-to-Female, Cat-to-dog, etc. Do the authors expect different performance for the latter case? I believe it would be helpful to include benchmarks for at least some of these additional datasets to understand the strengths and the limitations of the proposed method.

Questions

1. The description of the Stochasticity analysis (Fig. 7) is not clear to me. Could the author describe this ablation experiment and what does the generated map tell us? I understand that high pixel-wise std implies greater diversity, but why is it surprising that the std is greater for the foreground pixels? What is the scale of the std?

Rating

6: marginally above the acceptance threshold

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

2 fair

Summary

The paper proposes a method for unpaired image-to-image translation based on Schrödinger bridge. The authors combine several ideas to obtain a practically-sounding method. At first, they propose to learn their model in a multistep manner, which is a natural and worth-considering given the recent success of diffusion models. Secondly, they propose to use the additional GAN loss to better match the target distribution. Additionally, the authors report the usefulness of the additional regularization applied at the training stage. Overall, the proposed approach demonstrates good practical results and has clear theoretical motivation. **Update**: My concerns seems to be addressed. I rise my score to 6

Strengths

The paper is overall-well written. It has clear motivation (unpaired image-to-image translation via entropic optimal transport (or Schrödinger bridge) problems. The idea of multistep approach, where the timeline of the Schrödinger process is discretized into several steps and a practitioner learn a kind of “denoising” model (an intuition from diffusion models) is fresh and worth to be considering. Moreover, it seems that the considering several steps indeed improves the qualitative performance, as shown by the practical experiments. The practical results are also encouraging.

Weaknesses

My concerns are covered in the questions section (see below). I expect the authors to address them.

Questions

- I have some doubts regarding the statement (end of page 1) that “no work has successfully trained SBs for direct translation between high-res images” and (end of page 2) “our work represents the first endeavor on this problem”. For example, [1], [2] deals with $64 \times 64$ images and tackles unpaired image-to-image problems. I am not sure, if the methods from [1], [2] works with $256 \times 256$, but, anyway, [1] and [2] should be covered as the related works and the competitive methods. - Regarding the competitive methods (GANs), the authors miss the StarGAN and StarGAN v2 [3]. The latter model demonstrates a good quality on the unpaired image2image and I think this model is a good baseline method. Please, add it to the comparison. - I have one concern about objective (14). The transition from constrained problem (9)-(10) to the combination of losses (14) is non-trivial and not fully theoretically justified. I am not sure, that the objective (14) indeed solves the SB problem. Some comments/clarifications will contribute to the clearness of the manuscript flow. - I ask the authors to add the clarifications how to estimate the entropy with samples, a technique introduced at the beginning of page 6. It will facilitate the comprehension and smoothness of the reading. - I strongly encourage the authors to evaluate how does their approach solve SB problem on a nontrivial setup beyond gaussian2gaussian case. In particular, there is a recent paper [4] which proposes benchmark pairs of distributions with known GT EOT/SB solutions, even for images data. The validation on such a benchmark will highly contribute to the quality of the paper and will probably answer the question (3). - I have one more question regarding the baseline methods. The authors choose the work [5] (called them as NOT) “as the representative of OT [methods]”. This choice really surprised me because [5] also deals with SB problem, and, as I understand, their method is not designed for image data (at least, the authors of [5] does not consider image data use cases). That is why comparing with [5] seems strange. I recommend authors take [6, 7, 8] as the representative of OT methods instead. These methods also has a possibility to generate the samples of controllable diversity, have GAN-resembling objective. Moreover, they are designed specifically for unpaired image2image problems. [1] Diffusion Schrödinger Bridge Matching, NeurIPS’2023 [2] Entropic Neural Optimal Transport via Diffusion Processes, NeurIPS’2023 [3] StarGAN v2: Diverse Image Synthesis for Multiple Domains, CVPR’2020 [4] Building the Bridge of Schrödinger: A Continuous Entropic Optimal Transport Benchmark, NeurIPS’2023 [5] Transport with support: Data-conditional diffusion bridges, arxiv [6] Neural Optimal Transport, ICLR’2023 [7] Neural Monge Map estimation and its applications, TMLR [8] Kernel Neural Optimal Transport, ICLR’2023

Rating

6: marginally above the acceptance threshold

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

2 fair

Ethics concerns

No

Summary

This paper aims to address the problem that the diffusion model based unpaired image translation tasks rely on the Gaussion prior assumption. It resorts to Schrödinger Bridge and proposes the Unpaired Neural Schrödinger Bridge (UNSB) to express SB problem as a sequence of adversarial learning problems. By this formulation, SB can regarded as Lagrangian formulation under the constraint on the KL divergence between the true target distribution and the model distribution. This formulation leads to the composition of generators learned via adversarial learning that overcome the curse of dimensionality with advanced discriminators. Experiments demonstrated the effectiveness of the proposed method compared to its competitors.

Strengths

The paper experimentally identifies the cause behind the failure of previous SB methods for image-to-image translation as the curse of dimensionality, by a toy task. The proposed method gives a new formulation of SB, which benefits on addressing the identified curse of dimensionality. The proposed method shows better results than existing methods.

Weaknesses

The paper claimed that "none of SB models so far have been successful at unpaired translation between highresolution images", however, the experiments in this paper do not include highresolution images yet (256x256 can not be regarded as highresolution). How will the proposed method perform on this setting is still unclear. It is unclear that the findings by ablation study are still valid for other datasets? When presenting the stochasticity analysis, it only gives results for one image. It would be better if the paper can provide statistical numbers across many different images/cases.

Questions

no

Rating

8: accept, good paper

Confidence

1: You are unable to assess this paper and have alerted the ACs to seek an opinion from different reviewers.

Soundness

3 good

Presentation

3 good

Contribution

3 good

General Reply to All Reviewers

We sincerely thank all Reviewers for the constructive feedback, which was incorporated into the revised version of our paper. Here, we summarize changes to our paper, such as additional experiments or clarifications. Modified contents have been highlighted red in the revised paper. **Section 2.** We added missing references [1,2]. [1] Entropic neural optimal transport via diffusion processes, Gushchin et al., NeurIPS, 2023. [2] Diffusion Schrödinger Bridge Matching, Shi et al., NeurIPS, 2023. **Section 5.** We fixed a wrong reference for Neural Optimal Transport [3]. [3] Neural Optimal Transport, Korotin et al., ICLR, 2023. **Appendix B.** We added more description on entropy approximation. **Appendix C.2.** We added more results on stochasticity analysis. **Appendix C.3.** We added a result on Male2Female translation, with baselines EGSDE, StarGAN v2, and NOT. **Appendix C.4.** We added a result on the entropic OT benchmark [4]. [4] Building the bridge of Schrödinger: A continuous entropic optimal transport benchmark, Gushchin et al., NeurIPS Track on Datasets and Benchmarks, 2023. **Appendix C.5.** We added a result on Zebra2Horse translation. **Appendix C.6.** We added additional ablation study on Summer2Winter. **Appendix D.** We discuss why diffusion-based methods perform well on some tasks (e.g., Male2Female, Cat2Dog), but not so well on others (e.g., Horse2Zebra, Summer2Winter).

Result on $512 \times 512$ Resolution Images

We would like to let the Reviewers know that we obtained encouraging results on even higher resolution images. In Figure 15 of our revised paper, we added Cat2Dog translation samples on AFHQ $512 \times 512$ resolution images. **We emphasize that this is a preliminary result with only 100 epochs of UNSB training, and we may obtain better results with longer training.**