Dense Gaussian Processes for Few-Shot Segmentation

1 INTRODUCTION . Image few-shot segmentation ( FSS ) of semantic classes ( Shaban et al. , 2017 ) has received increased attention in recent years . The aim is to segment novel query images based on only a handful annotated training samples , usually referred to as the support set . The FSS method thus needs to extract information from the support set in order to accurately segment a given query image . The problem is highly challenging , since the query image may present radically different views , contexts , scenes , and objects than what is represented in the support set . The core component in any FSS framework is the mechanism that extracts information from the support set to guide the segmentation of the query image . However , the design of this module presents several challenges . First , it needs to aggregate detailed yet generalizable information from the support set , which requires a flexible representation . Second , the FSS method should effectively leverage larger support sets , achieving scalable segmentation performance when increasing its size . While perhaps trivial at first glance , this has proved to be a major obstacle for many state-of-the-art methods , as visualized in Fig . 1 . Third , the method is bound to be queried with appearances not included in the support set . To achieve robust predictions even in such common cases , the method needs to assess the relevance of the information in the support images in order to gracefully revert to e.g . learned segmentation priors when necessary . We address the aforementioned challenges by densely aggregating information in the support set using Gaussian Processes ( GPs ) . Specifically , we use a GP to learn a mapping between dense local deep feature vectors and their corresponding mask values . The mask values are assumed to have a jointly Gaussian distribution with covariance based on the similarity between the corresponding feature vectors . This permits us to extract detailed relations from the support set , with the capability of modeling complex , non-linear mappings . As a non-parametric model , the GP further effectively benefits from additional support samples , since all given data is retained . As shown in Fig . 1 , the segmentation accuracy of our approach improves consistently with the number of support samples . Lastly , the predictive covariance from the GP provides a principled measure of the uncertainty based on the similarity with local features in the support set . Our FSS approach is learned end-to-end through episodic training , treating the GP as a layer in a neural network . This further enables us to learn the output space of the GP . To this end , we encode the given support masks with a neural network in order to achieve a multi-dimensional output representation . In order to generate the final masks , our decoder module employs the predicted mean query encodings , together with the covariance information . Our decoder is thus capable of reasoning about the uncertainty when fusing the predicted mask encodings with learned segmentation priors . Lastly , we further improve our FSS method by integrating dense GPs at multiple scales . We perform comprehensive experiments on two benchmarks : PASCAL-5i ( Shaban et al. , 2017 ) and COCO-20i ( Nguyen & Todorovic , 2019 ) . Our proposed DGPNet outperforms existing methods for 1-shot and 5-shot by a large margin , setting a new state-of-the-art on both benchmarks . When using the ResNet101 backbone , our DGPNet achieves an absolute gain of 14.9 for 5-shot segmentation on the challenging COCO-20i benchmark , compared to the best reported results in the literature . We further demonstrate the cross-dataset transfer capabilities of our DGPNet approach from COCO-20i to PASCAL and perform detailed ablative studies to probe the effectiveness of our contributions . 2 RELATED WORK . Few-Shot Segmentation The earliest work in few-shot segmentation ( FSS ) , by Shaban et al . ( 2017 ) , proposed a method for predicting the weights of a linear classifier based on the support set , which was further built upon in later works ( Siam et al. , 2019 ; Liu et al. , 2020a ; Boudiaf et al. , 2021 ) . Instead of learning the classifier directly , Rakelly et al . ( 2018 ) proposed to construct a global conditioning prototype from the support set and concatenate it to the query representation , with several subsequent works ( Dong & Xing , 2019 ; Zhang et al. , 2020 ; Wang et al. , 2019 ; Nguyen & Todorovic , 2019 ; Zhang et al. , 2019b ; Liu et al. , 2020c ; Azad et al. , 2021 ; Liu et al. , 2020b ; Xie et al. , 2021 ; Wang et al. , 2021 ) . A major limitation of these methods is the unimodality assumption . To alleviate this problem , Zhang et al . ( 2021 ) construct additional prototypes by a self-guided module , while Yang et al . ( 2020a ) ; Liu et al . ( 2020c ) ; Li et al . ( 2021 ) instead cluster multiple prototypes to create a richer representation . However , clustering introduces extra hyperparameters , such as the number of clusters , as well as optimization difficulties . In contrast , our method is not restricted in any such sense , and only requires us to choose an appropriate kernel . Some recent works consider pointwise correspondences between the support and query set . These works have mostly focused on attention or attention-like mechanisms ( Zhang et al. , 2019a ; Yang et al. , 2020b ; Hu et al. , 2019 ; Tian et al. , 2020 ; Wang et al. , 2020 ) . In contrast with these methods , we construct a principled posterior over functions , which greatly aids the decoder . Combining GPs and Neural Networks While early work focused on combining GPs and neural networks in the standard supervised classification setting ( Salakhutdinov & Hinton , 2009 ; Wilson et al. , 2016 ; Calandra et al. , 2016 ) , there has recently been an increased interest in utilizing Gaussian processes in the context of few-shot classification ( Patacchiola et al. , 2020 ; Snell & Zemel , 2021 ) . Previous works employ the GP in the classification setting and as the final output layer in the network , and optimize proxies of either the predictive or marginal log likelihood directly . Note that the classification likelihood is non-gaussian , and hence computing the exact posterior and marginal likelihoods of the model becomes intractable . Here , we go beyond this limitation and propose an internal dense GP model of the support features , where the posterior predictive distribution is used as input to a CNN decoder . Moreover , this allows us to learn the output space of the GP to further increase its expressive power . To the best of our knowledge , we are the first to introduce a dense GP approach for the challenging dense prediction task of few-shot segmentation . 3 METHOD . 3.1 FEW-SHOT SEGMENTATION . Few-shot segmentation is a dense few-shot learning task ( Shaban et al. , 2017 ) . The aim is to learn to segment objects from novel classes , given only a small set of annotated images . A single instance of this problem , referred to as an episode , comprises a small set of annotated samples , called the support set , and a set of samples on which prediction is to be made , the query set . Formally , we denote the support set as { ( ISk , MSk ) } Kk=1 , comprising K image-mask pairs ISk 2 RH0⇥W0⇥3 and MSk 2 { 0 , 1 } H0⇥W0 . A query image is denoted as IQ 2 RH0⇥W0⇥3 and the aim is to predict its corresponding segmentation mask MQ 2 { 0 , 1 } H0⇥W0 . To develop our approach , we first provide a general formulation for addressing the FSS problem , which applies to several recent methods , including prototype-based ( Li et al. , 2021 ; Liu et al. , 2020a ) and correlation-based ( Tian et al. , 2020 ; Wang et al. , 2020 ) ones . Our formulation proceeds in three steps : feature extraction , few-shot learning , and prediction . In the first step , deep features are extracted from the given images , x = F ( I ) 2 RH⇥W⇥D . ( 1 ) These features provide a more disentangled and invariant representation , which greatly aids the problem of learning from a limited number of samples . The main challenge in FSS , namely how to most effectively leverage the annotated support samples , is encapsulated in the second step . As a general approach , we consider a learner module ⇤ that employs the support set in order to find a function f , which associates each query feature xQ 2 RD with an output yQ 2 RE . The goal is to achieve an output yQ that is strongly correlated with the ground-truth query mask , allowing it to act as a strong cue in the final prediction . Formally , we express this general formulation as , f = ⇤ ( { xSk , MSk } k ) , yQ = f ( xQ ) . ( 2 ) The learner ⇤ aggregates information in the support set { xSk , MSk } k in order to predict the function f . This function is then applied to query features equation 2 . In the final step of our formulation , output from the function f on the query set is finally decoded by a separate network as cMQ = U ( yQ , xQ ) to predict the segmentation cMQ . The general formulation in equation 2 encapsulates several recent approaches for few-shot segmentation . In particular , prototype-based methods , for instance PANet ( Wang et al. , 2019 ) , are retrieved by letting ⇤ represent a mask-pooling operation . The function f thus computes the cosine-similarity between the pooled feature vector and the input query features . In general , the design of the learner ⇤ represents the central problem in few-shot segmentation , since it is the module that extracts information from the support set . Next , we distinguish three key desirable properties of this module . 3.2 MOTIVATION . As discussed above , the core component in few-shot segmentation is the few-shot learner ⇤ . Much research effort has therefore been diverted into its design ( Nguyen & Todorovic , 2019 ; Liu et al. , 2020a ; Wang et al. , 2020 ; Liu et al. , 2020c ; Yang et al. , 2020a ; Li et al. , 2021 ) . To motivate our approach , we first identify three important properties that the learner should possess . Flexibility of f The intent in few-shot segmentation is to be able to segment a wide range of classes , unseen during training . The image feature distributions of different unseen classes are not necessarily linearly separable ( Allen et al. , 2019 ) . Prototypical few-shot learners , which are essentially linear classifiers , would fail in such scenarios . Instead , we need a mechanism that can learn and represent more complex functions f . Scalability in support set size K An FSS method should be able to effectively leverage additional support samples and therefore achieve substantially better accuracy and robustness for larger support sets . However , many prior works show little to no benefit in the 5-shot setting compared to 1-shot . As shown by Li et al . ( 2021 ) and Boudiaf et al . ( 2021 ) , it is crucial that new information is effectively incorporated into the model without averaging out useful cues . Uncertainty Modeling Since only a small number of support samples are available in FSS , the network needs to regularly handle unseen appearances in the query images . For instance , the query may include novel backgrounds , scenes , and objects . Since the function f predicted by the learner is not expected to generalize to such unseen scenarios , the network should instead utilize neighboring predictions or learned priors . However , this can only be achieved if f models and communicates the uncertainty of its prediction to the decoder .

Summary: Authors propose a novel few-shot segmentation method by adopting dense Gaussian process (GP) regression to capture complex appearance distributions. To boot the performance, authors consider the uncertainty in the final segmentation. Authors exploit the end-to-end learning capabilities of the proposed method to learn a high-dimensional output space for the GP.  Authors report state-of-the-art results in two public few shot segmentation benchmarks.