Source-linked AI summary
Generating Images with Sparse Representations
Charlie Nash, Jacob Menick, Sander Dieleman, Peter W. Battaglia
TL;DR
Likelihood-based image generation offers stable training and full-distribution modeling, but image dimensionality makes efficient architectures and sampling difficult. DCTransformer models sparse quantized DCT sequences autoregressively, achieving strong quality and diversity while supporting colorization and super-resolution; complex high-resolution datasets still require substantial computation.
Problem
High-dimensional images make likelihood-based generative models challenging to design and sample efficiently, despite their stable training and emphasis on full data coverage.
Method
DCTransformer converts images into sparse sequences of quantized DCT channel, spatial-location, and coefficient tuples, then autoregressively predicts each sequence element with a Transformer.
Results
DCTransformer achieves strong sample quality and diversity, typically the best recall, and the best spatial sFID scores on three of five datasets, while also supporting colorization and super-resolution.
Takeaways & Limitations
Sparse DCT representations provide a practical basis for neural generative modeling across varied image resolutions and tasks.
Takeaways & Limitations
Complex high-resolution datasets require large models and substantial computational resources, exceeding GAN requirements in particular.
Abstract
from arXiv · showhide
The high dimensionality of images presents architecture and sampling-efficiency challenges for likelihood-based generative models. Previous approaches such as VQ-VAE use deep autoencoders to obtain compact representations, which are more practical as inputs for likelihood-based models. We present an alternative approach, inspired by common image compression methods like JPEG, and convert images to quantized discrete cosine transform (DCT) blocks, which are represented sparsely as a sequence of DCT channel, spatial location, and DCT coefficient triples. We propose a Transformer-based autoregressive architecture, which is trained to sequentially predict the conditional distribution of the next element in such sequences, and which scales effectively to high resolution images. On a range of image datasets, we demonstrate that our approach can generate high quality, diverse images, with sample metric scores competitive with state of the art methods. We additionally show that simple modifications to our method yield effective image colorization and super-resolution models.
1. Introduction
Likelihood-based image generation offers stable training and broad data coverage, but raw-pixel modeling is difficult because images are high-dimensional and autoregressive sampling is costly. DCTransformer addresses this by modeling sparse DCT representations and predicts both where and what content to add.
- Likelihood-based objectives encourage full data-distribution coverage, provide more stable training than adversarial alternatives, and support held-out overfitting detection.
- Raw-pixel likelihood modeling is difficult because typical images contain hundreds of thousands of pixels, making autoregressive conditioning and sampling potentially prohibitive.
- Natural images contain substantial redundancy, motivating compressed representations for generative modeling as well as storage, transmission, and processing.
- DCTransformer models quantized DCT coefficients sparsely as sequences of channel, spatial-location, and coefficient triples rather than modeling pixels directly.
- The autoregressive model sequentially predicts the next sequence element, thereby predicting both where to place content and what content to add.
2. DCT-based sparse image representations
The paper converts images into quantized block-DCT representations and serializes only non-zero coefficients, exploiting frequency structure and spatial sparsity. This representation supports variable-length sequences and progressive reconstruction while reducing representation size.
- 2.1. Block DCT: The two-dimensional DCT transforms zero-centered B×B pixel blocks into coefficients indexed by horizontal and vertical spatial frequencies.The normalizing factor α enforces orthonormality.
- 2.2. Quantization: JPEG-style quantization divides DCT coefficients by a quality-parameterized matrix and rounds them, suppressing high-frequency components more strongly.
- 2.3. Sparsity: Quantized DCT blocks are flattened, reassembled into a 64-channel DCT image, and converted into non-zero channel-position-value tuples.
- 2.3. Sparsity: The sparse representation concatenates luminance and downsampled chroma tuples, then appends a stopping token to mark each variable-length sequence's end.
- 2.3. Sparsity: Low-to-high frequency ordering places low-frequency content first and progressively adds high-frequency content, producing a natural upsampling structure.An alternative luma-before-chroma ordering supports colorization.
- 2.3. Sparsity: Sparse representations are substantially smaller than dense representations for all but the highest DCT quality setting.
3. DCTransformer
DCTransformer autoregressively models sparse DCT coordinate lists with stacked channel, position, and value decoders, while chunked context processing keeps computation manageable for long, variable-resolution sequences.
- Architecture: DCTransformer predicts each sparse DCT tuple autoregressively by ordering channel, spatial position, and quantized value predictions.Samples are decoded from the resulting sparse sequence into images.
- Architecture: Sparse DCT sequences can exceed 100k tuples, making standard self-attention impractical because memory scales quadratically with sequence length.The issue is especially acute for high-resolution datasets.
- Training: Training randomly selects fixed-size target chunks and represents preceding context as a dense DCT image, optionally downsampled before Transformer encoding.This conditions chunk prediction on prior sequence content without processing the entire sequence densely.
- Architecture: The architecture uses separate channel, position, and value decoders, with spatially gathered DCT embeddings supplied to value prediction.Each decoder’s hidden states is projected to logits for its associated categorical distribution.
- Datasets: DCTransformer is demonstrated on datasets including plant leaves at long-side resolution 2048, diabetic retinopathy at 1024, CLEVR at 480, and FFHQ at 1024.Selected samples are presented as examples of achievable image quality.
- Sampling: Sampling repeats fixed-size autoregressive chunk continuations, adding completed values to the input DCT image until a stopping condition is reached.Computation and memory remain constant for each block rather than expanding with sequence length.
4. Experiments
Experiments evaluate DCTransformer with sample-based metrics across standard image-generation benchmarks and illustrate its use for high-resolution generation, upsampling, and colorization.
- Evaluation: Sample-based FID, sFID, precision, and recall are used because likelihoods are not directly comparable across model classes and representations.sFID uses intermediate spatial features to assess spatial distributional similarity.
- Image generation benchmarks: GANs achieve the strongest precision and FID overall, while DCTransformer typically achieves the best recall scores.The reported pattern reflects stronger data coverage for the likelihood-trained model.
- Image generation benchmarks: DCTransformer closes the FID gap substantially on spatial sFID and achieves the best scores on three of five studied datasets.Its samples are described as diverse with realistic textures and structures, but less reliably coherent than GAN samples.
- Diverse datasets: DCTransformer produces high-quality unconditional generations on varied datasets, including plant leaves at long-side resolution 2048.The paper identifies these as the first high-quality unconditional neural-network generations at that resolution to its knowledge.
- Upsampling and colorization: Reordering DCT channels enables class-conditional upsampling by conditioning on low-frequency components and colorization by conditioning on luminance information.The paper presents example figures for both tasks, with colorization trained on OpenImages V4.
5. Related Work
Related work includes autoregressive models over compressed image representations, with DCTransformer differing from VQ-VAE2 through its dynamically sized sparse representation.
- Compressed representations: VQ-VAE2 autoregressively models a lossy-compressed image representation with fewer dimensions than pixels but a fixed size.Its representation uses a hierarchy of latent codes.
- Compressed representations: DCTransformer instead represents images with a sparse sequence whose size varies with image content.This contrasts the fixed-size latent representation used by VQ-VAE2.
- Related approaches: Other approaches reduce image representation complexity through lower bit-depth values, while Subscale Pixel Network performs autoregressive spatial upsampling with Transformers.SPN also subsamples target sequences during training and uses a separate encoder for past context.
6. Discussion
The paper uses sparse DCT-based image representations with a Transformer autoregressive model, achieving strong sample quality and diversity while supporting super-resolution and colorization. High-resolution and complex datasets still require large models and substantial computational resources.
- Discussion: DCTransformer models sparse DCT-based image sequences with a Transformer autoregressive architecture and chunked training for long sequences.The approach addresses sequence-length challenges while retaining likelihood-based modeling.
- Discussion: DCTransformer achieves strong performance on sample quality and diversity benchmarks.The paper also reports support for super-resolution upsampling and colorization tasks.
- Discussion: Good results on complex and high-resolution datasets require large models and substantial computational resources.The paper contrasts this resource demand with GANs, which can achieve high-quality results using fewer computational resources.
A. Quantization matrices
The method follows JPEG-style quantization for DCT pixel blocks, adapting the quantization matrix to image quality, chrominance, and block size.
- A. Quantization matrices: The method uses the Independent JPEG Group’s quality-parameterized quantization matrix Q for DCT pixel blocks.The quality parameter q ranges from 1 to 100.
- A. Quantization matrices: Chrominance components use a separate base matrix Tchroma that applies stronger quantization than the luma matrix.This replaces Tluma for chrominance components.
- A. Quantization matrices: For block sizes other than 8, the base matrices Tluma and Tchroma are resized using nearest-neighbour interpolation.
B. Architecture details
DCTransformer combines a Transformer encoder, three stacked decoders, and architectural modifications intended to improve stability, training speed, inference speed, and memory use.
- B. Architecture details: DCTransformer uses a Transformer encoder for partial DCT images and three stacked decoders for DCT coordinate-list slices.
- B. Architecture details: Layer normalization is placed inside the residual path rather than applied outside it.This placement follows prior Transformer designs cited by the paper.
- B. Architecture details: ReZero multiplies each residual connection by a jointly optimized scalar initialized to zero.The paper reports a small improvement in training speed and stability.
- B. Architecture details: A PAR Transformer-style architecture places a series of fully connected layers after each self-attention layer to improve inference speed and reduce memory use.DCTransformer uses this style in both its encoder and decoders.
- B. Architecture details: Table 3 specifies the architecture configurations used in the main experiments.
C. Training details
Training uses token-based optimization, biased chunk selection, and sequence-length correction to train long DCT sequences efficiently. Experiments also examine scaling with model size on LSUN bedrooms.
- C. Training details: Models are optimized with Adam for a fixed number of tokens, using linear warmup followed by cosine learning-rate decay.Training uses Google Cloud TPUv3 hardware.
- C. Training details: Target chunks are biased toward the beginning of sequences because these chunks contain more low-frequency information and improve sample quality.Uniform chunk sampling would produce an unbiased Monte Carlo gradient estimate, but the biased policy was empirically advantageous.
- C. Training details: Figure 9 reports total bpd and channel, position, value, and sequence-chunk contributions as model size ranges from roughly 3 million to 448 million parameters.The scaling experiment trains DCTransformer models on LSUN bedrooms.
- C. Training details: Chunk-based training can bias unconditional generation toward chunks from shorter sequences.Initial chunks from long sequences are selected less frequently than those from short sequences.
- C. Training details: The method counters sequence-length bias by randomly filtering sequences with probability inversely proportional to sequence length.A maximum filtering sequence length Lmax is used.
D. Model scaling properties
The paper examines how DCTransformer scales with model size and analyzes how different sequence components contribute to bits per dimension (bpd).
- The scaling analysis is motivated by prior findings that Transformer performance improves with model size under sufficient compute and data.The cited passage frames the investigation around whether similar scaling behavior applies to DCTransformer.
- Value predictions contribute most to total bpd, followed by positions, while channels contribute very little.
- Per-chunk total bpd decreases with chunk position, likely because value bpd declines as quantization becomes heavier for higher-frequency components.
E. Additional samples
The paper presents uncurated comparisons with real images and baselines, alongside high-resolution upsampling and colorization samples across several datasets.
- Uncurated samples: Figures 10 and 11 compare uncurated DCTransformer samples with real data and baseline samples on FFHQ and LSUN subsets.Figure 10 uses 1024x1024 resolution for FFHQ, while Figure 11 covers bedrooms, towers, and church-outdoor images.
- Uncurated samples: DCTransformer produces variable-aspect-ratio LSUN samples with 384-pixel long-side resolution, whereas BigGAN and VQ-VAE use fixed 256x256 outputs.The fixed-resolution outputs correspond to resized long-side crops of the input images.
- Configuration: Table 3 lists the model and training hyperparameters, including Transformer encoder and decoder layer specifications.Its layer specification represents self-attention and fully connected layers within Transformer blocks.
- Image restoration: Figure 12 shows three uncurated DCTransformer samples for 8x image upsampling on the ImageNet validation set.The figure places the downsampled input on the left, generated samples in the middle, and the original image on the right.
- Image restoration: Figure 13 shows three uncurated DCTransformer samples for image colorization on the OpenImagesV4 validation set.The grayscale input appears on the left, generated colorizations in the middle, and the original image on the right.