Attentional Prototype Inference for Few-Shot Segmentation

We introduce a latent variable $\mathbf{z}$ to represent the class prototype, which is conditioned on the corresponding support set. By incorporating the latent variable, we have the conditional predictive log-likelihood as follows

where $p(\mathbf{z}|\mathcal{S})$ is a conditional prior. The model in (1) provides a probabilistic modeling of prototypes for semantic segmentation, which was introduced in our preliminary work [28]. In this way, our prototype serves as a global representation of an object category while previous ones do not take into account the local spatial structure of the image [9].

To further enhance the model in (1), we introduce a local latent variable $\mathbf{m}_{q}$ to represent the attention map associated with each image, which highlights the foreground object. The conditional predictive log-likelihood with respect to the two latent variables takes the following form

where we also deploy a conditional prior $p(\mathbf{m}_{q}|\mathbf{x}_{q})$ for $\mathbf{m}_{q}$, since the attention maps should be specific for each individual query image $\mathbf{x}_{q}$.

However, these posteriors are intractable in practice. Thus, we introduce the variational posterior to approximate the true posteriors by minimizing their Kullback-Leibler (KL) divergence. We employ the variational distributions $q_{\phi_{1}}(\mathbf{z}|\mathbf{x}_{q},\mathbf{y}_{q},\mathcal{S})$ and $q_{\phi_{2}}(\mathbf{m}_{q}|\mathbf{x}_{q},\mathbf{y}_{q})$ for the prototype $\mathbf{z}$ and the attention map $\mathbf{m}_{q}$, respectively. By incorporating the variational posteriors into the conditional log-likelihood $\log p(\mathbf{y}_{q}|\mathbf{x}_{q},S)$ of (2), we arrive at

Applying Jensen’s inequality gives rise to the ELBO as follows

The first term of the ELBO is the expectation of the log-likelihood of the conditional generative distribution $p(\mathbf{y}_{q}|\mathbf{z},\mathbf{m}_{q},\mathbf{x}_{q})$ based on the inferred prototypes $\bf{z}$ and attention maps $\mathbf{m}_{q}$. The second term is the KL divergence between the estimated posterior distribution $q_{\phi_{1}}(\mathbf{z}|\mathbf{x}_{q},\mathbf{y}_{q},\mathcal{S})$ and the prior distribution $p(\mathbf{z}|\mathcal{S})$. Minimizing this term encourages the model to leverage the object information for the segmentation of the query image. Minimizing the third term of the KL divergence enables the model to generate attention maps that highlight the foreground object. We derive the optimization objective based on the ELBO.

Maximizing the ELBO can yield accurate predictions for the segmentation masks and narrow the gap between the posterior and the prior distributions. This encourages 1) the inferred prototype from the support dataset to match the full dataset by minimizing the first KL term; and 2) the inferred map from the query set to approach the one based merely on the query image by minimizing the second KL term. Based on the ELBO, we define the empirical objective function for optimization.

Given a batch of episodes, the empirical objective for stochastic optimization with the Monte Carlo estimation of expectations is as follows

where $i$ indexes over the sampled episode in the meta-training set $\mathcal{D}_{\text{train}}$, $\mathbf{z}^{(l)}\sim q_{\phi_{1}}(\mathbf{z}|\mathbf{x}^{(i)}_{q},\mathbf{y}^{(i)}_{q},\mathcal{S}^{(i)})$ and $\mathbf{m}_{q}^{(j)}\sim q_{\phi_{2}}(\mathbf{m}_{q}|\mathbf{x}^{i}_{q},\mathbf{y}^{i}_{q})$ are the variables sampled from their variational distribution. $L,M$ are the number of samples. Generally, the parameters of a neural network are optimized jointly with stochastic gradient descent. However, it is usually intractable to calculate the gradient of the sampling operation. Therefore, we deploy the reparameterization trick [24] to handle the non-differentiable problem of the sampling process. Specifically, supposing the two variational posterior distributions take the form of a multivariate Gaussian with a diagonal covariance, the sampling process is formulated as

During training, the samples of the class prototype $\mathbf{z}$ are obtained by:

where $\odot$ denotes an element-wise multiplication and $\bm{\epsilon}^{(l)}\sim N(\bm{\epsilon};\bm{0},\bm{1})$. The same operation is also deployed for sampling $\mathbf{m}_{q}^{(j)}$.

The first term of the empirical loss $\mathcal{L}$ in (5) is implemented as a least square loss in [24]; we generalize this to a pixel-wise cross-entropy loss to penalize the difference between the predicted segmentation map $\tilde{\mathbf{y}}_{q}$ and the ground truth $\mathbf{y}_{q}$. The number of samples $L$ and $M$ are set to 1 during training to speed up the learning process. Since the KL terms minimize the discrepancy between two distributions, the prior networks can mimic the behavior of the posterior networks that produce effective prototypes or attention maps at training time.

