A Transformer-Based Adaptive Prototype Matching Network for Few-Shot Semantic Segmentation

1 Introduction

In recent years, traditional semantic segmentation [Wang et al., 2019] has made significant progress due to the rapid advancements in deep learning within the field of computer vision [Chen et al., 2022b; Ye et al., 2021]. However, this task has long struggled with challenges such as dense annotation requirements and limited generalization. In such case, fewshot semantic segmentation (FSS) [Tian et al., 2020] has been proposed to simulate real-world scenarios with limited data and multiple categories.

FSS follows the framework of meta-learning, conducting training in the form of episodes that consist of a support set and a query set. The model execution process can be divided into three stages. First, both the support set and the query set are sent synchronously to the parameter-sharing backbone network for feature extraction. Next, in the feature matching stage, interaction occurs between the annotated support features and the query features using either a prototypepixel mechanism [Zhang et al., 2020; Zhang et al., 2021a; Fan et al., 2022; Liu et al., 2022c; Cao et al., 2022] or a pixel-pixel mechanism [Xie et al., 2021; Min et al., 2021; Zhang et al., 2021b; Shi et al., 2022]. Finally, in the classification stage, FSS predicts the segmentation mask for the target category in the query image.

∗Corresponding author.

Existing FSS models [Cao et al., 2022; Zhang et al., 2021b; Liu et al., 2022c; Chen et al., 2024] have shown impressive results. However, as illustrated in Fig.5, previous works still face challenges due to background interference, including adjacent regions, seen classes, and analogs. The reasons can be explained from three aspects: (i) in the feature extraction stage, previous works [Tian et al., 2020; Chen et al., 2022a; Cao et al., 2022] have relied on features directly extracted from the pretrained backbone networks. However, these pretrained backbones exhibit inherent bias. As shown in Fig.1(a), they tend to prioritize extracting features related to "person" rather than the specific target category "bottle" for the current task. (ii) In the feature matching stage, previous works [Cao et al., 2022; Liu et al., 2022c; Shi et al., 2022] have often utilized attention mechanisms to establish the single-layer relationship between the support set and the query set for category information transfer. However, this subtle relationship proves to be insufficient when there exists significant intra-class variability within the target category, leading to attention bias (see in Fig.1(b)). (iii) In the classification stage, existing methods [Tian et al., 2020; Zhang et al., 2021a; Liu et al., 2022b] predominantly rely on semantic relevance to make predictions, overlooking the exploration of spatial information, which we refer to as spatialaware bias. As depicted in Fig.1(c), when faced with complex scenes, relying solely on semantic consistency may fail to accurately segment the target category "bottle."

The aforementioned concerns motivate us to introduce a novel Transformer-Based Adaptive Prototype Matching Network. This model mitigates background interference in FSS by strategically and efficiently interacting during the three stages of model execution. The central concept involves leveraging both semantic and spatial perceptions of query features to provide a more comprehensive understanding of category information, ultimately enhancing the robustness of the model.

Firstly, in the feature extraction stage, we draw inspiration from [Lin et al., 2023] and introduce the target enhancement module (TEM) to alleviate inherent bias. TEM focuses on exploring the relevance of multi-scale local context in an computationally efficient manner to enhance foreground features. Secondly, in the feature matching stage, to address attention bias, we devise a dual constraint aggregation module (DCAM). This module emulates the human visual matching process, identifying the most similar target regions in the image based on prior information and using these regions as references for self-retrieval. Finally, in the feature classification stage, addressing the spatial-aware bias, we propose a dual classification module (DCM). This module aims to decouple the segmentation task into two subtasks: semantic alignment and spatial alignment. It leverages semantic consistency for target category identification and spatial consistency for precise localization, collectively improving performance.

In summary, our contributions are as follows:

• We propose a TEM to effectively mitigate inherent bias.
• We propose a DCAM to alleviate attention bias.
• We propose a DCM to individually achieve semantic alignment and spatial alignment, addressing the spatialaware bias.
• Extensive experiments on two benchmark datasets, PASCAL-5i and COCO-20i , demonstrate that the proposed model outperforms other existing competitors using the same metrics.

