# Impact of novel aggregation methods for flexible, time-sensitive EHR
prediction without variable selection or cleaning

Jacob Deasy, Ari Ercole and Pietro Liò
University of Cambridge
{jd645, ae105<EMAIL_ADDRESS>
###### Abstract

Dynamic assessment of patient status (e.g. by an automated, continuously
updated assessment of outcome) in the Intensive Care Unit (ICU) is of
paramount importance for _early alerting, decision support_ and _resource
allocation_. Extraction and cleaning of expert-selected clinical variables
discards information and protracts collaborative efforts to introduce machine
learning in medicine. We present improved aggregation methods for a _flexible
deep learning architecture_ which learns a joint representation of patient
chart, lab and output events. Our models outperform recent deep learning
models for patient mortality classification using ICU timeseries, by embedding
and aggregating all events with _no pre-processing or variable selection_. Our
model achieves a strong performance of AUROC 0.87 at 48 hours on the MIMIC-III
dataset while using 13,233 unique un-preprocessed variables in an
interpretable manner via hourly softmax aggregation. This demonstrates how our
method can be easily combined with existing electronic health record systems
for automated, dynamic patient risk analysis.

##  1 Introduction

In the United States, over a third of hospitals now utilize Electronic Health
Record (EHR) databases that are considered broad enough to be comprehensive
[1]. The comprehensive, but complex, information contained in such EHRs gives
a full description of the patient’s clinical journey and is a promising source
of data not only for retrospective studies, but also for building
sophisticated decision support systems to improve care delivery [2].

However, a systematic review of 107 predictive models built with EHR data
found that only 34.6% (37/107) used longitudinal data, and the median number
of variables used was only 27 [3, 4]. The Intensive Care Unit (ICU) has
particularly high longitudinal data density since patients’ conditions may
change on timescales of minutes. Despite this, current mortality risk
estimates are often solely based on acute physiology scores [5]. These metrics
are static and typically reliant upon logistic regression of specific markers
of patient physiology recorded during the first hours after ICU admission.
Such approaches are likely to both be sub-optimal in terms of predictive power
and unsuitable for decision support as they cannot be updated in real-time.

This work introduces novel aggregation methods for a flexible method of
learning from _all chart, lab and output events_ , regardless of type,
frequency or cardinality. Our deep learning architecture achieves strong in-
hospital mortality prediction at 48 hours of AUROC 0.878 (95% CI: 0.871-0.883)
on patients in the MIMIC-III dataset [6], without performing any variable
selection or pre-processing. Instead, each model learns a powerful method of
embedding and aggregating discrete variable categories from raw EHR data. We
assess an aggregation variant which builds in model interpretability and find
that the predictive performance cost is negligible—with models still
performing in line with recent strong deep learning architectures. Any
sequential EHR data can be incorporated into our model to provide more
interpretable and personalized dynamic risk prediction, with reliable dynamic
performance exceeding traditional severity scores after only a few hours in
the ICU.

##  2 Related Work

Recurrent Neural Networks (RNNs) [7, 8] have been shown to outperform
traditional heuristic ICU scores for prediction of mortality and long Length
Of Stay (LOS) risk [2, 4]. Meanwhile, state-of-the-art models using sequential
data have grown in size and complexity and been adapted to address the unique
challenges of the ICU [4, 9]. In all of these cases, the dominant approach to
deep learning with EHR data remains reliant on an initial stage of variable
selection involving the use of expert knowledge to hand-pick a subset of
clinically relevant variables [9, 10]. Even on the broad MIMIC-III dataset,
clinical variable selection has recently been somewhat standardized by the
introduction of a benchmark [10]. Larger studies have managed to avoid
variable selection by building pipelines which convert their standard data
type into a format suitable for machine learning [4, 11]. Although this has
led to strong prediction results, these architectures do not generalize well
to EHRs and regions where little effort is being made to produce a
standardized EHR database.

On the other hand, previous research has shown that using sparsity-inducing
priors in deep Bayesian neural networks can automatically reduce the number of
clinical variables attended to by an _AI clinician_ [12], with little impact
on performance [13]—calling into question the need for variable selection in
the first place. This has lead to the recent introduction of the first deep
learning pipeline which only requires that EHR time series contain a sequence
of timestamps and a sequence of variable readings [14]. We extend this work by
exploring embedding aggregation variants that can process arbitrarily many
patient readings per hour. We create a pipeline and dynamic model that perform
zero variable selection, data processing or model ensembling with performance
boosting aggregation variants. Our models aggregate varied patient demographic
and physiological data in an interpretable manner to provide a more
clinically-relevant and accurate outcome prediction method in the ICU. As well
as exploring the power of embedding aggregation functions, by producing a
time-sensitive model, our approach tackles the inherent and oft-neglected need
for dynamic assessment in the ICU [15].

##  3 Model and inputs

###  3.1 Dataset

Table 1: Mortality class distribution and ICU lengths of stay information for our MIMIC-III subset. Model | Patients | Mean LOS | Median LOS
---|---|---|---
Died | 2,797 (13.2%) | 8.11 | 5.35
Survived | 18,342 (86.8%) | 5.65 | 3.50
Total | 21,139 | 5.97 | 3.72

We used three types of data from the MIMIC-III dataset for classification: (a)
_chart event_ recordings of routine vital signs from the electronic chart
forming the bulk of patient information, (b) _lab event_ results of laboratory
tests on blood and urine, and (c) _output event_ recordings of admission
information and treatments. The objective is to dynamically assign a risk of
mortality to each patient based on their event time series during the first 48
hours after their admission to the ICU.

The MIMIC-III database contains high-resolution patient data, including:
demographics, vital sign time-series, laboratory tests, illness severity
scores, medications and procedures, fluid intake and outputs, clinician notes,
and diagnostic coding. The median age of adult patients is 65.8 years
($Q_{1}$–$Q_{3}$: 52.8–77.8), 55.9% of patients are male, and in-hospital
mortality is 11.5% [6]. Between the two EHR systems that comprise MIMIC-III,
CareVue and MetaVision, in total the overall dataset contains 330,712,483
chart events, 27,854,055 lab events and 4,349,218 output events.

The architecture was trained using a subset of patient stays where the
proportion of in-hospital deaths was 13.2%. The proportion of long ICU stays
(greater than 7 days in the ICU) was 23.0% (4,868/21,139)—see Table 1 for a
summary of mortality and LOS statistics.

###  3.2 Data pipeline

Unlike traditional approaches, we retain all of the chart, lab and output
events for each stay without any data cleaning, outlier removal or domain-
specific knowledge. The processing we perform is enough to assign a patient
ID, a stay ID and a timestamp to each event—a pipeline which is independent of
EHR data formatting or structure. Our model takes the entire patient
timeseries as input, regardless of event type, frequency or cardinality and
discretizes each event (see Table 2). We note that, because we do not select
for clinical variables, after event association with patient stays, our EHR
dataset contains 208,572,237 events instead of the 31,868,114 employed in [10]
and all subsequent papers relying on the MIMIC-III benchmark. Due to the
increased number of variables used by our model, after processing we also have
a higher number of patients and stays available, supporting the evidence in
favour of models which can incorporate broad EHR data.

Table 2: Examples of discrete token creation alongside percentile-based quantization and tokenization of continuous variable. Discrete variables remain discrete, while values that can be converted to floating point numbers are considered continuous and separated into 20 percentile-based bins. Label | Value | Percentile | Token
---|---|---|---
Blood pH | 7.41 | 55-60% | Blood pH_12
Heart Rate | 69 | 35-40% | Heart Rate_8
Code Status | Full Code |  | Code Status Full Code

###  3.3 Model architecture

We use an embedding layer followed by an aggregation function to reduce an
arbitrary number of inputs to a fixed size representation (see Figure 1),
before passing this to a recurrent neural network to classify patient
mortality. All chart, lab and output events are tokenized and embedded; the
resulting set of vectors are aggregated over each hour by one of the
following: summation, average, weighted average or a softmax (masked by hour).
Dynamic classification at each time step is achieved by a dense layer with
sigmoid activation applied to each hidden state.

012Hours since
admission00:1400:2300:4701:1301:13$\vdots$01:3801:5902:03$\vdots$743002869674243$\vdots$957374447$\vdots$$e_{6742}$$e_{43}$$\vdots$$e_{9573}$$e_{74}$$e_{74}$$e_{3002}$$e_{869}$EmbeddingAggregation$f_{\textrm{agg}}\left(\\{e_{i_{t}}\\}_{0\leq
t<1}\right)$$f_{\textrm{agg}}\left(\\{e_{i_{t}}\\}_{1\leq
t<2}\right)$$f_{\textrm{agg}}\left(\\{e_{i_{t}}\\}_{2\leq t<3}\right)$$\vdots$
Figure 1: Example of our flexible EHR embedding for a single patient time
series. EHR token indices are embedded and reduced to a fixed size hourly
representation by aggregation function $f_{\textrm{agg}}$.

##  4 Experiments and results

###  4.1 Evaluation procedure

We used $k$-fold cross-validation with $k=10$ to assess model performance on a
held out test set of 2,114 patients (10,000 bootstrap samples were used to
construct confidence intervals). All models were trained for 50 epochs with an
early stopping threshold of 5 epochs with no increase in AUROC on the
validation set of 1,000 patients. All models were trained with a batch size of
128, using the Adam optimizer [16], with a learning rate of 0.001 and default
hyperparameters otherwise. As the distribution of mortality across the MIMIC-
III dataset is imbalanced (see Table 1), accuracy is an unsuitable metric so
we employ AUROC score to evaluate model performance.

###  4.2 Model parameters

Our baseline model contains 38,780 embeddings in a 32D space, generated by
discretizing 13,233 unique variable labels, with 20 bins per continuous
variable. We utilize a GRU [17] of depth 1 with a sigmoid output activation,
layer normalization applied to all linear projections [18], dropout
regularization [19] ($p=0.5$) on both the aggregated embedding and the hidden
states, as well as a weight decay coefficient of 0.001 [20]. We assess the
effect of using summation, simple averaging and a masked softmax (over a
separate token weight vector) as aggregation functions on model performance.
We note that the risk of unnormalized aggregation, such as summation, is that
the model learns to count the number of readings taken in an hour and
correlate this with patient outcome.

###  4.3 Results

Table 3 demonstrates that integrating events in a unified manner results in
strong early classification of in-hospital mortality (AUROC $0.88$)—stronger
than both traditional ICU scores and recent deep learning models despite no
feature engineering. AUROC of $0.88$ means that there is an $88$% chance of
assigning higher risk to a random patient destined to die than a random
patient destined to live. If clinical resource allocation were based on our
model rather than SAPS II, 25% more patients would be correctly prioritized.
In particular, the softmax aggregation model variant automatically learns
comparative hourly weights for each EHR token, building in interpretability
with negligible cost to performance. Exploratory work on this weighting system
has validated the association of well-known mortality indicators, such as high
respiratory rate, resuscitation instructions and even documentation of visits
by the priest, and opens avenues for future research on patient-specific risk
variation.

As such, this research presents a model that is flexible enough to reduce the
requirements on sequential EHR data to simply timestamps alongside variable
readings. Percentile-based quantization followed by embedding aggregation also
means our model automatically learns how _outliers_ and _missing values_
influence patient mortality risk. By eliminating the need for variable
selection and generating comparative weights while outperforming a strong deep
learning baseline, our aggregation improvements allow machine learning
specialists to consult clinicians about highly weighted variables _after_
model development, rather than before. Future work will focus on building
models that learn to convert any of the diverse range of design principles and
computational practices currently employed across EHR databases to the form
used by our model.

Table 3: Cross-validation results for prediction of mortality at 12 and 48 hours for the embedding aggregation functions, compared to models that make use of carefully selected predictor variables. Model/Aggregation | # curated features | 12h AUROC (95% CI) | 48h AUROC (95% CI)
---|---|---|---
OASIS [21] | 10 | — | 0.663
SAPS II [22] | 17 | — | 0.705
Benchmark [10] | 76 | — | 0.870 (0.852-0.887)
Weighted average [14] | — | 0.798 (0.794-0.802) | 0.856 (0.851-0.861)
Masked softmax | — | 0.805 (0.796-0.814) | 0.872 (0.862-0.881)
Average | — | 0.811 (0.807-0.815) | 0.875 (0.873-0.877)
Summation | — | 0.810 (0.804-0.815) | 0.878 (0.871-0.884)

## References

  * [1] Adler-Milstein J, DesRoches CM, Kralovec P, Foster G, Worzala C, Charles D, et al.  Electronic health record adoption in US hospitals: progress continues, but challenges persist.  Health affairs. 2015;34(12):2174–2180.
  * [2] Johnson AE, Ghassemi MM, Nemati S, Niehaus KE, Clifton DA, Clifford GD.  Machine learning and decision support in critical care.  Proceedings of the IEEE Institute of Electrical and Electronics Engineers. 2016;104(2):444.
  * [3] Goldstein BA, Navar AM, Pencina MJ, Ioannidis J.  Opportunities and challenges in developing risk prediction models with electronic health records data: a systematic review.  Journal of the American Medical Informatics Association. 2017;24(1):198–208.
  * [4] Rajkomar A, Oren E, Chen K, Dai AM, Hajaj N, Hardt M, et al.  Scalable and accurate deep learning with electronic health records.  NPJ Digital Medicine. 2018;1(1):18.
  * [5] Desai N, Gross J.  Scoring systems in the critically ill: uses, cautions, and future directions.  BJAE. 2019;19(7):212–218.
  * [6] Johnson AE, Pollard TJ, Shen L, Li-wei HL, Feng M, Ghassemi M, et al.  MIMIC-III, a freely accessible critical care database.  Scientific data. 2016;3:160035.
  * [7] Elman JL.  Finding structure in time.  Cognitive science. 1990;14(2):179–211.
  * [8] Rumelhart DE, Hinton GE, Williams RJ, et al.  Learning representations by back-propagating errors.  Cognitive modeling. 1988;5(3):1.
  * [9] Che Z, Purushotham S, Cho K, Sontag D, Liu Y.  Recurrent neural networks for multivariate time series with missing values.  Scientific reports. 2018;8(1):6085.
  * [10] Harutyunyan H, Khachatrian H, Kale DC, Galstyan A.  Multitask learning and benchmarking with clinical time series data.  arXiv preprint arXiv:170307771. 2017;.
  * [11] Tabak YP, Sun X, Nunez CM, Johannes RS.  Using electronic health record data to develop inpatient mortality predictive model: Acute Laboratory Risk of Mortality Score (ALaRMS).  Journal of the American Medical Informatics Association. 2013;21(3):455–463.
  * [12] Komorowski M, Celi LA, Badawi O, Gordon AC, Faisal AA.  The Artificial Intelligence Clinician learns optimal treatment strategies for sepsis in intensive care.  Nature Medicine. 2018;24(11):1716.
  * [13] Popkes AL, Overweg H, Ercole A, Li Y, Hernández-Lobato JM, Zaykov Y, et al.  Interpretable Outcome Prediction with Sparse Bayesian Neural Networks in Intensive Care.  arXiv preprint arXiv:190502599. 2019;.
  * [14] Deasy J, Liò P, Ercole A.  Dynamic survival prediction in intensive care units from heterogeneous time series without the need for variable selection or pre-processing.  arXiv preprint submission identifier: submit/2843514. 2019;.
  * [15] Meiring C, Dixit A, Harris S, MacCallum NS, Brealey DA, Watkinson PJ, et al.  Optimal intensive care outcome prediction over time using machine learning.  PloS one. 2018;13(11):e0206862.
  * [16] Kingma DP, Ba J.  Adam: A method for stochastic optimization.  arXiv preprint arXiv:14126980. 2014;.
  * [17] Cho K, Van Merriënboer B, Gulcehre C, Bahdanau D, Bougares F, Schwenk H, et al.  Learning phrase representations using RNN encoder-decoder for statistical machine translation.  arXiv preprint arXiv:14061078. 2014;.
  * [18] Lei Ba J, Kiros JR, Hinton GE.  Layer normalization.  arXiv preprint arXiv:160706450. 2016;.
  * [19] Srivastava N, Hinton G, Krizhevsky A, Sutskever I, Salakhutdinov R.  Dropout: a simple way to prevent neural networks from overfitting.  The Journal of Machine Learning Research. 2014;15(1):1929–1958.
  * [20] Krogh A, Hertz JA.  A simple weight decay can improve generalization.  In: Advances in neural information processing systems; 1992. p. 950–957.
  * [21] Johnson AEW, Kramer AA, Clifford GD.  A new severity of illness scale using a subset of acute physiology and chronic health evaluation data elements shows comparable predictive accuracy.  Critical care medicine. 2013;41(7):1711–1718.
  * [22] Le Gall JR, Lemeshow S, Saulnier F.  A new simplified acute physiology score (SAPS II) based on a European/North American multicenter study.  Jama. 1993;270(24):2957–2963.


# Re-ViLM: Retrieval-Augmented Visual Language Model for
Zero and Few-Shot Image Captioning

Zhuolin Yang  Wei Ping  Zihan Liu  Vijay Korthikanti  Weili Nie  De-An Huang
Linxi Fan  Zhiding Yu  Shiyi Lan  Bo Li  Ming-Yu Liu  Yuke Zhu  Mohammad
Shoeybi  Bryan Catanzaro  Chaowei Xiao  Anima Anandkumar

###### Abstract

Augmenting pretrained language models (LMs) with a vision encoder (e.g.,
Flamingo) has obtained state-of-the-art results in image-to-text generation.
However, these models store all the knowledge within their parameters, thus
often requiring enormous model parameters to model the abundant visual
concepts and very rich textual descriptions. Additionally, they are
inefficient in incorporating new data, requiring a computational-expensive
fine-tuning process. In this work, we introduce a Retrieval-augmented Visual
Language Model, Re-ViLM, built upon the Flamingo, that supports retrieving the
relevant knowledge from the external database for zero and in-context few-shot
image-to-text generations. By storing certain knowledge explicitly in the
external database, our approach reduces the number of model parameters and can
easily accommodate new data during evaluation by simply updating the database.
We also construct an interleaved image and text data that facilitates in-
context few-shot learning capabilities. We demonstrate that Re-ViLM
significantly boosts performance for image-to-text generation tasks,
especially for zero-shot and few-shot generation in out-of-domain settings
with $4\times$ less parameters compared with baseline methods.

Machine Learning, ICML

##  1 Introduction

Image-to-text generation, also known as image captioning (e.g., Karpathy &
Fei-Fei, 2015), plays a vital role in understanding visual information and
enhancing human-AI interaction. This task has a wide range of practical
applications such as gaming, virtual reality, and robotics (Luo et al., 2019;
Zhao et al., 2021; Liu et al., 2021). To address this problem, in recent
years, numerous methods have been proposed and have obtained great success
(e.g., Alayrac et al., 2022; Wang et al., 2022c; Hu et al., 2022; Chen et al.,
2022c, a; Wang et al., 2022b).

Among them, visual language models (LMs) build on top of pretrained
autoregressive LMs (e.g., GPT-3, Brown et al., 2020), and inherit its powerful
text generation ability. In particular, the pretrained LM parameters are
usually frozen and only some trainable layers (e.g., adaptor) are added into
the large LM during multimodal pretraining (Eichenberg et al., 2022; Mokady et
al., 2021; Tsimpoukelli et al., 2021; Alayrac et al., 2022). This frozen LM
strategy can avoid catastrophic forgetting when the visual LM is trained on
paired $\langle$_image, text_ $\rangle$ data, where the text quality is
usually lower than the text-only corpus to pretrain LM. In addition, it
enables the compelling zero-shot or few-shot capability of pretrained LM
(e.g., Flamingo, Alayrac et al., 2022).

However, such methods have two major limitations: 1) They store all acquired
knowledge within the model parameters, making them parameter inefficient in
modeling the abundant visual concepts (e.g., uncommon objects) and rich
textual descriptions (e.g., alternative descriptions for the same scene). 2)
They are inefficient in incorporating new data, typically requiring
computationally expensive fine-tuning (Chen et al., 2022a) or pertaining on
increasingly more parameters and interleaved $\langle$_image, text_ $\rangle$
data (Alayrac et al., 2022).

