Source-linked AI summary
Thinking Fast and Slow: Efficient Text-to-Visual Retrieval with Transformers
Antoine Miech, Jean-Baptiste Alayrac, Ivan Laptev, Josef Sivic, Andrew Zisserman
TL;DR
Large-scale image and video retrieval needs the scalability of dual encoders without sacrificing the accuracy of cross-attention models. The paper combines fine-grained cross-attention with distillation and re-ranking, achieving better results than the Slow model while reducing inference time by several orders of magnitude.
Problem
Dual encoders scale efficiently but have limited interaction modeling, whereas cross-attention improves accuracy but is costly for large-scale retrieval.
Method
The paper combines a fine-grained cross-attention model with a Fast dual encoder through knowledge distillation and candidate re-ranking.
Results
The combined Fast and Slow approach achieves better results than the Slow model while significantly reducing inference time by several orders of magnitude on large datasets.
Takeaways & Limitations
The approach provides a model-agnostic way to scale vision-text Transformer retrieval using dual-encoder indexing and Slow-model re-ranking.
Takeaways & Limitations
The Slow cross-attention model is expensive to run, and its encoded image representation is too large to keep an entire large dataset in memory.
Abstract
from arXiv · showhide
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.
1. Introduction
Large-scale text-to-image retrieval must balance dual encoders’ scalability with cross-attention models’ accuracy. The paper combines them through distillation, fine-grained cross-attention, and re-ranking.
- Dual encoders separately map images and text into a joint space, enabling dot-product similarity and approximate nearest-neighbour search at billion-image scale.
- Cross-attention models improve retrieval by grounding words to image locations, but their per-sample test-time cost makes large-scale search impractical.
- The proposed system distills the Slow cross-attention model into an indexable Fast dual encoder and uses the Slow model to re-rank a small candidate set.
- A fine-grained cross-attention architecture improves vision-text interactions through feature-map upsampling, while preserving scalability.
- The combined Fast and Slow approach is designed to improve retrieval accuracy while reducing inference cost for large-scale search.
2. Related work
Related work applies Transformer architectures and other neural models to vision-and-language tasks, often combining visual features with text objectives. Distillation is established for compressing expensive models, motivating its use here for non-indexable vision-text retrieval.
- Vision-language Transformers use multimodal objectives such as cross-modal matching, masked language modelling, and masked region modelling.
- Some vision-language systems operate directly on dense feature maps instead of relying on pretrained object-detector regions.
- Earlier vision-and-text approaches also use recurrent neural networks, multilayer perceptrons, or bag-of-words text models.
- Distillation has been used to compress computationally expensive models in language analysis, detection, classification, and speech recognition.
3. Thinking Fast and Slow for Retrieval
The paper combines fast, indexable dual encoders with accurate but costly cross-attention models for scalable language-based image retrieval. It introduces fine-grained visual-text attention, distillation, and re-ranking to preserve accuracy while improving retrieval efficiency.
- Fast and Slow retrieval: Dual encoders compute image-text similarity with a dot product, enabling approximate nearest neighbour search for efficient large-scale retrieval.Their simple interaction limits accuracy compared with richer cross-attention models.
- Fast and Slow retrieval: Cross-attention models jointly process visual and textual representations, improving scoring through richer interactions but incurring higher computational cost.The model computes similarity through cross-modal attention rather than a decomposable dot product.
- Thinking Slow with cross-attention: Fine-grained cross-attention increases visual feature resolution by gradually upsampling convolutional features while incorporating earlier high-resolution maps.The architecture avoids visual self-attention, allowing feature maps with thousands of vectors without substantially increasing backbone cost.
- Thinking Slow with cross-attention: The Slow model scores an image-text pair using forward and backward caption log-likelihoods from Transformer decoders conditioned on encoded image features.Captioning models can perform on-par with models trained using conventional contrastive image-text matching loss.
- Thinking Faster and better for retrieval: Scaling the Slow model is constrained because its network is expensive to run and its encoded image representation is too large to store for an entire dataset.The proposed Fast-and-Slow approach addresses these large-scale retrieval constraints through indexing, distillation, and selective re-ranking.
- Thinking Faster and better for retrieval: Distillation transfers Slow-model knowledge to an indexable Fast dual encoder using teacher distributions and cross-entropy, while re-ranking combines Fast retrieval with Slow scoring.Re-ranking fewer than ten examples out of thousands can recover Slow-model performance, and the combined approach reports more than 100× speed-up with improved retrieval accuracy.
4. Experiments
The experiments show that cross-attention improves retrieval, while distillation and re-ranking combine its accuracy with fast dual-encoder retrieval. The approach improves efficiency and achieves strong image and video retrieval results.
- Cross-attention retrieval: Cross-attention models generally outperform dual encoders on COCO and Flickr30K retrieval.The comparison includes PixelBERT and VirTex against NCE BoW and BERT.
- Cross-attention retrieval: Captioning objectives can outperform image-text matching and masked-language-modelling objectives for retrieval.Backward captioning further improves retrieval performance.
- Fast and Slow retrieval: Distillation improves the Fast dual encoder by over 10% on R@1 when training on COCO.Improvements are more moderate on Conceptual Captions, where distillation often takes longer to converge.
- Fast and Slow retrieval: Re-ranking only 10 candidates can recover or exceed the Slow model’s R@1 while substantially reducing query time.On COCO, combining distillation with re-ranking reduces the required candidate count from K = 50 to K = 10.
- Fast and Slow retrieval: Re-ranking five of 5,000 images from the distilled Fast model reaches the Slow model’s COCO R@1 performance.Figure 3 examines retrieval as the number of top-K candidates and β vary.
- State-of-the-art comparisons: The method outperforms selected state-of-the-art image-retrieval baselines but remains below UNITER and OSCAR, while achieving a new state of the art on VATEX video retrieval.The reported image comparisons use Flickr30K, and the video result uses VATEX.
5. Conclusion
The paper scales vision-text transformer retrieval by combining a fine-grained Slow cross-attention model with a Fast dual encoder through distillation and re-ranking. This combined approach improves accuracy while reducing inference time by several orders of magnitude and is model agnostic.
- 5. Conclusion: The combined Fast and Slow approach improves retrieval accuracy while reducing inference time by several orders of magnitude on large datasets.It combines distillation with re-ranking to make cross-attention retrieval scalable.
- 5. Conclusion: The approach can be applied to any vision-text Transformer Slow model and dual-encoder Fast retrieval model.The supplementary material provides additional ablations, hyperparameters, architecture details, video extension details, and qualitative results.
B. Additional quantitative results
Additional experiments show that distillation quality depends on the student text model, temperature, and loss combination, while re-ranking only a few Fast-model candidates can recover Slow-model performance.
- Distillation ablations: The BoW student text model performed best for distillation, suggesting complex language models were unnecessary for this task.This finding aligns with the reported lack of improvement from more complex language models for NCE.
- Distillation ablations: Distillation required careful temperature selection and combining the distillation loss with the original loss to ensure improvements.The experiments used τ = 10 and ατ^2 = 0.001 for subsequent distillation experiments.
- Re-ranking: Re-ranking four images out of one thousand retrieved by the distilled Fast model reached the Slow model’s Flickr validation R@1 performance.Figure 4 examines how re-ranking performance changes with the number of retrieved candidates K.
C. Experiment details
The experiments use standard image augmentation and specified optimization settings across datasets and model families, with separate preprocessing for videos and modified PixelBERT inputs.
- Data augmentation: Training crops are resized to 224 × 224 and horizontally flipped with probability 0.5; video training additionally samples 32-frame clips at 10 frames per second.The video clips span 3.2 seconds from the original 10-second VATEX clips.
- Test-time processing: At test time, image retrieval uses a central crop, while video scores average four uniformly sampled 32-frame clips.Each video clip is centrally cropped before visual-text scores are averaged.
- Optimization: All models use ADAM with cosine learning-rate decay and 5k warm-up steps, while training schedules vary by model and dataset.The listed configurations include distinct learning rates, batch sizes, step counts, and gradient clipping values for NCE, VirTex, Slow, and PixelBERT.
- PixelBERT: PixelBERT removes its downsampling and random-pixel-sampling modules and feeds all 7 × 7 visual features to the transformer.The modification addresses the use of smaller image resolutions than in the original PixelBERT implementation.
D. Architecture details
The upsampling block fuses the input feature map with a resized previous ResNet feature map using nonnegative learned weights and a lightweight convolutional refinement.
- Feature fusion: Fast normalized fusion combines P_in with Resize(P_prev) using nonnegative learned weights and a small stabilizing ε.The weights are constrained by applying ReLU, and Resize performs 2 × nearest-neighbour upsampling.
D.2. Video architecture with upsampling
For VATEX, the image architecture is adapted to video by starting from TSM ResNet-50 and adding spatio-temporal processing through temporal feature shifts.
- Video adaptation: The video model starts from TSM ResNet-50, whose Temporal Shift Modules move features forward and backward along the time dimension.This adapts the image-only architecture to handle spatio-temporal attention.
E.1. Retrieval results
The qualitative Flickr examples compare Fast and Slow retrieval, including Fast candidates re-ranked by Slow. Distillation improves the Fast model’s retrieval of the annotated groundtruth.
- Retrieval comparison: Four retrieval settings are compared on Flickr: Fast without distillation, Fast, Slow, and Fast and Slow with K=50 re-ranking.The models are trained on COCO and evaluated on Flickr zero-shot.
- Retrieval comparison: The Slow model is used to re-rank the most promising Fast candidates.The final row shows the top K=50 Fast candidates re-ranked with the Slow model.
- Qualitative example: The groundtruth image is absent from the Fast model’s top-5 without distillation but appears fourth after distillation.
E.2. Attention maps
The attention visualizations show how the Slow model relates text tokens to image regions across transformer layers. Its 56 × 56 feature map yields finer localization than VirTex’s 7 × 7 map, while attention becomes more focused in higher layers.
- Attention extraction: Each text token produces a query that scores an H × H visual feature map, whose softmax weights aggregate visual values.
- Visualization: The visualizations overlay layer-specific attention maps with image tokens shifted backward to show the regions used to predict each word.
- Observed patterns: Attention regions show coherence with predicted tokens, such as the mouth for “smilling,” the shirt for “peach,” and the bicycle for the final word.
- Observed patterns: The Slow model produces finer-grained attention than VirTex because its 56 × 56 feature map is higher resolution than VirTex’s 7 × 7 map.
- Observed patterns: Attention becomes more focused in higher layers, progressing from whole-image coverage in the first layer to a specific region.
- Caveat: The authors consider these visualizations valuable for understanding text–vision relations despite imperfect interpretability.