[Paper Review] BayeSeg: Bayesian Modeling for Medical Image Segmentation with Interpretable Generalizability

Shangqi Gao, Hangqi Zhou+2 more|arXiv (Cornell University)|Mar 3, 2023

Domain Adaptation and Few-Shot Learning4 citations

TL;DR

BayeSeg proposes a deep Bayesian segmentation framework that decomposes medical images into domain-invariant shape and domain-specific appearance components using hierarchical Bayesian priors, enabling interpretable, generalizable segmentation. By modeling label statistics via variational inference, it achieves state-of-the-art generalization on prostate and cardiac MRI segmentation under cross-sequence, cross-modality, and cross-site distribution shifts.

ABSTRACT

Due to the cross-domain distribution shift aroused from diverse medical imaging systems, many deep learning segmentation methods fail to perform well on unseen data, which limits their real-world applicability. Recent works have shown the benefits of extracting domain-invariant representations on domain generalization. However, the interpretability of domain-invariant features remains a great challenge. To address this problem, we propose an interpretable Bayesian framework (BayeSeg) through Bayesian modeling of image and label statistics to enhance model generalizability for medical image segmentation. Specifically, we first decompose an image into a spatial-correlated variable and a spatial-variant variable, assigning hierarchical Bayesian priors to explicitly force them to model the domain-stable shape and domain-specific appearance information respectively. Then, we model the segmentation as a locally smooth variable only related to the shape. Finally, we develop a variational Bayesian framework to infer the posterior distributions of these explainable variables. The framework is implemented with neural networks, and thus is referred to as deep Bayesian segmentation. Quantitative and qualitative experimental results on prostate segmentation and cardiac segmentation tasks have shown the effectiveness of our proposed method. Moreover, we investigated the interpretability of BayeSeg by explaining the posteriors and analyzed certain factors that affect the generalization ability through further ablation studies. Our code will be released via https://zmiclab.github.io/projects.html, once the manuscript is accepted for publication.

Motivation & Objective

• Address the challenge of distribution shift in medical image segmentation across different imaging sequences, modalities, and sites.
• Improve model generalization on unseen domains without relying on domain-specific data augmentation or large labeled datasets.
• Enhance interpretability by explicitly modeling anatomical shape as a domain-invariant factor through statistical priors.
• Develop a joint image and label statistical modeling framework that supports uncertainty estimation and causal-aware representation learning.
• Enable clinical trust and regulatory compliance by providing explainable segmentation predictions through posterior interpretation.

Proposed method

• Decompose each medical image into two latent variables: a spatially correlated shape variable (domain-invariant) and a spatially variant appearance variable (domain-specific).
• Assign hierarchical Bayesian priors to the shape and appearance components, enforcing the shape to capture stable anatomical structures across domains.
• Model segmentation masks as a locally smooth function of the shape variable, ensuring spatial consistency and interpretability.
• Formulate a joint probabilistic model over image, appearance, shape, and segmentation label using a generative process based on Gaussian Markov Random Fields.
• Implement posterior inference via a variational Bayesian framework with neural networks, enabling end-to-end training and uncertainty quantification.
• Use t-SNE visualization and ablation studies to validate the domain-invariance of the learned shape representations and their impact on generalization.

Experimental results

Research questions

• RQ1 Can explicit statistical modeling of shape as a domain-invariant factor improve generalization in cross-domain medical image segmentation?
• RQ2 How does hierarchical Bayesian prior modeling of image and label statistics enhance interpretability and robustness to distribution shift?
• RQ3 To what extent does separating appearance and shape components improve generalization compared to standard deep learning baselines?
• RQ4 Can the posterior distributions of shape and appearance components be meaningfully interpreted to explain model predictions?
• RQ5 What factors—such as anatomy-relevant vs. anatomy-independent shifts—influence the generalization performance of shape-based segmentation?

Key findings

• BayeSeg achieved 80.9 ± 6.8% Dice score on the LGE sequence (source) and 79.4 ± 10.2% on the T2 sequence (target), outperforming all ablation baselines in cross-sequence generalization.
• The method demonstrated superior generalization on the T2 sequence (79.4 ± 10.2% Dice) compared to ablation #6 (79.1 ± 10.9%), confirming the benefit of joint image-label modeling and explicit shape priors.
• t-SNE visualizations showed that shape representations from BayeSeg had significantly reduced domain shift compared to raw image features, especially between LGE and T2 sequences.
• Ablation studies confirmed that joint modeling of image and label statistics improved generalization, with the full model outperforming variants that omitted shape priors or label smoothing.
• Interpretation of posterior distributions revealed that shape components were stable across domains, while appearance components captured modality-specific variations.
• Despite strong shape invariance, the worst performance on T2 was attributed to high anatomy-relevant shift, indicating that current shape modeling does not fully eliminate anatomical variability.