PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2311.14097v3cs.CV

TL;DR

Diffusion models generate slowly because they iteratively denoise, while consistency training enables faster sampling but can require large models and batches. ACT uses a discriminator to optimize distribution distance at each timestep, and outperforms consistency training on three image benchmarks with substantially lower training resource requirements.

  • Problem

    Diffusion models’ lengthy iterative sampling limits practicality, while consistency training’s accumulated losses help explain its need for large batches and models.

    Source: §1 · p. 1, §3.2 · p. 5

  • Method

    ACT incorporates a discriminator into consistency training to directly reduce the Jensen-Shannon divergence between generated and target distributions at each timestep.

    Source: §1 · p. 2

  • Results

    ACT significantly surpasses consistency training on CIFAR10, ImageNet 64×64, and LSUN Cat 256×256 using less than 1/6 of its batch size and less than 1/2 of its model parameters and training steps.

    Source: §1 · p. 2

  • Takeaways & Limitations

    ACT retains consistency models’ inpainting capability while reducing resource consumption.

    Source: §5 · p. 8

    The interaction between consistency-training and generator losses remains to be explored, and the conclusion identifies other distribution distances as future directions.

    Source: §5 · p. 8

Abstract

from arXiv · show

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

1. Introduction

Diffusion models generate high-quality images but require lengthy iterative denoising, motivating consistency models that can sample in one or two steps. ACT adds a discriminator to consistency training, guided by analysis of its accumulated loss and distribution-distance bound.

1. Introduction

  • Consistency training can generate data in one or two steps, but approximating the trajectory conditionally yields lower sample quality than consistency distillation.A consistency function maps points along a trajectory to a consistent output; training fits its outputs at adjacent timesteps.Source: §1 · p. 1
  • The paper shows consistency training optimizes an upper bound on the Wasserstein distance, whose accumulated losses and errors increase with timestep and help explain its large-batch requirement.The bound includes current and previous consistency-training losses, as well as distance between sampling distributions and accumulated error.Source: §1 · p. 2
  • ACT significantly surpasses consistency training on CIFAR10, ImageNet 64×64, and LSUN Cat 256×256 while using less than 1/6 of its batch size and less than 1/2 of its model parameters and training steps.The reported setups use 1 RTX 3090 for ACT on CIFAR10, 4 A100s for ImageNet, and 8 A100s for LSUN Cat, versus 8, 64, and 64 A100s for consistency training, respectively.Source: §1 · p. 2
  • ACT incorporates a discriminator to directly reduce the Jensen-Shannon divergence between generated and target distributions at each timestep.The discriminator makes this reduction independent of the model output at the previous timestep.Source: §1 · p. 2

2. Related works

Related work contrasts diffusion models with GANs and describes consistency methods as a route to faster sampling. It also situates the paper among techniques for stabilizing GAN training and prior consistency-model approaches.

2. Related works

  • GAN training challenges include instability and mode collapse; prior stabilizers include spectral normalization, gradient penalties, and differentiable data augmentation.StyleGAN2-ADA adapts differentiable augmentation for small datasets.Source: §2 · p. 2
  • Diffusion models generate samples by reversing a noise process and are described as more stable to train than GANs, though their sampling takes longer.Prior acceleration approaches include DDIM and consistency models.Source: §2 · p. 2
  • Consistency models learn mappings from trajectory points to the endpoint for one-step sampling; consistency training approximates non-conditional trajectories with conditional ones, while distillation obtains trajectories from a score-based model.The related-work discussion also distinguishes prior consistency methods and notes a concurrent approach that adds a discriminator after multi-step iteration.Source: §2 · pp. 2–3

3. Method

The method develops ACT by analyzing the accumulated Wasserstein upper bound of consistency loss and adding a discriminator for direct JS-divergence minimization. It also balances the training losses and adjusts augmentation probability using gradient-penalty values over time.

3.1.1 Score-Based Generative Models

  • Score-based models describe sampling from noise to data as an ODE trajectory when deterministic sampling is used.The process starts from a normal distribution and proceeds toward the sample distribution.Source: §3.1.1 · p. 3

