Self-Supervised Tuning for Few-Shot Segmentation

1 Introduction

Recently, semantic segmentation models [Long et al., 2015] have made great progress under full supervision, and some of them have even surpassed the level of human recognition. However, when the learned model is applied to a new segmentation task, it takes great cost to collect a large amount of full annotated data in pixel level. Furthermore, samples are not available in large quantities in some areas such as health care, security and so on. To address this problem, various few-shot segmentation methods are proposed.

∗Corresponding author

One solution for solving few-shot segmentation [Shaban et al., 2017] is meta-learning [Munkhdalai and Yu, 2017], whose general idea is to utilize a large number of episodes similar to target task to learn a meta learner that generate an initial segmentation model, and a base learner that quickly tunes the model with few samples. In most methods, a powerful feature extractor with good migration ability is provided for the meta learner to map the query and support images into a shared embedding space. And the base learner generates a category-specific descriptor with the support set. The similarity between the query's feature maps and the descriptor is measured under a parametric or nonparametric metric, and leveraged as guidance for dense prediction of the query branch.

However, when the low-level visual features of foreground objects extracted from the support images are too marginalized in the embedding space [Zhang et al., 2014], the generated descriptor by base learner is not category-specifically discriminative. In this case, the regions to be segmented in the query image may be ignored or even confused with other categories in the background. For example, as shown in Fig. 1, the bottle-specific descriptor generated from the support image is unable to identify the bottle itself in the same image, and therefore inapplicable for bottles in other complicated scenarios to be segmented.

To address the issue, this paper presents an adaptive tuning framework for few-shot segmentation, in which the marginalized distributions of latent features are dynamically adjusted by a self-supervised scheme. The core of the proposed framework is the base learner driven by a self-segmentation task, which augments category-specific constraint on each new episode and facilitate successive labeling. Specifically, the base learner is designed as an inner-loop task, where the support images are segmented under the supervision of the presented masks of input support images. By back-propagating self-supervised loss through the support feature map, the corresponding gradients are calculated and used as guidance for augmenting each element in the support embedding space.

Moreover, the self-segmentation task can be considered as a special case of few-shot segmentation with the query and support image to be the same. Therefore, we also utilize the resulting loss of the task (we call it auxiliary loss in this paper) to promote the training process. Since auxiliary loss is equivalent to performing data enhancement to annotated samples, it can involve more information for training with the same number of iterations. The evaluation is performed on the prevailing public benchmark datasets PASCAL-5i and COCO20i , and the results demonstrate above-par segmentation performance and generalization ability of our proposed method.

Our main contributions are summarized as follows:

1. An adaptive tuning framework is proposed for fewshot segmentation, in which the marginalized distributions of latent category features are dynamically adjusted by a selfsupervised scheme.
2. A novel base learner driven by a self-segmentation task is proposed, which augments category-specific feature description on each new episode, resulting in better performance on label prediction.
3. Experimental results on two public benchmark datasets PASCAL-5i and COCO-20i demonstrate the superiority of our proposed method over SOTA.

2 RelatedWork

Few-shot learning. Few-shot learning has recently received a lot of attention and substantial progress has been made based on meta-learning. Generally, these methods can be divided into three categories. Metric-based methods focus on the similarity metric function over the embeddings [Snell et al., 2017]. Model-based method mainly utilizes the internal architecture of the network (such as memory module [Santoro et al., 2016], etc.) to realize the rapid parameter adaptation in new categories. The optimization-based method aims at learning an update scheme for base learner [Munkhdalai and Yu, 2017] in each episode. In the latest study, [Lee et al., 2019] and [Bertinetto et al., 2018] introduce machine learning methods such as SVM and ridge regression into the inner loop of the base learner, and [Rusu et al., 2018] directly replaces the inner loop with an encoded-decode network. These methods have achieved state-of-the-art performance in fewshot classification task. Our model also takes inspiration of them.

Semantic segmentation. Semantic segmentation is an important task in computer vision and FCNs [Long et al., 2015] [Zhang et al., 2019b] have greatly promoted the development of the field. After that, DeepLabV3 [Chen et al., 2017] and PSPNet [Zhao et al., 2017] propose different global contextual modules, which pay more attention to the scale change and global information [Chen et al., 2019] in the segmentation process. [Zhu et al., 2019] considers the full-image dependencies from all pixels based on Non-local Networks, which shows superior performance in terms of reasoning [Qiao et al., 2019].