2 RelatedWork

2.1 Few-Shot Semantic Segmentation

FSS predicts the mask for an unseen category by using a small number of annotated support images. Metric learning-based FSS can be divided into two main categories: prototype-based methods and pixel-based matching methods.

The prototype-based mechanism has emerged as a predominant method in FSS. The pioneering work by [Dong and Xing, 2018] introduced the prototype concept to FSS, utilizing a prototype to represent information about the target class in the support set and performing matching on query features to predict the segmentation mask. Subsequent studies extended this approach from different perspectives. Firstly, some studies [Yang et al., 2020; Li et al., 2021; Yang et al., 2021; Zhang et al., 2021a; Liu et al., 2022a] proposed generating multiple foreground prototypes to fully utilize the support foreground information. Secondly, considering the potential presence of new classes in the background, [Yang et al., 2021; Chen et al., 2022a; Liu et al., 2022b] generated one or more background prototypes by mining the query background.

Due to the inevitable loss of spatial information associated with the prototype-based approach, recent studies [Zhang et al., 2019; Min et al., 2021; Shi et al., 2022; Zhang et al., 2021b; Peng et al., 2023] proposed a pixel-based matching strategy. Unlike the prototype-based method, pixel-based matching seeks to establish a dense association between the support pixels and the query pixels.

Despite the commendable performance achieved by the pixel-based matching method, it is prone to overfitting on the training set. Furthermore, the computational demands of the pixel-based matching method surpass those of the prototypebased approach. After comprehensive consideration, our approach selects a prototype-based matching scheme.

2.2 Transformer

Transformers have shown incredible success from the field of natural language processing (NLP) [Vaswani et al., 2017] to computer vision (CV) [Huang et al., 2022; Ouyang et al., 2023] due to their ability to capture long-range correlations. Some recent works [Lu et al., 2021; Zhang et al., 2021b; Shi et al., 2022; Cao et al., 2022; Liu et al., 2022c] have explored the use of transformers in FSS. [Lu et al., 2021] introduced the classifier weight converter to dynamically adjust classifier weights for each query image. [Zhang et al., 2021b] proposed a cyclic consistent attention mechanism to filter out task-irrelevant pixels from the support set. [Shi et al., 2022] proposed dense pixel cross-query and support attention-weighted mask aggregation to predict the segmentation mask by aggregating multi-level support masks weighted by pixel relevance. [Cao et al., 2022] utilized traditional vanilla attention to strengthen the discriminant of class prototypes. While these methods have been successful, they have certain limitations. The first issue is inherent bias of the backbone. Despite using self-alignment in [Zhang et al., 2021b] to enhance target category features, the quadratic complexity of the input length remains a concern. The second issue is attention bias caused by insufficient constraints. [Liu et al., 2022c] employed mask attention to filter out background noise. However, in cases where there are significant variations within the same class, the accuracy of the mask is worrying.

In our work, we pioneer the integration of the convolutional transformer [Lin et al., 2023] into FSS and enhance foreground features through a local adaptive strategy. Additionally, we introduce a dual semantics-aware attention mechanism to explicitly model the cross-consistency between the support set and the query set as well as the self-consistency within the query set, in order to achieve sufficient mining of the query target class information.

3 Method

3.1 Overview

In this work, we aim to mitigate the susceptibility to background noise caused by inherent bias, attention bias, and spatial-aware bias. To this end, we propose a novel Transformer-Based Adaptive Prototype Matching Network that establishes robust matching relationships between the support set and the query set during the three model execution phases. In the initial feature extraction stage, we introduce a target enhancement module to alleviate inherent bias (see in Section 3.2). Subsequently, in the feature matching stage, to address attention bias, we design a dual constraint aggregation module (see in Section 3.3). Lastly, in the feature classification stage, to tackle spatial-aware bias, we propose a dual classification module (see in Section 3.4). Without loss of generality, we demonstrate the entire network architecture in a 1-shot setting (see in Fig.2).

