Robust Prototypical Few-Shot Organ Segmentation with Regularized Neural-ODEs

Deep learning models such as ResNets [38] learn a sequence of transformations by mapping input x to output y by composing a sequence of transformations to a hidden state. In a ResNet block, computation of a hidden layer representation can be expressed using the following transformation:

where $\mathrm{t\in\{0,\ldots,T\}}$ and $\mathrm{\textbf{h}}:[0,\infty]\rightarrow\mathbb{R}^{n}$. As the number of layers is increased and smaller steps are taken, in the limit, the continuous dynamics of the hidden layers are parameterized using an ordinary differential equation (ODE) [25] specified by a neural network,

where $f:\mathbb{R}^{n}\times[0,\infty]\rightarrow\mathbb{R}^{n}$ denotes the non-linear trainable layers parameterized by weights $\mathbf{\theta}$. These layers define the relation between the input $\mathrm{h(0)}$ and output $\mathrm{h(T)}$, at time $\mathrm{T}>0$, by providing a solution to the ODE initial value problem at terminal time T. Neural-ODEs are the continuous equivalent of ResNets where the hidden layers can be regarded as discrete-time difference equations. The time here is a dummy variable loosely corresponding to the equivalent of layer number in ResNets, though it may not have a physical interpretation. The curve h(t), obtained by solving Eq. 2, is termed as an integral curve and denotes the $n$-dimensional features at any intermediate time t. Different initial values h(0) lead to different particular solutions which can be thought of as different integral curves obtained as different inputs are passed through the Neural-ODE.

Multiple works have tried to explain the intrinsic robustness of Neural-ODEs. [25] argues that the non-intersecting property of the Neural-ODE’s integral curves allows the model to be intrinsically robust. Because the integral curves are governed by differential equations of the form shown in Eq. 2, it follows that if two integral curves intersect, their slopes at the point of intersection must be identical. Following the curves back to the y-axis (Fig. 2), it can be proved that any two intersecting integral curves have to be identical, or equivalently, distinct integral curves cannot intersect. As can be seen in Fig. 2, this leads to restrictions on the possible integral curves and this consequently improves robustness as illustrated in Fig. 1. We elaborate on this property, its proof and its consequences on robustness with a supplementary video presentation.

As an alternate theory to account for this intrinsic robustness, [39] shows that adaptive stepsize ODE solvers commonly used in Neural-ODEs tend to have a gradient-masking effect, and this allows the model to be robust to gradient-based attacks. Multiple studies [25, 40, 41] have applied Neural-ODEs to defend against adversarial attacks. [42] proposes time-invariant steady Neural-ODE that is more stable than conventional convolutional neural networks (CNNs) in the classification setting. [41] design a Neural-ODE such that the equilibrium points of the ODE solution have Lyapunov-stability, thereby suppressing input perturbations. [40] explores the equivalents of different regularising noise injection techniques like dropout and gaussian smoothing in Neural-ODEs, and argues how these techniques lead to better, more stable equilibrium points. [43] describe a framework that guarantees exponential convergence and improved adversarial robustness by using a novel Lyapunov loss, which they approximately optimise using a practical Monte Carlo approach.

Few-shot segmentation (FSS) [44, 4, 45, 12, 8, 10, 11] aims to perform pixel-level classification for novel classes in a query image when trained on only a few labelled support images. A commonly adopted approach for FSS is based on prototypical networks [5, 7, 4, 46] that employ prototypes to represent typical information for foreground objects present in the support images. The prototype is subsequently compared with the query images pixel-wise via cosine similarity, leading to the segmentation mask. In addition to the prototype-based setting, [44] incorporates ‘squeeze & excite’ blocks that avoid the need for pre-trained models for medical image segmentation. [45] uses a relation network and introduced the FSS-1000 dataset that is significantly smaller as compared to contemporary large-scale datasets for FSS. The data-scarce medical domain has a unique requirement for specialized FSS algorithms [47, 46]. One common problem setting in this domain is that of organ segmentation, with openly available organ segmentation datasets [35, 36, 37] and specialized methods [48, 6] addressing their unique challenges. [48] used a bidirectional gated recurrent unit (GRU) to capture relationships of features across slices. [6] uses an interclass segmentation network for organ segmentation on a multi-institution dataset from prostate cancer patients. Many other methods [49, 50, 51] addressing different unique challenges can be found in literature, but we focus more on the commonly used prototypical networks. MetaNODE [52] was one of the first works to use Neural-ODEs in the Few-shot domain. Our work is quite different to MetaNODE as the latter addresses prototype bias reduction. Also, the core technique varies in these two cases. MetaNODE uses ODEs for meta-optimization on mean prototypes while we use ODEs for generating a representative prototype for each class.

