SimPropNet: Improved Similarity Propagation for Few-shot Image Segmentation

1 Introduction

Semantic image segmentation assigns class labels to image pixels. It finds applications in image editing [1, 20, 21], medical diagnosis [8, 13, 18], automated driving [7] etc. Supervised deep neural network methods such as [2, 3, 4, 5, 6, 23, 27] enable highly accurate image segmentation. However, they work for only a small number of fixed object classes, and require a large number of image-mask pairs for training which are hard to manually annotate. In several practical scenarios, including online commerce and design, the images may exist in a large number of sparsely populated classes (for instance, images of products). In such cases obtaining image-mask pairs for all possible classes to train a supervised method may be infeasible. Thus, segmentation methods that generalize to new classes with scant training data are of significance. Few-shot image segmentation methods, like [15, 24, 26], are class-agnostic and alleviate the need for a large number of image-mask pairs. These methods utilize additional images and their masks of a particular class in predicting the binary segmentation mask for a given image of the same class. This work proposes a new few-shot segmentation framework that seeks to alleviate a fundamental limitation in existing techniques and significantly improves upon the stateof-the-art.

∗work done as part of Adobe MDSR internship program

† equal contribution

Few-shot segmentation (FSS) methods typically take as input - a query image for which the segmentation mask is required, a set of support images, and their segmentation masks (support masks). One-shot segmentation is the extreme setting where only a single support image-mask pair is available. Our method achieves state-of-the-art performance in both one-shot and few-shot settings (about 5% and 4% gain respectively).

Recent deep neural network based FSS methods (like [16, 24, 25, 26]) employ variations of the following broad process (Figure 1):

1. Query and support image features are extracted using

• shared, pre-trained (on ImageNet [14]) network layers.
• A probe representation for regions of attention in the support image is obtained from the support features and mask(s).
• The attention representation is fused with the query features.
• The fused features are decoded to obtain the querymask. The attention extraction and fusion modules leverage the high-dimensional feature similarity between the query and the support images to selectively decode the query features to output the segmentation mask. The focus of this work is to demonstrate gaps in the propagation of this similarity information in existing methods and to bridge those gaps.

Experiments with two state-of-the-art methods - [16] and [26] (see Section 3.1) reveal that FSS methods make errors in visually similar regions, across image pair samples of identical classes. They perform poorly when support is identical to query inputs as well. These results indicate that the class and visual similarity information is not propagated optimally across the support and query branches .

We predict the support mask from the support features to endow the features with specificity with respect to the target class, which in turn aids in similarity propagation between the support and query. We also leverage the similarities in the background regions of the support and query images through a background attention vector computed from the support features and inverse mask, and fuse it with the query features. Finally, to prevent the network from overfitting on training class-conditional similarities, we employ an input channel averaging input augmentation for the query input. With these improvements we achieve state-of-the-art performance on PASCAL 5i dataset for both one-shot and five-shot segmentation tasks.

The contributions of this work can be summarized as follows:

1. It highlights gaps in existing FSSmethods in fully utilizing the similarity between the support and query images.
2. It introduces SimPropNet, a few-shot segmentation framework that bridges those gaps and achieves state-ofthe-art performance on the PASCAL 5i dataset in both 1-shot and 5-shot evaluation settings. The framework employs:
3. A dual-prediction scheme (DPr) where the query and support masks are jointly predicted using a shared decoder, which aids in similarity propagation between the query and support features.
4. A novel foreground-background attentive fusion (FBAF) mechanism for utilizing cues from background feature similarities between the query and support images.

The next section places the proposed method in the context of previous work related to the problem.

2 RelatedWork

Recent methods for few-shot semantic segmentation (FSS) build a framework for one-shot segmentation and subsequently construct methods to use the framework for k-shot segmentation (k > 1). The proposed work follows the same methodology.

As described in Section 1, most FSS methods employ a dual branched neural network model with a support branch and a query branch. This model was introduced by Shaban et al. [15] where the support branch is conditioned on the support input to predict the weights of the last layer in the query branch which then predicts the query mask.

