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

In plain words

REVEAL is a model that answers questions about pictures and writes captions. Instead of packing every fact into its weights, it keeps a large lookup store built from several sources: image-caption pairs, Wikipedia-style passages, question-answer pairs and knowledge-graph facts. For each query it searches that store for the most useful entries, reads them together with the image and the question, and writes the answer.

The unusual part is how it is trained. The search step and the writing step are trained together on one language-modelling objective, so the model learns which stored entries help. To make the search trainable, each retrieved entry's weight scales its tokens inside the attention layers, so the answer loss flows back into the search.

Because facts live in the store rather than in the weights, entries can be added or changed without retraining the network. The authors test this by removing knowledge and putting it back. They report top results on two knowledge-based visual question answering benchmarks and on image captioning, without depending on a very large language model.

TL;DR

• A visual-language model (T5 text side, ViT image side) that retrieves from a memory of encoded multimodal knowledge: K=10 entries in pre-training, K=50 for OK-VQA fine-tuning §3.2 §4.3 §5.1.
• The retriever gets no labels. It learns through the generation loss, via an attentive fusion layer, after a contrastive warm start on pseudo-ground-truth passages §3.4 §4.1.
• Author-reported: 59.1% on OK-VQA, 52.2% on A-OKVQA, 145.4 CIDEr on MSCOCO (2.1B model) tab. 3 tab. 4 tab. 5.
• Relevant if you need facts that change without retraining. The released code differs from the paper in several weights and settings; read "Paper vs code" before reimplementing.

Problem

Large vision-language and language models store world knowledge in their parameters, which needs very large scale and full retraining whenever the knowledge changes §1.

Earlier retrieval-augmented models mostly use one modality and a frozen retriever, so they miss useful evidence from other modalities §1 §2. End-to-end training of a multimodal retriever has no direct supervision: no label says which entries help §1. Training signals that score each retrieved entry on its own are costly with many entries and ignore how entries depend on each other §1 §2.

Method

REVEAL has four parts: a query encoder, a memory of encoded knowledge, a retriever, and a generator §3 fig. 2. The model learns the distribution of an output text given a multimodal query. Retrieval is a latent variable, and the marginal is approximated by its top-K entries eq. 1 §3.

p(y \mid x) = \sum_{M \subset \tilde{\mathcal{M}}} p(M \mid x)\, p(y \mid x, M) eq. 1

The paper says the sum over all size-K subsets is intractable, so it uses the K entries with the highest retrieval probability instead §3.

Query encoding §3.1. An image goes through a ViT, and a projection maps its tokens into the text token space. Text goes through the lower layers of a T5 encoder. The two token sequences are concatenated. A query head (the upper T5 layers) reads the first [CLS] output, applies a linear projection and L2-normalises it into a d-dimensional query. The key head has the same design and produces keys for memory entries §3.1.

Memory §3.2. Each knowledge item may be a paired image and caption, plain text, a bare image, or a knowledge-graph triplet, and becomes one key and one value. The key is a d-dimensional embedding from the key head. The value is a sequence of tokens. The value is compressed by a Perceiver-style head: a transformer decoder with c learned latent embeddings, where the paper uses c=32 §3.2. Keys and values are precomputed for the whole corpus and refreshed every 1000 training steps §3.2. This procedure follows REALM §3.2.

Two regularisers shape the compressed values §3.2:

• Decorrelation: a penalty on linear correlation between output tokens, summed over pairs of knowledge items. The formula is written as a sum of squared Frobenius norms of covariances, indexed over i, j = 1..K.
• Alignment: pulls the norm of compressed values towards the norm of the query encodings, as a ratio of sums over the batch.

\mathcal{L}_{\text{decor}} = \sum_{i,j=1}^{K} \Big\lVert \text{Cov}\big(\psi(b(z_i)), \psi(b(z_j))\big) \Big\rVert_F^2 §3.2

\mathcal{L}_{\text{align}} = \Big\lvert 1 - \frac{\sum_z \lVert \psi(b(z)) \rVert_2}{\sum_x \lVert b(x) \rVert_2} \Big\rvert §3.2