Adversarial attacks for natural image classification have been extensively explored, and interest is turning towards the effects of adversarial attacks in the medical domain [18, 53]. FGSM [13] and PGD [15] generate adversarial examples based on the CNN gradients. Besides image classification, several attack methods have also been proposed for semantic segmentation [54, 33, 34, 55]. [33] introduced Dense Adversary Generation (DAG) that optimizes a loss function over a set of pixels for generating adversarial perturbations. Recently, [34] proposed an adversarial attack (SMIA) for images in the medical domain that employs a loss stabilization term to exhaustively search the perturbation space. While adversarial attacks expose the vulnerability of deep neural networks, adversarial training [15, 13, 17] is effective in enhancing the target model by training it with adversarial samples. However, none of the existing methods has explored SAT procedure for few-shot semantic segmentation. In addition to the explicit use of adversarial samples during training, many methods have also explored the use of simple and fast Gaussian perturbations [26, 28, 27, 29, 30] as effective regularisation mechanisms for boosting adversarial robustness. Augmenting the training set with examples perturbed using Gaussian noise to increase adversarial robustness has been suggested in [56]. [26] uses Renyi divergence to illustrate the connection of robustness to random noise. [30] proposes a defence method against query-based black-box attacks, Random Noise Defense (RND), which adds proper Gaussian noise to each query. RND can be combined with existing defence methods like adversarial training to further boost the adversarial robustness.

Other than external procedures to make models robust, recent studies also propose models that are intrinsically robust. [57] uses a Feature Pyramid Decoder with denoising and image restoration that can be applied to any general CNN, improving robustness intrinsically. [58] uses blocks that denoise intermediate features, effectively restoring perturbed features back to the clean ones. The intrinsic adversarial robustness of Neural-ODEs [25, 41, 40, 43, 39] is being increasingly explored in recent years, making it an exciting field of research.

FSS setting includes train $\mathcal{D}_{\mathrm{train}}$ and test $\mathcal{D}_{\mathrm{test}}$ datasets having non-overlapping class sets. Each dataset consists of a set of episodes with each episode containing a $N$-way $K$-shot task $\mathcal{T}_{i}$ = ($\mathcal{S}_{i},\mathcal{Q}_{i}$) where $\mathcal{S}_{i}$ and $\mathcal{Q}_{i}$ is support and query sets for the $i^{th}$ episode having a class set $C_{i}$. Formally, $\mathcal{D_{\mathrm{train}}}=\{(\mathcal{S}_{i},\mathcal{Q}_{i})\}^{E_{\mathrm{train}}}_{i=1}$ and $\mathcal{D_{\mathrm{test}}}=\{(\mathcal{S}_{i},\mathcal{Q}_{i})\}^{E_{\mathrm{test}}}_{i=1}$ where $E_{\mathrm{train}}$ and $E_{\mathrm{test}}$ denote the number of episodes during training and testing. The support set $\mathcal{S}_{i}$ has $K$ image $(\mathcal{I}_{\mathcal{S}})$, mask $(L_{\mathcal{S}})$ pairs per class with a total of $N$ semantic classes i.e. $\mathcal{S}_{i}=\{(\mathcal{I}_{\mathcal{S}}^{k},L_{\mathcal{S}}^{k}(\eta))\}$ where $L^{k}_{\mathcal{S}}(\eta)$ is the ground-truth mask for $k$-th shot corresponding to class$~{}\eta\in C_{i}$, $|C_{i}|=N$ and $k=1,2,\ldots,K$. The query set $\mathcal{Q}_{i}$ has $N_{\mathcal{Q}}$ $\mathcal{\mathrm{image}~{}(I_{Q})},{\mathrm{mask}~{}(L_{\mathcal{Q}})}$ pairs. The FSS model $\mathcal{F}(.)$ is trained on $\mathcal{D}_{\mathrm{train}}$ across the episodes with support sets and query images as inputs, and predicts the segmentation mask $M_{\mathcal{Q}}$ = $\mathcal{F}(\mathcal{S}_{i},\mathcal{I}_{\mathcal{Q}})$ in the $i$-th episode for query image $\mathcal{I}_{\mathcal{Q}}$. During testing, the trained model $\mathcal{F}(.)$ is used to predict masks for unseen novel classes with the corresponding support set samples and query images as inputs from $\mathcal{D}_{\mathrm{test}}$. Further, the trained FSS model is adversarially attacked to record the drop in performance. An adversarial version of a clean sample can be generated by exploiting gradient information from the model $\mathcal{F}(.)$ employing [13]. Specific to the case of FSS, the prediction of query masks not only depends on the query image but also the information from the support set. This enables the attacks to be designed in such a way that either attacked query or support can deteriorate the query prediction. These perturbations are specifically chosen so that the loss between ground truth and the predicted masks of the query increases.

