Quantifying Model Uncertainty in Inverse Problems via Bayesian Deep Gradient Descent

Recent advances in reconstruction methods for inverse problems leverage powerful data-driven models, e.g., deep neural networks. These techniques have demonstrated state-of-the-art performances for several imaging tasks, but they often do not provide uncertainty on the obtained reconstructions. In this work, we develop a novel scalable data-driven knowledge-aided computational framework to quantify the model uncertainty via Bayesian neural networks. The approach builds on and extends deep gradient descent, a recently developed greedy iterative training scheme, and recasts it within a probabilistic framework. Scalability is achieved by being hybrid in the architecture: only the last layer of each block is Bayesian, while the others remain deterministic, and by being greedy in training. The framework is showcased on one representative medical imaging modality, viz. computed tomography with either sparse view or limited view data, and exhibits competitive performance with respect to state-of-the-art benchmarks, e.g., total variation, deep gradient descent and learned primal-dual.


I. INTRODUCTION
The task of reconstructing an unobservable signal or image x from a given collection of observations y that have undergone a corruption process (often following a complex forward transform) is ubiquitous in nearly all scientific disciplines, and represents an integral component of many scientific investigations. One notable feature of such a task is that the underlying problem is often ill-posed in the sense that small perturbations in the observational data y can lead to large deviations of the reconstructions, since relevant information might be damped / lost in the forward process.
The pipeline of most reconstruction procedures usually begins with deriving an accurate forward model. To cope with the inherent ill-posedness of the inverse problem, one fruitful idea is to employ regularisation techniques, either explicitly via variational regularisation or implicitly via iterative regularisation (Landweber, Gauss-Newton and EM). One prominent class of reconstruction algorithms is based on variational regularisation, which involves minimising a Tikhonov functional that consists of two terms, i.e., a fidelity term measuring the fitting quality of the model output to the observational data, and a penalty term which encodes a priori knowledge about the sought-for signal [1], [2], [3]. From a statistical standpoint, it may be viewed as the maximum a posteriori estimation of a certain posterior probability density function. During the past few decades, many hand-crafted § Equal contribution. Correspondence to riccardo.barbano.19@ucl.ac.uk. penalties / prior distributions, e.g., Sobolev smoothness, sparsity, total (generalised) variation, and anatomical priors, have been proposed, and have achieved impressive reconstructions (along with rigorous theoretical guarantees, e.g., stability, consistency and convergence rates). We refer to the class of methods that explicitly defines the forward operator and the probability distribution of the noise as knowledge-driven methods. However, for many applied inverse problems, the choice and design of a proper penalty term remains highly nontrivial and an unsuitable choice may greatly compromise the reconstruction accuracy. Furthermore, the model is often only an approximate description to the real-world physical process, and the high complexity associated with the minimisation problem can present the major computational bottleneck.
In recent years, deep learning methods have been widely employed for solving inverse problems [4], [5], [6], [7] and offer novel computational frameworks to tackle the aforementioned shortcomings of knowledge-driven approaches. There are several different strategies to use deep neural networks (DNNs) for image reconstruction, e.g., to reconstruct signals or images purely based on training data (i.e., approximating the inverse map [8]), to build effective penalty terms directly learned from the training data (or even without training data, e.g., deep image prior [9]), and to replace components of an established optimisation algorithm (e.g., gradient descent, proximal gradient iteration, primal-dual algorithm [7] and ADMM [4]) by DNNs; see the work [10] for a recent overview. Very encouraging empirical results have also been demonstrated for several classical and challenging inverse problems, e.g., image denoising [11], super-resolution [12], undersampled MRI [13], low-dose computed tomography (CT) [14] and photo-acoustic tomography [15]. Despite the extremely promising empirical results, purely data-driven approaches have some shortcomings, e.g., the need for a large amount of training data (which is infeasible to acquire at a large scale in medical applications), the lack of good interpretability of deep architectures, and the lack of robustness with respect to adversarial attacks. The unrolling approach (see e.g., [16], [17], [15], [18]) aims at combining the strengths of knowledgedriven and data-driven approaches, by explicitly incorporating the physical models (prior knowledge etc.) and exploiting the good approximation properties and expressivity of DNNs, and thus this class of methods is very promising for solving many challenging inverse problems, and has received enormous attention recently [19].
All the aforementioned works focus on producing one single reconstruction for a given set of corrupted observations, and do not provide relevant uncertainty estimates. Characterising uncertainty in DNN-based solutions to inverse problems is still in its infancy; nevertheless it remains an interesting open problem. In view of the lack of robustness of DNNs, especially in sensitive domains, e.g., medical imaging [20], uncertainty estimates provide valuable additional insights. In the context of deep learning, there are several different types of uncertainties, e.g. model uncertainty (also known as epistemic uncertainty), aleatoric uncertainty, and algorithmic uncertainty. The present work focuses on model uncertainty, which refers to the fact that for a given training dataset, there are a multitude of model configurations that can fit, each giving different predictions on the test dataset. In practice, this occurs frequently, due to the severe over-parameterisation of the DNN model. In this case, it is desirable to assess all different ways that can be used to explain the data, i.e., placing uncertainty over our model parameters [21]. Recovering such uncertainty enables conveying additional information on the confidence we have in the model prediction and is thus important for comprehensive down-stream decision making. Uncertainty quantification of inverse problems can be naturally formulated within a Bayesian framework [22]. This has also been extensively explored in recent years in deep learning, and more recently for learned inversion techniques [23], [24]. The work [23] proposes a data-driven sampling technique for exploring the posterior distribution of the inverse solution based on Wasserstein GANs. The work [24] proposes a probabilistic framework based on conditional VAEs, which potentially also allows incorporating physics, and demonstrates its performance on Gaussian and Poisson denoising. However, both works [23], [24] quantify the uncertainty inherent to the data, i.e., aleatoric uncertainty, but do not provide the uncertainty of the network that fits the data, i.e., epistemic uncertainty.
To quantify model uncertainty of a learned inversion model, one popular approach is to use Bayesian neural networks (BNNs), and to encode the model uncertainty in the network parameters (conditioned on the training data), through a probability distribution instead of a point estimator [25], [26]. However, the exact posterior and posterior predictive distributions are intractable due to the nonlinear dependency of network parameters in the likelihood function and the high-dimensionality of the parameter vector. Although Markov chain Monte Carlo (MCMC) [27] techniques are asymptotically exact, they often do not scale well on BNNs [28], and due to high parameter dimensionality, MCMC still cannot be routinely applied to inverse problems in imaging (note that however, there is important recent progress in this direction using tools from convex optimisation for the posterior distribution on the inverse solution [29]). In the literature, variational inference (VI) [30], [31] is often adopted as a practical approximate inference scheme for BNNs. Note that standard BNN implementations tend to double the number of parameters per layer [26] (e.g., mean and variance in a Gaussian mean field approximation), and consequently, VI often exhibits slow convergence which can potentially greatly compromise the network performances [32], [33]. Thus, there has been extensive research into designing practical approximate inference methods that allow easily scaling up to deep neural architectures. In particular, it has led to connections between VI and dropout [34], approaches which are referred to as Monte Carlo dropout (MCDO) schemes [35], [21]. They have shown promising empirical results in several applications but their drawbacks and poor performances were also noted in sequential decision problems [36], [37]. Thus, it is still of enormous importance to develop a scalable inference procedure providing epistemic uncertainty for the learning-based inversion techniques. In this work, we propose an efficient data-driven knowledge-aided computational framework for quantifying model uncertainty based on Bayesian DNNs for unrolling type learned inversion methods. In particular, it can be viewed as a probabilistic analogue of a recently proposed deep gradient descent [15], and hence the framework is termed as Bayesian Deep Gradient Descent (BDGD). Our main contributions are summarised as follows: • We introduce a tractable and statistically principled framework that provides the model uncertainty. This is achieved by integrating a data-driven knowledge-aided framework with the advances in BNNs and VI. • We propose a "greedy" training scheme (in a statistical context), which trains the framework block-wise, in a manner similar to deep gradient descent [15]. This allows for greatly reducing the training time. Further, we provide an interpretation of the procedure as performing a sequence of constrained variational inference problems. • Building upon [38], we further achieve computational efficiency for training a hybrid architecture where VI is only applied on a small portion of the whole network together with the greedy training scheme. The proposed framework is evaluated on CT reconstruction, from either sparse view or limited view data. The numerical results show that the approach is capable of delivering mean estimators competitive with benchmark algorithms while also providing useful uncertainty information.
The remainder of the paper is organised as follows. In Section II, we introduce the proposed framework and formalise the training and inference procedures. In Section III, we showcase the proposed framework on CT reconstruction. In Section IV, we provide concluding remarks.