Retrieval §3.3. A softmax gate over knowledge sources, computed from the query, weights each corpus. The relevance of an entry is the dot product of query and key. The probability of an entry is the corpus gate times a softmax over relevances within its corpus eq. 2 eq. 3:

p(m^j_i \mid x) = p(\mathcal{M}^j \mid x)\, p(m^j_i \mid x;\, \mathcal{M}^j) eq. 2

p(m^j_i \mid x) = \text{Gate}_{\mathcal{M}^j}(x) \cdot \frac{\exp\big(\text{Rel}(x, m^j_i)/\tau\big)}{\sum_{m^j_k \in \mathcal{M}^j} \exp\big(\text{Rel}(x, m^j_k)/\tau\big)} eq. 3

Here Gate is a softmax over corpora of a linear map of the query, Rel(x, m) is the query-key dot product, and τ is a temperature. The paper approximates the probability of a set M as a product over its entries §3.3.

The top-K entries are returned from memory. A small share of them, which the paper states as 10%, is re-encoded from raw data so that the encoders receive gradients. The rest use the cached key and value §3.3.

Generator and attentive fusion §3.4. The query tokens and the K compressed value sequences are concatenated, giving I + c·K tokens of width d. A soft mask made of the query's ones and each retrieved entry's probability is applied to the token sequence, so each entry's weight follows its retrieval probability §3.4. The fused sequence goes to a T5 decoder that generates the output text §3.4 fig. 3. This is how the retrieval score receives gradients from the generation loss §3.4.

# REVEAL forward pass for one query  [§3.1] [§3.3] [§3.4]
# Input:  image (B, H, W, 3); text_tokens (B, L)
# Output: decoder logits (B, T, vocab)

def reveal_forward(image, text_tokens, memory, K):
    # Encode the query: ViT tokens plus lower-T5 tokens  [§3.1]
    query_tokens = encode_tokens(image, text_tokens)   # (B, I, d)
    # Summarise with the query head; L2-normalised  [§3.1]
    query = l2_normalise(query_head(query_tokens))     # (B, d)
    # Rank entries by gated retrieval probability  [eq. 2] [eq. 3]
    gated_probability = score_entries(query, memory)   # (B, N)
    # Keep the K most probable entries  [§3.3]
    top_ids = topk(gated_probability, K)               # (B, K)
    # Re-encode a small share of the K entries from raw data  [§3.3]
    fresh_values = encode_entries(memory.raw[top_ids[:, :n_fresh]])
    # Use cached values for the rest; cache refreshed every 1000 steps  [§3.2]
    cached_values = memory.values[top_ids[:, n_fresh:]]
    values = concat(fresh_values, cached_values)       # (B, K, c, d)
    # Retrieval probabilities of the K entries  [§3.4]
    prior = gated_probability[top_ids]                 # (B, K)
    # Scale each entry's tokens by its prior, then fuse  [§3.4]
    fused = attentive_fusion(query_tokens, values, prior)
    # Decode the answer or caption from the fused sequence  [§3.4]
    return t5_decoder(fused)

Hyperparameters

Results

Author-reported. Setup: OK-VQA and A-OKVQA for knowledge-based VQA, with soft VQA accuracy §5.1. MSCOCO (Karpathy test split) and NoCaps (validation) for captioning, with CIDEr §5.2. Pre-training on Web-Image-Text, and knowledge from WIT, CC12M, VQA-v2 and WikiData §4.2. Hardware for pre-training: 256 TPU v4 chips §4.3.

Knowledge sources tab. 1. Counts and text lengths as listed; length unit not stated.

Model variants tab. 2.

OK-VQA tab. 3. Memory is given as learnable + frozen on-device + CPU/disk, in GB, where the paper gives it.

The text gives +6.0 over single KAT and +2.5 over single ReVIVE §5.1. The paper's footnote says the 993 GB of values sit on disk, the 10 GB of keys stay in TPU memory, and the frozen GPT-3 models need about 350 GB on device §5.1.

A-OKVQA tab. 4. Accuracy (%).

The text reports +3.6 over GPV-2 §5.1.

Image captioning, CIDEr tab. 5.

