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

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.

25 pages, 15 figures

References in corpus (14)

• Auto-Encoding Variational Bayes
• Natural Scales in Geographical Patterns
• What Uncertainties Do We Need in Bayesian Deep Learning for Computer Vision?
• Semi-Supervised Learning with Deep Generative Models
• Semantic Image Synthesis with Spatially-Adaptive Normalization
• Deep learning to achieve clinically applicable segmentation of head and neck anatomy for radiotherapy
• Superhuman Accuracy on the SNEMI3D Connectomics Challenge
• Why M Heads are Better than One: Training a Diverse Ensemble of Deep Networks
• Bridging Category-level and Instance-level Semantic Image Segmentation
• BIVA: A Very Deep Hierarchy of Latent Variables for Generative Modeling
• Do Deep Generative Models Know What They Don't Know?
• Instance Segmentation by Deep Coloring
• Hierarchical Autoregressive Image Models with Auxiliary Decoders
• UPSNet: A Unified Panoptic Segmentation Network

Cited by in corpus (8)

• A Review of Uncertainty Quantification in Deep Learning: Techniques, Applications and Challenges
• Stochastic Segmentation Networks: Modelling Spatially Correlated Aleatoric Uncertainty
• Attentional Prototype Inference for Few-Shot Segmentation
• Fully Unsupervised Diversity Denoising with Convolutional Variational Autoencoders
• Weakly- and Semi-Supervised Probabilistic Segmentation and Quantification of Ultrasound Needle-Reverberation Artifacts to Allow Better AI Understanding of Tissue Beneath Needles
• Uncertainty quantification in medical image segmentation with normalizing flows
• Safety-aware Motion Prediction with Unseen Vehicles for Autonomous Driving
• An Uncertainty-aware Hierarchical Probabilistic Network for Early Prediction, Quantification and Segmentation of Pulmonary Tumour Growth