II. BAYESIAN DEEP GRADIENT DESCENT A. Learning Gradient Descent
In practical reconstruction, typically iterative methods are employed to optimise the resulting Tikhonov functional. The associated minimisation problem for recovering the signal x from the measurements y is given by where D(y, Ax) is the data fidelity term, A the forward operator describing the data formation mechanism, Diagram of the CNN architecture representing one step update with the k th block. The deterministic layers with parameter φ k are colour-coded in gray. The Gaussian mean-field convolutional layer with parameter θ k is colour-coded in yellow. Note that a ReLU is applied after the summation between the skip connection and the output of the Gaussian mean-field convolutional layer.
penalty term expressing a priori knowledge on the signal x, C a constraint set, and λ ∈ R + the regularisation parameter that balances the data fidelity term and the regularisation term R(x). In unrolled iteration, one constructs a cascade of convolutional neural networks (CNNs) that mimic an iterative minimisation algorithm, and each block within the cascade learns the update for the next iteration: where f φ k denotes one block of the cascade, with φ k being the corresponding parameter vector. Note that f φ k can incorporate a feasibility projection operator P C . At each iteration, x k−1 is updated by the information passed through the gradient ∇D (y, Ax k−1 ). In this way, the reconstruction procedure recurrently incorporates the information about A and its adjoint, which often encodes various important physical laws of the inverse problem. Fig. 1 shows the network architecture of a single block, how it mimics the gradient update, and the overall cascade. The two branches within each block analyse the information of x k−1 and ∇D (y, Ax k−1 ) separately, and the analysis resulting from these two branches is then merged and further processed to give the update δx k−1 . With the skip connection and projection, the update step is done by where δx k−1 is the output of the block before the skip connection and mimics the product of step size and update direction in a gradient descent type algorithm.
By applying the blocks sequentially for K iterations, the reconstructed image is given by:

