Scaling Autoregressive Multi-Modal Models: Pretraining and Instruction Tuning

Introduction

Diffusion models have recently dominated image generation work due to their strong performance and relatively modest computational cost (Saharia et al., 2022; Chen et al., 2022; Rombach et al., 2022). In contrast, token-based autoregressive models (Ramesh et al., 2021; Yu et al., 2022) are known to also produce strong results, with even better global image coherence in particular, but are much more expensive to train and use for inference. In this paper, we show that it is possible to extend training and inference ideas originally developed for text-only models to flip this narrative; autoregressive models can be efficient and performant while also generalizing beyond the strict text-to-image format to be tuneable for a wide range of image and text generation tasks.

More specifically, we introduce CM3Leon (pronounced “Chameleon”), a retrieval-augmented, token-based, decoder-only multi-modal language model capable of generating and infilling both text and images. CM3Leon uses the CM3 multi-modal architecture (Aghajanyan et al., 2022), but additionally shows the extreme benefits of scaling up and training on more diverse data. It is the first multi-modal model trained with a recipe adapted from text-only language models, including a large-scale retrieval-augmented pretraining stage and a second multi-task supervised fine-tuning (SFT) stage. The pretraining is efficient because it follows the retrieval-augmented CM3 approach (Yasunaga et al., 2022) but uses a new large-scale Shutterstock dataset that includes only licensed image and text data. The SFT stage follows multi-task instruction tuning for text-only models Iyer et al. (2022), which allow arbitrary mixtures of image and text tokens in both the inputs and outputs. The generality of CM3Leon also supports the introduction of an improved, self-contained contrastive decoding method Li et al. (2022), which can provide self-guidance to improve both text and image generation.

CM3Leon achieves state-of-the-art performance in text-to-image generation with 5x less training compute than comparable methods (zero-shot MS-COCO FID of 4.88). It can also do non-trivial image-to-text generation, even though it was trained on only 3B Shutterstock text tokens. After SFT, CM3Leon demonstrates unprecedented levels of controllability in tasks ranging from language-guided image editing to image-controlled generation and segmentation. We also show that retrieval augmentation is key for efficient training, and our new contrastive decoding method enables much higher quality generation overall. These results strongly suggest that autoregressive models are worth significantly more study for any text and image task.

Pretraining

We explore the potential of token-based decoder-only models in the text-to-image domain by building upon the foundation laid by RA-CM3 Yasunaga et al. (2022). We simplify the original settings in RA-CM3 by streamlining the objective, modifying the dataset, and incorporating insights from multi-modal scaling laws presented by Aghajanyan et al. (2023).

The ethical implications of image data sourcing in the domain of text-to-image generation have been a topic of considerable debate. In this study, we use only licensed images from Shutterstock. As a result, we can avoid concerns related to images ownership and attribution, without sacrificing performance.

We use the image tokenizer from Gafni et al. (2022a), which encodes a $256\times 256$ image into $1024$ tokens from a vocabulary of $8192$. For text, we train a custom tokenizer over the Zhang et al. (2022) data with a vocabulary size of $56320$. Additionally, we introduce a novel special token, denoted as , which serves to indicate a transition between modalities. A visualization of one caption-image pair after tokenization and formatting with our special tokens is available in § B.1(Figure 8).

Retrieval Augmentation

Our retrieval approach aims to retrieve relevant and diverse multi-modal documents from a memory bank, given an input sequence (Yasunaga et al., 2022). It includes both a dense retriever and a retrieval strategy.

The dense retriever takes a query $q$ (e.g., the input sequence $x$) and a candidate document $m$ from the memory bank $\mathcal{M}$ and returns a relevance score $r(q,m)$. We adopt the dense retrieval method from Karpukhin et al. (2020), which uses a bi-encoder architecture. The encoder is CLIP-based. We split the multi-modal document into a text part and an image part, encode them separately using off-the-shelf frozen CLIP text and image encoders, and then average the two as a vector representation of the document (Radford et al., 2021). We use the ViT-B-32 model and normalize the image/text embeddings. The final retrieval is done with Maximum Inner Product Search (MIPS) over the memory bank using the dense retriever to obtain a list of candidate documents sorted by relevance score (Tiwari et al., 2022).

