Regularized Distribution Matching Distillation for One-step Unpaired Image-to-Image Translation

We thank the reviewer for valuable feedback and discussion. Below, we address weaknesses and raised questions.

(1) "Could you provide details for the assumptions to guarantee the bijectivity of the Monge optimal transport problem in Line 179?"

We would like to note that the formula that is referred to does not directly relate to our work since we consider a more general soft-constrained optimal transport setting.

One can make the following assumptions to guarantee the bijectivity of $G^*$:

1. $\mathcal{X} = \mathcal{Y} = \Omega \subset \mathbb{R}^d$ and $\Omega$ is compact. In the case of image-to-image translation, we can take $\Omega = [0, 1]^d$;

2. $p^\mathcal{S}$ and $p^{\mathcal{T}}$ have densities with respect to the Lebesgue measure, and the densities are bounded away from zero and infinity. Both $p^\mathcal{S}$ and $p^{\mathcal{T}}$ are supported on $\Omega$;

3. $c(x, y) = h(x - y)$, where $h(x - y)$ is strictly convex.

From these assumptions, it follows that $G^*$ exists and has a particular form:

$$G^*(x) = x - (\nabla h)^{-1} \circ \nabla \phi(x),$$

where $\phi$ denotes Kantorovich potential. Moreover, it is invertible. For more detailed information, please refer to [2], Theorem 1.17 and Remark 1.20

(2) "Could you provide the quantitative results on the toy experiments in Figure 2?"

Let us denote the source and target distributions as $p^\mathcal{S}$ and $p^{\mathcal{T}}$ and $p^{G_{\theta}}$ as the distribution of the generator. $\mathcal{W}^2$ denotes squared Wasserstein Distance computed between given distributions via the POT library. We compute $\mathcal{W}^2$ (discrete version) between the corresponding empirical distributions, since for dimension two the difference between discrete and continuous settings is negligible. $L_2(x, G_\theta(x))$ denotes the squared pairwise $L_2$ distance between inputs and the corresponding generated images.

Please see the table below for the quantitative results on Gaussian $\rightarrow$ Swiss Roll.

Let us look at the $\mathcal{W}^2(p^{G_{\theta}}, p^{\mathcal{T}})$ denoting the "closeness" of the generated distribution to the target distribution. For lower coefficients of $\lambda$ we fit the target distribution better (as can be seen in the column 3 of the table) which is expected because the KL part of the loss begins to dominate. At the same time, for $\lambda$ close to zero, but not equal to zero the pairwise loss $L_2(x, G_\theta(x))$ is fairly close to $\mathcal{W}^2(p^\mathcal{S}, p^{G_{\theta}})$ (columns 1 and 2) and for larger $\lambda$ they are almost equal. For $\lambda = 0$ the pairwise distance $L_2(x, G_\theta(x))$ gets huge as is qualitatively seen in numerous intersections in Figure 2 for $\lambda = 0$. Therefore, we are interested in situation where we have both a good target fit $\mathcal{W}^2(p^{G_{\theta}}, p^{\mathcal{T}})$ and the the pairwise distance $L_2(x, G_\theta(x))$ is close to $\mathcal{W}^2(p^\mathcal{S}, p^{G_{\theta}})$, which means that the map is close to optimal. This holds true for small $\lambda > 0$.

(3) "In Table 1, SDEdit and DDIB achieve competitive FID results and inferior transport costs compared to RDMD. Since the transport cost and FID metric should be considered together, could you provide the translation examples for these models?"

Please, see Figures 8 and 9 in Appendix for a visual comparison of our method against the baselines.

(4) Do all the models in Tables 1 and 2 share the same backbone networks, except for the pixel-space EGSDE?

Yes, all models in Tables 1 and 2 share the same backbone DDPM++ with 55M parameters. Please see Appendix sections D.4 and D.5 for further information on backbone specifics.