Rakelly et al. [12] improve upon [15] by employing a late fusion strategy for the support mask and support feature maps to segment the query image. Late fusion adopts a parametric module in a two branch setting, which fuses information extracted from the support set with the features of the query image to produce the segmentation mask. Dong and Xing [9] combine the late fusion methodology with the idea of prototypical networks introduced in [19] to learn a metric space where distances to prototypes represent class-level similarity between images. Their method uses the prototypes as guidance to fix the query set segmentation rather than obtaining segmentation directly from metric learning. Zhang et al. [25] take a different approach to late fusion. They introduce masked average pooling (MAP) that pools support features of the regions of interest in the support images, and fuses them with the query features using vector cosine-similarity. This attention map is then decoded to predict the query's segmentation maps. Building on [9], Wang et al. [24] combine Prototypical Learning [19], and MAP [25] to incorporate the support information into the segmentation process. Zhang et al. [26] adopt a learnable method through an attention mechanism to fuse support information from multiple support examples along with iterative refinement. Their method also uses MAP, but instead of fusing its output with the query features using cosine-similarity, they concatenate it with the query features and then decode the output.

These methods for few-shot segmentation depend on the support set to produce the segmentation for the query image, but fail to fully utilize the support information efficiently. We present evidence for this gap in Section 3.1 and subsequently propose SimPropNet, a few-shot segmentation framework that seeks to bridge this gap and improve performance.

3 Method

In this section, we first establish firm ground to motivate our work, and then elucidate the proposed method in the one-shot segmentation setting. Subsequently, we demonstrate how we adapt the method for k-shot segmentation setting.

3.1 Premise Validation

We first present experimental evidence validating the premise that there are gaps in similarity propagation between the support and the query images. This argument is expressed as follows:

1. FSS methods make errors for image regions that may not be inherently hard to segment.
2. The regions of error in the support and query images are class-conditionally similar.
3. Even with maximally similar support and query images, current FSS methods are unable to produce good predictions.

We conduct experiments to corroborate this argument which indicates an opportunity to improve similarity sharing between the support and query images.

We also provide evidence to show that the background regions of the support and query images may be similar and could hold cues for improved segmentation of the query. The experiments use author-provided implementations of the recent state-of-the-art methods [16] and [26].

Errors by FSS versus Supervised Methods We measure the overlap (Table 1) between the error regions of output masks from two FSS methods ([16] and [26]) with the error regions of masks produced by DeepLab v3+ [6] (a state-ofthe-art supervised method). The overlap between the error regions for FSS methods and DeepLab v3+ is very small, while the gap in the mean-IoU for correct predictions (DeepLab v3+ versus FSS methods) is large. This indicates that there are image regions where supervised methods like DeepLab v3+ do not fail (but FSS methods do), and that they are not characteristically difficult to segment.

Similarity of mispredicted regions To determine the similarity between regions of the query and support images where [26] makes errors, we compute the Masked Average Pooling vectors (as described in [25]) using pre-trained VGG-16 [17] features. This is done for several pairs of images of identical classes. for the following two regions of each image:

1. the pixels covered by the ground-truth masks (Z g ), and
2. the pixels present in the regions mispredicted by the FSS method (Z e ).

For each image pair (A,B) of a class, the ratio of the innerproducts of the MAP vectors for the error regions and the ground-truth mask regions - (Z e (A)·Z e (B))/Z g (A)·Z g (B) - is a measure of the similarity of the corresponding error regions relative to the similarity of the ground-truth mask regions. We measure the ratio over a 1000 pairs of images from the PASCAL VOC dataset [10] and find the average to be 0.87 (std. dev. = 0.25). The high value of relative similarity of error regions substantiates the claim that errors are committed in regions of similarity that could have possibly been propagated from the support to the query.

Identical Inputs: (Support = Query) Table 2 reports results from a third experiment with [16] and [26] where the same image is given as input for both support and query (including the ground-truth mask for the support). This constitutes a basic test for similarity propagation in the network. Ideally, the network must produce the exact mask as provided in the support input, because the query and the support are as similar as possible. However, both methods ([16] and [26]) perform poorly for these inputs indicating the loss of similarity information in the networks.

Cues from Background Similarity Figure 3 indicates the degree of similarity measured as the cosine-similarity between the MAP vectors of foreground and background regions for the pairs of images. The figure indicates that images of identical object classes have higher degree of feature similarity in their background features, than in the actual foreground features. This represents an opportunity to collect more information for accurate segmentation.

