[Paper Review] Thinking Fast and Slow: Efficient Text-to-Visual Retrieval with Transformers

Antoine Miech, Jean-Baptiste Alayrac+3 more|arXiv (Cornell University)|Mar 30, 2021

Multimodal Machine Learning ApplicationsComputer Science89 references130 citations

TL;DR

The paper combines fast dual-encoder models with slow cross-attention transformers via distillation and re-ranking to achieve scalable, high-accuracy text-to-vision retrieval. It introduces fine-grained cross-attention with feature upsampling and a bi-directional captioning objective, validated on COCO, Flickr30K, and VATEX datasets.

ABSTRACT

Our objective is language-based search of large-scale image and video datasets. For this task, the approach that consists of independently mapping text and vision to a joint embedding space, a.k.a. dual encoders, is attractive as retrieval scales and is efficient for billions of images using approximate nearest neighbour search. An alternative approach of using vision-text transformers with cross-attention gives considerable improvements in accuracy over the joint embeddings, but is often inapplicable in practice for large-scale retrieval given the cost of the cross-attention mechanisms required for each sample at test time. This work combines the best of both worlds. We make the following three contributions. First, we equip transformer-based models with a new fine-grained cross-attention architecture, providing significant improvements in retrieval accuracy whilst preserving scalability. Second, we introduce a generic approach for combining a Fast dual encoder model with our Slow but accurate transformer-based model via distillation and re-ranking. Finally, we validate our approach on the Flickr30K image dataset where we show an increase in inference speed by several orders of magnitude while having results competitive to the state of the art. We also extend our method to the video domain, improving the state of the art on the VATEX dataset.

Motivation & Objective

• Motivate language-based search over large image/video datasets using scalable retrieval.
• Develop a fast, indexable dual-encoder that benefits from accurate cross-attention through distillation.
• Propose re-ranking of fast results with a slow cross-attention model to boost accuracy.
• Introduce a fine-grained cross-attention architecture and a captioning-based training objective.
• Demonstrate scalability and state-of-the-art-like performance on image and video retrieval datasets.

Proposed method

• Define fast dual-encoder (DE) and slow cross-attention (CA) models for image-text retrieval.
• Introduce a novel gradual feature upsampling architecture to enable fine-grained vision-language cross-attention.
• Use a bi-directional captioning loss to train the slow CA model for retrieval.
• Distill knowledge from the slow CA model into the fast DE model using a teacher-student distillation objective over sampled image-text pairs.
• Combine fast DE retrieval with slow CA re-ranking by selecting top-K candidates from the fast model and re-ranking with the slow model (Eq. 9).
• Optionally improve CA efficiency by applying a re-ranking step at query time rather than for all candidates.

Experimental results

Research questions

• RQ1 Can cross-attention models improve text-to-vision retrieval while remaining scalable with indexing?
• RQ2 Can knowledge distilled from a slow cross-attention model improve a fast dual-encoder without sacrificing indexability?
• RQ3 Does a re-ranking step using a slow model on top-K fast results yield substantial accuracy gains with modest latency?
• RQ4 Does a captioning-based training objective perform competitively for retrieval compared to traditional contrastive/objective losses?
• RQ5 How does gradual feature upsampling impact fine-grained grounding in vision-language retrieval?],
• RQ6 key_findings:[

Key findings

• Cross-attention models generally outperform dual encoders on retrieval benchmarks, but dual encoders with distillation can close the gap.
• Bi-directional captioning objectives yield retrieval-competitive results, sometimes surpassing standard cross-modal matching losses.
• Gradual upsampling of high-resolution visual features significantly improves recall at 1 (R@1) and recall at 5 (R@5) over baseline cross-attention models.
• Distilling Slow CA into Fast DE improves Fast DE performance by over 10 percentage points in R@1 when trained on COCO, and reduces the gap to Slow models.
• Re-ranking top-K fast results with the Slow model achieves state-of-the-art-like performance with substantial speedups (e.g., ~100x faster than Slow alone on COCO).
• The combined Fast+Slow approach yields strong results on image datasets (COCO, Flickr30K) and extends effectively to video (VATEX).