PaperlayerView exact revision on arXiv

Source-linked AI summary

Truncated Diffusion Probabilistic Models and Diffusion-based Adversarial Auto-Encoders

Huangjie Zheng, Pengcheng He, Weizhu Chen, Mingyuan Zhou

arXiv:2202.09671v4stat.MLcs.LG

TL;DR

Diffusion generation can be costly because it uses many forward and reverse steps, while reducing steps can undermine the Gaussian assumptions behind denoising. The paper truncates the forward chain and learns an implicit distribution to initialize reverse diffusion; experiments report faster generation with fewer steps and comparable or improved quality.

  • Problem

    Diffusion models require many steps for Gaussian denoising assumptions to hold, making training and sampling computationally costly; prior step-reduction methods are reported to risk generation quality.

    Source: §1 · pp. 1–2

  • Method

    TDPM truncates the forward diffusion chain and learns an implicit distribution to initialize reverse sampling, while TLDM applies truncation to latent diffusion for text-guided generation.

    Source: §1 · p. 2

  • Results

    Truncation can use fewer diffusion steps without sacrificing or even enhancing generation quality, and TLDM can generate more than 50 images in the time LDM takes one.

    Source: §1 · p. 2

  • Takeaways & Limitations

    A learnable implicit prior can cooperate with diffusion to provide better reverse-process starting samples, while truncation can be adjusted to balance generation quality and efficiency.

    Source: §1 · p. 2

    The likelihood in the Diffusion-AAE formulation is generally intractable, making its first loss term intractable.

    Source: §3.3 · p. 5

Abstract

from arXiv · show

Employing a forward diffusion chain to gradually map the data to a noise distribution, diffusion-based generative models learn how to generate the data by inferring a reverse diffusion chain. However, this approach is slow and costly because it needs many forward and reverse steps. We propose a faster and cheaper approach that adds noise not until the data become pure random noise, but until they reach a hidden noisy data distribution that we can confidently learn. Then, we use fewer reverse steps to generate data by starting from this hidden distribution that is made similar to the noisy data. We reveal that the proposed model can be cast as an adversarial auto-encoder empowered by both the diffusion process and a learnable implicit prior. Experimental results show even with a significantly smaller number of reverse diffusion steps, the proposed truncated diffusion probabilistic models can provide consistent improvements over the non-truncated ones in terms of performance in both unconditional and text-guided image generations.

1 INTRODUCTION

Diffusion models generate images by corrupting data with noise and learning a reverse denoising process, but the many steps needed can be computationally costly. The paper proposes truncating that process and learning an implicit starting distribution, with reported benefits for generation speed and quality.

1 INTRODUCTION

  • The paper situates diffusion models among image-generation methods, reporting that they had achieved strong benchmark results and surpassed many existing generative-model families.The introduction notes applications including data augmentation, image editing, and style transfer.Source: §1 · p. 1
  • Diffusion models corrupt data with noise and train a denoising network to reverse that corruption, but typically require hundreds or thousands of steps because large steps can violate the Gaussian assumption.Prior step-reduction approaches are reported to remain unable to achieve significant speedups without sacrificing generation quality.Source: §1 · pp. 1–2
  • TDPM truncates the forward chain and learns an implicit distribution to initialize reverse diffusion; TLDM applies this approach to latent diffusion for text-to-image generation.The implicit distribution can use a separate generator, which offers more flexibility and can improve quality, or reuse the denoising network with no added parameters and comparable results.Source: §1 · p. 2
  • TDPM has a diffusion-based adversarial auto-encoder formulation with a fixed encoder, learnable decoder, truncated diffusion process, and learnable implicit prior.The paper relates this structure to DDPM's analogy with a VAE using a fixed encoder, learnable decoder, and predefined prior.Source: §1 · p. 2
  • The method can generate images faster with fewer diffusion steps without sacrificing, and sometimes enhancing, generation quality; TLDM can generate more than 50 images in the time LDM takes one.The truncation level can be adapted to data complexity and computational resources.Source: §1 · p. 2

2 PRELIMINARIES ON DIFFUSION MODELS

Diffusion models gradually add Gaussian noise to data, then learn a reverse process to recover samples. Their standard formulation relies on many small steps so reverse transitions are Gaussian-like and the final distribution approaches isotropic Gaussian noise.