Few-shot semantic segmentation. While the work on fewshot learning is quite extensive, the research on few-shot segmentation [Hu et al., 2019] has been presented only recently. [Shaban et al., 2017] first proposes the definition and task of one-shot segmentation. Following this, [Rakelly et al., 2018] solves the problem with sparse pixel-wise annotations, and then extends their method to interactive image segmentation and video object segmentation. [Dong and Xing, 2018] generalizes the few-shot semantic segmentation problem from 1way (class) to N-way (classes). [Zhang et al., 2019a] introduces an attention mechanism to effectively fuse information from multiple support examples and proposes an iterative optimization module to refine the predicted results. [Nguyen and Todorovic, 2019] and [Wang et al., 2019] leverage the annotations of the support images as supervision in different ways. Different from [Tian et al., 2019] which employs the base learner directly with linear classifier, our method devises a novel self-supervised base learner which is more intuitive and effective for few-shot segmentation. Compared to [Nguyen and Todorovic, 2019], the support images are used for supervision in both training and test stages of our framework.

3 ProblemDescription

Here we define an input triple Tri = (Qi, Si, T iS ), a label T iQ and a relation function F : Ai = F (Qi, Si, T iS ; θ), where Qi and Si are the query and support images containing objects of the same class i, correspondingly. T iS and T iQ are the pixel-wise labels corresponding to the ith class objects in Si and Qi. Ai is the actual segmentation result, and θ is all parameters to be optimized in function F. Our task is to randomly sample triples from the dataset, train and optimize θ, thus minimizing the loss function L:

We expect that the relationship function F can segment object regions of the same class in another target image each time it sees few support images belonging to a new class. This is the embodiment of the meaning of few-shot segmentation. It should be mentioned that the classes sampled by the test set are not present in the training set, that is, Utrain ⋂ Utest = Ø. The relation function F in this problem is implemented by the model detailed in Sec. 4.3.

4 Method

4.1 Model Overview

Different from the existing methods, this paper proposes a novel self-supervised tuning framework for few shot segmentation, which is mainly composed of an adaptive tuning mechanism, a self-supervised base learner and a meta learner. These three components will be illustrated in the next three subsections. In this subsection, we mainly present the description of the whole framework mathematically, and the symbolic representation is consistent with that in Sec. 3.

As shown in Fig. 2, a Siamese network fe [Koch et al., 2015] is first proposed to extract the features of input query and support images [Zhang et al., 2012]. The mechanism of parameters sharing not only promotes optimization but also reduces the amount of calculation. In this step, the latent visual feature representations Rq and Rs are obtained as follows:

Here θe is the learnable parameter of the sharing encoder. To dynamically adjust the latent features, we firstly devise a novel self-supervised inner-loop as the base learner (fb) to exploit the underlying semantic information (θs).

Then the distribution of low-level visual features is tuned (ft) according to the high-level category-specific cues obtained above:

Inspired by Relation Network [Sung et al., 2018], a deep nonlinear metric fm is introduced into our meta learner. It measures the similarity between the feature map of query image and the tuned feature, and accordingly determines regions of interest in query images. Finally, a segmentation decoder fd is utilized to refine the response area to the original image size:

Similarly, θm and θd stand for the parameters of metric and decoding part. With the ability to continuously learn from different episodes, our meta learner gradually optimizes the whole process above (outer loop) and improves the performance of base learner, metric and decoder.

4.2 Self-Supervised Tuning Scheme

The self-supervised tuning scheme is the core of our proposed method, which is implemented by a base learner driven by self-segmentation and an adaptive tuning module. Different from the existing base learner used in [Lee et al., 2019] [Bertinetto et al., 2018], our proposed base learner is generated under the supervision of the presented masks of input support images. This method origins from an intuitive idea, that is, the premise of identifying the regions with the same category as target objects is to identify the object itself first. Therefore, we duplicate the features of the support image as the two inputs of Eq. 6 in Sec. 4.1 and calculate the standard cross-entropy loss (Lcross) with the corresponding support mask:

To exploit the above information to better implement the segmentation task of the query branch, the marginalized distribution is adjusted in embedding space based on the category-specific semantic constraint. Specifically, a gradient map that is calculated by back-propagating self-supervised

(a) Results for 1-shot segmentation.

(b) Results for 5-shot segmentation.

Table 1: Ablation study on PASCAL-5i dataset under the metric of mean-IoU. Bold fonts represent the best results.

loss through the support feature map is used as a guidance for augmenting each element in the support embedding space. Mathematically,

Note that only the feature representation is updated here and the network parameters are unchanged.

4.3 Deep Learnable Meta Learner

Inspired by the Relation Network, we apply a deep non-linear metric to measure the similarity between the feature map of query image and the descriptor generated by base learner. As in [Rakelly et al., 2018], we also use the deep learnable metric with late fusion as the main component of meta learner. First, we multiply the features by the downsampling mask and aggregate it to obtain the latent features of the foreground. This feature is then tiled to the original spatial scale, so that each dimension of the query feature is aligned with the representative feature (Rr s ):

Through the Relation Network comparator, the response area of the query image is obtained:

Finally, we feed it into the segmentation decoder, refining and restoring the original image size to get accurate segmentation results.

4.4 Loss and Generalization to 5-Shot Setting

In addition to the cross-entropy loss function of segmentation obtained by query-support set commonly used in other methods (main loss), we also include the cross-entropy loss of support-support segmentation from base learner (auxiliary loss) into the final training loss as an auxiliary. In our method, the auxiliary loss itself is created as an intermediate process, so there is no much extra computation required. It can be seen from the experiment part that the auxiliary loss can accelerate the convergence and improve the performance.

When generalizing to 5-shot segmentation, the main difference is the base learner part. Considering that the number of samples for 5-shot segmentation is still small, if we calculate the gradient optimization together, it is easy to produce large discrepancy. Therefore, we calculate the gradient value separately, getting 5 separate response areas, and then take the weighted summation according to the self-supervised scores to get the final result, that is:

where fIoU represents the function which is used to calculate the IoU scores.

5 Experiment

5.1 Dataset and Settings

Dataset. To evaluate the performance of our model, we experiment on PASCAL-5i and COCO-20i datasets. The former is first proposed in [Shaban et al., 2017] and is recognized as the standard dataset in the field of few-shot segmentation in subsequent work. That is, from the set of 20 classes on PASCAL dataset, we sample five and consider them as the test subsets Utesti = {4i +1, · · · 4i +5}, with i being the fold number (i = 0, 1, 2, 3), and the remaining 15 classes form the training set Utraini . COCO-20i dataset is proposed in recent work and the division is similar to PASCAL-5i dataset. In the test stage, we randomly sample 1000 pairs of images from the corresponding test subset.

Settings. The backbone of the existing methods are different, mainly VGG-16 [Simonyan and Zisserman, 2014] and ResNet-50 [He et al., 2016]. To make a fair comparison, we separately train two models with different backbones for testing. As adopted in [Shaban et al., 2017], we choose the perclass foreground Intersection-over-Union (IoU) and the average IoU over all classes (mean-IoU) as the main evaluation indicator of our task. While the foreground IoU and background IoU (FB-IoU) is a commonly used indicator in the field of binary segmentation, it is used by few papers of fewshot segmentation task. Because mean-IoU can better measure the overall performance of different classes and ignore the proportion of background pixels, we show the results of mean-IoU in all experiments. Our model uses the SGD optimizer during the training process. The initial learning rate is set to 0.0005 and the attenuation rate is set to 0.0005. The model stops training after 200 epochs. All images are resized to 321 × 321 size and the batch size is set to 16.

5.2 Ablation Study