(5) "Could you provide a comparison with optimal transport models trained from scratch on the image-to-image translation task? [1] reported that NOT and OTM can achieve competitive FID results with the large backbone network, DDPM++. Could you clarify the advantages of RDMD over these approaches?"

We thank the reviewer for the valuable suggestion. In the rebuttal, we add a comparison with five methods suggested by the reviewers. We ran RDMD on $64\times64$ Wild $\rightarrow$ Cat translation problem to compare with the results mentioned in the DIOTM [1] paper. Here, we consider DIOTM, UOTM [3] and DSBM [4] (mentioned by the reviewer 4Uuc) as the baselines. Moreover, we compare with the OTCS [10] method (suggested by the reviewer v4Tf). We exclude OTM [5] and NOT [6] since [1] and [3] indicate their inferior performance compared to UOTM and DIOTM. We also compare our method with ASBM on CelebA $64\times64$ Male $\rightarrow$ Female translation problem.

Since UOTM is not applied for I2I tasks in the original paper, we modify its official implementation for this scenario and run with the hyperparameters suggested for the larger-scale experiments. We use the official DIOTM implementation from OpenReview and train it with the default Wild $\rightarrow$ Cat hyperparameters to obtain similarity metrics. We train OTCS on Wild $\rightarrow$ Cat with the hyperparameters used for the main unpaired I2I experiment from the paper (unpaired CelebA deblurring). We train ASBM with the hyperparameters suggested for the $128 \times 128$ CelebA experiment. We do not train DSBM but report its performance measured in [7].

We note that the FID (test $\times 10$) corresponds to the FID calculation protocol used in DIOTM, where the authors generate ten samples per each source (we take this result from the paper and mark it with * ). For a more fair comparison with our one-to-one implementation of RDMD, we adapt this calculation by sampling ten augmentations for each source test sample. We present the results in the table below.

Here, we outperform all the mentioned baselines in terms of all metrics (except OTCS [10], mentioned by the reviewer v4Tf, which has a lower $L_2$, but does not fit the target distribution). We thus obtain a method which not only has a better faithfulness-realism trade-off but which is strictly better than the baselines.

We stress that the main advantage of RDMD is the usage of the DMD [11] objective for optimisation. DMD, being a method of pushing an arbitrary generator towards the target distribution, proved to be one of the SOTA methods in unconditional generation. This highly motivates its usage in problems beyond unconditional generation in any case where such distribution matching objective is needed. Methods from this family proved useful in text-to-3D generation [8, 9]. In our paper and the rebuttal, we empirically demonstrate the method's superior performance compared to the diffusion-based and OT-based baselines.

References.

• [1] Improving Neural Optimal Transport via Displacement Interpolation
• [2] Santambrogio, F. (2015). Optimal Transport for Applied Mathematicians.
• [3] Generative Modeling through the Semi-dual Formulation of Unbalanced Optimal Transport
• [4] Diffusion Schrödinger Bridge Matching
• [5] Scalable computation of monge maps with general costs
• [6] Neural Optimal Transport
• [7] Schrödinger Bridge Flow for Unpaired Data Translation
• [8] DreamFusion: Text-to-3D using 2D Diffusion
• [9] ProlificDreamer: High-Fidelity and Diverse Text-to-3D Generation with Variational Score Distillation
• [10] Optimal Transport-Guided Conditional Score-Based Diffusion Model
• [11] One-step Diffusion with Distribution Matching Distillation

References.

