{'fraction_non_alphanumeric': 0.04060305883633593, 'fraction_numerical': 0.011756381646981985, 'mean_word_length': 4.833015873015873, 'pattern_counts': {'":': 0, '<': 0, '<?xml version=': 0, '>': 0, 'https://': 0, '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}}
Semi-supervised Meta-learning with Disentanglement for Domain-generalised Medical Image Segmentation

Generalising deep models to new data from new centres (termed here domains) remains a challenge. This is largely attributed to shifts in data statistics (domain shifts) between source and unseen domains. Recently, gradient-based meta-learning approaches where the training data are split into meta-train and meta-test sets to simulate and handle the domain shifts during training have shown improved generalisation performance. However, the current fully supervised meta-learning approaches are not scalable for medical image segmentation, where large effort is required to create pixel-wise annotations. Meanwhile, in a low data regime, the simulated domain shifts may not approximate the true domain shifts well across source and unseen domains. To address this problem, we propose a novel semi-supervised meta-learning framework with disentanglement. We explicitly model the representations related to domain shifts. Disentangling the representations and combining them to reconstruct the input image allows unlabeled data to be used to better approximate the true domain shifts for meta-learning. Hence, the model can achieve better generalisation performance, especially when there is a limited amount of labeled data. Experiments show that the proposed method is robust on different segmentation tasks and achieves state-of-the-art generalisation performance on two public benchmarks.


Introduction
Despite recent progress in medical image segmentation [3,8,18], inference performance on unseen datasets, acquired from distinct scanners or clinical centres, is known to decrease [5,38]. Such reduction is mainly caused by shifts in data statistics between different clinical centres i.e. domain shifts [45], due to variation in patient populations, scanners, and scan acquisition settings [42]. The variation in population impacts the underlying anatomy and pathology due to factors such as gender, age, ethnicity, which may differ for patients in different locations [28,39,43]. The variation in scanners and scan acquisition settings impacts the characteristics of the acquired image, such as brightness and contrast [45].
The naive approach to handling domain shift is to acquire and label as many and diverse data as possible, the cost implications and difficulties of which are known to this community. Alternatively one can train a model on source domains to generalise for a target domain with some information on the target domain available i.e. domain adaptation [4] such as cross-site MRI harmonisation [37] to enforce the source and target domains to share similar image-specific characteristics [9]. A more strict alternative is to not use any information for the target domain, known as domain generalisation [22]. Herein, we focus on this more challenging and more widely applicable approach.
In domain generalisation, the overarching goal is to identify suitable representations that encode information about the task at hand whilst being insensitive to domain-specific information. There are several active research directions aiming to address this goal, including: direct augmentation of the source domain data [45], feature space regularisation [6,14,27,35,46], alignment of the source domain features or output distributions [25], and learning domain-invariant features with gradient-based meta-learning [11,24,29]. Of the above, gradient-based meta-learning methods have the advantage of not overfitting to dominant source domains which account for the more populous data in the training dataset [11]. Gradient-based meta-learning [22,26] exploits an episodic training paradigm [23] by splitting the source domains into meta-train and meta-test sets at each iteration. The model is trained to handle domain shift by simulating it during training. By using constraints to implicitly eliminate the information related to the simulated domain shifts, the model can learn to extract domain-invariant features. Previous work introduced different constraints in a fully supervised setting e.g. global class alignment and local sample clustering [11], shape-aware constraints [29] or simply the task objective [20,22], where [20] extends [22] to medical image segmentation. 1 These approaches do not scale in medical image segmentation as pixel-wise annotation is time-consuming, laborious, and requires expert knowledge. Meanwhile, in a low data regime where centres only provide a few labeled data samples, these methods may only learn to extract domaininvariant features from an under-represented data distribution [20,45]. In other words, the simulated domain shifts may not well approximate the true domain shifts between source and unseen domains.
To address this problem, we propose to explicitly disentangle the representations related to domain shifts for meta-learning as we illustrate in Fig. 1. Learning these complete and sufficient representations [1] via reconstruction brings the benefit of unsupervised learning, thus we can better simulate the domain shifts by also using unlabeled data from any of the source domains. We consider two sources of shifts: one due to scanner and scan acquisition setting varia- tion, and one due to population variation. Because our task is segmentation, we want to be sensitive to changes in anatomy but insensitive to changes in imaging characteristics be it some common across domains or domain-specific. We use spatial (grid-like) features as a representation of anatomy (Z) and two vectors (s,d) to encode common or domain-specific imaging characteristics. We apply specific design and learning biases to disentangle the above. For example, a spatial Z is equivariant to segmentation and this has been shown to improve performance [7,16]. We further encourage Z to be disentangled from s and d by exploiting also low-rank regularisation [25]. Gradient-based meta-learning also encourages Z, s, and d to generalise well to unseen domains whilst at the same time improves (implicitly) their disentanglement. Our contributions are summarised as follows: 1. We propose the first, to the best of our knowledge, semi-supervised domaingeneralisation framework combining meta-learning and disentanglement.
2. Use of low-rank regularisation as a learning bias to encourage better disentanglement and hence improved generalisation performance.

3.
Extensive experiments on cardiac and gray matter datasets show improved performance over several baselines especially for the limited annotated data case.

Proposed method
Consider a multi-domain training dataset K} that is defined on a joint space X × Y, where X k i is the i th training datum from the k th source domain with corresponding ground-truth segmentation mask Y k i , and N k denotes the number of training samples in the k th source domain. We aim to learn a model containing a feature network F ψ : X → Z to extract the anatomical features Z and a task network T θ : Z → Y to predict the segmentation masks, where ψ and θ denote the network parameters.

Gradient-based meta-learning for domain generalisation
In gradient-based meta-learning for domain generalisation, the domain shift is simulated by training the model on a sequence of episodes [23,24]. Specifically, the meta-train set D tr and the meta-test set D te are constructed by randomly splitting the source domains D for each iteration of training. Each iteration comprises a meta-train step followed by a meta-test step. For the meta-train step, the parameters ψ and θ of F ψ and T θ are calculated by optimising the meta-train loss L meta−train with data from D tr (inner-loop update), as defined by: where α is the learning rate for the meta-train update step. Typically, L meta−train is the task objective, e.g. the Dice loss [10] for a segmentation task. This step rewards accurate predictions on the meta-train source domains. For the meta-test step, the meta-test source domains D te are processed by the updated parameters (ψ , θ ) and the model is expected to contain certain properties quantified by the L meta−test loss. L meta−test is computed using the updated parameters (ψ , θ ), whilst the gradients are computed towards the original parameters (ψ, θ). The final objective is defined as: The intuition behind this scheme is that the model should not only perform well on the source domains, but its future updates should also generalise well to unseen domains. Below, we will describe our meta-train and meta-test objectives but first we present how we disentangle representations related to domain shifts.

Learning disentangled representations
To model appearance in a single-domain setting, typically a single vector-based variational representation is used [7]. Here, due to our multi-domain setting, inspired by [17,44], we separately encode domain-specific imaging characteristics as an additional vector-based variational representation. Hence, we aim to learn two independent vector representations, where one (s) captures common imaging characteristics across domains and the other one (d) captures specific imaging characteristics for each domain. In addition, we encode spatial anatomy information in a separate representation Z, which we encourage to be disentangled from s and d.
In particular, the input image X is first encoded in a common (appearance) representation s = E S (X), and a domain representation d = E D (X) that is followed by a shallow domain classifier T C (d) which predicts the source domain (ĉ) label of X. Then, a decoder DE combines the extracted features Z = F ψ (X) and the representations s and d to reconstruct the input image, i.e. X = DE(Z, s, d). Note that DE combines Z and s, d using adaptive instance normalisation (AdaIN) layers [15]. As shown in [16], AdaIN improves disentanglement and encourages Z to encode spatially equivariant information, i.e. anatomical information useful for segmentation, and s, d to only encode common or domain-specific appearance.
To achieve such "triple" disentanglement we consider several losses: 1) KL divergences L KL (s, N (0, 1)), L KL (d, N (0, 1)) to induce a Gaussian N (0, 1) prior in s and d, encouraging the representations to be robust on unseen domains [13]; 2) Hilbert-Schmidt Independence Criterion (HSIC) loss L HSIC (s, d), to force s and d to be independent from each other [34]; 3) a classification loss L cls (c,ĉ) such that the domain representation d is highly correlated with the domainspecific information [17]; and 4) a reconstruction loss L rec (X,X), defined as the 1 distance between X andX, to learn representations without supervision [7,16].
We further encourage the extracted features Z to be equivariant across the meta-train source domains and improve disentanglement between Z and s, d by applying rank regularisation [25]. Specifically, consider a batch } from K tr meta-train source domains, and K tr features } extracted using the feature network F ψ . By flattening and concatenating these features, we end up with a matrix Z with dimensions [C, K tr × H × W ], where C, H, W denote the number of channels, height, and width of Z. Then, by forcing the rank of Z to be m (i.e. the number of the segmentation classes), Z is encouraged to encode only globally-shared information across K tr source domains in order to predict the segmentation mask as discussed in [25]. We achieve that by minimising the (m + 1) th singular value σ m+1 of Z, i.e. L rank = σ m+1 . Overall, L DT is defined as: L DT =λ rank L rank (Z) + λ KL (L KL (s, N (0, 1)) + L KL (d, N (0, 1))) + λ rec L rec (X,X) + λ HSIC L HSIC (s, d) + λ cls L cls (c,ĉ), where c is the domain label. We adopt hyperparameter values according to our extensive early experiments and discussion from [7,25] as λ rank = 0.1, λ KL = 0.1, λ rec = 1 and λ cls = 1. Note that all the losses do not need ground-truth masks. The domain class label is available, as we know the centre where the data belong.