To prove the effectiveness of our architecture, We conduct several ablation experiments on PASCAL-5i dataset as shown in Table 1. To improve efficiency, we only choose ResNet50 as backbone for all models in the ablation study for fair comparison. The performance of our network is mainly attributed to two prominent components: SSM and auxiliary loss. Note that we progressively add addition components to the baseline, which enables us to gauge the performance improvement obtained by each of them. Due to the fact that the self-supervised tuning mechanism dynamically adjusts the marginalized distributions of latent features at each stage, SSM brings about a 1.2 and 2.9 percent mean-IoU increase in 1-shot and 5-shot settings, respectively. At the same time, we can see that auxiliary loss boosts the overall performance, resulting in a 2.3 and 3.8 percent improvement.

5.3 Analysis

As PASCAL-5i is the most commonly used dataset by all few-shot segmentation methods, the main analysis parts of our experiment are accomplished on PASCAL-5i dataset.

Effect of the Self-Supervised Module

To clarify the actual function of the self-supervised module for few-shot segmentation task, we show the following visualization results. As shown in Fig. 3, we feed the feature representations generated before and after the tuning process to Relation Network to obtain two sets of segmentation results. Note that the original network often segments other objects by mistakes that are easily confused in the background. After adding the SSM module, the category-specific semantic constraint is introduced to help query images to correct the prediction results. To demonstrate the meaning of this semantic constraint, the calculated gradient is visualized in the last column. We can see that it strengthens the focus on the regions of target categories.

Importance of Auxiliary Loss

To verify the role of auxiliary loss in the training stage, we designed the following experiment. First, we show the meanIoU curves generated with and without our auxiliary loss during test stage in Fig. 4. It can be clearly seen that the convergence is faster and the mean-IoU result is better after the loss is added. Then, we visualize the segmentation results of two models trained with and without auxiliary loss. As shown in Fig. 5, the model without auxiliary loss can not segment the support images themselves, therefore inapplicable for other images of the same category. The proposed auxiliary loss improves the self-supervised capability as well as the performance of few-shot segmentation.

Comparison among Different 5-Shot Fusion Methods

To prove the superiority of our weighted fusion in 5-shot setting, we compare the 5-shot segmentation results with the average and maximum fusion methods, in which the average and maximum segmentation results of 5 support images are computed, respectively. It can be seen in Table 2 and Fig. 6 that our weighted fusion strategy achieves the best. The samples in the PASCAL-5i dataset are relatively simple, and most of the weights are close to one fifth, so the result of average fusion method is similar to weighted fusion in general.

5.4 Comparison with SOTA

To better assess the overall performance of our network, we compare it to other methods (OSLSM [Shaban et al., 2017], SG-One [Zhang et al., 2018], PAnet [Wang et al., 2019], FWBFS [Nguyen and Todorovic, 2019] and CAnet [Zhang et al., 2019a]) on PASCAL-5i and COCO-20i datasets.

PASCAL-5i

We train two types of SST (Self-Supervised Tuning) models with VGG-16 and ResNet-50 backbones (we call them SSTvgg and SST-res model) on PASCAL-5i dataset. It can be seen in Table 3 that our SST-vgg model surpasses the best existing method over two percentage points in 1-shot setting, and the SST-res model yields 2.2 points improvement, which is even 1.4 points higher than the method with ResNet101 backbone. Under the setting of 5-shot, our SST-vgg and SST-res models significantly increase by 1.9 and 1.5 points as shown in Table 4, respectively. These comparisons indicate that our method boosts the recognition performance of few-shot segmentation.

COCO-20i

To prove that our method also has good generalization performance on larger datasets, we compare our method with others which recently report results on the COCO-20i dataset as shown in Table 5. Obviously, the average results of our method surpass other best methods by 1 and 1.6 points under 1-shot and 5-shot setting, respectively.

6 Conclusion

In this paper, a self-supervised tuning framework is proposed for few-shot segmentation. The category-specific semantic constraint is provided by the self-supervised inner loop and utilized to adjust the distribution of latent features across different episodes. The resulting auxiliary loss is also introduced into the outer loop of training process, achieving faster convergence and higher scores. Extensive experiments on benchmarks show that our model is superior in both performance and adaptability compared with existing methods.

Acknowledgments

This work was supported by the National Key R&D Program of China under Grant 2017YFB130092, the National Natural Science Foundation of China (NSFC) under Grants 61872327, the Fundamental Research Funds for the Central Universities under Grant WK2380000001 as well as Huawei Technologies Co., Ltd.