The inference of segmentation maps varies between the learning and inference stages. At test time, instead of sampling from the variational posterior distributions, we draw $L$ samples of prototypes $\{\mathbf{z}^{(l)}\}_{l=1}^{L}$ from the prior $p_{\theta_{1}}(\mathbf{z}|\mathcal{S})$ and $M$ samples of attention vectors from $p(\mathbf{m}_{q}|\mathbf{x}_{q})$. $\tilde{\mathbf{y}}_{q}$ is obtained by taking the average of $LM$ segmentation maps from these samples

where

and

PASCAL-${5}^{i}$ originates from PASCAL VOC12 and extends annotations. We follow the settings in [17], splitting the 20 original classes into four folds and conducting cross-validation among them. Specifically, we select 15 classes for training, while the remaining 5 classes are for testing. For a fair comparison, we adopt the same strategy as [17], randomly sampling 1,000 episodes of support-query pairs for evaluation.

COCO-${20}^{i}$ is a challenging dataset built upon MS-COCO with 80 object categories. We also divide the 80 classes in MS-COCO into four folds and conduct four-fold cross-validation. Under the same settings as PASCAL-${5}^{i}$, 60 object categories are selected for training, while the remaining 20 categories are used for testing. In each fold, we sample 1000 support-query pairs from the 20 testing classes for evaluation, following [17].

FSS-1000 is a specialized few-shot segmentation dataset. It contains 1,000 object categories including 520 classes for training, 240 classes for validation, and 240 classes for testing. Following [35], we choose the same 240 categories for testing and train the model on the specified 520 classes with the support of the validation set. The number of testing episodes is 1,000.

LIDC-IDRI is a dataset to represent typical ambiguities in vision signals. It contains 1,018 thoracic CT cases from 1,010 lung patients. The task is to segment lesions of lung nodules given a CT slice. Each CT case has been annotated by 4 individual radiologists that provide segmentation masks for independently detected lesions. We obtain 2,630 sub-cases by extracting a series of slices from 3D volumes and cropping 2D slices to 128 $\times$ 128 pixels centered at the lesion positions. Each sub-case contains a set of 2D images with different cardinalities. Each image corresponds to four annotated masks from independent radiologists (See Figure 1). For the setting of few-shot segmentation, we consider a sub-case as one episode, i.e., randomly sampling support and query images from one sub-case without replacement. For a total of 2,630 sub-cases, we choose the first 2,030 sub-cases for training and testing the model on the last 300 sub-cases. The remaining 300 sub-cases are for validation. The total numbers of images for training, testing, and validation are 11,631, 1,943, and 1,974.

We adopt a ResNet101 backbone pre-trained on ImageNet as the encoder. The decoder is designed as a skip-connection structure [30], which is composed of three convolutional blocks to generate segmentation maps. Each block receives the input of the concatenation with the corresponding encoded feature through the skip connections and the decoding features. The structure of the decoder, prior, and posterior networks are listed in Figures 3, 4, and 5. We choose Adam as the optimizer and train the model on four NVIDIA Tesla V100 with around 30 epochs. The learning rate is fixed to $5e-7$ for the backbone and $5e-5$ for other layers, and the batch normalization (BN) layers are frozen during training. The numbers of samples $L$ and $M$ are set to 10 during the test phase, which is analyzed in detail by our ablation study in Section 4.3.1.

We adopt the same metrics as [14] for evaluation, i.e. Class-IoU (C-IoU) and Binary-IoU (B-IoU). Class-IoU measures the intersection-over-union

for each class, where TP, FP and FN are the number of pixels that are true positives, false positives and false negatives of the predicted segmentation masks for each foreground category $c$. Binary-IoU measures the IoU between the foreground and background pixels, where all object classes are treated as foreground.

We generalize the IoU metric to the cross energy distance (CED) on the LIDC-IDRI dataset to measure the distance between two distributions. The core idea of CED is leveraging distances between observations. Let $\mathbf{y},\hat{\mathbf{y}}$ denote the annotation and the prediction, respectively. $d(\cdot)$ is the distance between two observations. For $m$ independent annotation samples and $n$ prediction samples, we have

Here, we choose $d(\mathbf{y},\hat{\mathbf{y}})=1-\text{IoU}(\mathbf{y},\hat{\mathbf{y}})$. A smaller CED value indicates that the prediction distribution is close to the ground truth distribution and can characterize the ambiguity that appears in images. In our LIDC-IDRI experiments, we set $m=4$ since there are four annotation masks for one image.

