Adversarial Schrödinger Bridge Matching

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

(1) I strongly believe the performance gains are mostly from applying GAN architecture, since D-IMF and IMF indicates essentially the same learning scheme. Therefore, it is indeed an improvement from DSBM, but the novelty of this work as a distinct SB study might not be significant from the perspective of diffusion models.

It is not possible to just apply GAN architecture in the original continuous IMF framework [2] since their Markovization objective is based on bridge-matching and implies learning a drift function of SDE and not transitional densities of a discrete process. The main cause of it is that the continuous IMF framework is based on the known property that SB is a unique continuous Markovian and reciprocal process.

The novelty of our work as a distinct SB study comes from our main theorem (Theorem 3.1). We show that to solve the SB problem, it is enough to find a probability distribution $\pi(x_0, x_{t_{1}}, \dots, x_{t_{N}}, x_1)$ with any number $N\geq 1$ of intermediate points, which is Markovian, discrete reciprocal and shares the marginals $p(x_0)$ and $p(x_1)$. Even $N=1$ is enough in our framework, i.e., it is enough to find a joint distribution with three variables $\pi(x_0, x_t, x_1)$ and $t \in (0, 1)$, which is Markovian and discrete reciprocal to solve SB.

We believe that our main theorem opens the door to a novel family of algorithms (including ASBM) for solving SB, based on finding a discrete probability distribution $\pi(x_0, x_{t_{1}}, \dots, x_{t_{N}}, x_1)$ that is Markovian and discrete reciprocal to solve the SB problem.

We also believe that this theorem may facilitate the theoretical study of this problem, including the study on the convergence speed of D-IMF and IMF algorithms since it allows us to work with a more simple object — joint probability distribution instead of continuous stochastic processes. For instance, we show that it is possible to derive closed-form Markovian and reciprocal projections for a multivariate Gaussian case in our framework. In contrast, the continuous counterpart, form solution, is derived only for a 1-dimensional case [2].

(2) This model does not always yield state-of-the-art scores for syntactic datasets (see Table 1&2).

Our algorithm does not beat all other competitors on all the setups. However, it provides the best results on most considered setups by both plan and target metrics (Tables 1 and 2).

(3) While I understand this work for some extent, I do not clearly understand the motivation why do we need a sampling-based approach when we have sampling-free method such as SF$^2$M [3] and LightSB [1] models.

In the limitation section, the LightSB [1] authors state that their algorithm is not designed for large-scale generative problems. Furthermore, they do not provide experiments in the image space, and they only consider image-to-image translation in the latent space of the autoencoder. Other experiments with real data include only single-cell data, which has lower dimensionality than images and does not lie on the low-dimensional manifold-like image distribution.

The authors of SF2M [3] also consider only single-cell problems. Furthermore, while SF$^2$M [3] is indeed simulation-free, it theoretically guarantees convergence to SB only if samples from the ground-truth static SB solution $\pi^*(x_0, x_1)$ are used for training. However, if one has access to the $\pi^*(x_0, x_1)$, the problem itself is no longer unpaired domain translation since there is access to the ground-truth paired data. The authors of [3] address this problem by approximation of $\pi^*(x_0, x_1)$ by minibatch Optimal Transport techniques, which seems to provide a good approximation in the considered single-cell setups but is known to have a bias growing with the dimensionality of the considered space. For instance, no experiments show that this approximation works well in the case of high-dimensional unpaired image-to-image setups.

Thus, while each of these simulation-free algorithms has its advantages and works well in the setups for which they were designed, there is still no simulation-free algorithm that can demonstrate good quality on the unpaired image-to-image translation without using any autoencoders.

(4) How do the authors suspect that the FID score is very low for ASBM? Following the arguments of [1], could the authors report CMMD scores or other metrics for results in Fig. 3?

Since it is hard to judge quality given only by the absolute values of FID, we compare our method with the DSBM and see that it provides better FID. As per your request, we measure the CMMD score and present the results below:

