Zero-Shot Text-to-Image Generation

In plain words

Most earlier text-to-image systems were built for one small dataset and used special tricks: extra labels, object masks, or custom losses. This paper asks a simpler question: what happens if you use a very large model and a very large pile of web image-caption pairs?

The method has two parts. First, a small image compressor turns each picture into a short grid of codes picked from a fixed dictionary of 8192 entries. The picture is no longer a million pixel values; it is a 32-by-32 grid of dictionary entries that a decoder can turn back into a recognisable image. Second, one large transformer reads the caption words followed by those grid codes, as if they were one sentence, and learns to predict the next item. To make a picture, you type a caption, let the model keep predicting codes, and decode the result. It draws several candidates and keeps the ones a second model says best match the caption.

The authors report that this works on a dataset it never trained on for those captions, and that human raters often preferred its pictures to a strong earlier model. It is weaker on specialised image sets, and its image-to-image edits are only partly reliable.

TL;DR

• Two stages: a discrete VAE compresses each 256x256 RGB image into 32x32 tokens from an 8192-entry codebook; then a 12-billion-parameter sparse transformer models caption BPE tokens followed by image tokens §2, §2.1, §2.2.
• Training uses about 250 million web image-text pairs; MS-COCO is not in the training set for its captions, though some of its validation images overlap §2.3, §3.1.
• At sampling time 512 candidates are drawn and reranked by a pretrained contrastive model §2.6.
• Author-reported: 90% and 93% majority preference over DF-GAN for realism and caption match on MS-COCO; FID within about 2 points of the best prior method fig. 7, §3.1.
• The 12B transformer and training code are not released; only the dVAE is in the code repository README.md.

Problem

Earlier text-to-image models were trained on small datasets such as MS-COCO and CUB and often needed extra supervision such as part labels or masks §1. Modelling raw pixels with a transformer needs far too much memory, and likelihood training spends capacity on high-frequency detail rather than on the low-frequency structure that makes objects recognisable §2. The paper assumes that a large autoregressive model on a large, noisy dataset can be competitive zero-shot without task-specific design §1, §4.

Method

Overview. The model maximises a two-stage variational bound. A discrete VAE (dVAE) first learns a visual codebook. A transformer then learns the joint prior over caption tokens and image tokens §2, eq. 1. The paper states that this bound holds exactly only for a KL weight of 1, and that larger weights work better in practice §2.

Stage 1: dVAE §2.1.

1. Encode a 256x256 RGB image to a 32x32 grid of categorical logits over K = 8192 codebook entries; this gives a 192-fold reduction in the number of positions §2.1, fig. 1.
2. Train the dVAE on images alone, with a relaxed categorical posterior (Gumbel-softmax) so that gradients pass through the discrete sample. The temperature is annealed; the reconstruction term is a logit-Laplace likelihood §2.1, app. A.3.
3. Decode tokens to images with a dVAE decoder that outputs six maps: three means and three log-scales of the logit-Laplace distribution for R, G and B app. A.3.

Stage 2: prior §2.2.

1. Lowercase the caption and BPE-encode it, keeping at most 256 tokens (vocabulary 16384). During training, apply 10% BPE dropout §2.2, app. B.2.
2. Encode the image with the frozen dVAE encoder and take the argmax over codebook logits, not a sample, giving 32x32 = 1024 tokens (vocabulary 8192) §2.2.
3. Concatenate text tokens and image tokens into one sequence. Text is padded to 256 positions with a separately learned padding token at each position §2.2.
4. Run a decoder-only sparse transformer with 64 layers. Text-to-text attention is causal. Image positions attend to all text tokens. Image-to-image attention uses row, column or convolutional masks §2.2, app. B.1, fig. 11.
5. Compute cross-entropy on text and image tokens, normalised by their counts in the batch. Weight text by 1/8 and image by 7/8 §2.2.

Architecture details of the transformer app. B.1.

• Each image position gets a row embedding and a column embedding added to its token embedding. Text positions get their own embeddings fig. 10.
• Layer i (indexed from 1 to 63) uses a column mask if (i - 2) mod 4 = 0, and a row mask otherwise. The first four layers are row, column, row, row. The last layer uses an 11x11 convolutional mask app. B.1, fig. 11.
• 62 heads with head size 64 give model width d_model = 3968 app. B.1, tab. 2.
• The paper says the model follows the sparse transformer of Child et al. (bibliography key child2019generating) with these changes app. B.1.

Training systems §2.4, §2.5, app. D, app. E.

• Most parameters, Adam moments and activations are stored in 16-bit. Activation checkpointing recomputes activations in the backward pass §2.4.
• Per-resblock gradient scaling replaces standard loss scaling. There is one scale per resblock (128 in total). The scales are updated as described in the Hyperparameters section §2.4, app. D, fig. 4.
• Gradient all-reduce is compressed with PowerSGD at rank 896 in total. Each machine has 8 GPUs and parameters are sharded across them. The right factor Q is a fixed random Gaussian matrix, not updated §2.5, app. B.2, app. E.2, tab. 1.
• Total training used 1024 V100 GPUs (128 machines) app. B.2, app. E.2.

