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

Authors: Denis Rakitin, Ivan Shchekotov, Dmitry Vetrov
Primary area: generative models
Keywords: diffusion distillation, distribution matching distillation, optimal transport, image-to-image translation

Abstract:
Diffusion-based generative models achieve SOTA results in mode coverage and generation quality but suffer from inefficient sampling. Recently introduced diffusion distillation techniques approach this issue by transforming the original multi-step model into a one-step generator with approximately the same output distribution. Among these methods, Distribution Matching Distillation (DMD) offers a suitable framework for training general-form one-step generators, applicable beyond unconditional generation. In this paper, we propose a modification of DMD, called Regularized Distribution Matching Distillation (RDMD), which applies to the unpaired image-to-image (I2I) translation problem. To achieve this, we regularize the generator objective from DMD with the transport cost between its input and output. We validate the method's applicability in theory by establishing its connection with optimal transport. Moreover, we demonstrate its empirical performance in application to several translation tasks, including 2D examples and I2I between different image datasets, where it performs on par or better than multi-step diffusion baselines.

Decision: reject
Reviewer scores: 5, 5, 3, 3

Reviews and discussion:
<thread>
creation date:
1730388107985

thread_type:
official_review

<thread_object>
creator: reviewer_q6cb

<content>
summary:
- This paper proposes a diffusion model distillation method for a one-step generator, called Regularized Distribution Matching Distillation (RDMD). RDMD introduces an additional cost regularizer into the previous Distribution Matching Distillation approach. The proposed method is evaluated on the image-to-image translation task, where this cost-minimization regularizer is desirable.

strengths:
- This paper provides the relationship between RDMD and the optimal transport map in Thm 1.
- This paper is easy to follow.

weaknesses:
- The technical novelty of the proposed method is incremental.
- The proposed method is not compared with the optimal transport models.

questions:
- Could you provide details for the assumptions to guarantee the bijectivity of the Monge optimal transport problem in Line 179?
- Could you provide the quantitative results on the toy experiments in Figure 2? Evaluating the transport cost and Wasserstein distance between the generated and target distribution would strengthen the experimental results.
- 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?
- Do all the models in Tables 1 and 2 share the same backbone networks, except for the pixel-space EGSDE?
- 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?

[1] Choi, Jaemoo, Yongxin Chen, and Jaewoong Choi. "Improving Neural Optimal Transport via Displacement Interpolation." arXiv preprint arXiv:2410.03783 (2024).
</content>
</thread_object>

</thread>
<thread>
creation date:
1730542355333

thread_type:
official_review

<thread_object>
creator: reviewer_4uuc

<content>
summary:
This paper perform unpaired image-to-image translation through adding quadratic regularization term on the existing distribution matching distillation method.

strengths:
- Experiments were conducted on various data and resolutions in both pixel-space and latent diffusion.

weaknesses:
I find the contribution of this paper to be somewhat limited. 

- The method mainly introduces a quadratic regularization term, $\\| x - G_\theta (x) \\|^2 $, into existing distillation approaches. This is a very simple method. Moreover, in Equations 9-11, the paper replaces the Gaussian input $p_{noise}$ with $p_{source}$, and re-train $G_\theta(x)$ with $x\sim p_{source}$. Since they no longer put noise into $G_\theta$, this approach seems more like re-training $G_\theta$ by leveraging a well-pretrained diffusion model than a distillation method. So, this method should be seen as a fully-retrained model, not a distillation method. However, the paper only compares with zero-shot diffusion-based methods or classifier-learned-based methods. 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].

- 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.

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

- CycleDiff, DDIB needs to be defined and cited in somewhere in the paper.


[1] T. Park et al., Contrastive Learning for Unpaired Image-to-Image Translation, ECCV, 2020.

[2] J. Fan et al, Scalable computation of monge maps with general costs, NeurIPS, 2021.

[3] J. Choi et al, Generative Modeling through the Semi-dual Formulation of Unbalanced Optimal Transport, NeurIPS, 2023.

[4] Y. Shi et al., Diffusion Schrodinger Bridge Matching, NeurIPS, 2023.

[5] N. Gushchin et al., Adversarial Schrodinger Bridge Matching, NeurIPS, 2024.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730635331937

thread_type:
official_review

<thread_object>
creator: reviewer_g3fx

<content>
summary:
This study proposes distilling a pre-trained diffusion model into a one-step generative model to efficiently address unpaired Image-to-Image (I2I) translation. The method uses distribution matching distillation, effective in image generation, and regularizes it for unpaired I2I tasks. The proposed approach achieves good results and is significantly efficient in the sampling phase, requiring only one generation step.

strengths:
- Applying distribution matching distillation (DMD) to efficiently tackle unpaired I2I translation is an interesting approach.
- Empirical results indicate that the proposed method is effective.

weaknesses:
- The contribution feels incremental, as it primarily extends existing distribution matching distillation from text-to-image generation (where the prior distribution is Gaussian) to unpaired I2I by simply substituting the Gaussian prior with a source image distribution and adding the transport constraint. Please provide deeper justifications for this approach if any exist. 

- 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 (potentially high FID), as observed in your experiments. 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.

- 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. 

- 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?
 
- 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?