2 PRELIMINARIES ON DIFFUSION MODELS

  • The forward process starts from data x0 and repeatedly samples xt from a Gaussian transition with variance βt, gradually adding noise.The transition mean is √(1−βt)xt−1 and its covariance is βtI.Source: §2 · p. 3
  • The variational objective expresses the negative ELBO through a reconstruction term, intermediate reverse-transition KL terms, and a final-time term.The intermediate terms compare q(xt−1 | xt, x0) with the learned reverse distribution pθ(xt−1 | xt).Source: §2 · p. 3
  • With sufficiently many steps and small noise rates, the final forward-chain distribution approaches isotropic Gaussian noise, allowing reverse sampling to start from N(0, I).The reverse distribution at each step is approximated with a neural-network-parameterized Gaussian.Source: §2 · p. 3
  • The standard formulation assumes enough steps for Gaussian reverse transitions and for the forward process to make q(xT | x0) approximately N(0, I), leaving the final ELBO term approximately zero.The KL divergence DKL(q||p) is defined as the expectation under q of log q − log p.Source: §2 · p. 3
  • Reducing the number of steps can make true reverse posteriors non-Gaussian and intractable, challenging Gaussian denoisers, especially when noise steps are large.Xiao et al. use an implicit generator and an adversarial objective, but their models work best with as few as four steps and deteriorate when the step count increases.Source: §2 · pp. 3–4

3 TRUNCATED DIFFUSION AND ADVERSARIAL AUTO-ENCODING

TDPM shortens the forward diffusion chain and learns an implicit distribution to initialize the reverse chain; the paper also frames TDPM as a diffusion-based adversarial auto-encoder. Its objective matches the learned prior to the noisy-data distribution while retaining the DDPM loss terms before truncation.

3 TRUNCATED DIFFUSION AND ADVERSARIAL AUTO-ENCODING

  • TDPM can be viewed as a hierarchical adversarial auto-encoder with a fixed encoder, a learnable implicit prior, and a reverse-diffusion decoder.The encoder is the forward diffusion process, and the decoder has Ttrunc stochastic layers parameterized by θ.Source: §3 · p. 4, §3.3 · p. 5

3.1 MOTIVATION AND TECHNICAL CHALLENGES

  • Truncating the forward chain leaves data partially corrupted, but requires learning the unknown noisy-data distribution at Ttrunc to initialize the reverse chain.The implicit generator approximates this distribution; the retained diffusion rates before Ttrunc remain unchanged.Source: §3.1 · p. 4

3.2 HAND-CRAFTED TDPM OBJECTIVE FUNCTION

  • TDPM adds a divergence term matching pψ(xTtrunc) to q(xTtrunc) and combines it with the DDPM loss terms before Ttrunc.The marginal q(xTtrunc) has a semi-implicit form whose density is often intractable.Source: §3.2 · p. 5

3.3 TDPM AS DIFFUSION-BASED ADVERSARIAL AUTO-ENCODER

  • Diffusion-AAE treats q(xTtrunc | x0) as a fixed encoder and the reverse diffusion chain as a decoder, while learning the prior to match the aggregated posterior.Its loss is upper bounded by the TDPM loss, although its reconstruction-likelihood term is generally intractable.Source: §3.3 · p. 5

3.4 MATCHING THE PRIOR TO AGGREGATED POSTERIOR

  • The prior-matching divergence can be estimated from samples, and the paper reports good performance with either of the explored statistical distances.One explored approach uses a GAN generator and discriminator to distinguish corrupted-data samples from implicit-prior samples.Source: §3.4 · p. 5

3.5 TRAINING ALGORITHM

  • Training separates denoising updates from prior-matching updates, and the implicit generator can share the reverse model’s parameters without clearly hurting empirical performance.A separate generator such as StyleGAN2 is also explored and provides better performance while increasing model size.Source: §3.5 · p. 6

3.6 RELATED WORK

  • The paper refers readers to Appendix B for its detailed discussion of related work.Source: §3.6 · p. 6

4 EXPERIMENTS

Experiments test TDPM with different implicit-generator configurations on unconditional and text-to-image tasks, using FID and Recall to assess image quality and diversity. Across image resolutions and tasks, results indicate that truncated chains can retain generation quality with fewer reverse steps.

4 EXPERIMENTS

  • The experiments use FID and Recall for fidelity and diversity, with two implicit-generator options: reusing the denoising network or using a separate network.Unconditional experiments include CIFAR-10, LSUN-bedroom, and LSUN-Church; text-to-image experiments use CUB-200 and MS-COCO.Source: §4 · p. 6

