ReVeaL: Retrieval-Augmented Visual-Language Pre-Training with
Multi-Source Multimodal Knowledge Memory

Figure 2 (a) depicts how the input image-text query is encoded. We use a base visual-language encoder $b(\cdot)$ to turn the query input and each knowledge item (with potentially different modalities e.g. text-only, image-only or image-text pairs) into a sequence of embeddings (tokens). We adopt a Vision Transformer (ViT) [10] to encode the images and we use a lower-layer¹¹1We denote the last $l$ layers of a T5 encoder as ‘upper-layer’, and the remaining ones including the token embedding layer as ‘lower-layer’. T5 encoder [33] to encode the texts. We add a projection layer on top of the ViT model to map the image tokens into the same space as the text tokens. We then concatenate the two modalities together. We use an upper-layer T5 module as both the query Head $\phi_{\texttt{Query}}(\cdot)$ and the key Head $\phi_{\texttt{Key}}(\cdot)$ to compute the query embedding and memory keys. We take the output of the first [CLS] tokens followed by a linear projection and L2-normalization to summarize the input into a $d$-dimensional embedding.

Figure 2 (b) shows how memory is constructed and updated by encoding knowledge items. Our approach differs from previous works primarily by leveraging a diverse set of multimodal knowledge corpora (WikiData knowledge graph, Wikimedia passages and images, Web image-text pairs). Throughout the paper, we denote each corpus as $\mathcal{C}^{j}=\{z^{j}_{1},\dots,z^{j}_{N}\}$, in which each $z^{j}_{i}\in\mathcal{C}^{j}$ is a knowledge item that could be an image-text pair, text only, image only, or a knowledge graph triplet. We denote the unified knowledge corpus as $\tilde{\mathcal{C}}=\mathcal{C}^{1}\cup\mathcal{C}^{2}\dots\cup\mathcal{C}^{S}$ that combines $|\tilde{\mathcal{C}}|=S$ different knowledge corpora. We encode the external knowledge corpora into a unified memory $\tilde{\mathcal{M}}=[\mathcal{M}^{1},\dots,\mathcal{M}^{|\tilde{\mathcal{C}|}}]$. Each knowledge item $z_{i}$ is encoded into a key/value pair $m_{i}=(\texttt{Emb}_{\texttt{Key}}(z_{i}),\texttt{Emb}_{\texttt{Value}}(z_{i}))$ in memory. Each key $\texttt{Emb}_{\texttt{Key}}(z)=\phi_{\texttt{Key}}\big{(}b(z)\big{)}\in\mathbb{R}^{d}$ is a $d$-dimensional embedding vector encoded via Key Head. Each value is a sequence of token embeddings representing the full information of knowledge item $z$. We follow a similar procedure as in [14] to precompute key/value embeddings of knowledge items from different sources and index them in a unified knowledge memory. We continuously re-compute the memory key/value embeddings as the model parameters get updated during the pre-training phase. We update the memory $\tilde{\mathcal{M}}$ asynchronously at every 1000 training steps.

Scaling Memory by Compression A naive solution for encoding the memory value is to keep the whole sequence of tokens for each knowledge item. Then, the generator could fuse the input query and the top-K retrieved memory values by concatenating all their tokens together and feeding them into a Transformer Encoder-Decoder pipeline [23]. This approach has two issues: (1) storing hundreds of millions of knowledge items in memory is impractical given that each memory value would consist of hundreds of tokens; (2) transformer encoder has quadratic complexity with respect to the total number of tokens times $K$ for self-attention.

