Improved Diffusion-based Generative Model with Better Adversarial Robustness

TL;DR

This paper addresses distribution mismatch in diffusion probabilistic models by applying adversarial training, improving their robustness and generative quality, supported by theoretical analysis and extensive empirical validation.

Contribution

It introduces a novel application of adversarial training to diffusion models and consistency models to mitigate distribution mismatch, backed by theoretical proofs and empirical results.

Findings

Adversarial training improves diffusion model robustness.

Theoretical analysis links DRO to adversarial training in DPMs.

Empirical results show enhanced generative performance.

Abstract

Diffusion Probabilistic Models (DPMs) have achieved significant success in generative tasks. However, their training and sampling processes suffer from the issue of distribution mismatch. During the denoising process, the input data distributions differ between the training and inference stages, potentially leading to inaccurate data generation. To obviate this, we analyze the training objective of DPMs and theoretically demonstrate that this mismatch can be alleviated through Distributionally Robust Optimization (DRO), which is equivalent to performing robustness-driven Adversarial Training (AT) on DPMs. Furthermore, for the recently proposed Consistency Model (CM), which distills the inference process of the DPM, we prove that its training objective also encounters the mismatch issue. Fortunately, this issue can be mitigated by AT as well. Based on these insights, we propose to conduct efficient AT on both DPM and CM. Finally, extensive empirical studies validate the effectiveness of AT in diffusion-based models. The code is available at https://github.com/kugwzk/AT_Diff.