Interclass Prototype Relation
for Few-Shot Segmentation

As mentioned in the introduction, the target class, which has a few shots in few-shot segmentation, is difficult to classify similar classes because the data plotted in the feature space is sparse and has low coverage. To address this problem, we propose the Interclass Prototype Relation Network (IPRNet). IPRNet has two modules: the Interclass Prototype Relation Module (IPRM) and the Respective Classifier Module (RCM). These modules aim to improve the identification performance of similar classes by reducing the similarity between prototypes and extracting the differences between classes. An overview of the network is shown in Figure 2.

This section describes the overall flow of IPRNet. First, the support and query images were fed into pretrained shared CNNs (pretrained by ImageNet [7]) to extract features. Next, we passed the support features with the support masks through the IPRM. Through IRPM, we obtain prototypes that represent feature vectors for each class, and the value $L_{r}$ that indicates the similarity between each prototype. Then, for more accurate pixel-by-pixel matching, the matching between prototypes and the query feature is performed by calculating the cosine similarity in map process. The two similarity maps can be obtained by matching with the query feature for the target class prototype and the background prototype, respectively. This process was inspired by earlier research on MLC [42]. The input to the multi-scale network is the concatenation of the support features, the query feature, and two similarity maps. The output is a relation feature valid for classification and $L_{m}$ which is a multi-scale loss introduced in PFENet [34]. The multi-scale network sets up the top-down structure of FPN-like [17] by using the feature enrichment module introduced in PFENet [34] to obtain multi-scale information. This structure enables fast multi-scale aggregation, by transferring features from finer to coarser and by easing feature interaction. Each scale yielded segmentation results for calculating the loss. The average loss value for each scale was $L_{m}$. Finally, the relation feature, a multi-scale information-intensive feature map, is generated by fusing all the different scales into a concatenated feature map by convolution.

The relation feature is the input to the RCM. The RCM performs a discriminative process to classify the foreground and background more and obtains the final inference result $\tilde{M^{Q}}$. Then, the loss $L_{p}$ between the inference result and ground-truth mask $M^{Q}$ is calculated. The loss function, which is the cost function of training, is given by equation (1).

$w_{1}$, $w_{2}$, and $w_{3}$ are the weight coefficients, which were trained with $w_{1}=0.4$, $w_{2}=0.2$, and $w_{3}=0.4$, respectively. Here, we describe the details of the proposed IPRM and RCM.

The Interclass Prototype Relation Module (IPRM) was proposed to reduce the similarity between classes. The IPRM has a prototype generation process and relation loss calculation process. The prototype generation process calculates the prototypes of all classes present in the batch and the support images. This process obtains a prototype, which is a global descriptor of a particular class in an image by taking as input the feature map extracted from the support image and the segmentation mask paired with it. We employed the masked average pooling strategy [30] to compute the prototype for each class. The prototype $P_{i}^{c}$ of the $c$ th class in the $i$ th support image is computed using equation (2).

$F_{i}$ is the feature map extracted via the backbone network with the support image $I_{i}^{S}$ as the input. The subscripts $x$ and $y$ indicate the horizontal and vertical spatial positions of the feature map, respectively. $M_{i}^{x,y}$ is the segmentation mask. By applying this, we can eliminate the regions of the feature map other than the specified class. Equation (2) was calculated for all the support images in the batch. The maximum number of prototypes $n$ obtained is the number of all classes in which $S_{train}$ has $c\in\{0,1,2,\ldots,n\}$.

Next, the relation loss calculation process was performed using the obtained prototypes. This process calculates the average value $L_{r}$ of the cosine similarity between different classes. The relation loss $L_{r}$ is calculated using equation (3,4).

The $c_{s}$ and $c_{t}$ refer to class numbers and the difference between the prototypes of two different classes is measured by the cosine similarity expressed in the equation (4). The similarity between the prototypes of two different classes $c_{s}$,$c_{t}$ computed in equation (4) is calculated for all combinations switching between pairs of prototypes to be measured using the equation (3). The average of these prototype similarity values are the relation loss $L_{r}$. The relation loss $L_{r}$ is designed to improve the separation performance between each class by training the network such that the similarity between each class is low.

Further, among the prototypes calculated by the equation (2), the prototype of the target class and the prototype extracted from the background region is selected. The prototypes were compared with the query feature and their respective similarity maps were computed. These similarity maps were used as input to the multi-scale network.

An overview of the Respective Classifier Module (RCM) is shown in Figure 3. The RCM takes as input the relation feature, which is the output of the multi-scale network shown in Figure 2. It was designed to improve the separation performance between the target class of objects and the background by estimating each of them independently, reintegrating the results.

