Q-Net: Query-Informed Few-Shot Medical Image Segmentation

Deep learning with CNNs has become the dominated approach to MIS on various tissues, anatomical structures and lesions. Fully Convolutional Networks (FCNs) stand out to be a powerful deep net architecture for semantic segmentation Long, Shelhamer, and Darrell (2015). FCNs replace the fully connected layers of standard CNNs with fully convolutional layers. Given an arbitrary-sized input image, FCNs can produce pixel-wise output of consistent size. Afterwards, the encoder-decoder networks become the major architecture for semantic segmentation, among which U-Net is well-recognized as a SOTA network for MIS Ronneberger, Fischer, and Brox (2015). U-Net adopts a symmetrical encoder-decoder architecture infused with skipped connections. In the downsampling path, it applies CNN, while in the upsampling path, it enhances the feature map, making its size be consistent with that of the input.

Following U-Net, several FCNs have been developed, such as U-Net 3D Çiçek et al. (2016), V-Net Milletari, Navab, and Ahmadi (2016), Y-Net Mehta et al. (2018), attention U-Net Oktay et al. (2018) and nnU-Net Isensee et al. (2018), which can be regarded as different variants of U-Net. All these CNN type models require abundant expert-annotated data. When only a few labeled images of an “unseen” class are available for segmenting an object of this class, these methods fail to provide good performance. When facing data-scarce cases, one needs to integrate data augmentation or transfer learning techniques to such methods to make them work.

Different from supervised learning methods that require a large number of labeled samples to work, few-shot learning (FSL) methods aim to learn with few labeled samples. Thus, it is natural to adapt FSL methods to deal with MIS, leading to FSS methods. Ouyang et al. (2020) propose an adaptive local prototype pooling module to add extra focus on the background. Hansen et al. (2022) use a single foreground prototype, and view image segmentation from an anomaly detection perspective. Rakelly et al. (2018) use weak annotations by replacing the pre-defined binary mask with a few selected landmarks of foreground and background. Most of these FSS methods aim to fully exploit information covered in the support set and assume that the labeled training data is abundant.

The prototypical network, which uses masked average pooling to extract class-wise features from the support set, is another major category of FSS methods featured by its strong interpretability and robustness against noises. Dong and Xing (2018) adopt the meta-learning mechanism in prototypical networks for semantic segmentation. The resulting model contains a prototype network for learning class-specific prototypes and a segmentation network for making predictions on query images. Wang et al. (2019) formulate image segmentation as an non-parametric prototype matching process, then propose prototype alignment network (PANet), which inversely predicts labels of the support images by using query images as the support set. Liu et al. (2020) improve PANet by introducing additional prototypes to capture more diverse features of the semantic classes. Base on PANet, Ouyang et al. (2020) and Hansen et al. (2022) employ superpixel/supervoxel segmentation for a self-supervised prototypical FSS on medical images. Tang et al. (2021) propose a context relation encoder and mask refinement module to iteratively refine the segmentation.

Another branch of medical FSS methods follows a two-arm architecture of Shaban et al. (2017), which consists of a conditioner arm and a segmenter arm, corresponding to the support set and the query data, respectively. Roy et al. (2020) employ dense connections with squeeze and excitation blocks to strengthen the interaction between the conditioner arm and the segmentor arm. Sun et al. (2022) present a global correlation network with discriminative embedding, obtaining an improved performance.

The aforementioned works do not explicitly take into account possible distribution shifts between the query image and the support set, while such shifts indeed exist in some practical MIS scenarios. Our proposed Q-Net learns to borrow information from the query image based on a meta-learning setting to remove or at least mitigate such distribution shifts.

We employ a shared feature extractor $f_{\theta }$ to extract multi-scale features from both the support and the query images in the embedding space. We denote the support features as ${\text{F}}^s=f_{\theta }\left({\text{x}}^s\right),{\text{F}}^s\in R^{H^{\prime }\times W^{\prime }\times Z}$, where $H^{\prime },W^{\prime }$ denote the height and the width of the features and $Z$ the channel depth. The query features ${\text{F}}^q=f_{\theta }\left({\text{x}}^q\right)$ is defined in the same way.

Following the common practice, we use ResNet-101 with pre-trained MS-COCO weights as the backbone net for feature extraction. It is composed of a $7\times 7$ convolutional layer with stride of 2, a max pooling layer, and four residual blocks with multiple 3-layer bottleneck blocks. The output features of the first residual block in the network are 4 times smaller than the input resolution. The network can provide another 3 times down-sampling in the last 3 residual blocks with dilation. In contrast to previous works Wang et al. (2019); Ouyang et al. (2020); Hansen et al. (2022), where the network outputs features that are 8 times smaller than input resolution by using dilation in the second residual block, we use dilation in the last residual block and extract features with two different scales from the last two residual blocks. This results in the resolution of multi-scale features being 1/4 and 1/8 of the image resolution $(H^{\prime }=H/4orH/8,W^{\prime }=W/4orW/8)$. The multi-scale feature extraction can better capture class prototypes of tissues of different sizes and shapes.

