Few-Shot Segmentation via Cycle-Consistent Transformer

1 Introduction . Recent years have witnessed great progress in semantic segmentation [ 19 , 4 , 47 ] . The success can be largely attributed to large amounts of annotated data [ 48 , 17 ] . However , labeling dense segmentation masks are very time-consuming [ 45 ] . Semi-supervised segmentation [ 15 , 39 , 38 ] has been broadly explored to alleviate this problem , which assumes a large amount of unlabeled data is accessible . However , semi-supervised approaches may fail to generalize to novel classes with very few exemplars . In the extreme low data regime , few-shot segmentation [ 26 , 35 ] is introduced to train a segmentation model that can quickly adapt to novel categories . ∗Part of this work was done when Gengwei Zhang was an intern at Baidu Research . †Corresponding author . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Most few-shot segmentation methods follow a learning-to-learn paradigm where predictions of query images are made conditioned on the features and annotations of support images . The key to the success of this training paradigm lies in how to effectively utilize the information provided by support images . Previous approaches extract semantic-level prototypes from support features and follow a metric learning [ 29 , 7 , 35 ] pipeline extending from PrototypicalNet [ 28 ] . According to the granularity of utilizing support features , these methods can be categorized into two groups , as illustrated in Figure 1 : 1 ) Class-wise mean pooling [ 35 , 46 , 44 ] ( Figure 1 ( a ) ) . Support features within regions of different categories are averaged to serve as prototypes to facilitate the classification of query pixels . 2 ) Clustering [ 18 , 41 ] ( Figure 1 ( b ) ) . Recent works attempt to generate multiple prototypes via EM algorithm or K-means clustering [ 41 , 18 ] , in order to extract more abundant information from support images . These prototype-based methods need to “ compress '' support information into different prototypes ( i.e . class-wise or cluster-wise ) , which may lead to various degrees of loss of beneficial support information and thus harm segmentation on query image . Rather than using prototypes to abstract the support information , [ 43 , 34 ] ( Figure 1 ( c ) ) propose to employ the attention mechanism to extract information from support foreground pixels for segmenting query . However , such methods ignore all the background support pixels that can be beneficial for segmenting query image , and incorrectly consider partial foreground support pixels that are quite different from the query ones , leading to sub-optimal results . In this paper , we focus on equipping each query pixel with relevant information from support images to facilitate the query pixel classification . Inspired by the transformer architecture [ 32 ] which performs feature aggregation through attention , we design a novel Cycle-Consistent Transformer ( CyCTR ) module ( Figure 1 ( d ) ) to aggregate pixel-wise support features into query ones . Specifically , our CyCTR consists of two types of transformer blocks : the self-alignment block and the cross-alignment block . The selfalignment block is employed to encode the query image features by aggregating its relevant context information , while the cross-alignment aims to aggregate the pixel-wise features of support images into the pixel-wise features of query image . Different from self-alignment where Query3 , Key and Value come from the same embedding , cross-alignment takes features from query images as Query , and those from support images as Key and Value . In this way , CyCTR provides abundant pixel-wise support information for pixel-wise features of query images to make predictions . Moreover , we observe that due to the differences between support and query images , e.g. , scale , color and scene , only a small proportion of support pixels can be beneficial for the segmentation of query image . In other words , in the support image , some pixel-level information may confuse the attention in the transformer . Figure 2 provides a visual example of a support-query pair together with the label 3To distinguish from the phrase `` query '' in few-shot segmentation , we use `` Query '' with capitalization to note the query sequence in the transformer . masks . The confusing support pixels may come from both foreground pixels and background pixels . For instance , point p1 in the support image located in the plane afar , which is indicated as foreground by the support mask . However , the nearest point p2 in the query image ( i.e . p2 has the largest feature similarity with p1 ) belongs to a different category , i.e . background . That means , there exists no query pixel which has both high similarity and the same semantic label with p1 . Thus , p1 is likely to be harmful for segmenting `` plane '' and should be ignored when performing the attention . To overcome this issue , in CyCTR , we propose to equip the cross-alignment block with a novel cycle-consistent attention operation . Specifically , as shown in Figure 2 , starting from the feature of one support pixel , we find its nearest neighbor in the query features . In turn , this nearest neighbor finds the most similar support feature . If the starting and the end support features come from the same category , a cycle-consistency relationship is established . We incorporate such an operation into attention to force query features only attend to cycle-consistent support features to extract information . In this way , the support pixels that are far away from query ones are not considered . Meanwhile , cycle-consistent attention enables us to more safely utilize the information from background support pixels , without introducing much bias into the query features . In a nutshell , our contributions are summarized as follows : ( 1 ) We tackle few-shot segmentation from the perspective of providing each query pixel with relevant information from support images through pixel-wise alignment . ( 2 ) We propose a novel Cycle-Consistent TRansformer ( CyCTR ) to aggregate the pixel-wise support features into the query ones . In CyCTR , we observe that many support features may confuse the attention and bias pixel-level feature aggregation , and propose incorporating cycle-consistent operation into the attention to deal with this issue . ( 3 ) Our CyCTR achieves state-ofthe-art results on two few-shot segmentation benchmarks , i.e. , Pascal-5i and COCO-20i . Extensive experiments validate the effectiveness of each component in our CyCTR . 2 Related Work . 2.1 Few-Shot Segmentation . Few-shot segmentation [ 26 ] is established to perform segmentation with very few exemplars . Recent approaches formulate few-shot segmentation from the view of metric learning [ 29 , 7 , 35 ] . For instance , [ 7 ] first extends PrototypicalNet [ 28 ] to perform few-shot segmentation . PANet [ 35 ] simplifies the framework with an efficient prototype learning framework . SG-One [ 46 ] leverage the cosine similarity map between the single support prototype and query features to guide the prediction . CANet [ 44 ] replaces the cosine similarity with an additive alignment module and iteratively refines the network output . PFENet [ 30 ] further designs an effective feature pyramid module and leverages a prior map to achieve better segmentation performance . Recently , [ 41 , 18 , 43 ] point out that only a single support prototype is insufficient to represent a given category . Therefore , they attempt to obtain multiple prototypes via EM algorithm to represent the support objects and the prototypes are compared with query image based on cosine similarity [ 18 , 41 ] . Besides , [ 43 , 34 ] attempt to use graph attention networks [ 33 , 40 ] to utilize all foreground support pixel features . However , they ignore all pixels in the background region by default . Besides , due to the large difference between support and query images , not all support pixels will benefit final query segmentation . Recently , some concurrent works propose to learn dense matching through Hypercorrelation Squeeze Networks [ 22 ] or mining latent classes [ 42 ] from the background region . Our work aims at mining information from the whole support image , but exploring to use the transformer architecture and from a different perspective , i.e. , reducing the noise in the support pixel-level features . 2.2 Transformer . Transformer and self-attention were firstly introduced in the fields of machine translation and natural language processing [ 6 , 32 ] , and are receiving increasing interests recently in the computer vision area . Previous works utilize self-attention as additional module on top of existing convolutional networks , e.g. , Nonlocal [ 36 ] and CCNet [ 14 ] . ViT [ 8 ] and its following work [ 31 ] demonstrate the pure transformer architecture can achieve state-of-the-art for image recognition . On the other hand , DETR [ 3 ] builds up an end-to-end framework with a transformer encoder-decoder on top of backbone networks for object detection . And its deformable vairents [ 51 ] improves the performance and training efficiency . Besides , in natural language processing , a few works [ 2 , 5 , 27 ] have been introduced for long documents processing with sparse transformers . In these works , each Query token only attends to a pre-defined subset of Key positions . 2.3 Cycle-consistency Learning . Our work is partially inspired by cycle-consistency learning [ 50 , 9 ] that is explored in various computer vision areas . For instance , in image translation , CycleGAN [ 50 ] uses cycle-consistency to align image pairs . It is also effective in learning 3D correspondence [ 49 ] , consistency between video frames [ 37 ] and association between different domains [ 16 ] . These works typically constructs cycle-consistency loss between aligned targets ( e.g. , images ) . However , the simple training loss can not be directly applied to few-shot segmentation because the test categories are unseen from the training process and no finetuning is involved during testing . In this work , we incorporate the idea of cycle-consistency into transformer to eliminate the negative effect of confusing or irrelevant support pixels . 3 Methodology . 3.1 Problem Setting . Few-shot segmentation aims at training a segmentation model that can segment novel objects with very few annotated samples . Specifically , given dataset Dtrain and Dtest with category set Ctrain and Ctest respectively , where Ctrain ∩ Ctest = ∅ , the model trained on Dtrain is directly used to test on Dtest . In line with previous works [ 30 , 35 , 44 ] , episode training is adopted in this work for few-shot segmentation . Each episode is composed of k support images Is and a query image Iq to form a k-shot episode { { Is } k , Iq } , in which all { Is } k and Iq contain objects from the same category . Then the training set and test set are represented by Dtrain = { { Is } k , Iq } Ntrain and Dtest = { { Is } k , Iq } Ntest , where Ntrain and Ntest is the number of episodes for training and test set . During training , both support masks Ms and query masks Mq are available for training images , and only support masks are accessible during testing . 3.2 Revisiting of Transformer . Following the general form in [ 32 ] , a transformer block is composed of alternating layers of multi-head attention ( MHA ) and multi-layer perceptron ( MLP ) . LayerNorm ( LN ) [ 1 ] and residual connection [ 12 ] are applied at the end of each block . Specially , an attention layer is formulated as Atten ( Q , K , V ) = softmax ( QKT√ d ) V , ( 1 ) where [ Q ; K ; V ] = [ WqZq ; WkZkv ; WvZkv ] , in which Zq is the input Query sequence , Zkv is the input Key/Value sequence , Wq , Wk , Wv ∈ Rd×d denote the learnable parameters , d is the hidden dimension of the input sequences and we assume all sequences have the same dimension d by default . For each Query element , the attention layer computes its similarities with all Key elements . Then the computed similarities are normalized via softmax , which are used to multiply the Value elements to achieve the aggregated outputs . When Zq = Zkv , it functions as self-attention mechanism . The multi-head attention layer is an extention of attention layer , which performs h attention operations and concatenates consequences together . Specifically , MHA ( Q , K , V ) = [ head1 , ... , headh ] , ( 2 ) where headm = Atten ( Qm , Km , Vm ) and the inputs [ Qm , Km , Vm ] are the mth group from [ Q , K , V ] with dimension d/h .

Summary: The paper introduces a new transformer architecture designed to tackle image segmentation in the few-shot setting. The new architecture seems to make it easier for the pre-trained model to generalise to new classes and produce more accurate segmentation masks. The comparison with benchmarks shows incremental improvements on PASCAL-5i and COCO-20i in 1-shot and 5-shot setting, with two different convolutional backbones.  Due to minor yet numerous issues the paper is currently below the acceptance bar, but I am willing to improve my score if these are addressed.