The proposed R-PNODE is based on existing prototypical few-shot segmentation models [7, 4]. Given an episode $i$ with task $\mathcal{T}_{i}$ = ($\mathcal{S}_{i},\mathcal{Q}_{i}$), the feature extractor $f_{\theta}$ generates intermediate feature representations ${Z}^{k}_{\mathcal{S}}$ and ${Z}_{\mathcal{Q}}$ for the support and query images $\mathcal{I}^{k}_{\mathcal{S}}$, $\mathcal{I}_{\mathcal{Q}}$. The outputs from the feature extractor $f_{\theta}$ are considered as initial states for the Neural-ODE block at time $t=0$, denoted as ${Z}^{k}_{\mathcal{S}}(0)$ and ${Z}_{\mathcal{Q}}(0)$ respectively. From each support and query image, an additional image is generated by applying multiplicative H$\times$W dimensional i.i.d. Gaussian noise, denoted mathematically by:

where $\odot$ is the Hadamard operator, for element-wise multiplication. The standard deviation $\sigma$ is a hyperparameter. Similar to the clean images, their corresponding Gaussian samples are also passed through the feature extractor to obtain intermediate features ${Z}^{kG}_{\mathcal{S}}(0)$ and ${Z}^{G}_{\mathcal{Q}}(0)$. The Neural-ODE block consists of hidden layers $h_{\phi}$ parameterized by $\phi$ and its dynamics are governed by $h_{\phi}$ which control how the intermediate state changes at any given time $t$. The output representation at fixed terminal time $T(T>0)$ for query features $Z_{\mathcal{Q}}$ is given by:

Similarly, the output representation at fixed terminal time $T(T>0)$ for support features $Z^{k}_{\mathcal{S}}$ are generated. The support feature maps $Z^{k}_{\mathcal{S}}(T)$ from the Neural-ODE block of spatial dimensions ($H^{\prime}\times W^{\prime}$) are upsampled to the same spatial dimensions of their corresponding masks ${L}_{\mathcal{S}}$ of dimension $(H\times W$). Inspired by late fusion [4] where the ground-truth labels are masked over feature maps, we employ Masked Average Pooling (MAP) between $Z^{k}_{\mathcal{S}}(T)$ and ${L}^{k}_{\mathcal{S}}(\eta)$ to form a $d$-dimensional prototype $p(\eta)$ for each foreground class $\eta\in C_{i}$ as shown:

where $(x,y)$ are the spatial locations in the feature map and $\mathbbm{1}(.)$ is an indicator function. The background is also treated as a separate class and the prototype for it is calculated by computing the feature mean of all the spatial locations excluding the ones that belong to the foreground classes. The probability map over semantic classes $\eta$ is computed by measuring the cosine similarity (cossim) between each of the spatial locations in $Z_{\mathcal{Q}}(T)$ with each prototype $p(\eta)$ as given by:

The denominator involves a summation over all foreground classes $\eta^{\prime}\in C_{i}$. The predicted mask $M_{\mathcal{Q}}$ is generated by taking the argmax of $M_{\mathcal{Q}}(\eta)$ across semantic classes. We use Binary Cross Entropy loss $\mathcal{L}_{CE}$ between $M_{\mathcal{Q}}$ and the ground-truth mask for training.

