PaperlayerView exact revision on arXiv

Source-linked AI summary

Wavelet Diffusion Models are fast and scalable Image Generators

Hao Phung, Quan Dao, Anh Tran

arXiv:2211.16152v2cs.CVeess.IV

TL;DR

Diffusion models provide high-fidelity images but remain too slow for many real-time uses, even after DiffusionGAN reduces sampling steps. Wavelet Diffusion applies wavelet decomposition at image and feature levels to accelerate processing while preserving quality. Experiments across standard benchmarks report state-of-the-art speed, comparable generation quality, and faster convergence than DDGAN.

  • Problem

    Diffusion models’ slow training and inference, including a persistent speed gap versus GANs, limits their use in large-scale or real-time applications.

  • Method

    Wavelet Diffusion decomposes images and feature maps into frequency subbands, using dimensional reduction and frequency-aware processing to accelerate diffusion.

  • Results

    Wavelet Diffusion achieves state-of-the-art training and inference speed while maintaining high generative quality across CIFAR-10, STL-10, CelebA-HQ, and LSUN-Church.

  • Takeaways & Limitations

    The method closes the speed gap with StyleGAN models while retaining comparable image-generation quality and converging faster than DDGAN.

Abstract

from arXiv · show

Diffusion models are rising as a powerful solution for high-fidelity image generation, which exceeds GANs in quality in many circumstances. However, their slow training and inference speed is a huge bottleneck, blocking them from being used in real-time applications. A recent DiffusionGAN method significantly decreases the models' running time by reducing the number of sampling steps from thousands to several, but their speeds still largely lag behind the GAN counterparts. This paper aims to reduce the speed gap by proposing a novel wavelet-based diffusion scheme. We extract low-and-high frequency components from both image and feature levels via wavelet decomposition and adaptively handle these components for faster processing while maintaining good generation quality. Furthermore, we propose to use a reconstruction term, which effectively boosts the model training convergence. Experimental results on CelebA-HQ, CIFAR-10, LSUN-Church, and STL-10 datasets prove our solution is a stepping-stone to offering real-time and high-fidelity diffusion models. Our code and pre-trained checkpoints are available at \url{https://github.com/VinAIResearch/WaveDiff.git}.

1. Introduction

Diffusion models offer high-quality image generation but remain slower to train and sample than GANs. Wavelet Diffusion targets this gap by processing frequency components at image and feature levels while preserving generative quality.

  • Motivation: Diffusion models can surpass GANs in generation quality but require substantially slower training and inference, limiting real-time adoption.DDPMs may require thousands of sampling steps, while DiffusionGAN remains slower than StyleGAN and converges slowly.
  • Proposed approach: Wavelet Diffusion decomposes inputs into low- and high-frequency subbands at both image and feature levels.The design uses wavelet dimensional reduction for speed and frequency information for visual quality.
  • Proposed approach: The method reduces spatial resolution fourfold at the image level and emphasizes wavelet information across generator blocks with marginal computational overhead.These choices are intended to reduce training and inference time while keeping output quality relatively unchanged.
  • Results: Experiments on CIFAR-10, STL-10, CelebA-HQ, and LSUN-Church report state-of-the-art training and inference speed while maintaining high generative quality.The authors position Wavelet Diffusion as a stepping-stone toward real-time, high-fidelity diffusion models.

2. Related work

Prior diffusion research improves quality, sampling, or convergence, but existing approaches do not balance speed and quality. Wavelet Diffusion addresses this by exploiting frequency sparsity and applying wavelet awareness at image and feature levels.

  • Diffusion models: Diffusion models use sequential denoising and generally face a trade-off between sampling speed and generated-image quality.Prior methods reduce steps, alter the sampling process, or move computation to latent space.
  • DiffusionGAN: DiffusionGAN reduces denoising to a few large steps through adversarial modeling but remains slower than GAN competitors.Its large-step generator models complex multimodal distributions that are difficult to represent with Gaussian transitions.
  • Wavelet-based approaches: Wavelet transforms have been used in compression, representation learning, restoration, style transfer, and GAN-based image generation.These applications exploit spatial and frequency information to support visual modeling.
  • Wavelet-based approaches: Earlier wavelet diffusion methods improve selected aspects but do not balance diffusion-model quality and speed.The cited approaches include multiscale wavelet score modeling and score matching for colorization.
  • Wavelet Diffusion: Wavelet Diffusion uses frequency sparsity and dimensionality reduction for efficient sampling, with frequency awareness at both image and feature levels.The approach is inspired by wavelet-based GANs and targets high-fidelity generation.

3. Background