Therefore, we propose to use the Perceiver architecture [19] as the Value Head to encode and compress knowledge items. The Perceiver model uses a transformer decoder $\psi(\cdot)$ with learnable $c$-length latent embeddings to compress the full token sequence into an arbitrary length $c$, such that $\texttt{Emb}_{\texttt{Value}}(z)=\psi(b(z))\in\mathbb{R}^{c\times d}$ (In our experiments we use $c=32$). This lets us retrieve top-K memory entries for K as large as a hundred. To make the compressed embeddings generated by Perceiver more expressive, we add two additional regularizations. The first one is a disentangled regularization [16] that forces every two output tokens to be linearly de-correlated $\mathcal{L}_{\text{decor}}=\sum_{i,j=1}^{K}\Big{\lVert}\text{Covariance}\big{(}\psi\big{(}b(z_{i})\big{)},\psi\big{(}b(z_{j})\big{)}\big{)}\Big{\rVert}_{F}^{2}$, and the second one is an alignment regularization that minimizes the distance of L2-Norm between the query and compressed knowledge embedding: $\mathcal{L}_{\text{align}}=\bigg{\lvert}1-\frac{\sum_{z}\lVert\psi\big{(}b(z)\big{)}\rVert_{2}}{\sum_{x}\lVert b(x)\rVert_{2}}\bigg{\rvert}$.

Figure 2 (c) shows ReVeaL’s retrieval procedure. Given the input query $x$, the retriever’s task is to find top-K memory entries $M$ with the highest probability $p(M\mid x)$ which we approximate as $p(M\mid x)=\prod_{m\in M}p(m\mid x)$ by retrieving each entry independently. Note that we retrieve from a large-scale unified memory $\tilde{\mathcal{M}}=[\mathcal{M}^{1},\dots,\mathcal{M}^{|\tilde{\mathcal{C}|}}]$ that is constructed from a diverse set of knowledge sources. To help the query to better choose the most appropriate knowledge sources, we learn a gating function that models the probability of retrieving from each memory corpus. With the corpus gating, for $m^{j}_{i}\in\mathcal{M}^{j}$ we re-weight $p(m^{j}\mid x)$ by the computed corpus gating score:

where $Gate_{\mathcal{M}^{j}}(x)=\text{Softmax}(W\cdot\texttt{Emb}_{\texttt{Query}}(x)+b)[j]$ is a softmax gating that assigns a score to each memory corpus $M^{j}$, with $W$ and $b$ as function parameters. $Rel(x,m^{j}_{i})$ models relevance score between query $x$ and each memory entry via embedding dot product, such that $Rel(x,m^{j}_{i})=\texttt{Emb}_{\texttt{Query}}(x)^{T}\cdot\texttt{Emb}_{\texttt{Key}}(z^{j}_{i})$. where $z_{i}$ is the knowledge item corresponding to the memory entry $m_{i}$ and $\tau$ is the temperature parameter.

After identifying the top-K memory entries, the retriever passes the pre-computed in-memory key and value embeddings to the generator. In the meantime, to support end-to-end training of the encoders, we also re-encode a small portion (i.e., 10$\%$) of the retrieved knowledge items $z_{i}$ from scratch. In this way, the memory encoders could be updated with moderate computational cost. We concatenate the re-encoded knowledge with in-memory ones to construct the final top-K retrieved key/value embeddings.

Figure 2 (d) shows how the query and the retrieved knowledge items are fused to generate the output answer. All $K$ retrieved memory values are concatenated with the query embedding, which is feasible due to the Perceiver module utilized as the value head $\psi(\cdot)$, compressing each knowledge item into a short sequence. We denote the concatenated query embedding and memory values as $X=[b(x),\psi(b(z_{1})),\dots,\psi(b(z_{K}))]\in\mathbb{R}^{(I+c\cdot K)\times d}$, where $I$ is the number of tokens of the input query $x$ and $c$ is the number of compressed tokens. To guide the generator towards attending to the most important items in $X$ and facilitate backpropagation of gradients to the retriever, we propose an attentive fusion module $f(\cdot)$ capable of incorporating the retriever score as a prior for calculating cross-knowledge attention. The detailed procedure is illustrated in Figure 3. We firstly compute a latent soft attention mask over $X$ as $\texttt{Mask}_{\texttt{att}}=[1,p(z_{1}|x),\dots,p(z_{K}|x)]$. Finally, we pass the fused representation $f(X,\texttt{Mask}_{\texttt{att}})$ into a T5 decoder module $g(\cdot)$ to generate the textual output.

