Doubly Deformable Aggregation of Covariance Matrices for Few-shot Segmentation

Semantic segmentation at the pixel level [47, 42, 19, 6] is one of the fundamental tasks in computer vision and has been extensively studied for decades. In recent years, the performance of semantic segmentation tasks has been significantly improved due to the substantial progress in deep learning techniques. Large-scale convolutional neural networks (CNNs) [28, 8], vision Transformers [3, 18], and MLP-based deep networks [15] have greatly improved the ability of visual representation learning, and significantly enhanced the performance of downstream tasks such as semantic segmentation [48, 47].

However, these advances currently rely heavily on large-scale annotated datasets that require extensive manual annotation efforts. This is not conducive to real-world computer vision applications, as obtaining pixel-level annotations is expensive and time-consuming. To remedy this issue, the few-shot segmentation task has emerged and attracted more and more research attention [41, 37, 17, 43, 10, 44]. For few-shot segmentation, only a handful of annotated samples (support samples) for each class are provided for the training of the model, which largely mitigates the reliance on large-scale manual annotations. Although an attempt has been made [1] to handle few-shot segmentation in a transductive way, most existing works model the few-shot segmentation in a meta-learning setting [26, 34] to avoid overfitting brought by insufficient training data. As support samples are the most important guidance for the prediction of query samples, the key to few-shot segmentation is to effectively exploit information from the support set.

Prototype-based few-shot segmentation methods attempt to extract a prototype vector to represent the support sample [35, 31, 41]. Despite its effectiveness, compressing the features of support samples into a single vector is prone to noise and neglects the spatial structure of objects in the support image. To maintain the geometric structure of objects, visual correspondence-based methods [20, 10, 44] opt to utilize dense, pixel-wise semantic correlations between support and query images. Sophisticated pixel-wise attention-based methods [45] have also been devised to leverage dense correlations between support and query samples. However, directly using pixel-level dense information is a double-edged sword, which brings two new challenges: feature ambiguities caused by repetitive patterns, and background noise. To handle these issues, recent work [20] shows that aggregating dense correlation scores (also called cost volume) with 4D convolution can attain outstanding performance for few-shot segmentation. To better model the interaction among correlation scores, [9] proposes to combine 4D convolutions and 4D swin Transformer [18] in a coarse-to-fine manner. Although recent cost aggregation-based methods [20, 9] attain state-of-the-art performance, they still suffer from two limitations: 1) lack of flexibility in computing correlation scores, which results in an extremely slow convergence speed; 2) lack of flexibility in aggregating high-dimensional correlation scores, which leads to limited performance with high computational cost.

To tackle the first limitation, we propose a Gaussian process (GP) based kernel learning method to learn flexible kernel functions for a more accurate similarity measurement. We find that existing state-of-the-art methods [20, 9] mainly focus on the cost volume aggregation module and directly rely on the hypercorrelation tensors constructed by the cosine similarity. However, we argue that directly using the cosine similarity lacks flexibility and cannot faithfully reflect the semantic similarities between support and query pixels. As shown in Fig. 1 (a), the match between the cosine similarity map and the ground truth label is poor. To handle this problem, we design a hard example mining mechanism to dynamically sample hard samples to train the covariance kernel functions. As shown in Fig. 1 (b), using the learned covariance kernels can generate a more reasonable similarity map and greatly improve the convergence speed.

Another limitation is mainly caused by the cost volume aggregation module. Prior works have shown that a better cost volume aggregation method can attain superior few-shot segmentation performance. However, the existing 4D convolution-based method [20] lacks the ability to model longer distance relations between elements in the cost volume. VAT [9] proposes to combine 4D convolutions with 4D swin Transformers [18] for cost volume aggregation. Although outstanding performance can be achieved, it requires a notable GPU memory consumption and suffers from a slow convergence speed. Considering this, we propose a doubly deformable 4D Transformer based aggregation method, with deformable attention mechanisms on both the support and query dimensions of the 4D cost volume input. Compared with 4D convolutions, it can not only model longer-distance interactions between pixels owing to the Transformer network design, but also learn to selectively attend to more informative regions of the support and query information.

To sum up, we propose a novel few-shot segmentation framework using doubly deformable aggregation of covariance matrices (DACM), which aggregates learnable covariance matrices with a doubly deformable 4D Transformer to predict segmentation results effectively. The proposed GP-based kernel learning can learn more accurate similarity measurements for cost volume generation. In what follows, the designed doubly deformable 4D Transformer can effectively and efficiently aggregate the multi-scale cost volume into the final segmentation result. Specifically, our contributions can be summarized as follows:

• •

  Towards a more accurate similarity measurement, a GP-based kernel learning method is proposed to learn flexible covariance kernel functions, with a novel hard example mining mechanism. To our knowledge, we are the first to use learnable covariance functions instead of cosine similarity for the few-shot segmentation task.

• •

  Towards a more flexible cost volume aggregation, we propose a doubly deformable 4D Transformer (DDT) module, which utilizes deformable attention mechanisms on both the support and query dimensions of the 4D cost volume input. DDT can enhance representation learning by selectively attending to more informative regions.

• •

  By combining these two modules, the proposed DACM method can attain new state-of-the-art performance on three public datasets with an extremely fast convergence speed. We also provide extensive visualization to better understand the proposed method.

The whole architecture of the proposed DACM framework is illustrated in Fig. 2. Given the input support and query images, multi-level deep features are first extracted by the fixed backbone network pre-trained on ImageNet [25]. Three levels of deep features with different spatial resolutions form a pyramidal design. For each level, average pooling is used to aggregate multiple features into a single one. Then, three different GP models are trained for computing the covariance matrices (4D cost volume) for each level. Next, the proposed Deformable 4D Transformer Aggregator (DTA) is combined with the weight-sparsified 4D convolution for cost volume aggregation. Finally, a decoder is used to predict the final segmentation result for the input query image.

In this paper, we present a few-shot segmentation method by doubly deformable aggregating covariance matrices. The proposed method aggregates learnable covariance matrices with a doubly deformable 4D Transformer to predict segmentation results effectively. Specifically, we make the following contributions. 1) We devise a novel hard example mining mechanism for learning covariance kernels of the Gaussian process to enable a more accurate correspondence measurement. 2) We design a doubly deformable 4D Transformer to effectively and efficiently aggregate the multi-scale cost volume into the final segmentation results. 3) By combining these two modules, the proposed method can achieve state-of-the-art few-shot segmentation performance with a fast convergence speed.

This work is jointly supported by the European Research Council (ERC) under the European Union’s Horizon 2020 research and innovation program (grant agreement No. [ERC-2016-StG-714087], Acronym: So2Sat), by the Helmholtz Association through the Framework of Helmholtz AI (grant number: ZT-I-PF-5-01) - Local Unit “Munich Unit @Aeronautics, Space and Transport (MASTr)” and Helmholtz Excellent Professorship “Data Science in Earth Observation - Big Data Fusion for Urban Research”(grant number: W2-W3-100), by the German Federal Ministry of Education and Research (BMBF) in the framework of the international future AI lab ”AI4EO – Artificial Intelligence for Earth Observation: Reasoning, Uncertainties, Ethics and Beyond” (grant number: 01DD20001) and by German Federal Ministry of Economics and Technology in the framework of the ”national center of excellence ML4Earth” (grant number: 50EE2201C).

Supplementary Material Zhitong Xiong Haopeng Li Xiao Xiang Zhu

Three datasets for few-shot segmentation are exploited to evaluate the proposed method, including the PASCAL-$5^{i}$ [26], COCO-$20^{i}$ [14], and FSS-1000 [13]. As we can see in Fig. 10, there exist large variations between the support images and the query images in scale, pose and appearance.

Although only 50 epochs are used for training the proposed DACM method, the performance of our model is clearly better than that of the baseline method for some difficult test samples. Compared with HSNet [20], qualitative results in Fig. 10 show that DACM can obtain more complete segmentation maps owing to larger receptive fields of the Transformer architecture. In Fig. 11, we also visualize more examples of sampled locations during the training process of our DACM method.

Different from HSNet [20], we use the learned covariance kernel functions to measure the similarity between the support features and the query features instead of a fixed cosine similarity. The fixed cosine similarity is computed as follows:

where $\bm{x}_{q}$ is the query feature, and $\bm{z}_{s}$ is the masked support feature.

Theoretically, the cosine similarity used in [20] can be viewed as a special case of our method, i.e., using a fixed linear kernel. For a more effective similarity measurement, we target at learning the kernel functions $k(\cdot,\cdot)$ on different datasets using Gaussian process. The covariance matrices can be computed as follows:

To better understand the learned kernel functions in GP, we have visualized the covariance matrices in Fig. 12. For the cosine similarity map $\bm{C}_{1}\in\mathbb{R}^{H_{q}\times W_{q}\times H_{s}\times W_{s}}$, we reshape it into $\bm{C}_{1}\in\mathbb{R}^{H_{q}\times W_{q}\times H_{s}W_{s}}$ and sum it up along the third dimension. Then we can get 2D similarity maps of the query image. The 2D similarity maps are shown in Fig. 12. For the covariance matrices $\bm{C}_{2}\in\mathbb{R}^{H_{q}\times W_{q}\times H_{s}\times W_{s}}$, we conduct the similar computation as the cosine similarity map $\bm{C}_{1}$.

