Autoregressive Model Beats Diffusion: Llama for Scalable Image Generation

Introduction

Built upon autoregressive models, large language models (LLMs) [Vaswani et al. 2017; Devlin et al. 2018; Radford et al. 2018; Raffel et al. 2020; Radford et al. 2019; Brown et al. 2020; Zhang et al. 2022] generate the text by predicting the next token in a sequence. This “next-token prediction” paradigm presents unprecedented capabilities in solving language tasks in a human-like conversational manner [Ouyang et al. 2022; OpenAI 2022, 2023b; Google 2023; Anthropic 2023; Workshop et al. 2022; Touvron et al. 2023a, b; Bai et al. 2023a; Yang et al. 2023; Team 2023; Bi et al. 2024] and incredible scalability [Kaplan et al. 2020; Henighan et al. 2020; Hoffmann et al. 2022; Wei et al. 2022; Alabdulmohsin et al. 2022; Chowdhery et al. 2023; Anil et al. 2023], demonstrating a promising path toward general-purpose artificial intelligence models.

Witnessed the scalability of autoregressive models on large language models, pioneering works attempt to explore autoregressive models in image generation, for example, VQVAE [Van Den Oord et al. 2017; Razavi et al. 2019], VQGAN [Esser et al. 2021; Lee et al. 2022], DALL-E [Ramesh et al. 2021], Parti [Yu et al. 2021, 2022]. They introduce image tokenizers to convert continuous images to discrete tokens, and apply autoregressive models to generate image tokens in the way of next-token prediction. They demonstrate strong performance among their contemporaries [Brock et al. 2018; Ho et al. 2020; Dhariwal & Nichol 2021] in the year before 2022. However, their open-source communities are not well developed, which largely limits their further improvements.

At the same period, another image generation method, diffusion models [Song & Ermon 2019; Ho et al. 2020; Song et al. 2020; Dhariwal & Nichol 2021; Nichol et al. 2021; Lu et al. 2022a; Ho et al. 2022a; Ho & Salimans 2022; Rombach et al. 2022; Ramesh et al. 2022; Saharia et al. 2022; Rombach et al. 2022] develop rapidly. Along with their open-source communities, they dominate the field of visual generation up to today. However, diffusion models share distinct paradigms with autoregressive language models, which poses a huge challenge to building a unified model between language and vision.

In this work, we are committed to pushing the envelope of autoregressive models on image generation further: continuing its research methodology and contributing to open-source community. Reviewing the literature on image generation in the year before 2024, we identify three keys to existing advanced models [Peebles & Xie 2023; Podell et al. 2023; Xue et al. 2023; Chen et al. 2023b, c; Betker et al. 2023; Li et al. 2024; Esser et al. 2024]: 1) well-designed image compressors, 2) scalable image generation models and 3) high-quality training data. Motivated by this, we reexamine the designs of image tokenizers (image compressors for autoregressive models), the scalability properties of image generation models, and the effects of training data.

Towards a potential unified model between language and vision, our design is reducing the inductive biases on visual signals and adopting the same architecture as LLM. This belongs to a different research philosophy with recent works [Chang et al. 2022; Yu et al. 2023b; Tian et al. 2024] that modify the architectures under the guidance of vision-oriented designs. For example, MaskGIT [Chang et al. 2022], MAGVIT [Yu et al. 2023a, b] adopt the masked image modeling strategy, VAR [Tian et al. 2024] uses hierarchical multi-scale property. Although they have succeeded in achieving leading image generation performance, and even better than diffusion models, it is still not clear whether the original language model architectures are capable of this. Instead, our work reveals that vanilla autoregressive models that apply the exactly same “next-token prediction” as language models are also able to achieve state-of-the-art image generation performance. As a bonus, we can leverage the techniques [Dao et al. 2022; Rasley et al. 2020; Shoeybi et al. 2019; Zhao et al. 2023; Kwon et al. 2023; Chen et al. 2023a; Dettmers 2022] developed in LLM community to optimize the training recipes and inference speeds of our models.

In summary, our contributions to the community include:

Image tokenizer: An image tokenizer with downsample ratio of 16, achieves reconstruction quality of 0.94 rFID and codebook usage of 97% on ImageNet benchmark. With the downsample ratio of 8, our tokenizer is competitive or even better than continuous VAE [Rombach et al. 2022; Podell et al. 2023; OpenAI 2023a] used in diffusion models. This shows that discrete representation in image tokenizers is no longer the bottleneck of the image reconstruction.

