SlowGuess/NAACL_2021_OCR huggingface dataset:
Siqi Sun; Yen-Chun Chen
Linjie Li, Shuohang Wang, Yuwei Fang, Jingjing Liu
{siqi.sun, yen-chun.chen, lindsey.li, shuohang.wang, yuwfan, jingjl} @microsoft.com
Multimodal pre-training has propelled great advancement in vision-and-language research. These large-scale pre-trained models, although successful, fatefully suffer from slow inference speed due to enormous computation cost mainly from cross-modal attention in Transformer architecture. When applied to real-life applications, such latency and computation demand severely deter the practical use of pre-trained models. In this paper, we study Image-text retrieval (ITR), the most mature scenario of  $\mathrm{V + L}$  application, which has been widely studied even prior to the emergence of recent pre-trained models. We propose a simple yet highly effective approach, LightningDOT that accelerates the inference time of ITR by thousands of times, without sacrificing accuracy. LightningDOT removes the time-consuming cross-modal attention by pre-training on three novel learning objectives, extracting feature indexes offline, and employing instant dot-product matching with further re-ranking, which significantly speeds up retrieval process. In fact, LightningDOT achieves new state of the art across multiple ITR benchmarks such as Flickr30k, COCO and Multi30K, outperforming existing pre-trained models that consume  $1000\times$  magnitude of computational hours.
Image-text retrieval (ITR) has been widely studied as a staple benchmark task in both NLP and computer vision communities. Traditional ITR search engines typically deploy ranking-based models built upon visual-semantic embedding matching (Faghri et al., 2017; Huang et al., 2018) or deep cross-modal fusion with attention mechanism (Lee et al., 2018; Li et al., 2020a,b). Earliest works (Kiros et al., 2014; Faghri et al., 2017;
Figure 1: Evolution of Image-Text Retrieval (ITR) paradigm. (a) Early work (Faghri et al., 2017) using dot product to learn the similarity between global image features and global text features. (b) Later study (Lee et al., 2018) applying cross-attention between the features of each region and each word. (c) Pre-trained  $\mathrm{V + L}$  models (Chen et al., 2020) with deep Transformer. (d) LightningDOT without cross-attention. CMR, SMRM and VMLM refer to different pre-training tasks, which will be introduced later in method section.
Wang et al., 2018) employ separate image encoder (e.g., CNN) and text encoder (e.g., RNN), the embeddings from which are then measured by doc product for similarity matching (Figure 1(a)). Later studies (Lee et al., 2018, 2019; Wang et al., 2019; Zhang et al., 2020) improve this paradigm by employing advanced region-level visual encoder (e.g., Faster-RCNN) and applying cross-attention between word features and region features for multimodal fusion (Figure 1(b)).
With the advent of Transformer (Vaswani et al., 2017) and BERT (Devlin et al., 2019), cross-modal retrieval tasks are more recently dominated by vision-and-language  $(\mathrm{V} + \mathrm{L})$  pre-trained models, such as ViLBERT (Lu et al., 2019), UNITER (Chen et al., 2020), OSCAR (Li et al., 2020b), and VILLA (Gan et al., 2020). Large-scale pre-trained models learned from massive corpus of image-text pairs can power heterogeneous downstream tasks that take diverse modalities as inputs (e.g., text, image, video, audio). These models benefit from
1Code and pre-training checkpoints are available at https://github.com/intersun/LightningDOT.
Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 982-997 June 6-11, 2021. ©2021 Association for Computational Linguistics
the self-attention mechanism in Transformer architecture, learning joint image+text embeddings through pre-training objectives such as masked language modeling (MLM) and masked region modeling (MRM) (Figure 1(c)).
However, the very ingredient that engenders the success of these pre-trained models, cross-modal attention between two modalities (through self-attention), also destines the inevitable latency and huge computation cost in training and deploying such massive-scale models. For example, UNITER (Chen et al., 2020) builds upon 12/24 Transformer layers, and trains over 10 million image+text pairs. The inference time of such large models with 110 million parameters is 48 seconds on average for text query from COCO dataset (Chen et al., 2015), not scalable in real-life applications serving millions of queries per second.
To make real-time ITR possible with low latency, we ask a bold question: can we go back to the beginning, reverting to simple dot product for efficient cross-modal retrieval? To make this retro experiment feasible, we rely on Transformer to pre-train high-quality image and text encoders, but use efficient dot product for multimodal fusion instead of computationally heavy self-attention. To still facilitate effective cross-modal embedding learning, we use a special [CLS] token on both encoders, which transfers the learned embedding from the other modality (Figure 1(d)). We name this new paradigm LightningDOT, for its lightening speed benefiting from dot product computation.
By removing the time-consuming cross-attention between modalities, the model can learn visualsemantic embeddings without extensive matching between each image-text pair during inference, as used in existing pre-trained models (Chen et al., 2020; Li et al., 2020b; Lu et al., 2019). Further, by eliminating the dependency on real-time computation over image-text pairs, we can compute all image and text embeddings independently offline just for once, and reuse these embeddings as cached indexes for new queries on the fly (Figure 2).
For model training, we propose three learning objectives to jointly train two Transformer blocks: Image Encoder and Language Encoder. Specifically, Visual-embedding fused MLM (namely VMLM) and Semantic-embedding fused MRM (namely SMRM) ensure cross-modal information is harnessed even without cross-modality self-attention. A cross-modal retrieval objective (namely CMR)
encourages the model to learn multimodal fusion through pre-training. To maintain competitive model performance, we further introduce a re-ranking mechanism to bring back the benefit of cross-attention methods.
In summary, LightningDOT is designed with late fusion to learn visual-semantic embeddings. Experiments on popular ITR benchmarks show that LightningDOT is 600/1900 times faster than existing pre-trained models on Flickr30k/COCO, while achieving new state-of-the-art results. When retrieving from larger candidate pool ( $>120\mathrm{K}$  images), LightningDOT is 23,000 times faster. To the best of our knowledge, this is the first known effort on improving  $\mathrm{V + L}$  model efficiency.
V+L Pre-training Inspired by the success of Transformer-based (Vaswani et al., 2017) language model pre-training (Devlin et al., 2019; Liu et al., 2019; Yang et al., 2019; Raffel et al., 2020; Lan et al., 2020; Clark et al., 2020), vision-and-language pre-training (Huang et al., 2020b; Su et al., 2020; Li et al., 2020b, 2019a) has become the prevailing paradigm in learning multimodal representations, with strong results on tasks such as image-text retrieval (Kiros et al., 2014), visual question answering (Antol et al., 2015) and referring expression comprehension (Yu et al., 2016). Exemplary works include two-stream (Tan and Bansal, 2019; Lu et al., 2019) and single-stream models (Chen et al., 2020; Li et al., 2020a; Zhou et al., 2020). Multi-task learning (Lu et al., 2020) and adversarial training (Gan et al., 2020) are also explored. This family of pre-training methods aims for general-purpose V+L without computation cost consideration. To the best of our knowledge, our work is the first known effort on pre-training visual-semantic embedding that enables low-latency real-time cross-modal retrieval. Ours is concurrent work with CLIP (Radford et al., 2021).
Image-Text Retrieval Early cross-modal embedding works (Kiros et al., 2014; Wang et al., 2018; Faghri et al., 2017) focus on using a two-stream model to learn a unified visual-semantic embedding, with progressive improvement on two popular benchmarks: Flickr30K (Plummer et al., 2015) and COCO (Chen et al., 2015). Later methods with cross-attention (Lee et al., 2018, 2019; Wang et al., 2019; Zhang et al., 2020) become more popular, with significant performance gain.
(a) Pre-training Tasks
(b) Image Retrieval Pipeline
Figure 2: An overview of our proposed framework. (a) LightningDOT is pre-trained with Sementic-embedding Fused Mask Region Modeling (SMRM), Visual-embedding Fused Mask Language Modeling (VMLM) and Cross-modal Retrieval (CMR). (b) LightningDOT ITR pipeline (image retrieval as an example). Similarities between input textual query and image candidates are computed via dot product. During inference, image representations can be computed offline, and a re-ranker can be applied for better accuracy, still with significant speedup.
Pre-trained  $\mathrm{V + L}$  models also fall into this category. By exploiting large-scale image-text datasets, pretrained  $\mathrm{V + L}$  models further push the performance on Flickr30K and COCO. Although achieving high recall, cross-attention requires excessive computation cost during inference that cannot be overlooked. In this work, inspired by dense retrieval in text retrieval domain (Guu et al., 2020; Karpukhin et al., 2020; Xiong et al., 2020; Mao et al., 2020; Lewis et al., 2020), we propose a more efficient attention-less framework. With pre-training, our model achieves better performance while being significantly faster than cross-modal attention methods. Note that the proposed approach is orthogonal to model compression techniques that reduce the number of layers/parameters (Sun et al., 2019; Jiao et al., 2020), since we do not reduce the number of parameters from the UNITER baseline. These two approaches can be combined to further boost the speed, which is an interesting future work direction.
3 LightningDOT Framework
In this section, we present the proposed LightningDOT framework, which consists of two deep Transformers as image and language encoders. We first introduce three tasks designed to pre-train the model, then present our inference pipeline from offline feature extraction to online instant retrieval.
3.1 Model Pre-training
We denote the Transformer-based (Vaswani et al., 2017) image encoder and language encoder by
$f_{\theta_V}$  and  $f_{\theta_L}$ , respectively ( $\theta_V, \theta_L$  are learnable parameters). Given a dataset of paired image and text  $\{(i,t)\}$ , we first extract region features  $\mathbf{v} = \{\mathbf{v}_0,\mathbf{v}_1,\dots,\mathbf{v}_N\}$  ( $\mathbf{v}_j \in \mathbb{R}^{d_v}$ ,  $N$  is the number of regions) for image  $i$ , along with bounding box positions of regions via a pre-trained Faster-RCNN (Ren et al., 2015; Anderson et al., 2018). The image encoder  $f_{\theta_V}$  encodes this sequence of image regions into a  $d$ -dimensional space  $f_{\theta_V}(\mathbf{v}) = \mathbf{h} = \{\mathbf{h}_0,\dots,\mathbf{h}_N\}$  ( $\mathbf{h}_j \in \mathbb{R}^d$ ). The corresponding text  $t$  is tokenized into sub-word units and projected into high-dimensional feature vectors  $\mathbf{w} = \{\mathbf{w}_0,\mathbf{w}_1,\dots,\mathbf{w}_T\}$  ( $\mathbf{w}_j \in \mathbb{R}^{d_w}$ ,  $T$  is the number of tokens) following Devlin et al. Similarly, the text encoding process can be written as  $f_{\theta_L}(\mathbf{w}) = \mathbf{z} = \{\mathbf{z}_0,\dots,\mathbf{z}_T\}$  ( $\mathbf{z}_j \in \mathbb{R}^d$ ). We regard the output [CLS] embedding  $\mathbf{h}_0$  as global image representation, and  $\mathbf{z}_0$  as global text representation. Following sections discuss how to jointly train these two encoders to learn strong visual-semantic embeddings, through three pre-training objectives.
Visual-embedding Fused Masked Language Modeling (VMLM) Masked Language Modeling (MLM) pre-training is first proposed by Devlin et al. (2019), where  $15\%$  of the words are masked and the model is trained to reconstruct the masked words. Formally, we denote  $\mathbf{w}_{\mathbf{m}} = \{\mathbf{w}_{\mathbf{m}_1},\dots ,\mathbf{w}_{\mathbf{m}_M}\}$  as masked tokens, where  $\mathbf{m}\in \mathbb{N}^{M}$  is the set of masked indices of size  $M$ , randomly sampled from a natural number  $\mathbb{N}$ .  $\mathbf{w}_{\backslash \mathbf{m}}$  are
${}^{3}\mathbf{v}_{0}$  is a special [CLS] embedding.
4A 30k BPE (Sennrich et al., 2016) vocabulary (bert-base-cased) is used to tokenize the text. A special [CLS] token is also prepended following the common practice  $(\mathbf{w}_0)$ .
In practice, this  $15\%$  is further decomposed into  $10\%$  random words,  $10\%$  unchanged, and  $80\%$  [MASK].
2The total inference time is quadratic to the dataset size with cross-attention for image-text retrieval task.
the unmasked words. MLM can be optimized by minimizing the negative log-likelihood:
$$
\begin{array}{l} \mathcal {L} _ {\mathrm {M L M}} (t) = - \log P _ {\theta_ {L}} (\mathbf {w} _ {\mathbf {m}} | \mathbf {w} _ {\backslash \mathbf {m}}) \\ = - \frac {1}{M} \sum_ {k = 1} ^ {M} \log P _ {\theta_ {\mathrm {m l m}}} \left(\mathbf {w} _ {\mathbf {m} _ {k}} \mid \mathbf {z} _ {\mathbf {m} _ {k}}\right), \tag {1} \\ \end{array}
$$
where  $\theta_{\mathrm{mlm}}$  is the additional parameters introduced to map hidden states  $\mathbf{z}$  to word probabilities.
Under the  $\mathrm{V + L}$  setting, the textual input is usually highly correlated with the image. To leverage this cross-modal relation, we propose visual-embedding fused MLM (VMLM), in which the paired image  $i$  is considered as additional input when training the model to reconstruct masked tokens in sentence  $t$ . The loss function of MLM can be formulated as:
$$
\begin{array}{l} \mathcal {L} _ {\mathrm {V M L M}} (t, i) = - \log P _ {\theta} (\mathbf {w} _ {\mathbf {m}} | \mathbf {w} _ {\backslash \mathbf {m}}, i) \\ = - \frac {1}{M} \sum_ {k = 1} ^ {M} \log P _ {\theta_ {\mathrm {m l m}}} \left(\mathbf {w} _ {\mathbf {m} _ {k}} \mid \mathbf {z} _ {\mathbf {m} _ {k}} + \mathbf {h} _ {0}\right), \tag {2} \\ \end{array}
$$
where  $\theta = \{\theta_V,\theta_L\}$  and the word probabilities  $P_{\theta}$  are conditioned on the corresponding image  $i$  via the global image representation  $\mathbf{h}_0$ . Although VMLM takes a similar mathematical form to the MLM task proposed in UNITER, they differ in two main aspects: 1) LightningDOT uses two separate encoders ( $\mathbf{h}_0$  is computed by  $f_{\theta_V}$ ); and 2) visual dependency is explicitly injected to text representations ( $\mathbf{z}_{\mathbf{m}_k} + \mathbf{h}_0$ ), instead of implicitly learned through cross-modal attention.
Semantic-embedding Fused Masked Region Modeling (SMRM) Recent works on  $\mathrm{V + L}$  pretraining (Lu et al., 2019; Tan and Bansal, 2019) have shown that mask-then-reconstruct pre-training on image regions also helps image+text embedding learning. Similar to MLM, Masked Region Modeling (MRM) is supervised by:
$$
\begin{array}{l} \mathcal {L} _ {\mathrm {M R M}} (i) = \mathcal {D} _ {\theta_ {\mathrm {m r m}}} (\mathbf {v} _ {\mathbf {m}}, f _ {\theta_ {V}} (\mathbf {v} _ {\backslash \mathbf {m}})) \\ = \frac {1}{M} \sum_ {k = 1} ^ {M} \mathcal {D} _ {\theta_ {\mathrm {m r m}}} \left(\mathbf {v} _ {\mathbf {m} _ {k}}, \mathbf {h} _ {\mathbf {m} _ {k}}\right), \tag {3} \\ \end{array}
$$
where  $\mathcal{D}$  can be any differentiable distance function. Among the variants of MRM, we consider Masked Region Feature Regression (MRFR) with L2 distance and Masked Region Classification with KL-Divergence (MRC-kl), due to their proven success in learning V+L representations (Chen et al.,
2020). In MRFR, the  $L_{2}$  distance between two feature vectors  $\mathbf{x}$  and  $\mathbf{y}$  is defined as:
$$
\mathcal {D} _ {\theta_ {\mathrm {f r}}} (\mathbf {x}, \mathbf {y}) = \sum_ {k} \| \mathbf {x} _ {k} - g _ {\theta_ {\mathrm {f r}}} (\mathbf {y} _ {k}) \| _ {2} ^ {2},
$$
where  $\| \cdot \| _2$  denotes  $L_{2}$ -norm, and  $g_{\theta_{\mathrm{fr}}}(\cdot)$  is a learnable Multi-layer Perceptron (MLP) with parameters  $\theta_{\mathrm{fr}}$ . The KL-divergence  $\mathcal{D}_{\mathrm{KL}}$  in MRC-kl measures distance between two probability distributions:
$$
\mathcal {D} _ {\theta_ {\mathrm {m r c}}} (\mathbf {x}, \mathbf {y}) = \sum_ {k} \mathcal {D} _ {\mathrm {K L}} (c (\mathbf {x} _ {k}) | | g _ {\theta_ {\mathrm {m r c}}} (\mathbf {y} _ {k})),
$$
where  $\theta_{\mathrm{mrc}}$  is the parameters of a trainable MLP that maps feature vector  $\mathbf{x}_k$  to the object class distribution  $c(\mathbf{x}_k)$  predicted by Faster R-CNN.
To incorporate language information encoded in the paired text, we extend MRM to Semantic-embedding fused MRM (SMRM), where the global text representation  $\mathbf{z}_0$  is exploited when reconstructing masked regions.
$$
\begin{array}{l} \mathcal {L} _ {\mathrm {S M R M}} (i, t) = \mathcal {D} _ {\theta_ {\mathrm {m r m}}} (\mathbf {v} _ {\mathbf {m}}, f _ {\theta_ {V}} (\mathbf {v} _ {\backslash \mathbf {m}}), t) \\ = \frac {1}{M} \sum_ {k = 1} ^ {M} \mathcal {D} _ {\theta_ {\mathrm {m r m}}} \left(\mathbf {v} _ {\mathbf {m} _ {k}}, \mathbf {h} _ {\mathbf {m} _ {k}} + \mathbf {z} _ {0}\right). \tag {4} \\ \end{array}
$$
The specific variants SMRFR and SMRC-kl can be derived using the corresponding distance function, which is omitted for simplicity. Note that both the cross-modal fusion introduced in Eqn. (2) and Eqn. (4) uses simple addition without introducing extra parameters from their uni-modal counterpart. Moreover, the extra parameters  $\theta_{\mathrm{mll}}$  and  $\theta_{\mathrm{mrm}}$  is not needed at downstream inference so will not slow down the retrieval.
Cross-modal Retrieval Objective (CMR) Beyond image or text focused reconstructive objectives, we also propose a new pre-training task, Cross-modal Retrieval (CMR), to leverage the paired information between image and text. With this learning objective, the model is optimized to promote high similarity score for a matched image-sentence pair  $(i,t)$  and vice versa. The similarity score between query  $t$  and image  $i$  is defined as:
$$
S (t, i) = \left\langle \mathbf {z} _ {0}, \mathbf {h} _ {0} \right\rangle , \tag {5}
$$
where  $\langle \cdot ,\cdot \rangle$  denotes the inner product between two vectors, and  $\mathbf{h}_0$  and  $\mathbf{z}_0$  are the output [CLS] embeddings from image encoder  $f_{\theta_V}$  and language encoder  $f_{\theta_L}$ , respectively.
In our implementation, no textual inputs are directly concatenated with image regions due to separate encoding of image and text.
Figure 3: An illustration of the proposed CMR Loss. Note that positive pairs lie in the diagonal of the matrix.
In order to capture both image-retrieval and text-retrieval supervision signals in a single forward-backward pass, we propose a bi-directional variant of contrastive loss. Given any matched image-text pair  $(i,t)$ , we treat text  $t$  as the query, sample  $n - 1$  negative images  $\{i_2,i_3,\dots ,i_n\}$ , and then compute the objective function as:
$$
\mathcal {L} _ {\mathrm {I R}} ^ {(t)} = - \log \frac {e ^ {S (t , i _ {1})}}{\sum_ {k = 1} ^ {n} e ^ {S (t , i _ {k})}},
$$
where  $t_1 \coloneqq t$ . Similarly, we take image  $i$  as query  $(i_1 \coloneqq i)$ , sample  $n - 1$  negative text, and compute:
$$
\mathcal {L} _ {\mathrm {T R}} ^ {(i)} = - \log \frac {e ^ {S (i , t _ {1})}}{\sum_ {k = 1} ^ {n} e ^ {S (i , t _ {k})}}
$$
to optimize for text retrieval.
Following Henderson et al. (2017); Gillick et al. (2019); Karpukhin et al. (2020), we use in-batch negatives to avoid the actual sampling of a negative image or text: given a batch of  $n$  positive image-text pairs  $B = \{(i_1, t_1), \ldots, (i_n, t_n)\}$ , we use all other images from within the batch as negatives ( $\{i_j\}$ , where  $j \in \{1, 2, \ldots, n\}$  and  $j \neq k$ ) for every positive pair  $(i_k, t_k)$ , and vice versa for negative text. The final CMR loss for batch  $B$  is:
$$
\mathcal {L} _ {\mathrm {C M R}} (B) = \frac {1}{2 n} \sum_ {k = 1} ^ {n} \mathcal {L} _ {\mathrm {T R}} ^ {(i _ {k})} + \mathcal {L} _ {\mathrm {I R}} ^ {(t _ {k})}. \tag {6}
$$
An illustration of  $\mathcal{L}_{\mathrm{CMR}}$  is presented in Figure 3.7 Through joint pre-training with CMR, VMLM and SMRM, the visual-semantic embeddings learned from image encoder and language encoder can be readily applied to downstream tasks. During finetuning stage, we directly adopt CMR loss to supervise the training process.
3.2 Real-time Inference
For simplicity, we take text-to-image retrieval as an example to introduce the real-time inference pipeline (Figure 2(b)): (i) Offline image feature extraction and encoding; (ii) Online retrieval with text query; and (iii) Online re-ranking with top-retrieved images. Text retrieval is conducted in a symmetric manner.
Offline Feature Extraction Image retrieval task requires the model to rank every image  $i$  in an image database  $I$  based on its similarity to a text query  $t$ . In LightningDOT, we first apply the image encoder  $f_{\theta_V}$  to all images in  $I$ , and cache the resulting global image representations  $\{\mathbf{h}_0^{(i)} \in \mathbb{R}^d | i \in I\}$  into an index (Johnson et al., 2019) in memory for later use. Note that the entire image-to-index process, including Faster-RCNN feature extraction and Transformer encoding, can all be conducted offline. Therefore, for every new query  $t$  at real time, the cached index can be reused for maximum inference time saving.
Online Retrieval During inference, given a text query  $t$ , we encode it with the language encoder  $\theta_{L}$  and then compute its similarity score to the embedding of every image in  $I$  (stored in memory index) via Eqn (5). Finally, the images will be ranked by their similarity scores, from the highest to lowest. In practice, people are more interested in top-  $K$  retrieval, with a list of  $K$  images  $I_{t}$  satisfying:
$$
I _ {t} := \left\{i _ {m _ {k}} \right\} _ {k = 1} ^ {K}, \text {w h e r e}
$$
$$
S (t, i _ {m _ {1}}) \geq S (t, i _ {m _ {2}}) \geq \dots \geq S (t, i _ {m _ {K}}) \quad \text {a n d}
$$
$$
S \left(t, i _ {m _ {K}}\right) \geq S (t, i) \quad \forall i \in \left(I \backslash I _ {t}\right). \tag {7}
$$
This optimization problem has been well studied, and we use FAISS (Johnson et al., 2019) to solve it in our implementation. It is worth noting that in order to apply fast search, the similarity function has to be decomposable. Therefore, we choose the simple dot product as  $S$  instead of a more complicated neural network function. Similarly, for text retrieval, the same architecture can be applied by simply pre-computing the embedding for all sentences and using an image as query instead.
Re-ranking To further improve retrieval accuracy, we propose a two-stage approach by adopting an optional re-ranking model. In the first stage, we use LightningDOT to retrieve top-  $M$  images (or texts), where  $M$  is an integer much smaller
The whole similarity matrix can be computed efficiently with one batched matrix multiplication call. This operation can take advantage of GPU hardware with Tensor Cores for faster training.