R-PNODE is further optimized using two additional losses. Eq. 5 and Eq. 7 hold for the corresponding features of the gaussian perturbed images as well to get $Z_{S}^{kG}(T)$ and ${M}^{G}$ for the support and query respectively. The cluster loss ($\mathcal{L}_{CL}$) maximises the cosine similarity between each feature $Z$ extracted from clean support samples and their corresponding Gaussian samples as shown below:

The cluster loss helps to bring the support features and their corresponding gaussian perturbations closer in the representation space which generates robust prototypes that are more representative of the class. This helps to learn better relationships between support and query examples which results in better mask predictions. This technique of augmenting the support images with different gaussian perturbations is not explored by the previous methods. It improves both generalizability and robustness as observed from the results in Table VIII. Further, to make the query features robust, we employ consistency loss ($\mathcal{L}_{CON}$) to minimize the difference between Gaussian perturbed query images (${M}^{G}_{\mathcal{Q}}$) and their corresponding ground-truth labels ($L_{\mathcal{Q}}$). For the sake of uniformity, we extend the same loss $\mathcal{L}_{CE}$ as the objective function.

$\mathcal{L}_{CL}$ helps to predict more realistic masks as predicted masks are dependent on the quality of the prototypes and aids $\mathcal{L}_{CON}$ in matching the predicted masks from gaussian perturbed query images with the ground truth masks which further improves robustness and generalizability.

For the overall procedure, refer to Algorithm 1 and Fig. 3. R-PNODE’s superior performance for few-shot segmentation is attributed primarily to two factors. First is the fact that the integral curves learnt by a Neural-ODE are always non-intersecting. We expect a good model to be such that the query and support features of the same class are close to each other. The non-intersecting properties of Neual-ODEs essentially lead to explicit constraints on the features learnt. Algorithm 1 Regularized Prototypical Neural-ODEs for FSS 0:  Clean training data $\mathcal{D_{\mathrm{train}}}=\{(\mathcal{S}_{i},\mathcal{Q}_{i})\}^{E_{train}}_{i=1}$, segmentation network $\mathcal{F}(.)$. 1:  for $\mathrm{i}\in\{1,\ldots,E_{train}\}$ do 2:     Sample episode $\mathcal{E}^{\mathrm{orig}}=\{(\mathcal{S}_{\mathrm{i}},\mathcal{Q}_{\mathrm{i}})\}$ from $\mathcal{D_{\mathrm{train}}}$. 3:     Get gaussian perturbed samples $\mathcal{S}_{\mathrm{i}}^{G}$ and $\mathcal{Q}_{\mathrm{i}}^{G}$. 4:     Get prototypes $p({\eta})$ and features $Z_{S}^{k}(T)$ from $\mathcal{S}_{\mathrm{i}}$. Use prototypes to get predictions $M_{\mathcal{Q}_{\mathrm{i}}}$ and $M_{\mathcal{Q}_{\mathrm{i}}}^{G}$. 5:     Apply $\mathcal{L}_{CE}$ on mask predictions $M_{\mathcal{Q}_{\mathrm{i}}}$ and $\mathcal{L}_{CON}$ on $M_{\mathcal{Q}_{\mathrm{i}}}^{G}$ using Eq. 9. 6:     Extract features $Z_{S}^{kG}(T)$ for each perturbed support sample and compare with clean $Z_{S}^{k}(T)$ to get $\mathcal{L}_{CL}$ by Eq. 8. 7:     Backpropagate losses $\mathcal{L}_{CE}+\alpha.\mathcal{L}_{CON}+\beta.\mathcal{L}_{CL}$ and update parameters of $\mathcal{F}(.)$. 8:  end for

The query features are constrained by the surrounding integral curves of those support prototype features that belong to the same class. Thus, the query features tend to be closer to the prototype features in the representation space. This leads to the accurate classification of query features which ultimately leads to the accurate segmentation of the query images. Fig. 1 tries to demonstrate this very process.

The second factor contributing to the performance of R-PNODE is the addition of cluster and consistency losses. As shown in Fig. 1, the gaussian samples serve as additional reference points in the feature space, and by bringing the final features of these closer to the original, the final features of all samples within this larger space are essentially constrained to be very close to each other.