To sample informative retrieved documents for the generator during training, we consider three key factors: relevance, modality, and diversity. First, the retrieved documents should be relevant to the input sequence, captured by the dense retriever score based on CLIP. Second, retrieving a multi-modal document consisting of images and text leads to better generator performance than retrieving either image or text. Third, diversity is essential to avoid redundancy in the retrieved documents. Simply taking the top $K$ documents based on relevance score can result in duplicates or highly similar documents, hurting downstream pretraining. We skip a candidate document if it is too similar to the query or if the documents have already been retrieved. In practice, we only use retrieved documents with relevance score $\leq 0.9$. Additionally, we use query dropout, which drops some tokens of the query used in retrieval (20% of tokens) to encourage diversity and serve as regularization for training.

Throughout our work, we retrieve two documents each, based on image and text, respectively. In training, we randomly select three retrieved samples for every caption-image pair in our dataset, effectively 4x the number of tokens available in the pretraining. A visualization of a single training example can be found in § B.1(Figure 9).

2 Objective Function

The CM3 objective accepts multi-modal inputs (e.g., $x_{\text{input}}=$ "Image of a chameleon: [image]") and transforms them into an infilling instance by masking specific spans and relocating them to the end (e.g., $x_{\text{input}}=$ "Image of : [image] a chameleon"). It uses a standard next token prediction loss, $-\log p(x_{\text{input}})$. This results in a versatile model capable of infilling and autoregressive generation tasks for both images and text. In the case of caption-to-image generation, CM3 creates a continuation from the prompt "Image of a chameleon:". For image-to-caption generation, CM3 utilizes the prompt "Image of : [image] ".

Yasunaga et al. (2022) built upon the original CM3 by including retrieved multi-modal documents in the context for each training example and up weighting the query image-caption pair loss, as illustrated in the last image-caption pair in Figure 9. This approach encourages the model to concentrate more on using retrieved samples during the generation process. However, this method adversely affects the zero-shot scenario, where the goal is to generate an image without retrieval, such as predicting a continuation from text . We remove this weighting in our setting and make a minor modification to the CM3 objective by preventing masking across tokens. This adjustment is justified by the fact that allowing masking across tokens may lead to the model generating image content from an arbitrary midpoint, which is not a desirable outcome.

3 Model

The CM3Leon models follow a decoder-only transformer architecture, similar to Zhang et al. (2022) and Brown et al. (2020). Compared to Zhang et al. (2022), we remove bias terms, dropout, and learnable parameters for layer norms and use a sequence length of 4096 instead of 2048. For weight initialization, we use a truncated normal distribution with a mean of 0 and a standard deviation of 0.006, truncated to 3 standard deviations. Output layers are initialized as 0, and the learned absolute positional embedding is initialized near zero with a standard deviation of 0.0002. The models were trained with Metaseqhttps://github.com/facebookresearch/metaseq, with experiment tracking done with Aim Arakelyan et al. (2020).

4 Training

Our models are trained across three distinct sizes, with the corresponding parameters and training setup detailed in Table 3. The major hyperparameters, such as the learning rate and batch size, are adopted from prior work in multi-modal scaling laws, creating a stable and smooth training progression as illustrated in Figure 3 (Aghajanyan et al., 2023). The 350 Million (350M), 760 Million (760M), and 7 Billion (7B) models are trained to 1.4 Trillion (T), 1.9T, and 2.4T tokens, respectively. The losses for all three models decrease steadily throughout training, strongly suggesting they have not saturated.

Text-To-Image Results

There has been significant work on developing decoding algorithms for autoregressive text-to-image models, such as DALL-E Ramesh et al. (2021), which can have a large effect on the quality of the final outputs. DALL-E employs temperature sampling and a re-ranking stage via CLIP over 512 prompt candidates. Models like PARTI and Make-A-Scene user token-based classifier-free guidance, significantly reducing the number of candidates required for re-ranking to just 16 samples (Yu et al., 2022; Gafni et al., 2022a). Our experiments show that different approaches offer complementary benefits, as decribed in this section. We compare the following options.

