PaperlayerView exact revision on arXiv

Source-linked AI summary

Zero-Shot Text-to-Image Generation

Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, Ilya Sutskever

arXiv:2102.12092v2cs.CVcs.LG

TL;DR

Text-to-image generation has often relied on specialized models trained on relatively small, fixed datasets, while generated samples can still contain severe visual artifacts. This paper trains a large autoregressive transformer on text and image tokens from internet-scale data, achieving strong zero-shot generalization and capabilities beyond its training task.

  • Problem

    Existing text-to-image methods are typically evaluated on small datasets, while generated samples can still exhibit severe artifacts such as object distortion and illogical placement.

  • Method

    The paper trains an autoregressive transformer on text and compressed image tokens as a single stream, using large-scale image-text data.

  • Results

    The model achieves strong zero-shot performance relative to prior domain-specific approaches and supports a broader range of capabilities from a single generative model.

  • Takeaways & Limitations

    The findings suggest that improving generalization with scale may be a useful direction for progress in text-to-image generation.

  • Takeaways & Limitations

    The model performs significantly worse on the specialized CUB dataset, with a nearly 40-point FID gap versus the leading prior approach.

Abstract

from arXiv · show

Text-to-image generation has traditionally focused on finding better modeling assumptions for training on a fixed dataset. These assumptions might involve complex architectures, auxiliary losses, or side information such as object part labels or segmentation masks supplied during training. We describe a simple approach for this task based on a transformer that autoregressively models the text and image tokens as a single stream of data. With sufficient data and scale, our approach is competitive with previous domain-specific models when evaluated in a zero-shot fashion.

1. Introduction

Text-to-image synthesis has advanced through architectural, loss, attention, conditioning, and pretrained-model innovations, but generated samples can still contain severe visual artifacts. This work shows that scaling an autoregressive transformer to 12 billion parameters and 250 million internet-collected image-text pairs yields a flexible, high-fidelity, natural-language-controllable model with strong zero-shot performance.

  • Related work: Prior text-to-image methods improved fidelity through multi-scale generators, attention, auxiliary losses, additional conditioning information, and pretrained discriminative or cross-modal models.These approaches followed early DRAW- and GAN-based systems for caption-conditioned image generation.
  • Limitations: Despite substantial progress, generated samples can still exhibit object distortion, illogical object placement, and unnatural foreground-background blending.These artifacts motivate approaches that can improve visual fidelity and compositional plausibility.
  • Approach: 12-billion parameters and 250 million image-text pairs enable training a flexible, high-fidelity image generator controllable through natural language.The model uses an autoregressive transformer and internet-collected image-text pairs rather than the relatively small datasets typical of prior evaluation.
  • Results: 90% of the time, human evaluators prefer the model’s zero-shot MS-COCO generations over prior work trained on that dataset.The model uses none of the MS-COCO training labels.
  • Emergent capabilities: The single large generative model also exhibits rudimentary image-to-image translation and can combine concepts, anthropomorphize animals, and render text.These capabilities emerge without custom approaches for each task, though concept combination is achieved with varying reliability.

2. Method

The method trains a transformer to autoregressively model text and image tokens as one stream, using a two-stage procedure to avoid the memory and modeling costs of pixel-level image tokens. A dVAE first compresses images into discrete tokens, after which a 12-billion-parameter transformer learns their joint distribution with text.

  • Two-stage training: The method uses two-stage training: a dVAE compresses images, then an autoregressive transformer models the joint text-image token distribution.This addresses the memory burden of using pixels directly as image tokens.
  • Stage 1: dVAE: A 256×256 RGB image becomes a 32 × 32 grid of tokens with 8192 possible values, reducing transformer context size by a factor of 192 without large visual-quality degradation.The compressed representation contains 1024 image tokens.
  • Stage 2: autoregressive transformer: The second stage concatenates up to 256 BPE-encoded caption tokens with 1024 image tokens and trains a 12-billion-parameter sparse transformer on their joint distribution.Captions use a vocabulary of 16,384, while image tokens use a vocabulary of 8192 and are obtained by argmax sampling from encoder logits.
  • Training objective: The model maximizes an evidence lower bound using pθ,ψ(x, y, z) = pθ(x | y, z)pψ(y, z), with qφ encoding images, pθ decoding them, and pψ modeling text and image tokens.In stage one, φ and θ are optimized on images; in stage two, they are fixed while ψ is learned.
  • Stage 2: autoregressive transformer: The decoder-only transformer lets image tokens attend to all text tokens across 64 self-attention layers, while image-to-image attention uses row, column, or convolutional masks.Text-to-text attention uses the standard causal mask.

3. Experiments

