Anti-aliasing Semantic Reconstruction for Few-Shot Semantic Segmentation

Few-shot semantic segmentation aims to learn a model ($e.g.$, a network) which can generalize to previously unseen classes. Given two image sets $\mathcal{D}_{base}$ and $\mathcal{D}_{novel}$, classes in $\mathcal{D}_{novel}$ do not appear in $\mathcal{D}_{base}$, it requires to train the feature representation on $\mathcal{D}_{base}$ (which has sufficient data) and test on $\mathcal{D}_{novel}$ (which has only a few annotations). Both $\mathcal{D}_{base}$ and $\mathcal{D}_{novel}$ contain several episodes, each of which consists of a support set ${(\mathbf{A}_{i}^{s},\mathbf{M}_{i}^{s})}_{i=1}^{K}$ and a query set $(\mathbf{A}^{q},\mathbf{M}^{q})$, where $K,\mathbf{A}_{i}^{s},\mathbf{M}_{i}^{s},\mathbf{A}^{q}$ and $\mathbf{M}^{q}$ respectively represent the shot number, the support image, the support mask, the query image, and the query mask. For each training episode, the model is optimized to segment $\mathbf{A}^{q}$ driven by the segmentation loss $\mathcal{L}_{seg}$. Segmentation performance is evaluated on $\mathcal{D}_{novel}$ across all the test episodes.

We propose a semantic reconstruction framework, where the semantics of novel classes are explicitly reconstructed by those of base classes, Fig. 2. Given support and query images, after extracting convolutional features through a CNN, the ground-truth mask is multiplied with support features in a pixel-wised fashion to filter out background features [47, 43, 22, 37, 24]. With a convolutional block we reduce the number of feature channels and obtain support features $\mathbf{F}_{c}^{s}\in{\mathbb{R}^{H\times W\times(B\times D)}}$ and query features $\mathbf{F}^{q}\in{\mathbb{R}^{H\times W\times(B\times D)}}$, where $H\times W$, $B$, and $D$ respectively denote the size of feature maps, base class number, and feature channel number. During training phase, $c$ denotes the base class. And $c$ denotes the novel class during testing phase. The convolutional block consists of pyramid convolution layers, which captures features from coarse to fine. To explicitly encode class-related semantics, we averagely partition the feature channels to $B$ groups, corresponding to $B$ base classes. The grouped features $\mathbf{F}_{c}^{s}$ and $\mathbf{F}^{q}$ are further spatially squeezed into two vectors $\mathbf{v}_{c}^{s}$ and $\mathbf{v}^{q}$, termed semantic vectors, by global average pooling, Fig. 2.

Corresponding to $B$ base classes, the semantic vectors $\mathbf{v}_{c}^{s}$ and $\mathbf{v}^{q}$ consists of $B$ sub-vectors $\{\mathbf{v}^{s}_{c,b}\}_{\{b=1,2,...,B\}}\in{\mathbb{R}^{D}}$ and $\{\mathbf{v}^{q}_{b}\}_{\{b=1,2,...,B\}}\in{\mathbb{R}^{D}}$. During the training phase, the sub-vectors are used to construct basis vectors in the $B$-dimensional class-level semantic space by the semantic span module, as explained in Section 3.3, Fig. 2. In the space, a basis vector ($\mathbf{v}_{b}$) corresponding to the $b$-th base class is defined as $\mathbf{v}_{b}=\mathbf{v}^{s}_{c,b}/||\mathbf{v}^{s}_{c,b}||=\mathbf{v}^{q}_{b}/||\mathbf{v}^{q}_{b}||$. In the inference phase, the semantic vector for the $c$-th class in support branch can be linearly reconstructed [16], as

where $\hat{\mathbf{v}}^{s}_{c}$ denotes the reconstructed support semantic vector (reconstructed support vector for short), and $w_{c,b}^{s}$ is the $b$-th element of the weight vector $\mathbf{w}_{c}^{s}=softmax([||\mathbf{v}_{c,1}^{s}||,||\mathbf{v}_{c,2}^{s}||,\cdots,||\mathbf{v}_{c,B}^{s}||])$. Large $w_{c,b}^{s}$ indicates that the basis vector, whose corresponding base class have strong similarity with the $c$-th class, contributes much to the reconstruction.