The relation feature $F_{r}$ is branched for foreground target class prediction and background prediction. Then, it is transformed into a probability distribution representation $V_{1}$ and $V_{0}$ through two convolutional layers, the activation layer, and the softmax layer respectively. $V_{1}$ indicates the probability distribution of a single object in the target class, and $V_{0}$ indicates the probability distribution of the background region only. The loss calculation between the probability distribution representation $V$ and the ground-truth mask $M^{Q}$ can be expressed by equation (5).

The loss is calculated by cross-entropy error. $V_{1}$, the calculated loss value $L_{1}$ is the loss function of a single object of the target class, and given probability distribution $V_{0}$, the calculated loss value $L_{0}$ is the loss function of the background region only. Equation (6) is a function that, according to the given class ID $c$, returns 255 pixel positions other than those corresponding to the class ID in the ground-truth mask $M^{Q}$. Two hundred and fifty-five means ignore the label, an area that is not used when calculating loss, thus eliminating the relationship between the background and foreground. This mechanism allows the RCM to acquire a discriminator through training that can make inferences about its pixels based on information about the object itself only.

The final estimation result $V_{f}$ can be obtained by fusing the information from the concatenation of the foreground probability distribution $V_{1}$, background probability distribution $V_{0}$, and relation feature $F_{r}$ with the two convolution layers, the activation layer, and the softmax layer. The cross-entropy error $L_{f}$ was also obtained between the prediction result $V_{f}$ and the ground-truth mask ${M^{Q}}$.

The final loss value $L_{p}$ output by the RCM is obtained by the weighted addition of the loss value $L_{1}$ to the prediction result of the foreground object alone, the loss value $L_{0}$ to the prediction result of the background region only, and the loss value $L_{f}$ to the final prediction result. This can be defined by the following equation (7).

$\alpha$, $\beta$, and $\gamma$ are the weight coefficients, which were trained by setting $\alpha=0.15$, $\beta=0.15$, and $\gamma=0.7$, respectively.

To analyze the performance of IPRNet, we selected two datasets that are widely used for few-shot segmentation, Pascal-$5^{i}$[29] and COCO-$20^{i}$ [18]. Pascal-$5^{i}$ [29] contains images from PASCAL VOC 2012 [9] and additional annotations from SBD [12]. For training, a total of twenty class categories were evenly divided into four splits, and model training was performed using a cross-validation method. Specifically, three splits were selected as the training data during the training process, and the remaining splits were used for testing. During testing, one thousand support-query pairs were randomly sampled and evaluated [29]. COCO-$20^{i}$ [18], unlike Pascal-$5^{i}$[29], is an incredibly challenging dataset. This is because it is a large dataset having 82,081 images, and many objects are included in the images of realistic scenes. Following the FWB [26], the eighty classes of COCO-$20^{i}$ [18] were evenly divided into four splits, and the same cross-validation scheme was used. To obtain more stable results, we randomly sampled 20,000 pairs of during the testing [34]. As an evaluation metric, we used the mean Intersection-over-Union (mIoU), which is commonly used in semantic segmentation.

An ablation study was conducted to verify the influence of the proposed IPRM and RCM. To verify the separation performance between similar classes, we compared the per-class IoU results of IPRNet and the baseline without our IPRM and RCM.

ResNet [13] is employed as the backbone network, and block2 and block3 are concatenated to generate the feature map [46]. The input support and query images were cropped to an image size of 400×400 pixels and fed into the backbone network. The initial learning rate was set to 0.05, momentum and weight decay to 0.9 and 0.0001, respectively, and the optimizer was trained with a batch size of thirty-two using the SGD optimizer. We also adapted the poly method [3], where the decay of the learning rate is achieved by multiplying by $(1-current\_iter)^{power}$, and the $power$ is set to 0.9. The pretrained backbone network is frozen so that it does not learn class-specific representations of the training data. We implemented it using Pytorch experimented with an NVIDIA A10G GPU.

In the ablation study, for the IPRM deletion, we also modified a mechanism that uses only the conventional masked average pooling strategy [30] to acquire the target class prototype and background class prototype. The RCM is completely removed and replaced with a mechanism that calculates the loss between the output of the multi-scale network and the ground truth.

The effectiveness of our method was evaluated on two benchmark datasets Pascal-$5^{i}$ [29] and COCO-$20^{i}$ [18]. We extensively experimented with the $1$-shot and $5$-shot on $1$-way problems in various few-shot split settings using the widely used encoder networks ResNet-50 and ResNet-101. Here, an n-way k-shot implies that k samples are given for each class among the n classes. Extensive experiments show that the mIoU is improved over the conventional method in all the cases. The ablation study confirmed that removing the IPRM and RCM reduced the mIoU.