3.2 Target Enhancement Module

Alleviating the inherent bias of the backbone is crucial for FSS. Despite significant progress in recent work [Zhang et al., 2021b], challenges persist due to the quadratic complexity of the input length. In TEM, to reduce computational costs, we introduce a Multi-Scale Local-Aware Modulation Transformer based on the convolutional transformer architecture [Lin et al., 2023] for performing multi-scale feature extraction. Differently, we employ multi-scale self-adaptive local attention to enhance foreground information and mitigate background interference. Furthermore, we replace the standard multi-layer perception (MLP) [Vaswani et al., 2017] with an invertible neural network (INN) [Dinh et al., 2016] to preserve more fine-grained features in the feedforward process.

Specifically, we first follow previous work [Tian et al., 2020] to get the initial activated query feature F act q using the middle-level query feature F m q ∈ RH×W×C , the support prototype feature F s p ∈ RH×W×C , and the prior mask M 0 as inputs. F s p is obtained from the support prototype extension, where the support prototype ps ∈ R1×C is obtained by applying masked average pooling(MAP) on the middle-level of the support feature F m s ∈ RH×W×C . M 0 is obtained from high-level support and query features. Then, as shown in Fig.3, F act q is treated as the input to theMulti-Scale LocalAware Modulation Transformer. We follow the paradigm of the multi-head attention mechanism, where features are divided into N groups F act q = [F act q1 , F act q2 , ..., F act qN ] by the channel dimension, and a convolutional layer ConvK×K with a kernel size of K × K is applied to each group of features F act qi ∈ RH×W×(C/N) to generate local attention weights Attn(F act qi ) ∈ RH×W×(K×K) . The weights are then normalized by a softmax function. In turn, the weighted features ̂ F act qi ∈ RH×W×(C/N) are integrated by a convolutional layer ConvK×K with the same kernel size of K × K. We then concatenate each group of the weighted features along the channel dimension to obtain the convolutional modulator M ∈ RH×W×C . We initialize the kernel size with 1 × 1 and gradually increase it by 2 per head. The process is described by the following equation:

where Γ1(·) denotes the transformation of Attn(F act qi ) to RH×W×1×(K×K) , and Γ2(·) denotes the transformation of F act qi to RH×W×(C/N)×(K×K) . ⊙ represents the elementwise product.

With the convolutional modulator, we obtain features F tmp q ∈ RH×W×C fused by spatial dimension. Finally, for preserving more detailed features, we employ the invertible feedforward network for inter-channel information fusion. This process is described as:

where Wv denote parameters of the linear mapping. Conv1×1 is a convolutional operation with a kernel size of 1×1. ˜ Fq ∈ RH×W×C indicates the output feature. F tmp q [1 : c] ∈ RH×W×c is the 1st to the cth channels of the input feature F tmp q . Si is the residual function in th INN layer.

3.3 Dual Constraint AggregationModule

To enhance the discriminant of category cues, existing approach [Cao et al., 2022] mines more category information by utilizing vanilla attention to explore the correlation between the support prototype and query features. However, in scenarios with significant intra-class variations of the target category, this single-layer constraint proves to be insufficient, resulting in the issue of attention bias (shown in Fig.4(c)). To address this limitation, we propose a DCAM as illustrated in Fig.4(b) , which consists of two key components: intra-class difference representation and a dual semantic-aware attention mechanism.

Intra-Class Difference Representation. To alleviate intraclass variations, we propose utilizing a set of learnable vectors, namely intra-class difference representation, to model the variance between the support set and the query set. Specifically, we expand the support prototype ps to a size of N × C, and assign an intra-class difference representation ∆q to each of them separately. The generated prototype can be denoted as pini s . To explain the validity of the proposed intra-class difference representation, we formulated its role in the subsequent attention process. In the role of the intraclass difference representation, the computational process of the attention map can be extended as follows:

