Uncertainty Estimation in Medical Image Denoising with Bayesian Deep Image Prior

Abstract

Uncertainty quantification in inverse medical imaging tasks with deep learning has received little attention. However, deep models trained on large data sets tend to hallucinate and create artifacts in the reconstructed output that are not anatomically present. We use a randomly initialized convolutional network as parameterization of the reconstructed image and perform gradient descent to match the observation, which is known as deep image prior. In this case, the reconstruction does not suffer from hallucinations as no prior training is performed. We extend this to a Bayesian approach with Monte Carlo dropout to quantify both aleatoric and epistemic uncertainty. The presented method is evaluated on the task of denoising different medical imaging modalities. The experimental results show that our approach yields well-calibrated uncertainty. That is, the predictive uncertainty correlates with the predictive error. This allows for reliable uncertainty estimates and can tackle the problem of hallucinations and artifacts in inverse medical imaging tasks.

A Appendix

1.1 A.1 Additional Figures

(See Figs. 6, 7, 8 and 10)

Fig. 6.

Denoised images after convergence.

Fig. 7.

Denoised images with early-stopping applied.

Fig. 8.

MSE (top row) between denoised $\hat{\textit{\textbf{x}}}$ image and noisy image $\tilde{\textit{\textbf{x}}}$ and SSIM (bottom row) between denoised $\hat{\textit{\textbf{x}}}$ image and ground truth $\textit{\textbf{x}}$ vs. iteration. Only MCDIP does not overfit the noisy image and converges with highest similarity to the ground truth. Despite the claim of the authors, SGLD suffers from overfitting and creates the need for carefully applied early stopping [4]. Note: We compared both our own implementation of SGLD and the original code provided by the authors. The plots show means from 3 runs with different random initialization.

1.2 A.2 Additional Tables

(See Table 2, 3 and 4)

Table 2. PSNR with early-stopping.

Table 3. SSIM after convergence.

Fig. 9.

Calibration diagrams and uncertainty maps for SGLD+NLL and MCDIP after convergence (best viewed with digital zoom). (Left) The calibration diagrams show MSE vs. uncertainty and provide mean uncertainty (U) and UCE values. (Right) Uncertainty maps show per-pixel uncertainty. Due to overfitting, the MSE and uncertainty from SGLD+NLL concentrates around 0.0.

Fig. 10.

Calibration diagrams and uncertainty maps for SGLD+NLL after early stopping and MCDIP after convergence (best viewed with digital zoom). (Left) The calibration diagrams show MSE vs. uncertainty and provide mean uncertainty (U) and UCE values. (Right) Uncertainty maps show per-pixel uncertainty.

Table 4. SSIM with early-stopping.

1.3 A.3 SGLD with Step Size Decay

Additionall, we implement SGLD with step size decay as described by Welling et al. [26]. The step size $\epsilon$ is used to scale the parameter update in the SGD step (i.e. the learning rate) and defines the variance of the noise that is injected into the gradients. Here, we reduce the step size at each step t exponentially with $\epsilon _{t} = 0.999^{t} \epsilon _{0}$. To satisfy the step size property (Eq. (2) in [26]), we fix the step size once it decreases below 1e-8. We observe no overfitting of the noisy image with step size decay (see Fig. 11). However, the quality of the resulting denoised image is very sensitive to the decay scheme. Choosing a decrease that is too low (i.e. $\epsilon _{t} = 0.9999^{t} \epsilon _{0}$) results in overfitting; a decrease that is too high (i.e. $\epsilon _{t} = 0.99^{t} \epsilon _{0}$) results in convergence to a subpar reconstruction. This is equivalent to carefully applied early stopping and therefore nullifies the advantage of SGLD for denoising of medical images.

Fig. 11.

Comparison of SGLD and SGLD+LR (with step size decay). Carefully chosen step size decay impedes overfitting the noisy image. (Right) Reconstruction of SGLD+LR after convergence (no early stopping applied).

1.4 A.4 Downsampling

Here, we provide justification why downsampling of an image by averaging neighboring pixels reduces the noise level and can be used as an approximation to a ground truth noise-free image (by sacrificing image resolution).

Proposition 1

Downsampling of an image reduces the observation noise.

Proof

Let $X = \mu _{x} + \varepsilon _{x}$ and $Y = \mu _{y} + \varepsilon _{y}$ be two neighboring pixels affected by additive i.i.d. noise $\varepsilon _{x} , \varepsilon _{y} \sim \mathcal {N}(0, \sigma ^{2})$. The pixels are assumed to be uncorrelated to noise. Pixels in a local neighborhood are highly correlated and assumed to be of high similarity $\mu _{x} \approx \mu _{y} = \mu$. Let $Z = \tfrac{1}{2} \left( X + Y \right)$ be the average of two neighboring pixels (i.e. the result of downsampling). The expectation is given by