Figure 1: Examples of input images and output captions from 2.1B Flamingo (re-
implemented) and 2.4B Re-ViLM. Re-ViLM can utilize the retrieved captions to
generate more informative and accurate captions.

In the past few years, retrieval-augmented LMs (Guu et al., 2020; Lewis et
al., 2020; Karpukhin et al., 2020; Borgeaud et al., 2022) have shown notable
success in improving accuracy while reducing model parameters by retrieving
large-scale text corpus. Despite their success, there are several issues to be
addressed before we apply retrieval technique to visual LMs for image-to-text
generation: _i)_ The visual LM needs to seamlessly retrieve and encode
external knowledge at the beginning of multi-modal pretraining. Otherwise, the
powerful pretrained autoregressive LM tends to ignore the poorly encoded
external knowledge. _ii)_ In multi-modal datasets, there are cases where
multiple captions describing the same image (e.g., from different annotators
(Lin et al., 2014)), and multiple images having the same caption (e.g., see
images in Figure 3 in Appendix). Thus, simply performing standard nearest
neighbor retrieval tends to make the model take a shortcut and copy-paste
retrieval examples during training. _iii)_ Training the model on large-scale
interleaved $\langle$_image, text_ $\rangle$ dataset facilitates few-shot
learning capability (e.g., M3W in Alayrac et al. (2022)), but it is really
expensive to collect such dataset.

