Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

ACT-Diffusion: Efficient Adversarial Consistency Training for One-step Diffusion Models

Fei Kong, Jinhao Duan, Lichao Sun, Hao Cheng, Renjing Xu, Hengtao Shen, Xiaofeng Zhu, Xiaoshuang Shi, Kaidi Xu

TL;DR

This paper proposes Adversarial Consistency Training (ACT), which directly minimizes the Jensen-Shannon divergence between distributions at each timestep using a discriminator, and enhances generation quality, and convergence.

Abstract

Though diffusion models excel in image generation, their step-by-step denoising leads to slow generation speeds. Consistency training addresses this issue with single-step sampling but often produces lower-quality generations and requires high training costs. In this paper, we show that optimizing consistency training loss minimizes the Wasserstein distance between target and generated distributions. As timestep increases, the upper bound accumulates previous consistency training losses. Therefore, larger batch sizes are needed to reduce both current and accumulated losses. We propose Adversarial Consistency Training (ACT), which directly minimizes the Jensen-Shannon (JS) divergence between distributions at each timestep using a discriminator. Theoretically, ACT enhances generation quality, and convergence. By incorporating a discriminator into the consistency training framework, our method achieves improved FID scores on CIFAR10 and ImageNet 64$\times$64 and LSUN Cat 256$\times$256 datasets, retains zero-shot image inpainting capabilities, and uses less than $1/6$ of the original batch size and fewer than $1/2$ of the model parameters and training steps compared to the baseline method, this leads to a substantial reduction in resource consumption. Our code is available:https://github.com/kong13661/ACT

ACT-Diffusion: Efficient Adversarial Consistency Training for One-step Diffusion Models

TL;DR

This paper proposes Adversarial Consistency Training (ACT), which directly minimizes the Jensen-Shannon divergence between distributions at each timestep using a discriminator, and enhances generation quality, and convergence.

Abstract

Though diffusion models excel in image generation, their step-by-step denoising leads to slow generation speeds. Consistency training addresses this issue with single-step sampling but often produces lower-quality generations and requires high training costs. In this paper, we show that optimizing consistency training loss minimizes the Wasserstein distance between target and generated distributions. As timestep increases, the upper bound accumulates previous consistency training losses. Therefore, larger batch sizes are needed to reduce both current and accumulated losses. We propose Adversarial Consistency Training (ACT), which directly minimizes the Jensen-Shannon (JS) divergence between distributions at each timestep using a discriminator. Theoretically, ACT enhances generation quality, and convergence. By incorporating a discriminator into the consistency training framework, our method achieves improved FID scores on CIFAR10 and ImageNet 6464 and LSUN Cat 256256 datasets, retains zero-shot image inpainting capabilities, and uses less than of the original batch size and fewer than of the model parameters and training steps compared to the baseline method, this leads to a substantial reduction in resource consumption. Our code is available:https://github.com/kong13661/ACT
Paper Structure (24 sections, 2 theorems, 25 equations, 11 figures, 7 tables, 2 algorithms)

This paper contains 24 sections, 2 theorems, 25 equations, 11 figures, 7 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Related works
  3. Method
  4. Preliminary
  5. Score-Based Generative Models
  6. Consistency Training
  7. Generative Adversarial Networks
  8. Analysis the Loss Function
  9. Enhancing Consistency Training with Discriminator
  10. Gradient Penalty Based Adaptive Data Augmentation
  11. Experiments
  12. Generation Performance
  13. Resource Consumption
  14. Ablation Study
  15. Impacts of $\lambda_N$
  16. ...and 9 more sections

Key Result

Theorem 3.1

If the consistency model satisfies the Lipschitz condition: there exists $L>0$ such that for all $\boldsymbol x$, $\boldsymbol y$ and $t$, we have $\Vert \boldsymbol f(\boldsymbol x,t,\boldsymbol \theta)-\boldsymbol f(\boldsymbol y,t,\boldsymbol \theta)\Vert_2\leq L\Vert \boldsymbol x-\boldsymbol y\ where the definition of $p_t$,$\boldsymbol f$, $\mathcal{L}_{CT}^{t_k}$ and $\boldsymbol g$ is cons

Figures (11)

  • Figure 1: Generated samples on ImageNet 64$\times$64 (top two rows) and LSUN Cat 256$\times$256 (the third row).
  • Figure 2: $\mathcal{L}_{gp}$, $\mathcal{L}_{CT}$, and FID of ACT on ImageNet 64x64 ($\lambda_N\equiv 0.3$, an overly large $\lambda_N$ leads to training collapse. Additionally, drastic changes in $\mathcal{L}_{gp}$ closely follow changes in $\mathcal{L}_{CT}$).
  • Figure 3: $\mathcal{L}_{gp}$, $\mathcal{L}_{CT}$, and FID of ACT on CIFAR10 ($\lambda_N\equiv 0.3$, an appropriate $\lambda_N$. In the later stages of training, without data augmentation, $\mathcal{L}_{CT}$, $\mathcal{L}_{gp}$, and FID all show relatively large increases).
  • Figure E1: $\mathcal{L}_{gp}$, $\mathcal{L}_{CT}$, and FID of ACT on ImageNet 64x64 ($w_{mid=0.2}, w=0.6$, a suitable parameter set. Under these parameters, all three metrics demonstrate stability).
  • Figure E2: $\mathcal{L}_{gp}$, $\mathcal{L}_{CT}$, and FID of ACT-Aug on CIFAR10 ($\lambda_N\equiv 0.3$, a suitable parameter set. Under these parameters, all three metrics demonstrate stability).
  • ...and 6 more figures

Theorems & Definitions (4)

  • Theorem 3.1
  • proof
  • Theorem C.1
  • proof