Probabilistic Modeling of Multi-rater Medical Image Segmentation for Diversity and Personalization

Technical Summary: Disentangling Preference and Ambiguity for Controllable Medical Lesion Segmentation

Problem Statement
Medical lesion segmentation is inherently challenged by imaging uncertainty, stemming from ill-defined boundaries, artifacts, and irregular shapes of medical targets. A critical manifestation of this uncertainty is inter-observer variability, where different experts provide divergent annotations for the same image. Existing approaches to multi-rater medical image segmentation generally fall into two limited categories:

1. Generation methods: These produce diverse segmentation results but fail to capture expert-specific personalization, often generating variations that do not resemble realistic expert behavior.
2. Personalization methods: These replicate individual annotators' outputs but lack the ability to generate diverse segmentation, effectively hard-coding expert preferences and ignoring inherent image ambiguity.
3. Crowdsourcing methods: These aggregate multiple annotations into a single meta-segmentation, assuming a single ground truth and failing to model the distribution of expert variability.

The paper identifies a gap in the literature: there is no effective approach capable of simultaneously generating segmentation results that are both diverse (capturing the range of expert disagreement) and personalized (aligning with specific expert annotations) within a unified framework.

Methodology: ProSeg
The authors propose ProSeg (Probabilistic modeling of multi-rater Segmentation), a probabilistic graphical model (PGM) designed to unify diversity and personalization. The core innovation lies in the introduction of two continuous latent variables to disentangle the sources of uncertainty:

• $\tau$ (Tau): Models the subjective variations among expert annotators (inter-observer variability). It is inferred from the specific expert annotator $R$.
• $Z$: Models the ambiguity in medical scans (intra-image uncertainty). It is inferred from the observed image data $X$.

Probabilistic Formulation
ProSeg models the joint distribution $p(Y, x, R)$ of annotations ($Y$), images ($x$), and experts ($R$). The model assumes that annotations are generated conditionally based on the latent variables $\tau$ and $Z$.

• Generative Process: The model defines conditional distributions $p(y_{ri} | \tau_i, z_i)$, $p(x | z_i)$, and $p(r_i | \tau_i)$.
• Variational Inference: Since the posterior distributions of the latent variables are intractable, the model employs variational inference to approximate them using $q(z_i | x)$ and $q(\tau | R)$.
• Training Objective: The model is trained by maximizing the Evidence Lower Bound (ELBO), which minimizes the negative log-likelihood of the observed data (reconstruction of images, classification of experts, and segmentation accuracy) while regularizing the posterior distributions of $\tau$ and $Z$ to match their priors (Dirichlet for $\tau$ and Gaussian for $Z$).

Model Architecture
The architecture consists of five neural network modules:

1. Image Encoder: Infers the ambiguity latent variable $z$ from the image $x$.
2. Image Decoder: Reconstructs the image from $z$.
3. Class Embedding: Infers the expert preference latent variable $\tau$ from the expert annotator $R$.
4. Classifier: Predicts the expert annotator from $\tau$.
5. Segmentation Predictor: Generates the segmentation $y_{ri}$ conditioned on both $\tau_i$ and $z_i$.

Control Mechanisms

• Personalization: By specifying a particular expert $r$, the model samples $\tau$ from $q(\tau | r)$, generating segmentation aligned with that specific expert's style while maintaining diversity via sampling $z$.
• Diversity: By sampling $\tau$ from the prior distribution $p(\tau)$ (Dirichlet) rather than a specific expert's posterior, the model generates diverse segmentation results representing the full spectrum of expert variability.

Experimental Results
The authors evaluated ProSeg on two benchmark datasets:

1. LIDC-IDRI: A lung nodule dataset with four expert annotations per image.
2. NPC: A nasopharyngeal carcinoma dataset with four real radiologists, exhibiting higher variability in preferences.

Key Metrics:

• Diversity: Measured by Generalized Energy Distance (GED) and Soft Dice ($D_{soft}$).
• Personalization: Measured by Match Dice ($D_{match}$), Maximum Dice ($D_{max}$), and average Dice scores against specific experts ($D_{mean}$).

Findings:

• ProSeg achieved State-of-the-Art (SOTA) performance on both datasets.
• On LIDC-IDRI, ProSeg achieved the lowest GED (0.1152) and highest $D_{soft}$ (91.53%), outperforming generation methods (e.g., Probabilistic U-Net) and personalization methods (e.g., TAB, Pionono).
• On the more challenging NPC dataset, ProSeg improved diversity performance by 20.71% over previous SOTA personalization methods in terms of GED, while maintaining high personalization accuracy ($D_{mean} = 82.07\%$).
• Ablation studies confirmed that both latent variables are crucial; removing either $\tau$ or $Z$ resulted in significant performance degradation.
• Visual analysis demonstrated that ProSeg generates segmentations that are both diverse and closely aligned with the "Gold Standard" distribution of expert annotations, whereas other methods either lack diversity or fail to match specific expert preferences.

Significance and Claims
The paper claims that ProSeg is the first probabilistic modeling framework to unify diversity and personalization modeling for multi-rater lesion segmentation. By jointly learning these aspects within a single probabilistic framework, ProSeg addresses the limitations of existing methods that treat diversity and personalization as mutually exclusive or require multi-stage training. The authors assert that ProSeg provides a generalizable framework that can be extended to other multi-rater medical image segmentation tasks, offering a robust solution for handling the inherent uncertainty and variability in clinical diagnosis. The work highlights that capturing both the ambiguity of the image and the subjectivity of the annotator is essential for generating reliable and clinically relevant segmentation outputs.