Meta-train and meta-test objectives
Our meta-train objective contains two components: where λ Dice = 5 when labeled data are available.
For the meta-test step, the model is expected to: 1) accurately predict segmentation masks (by applying the task objective), and 2) disentangle Z and s, d to the same level as meta-train sets. A naive strategy for the latter is to use L DT for meta-test sets. However, as analysed in [2,29], the meta-test step is unstable to train: the gradients from the meta-test loss are second-order statistics of ψ and θ. Our experiments revealed that including the unsupervised losses L KL and L HSIC make training even more unstable (even leading to model collapse). In addition, we use one domain for meta-test in experiments, while L rank requires multiple domains. According to [31,33], considering fixed learning and design biases, the level of disentanglement can be proxied by the reconstruction quality (with ground-truth image X) and the domain classification accuracy (with ground-truth label c). Hence, we adopt as the meta-test loss: Note that for unlabeled data, L rec and L cls do not need ground-truth masks.

Baseline models
nnUNet [19]: is a self-adapting framework based on 2D and 3D U-Nets [40] which does not specifically target domain generalisation. Given a labelled training dataset, nnUNet automatically adapts its model design and hyperparameters to obtain optimal performance. In the M&Ms challenge, methods based on nnUNet achieved the top performance [5].
SDNet+Aug. [30]: disentangles the input image to a spatial anatomy and a non-spatial modality factors. Here we use intensity-and resolution-augmented data in a semi-supervised setting. Compared to our method, "SDNet+Aug." only poses disentanglement to the latent features without meta-learning. LDDG [25]: is the latest state-of-the-art model for domain-generalised medical image analysis. It also uses a rank loss and when applied in a fully supervised setting, LDDG achieved the best generalisation performance on SCGM. SAML [29]: is another gradient-based meta-learning approach. SAML proposed to constraint the compactness and smoothness properties of segmentation masks across meta-train and meta-test sets in a fully supervised setting.

