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: The paper proposes a cost-volume refinement method using transformers to be used for a cost-volume-based semantic correspondence framework. The main innovations of the method are the use of transformers, how the cost volume is aggregated with appearance features, and the serial way in which appearance features are used. The method is tested on SPair-71K, PF-Pascal, and PF-Willow datasets.