is a probabilistic technique used in autoregressive models, such as Ramesh et al. (2021). The method involves modifying the softmax temperature during the sampling stage to control the randomness of predictions. We pair this with Classifier Free Guidance in all of our experiments.

TopP Sampling

also known as nucleus sampling, involves sampling from the smallest set of top-ranked tokens with a cumulative probability exceeding a predefined threshold (Holtzman et al., 2020). We pair this with Classifier Free Guidance in all of our experiments.

Classifier Free Guidance (CFG)

Classifier-free guidance refers to directing an unconditional sample towards a conditional sample (Gafni et al., 2022a). We replace the text with the mask token from the CM3 objective to facilitate unconditional sampling. This is one of the core benefits of training with the CM3 objective, allowing us to do classifier-free guidance without the need for finetuning. During the inference stage, two concurrent token streams are generated: a conditional token stream, which is contingent on the input text, and an unconditional token stream, which is conditioned on a mask token. Borrowing the notation from Gafni et al. (2022a):

where $T$ denotes the transformer, $t_{y}$ is the output tokens and $t_{x}$ is the conditional input text, represents the absence of input text (and replacement with a mask token), and $\alpha_{c}$ is a scaling factor. The classifier-free guidance effectively blends the unconditional and conditional logits, influencing the model’s output towards a more desired conditional output.

Contrastive Decoding TopK (CD-K)

A key insight is that the logit subtraction in Equation 2 resembles the log probability subtraction in contrastive decoding methods in text (Li et al., 2022). This leads us to propose a variant of the contrastive decoding (CD) algorithm, originally proposed by Li et al. (2022), as an alternative to CFG.

Recall that CD defines a score per token:

Here, $\mathcal{V}({t_{y}}_{<i})$ represents the set of potential subsequent tokens whose probabilities are at least $\alpha$ times the maximum probability value:

Traditionally $p_{\textsc{exp}}$ and $p_{\textsc{ama}}$ in the CD decoding algorithm represent a strong and weak model where the strong model was trained with more compute (or larger model size) compared to the weak model. Instead we select $p_{\textsc{exp}}$ having text conditioning and $p_{\textsc{ama}}$ has no text conditioning. Additionally we saw that the $\mathcal{V}({t_{y}}_{<i})$ constraint was too strict, and would consistently become greedy decoding. Therefore we propose a slight modification of CD we call CD-K that alters $\mathcal{V}({t_{y}}_{<i})$ to:

where instead of taking the largest probability we take the $k$-th largest probability.

Ablation

In Figure 4 we show that CD-K is competitive with standard CFG based sampling while providing a complementary set of generations to CFG allowing us to continue minimizing FID as we increase number of generations (while both CD-K and CFG independently stagnate).

2 Quantitative Evaluations

Table 1 and Figure 3 provide a comparative overview of CM3Leon and state-of-the-art text-to-image models, evaluated based on the zero-shot MS-COCO (30K) task using the Fréchet Inception Distance (FID) metric (Seitzer, 2020). CM3Leon-7B model set’s a new state-of-the-art FID score of 4.88, while only using a fraction of the training data and compute of other models such as PARTI.

This observation underlines the effectiveness of retrieval-augmented decoder-only models like CM3Leon. In particular, the CM3Leon-7B model, when operated with one or two retrieved examples during inference, records superior FID scores. This result demonstrates the crucial role retrieval plays in expanding the world knowledge provided to the model and its capacity to generate high-quality images. CM3Leon surpasses all other retrieval-augmented models, including KNN-diffusion and RE-IMAGEN.

Supervised Fine-Tuning

Supervised fine-tuning (SFT) is critical in training large language models (LLMs) like ChatGPT. Despite this, its application in multi-modal settings remains largely unexplored. SFT trains a model to better understand of future instructions or prompts, enhancing its performance in novel and even zero-shot tasks. We have found that instruction tuning notably amplifies multi-modal model performance across various tasks such as image caption generation, visual question answering, text-based editing, and conditional image generation.