3.1.2 Consistency Training

  • Consistency training fits a function that maps points on a shared ODE trajectory to the same endpoint, supporting faster sampling.Training divides the time interval into discrete steps and uses an ODE solver to obtain trajectory points.Source: §3.1.2 · p. 3

3.1.3 Generative Adversarial Networks

  • The GAN formulation alternates discriminator and generator training, and its minimax objective is equivalent to minimizing JS divergence between generated and target distributions.The generator learns to deceive the discriminator, which distinguishes generated from target samples.Source: §3.1.3 · p. 4

3.2. Analysis the Loss Function

  • Under a Lipschitz condition, minimizing consistency loss reduces an upper bound on Wasserstein distance, which includes accumulated losses from earlier timesteps.The accumulated losses and approximation errors grow with timestep, helping explain why consistency training can require large batches and models.Source: §3.2 · pp. 4–5

3.3. Enhancing Consistency Training with Discriminator

  • ACT adds a timestep-conditioned discriminator to directly minimize JS divergence between generated and target distributions at each timestep, without dependence on previous-timestep loss.The discriminator loss is equivalent to minimizing the JS divergence between the generated distribution at that timestep and the target distribution.Source: §3.3 · p. 5
  • When consistency loss is large, its optimization direction can conflict with the discriminator loss; ACT balances them and increases discriminator-loss weight as consistency loss rises.The paper says the conflict is possible when consistency loss is relatively large and describes the weighting scheme in terms of timestep.Source: §3.3 · p. 5

3.4. Gradient Penalty Based Adaptive Data Augmentation

  • For smaller datasets, ACT uses gradient-penalty values to adjust differentiable augmentation probability over time to mitigate discriminator overfitting; performance on larger datasets remains unexplored.The reported augmentation experiments use CIFAR10.Source: §3.4 · pp. 5–6

4. Experiments

Experiments compare ACT with consistency training on CIFAR10, ImageNet 64×64, and LSUN Cat 256×256. ACT improves reported sample quality, while ablations examine λN, training stability, and discriminator design.

4.1. Generation Performance

  • FID improves from 8.7 to 6.0 on CIFAR10, 13 to 10.6 on ImageNet 64×64, and 20.7 to 13.0 on LSUN Cat 256×256.The CIFAR10 table also reports the best FID and IS among one-step diffusion models; the LSUN Cat table compares ACT with consistency training.Source: §4.1 · p. 6, Table 3 · p. 7, Table 4 · p. 7

4.2. Resource Consumption

  • ACT uses 27.5M parameters on CIFAR10 (41.6M with discriminator) versus consistency training’s 63.8M, and 107M on ImageNet (161M with discriminator) versus 282M.CIFAR10 experiments used 1 RTX 3090 versus 8 A100 GPUs; ImageNet used 4 versus 64 A100 GPUs.Source: §4.2 · pp. 6–7

4.3.1 Impacts of λN

  • Performance initially improves and then declines as λN increases; λN≈1 causes CIFAR10 mode collapse, while λN=0.3 destabilizes ImageNet LCT near 150k steps.On ImageNet, LCT collapses around 170k training steps.Source: §4.3.1 · p. 7

4.3.2 Connection between gradient penalty and training stability

  • On CIFAR10, Lgp increases alongside LCT and FID; ACT-Aug adjusts data-augmentation probability using Lgp to stabilize training.The reported relationship between Lgp and LCT also appears in ImageNet training.Source: §4.3.2 · p. 7

4.3.3 Discriminator

  • On CIFAR10, SiLU sped convergence and improved final results; DDPM outperformed DDPM-res with better stability, and timestep embeddings improved results.The authors suggest LCT supplies an additional gradient direction that may mitigate discriminator overfitting.Source: §4.3.3 · p. 8

5. Conclusion

ACT improves consistency training by directly using Jensen-Shannon divergence to reduce the distance between generated and target distributions. The paper reports better generation quality with lower resource use, while identifying further loss-interaction and divergence choices to explore.