Sampling and reranking §2.6. Draw N candidate token grids from the transformer at temperature 1 (no temperature reduction for the main results), decode each with the dVAE decoder, score each (caption, image) pair with a pretrained contrastive model, and keep the top k. The main results use N = 512. The paper does not state k for the main results §2.6, §3.1, fig. 6.

# Stage 2 training step  [§2.2] [app. B.2]
# Input:  image (256, 256, 3); caption (string)
# Output: scalar loss

def training_step(image, caption):
    # Augment the image with random crops and a resize  [app. B.2]
    image = augment(image)
    # Encode to 32 * 32 image tokens by taking the argmax  [§2.2]
    image_tokens = argmax(dvae_encoder(image), axis=channels)   # (1024,)
    # Lowercase, BPE-encode with 10% dropout, keep <= 256 tokens  [§2.2]
    text_tokens = bpe_encode_with_dropout(lowercase(caption), max_len=256)
    # Concatenate into one stream and pad text to 256 positions  [§2.2]
    stream = concat(text_tokens_padded_to_256, image_tokens)     # (1280,)
    # Predict each token from the tokens before it  [§2.2] [app. B.1]
    logits = transformer(stream)
    # Cross-entropy, text weighted 1/8 and image weighted 7/8  [§2.2]
    loss = (1/8) * ce(text_logits, text_targets)
    loss = loss + (7/8) * ce(image_logits, image_targets)
    return loss

# Sampling with contrastive reranking  [§2.6] [app. A.3]
# Input:  caption; N candidates; k to keep
# Output: top-k images

def generate(caption, num_candidates, top_k):
    # Build the text prefix as in training  [§2.2]
    text_tokens = bpe_encode(lowercase(caption), max_len=256)
    # Draw candidate image-token grids at temperature 1  [§2.6]
    grids = [sample_image_tokens(text_tokens) for _ in range(num_candidates)]
    # Decode each grid to pixels: sigmoid of the means, then un-map  [app. A.3]
    images = [dvae_decode_to_pixels(grid) for grid in grids]
    # Score each (caption, image) pair with the contrastive model  [§2.6]
    scores = [contrastive_score(caption, image) for image in images]
    # Keep the k best; k is not stated for the main results  [§2.6]
    return best_k(images, scores, k=top_k)

Key equations:

\ln p_{\theta,\psi}(x,y) \geq \mathbb{E}_{z \sim q_\phi(z\mid x)}\big[\ln p_\theta(x\mid y,z) - \beta\,D_{\mathrm{KL}}(q_\phi(y,z\mid x),\,p_\psi(y,z))\big] eq. 1

f(x\mid\mu,b) = \frac{1}{2b\,x(1-x)}\exp\!\Big(-\frac{|\operatorname{logit}(x)-\mu|}{b}\Big) eq. 2

\varphi(x) = \frac{1-2\epsilon}{255}\,x + \epsilon, \qquad \epsilon = 0.1 eq. 3

Hyperparameters

Results

Author-reported. Setup: MS-COCO captions (zero-shot for the model; 30,000-caption validation subset for the FID and IS curves), CUB test captions, baseline DF-GAN §3.1, fig. 7, fig. 9. Hardware is in the Hyperparameters section.

Overlap removal did not change the results significantly; dashed lines in fig. 9 are the same models scored against overlap-free references §3.2.

Paper vs code

The official code release at code_commit covers only the dVAE, its weight loader and a usage notebook. The transformer, the dVAE training code, the contrastive reranker and the dataset are not in the repository README.md.

Checked and consistent:

• First encoder convolution is 7x7 app. A.1 encoder.py.
• Last encoder convolution is 1x1 with 8192 output channels app. A.1 encoder.py.
• Decoder first convolution is 1x1 from the vocabulary size app. A.1 decoder.py.
• Three max-pool layers in the encoder app. A.1 encoder.py encoder.py encoder.py.
• Nearest-neighbour upsampling in the decoder app. A.1 decoder.py decoder.py decoder.py.
• Decoder output has 6 channels (3 means, 3 log-scales) app. A.3 decoder.py.
• Pixel mapping with epsilon 0.1 app. A.3 utils.py utils.py.
• Tokens from argmax over logits, not sampled §2.2 usage.ipynb.
• Reconstruction as sigmoid of the mean, then un-mapping app. A.3 usage.ipynb.

Code-only details not stated in the paper:

• Residual gain 1/(n_layers^2), where n_layers = 4 groups x 2 blocks = 8 encoder.py decoder.py.
• Weight initialisation with standard deviation 1 / sqrt(n_in * k^2) utils.py.
• Encoder's final and decoder's first 1x1 convolutions run in 32-bit on GPU; others may use 16-bit encoder.py decoder.py utils.py.
• Default block layout (4 groups x 2 blocks, width 256) encoder.py. The released weights are pickled module objects loaded with torch.load [code: dall_e/init.py]; the code cannot show that these defaults match them.