The results of these experiments indicate the gaps in similarity propagation in existing FSS methods that we propose to exploit. Next, we describe the organization of the proposed network and how it addresses the similarity propagation problem.

3.2 Network Organization

Figure 2 illustrates the architecture of SimPropNet. The network is organized into two branches (support and query) with a shared encoder, a fusion module and a shared decoder. For our experiments, we use a ResNet-50 [11] based feature extractor, and atrous spatial pyramid pooling based decoder as in [26]. The encoder comprises of three layers from a ResNet50 network pre-trained on ImageNet [14], and a single 3 × 3 dilated (rate = 2) convolutional layer with 256 filters. The ResNet-50 layers are frozen during training. The decoder comprises of an atrous spatial pyramid pooling layer (as introduced in [4]), and two convolutional layers. The last layer has linear activation and produces a 2-channel output (1/8)th the input size. The output from the last layer is resized to the expected mask size with bilinear interpolation. The predicted query and support masks are compared to their respective ground-truths using the cross-entropy loss.

3.3 Support and QueryMask Prediction (DPr)

The network is trained to predict both the support and the query masks using the same encoder and decoder. We submit that this dual prediction requirement forces the query and support features from the shared encoder to share greater and more target-specific similarity. For instance, if the support image (and mask) is of a car, the network must be able to recover back the support mask (i.e. the entirety of the car) from the foreground support features. Because the encoder is shared between the support and the query images, the query features share this characteristic with the support features.

This is reflected in the effective gain in the mIoU as discussed in Section 5.1.

3.4 Foreground/Background Attentive Fusion (FBAF)

Late fusion methods (like [12, 24, 25, 26]) fuse the foreground features from the support branch into the query branch to locate the regions of class-conditional similarity. The foreground features are obtained using the masked average pooling (MAP) operation that computes a channel-wise weighted average of the support features where the weights are the support mask values at each position. We submit that fusing the background feature information from the support branch has the effect of suppressing similar background features in the query features. The fusion process can be concisely stated using the following three equations, each representing a step of the process:

1. Dual Probe Feature Extraction: The foreground and the background MAP vectors are computed using the support mask and the support features.

Here, µc is the average pooling operation with a kernel size equal to the feature map size, Fs are the support features, and ˘Ms is the support mask down-sampled to the height and width of Fs.

1. FG/BGAttentionMap Generation: Four attentionmaps, two (foreground and background) each from the support and query features are computed.

Here cossim is the cosine-similarity measure and Fq are the query features.

1. Feature and Attention Map Fusion: The attention maps and the features for the query and support are fused to two feature vectors that are finally decoded into the query and support mask predictions respectively.

Here G∗ is for both support (Gs) and query (Gq) features, and ⊕ is the concatenation operation and IN is the Instance-Normalization operation [22].

Because the fusion process combines information from both the foreground and the background support features, we term this FG/BG Attentive fusion (FBAF). The analysis in Section 5.1 demonstrates the effective increase in prediction accuracy by employing FBAF.

3.5 k-Shot Inference

The network is not specifically trained for k-shot training (k > 1). To incorporate more than one support example pairs in inference, the MAP vectors computed in the Dual Probe Feature Extraction step are averaged over the support pairs:

These MAP vectors are used to compute the foreground and background attention maps, and are fused with the query features (Section 3.4) to predict the query segmentation mask.

3.6 Implementation Details

The training is done on virtual machines on Amazon AWS with four Tesla V100 GPUs and 16-core Xeon E5 processors. The standard SGD optimization algorithm is used, learning rate is kept at (2.5×10-3 ) and the batch size is 8 for all training. Training with still higher learning rate yields even better results than reported in the paper, but the training is not always stable and may decay considerably in later training steps. Training for each split is run for 180 epochs and the checkpoint with the highest validation score (mIoU) is retained. To prevent the network from overfitting on the training classes, we also use an input regularization technique called Input Channel Averging (ICA) where the query RGB image is mapped to a grayscale input (after normalizing) with a switch probability (initialized at 0.25 for our experiments) that decays exponentially as training progresses. The particular benefit of using ICA is discussed in Section 5.1.

