Improved Diffusion-based Generative Model with Better Adversarial Robustness
This addresses a key training bottleneck for diffusion-based generative models, offering incremental improvements in robustness and performance.
The paper tackles the distribution mismatch problem in Diffusion Probabilistic Models (DPMs) and Consistency Models (CMs) during training and inference, proposing to use Adversarial Training (AT) based on Distributionally Robust Optimization, which improves generative accuracy and adversarial robustness as validated empirically.
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.