Training Vision Transformers for Image Retrieval

Introduction

One of the fundamental skills in reasoning is the ability to predict similarity between entities even if such entities have not been observed before. In the context of computer vision, learning similarity metric has many direct applications such as content-based image retrieval, face recognition and person re-identification. It is also a key component of many other computer vision tasks like zero-shot and few-shot learning. More recently, advances in metric learning have been essential to the progress of self-supervised learning, which relies on matching two images up to data augmentation as a learning paradigm.

Modern methods for image retrieval typically rely on convolutional encoders and extract compact image-level descriptors. Some early approaches used activations provided by off-the-shelf pre-trained models Babenko et al. (2014). However, models trained specifically for the image retrieval task achieve better performance Radenović et al. (2018b); Teh et al. (2020); Wang et al. (2020). A number of efficient objective functions have been proposed to penalize the discrepancy between computed similarities and the ground truth. In addition, research has been focused on improvements of sampling methods and data augmentation strategies.

The transformer architecture by Vaswani et al. (2017) has been successfully used for a number of NLP tasks Devlin et al. (2018); Radford et al. (2018), and more recently in the core computer vision task of image classification Dosovitskiy et al. (2021); Touvron et al. (2020). This is an interesting development as transformer-based models adapted for computer vision come with a different set of inductive biases compared to the currently dominant convolutional architectures. This suggests that such models may find alternative solutions and avoid errors that are typical for convolutional backbones. While there have been some efforts exploring attention-based metric learning for images Kim et al. (2018); Chen & Deng (2019), to our knowledge the adoption of a plain transformer has not been studied in this context.

In this work, we introduce and study Image Retrieval Transformers (IRT). As illustrated in Figure 1, our IRT model instantiates a Siamese architecture with a transformer backbone. We investigate the adaptation of metric learning techniques and evaluate how they interplay with transformers. In particular, we adopt a contrastive loss Hadsell et al. (2006), which has recently been reaffirmed as a very effective metric learning objective Musgrave et al. (2020); Wang et al. (2020). We also employ a differential entropy regularization that favors uniformity over the representation space and improves performance.

We perform an extensive experimental evaluation and validate our approach by considering two image retrieval tasks. First, we investigate the task of category-level image retrieval, which is often used to measure the progress in metric learning Teh et al. (2020); Musgrave et al. (2020). We also explore retrieval of particular objects, and compare our method to convolutional baselines in similar settings (same resolution and similar complexity).

The main contributions of this work are listed below.

We propose a simple way to train vision transformers both for category-based level and particular object retrieval, and achieve competitive performance when compared to convolutional models with similar capacity.

As a result, we establish the new state of the art on three popular benchmarks for category-level retrieval.

For particular object retrieval, in the regime of short-vector representation (128 components), our results on ${\mathcal{R}}$Oxford and ${\mathcal{R}}$Paris are comparable to those of convnets operating at a much higher resolution and FLOPS.

We show that the differential entropy regularizer enhances the contrastive loss and improves the performance overall.

Related Work

The transformer architecture was introduced by Vaswani et al. (2017) for machine translation. It solely relies on self-attention and fully-connected layers, and achieving an attractive trade-off between efficiency and performance. It has subsequently provided state-of-the-art performance for several NLP tasks Devlin et al. (2018); Radford et al. (2018). In computer vision, several attempts have been devoted to incorporate various forms of attention, for instance in conjunction Wang et al. (2018) or as a replacement to convolution Ramachandran et al. (2019). Other methods utilize transformer layers on top of convolutional trunks Carion et al. (2020) for detection.

More recently, convolution-free models that only rely on transformer layers have shown competitive performance Chen et al. (2020); Dosovitskiy et al. (2021); Touvron et al. (2020), positioning it as a possible alternative to convolutional architectures. In particular, the Vision Transformers (ViT) model proposed by Dosovitskiy et al. (2021) is the first example of a transformer-based method to match or even surpass state-of-the-art convolutional models on the task of image classification. Touvron et al. (2020) subsequently improved the optimization procedure, leading to competitive results with ImageNet-only training Deng et al. (2009).