where Wsupp is the attention weight obtained from query Q and key K, and Wintra is the weight obtained from the intraclass difference representation ∆q and key K. Wintra can be considered as a factor to adjust Wsupp according to the intra-class variation.

Dual Semantic-Aware Attention Mechanism. As illustrated in Fig.4(d), our proposed dual semantic-aware attention mechanism consists of two layers of constraints. In the first layer, we use the support prototype as a reference to select points with high matching confidence in the query feature. These selected points are then used as guidance in the second layer to find points with high feature similarity in the entire query feature map. Throughout this process, we refer to pini s as Q1, ˜ Fq as Q2, K, and V .

To begin, we calculate the matching confidence map Smat ∈ RN×HW between Q1 and K. The matching score can be obtained by

Then we use the indices of the M points with the highest matching confidence in Smat to guide the selection of corresponding points TopM(Q2) in Q2 as references for the second layer constraints. It is worth noting that we employ a soft conditioning approach, where the number of M is set to be 1/4 of the average size of all support foreground regions in each batch.

In the second layer, we calculate the feature similarity matrix between the guide points TopM(Q2) and K, and perform the maximization operation max(˙) along the guide points dimension. The process can be described as:

We then use the obtained attention map Ssim ∈ R1×HW to weigh the aggregation of the category prototype from the query feature, and send it to the feedforward layer (FFN) to generate the robust support category prototype ˜ ps.

3.4 Dual ClassificationModule

Existing methods [Tian et al., 2020; Zhang et al., 2021a; Liu et al., 2022b] predominantly predict based on semantic consistency, often neglecting the spatial consistency of the target object, leading to failures in locating target categories. In the DCM, our objective is to decouple the segmentation task into two subtasks: semantic alignment and spatial alignment. The semantic similarity-based mask serves to identify the target category, while the spatial distribution probability-based mask assists in precise localization, synergizing for improved performance. Firstly, we optimize ˜ Fq and ˜ ps using TEM and DCAM. Then, we generate the semantic similarity-based map Msim by element-wise product of refined prototype ˜ prs and ˜ F r q . Moreover, taking into account the intrinsic guidance provided by query features, we exploit the spatial consistency within the target object itself in the query set to obtain the spatial distribution probability-based mask. Specifically, we apply spatial attention [Woo et al., 2018] to ˜ Fq and ˜ F r q , obtaining spatial distribution guidences, Sq and Sr q , reflecting foreground position information. We then concatenate F m q , ˜ Fq, ˜ F r q , Sq, Sr q and the prior maskM 0 along the channel dimension and feed them into a decoder FEM(·)[Tian et al., 2020] to estimate the spatial distribution probability-based query foreground mask Mpro. The process can be described as:

Finally, we combin Msim and Mpro through simple addition to obtain the ultimate query foreground segmentation map Mf inal :

where M f pro denotes the foreground prediction map in Mpro.

3.5 Total Training Loss

Our training loss is composed of three parts. Firstly, we employ two dice losses, Lfinal and Lsim, to supervise the training of the final prediction map Mfinal and the semantic similarity-based mask Msim. Second, we utilize the binary cross-entropy loss Lpro to supervise the training of the spatial distribution probability-based map Mpro. Finally, we utilize KL(Kullback-Leibler) divergence loss LKL to distill the foreground distribution information of the target obeject in the query to students Sr q and Sq using the query ground truthMq and Sr q as the teacher, respectively. In summary, our overall objective function is:

where epoch represents the total number of training rounds, e represents the current round, and λ is a adjustable loss weights, here we set λ to 10.

4 Experiments

In the experiments, we leverage two popular FSS benchmarks, i.e., PASCAL-5i [Shaban et al., 2017] and COCO-20i [Nguyen and Todorovic, 2019], to evaluate the proposed approach. We adopt mean intersection-over-union (mIoU) as the evaluation metric for experiments.

4.1 Implementation Details