In this work, we propose a Retrieval-augmented Visual Language Model, Re-ViLM,
which enhances the state-of-the-art visual language model, Flamingo for zero-
shot and in-context few-shot image captioning,111Note that, no official
implementation of _Flamingo_ is available, which is trained on large-scale in-
house dataset (Alayrac et al., 2022). We re-implement the model on public
available dataset. by seamlessly incorporating a multimodal retriever and
retrieval-augmented LM layers that cross-attend to a text encoder (see
selected samples in Figure 1, and model framework in Figure 2). Specifically,
we make the following contributions:

  1. 1.

In contrast to previous work, we initialize Re-ViLM with RETRO, a pretrained
retrieval-augmented LM (Borgeaud et al., 2022), thus it can seamlessly
integrate the retrieval capability at the beginning of multimodal pretraining
and result in improved performance.

  2. 2.

We investigate the retrieval strategy to build the multimodal retriever. At
multimodal pretraining, we find the best performance is obtained by retrieving
_k_ -nearest neighbor captions based on _cosine_ similarity between image CLIP
(Radford et al., 2021) embeddings, while circumventing “copy-and-paste”
behavior in training by filtering out retrieved candidates with the same
caption as the training instance.

  3. 3.

We construct both pretraining and evaluation datasets consisting of
interleaved $\langle$_image, text_ $\rangle$ pairs for multimodal pretraining,
using existing public datasets. This facilitates in-context learning where
few-shot examples are given as interleaved $\langle$_image, text_ $\rangle$
pairs.

  4. 4.

We conduct extensive experiments for image-to-text generation under zero-shot,
few-shot, and fine-tuning settings on various benchmarks including MSCOCO (Lin
et al., 2014), Flickr30k (Plummer et al., 2015), and NoCaps (Agrawal et al.,
2019). Re-ViLM consistently outperforms the baseline Flamingo model across all
settings. The improvements are particularly notable in zero-shot and few-shot
settings, e.g., our Re-ViLM can outperform the Flamingo model containing even
$4\times$ more parameters in zero-shot evaluation.

We organize the rest of the paper as follows. In § 2, we discuss related work.
We introduce Re-ViLM model in § 3 and multimodal dataset for pretraining and
retrieval in § 4. We present our experimental results in § 5 and conclude the
paper in § 7.

Figure 2: The framework of Re-ViLM. The model first extracts CLIP embedding of
the input image, and use it to retrieve similar image-text pairs from the
database. Within some predetermined layers, the _retrieval-augmented LM_ will
cross-attend the visual representation from the image encoder, and the textual
representation from text encoder, which encodes the retrieved captions.

##  2 Related Work

Our work builds upon the state-of-the-art in visual language models (LMs) and
retrieval-augmented LMs.

##### Visual Language Models

Many recent work tackle the problem of generating text captions for given
images (e.g., Wang et al., 2022c; Alayrac et al., 2022; Aghajanyan et al.,
2022; Wang et al., 2022b; Hu et al., 2022; Li et al., 2022; Chen et al.,
2022c; Yu et al., 2022). Among these work, visual language models
(Tsimpoukelli et al., 2021; Alayrac et al., 2022) directly augment pretrained
LMs with visual component, achieving strong results in both zero-shot and few-
shot generation.

Retrieval-augmented Models Retrieval has been successfully applied in various
NLP tasks, including question answering (Guu et al., 2020; Karpukhin et al.,
2020), autoregressive language modeling (Borgeaud et al., 2022), and other
knowledge-intensive tasks (e.g., Lewis et al., 2020). In computer vision,
retrieval has also been applied for image recognition with long-tail
distribution of classes (Long et al., 2022). In this work, we apply retrieval
for image captioning.

There two concurrent work (Yasunaga et al., 2022; Chen et al., 2022b) that
apply retrieval for multi-modal generation. Chen et al. (2022b) propose MuRAG,
a retrieval-augmented multi-modal model for open question answering over
images and text. They evaluate their method on two QA datasets: WebQA, and
MultimodalQA.

RA-CM3 (Yasunaga et al., 2022) augments the CM3 backbone (Aghajanyan et al.,
2022) with retrieval, which can perform both image-to-text generation and
text-to-image synthesis. In contrast to our ReViLM, there are the following
differences: 1) We investigate the “copy-and-paste” behavior of retrieval-
augmented model during training, and propose a simple filtering strategy
during retrieval. In contrast, RA-CM3 proposes a query-dropout strategy that
drops some tokens of the query caption used in retrieval. In our ablation
study, we find our simple strategy works better than the dropout
regularization as shown in § 5.5.1. 2) Our Re-ViLM uses retrieval-augmented LM
decoder layer with cross-attention module to attend to the retrieved similar
captions. In contrast, RA-CM3 appends the retrieved captions as the prefix
context on the decoder side. In our ablation study, we find ours are more
effective, as shown in § 5.5.2. 3) For image-to-text generation, RA-CM3 only
provides the result of 2-shot in-context image captioning on MSCOCO (Lin et
al., 2014). In contrast, we perform extensive evaluations of our Re-ViLM on
various benchmarks under zero-shot, few-shot and fine-tuning settings.

##  3 Re-ViLM Architecture

In this section, we begin by outlining the framework of Re-ViLM. After that,
we delve into the details of each component in depth.

###  3.1 Framework

We illustrate our Re-ViLM framework for image captioning in Figure 2. It
consists of three essential components:

  * •

