A Hierarchical Probabilistic U-Net for Modeling Multi-Scale Ambiguities

TL;DR

This paper introduces a Hierarchical Probabilistic U-Net that models multi-scale ambiguities in medical image segmentation, enabling high-fidelity sampling and capturing complex distributions across different image scales.

Contribution

It proposes a novel hierarchical latent space within a probabilistic U-Net to better model multi-scale ambiguities in segmentation tasks.

Findings

Effective in segmenting ambiguous medical scans

Capable of detailed and flexible sample generation

Separates independent factors across scales

Abstract

Medical imaging only indirectly measures the molecular identity of the tissue within each voxel, which often produces only ambiguous image evidence for target measures of interest, like semantic segmentation. This diversity and the variations of plausible interpretations are often specific to given image regions and may thus manifest on various scales, spanning all the way from the pixel to the image level. In order to learn a flexible distribution that can account for multiple scales of variations, we propose the Hierarchical Probabilistic U-Net, a segmentation network with a conditional variational auto-encoder (cVAE) that uses a hierarchical latent space decomposition. We show that this model formulation enables sampling and reconstruction of segmenations with high fidelity, i.e. with finely resolved detail, while providing the flexibility to learn complex structured distributions across scales. We demonstrate these abilities on the task of segmenting ambiguous medical scans as well as on instance segmentation of neurobiological and natural images. Our model automatically separates independent factors across scales, an inductive bias that we deem beneficial in structured output prediction tasks beyond segmentation.