PaperlayerView exact revision on arXiv

Source-linked AI summary

Simple diffusion: End-to-end diffusion for high resolution images

Emiel Hoogeboom, Jonathan Heek, Tim Salimans

arXiv:2301.11093v2cs.CVcs.LGstat.ML

TL;DR

High-resolution pixel-space diffusion is difficult because common latent-space and multistage alternatives add complexity and usually do not support single end-to-end training. The paper introduces simple changes to the noise schedule, architecture scaling, dropout placement, and downsampling, achieving state-of-the-art ImageNet diffusion performance without sampling modifiers and high-quality single-stage text-to-image generation.

  • Problem

    High-resolution diffusion commonly relies on latent spaces or multiple generation stages, creating additional complexity and limiting single end-to-end training setups.

  • Method

    Simple diffusion combines a resolution-adjusted noise schedule, selective low-resolution architecture scaling and dropout, multiscale loss, and downsampling in one end-to-end model.

  • Results

    Simple diffusion achieves state-of-the-art ImageNet FID without sampling modifiers and is the first reported single-stage text-to-image model to generate images of such high visual quality.

  • Takeaways & Limitations

    The paper shows that high-resolution diffusion can retain a single-stage, end-to-end pixel-space design rather than requiring cascades or mixtures of experts.

  • Takeaways & Limitations

    Architecture scaling is evaluated empirically at the 16 × 16 level and may not be the ideal way to scale the network.

Abstract

from arXiv · show

Currently, applying diffusion models in pixel space of high resolution images is difficult. Instead, existing approaches focus on diffusion in lower dimensional spaces (latent diffusion), or have multiple super-resolution levels of generation referred to as cascades. The downside is that these approaches add additional complexity to the diffusion framework. This paper aims to improve denoising diffusion for high resolution images while keeping the model as simple as possible. The paper is centered around the research question: How can one train a standard denoising diffusion models on high resolution images, and still obtain performance comparable to these alternate approaches? The four main findings are: 1) the noise schedule should be adjusted for high resolution images, 2) It is sufficient to scale only a particular part of the architecture, 3) dropout should be added at specific locations in the architecture, and 4) downsampling is an effective strategy to avoid high resolution feature maps. Combining these simple yet effective techniques, we achieve state-of-the-art on image generation among diffusion models without sampling modifiers on ImageNet.

1. Introduction

Diffusion models work well for generation, but high-resolution applications commonly use latent spaces or multistage systems that add complexity. This paper investigates simpler end-to-end pixel-space diffusion and identifies four architectural and training modifications for high-resolution images.

  • Diffusion models generate data by learning a neural-network approximation to the reverse of a noise-adding destruction process.
  • The approach trains a single model end-to-end and generates images directly in full pixel space without cascades or mixtures of experts.
  • High-resolution diffusion typically uses latent spaces, cascades, or mixtures of denoising experts, which add complexity and usually prevent single end-to-end training.
  • Simple diffusion adjusts the noise schedule, scales the U-Net at 16 × 16 resolution, places dropout below the highest-resolution maps, and downsamples at higher resolutions.
  • After distillation, the single-stage model can generate an image in 0.4 seconds.

2. Background: Diffusion Models

Diffusion models define a forward process that progressively adds Gaussian noise and a learned reverse process that denoises it. The background introduces variance-preserving schedules, signal-to-noise ratios, denoising parameterizations, and the standard training loss.

  • A diffusion process gradually destroys information by adding Gaussian noise, while a neural network approximates the reverse denoising process.
  • In a variance-preserving process, α_t decreases and σ_t increases over time while determining the balance between signal and noise.
  • The α-cosine schedule uses α_t = cos(πt/2) and σ_t = sin(πt/2), with signal-to-noise ratio α_t/σ_t governing the schedule.
  • Denoising can estimate x from a noisy state, and the paper uses v prediction because initial experiments found it more reliable at larger resolutions.
  • Training uses the standard epsilon loss, which is related through variational inference to a lower bound on model log-likelihood.

3. Method: simple diffusion