_Image encoder_ begins with a pretrained vision transformer from CLIP (Radford
et al., 2021) to extract visual features from the input images. These features
are subsequently fed into a trainable perceiver resampler (Jaegle et al.,
2021) to unify the image features into the textual representation used in the
retrieval-augmented LM.

  * •

_Retrieval-augmented LM_ is initialized with a pretrained RETRO model
(Borgeaud et al., 2021) to generate corresponding captions based on the image
features and retrieved evidence. Among some LM layers, the module cross-attend
the hidden representations from the image encoder or the shared text encoder.
The bidirectional text encoder encodes the retrieved captions obtained from
the multimodal retriever.

  * •

_Multimodal retriever_ consists of a retrieval database storing
$\langle$_image, text_ $\rangle$ pairs indexed by Faiss, a fast similarity
search library (Johnson et al., 2019) using their CLIP embeddings. Given a
query image, the retriever extracts its embedding using the CLIP-ViT module
within the image encoder. It then returns the top-_k_ $\langle$_image, text_
$\rangle$ pairs, measured by the cosine similarity of embeddings between the
query image and the retrieved images. After that, the retrieved captions are
encoded by the LM to generate relevant captions.

In the following subsections, we provide more details about each component of
our model.

###  3.2 Image Encoder

To fully leverage the existing pretrained model, we initialize the image
encoder with CLIP-ViT (Radford et al., 2021), a pretrained vision transformer
that processes images by dividing them into a grid of patches and then
processing each patch with a transformer encoder. In our experiments, we used
two different sizes of the CLIP-ViT model: ViT-B/32 and ViT-L/14. We freeze
the CLIP-ViT part during multimodal pretraining to avoid catastrophic
forgetting while making it trainable during fine-tuning for better results. We
then use a perceiver resampler to obtain fixed-length hidden representations
from the CLIP-ViT token embeddings. The perceiver is trainable during both the
multimodal pretraining and future fine-tuning, allowing it to adapt the visual
representations for the text decoder and connects the two modalities.

###  3.3 Multimodal Retriever

Our Multimodal Retriever consists of three steps: (1) building the database,
(2) retrieval and (3) filtering.

##### Building Database.

The retrieval database is built upon a image-text paired dataset and is
structured as a key-value map, where the keys are CLIP ViT-B/32 image
embeddings for images, and the values are the corresponding text descriptions.
The database is indexed by Faiss library (Johnson et al., 2019).

Retrieval  Given an input query image $I$, we perform $k$-nearest neighbor
retrieval with cosine similarity of embeddings between query image and
database images.222We also tried to use embeddings of query image and database
captions with the same CLIP model, and the empirical results are very similar.
The retrieval results are denoted as
$\mathcal{R}(I)=\\{(i_{1},c_{1}),\cdots,(i_{k},c_{k})\\}$, where $i_{j}$ and
$c_{j}$ with $j\in[1,k]$ represents the retrieved image and caption,
respectively.

Filtering strategy  There can be multiple captions for the same image from
different annotators in some datasets (e.g., MSCOCO (Lin et al., 2014)). In
this case, these retrieved captions, which are highly correlated (even near-
duplicate) with the ground-truth caption, could give a false sense of
retrieved evidence quality to the model, resulting in potentially degraded
test performance due to the discrepancy between the training and evaluation
setting. To avoid it, we filter out the retrieved image-text pairs if the
retrieved image is identical to the query image $I$ (e.g., $i_{1}=I$) during
both training and inference.

Furthermore, in the image-text datasets (e.g., Conceptual Captions (Sharma et
al., 2018; Changpinyo et al., 2021)), multiple images can have identical
captions. For example, one annotator may provide the same caption or alt-text
for similar images (see Figure 3 in Appendix for examples). If the retrieved
caption from database is the same as the ground-truth caption during training,
it will encourage the model to take a short path and simply copy-and-paste the
retrieved caption to the model output, hindering the training of Re-ViLM. To
address this issue, we employ a filtering strategy that filters out the
retrieved image-text pair if its text is identical to the training image’s
caption that is used as the teacher-forced input at the LM decoder layer.

Thus, if the corresponding ground-truth caption of $I$ is $C$, the filtered
retrieval results $\mathcal{R}(I)$ is,

$\\{(i_{1},c_{1}),\cdots,(i_{k},c_{k})\mid c_{j}\neq C,i_{j}\neq I,\forall
j\in[1,k]\\}$

In § 5.5.1, we show that Re-ViLM can be largely improved by using this simple
and effective filtering strategy.

###  3.4 Retrieval-augmented LM

To facilitate the model using the retrieved captions, the visual LM needs to
seamlessly retrieve and encode the external knowledge at the beginning of
multimodal pretraining. Thus, we initialize our text encoder and LM decoder
layer with pretrained RETRO (Borgeaud et al., 2021), a state-of-the-art
retrieval-augmented LM. In our experiments, we use three different sizes of
RETRO models: RETRO${}_{\textrm{\small base}}$ with 148M parameters,
RETRO${}_{\textrm{\small medium}}$ with 405M, and RETRO${}_{\textrm{\small
large}}$ with 1.5B parameters. To generate captions conditioned on visual
input, we interleave the LM decoder layers with gated cross-attention dense
layers (gated xttn-dense) as in _Flamingo_ (Alayrac et al., 2022), which take
the output of perceiver resampler as the _key_ and _value_ for cross-
attention. To incorporate the retrieved captions as evidence, we interleave
the LM decoder layers with retrieval-augmentation layers, which take text
encoder output as _key_ and _value_ for cross-attention. Note that, the text
encoder is shared across retrieval-augmentation layers. We freeze the
retrieval-augmented LM at multimodal pretraining, and make it trainable at
fine-tuning.

Text encoder  Given the retrieved $k$-nearest neighbor captions, we use a
transformer-based bidirectional encoder model, to obtain the hidden
representations. Specifically, our text encoder shares the subword embedding
table with the LM decoder. We concatenate $k$ embeddings along with the length
dimension to form the retrieval raw embedding tensor
$\mathcal{E}\in\mathbb{R}^{(k\times m)\times d}$, where $m$ is sequence
length, $d$ is hidden dimension. After applying transformer layers, the
encoder output is cross-attended by the LM decoder layers. We initialize the
text encoder with a pretrained RETRO language model (Borgeaud et al., 2021)
instead of training from scratch to improve model performance.

###  3.5 Trainable Modules at Pretraining and Finetuning

Similar to Flamingo, our training objective is to maximize the conditional
likelihood of the captions given the images. At pretraining, we follow
Flamingo’s strategy by freezing the pretrained components, including CLIP-ViT
in image encoder and retrieval-augmented LM, training only the perceiver
resampler from scratch. At finetuning, we unfreeze all the pretrained
components and increase image resolution from $224\times 224$ to $480\times
480$, as suggested in (Alayrac et al., 2022). This has been shown to improve
overall performance.

##  4 Multimodal Data for Pretraining and Retrieval

###  4.1 Image-Text Pair Data

In this work, we pretrain our models using two multi-modal datasets: 1) CC3M +
CC12M + SBU, which consists of overall 15 million high-quality image-text
pairs from the Conceptual Dataset (Sharma et al., 2018; Changpinyo et al.,
2021) and SBU Captions (Ordonez et al., 2011); 2) COYO-700M (Byeon et al.,
2022), which contains 747 million image-text pairs after filtering out low-
quality samples from a collection of 10 billion web image-text sources. In our
experiment, we find that the high-quality captions are essential for
pretraining in image captioning task. Thus, we further filter out instances
with irregular textual tokens or low CLIP similarity scores between image and
text, and obtain 104M high-quality image-text pairs. For simplicity, we refer
to the CC3M + CC12M + SBU dataset as CCS and the COYO-104M dataset as COYO.

At retrieval, we use CCS and COYO as the main sources for our retrieval
database, and utilize the Faiss library (Johnson et al., 2019) to support fast
similarity-based retrieval.

###  4.2 Interleaved Image-Text Data

In this subsection, we discuss our proposed pretraining method for enhancing
the in-context few-shot ability of our model. In this scenario, the model
needs to be highly conditioned on the previous few-shot samples (image-text
pairs) to effectively generate captions of test images. However, existing
multimodal models generally do not use multiple image-text pairs as inputs for
pretraining (Tsimpoukelli et al., 2021; Yasunaga et al., 2022). This makes the
in-context few-shot learning at inference time challenging, as there is no
such supervision during pretraining. While Alayrac et al. (2022) built an in-
house large-scale multimodal corpora with interleaved images and text,
collecting such a dataset is expensive.

We construct our image-text interleaved datasets using publicly available
image-text pair datasets. We make the image-text pairs in each interleaved
sample relevant, in order to explicitly teach the model how to condition on
previous data samples for generating the caption of the current image.

