CATrans: Context and Affinity Transformer for Few-Shot Segmentation

1 Introduction

Fully supervised semantic segmentation has made tremendous progress in recent years [Long et al., 2015; Lin et al., 2017]. However, these methods rely heavily on a large amount of pixel-wise annotations which requires intensive manual labor, and are incapable of generalizing to new classes with a handful of annotations. In contrast, humans can recognize a new category even with little guidance. Inspired by this, Few-shot segmentation (FSS) has recently received a growing interest in the computer vision community [Wang et al., 2019; Zhang et al., 2019b; Zhang et al., 2020].

∗ Interns at the Institute of Deep Learning, Baidu Research

†Corresponding author

The goal of FSS is to segment the novel class in the query image, conditioned on the given support set which contains only a few support images and the corresponding groundtruth masks. A fundamental challenge is the large intra-class appearance and geometric variations between support-query pairs, so the key issue is how to effectively reason the relationships of paired samples.

Most FSS methods [Cho et al., 2021; Zhang et al., 2021; Lu et al., 2021] follow a learning-to-learn paradigm. Specifically, features are extracted from both query and support images, and then passed through a feature matching procedure to transfer the support mask to the query image. The convolutional neural network (CNN)-based approaches [Wang et al., 2019; Yang et al., 2020; Liu et al., 2020] condense the masked object features in the support image into a single or few prototypes. Recently, some approaches introduce Transformer-based architecture to establish pixel-wise matching scores between support-query pairs, containing two main technical routes, ie., the context and affinity.

As the dense context information is beneficial to FSS task, especially when the large intra-class variances exist in the support and query set, [Lu et al., 2021] propose the Classifier Weight Transformer (CWT) to dynamically adapt the classifier's weights trained on support set to each query image as shown in Figure 1 (a). [Zhang et al., 2021] aggregates the context information within query images and between support-query pairs via transformer blocks for self- and crossalignment, as shown in Figure 1 (b). This method, however, suffers from unrepresentative of support feature which stimulates us to propose the Relation-guided Context Transformer (RCT), in which the global context information of support can be considered with a mask encoder [Johnander et al., 2021] and a transformer block. The RCT takes as input discriminative self-feature to build more accurate cross-correlation of context, as briefly shown in Figure 1 (d).

The attention-aware affinity is globally constructed between support and query features as another guidance for query segmentation. The Cost Aggregation with Transformers (CATs) [Cho et al., 2021] builds the cross-affinity between support and query features followed by transformer blocks, as shown in Figure 1 (c). However, this method does not incorporate individual self-affinity for support object or query image to disambiguate noisy correlations, which measures pixel-wise correspondences within itself, enabling each spatial fiber to match itself and other tokens. We, thus, design a Relation-guided Affinity Transformer (RAT) for generating a reliable cross-affinity map inherited from the self-affinity. The illustration is schematically depicted in the Figure 1 (d).

Additionally, we explore how to utilize both context and affinity guidance simultaneously. To be specific, we develop a hierarchical CATrans: Context and Affinity Transformer, where we leverage a stack of multi-level correlation maps related to both context and affinity. Moreover, following by [Johnander et al., 2021], we also concatenate the query embedding with high resolution with those low-resolution correlation maps to guide the decoder.

Overall, our contributions are : i) we design a Relationguided Context Transformer (RCT) with the enhanced support features to propagate informative semantic information from support to query images. ii) we develop a Relationguided Affinity Transformer (RAT) to measure the reliable cross correspondences by considering the auxiliary selfaffinity of both support object and query images. iii) we propose Context and Affinity Transformer, dubbed as CATrans, in a hierarchical architecture to aggregate the context and affinity together, resulting in discriminative representations from support to query mask, enhancing robustness to intraclass variations between support and query images. Our CATrans outperforms the state-of-the-art methods on two benchmarks, Pascal-5i and COCO-20i .

2 RelatedWork

2.1 Semantic Segmentation

Semantic segmentation is a fundamental problem in computer vision, which aims to classify each pixel of an image into predefined categories. Most existing semantic segmentation methods are based on fully convolutional networks (FCNs)

[Long et al., 2015], that replaces the fully connected layer with fully convolutional ones for pixel-level prediction. Recent breakthroughs in semantic segmentation have mainly come from multi-scale feature aggregation or attention mechanisms. However, the traditional fully supervised segmentation methods require large amounts of image-mask pairs for training, which are very expensive and time consuming. Additionally, it cannot extend model's generalizability to unseen classes with only a few well-labeled samples.

2.2 Few-shot Segmentation