The method adapts standard pixel-space diffusion to high resolutions by adjusting its noise schedule, training loss, and architecture while avoiding high-resolution feature maps.

  • 3.1. Adjusting Noise Schedules: The shifted schedule reduces excessive signal preservation at high resolution, giving the denoiser more time to determine global image structure.Under the standard cosine schedule, global structure is already defined across a wide time range for high-resolution images.
  • 3.1. Adjusting Noise Schedules: High-resolution diffusion schedules are shifted relative to a reference resolution because downsampling increases the signal-to-noise ratio.Average pooling over an s × s window increases α_t/σ_t by s and SNR by s^2, motivating a resolution-dependent schedule shift.
  • 3.2. Multiscale training loss: A multiscale loss evaluates denoising at downsampled resolutions and weights lower-resolution terms more heavily to counter high-frequency-dominated training losses.The authors report quicker convergence especially beyond 256 × 256, while decreasing relative weights as resolution increases.
  • 3.3. Scaling the Architecture: Scaling primarily at the 16 × 16 feature-map resolution improves performance while keeping feature maps and parameters more manageable.The paper motivates lower-resolution scaling through accelerator utilization and memory constraints, and reports that 16 × 16 scaling works empirically.
  • 3.3. Scaling the Architecture: Downsampling at the network input and upsampling at the output avoids computation on the highest resolutions and empirically improves performance while reducing cost.One implementation uses an invertible linear 5/3 wavelet transform to produce lower-resolution frequency responses.
  • 3.3. Scaling the Architecture: The method further replaces convolutional layers with MLP blocks at resolutions already using self-attention, producing a U-Net with a transformer backbone.This U-ViT design combines U-Net structure with transformer blocks at selected resolutions.

4. Related Work

Diffusion models are established generative models for multiple modalities, but high-resolution image generation commonly relies on lower-dimensional latent spaces or multi-stage processes that add complexity and prevent single end-to-end training.

  • Diffusion models learn generation by approximating the reverse of a predefined stochastic information-destruction process with neural networks.
  • Diffusion models have been successfully applied to image, speech, and video generation.

5. Experiments

Experiments show that targeted changes to noise scheduling, regularization, architecture scaling, downsampling, and loss design improve high-resolution diffusion, supporting strong single-stage ImageNet results without sampling modifications.

  • Noise schedule: 7.65 versus 3.76 FID on 256 × 256 ImageNet train data shows that shifting the noise schedule substantially improves performance over the original cosine schedule.The authors generally recommend the shift 64 schedule because shifts toward 64 and 32 perform similarly, with 64 converging slightly better early.
  • Dropout: Dropout from resolutions 64, 32, or 16 performs comparably on ImageNet 128, while no dropout or dropout from the data resolution should be avoided.The 16 × 16 setting performed slightly worse but converged faster in early iterations and was used subsequently.
  • Architecture scaling: Increasing the number of 16 × 16 U-Net modules generally improves performance, although eval FID decreases slightly when scaling from 8 to 12 blocks.The authors suggest the exception may indicate that larger networks require stronger regularization.
  • Downsampling: Downsampling 512-resolution inputs to 256 or 128 produces faster models with better sample quality and avoids excessively large feature maps without performance degradation.The authors suggest that downsampling may help the network optimize for sample quality at high resolution.
  • Comparison with literature: Simple diffusion achieves state-of-the-art FID for class-conditional ImageNet generation without guidance or other sampling modifications and operates at resolutions up to 512 × 512.Larger U-ViT models excel on train FID and Inception Score, while U-Net models perform better on eval FID.
  • Text to image: 9.57 FID@30K at 512 × 512 is reported for text-to-image generation using a single diffusion model trained end-to-end.

6. Conclusion

The paper introduces simple modifications to denoising diffusion that work for high-resolution images and reports state-of-the-art ImageNet FID without sampling modifiers.

  • Simple diffusion achieves state-of-the-art ImageNet performance in FID score without sampling modifiers.The model can also be trained in an end-to-end setup.

A. Additional Background Information on Diffusion Models

This section outlines diffusion training and sampling, including shifted noise schedules, U-ViT denoising, classifier-free guidance, clipping, and optional distillation.

  • Diffusion destroys information by adding noise, then learns the reverse denoising process with a neural network.
  • The shifted cosine log-SNR schedule offsets the standard schedule using image and noise dimensions.Its endpoints must shift consistently when they define architectural embeddings.
  • Sampling starts from Gaussian noise and repeatedly applies denoiser predictions through DDPM sampler steps.The final prediction uses the mean estimate and clips x-space values to the range -1 to 1.
  • Classifier-free guidance is enabled by dropping the conditioning signal during training and recombining conditional and unconditional predictions.The reported guidance scale is expressed as (1 + η).
  • Distillation reduces sampling steps and neural-network evaluations; a distilled U-ViT generates one image in 0.42 seconds on a TPUv4.A batch of 8 images takes 2.00 seconds.

B. Experimental details

The experiments section introduces the specific experimental details, beginning with standard optimizer settings for the U-Net experiments.

  • The experimental details begin by specifying standard optimizer settings for the U-Net experiments.

