Image Retrieval with Self-Supervised Divergence Minimization and Cross-Attention Classification

1 Introduction

Representing similarity between entities is critical to many tasks such as content-based image retrieval, cross-modal information retrieval, face recognition, and person reidentification. A key driver behind these methods is metric learning where the task is to learn embeddings such that their distances preserve a notion of semantic similarity. There are several variations on metric learning losses, which include ideas such as sampling [Yu et al., 2019; Robinson et al., 2021] and different schemes for grouping examples together such as pair-based [Khosla et al., 2021; Hoffer and Ailon, 2018] and proxy-based [Kim et al., 2020a; MovshovitzAttias et al., 2017b]. The overall goal behind metric learning methods is to train an embedding model to pull similar objects together in the embedding space while (optionally) pushing dissimilar objects apart. Similarity is often given by a class label where all examples in a dataset of class c are pulled together and examples of other classes are pushed apart. Other options for determining similarity include selfsupervised learning [Caron et al., 2021b; Caron et al., 2019; Caron et al., 2021a] where an input image is transformed into a pair or set of views via data augmentation. Here, all views of the same image are "similar." In the metric learning for image retrieval setting, the encoder often takes one of two forms: a convolutional neural network (CNN) [He et al., 2015; Szegedy et al., 2014] or more recently a Vision Transformer [Dosovitskiy et al., 2021; Touvron et al., 2021]. The goal of the encoder is to transform each input image into a learned embedding optimized for the downstream task of retrieval from an external set, usually called a database. An important point to note is that the training setting is decoupled from the retrieval setting in that there is no database present during training so queries are not conditioned over a database. Instead there is an auxiliary task such as a metric learning objective which is then used for retrieval during evaluation. Here we propose an approach that instead jointly learns query and database representations.

We offer a novel framework for training image retrieval systems which we call DMCAC (Divergence Minimization with Cross-Attention Retrieval). It offers a few novel insights. Firstly, we use a query set and database set directly during training. We generate and embed several views of an input query via data augmentation as is common in the selfsupervised regime. But rather than minimizing the distance between each query view, we first compute a similarity distribution between the query views and a separate database set prodcing a similarity distribution for each view over the database. Our method is motivated by distribution matching which we empirically find to be stronger than point-wise objectives like triplet loss. For our loss, we minimize the divergence in this similarity distribution between views so that our loss is conditioned from a query set over a database set, rather than over a query set alone. This formulation closely resembles the retrieval test setting where a query set is used to retrieve similar sets of items from the database.

Importantly, we note that the self-supervised method is prone to collapse [Caron et al., 2021a] where the encoder "cheats" by simply embedding all views to the same point because there is no grounding signal of "correctness" such as a class label. We overcome this issue by using the retrieved database embeddings to also directly predict the query class label via cross-attention. In this way, the retrieved examples must learn some notion of semantic meaning about the query. Our key contributions are:

• We present a novel retrieval framework that applies

self-supervised learning to directly learn a joint representation between a query set and a database set. This approach offers a competitive and often improved alternative to common metric learning retrieval frameworks.

• We introduce a novel loss for retrieval. We minimize divergence between query views via a Frobenius norm and classification error via cross-attention with crossentropy using query images conditioned over a database set.
• We investigate the tradeoffs between end-to-end differentiability and efficient approximate nearest neighbor (ANN) search alongside a thorough ablation of hyperparameters. We show that our method is competitive across both settings.

2 RelatedWorks

2.1 Transformers and Vision Transformers

