Shape-aware Meta-learning for Generalizing
Prostate MRI Segmentation to Unseen Domains

The foundation of our learning scheme is the gradient-based meta-learning algorithm [17], to promote robust optimization by simulating the real-world domain shifts in the training process. Specifically, at each iteration, the source domains $\mathcal{D}$ are randomly split into the meta-train $\mathcal{D}_{tr}$ and meta-test $\mathcal{D}_{te}$ sets of domains. The meta-learning can be divided into two steps. First, the model parameters $\theta$ are updated on data from meta-train $\mathcal{D}_{tr}$, using Dice segmentation loss $\mathcal{L}_{seg}$:

where $\alpha$ is the learning-rate for this inner-loop update. Second, we apply a meta-learning step, aiming to enforce the learning on meta-train $\mathcal{D}_{tr}$ to further exhibit certain properties that we desire on unseen meta-test $\mathcal{D}_{te}$. Crucially, the meta-objective $\mathcal{L}_{meta}$ to quantify these properties is computed with the updated parameters $\theta^{\prime}$, but optimized towards the original parameters $\theta$. Intuitively, besides learning the segmentation task on meta-train $\mathcal{D}_{tr}$, such a training scheme further learns how to generalize at the simulated domain shift across meta-train $\mathcal{D}_{tr}$ and meta-test $\mathcal{D}_{te}$. In other words, the model is optimized such that the parameter updates learned on virtual source domains $\mathcal{D}_{tr}$ also improve the performance on the virtual target domains $\mathcal{D}_{te}$, regarding certain aspects in $\mathcal{L}_{meta}$.

In segmentation problems, we expect the model to well preserve the complete shape (compactness) and smooth boundary (smoothness) of the segmentations in unseen target domains. To achieve this, apart from the traditional segmentation loss $\mathcal{L}_{seg}$, we further introduce two complementary loss terms into our meta-objective, $\mathcal{L}_{meta}=\mathcal{L}_{seg}+\lambda_{1}\mathcal{L}_{compact}+\lambda_{2}\mathcal{L}_{smooth}$ ($\lambda_{1}$ and $\lambda_{2}$ are the weighting trade-offs), to explicitly impose the shape compactness and shape smoothness of the segmentation maps under domain shift for improving generalization performance.

Traditional segmentation loss functions, e.g., Dice loss and cross entropy loss, typically evaluate the pixel-wise accuracy, without a global constraint to the segmentation shape. Trained in that way, the model often fails to produce complete segmentations under distribution shift. Previous study have demonstrated that for the compact objects, constraining the shape compactness [9] is helpful to promote segmentations for complete shape, as an incomplete segmentation with irregular shape often corresponds to a worse compactness property.

Based on the observation that the prostate region generally presents a compact shape, and such shape prior is independent of observed domains, we propose to explicitly incorporate the compact shape constraint in the meta-objective $\mathcal{L}_{meta}$, for encouraging the segmentations to well preserve the shape completeness under domain shift. Specifically, we adopt the well-established Iso-Perimetric Quotient [19] measurement to quantify the shape compactness, whose definition is $C_{IPQ}={4\pi A}/{P^{2}}$, where $P$ and $A$ are the perimeter and area of the shape, respectively. In our case, we define the shape compactness loss as the reciprocal form of this $C_{IPQ}$ metric, and expend it in a pixel-wise manner as follows:

where $p$ is the prediction probability map, $\Omega$ is the set of all pixels in the map; $\nabla p_{u_{i}}$ and $\nabla p_{v_{i}}$ are the probability gradients for each pixel $i$ in direction of horizontal and vertical; $\epsilon$ ($1e^{-6}$ in our model) is a hyperparameter for computation stability. Overall, the perimeter length $P$ is the sum of gradient magnitude over all pixels $i\in\Omega$; the area $A$ is calculated as the sum of absolute value of map $p$.

Intuitively, minimizing this objective function encourages segmentation maps with complete shape, because an incomplete segmentation with irregular shape often presents a relatively smaller area $A$ and relatively larger length $P$, leading to a higher loss value of $\mathcal{L}_{compact}$. Also note that we only impose $\mathcal{L}_{compact}$ in meta-test $\mathcal{D}_{te}$, as we expect the model to preserve the complete shape on unseen target images, rather than overfit the source data.