Consistently, given a query image, the corresponding query features $\mathbf{F}^{q}$ are reconstructed by regarding each location on the feature maps as a feature vector. Each location of feature map is reconstructed as $\hat{\mathbf{F}}^{q}(x,y)=\sum_{b=1}^{B}{\mathbf{W}_{b}^{q}(x,y)\cdot\mathbf{v}_{b}}$, where $(x,y)$ denotes the coordinates of pixels on the feature map, and $\mathbf{W}_{b}^{q}(x,y)$ is defined as the norm of sub-vector $\mathbf{F}^{q}_{b}(x,y)$. Considering that the query image contains objects not only belonging to the target class but also other classes, we exploit the semantic filtering module, as illustrated in Section 3.4, to filter out the interfering components in the reconstructed query features for the $c$-th target class semantic segmentation.

Within the origin feature space, when base class features are close to each other, there could be semantic aliasing among novel classes. To minimize semantic aliasing, we propose to span a class-level semantic space in the training phase. To construct a group of basis vectors which tends to be orthogonal and representative, we propose the semantic span module (semantic span for short). As shown in Fig. 3, the semantic span is driven by two loss functions, $i.e.$, semantic decoupling and contrastive losses.

On the one hand, the semantic span targets at constructing basis vectors by regularizing the feature maps so that each group of features is correlated to a special object class. To fulfill this purpose, we propose the following semantic decoupling loss, as

where $\mathbf{w}_{c}^{s}$ denotes the reconstruction weight vector, and $\mathbf{y}\in\mathbb{R}^{B}$ denotes the one-hot class label of a support image. Obviously, minimizing $\mathcal{L}_{dec}$ is equivalent to maximize the reconstruction weights related to the specific class ($e.g.$, the $c$-th class), while minimizing those unrelated to it. This defines a soft manner converting a group of features correlated to the semantics of the specific class ($e.g.$, the $c$-th class) to its corresponding basis vector in the class-level semantic space.

On the other hand, the semantic span targets at further enhancing orthogonality of basis vectors, which improve the quality of novel class reconstruction. In details, sub-vectors in $\{\mathbf{v}^{s}_{c,b}\}_{\{b=1...B\}}\cup\{\mathbf{v}^{q}_{b}\}_{\{b=1...B\}}$ belonging to different classes are expected to be orthogonal to each other while those corresponding to the same base classes, $e.g.$, $\mathbf{v}^{s}_{c,b}$ and $\mathbf{v}^{q}_{b}$, are expected to have a small vector angle. These two objectives are simultaneously achieved by minimizing the contrastive loss defined as

where $cos<\cdot>$ denotes the $Cosine$ distance metric of two vectors. In summary, the final loss of the semantic reconstruction framework is defined as:

where $\alpha$, $\beta$ and $\gamma$ are weights of the loss functions. Note that $\mathcal{L}_{con}$ is calculated during the later stage of training phase.

When multiple objects from different classes exist in the same query image, the reconstructed features of the query image contains components of all these classes. To pick out objects belonging to the target class and suppress interfering semantics, $i.e.$, divorcing the semantics related to the background or objects from other classes, we propose a semantic filtering module. Moreover, owing to that the reconstructed vectors of different classes are non-collinear, the semantic filtering module is implemented by projecting query feature vectors to the reconstructed support vector, as shown in Fig. 4. This is also based on the fact that the reconstructed support vector has precise semantics because the corresponding features have been multiplied with the ground-truth mask, as shown in Fig. 2.

On the support branch, we follow Eq. 1 to reconstruct the support features using basis vectors and obtain the reconstructed support vector $\hat{\mathbf{v}}_{c}^{s}$. On the query branch, we reconstruct each feature vector $\mathbf{F}^{q}(x,y)$ on the feature maps in the same way and obtain the reconstructed features $\hat{F}^{q}$. We then project $\hat{\mathbf{F}}^{q}$ to $\hat{\mathbf{v}}_{c}^{s}$ to calculate filtered features as

where $(x,y)$ denotes the coordinates of pixels on the feature maps. The intuitive effects of the filter operation are displayed in Fig. 4, which illustrates that the support branch guides the query branch more effectively. The filtered query features $\hat{\mathbf{F}}_{c}^{q}$ are further enhanced by a residual convolutional module with iterative refinement optimization and fed to Atrous Spatial Pyramid Pooling (ASPP) to predict the segmentation mask, Fig. 2. For the residual convolutional module, we replace the history mask in CANet [43] with the squeezed $\hat{\mathbf{F}}_{c}^{q}$.

