SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Weizhe Lin, Jinghong Chen,\* Jingbiao Mei, Alexandru Coca, Bill Byrne
Cambridge, United Kingdom CB2 1PZ
{wl356，jc2124，jm2245，ac2123，wjb31}@cam.ac.uk
Knowledge-based Visual Question Answering (KB-VQA) requires VQA systems to utilize knowledge from external knowledge bases to answer visually-grounded questions. Retrieval-Augmented Visual Question Answering (RA-VQA), a strong framework to tackle KB-VQA, first retrieves related documents with Dense Passage Retrieval (DPR) and then uses them to answer questions. This paper proposes Fine-grained Late-interaction Multi-modal Retrieval (FLMR) which significantly improves knowledge retrieval in RA-VQA. FLMR addresses two major limitations in RA-VQA's retriever: (1) the image representations obtained via image-to-text transforms can be incomplete and inaccurate and (2) relevance scores between queries and documents are computed with one-dimensional embeddings, which can be insensitive to finer-grained relevance. FLMR overcomes these limitations by obtaining image representations that complement those from the image-to-text transforms using a vision model aligned with an existing text-based retriever through a simple alignment network. FLMR also encodes images and questions using multi-dimensional embeddings to capture finer-grained relevance between queries and documents. FLMR significantly improves the original RA-VQA retriever's PRRecall@5 by approximately  $8\%$ . Finally, we equipped RA-VQA with two state-of-the-art large multi-modal/language models to achieve  $\sim 61\%$  VQA score in the OK-VQA dataset.
Knowledge-based Visual Question Answering (KB-VQA) is a challenging problem that lies at the intersection of Computer Vision, Natural Language Processing, and Information Retrieval. The objective of VQA is to read an image and answer a question related to the image content. KB-VQA poses an additional challenge: in order to answer the question correctly, the system needs to draw on relevant information from an external knowledge source, such as a knowledge graph or a database. Therefore, tackling KB-VQA tasks crucially depends on the ability to retrieve relevant information and to ground the answer generation process in the retrieved knowledge.
Retrieval Augmented Visual Question Answering (RA-VQA) is a framework designed to answer difficult KB-VQA questions [Luo et al., 2021, Gao et al., 2022, Lin and Byrne, 2022], with the most recent version from Lin and Byrne [2022] achieving performance close to large models (such as GPT-3 [Brown et al., 2020]) while using much simpler models. RA-VQA first retrieves  $K$  documents relevant to the image and the question from an external knowledge base, and then generates the answer using a Large Language Model (LLM) grounded in the retrieved passages.
*Equally contributed as the first author
We focus on two major limitations in RA-VQA's retriever. (1) Incomplete image understanding: image representations are obtained via image-to-text transforms such as captioning and object detection. While effective, this approach can result in incomplete image understanding, which hinders the retrieval of relevant knowledge. This is a common issue for retrieval-based KB-VQA systems in the literature. (2) Lossy compression of visual scenes and questions to a single embedding: the Dense Passage Retrieval (DPR) [Karpukhin et al., 2020] retriever, widely used in current retrieval-based QA systems, computes relevance scores between queries and documents with their respective, one-dimensional embeddings. However, compressing complex visual scenes and questions into a single embedding can be lossy. This is especially problematic in KB-VQA, where queries and visual elements are more diverse than in other Open-domain QA tasks. DPR could overlook finer-grained relevance, resulting in degraded retrieval performance.
To address these two limitations we propose an enhanced knowledge retrieval approach called Fine-grained Late-interaction Multi-modal Retrieval (FLMR). FLMR incorporates finer-grained, token-level visual and textual features into multi-dimensional embeddings. When computing relevance scores, FLMR considers the interaction between every pair of token embeddings, including cross-modality interaction between texts and images, enabling a finer-grained assessment of relevance. We also introduce large vision models such as ViT [Dosovitskiy et al., 2021] to produce visual tokens that complement text-based image representations for more complete image understanding. To ensure that the interactions between visual and text tokens are well-defined, we align the vision model with the text-based retriever with a simple yet effective alignment training procedure. We also find that FLMR is able to make use of finer-grained regions of interest, leading to better recall rate, whereas DPR's recall rate degrades when these finer-grained features are incorporated. Our FLMR retriever achieves a significant increase of approximately  $8\%$  in PRRecall@5 for knowledge retrieval, and a competitive VQA score of  $61\%$ , surpassing the state-of-the-art models with the same scale of parameters.
- We introduce FLMR, the first-of-its-kind to leverage Late Interaction $^2$  and multi-dimensional representations to capture fine-grained, cross-modality relevance that significantly improve retrieval performance over existing state-of-the-art KB-VQA retrievers;
- We show that introducing image representations from large vision model after a simple yet effective alignment procedure can complement image representations obtained via image-to-text transforms, leading to more complete image understanding, better knowledge retrieval and higher VQA accuracy. This offers improvements to current VQA systems as many systems have only a single mode of image understanding that relies on either image-to-text transforms or vision models;
- We achieve a substantial improvement of approximately  $8\%$  in knowledge PRRecall@5 over other state-of-the-art retrievers in the OK-VQA dataset, with an accuracy of  $61\%$  that surpasses other systems with similar parameter sizes.
Visual Question Answering Systems. Recent work in VQA can be roughly divided into four categories with respect to multi-modal modeling: (1) Visual and textual features can be fused via cross-modality fusion [Yu et al., 2018, Singh et al., 2019, Yu et al., 2019, Jiang et al., 2020, Guo et al., 2021]; (2) Multi-modal models can be trained from scratch to jointly understand vision and language before they are fine-tuned to perform VQA tasks [Tan and Bansal, 2019, Chen et al., 2020, Gan et al., 2020, Li et al., 2020a, Wang et al., 2022, Zhang et al., 2021, Li et al., 2021]; (3) Vision model and language model that have been pre-trained on uni-modal corpora can be aligned to avoid expensive multi-modal pre-training [Guo et al., 2023, Dai et al., 2022, Singh et al., 2022]. (4) Image-to-text transforms such as captioning can be used to transform images into texts to enable the use of text-only reasoning pipelines [Lin and Byrne, 2022, Gui et al., 2021, Lin et al., 2022, Luo et al., 2021, Yang et al., 2022, Gao et al., 2022, Hu et al., 2022a]. Building on these Vision-and-Language modeling techniques, our work shows that image-to-text transforms and aligned vision models can complement each other to provide more complete visual information, leading to improved performance in both knowledge retrieval and VQA.
$^{2}$ Dual encoder architecture where the queries and documents are first encoded into token-level embeddings and these embeddings are then aggregated to compute final relevance scores
Figure 1: Overview of RA-VQA-v2. The system consists of two steps: (A) Knowledge Retrieval and (B) Answer Generation. (A.1) A text retriever is used to obtain token-level embeddings of text-based vision (obtained by captioning and object detection) and text documents in the database. (A.2) Visual tokens are obtained from the image and the region-of-interest patches using a vision model and a mapping network. (A.3) Relevance score between the query and the document is computed by aggregating the fine-grained relevance at token level with late interaction mechanism (Eq. 12). (B.1) The answer generator takes the text query, the image, and the retrieved documents as input, generating one candidate answer per retrieved document. (B.2) The answer with the highest joint probability is selected.
Knowledge-based VQA Systems. Recent KB-VQA systems can access both structured data, such as ConceptNet and other KGs [Narasimhan et al., 2018, Garderes et al., 2020, Li et al., 2020b, Wu et al., 2022, Marino et al., 2021, Chen et al., 2023a], as well as unstructured data such as Wikipedia passages [Wu et al., 2022, Gao et al., 2022, Gui et al., 2021] for knowledge retrieval. LLMs can also be a source of "implicit world knowledge": KAT [Gui et al., 2021] and REVIVE [Lin et al., 2022] prompt GPT-3 to generate potential answer candidates. RA-VQA [Lin and Byrne, 2022] and its prior works [Luo et al., 2021, Qu et al., 2021, Gao et al., 2022] ground answer generation in the retrieved knowledge from external KBs to achieve excellent VQA performance. Our work improves this retriever-reader pipeline with a novel knowledge retriever which significantly improves the recall rate of knowledge retrieval as well as the final VQA performance.
Knowledge Retrieval. Most knowledge retrievers in QA systems are based on DPR and its variants [Karpukhin et al., 2020, Gui et al., 2021, Luo et al., 2021, Gui et al., 2021, Lin and Byrne, 2022, Wu and Mooney, 2022]. These mainly use one-dimensional embeddings and contrastive learning for training. Late Interaction models [Khattab and Zaharia, 2020, Santhanam et al., 2022a] have recently achieved state-of-the-art performance on QA knowledge retrieval. Our FLMR extends this paradigm to work with multi-modal features and shows that incorporating finer-grained visual features, such as regions-of-interest, leads to superior retrieval performance. EnFoRe [Wu and Mooney, 2022] retrieves a list of entities from the image, the query, and the answer candidates, and then explicitly learns scores to indicate the importance of each fine-grained entity. FILIP [Yao et al., 2022] has a similar late-interaction setting but it focuses on single modal query (image-text retrieval). To the best of our knowledge, FLMR is also the first to introduce cross-modality, token-level late interactions to compute relevance scores for KB-VQA knowledge retrieval. We also propose a light-weight method that aligns a vision model with a text-based retriever to incorporate more complete multi-modal information in retrieval queries. Compared to previous approaches that rely on expensive pre-training on multi-modal datasets [Chen et al., 2022a, Yao et al., 2022], FLMR's vision-language alignment process is efficient and can be done in 4 hours with one A-100 GPU.
In this section, we introduce RA-VQA-v2, which builds upon the original RA-VQA framework [Lin and Byrne, 2022] but is equipped with Fine-grained Late-interaction Multi-modal Retriever (FLMR) to enhance knowledge retrieval. As illustrated in Fig. 1, the framework consists of two stages: Knowledge Retrieval (Sec. 3.1) and Answer Generation (Sec. 3.2).
3.1 Knowledge Retrieval
The FLMR system consists of two encoders: a vision model  $F_V$  and a language model  $F_L$  that encode image and textual features, respectively.
Visual Features. We utilize two types of visual representations: (1) text-based vision representations (textual description of visual elements) obtained by image-to-text transforms and (2) feature-based vision representations obtained by large vision models.
For text-based vision representations, to allow a direct comparison, we follow Lin and Byrne [2022] to extract objects and their attributes using VinVL [Zhang et al., 2021] and generate image captions using Oscar [Li et al., 2020a]. For each image  $I$ , we obtain a textual description that contains serialized object names, attributes, and descriptive captions [Lin and Byrne, 2022]. The sequence is appended to the question  $q$  to form the query. For simplicity of notation, the question  $q$  always includes text-based vision unless otherwise specified.
For feature-based vision representations, we use the vision model  $F_V$  to extract both global and regional image feature representations. For regional image feature representations, we further use the object detection results of VinVL to locate  $N_{ROI}$  (Region-of-Interest) bounding boxes. To filter bounding box proposals from VinVL, we use the predicted class name associated with each box to select objects explicitly mentioned in the question  $q$ , and then prioritize bounding boxes with larger areas. Using the vision model  $F_V$ , we then obtain one global image representation  $g = F_V(I) \in \mathcal{R}^{d_V}$  from the image  $I$  and ROI-based regional representations  $\{r_i = F_V(I_i^p) \in \mathcal{R}^{d_V}\}_{i=1,\dots,N_{ROI}}$  from the image ROI patches  $\{I_i^p : i = 1,\dots,N_{ROI}\}$  which contain finer-grained details.
Token-Level Embeddings. Compared with DPR's compressed, one-dimensional representation of queries and documents, FLMR preserves richer information by employing token-level, multidimensional embeddings to improve retrieval. We obtain token-level embeddings for both textual input and visual input. These are concatenated to form the final embeddings of queries and documents.
To align the vision and text modalities, we train a mapping network  $\mathcal{F}_M$  that learns to project visual features from vision model  $\mathcal{F}_V$  with hidden size  $d_V$  into the latent space of the language model  $\mathcal{F}_L$  with hidden size  $d_L$ . The mapping network, a 2-layer multi-layer perceptron, projects each visual representation into  $N_{vt}$  visual tokens, i.e.  $\mathcal{R}^{d_V} \to \mathcal{R}^{N_{vt}d_L/2} \to \mathcal{R}^{N_{vt}d_L}$  and finally reshaped into  $\mathcal{R}^{N_{vt} \times d_L}$ .
Formally, the final query embeddings  $\mathbf{Q}$  are:
$$
\mathbf {Q} = \left[ \mathcal {F} _ {L} (q), \mathcal {F} _ {M} \left(\left[ g, r _ {1}, r _ {2}, \dots , r _ {N _ {R O I}} \right]\right) \right] \in \mathcal {R} ^ {l _ {Q} \times d _ {L}}, \tag {1}
$$
where  $l_{Q} = l_{q} + (N_{ROI} + 1) \times N_{vt}$ .  $l_{q}$  is the length of the question  $q$ .  $[v_{1},\dots,v_{N}]$  denotes the concatenation of  $N$  embeddings  $v_{1}$  to  $v_{N}$ .
The documents in the knowledge base are represented by embeddings  $\mathbf{D}$  obtained from the document content  $d$  of length  $l_{D}$ :
$$
\mathbf {D} = \mathcal {F} _ {L} (d) \in \mathcal {R} ^ {l _ {D} \times d _ {L}}, \tag {2}
$$
Multi-Modal Late Interaction. We compute the relevance score between a question-image pair  $\bar{\mathbf{q}} = (q, I)$  and a document  $d$  by a late interaction formula similar to that in ColBERT but under the multi-modal context:
$$
r (\bar {\mathbf {q}}, d) = r ((q, I), d) = \sum_ {i = 1} ^ {l _ {Q}} \max  _ {j = 1} ^ {l _ {D}} \mathbf {Q} _ {i} \mathbf {D} _ {j} ^ {\top} \tag {3}
$$
For each query token, the MAX operation selects the highest relevance score over all document tokens. In preliminary experiments, other operations (e.g. MEAN or SUM) were found to be overly sensitive to the length of documents, which can be as short as a single sentence. We note that [PAD]
may dominate in the final score for short documents, whereas longer documents have an inherent advantage due to their greater number of meaningful tokens.
In contrast to DPR, FLMR allows full interactions between every query embedding vector  $\mathbf{Q}_i$  and every document embedding vector  $\mathbf{D}_j$ . FLMR retriever also supports retrieving multi-modal documents. We leave the formulation and results to Appendix H.
Training and Inference. To train the model, we treat documents  $d^{*}$  that contain the ground-truth answer to question  $q$  as gold (positive) documents. We use in-batch negative sampling for training following Karpukhin et al. [2020]. All documents in a training batch other than  $d^{*}$  are considered negative for  $q$ , denoted as  $\mathcal{N}(q)$ . We train with the contrastive loss  $\mathcal{L}_{CL}$  over the dataset  $\mathcal{D}$ :
$$
\mathcal {L} _ {C L} = - \sum_ {(q, d ^ {*}) \in \mathcal {D}} \log \frac {\exp (r (q , d ^ {*}))}{\exp (r (q , d ^ {*})) + \sum_ {z \in \mathcal {N} (q)} \exp (r (q , z))} \tag {4}
$$
After training, all documents are indexed using PLAID [Santhanam et al., 2022b], which enables fast late-interaction retrieval with a time cost similar to that of DPR.
Training the Mapping Network for Vision-Language Alignment. Directly fine-tuning the two models  $\mathcal{F}_V$  and  $\mathcal{F}_L$  on the retrieval task leads to performance degradation at the start of training, as the models are not yet aligned. Inspired by CLIP [Radford et al., 2021], where a language model is trained to align with a vision model, we align  $\mathcal{F}_V$  and  $\mathcal{F}_L$  in the context of knowledge retrieval by pre-training the parameters of the mapping network  $\mathcal{F}_M$  with a retrieval task.
Given ground-truth image-document pairs  $\{(I_p, d_p)\}$ , which can be Wikipedia images and their accompanying texts, the system is trained to retrieve the document  $d_p$  associated with the input image  $I_p$ . The relevance between the input image  $I$  and a document  $d$  is formulated as
$$
\mathbf {Q} = \mathcal {F} _ {M} \left(F _ {V} (I)\right) \in \mathcal {R} ^ {N _ {v t} \times d _ {L}}; \quad \mathbf {D} = \mathcal {F} _ {L} (d) \in \mathcal {R} ^ {l _ {D} \times d _ {L}}; \quad r (I, d) = \sum_ {i = 1} ^ {N _ {v t}} \max  _ {j = 1} ^ {l _ {D}} \mathbf {Q} _ {i} \mathbf {D} _ {j} ^ {\top} \tag {5}
$$
where only the parameters of the mapping network  $\mathcal{F}_M$  are trained with the contrastive loss in Eq. 4. We provide details of pre-training in Appendix E and discuss its effectiveness in Sec. 5.2.
Knowledge Retrieval. We extract top- $K$  documents from the knowledge base as relevant knowledge. The retrieval probability is defined below following the notation of Lin and Byrne [2022] and Lewis et al. [2020]:
$$
p _ {\theta} \left(d _ {k} \mid \bar {\mathbf {q}}\right) = \frac {\exp \left(r \left(\bar {\mathbf {q}} , d _ {k}\right)\right)}{\sum_ {j = 1} ^ {K} \exp \left(r \left(\bar {\mathbf {q}} , d _ {j}\right)\right)} \tag {6}
$$
where  $\theta$  denotes the model parameters of  $\mathcal{F}_V, \mathcal{F}_L$ , and  $\mathcal{F}_M$ .
3.2 Answer Generation
In principle, the knowledge retrieved by FLMR can be used by any VQA answer generator. We denote the answer generator as  $\mathcal{F}_A$  with parameters  $\phi$ . Following Lin and Byrne [2022], RA-VQA-v2 generates an answer for each retrieved document and selects the best candidate by the joint probability of retrieval and answer generation:
$$
\left\{d _ {k} \right\} _ {k = 1} ^ {K} = \operatorname {t o p} \mathrm {K} _ {d} \left(p _ {\theta} (d | \bar {\mathbf {q}})\right); \quad \widehat {y}, \widehat {d} = \underset {y, d _ {k}} {\arg \max } p (y, d _ {k} | \bar {\mathbf {q}}) = \underset {y, d _ {k}} {\arg \max } p _ {\phi} (y | \bar {\mathbf {q}}, d _ {k}) p _ {\theta} (d _ {k} | \bar {\mathbf {q}}) \tag {7}
$$
The training loss of the answer generator follows that of the underlying model. For example, when using BLIP 2 [Li et al., 2023], we use the cross-entropy loss of the generated sequences:
$$
\mathcal {L} = \sum_ {(\bar {\mathbf {q}}, S) \in \mathcal {T}} \sum_ {k = 1} ^ {K} \log p _ {\phi} \left(s _ {k} ^ {*} \mid \bar {\mathbf {q}}, d _ {k}\right) \tag {8}
$$
where  $\mathcal{T}$  is the whole dataset.  $S$  is the set of human responses.  $s_k^* \in S$  is the answer string that appears in the retrieved document  $d_k$ , or the most popular answer string if an exact match cannot be found in the document.
3The most popular answer is the one chosen by most annotators.
Datasets. We focus on the OK-VQA dataset where a large portion of questions requires external knowledge (either commonsense or domain-specific) to answer. There are no annotations of 'ground-truth' documents for OK-VQA questions. We follow the literature to use pseudo-relevance labels (a binary indicator of whether a document contains the answer string) as document annotations. We do not evaluate A-OKVQA [Schwenk et al., 2022], a successor of OK-VQA, as it emphasizes visually-grounded reasoning rather than knowledge retrieval. To validate the effectiveness of our proposed approach, we test the retrieval abilities using 2 different corpora, whose statistics can be found in Appendix C:
(1) Google Search Corpus for OK-VQA [Luo et al., 2021]: a passage corpus collected for answering OK-VQA questions. Previous work has shown that the corpus is effective for OK-VQA [Luo et al., 2021, Lin and Byrne, 2022]. We use this corpus in evaluating VQA performance since it covers more knowledge for answering the OK-VQA questions.
(2) Wikipedia Corpus for OK-VQA: we collect this corpus by gathering all Wikipedia passages on common objects and concepts (e.g. umbrella, dog, hat) and those containing any of the potential answers in OK-VQA training set. This ensures the corpus covers useful knowledge for answering OK-VQA questions. We note that the collected corpus encompasses a multitude of semantically-diverse documents (>100,000) that challenge the retrieval system to identify actually useful documents. For example, all Wikipedia documents with the word 'party' are included in the corpus, ranging from descriptions of fairy tales to political parties. We also use  $10\%$  of the WIT dataset [Srinivasan et al., 2021], a corpus based on Wikipedia with image-text pairs, to train the mapping network for multi-modal alignment.
We evaluate on two additional KB-VQA datasets to demonstrate FLMR's generalizability.
(1) FVQA [Wang et al., 2017]: We follow RAVQA [Lin and Byrne, 2022] to preprocess the data. All knowledge triplets are serialized into text sequences to form the knowledge base for retrieval. The average of 5 cross-validation splits is reported.
(2) Infoseek [Chen et al., 2023b]: Infoseek is a newly proposed KB-VQA dataset that provides Wikipedia documents that can be used in answering its questions. We follow Chen et al. [2023b] in preprocessing. First, we remove questions whose answers cannot be found in the provided Wikipedia passages. Second, in additional to the documents covered in the dataset ( $\sim$ 60,000), we include less relevant passages to form a knowledge base for retrieval ( $\sim$ 100,000 documents). The test set annotation has not been released, and so we split the official validation set again into validation and test sets ( $\sim$ 5200 questions).
Training Setup. We use Co1BERTv2 [Santhanam et al., 2022a] and CLIP ViT-base [Radford et al., 2021] to initialize the text-based retriever and vision encoder. For the DPR baseline, we use the official DPR checkpoints to initialize the retriever. In answer generation, we use T5-large [Raffel et al., 2020] and BLIP2-Flan-T5-XL. We use 1 Nvidia A100 (80G) for all experiments. We give detailed training hyperparameters in Appendix E. We use LoRA [Hu et al., 2022b] to fine-tune RA-VQA-v2 (BLIP 2) on 1 single GPU. The vision model is frozen throughout all experiments. During vision-language alignment training, only the mapping network is trainable. In training the answer generator, the retriever is frozen. Our implementations are released at https://github.com/LinWeizheDragon/Retrieval-Augmented-Visual-Question-Answering.
Evaluation. We present the metrics used to assess the generated answer and the performance of our knowledge retriever below. All reported numbers are averaged from 3 runs with different seeds. We verified the significance of all mentioned improvements with scipy.stats.ttest_ind  $(p < 0.05)$ .
(1) VQA Score: To evaluate VQA performance, we use the official VQA Score [Marino et al., 2019] which assigns score to the generated answer based on its exact occurrence count in the set of human responses  $S$ :
$$
\operatorname {V Q A S c o r e} (y, \mathcal {S}) = \min  \left(\frac {\# \mathcal {S} (y)}{3}, 1\right), \tag {9}
$$
where  $\#_{S}(y)$  is the occurrence of  $y$  in human responses  $S$ . This score ensures that a model is partially rewarded even if it generates a less popular answer among the human responses [Luo et al., 2021].