We fine-tune CM3Leon on a wide array of mixed image and text tasks. We organized each task as a series of interleaved text and image examples, as shown in Figure 5. The fine-tuning process follows the pretraining stage, employing the same CM3 objective by combining the task instruction with the output. Further details about the hyperparameters and scale of the SFT can be found in Section E.1.

allows the modification of an initial image based on text instructions, with changes such as seasonal and weather adjustments, background changes, and material alterations. We used InstructPix2Pix methodology and proprietary face-filtering techniques on their data, yielding around 600,000 examples (Brooks et al., 2023).

Image-to-Image Grounded Generation

involves producing grounding images with various features and text prompts. Features like edge maps, segmentation maps, key points, and human poses can be derived from user-uploaded images or sketches. We used ControlNet processing code on Shutterstock datasets to curate 7 million examples with features like canny edge, hed boundary, user sketching, human pose, and more (Zhang & Agrawala, 2023).

Spatially Grounded Image Generation

allows the user to integrate spatial information into text prompts for image generation, with each object represented by discrete tokens. We used object detection datasets like MS-COCO, Openimage, and Object365 to compile 3 million training examples(Lin et al., 2014; Kuznetsova et al., 2020; Shao et al., 2019).

How-to-write

task enables users to request the model to create signs or logos based on text prompts. We used an OCR detector to find suitable examples from Shutterstock datasets, resulting in 200,000 examples.

Results:

We showcase qualitative examples of images produced by a fine-tuned CM3Leon-7B model, as depicted in Figure 6. All instances in text-guided editing and image-image-grounded generation utilize a task prefix. For instance, we precede every text-guided editing example with the phrase, "Edit the image following the text instruction," and every scribble generation example with "Create a high-quality image from children’s scribble and text description," amongst others. The top row of Figure 6 presents text-guided image examples. We employ separate image CFG (1.5) and text CFG (7.5) values during decoding. This approach is crucial for producing edited images that mirror the original image and closely align with the text editing instruction. The second row in Figure 6 show Structure-Guided Image Editing examples. For decoding, we utilized a single CFG value of 3. Given identical input open pose features, our model can generate markedly distinct images that follow different text prompts while maintaining the same pose as in the input image. More examples in 15

2 Conditional Text Generation

We also include several vision-language tasks to teach CM3Leon to respond in text to various kinds of textual prompts conditioned on an image, such as visual question answering, long-form captioning, etc. We use the following 8 vision-language tasks: MS-COCO (Chen et al., 2015), Flickr30k (Young et al., 2014), Image Paragraph (Krause et al., 2017), Localized Narratives (Pont-Tuset et al., 2020), VQA2 Goyal et al. (2017), VizWiz (Gurari et al., 2018), OKVQA (Marino et al., 2019), and ScienceQA (Lu et al., 2022). We use multiple prompt templates for each task to make the model robust to prompt variations (more details on the templates in Table 5 of the Appendix).

Table 2 presents the performance comparison of our SFT-CM3Leon model w.r.t. previous state-of-the-art (SoTA) such as Flamingo (Alayrac et al., 2022) and OpenFlamingohttps://laion.ai/blog/open-flamingo/. We show that our SFT-CM3Leon model achieves strong zero-shot performance on several vision-language tasks even though they saw significantly fewer text data ($\approx$ 3B tokens) compared to Flamingo (100B tokens) and OpenFlamingo (40B tokens). Notably, SFT-CM3Leon even beats Flamingo on the VizWiz task. Figure 16 presents our SFT-CM3Leon-7B model generations, given an image context and an instruction. The model is quite flexible with the instruction and can generate captions or answer a variety of questions. Further, the ability of to follow instructions is more evident in Figure 7 where the model can generate very long captions or reason over an image based on the given instruction.

Related Work

Significant progress in the domain of text-to-image generation has been achieved through the use of diffusion models (Rombach et al., 2022; Nichol et al., 2021; Ramesh et al., 2022). The underlying mechanism involves sequentially adding noise to an image and then learning to reverse the noise based on provided text inputs or features (Luo, 2022). Diffusion models generally incorporate pretrained text or language representations such as the text encoder of the CLIP (Radford et al., 2021) image-text model or text encoders like T5 (Raffel et al., 2020). The recursive application of multi-resolution diffusion model (by employing multiple steps of super-resolution) has further enhanced their capability to generate high-quality images from text prompts, leading to state-of-the-art zero-shot non-retrieval based MS-COCO FID scores

