Contextual Similarity Aggregation with Self-attention for Visual Re-ranking

1 Introduction . In instance image retrieval , the goal is to efficiently identify images containing the same object or describing the same scene with the query image from a large corpus of images . Towards this goal , many works have emerged in recent years [ 42 , 21 , 36 , 26 , 24 ] . With the development of deep learning , a great number of methods leverage convolutional neural network ( CNN ) as feature extractor [ 38 , 3 , 44 , 4 , 14 , 34 , 25 , 35 , 45 ] to replace or combine with the classic SIFT feature [ 23 ] . Generally , the performance of the initial retrieval results by simple comparison of visual feature may not be satisfactory . To refine it , many visual re-ranking techniques have been proposed [ 9 , 19 , 48 ] . Popular visual re-ranking techniques include query expansion , geometric context verification , kNNbased re-ranking , and diffusion-based methods , etc . Query Expansion ( QE ) [ 9 , 10 , 39 , 34 , 30 ] computes an average feature of the query and the top-ranked images to update the original query * The first and second authors contribute equally to this work . †Corresponding authors : Min Wang and Wengang Zhou Code will be released at https : //github.com/MCC-WH/CSA . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . feature , which is then used for a second-round retrieval . Geometric Context Verification ( GCV ) [ 30 , 12 , 20 , 2 ] leverages the geometric context of local features to remove the false matches in the original results . Besides , if two images belong to the k-reciprocal nearest neighbors ( kNN ) of each other , these two images have a high probability of being relevant . Based on this observation , lots of k-Nearest Neighbors ( kNN ) based re-ranking methods thus emerge and achieve outstanding performance [ 19 , 48 , 50 ] . Diffusion-based methods [ 51 , 11 , 18 ] consider the global similarity of all images and perform similarity propagation iteratively . Most of the above methods do not involve training , thus can be quickly equipped with various features . The recent learning-based methods such as LAttQE [ 16 ] and GSS [ 22 ] , achieve better performance but require training a specific model for each kind of feature . Different from the methods mentioned above , in this paper , we propose a novel visual re-ranking method by contextual information aggregation . The initial retrieval results generated by feature comparison only consider pairwise image similarity , but ignore the rich contextual information contained in the ranking list . As illustrated in [ 29 ] , if two images are mutually relevant , they share similar distances from a set of anchor images . Thus , we directly select the top-L retrieval results as anchors for each query to fully model the contextual information in the ranking list . We define an affinity feature for each image in top-K candidates by computing the similarity between it and the anchor images . To further promote the re-ranking quality , we propose a new data augmentation method . Moreover , inspired by Query Expansion ( QE ) which employs the information of top-ranked images to update the query feature , we propose to update affinity features of each top-K candidates by aggregating the affinity features of other candidates to promote the re-ranking performance with transformer encoder [ 46 ] . For each image in the top-K candidates , our re-ranking model dynamically aggregates the affinity features of other candidates based on the importance between candidates by computing their similarities using the affinity features . The output of our re-ranking model can be regarded as the refined affinity features for top-K candidates which contain more contextual information . To train the re-ranking model , two loss functions are introduced . Firstly , we use a contrastive loss to restrain the updated affinity features so that the relevant images have large similarity , and vice versa . Besides , we exploit a Mean Squared Error ( MSE ) loss to reserve the information in original affinity features by restricting the difference between the original affinity features and the refined affinity features . During the inference time , we compute the affinity features for the top-K candidates and utilize the transformer encoder to refine the features . Then we re-rank these candidates by computing the cosine similarity between the refined affinity features . The rank of images outside the top-K ranking list remains unchanged . Note that our re-ranking model is not directly involved with the original visual feature . Instead , it computes the affinity features for top-K images , which serve as the input in our method . Therefore , it can be combined with various existing image retrieval algorithms including representation learning and re-ranking methods , and further enhance the retrieval performance with low computational overhead . We conduct comprehensive experiments on four benchmark datasets to prove the generality and effectiveness of our proposed visual re-ranking method . Besides , the time and memory complexity of our re-ranking method is lower than the state-of-the-art re-ranking methods . 2 Related Work . In this section , we review the related works on visual re-ranking including query expansion , geometric context verification , kNN based methods , and diffusion-based methods . Query Expansion . Query Expansion ( QE ) uses the information of the top-ranked images obtained from the initial retrieval to refine the representation of the original query for further retrieval . Average query expansion ( AQE ) [ 9 ] aggregates the features of top-K retrieval results by average pooling . Average query expansion with decay ( AQEwD ) [ 15 ] proposes a weighted average aggregation , where the weights decay over the rank of retrieved images . Discriminative query expansion ( DQE ) [ 1 ] regards the top-ranked images as positive images , and regards the bottom-ranked images as negative images to train a linear SVM . Then DQE sorts images according to their signed distance from the decision boundary . In [ 34 ] , α-weighted query expansion ( αQE ) weights the top-K images by their Geometric Context Verification . Geometric Context Verification uses the geometric context of local features to remove the false matches . Some approaches [ 9 , 30 , 37 ] estimate the transformation model to verify the local correspondences by RANSAC-based methods . RANSAC [ 12 ] generates hypotheses on random sets of correspondences and then identifies a geometric model with maximum inliers . Spatial coding [ 52 ] proposes to represent the relative spatial locations among local features into binary maps for efficient local matching verification . DSM [ 40 ] leverages the MSER regions which are detected on the activation maps of the convolutional layer . Then DSM regards these regions as local features for spatial verification . k-NN-based Re-ranking . The k-reciprocal nearest neighbors of an image are considered as highly relevant candidates [ 49 , 32 , 39 ] . Contextual Dissimilarity Measure ( CDM ) [ 19 ] iteratively regularizes the average distance of each point to its neighbors to update the similarity matrix . Visual Rank [ 48 ] determines the centrality of nodes on a similarity graph according to the link structures among images to rank images . In CRL [ 27 ] , a lightweight CNN model is trained to explore the contextual information and learn the relevance between images . GSS [ 22 ] proposes an unsupervised method for re-ranking based on graph neural network neighbor encoding and pairwise similarity separation loss . In [ 50 ] , the k-reciprocal feature is designed by encoding the k-reciprocal nearest neighbors into a single vector to perform re-ranking . Diffusion-based Re-ranking . Diffusion technique is initially designed for ranking on manifolds [ 51 ] , and it has been successfully applied to many computer vision tasks , such as image classification , object detection , image retrieval , etc . In image retrieval , a lot of research concentrates on propagating similarity through the kNN graph [ 11 ] . Diffusion is usually used as a re-ranking method , and it has achieved state-of-the-art performance on many benchmarks [ 28 , 47 , 5 , 18 , 33 , 8 , 7 , 6 , 13 , 28 ] . Diffusion considers the underlying manifold structure and explores the internal relationships between images . As one of the representative methods , DFS [ 18 ] proposes a regional diffusion mechanism to further improve the recall of small objects based on the neighborhood graphs . Then it employs conjugate gradient method to compute the closed-form solution of diffusion as well as maintaining time efficiency and accuracy . 3 Method . 3.1 Framework Overview . In image retrieval , the database is defined asD = { I1 , I2 , · · · , IN } , where Ii denotes the ith database image and N is the database size . Usually , images are first mapped to high-dimensional vectors through a feature extractor φ ( · ) . We define the feature representation of images as f = φ ( I ) ∈ Rd . Given a query , we obtain an initial ranking list of the top-K images R = [ r1 , r2 , · · · , rK ] by a retrieval algorithm , where ri denotes the ID of the ith image . For convenience of the following discussion , we assume the query image is returned at the first position in the ranking list , otherwise we directly insert it in the front of the ranking list . The features of top-K images in the ranking list are described as FK = [ fr1 , fr2 , · · · , frK ] ∈ Rd ×K . Our framework is shown in Fig . 1 . With the initial retrieval ranking list obtained by a query , we propose to represent each of the top-K candidates into an affinity feature vector by computing the cosine similarity between it and a set of anchor images . Then the affinity features are refined by a transformer encoder to aggregate the contextual information in the ranking list . Affinity Feature . Generally , if two images are relevant to the query , they shall be relevant to each other [ 43 ] . Moreover , as revealed in [ 29 ] , if two images are mutually relevant , their distances to a set of pre-defined anchor images shall be similar . Actually , it is a non-trivial issue to select a proper set of images as anchors . Generally , if the selected anchor images are far from the top-K candidates in the ranking list , the variance of its distance with candidates is relatively small , which is difficult to provide useful information to distinguish the candidates . In other words , a better alternative is to choose the top-L images in ranking list as the anchor images . With anchor images , we define the affinity features for the ith image in the rank list R as follows : ai = [ f T rifr1 , f T rifr2 , · · · , f T rifrL ] T , 1 ≤ i ≤ K. ( 1 ) After that , we take the affinity features of the K images in R as a sequence , which is fed into our re-ranking model for refinement . Note that the input of our proposed method is affinity feature sequence and is not directly related to the original feature , which endows our method with good generality . Contextual Similarity Aggregation with Transformer . Although the affinity features have taken into account the contextual information in the ranking list , there may be inconsistency in the initial affinity features obtained by directly comparing with anchor images . In order to make the affinity features of related images more consistent , we propose a new re-ranking model by contextual information aggregation to dynamically refine the affinity features of different images . Specifically , with transformer encoder as the core component of our re-ranking model , we first map the affinity feature ai ∈ RL for the ith image to a′ ∈ RL ′ using a learnable projection matrix Wp ∈ RL×L ′ , and then feed the sequence a′i , i = 1 , 2 , · · · , K into a transformer encoder consisting of Multi-Head Attention ( MHA ) and Feed-Forward Networks ( FFNs ) . In MHA , each sequence element is updated by a weighted sum of all other elements based on the scaled dot-product similarity . The scaled dot-product attention first maps affinity features a′i ( i = 1 , 2 , · · · , K ) to Queries ( Q ∈ RK×ds ) , Keys ( K ∈ RK×ds ) and Values ( V ∈ RK×ds ) with three learnable projection matrices . After that , the similarity between Q and K is aggregated with V together , Attention ( Q , K , V ) = Softmax ( QKT√ ds ) V , ( 2 ) which can be seen as using the similarity between images as weights to aggregate affinity features . The multi-head structure concatenates and fuses the outputs of multiple scaled dot-product attention modules using the learnable projection WM . MHA is defined as follows : MHA ( Q , K , V ) = Concatenation ( h1 , h2 , · · · , hNH ) WM , ( 3 ) where hi = Attention ( Qi , Ki , Vi ) , i = 1 , 2 , · · · , NH , ( 4 ) andNH is the head number . Note that to update the affinity feature of an image , the affinity features of all the relevant candidates should be considered as contextual information . A single-head self-attention layer limits the ability of the model to focus on one or more relevant candidates simultaneously without affecting other equally important candidates . This can be achieved by projecting the original affinity features into different representation subspaces . Specifically , in MHA , different projection matrices for Q , K , and V are used for different heads , and these matrices can project affinity features into different subspaces . The output is normalized via Layer Normalization ( LN ) and finally added to S′ to form a residue connection , S′ = S′ + LN ( MHA ( Q , K , V ) ) . ( 5 ) In addition to MHA , the transformer encoder layer has a position-wise feed-forward network . We stack n transformer encoder layers to obtain more consistent affinity features . The refined affinity features yi , i = 1 , 2 , · · · , K come from the output of the last layer of our re-ranking model . Data Augmentation . According to the description above , we represent the top-K images into affinity features , which serve as the input to our model . To improve the robustness of the re-ranking model and avoid overfitting , it is necessary to increase the size of the training sets . To this end , we design a data augmentation scheme oriented to our framework . Considering that our training data comes from retrieval results , we propose to use different visual features to obtain the ranking lists of the same query . With different types of features to create the affinity features , the network will learn more information about the contextual information . For a query , we employ M different visual feature extractors { φ ( · ) 1 , φ ( · ) 2 , · · · , φ ( · ) M } to obtain different ranking lists , and compute their corresponding affinity features according to the previous section . The obtained M affinity feature sequences for the same query are used as training samples . In this way , we enlarge the training set size to M times of the original size and improve the robustness of our model .

Summary: The paper proposes a re-ranking algorithm for image retrieval, capable of improving the ranking results of an off-the-shelf similarity-based retrieval system, by representing each query by an affinity vector that measures similarities to a set of anchor images, and then refine this representation by aggregating it with contextual information obtained from other top-ranked retrieval results using a transformer architecture. The refined representations of the top-K images can then be used for re-ranking.   The paper shows improvements both in precision and time complexity over previous re-ranking algorithms. The time complexity improvements are mostly due to the introduction of a fixed set of anchor images, which avoids context affinity features that scale linearly with the number of top-k retrievals considered. The precision improvements are a result of a learnable context aggregation module (in the form of a transformer) and the data augmentations used for training it.  Experiments on several benchmarks show healthy improvements over prior re-ranking procedures.