Retrieval training methods tab. 6. The REVEAL-Base generator is frozen, the retriever is trained from random initialisation on text-only WIT passages, and OK-VQA accuracy is measured after fine-tuning §5.3.

Ablations, from the text and figures §5.3 fig. 5 fig. 6 fig. 7. These use REVEAL-Base on the OK-VQA validation set:

• WIT alone gives the highest single-source accuracy, 53.1%.
• Leaving out CC12M, VQA-v2 or WikiData lowers accuracy by 1.3, 0.6 and 1.1 points.
• Pairing WIT with any second source improves on WIT alone fig. 6.
• Knowledge update: after fine-tuning with 10% of the knowledge, putting the removed knowledge back into memory (no retraining) gives 51.8%, which is +6.7 over the version without it §5.3 fig. 7.

Paper vs code

Code: the official code is the Scenic project folder scenic/projects/knowledge_visual_language, linked from the project page at reveal-cvpr.github.io, in the google-research/scenic repository at code_commit (Apache-2.0). Citations below are relative to that folder. The released configs are for pre-training (configs/wit_memory_*.py), warm start (configs/wit_retrieval_soft_*.py) and OK-VQA fine-tuning (configs/finetune_okvqa_base.py). Model code is in models/knowledge_fid.py, models/fusion_in_decoder_soft.py and models/layers.py.

Discrepancies found, each with both citations:

1. Share re-encoded. Paper: 10% §3.3. Pre-training configs use 0.2, so with K=10 the code re-encodes ceil(K × 0.2) = 2 entries, not 1 wit_memory_base.py knowledge_fid.py. Fine-tuning uses 0.1, which matches finetune_okvqa_base.py.
2. Loss weights. Paper: L_PrefixLM plus 0.01 times each of L_contra, L_decor and L_align §4.3. Code: 0.99 times the generation loss plus 0.01 times the contrastive loss, using the config's retrieval_ratio knowledge_fid.py wit_memory_base.py. The decorrelation term is added with weight 1e-2 knowledge_fid.py. The alignment term is added with weight 1e-4, not 0.01 knowledge_fid.py wit_memory_base.py.
3. Decorrelation form. The paper's formula sums over pairs of knowledge items i, j in 1..K §3.2. The code decorrelates the c output tokens of each entry: it centres them across tokens, normalises each token, squares the pairwise cosines and averages layers.py. The regulariser acts within each entry, not across entries.
4. Perceiver depth and scale. Code uses num_fusion_layers (6 for Base) as the Perceiver depth too, and multiplies its output by 4 fusion_in_decoder_soft.py layers.py. Neither is stated in the paper.
5. Corpus gate placement. The paper multiplies the gate outside the softmax eq. 3. Code multiplies each retrieved key by its entry's gate before fusion, so the gate acts on the logit knowledge_fid.py. Top-K selection ranks by gate times dot product, not by the probability of eq. 3 knowledge_fid.py.
6. Fusion prior. Code computes softmax over an affine transform of the dot products, with a learned scale (sigmoid times 5) and bias, and then multiplies by K fusion_in_decoder_soft.py layers.py. The paper's mask is p(z|x) with no factor K and a fixed temperature τ §3.4.
7. Pseudo-ground-truth replacement (code only). During training, with probability 0.02 per example, retrieved keys and values are replaced by the pseudo-ground-truth entry knowledge_fid.py, enabled by use_psudo_retr wit_memory_base.py. The paper does not mention this.
8. Knowledge sources. The paper uses four corpora including WikiData §4.2. Released configs list three: wit_table, cc12m_table and vqa_table wit_memory_base.py finetune_okvqa_base.py. A WikiData loader and linearisation script exist in data/wikidata/. Whether the published run used them cannot be determined from the configs.
9. Pre-training hyperparameters. The paper gives peak learning rate 1e-3, batch 4096, 10,000 warm-up steps and 256 chips §4.3. No config matches all four. The base config has learning rate 1e-4, batch 3072 on 128 devices, and warm-up 10,000 wit_memory_base.py wit_memory_base.py wit_memory_base.py. The 256-device config for the ViT-g model, wit_memory_g.py, has learning rate 1e-3, batch 2048, warm-up 20,000 and 2 epochs wit_memory_g.py wit_memory_g.py wit_memory_g.py wit_memory_g.py.
10. Warm-start settings. The paper gives 10 epochs and weight decay 0.001 §4.3. The closest config, wit_retrieval_soft_base_config.py, runs 100 epochs with weight decay 1e-2 and a retrieval ratio of 1/2 wit_retrieval_soft_base_config.py wit_retrieval_soft_base_config.py wit_retrieval_soft_base_config.py. Which config produced the paper's warm start is not clear.
11. Fine-tuning settings. Paper: batch 128, peak learning rate 1e-4 §5.1. Code: batch 1536 finetune_okvqa_base.py, learning rate 1e-5 finetune_okvqa_base.py, 1000 epochs finetune_okvqa_base.py and no weight decay finetune_okvqa_base.py.
12. Memory refresh during fine-tuning. Paper: every 1000 steps §3.2. The fine-tuning config evaluates every 500 steps, and the memory is refreshed at those evaluation steps finetune_okvqa_base.py trainer_memory.py.
13. Dataset size for schedules. Configs assume 1e9 training examples wit_memory_base.py. The paper says about 1.3B filtered pairs §4.1. Step counts in the configs will not match the paper's 5-day run.