• [1] One-Step Diffusion with Distribution Matching Distillation
• [2] Improved Distribution Matching Distillation for Fast Image Synthesis
• [3] Score identity Distillation: Exponentially Fast Distillation of Pretrained Diffusion Models for One-Step Generation
• [4] DreamFusion: Text-to-3D using 2D Diffusion
• [5] ProlificDreamer: High-Fidelity and Diverse Text-to-3D Generation with Variational Score Distillation
• [6] Diffusion Schrödinger Bridge Matching
• [7] Generative Modeling through the Semi-dual Formulation of Unbalanced Optimal Transport
• [8] Improving Neural Optimal Transport via Displacement Interpolation
• [9] Optimal Transport-Guided Conditional Score-Based Diffusion Model
• [10] Adversarial Schrödinger Bridge Matching
• [11] Scalable computation of monge maps with general costs
• [12] Contrastive Learning for Unpaired Image-to-Image Translation
• [13] EGSDE: Unpaired Image-to-Image Translation via Energy-Guided Stochastic Differential Equations
• [14] Schrödinger Bridge Flow for Unpaired Data Translation

(1) "Overall, it currently lacks a theoretical foundation and does not appear to bring substantial methodological improvements. Moreover, the comparisons are extremely weak. Even with this weak comparison, the performance is not convincing. I believe there needs big improvements in concept, literature, experimental design, and comparison groups to strengthen the paper."

In terms of theoretical foundations, we verify our method by establishing its connection with the optimal transport map. In terms of methodology, we stress that DMD [1] and concurrent distribution matching methods (iDMD [2], SiD [3]) achieve SOTA in a one-step generation. This highly motivates adapting these methods for other problems in generative modelling, such as unpaired I2I, and hopefully, achieve SOTA results (e.g. DreamFusion [4] and ProlificDreamer [5] achieve remarkable results by applying distribution matching in text-to-3D).

In addition, the results presented in Figures 4, 7, 10 and Tables 1, 2, suggest that the proposed method achieves better faithfulness-realism tradeoff than the multi-step baselines while being one-step.

(2) "I suggest authors to compare with other methods that fully train a new I2I model with the help of diffusion models. Or I suggest the authors to compare with the GAN-based methods [1] and some of the optimal-transport-based methods [2,3,4,5]."

We thank the reviewer for the valuable suggestion. Below, we perform the additional comparisons with most of the mentioned baselines. Specifically, we compare RDMD with DSBM [6], UOTM [7], DIOTM [8] and OTCS [9] on $64 \times 64$ Wild $\rightarrow$ Cat and with ASBM [10] on $64 \times 64$ Male $\rightarrow$ Female. We exclude OTM [11] and CUT [12] from comparisons, since OTM was shown [7, 8] to have inferior performance than UOTM and DIOTM, and CUT performs worse than EGSDE [13], with which we compared.

Since UOTM is not applied for I2I tasks in the original paper, we modify its official implementation for this scenario and run with the hyperparameters suggested for the larger-scale experiments. We use the official DIOTM implementation from OpenReview and train it with the default Wild->Cat hyperparameters to obtain similarity metrics. We train OTCS on Wild $\rightarrow$ Cat with the hyperparameters used for the main unpaired I2I experiment from the original paper (unpaired CelebA deblurring). We train ASBM with the hyperparameters suggested for the $128 \times 128$ CelebA experiment. We do not train DSBM but report its performance measured in [14]. We obtain the following results:

Here, we outperform all the mentioned baselines in terms of all metrics (except OTCS [9], mentioned by the reviewer v4Tf, which has a lower $L_2$, but does not fit the target distribution). We thus obtain a method, which not also has a better faithfulness-realism trade-off, but which is strictly better than both the diffusion-based and OT-based baselines. For further qualitative comparisons, please see Appendix E of the revised paper and Figures 11, 12, 13.

(3) "In the implementation, do this method use LPIPS as a loss term in training process as Distribution Matching Distillation (DMD) does?"

In DMD, using the additional LPIPS loss is possible due to the ability to generate deterministic (noise, data) pairs from the pre-trained diffusion model. Our method, in contrast, would require (source, target) pairs, which requires a paired data set or an already trained I2I model, which is why we do not use it.

(4) "CycleDiff, DDIB needs to be defined and cited in somewhere in the paper."

We cite both methods in Line 298. We also added the methods names in the text, thanks for pointing that out.

We thank the reviewer for the valuable discussion and comments. Below, we address the aforementioned concerns.