5. Conclusion

  • ACT directly uses Jensen-Shannon divergence to reduce the distance between generated and target distributions.The conclusion describes this as an improvement over consistency training.Source: §5 · p. 8
  • The method reports superior generation quality with less than 1/6 of the original batch size and approximately 1/2 of the original model parameters and training steps.The conclusion also reports retained inpainting capability and smaller resource consumption.Source: §5 · p. 8
  • The interaction between L_CT and L_G remains open for further exploration, as do distance choices beyond JS divergence.The authors identify both directions as future work to improve performance.Source: §5 · p. 8

A. Architecture and Experiment settings

The appendix describes the consistency-model and discriminator architectures, experimental settings, and evaluation metrics. It points to tables for model parameters and setup details.

A. Architecture and Experiment settings

  • The consistency model follows a DDPM-like structure with embeddings adapted to continuous time, while the discriminator uses DDPM downsampling through the mid-block and a final linear layer mapping to R.The discriminator's layers-per-block parameter is 150% of the consistency model's; other parameters are the same.Source: A. Architecture and Experiment settings · p. 11
  • Table A1 lists the UNet2DModel parameters, with unlisted settings using diffusers defaults.The passage identifies the table's scope but does not provide its parameter values.Source: Table A1 · p. 11, A. Architecture and Experiment settings · p. 11
  • Table A2 summarizes the experimental setup; unless otherwise specified, the consistency model and discriminator use the same learning rate.The ProjectedGAN discriminator instead uses learning rate 0.002, with w and wmid set to 0.1.Source: A. Architecture and Experiment settings · p. 11, Table A2 · p. 11
  • Evaluation uses IS, FID, Improved Precision, and Improved Recall, with precision measuring generated-sample fidelity and recall measuring represented real-sample diversity.The passage notes that IS does not incentivize full-distribution or within-class diversity.Source: A. Architecture and Experiment settings · p. 11

B. Details of the Proof for Theorem 3.1

The proof details state why two inequalities used in the theorem hold: one from a Lipschitz condition and another from consistency, Lipschitz continuity, and Euler-solver truncation error.

B. Details of the Proof for Theorem 3.1

C. Conditional Discriminator

The conditional-discriminator theorem considers generated samples conditioned on a diffusion state and states that the optimal generator distribution matches the target conditional distribution. Its proof expresses the objective in integral form and relates the resulting expression to Jensen-Shannon divergence.

C. Conditional Discriminator

  • The theorem considers G(z, x_t, t) with z drawn from a normal distribution conditioned on x_t, alongside x_t from the diffusion marginal and x_0 from the target distribution.Source: C. Conditional Discriminator · p. 12
  • At the optimum of Eq. (11), the generated conditional distribution p_g(·|x_t) equals the target conditional distribution p(·|x_t).The proof identifies Jensen-Shannon divergence and states the equality condition for its expression.Source: C. Conditional Discriminator · p. 12
  • The proof rewrites Eq. (11) in integral form, identifies the optimal discriminator, and substitutes it into the objective.Source: C. Conditional Discriminator · p. 12

D. ACT-Aug

This section introduces ACT-Aug, distinguishes it from ACT, and points to its algorithm.

D. ACT-Aug

  • ACT-Aug’s differences from ACT are highlighted in red.The section directs readers to Algorithm 2 for the full procedure.Source: D. ACT-Aug · p. 12
  • Algorithm 2 lists the ACT-Aug procedure.Source: D. ACT-Aug · p. 12

E. More Experiment Results

The experiments report that ACT retains zero-shot inpainting and conditional-trajectory properties. Training metrics are stable under suitable parameter choices, while ACT-Aug shows decreasing consistency loss and FID on CIFAR10.

Zero-shot Image Inpainting

ACT-Aug algorithm

  • ACT-Aug trains the consistency model with a weighted combination of consistency and generator losses, and trains a discriminator with adversarial and gradient-penalty losses.Algorithm 2 also includes augmentation in the discriminator and an augmentation-probability update.Source: E. More Experiment Results · p. 13

Generation Visualization on Conditional Trajectory

More samples

  • λN ≈ 1 is associated with mode collapse in failed CIFAR10 generations.Figure E6 shows the failed generations.Source: Figure E6 · p. 14

Examples of proper λN

More samples

Loading 2311.14097v3…