Deep MMD Gradient Flow without adversarial training

Authors: Alexandre Galashov, Valentin De Bortoli, Arthur Gretton
Primary area: generative models
Keywords: Generative modeling, diffusion models, Wasserstein gradient flows, generative adversarial networks, discriminator flow

Abstract:
We propose a gradient flow procedure for generative modeling by transporting particles from an initial source distribution to a target distribution, where the gradient field on the particles is given by a noise-adaptive Wasserstein Gradient of the Maximum Mean Discrepancy (MMD). The noise adaptive MMD is trained on data distributions corrupted by increasing levels of noise, obtained via a forward diffusion process, as commonly used in denoising diffusion probabilistic models. The result is a generalization of MMD Gradient Flow, which we call Diffusion-MMD-Gradient Flow or DMMD. The divergence training procedure is related to discriminator training in Generative Adversarial Networks (GAN), but does not require adversarial training. We obtain competitive empirical performance in unconditional image generation on CIFAR10, MNIST, CELEB-A (64 x64) and LSUN Church (64 x 64). Furthermore, we demonstrate the validity of the approach when MMD is replaced by a lower bound on the KL divergence.

Decision: accept (poster)
Reviewer scores: 6, 5, 6, 6

Reviews and discussion:
<thread>
creation date:
1730063242418

thread_type:
official_review

<thread_object>
creator: reviewer_vkum

<content>
summary:
This paper proposes a generative model based on the Wasserstein gradient flow of the Maximum Mean Discrepancy (MMD), called DMMD. Given the forward process in diffusion models, the proposed method learns a discriminator between clean data and noisy data. The Wasserstein gradient flow is represented via this discriminator. Therefore, DMMD generates samples by simulating particle trajectories through this flow. DMMD achieves superior performance compared to other discriminator flow and MMD flow baselines in image generation across various datasets.

strengths:
- This paper is easy to follow.
- This paper suggests an efficient approximate sampling procedure for the linear kernel.
- DMMD is evaluated on various image datasets, such as CIFAR-10, MNIST, CelebA (64x64), and LSUN-Church (64x64).

weaknesses:
- My main concern is whether the trajectories of the probability distributions for the forward process $\\{ p^{1}\_{t} \\}\_{t \geq 0}$ in the diffusion model and the MMD gradient flow $\\{ p^{2}\_{t} \\}\_{t \geq 0}$ coincide. If these trajectories are different, DMMD learns the Wasserstein gradient flow for minimizing $MMD (p^{1}\_{t}, p_{data})$. However, during the generation process, the particles follow $\\{ p^{2}\_{t} \\}$ at $ t - \triangle t$. This gradient flow mismatch cannot guarantee that the particles correctly generate the target distribution. Could you clarify this concern?
- The generation process is computationally expensive. How many $N_{s}$ steps are required for each time $t$ during sample generation (Eq 9)?
- For a general kernel, the sampling from DMMD requires access to the training data (Eq 10). Only for linear kernel, this issue can be avoided by saving the average features for each time $t$ (Eq. 12).

questions:
- In Table 1, do all the other baselines use the same backbone network?
- Could you provide the wall-clock time comparison to demonstrate the computational benefits of approximate sampling in Table 2?
- Is this method applicable to higher-dimensional datasets, such as CelebA-HQ (256x256)?
- Could you provide the potential advantages of this approach compared to other non-adversarial dynamic generative models, such as Flow Matching [1]?

[1] Lipman, Yaron, et al. "Flow matching for generative modeling." ICLR 2023.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730555791732

thread_type:
official_review

<thread_object>
creator: reviewer_xqpy

<content>
summary:
The article proposes a gradient flow procedure for generative modeling by transporting particles from an initial source distribution to a target distribution, where the gradient field on the particles is given by a noise-adaptive Maximum Mean Discrepancy (MMD) divergence. The noise adaptive MMD is trained on a data distribution corrupted by forward diffusion process. The divergence training procedure is related to discriminator training in Generative Adversarial Networks (GAN), but does not require adversarial training. The article demonstrates competitive empirical performance of the method in unconditional image generation on CIFAR10 dataset.

strengths:
The article explores the integration of a diffusion process into models based on Maximum Mean Discrepancy (MMD) gradient flow, offering a theoretical foundation for both the discriminator and the sampling process used in MMD-GAN with this diffusion approach. Additionally, it employs a linear kernel for the scalable MMD-GAN to reduce computational complexity. The authors also conduct experiments using other forms of KL divergence, such as KALE divergence, to demonstrate its effectiveness.

weaknesses:
1. I believe the contribution of this article is inadequate. Previous research has utilized the diffusion process in the discriminator, as noted in this work [1]. However, this article does not provide theoretical proof demonstrating that the MMD GAN can converge to more optimal points when using the diffusion process.