We observe that CMMD scores align with FID scores presented in our work.

(5) Typo L213: Denoising Denoising GANs

Thank you for pointing out this. We will fix it in the final version.

Concluding remarks. We would be grateful if you could let us know if the explanations we gave have been satisfactory in addressing your concerns 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] Alexander Korotin. Light schrödinger bridge. ICLR, 2024.

[2] Shi et al. (2023) -- Diffusion Schrodinger Bridge Matching

[3] Tong A. Y. et al. Simulation-Free Schrödinger Bridges via Score and Flow Matching AISTATS 2024.

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

(1) Computational burden and stability study.

As requested, we provide a study on the effectiveness of ASBM and compare it with the DSBM.

Number of parameters. Our ASBM generator has about 42 million parameters, while each discriminator has 27 million. The generator of DSBM has 38 million parameters. Thus, generators have similar capacities.

Training time. We mentioned (lines 714-716) that for Celeba setups, our ASBM uses 7 days of training on a A100 GPU, while the DSBM uses 5 days of training. The time is comparable since both methods use the same number of IMF and D-IMF iterations.

Please note that we provide an analysis of ASBM (ours) and DSBM algorithms in Appendix E, where we show that additional iterations only slightly change the results.

Inference time. We measure the inference time of both ASBM (ours) and DSBM algorithms on the same hardware with one A100 GPU. We measure the mean and standard deviation over 16 runs. For a batch of 100 images, ASBM requires about 842 ms (std 97 ms), while DSBM requires 33870 ms (std 47 ms). Thus, our algorithm provides up to 40x speedup compared to DSBM.

GPU memory consumption. On the Celeba setup, our algorithm requires about 25 GB of GPU for training with a batch size of 32 and about 12 GB on the inference stage to translate a batch of 100 images. Meanwhile, DSBM requires 34 GB of GPU for training with batch size 32 and about 26.7 GB on the inference stage for a batch of 100. Thus, our algorithm utilizes significantly less GPU memory.

Stability of training. To show that our procedure is stable, we provide the plot of FID during the training of all D-IMF iterations in Figure 3 of the attached pdf. On the plot, we provide FID vs the number of generator gradients updated during the training. We highlight the starts of each D-IMF iteration by the vertical lines. We observe stable decreases of FID.

(2) Lack of empirical and ablation studies.

The number of phases K. We provide an analysis of D-IMF/IMF iterations (K) in Appendix E.1. of our work.

Independent vs minibatch coupling. As per request, we provide a study on running ASBM with different coupling. We provide the additional results for the Colored MNIST setup with $\epsilon=10$ (which we discuss in Appendix C.3 of our paper). In Figure 1 of attached pdf, we provide qualitative generation results for our ASBM model with independent and minibatch couplings after 3 D-IMF iterations, which we used in our work for this setup. We also provide FID values in Table 1 of attached pdf. We see comparable visual quality as well as similarity for both results; the FID metric is slightly better for minibatch case.

Study on different NFE. As requested, we present the study's results on different NFE of our ASBM model. In our study, we use our ASBM model trained with $4$ steps for Celeba male to female with $\epsilon=1.0$. We test this model with different numbers of NFE $\in [1, 2, 3, 4, 8, 16, 32]$. In all the cases, we use a uniform schedule. (See answer 4 for details). We present qualitative results of generation for several inputs in Figure 2 of the attached pdf. We also present the FID values computed on the test set and average $L_2$ cost $c(x_0,x_1) = \frac{||x_0-x_1||^2}{2}$ between the input image $x_0$ and generated image $x_1$ to assess similarity in Table 2 of the attached pdf. Interestingly, our model with $\text{NFE}=3$ provides almost the same FID value as for $\text{NFE}=4$, which we use at the training, but significantly lower $L_2$ cost.