The model performs competitively zero-shot on MS-COCO, with strong human-evaluation results and FID within 2 points of the best prior approach, while remaining substantially weaker on specialized CUB data. Qualitative evaluations also show unusual-concept composition, combinatorial generalization, and limited language-controlled image transformations, though these abilities are inconsistent.

  • Quantitative evaluation: 93% of the time, human evaluators judged our model’s sample to better match the caption than DF-GAN, and 90% judged it more realistic.The comparison used a human evaluation similar to Koh et al. (2021).
  • Quantitative evaluation: Within 2 points of the best prior approach, our model’s MS-COCO FID remained competitive despite never being trained on the captions.About 21% of MS-COCO validation images appeared in the filtered YFCC100M training data, but removing them caused no significant change in FID results.
  • Quantitative evaluation: Nearly 40 points separated our model’s CUB FID from the leading prior approach, while removing the dataset’s 12% training-overlap images caused no significant change.The authors speculate that zero-shot generation compares less favorably on specialized distributions such as CUB.
  • Quantitative evaluation: Increasing the MS-COCO reranking sample size clearly improved FID and IS through 32 samples, after which gains diminished.The reranking uses a contrastive model.
  • Qualitative evaluation: The model demonstrated unusual-concept composition, combinatorial generalization, and limited language-controlled image-to-image translation, including transformations such as recoloring, grayscale conversion, flipping, and style transfer.Its composition and variable-binding behavior were inconsistent, sometimes assigning Christmas sweaters incorrectly or producing the wrong animal relationship.

4. Conclusion · A. Details for Discrete VAE · A.1. Architecture

The paper concludes that scaling an autoregressive transformer improves zero-shot generalization and enables broader capabilities in a single text-to-image model. Its dVAE uses convolutional ResNet encoder and decoder architectures with bottleneck-style residual blocks and specified convolutional projections.

  • 4. Conclusion: Scale improves zero-shot performance relative to previous domain-specific approaches.The conclusion identifies improved generalization as a benefit of executing the autoregressive transformer at scale.
  • 4. Conclusion: Scale broadens the range of capabilities emerging from a single generative model.The paper links this broader capability range to improved generalization as a function of scale.
  • 4. Conclusion: The approach models text-to-image generation with an autoregressive transformer executed at scale.The conclusion characterizes the method as a simple approach based on an autoregressive transformer.
  • 4. Conclusion: Improving generalization as a function of scale may drive progress on text-to-image generation.This is presented as a suggestion arising from the paper’s findings.
  • A.1. Architecture: The dVAE encoder and decoder are convolutional ResNets with bottleneck-style resblocks.The architecture follows LeCun et al. (1998) and He et al. (2016).
  • A.1. Architecture: The models primarily use 3 × 3 convolutions, with 1 × 1 skip-connections convolutions when feature-map counts change.The first encoder convolution is 7 × 7.

A.2. Training · A.3. The Logit-Laplace Distribution · B. Details for Transformer

The dVAE training procedure relies on cosine-decayed schedules for optimization stability, while the logit-Laplace distribution matches the bounded range of image pixels. Transformer details specify padding for fixed-length text inputs and additive row-column embeddings for image tokens.

  • A.2. Training: The dVAE uses cosine schedules for several quantities during training, including the KL weight, relaxation temperature, and step size.The dVAE is trained on the same dataset as the transformer with the augmentation code from Listing 1.
  • A.2. Training: The KL weight β increases from 0 to 6.6 over the first 5000 updates.Bowman et al. (2015) use a similar sigmoid-based schedule.
  • A.2. Training: The relaxation temperature τ anneals from 1 to 1/16 over the first 150,000 updates, whereas linear annealing typically led to divergence.The relaxation-temperature and step-size schedules are especially important for stability and successful optimization.
  • A.2. Training: The step size anneals from 1 · 10^-4 to 1.25 · 10^-6 over 1,200,000 updates.Parameters are updated with AdamW using β1 = 0.9, β2 = 0.999, ϵ = 10^-8, and weight decay multiplier 10^-4.
  • A.3. The Logit-Laplace Distribution: Standard ℓ1 and ℓ2 reconstruction objectives correspond to Laplace and Gaussian likelihoods, but both distributions assign likelihood outside the bounded pixel-value interval.This creates a mismatch between the modeled pixel range and the distributions’ support over the entire real line.
  • B. Details for Transformer: For a transformer with maximum text length 6, a caption of length 4 uses 2 padding tokens, and each image vocabulary embedding is summed with row and column embeddings.Each illustrated vector has size dmodel = 3968.
  • A.3. The Logit-Laplace Distribution: The logit-Laplace distribution applies a sigmoid to a Laplace-distributed variable, producing a probability density supported on (0, 1).Its log-density is used as the dVAE reconstruction term.
  • A.3. The Logit-Laplace Distribution: The dVAE decoder outputs six feature maps: three for the RGB µ parameters and three for ln b.Input pixels are transformed from [0, 255] to (ϵ, 1 − ϵ) with ϵ = 0.1; reconstruction uses ϕ^-1(sigmoid(µ)).