Few-shot semantic segmentation has attracted lots of research attentions after that [Shaban et al., 2017], which first dealt with this issue by proposing to adapt the classifier for each class, conditioned on the support set. Recent approaches formulate few-shot segmentation from the view of metric learning. [Dong and Xing, 2018] learned prototypes for different classes and the segmentation results are made by cosine similarity between the features and the prototypes. [Wang et al., 2019] developed an efficient prototype learning framework to build consistent prototypes. PFENet [Tian et al., 2020] made progress by further designing an effective feature pyramid module and leveraged a prior map to achieve a better segmentation performance. Recently, [Liu et al., 2020; Yang et al., 2020] found that it is insufficient to represent a category with a single support prototype. Therefore, they used multiple prototypes to represent the support objects via the EM algorithm or K-means clustering. However, these methods disregard the pixel-wise relationships of spatial structure in feature maps.

Recent works [Zhang et al., 2019a; Zhang et al., 2021; Cho et al., 2021; Lu et al., 2021] attempted to fully utilize a correlation map to leverage the pixel-wise relationships between support and query features. Specially, [Zhang et al., 2019a] used graph attention networks to propagate information from the support image to query images, and [Zhang et al., 2021] utilized cycle-consistent transformer to aggregate the pixel-wise support features into the query. [Lu et al., 2021] proposed a classifier weight transformer where the transformer is applied to adapt the classifier solely by freezing the encoder and decoder. However, all these dense matching methods focus on either context correspondences or affinity maps only. This is no study about whether these two measures are complementary, and can be integrate to achieve a better performance.

2.3 Transformers in Vision

Recently, transformers, first introduced in natural language processing [Vaswani et al., 2017], and are receiving increasing interests in the computer community. Since the pioneer works such as ViT [Dosovitskiy et al., 2021], it demonstrates the pure transformer architecture can achieve the state-of-theart for image recognition. On the other hand, DETR [Carion et al., 2020] built up an end-to-end framework with a transformer encoder-decoder on top of backbone networks for object segmentation. And its deformable variants [Zhu et al., 2021] improved the performance and training efficiency. However, there are few studies that compute both context and affinity.

3 Preliminaries

3.1 ProblemFormulation

The goal of few-shot segmentation is to segment novel objects with very few annotated samples. Specifically, all classes are divided into two disjointed class set Ctrain and Ctest. To mitigate the overfitting caused by insufficient training data, we follow the common protocol called episodic training. Under K-shot setting, each episode is composed of a support set S = {(Is,Ms)}K , where Is,Ms are support image and its corresponding mask, and a query sample Q = (Iq,Mq), where Iq,Mq are the query image and mask, respectively. In particular, given dataset Dtrain = {S,Q}Ntrain and Dtest = {S,Q}Ntest with category set Ctrain and Ctest, respectively, where Ntrain and Ntest is the number of episodes for training and test sets. During training, our model takes a sampled episode from both support masks Ms and query masks Mq , and only use support masks for predicting the query segmentation map ˆMq during testing.

3.2 Revisiting of Transformer

Transformer block [Vaswani et al., 2017] consists of multihead attention (MHA) and multi-layer perception (MLP) with inputs of a set of Query (Q), Key (K) and Value (V) elements. In addition, LayerNorm (LN) and residual connection are available at the end of each block. Specially, an attention layer is formulated as Architecture Overview:

where [Q;K; V ] = [WqFq;WkFk;WvFv], in which Fq is the input query sequence, Fk/Fv is the input key/value sequence, Wq,Wk,Wv ∈ Rc×c are learnable weights, c is the channel dimension of the input sequences.

The multi-head attention layer is an enhancement of attention layer, where h attention units are applied and then concatenated together. Concretely, this operation splits input sequences along the channel dimension c into h groups:

where headm = Atten(Qm,Km, Vm) and the inputs [Qm,Km, Vm] are the mth group from [Q,K, V ] with dimension c/h.

4 Methodology

Below we present the overall architecture of CATrans followed by a description of its individual components.

4.1 Architecture Overview

The overall architecture of our Context and Affinity Transformer (CATrans) is illustrated in Figure 2, which consists of Image Encoder, Mask Encoder, Relation-guided Context Transformer (RCT), Relation-guided Affinity Transformer (RAT) and the Decoder. Specifically, the input supportquery images {Is, Iq} are passed through the Image Encoder to extract multi-scale features {F l s , F l q ∈ RHl×Wl×Cl }4l=1 , where l denotes the scale-level, and pyramid mask features

