Source-linked AI summary
BayeSeg: Bayesian Modeling for Medical Image Segmentation with Interpretable Generalizability
Shangqi Gao, Hangqi Zhou, Yibo Gao, Xiahai Zhuang
TL;DR
Distribution shifts across medical imaging systems challenge segmentation on unseen data, and domain-invariant representations remain difficult to interpret. BayeSeg uses Bayesian modeling of image and label statistics to separate shape and appearance and infer their posteriors. Experiments on prostate and cardiac segmentation show effectiveness, while analysis identifies a low-generalization case on T2.
Problem
Distribution shifts across imaging systems challenge segmentation on unseen data, and the interpretability of domain-invariant features remains difficult.
Source: §1 · pp. 1–2Method
BayeSeg decomposes images into shape and appearance, assigns hierarchical Bayesian priors, models segmentation from shape, and uses variational inference implemented with neural networks.
Source: §1 · p. 2, §3 · p. 4, §3.2 · p. 5Results
Experiments on prostate and cardiac segmentation show BayeSeg’s effectiveness, and the authors report its lowest generalization performance on T2 in the discussed cardiac setting.
Source: §1 · p. 2, §6 · p. 13, §7 · p. 14Takeaways & Limitations
The authors conclude that joint modeling of image and label statistics can promote interpretability and generalization in medical image segmentation.
Source: §7 · p. 14The discussed cardiac evaluation identifies T2 as BayeSeg’s lowest-generalization case, associated with its largest anatomy-relevant shift.
Source: §6 · p. 13
Abstract
from arXiv · showhide
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.
1. Introduction
Medical imaging systems create distribution shifts across sequences, modalities, and sites, while labeling medical images is costly. BayeSeg addresses the need for generalizable segmentation by explicitly modeling image and label statistics, with experiments and interpretability analyses across cardiac and prostate tasks.
1. Introduction
- Cross-sequence, cross-modality, and cross-site imaging differences can shift image distributions, and deep networks often generalize poorly to unseen datasets.Obtaining the labeled medical images needed to address overfitting is also expensive because it requires qualified expert labor.Source: §1 · pp. 1–2
- BayeSeg jointly models image and label statistics, separating spatial-correlated shape from spatial-variant appearance and assigning hierarchical Bayesian priors to these variables.Variational inference implemented with neural networks estimates their posterior distributions.Source: §1 · p. 2
- The paper extends its MICCAI 2022 framework by statistically describing the relationship between segmentation and its ground-truth.The earlier framework jointly modeled image and label statistics and used neural-network-implemented variational Bayesian inference.Source: §1 · p. 2
- BayeSeg’s experiments cover cardiac and prostate segmentation across cross-modality, cross-sequence, and cross-site settings.The paper also clarifies clinical distribution-shift background and reviews domain generalization.Source: §1 · p. 2
- The paper probes interpretability by visualizing and explaining BayeSeg posteriors, and studies generalizability through prior selection, probabilistic-model pruning, and data-distribution analysis.The paper also describes its limitations and future work in a later section.Source: §1 · p. 2
2. Related works
Prior image-statistics work spans Markov models, spatial correlation, and transform-domain sparsity, while domain-generalization work spans data, learning, and representation strategies. BayeSeg combines explicit statistical modeling with image decomposition to pursue interpretable, generalizable segmentation.
2.1. Image statistical modeling
- Image-statistics research includes Gaussian priors, MRF/Gibbs models, SAR/CAR spatial-correlation models, and transform-domain sparsity models.Pixel dependencies make some field models computationally challenging; later methods pursue tractability or flexible learned priors.Source: §2.1 · pp. 2–3
- Medical image intensities combine spatially independent and correlated information, motivating spatially variant Gaussian priors for independence and SAR priors for spatial correlation.Earlier fixed-neighbor operators can limit modeling flexibility, while some prior work uses learnable high-order filters.Source: §2.1 · p. 3
2.2. Domain generalization for medical image segmentation
- Domain generalization targets unseen domains without target-domain training data, using data-based, learning-based, or representation-based approaches.Data-based approaches enrich training diversity; learning-based approaches combine source-domain knowledge; representation-based approaches seek domain-invariant features.Source: §2.2 · pp. 3–4
- BayeSeg explicitly models shape as domain-invariant organ structure, using hierarchical Bayesian priors rather than relying on multiple source domains for feature alignment.Unlike DSU’s uncertain-feature-statistics augmentation and DIVA’s low-dimensional classification representation, it aims for interpretable statistical representations for medical segmentation.Source: §2.2 · p. 4
3. Methodology
BayeSeg jointly models image and label statistics, then uses variational inference implemented with neural networks to estimate posteriors for shape, appearance, and segmentation.
3. Methodology
- BayeSeg uses deep learning to infer posteriors because conventional variational Bayesian inference is computationally expensive for high-dimensional variables.The methodology combines statistical modeling of images and labels with deep inference.Source: §3 · p. 4
3.1. Statistical modeling of image and label
- The image is decomposed into spatial-correlated shape and spatial-variant appearance, while segmentation is modeled as locally smooth and dependent on shape.Hierarchical priors characterize image and label statistics; the model assigns Gaussian, Gamma, and Beta priors to specified variables.Source: §3 · p. 4, §3.1 · p. 5
3.2. Variational inference of image and label
- Variational Bayesian inference approximates the intractable posterior over image and label variables with a variational distribution whose variables are generally assumed independent.The objective includes a cross-entropy term between predicted segmentation and manual ground-truth labels.Source: §3.2 · pp. 5–6
3.3. Unfolding of the variational loss
- The variational objective is unfolded into semantically interpretable terms, and minimizing the unfolded loss is equivalent to minimizing the original objective without manually balancing its terms.The unfolded expression supports implementation and interpretation of the variational loss.Source: §3.3 · p. 6
3.4. Neural networks and training strategy
- Two ResNets infer shape and appearance posteriors, while a segmentation network infers the segmentation posterior; training balances variational and cross-entropy losses, and testing uses the posterior mean of segmentation.The variational-loss balancing weight is 100 in the experiments.Source: §3.4 · p. 7
4. Experiments
Experiments evaluate BayeSeg on prostate cross-site and cardiac cross-sequence, cross-site, and cross-modality segmentation, alongside component ablations. BayeSeg outperforms comparison methods in the reported settings, while the ablation identifies variational loss as key to generalizability.
4.1. Datasets
- The evaluation spans three prostate datasets and six domains, plus six cardiac domains from five datasets, covering cross-site, cross-sequence, and cross-modality settings.Prostate images were resampled and converted to 2D slices; cardiac slices were aligned and resampled, with CT contrast enhanced before normalization.Source: §4.1 · p. 7
4.2. Prostate segmentation
- BayeSeg averages 77.5 Dice on prostate segmentation, exceeding the second-best method by 7.7; it also retains connected outputs in the reported worst cases.On BIDMC, it exceeds the second-best successful method by over 20 and has the least performance drop, 24.4; on HK, its drop is 4.1, also the least. In the worst I2CVB case, DSU produces a disconnected segmentation while BayeSeg remains connected.Source: §4.2 · p. 8, Table 5 · p. 8
4.3. Cardiac segmentation
- On cardiac target domains, BayeSeg has the smallest reported source-target Dice drop, 5.5, and beats the second-best method by 34.6 on MSCMR T2, 10.2 on EMIDEC LGE, and 3.6 on CASDC CT.Qualitative results are comparable or superior in five of six cases, but every method fails on the worst EMIDEC LGE case.Source: §4.3 · pp. 8–9, Table 6 · p. 8, Figure 3 · p. 9
4.4. Ablation study on basic components
- Variational loss is the key factor in BayeSeg’s generalizability; stochastic mapping alone brings no remarkable improvement but supports unseen-domain performance under that loss.With variational-loss guidance, stochastic mapping provides infinite shape samples for a given image.Source: §4.4 · p. 9
5. Study of interpretability and generalizability
BayeSeg’s posterior visualizations give shape, appearance, and segmentation variables interpretable roles, while ablations examine how priors and joint image-label modeling relate to generalization. The full probabilistic graphical model performs better than simplified configurations, and suitable priors matter especially on difficult target domains.
5. Study of interpretability and generalizability
- Ablations of ρ, υ, and ω test how statistical priors affect posterior variables, while simplified PGM configurations test the contributions of image and label statistics.The section frames explicit prior modeling and joint image-label modeling as two directions for examining generalization.Source: §5 · p. 9, §5.1 · p. 11, §5.2 · p. 11
5.1. Interpretation of statistical representations
- BayeSeg’s posterior visualizations assign semantic roles to appearance, shape, and segmentation variables, supporting interpretation across input distributions rather than only individual samples.Shape is spatially correlated and captures structure; appearance is spatially variant and approximates the image, while υ and ω indicate shape and segmentation boundaries.Source: §5.1 · p. 10, Figure 4 · p. 10
5.2. Generalizability of statistical prior modeling
- The proposed ρ prior achieves superior average performance on three of five target domains, while overly large or small ρ can respectively constrain appearance detail or hinder convergence.A scalar ρ also restricts spatial variation in appearance; the proposed prior helps capture inter-structure details.Source: §5.2 · p. 11, Figure 5 · p. 11, Table 8 · p. 11
- Sharper υ priors can neglect important edges and oversmooth shape, whereas flatter priors can add noise; the sharper example performs poorly on the hardest T2 target.The proposed priors support decomposition into shape and appearance for subsequent segmentation.Source: §5.2 · pp. 11–12, Figure 6 · p. 12, Table 9 · p. 12
- Suitable priors improve generalization, with stronger effects on difficult targets; changing ρ and υ has more obvious quantitative effects than changing ω.The discussion attributes this difference to the variables’ respective roles in the statistical model.Source: §5.2 · p. 12
5.3. Generalizability of the joint modeling
- The full PGM outperforms configurations that make appearance or segmentation deterministic, which reduces generalization ability by removing appearance-related details or label-statistics modeling.The reported losses are especially notable on T2 of MSCMR when appearance is deterministic.Source: §5.3 · p. 12
6. Discussion
The extracted shape improves generalizability but is not fully domain-invariant. The authors explain the weak T2 result by distinguishing anatomy-independent from anatomy-relevant shifts.
6. Discussion
- The extracted shape is not completely domain-invariant, and BayeSeg has its worst generalization ability on T2 despite T2’s comparatively small t-SNE shift from source LGE.The paper attributes the mismatch to T2’s largest anatomy-relevant shift; t-SNE also reflects anatomy-independent shift, whereas segmentation relates only to anatomy-relevant shift.Source: §6 · p. 13
- Only source LGE and target T2 orientations were aligned, which greatly eliminated T2’s anatomy-independent shift but not its largest anatomy-relevant shift.The authors identify extracting domain-invariant, particularly anatomy-relevant, features as an open challenge.Source: §6 · p. 13
7. Conclusion
The conclusion presents BayeSeg as a Bayesian segmentation framework that jointly models image and label statistics to support interpretability and generalization. It reports effectiveness on prostate and cardiac segmentation and uses posterior interpretation and ablations to examine model behavior.
7. Conclusion
- BayeSeg decomposes images into appearance and shape, assigns hierarchical priors to model domain-specific appearance and domain-stable shape, and predicts segmentation from shape through label statistics.Neural-network-implemented variational inference estimates the variables’ posteriors.Source: §7 · p. 14
- Experiments on prostate and cardiac segmentation report effectiveness, while posterior interpretation and ablations examine explicit priors and joint image-label modeling.The ablations validate benefits for generalization capability, as reported by the authors.Source: §7 · p. 14