ASR can be analyzed from the perspectives of vector orthogonality and sparse reconstruction. Without loss of generality, we take the two-dimensional space as an example. Denote ${\mathbf{v_{1}},\mathbf{v}_{2}}\in\mathbb{R}^{2}$ two unit basis vectors ($\left\|\mathbf{v}_{1}\right\|=1,\left\|\mathbf{v}_{2}\right\|=1$), which span the space. Denote $\theta$ as the angle between $\mathbf{v}_{1}$ and $\mathbf{v}_{2}$, and $cos\,\theta=\mathbf{v}_{1}\mathbf{v}_{2}$. According to the properties of linear algebra [16], any vectors, $e.g.,{\mathbf{u}_{1},\mathbf{u}_{2}}\in\mathbb{R}^{2}$, in the spanned space can be linearly reconstructed as

where $w_{11},w_{12},w_{21},w_{22}\in[0,1]$ are reconstruction weights which feed the linear constraints: $w_{11}+w_{12}=1.0$ and $w_{21}+w_{22}=1.0$. $C_{1}$ and $C_{2}$ are scaling constants. The cosine similarity between $\mathbf{u}_{1}$ and $\mathbf{u}_{2}$ is computed as

where $(w_{11}+w_{21}-2w_{11}w_{21})\in[0,1]$ and $(cos\,\theta-1)\in[-1,0].$

Orthogonality. To reduce semantic aliasing of any two novel classes, the angle between their semantic vectors, $\mathbf{u}_{1}$ and $\mathbf{u}_{2}$, should be large, known as to reduce $cos<\mathbf{u}_{1},\mathbf{u}_{2}>$. According to the last line of Eq. 7, to obtain a small $cos<\mathbf{u}_{1},\mathbf{u}_{2}>$, the term $cos\,\theta$ should approach to $0$, which means that the angle between the basis vectors $\mathbf{v}_{1}$ and $\mathbf{v}_{2}$ is large, which implies the orthogonality of basis vectors. The proposed ASR approach satisfies the orthogonality by introducing the semantic span module. As shown in Fig.5, the statistical visualization results over base classes validate the orthogonality.

Sparse Reconstruction. Refer to the last line of Eq. 7, the another manner to reduce the $cos<\mathbf{u}_{1},\mathbf{u}_{2}>$ is enlarging the term $(w_{11}+w_{21}-2w_{11}w_{21})$. According to the function characteristics, $(w_{11}+w_{21}-2w_{11}w_{21})$ reaches its maximum when $|w_{11}$-$w_{21}|$ approaches to $1.0$, which contains underlying conditions that $|w_{11}$-$w_{12}|$ and $|w_{21}$-$w_{22}|$ approach to $1.0$ due to the linear constraints. This illustrates that to further aggregate the capability of anti-aliasing and guarantee the discrimination of novel classes, the reconstruction weights for novel classes should be differential and sparse. ASR satisfies these requirements due to the potential class-level semantic similarity according to the statistical results shown in Fig. 6. Meanwhile, nonzero weights over multiple basis classes other than the dominate one enable ASR to distinguish novel classes from base classes.

Datasets. The experiments are conducted on PASCAL VOC 2012 [7] and MS COCO [19] datasets. We combine the PASCAL VOC 2012 with SBD [12] and separate the combined dataset into four splits. The cross-validation method is used to evaluate the proposed approach by sampling one split as test categories $\mathcal{C}_{test}={4i+1,…,4i+5}$, where $i$ is the index of a split. The remaining three splits are set as base classes for training. The reorganized dataset is termed as Pascal-$5^{i}$ [33, 37]. Following the settings in [25, 33, 37] we construct the COCO-$20^{i}$ dataset. MS COCO is divided into four splits, each of which contains 20 categories. We follow the same scheme for training and evaluation as on the Pascal-$5^{i}$. The category labels for the four splits are included in the supplementary material. For each split, 1000 pairs of support and query images are randomly selected for performance evaluation.

Training and Evaluation. We use CANet [43] without attention modules as the baseline. In training, we set the learning rate as 0.00045. The segmentation model (network) is trained for 200000 steps with the poly descent training strategy and the stochastic gradient descent (SGD) optimizer. Several data augmentation strategies including normalization, horizontal flipping, gaussian filtering, random cropping, random rotation and random resizing are used. We adopt both the single-scale and multi-scale [43, 42, 22] evaluation strategies during testing. Our approach is implemented upon the PyTorch 1.3 and run on Nvidia Tesla V100 GPUs.