4 Results

In this section, we present the evaluation metrics and reporting criteria for few-shot segmentation, the quantitative results for SimPropNet, and comparison with other state-of-the-art methods. For brevity, we include very few qualitative results highlighting the key benefits our work.

4.1 Metric and Reporting

For consistency of comparison with baseline techniques, we follow the evaluation protocol established by Shaban et al. [15]. Different instances of the network are trained with the different training splits of the PASCAL-5i dataset, each with 15 classes of objects and their masks. Testing is done on the images with objects of the 5 withheld classes. The mean intersection-over-union (mIoU) metric for output masks with respect to the ground-truth mask is computed. We report the mean IoU values for random support and query pairs (images of the same class) in the standard 1-shot and 5-shot (k = 5 support images) settings. To ensure fairness in comparison, we use author-provided code for baselines and evaluate performance for three sets of random support-query test pairs for each of the splits and report the average performance. These splits will be released along with the paper upon acceptance.

4.2 SimPropNet Results

Quantitative Results We compare our performance using the method described in [15] that employs random query and support image pairs (partitions) from each test class in PASCAL-5i splits. The mean intersection-over-union (mIoU) values are reported in Table 3. Our method, SimPropNet, achieves state-of-the-art performance in both the 1-shot and the 5-shot setting. SimPropNet outperforms CANet [26], the current state-of-the-art in 1-shot setting, by 5.44% in 1-shot and 6.23% in the 5-shot evaluation. Further, SimPropNet outperforms PANet [24], the current state-of-the-art in 5-shot setting, by 9.06% in the 1-shot evaluation and 4.34% in the 5-shot evaluation task.

Qualitative Results Qualitative one-shot segmentation results also indicate significant improvements in output. Figure 4 highlights how providing a more similar support image and mask helps in improving the segmentation, and is expected in a practical scenario. Figure 5 presents two sample results for comparison with CANet [26]. The first row illustrates how CANet may overfit on training classes (target test class (person) is from split-2, and bicycle is in the training set) or its features lack specificity to the target class, and how SimPropNet overcomes this issue. Row 2 of the figure illustrates the benefit of providing support images with similar backgrounds.

5 Analysis

In this section, we first analyse the particular benefit of each component contribution of SimPropNet individually. Next, we probe the effectiveness of SimPropNet in improving similarity propagation. We present these evaluations in the oneshot setting and compare against CANet [26], the existing state-of-the-art.

5.1 Ablation Study of Components

We study the effectiveness of each of our contributions individually - dual prediction (DPr) of the support and query masks, and foreground-background attentive fusion (FBAF). Table 4 reports the mIoU values over the different splits of the PASCAL 5i dataset. As highlighted by the results, both DPr and FBAF used individually demonstrate clear gains over the baseline (CANet [26]) of 4.14% and 3.75% in mIoU respectively. FBAF alone performs better than DPr in three of the four splits, but has a slightly worse mean because of its sharp decline in performance in split-3. The combination of DPr and FBAF achieves an improvement of 5.34% over [26]. Additionally using ICA during training further improves mIoU on three of the splits and increases mean mIoU to 57.19%.

5.2 Measuring Gain in Similarity Propagation

We evaluate the performance of the proposed network on identical support and query images (as reported for [26] and [16] in Table 2 in Section 3.1). The performance of a oneshot segmentation method on identical inputs for the query and the support is the upper bound on the degree of similarity propagation in the network. Results of this experiment for SimPropNet (Table 5) show an average gain of 21.5% over CANet [26], and 14.23% over AMP [16], over all splits. This indicates clearly that the network is utilizing the similarity information between the query and support images better.

6 Conclusions and FutureWork

The paper presents a rigorous argument that similarity propagation in existing few-shot image segmentation networks is sub-optimal. It proposes SimPropNet, a deep neural network with two strategies for bridging this gap - predicting the support mask besides the query mask with a shared decoder, and fusing background features into the feature fusion module. The network achieves a new state-of-the-art as shown by a comprehensive set of experiments. Class-conditional similarity matching can only match pixels with a similar class-mix between the query and the support images. In future work, we focus on exploring the "objectness" aspect of the target image to improve few-shot segmentation.