Spatial-Frequency U-Net for Denoising Diffusion Probabilistic Models

TL;DR

This paper introduces SFUNet, a novel architecture that models images in wavelet space for denoising diffusion probabilistic models, leading to higher quality image generation by jointly capturing spatial and frequency domain information.

Contribution

The paper proposes a wavelet space-based U-Net architecture with spatial frequency-aware modules, improving image synthesis quality over pixel-based models.

Findings

Higher quality images on multiple datasets

Effective joint modeling of spatial and frequency information

Compatible with existing DDPM training processes

Abstract

In this paper, we study the denoising diffusion probabilistic model (DDPM) in wavelet space, instead of pixel space, for visual synthesis. Considering the wavelet transform represents the image in spatial and frequency domains, we carefully design a novel architecture SFUNet to effectively capture the correlation for both domains. Specifically, in the standard denoising U-Net for pixel data, we supplement the 2D convolutions and spatial-only attention layers with our spatial frequency-aware convolution and attention modules to jointly model the complementary information from spatial and frequency domains in wavelet data. Our new architecture can be used as a drop-in replacement to the pixel-based network and is compatible with the vanilla DDPM training process. By explicitly modeling the wavelet signals, we find our model is able to generate images with higher quality on CIFAR-10, FFHQ, LSUN-Bedroom, and LSUN-Church datasets, than the pixel-based counterpart.