Metric Learning.

A first class of deep metric learning methods is based on classification: these approaches represent each category using one Movshovitz-Attias et al. (2017); Teh et al. (2020); Zhai & Wu (2018); Boudiaf et al. (2020) or multiple prototypes Qian et al. (2019). The similarity and dissimilarity training signal is computed against the prototypes rather than between individual instances. Another class of methods operate on pairs methods: the training signal is defined by similarity/dissimilarity between individual instances directly. A contrastive loss Hadsell et al. (2006) aims to push representations of positive pairs closer together, while representations of negative pairs are encouraged to have larger distance. The triplet loss Weinberger & Saul (2009) builds on the same idea but requires the positive pair to be closer than a negative pair by a fixed margin given the same anchor. Wu et al. (2017) proposes negative sampling weighted by pair-wise distance to emphasize harder negative examples. Other pair-based losses rely on the softmax function Goldberger et al. (2005); Sohn (2016); Wu et al. (2018); Wang et al. (2019), allowing for more comparisons between different positive and negative pairs.

While a vanilla contrastive loss has been regarded to have a weaker performance when compared to its successors like triplet and margin Wu et al. (2017) losses, recent efforts Musgrave et al. (2020) showed that a careful implementation of the contrastive loss leads to results outperforming many more sophisticated losses. Additionally, Wang et al. (2020) showed that when augmented with an external memory to allow sampling of a sufficient number of hard negatives, contrastive loss achieves a state-of-the-art performance on multiple image retrieval benchmarks.

Particular Image Retrieval

has progressively evolved from methods based on local descriptors to convolutional encoders. In this context, an important design choice is how to compress the spatial feature maps of activations into a vector-shaped descriptor Babenko & Lempitsky (2015); Tolias et al. (2015). Subsequent works have adopted end-to-end training Gordo et al. (2016); Radenović et al. (2018b); Revaud et al. (2019) with various forms of supervision. In a concurrent work, Gkelios et al. (2021) investigated off-the-shelf pre-trained ViT models for particular image retrieval.

Differential Entropy Regularization.

Zhang et al. (2017) aim a better utilization of the space by spreading out the descriptors through matching first and second moments of non-matching pairs with points uniformly sampled on the sphere. Wang & Isola (2020) provide a theoretical analysis for contrastive representation learning in terms of alignment and uniformity on the hypersphere. In the context of face recognition, Duan et al. (2019) argue for spreading the class centers uniformly in the manifold, while Zhao et al. (2019) minimize the angle between a class center and its nearest neighbor in order to improve inter-class separability.

We focus our study on pairwise losses, the contrastive loss in particular, aiming to prevent the collapse in dimensions that happens as a byproduct of adopting such an objective. Most related to out method, Sablayrolles et al. (2019) propose a differential entropy regularization based on the estimator by Kozachenko & Leonenko (1987), in order to spread the vectors on the hypersphere more uniformly, such that it enables improved lattice-based quantization properties. Bell et al. (2020) adopted it as an efficient way to binarize features output by convnets in commerce applications.

Methods

In this section, after reviewing the transformer architecture, we detail how we adapt it to the category-level and particular object retrieval. Note, in the literature these tasks have been tackled by distinct techniques. In our case we use the same approach for both of these problems. We gradually introduce its different components, as follows:

IRT${}_{\text{O}}$– off-the-shelf extraction of features from a ViT backbone, pre-trained on ImageNet;

IRT${}_{\text{L}}$– fine-tuning a transformer with metric learning, in particular with a contrastive loss;

IRT${}_{\text{R}}$– additionally regularizing the output feature space to encourage uniformity.