In Table 1, we compare the performance of API with typical methods on PASCAL-${5}^{i}$ in terms of the Class-IoU metric and Binary-IoU. API outperforms those prototype-based methods by good margins under both the 1-shot and 5-shot settings (57.4%, 60.7%). The 1-shot setting is more challenging than the 5-shot setting due to the much larger intra-class variation. The Monte Carlo estimation in our probabilistic model serves as an ensemble of the prediction results. Specifically, API outperforms PMM which computes segmentation masks with multiple prototypes. This accounts for the robustness of our API in the 1-shot case. We also evaluate the model in terms of Binary-IoU. Our model again yields comparable performance under both the 1-shot and 5-shot settings of $70.3\%$ and $72.1\%$.

However, due to the intrinsic frailty of prototype-based methods, API has limited capability of handling objects with small sizes or complex shapes. Recall the computation of the prototype is from feature maps at the high level, discriminative information of objects with small sizes may vanish as downsampling operations, especially for small objects with complex backgrounds. This challenge also exists when a large object has complex shapes, i.e.some parts of the object, such as table legs, might be omitted. Hierarchical prototypes would be potential to tackle it as we discussed in Conclusion. Therefore, we remove one class with the lowest confidence in each fold and demote it as API*. As shown in the bottom part of Table 1, API* can achieve considerable performance compared with the state-of-the-art, e.g., $64.7\%$ for the setting of 1-shot. Compared with other methods, our prototype-based model is more efficient and easy to be implemented. Because of the advanced architecture of transformers, VAT [22] obtained the highest IoU metric for most cases on PASCAL-${5}^{i}$.

Some qualitative results on PASCAL-${5}^{i}$ are visualized in Figure 6. The proposed API is capable of producing more accurate segmentation under various challenging scenarios, where the query images vary in appearance and object size from the associated support images. For instance, in the fifth column, the size and viewpoint of the bus in the query image is significantly different from the annotated plane in the support image; in the eighth column, the annotated boat in the support image is much smaller than the one in the query image.

COCO-${20}^{i}$ is more challenging than PASCAL-${5}^{i}$ since the scenes in COCO-${20}^{i}$ are more complex with more intra-class diversity. Therefore, few-shot segmentation on COCO-${20}^{i}$ has more ambiguity and it is difficult to acquire an effective class-specific prototype. As can be seen in Table 2, our method outperforms the homogeneous method PMM [17] by $5.7\%$ and $5.5\%$ in terms of Class-IoU under the 1-shot and 5-shot settings. Since the COCO-${20}^{i}$ dataset contains more complex objects than PASCAL-${5}^{i}$, we drop five classes with the lowest confidence in each fold. As shown in Table 2, we obtain reasonable results on the rest of 15 classes in each fold, e. g., $41.8\%$ of Class-IoU. As DGPNet [23] applied the powerful probabilistic model of Gaussian Process and introduced extra mask encoders, it non-trivially outperforms those recent works. The qualitative results are provided in Figure 7. Our method successfully predicts the segmentation maps for query images, though the objects in the query image are significantly different from those in the support images in terms of appearance, size, and viewpoints. As shown in the third column, the object can also be successfully segmented even with serve occlusions.

We evaluate our method following the official evaluation protocols in [35]. The evaluation metric used for FSS-1000 is the IoU of positive labels in a binary segmentation map. Since this dataset contains the validation set, we select the prediction model with the support of the validation set (i. e., the model for the best validation class-IoU). Then, we report the result on the testing set. A performance comparison with other models in terms of Positive-IoU (P-IoU) is provided in Table 3. Our method improves over the state-of-the-art set by Wei et al. [35] by $12.1\%$ and $7.9\%$ in the 1-shot and 5-shot settings, demonstrating the effectiveness of our proposal for few-shot segmentation across large-scale semantic categories. Figure 8 visualizes segmentation results on the FSS-1000 dataset, where API produces accurate segmentation maps close to the ground truth. As shown in Figure 8, the foreground object is usually located at the center of the image and distinct from the background. This could lead to considerable performance for our model on FSS-1000.

