Wavelet Diffusion Models are fast and scalable Image Generators (2211.16152v2)

Abstract: 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}.

• The paper introduces a wavelet-based diffusion scheme using DWT to decompose images, reducing computational costs by a factor of four while retaining high-frequency details.
• It embeds wavelet transformations into the network architecture, enabling efficient upsampling and downsampling and achieving a 2.5x speedup on CIFAR-10 with competitive FID scores.
• A reconstruction term in the loss enhances model convergence and robustness, paving the way for real-time diffusion model applications.

Wavelet Diffusion Models: Efficient Image Generation

The paper "Wavelet Diffusion Models are Fast and Scalable Image Generators" presents a significant advancement in the efficiency of diffusion models for image generation. The authors introduce a novel wavelet-based diffusion framework aimed at reducing the extensive computational demands of traditional diffusion models, thus bridging the speed gap with prominent generative adversarial networks (GANs) such as StyleGAN.

Key Contributions

1. Wavelet-Based Diffusion Scheme: The core of the proposed approach leverages discrete wavelet transform (DWT) to decompose images into low- and high-frequency components. This strategic decomposition serves two purposes: it emphasizes high-frequency details essential for image fidelity while concurrently reducing spatial dimensions by a factor of four, considerably enhancing computational efficiency.
2. Wavelet Embedded Networks: By integrating wavelet transformations at both image and feature levels, the authors design networks capable of maintaining high-quality output with reduced processing times. The generator employs frequency-aware mechanisms, including novel downsampling and upsampling techniques, to preserve crucial frequency information without significantly increasing computational complexity.
3. Reconstruction Enhancements: A reconstruction term is introduced into the training objective, adding robustness and ensuring fidelity in the generated outputs. This term aids the model's convergence, allowing it to learn both low-resolution approximations and the intricate high-frequency components effectively.

Experimental Validation

The authors performed extensive testing across several datasets, namely CIFAR-10, STL-10, CelebA-HQ, and LSUN-Church, demonstrating the model's superior speed without compromising image quality.

• CIFAR-10 Results: The proposed model achieved a 2.5x speedup compared to the fastest previous diffusion approach (DDGAN), maintaining similar Frechet Inception Distance (FID) scores indicative of high visual quality.
• STL-10 and CelebA-HQ: On larger images, the model continued to outperform both in terms of speed and quality, sometimes lowering FID significantly compared to baseline models. The introduction of a wavelet-aware generator showed additional performance gains.
• Real-Time Capabilities: Importantly, the model shows a trajectory toward real-time performance, a crucial factor for embedding diffusion models into interactive applications.

Implications and Future Directions

This work presents a foundational improvement in computational efficiency for diffusion models, making them more viable for real-time applications where rapid generation is essential. The utilization of wavelet transformation offers a generalizable enhancement applicable to various generative tasks, potentially paving the way for further exploration into hybrid models combining frequency-domain analysis with spatial processing in deep learning.

The success of this approach suggests promising avenues for future research, including more sophisticated frequency decomposition techniques and their applications in other domains like video or 3D model generation. Exploring these methods could further diminish current limitations, such as the balance between speed and high-resolution fidelity, solidifying diffusion models' place next to GANs in the domain of high-performance generative models.

Overall, the paper makes a substantial contribution by addressing a critical bottleneck in diffusion models, offering an efficient pathway toward scalable and practical generative systems.