Let us review the main building blocks for transformer-based models, and more specifically of the recently proposed ViT architecture by Dosovitskiy et al. (2021). The input image is first decomposed into $M$ fixed-sized patches (e.g. 16$\times$16). Each patch is linearly projected into $M$ vector-shaped tokens and used as an input to the transformer in a permutation-invariant manner. The location prior is incorporated by adding a learnable 1-D positional encoding vector to the input tokens. An extra learnable CLS token is added to the input sequence such that its corresponding output token serves as a global image representation.

The transformer consists of $L$ layers, each of which is composed of two main blocks: a Multi-Headed Self Attention (MSA) layer, which applies a self-attention operation to different projections of the input tokens, and a Feed-Forward Network (FFN). Both the MSA and FFN layers are preceded by layer normalization and followed by a skip connection. We refer the reader to Dosovitskiy et al. (2021) for details.

Table 1 presents the neural networks models used through this paper. They are all pre-trained on ImageNet1k Deng et al. (2009) only. In order to have a fair comparison with other retrieval methods, we choose to use the DeiT-Small variant of the ViT architecture introduced by Touvron et al. (2020) as our primary model. The DeiT-Small model has a relatively compact size which makes it comparable to the widely adopted ResNet-50 convolutional model He et al. (2016) in terms of parameters count and FLOPS, as shown in Table 1. Additionally, we provide some analysis and results of larger models like ResNet-101 and DeiT-Base, as well as DeiT variants with advanced pre-training.

2 IRTOO{}_{\text{O}}: off-the-shelf features with Transformers

We first consider the naive approach IRT${}_{\text{O}}$, where we extract features directly from a transformer pre-trained on ImageNet. This strategy is in line with early works on image retrieval with convolutional networks Babenko et al. (2014), which were featurizing activations.

We extract a compact vector descriptor that represents the image globally. In the ViT architecture, pre-classification layers output $M+1$ vectors corresponding to $M$ input patches and a class (CLS) embedding.

In our referent pooling approach, CLS, we follow the spirit of BERT Devlin et al. (2018) and ViT models, and view this class embedding as a global image descriptor. In addition, we investigate performance of global pooling methods that are typically used by convolutional metric learning models, including average, maximum and Generalized Mean (GeM) pooling, and apply them to the $M$ output tokens.

We follow the common practice of projecting the descriptor vector into a unit ball after pooling. In the case when the target dimensionality is smaller than that provided by the architecture, we optionally reduce the vector by principal component analysis (PCA) before normalizing it.

3 IRTLL{}_{\text{L}}: Learning the Metric for Image Retrieval

We now consider a metric learning approach for image retrieval, denoted by IRT${}_{\text{L}}$. It is the dominant approach to both category-level and particular object retrieval. In our case we combine it with transformers instead of convolutional neural networks. We adopt the contrastive loss with cross-batch memory by Wang et al. (2020) and fix the margin $\beta$ = $0.5$ by default for our metric learning objective.

The contrastive loss maximizes the similarity between encoded low-dimensional representations $z_{i}$ of samples with the same label $y$ (or any other pre-defined similarity rule). Simultaneously, it minimizes the similarity between representations of samples with unmatched labels which are referred to as negatives. For the contrastive loss, only negative pairs with a similarity higher than a constant margin $\beta$ contribute to the loss. This prevents the training signal from being overwhelmed by easy negatives. Formally, the contrastive loss over a batch of size $N$ is defined as:

The representations $z_{i}$ are assumed to be $l_{2}$-normalized, therefore the inner product is equivalent to cosine similarity.

4 IRTRR{}_{\text{R}}: Differential Entropy Regularization

Recently, Boudiaf et al. (2020) studied connections between a group of pairwise losses and maximization of mutual information between learned representations $Z=\{z_{i}\}$ and corresponding ground-truth labels $Y=\{y_{i}\}$. We are interested in the particular case of the contrastive loss. The mutual information is defined as

The positive term of the contrastive loss leads to minimization of the conditional differential entropy $\mathcal{H}(Z|Y)$, where intuitively, samples representations belonging to the same category are trained to be more similar:

On the other hand, the negative term of this loss is responsible for preventing trivial solutions where all sample representations are collapsed to a single point. Therefore, it maximizes the entropy of the learned representations:

The margin $\beta$ plays an important role in the training dynamics. Low values of $\beta$ allow exploration of a larger number of negative samples. Yet in this case easy negatives can dominate the training and cause the performance to plateau. In contrast, higher values of $\beta$ would only accept hard negatives, possibly leading to noisy gradients and unstable training Wu et al. (2017).

Motivated by the entropy maximization view of the negative contrastive term in Equation 4, we add an entropy maximization term that is independent of the negative samples accepted by the margin. In particular, we use the differential entropy loss proposed by Sablayrolles et al. (2019). It is based on the Kozachenko & Leonenko (1987) differential entropy estimator:

Intuitively, the different entropy regularization prevents the representations of different samples from lying too close on the hypersphere, by increasing their distance from positive examples, and the hard negatives as well. Having hard negatives with extremely small distances is a main source of noise in the training signal, as identified by Wu et al. (2017).

5 Analysis

We study the behaviour of the output representation space when training with a contrastive loss, and how augmenting this loss with a differential entropy regularization impacts the space properties and the model performance.

In Figure 4 we examine the cumulative energy of the principle components for features from an off-the-shelf, ImageNet pre-trained model, as well as models trained using contrastive loss. We observe that the features after training with the contrastive loss suffer from a collapse in dimensions compared to an untrained model. This suggests an ineffective use of the representational capacity of the embedding space, as alignment is favored over uniformity while both are necessary for good representations Wang & Isola (2020). As we augment the contrastive loss with the differential entropy regularization, the cumulative energy spreads across more dimensions (see Figure 4 with non-zero values of $\lambda$). Higher values of $\lambda$ alleviate the dimensionality collapse problem.

Another observation is that the transformer-based architecture is less impacted than convnets by the collapse (see Figure 4). Despite having a lower extrinsic dimensionality compared to the ResNet-50 model, the DeiT-Small features are more spread over principle components. A possible reason for that behavior is that in multi-headed attention, each input feature is projected to different sub-spaces before the attention operation, reducing the risk of collapse.

Gradient Analysis.

As pointed out by Wu et al. (2017), very hard negatives can lead to noisy gradients. We examine the nuclear norm $\|\cdot\|_{*}$ associated with the covariance matrix of the gradients directions $\gamma=\|\text{Cov}(\nabla_{z}\mathcal{L_{\text{contr.}}})\|_{*}$, averaged over all training iterations (see Figure 4). Higher values of $\gamma$ could indicate noisy gradients. We observe them for both very high and very low values of margin $\beta$ which aligns with our understanding that very easy and very hard negatives lead to less informative and less stable training signal. Moreover, we observe a decrease in the $\gamma$ values after the addition of the entropy regularization term.

Experiments & Ablation Studies

We first describe datasets and implementation details, and then proceed with discussions of empirical results.

We report performance on three popular datasets commonly used for category-level retrieval. Stanford Online Products (SOP) Oh Song et al. (2016) consists of online products images representing 22,634 categories. Following the split proposed by Oh Song et al. (2016), we use first 11,318 categories for training and the remaining 11,316 for testing. CUB-200-2011 Wah et al. (2011) contains 11,788 images corresponding to 200 bird categories. Following Wah et al. (2011), we split this dataset into two class-disjoint sets, each with 100 categories for training and testing. In-Shop Liu et al. (2016) contains 72,712 images of clothing items belonging to 7,986 categories, 3,997 of which used for training. The remaining 3,985 categories are split into 14,218 query and 12,612 gallery images for testing. We compute the Recall@K evaluation metric for a direct comparison with previous methods.

Particular Object Retrieval.

For training, we use the SFM120k dataset Radenović et al. (2018b) which is obtained by applying structure-from-motion and 3D reconstruction to large unlabelled image collections Schonberger et al. (2015). The positive images are selected such that enough 3D points are co-observed with the query image, while negative images come from different 3D models. We use 551 3D models for training and 162 for validation.