Scalable image generation model: A series of class-conditional image generation models, ranging from 111M to 3.1B parameters, are developed based on Llama architecture [Touvron et al. 2023a, b]. The largest model realizes 2.18 FID on ImageNet 256$\times$256 benchmarks, outperforming the popular diffusion models such as LDM [Rombach et al. 2022], DiT [Peebles & Xie 2023]. This shows that vanilla autoregressive models without inductive biases on visual signals can serve as the basis of image generation systems.

Hiqh-quality training data: A text-conditional image generation model with 775M parameters, is firstly trained on a 50M subset of LAION-COCO [LAION 2022] and then fine-tuned on 10M internal high aesthetics quality images. It demonstrates competitive performance of visual quality and text alignment.

Optimized inference speed: We adopt vLLM [Kwon et al. 2023], one of the most popular LLM serving frameworks, to optimize the inference speed of our image generation models, and remarkable 326% - 414% speedup is achieved.

We release all models and codes to facilitate the open-source community of visual generation and multimodal foundation models. It is worth noticing that our released models are still behind state-of-the-art visual generation models based on diffusion models [Alpha-VLLM 2024; Esser et al. 2024; Brooks et al. 2024]. When more training data and computation resources are available in the future, large-scale AR-based visual generation models, e.g., above 7B parameters, will be explored.

Autoregressive Models for Image Generation

During image generation, image tokens $(q_{1},q_{2},\dots,q_{h\cdot w})$ are generated by autoregressive models [Radford et al. 2018, 2019; Brown et al. 2020; Touvron et al. 2023a] in the way of next-token prediction $\prod_{t=1}^{h\cdot w}p(q_{t}\mid q_{<t},c)$, where $c$ is class label embedding or text embedding. Finally, these image tokens are converted to image pixels by the image tokenizer decoder.

2 Image Tokenizer

Training losses.

Since quantization is a non-differentiable operation, a straight-through gradient estimator [Bengio et al. 2013] is used to preserve the gradient from the decoder to the encoder $z=\text{sg}[z-f]+f$, $\text{sg}[\cdot]$ is stop-gradient operation. For codebook learning, $\mathcal{L}_{\text{VQ}}=\|\text{sg}[f]-z\|_{2}^{2}+\beta\|f-\text{sg}[z]\|_{2}^{2}$, where the second term is commitment loss [Van Den Oord et al. 2017] to force feature vectors extracted from the encoder to be close to codebook vectors, $\beta$ is commitment loss weight. For simplicity, we don’t add entropy loss [Yu et al. 2023a; Chang et al. 2022] in codebook learning.

3 Image Generation by Autoregressive Models

Our model architecture is largely based on Llama [Touvron et al. 2023a, b], applying pre-normalization using RMSNorm [Zhang & Sennrich 2019], SwiGLU activation function [Shazeer 2020], and rotary positional embeddings [Su et al. 2024]. Specifically, we use 2D RoPE in at each layer of our model, following the implementation of [Lu et al. 2023; Fang et al. 2023]. We do not use the technique of AdaLN [Peebles & Xie 2023] to keep our structure the same as LLM.

Class-conditional image generation.

The class embedding is indexed from a set of learnable embeddings [Peebles & Xie 2023; Esser et al. 2021] and is used as the prefilling token embedding. Starting from this token embedding, the model generates the sequence of image tokens by next-token prediction way, and stops at the location of the pre-defined maximum length.

Text-conditional image generation.

To integrate the text condition into autoregressive models, we use FLAN-T5 XL [Chung et al. 2024] as the text encoder, the encoded text feature is projected by an additional MLP [Chen et al. 2023b, c] and is used as prefilling token embedding in autoregressive models. We note that this design is not an ultimate design for multimodal foundation models, where a unified vocabulary is established between language and vision [Lu et al. 2023; Team et al. 2023]. We leave it for future research.

Classifier-free guidance.

It is worth noting that all design choices discussed so far are largely inspired by previous works, for example, image tokenizer is borrowed from [Rombach et al. 2022; Yu et al. 2021], image generation is from [Peebles & Xie 2023; Chen et al. 2023b; Esser et al. 2021]. A large portion of these techniques are well studied in diffusion models but little in AR models. Our work adapts these advanced designs collectively to AR-based visual generation models.

4 Scale Up