[1] Lu et al., 2022, "DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps", NeurIPS'22.
[2] Zhang et al., 2023, "Fast Sampling of Diffusion Models with Exponential Integrator", ICLR'23.
[3] Wizadwongsa et al., 2023, "Accelerating Guided Diffusion Sampling with Splitting Numerical Methods", ICLR'23.
[4] Lee et al., 2023, "Minimizing Trajectory Curvature of ODE-based Generative Models", ICML'23.
[5] Liu et al., 2023, "Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow", ICLR'23.
[6] Min et al., 2022, "EGSDE: Unpaired Image-to-Image Translation via Energy-Guided Stochastic Differential Equations", NeurIPS'22.

questions:
Please refer to the Weaknesses section for my concerns and questions.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730710883124

thread_type:
official_review

<thread_object>
creator: reviewer_v4tf

<content>
summary:
Authors proposed a modified extension of DMD that includes transport cost for unpaired I2I task. They established connection to optimal transport and showed that solution of the soft-constrained RDMD converges to that of the hard-constrained Monge problem. Experimental results on several datasets also prove the effect of the model.

strengths:
Authors prove that solution of the soft-constrained RDMD converges to that of the hard-constrained Monge problem. The proof seems complicated and unfortunately I do not have the mathematical background to fully verify all steps.

weaknesses:
Authors fail to mention related works such as OTCS (Optimal Transport-Guided Conditional Score-Based Diffusion Model). Experiments are also simple and did not outperform baseline in some metrics. Qualitative evaluations are very hard to judge given the 64x64 generated image resolution. 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.

questions:
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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732802234930

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Added comparisons with the suggested baselines

comment:
We sincerely thank the reviewers for their valuable discussion, comments and suggestions. A large part of the remarks is related to the comparison with OT-based methods. Thus, we write a general comment, in which we present comparisons with the suggested baselines. 

Among the suggested baselines, we compare with UOTM [1], DIOTM [2], OTCS [3], DSBM [4] and ASBM [5]. We do not directly compare with OTM [6] and NOT [7] since the authors of [1] and [2] demonstrate their inferior performance compared to UOTM and DIOTM. **All the comparisons, including generated samples, can be found in Appendix E of the revised paper.**

We choose two I2I tasks for comparison. First is $64\times64$ AFHQv2 Wild $\rightarrow$ Cat, which is widely studied as a benchmark for these methods (e.g. for DSBM in [8] and DIOTM). 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 also train OTCS on Wild $\rightarrow$ Cat with the hyperparameters used for the main unpaired I2I experiment from the paper (unpaired CelebA deblurring). We do not train DSBM but report its performance measured in [8]. 

The second I2I task is $64 \times 64$ CelebA Male $\rightarrow$ Female studied, for example, in [2, 5]. Here, we train ASBM with the hyperparameters suggested for the $128 \times 128$ CelebA experiment. 

We compare RDMD with OTCS, DSBM, UOTM and DIOTM on **Wild $\rightarrow$ Cat** in the following table, where we beat all the baselines in all the metrics (in both faithfulness and realism), except OTCS, which has a lower $L_2$, but does not fit the target distribution. Figures 11 and 12 from the revised paper version validate the performance qualitatively.

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.

|  Wild $\rightarrow$ Cat        | FID (train) | FID (test) | FID (test $\times 10$) | $\sqrt{L_2}$ | LPIPS     | PSNR      | SSIM      |
| -------- | ----------- | ---------- | ---------------------- | ------------ | --------- | --------- | --------- |
| OTCS     | 54.50       | 65.01      | ---                    | **13.71**        | 0.508     | 18.38     | 0.468     |
| DSBM     | ---         | 25.41      | ---                    | ---          | 0.485     | ---       | ---       |
| UOTM     | 14.85       | 26.7       | ---                    | 27.39        | 0.509     | 12.14     | 0.250     |
| DIOTM    | 8.94        | 20.28      | 10.72*                 | 18.69        | 0.465     | 15.66     | 0.496     |
| **RDMD** | **7.87**    | **18.18**  | **9.24**               | **15.59**    | **0.363** | **19.22** | **0.594** |

Next, we compare RDMD with ASBM on CelebA **Male $\rightarrow$ Female**. Here, we also beat the baseline in all metrics (and obtain higher-quality samples, which can be verified by Figure 13).

|     Male $\rightarrow$ Female     | FID (train) | FID (test) | $\sqrt{L_2}$ | PSNR      | SSIM      |
| -------- | ----------- | ---------- | ------------ | --------- | --------- |
| ASBM     | 22.94       | 31.99      | 15.32        | 17.40     | 0.524     |
| **RDMD** | **12.04**   | **25.6**   | **11.88**    | **20.97** | **0.701** |

**References.**

* [1] Generative Modeling through the Semi-dual Formulation of Unbalanced Optimal Transport
* [2] Improving Neural Optimal Transport via Displacement Interpolation
* [3] Optimal Transport-Guided Conditional Score-Based Diffusion Model
* [4] Diffusion Schrödinger Bridge Matching
* [5] Adversarial Schrödinger Bridge Matching
* [6] Scalable computation of monge maps with general costs
* [7] Neural Optimal Transport
* [8] Schrödinger Bridge Flow for Unpaired Data Translation
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper extends the diffusion distillation methods, specifically, distribution matching distillation to the unpaired image-to-image translation task. The modification mainly consists of 1) replace the gaussian noise by the source domain distribution, 2) add regularisation related to optimal transport. As pointed out by several reviewers, the novelty of the proposed work is limited. Also, the L2 loss might not be able to capture the semantic difference between domains. Finally, there are several baseline methods that were not included in the comparison. Given these concerns, I would recommend rejecting this paper.