The code detaches the query before the top-K search, so gradients reach the retriever only through the fusion attention and the gate knowledge_fid.py. The paper's text is consistent with this. The query and key head internals were not checked line by line.

Pitfalls

• The appendix is not in the compiled paper. In main.tex the appendix include is commented out. The appendix's K and c sensitivity, corpus-size ablation, visual-only retrieval baseline, Perceiver algorithm and fusion algorithm are therefore not part of the compiled paper. This card does not cite them. The Perceiver and fusion designs are described here from the main text and the code.
• The contrastive warm start is underspecified. The paper writes one log-softmax over a single dot product with the pseudo-ground-truth key, without naming negatives §4.1. The code uses in-batch negatives across all devices knowledge_fid.py. Implement the code's version or state the choice.
• The warm start is necessary. With random retriever weights, the retriever returns irrelevant entries that give no useful training signal §4.1. The compiled paper gives no number for this, because the ablation sits in the commented-out appendix.
• Stale memory. Most of the K entries use values cached up to 1000 steps earlier. Only the re-encoded share follows the current encoder §3.2 §3.3. The paper does not measure the effect of this lag.
• Retrieval cost is high. The paper uses distributed approximate MIPS across all devices (TPU-KNN) §4.3. Values for the full memory occupy 993 GB on disk for the largest model §5.1. Pre-training took about 5 days on 256 chips §4.3.
• Temperature and dimension are not stated. τ in eq. 3 and the key dimension d are not given in the compiled paper eq. 3 §3.2. The code's values are learned or set in configs (see Hyperparameters), so the paper's model cannot be reproduced exactly from the text alone.
• Fusion layer count is not stated. The compiled text does not give the number of fusion layers. The code uses 6 for Base and 8 for Large and REVEAL wit_memory_base.py wit_memory_g.py.
• Figure and table numbers follow source order in the LaTeX. Cite the section if the numbering in the rendered paper differs.
• Results are author-reported. Several rows (GPT-3-based models) use different knowledge sources and far larger frozen models, so per-parameter comparisons are not like-for-like §5.1 tab. 3.
• Wikidata linearisation follows a cited procedure, not one given in the paper §4.2. Reimplementing it needs that reference.

Lineage

• Builds on: REALM (arXiv:2002.08909) for the precomputed key-value memory procedure §3.2; ViT and T5 for the encoders and checkpoints §3.1 §4.3; Perceiver for value compression §3.2; Web-Image-Text and WIT for pre-training and warm start §4.1 §4.2.
• Compared against: KAT, ReVIVE, PICa, Flamingo, CoCa, SimVLM, GPV-2, MAVEx and KRISP on VQA and captioning tab. 3 tab. 4 tab. 5. Retrieval-training methods: Attention Distill, EMDR², and Perplexity Distill (Atlas, arXiv:2208.03299) tab. 6.
• Superseded by: not known from content already in this wiki.