4.1 EFFICIENCY IN BOTH TRAINING AND SAMPLING

  • On CIFAR-10, TDPM obtains good FID with fewer NFE, while TDPM+ improves FID further and is best at NFE=100; TDPM with reverse steps also has better Recall than the GAN comparisons and FID as good as StyleGAN2-ADA.The table reports unconditional-generation results; Figure 2 separately shows ImageNet-64×64 samples at 5, 50, and 100 NFE.Source: §4.1 · p. 7, Table 1 · p. 7, Figure 2 · p. 7
  • With 4 reverse steps, TDPM has much lower FID than with 0 steps at slightly longer sampling time; increasing steps lowers FID slowly while sampling time rises linearly.At 99 reverse steps, TDPM's FID is better than DDPM with 1000 steps; the authors suggest using 4 or more steps to balance quality and speed.Source: §4.1 · pp. 7–8, Figure 4 · p. 8
  • Fewer reverse steps reduce the diffusion component's training time but increase the implicit prior's training time; at 99 steps, their training times are similar and the whole model trains faster than both GAN and DDPM.Figure 3 marks iterations for t < TTrunc in red and for t = TTrunc in blue.Source: §4.1 · p. 7, Figure 3 · p. 8

4.2 RESULTS ON HIGHER-RESOLUTION AND MORE DIVERSE IMAGE DATASETS

  • On 256 × 256 LSUN-Church and LSUN-Bedroom, TDPM generates images of similar quality with much smaller truncation steps, and therefore runs significantly faster than its backbone models.The comparison evaluates DDPM and ADM backbones; Figure 5 shows implicit-prior samples and their refined outputs.Source: Table 2 · p. 7, §4.2 · p. 8, Figure 5 · p. 8
  • On ImageNet-1K at 64×64, TDPM retains good generation quality with small truncation steps, and a properly trained implicit model can further improve backbone performance.The TDPM+ configuration uses a pre-trained ADM and a StyleGAN-XL-based implicit model.Source: Table 3 · p. 7, Figure 2 · p. 7, §4.2 · p. 8

4.3 TEXT-TO-IMAGE GENERATION

  • TLDM achieves comparable or better performance than LDM with many fewer reverse steps; at NFE=5, sampling is 10 and 50 times faster than LDM with 50 and 250 steps, respectively.At NFE=5, FID is higher, but generated images remain visually appealing.Source: §4.3 · p. 9, Figure 6 · p. 9, Figure 7 · p. 9

5 CONCLUSION

The conclusion presents TDPM as a way to shorten diffusion trajectories by learning an implicit starting distribution for reverse denoising. It reports effectiveness in sample quality and efficiency across synthetic and real image data.

5 CONCLUSION

  • TDPM shortens the diffusion chain by learning an implicit generative model for the otherwise intractable distribution at the truncation point.The implicit model can use the same U-Net as the denoiser, and the paper reports multiple ways to learn this distribution.Source: §5 · p. 9
  • TDPM can be formulated as an adversarial auto-encoder with a learnable implicit prior.Source: §5 · p. 9
  • Experiments on synthetic and real image data demonstrate sample-quality and efficiency effectiveness, with diffusion chains shortened to only a few steps.Source: §5 · p. 9

A PROOF

The proof excerpt invokes the DDPM ELBO derivation in establishing its inequality. The related-work discussion places the paper among approaches to sampling efficiency and hybrid generative models.

A PROOF

  • The proof reduces its task to showing that the first term in equation (11) is no greater than L0 + PTtrunc.The excerpt says the second-to-last equality follows the ELBO derivation of Ho et al. (2020).Source: A PROOF · p. 15

B RELATED WORK

  • Prior sampling-efficiency approaches include distillation, adaptive noise schedules, non-Markovian processes, and improved SDE solvers.Source: B RELATED WORK · p. 15
  • Related hybrid approaches use implicit models at each denoising step or use diffusion to refine VAE outputs; TDPM instead places diffusion on an implicit model.The implicit model in TDPM can use a U-Net or a separate network.Source: B RELATED WORK · p. 15

C DISCUSSION

The discussion describes TDPM as an efficient diffusion-based model while noting risks from misuse and open directions for extending the method. The appendix section announces further algorithm details and complementary results.

C DISCUSSION

  • The discussion describes TDPM’s capacity as competitive with existing diffusion-based models and its efficiency as largely improved.Source: C DISCUSSION · p. 15
  • TDPM’s truncated component could be misused to unlock existing diffusion models’ capacity to generate private data or to generate ill-intentioned images at lower cost.The discussion describes providing intermediate noisy images or fine-tuning early steps as possible ways to unlock that capacity.Source: C DISCUSSION · p. 16
  • The work focuses on finite-time diffusion models; extending TDPM to continuous-time models remains non-trivial, and integration with other chain-shortening strategies and advanced architectures is left for future work.The discussion also notes that TDPM can be built on diffusion models beyond DDPM and that the truncated distribution could use a semi-implicit formulation.Source: C DISCUSSION · p. 16