Following [Lang et al., 2022], we first train the PSPNet to obtain a backbone based on the seen training classes for each fold, i.e., 16/61 training classes (including background) for PASCAL-5i /COCO-20i . Subsequently, the parameters of the trained backbone are frozen, and a meta-learning strategy is employed to train the remaining structures. Optimization of these structures is conducted using the Adam optimizer with a learning rate of 10e-3, involving 50 epochs on PASCAL-5i and 100 epochs on COCO-20i . All images are resized directly to 473×473, and the channel dimension of the image features is set to 64. The training batch size is configured as 20 for the 1-shot setting and 15 for the 5-shot setting. No data augmentation strategies are applied during training. All experiments are executed on a single 24GB RTX3090 GPU.

4.2 ComparisonWith State-of-the-Art Methods

PASCAL-5i Results. Table 1 presents a performance comparison of mIoU on the PASCAL-5i dataset between our method and several representative models. It is evident that (1) our method outperforms the previous state-of-the-art [Min et al., 2021] by 2.4% and 0.6% in the 1-shot and 5-shot settings with VGG16 as the backbone network, respectively. (2) On ResNet50, our model achieves state-of-the-art performance [Bao et al., 2023] with just 1/7 of the parameters. With a minimal increase of 0.5M parameters, we surpass the baseline [Cao et al., 2022] by 3.9% and 3.1% in the 1-shot and 5-shot settings. Notably, in the 1-shot and 5-shot settings of fold2, we outperform the baseline by 5.2% and 7.0%, and in fold3, we achieve improvements of 6.5% and 3.4%, respectively. These results further emphasize the effectiveness of our method in mitigating background interference and achieving accurate query segmentation mask.

COCO-20i Results. COCO-20i is a more challenging dataset with a diverse range of categories and intricate backgrounds. Table 2 illustrates the mIoU performance comparison on the COCO-20i benchmark. It can be seen that (1) our method built on VGG16 surpass the previous state-of-the-art [Liu et al., 2022a] by 2.3% and 1.5% in the 1-shot and 5-shot settings, respectively. (2) Using ResNet50 as backbone, our image-only model has competitive performance with stateof-the-art [Yang et al., 2023] that utilizes both text and image information under the 1-shot setting. Furthermore, in the 5shot setting, we surpass [Yang et al., 2023] by 3.0% using only 1.1M parameters. This underscores the robust generalization capability of our model to effectively cope with in handling bias problems.

Qualitative Results. We present qualitative results comparing our method with previous works, including CYCTR [Zhang et al., 2021b], IPMT [Liu et al., 2022c], and Protoformer [Cao et al., 2022], on the PASCAL-5i and COCO-20i benchmarks. Our method demonstrates several advantages over previous works, as depicted in Fig.5. (1) Our approach successfully mitigates background interference from adjacent regions, a contrast to previous works that erroneously treat the background surrounding the foreground as the region for segmentation(see 1st to 3rd column). (2) Our method rectifies the misconception of considering known classes in the background as foreground, concentrating attention on the current target category(see 4th to 7th columns). (3) Our method effectively suppresses interference from analogs in the background, facilitating precise target category localization (see 8th to 10th columns).

4.3 Ablation Experiments

We conduct following ablation studies with ResNet-50 backbone under the 1-shot setting on PASCAL-5i dataset.

Components Analysis. Our approach comprises three main modules: target enhancement module (TEM), dual constraint aggregation module (DCAM), and dual classification module (DCM). Table 3 presents our validation on the effectiveness of each component. Compared to the baseline, using TEM alone to enhance query foreground features and using DCAM alone to enhance the discriminant of category prototypes results in a 0.9% and 2.2% improvement, respectively. The synergistic effect of TEM and DCAM lead to a 2.6% improvement. Employing DCM to achieve semantic alignment and spatial alignment provides an extra growth of 1.3%. The results reveal a 3.9% improvement of our model over the baseline, indicating that the introduced modules effectively address three issues-namely, inherent bias, attention bias, and spatial-aware bias. This ultimately reduces background interference, leading to precise segmentation.