2. Additionally, the effectiveness of MMD Gradient Flow has only been tested on low-resolution datasets, which does not provide sufficient evidence to confirm its overall efficacy. I recommend that the author conduct additional experiments using high-resolution datasets at a resolution of 256x256, specifically on the LSUN and CelebA datasets. These experiments should include evaluations based on the number of metric sampling steps (NFE) and diversity (FID).

[1] Xiao Z, Kreis K, Vahdat A. Tackling the generative learning trilemma with denoising diffusion

questions:
1. I recommend that the author conduct additional experiments using a variety of datasets to demonstrate the effectiveness of the DMMD. The results on CIFAR10 are not as impressive as those achieved by state-of-the-art generative models. I also suggest that the author include more comparative experiments with additional relevant works.

2. I am wondering if using vector Z in Eq.(15) to train a generator will result in better outcomes than using it as a sampling process.

3. I am curious about the time required to train the MMD Gradient Flow on the CIFAR-10 dataset. In my experience, training a diffusion-based discriminator takes significantly longer than training the original models. I find that incorporating a diffusion process in gradient flow-based models can reduce the number of accumulated gradient steps. I would like to know if the sampling steps of the DMMD model are influenced by the diffusion process.

4. I wonder how the sampling time changes with increasing data dimensions, such as in 256x256 high-dimensional datasets.

5. Can the MMD Gradient Flow be applied to other tasks such as Super-Resolution or Inpainting? These tasks also involve particle transportation from an initial source to a target distribution.

6. There are writing errors in the article, particularly in Table 4 and Algorithm 2.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730686257016

thread_type:
official_review

<thread_object>
creator: reviewer_kt31

<content>
summary:
The paper integrates MMD GAN with the diffusion forward process, introducing a novel generative model framework called DMMD. Through the design of a noise-adaptive MMD gradient flow, this framework aims to reduce the challenges of adversarial training and address the singularity issues found in score-based methods.

strengths:
S1 - The paper provides a well-formulated background and problem statement, with a theoretically motivated and well-grounded DMMD framework.

S2 - The idea of using an adversarial training-free discriminator based on the diffusion forward process could offer valuable insights to the community.

weaknesses:
W1 - The framework's absolute performance is a concern, as DMMD shows a significant performance gap compared to DDPM and more modern methods on the selected image generation benchmarks.

W2 - Its broader application potential is limited, with empirical evaluation restricted to small datasets like MNIST and CIFAR.

W3 - The sampling method appears restrictive, requiring reference features from the ground truth dataset to formulate the witness function.

questions:
Q1 - The need for dataset features during sampling seems counterintuitive. Do the authors have any insights into potential solutions for addressing this limitation?
</content>
</thread_object>

</thread>
<thread>
creation date:
1731106759425

thread_type:
official_review

<thread_object>
creator: reviewer_qpwn

<content>
summary:
The paper presents a novel approach, "Diffusion Maximum Mean Discrepancy Gradient Flow" (DMMD), to improve generative modeling by combining Maximum Mean Discrepancy (MMD) with a noise-adaptive gradient flow mechanism. Unlike GANs, DMMD eliminates adversarial training by utilizing a noise-conditional MMD discriminator. DMMD introduces a particle transport technique, adapting MMD as the divergence metric to transport particles from a source distribution to a target distribution.

strengths:
1. The paper is clearly written and well-organized, tackling a genuine problem and effectively presenting its contributions and findings.

2. The paper establishes a solid mathematical foundation, rigorously linking each section to prior research.

3. The results presented are sufficient to validate the theoretical findings and showcase the effectiveness of the proposed approach.

weaknesses:
1. While the paper shows promising results, it is still outperformed by standard diffusion models, especially in terms of FID scores. Further work might be necessary to reach SOTA performance on larger datasets like ImageNet.

2. Related to the previous point, the experimental results are primarily limited to smaller datasets (CIFAR10, MNIST, CELEB-A, and LSUN Church), which may not reflect the potential scalability of DMMD to more complex, high-resolution datasets.

3. Although the method avoids adversarial training, the noise-adaptive MMD flow still introduce complexity, which may limit reproducibility.

questions:
1. As mentioned above, what challenges do you consider in scaling DMMD to larger datasets like ImageNet?
2. What factors limit DMMD's FID performance relative to state-of-the-art diffusion models?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732223286967

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Joint answer

comment:
Please find below a common answer to the occurring points raised during the review process.

# DMMD does not achieve SOTA on Image Generation datasets

Please note that the focus of our work was to demonstrate that it is possible to learn a Wasserstein gradient flow based on a GAN-like discriminator, without adversarial training, and to use it to produce samples with this flow. As it stands now, our method delivers the best empirical results in the class of methods which rely on Wasserstein gradient flows (see our related work and experiments section), and are the first such results to come close to score-based diffusion approaches.

As noted by reviewers, DMMD nonetheless still underperforms compared to the standard diffusion models and these limitations will be the focus for the future work.