Not in code, so not verifiable here: the relaxed ELB, the gumbel-softmax and temperature schedule, the KL and step-size schedules, all transformer code, the sparse attention masks, the per-resblock gradient scaling, PowerSGD, the reranking model, the data pipeline and filters, and the human-evaluation setup.

No discrepancies found in the parts that can be checked.

Pitfalls

• Transformer not released. Nothing in the code can produce a caption-to-image sample. An implementer must write the 64-layer sparse transformer, its masks and its embeddings from the text and figures alone. The paper points to Child et al. for the base attention scheme, so details such as the exact mask extents for a 32x32 grid are not fully specified here app. B.1 fig. 11.
• Row-mask extent for the real grid is derived, not stated. The illustration uses a 4x4 grid with an extent of 5. Extending the same rule to 32 columns gives 33 previous tokens; this is an inference from the caption of fig. 11, not a number the paper gives.
• How text and image vocabularies share one stream is not stated. The paper gives two vocabulary sizes, 16384 and 8192, but not whether they share one embedding table or use index offsets §2.2.
• Padding is learned per position. Caption padding positions get their own learned tokens, one per text position up to 256. Positions after a short caption are not masked out §2.2.
• Argmax, not sampling, when tokenising. The paper notes that the ELB strictly calls for sampling from the encoder's categorical distribution. The authors used argmax, reporting a regularising effect only in preliminary ImageNet runs §2.2 fn. usage.ipynb.
• ELB is only an exact bound at beta = 1. The model uses larger beta values (6.6 in the dVAE), so the training objective is no longer a strict bound §2.1 app. A.2.
• Optimiser mismatch. The main text says Adam for the dVAE §2.1, while the appendix specifies AdamW with weight decay 1e-4 app. A.2. Use the appendix values.
• Temperature schedule matters. Linear annealing of the relaxation temperature reportedly diverged; cosine annealing worked app. A.2. Changing the schedule can break stability.
• 16-bit training was the hard part. The authors report underflow in 16-bit gradients as the main obstacle past one billion parameters. Standard loss scaling was not enough, and per-resblock scaling was used §2.4 app. D fig. 4.
• Skipped updates are routine. About 5% of updates are skipped because a non-finite gradient appeared app. E.2.
• Fixed random Q in PowerSGD. Resampling the right factor at every update caused a large performance drop. Warm-starting was not needed §2.5 fn.
• Gradient clipping only at warm-up. Clipping at norm 4 is applied only during warm-up app. B.2.
• Learning-rate schedule is manual. The rate is halved each time the training loss appears to plateau, five times in total. The plateau test is not defined app. B.2.
• No horizontal flips in transformer augmentation. Flips are excluded because images can contain text. The dVAE augmentation does use flips app. B.2 [listing].
• Reranking is costly. The main results sample 512 candidates per caption and rerank them with a second model. The reranker is not released; the bibliography entry (key radford2021learning, titled "Learning Transferable Visual Models From Natural Language Supervision") is the only identification §2.6 [bibliography].
• Dataset is not released and only partly described. The filters are a subset of those in Conceptual Captions. Deduplication thresholds were chosen by hand to minimise false negatives app. C §3.2.
• Overlap with evaluation images. About 21% of MS-COCO validation images overlap the training data by the authors' de-duplication. The paper says removing them changed results only slightly §3.1 §3.2.
• Weak where the data is narrow. The gap on CUB is about 40 FID points worse than the best prior method, which the authors attribute to specialised distributions §3.1.
• Image-to-image and text rendering are qualitative. Image-to-image translation is described as rudimentary and inconsistent; variable binding fails sometimes. No metric is reported for these §3.3 app. G.
• Human evaluation details. One worker's answers were disqualified for disagreement and fast responses. The rest were kept app. F.
• Code license conditions. The dVAE code is under a "Modified MIT" license with an extra request that users apply the model responsibly and say it was used. It is not a standard SPDX licence. This card copies no code verbatim.

Lineage

• Builds on: discrete latent codes and two-stage training of VQ-VAE and VQ-VAE-2 (bibliography keys oord2017neural, razavi2019generating); sparse transformer (child2019generating); PowerSGD gradient compression (vogels2019powersgd); Gumbel-softmax relaxation (jang2016categorical, maddison2016concrete).
• Compared against: AttnGAN (xu2018attngan), DM-GAN (zhu2019dm), DF-GAN (tao2020df). The paper says DF-GAN had the best prior Inception score and FID on MS-COCO §3.1.
• Reranking follows the earlier contrastive approach in razavi2019generating and radford2021learning, and resembles the auxiliary text-image matching loss of AttnGAN §2.6.