SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Jing Yu Koh
jingyuk@cs.cmu.edu
Daniel Fried
Carnegie Mellon University dfried@cs.cmu.edu
Ruslan Salakhutdinov
Carnegie Mellon University rsalakhu@cs.cmu.edu
We propose a method to fuse frozen text-only large language models (LLMs) with pre-trained image encoder and decoder models, by mapping between their embedding spaces. Our model demonstrates a wide suite of multimodal capabilities: image retrieval, novel image generation, and multimodal dialogue. Ours is the first approach capable of conditioning on arbitrarily interleaved image and text inputs to generate coherent image (and text) outputs. To achieve strong performance on image generation, we propose an efficient mapping network to ground the LLM to an off-the-shelf text-to-image generation model. This mapping network translates hidden representations of text into the embedding space of the visual models, enabling us to leverage the strong text representations of the LLM for visual outputs. Our approach outperforms baseline generation models on tasks with longer and more complex language. In addition to novel image generation, our model is also capable of image retrieval from a prespecified dataset, and decides whether to retrieve or generate at inference time. This is done with a learnt decision module which conditions on the hidden representations of the LLM. Our model exhibits a wider range of capabilities compared to prior multimodal language models. It can process image-and-text inputs, and produce retrieved images, generated images, and generated text — outperforming non-LLM based generation models across several text-to-image tasks that measure context dependence.
Autoregressive language models (LMs) and large language models (LLMs) trained on text corpora have shown impressive abilities to efficiently adapt to other modalities. Prior work showcased the effectiveness of grounding text-only LMs to images for vision-and-language tasks [56, 4, 29, 33, 31, 35], to embodied settings for robotics [3, 18], offline reinforcement learning [48], and more. These methods typically keep most of the LLM weights frozen. This allows them to leverage the capabilities that the LLM learns during large scale text-only pretraining, such as the ability to learn from in-context examples [9], more effectively process longer context, and condition on inputs more strongly.
In this work, we tackle the task of extending multimodal language models to generate novel images. Our approach, Generating Images with Large Language Models (GILL), is capable of processing arbitrarily interleaved image-and-text inputs to generate text, retrieve images, and generate novel images (Fig. 1). Our findings show that it is possible to efficiently map the output embedding space of a frozen text-only LLM to that of a frozen generation model (in this work, Stable Diffusion [49]) despite both models using entirely different text encoders. We achieve this by finetuning a small number of parameters on image-caption pairs [52], in contrast to other methods which require interleaved image-text data [4, 2]. Our approach is computationally efficient and does not require running the image generation model at training time. To achieve strong image generation performance, we propose efficient architectural changes to learn the LLM-to-generation mapping effectively with the GILLMapper module. GILLMapper is a lightweight Transformer [57] conditioned on special
Figure 1: Our model is capable of generating text, retrieving images, generating novel images, and interleaving results into coherent multimodal dialogue.
learnt text tokens. We train it by minimizing the  $l_{2}$  distance between its outputs and the outputs of the text encoder of a text-to-image generation model. This distillation training allows us to use the image decoder of the text-to-image model at inference time. Despite its simplicity, we show that this allows us to outperform the baseline text-to-image generation model on several tasks that measure language context dependence. Finally, to decide whether to produce a retrieved image or a generated one at inference time, we train a decision model that outputs a decision conditioned on the LM hidden representations. This allows us to both generate and retrieve in output sequences, as shown in Fig. 1.
Our experimental results demonstrate that GILL is more effective than Stable Diffusion at processing longer-form text, including dialogue and discourse. We show on dialogue-conditioned image generation that GILL can outperform non-LLM based generation models, and benefit from multimodal context: generating images that match text better than the backbone generation models that we distill from. In addition, GILL can process arbitrarily interleaved image-text inputs, unlike typical text-to-image models which only process text. GILL is the first model capable of outputting retrieved images, novel images, and text — interleaving these for coherent multimodal dialogue generation. $^{1}$
Multimodal Language Models Several prior works have developed multimodal language models which process image and text inputs to generate text outputs. Frozen [56] showed that it is possible to finetune a visual encoder to map images into the hidden space of a text-only LLM, and that this exhibits compelling few-shot, captioning, and question answering abilities. Other methods improve upon this approach by introducing adapters [19], scaling up model and data sizes [4, 64], improving the visual encoder [4, 33], finetuning on instructions [35], and training unified models on multi-task objectives [36, 63, 42]. CM3 [2, 62] trained multimodal LMs on HTML webpages consisting of interleaved images and text. Many state-of-the-art models also require significant computational resources to train. For example, Flamingo [4] is trained on 1535 TPUs for 15 days, while RA-CM3 [62] use 256 GPUs for 5 days. In contrast, our efficient adaptation method is trained on 2 GPUs for 2 days. The most similar work to our approach is FROMAGe [31], which trains a multimodal language model capable of processing arbitrarily interleaved image and text inputs to
<sup>1</sup>Our code and pretrained models are publicly released at https://github.com/kohjingyu/gill.
Figure 2: GILL model architecture overview. It is trained with a captioning loss to learn to process images (left), and losses for image retrieval and image generation to learn to produce images (right).
generate text interleaved with retrieved images. While FROMAGe can only retrieve images in their outputs, GILL is capable of both image retrieval and image generation, which allows it to outperform retrieval-only models when they are limited by their candidate retrieval set (Fig. 5).
Large Language Models Our work leverages recent advances in Transformer-based [57] LLMs. When trained at large enough scale, LLMs exhibit compelling properties, such as the ability to learn from few-shot in-context examples [9, 11] and generate and process long text inputs [61, 59, 53, 7]. Our approach also builds upon recent efforts on open sourced LLM weights [69, 55].
Text-to-Image Generation Text-to-image generation is the task of synthesizing a realistic image conditioned on natural language descriptions. [47] was one of the first to tackle this with a conditional GAN [23]. Later work improved upon this by introducing multi-stage models [67], attention mechanisms [60], and contrastive methods [73, 66]. Several recent papers also formulate the text-to-image generation task as a sequence modeling problem [45, 17, 13], training large Transformer [57] models on discretized image tokens [46]. [20, 65] improved upon this approach by introducing stronger image quantizers and scaling up model parameters. Several recent methods [38, 44, 49] apply diffusion models [26] to improve generated image quality. [50, 65] scale up text encoder models to achieve significant gains in generating relevant images. In contrast with computationally intensive methods that train end-to-end, GILL does not require running the image generation model during training.
We efficiently adapt a pretrained autoregressive language model of text, to process image and text inputs and produce image and text outputs. Most of the model weights (including those of the base LLM and image generator) are kept frozen, and we finetune a small number of parameters on image caption data (Fig. 2) to achieve a wide range of capabilities (Fig. 5). There are several challenges that need to be resolved. The model needs to learn to process image-and-text content (Sec. 3.1). It also needs to learn to produce images (either retrieved or generated), and determine whether to produce text or images at each step (Sec. 3.2). Finally, whenever an image is produced, the model needs to decide whether image retrieval (from a candidate set) or generation is more appropriate (Sec. 3.3).
3.1 Learning to Process Images
Given an image  $x$  and its text caption  $y$  (tokenized as  $(s_1, \ldots, s_T)$ ), our goal is to adapt a frozen LLM to enable it to complete any sequence of arbitrarily interleaved image and text inputs. For example, inputs for the Visual Storytelling dataset [28] consist of 5 images and 5 text descriptions, interleaved in a manner such as  $(x_1, y_1, \ldots, x_5, y_5)$ . We follow prior work [56, 19, 35, 31] in learning translation parameters that map from image features to text embedding space.
We first extract visual embeddings  $v_{\phi}(x) \in \mathbb{R}^d$  with a pretrained visual backbone (its weights  $\phi$  and LLM weights  $\theta$  are kept frozen). We learn a linear mapping  $\mathbf{W}_{\mathrm{cap}} \in \mathbb{R}^{d \times ke}$  which maps  $v_{\phi}(x)$  into a sequence of  $k$ $e$ -dimensional vectors that we use as inputs to the LLM (Fig. 2, left), where  $e$  is the
Figure 3: Inference time procedure for GILL. The model takes in image and text inputs, and produces text interleaved with image embeddings. After deciding whether to retrieve or generate for a particular set of tokens, it returns the appropriate image outputs.
Figure 4: GILLMapper model architecture. It is conditioned on the hidden [IMG] representations and a sequence of learnt query embedding vectors.
LLM input embedding dimension. We train  $\mathbf{W}_{\mathrm{cap}}$  on image-caption pairs (details in Sec. 3.4), by minimizing the negative log-likelihood loss of the token sequence  $(s_1, \ldots, s_T)$ :
$$
l _ {c} (x, y) = - \sum_ {t = 1} ^ {T} \log p _ {\theta} \left(s _ {t} \mid v _ {\phi} (x) ^ {T} \mathbf {W} _ {\text {cap}}, s _ {1}, \dots , s _ {t - 1}\right) \tag {1}
$$
Intuitively, this objective trains a mapping  $\mathbf{W}_{\mathrm{cap}}$  that allows us to translate images into embedding vectors in the token embedding space of the LLM (illustrated in Fig. 2, left).
3.2 Learning to Produce Images
In order to enable the model to produce image outputs, we add special [IMG] tokens to the vocabulary of the LLM, similar to [71, 31] which introduce special tokens correspond to images that should be output by the model. The hidden states that the LLM produces for these tokens will be used to retrieve or generate images. While [31] use a single token for their image retrieval model, we observed in our experiments that image generation requires much more finegrained textual information (Sec. 5). In order to improve the expressivity of the frozen LLM for novel image generation, we generalize to use  $r$  tokens [IMG1], ..., [IMG{r}] for representing visual outputs.
Concretely, we add a trainable matrix  $\mathbf{E}_{\mathrm{img}} \in \mathbb{R}^{r \times e}$  to the embedding matrix of the frozen LLM, which represents the  $r$  [IMG] token embeddings. We wish to train the model to learn when it should produce [IMG] tokens. This is done by minimizing the negative log-likelihood of producing the first [IMG] token conditioned on previously generated tokens:
$$
l _ {p} (y) = - \log p _ {\left\{\theta \cup \mathbf {E} _ {\mathrm {i m g}} \right\}} ([ \mathrm {I M G 1} ] \mid s _ {1}, \dots , s _ {t}) \tag {2}
$$
The LLM weights  $\theta$  are kept frozen, and we only update  $\mathbf{E}_{\mathrm{img}}$ . During inference, we always generate the [IMG2], ..., [IMG{r}] tokens whenever the first [IMG1] token is produced. During training, the [IMG] tokens are appended to each caption (Fig. 2). The LLM hidden states of the [IMG] tokens are used for image retrieval and generation, as described in the following sections.
Novel Image Generation In order for our LLM to produce image outputs, the [IMG] tokens need to be mapped into a semantically meaningful region of the input space of an image generation model  $G_{\psi}$  (such as that of the Stable Diffusion [49] image decoder). In initial experiments, we found that training a simple linear mapping such as those used in previous work on retrieval [31] was insufficient, and that such a model was unable to handle more complex prompts (see Sec. 5 for analysis). Hence, we propose GILLMapper (Fig. 4), a lightweight 4-layer encoder-decoder transformer model with trainable weights  $\omega$ . The GILLMapper module  $f_{\omega}$  conditions on  $h_{\{\theta \cup \mathbf{E}_{\mathrm{img}}\}}(y, [\mathrm{IMG}])$  (the [IMG] representations from the last hidden layer of the LLM) and  $L$  learnt query embeddings  $(q_1, \dots, q_L) \in \mathbb{R}^{L \times m}$  (where  $L$  is the maximum input sequence length of the text-to-image generation backbone  $G_{\psi}$ ).
The purpose of introducing learnable query embeddings is to enable GILLMapper to extract sequences of  $L$  features from the LLM [IMG] hidden states. This is similar to the queries introduced in DETR [10] for object detection and BLIP-2 [33] for extracting image features. We optimize the GILL trainable weights  $(q_{1},\ldots ,q_{L}$  and  $\omega$ ) by minimizing the MSE loss of the GILLMapper model outputs against the embeddings produced by the text encoder  $(T_{\psi})$  of a frozen text-to-image generation model:
$$
l _ {g} (y) = \left\| f _ {\omega} \left(h _ {\{\theta \cup \mathbf {E} _ {\mathrm {i m g}} \}} (y, [ \mathrm {I M G} \{1 \} ]), \dots , h _ {\{\theta \cup \mathbf {E} _ {\mathrm {i m g}} \}} (y, [ \mathrm {I M G} \{\mathbf {r} \} ]), q _ {1}, \dots , q _ {L}\right) - T _ {\psi} (y) \right\| _ {2} ^ {2} \tag {3}
$$
This is essentially distilling from  $T_{\psi}$  to learn a valid mapping from the output representations of our frozen LLM to the input space of  $G_{\psi}$ . Note that this does not require  $G_{\psi}$  during training, so we can precompute  $T_{\psi}(y)$  ahead of time, making training highly efficient. During inference, when [IMG] tokens are generated, we can synthesize an image by applying GILLMapper and the decoder  $G_{\psi}$ :
Generated Image  $= G_{\psi}(f_{\omega}(h_{\{\theta \cup \mathbf{E}_{\mathrm{img}}\}}(y,[\mathrm{IMG}\{1\} ]),\ldots ,h_{\{\theta \cup \mathbf{E}_{\mathrm{img}}\}}(y,[\mathrm{IMG}\{\mathbf{r}\} ]),q_1,\ldots ,q_L))$
where  $h_{\{\theta \cup \mathbf{E}_{\mathrm{img}}\}}(y, [\mathrm{IMG}\{\mathbf{i}\}])$  represents the hidden states from the last hidden layer of the modified LLM corresponding to the  $i^{th}$  [IMG] token. The learnt query embeddings  $(q_1, \ldots, q_L)$  are part of the GILLMapper model weights, and are hence kept fixed during inference.
Image Retrieval Similar to [31], we learn a linear mapping  $\mathbf{W}_{\mathrm{t2i}} \in \mathbb{R}^{e \times p}$  that maps the first token ([IMG1] to a  $p$ -dimensional vector. We also learn a linear mapping  $\mathbf{W}_{\mathrm{i2t}} \in \mathbb{R}^{d \times p}$  that map the pooled visual output of the image encoder  $v_{\phi}(x)$  to a  $p$ -dimensional space. These represent image and text embeddings, and we train the model by minimizing the InfoNCE loss [39]:
$$
l _ {r} \left(\mathbf {x} _ {i}, \mathbf {y} _ {i}\right) = - \log \frac {\exp \left(\sin \left(\mathbf {x} _ {i} , \mathbf {y} _ {i} , \mathbf {W} _ {\mathrm {t} 2 \mathrm {i}}\right) / \tau\right)}{\sum_ {j = 1} ^ {N} \exp \left(\sin \left(\mathbf {x} _ {j} , \mathbf {y} _ {i} , \mathbf {W} _ {\mathrm {t} 2 \mathrm {i}}\right) / \tau\right)} - \log \frac {\exp \left(\sin \left(\mathbf {x} _ {i} , \mathbf {y} _ {i} , \mathbf {W} _ {\mathrm {i} 2 \mathrm {t}}\right)\right) / \tau\left. \right)}{\sum_ {j = 1} ^ {N} \exp \left(\sin \left(\mathbf {x} _ {i} , \mathbf {y} _ {j} , \mathbf {W} _ {\mathrm {i} 2 \mathrm {t}}\right)\right) / \tau)} \tag {4}
$$
where the similarity is computed as
$$
\operatorname {s i m} (x, y, \mathbf {W}) = \frac {\left(\mathbf {W} ^ {T} v _ {\phi} (x)\right) ^ {T} \left(\mathbf {W} ^ {T} h _ {\{\theta \cup \mathbf {E} _ {\mathrm {i m g}} \}} (y , [ \mathrm {I M G 1} ])\right)}{\| \mathbf {W} ^ {T} v _ {\phi} (x) \| \| \mathbf {W} ^ {T} h _ {\{\theta \cup \mathbf {E} _ {\mathrm {i m g}} \}} (y , [ \mathrm {I M G 1} ]) \|}
$$
During inference, we follow standard procedure [43] in retrieving the image with the highest cosine similarity (between image embeddings and the [IMG] tokens) from a candidate set of images.
3.3 Deciding to Generate or Retrieve
While learning to produce [IMG] tokens allows us to decide when to interleave images in text, the task of deciding whether to retrieve or generate from [IMG] tokens remains. Intuitively, for a given prompt, we would like to retrieve when there is a strong match from our set of candidate images, and generate otherwise. In order to evaluate this, we collect human annotations on PartiPrompts (P2) [65], a collection of prompts used to benchmark image generation models. P2 contains some prompts that are well-represented by naturally occurring images, but others that are unlikely to occur in natural image sets, making it a test of generative models. For each of the 1,632 examples in P2, we generate an image with the text-to-image generation model  $G_{\psi}$ , and use the CLIP ViT-L [43] model to retrieve the top ranked image from CC3M [52] according to the cosine similarity of image embeddings  $v_{\phi}$ .
We have 5 independent human annotators (details in the appendix) select which of the two images for each prompt, retrieved or generated, is better matched to the prompt. We labeled the examples where the generated image was selected as 'gen' (indicating prompts which we should generate an image for) and 'ret' for prompts that should have an image retrieved. We extract the most confident set of these annotations (retaining roughly 900 examples with an inter-annotator agreement of at least  $4/5$ ), and split them into a  $67\%$  train (600) and  $33\%$  test (300) split. We use this to train a linear classifier on the LLM [IMG] hidden states as a decision model for deciding when to retrieve or generate (more details and baselines are provided in the appendix). Although these annotations of retrieving versus generating are somewhat model dependent, we believe that this data is still a valuable metric during model development. We will release our annotations to encourage future work in this space.
3.4 Data and Implementation Details
The final training objective for a batch of image-text pairs  $(\mathbf{x},\mathbf{y})$  is the sum of the captioning loss  $l_{c}$  (Eq. 1), image token prediction loss  $l_{p}$  (Eq. 2), generation loss  $l_{g}$  (Eq. 3) and retrieval loss  $l_{r}$  (Eq. 4):
$$
\min  _ {\mathbf {W} _ {i 2 t}, \mathbf {W} _ {i 2 i}, \mathbf {W} _ {\text {cap}}, \mathbf {E} _ {\text {i m g}}, \omega , q _ {1: L}} \frac {1}{N} \sum_ {i = 1} ^ {N} \left(l _ {c} (\mathbf {x} _ {i}, \mathbf {y} _ {i}) + l _ {p} (\mathbf {y} _ {i}) + l _ {g} (\mathbf {y} _ {i}) + l _ {r} (\mathbf {x} _ {i}, \mathbf {y} _ {i})\right) \tag {5}
$$
The decision model is trained separately after convergence of the other components. The multitask loss (Eq. 5) trains GILL to process images  $(l_{c})$ , produce [IMG] tokens  $(l_{p})$ , generate images  $(l_{g})$ , and retrieve images  $(l_{r})$ . This enables it to generalize to a wide range of vision-and-language tasks.
We train on Conceptual Captions (CC3M) [52], which consists of 3.3M image-text pairs. Following [31], we pack two random examples together during training with probability 0.5 (i.e.,  $50\%$  of the time, the input is a single image and caption example, while the other  $50\%$  of the time the input consists of a sequence consisting of two interleaved images and captions). We use the OPT-6.7B [69] model as the LLM backbone (which produce hidden states  $h_\theta$  with embedding dim  $e = 4096$ ). For the visual model used to extract features  $v_\phi$  for captioning and retrieval, we use the CLIP [43] ViT-L model. For our text-to-image generation backbone  $G_\psi$ , we use the Stable Diffusion [49] v1.5 model (with  $L = 77$  input vectors). We use  $k = 4$  visual tokens, and  $r = 8$  learnt [IMG] tokens. We set the GILLMapper query embedding dimension  $m = 512$ . For retrieval, we use an embedding dimension  $p = 256$ . All pretrained model weights are kept frozen, and we only train the linear layers  $\mathbf{W}_{\mathrm{i2t}}$ ,  $\mathbf{W}_{\mathrm{t2i}}$ ,  $\mathbf{W}_{\mathrm{cap}}$ , the [IMG] embedding matrix  $\mathbf{E}_{\mathrm{img}}$ , and the GILLMapper parameters  $\omega$  and query vectors  $q_{1:L}$ . In total, there are 50M trainable parameters, significantly fewer than in the frozen LLM and visual models (which total approximately 8B parameters). We use bfloat16 precision [1], and optimize using Adam [30] ( $\beta_1 = 0.9$ ,  $\beta_2 = 0.95$ ) with a learning rate of 0.001. We train with a batch size of 200 for 20K iterations, which takes 2 days on 2 A6000 GPUs. We follow [31] and concatenate captions to encourage the model to attend to relevant images within an image-text sequence.
GILL is the first multimodal language model capable of conditioning on image-and-text inputs to generate meaningful images interleaved with text. Hence, our experiments primarily focus on evaluating its ability to produce novel images (Sec. 4.1). Our results show that GILL improves over Stable Diffusion [49] on tasks that require processing long-form text such as dialogue and discourse. We also benchmark the performance of models in deciding whether to retrieve or generate (see appendix). GILL is capable of generating text, retrieving images, and generating images. Despite being more general than prior work [56, 4, 31], we find that GILL performs comparably to or better than existing multimodal LMs on contextual image retrieval and text generation tasks (see Sec. 5).
4.1 Contextual Image Generation
To test the ability of our model against baseline methods for novel image generation, we run experiments on the VIST [28] and VisDial [16] datasets. These are the same datasets used in prior work [31] for benchmarking image retrieval conditioned on multimodal text-and-image context.
Evaluation Metrics The focus of our evaluation is on the ability of generative models to handle complex language descriptions. Hence, we compute metrics which measure the relevance of the generated image content. We evaluate models with two metrics:
1. CLIP Similarity: We use the CLIP [43] ViT-L image encoder to produce pooled representations of generated images and the corresponding real images, and report their cosine similarity. A higher score indicates that a generated image is more similar to the real image.
2. Learned Perceptual Image Patch Similarity (LPIPS): LPIPS [68] evaluates the distance between image patches. We measure LPIPS between real and generated images. A lower value indicates that two images are closer in perceptual space (i.e., more similar), while a higher value indicates that two images are more dissimilar.
2These models were selected for their strong performance and open source availability, but in principle our approach can be applied with any other pretrained models with similar function calls.
Image + Text Input
FROMAGE (retrieved)
Ours (generated)
rstFROMAGe
Multimodal Dialogue
Figure 5: Qualitative results over various input and output modalities. GILL is able to process contextual multimodal cues to retrieve and generate appropriate image and text outputs.