Training with different NFE. Since the time of the rebuttle and our resources are limited we train ASBM with $NFE=2$ and $NFE=8$ with $K=2$ D-IMF iterations and $\epsilon=1.0$ using the same procedure as before. We present samples in the Figure 5 together with ASBM model with $NFE=4$ after $K=2$ D-IMF iteration trained before. Male-to-female metrics:

We observe better results for lower NFE, possibly because the model needs to learn fewer transitional distributions. We will add all the results to the final version.

(3) The FID of the reverse direction.

Thank you for this suggestion. The values are given below.

Thus, the FID values for the reverse direction show the same behavior as for the forward direction.

(4) Measurement of perceptual quality

As requested, we measure perceptual similarity between input and translated images using standard LPIPS metric [1] (as used in DSBM). For our Celeba male-to-female setup, we measure LPIPS between input and translated images from the test. We present the results below (the lower the better):

We observe almost the same similarity in the case of $\epsilon=1$ and better similarity for ASBM (ours) in the case of $\epsilon=10$. We will add these results in the final version.

(5) Time discretization.

We use uniform time discretization in all our experiments, i.e., for the number of inner times points $N$, we use time discretization $t_{n} = \frac{n}{N+1}$ for $n \in [0, N+1]$.

(6) Did you use any pretrained model?

No, we do not use any pretrained models in our experiments.

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] Zhang R. The unreasonable effectiveness of deep features as a perceptual metric. CVPR, 2018.

[2] Shi Y. Diffusion Schrödinger bridge matching. NeurIPS, 2024.

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

(1) I enjoyed reading the majority of the paper—except, perhaps oddly, the paper title. Over 85% of the main technical contributions (Sec 3) stem from the construction of the discrete IMF formulation. The concept of "adversarial" appears only as one (of many plausible) implementation of Eq (15). I feel like the "discrete" aspect isn't highlighted enough, or at all, in the title. This is, of course, just a personal comment rather than a criticism.

The title of the work which introduces DSBM algorithm is "Diffusion Schrödinger Bridge Matching". Our work is based on it but, at the same time, proposes to use adversarial training as a current main competitor of diffusion models in the field of generative models. To highlight both the relation with the prior work and the usage of adversarial training in our approach, we decide to use the title "Adversarial Schrödinger Bridge Matching".

(2) In L220, the non-saturating GAN loss should be provided in the main paper rather than being postponed to the Appendix.

Thank you for this suggestion. We initially placed the non-saturating GAN loss in the Appendix to keep the main text focused and avoid introducing technical details too early. Furthermore, the non-saturating GAN loss is rather standard approach in GANs, so we did not pay much attention to it. However, we agree that it would be beneficial to mention the non-saturating GAN loss in the main paper for clarity. We will move it to the main text in the final version and refer readers to the Appendix for a more detailed description of our adversarial procedure.

(3) The notions of T and M are slightly confusing. For example, In L73 and L76, both represent Markovian processes. Is there a typo?

In this context, the process $T$ and $M$ are the same Markovian process. We will replace $T$ with $M$ in the final version. Thank you for mentioning it.

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. We are also open to discussing any other questions you may have.

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

(1) Part of the theory was already done in [1]. ...

Indeed, the authors of [1] analyze the discrete-time Markovian projection and show that it converges to continuous-time projection in certain limiting cases. We cite this result in line 260. We will add additional citations and discussion in Section 3.3 after defining Markovian projection to highlight some aspects that have already been introduced in [1].

Please note that our Proposition 3.5 on KL-minimization extends the results from [1, Appendix E]. The authors of [1] considered the Markovian projection for a mixture of Markovian bridges. In turn, in our Proposition 3.5 we consider arbitrary distribution $\pi(x_0, x_{t_{1}}, \dots, x_{t_{N}}, x_1)$ without any assumptions on its conditional distribution $\pi(x_{t_{1}}, \dots, x_{t_{N}}|x_0, x_1)$. It makes the proof technically more complex but provides a more general result.

