Megapixel Image Generation with Step-Unrolled Denoising Autoencoders

TL;DR

This paper introduces a scalable, efficient framework combining VQ-GAN, hourglass transformers, and SUNDAE autoencoders to generate high-resolution, diverse, and realistic megapixel images rapidly on consumer hardware.

Contribution

It presents a novel combination of techniques and modifications that enable megapixel image generation with fast sampling, high scalability, and flexible capabilities, surpassing prior methods in efficiency and resolution.

Findings

Achieved high-resolution ($1024 imes 1024$) image generation in 2-4 days.

Produced diverse, realistic megapixel samples in about 2 seconds on a GTX 1080Ti.

Obtained competitive FID scores close to state-of-the-art with fewer sampling steps.

Abstract

An ongoing trend in generative modelling research has been to push sample resolutions higher whilst simultaneously reducing computational requirements for training and sampling. We aim to push this trend further via the combination of techniques - each component representing the current pinnacle of efficiency in their respective areas. These include vector-quantized GAN (VQ-GAN), a vector-quantization (VQ) model capable of high levels of lossy - but perceptually insignificant - compression; hourglass transformers, a highly scaleable self-attention model; and step-unrolled denoising autoencoders (SUNDAE), a non-autoregressive (NAR) text generative model. Unexpectedly, our method highlights weaknesses in the original formulation of hourglass transformers when applied to multidimensional data. In light of this, we propose modifications to the resampling mechanism, applicable in any task applying hierarchical transformers to multidimensional data. Additionally, we demonstrate the scalability of SUNDAE to long sequence lengths - four times longer than prior work. Our proposed framework scales to high-resolutions ($1024 \times 1024$) and trains quickly (2-4 days). Crucially, the trained model produces diverse and realistic megapixel samples in approximately 2 seconds on a consumer-grade GPU (GTX 1080Ti). In general, the framework is flexible: supporting an arbitrary number of sampling steps, sample-wise self-stopping, self-correction capabilities, conditional generation, and a NAR formulation that allows for arbitrary inpainting masks. We obtain FID scores of 10.56 on FFHQ256 - close to the original VQ-GAN in less than half the sampling steps - and 21.85 on FFHQ1024 in only 100 sampling steps.