State of the art diffusion models have benefited from an intensive search effort over parameters and architectures by the research community, focusing on optimizing their performance on datasets like CIFAR-10. In fact, it took the community 5 years to improve the originally proposed diffusion model [1] to achieve good performance (see DDPM [2]). We believe pushing DMMD’s FID to be closer to diffusion models would require to re-evaluate the design choices made in DMMD (such as architecture, forward noising process, etc) and carefully tune these to the best possible performance on CIFAR-10. We presently borrow settings and architectures  which may not be optimal for DMMD, yet the performance of DMMD is already better than all existing methods employing Wasserstein gradient flows. Further work on tuning parameters and architectures is an important direction for future work.

# Scaling DMMD to high resolution datasets

In our experiments, we considered CIFAR-10 (50k samples) and MNIST (60k samples) to be 32x32 datasets and CELEB-A (162k samples) and LSUN Church (126k samples) to be 64x64 datasets. In principle, scaling to larger image sizes will entail the same increase in complexity as for diffusion models. In fact, in the literature it is common [3-4] to use latent diffusion models for large image sizes, effectively reducing the size of the images to a much smaller, latent dimensionality. The implementation of the Diffusion MMD in latent space for larger image sizes is an interesting topic for future work.

For the purpose of this rebuttal, we scaled DMMD to CELEB-A-HQ (30k examples) of size 128x128. We used the same backbone as in the DDPM paper [4] for the CELEB-A-HQ-256x256 dataset.

We train both DDPM and DMMD on these datasets. Both methods use unconditional sampling with 100 NFEs. For DMMD at sampling time, we use $T=10$ noise levels with $N_{s}=10$ and $\eta=0.001$ gradient flow learning rate and with $N_p=50$ noisy particles. See Alg. 2 for details.

Our results:

* DDPM -- FID=11.67, Inception Score (IS)=2.87
* DMMD -- FID=14.09, IS=2.57

We provide samples in the anonymized link  https://tinyurl.com/iclr2025images.

The results suggest that DMMD still have a performance gap compared to DDPM as observed for lower resolution datasets. Training on yet higher resolution images (256x256) currently presents memory challenges for our training hardware, which we were not able to overcome in time for the rebuttal. We are working to obtain higher resolution results, but we believe that the 128x128 results are already representative of expected performance.

We believe that future improvements in parameter and architecture refinement for DMMD (see our point about SOTA) will translate into improved performance of DMMD in these high resolution datasets. This is left for future work.

We will add these results to the camera-ready version of our paper.

# Computational complexity

We add more details about the computational complexity of DMMD during training and sampling time. We will add this information in the Appendix.

## Training-time complexity

Computing total loss in training iteration (see Alg. 1) on a batch of B samples is $O(B^2)$for an arbitrary kernel and $O(B)$ for a linear kernel (see Appendix B.3). We compute it for $N_n$ noise levels meaning that the cost scales as $O(N_nB^2)$ for an arbitrary kernel and $O(N_nB)$ for a linear kernel. This is more expensive than DDPM [2], which scales as $O(B)$.

## Sampling-time complexity

During sampling time (see Alg. 2), we use $T$ noise levels and do $N_s$ steps per noise level with $N_p$ noisy particles. Let $O(F)$ be the cost of forward pass for the backbone network. The gradient takes $O(2F)$ time to compute. Each sampling step for DDPM costs $O(F)$, whereas for DMMD it costs $O(FN_p^2)$ for an arbitrary kernel, $O(FN_p)$ for linear kernel and $O(F)$ when using approximate sampling procedure.

**References**:

[1] Deep Unsupervised Learning using Nonequilibrium Thermodynamics, Jascha Sohl-Dickstein, Eric A. Weiss, Niru Maheswaranathan, Surya Ganguli

[2] Denoising Diffusion Probabilistic Models, Jonathan Ho, Ajay Jain, Pieter Abbeel

[3] Imagen 3

[4] Dalle-2
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper presents an approach to generative modeling by combining Maximum Mean Discrepancy (MMD) with diffusion models, avoiding the complexities of adversarial training. Reviewers appreciated the paper’s mathematical rigor, particularly its noise-adaptive gradient flow mechanism and efficient approximation of sampling using linear kernels. The framework offers valuable theoretical insights, and its empirical validation on datasets such as CIFAR-10, MNIST, CelebA, and LSUN Church showcases competitive results against some baseline methods. These strengths underline the potential of DMMD in addressing challenges in adversarial training. However, reviewers expressed concerns regarding the framework’s scalability and performance relative to state-of-the-art diffusion models. The DMMD framework has yet to demonstrate comparable results on larger or high-resolution datasets, such as ImageNet, and its computational cost during the generation process remains significant, especially when using non-linear kernels. Experiments primarily focused on lower-resolution datasets, limiting the evidence of scalability to more complex datasets. Additionally, the framework requires access to the training dataset for certain kernels, which might hinder general applicability. These issues highlight the need for further work to validate its robustness and efficiency on high-dimensional data. After the rebuttal phase the work remains borderline. With moderate revisions and a focus on addressing scalability and methodological concerns, the DMMD paper has the potential to make a significant contribution to generative modeling research.