Traditional diffusion gradually adds noise and learns a reverse denoising process, while DiffusionGAN enables larger sampling steps with adversarial training. Wavelet transforms separate frequency components and reduce spatial dimensions for more efficient processing.

  • Diffusion models: Traditional diffusion adds small amounts of noise over thousands of timesteps, then approximates the reverse process with a parameterized denoising distribution.The reverse model is commonly represented as a Gaussian with learned mean and variance and trained using KL divergence.
  • Diffusion models: The forward diffusion posterior has a closed form in which x_t is distributed around a scaled clean sample with schedule-dependent variance.The variance schedule uses small β_t values, which determine the noise added at timestep t.
  • DiffusionGAN: DiffusionGAN uses a conditional generator and discriminator to model large-step transitions that are no longer well described by Gaussian distributions.The generator receives a noisy sample, latent noise, and timestep, while adversarial training distinguishes real and generated pairs.
  • Wavelet transform: Wavelet transform separates an image into low-frequency approximation and high-frequency horizontal, vertical, and diagonal edge details.The Haar wavelet implements this through discrete wavelet transform and inverse wavelet transform operations.
  • Wavelet transform: Wavelet filters decompose an H × W input into four H/2 × W/2 subbands that can be inverted to reconstruct the original signal.The subbands are LL, LH, HL, and HH, formed by pairwise orthogonal filters.
  • Wavelet transform: In Wavelet Diffusion, the transform is applied to images and feature maps to emphasize high frequencies and reduce spatial dimensions fourfold for efficient sampling.This provides the frequency-aware basis for the proposed acceleration strategy.

4. Method

Wavelet Diffusion performs diffusion denoising in a wavelet representation, reducing spatial computation while retaining high-frequency information. Its wavelet-embedded generator uses frequency-aware components and a reconstruction objective to improve efficient, high-quality generation.

  • Wavelet-based diffusion scheme: The framework decomposes images into four low- and high-frequency wavelet subbands, concatenates them, and performs diffusion denoising in wavelet space.The final image is recovered through inverse wavelet transformation after the denoising steps.
  • Wavelet-based diffusion scheme: Most network layers operate at 4× smaller spatial dimensions, significantly reducing computation while preserving the network width used by DDGAN.
  • Wavelet-based diffusion scheme: At each timestep, the generator uses a noisy wavelet sample, latent vector, and timestep embedding to approximate the original signals, while a discriminator distinguishes real and generated sample pairs.The generated previous-step sample is then drawn from a tractable posterior distribution.
  • Training objective: The generator combines adversarial loss with an L1 reconstruction term between generated and ground-truth clean images to preserve frequency information and wavelet-subband consistency.The reconstruction term is weighted by λ, whose default value is 1.
  • Wavelet-embedded generator: The wavelet-embedded generator follows a UNet with frequency-aware sampling blocks, frequency residual connections, and frequency bottleneck blocks that separately process low- and high-frequency components.Wavelet downsampling maps residual shortcuts to feature pyramids, while inverse wavelet transformation returns bottleneck features to the original space.

5. Experiments

Experiments across CIFAR-10, STL-10, CelebA-HQ, and LSUN-Church evaluate fidelity, diversity, convergence, resource use, and generation speed. Wavelet Diffusion generally combines faster sampling and convergence with competitive image quality, while ablations identify reconstruction and wavelet-generator components as beneficial.

  • Experimental setup: The evaluation covers CIFAR-10, STL-10, CelebA-HQ at 256, 512, and 1024 resolutions, and LSUN-Church.FID measures fidelity, Recall measures diversity, and inference time is averaged over repeated trials.
  • Experimental setup: Table 1 reports comparable parameter counts but lower FLOPs and memory usage than DDGAN.
  • CIFAR-10: 0.08(s) inference time on CIFAR-10 is 2.5× faster than DDGAN and reaches real-time performance alongside StyleGAN2.
  • STL-10: STL-10 achieves FID 12.93 and sampling time 0.38(s), with faster convergence than DDGAN, especially during early epochs.
  • CelebA-HQ: CelebA-HQ 256 achieves FID 5.94 and Recall 0.37 while exceeding 2× faster sampling than DDGAN; at 512, FID is 6.40 versus 8.43 and time is 0.59 versus 1.49.
  • LSUN-Church: LSUN-Church reaches FID 5.06, comparable image quality with GANs, 2× faster inference than DDGAN, and Recall 0.40 exceeding StyleGAN2.
  • Ablation studies: The reconstruction term reduces FID by around 0.6 points to 5.94, while the full wavelet generator reaches FID 5.94 with a small running-speed cost.
  • Running time: Single-image generation reaches up to 1024 × 1024 in 0.1s, while wavelet decomposition separates coarse and detailed information across feature-space scales.