D ALGORITHM DETAILS AND COMPLEMENTARY RESULTS

D.1 ADDITIONAL ANALYSIS ON THE PARAMETERIZATION OF THE IMPLICIT GENERATOR

The appendix examines shared and separate parameterizations for TDPM’s implicit generator, an alternative conditional-transport objective, and a conditional TDPM formulation. It also reports that conditional transport has no significant performance difference from the GAN objective when the generator is well trained.

D.1 ADDITIONAL ANALYSIS ON THE PARAMETERIZATION OF THE IMPLICIT GENERATOR

  • A shared U-Net generates the truncated noisy image and performs subsequent denoising, but having it play both roles makes training more challenging.As Ttrunc grows, the truncated distribution becomes more noise-like, making noisy-image generation and noise prediction easier for the generator.Source: §D.1 · p. 16
  • A separate generator for the truncated step increases model parameters but offers greater training flexibility, including parallel training or use of a pre-trained model.The experiments use StyleGAN2 for this generator, adding 19M and 28M parameters for 32 × 32 and 256 × 256 images, respectively.Source: §D.1 · p. 16

D.2 ALTERNATIVES OF LEARNING THE IMPLICIT DISTRIBUTION

  • Conditional transport replaces the discriminator with a learned conditional distribution and a critic measuring feature-space point-to-point cost; its performance does not significantly differ from the GAN objective when the generator is well trained.The approach uses the same generator as the GAN-based alternative.Source: §D.2 · p. 17

D.3 CONDITIONAL TRUNCATED DIFFUSION PROBABILISTIC MODELS

  • The appendix derives a conditional version of TDPM and states that its truncated distribution can be fitted with either GAN or conditional transport.Source: §D.3 · p. 17

D.4 ANALYSIS ON TOY EXPERIMENTS

On toy distributions, TDPM fits short diffusion chains by learning the non-Gaussian truncated distribution, whereas DDPM’s pure-noise starting point can leave a gap when few steps are used. Results vary with the toy distribution and the implicit generator’s ability to fit it.

D.4 ANALYSIS ON TOY EXPERIMENTS

  • The toy study evaluates Swiss Roll, Double Moons, 8-modal Gaussian, and 25-modal Gaussian mixtures using equal-weight components and reports forward KL divergence as a quantitative metric.It uses 2,000 samples and illustrates generated samples after 5,000 training epochs.Source: §D.4 · pp. 17–18, Figure 11 · p. 20
  • TDPM-GAN and TDPM-CT fit the truncated Swiss Roll distribution in short chains, while DDPM struggles to recover the data from pure noise when the chain is too short.The comparison covers T = 2 and T = 5; TDPM-CT fits slightly better at length 2, while GAN fitting becomes easier at length 5.Source: Figure 8 · p. 18, §D.4 · p. 18
  • On 25-modal Gaussian data, GAN avoids mode collapse and approximates the distribution better than conditional transport, yielding better final data recovery.The passage also reports that DDPM can recover a distribution similar to the data on 8-modal Gaussian data with five steps because its diffused distribution is more isotropic.Source: §D.4 · p. 18, Figure 10 · p. 19, Figure 9 · p. 19

D.5 ADDITIONAL ABLATION STUDIES

The appendix tests pretrained diffusion backbones, noise schedules, truncation choices, and model efficiency, then describes the experiment architectures. It reports comparable performance with a pretrained DDPM backbone and end-to-end TDPM, and insensitivity to the tested noise schedules.

D.5 ADDITIONAL ABLATION STUDIES

  • TDPM using a separately trained implicit generator with a pretrained DDPM has generation performance comparable to end-to-end TDPM on CIFAR-10.This two-step procedure trains only the implicit model rather than pretraining the DDPM in that procedure.Source: Table 5 · p. 21, §D.5 · p. 21
  • TDPM is not sensitive to the choice between linear and cosine noise schedules on CIFAR-10.Table 6 reports FID values for the two schedules.Source: Table 6 · p. 21, §D.5 · p. 21
  • TDPM adds parameters for a training discriminator that is not used during generation, while the appendix reports sampling-efficiency results and model size.Source: Table 7 · p. 22, §D.5 · p. 22
  • SNR rapidly decays during the first 100 diffusion steps, and the authors suggest noise in that interval may more significantly improve GAN training.For steps greater than 200, SNR changes little and is close to zero; the authors report that GAN training at Ttrunc = 99 is significantly facilitated, while Ttrunc = 4 requires iterations similar to training on clean data.Source: §D.5 · pp. 21–22, Figure 12 · p. 22