Target Enhancement Module. TEM aims to mitigates the inherent bias of the backbone and enhance the query foreground regions. To evaluate the performance of our proposed method, we perform experiments with other methods in terms COCO-20!

of computational effort and accuracy, respectively. We modified our model by (1) employing the self alignment module in [Zhang et al., 2021b] as the attention mechanism for TEM (referred to as SA), (2) replacing SA with the convolutional transformer architecture [Lin et al., 2023] (referred to as SAM), (3) substituting it with our Multi-Scale SelfAdaptive Local Attention (referred to as MSLA+MLP), and (4) utilizing INN[Dinh et al., 2016] instead of MLP as feedforward network (referred to as MSLA+INN). As illustrated in Figure 6, our approach maintains a high level of accuracy while reducing computational complexity. Moreover, our feedforward network retains more feature details at a slightly increased computational cost.

Dual Constraint Aggregation Module. We provide a comprehensive analysis of a crucial component in DCAM. We modified our model by (1) employing the original Vanilla Attention[Cao et al., 2022] as the attention mechanism for DCAM (referred to as VA), (2) replacing VA with mask attention [Cheng et al., 2022] (referred to as MA), (3) substituting it with our dual semantic-aware attention (referred to as DSAA), and (4) using intra-class difference representation (referred to as IDR). The results in Table 4 suggest that employing mask attention to mitigate background noise interference has negligible impact on performance enhancement. We attribute this to the fact that the mask derived from the similarity between support and query set suffers from the challenge of accuracy when there exists significant intra-class differences. In contrast, our dual semantic-aware attention mechanism copes with sensitivity to intra-class differences by mitigating background interference in a learnable manner(see in Fig. 4(c)). Moreover, our intra-class difference representation proves beneficial across three different attention mechanisms. Dual Classification Module. To assess different DCM components, ablation experiments were conducted. Table 5 shows that using only the semantic similarity-based mask improves the model's performance by 1.5%, demonstrating the necessity of optimizing class prototypes and query features. However, when using only the spatial distribution probabilitybased segmentation map, performance decreases by 2.3%. We consider this because relying only on the foreground distribution of the query image itself causes the model to bias towards the areas of known classes, resulting in a failure of segmentation for unseen classes.

5 Conclusion

We introduce a novel transformer-based adaptive prototype matching network to counteract background interference arising from inherent bias, attention bias, and spatial-aware bias. Our method includes three modules: Target Enhancement Module (TEM) addresses inherent bias by leveraging multiscale local context relevance to enhance foreground features. Dual Constraint Aggregation Module (DCAM) handles attention bias through a dual semantic-aware attention mechanism to reinforce constraints. Dual Classification Module (DCM) decouples the segmentation task into semantic alignment and spatial alignment to alleviate spatial-aware bias. Our experiments demonstrate that our method achieves the state-of-the-art performance with minimal parameters.

Acknowledgements

This work was supported in part by the National Natural Science Foundation of China under Grant U20B2065, Grant U22B2056, Grant 62272468, Grant 62332015, and Grant 62072449; in part by the Natural Science Foundation of Jiangsu Province under Grant BK20211539; in part by the Science and Technology Development Fund, Macau, SAR, under Grant 0075/2023/AMJ, Grant 0003/2023/RIB1, and Grant 001/2024/SKL; in part by the Zhuhai Science and Technology Innovation Bureau under Grant ZH2220004002524; and in part by the University of Macau under Grant MYRG2022-00059-FST and Grant MYRG-GRG2023-00237-FST-UMDF.

Contribution Statement

The contributions of each author to the manuscript are detailed below:

• Sihan Chen and Yadang Chen (Co-First Authors): Both authors contributed equally to the Conceptualization, Methodology, Investigation, Visualization, and Writing - Original Draft.
• Yuhui Zheng (Corresponding Author): Conceptualization, Funding Acquisition, Supervision, Writing - Review & Editing.
• Zhi-Xin Yang: Funding Acquisition, Resources.
• EnhuaWu: Resources, Supervision.