Our model architecture is almost the same as Llama, which allows us to seamlessly adopt optimization techniques [Zhang & Sennrich 2019; Shazeer 2020; Su et al. 2024] and training recipes [Dao et al. 2022; Rasley et al. 2020; Shoeybi et al. 2019] in LLM community. As shown in Table 1, we scale the model size up to 3.1B parameters in this work. All models are implemented with PyTorch 2 [Ansel et al. 2024] and trained on 80GB A100 GPUs. For training the models with parameters below 1.4B, we directly use DDP, otherwise, we adopt PyTorch FSDP [Zhao et al. 2023] to optimize GPU memory usage.

5 Serving

Autoregressive models have always suffered from its low inference speed. With the rapid development of large language models, advanced inference techniques [Kwon et al. 2023; Chen et al. 2023a; Dettmers 2022] are proposed in the LLM community to optimize the inference speed.

Similar to training, inference techniques developed in the LLM community can also be adopted to optimize our models. We verify the effectiveness of vLLM [Kwon et al. 2023], one of the most popular LLM serving frameworks, on our image generation methods. As shown in Table 7, 326% - 414% speedup is achieved compared to the baseline setting.

Experiments

The training is on ImageNet [Deng et al. 2009] train set, using the resolution of 256$\times$256 and random crop data augmentation. The image tokenizer model size is 72M and 70M when the downsample ratio is 16 and 8, respectively. All models are trained with the same settings: constant learning rate of $10^{-4}$, AdamW optimizer with $\beta_{1}=0.9$, $\beta_{2}=0.95$, weight decay$=0.05$, batch size of 128 and training epochs of 40. For the training losses, commitment loss weight is $0.25$ and adversarial loss weight is $0.5$. The adversarial loss is enabled after 20k training iterations.

Evaluation metrics.

We use the popular ImageNet benchmark under the image resolution of 256 $\times$ 256. The image reconstruction quality is measured by r-FID, reconstruction-FID on 256$\times$256 ImageNet 50k validation set. The codebook usage is calculated as the percentage of used codes in the queue of size 65536 over the whole codebook size. We also report PSNR and SSIM as the metrics of reconstruction quality, following SDXL [Podell et al. 2023].

Effect of image codebook designs.

As shown in Table 2, when the codebook vector dimension is reduced from 256 to 32 to 8, much better reconstruction quality and codebook usage are consistently achieved. For codebook size, a larger size from 4096 to 16384 benefits the overall performance. These observations are consistent with previous works [Yu et al. 2021, 2023b].

Effect of number of tokens to represent the image.

Table 3 studies the effect of image token number on image reconstruction quality. Using the same image tokenizer, for example, downsample ratio as 16, representing an image with only 256 tokens (16$\times$16) is not sufficient for good reconstruction quality, and increasing the number of tokens to 576 (24$\times$24) could largely improve the image quality from 2.43 to 0.99 rFID.

Comparisons with other image tokenizers.

We compare with other image tokenizers, including VQGAN [Esser et al. 2021], MaskGIT [Chang et al. 2022], ViT-VQGAN [Yu et al. 2021]. As shown in Table 4, our tokenizer outperforms previous image tokenizers. We also evaluate our tokenizer on COCO val2017 [Lin et al. 2014] of 256 $\times$ 256 image resolution to verify the image reconstruction quality, since COCO images contain more complex scenes. The comparison results are consistent with those in ImageNet validation set. This shows our tokenizer is a generalizable image tokenizer for both object-centric and scene-centric images.

Importantly, our tokenizer is competitive to continuous latent space representation, such as SD VAE [Rombach et al. 2022], SDXL VAE [Podell et al. 2023], and Consistency Decoder from OpenAI [OpenAI 2023a], which are widely used in diffusion models. This shows that discrete representation in the image tokenizer is no longer the bottleneck of the image reconstruction.

2 Class-conditional Image Generation

Our benchmark is the popular 256 $\times$ 256 ImageNet. All models are trained with the similar settings: base learning rate of $10^{-4}$ per 256 batch size, AdamW optimizer with $\beta_{1}=0.9$, $\beta_{2}=0.95$, $\text{weight decay}=0.05$, gradient clipping of 1.0. The dropout is always 0.1 for input token embedding, attention module and FFN module. The class condition embedding dropout for classifier-free guidance is 0.1.

Precomputing image codes. To accelerate the model training, we use the image tokenizer to precompute image codes before training. To achieve the similar effect of random crop data augmentation, we extract image codes of ten crops of the original image. During training, we randomly select one copy code from the ten augmentations.

Evaluation metrics.

