Improving Aleatoric Uncertainty Quantification in Multi-annotated Medical Image Segmentation with Normalizing Flows

Abstract

Quantifying uncertainty in medical image segmentation applications is essential, as it is often connected to vital decision-making. Compelling attempts have been made in quantifying the uncertainty in image segmentation architectures, e.g. to learn a density segmentation model conditioned on the input image. Typical work in this field restricts these learnt densities to be strictly Gaussian. In this paper, we propose to use a more flexible approach by introducing Normalizing Flows (NFs), which enables the learnt densities to be more complex and facilitate more accurate modeling for uncertainty. We prove this hypothesis by adopting the Probabilistic U-Net and augmenting the posterior density with an NF, allowing it to be more expressive. Our qualitative as well as quantitative (GED and IoU) evaluations on the multi-annotated and single-annotated LIDC-IDRI and Kvasir-SEG segmentation datasets, respectively, show a clear improvement. This is mostly apparent in the quantification of aleatoric uncertainty and the increased predictive performance of up to 14%. This result strongly indicates that a more flexible density model should be seriously considered in architectures that attempt to capture segmentation ambiguity through density modeling. The benefit of this improved modeling will increase human confidence in annotation and segmentation, and enable eager adoption of the technology in practice.

Appendices

Appendices

A Probabilistic U-Net Objective

The loss function of the PU-Net is based on the standard ELBO and is defined as

$$\begin{aligned} \begin{aligned} \mathcal {L}&= -\mathbb {E}_{q_\phi (\mathbf {z}\vert \mathbf {s},\mathbf {x})}[\,{\text {log}}p(\mathbf {s}\vert \mathbf {z},\mathbf {x})\,]+{\text {KL}}\left( \,q_\phi (\mathbf {z}\vert \mathbf {s},\mathbf {x})\vert \vert p_\psi (\mathbf {z}\vert \mathbf {x})\,\right) , \end{aligned} \end{aligned}$$

(4)

where the latent sample $\mathbf {z}$ from the posterior distribution is conditioned on the input image $\mathbf {x}$, and ground-truth segmentation $\mathbf {s}$.

B Planar and Radial Flows

Normalizing Flows are trained by maximizing the likelihood objective

$$\begin{aligned} \begin{aligned} \log p(\mathbf {x})=\log p_{0}\left( \mathbf {z}_{0}\right) -\sum _{i=1}^{K} \log \left( \left| {\text {det}} \frac{d f_{i}}{d \mathbf {z}_{i-1}}\right| \right) . \end{aligned} \end{aligned}$$

(5)

In the PU-Net, the objective becomes

$$\begin{aligned} \begin{aligned} \log q(\mathbf {z}\vert \mathbf {s},\mathbf {x})=\log q_{0}(\mathbf {z}_0\vert \mathbf {s},\mathbf {x})-\sum _{i=1}^{K} \log \left( \left| {\text {det}} \frac{d f_{i}}{d \mathbf {z}_{i-1}}\right| \right) , \end{aligned} \end{aligned}$$

(6)

where the i-th latent sample $\mathbf {z}_i$ from the Normalizing Flow is conditioned on the input image $\mathbf {x}$, and ground-truth segmentation $\mathbf {s}$.

The planar flow expands and contracts distributions along a specific directions by applying the transformation

$$\begin{aligned} f(\mathbf {x})=\mathbf {x} + \mathbf {u}h(\mathbf {w}^T\mathbf {x}+\mathbf {b}), \end{aligned}$$

(7)

while the radial flow warps distributions around a specific point with the transformation

$$\begin{aligned} f(\mathbf {x})=\mathbf {x} + \frac{\beta }{\alpha \left| \mathbf {x}-\mathbf {x}_0\right| }(\mathbf {x}-\mathbf {x}_0). \end{aligned}$$

(8)

C Dataset Images

Here example images from the datasets used in this work can be seen. Figure 4 depicts four examples from the LIDC dataset. On the left in the figure the 2D CT image containing the lesion, followed by the four labels made by four independent annotators is shown. In Fig. 5, eight examples from the Kvasir-SEG dataset is depicted. An endoscopic image with its ground truth label can be seen.

Fig. 4.

Example images from the LIDC dataset.

Fig. 5.

Example images from the Kvasir-SEG dataset.

D Sample Size Dependent GED

The GED evaluation is dependent on the number of reconstructions sampled from the prior distribution. Figure 6 depicts this relationship for the vanilla, 2-planar and 2-radial posterior models. The uncertainty in the values originate from the changing results when training with ten-fold cross validation. One can observe that with increasing sample size, the GED as well as the associated uncertainty decrease. This is also the case when the posterior is augmented with a 2-planar or 2-radial flow. Particularly, the uncertainty in the GED evaluation significantly decreases.

Fig. 6.

The GED based on sample size evaluated on the vanilla, 2-planar and 2-radial models.

E Prior Distribution Variance

We investigated whether the prior distribution captures the degree of ambiguity in the input images. For every input image $\mathbb {X}$, we obtain a latent L-dimensional mean and standard deviation vector of the prior distribution $P(\boldsymbol{\mu }, \boldsymbol{\sigma }\vert \mathbb {X})$. The mean of the latent prior variance vector $\mu _{LV}$, is obtained from the input images in an attempt to quantify this uncertainty. Figure 7 shows this for several different input images of the test set. As can be seen, the mean variance over the latent prior increases along with a subjective assessment of the annotation difficulty.

Fig. 7.

Depicted in the CT image is the mean of the prior distribution variance of the 2-planar model. We show the input CT image, its average segmentation prediction (16 samples) and ground truth from four annotators.