Implementation details
Models are trained using the Adam optimiser [21] with a learning rate of 2e −5 for 50K iterations using batch size 4. Images are cropped to 224 × 244 for M&Ms and 144 × 144 for SCGM. F ψ is a 2D UNet [40] to extract Z features with 8 channels of same height and width as input image. We follow the designs of SDNet [7] for E S , T θ and DE. E D has the same architecture as E S . Both s and d have 8 dimensions. T C is a single fully-connected layer. All models are implemented in PyTorch [36] 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 use Dice (%) and Hausdorff Distance [12] as the evaluation metrics.

Results and discussion
Tables 1, 2, 3 and 4 show that we consistently achieve the best generalisation performance on cardiac and gray matter segmentation. Particularly in the low data regime we improve Dice by ≈ 5% on M&Ms and ≈ 3% on SCGM compared to the best performing baseline. For 100% annotations in M&Ms (see Appendix), our model still outperforms the baselines. We also show the qualitative results in Appendix, where the improved performance is visually observed. M&Ms: Compared to "SDNet+Aug." which can also use (due to disentanglement) unlabeled data, our model performs consistently better. The results agree with the conclusion in [32]: without specific designs tuned to the tasks, disentanglement can not provide guaranteed generalisation ability. For LDDG and SAML, the generalisation performance significantly drops with small amounts of labeled data. Note that nnUNet adapts the model design per each run/training set. However, adapting the model design for different training data limits the scalability of nnUNet. In the Appendix, we also show that nnUNet possibly overfits the source domains in some cases. SCGM: We obtain consistent improvements also on SCGM, demonstrating application in other organs. Our model benefits from the additional 10 unlabeled subjects of each domain leading to better performance overall.

Ablation study
Here we conduct ablations on key losses crucial to disentanglement and the extraction of good anatomical features for good generalisation performance. We omit ablations on the KL losses as [13,17] showcase that variational encoding

Conclusion
We have presented a novel semi-supervised meta-learning framework for domain generalisation. Using disentanglement our approach models domain shifts, and thanks to our reconstruction approach to disentanglement, our model can be trained also with unlabeled data. By applying the designed constraints (including the low-rank regularisation) to the gradient-based meta-learning approach, the model extracts robust anatomical features useful for predicting segmentation masks in a semi-supervised manner. Extensive quantitative results, especially when insufficient annotated data are available, indicate remarkable improvements compared to previous state-of-the-art approaches. The performance of our method might improve with the use of additional unlabeled data from other domains, which we leave as future work. In Table A5 and Table 2, we report the Dice (%) and Hausdorff Distance results on the cases of giving 100% labeled data for each model. For M&Ms, the unlabeled data (of labeled subjects) from phases between end-systole and enddiastole phases still gives the proposed model slightly better performance i.e. 0.65% of improvement compared with the best baseline model.

Disentanglement networks
We use AdaIN module in the decoder DE to combine Z and s, d. Each AdaIN layer performs the following operation: where each feature map Z i is first normalized separately, and then is scaled and shifted based on the style representation and domain representation s i , d i mean and standard deviation scalars.  Fig. A2. We show the example images and predicted segmentation masks of each model for different cases. When training the baseline models with less labeled data, the performance drops significantly. In contrast, our model can produce satisfactory masks in every case.