We use Fréchet inception distance (FID) [Heusel et al. 2017] as the main metric. We also report Inception Score (IS) [Salimans et al. 2016], sFID [Nash et al. 2021] and Precision/Recall [Kynkäänniemi et al. 2019] as secondary metrics. All evaluations are implemented using ADM’s TensorFlow scripts [Dhariwal & Nichol 2021] for fair comparisons.

Effect of image tokens.

Although increasing the image tokens brings better image reconstruction quality, it is not strongly correlated to image generation quality. As shown in Table 5, when the model parameter is smaller than 1B, 256 (16$\times$16) tokens bring better image generation performance than 576 (24$\times$24). This shows the synergistic effect of scaling up model parameters and token numbers. Nevertheless, fewer image tokens would limit the image generation performance, for example, 256 (16$\times$16) tokens limit the FID at 3.06 FID, while 576 (24$\times$24) could further improve the FID to a lower value.

Effect of model size.

We train our models across five model sizes (B, L, XL, XXL, 3B) and evaluate their performance with and without classifier-free guidance. Figure 2 illustrates how FID changes as both the model sizes and the training epochs increase. Notable improvements in FID are observed when scaling the model from LlamaGen-B to LlamaGen-XXL. Further scaling to 3B yields only marginal improvements. A plausible explanation for this phenomenon could be the limitation in dataset size: ImageNet [Deng et al. 2009] comprises approximately only 1 million images, expanding the dataset or using stronger data augmentation could potentially lead to further improvements.

Effect of classifier-free guidance (CFG).

First, as shown in Figure 2, using classifier-free guidance can significantly enhance the visual quality across all model sizes. Moreover, Figure 3(a) illustrates that the model achieves optimal FID at CFG = 2.0 and further increasing CFG would deteriorate FID, which is consistent with previous findings [Dhariwal & Nichol 2021]. Additionally, the increment in CFG results in a trade-off between diversity and fidelity, as evidenced by increased precision and decreased recall, demonstrated in Table 10.

Effect of top-k sampling.

As shown in Figure 3(b), a small top-k value is not beneficial for FID and IS. Increasing top-k continuously improves FID but decreases IS, which trades off fidelity for diversity. We observe a similar trend when changing the parameter of top-p and temperature in sampling. Since FID is our main metric, we use maximum value as the default top-k value, which is the whole codebook size.

Comparisons with other image generation methods.

In Table 6, we compare with popular image generation models, including GAN [Brock et al. 2018; Kang et al. 2023; Sauer et al. 2022], Diffusion models [Dhariwal & Nichol 2021; Ho et al. 2022b; Rombach et al. 2022; Peebles & Xie 2023], and masked-prediction models [Chang et al. 2022]. Our models exhibit competitive performance in all metrics of FID, IS, Precision and Recall. Notably, our 3B model outperforms the popular diffusion models LDM [Rombach et al. 2022], DiT [Peebles & Xie 2023]. This shows that vanilla autoregressive models can serve as the basis of advanced image generation systems.

When comparing with autoregressive models [Esser et al. 2021; Yu et al. 2021; Lee et al. 2022], our model outperforms all previous models at different levels of model parameters. This benefits from better designs of image tokenizers and better scalability of image generation models. We hope our simple and effective implementation will serve as a solid baseline and help facilitate future research in autoregressive models for image generations.

3 Text-conditional Image Generation

We adopt a two-stage training strategy. In stage I, the model is trained on a 50M subset of LAION-COCO [LAION 2022] with the image resolution 256$\times$256. In Stage II, the model is fine-tuned on 10M internal high aesthetic quality images with the image resolution 512$\times$512. Examples of training data are shown in the Appendix. The maximum length of text token embedding is set to 120, and left padding is used to enable batch processing. The text condition embedding dropout for classifier-free guidance is 0.1. All models are trained with similar settings: model parameters of 775M, base learning rate of $10^{-4}$ per 256 batch size, AdamW optimizer with $\beta_{1}=0.9$, $\beta_{2}=0.95$, $\text{decay}=0.05$, gradient clipping of 1.0.

Precomputing image codes and text embeddings. We use pre-trained FLAN-T5 XL [Chung et al. 2024] to precompute text embedding of the image captions. For image code, we only extract image codes of the original image center crop in text-conditional models training.

Fine-tune image tokenizer.

Before two-stage training for text-conditional image generation models, we first fine-tune the image tokenizer on the joint of 50M LAION-COCO and 10M internal high aesthetic quality data.