D.6.1 MODEL ARCHITECTURE

  • The image-generation backbone follows a DDPM-style U-Net; TDPM+ instead uses a separate StyleGAN2 generator, with a discriminator for the truncated distribution and a navigator for conditional transport.Text-to-image experiments use a pretrained 1.45B LDM and fine-tune only its latent model.Source: §D.6.1 · pp. 22–23

D.6.2 TRAINING CONFIGURATIONS

The configurations cover unconditional and text-to-image datasets, use truncated versions of the backbone diffusion schedules, and specify optimization and evaluation settings. For ImageNet1K-64×64, the implicit model is trained directly at 64×64, and TDPM brings its final performance to a competitive level despite weaker standalone results.

Training configurations

  • Experiments use CIFAR-10, LSUN-bedroom, LSUN-Church, CelebA, and CelebA-HQ for unconditional generation, plus CUB-200 and MS-COCO for text-to-image generation.CIFAR-10 uses 32×32 images; CelebA uses 64×64, and the other listed datasets use 256×256 after center-cropping and resizing.Source: §D.6.2 · p. 23
  • Truncated schedules take the first TTrunc diffusion rates from the backbone schedule; the NFE=50 example uses the first 49 steps.The example starts from the DDPM schedule with β1 = 10−4, βT = 0.02, and T=1000.Source: §D.6.2 · p. 23
  • Training uses Adam with most hyperparameters matching Xiao et al. (2022), and a slightly modified generator learning rate to match Ho et al. (2020).Table 8 is identified as the optimization hyper-parameters table; its values are not given in the supplied passage.Source: §D.6.2 · p. 23, Table 8 · p. 23
  • TLDM uses a base learning rate of 10−5 and batch size 64; TDPM+ with StyleGAN2 follows the original hyperparameters and trains in parallel with the diffusion model.ImageNet1K-64×64 uses StyleGAN-XL with default training hyperparameters and direct 64×64 training instead of progressive growing.Source: §D.6.2 · p. 23
  • Without progressive growing, StyleGAN-XL’s result is clearly worse than the paper-reported FID 1.51 result, but as TDPM’s implicit model, TDPM’s final performance becomes competitive with this result.The passage also mentions images corrupted at TTrunc.Source: §D.6.2 · p. 24
  • For sampling-time evaluation, CIFAR-10 uses batches of 128; FID and recall use 50k samples for CIFAR-10 and LSUN, and 30k for CelebA-HQ and text-to-image datasets.Recall follows Kynkäänniemi et al. (2019); LDM uses classifier-free guidance with scale 1.5, and TDLM uses no DDIM steps.Source: §D.6.2 · p. 24

D.7 ADDITIONAL RESULTS ON UNCONDITIONAL GENERATION

The additional unconditional-generation materials provide a CIFAR-10 comparison table and qualitative samples across several datasets and truncation settings. The text-to-image figures show further TLDM samples at two truncation lengths.

D.7 ADDITIONAL RESULTS ON UNCONDITIONAL GENERATION

  • Table 9 gives a full CIFAR-10 unconditional-generation comparison grouped by sampling-step order; TDPM with NFE=1 is equivalent to a GAN using the DDPM architecture as generator.The caption says best FID and Recall in each group are bold but supplies no values in the provided passage.Source: Table 9 · p. 24
  • Figure 13 presents qualitative unconditional-generation results on CIFAR-10 and CelebA at 64×64.Source: Figure 13 · p. 25
  • Figure 14 compares TDPM samples from pθ(x0) and pθ(xTtrunc) on LSUN-Church at Ttrunc values 99, 49, and 4; NFE equals Ttrunc + 1.Figures 15 and 16 provide analogous TDPM samples for LSUN-Bedroom and CelebA-HQ.Source: Figure 14 · p. 25, Figure 15 · p. 25, Figure 16 · p. 25
  • Figures 17–19 provide analogous qualitative samples from TDPM-CT on LSUN-Church, LSUN-Bedroom, and CelebA-HQ.Source: Figure 17 · p. 26, Figure 18 · p. 26, Figure 19 · p. 26

D.8 ADDITIONAL RESULTS ON TEXT-TO-IMAGE GENERATION

  • Figure 20 shows additional TLDM text-to-image samples for different prompts at Ttrunc=49.Source: Figure 20 · p. 27
  • Figure 21 shows additional TLDM text-to-image samples for different prompts at Ttrunc=4.Source: Figure 21 · p. 28
Loading 2202.09671v4…