For evaluation, we report results using revisited benchmarks Radenović et al. (2018a) of the Oxford and Paris Philbin et al. (2007; 2008) datasets. These two datasets each contain 70 query images depicting buildings, and additionally include 4993 and 6322 images respectively in which the same query buildings may appear. The revisited benchmarks contain 3 splits: Easy (E), Medium (M) and Hard (H), grouped by gradual difficulty of query/database pairs. (E) ignores hard queries, (M) includes both easy and hard ones, while (H) considers hard queries only. We report the Mean Average Precision (mAP) for the Medium and Hard splits in all our experiments.

2 Implementation & Training Details

The transformer-based models and their pre-trained weights are based on the public implementationhttps://github.com/facebookresearch/deit of DeiT Touvron et al. (2020) built upon the Timm library by Wightman (2019). All models are optimized using the AdamW optimizer Loshchilov & Hutter (2017) with learning rate $3.10^{-5}$, weight decay $5.10^{-4}$ and batch size of 64. For all experiments, unless mentioned otherwise, the contrastive loss margin is set to $\beta=0.5$ and the entropy regularization strength is set to $\lambda=0.7$. We show later in ablation that the results are relatively stable (and not overfitted) to this hyper-parameter setting. We use standard data augmentation methods of resizing the image to 256$\times$256 and then taking a random crop of size 224$\times$224, combined with random horizontal flipping. Following Wang et al. (2020), we use a dynamic offline memory queue of the same size as the dataset (with the exception of In-Shop dataset for which the memory size is 0.2 of the dataset size). Additionally, for the In-Shop dataset we adopt a momentum encoder for the memory entries (similarly to He et al. (2020)) with momentum value of 0.999. We have found this was not necessary for SOP and CUB-200-2011. Finally, SOP and In-Shop models were trained for 35k iterations and the CUB-200-2011 model was trained for 2000 iterations.

Particular Object Retrieval.

For the particular object retrieval experiments, we build our implementation on top of the public codehttps://github.com/filipradenovic/cnnimageretrieval-pytorch associated with the work by Radenović et al. (2018b). We follow the same optimization and regularization procedure. All models, transformer-based and convolutional, are finetuned using the SFM120k dataset. The input images are resized to have the same fixed width and height. We report results for image sizes of 224$\times$224 and 384$\times$384. For finetuning, each batch consists of 5 tuples of (1 anchor, 1 positive, 5 negatives). For each epoch, we randomly select 2,000 positive pairs and 22,000 negative candidates (using hard-negative mining). We use the default hyper-parameters of Radenović et al. (2018b): the models are optimized using Adam Kingma & Ba (2015) with small learning rate of $5.10^{-7}$ and weight decay of $10^{-6}$. The contrastive loss margin is set to $\beta=0.85$. The models are finetuned for 100 epochs. All models with GeM pooling use a pooling exponent value of $p=3$. The dimensionality reduction is achieved using a PCA trained on the SFM120k dataset. For the evaluation, all the query and database images are resized into a square image with the same resolution as used during the finetuning stage.

3 Results

We present the Recall@K performance for three public benchmarks for category-level retrieval. For the SOP dataset, we can see in Table 2 that our IRT${}_{\text{R}}$ model with DeiT-S384 backbone achieves state-of-the-art performance for all values of K, outperforming previous methods by a margin of 2.6% absolute points for Recall@1. The DeiT-S$\dagger$ variant with distillation pre-training achieves the best results on this benchmark. Even when reducing the dimensionality to 128-D, our method outperforms all the convnets except at Recall@1000. On the CUB-200-2011 dataset, the DeiT-S384 model outperforms the current state of the art by 2.5% points at Recall@1. The distilled DeiT-S model provides an additional 1.9% improvement, achieving the best results for all values of K. The DeiT-S128 variant with compressed representation outperforms all previous methods except for the ProxyNCA++ model that uses 2048-D descriptors. Similarly, for In-Shop, the DeiT-S384 model and its distilled variant outperform all previous models at Recall@1 with a margin of 0.2% and 0.6% respectively.

