Quantifying Sources of Uncertainty in Deep Learning-Based Image Reconstruction
Abstract
Image reconstruction methods based on deep neural networks have shown outstanding performance, equalling or exceeding the state-of-the-art results of conventional approaches, but often do not provide uncertainty information about the reconstruction.
In this work we propose a scalable and efﬁcient framework to simultaneously quantify aleatoric and epistemic uncertainties in learned iterative image reconstruction.
We build on a Bayesian deep gradient descent method for quantifying epistemic uncertainty, and incorporate the heteroscedastic variance of the noise to account for the aleatoric uncertainty.
We show that our method exhibits competitive performance against conventional benchmarks for computed tomography with both sparse view and limited angle data.
The estimated uncertainty captures the variability in the reconstructions, caused by the restricted measurement model, and by missing information, due to the limited angle geometry.
Introduction
In the past few years deep learning (DL) based image reconstruction techniques have demonstrated remarkable empirical results.
As examples in medical imaging, we mention approaches that replace components of established optimisation algorithms (e g., gradient descent [1, 2, 3], primal-dual algorithm [4], ADMM [5]) by DNNs.
See [6, 7] for overviews.
The overwhelming majority of these techniques are deterministic, and there is a lack of solutions that provide mechanisms to estimate the uncertainty.
The absence of quantitative estimates of predictive uncertainty, bias, and robustness has greatly hindered the applicability of many deep learning methods in sensitive domains [8, 9], such as medical imaging and autonomous driving.
There are several types of uncertainty in the context of deep learning [9, 10].
These include aleatoric and epistemic uncertainties, which originate from the stochastic variability inherent in the data generating process, and from the uncertainty in the parameters of the model, respectively.
Uncertainty quantiﬁcation can be conveniently formulated within a Bayesian framework.
Although easy to formulate, such frameworks can be challenging to use since computing quantities from the posterior is computationally intractable.
In addition, learning-based methods that incorporate uncertainty quantiﬁcation often do not outperform, but rather are, at best, on par with classical non-Bayesian approaches [11].
There have been several recent attempts to characterise uncertainty in deep learning-based methods for solving inverse problems, see [12, 13] for aleatoric uncertainty and [14] for epistemic uncertainty, but the topic is still largely in its infancy.
In [14] a greedy, iterative, data-driven, knowledge-aided,
Bayesian approach, termed Bayesian deep gradient descent (BDGD), is developed to solve inverse problems in imaging.
In this work we extend BDGD so that aleatoric and epistemic uncertainties can both be quantiﬁed.
We develop a hybrid architecture (in an unrolled optimisation scheme) by appending a Gaussian mean-ﬁeld CNN to a deterministic CNN-based feature extractor.
Moreover, epistemic and aleatoric uncertainties are accounted for using a decomposition approach proposed in [15], which is achieved by bifurcating the network’s architecture.
Due to its hybrid nature, the resulting Bayesian framework is easy to train using stochastic variational inference (VI) [16], and allows us to improve the reconstruction quality over deterministic DNN techniques.
Our main contributions can be summarised as follows.
Firstly, we introduce a statistically principled framework to quantify aleatoric and epistemic uncertainties in image reconstruction by integrating a data-driven, knowledge-aided framework with the advances in BNNs and VI.
Secondly, we apply the framework to computed tomography (CT) reconstructions, where our results yield interpretable uncertainty.
Model Speciﬁcation
In image reconstruction we aim to recover an image x from (noisy) measurements y , given a forward operator A ,
such that Ax = y .
An established way to obtain a reconstruction is by solving the variational problem
[Equation]
[Equation] where D ( y, Ax ) is the data ﬁdelity and R ( x ) a regularising penalty term, C a constraint set, and λ ∈ R + balances the two terms.
The solution x ∗ is often found by iterative optimisation methods (with initial guess x 0 ).
This produces a sequence of iterates, and the iterations terminate once a stopping criterion is satisﬁed.
Unrolled optimisation [17, 1] is a prominent DL paradigm for image reconstruction.
It mimics an iterative method, but executes only a ﬁxed number of iterations, and at each iteration the update is computed using a DNN.
In this work we consider the iterates computed as residual update by
[Equation]
where the rectiﬁer linear unit, relu( x ) :=
max(0 , x ) , and the increment δx k − 1 is obtained by
[Equation]
To simplify the notation we overload the mappings f φ k and write the iterates as x k = f φ k ( ∇D , x k − 1 ) .
We refer to mappings (networks)
f φ k as blocks , and the collection of all the blocks deﬁning the iterations, as a cascade .
In a cascade K blocks are applied sequentially, and the reconstructed image for an initial guess x 0 is given by
[Equation]
Fig.
1 shows the overall framework.
Note that x k is at each iteration progressively reﬁned with the information passed through the gradient ∇D ( y, Ax k − 1 ) , thus incorporating the physical model underlying the inverse problem into the reconstruction process.
Speciﬁcs of the model and its training.
In our architecture each block consists of deterministic layers and one (mean-ﬁeld CNN)
Bayesian layer (see Fig. 1).
In the following, we will omit the deterministic parameters from the discussion, for the sake of clarity of notation.
Within a Bayesian framework the network parameters θ k are random variables, and we approximate the posterior with
VI.
In the k th block we compute an optimal approximate distribution q ∗ ψ k ( θ ) , parametrised by ψ k , such that it approximates the true posterior p ( θ | X, Y ) , where X, Y is our dataset consisting of ground truth data and corresponding observations.
Moreover, we train the cascade greedily, i.e., one block at a time.
Let Θ k :=
( θ 1 , . .
, θ k ) .
Assume the ﬁrst k − 1 blocks have been learnt, i.e., q ∗ Ψ k − 1 is known.
We consider the variational family Q k that consists of distributions of the form [Equation] with q ψ [Equation]
θ
[Equation]
Θ
[Equation]
=
[Equation]
where ψ k = { ( µ k,(cid:96) , σ 2 k,d ) } Dd =1 are the variational parameters, and D is the number of parameters per block.
In the k th block, we optimise for
[Equation]
where the loss function L k is given by
[Equation]
The prior and the likelihood remain to be speciﬁed.
The former is set recursively as
[Equation]
where p ( θ k ; Θ k − 1 ) =
N (0 , I ) .
Choosing the likelihood adequately allows the capture of either aleatoric or epistemic uncertainty, or both.
To do this, in BDGD [14] we take the likelihood as
[Equation]
Note that we model isotropic homoscedastic noise with variance σ 2 k (a trainable parameter).
Following [9], to capture aleatoric uncertainty, we can use input-dependent variance, and set the likelihood to be
[Equation]
Note that the network σ 2Θ k ( ∇D , x 0 ) is an a posteriori heteroscedastic model.
We refer to this model as BDGD+.
Once the model is trained, following [15], we can decouple aleatoric and epistemic uncertainties by decomposing the (entry-wise) predictive variance Var[ x ] at the K th step, using the law of total variance, and estimating with T ≥ 1 Monte Carlo samples as [Equation]
(cid:123)(cid:122)
[Equation]
epistemic . All the operations in (8) are meant entry-wise, and Θ tK ∼ q ∗ Ψ K
(Θ K ) for each t = 1 , . . .
, T .
We denote the Monte-Carlo estimate on the right-hand side, as (cid:100) Var[ x ] .
More details and practicalities regarding the training of our hybrid architecture, and the inference, can be found in Appendix A.1.
Experimental Results
We showcase the performance of BDGD+ for CT reconstructions.
Here the forward map A is given by the discrete Radon transform and we can compute a direct reconstruction by ﬁltered back-projection (FBP).
As the reconstruction problem is ill-posed, FBP reconstructions typically exhibit strong artefacts, especially if only a subset of data is available.
In the following, we study two cases of practical interest: sparse view and limited angle.
In the former, (cid:96) directions (here (cid:96) ∈ { 8 , 16 , 32 , 64 , 128 } ) are taken uniformly from 0 ◦ to 180 ◦ ; in the latter, the maximum range of available angle data is less than 180 ◦ .
We add 1% Gaussian noise to the sinogram and the data ﬁdelity is the squared L 2 norm 1 2 (cid:107)
Ax − y [Equation] 2 2 , so that
[Equation]
The depth K of the cascade in BDGD+ (and BDGD and DGD) depends on the problem difﬁculty.
For sparse view we use K = 20 for 8 and 16 directions (further increasing K does not lead to better reconstructions), and K = 10 for 32, 64, and 128 directions.
For limited angle we use K =for [0 ◦ , 90 ◦ ) and [0 ◦ , 120 ◦ ) , and K = 20 for [0 ◦ , 150 ◦ ) .
The initial guess x 0 is given by the FBP reconstruction, and we train each block for 150 epochs.
All the methods are trained on 4000 randomly generated ellipses (all of size 128 × 128 ), and are tested on 100 ellipses, and on the Shepp-Logan phantom.
Further details, including the benchmarks used, can be found in Appendix B.
The results are presented in Table 1 and Fig.
The values in Table 1 denote the peak signal to noise ratio (PSNR) of the reconstruction (i.e., the mean for BDGD and BDGD+).
Clearly, BDGD+ is competitive with benchmarks used: FBP+U-Net [18], LPD [4] and DGD [3], which are all far more accurate than non-deep learning based techniques, i.e., FBP and TV [19] (where the regulariser, in (1), is the total variation, R ( x ) =
(cid:107)∇ x (cid:107) 1 ).
Interestingly, BDGD+ consistently improves on BDGD in terms of PSNR, and the improvement can be nontrivial.
This shows the importance of modelling heteroscedastic variance in obtaining high-quality reconstructions.
limited angle, with respect to the available range of angles. The ﬁrst reported number is the mean PSNR over the ellipses, and the second number is the PSNR for the Shepp-Logan.
[Table]
Decomposition (8) allows us to separately quantify aleatoric and epistemic uncertainties.
In both sparse view and limited angle reconstructions, aleatoric uncertainty appears to dominate, with its overall shape close to the mean (but of a smaller magnitude).
In contrast, epistemic uncertainty is localised to certain regions (and is of a smaller magnitude), capturing the "out of distribution" (not in the training data) text in sparse view (cf.
Fig.
3 in the Appendix), or capturing artefacts due to limited angle data.
In particular, we stress that in the latter case, aleatoric uncertainty captures the limits of the geometry.
Thus, aleatoric and epistemic uncertainties provide complementary information about the reconstructions, and might shed different insights into their reliability.
Conclusion
In this work we design a hybrid framework for quantifying the uncertainty in image reconstruction for inverse problems.
The experimental results with sparse view/limited angle CT show that BDGD+ is competitive with conventional benchmarks, and can improve reconstruction quality in terms of PSNR over BDGD.
The uncertainty maps can capture variabilities in the reconstructions caused by different sources, e g., limited information or out of distribution data.
A Additional Materials for Bayesian Deep Gradient Descent (BDGD)
In this part, we provide further details about the proposed framework.
In Appendix A.1, we discuss the training and inference of BDGD and BDGD+.
In Appendix A.2, we review how we model aleatoric and epistemic uncertainties.
In A.3 we analyse a 2-block cascade, for the sake of clarity.
A.1 Practicalities in Training and Inference
Hereafter, we denote the deterministic parameters of the (deterministic) feature extractor as φ k , the random variable of the mean-ﬁeld CNN layer as θ k , and the variational parameters deﬁning the
Bayesian layer as ψ k .
In the proposed framework, we use a composition of maps f θ k ◦ f φ k to model the k th block, and perform VI on the parameters θ k ’s in a greedy manner.
During the training we optimise jointly with respect to φ k ’s as well.
In doing so, φ k ’s are in effect regarded as hyperparameters.
Hence, in the k th block (by abusing the notation) we solve for
[Equation]
where p Φ k ( X | Y, Θ k ) and q Ψ k , Φ k (Θ k ) are respectively the likelihood and the approximate posterior distribution parametrised by the hyperparameters Φ k .
Below we suppress the deterministic parameters Φ from notation, and focus only on the probabilistic ones.
Methodologically, this is equivalent to setting the variational family to be a delta approximation, i.e., mean ﬁeld with zero variance, on some parameters, but a Gaussian mean ﬁeld approximation on the remaining ones.
Note that we can add an L 2 penalty on deterministic parameters Φ K to the lower bound functional.
This does not change the statistical interpretation, since it does not penalise the random parameter θ , and we can add the penalty right after the KL divergence.
The hybrid approach greatly reduces the number of variational parameters, especially if the Bayesian component is only a small portion of the overall network.
In doing so, the resulting cascade has an overall complexity comparable with its deterministic counterpart, while retaining the ability to quantify uncertainty.
Once all the blocks in the cascade are trained, it can be used for inference.
Each sampling step amounts to a feed forward propagation through the network, which is computationally very efﬁcient (at least when compared with classical iterative reconstruction algorithms).
Recall and the likelihood for BDGD is
[Equation]
and for BDGD+ is
[Equation]
The approximate posterior distribution is given by
[Equation]
[Equation]
At test time, for each input y , we can use MC to estimate the statistics, i.e., the mean and the covariance, of the approximate predictive distribution
[Equation]
[Equation]
Speciﬁcally, E [ x ] can be estimated with the unbiased empirical estimator
[Equation]
with T samples of ˆΘ K from q ∗ Ψ K (Θ K ) , i [Equation] e., { ˆΘ t K } Tt =1 .
Moreover, the predictive uncertainty of q ∗ Ψ K ( x | y ) can be estimated with the sample covariance matrix
[Equation]
where x ⊗ 2 = xx (cid:62) .
Indeed
[Equation]
E
[Equation]
=
[Equation]
And
[Equation]
Note that the above derivations suggest an isotropic homoscedastic noise model.
It follows that ˆ E [ x ] and
(cid:100) Cov[ x ] are unbiased MC estimators of E [ x ] and Cov[ x ] =
E (cid:2) x ⊗ 2 (cid:3) − E [ x ] ⊗with T samples.
Note that we use the diagonal element of the sample covariance matrix to quantify the corresponding predictive uncertainty.
The training and inference procedures of our proposed framework are summarised by Algorithms 1 and 2, respectively.
After training all the blocks, we reconstruct the next update x k with 1 MC sample, and compute the gradient of the data ﬁdelity term D , which is problem speciﬁc, dependent on the forward operator and noise statistics.
At inference we use 100 MC samples to estimate the mean image and the pixel-wise variance (cid:100) Var[ x ] , which is a vector consisting of the diagonal elements in (cid:100) Cov[ x ] .
A.2 Modelling Aleatoric Uncertainty via a Heteroscedastic Noise Model
The framework in [14] can be viewed as an example of an isotropic homoscedastic, block-wise noise model, as the variance is not a function of the input, and it is ﬁxed at all spatial locations, i.e., σ 2 k I is background noise.
This is a restrictive and unrealistic assumption for most imaging modalities.
In this work we introduce an anisotropic heteroscedastic noise model, i.e., BDGD+.
For computational tractability, heteroscedastic anisotropic modelling assumes that the covariance matrix is diagonal, i.e., diag( σ 2Θ k ( ∇D , x 0 )) .
We model the likelihood p ( x | y, Θ k ) of the k th block as a Gaussian distribution with input-varying variance,
[Equation]
where f Θ k ( ∇D , x 0 ) is the mean (the reconstructed image), diag( σ 2Θ k ( ∇D , x 0 )) is the covariance matrix, and thus a function of the input x 0 , [ · ] (cid:96) denotes the (cid:96) th entry of a vector, and L is the number of pixels (the dimensionality) of x .
When jointly optimising the loss for the k th block, the negative log-likelihood, NLL( q Ψ k ; X, Y ) , over a mini-batch set B is deﬁned as
[Equation]
where c is an absolute constant, and x i, 0 is the FBP of the data sample x i .
The ﬁrst term
M ( q Ψ k ; X, Y ) is the squared Mahalanobis distance over the mini-batch, deﬁned as
[Equation]
The second term H ( q Ψ k ; X, Y ) is the mean differential entropy over the mini-batch B
[Equation]
When the covariance is diagonal, the ﬁrst term M ( q Ψ k ; X, Y ) corresponds to the mean squared error weighted by the inverse of the corresponding variance.
The second term H ( q Ψ k ; X, Y ) prevents the variance from growing too large.
Input: number of reconstruction steps K , dataset X, Y , initial guesses x i, 0 , batch-size |B| Compute FBPs x i, 0 of all data samples x i 1 for k ← 1 to K do Construct the block’s input:
D k − 1 = { x i,k − 1 , ∇D ( y i , Ax i,k − 1 ) } Ni =
Train the k th block f φ k ,θ k ( ∇D ( y i , Ax i,k − 1 ) , x i,k − 1 ) : // stochastic mini-batch optimisation ψ ∗ k , φ ∗ k ← argmin q Ψ k, Φ k ∈Q k ,φ k (cid:26) ˆ L
[Equation]
φ k , q Ψ
[Equation]
− N |B| (cid:88)
[Equation]
log
[Equation]
(cid:27) // update with ˆ θ k ∼ q ∗ ψ k , Φ k ( θ k ; Θ k − 1 ) x i,k ← f φ k , ˆ θ k ( ∇D ( y i ,
Ax i,k − 1 ) , x i,k − 1 )
Output: approximate posterior at each reconstruction step
Algorithm Input: observation y , trained cascade parameters ( Φ K , Ψ K ), the number T of MC samples 1 //we set T to2 for t ← 1 to T do //with ˆΘ tK ∼ q ∗ Ψ K , Φ K (Θ K )
Sample
[Equation]
=
[Equation]
) 5 Evaluate ˆ E [ x ] and (cid:100)
Var[ x ] with { x tK } Tt =
Output: ˆ E [ x ] and (cid:100)
Var[ x ]
Thus, we quantify the uncertainty of the reconstructed signal at the k th block by computing the predictive uncertainty as
[Equation]
) with Θ tK ∼ q ∗ Ψ K (Θ K ) .
As a reminder, note that (cid:100) Var[ x ] , in (8), is the diagonal of (cid:100) Cov[ x ] .
A.3 Example: 2-block Cascade
We will now analyse a fully-Bayesian 2-block cascade to explain our methodology.
For additional clarity we set D = 1 .
Let θ
and θ 2 be the parameters of the ﬁrst and second blocks, respectively, whereas ψ 1 and ψ 2 the variational parameters of the ﬁrst and second blocks, respectively.
As a reminder, through VI we aim to approximate the true posterior distribution p ( θ 1 , θ 2 | X, Y ) with a simpler distribution q Ψ 2 ( θ 1 , θ 2 ) .
When training the ﬁrst block, we minimise the following functional with respect to the variational parameters
[Equation]
where p ( θ 1 ) = 12 exp( − θ 21 / 2) is the prior distribution of θ 1 , and q ψ 1 ( θ 1 ) = 1 2 σ 21 exp (cid:16) − ( θ 1 − µ 1 ) 2 2 σ(cid:17) , with ψ 1 = ( µ 1 , σ 21 ) , is the approximate distribution in the mean ﬁeld Gaussian family.
Moreover, with a sample of θ 1 , our ﬁrst block f θ 1 ( ∇D , x 0 ) outputs the mean of p ( x | y, θ 1 ) .
Minimising
L 1 ( q ψ 1 ; X, Y ) , for given data X, Y , we obtain an (approximate) optimal posterior distribution of θ[-
-]
[Equation]
This is used to construct the (joint) prior distribution of ( θ 1 , θ 2 ) , which is deﬁned by p ( θ 1 , θ 2 ) =
q ∗ ψ 1 ( θ 1 ) p ( θ 2 ; θ 1 ) , where p ( θ 2 ; θ 1 ) = 12 exp( − θ 22 / 2) is the standard Gaussian distribution.
On the other hand, the approximate posterior distribution of ( θ 1 , θ 2 ) is pursued among distributions of the form q Ψ 2 ( θ 1 , θ 2 ) =
q ∗ ψ 1 ( θ 1 ) q ψ 2 ( θ 2 ; θ 1 ) , where q ψ 2 ( θ 2 ; θ 1 ) = 1 2 σ 2 2 exp (cid:16) − ( θ 2 − µ 2 ) 2 2 σ 2(cid:17) , and where ψ 2 = ( µ 2 , σ 22 ) .
In other words,
[Equation]
Alternatively, we can say that the approximately optimal joint distribution is computed as
[Equation]
over the constrained variational family
[Equation]
Note that during training for samples θ 1 ∼ q ∗ ψ 1 ( θ 1 ) and θ 2 ∼ q ∗ ψ 2 ( θ 2 ; θ 1 ) , the composite function f θ 2 ( f θ 1 ( ∇D , x 0 )) is outputting the mean of p ( x | y, θ 1 , θ 2 ) .
The derivation of the loss in (21) follows as
[Equation]
Let us now expand the negative Evidence
Lower BOund ELBO( q Ψ 2 ; X, Y ) term, L [Equation] Ψ
[Equation]
X, Y
[Equation]
Note that this can be interpreted as KL (cid:63) (
q Ψ k ( θ k ; Θ k − 1 ) (cid:107) p (Θ k
| X, Y )) with respect to a weighted measure ( q ∗ Ψ k (Θ k − 1 ) d
Θ k − 1 dθ k ) .
B Benchmarks Implementation Details
We consider four benchmark approaches:
• Total variation regularisation (TV) [20]; • FBP + U-Net [18]; • Deep gradient descent (DGD) [3]; • Learned primal-dual (LPD) [4].
Total variation [20] is an established image reconstruction technique, suitable for recovering piecewise constant like images.
The latter three are well-established deep unrolled iteration approaches.
Speciﬁcally, in FBP + U-Net, U-Net learns the identity map, as it fails to recover the missing information in the limited angle setting.
The latter could be overcome by either adding an additional loss term [3] or using extensive training data [21].
DGD can be considered a greedy version of LGS [2], and it is used here for consistency with BDGD.
TV reconstruction is computed with the Chambolle-Pock algorithm, with the regularisation parameter selected via grid search.
BDGD/BDGD+ and the benchmarks are all implemented in Python using the Operator Discretisation Library (ODL)
[22], PyTorch and TensorFlow.
To evaluate the operator
A and its adjoint, we use the GPU accelerated ASTRA backend [23].
For LPD we set the common parameters of training, e.g., number of iterations/epochs, batch size and number of batches, the same as DGD/BDGD to enforce comparability.
Note that LPD is an unrolled scheme trained end-to-end.
Hence, the parameters of the model are trained simultaneously instead of greedily.
C Additional Experimental Figures
In Fig.
3 we compare how the estimation of uncertainty, in the sparse view case, changes as we add out of distribution text.
We also include reconstructed SL phantoms with deep learning-based methods for sparse view and limited angle problems in Fig. 4, and Fig. 5, respectively.
D Additional Details on the Architectures
In Fig.
6 we report the architecture of the network representing one step update with the k th block.
Fig.
7 shows the architecture used for modelling the input-dependent variance in BDGD+.
x k −3 × 3 ×3 × 3 ×∇D 3 × 3 ×3 × 3 ×323 × 3 ×3 × 3 ×3 × 3 ×+ x k
[Figure]