Transformers [Vaswani et al., 2017] revolutionized natural language processing by using self-attention and blocks of fully-connected layers. Models such as BERT [Devlin et al., 2019] were highly performant across a broad spectrum of NLP tasks such as text classification and language translation. Transformers excel in capturing local dependencies and using them to produce global representations which made them an attractive candidate for computer vision. The key breakthrough came with Vision Transformers (ViTs) [Dosovitskiy et al., 2021] which had the insight that local tokens in an image are simply patches which can then be fed to the Transformer architecture. ViTs match or surpass convolutional models across tasks such as image classification [Caron et al., 2021b; Chen et al., 2021] and semantic segmentation [Strudel et al., 2021]. Because ViTs lack the traditional inductive biases of CNNs, methods arose to efficiently train ViTs [Steiner et al., 2022]. These methods include using data augmentation [Cubuk et al., 2019b], knowledge distillation [Touvron et al., 2021], pretraining with larger datasets [Koppula et al., 2022], and specialized architectural variants such as Swin Transformer [Liu et al., 2021] which reintroduce vision specific biases.

2.2 Metric Learning

There are many approaches to Metric Learning, but here we give an overview of pairwise and classification approaches. Contrastive loss [Carreira-Perpi˜n´an and Hinton, 2005; Hadsell et al., 2006] and triplet loss [Hoffer and Ailon, 2018] are examples of pairwise losses where the goal is to bring similar examples together while pushing dissimilar examples away in embedding space by generating groups of examples. Other variations such as Supervised Contrastive Loss [Khosla et al., 2021] introduce class information that helps create positive and negative pairs conditioned on a semantic label. These methods can be further improved by employing techniques such as Hard Negative Sampling [Robinson et al., 2021] and other sampling methods [Yu et al., 2019] to ensure that the network avoids collapse from an oversaturated signal of easy examples during optimization.

Methods such as [El-Nouby et al., 2021a; Song et al., 2023] show the effectiveness of combining contrastive methods with Vision Transformers for retrieval.

Classification based approaches represent categories with a single or multiple representative points [Movshovitz-Attias et al., 2017a; Teh et al., 2020a; Boudiaf et al., 2021]. In these methods, a new example is compared against the representative points, and a class is determined based on similarity to each class representative [Zhai and Wu, 2019a; Trivedy and Latecki, 2023]. All of these methods show that direct classification is a viable and performant technique for retrieval. We draw inspiration from these approaches in designing our novel classification approach which uses crossattention over a database set.

2.3 Self-Supervised Learning

Self-supervised learning (SSL) constructs its learning signal from the data itself where a model learns to predict a portion of the input data from other parts, compelling it to learn robust, high-level features. The key benefit here is that this requires no labeled data and is thus a good approach for large-scale pretraining [He et al., 2021; Baevski et al., 2022]. Common SSL approaches include solving pretext tasks, masking, and learning invariances to data augmentations. SSL pretext tasks include [Doersch et al., 2016] which predicts the relative position of image patches, and [Noroozi and Favaro, 2017a; Noroozi and Favaro, 2017b] which use jigsaw puzzles where a model learns to recognize the permutation of shuffled image patches. Masking methods such as [He et al., 2021] are trained to reconstruct masked portions of an image. In this work, we focus on augmentation based SSL methods such as SimCLR [Chen et al., 2020] which generate multiple views of an image via data augmentation and minimize the embedding distance between them. These methods are often paired with contrastive learning where positive pairs are two augmented views of a single image and negative pairs are views from two different images.

Self-supervised learning has benefited from work on selecting performant data augmentation strategies [Cubuk et al., 2019b; Cubuk et al., 2019a]. Reducing the embedding distance between variations of the same image (via cropping, color jittering, etc.), teaches a model to be invariant to these changes while hopefully preserving the semantic meaning of the image which leads to more robust and generalizable features [Park et al., 2023].

3 Method

In this section, we outline the details of our approach. In particular, we describe the building blocks for training our model, define our dataset contruction, highlight key design choices in our forward pass, and discuss the retrieval mechanism and loss function.

3.1 Dataset Construction

To motivate our approach we share the rationale and process for constructing the training and testing sets, following common practices of usual data splits. Because we want to directly train our model on the task of retrieval, we construct four sets which we describe here. Please see Section 4.1 for detailed information on each dataset.