Visualizations.

In Figure 4, we select text prompts from COCOPrompts [Lin et al. 2014] to generate images using models after stage I training and stage II training. After stage I training, the model captures the text-image alignment, while its ability to represent image details is not clear. Stage II training improves the visual aesthetic quality by a significant margin. We explain this improvement comes from two aspects: high aesthetic quality images shift the domain, and high image resolution brings better visual details. We notice that further increasing the image resolution to 1024$\times$1024 could bring better visual quality, and we leave it for future research.

More visualizations on PartiPrompts [Yu et al. 2022] are in Appendix. PartiPrompts have more longer captions than COCOPrompts, and our model demonstrates competitive performance in text-image alignment for long caption image generation tasks.

Limitation.

Due to the training data and model parameters, our text-conditional models have several limitations, such as text rendering errors, counting errors, and common misconceptions. These problems are promising to be mitigated when more training data and computation resources are available in the future.

4 Inference Speed

We verify the effectiveness of vLLM [Kwon et al. 2023] serving framework on our methods. Since our models use the same architecture as Llama, which is already supported by vLLM, we can seamlessly adopt its implementation. As shown in Table 7, we achieve 326% - 414% speedup compared to the baseline setting in the models from 111M to 1.4B parameters. Please note that the baseline setting has already integrated KV-Cache technique. In the 3B model, its head size 100 is not supported by PagedAttention in vLLM.

Related Work

Generative adversarial network (GAN) [Goodfellow et al. 2014; Brock et al. 2018; Karras et al. 2019; Kang et al. 2023] is the first representative visual generation method in deep learning era. To improve the distribution coverage, several likelihood-based methods are proposed. Diffusion models [Ho et al. 2020; Song & Ermon 2019; Song et al. 2020; Dhariwal & Nichol 2021] view image generation as the reverse diffusion process from noises to images. Masked-prediction models [Chang et al. 2022, 2023; Yu et al. 2023a, b] apply language model BERT-style [Devlin et al. 2018] by learning to predict masked tokens. Instead, autoregressive models [Esser et al. 2021; Ramesh et al. 2021; Yu et al. 2022] leverage GPT-style [Radford et al. 2018] to predict the next token in a sequence. To ease the modeling and improve the generation quality, these methods always introduce the image tokenization process [Kingma & Welling 2013; Van Den Oord et al. 2017] to convert pixel space to semantic space.

Multimodal foundation models.

Recently, vision-and-language models [Liu et al. 2024; Zhu et al. 2023; Dai et al. 2024; Peng et al. 2023; Zhang et al. 2023; Ma et al. 2024] have achieved versatile visual understanding through visual instruction tuning [Liu et al. 2024; Zhu et al. 2023]. However, unifying the understanding and generation in multimodal models is still in its early stages. Most existing methods [Sun et al. 2023b, a; Dong et al. 2024; Ge et al. 2023] try to collaborate a pre-trained diffusion model with existing models, rather than utilizing a unified next-token prediction paradigm. These methods need sophisticated designs to connect the two parts with distinct training paradigms, which makes scaling up challenging. Pioneering methods [Lu et al. 2022b, 2023; Bai et al. 2023b; Team et al. 2023; Team 2024] attempt to incorporate image generation into LLM using an autoregressive approach and achieve promising results. They do not specifically focus on demonstrating that a plain autoregressive approach can serve as a scalable image generator, which is our main argument in this work.

Conclusion

In this work, we delve into vanilla autoregressive models for scalable image generation. By reexamining their image tokenizers, image generation models and training data, our class-conditional models outperform the popular diffusion models, and our text-conditional models demonstrate competitive performance of visual quality and text alignment.

References

Appendix A Examples of Image-Text Pair Data

The original dataset has 600M image-text pair. We filter these images by valid image URL, aesthetic score, watermark score, CLIP image-text similarity score and image size. The remaining images are about 50M. Some examples are shown in Figure 5.

Training stage II: 10M internal high aesthetic quality images.

Each image is provided a long caption by LLaVA [Liu et al. 2024] using the prompt of “Describe this image in as much detail as possible”. Some examples are shown in Figure 6. We notice that the first sentence of the long caption is always a summary description of its image, so we use it as the short caption to augment the training of text-conditional image generation models.

Appendix B More Results on ImageNet Benchmark

We provide more detailed performance on ImageNet 256$\times$256 benchmark in Table 8 9 10. The generated image is always resized to 256$\times$256 when evaluating.