We pre-train our model on the Web-Image-Text dataset [51], a large-scale corpus containing 3 billion image alt-text caption pairs collected from the public Web. Since the dataset is noisy, we add a filter to remove data points whose captions are shorter than 50 characters. This yields roughly 1.3 billion image caption pairs for pre-training.

We denote the pre-training Web-Image-Text dataset [51] as $\mathcal{D}$. We use the text generation objective used in Wang et al. [45]) to pre-train our model on $\mathcal{D}$. Given an image-text example $x=$($\texttt{img},\texttt{txt}$) from $\mathcal{D}$, we randomly sample a prefix length $T_{p}$. We feed $x_{<T_{p}}$ that contains the text prefix and image to the model as input and our objective is to generate $x_{\geq T_{p}}$ containing the rest of the text as output. The training goal is to condition on $x_{<T_{p}}$ and autoregressively generate the remaining text sequence $x_{\geq T_{p}}$:

Warm Starting the Model In order to pre-train all components of our model end-to-end, we need to warm start the retriever at a good state. Otherwise, if starting with random weights, the retriever would often return irrelevant memory items that would never generate useful training signals.

To avoid this cold-start problem, we propose to construct an initial retrieval dataset with pseudo ground-truth knowledge to give the pre-training a reasonable head start. We create a modified version of the Wikipedia-Image-Text (WIT) [37] dataset for this purpose. Each image-caption pair in WIT also comes with a corresponding Wikipedia passage (words surrounding the text). We put together the surrounding passage with the query image and use it as the pseudo ground-truth knowledge that corresponds to the input query. As the passage provides rich information about the image and caption, it definitely is useful for initializing the model. To avoid the model from relying on low-level image features for retrieval, we apply random data augmentation to the input query image. Given this modified dataset that contains pseudo retrieval ground-truth, we train the query and memory key embeddings by optimizing the following contrastive loss:

where $\hat{z}$ represents the pseudo ground-truth knowledge entry corresponding to the input query $x$.

We use the following four sources of knowledge in our experiments: Wikipedia-Image-Text (WIT) [37] consists of the images in Wikipedia, as well as their alt-text captions and contextualized text passages. Conceptual (CC12M) [5] contains web images paired with alt-text captions. It includes many long-tail entities. VQA-v2 [12] is a visual question answering dataset. We merge all question-answer pairs per image into a single passage. WikiData [40] is a structural knowledge graph encoding relations between Wikipedia entities. We linearize all relational triplets per entity into a textual passage following the procedure of [32]. We have listed the statistical details of these knowledge sources in Table 1.

Incorporating all the components introduced above, ReVeaL can be directly pre-trained over large-scale image caption datasets after proper initialization. As our model architecture is based on T5 and ViT, we use pre-trained ViT checkpoints from [50] and pre-trained T5 checkpoints from [33] to initialize the encoder parameters. The query head, key head and attentive fusion layers are initialized from upper T5, while the base text encoder is initialized from lower T5. The combination of these modules can be found in Table 2 for three three model variants, ReVeaL-Base, ReVeaL-Large and ReVeaL, of which the largest ReVeaL model has around 2 billion parameters.

Distributed Online Retrieval. Finding the top-k most-relevant knowledge entries is a standard Maximum Inner Product Search (MIPS) problem. There are approximate search algorithms [36, 8] that scale sub-linearly with the size of the knowledge corpus $|C|$. We use TPU-KNN [8] to conduct distributed MIPS search, by splitting and storing the memory embeddings across all training devices. The query is synced to each device, which retrieves approximate top-K results from its own memory. Then these results are combined to compute the global top-K retrieved items.