• Training Query Set: DQ
• Training Database Set: DD
• Evaluation Query Set: D′Q
• Evaluation Database Set: D′D

Note that the training and evaluation sets have a disjoint set of classes. In this way, during training, we learn general representations that can be extended easily to new classes which is a common practice in retrieval evaluation.

3.2 Revisiting Vision Transformers

In our approach, we leverage Vision Transformers (ViTs). Here we provide an overview of their building blocks. A ViT operates over sequences like the original transformer [Vaswani et al., 2017] architecture. The input sequence can be described by

Here each xi represents a patch token with dimension F of which there are M, and the global representation for the image is captured by xcls. The most common approach for splitting an image into a set of patches is to create nonoverlapping patches of shape P × P to produce a sequence of patch tensors of shape P × P × C where C is the number of channels. Each patch is flattened to a vector of size P 2C and embedded to dimension F via a shared linear layer. This is equivalent to applying a square convolutional filter with kernel size P and stride P .

After generating the sequence of patches, learnable positional encodings are added so that the model is aware of spatial information instead of being a pure set-to-set model. The ViT then performs K blocks of computation consisting of multi-head self-attention (MHSA), layer normalization [Ba et al., 2016], feed-forward layers (FFN), and a skip connection. The input sequence length and dimensionality are preserved in each block. The output we use is an embedding zcls which captures the global representation of the input image.

3.3 Forward Pass

Here we present the data flow through our model and our augmentation setup. We define our training objective, namely to learn an updated encoder function ϕ via a base encoder, training query set, and training database set.

where initially ϕnew = ϕbase. In this case, our base encoder model, ϕbase is a flavor of ViT that comes initialized with pretrained weights. We first embed each example in our training database set DD ∈ RD×C×H×W via our base encoder, where D is the number of the training database images.

Here we produce ZD from our database where each of the D rows is an image embedding with output dimension F. This is achieved by outputting the CLS token for each image from the penultimate layer of the ViT and concatenating each of the tokens row-wise.

After constructing the database embeddings, we begin training the base encoder. We present our approach for a single query, but the method extends naturally to the minibatch setting. From our training query set, DQ, we select an image and perform data augmentation to generate A views of the image (one of them is the original image) and produce XA ∈ RA×C×H×W , where H and W are image dimensions and C is the number of channels. We then embed each image view with the encoder

where ZA is the concatenated set of CLS tokens for each of the views in the minibatch. As is common in other methods [Caron et al., 2021a; El-Nouby et al., 2021a], we project the output embedding for each view onto the unit ball for each view xi via ϕnew(xi)/∥ϕnew(xi)∥2. We compute the same projection for the database set so a dot product similarity computed between the embedded query views and database set is equivalent to cosine similarity.

3.4 Training Retrieval Mechanism

We begin with retrieval over all database image embeddings, ZD, using our query embeddings, ZA. Hence, the full retrieval computes the cosine similarity exhaustively over all database embeddings ZD

Full Retrieval Setting. We compute a similarity matrix P as a product between the embeddings of the query views, zA, and the transpose of ZD:

where σ is the softmax operator, so each row of P ′ can be interpreted as a probability distribution. We use P ′ to compute our loss as described in Section 3.5.

Note that in this approach we receive a full gradient signal from the database to the encoder. This setting acts as an implicit negative sampling where the encoder receives signals from both similar and dissimilar embeddings.

However, this simple setting only makes sense if all embeddings of the database images fit into GPU memory, which is an unrealistic assumption for larger databases. To address this limitation, we consider an approximate retrieval setting, where first retrieve embeddings of top-k images, and use them to construct the loss function defined in Section 3.5. These two methods differ in that approximate nearest neighbor retrieval uses FAISS [Johnson et al., 2019] with cosine similarity while full retrieval computes the cosine similarity exhaustively over all database embeddings ZD.