6. Conclusions

The conclusion presents Wavelet Diffusion as a wavelet-based scheme that improves diffusion-model speed and fidelity across benchmarks. It reports state-of-the-art diffusion running speed, comparable image quality, and faster convergence than DDGAN.

  • Wavelet transformations applied to image and feature space improve both image fidelity and sampling speed.
  • The method closes the speed gap with StyleGAN models while maintaining comparable image quality to StyleGAN2 and other diffusion models.
  • Wavelet Diffusion converges faster than the DDGAN baseline, supporting its efficiency for real-time and high-fidelity diffusion research.

A. Sensitivity analysis of training batch size

Training batch size materially affects final FID in the evaluated LSUN-Church and CelebA-HQ settings. Batch size 64 outperforms 128, while the 512-resolution comparison uses a larger batch than DDGAN because retraining was resource-limited.

  • Batch size is evaluated as a critical factor using FID on LSUN-Church and CelebA-HQ 256 × 256.
  • On LSUN-Church, batch size 64 reaches best FID 5.06, 0.6 points lower than batch size 128.
  • On CelebA-HQ 256, batch size 64 reaches minimum FID 5.93, 0.28 points lower than batch size 128 at FID 6.21.
  • For CelebA-HQ 512, the model uses batch size 32 versus DDGAN’s 16, and resource limits prevented retraining under the same batch size.

B. Experimental details

The implementation applies discrete wavelet transforms to images and feature maps within a UNet-like generator and discriminator framework. Configurations are adjusted for wavelet coefficients, with resource-saving batch and layer settings reported across datasets.

  • Discrete wavelet transform and inverse wavelet transform are applied to input images and feature maps in the Wavelet-based Diffusion framework.
  • The generator uses a UNet-like architecture based mainly on NCSN++, with dataset-specific configurations.
  • A100-40GB GPUs allow batch size 4 per GPU for the model instead of DDGAN’s batch size 2.
  • With batch size 64, LSUN-Church reaches minimum FID 5.06; the figure compares training curves for batch sizes 64 and 128.
  • The generator’s time-embedding channels are 4× larger than its base channels, and discriminator depth matches generator depth.

B.3. Training hyper-parameters

The appendix documents tuned hyper-parameters and training configurations across datasets, including GPU allocations and training durations.

  • Hyper-parameter choices are summarized in Table 10 for reproducibility.
  • Most settings match the DDGAN baseline, with differences mainly in epoch counts and allocated GPUs for specific datasets.
  • CIFAR-10 and STL-10 training take 1.6 and 3.6 days on one GPU, respectively.
  • CelebA-HQ 256 and LSUN-Church require 1.1 and 6.8 days on 2 and 4 GPUs, respectively.
  • CelebA-HQ 512 training takes 4.3 days on 8 GPUs, while training time depends mainly on denoising steps, network size, and image resolution.

C. More qualitative results

Additional qualitative samples span CIFAR-10, STL-10, CelebA-HQ, and LSUN-Church, with direct comparisons showing stronger visual results than DDGAN on selected datasets.

  • Additional samples are presented for CIFAR-10, STL-10, CelebA-HQ at three resolutions, and LSUN-Church.
  • On STL-10, Wavelet Diffusion produces better sample quality and plausible object appearances than DDGAN.
  • DDGAN fails to represent object-specific shapes in the STL-10 comparison.
  • On CelebA-HQ 512, Wavelet Diffusion produces clearer details such as eyebrows and wrinkles.

D. More discussion

The discussion explains how wavelet decomposition accelerates computation, motivates the framework’s novelty, and examines design choices affecting quality and speed.

  • Connection between wavelet transformation and speed: Wavelet transformation maps X ∈ R^C×H×W to Y ∈ R^4C×H/2×W/2, reducing spatial dimensions fourfold while preserving network width.
  • Connection between wavelet transformation and speed: The reduced spatial dimensions lower most network computation in FLOPs and allow a shallower network with fewer downsampling steps.
  • Increment novelty: The paper claims first use of wavelets in diffusion models with comprehensive incorporation at both pixel and feature levels.
  • Increment novelty: Feature-level components use low- and high-frequency subbands, while high-frequency components may bypass processing to preserve details.
  • Processing high- and low-frequency subbands on STL-10 produced nearly identical FID scores, 12.96 versus 12.93.
  • Progressive upsampling: A two-level progressive upsampling network reached FID 13.11 on CelebA-HQ 256, compared with 5.94 for the proposed approach.
Loading 2211.16152v2…