B. How to Learn in a Bayesian Framework
To provide model uncertainty for the reconstructed signal, we use Gaussian mean-field variational inference (MFVI) to Bayesian convolutional neural networks [26], and learn an approximate posterior distribution q(θ) over each block's parameters θ. Compared to vanilla deterministic DNNs, where the weights are represented by point estimates, BNNs place prior distributions over the weights, and obtain a posterior distribution via approximate inference. Thus, we replace the block's weights with distributions over the parameters [39]. To distinguish between deterministic and stochastic parameters, below we use φ and θ, respectively.
Given a training dataset where Y is the observation matrix and X is the ground truth matrix, we train the cascade using a greedy approach, i.e., block by block. One important issue is to interpret the resulting greedy scheme in a statistical context. To this end, we recursively define the prior distribution and likelihood function of Θ k := (θ 1 , . . . , θ k ) for training the k-th block. Specifically, when training the k-th block, we have already trained all blocks from f θ1 to f θ k−1 and computed the optimal approximate posterior distribution q * (Θ k−1 ). Then we define the joint prior distribution as p(Θ k ) = q * (Θ k−1 )p(θ k |Θ k−1 ), where the (conditional) prior p(θ k |Θ k−1 ) = N (0, I), and take the likelihood function to be p(x|y, Θ k ) = N (f Θ k (∇D, x 0 ), σ 2 k I), where σ 2 k is an additional trainable parameter. The variational family Q k for approximating the true posterior distribution p(Θ k |X, Y ) is chosen to be of the form q( are the variational parameters of the mean field approximation and D is the number of parameters per block. The optimal approximate posterior q * (Θ k ) is then learned by minimising the following constrained loss function over Q k , where KL(·||·) is the Kullback-Leibler divergence [40]. We employ the local reparametrization trick to compute the gradients [41], [42]. Mathematically minimising Eq. (4) is equivalent to the following constrained KL minimising problem min where p(Θ k |X, Y ) is given by Bayes' rule with likelihood function p(X|Y, Θ k ) and prior distribution p(Θ k ). It is worth noting that by its very construction, the approximate posterior distribution automatically admits a factorisation form with respect to the blocks, i.e., q * (Θ k ) = q * (θ 1 ) k i=2 q * (θ i |Θ i−1 ).

D. Practicalities in Training and Inference
Generally, VI methods need some tuning effort to perform well, especially on CNNs, otherwise they may exhibit slow convergence, which can significantly compromise accuracy [43]. Therefore, in the proposed framework, we use a composition of maps f θ k • f φ k to model the k-th block and perform variational inference on the parameters θ k 's in a greedy manner. Specifically, in the k-th block, we denote the parameters of the last layer by θ k , which is a random variable, and parameters of remaining layers in the block by φ k , which is regarded as deterministic variable, c.f., Fig. 1 (Bottom). It is worth noting that we are still optimising with respect to φ k 's. In doing so, we actually regard φ k 's as hyperparameters, use the likelihood functions with hyperparameters p Φ k (x|y, Θ k ) and for the k-th block solve for the following problem For notational simplicity, we omit the notation of Φ and, for instance, denote . Methodologically, it is equivalent to the variational family being a delta approximation (i.e., mean field with zero variance) on some network parameters but a Gaussian mean field approximation on the remaining ones. Naturally, the hybrid approach can greatly reduce the number of variational parameters, especially if the Bayesian component is only a small portion of the overall network, and the resulting cascade has an overall complexity comparable with the deterministic counterpart. The resulting construction retains the Bayesian strength for quantifying epistemic uncertainty with only slightly increased computational efforts and memory requirements. 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 efficient (especially when compared with classical iterative reconstruction algorithms). Recall that the likelihood function of Θ K is p(x|y, Θ K ) = N (f Θ K (∇D, x 0 ), σ 2 K I), and the approximate posterior distribution is given by q * (Θ K ) = q * (θ 1 ) K k=2 q * (θ k |Θ k−1 ). One can use Monte Carlo (MC) estimators to estimate the statistics of the distribution q * (x|y) = p(x|y, Θ K )q * (Θ K )dΘ K . (11) Specifically, E q * (x|y) [x] can be estimated with the unbiased empirical estimator Moreover, the predictive variance of q * (x|y) can be estimated by where x ⊗2 = xx T . Indeed, with T samples. The training and inference procedures of our proposed framework are summarised by Algorithms 1 and 2, respectively. At the end of the training of each block, we reconstruct the next update x k with 1 MC sample and compute the gradient of the data fidelity term D, which is problem specific. In our implementation, at the inference stage, we use 100 MC samples to estimate the mean image and pixel-wise variance to compute the summarising statistics.
Remark (On the Approximate Inference's Landscape). Apart from Mean Field Variational Inference (MFVI), several other Bayesian approximate schemes have been proposed for uncertainty quantification of neural networks, e.g., MCDO [35]. These Bayesian treatments can also be used within the proposed framework as alternatives to MFVI. Note that with other Bayesian treatments, the underlying choice of prior and Algorithm 1: BDGD (Training) Input: # reconstruction steps K, dataset D, initial guess x (i) 0 1 for k ← 1 to K do 2 Construct network's input: Train the k-th network f Φ k ,θ k (∇D(y (i) , Ax k−1 ) Output: approximate posterior at each reconstruction step Algorithm 2: BDGD (Inference) Input: # reconstruction steps K, observation y, initial guess x0, trained parameters (ΦK , ΘK ), # samples S 1 Construct ∇D(y, Ax0)

Output:Ê[x] and Var[x]
thus also the approximate posterior is different. We refer to [21,Section 3.2] for the case of MCDO, which has a VI reinterpretation [21], [35], [44]. It consists in training a DNN with dropout, and also applying dropout at test time, which can be seen as approximate marginalisation.

III. EXPERIMENTAL RESULTS
We showcase the performance of BDGD on CT reconstruction, which is one of the most prominent medical imaging modalities. In (2D) CT, observations are made in the space Y ⊂ H 1 2 (R × [0, π)) which is the range of the forward operator, i.e., Radon transform, R : X → Y and X ⊂ L 2 (R 2 ), consisting of line integrals through X in ray directionsω ∈ S 1 . While the complete data problem in CT is mildly ill-posed and can be exactly reconstructed by filtered back-projection (FBP), in practice often only a subset of data is available, which can be represented by the composition A = S • R with S a subsampling of the directions Ω ⊂ [0, π). We shall consider two different cases of practical interest in the parallel beam geometry, i.e., i) sparse view and ii) limited view. Throughout, we assume that the projections contain 1% Gaussian noise. Then the data fidelity term is given by the standard squared L 2 norm, and accordingly, where the notation A denotes the (unfiltered) back-projection operator (i.e., the adjoint of A). The depth K of the cascade in BDGD is set to either 10 or 20 depending on the problem setting, and the initial guess x 0 is set to be the FBP solution.
For benchmarks, we take three approaches, i.e., total variation regularisation (TV) [45], deep gradient descent (DGD) [15], and learned primal-dual (LPD) [7]. The latter two are well-established deep unrolled iteration approaches. TV reconstruction is computed with the Chambolle-Pock algorithm, with the regularisation parameter selected via grid search. We also report results for a variant of BDGD, which uses MCDO instead of MFVI. The number of parameters (per block) is 32833, 32978, and 32833 for DGD, BDGD-MFVI and BDGD-MCDO, respectively. Meanwhile, LPD has 253220 shared parameters, and the parameters are trained simultaneously instead of greedily.
All the methods are trained on 4000 randomly generated ellipses (all of size 128×128), with the background having the lowest value 0. We train each block for 150 epochs. BDGD and the benchmarks are all implemented in Python using Operator Discretisation Library (ODL) [46], PyTorch and TensorFlow. To evaluate the operator A and its adjoint, we use the GPU accelerated ASTRA backend [47].
For validation, we construct a second ellipse dataset (Ellipses Phantoms), and also test on the Shepp-Logan (SL) phantom, one widely used medical test image. For a quantitative comparison, we use peak signal-to-noise ratio (PSNR), which is computed by averaging the numerical results over five seeds.

A. Sparse View CT
In Table I, we present numerical results for sparse view CT with 30 directions uniformly taken from 0 to π, computed with K = 10 (further increasing the cascade's depth does not lead to better reconstructions). Table I indicates that BDGD-MFVI significantly outperforms all other methods on both datasets with the only exception of LPD on the ellipses datasets. In Fig. 2, we show the reconstructions of the SL phantom, and to shed insights into the working mechanisms of BDGD, we show the mean and pixel variance after the first, fifth, and tenth block, e.g.,Ê[x 10 ] and Var[x 10 ]. It is observed that both BDGD-MFVI and BDGD-MCDO show large pixel variance at k = 1, which is reduced at k = 5. However, BDGD-MFVI shows larger pixel variance at k = 10 than at k = 5, whereas BDGD-MCDO still presents low model uncertainty. The boxplot in the left panel of Fig. 3 (Left) shows the evolution of PSNR with the number k of blocks. Interestingly, during the first few blocks, BDGD-MFVI and BDGD-MCDO perform comparably, but BDGD-MFVI allows further improvement with additional blocks.
Numerically, in BDGD-MFVI, there appear two distinct phases from a variance standpoint. In the first phase, i.e., the reconstruction phase, it gradually improves the sample quality and reduces the overall model uncertainty level. For instance, in Fig. 2, BDGD-MFVI exhibits high variance up to the first block, i.e., the reconstructions disagree the most; but as the method progresses up to the fifth block, the reconstruction shows reduced model uncertainty. In the second phase, i.e., the fine-tuning phase, it improves the sample quality and gradually reveals more uncertainty for fine details. That is, the greedy training strategy actually incrementally improves the reconstruction. Overall, BDGD-MFVI performs better than BDGD-MCDO, which, in terms of PSNR, behaves nearly identical with the deterministic one on the ellipse phantom (we do also fail to observe a fine-tuning phase), whereas on the SL phantom, BDGD-MCDO outperforms DGD, showing again the benefit of being Bayesian.

B. Limited View CT
We show results for limited view [0, 2π/3), using K = 20; see Table II for PSNR and Fig. 4 for exemplary reconstructions. Like before, BDGD outperforms the other methods on both datasets in terms of PSNR. Fig. 4 shows the reconstruction of an ellipse phantom, along with the pixel variance for BDGD (for both MFVI and MCDO). In the phantom, a larger ellipse almost blends into the background, which is is very challenging to reconstruct for all the methods under consideration. However, the presence of the ellipse is well captured by the variance in BDGD-MFVI. Thus, being Bayesian indeed allows a more thorough analysis of the reconstructed image. Note that the uncertainties by BDGD-MFVI capture far more detail about the ellipse edges than those by BDGD-MCDO, and thus are potentially more informative than the latter. The right panel of Fig. 3 (Right) shows the boxplot of the evolution of PSNR with the number of blocks, which exhibits a similar behaviour to the case of sparse view.

C. Model Uncertainty with Different Geometries and Unseen Abnormalities
One of the serious issues in CT reconstructions is possible artefacts due to insufficient information in the data. Thus, it is of much interest to have indicators on artefacts. Interestingly, BDGD can provide indicative information by allocating higher variances in related regions, cf. Fig. 5, which also includes a more challenging scenario, i.e., limited view [0, π/3). In particular, as the view angle range decreases, the magnitude and significant areas of the variance increase, flaring up potential issues with the reconstruction in relevant regions. A very similar phenomenon can be observed, in an even more striking way, in the out-of-distribution (OOD) test data. In this test, we have added the text "ICPR 2020" in the ground truth phantom, following the work [20]. Being abnormal features in the ellipses patterns, large uncertainty quantities now concentrate on the area around the text, indicating potentially serious issues in the region. This shows the benefit of having the uncertainty available in addition to the point estimate.

IV. CONCLUSION AND FUTURE WORK
In this work, we have proposed a novel data-driven knowledge-aided framework, named as Bayesian Deep Gradient Descent (BDGD), for providing model uncertainty within the learning based inversion techniques. We adopted a hybrid model with both point estimation and probabilistic treatment of network parameters, and numerically realised it using an efficient greedy training strategy in a consistent probabilistic context. The numerical results with CT show that it is competitive with state-of-the-art benchmarks. In particular, the results indicate that being Bayesian, even if only a little (one layer per block), can actually be very beneficial to the mean estimate (in terms of PSNR), while also delivering useful uncertainty estimates. There are several avenues worth further pursuing. First, it is of great interest to explore how the predictive variance information can aid the performance of a downstream processing pipeline, e.g., segmentation of reconstructed images. Secondly, in view of the outstanding performance of BDGD, it is of much interest to evaluate the significant potential of the framework on more complex medical imaging settings, e.g., photo-acoustic tomography with limited view geometry, and positron emission tomography with a low count emission level. We will explore these possibilities in future works.