To study the capability of handling ambiguities in few-shot segmentation, we conduct experiments on LIDC-IDRI with typical image ambiguities. The evaluation metric for LIDC-IDRI is the cross energy distance (CED) rather than IoU values. We compared our API with two state-of-the-art deterministic models, i.e., FS-PARN [19] and HSNet [21]. Since it is unfeasible to directly obtain prediction samples, we adopt an ensemble strategy that trains a set of predictors. Considering there are 4 annotations, we initialize 15 models and select different combinations of annotations to train the model. Thus, the total number of those combinations is $C_{4}^{1}+C_{4}^{2}+C_{4}^{3}+C_{4}^{4}=15$. For each model, we train it with the data batch that randomly sampled annotation masks from the current combination. Once 15 predictors are obtained, we randomly select a subset of predictors to compute multiple segmentation maps for testing. As for our API model, we randomly choose one of 4 annotation masks at each update step. In this experiment, the sampling number $n$ in CED (Eq. 16) is set as 9 for all methods. We train the model on the training set of 2,030 sub-cases with the support of the validation set of 300 sub-cases. The results shown in Table 4 are CED values on the testing set. We test the model five times with different predictors selected randomly and report the average value for evaluation. Our API achieves non-trivial improvement in terms of CED compared with state-of-the-art methods. In contrast to evaluating the method with the class-IoU metric, the results of CED include all prediction samples and annotation masks, demonstrating the good statistical property of our proposed framework.

Different from previous deterministic models by learning a deterministic prototype vector, our probabilistic methods infers the distributions of the class prototype and the attention vector for each image. To demonstrate the advantage of the proposed probabilistic modeling, we implement a deterministic counterpart. We utilize the same network architecture for fair comparison and predict a deterministic class prototype vector or an attention map by the $\bm{\mu}$ branch and remove the KL divergence term during training. We implement both models with a VGG-16 , ResNet50, and ResNet101 backbone, which are commonly adopted in previous works [32, 37].

The results on PASCAL-$5^{i}$ are shown in Table 5. Our attentional prototype inference consistently achieves better performance than the deterministic models under both the $1$ and $5$-shot settings in terms of Class-IoU and the Binary-IoU metrics, because the proposed probabilistic modeling of prototypes and attention maps is more expressive of object classes and has a compelling capability of capturing the categorical concepts of objects. Therefore, the learned model is endowed with a stronger generalization ability for query images that usually exhibits large variations. The results also illustrate the advantage of probabilistic modeling for few-shot segmentation. As the ResNet101 backbone outperforms both VGG16 and ResNet50, we adopt ResNet101 as the backbone network in our experiments.

The newly introduced attention mechanism faithfully leverages the structure information that is neglected by the prototype, which is essential for mask prediction. The transformer architecture adopted in our model can utilize the context knowledge from pixels with high fidelity and enhances the foreground to achieve accurate prediction. We conduct extensive comparison experiments with the baseline variational prototype inference (VPI) [28] to show the effectiveness of the latent attention mechanism. As shown in Table 6, we evaluate their performance on three benchmarks of PASCAL-${5}^{i}$, COCO-${20}^{i}$ and FSS-1000 with three metrics of Class-IoU, Binary-IoU, and Positive IoU. API achieves considerable improvement compared with VPI in most cases. In particular, API improves VPI by $7.8\%$ under the 1-shot, and $8.1\%$ under the 5-shot settings for Class-IoU on COCO-20ⁱ. The performance advantage of API in Binary-IoU compared to VPI is relatively smaller than in Class-IoU on COCO-20ⁱ. This is due to the bias of Binary-IoU towards objects that cover a large part of the foreground and background areas.

We also visualize the attention maps computed by our API model in Figure 9. The pixel on the foreground object is highlighted to enhance the prediction of the prediction maps. The proposed latent attention module can capture the outline of the foreground objects. Besides, this module merely increases a little inference time. As evidenced in Figure 11, when we fix $L=10$ and increase $M$ from 1 to 15, the extra time introduced by the attention module is less than 0.1 seconds even under the 5-shot setting. These observations indicate the advantage of our strategy for computational efficiency.

The segmentation map is estimated by Monte Carlo sampling that obtains multiple prototypes z and attention maps m, and produces multiple outputs, then, all outputs are aggregated to produce the final segmentation map. We conduct a qualitative analysis of the effect of the Monte Carlo sampling on the segmentation results. As shown in Figure 11, the segmentation map for each sampled prototype is not always adequate. For example, in the first row, the segmentation map generated by 3 times sampling does not completely recover the object. By averaging the segmentation maps produced by the individual samples, the final segmentation map tends to be more complete and robust.

The quantitative results (Class-IoU of the sample number 5 & 15: 0.56 vs. 0.57) show that the prediction result turns to be better as the number of samples increases. Despite the fact that the segmentation results are more accurate given more samples, it will take more time for inference. We observe that the performance tends to saturate when $L,M$ reach $15$. Therefore, in our experiments, we set $L,M$ to 5 during inference to achieve precise segmentation maps with acceptable inference cost.