Uncertainty quantification in medical image segmentation with normalizing flows

Raghavendra Selvan,Frederik Faye,Jon Middleton,Akshay Pai

DOI: https://doi.org/10.48550/arXiv.2006.02683

2020-08-04

Abstract:Medical image segmentation is inherently an ambiguous task due to factors such as partial volumes and variations in anatomical definitions. While in most cases the segmentation uncertainty is around the border of structures of interest, there can also be considerable inter-rater differences. The class of conditional variational autoencoders (cVAE) offers a principled approach to inferring distributions over plausible segmentations that are conditioned on input images. Segmentation uncertainty estimated from samples of such distributions can be more informative than using pixel level probability scores. In this work, we propose a novel conditional generative model that is based on conditional Normalizing Flow (cFlow). The basic idea is to increase the expressivity of the cVAE by introducing a cFlow transformation step after the encoder. This yields improved approximations of the latent posterior distribution, allowing the model to capture richer segmentation variations. With this we show that the quality and diversity of samples obtained from our conditional generative model is enhanced. Performance of our model, which we call cFlow Net, is evaluated on two medical imaging datasets demonstrating substantial improvements in both qualitative and quantitative measures when compared to a recent cVAE based model.

Machine Learning,Computer Vision and Pattern Recognition

What problem does this paper attempt to address?

The problem that this paper attempts to solve is to quantify uncertainty in medical image segmentation. Specifically, medical image segmentation is an inherently ambiguous task, and factors such as partial volume effects and changes in anatomical definitions make the segmentation results uncertain. Moreover, in most cases, this uncertainty is concentrated near the boundaries of the structures of interest, but there are also significant inter - rater differences. The paper proposes a new conditional generative model based on conditional normalizing flows (cFlow) to improve the approximation ability of the underlying posterior distribution, thereby capturing more abundant segmentation variations. Through this method, the paper aims to provide more informative estimates of segmentation uncertainty than traditional pixel - level probability scores.

The paper mentions that traditional conditional variational auto - encoders (cVAE) can be used to infer the possible segmentation distributions based on input images, but their latent representation capabilities are limited. To overcome this limitation, the authors introduce a conditional normalizing flow step to enhance the model's ability to express complex latent posterior distributions, thereby capturing more diverse and realistic segmentation results at a single resolution. Experimental results show that, compared with existing cVAE models, the proposed cFlow Net model has significant improvements in both quantitative and qualitative metrics.