Particular Object Retrieval.

We present the mAP performance for the Medium and Hard splits of the revisited Oxford and Paris benchmarks in Table 7. First observe that for input images with size 224$\times$224, the DeiT-S$\dagger$ backbone outperforms its ResNet-50 counterpart with the same capacity, as well as the higher capacity ResNet-101 across all benchmarks and descriptor sizes. The larger DeiT-B$\dagger$ provides a significant gain in performance and achieves the best result among the reported models. Scaling up the image size to 384$\times$384 considerably improves the performance for all models with the DeiT-B$\dagger$ model retaining its position as the strongest model. In Table 9 we compare our model to strong state-of-the-art methods in particular object retrieval, following the standard extensive evaluation procedure. Revaud et al. (2019) use the original resolution of the dataset (i.e. 1024$\times$768), while Radenović et al. (2018a) utilizes multi-scale evaluation. Although these methods outperform our DeiT-B$\dagger$ model at resolution $384\times 384$ in mAP, they are approximately 248% and 437% more expensive w.r.t. FLOPS. Furthermore, we observe that for compressed representations of 128-D, our model closes the gap with Radenović et al. (2018a), achieving a higher mAP for $\mathcal{R}$Paris.

4 Ablations

We provide a comparison between different degrees of supervision corresponding to IRT${}_{\text{O}}$, IRT${}_{\text{L}}$ and IRT${}_{\text{R}}$ in Table 3. We observe that finetuning substantially improves performance over off-the-shelf features, especially for category-level retrieval. Augmenting the contrastive loss with differential entropy regularization further improves the performance across all benchmarks. Figure 5 demonstrates how the distribution of the cosine similarities between positive and negative pairs is impacted by the different variants we study. We notice that finetuning strongly helps to make the positive and negative distributions more separable. The entropy regularization term spreads the similarity values across a wider range.

Choice of Feature Extractor: Pooling Methods.

In Table 9, we study different feature aggregation methods, as described in Section 3.2. Both for category-level and particular object retrieval, we observe that utilizing the CLS token as the image-level descriptor provides the strongest performance (or at least on par) compared to other popular pooling methods such as average pooling, max pooling and GeM. This suggests that the transformer operates as a learned aggregation operator, thereby reducing the need for careful design of feature aggregation methods.

Performance across Objective Functions.

The choice of the objective function used to train image descriptors is crucially important and is the focus of the majority of the metric learning research. While we adopt the contrastive loss as our primary objective function, we additionally investigate two objective functions with different properties: (1) Normalized Softmax Zhai & Wu (2018) as a classification-based objective, and (2) Scalable Neighborhood Component Analysis (NCA) Wu et al. (2018), a pairwise objective with implicit weighting of hard negatives through temperature. Table 6 shows that DeiT-S outperforms its convolutional counterpart across all different choices of objective functions. This suggests that transformer-based models are strong metric learners and hence an attractive alternative to convolutional models for image retrieval.

Regularizing Hyper-parameter λ𝜆\lambda.

We explore the differential entropy regularization strength and its impact on the improvement of retrieval performance. First, we use the SOP dataset for our analysis and show how the Recall@1 performance changes with different margin values $\beta$ and entropy regularization strengths $\lambda$ in Table 6.

Conclusion

In this paper, we have explored how to adapt the transformer architecture to metric learning and image retrieval. In this context, we have revisited the contrastive loss formulation and showed that a regularizer based on a differential entropy loss spreading vectors over the unit hyper-sphere improves the performance for transformer-based models, as well as for convolutional models. As a result, we establish the new state of the art for category-level image retrieval. Finally, we demonstrated that, for comparable settings, transformer-based models are an attractive alternative to convolutional backbones for particular object retrieval, especially with short vector representations. Their performance is competitive against convnets having a much higher complexity.

References