B.1. Architecture

The model is a decoder-only sparse transformer with 64 attention layers and 62 heads per layer, using broadcasted row and column embeddings for image-token context. It alternates row and column sparse attention masks, reserving a convolutional mask for the final layer, where it provides a small performance boost.

  • Architecture: The decoder-only sparse transformer uses 64 attention layers, each with 62 attention heads and a per-head state size of 64.Broadcasted row and column embeddings are used for the image-token context.
  • Architecture: The model uses three sparse attention masks: row, column, and convolutional.For layers i ∈ [1, 63], column attention is used when i − 2 mod 4 = 0 and row attention otherwise; the first four layers are row, column, row, row.
  • Architecture: The convolutional attention mask is used only in the final self-attention layer and provides a small performance boost over row and dense causal masks.The comparison applies when the convolutional mask replaces the other masks in the final self-attention layer.

B.2. Training

Training uses augmented image and caption inputs, stabilized optimization techniques, and large-scale distributed training. The model is trained for 430,000 updates with a scheduled learning rate and validation data reserved.

  • Data and stabilization: Image inputs receive augmentation before dVAE encoding, while captions use 10% BPE dropout during BPE encoding.Training also uses per-resblock scaling and gradient compression with total compression rank 896.
  • Optimization: AdamW uses β1 = 0.9, β2 = 0.96, ϵ = 10−8, and weight decay multiplier 4.5 · 10−2.Decompressed gradients are norm-clipped at 4 only during warm-up, and most Adam moments use 16-bit formats.
  • Training schedule: 1024 NVIDIA V100 GPUs train the model with a total batch size of 1024 for 430,000 updates.The step size ramps linearly to 4.5 · 10−4 over 5000 updates, is halved five times after apparent plateaus, and ends 32 times below its initial value.
  • Training schedule: About 606,000 images are reserved for validation, with no observed overfitting reported in the supplied passage.

C. Details for Data Collection

The 12-billion-parameter transformer was trained on 250 million internet-collected text-image pairs assembled from several sources and filtered for caption and image quality. Training preprocessing used random square crops and resizing to a target resolution of 256 without horizontal flips because images may contain text.

  • Dataset construction: The dataset contains 250 million internet-collected text-image pairs, combining Conceptual Captions, Wikipedia pairs, and a filtered YFCC100M subset.It was created at a scale similar to JFT-300M for training the 12-billion-parameter transformer.
  • Filtering: Filters discard captions that are too short, non-English according to cld3, primarily boilerplate, or contain date-photography phrases.Instances with image aspect ratios outside [1/2, 2] are also discarded because square crops could exclude captioned objects.
  • Image preprocessing: Training preprocessing randomly crops images to a full square, resizes them to target_res = 256 with three channels, and applies a final random crop.Horizontal-flip augmentation is omitted because images may contain text.

D. Guidelines for Mixed-Precision Training · E. Details for Distributed Optimization

Training models beyond one billion parameters in 16-bit precision required eliminating underflow, motivating guidelines for gradient scaling, precision placement, and distributed gradient division. Distributed optimization uses PowerSGD compression, excluding selected parameter types and addressing billion-parameter performance regressions.

  • D. Guidelines for Mixed-Precision Training: Training beyond one billion parameters in 16-bit precision required detecting and eliminating underflow across training components.The authors developed mixed-precision guidelines from this process.
  • D. Guidelines for Mixed-Precision Training: The model uses 128 per-resblock gradient scales instead of standard loss scaling, initialized to M · 2^13 and updated according to gradient finiteness.Scales increase by 2^1/1000 when gradients are finite and otherwise divide by 2 and skip the update.
  • D. Guidelines for Mixed-Precision Training: Gradient scales cannot be divided consecutively within 125 updates and are clamped to [M · 2^7, M · 2^24].Figure 12 tracks these scales during early training of a 2.8-billion parameter model.
  • D. Guidelines for Mixed-Precision Training: The guidelines reserve 32-bit precision for gains, biases, embeddings, unembeddings, their gradients and Adam moments, while computing token logits in 32-bit precision.Gradient compression is disabled for these parameters.
  • D. Guidelines for Mixed-Precision Training: Dividing 16-bit gradients by the number of data-parallel workers M before all-reduce can cause underflow, whereas dividing after summation can cause overflow.The issue arises because gradients are commonly computed and stored in 16-bit precision.
  • D. Guidelines for Mixed-Precision Training: Fewer mixed-precision guidelines may be needed on TPUs because bfloat16 provides a larger 8-bit exponent range, making underflow less likely.This is presented as a hardware-dependent limitation.
  • D. Guidelines for Mixed-Precision Training: The solution divides the loss by the overall batch size, multiplies gradient scales by M, and applies a hand-tuned pre-all-reduce division to avoid underflow and overflow.This approach aims to minimize information loss before all-reduce without introducing overflow.
  • E. Details for Distributed Optimization: PowerSGD compresses gradients for all parameters except embeddings, unembeddings, gains, and biases, with adaptations addressing performance regressions at billion-parameter scale.The paper analyzes communication reduction by compression rank and model size and details the PowerSGD modifications in Section E.

