Hypercorrelation Squeeze for Few-Shot Segmentation
Juhong Min, Dahyun Kang, Minsu Cho
Introduction
The advent of deep convolutional neural networks has promoted dramatic advances in many computer vision tasks including object tracking , visual correspondence , and semantic segmentation to name a few. Despite the effectiveness of deep networks, their demand for a heavy amount of annotated examples from large-scale datasets still remains a fundamental limitation since data labeling requires substantial human efforts, especially for dense prediction tasks, e.g., semantic segmentation. To cope with the challenge, there have been various attempts in semi- and weakly-supervised segmentation approaches which in turn effectively alleviated the data-hunger issue. However, given only a few annotated training examples, the problem of poor generalization ability of the deep networks is yet the primary concern that many few-shot segmentation methods struggle to address.
In contrast, human visual system easily achieves generalizing appearances of new objects given extremely limited supervision. The crux of such intelligence lies at the ability in finding reliable correspondences across different instances of the same class. Recent work on semantic correspondence shows that leveraging dense intermediate features and processing correlation tensors with high-dimensional convolutions are significantly effective in establishing accurate correspondences. However, while recent few-shot segmentation research began active exploration in the direction of correlation learning, most of them neither exploit diverse levels of feature representations from early to late layers of a CNN nor construct pair-wise feature correlations to capture fine-grained correlation patterns. There have been some attempts in utilizing dense correlations with multi-level features, but they are yet limited in the sense that they simply employ the dense correlations for graph attention, using only a small fraction of intermediate conv layers.
In this work we combine the two of the most influential techniques in recent research of visual correspondence, multi-level features and 4D convolutions, and deign a novel framework, dubbed Hypercorrelation Squeeze Networks (HSNet), for the task of few-shot semantic segmentation. As illustrated in Fig. 1, our network exploits diverse geometric/semantic feature representations from many different intermediate CNN layers to construct a collection of 4D correlation tensors, i.e., hypercorrelations, which represent a rich set of correspondences in multiple visual aspects. Following the work of FPN , we adapt pyramidal design to capture both high-level semantic and low-level geometric cues for precise mask prediction in coarse-to-fine manner using deeply stacked 4D conv layers. To reduce computational burden caused by such heavy use of high-dimensional convs, we devise an efficient 4D kernel via reasonable weight-sparsification which enables real-time inference while being more effective and light-weight than the existing ones. The improvements on standard few-shot segmentation benchmarks of PASCAL-5i , COCO-20i , and FSS-1000 verify the efficacy of the proposed method.
Related Work
Semantic segmentation. The goal of semantic segmentation is to classify each pixel of an image into one of the predefined object categories. Prevalent segmentation approaches typically employ encoder-decoder structure in their architecture; the encoder aggregates features along deep convolutional pathways and provides high-dimensional feature map in low-resolution and the corresponding decoder takes the output to predict segmentation mask by reversing this process . Although the methods clearly show the effectiveness of the encoder-decoder architecture in the task of semantic segmentation, offering useful insights to our study, they still suffer apparent disadvantages of data-driven nature of neural networks: lack of generalizibility under insufficient training data.
Few-shot learning. To resolve the generalization problem, many recent approaches to image classification made various attempts in training deep networks with a few annotated examples . Vinyals et al. propose matching networks for one-shot learning; the method utilizes a special kind of mini-batches called episodes to match training and testing environments, facilitating better generalization on novel classes. Snell et al. introduce prototypical networks which compute distances between representative embeddings, i.e., prototypes, for few-shot classification. With the growing interests in few-shot learning in classification domain, the problem of few-shot segmentation has attracted a great deal of attention as well. Shaban et al. propose one-shot semantic segmentation networks which (meta-) learns to generate parameters of FCN . Inspired by the prototypical networks , utilizing prototype representations to guide mask prediction in a query image became a popular paradigm in few-shot segmentation literature . Witnessing the limitation of prototypical approaches, e.g., loss of spatial structure due to masked average pooling , work of build pair-wise feature correlations, e.g., graph attention, to retain the spatial structure of the images for fine-grained mask prediction. Note that both prototypical and graph-based methods fundamentally focus on learning to find reliable correspondences between support and query images for accurate mask prediction. In this work, we advance this idea and focus on learning to analyze correspondences using adequately designed learnable layers, e.g., 4D convolutions , for effective semantic segmentation.
Learning visual correspondences. The task of visual correspondence aims to find reliable correspondences under challenging degree of variations . Many methods typically built upon convolutional features pretrained on classification task , showing they serve as good transferable representations. Recent approaches to semantic correspondence show that efficiently exploiting different levels of convolutional features distributed over all intermediate layers clearly benefits matching accuracy. In wide-baseline matching literature, a trending choice is to employ 4D convolutions on dense feature matches to identify spatially consistent matches by analyzing local patterns in 4D space. The use of multi-level features and relational pattern analysis using 4D convs are the two widely adopted techniques in the field of visual correspondence.
In this paper we adapt the two most influential methodologies in visual correspondence to tackle few-shot segmentation: multi-level features and 4D convolutions. Inspired by the previous matching methods , which use multi-level features to build effective “appearance features”, we construct high-dimensional “relational features” using intermediate CNN features and process them with a series of 4D convolutions. However, their quadratic complexity still remains a major bottleneck in designing cost-effective deep networks, constraining many previous matching methods to use only a few 4D conv layers. To resolve the issue, we develop a light-weight 4D convolutional kernel by collecting only a small subset of vital parameters for effective pattern recognition, which eventually leads to an efficient decomposition into a pair of 2D conv kernels with a linear complexity. Our contributions can be summarized as follows:
We present the Hypercorrelation Squeeze Networks that analyze dense feature matches of diverse visual aspects using deeply stacked 4D conv layers.
We propose center-pivot 4D conv kernel which is more effective than the existing one in terms both accuracy and speed, achieving real-time inference.
The proposed method sets a new state of the art on three standard few-shot segmentation benchmarks: PASCAL-5i , COCO-20i , and FSS-1000 .
Problem Setup
Proposed Approach
2 4D-convolutional pyramid encoder
3 2D-convolutional context decoder
4 Center-pivot 4D convolution
Apparently, our network with such a large number of 4D convolutions demands a substantial amount of resources due to the curse of dimensionality, which constrained many visual correspondence methods to use only a few 4D conv layers. To address the concern, we revisit the 4D convolution operation and delve into its limitations. Then we demonstrate how a unique weight-sparsification scheme effectively resolves the issues.
where denotes a set of neighbourhood regions within the local 4D window centered on position , i.e., as visualized in Fig. 4. Although the use of 4D convolutions on a correlation tensor has shown its efficacy with good empirical performance in correspondence-related domains , its quadratic complexity with respect to the size of input features still remains a primary bottleneck. Another limiting factor is over-parameterization of the high-dimensional kernel: Consider a single activation in an D tensor convolved by D conv kernel. The number of times that the kernel processes this activation is exponentially proportional to . This implies some unreliable input activations with large magnitudes may entail some noise in capturing reliable patterns as a result of their excessive exposure to the high-dimensional kernel. The work of resolves the former problem (quadratic complexity) using spatially separable 4D kernels to approximate the 4D conv with two separate 2D kernels along with additional batch normalization layers that settle the latter problem (numerical instability). In this work we introduce a novel weight-sparsification scheme to address both issues at the same time.
Center-pivot 4D convolution. Our goal is to design a light-weight 4D kernel that is efficient in terms of both memory and time while effectively approximating the existing ones . We achieve this via a reasonable weight-sparsification; from a set of neighborhood positions within a local 4D window of interest, our kernel aims to disregard a large number of activations located at fairly insignificant positions in the 4D window, thereby focusing on a small subset of relevant activations only. Specifically, we consider the activations at positions that pivots either one of 2-dimensional centers, e.g., or , as the foremost influential ones as illustrated in Fig. 4. Given 4D position , we collect its neighbors if and only if they are adjacent to either or in its corresponding 2D subspace and define two respective sets as and . The set of center-pivot neighbours is defined as . Based on these two subsets of neighbors, center-pivot 4D convolution can be formulated as a union of two separate 4D convolutions:
which performs two different convolutions on separate 2D subspaces, having a linear complexity. In Sec. 5.2, we experimentally demonstrate the superiority of the center-pivot 4D kernels over the existing ones in terms of accuracy, memory, and time. We refer the readers to the Appendix A for a complete derivation of Eqn. 5.
5 Extension to KK-shot setting
Experiment
In this section we evaluate the proposed method, compare it with recent state of the arts, and provide in-depth analyses of the results with ablation study.
Implementation details. For the backbone network, we employ VGG and ResNet families pre-trained on ImageNet , e.g., VGG16, ResNet50, and ResNet101. For VGG16 backbone, we extract features after every conv layer in the last two building blocks: from conv4_x to conv5_x, and after the last maxpooling layer. For ResNet backbones, we extract features at the end of each bottleneck before ReLU activation: from conv3_x to conv5_x. This feature extracting scheme results in 3 pyramidal layers () for each backbone. We set spatial sizes of both support and query images to , i.e., , thus having , , and . The network is implemented in PyTorch and optimized using Adam with learning rate of 1e-3. We freeze the pre-trained backbone networks to prevent them from learning class-specific representations of the training data.
Datasets. We evaluate the proposed network on three standard few-shot segmentation datasets: PASCAL-5i , COCO-20i , and FSS-1000 . PASCAL-5i is created from PASCAL VOC 2012 with extra mask annotations , consisting of 20 object classes that are evenly divided into 4 folds: . COCO-20i consists of mask-annotated images from 80 object classes divided into 4 folds: . Following common training/evaluation scheme , we conduct cross-validation over all the folds; for each fold , samples from the other remaining folds are used for training and 1,000 episodes from the target fold are randomly sampled for evaluation. For every fold, we use the same model with the same hyperparameter setup following the standard cross-validation protocol. FSS-1000 contains mask-annotated images from 1,000 classes divided into training, validation and test splits having 520, 240, and 240 classes respectively.
We evaluate the proposed model on PASCAL-5i, COCO-20i, and FSS-1000 and compare the results with recent methods . Table 1 summarizes 1-shot and 5-shot results on PASCAL-5i; all of our models with three different backbones clearly set new state of the arts with the smallest the number of learnable parameters. With ResNet101 backbone, our 1-shot and 5-shot results respectively achieve 6.1%p and 4.8%p of mIoU improvements over and , verifying its superiority in few-shot segmentation task. As shown in Tab. 3, our model outperforms recent methods with a sizable margin on COCO-20i as well, achieving 2.7%p (1-shot) and 6.8%p (5-shot) of mIoU improvements over with ResNet101 backbone. Also on the last benchmark, FSS-1000, our method sets a new state of the art, outperforming as shown in Tab. 3.
We conduct additional experiments without support feature masking (Eqn. 1). Note that this setup is similar to co-segmentation problem with stronger demands for generalizibility since the model is evaluated on novel classes. As seen in the bottom row of Tab. 1, our model without support masking still performs remarkably well, achieving 1.4%p mIoU improvement over the previous best method in 1-shot setting whereas it rivals in 5-shot setting. This interesting result reveals that our model is also capable of identifying ‘common’ instances across different input images as well as predicting fine-grained segmentation masks.
Robustness to domain shift. To demonstrate the robustness of our method to domain shift, we evaluate COCO-trained HSNet on each fold of PASCAL-5i following the recent work of . We use the same training/test folds as in where object classes in training and testing do not overlap. As seen in Tab. 4, our model, which is trained without any data augmentation methods with 18 times smaller number of trainable parameters compared to (2.6M vs. 46.7M), performs robustly in presence of large domain gaps between COCO-20i and PASCAL-5i, surpassing by 1.0%p in 5-shot setting, and further improves with a larger backbone, e.g., ResNet101. The results clearly show the robustness of our method to domain shift, and may further increase when trained with data augmentations used in .
2 Ablation study
We conduct extensive ablation study to investigate the impacts of major components in our model: hypercorrelations, pyramidal architecture, and center-pivot 4D kernels. We also study how freezing backbone networks prevents overfitting and helps generalization on novel classes. All ablation study experiments are performed with ResNet101 backbone on PASCAL-5i dataset.
Ablation study on pyramid layers. To see the impact of hypercorrelation at each layer , we perform experiments in absence of each pyramidal layer. We train and evaluate our model using two different hypercorrelation pyramids, and , and compare the results with ours . Figure 6 summarizes the results; given hypercorrelation pyramid without geometric information (), our model fails to refine object boundaries in the final mask prediction as visualized in Fig. 7. Given a single hypercorrelation that only encodes semantic relations (), the model predictions are severely damaged, providing only rough localization of the target objects. These results indicate that capturing patterns of both semantic and geometric cues is essential for fine-grained localization.
Comparison between three different 4D kernels. We conduct ablation study on 4D kernel by replacing the proposed center-pivot 4D kernel with the original and spatially separable 4D kernels and compare their model size, per-episode inference time (1-shot), memory consumption, and floating point operations per second (FLOPs) with ours. Table 5 summarizes the results. The proposed kernel records the fastest inference time with the smallest memory/FLOPs requirements while being comparably effective than the other two. The results clearly support our claim that a large part of parameters in a high-dimensional kernel can safely be discarded without harming the quality of predictions; only a few relevant parameters are sufficient and even better for the purpose. While both the separable and our center-pivot 4D convolutions operate on two separate 2D convolutions, auxiliary transformation layers with multiple batch normalizations that make the separable 4D conv numerically stable in its sequential design result in twice larger number of parameters (4.4M vs. 2.6M) and slower inference time (28.48ms vs. 25.51ms) than ours.
The number of 4D layers in building blocks. We also perform experiments with varying number of 4D conv layers in the two building blocks: and . Figure 8 plots 1-shot and 5-shot mIoU results on PASCAL-5i with the model sizes. In the experiments, appending additional 4D layers (with a group norm and a ReLU activation) in the building blocks provides clear performance improvements up to three layers but the accuracy eventually saturates after all. Hence we use a stack of three 4D layers for both.
Finetuning backbone networks. To investigate the significance of learning ‘feature correlations’ over learning ‘feature representation’ in few-shot regime, we finetune our backbone network and compare learning processes of the finetuned model and ours (frozen backbone). Figure 9 plots the training/validation curves of the finetuned model and ours on every fold of PASCAL-5i. The finetuned model rapidly overfits to the training data, losing generic, comprehensive visual representations learned from large-scale dataset . Meanwhile, our model with frozen backbone provides better generalizibility with large trade-offs between training and validation accuracies. The results reveal that learning new appearances under limited supervision requires understanding their ‘relations’ to diverse visual patterns acquired from a vast amount of past experiences, e.g., ImageNet classification. This is quite analogous to human vision perspective in the sense that we generalize novel concepts (what we see) by analyzing their relations to the past observations (what we know) .
For additional experimental details, results and analyses, we refer the readers to the Appendix.
Conclusion
We have presented a novel framework that analyzes complex feature correlations in a fully-convolutional manner using light-weight 4D convolutions. The significant performance improvements on three standard benchmarks demonstrate that learning patterns of feature relations from multiple visual aspects is effective in fine-grained segmentation under limited supervision. We also demonstrated a unique way of discarding insignificant weights leads to an efficient decomposition of a 4D kernel into a pair of 2D kernels, thus allowing extensive use of 4D conv layers at a significantly small cost. We believe our investigation will further facilitate the use of 4D convolutions in other domains that require learning to analyze high-dimensional correlations.
Acknowledgements. This work was supported by Samsung Advanced Institute of Technology (SAIT), the NRF grant (NRF-2017R1E1A1A01077999), and the IITP grant (No.2019-0-01906, AI Graduate School Program - POSTECH) funded by Ministry of Science and ICT, Korea.
References
Appendix A Complete derivation of the center-pivot 4D convolution
where denotes a set of neighbourhood regions within the local 4D window centered on position , i.e., as visualized in Fig. 4. Now we design a light-weight, efficient 4D convolution via a reasonable weight-sparsification; from a set of neighborhood positions within a local 4D window of interest, our kernel aims to disregard a large number of activations located at fairly insignificant positions in the 4D window, thereby focusing only on a small subset of relevant activations for capturing complex patterns in the correlation tensor. Specifically, we consider activations at positions that pivots either one of 2-dimensional centers, e.g., or , as the foremost influential ones. Given 4D position , we collect its neighbors if and only if they are adjacent to either or in its corresponding 2D subspace and define two respective sets as
The set of center-pivot neighbours is defined as a union of the two subsets:
Based on this small subset of neighbors, center-pivot 4D (CP 4D) convolution can be formulated as a union of two separate 4D convolutions:
where and are 4D kernels with their respective neighbours and . Now consider below
which performs two different convolutions on separate 2D subspaces, having a linear complexity.
Appendix B Implementation details
For the backbone networks, we employ VGG and ResNet families pre-trained on ImageNet , e.g., VGG16, ResNet50, and ResNet101. For the VGG16 backbone, we extract features after every conv layer in the last two building blocks: from conv4_x to conv5_x, and after the last maxpooling layer. For the ResNet backbones, we extract features at the end of each bottleneck before ReLU activation: from conv3_x to conv5_x. This feature extracting scheme results in 3 pyramidal layers () for every backbone. We set spatial sizes of both support and query images to , i.e., , thus having , , and for both ResNet50 and ResNet101 bakcbones and , , and for the VGG16 backbone. The network is implemented in PyTorch and optimized using Adam with learning rate of 1e-3. We train our model with batch size of 20, 40, and 20 for PASCAL-5i, COCO-20i, and FSS-1000 respectively. We freeze the pre-trained backbone networks to prevent them from learning class-specific representations of the training data. The intermediate tensor dimensions, the number of parameters of each layers and other additional details of the network are demonstrated in Tab. A5, A6, and A7 for respective backbones of VGG16, ResNet50, and ResNet101.
Appendix C Additional results and analyses
Additional -shot results. Following the work of , we conduct -shot experiments with . Table A2 compares our results with the recent methods on PASCAL-5i and COCO-20i. The significant performance improvements on both datasets clearly indicate the effectiveness of our approach. Achieving 2.5%p and 4.6%p mIoU improvements over the previous best method on respective PASCAL-5i and COCO-20i, our model again sets a new state of the art in 10-shot setting as well, showing notable improvements with larger .
Numerical comparisons of ablation study. We tabularize Figures 5 and 6, e.g., ablation study on hypercorrelations and pyramidal layers, in Tables A3 and A4 respectively. Achieving 4.5%p mIoU improvements over , our method clearly benefits from diverse feature correlations from multi-level CNN layers () as seen in Tab. A3. A large performance gap between and in Tab. A4 (63.9 vs. 55.5) reveals that the intermediary second pyramidal layer () is especially effective in robust mask prediction compared to the first pyramidal layer ().
Evaluation results without using ignore_label on PASCAL-5i. The benchmarks of PASCAL-5i , COCO-20i , and FSS-1000 consist of segmentation mask annotations in which each pixel is labeled with either background or one of the predefined object categories. As pixel-wise segmentation near object boundaries is ambiguous to perform even for human annotators, PASCAL-5i uses a special kind of label called ignore_label which marks pixel regions ignored during training and evaluation to mitigate the ambiguity The use of ignore_label was originally adopted in PASCAL VOC dataset . The same evaluation criteria is naturally transferred to PASCAL-5i as it is created from PASCAL VOC..
Most recent few-shot segmentation work adopt this evaluation criteria but we found that some methods do not utilize ignore_label in their evaluations. Therefore, the methods are unfairly evaluated as fine-grained mask prediction near object boundaries is one of the most challenging part in segmentation problem. For fair comparisons, we intentionally exclude the methods of from Tab. 1 and compare the results of our model evaluated without the use of ignore_label with those methods . The results are summarized in Tab. A1. Even without using ignore_label, the proposed method sets a new state of the art with ResNet50 backbone, outperforming the previous best methods of and by (1-shot) 2.8%p and (5-shot) 5.4%p respectively. With VGG16 backbone, our method performs comparably effective to the previous best method while having the smallest learnable parameters.
Appendix D Qualitative results
Results without support feature masking. As demonstrated in Sec. 5.1, we conduct experiments without support feature masking (Eqn. 1), similarly to co-segmentation problem with stronger demands for generalizibility. Figure A1 visualizes some example results on PASCAL-5i dataset. Even without the use of support masks (in both training and testing), our model effectively segments target instances in query images. The results indicate that learning patterns of feature correlations from multiple visual aspects is effective in fine-grained segmentation as well as identifying ‘common’ instances in the support and query images.
Additional qualitative results. We present additional qualitative results on PASCAL-5i , COCO-20i , and FSS-1000 benchmark datasets. All the qualitative results are best viewed in electronic forms. Example results in presence of large scale-differences, truncations, and occlusions are shown in Fig. A2, A3, and A4. Figure A5 visualizes model predictions under large illumination-changes in support and query images. Figure A6 visualizes some sample predictions given exceptionally small objects in either support or query images. As seen in Fig. A7, we found that our model sometimes predicts more reliable segmentation masks than ground-truth ones. Some qualitative results in presence of large intra-class variations and noisy clutters in background are shown in Fig. A8 and A9. Given only a single support image-annotation pair, our model effectively segments multiple instances in a query image as visualized in Fig. A10. Figure A11 shows representative failure cases; our model fails to localize target objects in presence of severe occlusions, intra-class variances and extremely tiny support (or query) objects. As seen in Fig. A12, the model predictions become much reliable given multiple support image-mask pairs, i.e., .
The code and data to reproduce all experiments in this paper is available at our project page: http://cvlab.postech.ac.kr/research/HSNet/.