$$\begin{aligned} \mathbb {E}[Z]&= \frac{1}{2} \left( \mathbb {E}[X] + \mathbb {E}[Y] \right) \end{aligned}$$

(9)

$$\begin{aligned}&= \frac{1}{2} 2 \, \mathbb {E}[X] \end{aligned}$$

(10)

$$\begin{aligned}&= \mu \end{aligned}$$

(11)

and the variance is given by

$$\begin{aligned} \mathrm {Var}\left[ Z\right]&= \mathrm {Var}\left[ \frac{1}{2} \left( X + Y \right) \right] \end{aligned}$$

(12)

$$\begin{aligned}&= \frac{1}{2^{2}} \left( \mathrm {Var}\left[ X\right] + \mathrm {Var}\left[ Y\right] \right) \end{aligned}$$

(13)

$$\begin{aligned}&= \frac{1}{2^{2}} 2 \mathrm {Var}\left[ X\right] \end{aligned}$$

(14)

$$\begin{aligned}&= \frac{1}{2} \sigma ^{2} ~ . \end{aligned}$$

(15)

Thus, if the similarity of neighboring pixels is sufficiently high, downsampling reduces the variance of average pixel Z by a factor of 2.    $\square$

Naturally, two neighboring pixels are not exactly equal. However, downsampling can also be viewed as superposing two signals, each with a highly correlated and an uncorrelated part. Without providing proof, the amplitude of the addition of two signals can be viewed as vector addition. In the uncorrelated case, the two signals are perpendicular to each other and in the correlated case, the angle between the two signals is acute. Thus, the correlated parts of the two signals have a higher impact on the resulting addition than the uncorrelated (noise) parts. In the ideal case, where the noise is uncorrelated and the signals are in parallel, the same noise reduction as above follows.

1.5 A.5 Link Between Poisson Distribution and Normal Distribution

We approximate the Poisson noise to simulate a low-dose X-ray image with a Normal distribution. It is well-known that the limiting distribution of $\mathsf {Poisson}(\lambda )$ is Normal as $\lambda \rightarrow \infty$ [10]. For completeness, we list a common proof using the moment generating function of a standardized Poisson random variable:

Theorem 1

The Poisson($\lambda$) distribution can be approximated with a Normal distribution as $\lambda \rightarrow \infty$.

Proof

Let $X_{\lambda } \sim \mathsf {Poisson}(\lambda ), ~ \lambda \in \{ 1, 2, \ldots \}$. The probability mass function of $X_{\lambda }$ is given by

$$\begin{aligned} f_{X_{\lambda }}(x) = \frac{\lambda ^{x}e^{-\lambda }}{x!} \quad x \in \{ 0, 1, 2, \ldots \} ~ . \end{aligned}$$

(16)

The moment generating function is given by [10]

$$\begin{aligned} M_{X_{\lambda }}(t) = \mathbb {E} [ e^{t X_{\lambda }} ] = e^{\lambda (e^{t}-1)} ~ . \end{aligned}$$

(17)

The standardized Poisson random variable

$$\begin{aligned} Z = \frac{X_{\lambda } - \lambda }{\sqrt{\lambda }} \end{aligned}$$

(18)

has the limiting moment generating function

$$\begin{aligned} \lim _{\lambda \rightarrow \infty } M_{Z} (t)&= \lim _{\lambda \rightarrow \infty } \mathbb {E} \left[ \exp {\left( t \cdot \frac{X_{\lambda } - \lambda }{\sqrt{\lambda }} \right) } \right] \end{aligned}$$

(19)

$$\begin{aligned}&= \lim _{\lambda \rightarrow \infty } \exp { \left( -t \sqrt{\lambda } \right) } \mathbb {E} \left[ \exp {\left( \frac{t X_{\lambda }}{\sqrt{\lambda }} \right) } \right] \end{aligned}$$

(20)

$$\begin{aligned}&= \lim _{\lambda \rightarrow \infty } \exp { \left( -t \sqrt{\lambda } \right) } \exp {\left( \lambda \left( e^{t/\sqrt{\lambda }} - 1 \right) \right) } \end{aligned}$$

(21)

$$\begin{aligned}&= \lim _{\lambda \rightarrow \infty } \exp { \left( -t \sqrt{\lambda } + \lambda \left( t \lambda ^{-1/2} + t^{2} \lambda ^{-1}/2 + t^{3} \lambda ^{-3/2}/6 + \ldots \right) \right) } \end{aligned}$$

(22)

$$\begin{aligned}&= \lim _{\lambda \rightarrow \infty } \exp { \left( t^{2} / 2 + t^{3}\lambda ^{-1/2}/6 + \ldots \right) } \end{aligned}$$

(23)

$$\begin{aligned}&= \exp {\left( t^{2} / 2 \right) } \end{aligned}$$

(24)

which is the moment generating function of a standard normal random variable.    $\square$