Our interleaved dataset is constructed by using CCS. For each image-text pair
(query) in CCS, we select four relevant data pairs from the same corpus to
construct each interleaved sample, which results in five data pairs for each
interleaved sample. The data selection process for each query consists of two
steps. Step-1: We use $L_{2}$ metric to measure the distances between the CLIP
embeddings of the query image and the rest of images in CCS. After that, we
choose those data pairs where the images have a normalized distance score
between $0.4$ and $0.6$ to the query image. Step-2: To ensure the captions in
an interleaved sample are similar, we further use CLIP embeddings to calculate
the distances between the query caption and the captions from the selected
data pairs in Step-1. Then, we pick the top-4 data pairs where the captions
are most similar to the query caption. Note that we don’t choose the most
similar images in Step-1, because there are lots of near-duplicate images,
e.g., they have very few differences in terms of resizing, cropping, color,
rotation, watermark. Hence, we set thresholds (in practice, [0.4, 0.6]) to
filter out these cases, increase the diversity of the data pairs in one
interleaved sample, and still make sure the selected images are relevant.

##  5 Experiments

In this section, we evaluate the performance of Re-ViLM under three different
settings: zero-shot, few-shot and fine-tuning, on various image captioning
benchmarks. We aim to demonstrate the superiority of our retrieval
augmentation technique in improving the quality and relevance of generated
captions through retrieving relevant knowledge from external databases. We
compare our results to several widely-used image captioning models such as
Enc-Dec (Changpinyo et al., 2021), SimVLM (Wang et al., 2022c), and Flamingo
(Alayrac et al., 2022). Through extensive evaluation of Re-ViLM, we conclude
that Re-ViLM is compelling under zero-shot and few-shot settings.

###  5.1 Experimental setup

Evaluation Dataset. We conduct our image captioning evaluation on three multi-
modal datasets: 1) MSCOCO (Lin et al., 2014) is a dataset for image
captioning, object detection, and segmentation. We use the Karpathy split
(Karpathy & Fei-Fei, 2015), with 82k/5k/5k images for training, validation,
and testing respectively. Each image is annotated with at most 5 human-
generated captions. 2) Flickr30k (Plummer et al., 2015) is a standard
benchmark for sentence-based image captioning, which includes 29k/1k/1k images
in its Karpathy split. 3) NoCaps contains 15k images containing nearly 400
additional novel classes to original MSCOCO, which can be used to evaluate
novel object captioning performance after finetuning on MSCOCO.

For zero-shot setting, we focus on the evaluation under MSCOCO and Flickr30k
datasets. For fine-tuning setting, we evaluate on MSCOCO, Flickr30k and NoCaps
datasets. We conduct our few-shot experiments on MSCOCO dataset only to assess
Re-ViLM’s generalization and adaptability. Throughout our experiments, we
report BLEU@4, CIDEr, and SPICE scores (Lin et al., 2014) to measure the
quality and relevance of the generated captions given input images.

Implementation. We develop Re-ViLM with three different sizes based on
different size of CLIP-ViT and RETRO (Borgeaud et al., 2021), named Re-
ViLM${}_{\textrm{\small base}}$ (ViT-B/32, RETRO-148M), Re-
ViLM${}_{\textrm{\small medium}}$(ViT-L/14, RETRO-410M) and Re-
ViLM${}_{\textrm{\small large}}$(ViT-B/32, RETRO-1.5B). Compared to Flamingo
model with the same CLIP-ViT and comparable GPT-3 configuration (Brown et al.,
2020), Re-ViLM introduces up to $16\%$ additional parameters while largely
boosting the performance. We build our model with Megatron-LM infrastructure
to support large visual language model training and evaluation. We set global
batch size as $256$ and use Adam optimizer at training. We use beam search
with beam size as $3$, maximum generation length as 10 for inference.

Table 1: Zero-shot evaluation results on MSCOCO, Flickr30k benchmarks,
compared with different image captioning baselines. We report BLEU@4, CIDer,
SPICE scores for different methods. Note that MSCOCO, Flickr30k were excluded
from pretraining set in the following MSCOCO and Flickr30k results. We
replicate Flamingo models with the same image encoder and text decoder as Re-
ViLM based on original paper.

Method | Total | Trainable | MSCOCO karpathy test | Flickr30k karpathy test
---|---|---|---|---
params. | params. | BLEU@4 | CIDer | CIDer | SPICE
VL-T5 (Cho et al., 2021) | 224M | 224M | - | 4.9 | 2.6 | 2.0
Unfied VLP (Zhou et al., 2020) | 122M | 122M | - | - | 24.9 | 7.2
SimVLM${}_{\textrm{\small base}}$ (Wang et al., 2022c) | - | - | 9.5 | 24.0 | - | -
SimVLM${}_{\textrm{\small large}}$ | - | - | 10.5 | 24.9 | - | -
SimVLM${}_{\textrm{\small huge}}$ |  $\sim$ 1.4B |  $\sim$1.4B | 11.2 | 32.2 | - | -
Flamingo${}_{\textrm{\small base}}$(re-impl) | 364M | 102M | 12.4 | 39.6 | 42.2 | 7.9
Flamingo${}_{\textrm{\small medium}}$(re-impl) | 894M | 233M | 15.6 | 44.3 | 43.2 | 8.8
Flamingo${}_{\textrm{\small large}}$(re-impl) | 2.1B | 489M | 16.5 | 49.2 | 46.4 | 9.4
Re-ViLM${}_{\textrm{\small base}}$ | 420M | 158M | 17.0 | 51.2 | 45.2 | 9.2
Re-ViLM${}_{\textrm{\small medium}}$ | 1.0B | 347M | 17.9 | 53.6 | 52.0 | 9.8
Re-ViLM${}_{\textrm{\small large}}$ | 2.4B | 806M | 18.6 | 60.8 | 52.1 | 10.0

Table 2: Few-shot evaluation results on MSCOCO benchmarks, compared with
vanilla Flamingo models as our baseline. We report BLEU@4, CIDer scores for
different methods. We pretrain our Re-ViLM on constructed CCS interleaved
dataset and evaluate on constructed COCO interleaved dataset respectively. We
adopt CCS as our retrieval set during both pretraining and evaluate stage.

Method | Total | Trainable | 2 shots | 4 shots | 8 shots
---|---|---|---|---|---
params. | params. | BLEU@4 | CIDer | BLEU@4 | CIDer | BLEU@4 | CIDer
Flamingo-3B (Alayrac et al., 2022) | 3.2B | 1.3B | - | - | - | 85.0 | - | -
Flamingo-9B | 9.3B | 1.6B | - | - | - | 93.1 | - | -
Flamingo${}_{\textrm{\small base}}$(re-impl) | 364M | 102M | 13.7 | 53.9 | 19.5 | 66.0 | 22.1 | 71.8
Re-ViLM${}_{\textrm{\small base}}$ | 420M | 158M | 14.8 | 60.1 | 20.8 | 72.2 | 21.8 | 72.6
Flamingo${}_{\textrm{\small medium}}$(re-impl) | 894M | 233M | 17.9 | 69.0 | 23.3 | 80.2 | 23.1 | 76.8
Re-ViLM${}_{\textrm{\small medium}}$ | 1.0B | 347M | 18.2 | 73.6 | 24.0 | 84.5 | 24.1 | 81.0
Flamingo${}_{\textrm{\small large}}$(re-impl) | 2.1B | 489M | 18.2 | 71.6 | 25.7 | 89.2 | 26.3 | 89.1
Re-ViLM${}_{\textrm{\small large}}$ | 2.4B | 806M | 18.4 | 77.2 | 25.5 | 90.5 | 26.2 | 90.2

Table 3: Finetuning evaluation results on MSCOCO, Flickr30k, and NoCaps
benchmarks, compared with different image captioning baselines. Note that, for
NoCaps, we finetune on MSCOCO karpathy train, following prior works (Li et
al., 2022), while some work mentioning this setting as zero-shot evaluation.
We finetune our Re-ViLM on MSCOCO/Flickr30k karpathy train split respectively
for MSCOCO and Flick30k evaluation. We report BLEU@4, CIDer, SPICE scores for
different methods.

Method | params. | MSCOCO karpathy test | Flickr30k karpathy test | NoCaps validation
---|---|---|---|---
BLEU@4 | CIDer | BLEU@4 | SPICE | CIDer | SPICE
Enc-Dec (Changpinyo et al., 2021) | - | - | 110.9 | - | - | 90.2 | 12.1
VinVL (Zhang et al., 2021) | - | 38.2 | 129.3 | - | - | 92.5 | 13.1
VL-T5 (Cho et al., 2021) | 172M | 34.6 | 116.1 | - | - | 4.4 | 5.3
MetaLM (Hao et al., 2022) | 545M | 37.6 | 126.6 | - | - | 58.7 | 8.6
Unfied VLP (Zhou et al., 2020) | 122M | 36.5 | 116.9 | 30.1 | 17.0 | - | -
BUTD (Anderson et al., 2018) | - | 36.2 | 113.5 | 27.3 | 16.0 | - | -
NBT (Lu et al., 2018) | - | 34.7 | 107.2 | 27.1 | 15.6 | - | -
SimVLM${}_{\textrm{\small huge}}$ (Wang et al., 2022c) |  $\sim$1.4B | 40.6 | 143.3 | - | - | 110.3 | 14.5
BLIP (Li et al., 2022) | 252M | 38.6 | 129.7 | - | - | 105.1 | 14.4
BLIP${}_{\textrm{\small CapFilt-L}}$ (Li et al., 2022) | 252M | 40.4 | 136.7 | - | - | 113.2 | 14.8
Flamingo${}_{\textrm{\small base}}$(re-impl) | 364M | 37.0 | 128.0 | 30.4 | 16.5 | 102.8 | 14.0
Flamingo${}_{\textrm{\small medium}}$(re-impl) | 894M | 37.4 | 129.0 | 30.7 | 17.2 | 105.6 | 14.4
Flamingo${}_{\textrm{\small large}}$(re-impl) | 2.1B | 38.2 | 129.4 | 31.2 | 17.4 | 109.2 | 14.5
Re-ViLM${}_{\textrm{\small base}}$ | 420M | 37.8 | 129.1 | 30.6 | 17.3 | 105.2 | 14.2
Re-ViLM${}_{\textrm{\small medium}}$ | 1.0B | 38.2 | 131.2 | 31.0 | 17.5 | 106.8 | 14.4
Re-ViLM${}_{\textrm{\small large}}$ | 2.4B | 39.4 | 134.2 | 31.6 | 18.0 | 109.5 | 14.7

