The second stage of the training pipeline is called "learning the prior," and its core objective is to train a massive autoregressive Transformer to learn the joint probability distribution over text descriptions and image content.
In this stage, the parameters of the dVAE encoder and decoder trained in the first stage are frozen. Text descriptions are BPE-encoded into at most 256 tokens, while images are converted by the frozen dVAE encoder into 32×32=1,024 discrete image tokens. The text tokens and image tokens are then concatenated into a single sequence and fed into a Transformer with sparse attention mechanisms. The model is trained in an autoregressive manner (i.e., predicting the next token based on all preceding tokens), with the loss function being the cross-entropy loss over both text and image tokens, assigning a higher weight to the image portion to emphasize visual fidelity.
Ultimately, this trained Transformer serves as a powerful "prior" model that has mastered the regularities governing how plausible image token sequences are generated given a text condition. At inference time, it can autoregressively generate a complete image token sequence solely from a text input, which is then converted into a pixel image by the frozen dVAE decoder, achieving high-quality zero-shot text-to-image generation.