Extending the reasoning for the superior FSS performance of R-PNODE, the adversarial robustness is again attributed to the intrinsic property of non-intersecting curves of Neural-ODEs along with the proposed losses. As the integral curves of the support features imposed constraints on the features of the query, in a similar way, the clean features also impose constraints on the features of adversarially perturbed images. Thus, features for perturbed support and query tend to lie closer to the features of the clean ones giving rise to adversarial robustness.

An important difference from the usual adversarial perturbation methods is how the support is perturbed here. Although the ground-truth masks for support images are also available, it makes more sense in the FSS task to take the loss of query prediction using the corresponding support samples. Thus, in Eq. 10, the loss is between the query mask and query label, but the gradients are calculated w.r.t the support image. We take batches of either the query or the support perturbed, but not both together. Generating gaussian perturbations is not computationally expensive when compared to generating adversarial perturbations using gradient-based methods like PGD, SMIA or even FGSM. Thus, R-PNODE has a much smaller computational overhead than SAT. Note that the purpose of the Gaussian samples in R-PNODE is not the same as the adversarial samples in SAT. While SAT uses the perturbations to learn better features on the expected test adversarial distribution, R-PNODE uses gaussian perturbations to populate the training input feature space with more samples thereby having more constrained integral curves. Additionally, note that R-PNODE is attack-agnostic, and thus, its robustness should generalize well to multiple attack strategies. Thus, R-PNODE is able to overcome many of the limitations prevalent in adversarial training methods.

We experiment on three publicly available multi-organ segmentation datasets, BCV [35], CT-ORG [36], and Decathlon [37] to evaluate the generalisability of our method. For the in-domain setting, we split BCV into the train (seen) and test (novel) classes. Of the total BCV dataset available, we use half of the dataset for training (6506 slices), $\frac{1}{6}\mathrm{th}$ for validation (1779 slices) and $\frac{1}{3}\mathrm{rd}$ for testing (3873 slices). It is to be noted that the dataset splits are done according to the subjects. All slices of any subject will belong to either of the train, validation or test subset. So, the model will not be tested on a slice of any subject it has seen during training. For the cross-domain FSS, we train on the BCV dataset (with seen classes) and use novel classes from CT-ORG and Decathlon to test. To have a more uniform size of the test set, we sample 500 random slices per organ from the much larger CT-ORG dataset. We sample 500 and 325 slices from the Liver and Spleen test organs respectively, in the Decathlon dataset. For the 3D volumes in all three datasets, we extract slices with non-empty masks and divide them into the fixed train, test, and validation splits. [48] cropped the slices based on the ROI for each organ using available masks. We follow a more challenging setting, where no such cropping is applied (Fig. 5), thereby depriving any unfair advantage due to leakage of localisation information. Of the available organs, we report results on the Liver and Spleen (as novel classes) due to their medical significance and availability across datasets.

For baseline comparisons on few-shot organ segmentation, we experiment with traditional FSS methods, PANet [4], FSS-1000 [45] and SENet [44]. We compare with the recently proposed BiGRU [48], achieving state-of-the-art results on the selected datasets (in the ROI-cropped setting). Additionally, we compare with a recently proposed Neural-ODE based method SONet [39]^***SONet was requiring exorbitant amounts of time for each experiment, so we perform only one run for it.. To provide a fair comparison, we use skew-symmetric and input-output stable Neural-ODE blocks of SONet and extend the model to have VGG-like architecture. Finally, to demonstrate adversarial robustness, we compare with AT-PANet, which is PANet [4] trained with the modified adversarial training procedure described in section V-B.

We experiment on 6 different attacks, some traditional and some state-of-the-art. We first use FGSM[13], one of the most commonly used traditional single-step attack methods. While BIM [14] and PGD[15] are some basic iterative methods that extend single-step attacks, CW [31] is a more sophisticated one. Auto-Attack[32] is an adaptive method involving an ensemble of four diverse attacks to reliably evaluate robustness. While these attacks were originally proposed for classification, we extend them to segmentation in our experiments. DAG [33] was one of the first methods for adversarial attacks specific to segmentation. SMIA[34] was another segmentation attack proposed quite recently, that focuses on the medical domain, and is shown to perform better than many other methods. With this broad spectrum of attacks covering the most important categories, we are able to effectively demonstrate the efficacy of our proposed approach.