Pre-Training Pipeline. We first train the multimodal retriever on our modified version of the Wikipedia Image Text (WIT) dataset via $\mathcal{L}_{contra}$. We use the Adafactor optimizer without momentum ($\beta_{1}=0$, $\beta_{2}=0.999$), with weight decay of 0.001²²2The remaining experiments use the same optimizer configuration., and with a peak learning rate of $6e4$, to train for 10 epochs. We use this checkpoint to warm-start our generative pre-training. We set the number of retrieved knowledge entries as $K=10$ during pre-training, and use adafactor with a peak learning rate of $1e{-3}$ and inverse squared root learning rate scheduler with 10,000 linear warm-up steps. We use $\mathcal{L}_{PrefixLM}$ as the main objective, adding $\mathcal{L}_{contra}$, $\mathcal{L}_{decor}$ and $\mathcal{L}_{align}$ weighted by $0.01$. We use a batch size of 4096 across 256 CloudTPUv4 chips and train for about 5 days.

One of the most knowledge intensive visual-language tasks is knowledge-based visual question answering (VQA), exemplified by the OK-VQA [29] and A-OKVQA [35] benchmarks. To finetune our pre-trained model on these VQA tasks, we use the same generative objective where the model takes in an image question pair as input and generates the text answer as output. There are a few differences between the fine-tuning and the pre-training stages: 1) we set the number of retrieved knowledge entries to $K=50$, so the model is able to retrieve sufficient supporting evidence; 2) we freeze the whole base V-L encoder to stabilize training; and 3) we use a batch size of 128, with the Adafactor optimizier, a peak learning rate of 1e-4. We use the soft VQA accuracy metric [3] to evaluate the model’s generated answer.

Our results on OKVQA and A-OKVQA datasets are shown in Table 3 and Table 4 respectively. For OKVQA, earlier attempts that incorporate a fixed knowledge retriever report results that are below 45%. Recently a series of works utilize large language models (e.g. GPT-3) as implicit knowledge sources, which achieve much better performance with the trade-off of a huge computational cost. ReVeaL achieves higher performance than those methods without relying on such large language models³³3As shown in the last column of Table 3, ReVeaL stores external knowledge as value embeddings on disk, occupying 993GB of space. The key embeddings consume 10GB space and are kept in TPU memory for fast lookup. On the other hand, KAT and REVIVE need to load the entire 350GB GPT-3 model in the GPU/TPU memory. Furthermore, storing WikiData on disk consumes 500GB of disk memory.. Compared with the previous state-of-the-art, KAT and ReVIVE, which also utilizes T5-Large as a generator, ReVeaL achieves accuracy of 59.1%, which is +6.0% higher than the single KAT [13] model and +2.5% higher than ReVIVE [24].

On A-OKVQA, ReVeaL achieves 52.2% accuracy, which is +3.6% higher than the previous best, GPV-2 [21]. We also show two examples of these datasets in Figure 4. All these results show that, with proper end-to-end retrieval training and a diverse set of knowledge sources, ReVeaL can learn to retrieve meaningful knowledge entries, and achieve promising results without relying on a large language model.

We also evaluate ReVeaL on image captioning benchmarks: MSCOCO Captions [6] and NoCaps [1]. We follow the evaluation protocol used in [49]. We directly fine-tune our generator model on the MSCOCO training split via cross-entropy generative objective. We measure our performance on the MSCOCO test split and NoCaps val set with the CIDEr metric [39]. The results of these two datasets are shown in Table 5. Note that ReVeaL achieves better results than strong recent baselines such as SimVLM [45] and CoCa [49] on both benchmarks. Notably, ReVeaL $-{\text{Large}}$ with 1.4B parameters outperforms the 2.1B-parameter CoCa model and is significantly better than 80B-parameter Flamingo model [2].

In the following we study which design choices contribute most to the model’s performance. We focus on three research questions: (1) Does utilizing multiple knowledge sources enhance performance? (2) Does the proposed attentive fusion surpass existing end-to-end retrieval training methods? (3) Can we add knowledge by only updating the memory without modifying model parameters?