From Fig. 12 we can see that covariance matrices learned by the DACM model are more reasonable than the simple cosine similarity map. Then, the learned similarity maps can be processed by the proposed DDT model to further enhance the performance of few-shot segmentation.

In this subsection, we will visualize and analyze some failure cases of the proposed DACM method. In Fig. 13, we can see some failure cases of the learned covariance matrices. It can be seen that, the similarity between the support and the background of query samples is clearly low. However, in some cases, the similarity between the support and foreground objects of query samples is incorrectly high. From the visualization results, we can find that although the global representation learning of Transformer is beneficial in most cases, there still exist some situations where smaller receptive fields work better.

In Fig. 14, we also visualize some segmentation failures. As for the top-left image, the man holding the bottle is misclassified as a foreground object. Take the top-right result as an example, DACM incorrectly classify the reflection of the boat as a foreground object. Note that the baseline method obtains better mIoU results on these visualized samples than DACM. In the future work, we believe that studying how to effectively combine global and local representations can further enhance the performance of the few-shot segmentation task.

Gaussian process (GP) is a non-linear and non-parametric Bayesian model for regression and classification [38]. By exploiting the correlation between data samples, Gaussian process performs probabilistic non-linear prediction that is described by Gaussian distribution with estimated mean and covariance. We explain the details of general GP as follows. Formally, the data set consists of $N$ samples of dimension $D$, i.e., $\{X=\{{x}_{i}\}_{i=1}^{N},\bm{y}=\{y_{i}\}_{i=1}^{N}\}$, where ${x}_{i}\in\mathbb{R}^{D}$ is a data point and $y_{i}$ is the corresponding label. The GP regression model assumes that the outputs $y_{i}$ can be regarded as certain deterministic latent function $f({x}_{i})$ with zero-mean Gaussian noise $\varepsilon$, i.e., $y_{i}=f({x}_{i})+\varepsilon$, where $\varepsilon\sim\mathcal{N}(0,\sigma^{2})$. GP sets a zero-mean prior on $f$, with covariance $k({x}_{i},{x}_{j})$. The covariance function (kernel) $k$ reflects the smoothness of $f$. Commonly-used kernels can be found in Section 11. The hyper-parameters in the kernels are optimized by maximizing the marginal likelihood of the training data, which is given as follows,

where $\bm{f}=[f(x_{1}),f(x_{2}),\cdots,f(x_{N})]^{\mathrm{T}}$. Here, the term marginal likelihood refers to the marginalization over the function values $\bm{f}$. Setting the prior of the Gaussian process model $p(\bm{f}|X)$ to be a Gaussian distribution $\mathcal{N}(\bm{0},\bm{K})$, the marginal log-likelihood $\log p(\bm{y}|X)$ can be expressed by,

where $\bm{y}=[y_{1},y_{2},\cdots,y_{N}]^{\mathrm{T}}$, $\bm{K}_{ij}=k({x}_{i},{x}_{j})$, and $\bm{I}_{N}$ is the identity matrix of size $N$. For example, the Automatic Relevance Determination Squared Exponential (ARD SE) kernel function is defined as

where $\sigma^{2},\sigma_{0}^{2},\left\{l_{d}\right\}_{d=1}^{D}$ are hyper-parameters. By maximizing Eq. 12, the hyper-parameters in Eq. 13 for computing the covariance matrices can be optimized.

After the hyper-parameters are optimized, the predictive distribution for the test case $x^{*}$ can be calculated in a closed-form as follows:

where $\bm{k}=[k(x_{1},x^{*}),k(x_{2},x^{*}),\cdots,k(x_{N},x^{*})]^{\mathrm{T}}$ and $k^{*}=k(x^{*},x^{*})$.

In this section, we present the Pytorch-style implementation of the GP model. The most critical part of the GP model is the choice of covariance kernel functions. In algorithm 1, we show three types of kernel functions, including the linear kernel, the radial basis function (RBF) kernel and the additive mixed kernel.

For two vectors $\mathbf{x}_{\mathbf{1}},\mathbf{x}_{\mathbf{2}}$ in the feature map, the linear kernel, the RBF kernel, and the additive mixed kernel, with hyperparameters $\{v,\Theta\}$, are defined as follows:

Although the RBF kernel is used in this work, here we show that the proposed framework can be easily extended by using different kernel types or their combinations.