Following Hansen et al. (2022), we consider a single foreground prototype to avoid introducing additional background prototypes in each training episode. The prototypical prediction is undertaken by non-parametric metric learning. First, we compute the foreground prototype $\text{p}\in R^{1\times Z}$ from the support feature map ${\text{F}}^s$ via masked average pooling:

where the support feature map ${\text{F}}^s$ is resized to the mask size $(H,W)$, and $\odot$ denotes the Hadamard product for masked average pooling.

In order to perform soft thresholding on the predicted mask, we employ a shifted Sigmoid on the negative cosine similarity between the query feature vectors and the foreground prototype, as follows

where $(h,w)$ denotes the spatial locations of the foreground mask, $a=20$ a scaling factor Oreshkin, Rodríguez López, and Lacoste (2018), and $T$ the learned threshold. The subscripts ${}_f$ and ${}_b$ denote “foreground” and “background”, respectively.

In common practice of FSS, when $T$ is learned in the meta-training stage, it is then directly used in the meta-testing stage to segment unseen foreground objects. In the standard setting, the background may contain novel test objects, then the supervoxels overlapped with the test labels might be used as the training classes. Compared to the spatially heterogeneous background, the foreground objects, such as the right and the left kidneys in the abdominal MRI image, are homogeneous, thus can be easily clustered into the same supervoxel. The learned threshold used for anomaly detection has been optimized on supervoxels that share similar shape and size to the “unseen” objects. For such cases, using a fixed $T$ makes sense. However, if image slices containing the test classes are removed, then the test classes becomes completely unseen classes for the algorithm agent, then the threshold learned in the training process can not be adapted straightforward to such “unseen” classes.

Following the QIL strategy, we add two fully-connected layers $g_{\varphi }$ at the top of the feature extractor $f_{\theta }$ to learn a soft threshold from the query image instead of the predicted foreground masks on the seen classes:

where the fully-connected layer maps the number of channels of the flattened feature maps from 2048 to one. See more details in the Appendix Section.

Given class prototypes extracted from the support images, we use cosine similarity to make segmentation predictions on the foreground and background. This prediction process assumes that a query image holds the similar appearance with support images. During the training phase, the image slices adjacent to the support image are selected as query samples, while during the inference phase, the query image and the support set can be selected from different image volumes. Therefore, a query image can be quite different from the support images on intensity and appearance. Following the QIL strategy, We propose a prototype refinement module to update the class prototypes by borrowing information from the query image features in test time.

In prototype refinement, firstly we produce the predicted mask by calculating the cosine similarity between the query image features and the class prototypes. At each iteration, we apply element-wise multiplication with the query feature maps and the background mask ${\hat{\text{m}}}_b^q$ to obtain new feature maps ${\hat{\text{F}}}^q$. Then, we allocate the single foreground prototype p as a feature vector to the new feature maps.

We design this module to mimic the steps of an optimization algorithm in inference time. This module is trained to modify the foreground prototypes p gradually so that the final predicted mask ${\hat{\text{m}}}^q$ converges to an optimum solution.

where $L({\hat{\text{F}}}^q,{\text{F}}^q)$ is the cross-entropy loss computed for the new query features ${\hat{\text{F}}}^q$ to original query features ${\text{F}}^q$, and $N$ denotes the number of iterations of gradient back-propagation. ${\text{p}}_N$ denotes the optimal prototype after $N$ iterations, thus, an optimal predicted mask ${\tilde{\text{m}}}_f^q$ is generated.

Therefor, the final segmentation mask ${\tilde{\text{m}}}^q$ is generated by element-wise summation of predicted masks based on multi-scales in inference time, and $\alpha \in (0,1.0)$ is a coefficient for prediction based on $64\times 64$ features.

We train our network by minimizing the cross-entropy between the predicted masks and the ground-truth segmentation ${\text{m}}^q$. The dual-path predicted masks after anomaly threshold adaptation are up-sampled to $(H,W)$ and then combined through element-wise summation, as follows

where $\alpha$ is a balance factor, whose value is set empirically, see Table 3 for an ablation analysis.

We compute the binary cross-entropy loss as the segmentation loss for each training episode:

Following Wang et al. (2019); Ouyang et al. (2020), we inversely predict labels of the support images by using query images as the support set, then build a prototypical alignment regularization item as follows

Overall, the loss function for each training episode is defined to be

For more implementation details on prototype refinement, readers are referred to the Appendix Section.