Analyzing multiple knowledge sources. A major distinction of ReVeaL compared to previous retrieval-augmented approaches is its capacity to utilize a diverse set of knowledge sources during inference. To assess the relative importance of each data source and the efficacy of retrieving from various corpora, we conduct two ablation studies: 1) Only-One-Left: employing a single knowledge source to evaluate the outcomes; and 2) Leave-One-Out: excluding one knowledge source from the complete set $\mathcal{C}$. These ablation studies are executed using the ReVeaL ${}_{\text{Base}}$, evaluated on the OKVQA validation set under the aforementioned conditions. As shown in Figure 6, among the four knowledge sources utilized in this paper, WIT is the most informative, with the highest accuracy when used in isolation (53.1The remaining three corpora, CC12M, VQA-v2, and WikiData, do not offer the same level of informativeness as WIT when utilized independently. However, excluding any of these corpora from the complete dataset results in performance decreases of 1.3%, 0.6%, and 1.1%, respectively. This observation implies that these knowledge sources effectively complement one another, contributing valuable information to enhance performance. To further substantiate this hypothesis, we perform an additional experiment involving pairs of knowledge sources, as illustrated in Figure 6. Notably, even when paired with an informative knowledge source such as WIT, incorporating an extra corpus consistently leads to performance improvements.

Analyzing different retrieval training methods. Another core component of ReVeaL is the attentive fusion layer, which supports efficient joint training of the retriever and generator. We investigate its performance compared to two existing retrieval training method categories: 1) a frozen retriever based on ALIGN [20] representations ; 2) end-to-end retrieval training methods including Attention Distill [17], EMDR² [47], and Perplexity Distill [18].

We use the pre-trained ReVeaL-Base model, fix the generator and randomly initialize the retriever (query head and key head). We utilize our modified version of WIT dataset with pseudo ground-truth retrieved labels as the evaluation corpus. We evaluate retrieval performance by checking whether the correct passage appears in top-10/100 results. For the ALIGN model, we directly evaluate the retrieval results from the pre-trained checkpoint, while for other models, we perform retrieval-augmented training on the WIT dataset. To prevent the model from relying on image similarity for accurate results, we only use text passages as knowledge entries and discard images. Subsequently, we finetune the model on OKVQA and report its accuracy. The results are presented in Table 6. We observe that directly using pre-trained encoder does not perform well, even with a strong model like ALIGN. Moreover, among the various end-to-end retrieval training approaches, our attentive fusion method attains better accuracy in both retrieval and OKVQA tasks. Importantly, our technique exhibits a computational cost (quantified by GFLOPs) comparable to that of attention distillation, yet significantly lower than EMDR² and Perplexity distillation. This indicates that our proposed method is more efficient and effective for pre-training retrieval-augmented visual-language models.

Analyzing Knowledge Modification. One advantage of utilizing knowledge memory is that we could easily add or update knowledge entries without re-training model’s parameters. To validate this, we conducted ablation studies in which we removed a specific percentage of knowledge entries from the corpora and assessed the performance of the ReVeaL-Base model on the OKVQA dataset. Subsequently, we add the removed knowledge back into the corpora, allowing the trained model to make predictions using the complete set of corpora. This approach ensured that the removed knowledge was not seen by the model during fine-tuning, enabling us to test its ability to accurately retrieve and utilize that knowledge for problem-solving.

The results are illustrated in Figure 7, with the blue curves representing the inference outcomes without the removed knowledge and the orange curve depicting the results after adding the removed knowledge back. A notable performance improvement was observed upon reintroducing the knowledge (orange curve) compared to the outcomes with the removed knowledge (blue curve). Specifically, for the model fine-tuned with only 10% of the knowledge, the reintroduction of the removed knowledge resulted in an accuracy of 51.8 (+6.7 higher than when removed). This finding demonstrates that the ReVeaL model can swiftly adapt to new knowledge by merely updating the memory, obviating the need for re-training model parameters.