(1) "Please provide deeper justifications for this approach if any exist."

We agree that the idea of combining transport cost with an objective that pushes the distribution towards the target previously existed in literature [1, 2]. However, Distribution Matching Distillation [3] and the successive/concurrent works (e.g. [4], [5]) achieve SOTA results in training one-step generators by propagating knowledge from diffusion models. This motivates using the technique for a wider range of problems and hopefully achieving SOTA results in tasks beyond (un-)conditional generation. For instance, DreamFusion [6] and ProlificDreamer [7] use distribution matching loss for 3D generation and achieve remarkable results in text-to-3D.

(2) "In the experimental section, I suggest that the authors use FID computed based solely on the target distribution, rather than both the source and target as in the current design. This adjustment would better reflect the translation performance of the proposed method, given that other metrics do not adequately capture translation quality."

If we understood correctly, the remark regarding FID computation relates to the footnote "We measure FID between the outputs of the model on the train source data set and the train target data set." In order to avoid misunderstandings, here we mean that FID is measured between the train target data and the generator's outputs. Here, generator takes source training samples as input. We do not mix source and target data sets for FID computation. We will reformulate the footnote more clearly.

(3) "While the use of a transport constraint for regularization is sound, I am concerned about the choice of the squared difference norm for adapting DMD to unpaired I2I tasks. This transport cost may be too simple and fail to capture semantic details between the generated images from the one-step generator and the source images. A more robust cost function, such as the Energy function in EGSDE, could enhance translation capacity. This issue may lead to output images with low L2 loss (high faithfulness) but less realism."

While we agree that choosing the appropriate cost function is important for applications and we consider it as an important future work, we would like to note that even the direct usage of the simple $L_2$ cost frequently produces competitive results (e.g. in Neural Optimal Transport [2]).

Besides, recently popular Schrödinger bridge-based methods ([8, 9]) implicitly solve entropic OT with quadratic cost. We also note that in Figures 4, 7, and 10 we visualise the trade-off between FID and not only $L_2$ but also PSNR/SSIM/LPIPS and achieve results comparable to / better than the baselines (among which there is EGSDE with the aforementioned energy function).

In addition, the latent-space experiment in Section 5.3 could be seen as training a map between male and female faces with transport cost equal to the distance between latent representations. Results in Table 2 and Figure 10 suggest that this type of transport cost can indeed amplify preserving semantic properties: our model achieves better LPIPS but worse $L_2$, than the pixel-space EGSDE. If LPIPS is a representative semantic metric for the task, one may use RDMD in latent space for this purpose.

(4) "Although the sampling efficiency is evident with the one-step approach, the training phase requires an additional fake model. Could the authors provide an analysis of the training resources needed for the proposed framework?"

We provide all the details about computational resources in Appendix D.

(5) "Additionally, the paper lacks references to relevant work on accelerating diffusion models without additional training, such as [1, 2, 3]. There are also unpaired I2I translation methods, like [4, 5, 6], that could utilize these acceleration techniques to generate samples more efficiently, offering alternative baselines to your method. The paper could be more comprehensive if the study regarding these baselines is included."

We consider combining the existing methods of unpaired I2I with acceleration techniques as a relevant yet parallel line of work, which may itself have a significant novelty and give rise to new methods. Besides, accelerating diffusion models almost always results in a decrease in quality. Given that our method already outperforms EGSDE, one could expect it to beat the accelerated EGSDE as well. As for Rectified Flow (RF) [10], we add the comparison with DSBM [8], its SDE-based analogue (proven to have better quality than RF in [8]). The other three baselines are proposed by the reviewers. We run them on the unpaired $64 \times 64$ Wild$\rightarrow$Cat translation problem. We obtain the following results:

Here, we outperform all the mentioned baselines in terms of all metrics (except OTCS [11], mentioned by the reviewer v4Tf, which has a lower $L_2$, but does not fit the target distribution). We thus obtain a method which not only has a better faithfulness-realism trade-off but which is strictly better than the suggested baselines, including DSBM. For further qualitative comparisons, please see Appendix E and Figures 11, 12, 13.

As for [12], the original paper studies unconditional generation and does not include I2I experiments. At the same time, the unpaired I2I modification of the method does not seem straightforward: one should adapt the encoder to map $x_t$ to the source distribution and measure the corresponding KL divergence, which may require novel approaches in both architecture and the procedure itself.

"(6) Is it challenging for both models to converge with the proposed loss function? Could the authors provide training curves for the loss functions of both the one-step generator and the fake diffusion model?"

Due to the initialization with the pre-trained diffusion model the convergence is fast and stable. After a few hundred warm-up iterations, fake score loss monotonically converges. DMD loss quickly reaches its equilibrium and starts oscillating.

We thank the reviewer for the valuable feedback and comments. Below we address weaknesses and the raised questions.

(1) "Experiments are also simple and didn't outperform baseline in some metrics"

We would like to note that it is not fully correct to compare against all metrics simultaneously because the unpaired image-to-image translation inherently suggests a trade-off between the quality of generated samples (realism, i.e. measured by FID or other metrics) and the similarity between the generator input and output (i.e. measured by L2/SSIM/PSNR) because the two properties interfere with each other.

Therefore, we stress that it is more reasonable to compare the trade-off curves of the methods. We provide the corresponding graphs (FID vs L2/PSNR/SSIM/LPIPS) comparing our method to the baselines in Figure 4, Figure 7 and Figure 10. Here, RDMD achieves a better trade-off than other methods while being one-step (the other methods are not). At the same time, Tables 1 and 2 compare the similarity given the fixed realism (among all runs, we choose ones with the closest FID) and outperform other baselines in at least 2/3 or 3/4 metrics.

(2) "Qualitative evaluations are very hard to judge given the 64x64 generated image resolution"

We provide qualitative evaluations for latent space experiments for 256x256 images in Figure 5.

(3) "Overall it is hard to verify the effectiveness of the proposed method based on the limited experiment results, especially when very related work is missing in the comparison."

Our experimental results indicate that the method performs better than multi-step diffusion baselines in terms of quality as it gives a better realism-faithfulness tradeoff. We thank the reviewer for pointing to the OTCS [2] paper. In the rebuttal, we add a comparison with a total of 5 baselines, one of which is OTCS. We run it on the AFHQv2 Wild $\rightarrow$ Cat translation problem (along with 3 other baselines: DSBM [3], UOTM [4], DIOTM [5], suggested by the reviewers 4Uuc and q6cb), and obtain the following results:

While obtaining better $L_2$ and comparable (but still worse) values for other similarity metrics, OTCS fails at producing realistic samples, which is indicated by large FID values and low-quality samples (which we report in Figure 11 of the revised paper version along with the RDMD samples). We also beat all other mentioned baselines in all metrics, thus obtain a method with better faithfulness and realism at the same time.

For comparison with one more baseline, ASBM, please see the general comment or Appendix E of the revised paper (Table 4, Figure 13).

(4) "I wonder if there is potential limitations of the method that led to the use of 64x64 image resolution, which is much lower than the standard 256x256 for most of those dataset. Could there be efficiency or computation challenge for scaling up?"

We stick to 64x64 pixel and 256x256 latent-space experiments due to our constraints in computational resources. The computational complexity of the method is identical to DMD [1]. Therefore, there are no problems with higher dimensions, given the corresponding scale of resources.

References.

• [1] One-Step Diffusion with Distribution Matching Distillation
• [2] Optimal Transport-Guided Conditional Score-Based Diffusion Model
• [3] Diffusion Schrödinger Bridge Matching
• [4] Generative Modeling through the Semi-dual Formulation of Unbalanced Optimal Transport
• [5] Improving Neural Optimal Transport via Displacement Interpolation