A Probabilistic U-Net for Segmentation of Ambiguous Images

Citation

• Authors: Simon A. A. Kohl, Bernardino Romera-Paredes, Clemens Meyer, Jeffrey De Fauw, Joseph R. Ledsam, Klaus H. Maier-Hein, S. M. Ali Eslami, Danilo Jimenez Rezende, and Olaf Ronneberger
• Year: 2018
• Venue: NeurIPS 2018
• arXiv: arXiv:1806.05034
• URL: https://arxiv.org/abs/1806.05034

One-Line Takeaway

This paper is a strong conditional-VAE reference for genuinely ambiguous outputs: it combines a conditional latent variable model with a U-Net so one input image can produce multiple coherent segmentation hypotheses rather than one averaged prediction.

Technical Core

• The model is a conditional latent-variable segmentation model: $$p_{\theta,\psi,\omega}(S, z \mid X) = p_{\theta,\psi}(S \mid X, z)\, p_\omega(z \mid X),$$ where $X$ is an input image, $S$ is a segmentation mask, and $z \in \mathbb{R}^N$ indexes a segmentation mode.
• Training introduces a posterior network $$q_\nu(z \mid X, Y)$$ over latent codes given the image and one ground-truth mask $Y$.
• The conditional-VAE objective is the usual lower bound: $$\mathcal{L}(X,Y) = \mathbb{E}_{q_\nu(z \mid X,Y)}[\log p_{\theta,\psi}(Y \mid X,z)] - \mathrm{KL}\!\left(q_\nu(z \mid X,Y)\,\|\,p_\omega(z \mid X)\right).$$
• The key modeling point is that the latent variable parameterizes joint segmentation variants, not independent pixelwise uncertainty, so samples correspond to coherent full-image hypotheses.

Method

• Use a prior net to predict an axis-aligned Gaussian $$p_\omega(z \mid X) = \mathcal{N}(\mu_{\mathrm{prior}}(X), \operatorname{diag}(\sigma_{\mathrm{prior}}(X))).$$
• Use a posterior net to predict $$q_\nu(z \mid X,Y) = \mathcal{N}(\mu_{\mathrm{post}}(X,Y), \operatorname{diag}(\sigma_{\mathrm{post}}(X,Y))).$$
• Sample $z$, broadcast it spatially, concatenate it to the last U-Net activation map, and map the combined representation through several $1 \times 1$ convolutions to produce class logits: $$S = f_{\mathrm{comb}}\!\bigl(f_{\mathrm{U\mbox{-}Net}}(X), z\bigr).$$
• The decoder likelihood is implemented through segmentation cross-entropy, so the reconstruction term is the log-likelihood of the target mask under the latent-conditioned segmentation network.
• A notable architectural choice is to inject $z$ late, at the end of the U-Net, which keeps sampling cheap and empirically works better than early injection in their ablations.

Empirical Evaluation

• The paper evaluates on ambiguous lung-lesion segmentation and a modified Cityscapes setup with synthetic ambiguity.
• The headline result is that the model better captures both the set of plausible outputs and their frequencies than dropout, ensembles, multi-head predictors, and an image-to-image VAE baseline.
• Evaluation emphasizes distribution matching through generalized energy distance rather than single-best-mask accuracy alone.

Assumptions And Limits

The paper is about ambiguous supervised vision labels, not sequence generation. The output space is structured but non-autoregressive, and the ambiguity comes from multiple plausible annotations rather than from latent reasoning modes in a language model. It therefore speaks more to conditional multi-modality than to collapse in strong autoregressive decoders.

Relevance To This Project

• Similarities:
  • The paper uses the same core conditional-latent picture as the repo: $$p(z \mid x), \qquad p(y \mid x, z), \qquad q(z \mid x, y).$$
  • Its main motivation is also that one input can support multiple valid outputs, so the model should represent a distribution over latent modes rather than produce one averaged answer.
  • It is a useful example of a learned conditional prior over latents instead of a fixed unconditional prior.
• Differences:
  • The decoder is a segmentation network, not an autoregressive sequence model, so it avoids the exact posterior-collapse dynamics that dominate text VAEs and the repo's language-model setting.
  • The learned modes correspond to annotation ambiguity in images, not latent strategies for solving a reasoning problem.
  • The repo's methodology is more focused on factorizing an already-strong generator and on making latent structure semantically reusable, while this paper learns a conditional generator end-to-end from supervised examples.

Comparison Notes

Sohn et al. (2015) is the more classical CVAE formulation, while this paper is a particularly clean demonstration that a conditional latent can represent coherent one-to-many outputs in a modern structured predictor. Zhao et al. (2017) is the corresponding dialogue example where conditional multi-modality meets stronger sequence-generation challenges.

Related Notes

• Learning Structured Output Representation using Deep Conditional Generative Models
• Learning Discourse-level Diversity for Neural Dialog Models using Conditional Variational Autoencoders
• CVAE post-training methodology