Approximate Retrieval Setting. Let z ∈ R1×F denotes a deep embedding of query q, and let ZA ∈ RA×F denotes embeddings of the query views, where row i is an embedding of view qi with q1 = q. For each view qi of q, we retrieve database indices of its k nearest neighbors, denoted as Si. From here, we use set operations to generate a combined set of retrieved indices across all views. The goal is to find the full set of indices retrieved by all of the views together:

Now Sunion is a set containing all indices of the nearest neighbors across all A query views. We then retrieve the associated embeddings from the database to produce Z union S ∈ RT×F , where T is the cardinality |Sunion |. Finally, we compute a similarity matrix P as a product between the embeddings of the query views, zA, and the transpose of Z union S :

Discussion. Although the retrieval of top-k images and the set union operation on their indices are not differentiable, the encoder transformer weights can still be updated by the gradient backpropagation from this loss function. In other words, we use the nearest neighbor indices to obtain the associated embeddings from the database, compute the similarities between these selected images, and pass a gradient back to the encoder.

In a sense, we treat the embeddings of top-k database images as constants. However, since they change with each query, the gradient is informed by the dot product computed between the query views and their top-k database images. Since the encoder weights are updated after each minibatch, all query views in the next minibatch, which go through a forward pass, have new embeddings. In contrast, the database images get new embeddings every few epochs (15 epochs in our implementation). See supplementary materials for more details.

The key benefit of approximate retrieval is the ability to scale to much larger database sizes as is common in industrial settings. As our experiments in Section 4.3 demonstrate, the performance in comparison to full retrieval is competetive.

3.5 Loss over Query Views and Database

Here we describe our loss formulation for a given input embedding and its views. There are two main ideas behind our loss design: divergence minimization and crossattention based classification. Both of these losses compute a representation for the query conditioned on the database. In Section 4.3 we present a thorough ablation on the relative importance of each term in our loss function and also provide a discussion detailing the behavior that each term induces in the training process.

Frobenius Loss for Divergence Minimization Instead of directly minimizing the distance between A views of the query, we propose to minimize the divergence between the distributions over all pairwise views conditioned over the database. Our intuition is that different views of the same query should retrieve similar database images. As described in Section 3.4, we interpret a softmax normalized similarities P ′ as distributions, where each row gives a similarity distribution for each view over the retrieved image embeddings from the database. We compute a symmetric pairwise divergence matrix L ∈ RA×A using the JensenShannon (JS) Divergence:

where P and R are two distributions. Since different views of the same query should retrieve similar database images, each Lij ≥ 0 should be as small as possible. Hence, to compute an overall loss, we use the Frobenius Matrix Norm [Shen et al., 2013] of matrix L (using only the lower-triangular portion as the JS Divergence is symmetric):

Intuitively each entry, Lij, gives the distance between view i and view j conditioned on the distances to the retrieved database embeddings. So by minimizing Lfrob we want new embeddings of views i and j to retrieve the same set of knearest neighbors with corresponding similarity values being similar. We reiterate that this loss is computed in a selfsupervised way.

Classification Loss We use two modes of classification loss. The first computes a simple cross-entropy for each view embedding, z. Note that the original query is included as one of the views. Each embedding is passed through a linear layer and softmax layer to produce an output distribution, q over c classes. This is then passed through the cross-entropy loss with labels y.

CAC Loss The second approach is to classify a new representation of the input query embedding, z, which is obtained via cross-attention over the retrieved database embeddings, Z union S ∈ RT×F , where T is the total number of retrieved embeddings. We call this the cross-attention classification (CAC) loss. Intuitively the goal is to determine the class of z by attending on similar examples from the database and this setting closely resembles the category-level retrieval task where an input image of class ci returns other images of the same class from the database. We design our loss as follows

where Wq,Wk,Wv ∈ RF×F and Wc ∈ RF×c for c classes. Note that both Lce and Lcac are computed and summed over all query views z.

