CATs: Cost Aggregation Transformers for Visual Correspondence

1 Introduction . Establishing dense correspondences across semantically similar images can facilitate many Computer Vision applications , including semantic segmentation [ 46 , 54 , 36 ] , object detection [ 29 ] , and image editing [ 53 , 30 , 28 , 25 ] . Unlike classical dense correspondence problems that consider visually similar images taken under the geometrically constrained settings [ 16 , 19 , 50 , 18 ] , semantic correspondence poses additional challenges from large intra-class appearance and geometric variations caused by the unconstrained settings of given image pair . Recent approaches [ 42 , 43 , 45 , 34 , 37 , 39 , 31 , 58 , 47 , 57 , 51 , 35 ] addressed these challenges by carefully designing deep convolutional neural networks ( CNNs ) -based models analogously to the classical matching pipeline [ 48 , 41 ] , feature extraction , cost aggregation , and flow estimation . Several works [ 24 , 9 , 37 , 39 , 47 , 51 ] focused on the feature extraction stage , as it has been proven that the more powerful feature representation the model learns , the more robust matching is obtained [ 24 , 9 , 51 ] . However , solely relying on the matching similarity between features without any prior often suffers ∗Equal contribution †Corresponding author 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) from the challenges due to ambiguities generated by repetitive patterns or background clutters [ 42 , 24 , 26 ] . On the other hand , some methods [ 42 , 49 , 43 , 23 , 26 , 58 ] focused on flow estimation stage either by designing additional CNN as an ad-hoc regressor that predicts the parameters of a single global transformation [ 42 , 43 ] , finding confident matches from correlation maps [ 20 , 26 ] , or directly feeding the correlation maps into the decoder to infer dense correspondences [ 58 ] . However , these methods highly rely on the quality of the initial correlation maps . The latest methods [ 45 , 37 , 44 , 21 , 31 , 27 , 35 ] have focused on the second stage , highlighting the importance of cost aggregation . Since the quality of correlation maps is of prime importance , they proposed to refine the matching scores by formulating the task as optimal transport problem [ 47 , 31 ] , re-weighting matching scores by Hough space voting for geometric consistency [ 37 , 39 ] , or utilizing high-dimensional 4D or 6D convolutions to find locally consistent matches [ 45 , 44 , 27 , 35 ] . Although formulated variously , these methods either use hand-crafted techniques that are neither learnable nor robust to severe deformations , or inherit the limitation of CNNs , e.g. , limited receptive fields , failing to discriminate incorrect matches that are locally consistent . In this work , we focus on the cost aggregation stage , and propose a novel cost aggregation network to tackle aforementioned issues . Our network , called Cost Aggregation with Transformers ( CATs ) , is based on Transformer [ 61 , 10 ] , which is renowned for its global receptive field . By considering all the matching scores computed between features of input images globally , our aggregation networks explore global consensus and thus refine the ambiguous or noisy matching scores effectively . Specifically , based on the observation that desired correspondence should be aligned at discontinuities with appearance of images , we concatenate an appearance embedding with the correlation map , which helps to disambiguate the correlation map within the Transformer . To benefit from hierarchical feature representations , following [ 26 , 39 , 58 ] , we use a stack of correlation maps constructed from multilevel features , and propose to effectively aggregate the scores across the multi-level correlation maps . Furthermore , we consider bidirectional nature of correlation map , and leverage the correlation map from both directions , obtaining reciprocal scores by swapping the pair of dimensions of correlation map in order to allow global consensus in both perspective . In addition to all these combined , we provide residual connections around aggregation networks in order to ease the learning process . We demonstrate our method on several benchmarks [ 38 , 11 , 12 ] . Experimental results on various benchmarks prove the effectiveness of the proposed model over the latest methods for semantic correspondence . We also provide an extensive ablation study to validate and analyze components in CATs . 2 Related Work . Semantic Correspondence . Methods for semantic correspondence generally follow the classical matching pipeline [ 48 , 41 ] , including feature extraction , cost aggregation , and flow estimation . Most early efforts [ 7 , 30 , 11 ] leveraged the hand-crafted features which are inherently limited in capturing high-level semantics . Though using deep CNN-based features [ 5 , 24 , 42 , 43 , 23 , 49 , 26 ] has become increasingly popular thanks to their invariance to deformations , without a means to refine the matching scores independently computed between the features , the performance would be rather limited . To alleviate this , several methods focused on flow estimation stage . Rocco et al . [ 42 , 43 ] proposed an end-to-end network to predict global transformation parameters from the matching scores , and their success inspired many variants [ 49 , 23 , 25 ] . RTNs [ 23 ] obtain semantic correspondences through an iterative process of estimating spatial transformations . DGC-Net [ 34 ] , Semantic-GLU-Net [ 58 ] and DMP [ 15 ] utilize a CNN-based decoder to directly find correspondence fields . PDC-Net [ 59 ] proposed a flexible probabilistic model that jointly learns the flow estimation and its uncertainty . Arguably , directly regressing correspondences from the initial matching scores highly relies on the quality of them . Recent numerous methods [ 45 , 37 , 39 , 31 , 47 , 51 , 35 ] thus have focused on cost aggregation stage to refine the initial matching scores . Among hand-crafted methods , SCOT [ 31 ] formulates semantic correspondence as an optimal transport problem and attempts to solve two issues , namely many to one matching and background matching . HPF [ 37 ] first computes appearance matching confidence using hyperpixel features and then uses Regularized Hough Matching ( RHM ) algorithm for cost aggregation to enforce geometric consistency . DHPF [ 39 ] , that replaces feature selection algorithm of HPF [ 37 ] with trainable networks , also uses RHM . However , these hand-crafted techniques for refining the matching scores are neither learnable nor robust to severe deformations . As learningbased approaches , NC-Net [ 45 ] utilizes 4D convolution to achieve local neighborhood consensus by finding locally consistent matches , and its variants [ 44 , 27 ] proposed more efficient methods . GOCor [ 57 ] proposed aggregation module that directly improves the correlation maps . GSF [ 21 ] formulated pruning module to suppress false positives of correspondences in order to refine the initial correlation maps . CHM [ 35 ] goes one step further , proposing a learnable geometric matching algorithm which utilizes 6D convolution . However , they are all limited in the sense that they inherit limitation of CNN-based architectures , which is local receptive fields . Transformers in Vision . Transformer [ 61 ] , the de facto standard for Natural Language Processing ( NLP ) tasks , has recently imposed significant impact on various tasks in Computer Vision fields such as image classification [ 10 , 55 ] , object detection [ 3 , 62 ] , tracking and matching [ 52 , 51 ] . ViT [ 10 ] , the first work to propose an end-to-end Transformer-based architecture for the image classification task , successfully extended the receptive field , owing to its self-attention nature that can capture global relationship between features . For visual correspondence , LoFTR [ 51 ] uses cross and self-attention module to refine the feature maps conditioned on both input images , and formulate the hand-crafted aggregation layer with dual-softmax [ 45 , 60 ] and optimal transport [ 47 ] to infer correspondences . COTR [ 22 ] takes coordinates as an input and addresses dense correspondence task without the use of correlation map . Unlike these , for the first time , we propose a Transformer-based cost aggregation module . 3 Methodology . 3.1 Motivation and Overview . Let us denote a pair of images , i.e. , source and target , as Is and It , which represent semantically similar images , and features extracted from Is and It as Ds and Dt , respectively . Here , our goal is to establish a dense correspondence field F ( i ) between two images that is defined for each pixel i , which warps It towards Is . Estimating the correspondence with sole reliance on matching similarities betweenDs andDt is often challenged by the ambiguous matches due to the repetitive patterns or background clutters [ 42 , 24 , 26 ] . To address this , numerous methods proposed cost aggregation techniques that focus on refining the initial matching similarities either by formulating the task as optimal transport problem [ 47 , 31 ] , using regularized Hough matching to re-weight the costs [ 37 , 39 ] , or 4D or 6D convolutions [ 45 , 27 , 44 , 35 ] . However , these methods either use hand-crafted techniques that are weak to severe deformations , or fail to discriminate incorrect matches due to limited receptive fields . To overcome these , we present Transformer-based cost aggregation networks that effectively integrate information present in all pairwise matching costs , dubbed CATs , as illustrated in Fig . 1 . As done widely in other works [ 42 , 45 , 50 , 34 , 37 ] , we follow the common practice for feature extraction and cost computation . In the following , we first explain feature extraction and cost computation , and then describe several critical design choices we made for effective aggregation of the matching costs . 3.2 Feature Extraction and Cost Computation . To extract dense feature maps from images , we follow [ 26 , 37 , 39 ] that use multi-level features for construction of correlation maps . We use CNNs that produce a sequence of L feature maps , and Dl represents a feature map at l-th level . As done in [ 37 ] , we use different combination of multi-level features depending on the dataset trained on , e.g. , PF-PASCAL [ 12 ] or SPair-71k [ 38 ] . Given a sequence of feature maps , we resize all the selected feature maps to Rh×w×c , with height h , width w , and c channels . The resized features then undergo l-2 normalization . Given resized dense features Ds and Dt , we compute a correlation map C ∈ Rhw×hw using the inner product between features : C ( i , j ) = Dt ( i ) ·Ds ( j ) with points i and j in the target and source features , respectively . In this way , all pairwise feature matches are computed and stored . However , raw matching scores contain numerous ambiguous matching points as exemplified in Fig . 2 , which results inaccurate correspondences . To remedy this , we propose cost aggregation networks in the following that aim to refine the ambiguous or noisy matching scores . 3.3 Transformer Aggregator . Renowned for its global receptive fields , one of the key elements of Transformer [ 61 ] is the selfattention mechanism , which enables finding the correlated input tokens by first feeding into scaled dot product attention function , normalizing with Layer Normalization ( LN ) [ 1 ] , and passing the normalized values to a MLP . Several works [ 10 , 3 , 62 , 51 ] have shown that given images or features as input , Transformers [ 61 ] integrate the global information in a flexible manner by learning to find the attention scores for all pairs of tokens . In this paper , we leverage the Transformers to integrate the matching scores to discover global consensus by considering global context information . Specifically , we obtain a refined cost C′ by feeding the raw cost C to the Transformer T , consisting of self-attention , LN , and MLP modules : C′ = T ( C + Epos ) , ( 1 ) where Epos denotes positional embedding . The standard Transformer receives as input a 1D sequence of token embeddings . In our context , we reshape the correlation map C into a sequence of vectors C ( k ) ∈ R1×hw for k ∈ { 1 , ... , hw } . We visualize the refined correlation map with self-attention in Fig . 2 , where the ambiguities are significantly resolved . Appearance Affinity Modeling . When only matching costs are considered for aggregation , selfattention layer processes the correlation map itself disregarding the noise involved in the correlation map , which may lead to inaccurate correspondences . Rather than solely relying on raw correlation map , we additionally provide an appearance embedding from input features to disambiguate the correlation map aided by appearance affinity within the Transformer . Intuition behind is that visually similar points in an image , e.g. , color or feature , have similar correspondences , as proven in stereo matching literature , e.g. , Cost Volume Filtering ( CVF ) [ 16 , 50 ] . To provide appearance affinity , we propose to concatenate embedded features projected from input features with the correlation map . We first feed the features D into linear projection networks , and then concatenate the output along corresponding dimension , so that the correlation map is augmented such that [ C , P ( D ) ] ∈ Rhw× ( hw+p ) , where [ · ] denotes concatenation , P denotes linear projection networks , and p is channel dimension of embedded feature . Within the Transformer , self-attention layer aggregates the correlation map and passes the output to the linear projection networks to retain the size of original correlation C. Multi-Level Aggregation . As shown in [ 37 , 34 , 39 , 58 , 31 ] , leveraging multi-level features allows capturing hierarchical semantic feature representations . Thus we also use multi-level features from different levels of convolutional layers to construct a stack of correlation maps . Each correlation map Cl computed between Dls and Dlt is concatenated with corresponding embedded features and fed into the aggregation networks . The aggregation networks now consider multiple correlations , aiming to effectively aggregates the matches by the hierarchical semantic representations . As shown in Fig . 3 , a stack of L augmented correlation maps , [ Cl , P ( Dl ) ] Ll=1 ∈ Rhw× ( hw+p ) ×L , undergo the Transformer aggregator . For each l-th augmented correlation map , we aggregate with self-attention layer across all the points in the augmented correlation map , and we refer this as intra-correlation self-attention . In addition , subsequent to this , the correlation map undergoes intercorrelation self-attention across multi-level dimensions . Contrary to HPF [ 37 ] that concatenates all the multi-level features and compute a correlation map , which disregards the level-wise similarities , within the inter-correlation layer of the proposed model , the similar matching scores are explored across multi-level dimensions . In this way , we can embrace richer semantics in different levels of feature maps , as shown in Fig . 4 .

Summary: This paper addresses the problem of finding dense correspondences between semantically similar images. Given a feature extractor, the primary focus of this work is on the design of a matching algorithm to find the optimal matches for each spatial location. The proposed method is a learned transformer based architecture that takes advantage of multi-scale features to resolve ambiguous local matches. Results demonstrate that the performance of the proposed model is comparable to existing methods on multiple datasets.