###  5.2 Zero-shot Evaluation

We conduct zero-shot evaluation on MSCOCO and Flickr30k datasets. During
pretraining, we include both CCS and COYO as our retrieval database and report
the best number among all different settings. Results are shown in Table 1. We
find that, Re-ViLM could achieve significant boosts (around $10.0$ on CIDer
score) compared to the Flamingo model, by introducing up to $16\%$ additional
parameters. Even Re-ViLM${}_{\textrm{\small base}}$ outperforms the largest
SimVLM by a large margin. We leave the full results containing Re-ViLM’s
performance under different pretraining and retrieval database combination in
Appendix A.

###  5.3 Few-shot Evaluation

We evaluate the few-shot learning capability of Re-ViLM by pretraining it on
the constructed interleaved CCS dataset, and evaluating it under the
interleaved MSCOCO dataset, constructed by the same process described in §
4.2, with $\\{2,4,8\\}$-shots. Results are shown in Table 2.

While the significant improvements on $\\{2,4\\}$-shots setting compared with
the comparable size Flamingo model are clearly observed, we notice that the
retrieval augmentation benefits becomes less when the number of shots
increases (i.e., 8-shot). This is not surprising as the few-shot in-domain
examples from MSCOCO has more useful information to boost the model
performance on MSCOCO, than the out-of-domain samples from our retrieval
database, CCS and COYO. As the number of in-domain examples increases, the
benefit of retrieval from out-of-domain examples becomes marginal.

###  5.4 Fine-tuning Evaluation

We conduct fine-tuning evaluation of Re-ViLM on MSCOCO, Flickr30K and NoCaps
benchmarks. For evaluation on MSCOCO and Flickr30k, we fine-tune our
pretrained Re-ViLM with smaller learning rate and early-stop strategy on
MSCOCO and Flickr30k dataset respectively. For NoCaps evaluation, we fine-tune
our model on MSCOCO dataset, following prior works (Li et al., 2022). Results
are shown in Table 3. We observe that Re-ViLM still consistently outperforms
Flamingo, although the relative improvements becomes smaller compared to the
zero-shot and few-shot settings. We leave the full results containing Re-
ViLM’s performance under different pretraining and retrieval database
combination in Appendix A.

###  5.5 Ablation Study

Table 4: Comparison between ReViLM with simple filtering strategy, query-
dropout strategy, and without any filtering method during retrieval. Models
are pretrained on CCS dataset, and evaluated on MSCOCO, Flickr30k under zero-
shot setting. We report B@4: BLEU@4, C: CIDer, S: SPICE scores for different
methods.

Method | MSCOCO | Flickr30k
---|---|---
B@4 | C | C | S
Re-ViLM${}_{\textrm{\small base}}$ [No filtering] | 12.3 | 35.5 | 41.4 | 8.1
Re-ViLM${}_{\textrm{\small base}}$ [Query-dropout] | 16.5 | 48.6 | 43.4 | 9.1
Re-ViLM${}_{\textrm{\small base}}$ | 17.0 | 51.2 | 45.2 | 9.2
Re-ViLM${}_{\textrm{\small medium}}$ [No filtering] | 12.3 | 35.5 | 41.4 | 8.1
Re-ViLM${}_{\textrm{\small medium}}$ [Query-dropout] | 17.5 | 52.1 | 50.5 | 9.6
Re-ViLM${}_{\textrm{\small medium}}$ | 17.9 | 53.6 | 52.0 | 9.8

####  5.5.1 Filtering during Retrieval

There could exist two different types of duplication scenarios in multi-modal
datasets: Same image with multiple captions, which is commonly found in
MSCOCO, Flickr30k and NoCaps datasets, could lead to label leakage during
training. Multiple images with identical caption, which is common in
multimodal datasets such as Conceptual Captions, as shown in Figure 3 in
Appendix B. 333We find that the ratio of identical captions in Conceptual
Captions can be as high as $15.7\%$.  Both of these duplication can lead to a
severe issue that Re-ViLM can simply copy and paste the retrieved captions to
achieve $100\%$ match to the ground-truth captions. See Appendix B for more
in-depth discussion.

To mitigate these above issues, we develop a simple filtering strategy that
discards retrieved samples that matches the training query image $I$ or its
caption $C$ at training (i.e.
$\mathcal{R}(I)=\\{(i_{1},c_{1}),\cdots,(i_{k},c_{k})\mid c_{j}\neq
C,i_{j}\neq I,\forall j\in[1,k]\\}$). We notice that another concurrent work
RA-CM3, has also proposed the _query-dropout_ strategy which mitigates such
duplication issue by randomly dropping out retrieved caption tokens based on
their similarity to the query image and text. We conduct ablation study to
compare our simple filtering method with the query-dropout method. The
results, as shown in Table 4, indicates that our simple filtering strategy
leads to consistent improvement in the performance of Re-ViLM, while the
query-dropout strategy achieves competitive but slightly worse results than
simple filtering strategy.

####  5.5.2 Retrieval Augmentation as In-Context Prepending

Our Re-ViLM incorporate retrieved captions through the cross attention between
the bidirectional text encoder and LM decoder layers. A concurrent work, RA-
CM3, proposed an alternative retrieval augmentation method by appending the
retrieved captions as the prefix context on decoder side as a simpler way to
utilize retrieved evidence without introducing additional parameters. We
investigate this prompt-like augmentation method, and replicate it by
appending the top 2 retrieved evidence as prefix during pretraining and
inference of Flamingo model. We compare this retrieval augmentation method
with our retrieval-augmented LM layer approach. The results are shown in in
Table 5. We can observe that our retrieval-augmented design has better zero-
shot captioning performance than the retrieval augmentation method in Yasunaga
et al. (2022). It reveals the importance of our retrieval-based architecture
design.

Table 5: Comparison between different retrieval augmentation methods:
retrieval-augmented LM layers (Re-ViLM) and in-context prepending as prompt
(Flamingo + prepend), along with vanilla Flamingo model. Models are pretrained
on CCS dataset, and evaluated on MSCOCO, Flickr30k under zero-shot setting. We
report B@4: BLEU@4, C: CIDer, S: SPICE scores for different methods.

Method | MSCOCO | Flickr30k
---|---|---
B@4 | C | C | S
Flamingo${}_{\textrm{\small base}}$ | 12.4 | 39.6 | 42.2 | 7.9
Flamingo${}_{\textrm{\small base}}$ \+ prepend | 13.4 | 43.4 | 43.5 | 8.2
Re-ViLM${}_{\textrm{\small base}}$ | 17.0 | 51.2 | 45.2 | 9.2
Flamingo${}_{\textrm{\small medium}}$ | 15.6 | 44.3 | 43.2 | 8.8
Flamingo${}_{\textrm{\small medium}}$ \+ prepend | 16.4 | 45.6 | 46.6 | 9.2
Re-ViLM${}_{\textrm{\small medium}}$ | 17.9 | 53.6 | 52.0 | 9.8

##  6 Limitations

In this paper, we focus on exploring emergent zero-shot and in-context few-
shot image captioning. To achieve this, we designed our retrieval augmented
model mainly based on the Flamingo framework (Alayrac et al., 2022), and leave
the application of our retrieval design to other image-to-text frameworks (Bao
et al., 2021; Chen et al., 2022c; Wang et al., 2022a) as future work.
Furthermore, since there is no official implementation of Flamingo and its
training datasets, our framework is based on a reimplemented version of
Flamingo, trained on publicly available datasets and manually crafted
interleaved image-text datasets.

##  7 Conclusion

