Source-linked AI summary
Shape-aware Meta-learning for Generalizing Prostate MRI Segmentation to Unseen Domains
Quande Liu, Qi Dou, Pheng-Ann Heng
TL;DR
The paper addresses prostate MRI segmentation generalization across domain shifts when unseen target-domain data are unavailable. It combines gradient-based meta-learning with shape compactness and smoothness objectives, and reports consistent improvements across six unseen-domain settings.
Problem
Domain shifts across clinical sites, imaging protocols, and scanners can degrade segmentation, while target-domain data are unavailable for domain generalization training.
Method
SAML simulates domain shifts by splitting source domains into meta-train and meta-test sets, then adds shape compactness and smoothness objectives to meta-optimization.
Results
SAML outperforms state-of-the-art methods across all six unseen sites, exceeding DeepAll by 2.15% Dice and 0.60mm ASD.
Takeaways & Limitations
The study supports combining shape constraints with meta-learning for domain generalization in medical image segmentation.
Abstract
from arXiv · showhide
Model generalization capacity at domain shift (e.g., various imaging protocols and scanners) is crucial for deep learning methods in real-world clinical deployment. This paper tackles the challenging problem of domain generalization, i.e., learning a model from multi-domain source data such that it can directly generalize to an unseen target domain. We present a novel shape-aware meta-learning scheme to improve the model generalization in prostate MRI segmentation. Our learning scheme roots in the gradient-based meta-learning, by explicitly simulating domain shift with virtual meta-train and meta-test during training. Importantly, considering the deficiencies encountered when applying a segmentation model to unseen domains (i.e., incomplete shape and ambiguous boundary of the prediction masks), we further introduce two complementary loss objectives to enhance the meta-optimization, by particularly encouraging the shape compactness and shape smoothness of the segmentations under simulated domain shift. We evaluate our method on prostate MRI data from six different institutions with distribution shifts acquired from public datasets. Experimental results show that our approach outperforms many state-of-the-art generalization methods consistently across all six settings of unseen domains.
1 Introduction
Domain shifts across clinical sites can degrade prostate MRI segmentation, while domain generalization seeks direct transfer to unseen targets without target-domain data. The paper proposes shape-aware meta-learning that simulates such shifts and regularizes segmentation shape and boundaries.
- Performance can degrade across clinical sites using different imaging protocols or scanner vendors.
- Domain generalization learns from multiple source domains to generalize directly to unseen target domains without target-domain knowledge during training.
- Existing domain-generalization methods largely target natural-image classification, leaving shape-based regularization for dense medical-image segmentation underexplored.
- SAML uses episodic meta-learning, randomly splitting source domains into meta-train and meta-test sets to simulate domain shift during training.
- Two complementary objectives address incomplete shapes and ambiguous boundaries by promoting shape compactness and domain-invariant contour-background representations.
- The method is evaluated on prostate MRI data from six institutions with varied scanners and imaging protocols.
2 Method
The method simulates domain shifts through gradient-based meta-learning and augments meta-optimization with shape compactness and smoothness objectives for prostate MRI segmentation. These objectives target incomplete masks and ambiguous boundaries by preserving complete shapes and encouraging domain-invariant contour representations.
- Gradient-based Meta-learning Scheme: Source domains are split into meta-train and meta-test sets so parameter updates learned on virtual source domains are evaluated for generalization across simulated shifts.The meta-objective is computed using updated parameters but optimized toward the original parameters.
- Meta Shape-aware Objectives: The meta-objective combines Dice segmentation loss with compactness and smoothness terms: Lmeta = Lseg + λ1Lcompact + λ2Lsmooth.The weighting terms λ1 and λ2 control the trade-offs between the complementary shape objectives.
- Meta Shape Compactness Constraint: Compactness regularization uses the reciprocal isoperimetric quotient, with perimeter estimated from probability gradients and area from summed prediction probabilities.The prediction map p supplies pixel-wise quantities, while ε=1e-6 stabilizes computation.
- Meta Shape Compactness Constraint: The compactness loss is applied only on meta-test data to encourage complete shapes on unseen-target images rather than overfitting source data.Incomplete irregular segmentations tend to have smaller area and larger perimeter, producing higher compactness loss.
- Meta Shape Smoothness Enhancement: The smoothness objective promotes intra-class cohesion and inter-class separation for contour- and background-relevant embeddings across source domains.Contour and nearby-background embeddings are extracted from activation maps using morphological masks and weighted averaging, then regularized through contrastive learning.
- Meta Shape Smoothness Enhancement: Training objectives are optimized with respect to the original parameters θ, while Lsmooth is additionally optimized with respect to the embedding-network parameters Hφ.The smoothness loss is computed over embedding pairs in a mini-batch, with either contour- or background-relevant embeddings randomly selected per sample.
3 Experiments
Experiments evaluate prostate MRI segmentation across six shifted sites, comparing generalization methods, ablations, and varying numbers of source domains. The proposed approach consistently performs strongly, preserving segmentation shape while benefiting from broader source-domain coverage.
- Dataset and Evaluation Metric: Experiments use prostate T2-weighted MRI from six distribution-shifted sites, evaluating Dice and Average Surface Distance.The sites derive from NCI-ISBI13, I2CVB, and PROMISE12 datasets.
- Comparison Methods: The comparison includes DeepAll, BigAug, Epi-FCR, LatReg, MASF, and an Intra-site setting using an adapted Mix-residual-UNet backbone.DeepAll aggregates all source-domain data, while Intra-site trains and tests within the same domain.
- Generalization Results: Our approach outperforms the compared state-of-the-art methods across all six unseen sites, exceeding DeepAll by 2.15% on Dice and 0.60mm on ASD.The authors also report better preservation of complete shapes and smooth boundaries in qualitative comparisons.
- Intra-site Analysis: Intra-site training is relatively weaker for sites D, E, and F because these sites contain fewer samples, whereas aggregating shifted data from other sites can help.The proposed method outperforms Intra-site models at four of six sites and has superior overall Dice and ASD performance.
- Ablation Analysis: Adding shape compactness and then shape smoothness to meta-learning produces progressive improvements, with the complete SAML configuration improving consistently across all six sites.The plain meta-learning method already outperforms DeepAll, while the added smoothness objective encourages domain-invariant boundary embeddings.
- Influence of Training Domain Numbers: Generalization improves as the number of training source domains increases, while the proposed approach consistently outperforms DeepAll at every tested source-domain count.Training from a single source domain produces unsatisfactory direct transfer to the target domain.
4 Conclusion
The conclusion presents shape-aware meta-learning for prostate MRI segmentation domain generalization. It combines meta-learning with compactness and smoothness objectives, and reports effectiveness with potential extension to other shifted segmentation settings.
- The proposed shape-aware meta-learning scheme improves model generalization in prostate MRI segmentation.
- Two complementary objectives impose shape compactness and smoothness during meta-optimization to enhance segmentation outputs on unseen domains.
- Extensive experiments demonstrate the effectiveness of the proposed method, which the authors describe as the first medical-image-segmentation domain-generalization work combining shape constraints with meta-learning.
- The method can be extended to segmentation scenarios that suffer from domain shift.