Leaner and Faster: Two-Stage Model Compression for Lightweight Text-Image Retrieval

Dual-encoder architecture employs two separate neural networks to encode inputs from different modalities and map them to a shared space.

We denote the image encoder as $f_{v}$ and the text encoder as $f_{t}$ in the context of text-image retrieval. To train $f_{v}$ and $f_{t}$, it is common to adopt an objective that pushes the embeddings of matched text-image pairs closer while pushing those of non-matched text-image pairs apart. Specifically, Contrastive Language-Image Pretraining (CLIP) Radford et al. (2021) optimizes an InfoNCE van den Oord et al. (2018) loss:

Here, $f_{t}(x_{i})$ and $f_{v}(y_{j})$ are the L2-normalized embeddings of text in the $i$-th pair and image in the $j$-th pair. $N$ is the mini-batch size and $\tau$ is the temperature to scale the logits. The final objective is the sum of $L_{t2v}$ and its symmetric version $L_{v2t}$.

Despite good retrieval accuracy, models like CLIP still pose non-trivial memory footprint and inference time, which is undesirable for low-resource devices such as smart phones.

To tackle this issue, we propose a two-stage compression framework to make large dual-encoder model smaller and faster while retaining its accuracy. A schematic overview is illustrated in Figure 1. The first stage is task-agnostic, where we leverage massively available non-paired texts/images to separately compress the text/image encoder using an intra-modal contrastive knowledge distillation scheme. The second stage is task-specific, where we sequentially fine-tune the distilled image and text encoder using a combination of multiple techniques. We denote the image and text encoder of the large dual-encoder as $f_{v}^{T}$ and $f_{t}^{T}$ and those of the compressed model as $f_{v}^{S}$ and $f_{t}^{S}$.

Dataset. We use Conceptual Caption Sharma et al. (2018) for stage-1 compression. It consists of 3M noisy image alt-text pairs. However, we do not use the image-text alignment information but only treat it as a reservoir of general-domain images and texts. In stage-2, we use MSCOCO Lin et al. (2014) and Flickr30K Plummer et al. (2015) as the benchmarks. For MSCOCO, there are 113,287 images for training, 5,000 images for validation, and both 5K and 1K for testing. For Flickr30K, there are 28,783 images for training, 1,000 images for validation, and 1k for testing.

Implementation Detail. In stage-1, we train 1 epoch using AdamW Loshchilov and Hutter (2017) with a batch size of 84 for both images and texts, learning rate of 3e-4, and weight decay of 0.1. In stage-2, we use the same optimization setting except that we train with batch size 96 for 5 epochs. We employ a cosine learning rate scheduler with 10,000 warm-up steps for both stages. All reported results are calculated on the test set using checkpoints with the highest validation performance.

Main Results. Table 1 summarizes the main results. As can be observed, the pre-trained CLIP model can already deliver moderately good retrieval performance. The performance is further improved after fine-tuning. Fine-tuning pre-trained ViT-S/16 and TinyBERT underperforms the zero-shot CLIP, showing that training with inter-modal InfoNCE is not effective without extremely large-scale paired data. On most evaluation metrics, models compressed by our proposed two-stage pipeline perform on par with or better than the fine-tuned target model. We also found that the capacity of text encoder has limited affect on the performance.

Ablation Study. We perform extensive ablations to study the importance of each proposed technique. Due to the computational budget, we only conduct ablation on the image encoder and fix the text encoder as $f_{t}^{T^{\prime}}$. We evaluate w/o stage-1, stage-1_MSE (mean-square-error between $f_{v}^{S}$ and $f_{v}^{T}$), and stage-1_InfoNCE (identical to the loss in Section 3.1) for stage-1 ablation. We also study the effectiveness of KD/SF/HN by removing them separately or together. We made several observations based on Table 2: 1) SF makes fine-tuning stable and is essential for convergence. 2) both KD and HN improve retrieval accuracy and are complementary to each other. 3) intra-modal contrastive distillation helps when image-text pairs are noisy and outperforms inter-modal infoNCE loss.

Efficiency. In Table 3, we compare the disk space and QSP used by models on a RTX 2080Ti of 12GB memroy. The compressed image encoder $f_{v}^{S}$ takes 85MB disk space (39% of $f_{v}^{T}$) meanwhile being 1.51x times faster. Our compressed text encoder can achieve up to x2.77 inference speed-up and 40% size reduction (from 243MB to 146MB). We further benchmark models’ memory and run-time performance on a real iPhone X with 1,000 images in the gallery for testing. It takes 870MB and 295MB for loading CLIP and our compressed model into main memory respectively. After indexing, the response time for a single text query is 0.4s for CLIP while it is only 0.1s for our compressed model. Although the results are hardward-dependent, our compressed model still shows an evident improvement in efficiency.