In this work, we propose Re-ViLM, a retrieval-augmented image-to-text model,
with strong zero-shot and few-shot image captioning results. Our Re-ViLM
architecture is built on Flamingo, providing substantial reduction in the
number of parameters while obtaining compelling results across different
settings, as it does not need to store all knowledge within the parameters. In
particular, we fist initialize Re-ViLM with RETRO, a pretrained retrieval-
augmented language model, thus it can seamlessly integrate the retrieval
capability at the beginning of multimodal pretraining. We also propose a
simple yet effective filtering strategy at retrieval to circumvent the “copy-
and-paste” behavior of retrieval-augmented model. Furthermore, we construct an
interleaved image-text dataset for pretraining, which is crucial for in-
context few-shot learning. We conduct extensive experiments on diverse image-
captioning datasets and evaluate their zero-shot and in-context few-shot
performance. Our results show that Re-ViLM consistently outperform the
baseline Flamingo model across all settings. Additionally, we conduct
experiments on fine-tuning settings and show promising results.

## References

  * Aghajanyan et al. (2022) Aghajanyan, A., Huang, B., Ross, C., Karpukhin, V., Xu, H., Goyal, N., Okhonko, D., Joshi, M., Ghosh, G., Lewis, M., et al.  Cm3: A causal masked multimodal model of the internet.  _arXiv preprint arXiv:2201.07520_ , 2022.
  * Agrawal et al. (2019) Agrawal, H., Desai, K., Wang, Y., Chen, X., Jain, R., Johnson, M., Batra, D., Parikh, D., Lee, S., and Anderson, P.  Nocaps: Novel object captioning at scale.  In _Proceedings of the IEEE/CVF International Conference on Computer Vision_ , pp. 8948–8957, 2019.
  * Alayrac et al. (2022) Alayrac, J.-B., Donahue, J., Luc, P., Miech, A., Barr, I., Hasson, Y., Lenc, K., Mensch, A., Millican, K., Reynolds, M., et al.  Flamingo: A visual language model for few-shot learning.  In _NeurIPS_ , 2022.
  * Anderson et al. (2018) Anderson, P., He, X., Buehler, C., Teney, D., Johnson, M., Gould, S., and Zhang, L.  Bottom-up and top-down attention for image captioning and visual question answering.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pp. 6077–6086, 2018.
  * Bao et al. (2021) Bao, H., Dong, L., Piao, S., and Wei, F.  Beit: Bert pre-training of image transformers.  _arXiv preprint arXiv:2106.08254_ , 2021.
  * Borgeaud et al. (2021) Borgeaud, S., Mensch, A., Hoffmann, J., Cai, T., Rutherford, E., Millican, K., Driessche, G. v. d., Lespiau, J.-B., Damoc, B., Clark, A., et al.  Improving language models by retrieving from trillions of tokens.  _arXiv preprint arXiv:2112.04426_ , 2021.
  * Borgeaud et al. (2022) Borgeaud, S., Mensch, A., Hoffmann, J., Cai, T., Rutherford, E., Millican, K., Van Den Driessche, G. B., Lespiau, J.-B., Damoc, B., Clark, A., et al.  Improving language models by retrieving from trillions of tokens.  In _ICML_ , 2022.
  * Brown et al. (2020) Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al.  Language models are few-shot learners.  _NeurIPS_ , 2020.
  * Byeon et al. (2022) Byeon, M., Park, B., Kim, H., Lee, S., Baek, W., and Kim, S.  Coyo-700m: Image-text pair dataset.  https://github.com/kakaobrain/coyo-dataset, 2022.
  * Changpinyo et al. (2021) Changpinyo, S., Sharma, P., Ding, N., and Soricut, R.  Conceptual 12M: Pushing web-scale image-text pre-training to recognize long-tail visual concepts.  In _CVPR_ , 2021.
  * Chen et al. (2022a) Chen, J., Guo, H., Yi, K., Li, B., and Elhoseiny, M.  Visualgpt: Data-efficient adaptation of pretrained language models for image captioning.  In _CVPR_ , 2022a.
  * Chen et al. (2022b) Chen, W., Hu, H., Chen, X., Verga, P., and Cohen, W. W.  MuRAG: Multimodal retrieval-augmented generator for open question answering over images and text.  In _EMNLP_ , 2022b.
  * Chen et al. (2022c) Chen, X., Wang, X., Changpinyo, S., Piergiovanni, A., Padlewski, P., Salz, D., Goodman, S., Grycner, A., Mustafa, B., Beyer, L., et al.  Pali: A jointly-scaled multilingual language-image model.  _arXiv preprint arXiv:2209.06794_ , 2022c.
  * Cho et al. (2021) Cho, J., Lei, J., Tan, H., and Bansal, M.  Unifying vision-and-language tasks via text generation.  In _International Conference on Machine Learning_ , pp. 1931–1942. PMLR, 2021.
  * Eichenberg et al. (2022) Eichenberg, C., Black, S., Weinbach, S., Parcalabescu, L., and Frank, A.  Magma–multimodal augmentation of generative models through adapter-based finetuning.  In _EMNLP_ , 2022.
  * Guu et al. (2020) Guu, K., Lee, K., Tung, Z., Pasupat, P., and Chang, M.  Realm: Retrieval augmented language model pre-training.  In _ICML_ , 2020.
  * Hao et al. (2022) Hao, Y., Song, H., Dong, L., Huang, S., Chi, Z., Wang, W., Ma, S., and Wei, F.  Language models are general-purpose interfaces.  _arXiv preprint arXiv:2206.06336_ , 2022.
  * Hu et al. (2022) Hu, X., Gan, Z., Wang, J., Yang, Z., Liu, Z., Lu, Y., and Wang, L.  Scaling up vision-language pre-training for image captioning.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pp. 17980–17989, 2022.
  * Jaegle et al. (2021) Jaegle, A., Gimeno, F., Brock, A., Vinyals, O., Zisserman, A., and Carreira, J.  Perceiver: General perception with iterative attention.  In _ICML_ , 2021.
  * Johnson et al. (2019) Johnson, J., Douze, M., and Jégou, H.  Billion-scale similarity search with GPUs.  _IEEE Transactions on Big Data_ , 7(3):535–547, 2019.
  * Karpathy & Fei-Fei (2015) Karpathy, A. and Fei-Fei, L.  Deep visual-semantic alignments for generating image descriptions.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pp. 3128–3137, 2015.
  * Karpukhin et al. (2020) Karpukhin, V., Oğuz, B., Min, S., Lewis, P., Wu, L., Edunov, S., Chen, D., and Yih, W.-t.  Dense passage retrieval for open-domain question answering.  In _EMNLP_ , 2020.
  * Lewis et al. (2020) Lewis, P., Perez, E., Piktus, A., Petroni, F., Karpukhin, V., Goyal, N., Küttler, H., Lewis, M., Yih, W.-t., Rocktäschel, T., et al.  Retrieval-augmented generation for knowledge-intensive NLP tasks.  In _NeurIPS_ , 2020.
  * Li et al. (2022) Li, J., Li, D., Xiong, C., and Hoi, S.  Blip: Bootstrapping language-image pre-training for unified vision-language understanding and generation.  In _ICML_ , 2022.
  * Lin et al. (2014) Lin, T.-Y., Maire, M., Belongie, S., Hays, J., Perona, P., Ramanan, D., Dollár, P., and Zitnick, C. L.  Microsoft coco: Common objects in context.  In _European conference on computer vision_ , pp. 740–755. Springer, 2014.
  * Liu et al. (2021) Liu, A.-A., Zhai, Y., Xu, N., Nie, W., Li, W., and Zhang, Y.  Region-aware image captioning via interaction learning.  _IEEE Transactions on Circuits and Systems for Video Technology_ , 2021.
  * Long et al. (2022) Long, A., Yin, W., Ajanthan, T., Nguyen, V., Purkait, P., Garg, R., Blair, A., Shen, C., and van den Hengel, A.  Retrieval augmented classification for long-tail visual recognition.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pp. 6959–6969, 2022.
  * Lu et al. (2018) Lu, J., Yang, J., Batra, D., and Parikh, D.  Neural baby talk.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pp. 7219–7228, 2018.
  * Luo et al. (2019) Luo, R. C., Hsu, Y.-T., Wen, Y.-C., and Ye, H.-J.  Visual image caption generation for service robotics and industrial applications.  In _2019 IEEE International Conference on Industrial Cyber Physical Systems (ICPS)_ , pp. 827–832. IEEE, 2019.
  * Mokady et al. (2021) Mokady, R., Hertz, A., and Bermano, A. H.  Clipcap: Clip prefix for image captioning.  _arXiv preprint arXiv:2111.09734_ , 2021.
  * Ordonez et al. (2011) Ordonez, V., Kulkarni, G., and Berg, T.  Im2text: Describing images using 1 million captioned photographs.  _Advances in neural information processing systems_ , 24, 2011.
  * Plummer et al. (2015) Plummer, B. A., Wang, L., Cervantes, C. M., Caicedo, J. C., Hockenmaier, J., and Lazebnik, S.  Flickr30k entities: Collecting region-to-phrase correspondences for richer image-to-sentence models.  In _Proceedings of the IEEE international conference on computer vision_ , pp. 2641–2649, 2015.
  * Radford et al. (2021) Radford, A., Kim, J. W., Hallacy, C., Ramesh, A., Goh, G., Agarwal, S., Sastry, G., Askell, A., Mishkin, P., Clark, J., et al.  Learning transferable visual models from natural language supervision.  In _ICML_ , 2021.
  * Sharma et al. (2018) Sharma, P., Ding, N., Goodman, S., and Soricut, R.  Conceptual captions: A cleaned, hypernymed, image alt-text dataset for automatic image captioning.  In _ACL_ , 2018.
  * Tsimpoukelli et al. (2021) Tsimpoukelli, M., Menick, J. L., Cabi, S., Eslami, S., Vinyals, O., and Hill, F.  Multimodal few-shot learning with frozen language models.  _Advances in Neural Information Processing Systems_ , 34:200–212, 2021.
  * Wang et al. (2022a) Wang, J., Yang, Z., Hu, X., Li, L., Lin, K., Gan, Z., Liu, Z., Liu, C., and Wang, L.  Git: A generative image-to-text transformer for vision and language.  _arXiv preprint arXiv:2205.14100_ , 2022a.
  * Wang et al. (2022b) Wang, W., Bao, H., Dong, L., Bjorck, J., Peng, Z., Liu, Q., Aggarwal, K., Mohammed, O. K., Singhal, S., Som, S., et al.  Image as a foreign language: Beit pretraining for all vision and vision-language tasks.  _arXiv preprint arXiv:2208.10442_ , 2022b.
  * Wang et al. (2022c) Wang, Z., Yu, J., Yu, A. W., Dai, Z., Tsvetkov, Y., and Cao, Y.  SimVLM: Simple visual language model pretraining with weak supervision.  In _ICLR_ , 2022c.
  * Yasunaga et al. (2022) Yasunaga, M., Aghajanyan, A., Shi, W., James, R., Leskovec, J., Liang, P., Lewis, M., Zettlemoyer, L., and Yih, W.-t.  Retrieval-augmented multimodal language modeling.  _arXiv preprint arXiv:2211.12561_ , 2022.
  * Yu et al. (2022) Yu, J., Wang, Z., Vasudevan, V., Yeung, L., Seyedhosseini, M., and Wu, Y.  Coca: Contrastive captioners are image-text foundation models.  _arXiv preprint arXiv:2205.01917_ , 2022.
  * Zhang et al. (2021) Zhang, P., Li, X., Hu, X., Yang, J., Zhang, L., Wang, L., Choi, Y., and Gao, J.  Vinvl: Revisiting visual representations in vision-language models.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pp. 5579–5588, 2021.
  * Zhao et al. (2021) Zhao, R., Shi, Z., and Zou, Z.  High-resolution remote sensing image captioning based on structured attention.  _IEEE Transactions on Geoscience and Remote Sensing_ , 60:1–14, 2021.
  * Zhou et al. (2020) Zhou, L., Palangi, H., Zhang, L., Hu, H., Corso, J., and Gao, J.  Unified vision-language pre-training for image captioning and vqa.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 34, pp. 13041–13049, 2020.