Evaluation Metric. Following [34, 25, 43], we use the mean Intersection over Union (mIoU) and binary Intersection over Union (FB-IoU) as the performance evaluation metrics. The mIoU calculates the per-class foreground IoU and averages the IoU for all classes to obtain the final evaluation metric. The FB-IoU calculates the mean of foreground IoU and background IoU over all images regardless of category. For category $k$, IoU is defined as $IoU_{k}=TP_{k}/(TP_{k}+FP_{k}+FN_{k})$, where the $TP_{k},FP_{k}$ and $FN_{k}$ are the number of true positives, false positives and false negatives in segmentation masks. mIoU is the average of IoUs for all the test categories and FB-IoU is the average of IoUs for all the test categories and the background. We report the segmentation performance by averaging the mIoUs on the four cross-validation splits.

PASCAL VOC. In Table 1, we report the performance on Pascal VOC. ASR outperforms the prior methods with significant margins. Under 1-shot settings, with a VGG16 backbone, it respectively outperforms RPMMs [37] and SST [49] by 2.64% and 1.36%. Under the 1-shot settings, with a ResNet50 backbone, ASR outperforms CANet [43] and RPMMs [37] method by 2.76% and 1.82%. Under the 5-shot settings, ASR is comparable to the state-of-the-art method. It is worth mentioning that the SST and PPNet used additional $k$-shot fusion strategies while ASR uses a simple averaging strategy to get five-shot results. In Table 3, ASR is compared with state-of-the-art approaches with respect to FB-IoU. FB-IoU calculates the mean of foreground IoU and background IoU over images regardless of the categories, which reflects how well the full object extent is activated. ASR is on par with the compared methods, if not outperforms.

MS COCO. In Table 2, we report the segmentation performance on MS COCO. ASR outperforms the prior methods in most settings. Particularly under the 1-shot setting, it improves RPMMs [37] by 3.27%. Under 5-shot setting, it improves DAN [33] by 6.15%, which are significant margins. For the MS COCO dataset with larger semantic aliasing for the more object categories, semantic reconstruction demonstrated larger advantages. For the larger object category number, we construct a space using more orthogonal basis vectors, which have stronger ability of representation and discrimination. According to Section 3.5, semantic aliasing among novel classes is suppressed effectively. That is why ASR achieves larger performance gains on the MS COCO dataset.

We sampled 4000 images from 20 classes in PASCAL VOC, and drew the confusion matrix according to the segmentation results, Fig. 7. ASR effectively reduce semantic aliasing among classes. We further visualize segmentation results and compare them with baseline, Fig. 8. Based on the anti-aliasing representation of novel classes and semantic filtering, ASR reduces the false positive segmentation caused by interfering semantics within the query images.

Semantic Span. In Table 4, when simply introducing semantic reconstruction to the baseline method, the performance slightly drops. By using the semantic span module, we improved the performance from 53.26% to 55.98%, demonstrating the necessity of establishing orthogonal basis vectors during semantic reconstruction.

Semantic Filtering. As shown in Table 4, directly applying semantic filtering on the baseline method harms the performance because the support features contain aliasing semantics. By combining all the modules, ASR improves the mIoU by 2.66% (58.64% vs. 55.98%). In Table 5, four filtering strategies are compared. The vector projection strategy defined in Section 3.4 achieves the best result. Vector projection utilizes the characteristic of vector operations to retain semantics related to the target class and suppress unrelated semantics at utmost.

Channel Number ($D$). The channel number ($D$) of features to construct basis vectors is an important parameter which affects the orthogonality of basis vectors. From Fig. 9 we can see that the performance improves with the increase of $D$ and starts to plateau when $D=8$, where the orthogonality of different basis vectors is sufficient for novel class reconstruction. For the MS COCO dataset $D$ is set to 30.

Model Size and Efficiency. The model size of ASR is 36.7M, which is slightly larger than that of the baseline method [43] (36.3M) but much smaller than other methods, such as OSLSM [29] (272.6M) and FWB [25] (43.0M). With a Nvidia Tesla V100 GPU, the inference speed is 30 FPS, which is comparable with that of CANet (29 FPS).

Following the settings in [24], we conduct two-way one-shot segmentation experiments on PASCAL VOC. From Tab. 6 one can see that ASR outperforms PPNet [24] with a significant margin (53.13% vs. 51.65%). Because two-way segmentation requires not only to segment targets objects but also to distinguish different classes, the model is more sensitive to semantic aliasing. Our ASR approach effectively reduces the semantic aliasing between novel classes and thereby achieves superiors segmentation performance.