{'fraction_non_alphanumeric': 0.03297635457533665, 'fraction_numerical': 0.011769695684444916, 'mean_word_length': 4.856255821173549, 'pattern_counts': {'":': 0, '<': 0, '<?xml version=': 0, '>': 0, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
vMFNet: Compositionality Meets Domain-generalised Segmentation

Training medical image segmentation models usually requires a large amount of labeled data. By contrast, humans can quickly learn to accurately recognise anatomy of interest from medical (e.g. MRI and CT) images with some limited guidance. Such recognition ability can easily generalise to new images from different clinical centres. This rapid and generalisable learning ability is mostly due to the compositional structure of image patterns in the human brain, which is less incorporated in medical image segmentation. In this paper, we model the compositional components (i.e. patterns) of human anatomy as learnable von-Mises-Fisher (vMF) kernels, which are robust to images collected from different domains (e.g. clinical centres). The image features can be decomposed to (or composed by) the components with the composing operations, i.e. the vMF likelihoods. The vMF likelihoods tell how likely each anatomical part is at each position of the image. Hence, the segmentation mask can be predicted based on the vMF likelihoods. Moreover, with a reconstruction module, unlabeled data can also be used to learn the vMF kernels and likelihoods by recombining them to reconstruct the input image. Extensive experiments show that the proposed vMFNet achieves improved generalisation performance on two benchmarks, especially when annotations are limited. Code is publicly available at: https://github.com/vios-s/vMFNet.


Introduction
Deep learning approaches can achieve impressive performance on medical image segmentation when provided with a large amount of labeled training data [3,8]. However, shifts in data statistics, i.e. domain shifts [41,42], can heavily degrade deep model performance on unseen data [4,24,33]. By contrast, humans can learn quickly with limited supervision and are less likely to be affected by such domain shifts, achieving accurate recognition on new images from different The feature extractor The reconstruction module The segmentation module clinical centres. Recent studies [25,36] argue that this rapid and generalisable learning ability is mostly due to the compositional structure of concept components or image patterns in the human brain. For example, a clinical expert usually remembers the patterns (components) of human anatomy after seeing many medical images. When searching for the anatomy of interest in new images, patterns or combinations of patterns are used to find where and what the anatomy is at each position of the image. This process can be modeled as learning compositional components (learning), composing the components (composing), and performing downstream tasks (performing). This compositionality has been shown to improve the robustness and explainability of many computer vision tasks [16,21,36] but is less studied in medical image segmentation.
In this paper, we explore compositionality in domain-generalised medical image segmentation and propose vMFNet. The overall model design is shown in Fig. 1. Motivated by Compositional Networks [21] and considering medical images are first encoded by deep models as deep features, we model the compositional components of human anatomy as learnable von-Mises-Fisher (vMF) kernels (learning). 1 The features can be projected as vMF likelihoods that define how much each kernel is activated at each position. We then compose the kernels with the normalised vMF likelihoods to reconstruct the input image (composing). Hence, the kernels and the corresponding vMF likelihoods can also be learnt end-to-end with unlabeled data. In terms of a downstream segmentation task, the spatial vMF likelihoods are used as the input to a segmentation module to predict the segmentation mask (performing). In this case, the prediction is based on the activation of the components at each position of the image.
For data from different domains, the features of the same anatomical part of images from these domains will activate the same kernels. In other words, the compositional components (kernels) are learnt to be robust to the domain shifts. To illustrate such robustness, we consider the setting of domain generalisation [22]. With data from multiple source domains available, domain generalisation considers a strict but more clinically applicable setting in which no information about the target domain is known during training. We also consider test-time training [13,14,19,38] with our model on unseen test data at inference time, i.e. test-time domain generalisation [18]. We observe that by freezing the kernels and the segmentation module and fine-tuning the rest of the model to better reconstruct the test data, improved generalisation performance is obtained. This is due to the fact that medical images from different domains are not significantly different; new components do not need to be learned, rather the composing operations need to be fine-tuned.

Contributions:
-Inspired by the human recognition process, we design a semi-supervised compositional model for domain-generalised medical image segmentation. -We propose a reconstruction module to compose the vMF kernels with the vMF likelihoods to facilitate reconstruction of the input image, which allows the model to be trained also with unlabeled data. -We apply the proposed method to two settings: semi-supervised domain generalisation and test-time domain generalisation. -Extensive experiments on cardiac and gray matter datasets show improved performance over several baselines, especially when annotations are limited.

Related work
Compositionality: Compositionality has been mostly incorporated for robust image classification [21,36] and recently for compositional image synthesis [2,25]. Among these work, Compositional Networks [21] originally designed for robust classification under object occlusion is easier to extend to pixel-wise tasks as it learns spatial and interpretable vMF likelihoods. Previous work integrates the vMF kernels and likelihoods [21] for object localisation [40] and recently for nuclei segmentation (with the bounding box as supervision) in a weakly supervised manner [43]. We rather use the vMF kernels and likelihoods for domaingeneralised medical image segmentation and learn vMF kernels and likelihoods with unlabeled data in a semi-supervised manner. Domain generalisation: Several active research directions handle this problem by either augmenting the source domain data [7,42], regularising the feature space [5,15], aligning the source domain features or output distributions [23], carefully designing robust network modules [11], or using meta-learning to approximate the possible domain shifts across source and target domains [9,26,27,29]. Most of these methods consider a fully supervised setting. Recently, Liu et al. [29] proposed a gradient-based meta-learning model to facilitate semi-supervised domain generalisation by integrating disentanglement. Extensive results on a common backbone showed the benefits of meta-learning and disentanglement. Following [29], Yao et al. [39] adopted a pre-trained ResNet [12] as a backbone feature extractor and augmented the source data by mixing MRI images in the Fourier domain and employed pseudo-labelling to leverage the unlabelled data. Our proposed vMFNet aligns the image features to the same vMF distributions to handle the domain shifts. The reconstruction further allows the model to handle semi-supervised domain generalisation tasks.

Proposed method
We propose vMFNet, a model consisting of three modules; the feature extractor F ψ , the task network T θ , and the reconstruction network R ω , where ψ, θ and ω denote the network parameters. Overall, the compositional components are learned as vMF kernels by decomposing the features extracted by F ψ . Then, we compose the vMF kernels to reconstruct the image with R ω by using the vMF likelihoods as the composing operations. Finally, the vMF likelihoods that contain spatial information are used to predict the segmentation mask with T θ . The decomposing and composing are shown in Fig. 1 and detailed below.

Learning compositional components
The primal goal of learning compositional components is to represent deep features in a compact low dimensional space. We denote the features extracted by F ψ as Z ∈ R H×W ×D , where H and W are the spatial dimensions and D is the number of channels. The feature vector z i ∈ R D is defined as a vector across channels at position i on the 2D lattice of the feature map. We follow Compositional Networks [21] to model Z with J vMF distributions, where the learnable mean of each distribution is defined as vMF kernel µ j ∈ R D . To allow the modeling to be tractable, the variance σ of all distributions are fixed. In particular, the vMF likelihood for the j th distribution at each position i can be calculated as: where ||z i || = 1 and C(σ) is a constant. After modeling the image features with J vMF distributions with Eq. 1, the vMF likelihoods Z vM F ∈ R H×W ×J can be obtained, which indicates how much each kernel is activated at each position. Here, the vMF loss L vM F that forces the kernels to be the cluster centres of the features vectors is defined in [21] as This avoids an iterative EM-type learning procedure. Overall, the feature vectors of different images corresponding to the same anatomical part will be clustered and activate the same kernels. In other words, the vMF kernels are learnt as the components or patterns of the anatomical parts. Hence, the vMF likelihoods Z vM F for the features of different images will be aligned to follow the same distributions (with the same means). In this case, the vMF likelihoods can be considered as the content representation in the context of content-style disentanglement [6,30], where the vMF kernels contain the style information.

Composing components for reconstruction
After decomposing the image features with the vMF kernels and the likelihoods, we re-compose to reconstruct the input image. Reconstruction requires that com-plete information about the input image is captured [1]. However, the vMF likelihoods contain only spatial information as observed in [21], while the non-spatial information is compressed as varying kernels µ j , j ∈ {1 · · · J}, where the compression is not invertible. Consider that the vMF likelihood p(z i |µ j ) denotes how much the kernel µ j is activated by the feature vector z i . We propose to construct a new feature space Z with the vMF likelihoods and kernels. Let z vM F i ∈ R J be a normalised vector across Z vM F channels at position i. We devise the new feature vector z i as the combination of the kernels with the normalised vMF likelihoods as the combination coefficients, namely: After obtaining Z as the approximation of Z, the reconstruction network R ω reconstructs the input image with Z as the input, i.e.X = R ω ( Z).

Performing downstream task
As the vMF likelihoods contain only spatial information of the image that is highly correlated to the segmentation mask, we design a segmentation module, i.e. the task network T θ , to predict the segmentation mask with the vMF likelihoods as input, i.e.Ŷ = T θ (Z vM F ). Specifically, the segmentation mask tells what anatomical part the feature vector z i corresponds to, which provides further guidance for the model to learn the vMF kernels as the components of the anatomical parts. Then the vMF likelihoods will be further aligned when trained with multi-domain data and hence perform well on domain generalisation tasks.

Learning objective
Overall, the model contains trainable parameters ψ, θ, ω and the kernels µ. The model can be trained end-to-end with the following objective: where λ Dice = 1 when the ground-truth mask Y is available, otherwise λ Dice = 0. L Dice is the Dice loss [31], L rec is the reconstruction loss (we use L1 distance).

Datasets and baseline models
To make comparison easy, we adopt the datasets and baselines of [29], which we briefly summarise below. We also include [29] as a baseline.  [34], without any pre-training. nnUNet [17] is a supervised baseline. It adapts its model design and searches the optimal hyperparameters to achieve the optimal performance. SDNet+Aug. [28] is a semi-supervised disentanglement model, which disentangles the input image to a spatial anatomy and a non-spatial modality factors. Augmenting the training data by mixing the anatomy and modality factors of different source domains, "SDNet+Aug." can potentially generalise to unseen domains. LDDG [23] is a fully-supervised domain generalisation model, in which low-rank regularisation is used and the features are aligned to Gaussian distributions. SAML [27] is a gradient-based meta-learning approach. It applies the compactness and smoothness constraints to learn domain-invariant features across meta-train and metatest sets in a fully supervised setting. DGNet [29] is a semi-supervised gradientbased meta-learning approach. Combining meta-learning and disentanglement, the shifts between domains are captured in the disentangled representations. DGNet achieved the state-of-the-art (SOTA) domain generalisation performance on M&Ms and SCGM datasets.

Implementation details
All models are trained using the Adam optimiser [20] with a learning rate of 1 × e −4 for 50K iterations using batch size 4. Images are cropped to 288 × 288 for M&Ms and 144 × 144 for SCGM. F ψ is a 2D UNet [34] without the last upsampling and output layers to extract features Z. Note that F ψ can be easily replaced by other encoders such as a ResNet [12] and the feature vectors can be extracted from any layer of the encoder where performance may vary for different layers. T θ and R ω are two shallow convolutional networks that are detailed in Section 2 of Appendix. We follow [21] to set the variance of the vMF distributions as 30. The number of kernels is set to 12, as this number performed the best according to early experiments. For different medical datasets, the best number of kernels may be slightly different. All models are implemented in PyTorch [32] and are trained using an NVIDIA 2080 Ti GPU. In the semi-supervised setting, we use specific percentages of the subjects as labeled data and the rest as unlabeled data. We train the models with 3 source domains and treat the 4 th domain as the target one. We use Dice (%) and Hausdorff Distance (HD) [10] as the evaluation metrics.   16.56 HD, training with both losses gives better generalisation results when the model is trained to learn better kernels and with unlabeled data. When removing the two losses, the model can still perform adequately compared to the baselines due to the decomposing mechanism. Alignment analysis: To show that the vMF likelihoods from different source domains are aligned, for M&Ms 100% cases, we first mask out the non-heart part of the images, features and vMF likelihoods. Then, we train classifiers to predict which domain the input is from with the masked images X or masked features Z or masked vMF likelihoods Z vM F as input. The average cross-entropy errors are 0.718, 0.701 and 0.756, which means that it is harder to tell the domain class with the heart part of Z vM F , i.e. the vMF likelihoods for the downstream task are better aligned compared to the features Z from different source domains.

Visualisation of compositionality
Overall, the segmentation prediction can be interpreted as the activation of corresponding kernels at each position, where false predictions occur when the wrong kernels are activated i.e. the wrong vMF likelihoods (composing operations) are used to predict the mask. We show example images, reconstructions, predicted segmentation masks, and 8 most informative vMF likelihoods channels in Fig. 2. As shown, vMF kernels 1 and 2 (red box) are mostly activated by the left ventricle (LV) feature vectors and kernels 3 (blue box) and 4 (green box) are mostly for right ventricle (RV) and myocardium (MYO) feature vectors. Interestingly, kernel 2 is mostly activated by papillary muscles in the left ventricle even though no supervision about the papillary muscles is provided during training. This supports that the model learns the kernels as the compositional components (patterns of papillary muscles, LV, RV and MYO) of the heart.

Test-time domain generalisation
As discussed in Section 4.4, poor segmentation predictions are usually caused by the wrong kernels being activated. This results in the wrong vMF likelihoods being used to predict masks. Reconstruction quality is also affected by wrong vMF likelihoods. In fact, average reconstruction error is approximately 0.007 on the training set and 0.011 on the test set. Inspired by [13,37] we perform testtime training (TTT) to better reconstruct by fine-tuning the reconstruction loss L rec (X,X) to update F ψ and R ω with the kernels and T θ fixed. This should in turn produce better vMF likelihoods. For images of each subject at test time, we fine-tune the reconstruction loss for 15 iterations (saving the model at each iteration) with a small learning rate of 1 × e −6 . Out of the 15 models, we choose the one with minimum reconstruction error to predict the segmentation masks for each subject. The detailed results of TTT for M&Ms are included in Section 1 of Appendix. For M&Ms 2%, 5% and 100% cases, TTT gives around 3.5%, 1.4% and 1% improvements in Dice compared to results (without TTT) in Table 1.

Conclusion
In this paper, inspired by the human recognition process, we have proposed a semi-supervised compositional model, vMFNet, that is effective for domaingeneralised medical image segmentation. With extensive experiments, we showcased that vMFNet achieves consistently improved performance in terms of semi-supervised domain generalisation and highlighted the correlation between downstream segmentation task performance and the reconstruction. Based on the presented results, we are confident that freezing the kernels and fine-tuning the reconstruction to learn better composing operations (vMF likelihoods) at test time gives further improved generalisation performance.

More quantitative results
We report the results of test-time training on M&Ms dataset in Table A1. In Table A2, Table A3, Table A4 and Table A5, we report more detailed Dice (%) and Hausdorff Distance results of semi-supervised domain generalisation.

Implementation details
T θ and R ω have similar structures, where a double CONV layer (kernel size 3, stride size 1 and padding size 1) in UNet with batch normalisation and ReLU is first used to process the features. Then a transposed convolutional layer is used to upsample the features following with a double CONV layer with batch normalisation and ReLU. Finally, a output convolutional layer with 1 × 1 kernels is used. For T θ , the output of the last layer is processed with Sigmoid.