Appendix

##  Appendix A Re-ViLM with Increasing Size of Pretraining and Retrieval Data

In this section, we present detailed results of Re-ViLM in both zero-shot and
fine-tuning settings, as shown in Table 6, Table 7, and Table 8. The
performance of Re-ViLM continually improves with the increasing size of the
pretraining set and retrieval database. As the retrieval database grows, Re-
ViLM has a higher chance of obtaining more relevant evidence from the
retriever, resulting in better performance.

##  Appendix B Multiple Images with The Same Caption

In this section, we delve into the issue of data duplication, where multiple
images have the identical caption. As illustrated in Figure 3, there are three
distinct scenarios: #1) different views of the same event, #2) identical
scenes with different poses, and #3) different scenes with the same
description. In scenario 1), the provided caption describes an event with a
series of pictures taking at the same time, but it is not well related to a
specific picture, e.g., the first picture in this scenario. As a result, it is
undesired to directly copy the retrieved caption as the model output at both
training and test stage. If the retrieved similar image-text pair belongs to
scenario 2) and 3), the caption can be useful even with direct copy-and-paste
at test time. Specifically, for scenario 2), images depict the same scene
(e.g. a woman holding a cell phone in an office) and should have highly
correlated or even identical caption. However, if the large part of training
set consists of such examples 444For example, we find the ratio of training
examples that have the same captions as others can be as high as 15.1% in
Conceptual Captions (Sharma et al., 2018)., the retrieval-augmented models are
overly encouraged to copy-and-paste the retrieved caption as the output, which
can undermine their generalization ability at test stage, especially for out-
of-domain settings.

Figure 3: Set of images from conceptual captions dataset with the same
captions under (Top) Scenario #1): Different views from the same event.
(Middle) #2): Identical scenes with different poses. (Bottom) Scenario #3):
Different scenes with the same description. The example of images are from
Conceptual Captions (Sharma et al., 2018) Table 6: Full zero-shot evaluation
results on MSCOCO dataset. CCS refers to CC3M+CC12M+SBU dataset and COYO the
COYO104M.

Method | Pretraining | Evaluation
---|---|---
Training Set | Retrieval Database | Retrieval Database | BLEU@4 | CIDer
Re-ViLM${}_{\textrm{\small base}}$ | CCS | CCS | CCS | 17.6 | 49.4
Re-ViLM${}_{\textrm{\small base}}$ |  CCS+COYO |  CCS+COYO |  CCS+COYO | 17.0 | 51.2
Re-ViLM${}_{\textrm{\small medium}}$ | CCS | CCS | CCS | 17.3 | 52.8
Re-ViLM${}_{\textrm{\small medium}}$ |  CCS+COYO |  CCS+COYO |  CCS+COYO | 17.9 | 53.6
Re-ViLM${}_{\textrm{\small large}}$ | CCS | CCS | CCS | 19.2 | 59.6
Re-ViLM${}_{\textrm{\small large}}$ |  CCS+COYO |  CCS+COYO |  CCS+COYO | 18.6 | 60.8

Table 7: Full zero-shot evaluation results on Flickr30k dataset. CCS refers to
CC3M+CC12M+SBU dataset and COYO the COYO104M.

Method | Pretraining | Evaluation
---|---|---
Training Set | Retrieval Database | Retrieval Database | CIDer | SPICE
Re-ViLM${}_{\textrm{\small base}}$ | CCS | CCS | CCS | 45.0 | 9.2
Re-ViLM${}_{\textrm{\small base}}$ |  CCS+COYO |  CCS+COYO |  CCS+COYO | 45.2 | 9.2
Re-ViLM${}_{\textrm{\small medium}}$ | CCS | CCS | CCS | 50.8 | 9.5
Re-ViLM${}_{\textrm{\small medium}}$ |  CCS+COYO |  CCS+COYO |  CCS+COYO | 52.0 | 9.8
Re-ViLM${}_{\textrm{\small large}}$ | CCS | CCS | CCS | 51.5 | 9.7
Re-ViLM${}_{\textrm{\small large}}$ |  CCS+COYO |  CCS+COYO |  CCS+COYO | 52.1 | 10.0

Table 8: Full fine-tuning evaluation results on MSCOCO dataset. CCS refers to
CC3M+CC12M+SBU dataset and COYO the COYO104M.

Method | Pretraining | Evaluation
---|---|---
Training Set | Retrieval Database | Retrieval Database | BLEU@4 | CIDer
Re-ViLM${}_{\textrm{\small base}}$ | CCS | CCS | CCS | 37.0 | 127.5
Re-ViLM${}_{\textrm{\small base}}$ | CCS | CCS |  CCS+COYO | 37.3 | 128.0
Re-ViLM${}_{\textrm{\small base}}$ |  CCS+COYO | CCS | CCS | 37.2 | 128.0
Re-ViLM${}_{\textrm{\small base}}$ |  CCS+COYO | CCS |  CCS+COYO | 37.4 | 128.6
Re-ViLM${}_{\textrm{\small base}}$ |  CCS+COYO |  CCS+COYO |  CCS+COYO | 37.5 | 128.2
Re-ViLM${}_{\textrm{\small base}}$ |  CCS+COYO |  CCS+COYO |  CCS+COYO+COCO | 37.8 | 129.1
Re-ViLM${}_{\textrm{\small medium}}$ | CCS | CCS |  CCS+COYO | 37.5 | 129.0
Re-ViLM${}_{\textrm{\small medium}}$ |  CCS+COYO | CCS |  CCS+COYO | 37.7 | 129.1
Re-ViLM${}_{\textrm{\small medium}}$ |  CCS+COYO |  CCS+COYO |  CCS+COYO | 38.0 | 129.9
Re-ViLM${}_{\textrm{\small medium}}$ |  CCS+COYO |  CCS+COYO |  CCS+COYO+COCO | 38.2 | 131.2
Re-ViLM${}_{\textrm{\small large}}$ | CCS | CCS | CCS | 38.4 | 129.8
Re-ViLM${}_{\textrm{\small large}}$ | CCS | CCS |  CCS+COYO | 38.1 | 128.4
Re-ViLM${}_{\textrm{\small large}}$ |  CCS+COYO |  CCS+COYO |  CCS+COYO+COCO | 39.4 | 134.2