[Paper Review] Training Vision Transformers for Image Retrieval

Alaaeldin El-Nouby, Natalia Neverova+2 more|arXiv (Cornell University)|Feb 10, 2021

Advanced Image and Video Retrieval TechniquesComputer Science52 references122 citations

TL;DR

This paper shows that vision transformers can be effectively trained for image retrieval using a Siamese transformer architecture with contrastive loss and a differential entropy regularizer, achieving state-of-the-art results on category-level and strong results on particular object retrieval.

ABSTRACT

Transformers have shown outstanding results for natural language understanding and, more recently, for image classification. We here extend this work and propose a transformer-based approach for image retrieval: we adopt vision transformers for generating image descriptors and train the resulting model with a metric learning objective, which combines a contrastive loss with a differential entropy regularizer. Our results show consistent and significant improvements of transformers over convolution-based approaches. In particular, our method outperforms the state of the art on several public benchmarks for category-level retrieval, namely Stanford Online Product, In-Shop and CUB-200. Furthermore, our experiments on ROxford and RParis also show that, in comparable settings, transformers are competitive for particular object retrieval, especially in the regime of short vector representations and low-resolution images.

Motivation & Objective

• Demonstrate that plain vision transformers can be used for image retrieval with competitive accuracy
• Investigate the interaction of metric learning losses with transformer backbones
• Assess whether differential entropy regularization improves embedding space utilization
• Establish state-of-the-art results on SOP, CUB-200-2011, and In-Shop for category-level retrieval
• Evaluate performance on Oxford and Paris datasets for particular object retrieval

Proposed method

• Use a Siamese Vision Transformer (IRT) architecture to map image pairs into a common embedding space
• Apply a contrastive loss with cross-batch memory as the primary supervision signal
• Augment the contrastive loss with a differential entropy regularizer to encourage uniformity in the embedding space
• Experiment with off-the-shelf ViT features, fine-tuning with metric learning (IRT_O, IRT_L, IRT_R)
• Explore pooling variants (CLS token, average, max, GeM) and dimensionality reduction (PCA) to obtain compact descriptors
• Train and evaluate on SOP, CUB-200-2011, In-Shop (category-level) and on Oxford/Paris (particular object) using standard retrieval metrics

Experimental results

Research questions

• RQ1 Can a plain Vision Transformer backbone trained with metric learning achieve competitive or superior performance for category-level image retrieval compared to convolutional baselines?
• RQ2 Does fine-tuning ViT with a contrastive loss improve retrieval performance over off-the-shelf ViT features?
• RQ3 Does adding a differential entropy regularizer to the contrastive loss further improve embedding space utilization and retrieval accuracy?
• RQ4 How do transformer-based descriptors compare to convolutional descriptors in particular object retrieval at various descriptor sizes and image resolutions?

Key findings

• IRT_R with DeiT-S backbones achieves state-of-the-art Recall@1 on SOP, outperforming prior methods by a notable margin
• On CUB-200-2011, DeiT-S 384 with regularized training outperforms prior art at Recall@1
• For In-Shop, DeiT-S 384 yields superior Recall@1 versus prior convnet-based methods
• In particular object retrieval, DeiT-S and DeiT-B variants outperform ResNet-50/101 at 224x224 and scale well to 384x384, with competitive FLOPS
• Differential entropy regularization improves performance across benchmarks and mitigates feature collapse observed with plain contrastive loss
• Transformers show robustness to feature collapse and can match or exceed convnets at comparable capacity and resolutions