Overall Loss Our combined loss is a simple weighted sum over each of the individual losses.

We note that this loss performs well without any careful weighting and by simply allowing equal contribution for each term. In Section 4.3, we show experiments with different configurations for β including βi = 0, which removes the corresponding term.

4 Experiments

We follow a commonly used protocol for training and evaluation [Ermolov et al., 2022; El-Nouby et al., 2021a] and compare our method to state-of-the-art approaches across three datasets. We also share our implementation details and a thorough set of ablations on our architecture, hyperparameters, and loss design.

4.1 Datasets

In training, we use 80-20 stratified split over classes to create the Training Query and Training Database Sets, respectively. Our splits follow the common approaches used in [El-Nouby et al., 2021a; Ermolov et al., 2022].

CUB-200 [Wah et al., 2011] is a fine-grained dataset containing 11,788 images covering 200 sub-category classes of birds. We use the first 100 classes for training and the remaining for testing with no class overlap between train and test settings.

In-Shop Clothes Retrieval [Liu et al., 2016] is a clothing dataset containing 52,712 total images and 7,896 classes. We use the first 3997 classes for training and the remaining classes for testing.

Cars-196 [Krause et al., 2013] is a dataset of car models containing 16,185 total images and 196 classes. We split the dataset into 8,054 images for training (98 classes) and the remaining 98 classes for testing.

Stanford Online Products (SOP) [Song et al., 2016] contains 120,053 images of 22,634 products downloaded from eBay.com. We use the standard split of 59,551 images (11,318 classes) for training and 60,502 images (11,316 classes) for testing.

4.2 Implementation and Evaluation Details

We use two ViT variants as encoders across experiments, ViTS [Dosovitskiy et al., 2021] and DeiT-S [Touvron et al., 2021] following other transformer-based retrieval approaches [Ermolov et al., 2022; El-Nouby et al., 2021a]. Both are based on the ViT-S architectures and are readily comparable across other convolutional approaches on parameter counts (ViT-S 22M, ResNet-50 23M [He et al., 2015]). ViT-S is pretrained on ImageNet-21k [Deng et al., 2009] which contains 14M images and 21K classes while DeiT-S is pretrained using the ImageNet-1k subset containing 1.3M images and 1K classes. DeiT is trained via distillation using a ResNet teacher model, more details can be found in [Touvron et al., 2021]. Both encoders output embeddings of size 384. We use 224 × 224 input resolution across datasets. During training, we resize to 256 on the smaller size and take a random crop of size 224 × 224 while during testing we take a center crop of 224 × 224. We compute Recall@K as our evaluation metric during testing and use the embeddings z as described in Equation 5 projected to the unit ball. We choose RandAugment for query augmentation [Cubuk et al., 2019b] across all experiments. This is chosen over learnable augmentation methods [Cubuk et al., 2019a] for computational efficiency. Unless otherwise stated, we use A = 6 to generate A views of a query (which includes the query as one of the views) during training and use the approximate retrieval approach with k = 12 rather than the full retrieval approach. There is no augmentation applied during testing. Unless otherwise stated, we use βfrob = βce = βcac = 1 as described in Equation 20. We use AdamW [Loshchilov and Hutter, 2019] with learning rate 3 × 10-5 as the optimizer with weight decay 0.01. We use batch size of 256 for all datasets except SOP where we use 128 and number of steps as 250, 600, 2500, and 25000 for CUB, Cars, InShop, and SOP respectively.

4.3 Results