(2) ... the purpose of the Gaussian study. ...

There are two purposes: to experimentally study the rate of convergence and to demonstrate that a closed-form solution for the multivariate Gaussian case can be derived within the discrete-IMF framework. In the final version, we will shorten Section 3.4 by moving detailed technical definitions (such as matrices $U$ and $K$) to the Appendix but retain the key points in the main text.

(3) The claim of exponential convergence should be tamed down. ...

We agree that the phrasing in line 294 could be misleading. This observation is based solely on our experimental results, and the whole discussion on convergence rate is placed inside the experimental section (Section 4), while all of our theory is placed in Section 3. We also directly mentioned that there is no theoretical proof in the limitation section (lines 465-467). To avoid possible misunderstanding, we will revise the bolded text.

(4) ... remove this noise ... consider a deterministic last step as a postprocessing ...

We use the original published code from the DSBM authors [1] to ensure a fair comparison. In this code, noise is not added at the last sampling step. As you requested, we additionally used the drift function for the last time-moment $t=0.99$ to add an extra denoising step at the final step of the sampling process. It indeed removes the remaining noise, (see the samples in Figure 4 of the attached pdf). For the Celeba setup with $\epsilon=1.0$, we observe the improvement of the FID metric from $37.8$ to $29.7$. We will add this additional result in the experimental section.

(5) ... it seems that the diversity is much lower than DSBM. I would like to see a proper investigation of this phenomenon ...

As requested, we provide a study on the diversity of images produced. We measure the diversity of images produced by DSBM and ASBM on the Celeba setup by utilizing the LPIPS diversity as was done in [3, Table 1]. We take a subset of $500$ images from the test part of the Celeba dataset. We sample a batch of generated images of size $16$ for each input image, compute the average distance between all possible pairs of these images, and average these values over all inputs. We present the results below for DSBM and ASBM and different values of the coefficient $\epsilon=1$ and $\epsilon=10$.

Indeed, the LPIPS diversity is lower for ASBM. However, without a known ground truth for diversity, it is difficult to determine definitively whether ASBM provides insufficient diversity or if DSBM produces too much. We will include these findings in the final version of the paper and discuss this trade-off.

(6) .... However, why choose DD-GAN? ...

The mentioned adversarial distillation techniques [1,2] are based on tractable forward process in diffusion models for simulation-free sampling, which is not feasible for the unpaired Schrödinger Bridge (SB). Additionally, these methods require an already trained model for distillation, whereas our approach does not need a pretrained diffusion model.

Regarding consistency models, it is unclear how to use them in an unpaired image-to-image setup. Works [3, 4] do not address this. The work [3] discusses generative models and zero-shot translation with known degradations, applicable only to specific tasks like colorization. Similarly, [4] focuses on generative modeling. While consistency models can potentially accelerate arbitrary ODEs, the Schrödinger Bridge (SB) problem's solution is an SDE, which makes it challenging to use consistency models.

(7) ...the authors are quick to dismiss these limitations. ... (Appendix A)

We understand that adversarial training is a limitation. We intended to say that we use the DD-GAN approach, designed to address this limitation to some extent (mode coverage [2, Section 5.3], training stability [2, Appendix D]).

Furthermore, to show that our DDGAN-based ASBM procedure is stable, we additionally provide the plot of FID during the training of all D-IMF iterations in Figure 3 of the attached pdf. We provide FID vs the number of generator gradients updated during the training. We highlight the starts of each D-IMF iteration by the vertical lines. We observe the stable decrease of FID over the generator gradient update steps.

We will remove the phrase "we do not treat it to a be a serious limitation" in the final version.

Concluding remarks. We would be grateful if you could let us know if the explanations we gave have been satisfactory in addressing your concerns 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 are the same as you used in the corresponding weaknesses and questions.