In addition to promoting the complete segmentation shape, we further encourage smooth boundary delineation in unseen domains, by regularizing the model to capture domain-invariant contour-relevant and background-relevant embeddings that cluster regardless of domains. This is crucial, given the observation that performance drop at the cross-domain deployment mainly comes from the ambiguous boundary regions. In this regard, we propose a novel objective $\mathcal{L}_{smooth}$ to enhance the boundary delineation, by explicitly promoting the intra-class cohesion and inter-class separation between the contour-relevant and background-relevant embeddings drawn from each sample across all domains $\mathcal{D}$.

Specifically, given an image $x_{m}\in\mathbb{R}^{H\times W\times 3}$ and its one-hot label $y_{m}$, we denote its activation map from layer $l$ as $M^{l}_{m}\in\mathbb{R}^{H_{l}\times W_{l}\times C_{l}}$, and we interpolate $M_{m}^{l}$ into $T_{m}^{l}\in\mathbb{R}^{H\times W\times C_{l}}$ using bilinear interpolation to keep consistency with the dimensions of $y_{m}$. To extract the contour-relevant embeddings $E_{m}^{con}\in\mathbb{R}^{C_{l}}$ and background-relevant embeddings $E_{m}^{bg}\in\mathbb{R}^{C_{l}}$, we first obtain the binary contour mask $c_{m}\in\mathbb{R}^{H\times W\times 1}$ and binary background mask $b_{m}\in\mathbb{R}^{H\times W\times 1}$ from $y_{m}$ using morphological operation. Note that the mask $b_{m}$ only samples background pixels around the boundary, since we expect to enhance the discriminativeness for pixels around boundary region. Then, the embeddings $E_{m}^{con}$ and $E_{m}^{bg}$ can be extracted from $T_{m}^{l}$ by conducting weighted average operation over $c_{m}$ and $b_{m}$:

where $\Omega$ denotes the set of all pixels in $T_{m}^{l}$, the $E_{m}^{con}$ and $E_{m}^{bg}$ are single vectors, representing the contour and backgound-relevant representations extracted from the whole image $x_{m}$. In our implementation, activations from the last two deconvolutional layers are interpolated and concatenated to obtain the embeddings.

Next, we enhance the domain-invariance of $E^{con}$ and $E^{bg}$ in latent space, by encouraging embeddings’ intra-class cohesion and inter-class separation among samples from all source domains $\mathcal{D}$. Considering that imposing such regularization directly onto the network embeddings might be too strict to impede the convergence of $\mathcal{L}_{seg}$ and $\mathcal{L}_{compact}$, we adopt the contrastive learning [6] to achieve this constraint. Specifically, an embedding network $H_{\phi}$ is introduced to project the features $E\in[E^{con},E^{bg}]$ to a lower-dimensional space, then the distance is computed on the obtained feature vectors from network $H_{\phi}$ as $d_{\phi}(E_{m},E_{n})=\|H_{\phi}(E_{m})-H_{\phi}(E_{n})\|_{2}$, where the sample pair $(m,n)$ are randomly drawn from all domains $\mathcal{D}$, as we expect to harmonize the embeddings space of $\mathcal{D}_{te}$ and $\mathcal{D}_{tr}$ to capture domain-invariant representations around the boundary region. Therefore in our model, the contrastive loss is defined as follows:

where the function $\tau(E)$ indicates the class (1 for $E$ being $E^{con}$, and 0 for $E^{bg}$) $\zeta$ is a pre-defined distance margin following the practice of metric learning (set as 10 in our model). The final objective $\mathcal{L}_{smooth}$ is computed within mini-batch of $q$ samples. We randomly employ either $E^{con}$ or $E^{bg}$ for each sample, and the $\mathcal{L}_{smooth}$ is the average of $\ell_{contrastive}$ over all pairs of $(m,n)$ embeddings:

where $C(q,2)$ is the number of combinations. Overall, all training objectives including $\mathcal{L}_{seg}(\mathcal{D}_{tr};\theta)$ and $\mathcal{L}_{meta}(\mathcal{D}_{tr},D_{te};\theta^{\prime})$, are optimized together with respect to the original parameters $\theta$. The $\mathcal{L}_{smooth}$ is also optimized with respect to $H_{\phi}$.