B.1. U-Net settings

The U-Net experiments specify optimization, architecture, dropout, schedule, training-step, and downsampling settings across ImageNet resolutions.

  • U-Net settings: The default U-Net optimization uses Adam with batch size 512, learning rate 5e-5, 10,000 warmup steps, EMA decay 0.9999, and gradient clipping 1.0.
  • U-Net settings: ImageNet 128 and 256 U-Nets use cosine-shifted schedules with shift 64 and train for 1,500,000 and 2,000,000 steps, respectively.
  • U-Net settings: The configurations use dropout 0.1 from resolution 16 and attention at resolutions 8 and 16.
  • U-Net settings: The ImageNet 512 U-Net uses DWT-2 patching, a cosine-shifted schedule with shift 64, and 2,000,000 training steps.
  • U-Net settings: When downsampling skips high-resolution blocks, the corresponding residual blocks are added to lower-resolution levels to keep their total number constant.

B.2. U-ViT settings

The U-ViT closely follows the U-Net while using transformer blocks in self-attention modules. Its ImageNet settings are largely shared across resolutions, with resolution-dependent patching and v-based loss computation.

  • Architecture: U-ViT replaces convolutional layers in self-attention modules with MLP blocks, forming transformer blocks while retaining U-Net-like connections.Only transformer blocks in the middle use residual connections without skip connections.
  • Optimization: The default U-ViT training configuration uses Adam with a 1e-4 learning rate, 10,000 warmup steps, batch size 2048, and 500,000 training steps.The optimizer settings also include β1=0.9, β2=0.99, weight decay 0.0, EMA decay 0.9999, and gradient clipping at 1.0.
  • Resolution-dependent settings: Architecture settings are almost identical across 128, 256, and 512 resolutions, while patching changes from none to dwt_1 and dwt_2.The loss is computed on v rather than epsilon, with only minor observed performance differences between the two choices.
  • Architecture settings: The 512-resolution architecture uses base_channels=128, emb_channels=1024, channel_multiplier=[1, 2, 4, 16], and 36 transformer blocks with four heads.Transformer dropout is set to 0.2.

B.2.1. PSEUDO-CODE FOR U-VIT MODULES

The U-ViT combines a U-Net downsampling and upsampling path with a transformer middle section. Its transformer blocks apply MLP and self-attention updates, while residual blocks process features and incorporate embeddings and skip connections.

  • MLP block: The MLP block normalizes features, expands the channel dimension by an expansion factor of 4, applies embedding-based scale and shift, and can apply dropout.The dropout rate is controlled by the transformer_dropout configuration.
  • Self-attention: Self-attention projects normalized inputs into queries, keys, and values, computes softmax attention weights, and projects the attended values back to the channel dimension.The attention computation uses multiple heads.
  • Transformer blocks: A transformer block updates its input by adding an MLP block followed by a self-attention block.Both updates are residual additions.
  • Residual block: Each residual block uses normalization, swish activations, convolutions, embedding-conditioned scale and shift, and a residual output addition.The final convolution is initialized with zeros.
  • U-ViT pipeline: The down path applies residual blocks and stores intermediate features, then reshapes the deepest representation into a sequence for repeated transformer blocks with positional embeddings.After the transformer, the representation is reshaped back into spatial feature maps.
  • U-ViT pipeline: The up path upsamples features and applies residual blocks that consume stored skip features from the down path.This preserves the U-Net-style encoder-decoder structure around the transformer middle.

C. Additional Experiments

Additional experiments examine guidance sensitivity, scaling to 1024 resolution, sampling speed, and qualitative samples. The reported results emphasize the effect of downsample loss at 1024 resolution and the trade-off introduced by guidance.

  • Guidance scale: Guidance increases ImageNet IS at relatively small scales but especially raises evaluation FID, while the shifted schedule is sensitive to guidance.The text-to-image model is evaluated using CLIP versus MSCOCO FID30K across guidance scales.
  • Experiments on 1024: 6.06 versus 8.10 train FID is reported for the 1024 model with downsample loss versus without it.The experiment also notes that 4 × 4 patching produces 256-resolution feature maps and increases model cost.
  • Sampling speed: 0.42 seconds is reported for generating one distilled 512 × 512 image on a TPUv4, excluding the text encoder.A batch of eight images takes 2.00 seconds.
  • Qualitative results: Figure 10 presents random, non-cherry-picked 256 × 256 ImageNet samples from the U-ViT.The figure provides qualitative samples rather than a numerical comparison.
Loading 2301.11093v2…