{F lm ∈ RHlm×W lm×Clm}4l=1 are extracted via the Mask Encoder with the input of support binary mask Ms. The resulting triple {F l s , F l q , F lm}4l=3 is passed into the RCT and RAT, respectively. In practice, the pixel-wise context C l ∈ RHl×Wl×Clm provided by the RCT and dense affinity map A ∈ RHlWl×HlWl generated by RAT are concatenated for information aggregation. In the decoder, the fused representations associate with the high-resolution features of query image {F l s , F l q }2l=1 for predicting the query mask ˆMq.

4.2 Relation-guided Context Transformer

Inspirations: Current CNN-based prototypical learning approaches, ie. [Wang et al., 2019], condense the support features into a single or few context-wise prototypes. However, prototypical features inevitably drop spatial information and fail to discriminate incorrect matches due to limited receptive fields of CNNs. So recent research applied the Transformer-based architecture, such as [Zhang et al., 2021], to establish long-range and pixel-wise context relationships between paired support-query samples, outperforming previous CNN-based methods by a large margin. We conjecture that the representative features within the individual support and query image would help to aggregate more precise crossrelationships, being robust to large intra-class differences of paired support-query samples.

To this end, the RCT is designed and shown in Figure 2 (b). We flatten the given triple features (F l s , F l q , F lm) into 1D tokens as inputs for the following items.

Self-context. Self-context is separately employed for support objects and query features by aggregating its relevant context information, leading to more informative support and query features to be connected via the cross-context. The resulting contexts C lss and C lqq are designed as:

where MHA (·), LN(·) and MLP(·) are operations introduced in the section 3.

Cross-context. Considering that C lss , guided by mask features F lm, mainly focuses on the foreground, but background support pixels are beneficial for building the semantic relationships, we collaborate the enhanced self-context features of support and query with F l s to establish the pixel-wise crosscontext. The process of cross-context is formed as:

where C lsq ∈ RHlWl×Clm is spatially rearranged to the shape of Hl ×Wl × C lm.

4.3 Relation-guided Affinity Transformer

A huge feature distinction between the support and query images brings barriers for context knowledge transfer, which cripples the segmentation performance. We explore several attention-aware affinity maps that measure pixel-wise correspondences to facilitate the FSS task, as shown in Figure 2 (c). Overall, this module provides affinity guidance stemming from attention-aware features instead of semantics.

Why needs self-affinity? The training samples belonging to the same class always have features with large variations in appearance as these objects are taken in unconstrained settings. Take aeroplane as an example, all aeroplanes are made by metal and have wings. These features can be seen as intrinsic features. As the differences of shooting angle and lighting conditions, the shape and color of aeroplanes can be different. In few-shot segmentation, we need to enable each pixel-wise feature belonging to itself to match pixel-wise feature at the same position, making it robust to large variations in object appearance between the support and query images.

What is Self-affinity? For a support image, we utilize highdimension support mask features F lm concatenated with support image features F l s for estimating its affinity map by scaled-dot product followed by softmax function. The process for l-th features can be defined as:

where {·||·} denotes the concatenation operation, softmax (·) is a row-wise softmax function for attention normalization and two individual fc layers are applied to learn discriminative features by learnable parameters. Analogy to support feature, we formulate the self-affinity for query as:

Cross-affinity. The cross-affinity between support and query features F l ca is computed by:

However, solely relying on the cross-affinity between features often suffers from the challenges caused by large intra-class variations. We then embed the self-affinity of query Alqq into cross-affinity. The final cross-affinity is formulated as:

5 Experiments

In this section, we conduct extensive experiments for our CATrans on two widely-used few-shot segmentation benchmarks, PASCAL-5i and COCO-20i , to demonstrate the effectiveness of our method.

5.1 Datasets

PASCAL-5i [Shaban et al., 2017] is composed of PASCAL VOC 2012 with additional SBD [Hariharan et al., 2011] annotations, which contains 20 categories split into 4 folds (15/5 categories as base/novel classes).

COCO-20i [Lin et al., 2014] is created from MS COCO where the 80 object categories are divided into four splits (60/20 categories as base/novel classes).

5.2 Implementation Details

We conduct all experiments on 1 NVIDIA V100 GPU. The models are trained for 20k and 40k iterations on PASCAL-5i and COCO-20i , respectively, with AdamW as the optimizer. The initial learning rate is set to 5e-5 and decays at 10k iteration with a factor of 0.1. During training, we first resize the input images to 384×384 and 512×512 for PASCAL-5i and COCO-20i , respectively, and then perform the horizontal flip operation randomly. We simply use cross-entropy loss with a weight of 1 and 4 for background and foreground pixels, respectively. The BN layers of image encoder are frozen. For a fair comparision, we employ the widely-used ResNet50, ResNet-101 and Swin-Transformer as the image encoder. The mask encoder includes four light-weight layers, each of which is composed of 3 × 3 convolution, BatchNorm and ReLU.