E.1. Bandwidth Analysis

The method compresses gradients by factorizing them into two rank-r matrices and transmitting those factors with two smaller all-reduces. In the stated sharding setup, the bandwidth compression ratio is r(m + 2)/(2dm), or 5r/8d when m = 8, while several parameter types remain uncompressed.

  • Compression mechanism: Gradient compression factorizes G ≈ P Q^t with rank-r matrices P and Q, then transmits P and Q^t through two smaller all-reduces.Parameters are sharded along axis 1, except for the second MLP matrix.
  • Compression ratio: r(m + 2)/(2dm) is the compression ratio, reducing to 5r/8d for m = 8 GPUs per machine.The ratio is computed from the data sent between GPUs on different machines.
  • Uncompressed parameters: Embeddings, unembeddings, gains, and biases are not compressed, but their total bandwidth fraction decreases as model size increases.These parameters become a smaller share of total bandwidth for larger models.

E.2. Implementation Details

The implementation combines sharded training, compressed gradient communication, error correction, and carefully scheduled low-rank factorization across resblocks. Compute–communication overlap, operation throttling, and transpose-aware storage are used to improve efficiency and avoid instability or memory exhaustion.

  • Training and gradient compression: Parameter sharding and gradient compression overlap activation recomputation with parameter prefetching, then reduce-scatter each gradient shard while discarding the full gradient matrix.This setup is designed for efficient, reliable execution at billion-parameter scale.
  • Training and gradient compression: Nonfinite reduce-scatter results skip the update, while finite results are rescaled and accumulated in error buffers for subsequent correction.A single nonfinite gradient causes the entire update to be skipped, occurring about 5% of the time.
  • Low-rank factorization and updates: Grouped all-reduces aggregate low-rank P and Q factors, custom kernels clamp infinities, and Householder orthogonalization replaces numerically unstable Gram-Schmidt.The implementation adds ϵI_m×r to P with ϵ = 10^-6 before orthogonalization and computes decompressed gradients PQt in 32-bit precision.
  • Low-rank factorization and updates: After global norm computation, Adam applies updates using decompressed gradients, and finite low-rank factors update error buffers by subtracting machine-averaged decompressed gradients.Uncompressed parameters are reduced and updated separately.
  • Optimizations: Overlapping later resblock operations with current reduce-scatter, throttling prefetches and Adam updates, and using transpose-aware parameter storage improve performance while limiting memory pressure.The implementation avoids unnecessary transposes around NCCL communication by storing selected parameters in transposed formats.

F. Details for Human Evaluation Experiments · G. Zero-Shot Image-to-Image Translation

Human evaluation used controlled caption-based image comparisons on Amazon Mechanical Turk, while zero-shot image-to-image translation appeared unexpectedly without training modifications. The translation capability is illustrated through additional examples.

  • F. Details for Human Evaluation Experiments: 1000 captions generated one sample image per model per caption for each human-evaluation experiment.These captions and samples formed the evaluation tasks.
  • F. Details for Human Evaluation Experiments: 1000 image comparison tasks per experiment were submitted to Amazon’s Mechanical Turk.Each task compared model-generated images associated with a shared caption.
  • F. Details for Human Evaluation Experiments: Each task was answered by five distinct workers.Workers compared two images on realism and caption matching.
  • F. Details for Human Evaluation Experiments: Workers judged which image was most realistic and which best matched the shared caption.The two questions evaluated visual realism and semantic alignment with the caption.
  • F. Details for Human Evaluation Experiments: Other worker answers were retained despite some fast submissions, including many submission times under 4 seconds.Answers were kept unless excluded by the stated comparison with other workers combined with fast answer velocity.
  • G. Zero-Shot Image-to-Image Translation: Figure 14 presents further examples of zero-shot image-to-image translation discussed in Section 3.3.The examples provide additional evidence of the reported translation capability.
  • G. Zero-Shot Image-to-Image Translation: The authors did not anticipate zero-shot image-to-image translation and made no training-procedure modifications to encourage it.The capability emerged without an explicit intervention during training.
Loading 2102.12092v2…