Retrieval Augmented Diffusion Models

Conventional text-to-image models often struggle to capture the vast array of concepts and entities in the image domain. Methods like enabling retrieval during inference time can help address the complexity of these tail entities by delegating it to a retrieval step. Building on the work of Saharia et al. (2022), Chen et al. (2022) incorporates retrieval to enhance zero-shot MS-COCO FID scores, demonstrating further improvement in this area.

Autoregressive Token Models

Significant advancements have been made in the field by utilizing LLMs over tokenized image representations (Esser et al., 2020; Ramesh et al., 2021). A widely-used approach in the field (Van Den Oord et al., 2017; Razavi et al., 2019; Esser et al., 2021) involves an initial stage of converting images into discrete latent variables through tokenization, which transforms a text-to-image generation problem into a sequence-to-sequence problem, thereby enabling subsequent application of LLM techniques (Ramesh et al., 2021; Gafni et al., 2022b).

Non-Autoregressive Token Models

Although autoregressive models have benefited from extensive research in NLP, autoregressive decoding can be quite computationally expensive. Non-autoregressive models, such as Ghazvininejad et al. (2019), have been proposed in NLP and extended to text-to-image models, exemplified by Chang et al. (2023) which achieves state-of-the-art image generation performance and higher efficiency than diffusion or autoregressive models by employing masked modeling in discrete token space (non-autoregressively with iterative decoding).

Retrieval Augmented Autoregressive Token Models

Token-based models face challenges akin to those encountered by non-retrieval augmented diffusion models. To address these issues, Yasunaga et al. (2022) suggested prefixing decoder-only text-to-image models, such as Ramesh et al. (2021); Aghajanyan et al. (2022), with statically retrieved instances during training, resulting in significant efficiency gains during the training process.

Our paper primarily concentrated on scaling this strategy.

Conclusion

We presented CM3Leon, a retrieval-augmented, token-based, decoder-only multi-modal language model that efficiently and flexibly generates and infills text and images. Our approach extends the scope of autoregressive models, demonstrating their potential to compete with and exceed diffusion models in terms of cost-effectiveness and performance. By integrating a retrieval-augmented pretraining stage with a diverse, large-scale Shutterstock dataset and a second multi-task supervised fine-tuning stage, CM3Leon demonstrates the benefits of a comprehensive training approach. Further enhanced by an innovative, self-contained contrastive decoding method, our model offers improved text and image generation quality. Our results support the value of autoregressive models for a broad range of text and image tasks, encouraging further exploration for this approach.

References

Appendix A Showcase Prompts

Chameleon and octopus, side by side, high quality render, drawing, professional.

A plush toy koala bear relaxing on a lounge chair and working on a laptop. The chair is beside a rose flower pot. There is a window on the wall beside the flower pot with a view of snowy mountains.

A photo of an astronaut riding a horse in the forest. There is a river in front of them with water lilies.

A teddy bear wearing a motorcycle helmet and cape is riding a motorcycle in Rio de Janeiro with Dois Irmãos in the background. dslr photo.

A black german shepherd wearing a red beret

An Armenian church on the surface of Mars, with Astronaut walking into the church, in Focus. Photo. Fantasy. Dramatic.

Armenian khachkars surrounded by pomegranates in a bright green forest.

A small cactus wearing a straw hat and neon sunglasses in the Sahara desert.

A close up photo of a human hand, hand model. High quality

A raccoon main character in an Anime preparing for an epic battle with a samurai sword. Battle stance. Fantasy, Illustration

A stop sign in a Fantasy style with the text "1991"

Appendix B Pre-Training

B.2 Model Hyper-Parameters

Appendix C Inference Latency and Throughput

Appendix D Image Generation

Appendix E Supervised Fine Tuning

E.2 Training Data Breakdown

E.3 More Qualitative Samples