Deep EM with Hierarchical Latent Label Modelling for Multi-Site Prostate Lesion Segmentation

TL;DR

This paper introduces a hierarchical EM framework that models site-specific annotation variability to improve multi-site prostate lesion segmentation, achieving better cross-site generalisation and interpretable label-quality estimates.

Contribution

It proposes a novel hierarchical EM approach that explicitly accounts for site-dependent annotation differences, enhancing segmentation accuracy across diverse datasets.

Findings

Improved cross-site generalisation over state-of-the-art methods.

Statistically significant DSC improvements in pooled and leave-one-site-out evaluations.

Provides interpretable site-specific label-quality estimates.

Abstract

Label variability is a major challenge for prostate lesion segmentation. In multi-site datasets, annotations often reflect centre-specific contouring protocols, causing segmentation networks to overfit to local styles and generalise poorly to unseen sites in inference. We treat each observed annotation as a noisy observation of an underlying latent 'clean' lesion mask, and propose a hierarchical expectation-maximisation (HierEM) framework that alternates between: (1) inferring a voxel-wise posterior distribution over the latent mask, and (2) training a CNN using this posterior as a soft target and estimate site-specific sensitivity and specificity under a hierarchical prior. This hierarchical prior decomposes label-quality into a global mean with site- and case-level deviations, reducing site-specific bias by penalising the likelihood term contributed only by site deviations. Experiments on three cohorts demonstrate that the proposed hierarchical EM framework enhances cross-site generalisation compared to state-of-the-art methods. For pooled-dataset evaluation, the per-site mean DSC ranges from 29.50% to 39.69%; for leave-one-site-out generalisation, it ranges from 27.91% to 32.67%, yielding statistically significant improvements over comparison methods (p<0.039). The method also produces interpretable per-site latent label-quality estimates (sensitivity alpha ranges from 31.5% to 47.3% at specificity beta approximates 0.99), supporting post-hoc analyses of cross-site annotation variability. These results indicate that explicitly modelling site-dependent annotation can improve cross-site generalisation.