We present results across a variety of settings and compare to other state-of-the-art methods across each of our datasets. In Table 1 we compare our method, DMCAC, against other state-of-the-art approaches that use both convolutional and ViT based backbones. Across the four datasets, the ViT based methods outperform their convolutional counterparts despite using a lower embedding dimension. For our DMCAC-DeiT method, we consistently outperform the other DeiT methods across all four datasets and are pushing the state-of-the-art results even against the ViT (ImageNet-21K) models which are pre-trained on a much larger dataset. Using our DMCACViT method, we achieve several state-of-the-art results across the datasets and are comparable to existing ViT methods [Ermolov et al., 2022]. We perform especially well on CUB200, In-Shop, and Stanford Online Products. We hypothesize that this is because our approach can better distinguish between birds, clothing, and eBay products, which have more varied geometries and color patterns than the examples in Cars-196, though we are competitive there as well. We also make a similar observation to the one in [Ermolov et al., 2022] that there is a large gap in performance between DMCAC-DeiT and DMCAC-ViT on CUB-200, which is likely because ImageNet-21k contains bird classes, which yields better separability for DMCAC-ViT.

In Table 2 we present the effect of tuning our loss weights for each term in Ltotal, namely [βfrob, βce, βcac]. We highlight that our method works competitively out-of-the-box with the default configuration of equally weighting each loss term. This is clearly shown in Figure 3, where as training progresses, query examples learn to attend to the database examples of the same class. We find that with careful tuning of loss weights on the In-Shop dataset, we incrementally increase our already state-of-the-art retrieval performance.

We find that when we completely remove or down-weight Lce relative to βfrob and βcac, retrieval performance stays stable compared to the baseline of equal weighting. However, if we remove either βfrob or βcac as shown in the last two runs, performance drops drastically. This shows that both divergence minimization and cross-attention classification are necessary design choices in order to make our method work, and removing either significantly decrease retrieval performance. We hypothesize that this is because both terms have a synergistic relationship in the following way. As βfrob decreases, the encoder learns both robustness to augmentations and to retrieve similar embeddings from the database across each augmented view. Simultaneously, the cross-attention operation used in βcac re-encodes each input view as a linear combination of the retrieved embeddings which are treated as basis vectors. A correct classification can then reasonably be achieved via a simple linear combination only if the retrieved embeddings capture the semantic meaning of the query. This is the very goal of category level retrieval.

In Table 3 we perform an ablation on the number of views generated per image and how this affects Recall@1 for CUB-200. There is a tradeoff between the computational complexity of generating more views and the benefits of learning invariances over a greater number of augmentations. The effective batch size increases proportionally to the number of generated views which can be limiting depending on available GPU memory. However, we find that retrieval results stayed stable across the DMCAC-ViT settings and degraded for more views when using DMCAC-DeiT. Thus we find that our method is relatively stable to the number of augmentations used and thus is accessible and performs well even in low-memory environments.

In Table 4 we compare our approximate retrieval setting against full retrieval (FR) on CUB-200 and Cars-196. Due to the increase in memory and computational complexity of FR we reduce the relative size of the database compared to the query set during training. The testing setting remains unchanged. For CUB-200 we use a stratified 9010 query-database split during training and a stratified 955 split for Cars-196 during training. We see that the FR outperforms approximate retrieval across both datasets and models in most cases, but the difference is relatively minor. We hypothesize that this is because FR is end-to-end differentiable which means a gradient signal gets passed back to the encoder from all of the database embeddings rather than just from the subset of top-k examples. This acts as a sort of negative sampling where the encoder sees the full database distribution, including the long tail of low-similarity examples for a given query.

5 Conclusions

We offer a new perspective on image retrieval using selfsupervision with a novel loss formulation. We hypothesize that our distribution matching objective learns better query representations than point-wise objectives which can suffer from sampling noise. We instead learn a full mapping from the queries to the database via Lfrob. We also hypothesize that cross-attention classification (CAC) is better aligned to retrieval than cross-entropy classification (CEC). In CAC, each query is projected to a "database space" as an attentionweighted linear combination of the database images and is correctly classified if its "database space" representation accurately captures its class.

Acknowledgements

This work was supported in part by the NSF Grants III2107213 and CNS-2333611.