During evaluation, the results are averaged on the randomly sampled 5k and 20k episodes for each fold and 5 runs with different seeds. we report the mean-IoU (mIoU) under both 1-shot (given a single support example) and 5-shot (given five support examples).

5.3 Comparisons with the State-of-the-art

Results onPASCAL-5i . As shown in Table 1, our CATrans outperforms the previous best DGPNet by +2.8/+1.8% and +2.4/+1.1% for 1-shot/5-shot mIoU using ResNet-50 and ResNet-101 as the backbone, repectively. Rendering the Swin-T as the image encoder, our CATrans further achieves 67.6% and 77.3% mIoU for 1-shot and 5-shot segmentation.

Results onCOCO-20i . Table 3 reports the comparisons on the more challenging COCO-20i dataset. Compared with the previous best DGPNet, our CATrans surpasses it by +1.6/+2.0% 1-shot/5-shot mIoU using the ResNet-50. When using ResNet-101, our CATrans is +2.1% and +1.8% higher than DGPNet for 1-shot and 5-shot protocals, repectively. Equipped with the Swin-T, our CATrans achieves 49.4/60.1% 1-shot/5-shot mIoU.

5.4 Ablation Study

Here, we conduct extensive ablation studies with ResNet-50 on PASCAL-5i and COCO-20i to analyze the effect of the key components in our CATrans.

Effectiveness of RAT and RCT. We ablate our CATrans to observe the effectiveness of RAT and RCT modules, as shown in Table 2a. We define the baseline, without RAT and RCT modules, that is simply concatenating the support and query features along the channel mode. The variant with either RAT or RCT boosts baseline by up to +6.8/+7.7 on 1-shot protocol. Moreover, it can be seen that the support self-context branch of RCT provides additional ∼1% mIoU on 1-shot setting (RCT vs. RCT◦ ), endowing RCT with a powerful context aggregator. Considering that the cross-affinity between support and query images can serve as the additional guidance for FSS task, when the large feature distinctions impede context knowledge transfer, we verify how much benefit comes from the RAT. Table 2a shows the model equips with both RCT and RAT has ∼5% increase for FSS results.

Multi-scale Representations. We first verify the influence of fusing the high-resolution features of query image (l =1/2) on the top panel of Table 2b. Then on the bottom panel of Table 2b, it shows a comparison experiment between singlescale and multi-scale representations of CATrans (l = 3/4). The more levels of information used, the better the performance. The performance reaches the highest when these two levels are both leveraged, especially in the 1-shot setting where multi-scale guidance can extract more guiding information for query segmentation with 1.4%/1.8% mIoU gain on PASCAL-5i /COCO-20i .

Effect of Model Capacity. We stack more numbers of heads of attention layer to increase capacity of our CATrans and validate its effectiveness. It shows that our model performance is stable across different choices, particularly for #h ≥2.

Memory and Run-time. CATrans, with trivial computational overhead, performs the best over closely related transformer-based methods. CyCTR and CATs stack two successive transformer blocks whereas CATrans consists of RCT and RAT, each with one transformer block. The memory and run-time comparisons are below:

Qualitative results. To show the performance of CATrans intuitively, we visualize some final prediction masks produced by our method and its variants in Figure 3. The first column is support image and its ground truth, and the next three columns are query mask produced by RAT, RCT and CATrans, respectively. The last column is ground truth of query image. On the top row of Figure 3, as the large intraclass appearance variations between support and query images hidden context knowledge transferring, RCT fails to precisely segment one of its aerofoils, where RCT performs better by use of the self- and cross- affinity for query segmentation. On the bottom of Figure 3, the RAT mistakenly regards some plates as part of the target cup because of the similarity in the shape and color, measured by the attention-aware affinity. But RAT reduces the area of wrong segmentation significantly by successfully aggregating the context information. Overall, adopting both RAT and RCT performs the best.

6 Conclusion

We have proposed the novel Context and Affinity Transformer (CATrans) with RCT and RAT in a hierarchical architecture to deal with the large intra-class appearance and geometric variations in few-shot segmentation, so that CATrans can effectively incorporate both context similarity and affinity map for query segmentation. In addition, we consider pixelwise correspondences for individual support and query features to disambiguate noisy correlations.