Conditional Variational Autoencoder for Learned Image Reconstruction

Abstract

Learned image reconstruction techniques using deep neural networks have recently gained popularity and have delivered promising empirical results. However, most approaches focus on one single recovery for each observation, and thus neglect information uncertainty. In this work, we develop a novel computational framework that approximates the posterior distribution of the unknown image at each query observation. The proposed framework is very flexible: it handles implicit noise models and priors, it incorporates the data formation process (i.e., the forward operator), and the learned reconstructive properties are transferable between different datasets. Once the network is trained using the conditional variational autoencoder loss, it provides a computationally efficient sampler for the approximate posterior distribution via feed-forward propagation, and the summarizing statistics of the generated samples are used for both point-estimation and uncertainty quantification. We illustrate the proposed framework with extensive numerical experiments on positron emission tomography (with both moderate and low-count levels) showing that the framework generates high-quality samples when compared with state-of-the-art methods.

1. Introduction

Machine learning techniques, predominantly those using deep neural networks (DNNs), have received attention in recent years and delivered state-of-the-art reconstruction performance on many classical imaging tasks, including image denoising [1], image deblurring [2], and super-resolution [3], as well as on challenging applications such as low-dose and sparse-data computed tomography [4,5] and under-sampled magnetic resonance imaging [6], just to name a few.

Most existing works on DNN-based image reconstruction aim at providing a one-point estimate of the unknown image of interest for a given (noisy) observation; this implies that we regard DNNs as deterministic machineries. Nonetheless, the stochastic nature of practical imaging problems (e.g., a noisy observation process, or imprecise prior knowledge) implies that there actually exists an ensemble of plausible reconstructions, which are still consistent with the given data (although to various extents). This observation has motivated the need to develop a fully probabilistic treatment of the reconstruction task in order to assess the reliability of one specific reconstructed image so as to inform downstream decision making [7]. Unfortunately, information uncertainty/reliability is not directly available from most existing DNN-based image reconstruction techniques, and hence there is an imperative need to develop uncertainty quantification (UQ) frameworks that are capable of efficiently delivering uncertainty information along with the reconstruction [8].

The Bayesian framework provides a systematic yet very flexible way to address many UQ tasks by modelling the data and the unknown image probabilistically (i.e., as random variables); this has been popular for UQ in imaging inverse problems [9,10]. Furthermore, recent advances in Bayesian inference leverage powerful deep generative modelling tools such as the Wasserstein generative adversarial networks (GANs) [11] and variational autoencoders (VAEs) [12]. Deep generative models output one sample of the unknown image via feed-forward propagation, which is computationally feasible and can also be performed in parallel for multiple samples. These techniques hold enormous potential for the quantification of uncertainty in image reconstruction; nonetheless, developing rigorous data-driven image reconstruction techniques within a Bayesian framework remains challenging.

The first challenge stems from the high complexity of the posterior distribution of the unknown image (conditioned on a given observation). The “conventional” Bayesian setting often involves an explicit likelihood and the construction of a prior. In fact, in Bayesian inversion, both likelihood and prior are given explicitly (or hierarchically). The likelihood is derived from the statistical model of an observation process, under the assumption that both the noise statistics and the underlying physical principles of the imaging modality are well-calibrated. Nonetheless, deriving precise likelihoods can be highly nontrivial (e.g., due to a complex corruption process). Moreover, how to stipulate statistically meaningful yet explicit priors is a long-standing open problem in image reconstruction. In the context of learning-based approaches, one can only implicitly extract prior information in a “data-driven” way using deep neural networks from a set of training pairs instead of an explicit posterior distribution (e.g., up to a normalising constant), leading to an approximate posterior, which is “intrinsically implicit”. Furthermore, the UQ task is substantially complicated by the high-dimensionality of the reconstructed image, which renders the many conventional sampling type techniques not directly applicable.

The second challenge is related to the presence of physical laws. In many practical imaging applications, the data formation process itself is highly complex. To make matters even worse, the forward operator itself can be implicitly defined, for instance, via a system of differential equations in many medical imaging modalities such as the second-order elliptic differential equation in electrical impedance tomography [13], or the radiative transfer equation in diffuse optical tomography [14]. The forward maps often describe the fundamental physical laws and can be regarded as “established” prior knowledge. Therefore, it is useful to directly inform DNNs with the underlying physical laws (instead of having the DNNs learn them from the given training data); how to best use the physical laws represents an important issue in architectural design, which is currently being actively researched.

In this work (a preliminary version of the paper appeared as the arXiv preprint [15]), we develop a novel computational framework for uncertainty-aware learned image reconstruction with the use of a conditional variational autoencoder (cVAE) [16], implemented via algorithmic unrolling [17]. The framework recurrently refines the (stochastic) reconstruction that benefits directly from the physical laws as well as from the samples of a low-dimensional random variable, which is conditionally dependent on the observation. Furthermore, minimising the cVAE objective is equivalent to optimising an upper bound of an expected Kullback–Leibler (KL) divergence (see Proposition 1 for a precise statement), and the output samples (i.e., the plausible reconstructions) are actually drawn from an approximate posterior distribution of the unknown image. In sum, the proposed framework has the following features: (i) it serves as an efficient sampler from an approximate posterior in the sense of (generalised) variational inference [18]; (ii) it is scalable due to its encoder-decoder structure, more specifically due to the low-dimensionality of the latent variable. Note that the low-dimensional latent variable in cVAE is introduced in order to inject stochasticity into the reconstructions, while the heavy-duty reconstruction task is still carried out by the unrolling construction. Thus, the approach does not suffer from the over-smoothing effect typically observed for VAEs.

This work is performed within the context of uncertainty quantification for learned image reconstruction techniques. There are mainly two different types of uncertainties: aleatoric [7,19] and epistemic [7,20] uncertainties. See the review in [8] for UQ in medical image analysis, and the review in [21] for the UQ of DL techniques in general. We focus on aleatoric uncertainty associated with the reconstructed image. This arises from the intrinsic stochasticity with the given data, and it is irreducible even if one can collect more data. It differs from epistemic uncertainty, which is often a by-product of a Bayesian treatment for neural networks, that is, the Bayesian neural networks (BNNs) [20,22,23,24,25]). BNNs model the uncertainty within network weights, which are then inferred with approximate inference techniques [18,26,27]. The corresponding uncertainty arises from the fact that there is an ensemble of weight configurations that can explain the given training data equally well since there are generally far more network parameters than the available training data (also known as overparameterisation).

The rest of the paper is organised as follows. In Section 2, we describe the related works. In Section 3, we describe the problem formulation. Section 4 provides background information on the conditional variational autoencoder and describes our proposed framework, including the network architecture and the training and inference phases. In Section 5, we showcase our framework on an established medical imaging modality—positron emission tomography (PET) [28]—for which uncertainty quantification has long been desired yet is still very challenging to achieve [29,30,31], and confirm that the generated samples are indeed of high quality in terms of both point estimation and uncertainty quantification when compared with several state-of-the-art benchmarks. Finally, in Section 6, we conclude the paper with additional discussions.

2. Related Works

This work lies at the interface of the following three topics: learned image reconstruction, uncertainty quantification using generative models, and approximate inference. In this section, we briefly review related works.

Learned reconstruction techniques have received a lot of attention recently. Most of them are learned in a supervised manner [32]. One prominent idea is algorithmic unrolling, i.e., to unfold classical iterative optimization schemes and to replace certain components of the schemes with neural networks, whose parameters are then learned from training data. This construction brings good interpretability into the function of the algorithms, and allows for the incorporation of physical models, thus often delivering high-performance results without resorting to very large training datasets. This was first performed in [33], which unrolls the iterative shrinkage thresholding algorithm for sparse coding. In the context of inverse problems, an early work is [34]. Currently, there is a large body of works on unrolling; see [17] for a recent overview. These approaches focus mostly on point estimates and ignore the associated uncertainties. This work employs the unrolling idea to construct the backbone network, but also endows the reconstructions with uncertainty estimates using the low-dimensional latent space.

Over the last decade, there has been impressive progress on the UQ of DNNs, especially Bayesian neural networks [7,20,22,23,27] (see [21] for an overview). Nonetheless, high-dimensional inference tasks such as image reconstruction remains largely under-explored due to the associated computational cost [8]; thus, the UQ of learned image reconstruction techniques is just emerging. For epistemic uncertainty, the most common technique used in image reconstruction is the Monte Carlo dropout [20] due to its low computational cost, although it may sacrifice the accuracy [24]. One may further separate the sources of uncertainties into aleatoric and epistemic ones [25], if desired. This work addresses aleatoric uncertainty with a reconstructed image by approximating an implicitly defined posterior distribution, which differs markedly from existing works on epistemic uncertainty.

In the machine learning community, there is a myriad of different ways to approximate a target distribution using an approximate inference scheme, e.g., variational inference [35], variational Gaussian [36,37], expectation propagation [26,29], and Laplace approximation [38]. This work provides one way to approximate the distribution using cVAE. Note that to rigorously justify the distribution modelled by DNNs for aleatoric uncertainty, proper Bayesian interpretations of the loss function used in the training of DNNs (in connection with the target posterior distribution) are needed. The cVAE approach adopted in the present work does admit the interpretation as generalised variational inference (cf. Proposition 1). Very recently, a graph version of cVAE was employed for synthesizing cortical morphological changes in [39]. Of course, cVAE represents only one way of constructing the approach. Other approaches, e.g., the generative adversarial networks [11], also show potential by suitably adapting the corresponding loss function and equipping it with a Bayesian interpretation. For example, the interesting work performed in [40] constructs a sampler by approximating the transport map with DNNs. However, a thorough comparison of these different approaches with the use of deep generative models for the UQ of learned reconstruction remains missing.

3. Preliminaries

In this section we describe the problem setting and the variational autoencoder (VAE) framework.

3.1. Problem Formulation

First, we specify the setting of this work, i.e., finite-dimensional linear inverse problems. Let $x\in {\mathbb{R}}^n$ be the unknown image of interest, and $y\in {\mathbb{R}}^m$ be the observational data. The linear forward operator $\mathcal{A}:{\mathbb{R}}^n\to {\mathbb{R}}^m$ maps the unknown x to the data y. In practice, the observation y is a noisy version of the exact data $\mathcal{A}x$:

$$y=\eta \left(\mathcal{A}x\right),$$

where $\eta (·)$ denotes a general corruption process by (possibly unknown type) noise, e.g., Gaussian, Poisson, and Salt and Pepper, Cauchy, uniform, or the mixtures thereof. The image reconstruction task is to recover the unknown ground-truth x from the given noisy observation y. In the Bayesian framework, the corruption process $\eta (·)$ is encoded by a conditional distribution $p^{*}\left(y\right|x)$, and the unknown image x of interest is a random variable with a prior distribution $p^{*}\left(x\right)$. In the proposed framework, we only require (i) samples from the joint distribution $p^{*}(x,y):=p^{*}\left(y\right|x)p^{*}\left(x\right)$ (which is proportional to the posterior distribution $p^{*}\left(x\right|y)$, up to a normalizing constant $p\left(y\right)={\int }_{{\mathbb{R}}^n}{p}^{*}(x,y)\mathrm{d}x$), and (ii) the action of the maps $\mathcal{A}$ and ${\mathcal{A}}^{*}$ (the adjoint of $\mathcal{A}$). These two conditions hold for many practical imaging problems.

Note that DNN-based image reconstruction techniques employ paired training data ${\left\{(x_i,y_i)\right\}}_{i=1}^N$ from the underlying joint distribution $p^{*}(x,y)$, and also the associated forward operators, which may differ for each data pair. Thus, the training data tuple takes the form of $(x_i,y_i,{\mathcal{A}}_i)$, which is ideally an exhaustive representation of the image reconstruction problem. In particular, a closed-form expression of the posterior distribution $p^{*}\left(x\right|y)$ is not needed, hence the term “implicit” posterior distribution. In practice, one can collect measurements corresponding to samples of x drawn from the prior $p^{*}\left(x\right)$ (i.e., physically derived data) without explicitly knowing the corruption process $\eta (·)$. However, in many medical imaging applications, a dataset of ordered ground-truth images and observational data pairs is often expensive to acquire at large volumes, if not impossible at all. Note that the dimension of $(x_i,y_i,{\mathcal{A}}_i)$ can also vary from one sample to another, e.g., due to the discretization of different resolution levels. The explicit presence of the operator $\mathcal{A}$ is a noticeable difference from standard supervised learning in machine learning. The main goal of learned image reconstruction with aleatoric UQ is to provide a computational framework that gives an approximation to the posterior distribution $p^{*}\left(x\right|y)$ or the joint distribution $p^{*}(x,y)$ by suitably training on the given dataset ${\left\{(x_i,y_i,{\mathcal{A}}_i)\right\}}_{i=1}^N$.

Below, we use the notation $h(·)$ to denote a DNN, and the subscript to distinguish between DNNs. Furthermore, we abuse the subscript for a distribution and a DNN to reparameterise the corresponding random variable. For instance, $p_{\theta }\left(x\right)$ is a distribution of x and $h_{\theta }(·)$ is a DNN to reparameterise x, both with the parameter vector $\theta$.

3.2. Variational Inference and Variational Autoencoders

Now we describe the basic technique, variational inference (VI), and the building blocks of the proposed framework: variational autoencoders (VAEs).

First, we describe the idea of VI for posterior approximation. Let $p_{\theta }\left(x\right|y)$ be an intractable target distribution of interest, where the vector $\theta$ in $p_{\theta }\left(x\right|y)$ contains the hyperparameters of both prior $p_{\theta }\left(x\right)$ and likelihood $p_{\theta }\left(y\right|x)$, such as the prior belief strength (also known as the regularisation parameter) and noise precision. The intractability of the distribution $p_{\theta }\left(x\right|y)$ often arises from the high-dimensionality of the parameter space (i.e., n is large) and the non-Gaussian nature of the distribution (e.g., due to the complex data formation process).

To approximate the target distribution $p_{\theta }\left(x\right|y)$, we employ variational inference (VI), which is a popular approximate inference technique in machine learning [18,35]. It selects the best approximation $q_{\varphi }\left(x\right|y)$ from a candidate family $\mathcal{Q}$ (parameterised by the vector $\varphi$, commonly known as the variational parameters) by minimising a suitable divergence measure for probability distributions. The family $\mathcal{Q}$ is often taken to be an independent Gaussian distribution, which is fully characterised by the mean and (diagonal) covariance. This is commonly known as the mean field approximation. In VI, the most popular choice for the probability metric is the Kullback–Leibler (KL) divergence [41]. The KL divergence $D_{\mathrm{KL}}(q_{\varphi }\left(x\right|y)\left|\right|p_{\theta }\left(x\right|y))$ of the approximate posterior $q_{\varphi }\left(x\right|y)$ from the target posterior $p_{\theta }\left(x\right|y)$ is defined by the following formula:

$$D_{\mathrm{KL}}(q_{\varphi }\left(x\right|y)\left|\right|p_{\theta }\left(x\right|y))=\int q_{\varphi }\left(x\right|y)log\frac{q_{\varphi }\left(x\right|y)}{p_{\theta }\left(x\right|y)}\mathrm{d}x.$$

It follows directly from Jensen’s inequality that it is always non-negative and zero if and only if $q_{\varphi }\left(x\right|y)=p_{\theta }\left(x\right|y)$ almost everywhere. Since the target $p_{\theta }\left(x\right|y)$ is often only known up to a normalising constant, the problem

$$\underset{q_{\varphi }\in \mathcal{Q}}{min}{D}_{\mathrm{KL}}(q_{\varphi }\left(x\right|y)\left|\right|p_{\theta }\left(x\right|y))$$

is often cast into an equivalent evidence lower bound (ELBO) maximisation

$$\underset{\varphi }{max}\{\mathcal{L}(\varphi ;\theta ,y)={\mathbb{E}}_{q_{\varphi }\left(x\right|y)}[log{p}_{\theta }\left(y\right|x)]-D_{\mathrm{KL}}(q_{\varphi }\left(x\right|y)\left|\right|p_{\theta }\left(x\right))\},$$

(1)

where the notation ${\mathbb{E}}_q[·]$ denotes taking expectation with respect to the distribution q. In the ELBO, the first term enforces data consistency, whereas the second term represents a penalty, which is relative to the prior distribution $p_{\theta }\left(x\right)$.

Solving the optimization problem (1) requires evaluating the gradient of the loss $\mathcal{L}$ with respect to the variational parameters $\varphi$ (i.e., ${\nabla }_{\varphi }{\mathbb{E}}_{q_{\varphi }\left(x\right)}\left[f_{\theta }\left(x\right)\right]$) for a deterministic function $f_{\theta }$, parameterised by $\theta$. The challenge lies in the fact that the integral ${\mathbb{E}}_{q_{\varphi }\left(x\right)}\left[f_{\theta }\left(x\right)\right]$ is often not analytically tractable and can only be evaluated by means of the Monte Carlo methods. The reparameterisation trick [12,42] is useful for overcoming the challenges in directly back-propagating the gradients (with respect to $\varphi$) through the random variable x. It provides an unbiased estimate of the gradient of the ELBO with respect to the variational parameters $\varphi$. In fact, we assume that the conditional sampling of x depends on y and an easy-to-sample auxiliary variable $ϵ$ distributed according to $p\left(ϵ\right)$ (e.g., a Gaussian distribution):

$$x=g_{\varphi }(y,ϵ),\mathrm{with}ϵ\sim p\left(ϵ\right),$$

where $g_{\varphi }(·)$ is a deterministic function (in our case, a DNN). By estimating $\theta$ from the given data y while simultaneously optimising the variational parameter $\varphi$, one obtains the following Monte Carlo estimator of the loss in (1):

$$\mathcal{L}(\theta ,\varphi )=\frac{1}{L}\sum _{\ell =1}^L[log{p}_{\theta }\left(y_i\right|x_{(i,\ell )})]-D_{\mathrm{KL}}(q_{\varphi }\left(x\right|y_i)\left|\right|p_{\theta }\left(x\right)),$$

where ${\left\{x_{(i,\ell )}\right\}}_{\ell =1}^L$ are L samples generated with $y_i$, and using the variational encoder $q_{\varphi }\left(x\right|y)$: ${\left\{{ϵ}_{\ell }\right\}}_{\ell =1}^L$ are sampled from $p\left(ϵ\right)$ and $x_{(i,\ell )}=g_{\varphi }(y_i,z_{\ell })$. The KL term can often be evaluated in closed form since both factors therein are often taken to be Gaussian, otherwise it can also be approximated using the Monte Carlo. Note that in the preceding construction, we allowed the observation data y to vary with the unknown image x in order to accommodate the training data ${\left\{(x_i,y_i)\right\}}_{i=1}^N$. This differs from the standard approximate inference schemes, and it is often referred to as amortized variational inference.

In generative modelling, the considerations above give rise to a formalism very similar to the popular variational autoencoders (VAEs) [12,43], with x assuming the role of a latent variable, and y being the data to be modelled probabilistically; see the remark below for more details on the difference.

VAE is an unsupervised deep generative framework that learns stochastic mapping between the observed y-space and a latent x-space. The model learns a joint distribution $p_{\theta }(y,x)$ that factorises as $p_{\theta }(y,x)=p_{\theta }\left(x\right)p_{\theta }\left(y\right|x)$ with a stochastic decoder $p_{\theta }\left(y\right|x)$ and a prior distribution over a low-dimensional latent space $p_{\theta }\left(x\right)$. A stochastic encoder $q_{\varphi }\left(x\right|y)$ (also known as a parametric inference model) approximates the intractable posterior $p_{\theta }\left(x\right|y)$. The framework compresses the observed data into a constrained latent distribution (via the encoder) and reconstructs it faithfully (via the decoder). This process is carried out by two neural networks, an encoding network $h_{\varphi }(·)$ with the parameter $\varphi$ (i.e., the weights and the biases of the network), also called the variational parameters, and a decoding network $h_{\theta }(·)$ with the parameter $\theta$.

In practice, VAEs do not often use the encoding network to model the parameterisation function $g_{\varphi }(·)$ in an end-to-end way. Instead, VAEs take the observation y and outputs the coefficients to reparameterise the image x. For example, for a multivariate Gaussian $q_{\varphi }\left(x\right|y)$, the decoding network can output the mean $\mu$ and the Cholesky factor L of the covariance $\Sigma =L{L}^{\top }$:

$$\begin{array}{c}\hfill \left(\mu ,L\right)=h_{\varphi }\left(y\right),q_{\varphi }\left(x\right|y)=\mathcal{N}\left(x\right|\mu ,\Sigma ).\end{array}$$

We can generate samples from $q_{\varphi }\left(x\right|y)$ by sampling $ϵ$ from the standard Gaussian $p\left(ϵ\right)=\mathcal{N}\left(ϵ\right|0,I)$, followed by an affine transformation $x=\mu +Lϵ$.

VAE allows for the simultaneous performance of VI (with respect to $\varphi$) and the model selection (with respect to $\theta$), and the resulting VAE objective is given by the following equation:

$$\underset{\varphi ,\theta }{max}\left\{{\mathcal{L}}_{\mathrm{VAE}}(\theta ,\varphi ;y)={\mathbb{E}}_{q_{\varphi }\left(x\right|y)}[log{p}_{\theta }\left(y\right|x)]-D_{\mathrm{KL}}(q_{\varphi }\left(x\right|y)\left|\right|p_{\theta }\left(x\right))\right\}.$$

In practice, one may employ an identity variance Gaussian, with the mean being the decoder’s output as the likelihood $p_{\theta }\left(y\right|x)$, and a standard Gaussian distribution as the prior distribution $p_{\theta }\left(x\right)$.

Remark 1.

Note that the original formalism of VAE [12] is slightly different from the above. We briefly recall its derivation for the convenience of the readers. Given a dataset ${\left\{y_i\right\}}_{i=1}^N$ from an unknown probability function $p\left(z\right)$ and a multivariate latent encoding vector z, the objective of VAE is to model the data y as a distribution $p_{\theta }\left(y\right)$, i.e.,

$$p_{\theta }\left(y\right)=\int p_{\theta }\left(y\right|z)p_{\theta }\left(z\right)\mathrm{d}z,$$

where θ represents the network parameters. Thus, in VAE, if we assume that $p_{\theta }\left(y|z\right)$ is a Gaussian distribution, then $p_{\theta }\left(y\right)$ is a mixture of Gaussians. To solve for θ in a learning paradigm, one constructs an approximation $q_{\varphi }\left(z|y\right)\approx p_{\theta }\left(z|y\right)$ by means of VI with variational parameters ϕ. Repeating the preceding discussion on VI directly yields the following standard VAE loss:

$$\mathcal{L}(\theta ,\varphi )={\mathbb{E}}_{q_{\varphi }\left(z\right|y)}log{p}_{\theta }\left(y\right|z)-D_{\mathrm{KL}}(q_{\varphi }\left(z\right|y)\left|\right|p_{\theta }\left(z\right)).$$

The problem is then reduced to an autoencoder formalism: the conditional likelihood distribution $p_{\theta }\left(y\right|z)$ is the probabilistic decoder, and the approximation $q_{\varphi }\left(z\right|y)$ serves as the probabilistic encoder. Hence, the goal of VAE is to represent the given unlabelled data ${\left\{y_i\right\}}_{i=1}^N$ and to generate new data (from the latent variable z), which differs markedly from the task in learned image reconstruction, for which the image (represented as latent variable) is of the main objects of interest (and often of very high dimensionality, larger than that of y). Additionally, the problem of modelling $p_{\theta }\left(y\right|z)$ and $p_{\theta }\left(z\right)$ in the VAE framework is usually unidentifiable, in the sense that there may exist many different $(p_{\theta }\left(y\right|z),p_{\theta }\left(z\right))$ pairs that admit the same marginal distribution $p_{\theta }\left(y\right)$ [44]. Thus, the associated modelled approximate posterior $q_{\varphi }\left(z\right|y)$ is not unique.

4. Proposed Framework

We develop a computational UQ framework that learns a map from the observation y to a distribution, denoted by $p_{\varphi }\left(x\right|y)$, which approximates the true posterior distribution $p^{*}\left(x\right|y)$. The map is modelled with a recurrent network, and a probabilistic encoder therein allows for diverse reconstruction samples, and hence facilitating the quantification of the associated aleatoric uncertainty.

4.1. Conditional VAE as Approximate Inference

Note that a direct application of VAEs to image reconstruction is problematic: VAEs are unsupervised generative machineries and would only use noisy observations y, but not the ground-truth image x during the training process. To circumvent the issue, we employ the cVAE loss [16,45], a conditional variant of VAE:

$$\underset{\varphi ,\theta }{max}\left\{{\mathcal{L}}_{\mathrm{cVAE}}(\theta ,\varphi ;x,y)={\mathbb{E}}_{q_{\theta }\left(z\right|x,y)}[log{p}_{{\varphi }_2}\left(x\right|y,z)]-D_{\mathrm{KL}}(q_{\theta }\left(z\right|x,y)\left|\right|p_{{\varphi }_1}\left(z\right|y))\right\}.$$

(2)

In cVAEs, there are three distributions: a teacher encoder $q_{\theta }\left(z\right|x,y)$, a student encoder $p_{{\varphi }_1}\left(z\right|y)$, and a conditional decoder $p_{{\varphi }_2}\left(x\right|y,z)$. The vector $\varphi =({\varphi }_1,{\varphi }_2)$ assembles the parameters of $p_{{\varphi }_1}\left(z\right|y)$ and $p_{{\varphi }_2}\left(x\right|y,z)$. The approximate posterior $p_{\varphi }\left(x\right|y)$ obtained by cVAE is given by the following:

$$p_{\varphi }\left(x\right|y)=\int p_{{\varphi }_2}\left(x\right|y,z)p_{{\varphi }_1}\left(z\right|y)\mathrm{d}z.$$

The cVAE loss admits the following approximate inference interpretation in a generalised sense. Note that $J^{*}$ is a functional of the variational distribution $p_{\varphi }\left(x\right|y)$ and other auxiliary distributions involving z.

Proposition 1.

Minimising the loss ${\mathcal{L}}_{\mathrm{cVAE}}(\theta ,\varphi ;x,y)$ in (2) that is expected on the training data distribution $p^{*}(x,y)$ is equivalent to optimising an upper bound of the expected $\mathrm{KL}$ divergence:

$$J^{*}\left(p_{\varphi }\left(x\right|y)\right)={\mathbb{E}}_{p^{*}\left(y\right)}[D_{\mathrm{KL}}(p^{*}\left(x\right|y)\left|\right|p_{\varphi }\left(x\right|y)\left)\right].$$

Proof. 

By the definition of the KL divergence and the Fubini theorem,

$$\begin{array}{cc}\hfill J^{*}\left(p_{\varphi }\left(x\right|y)\right)& ={\mathbb{E}}_{p^{*}\left(y\right)}[D_{\mathrm{KL}}(p^{*}\left(x\right|y)\left|\right|p_{\varphi }\left(x\right|y)\left)\right]\hfill \\ \hfill & =\int p^{*}\left(y\right)\int p^{*}\left(x\right|y)log\frac{p^{*}\left(x\right|y)}{p_{\varphi }\left(x\right|y)}\mathrm{d}x\mathrm{d}y\hfill \\ \hfill & =\int p^{*}(x,y)[log{p}^{*}\left(x\right|y)-log{p}_{\varphi }\left(x\right|y)]\mathrm{d}(x,y)\hfill \\ \hfill & ={\mathbb{E}}_{p^{*}(x,y)}[log{p}^{*}\left(x\right|y)]+{\mathbb{E}}_{p^{*}(x,y)}[-log{p}_{\varphi }\left(x\right|y)].\hfill \end{array}$$

(3)

Next, we derive a lower bound for the $log{p}_{\varphi }\left(x\right|y)$ of the conditional distribution $p_{\varphi }\left(x\right|y)$ using the standard procedure. Since $p_{\varphi }(x,z|y)=p_{\varphi }\left(z\right|x,y)p_{\varphi }\left(x\right|y)$, we have the following:

$$p_{\varphi }\left(x\right|y)=\frac{p_{\varphi }(x,z|y)}{p_{\varphi }\left(z\right|x,y)},$$

and consequently, for any distribution $q\left(z\right|x,y)$, there holds the subsequent equation:

$$\begin{array}{cc}\hfill log{p}_{\varphi }\left(x\right|y)& =\int q\left(z\right|x,y)log{p}_{\varphi }\left(x\right|y)\mathrm{d}z=\int q\left(z\right|x,y)log\frac{p_{\varphi }\left(x\right|y)p_{\varphi }\left(z\right|x,y)}{p_{\varphi }\left(z\right|x,y)}\mathrm{d}z\hfill \\ \hfill & =\int q\left(z\right|x,y)log\frac{p_{\varphi }(x,z|y)}{p_{\varphi }\left(z\right|x,y)}\mathrm{d}z=\int q\left(z\right|x,y)log\frac{p_{\varphi }(x,z|y)}{q\left(z\right|x,y)}\frac{q\left(z\right|x,y)}{p_{\varphi }\left(z\right|x,y)}\mathrm{d}z\hfill \\ \hfill & =\int q\left(z\right|x,y)log\frac{p_{\varphi }(x,z|y)}{q\left(z\right|x,y)}\mathrm{d}z+\int q\left(z\right|x,y)log\frac{q\left(z\right|x,y)}{p_{\varphi }\left(z\right|x,y)}\mathrm{d}z.\hfill \end{array}$$

By the non-negativity of the KL divergence, the second term is non-positive, and then using the splitting $p_{\varphi }(x,z|y)=p_{{\varphi }_2}\left(x\right|y,z)p_{{\varphi }_1}\left(z\right|y)$, we can deduce the following:

$$\begin{array}{cc}\hfill log{p}_{\varphi }\left(x\right|y)& \ge \int q\left(z\right|x,y)log\frac{p_{\varphi }(x,z|y)}{q\left(z\right|x,y)}\mathrm{d}z=\int q\left(z\right|x,y)log\frac{p_{{\varphi }_1}\left(z\right|y)p_{{\varphi }_2}\left(x\right|y,z)}{q\left(z\right|x,y)}\mathrm{d}z\hfill \\ \hfill & =\int q\left(z\right|x,y)log\frac{p_{{\varphi }_1}\left(z\right|y)}{q\left(z\right|x,y)}\mathrm{d}z+\int q\left(z\right|x,y)log{p}_{{\varphi }_2}\left(x\right|y,z)\mathrm{d}z\hfill \\ \hfill & =-\mathrm{KL}\left(q\left(z\right|x,y)\right||p_{{\varphi }_1}\left(z\right|y))+{\mathbb{E}}_{z\sim q\left(z\right|x,y)}[log{p}_{{\varphi }_2}\left(x\right|y,z)],\hfill \end{array}$$

i.e.,

$$-log{p}_{\varphi }\left(x\right|y)\le D_{\mathrm{KL}}\left(q\left(z\right|x,y)\right||p_{{\varphi }_1}\left(z\right|y))+{\mathbb{E}}_{z\sim q\left(z\right|x,y)}[-log{p}_{{\varphi }_2}\left(x\right|y,z)].$$

(4)

Taking expectation of (4) with respect to $p^{*}(x,y)$, and then substituting it into identity (3) yields the following equation:

$$\begin{array}{ccc}\hfill J^{*}\left(p\left(x\right|y)\right)& \le & {\mathbb{E}}_{p^{*}(x,y)}[log{p}^{*}\left(x\right|y)]+{\mathbb{E}}_{p^{*}(x,y)}[D_{\mathrm{KL}}\left(q\left(z\right|x,y)\right||p_{{\varphi }_1}\left(z\right|y)\left)\right]\hfill \\ & & +{\mathbb{E}}_{p^{*}(x,y)}\left[{\mathbb{E}}_{z\sim q\left(z\right|x,y)}[-log{p}_{{\varphi }_2}\left(x\right|y,z)]\right].\hfill \end{array}$$

Since the term ${\mathbb{E}}_{p^{*}(x,y)}[log{p}^{*}\left(x\right|y)]$ is independent of the variational distribution $p_{\varphi }\left(x\right|y)$ and other auxiliary distributions involving z, minimising the cVAE loss in (2) that is expected on the training data distribution $p^{*}(x,y)$ is equivalent to minimising an upper bound of $J^{*}\left(p_{\varphi }\left(x\right|y)\right)$. This shows the desired assertion.    □

In view of Proposition 1, cVAEs can indeed learn an optimal map from the observation y to an approximate posterior distribution $p_{\varphi }\left(x\right|y)$ in the sense of minimising the expected loss of the KL divergence. This interpretation underpins the validity of the procedure for quantifying aleatoric uncertainty. The minimiser is indeed an approximation to the target posterior distribution $p^{*}\left(x\right|y)$, constructed by a generalised variational inference principle.

Example 1.

We briefly validate Proposition 1, which states that the cVAE framework can approximate the ground-truth distribution $p^{*}\left(x\right|y)$ in the sense of generalised variational inference. Note that it is notoriously challenging to numerically verify the accuracy of any approximate inference scheme for high-dimensional distributions. Nonetheless, in the low-dimensional case, the target distribution can be efficiently sampled through a Markov chain Monte Carlo [46]. To illustrate this, we take the ground-truth distribution $p^{*}\left(x\right|y)$ to be a two-dimensional multi-modal distribution, which consists of the mixture of seven Gaussians, as shown in Figure 1a, and approximate it by cVAE (with teacher encoder, student encoder, and decoder modelled by different three-layer neural networks, and with ReLu as the nonlinear activation function). We clearly observe from Figure 1b that the approximation $p_{\varphi }\left(x\right|y)$ is fairly accurate and can capture the multi-modality of the distribution well, thereby verifying Proposition 1.

We model the conditional encoder $p_{{\varphi }_2}\left(x\right|y,z)$ by a mean-field Gaussian, with a covariance $\beta I$, where $\beta >0$ is a hyperparameter. The DNN $h_{{\varphi }_2}$ with parameter ${\varphi }_2$ only outputs the mean of $p_{{\varphi }_2}\left(x\right|y,z)$, and on a mini-batch ${\left\{(x_i,y_i)\right\}}_{i=1}^M$, the objective function ${\widehat{\mathcal{L}}}_{\mathrm{cVAE}}(\varphi ,\theta )$ is given by:

$${\widehat{\mathcal{L}}}_{\mathrm{cVAE}}(\varphi ,\theta )=-\frac{1}{2M}\sum _{i=1}^M\frac{1}{L}\sum _{\ell =1}^L\parallel x_i-{\widehat{x}}_{(i,\ell )}{\parallel }^2-\frac{\beta }{M}\sum _{i=1}^M{D}_{\mathrm{KL}}(q_{\theta }\left(z\right|x_i,y_i)\left|\right|p_{{\varphi }_1}\left(z\right|y_i)),$$

where ${\widehat{x}}_{(i,\ell )}$ is the mean of the distribution $p_{{\varphi }_2}\left(x\right|y_i,z_{i,\ell })$, and ${\left\{z_{i,\ell }\right\}}_{\ell =1}^L$ are L samples drawn from the distribution $q_{\theta }\left(z\right|x_i,y_i)$, with the default choice $L=1$. Note that for special choices of $q_{\theta }\left(z\right|x,y)$ and $p_{{\varphi }_1}\left(z\right|y)$, the KL divergence term may be evaluated analytically and can be used, if available. The gradient of the loss ${\widehat{\mathcal{L}}}_{\mathrm{cVAE}}^L(\varphi ,\theta )$ is then evaluated by the reparameterisation trick.

Remark 2.

In VAEs, the approximate posterior of the image x (i.e., the latent variable) is modelled by $q_{\varphi }\left(x\right|y)$, whereas in cVAEs, it is modelled by $p_{\varphi }\left(x\right|y)=\int p_{{\varphi }_2}\left(x\right|y,z)p_{{\varphi }_1}\left(z\right|y)\mathrm{d}z.$ In both VAEs and cVAEs, the stochasticity of x comes from z: in VAEs, z is not dependent on the observation y, whereas in cVAEs, z is dependent on y. Since the distribution of z in cVAEs is learned, it can be more flexible than that in VAEs. Thus, even if $p_{{\varphi }_2}\left(x\right|y,z)$ is chosen to be a simple distribution (e.g., Gaussian distributions), $p_{{\varphi }_1}\left(x\right|y)$ can still model a broad family of distributions for a continuous unobservable x due to the presence of $p_{{\varphi }_1}\left(z\right|y)$, in a manner similar to the scale mixture of Gaussians.

4.2. cVAE for Learned Reconstruction

Probabilistic modelling consists of a learning principle given by a loss function with a proper probabilistic interpretation, and a graphical model describing probabilistic dependency between variables. In the proposed framework, we employ a cVAE-type loss:

$$\underset{\varphi ,\theta }{max}\{\mathcal{L}(\theta ,\varphi ;x,y,\mathcal{A})={\mathbb{E}}_{q_{\theta }\left(z\right|x,y,\mathcal{A})}[log{p}_{\varphi }\left(x\right|y,z,\mathcal{A})]-D_{\mathrm{KL}}(q_{\theta }\left(z\right|x,y,\mathcal{A})\left|\right|p_{\varphi }\left(z\right|y)\left)\right\}.$$

(5)

Its difference from the standard cVAE loss in (2) is that (5) also includes the forward map $\mathcal{A}$ (and its adjoint ${\mathcal{A}}^{*}$) as part of the training data. Here, $\mathcal{A}$ may have different realisations (e.g., corresponding to different levels of discretization) with varying dimensions. Nonetheless, it is viewed as a deterministic variable. The modelled approximate posterior distribution $p_{\varphi }\left(x\right|y)$ is then given by the following formula:

$$p_{\varphi }\left(x\right|y)=\int p_{\varphi }\left(x\right|y,z,\mathcal{A})p_{\varphi }\left(z\right|y)\mathrm{d}z.$$

The graphical model is given in Figure 2a. The learning algorithm learns a conditional sampler, in a manner similar to a random number generator (RNG) for a given probability distribution. Note that for an RNG, different runs lead to different samples; but with a fixed random seed, the path is the same for different runs. The auxiliary (low-dimensional) latent variable z (conditionally dependent on the observation y) is an analogue of the random seed in the RNG, and is introduced into the deterministic recurrent process (modelled by a network) to diversify the reconstruction samples. In particular, for a fixed z, the recursion process inputs the sample initialization $x^0$ and applies a recurrent refining step based on suitable sample quality measures and the information encoded in the variable z.

Three DNNs are employed to model the distributions in $\mathcal{L}(\theta ,\varphi ;x,y,\mathcal{A})$ and to carry out the conditional sampling process, including one for the recursive process (i.e., recurrent unit modelled by a neural network) and two for probabilistic encoders (i.e., teacher and student encoders). The observation y and forward map $\mathcal{A}$ constitute their inputs: y is fed into the two probabilistic encoders and each recurrence of the recurrent unit; $\mathcal{A}$ is fed into each recurrence of the recurrent unit and the teacher encoder. Subsequently, we explain how the three networks work separately.

The recurrent component is the (deep) network $h_{{\varphi }_2}(·)$; see Figure 2b. The network begins with an initial guess $x^0$ (default: back-projected data ${\mathcal{A}}^{*}y$) and outputs $x^K$ after K iterations as the mean of $p_{\varphi }\left(x\right|y,z,\mathcal{A})$, following the established idea of algorithmic unrolling [17,34], which mimics the iterations from standard optimization algorithms (e.g., (proximal) gradient descent, alternating direction method of multipliers (ADMM), and primal-dual hybrid gradient). At the k-th recursion, the network $h_{{\varphi }_2}$ takes one sample $x^{k-1}$ to refine and outputs an improved sample $x^k$. To incorporate the forward map $\mathcal{A}$ and the observation y, we employ a functional $E(\mathcal{A},y,x^{k-1})$. In the lens of variational regularisation [47], $E(\mathcal{A},y,x^{k-1})$ measures how well $x^{k-1}$ can explain the data y. To indicate how well $x^{k-1}$ fulfils the prior knowledge (e.g., sparsity and smoothness), we use the penalty $R\left(x^{k-1}\right)$ as part of the input. For the sample quality measure $E(\mathcal{A},y,x^{k-1})$ and the penalty $R\left(x^{k-1}\right)$, we use $\left|\right|y-\mathcal{A}\left(x^{k-1}\right){\left|\right|}^2$ (or its gradient), and $\left|\right|x^{k-1}{\left|\right|}_2^2$ or $|x^{k-1}{|}_{\mathrm{TV}}$ (total variation semi-norm), respectively. Besides the latest iterate $x^{k-1}$ and the quality measures E and R, the network $h_{{\varphi }_2}(·)$ also takes a memory variable $a^{k-1}$ and an auxiliary variable z. The memory variable $a^{k-1}$ plays the role of momentum in gradient-type methods and is to retain long-term information between recursions. The overall procedure resembles an unrolled momentum accelerated gradient descent. The auxiliary random variable z is low-dimensional and injects randomness into the iteration procedure. Since both $x^{k-1}$ and $x^k$ belong to the image space ${\mathbb{R}}^n$, we adopt CNNs without pooling layers to model the recurrent unit $h_{{\varphi }_2}$. Different inputs of $h_{{\varphi }_2}(·)$ are concatenated along the channel axis, and the outputs of $h_{{\varphi }_2}(·)$—that is, the updated $\delta x^k$ with $x^k=x^{k-1}+\delta x^k$ and the updated memory $a^k$—are also concatenated.

Remark 3.

At each step, the recurrent unit (neural network $h_{{\varphi }_2}$) reuses the observation data y and the map $\mathcal{A}$ for refinement, and the overall procedure differs from the deterministic mapping that serves as a post-processing step of back-projection. The latter is an end-to-end mapping that takes the back-projected data and outputs a refinement, but the proposed approach employs the current sample and quality measures and then decides the refinement strategy.

The framework employs two encoders of z: a teacher encoder $q_{\theta }\left(z\right|x,y,\mathcal{A})$ and a student encoder $p_{{\varphi }_1}\left(z\right|y)$; see Figure 3. The student encoder $p_{{\varphi }_1}\left(z\right|y)$ takes the observation y, and encodes the observation-based knowledge so as to inform the recurrent unit $h_{{\varphi }_2}$. Given one sample z from $p_{{\varphi }_1}\left(z\right|y)$, the network $h_{{\varphi }_2}$ gives one refining increment, and the distribution of z contributes to the diversity of the unknown image x. To help train the student encoder $p_{{\varphi }_1}\left(z\right|y)$, we provide the teacher encoder $q_{\theta }\left(z\right|x,y,\mathcal{A})$ with the ground-truth x and the forward map $\mathcal{A}$. The teacher encoder $q_{\theta }\left(z\right|x,y,\mathcal{A})$ is discarded once the training is finished. The encoders $p_{{\varphi }_1}\left(z\right|y)$ and $q_{\theta }\left(z\right|x,y,\mathcal{A})$ are modelled by two DNNs $h_{{\varphi }_1}(·)$ and $h_{\theta }(·)$, respectively, which reparameterise the auxiliary variable z. Since z is low-dimensional, CNNs with reduced mean layers and $1\times 1$ convolutional layers can guarantee the flexibility dimension of the input y. To input the ground-truth data x and y into $h_{\theta }(·)$, we use the forward map $\mathcal{A}$ and concatenate $\mathcal{A}\left(x\right)$ with y along the channel axis. This construction is very flexible with the dimension problem and allows for the training of $(x_i,y_i)$ of varying sizes (and the corresponding ${\mathcal{A}}_i$).

Now we can state the complete algorithms for the training and inference of the cVAE framework for uncertainty-aware learned image reconstruction in Algorithms 1 and 2, respectively. Here, M denotes the mini-batch size, T the maximum number of training batches, K the number of recurrences of $h_{{\varphi }_2}$ for one sample, and $(\widehat{\varphi },\widehat{\theta })$ the output of the training procedure (i.e., the learned parameters). There are many possible choices for the stochastic optimizer in line 11 of Algorithm 1 (e.g., ADAM [48], and SGD). We shall employ ADAM [48] in our experiment. The final sample from the recurrent process is regarded as the mean of the conditional distribution $p_{\varphi }\left(x\right|y,z,\mathcal{A})$. Thus, given the initial $x^0$, the iteration with different realisations of z leads to diverse samples of the unknown image x. Since each sample is the mean of $p_{\varphi }\left(x\right|y,z,\mathcal{A})$ rather than a direct sample from the approximate posterior $p_{\varphi }\left(x\right|y)$, the summarizing statistics also have to be transformed; see (6). Note that the posterior variance contains two components: one is due to the background (i.e., $\beta I$), and the other is due to the sample variance (i.e., $\frac{1}{S}{\sum }_{i=1}^S{x}_i{x}_i^t-{\widehat{\mathbb{E}}}_{p_{\varphi }\left(x\right|y)}\left[x\right]{\widehat{\mathbb{E}}}_{p_{\varphi }\left(x\right|y)}{\left[x\right]}^t$), as shown in the following result.

Algorithm 1 Training procedure

1: Input: Training data ${\left\{(x_i,y_i,{\mathcal{A}}_i)\right\}}_{i=1}^N$, $\beta$, T, K, M;   2: for$t=1,2,\dots ,T$do   3:  Randomly select a mini-batch training data ${\left\{(x_i,y_i,{\mathcal{A}}_i)\right\}}_{i=1}^M$;   4:  Sample ${\left\{z_i\right\}}_{i=1}^M$ from ${\left\{q_{\theta }\left(z\right|x_i,y_i)\right\}}_{i=1}^M$;   5:  Initialize ${\left\{{\widehat{x}}_i\right\}}_{i=1}^M$ with ${\left\{{\mathcal{A}}_i^{*}\left(y_i\right)\right\}}_{i=1}^M$, and ${\left\{a_i\right\}}_{i=1}^M$ with zeros;   6:  for $k=1,2,\dots ,K$ do   7:   Update ${\left\{{\widehat{x}}_i\right\}}_{i=1}^M$ and ${\left\{a_i\right\}}_{i=1}^M$ with ${\left\{h_{{\varphi }_2}({\widehat{x}}_i,E({\mathcal{A}}_i,y_i,{\widehat{x}}_i),R\left({\widehat{x}}_i\right),a_i,z_i)\right\}}_{i=1}^M$;   8:  end for   9:  Evaluate the KL divergence $\{D_{\mathrm{KL}}(q_{\theta }\left(z\right|x_i,y_i)\left|\right|p_{{\varphi }_1}\left(z\right|y_i){\left)\right\}}_{i=1}^M$; 10:  Compute the objective function $\widehat{\mathcal{L}}(\varphi ,\theta )$; 11:  Update the parameters $(\varphi ,\theta )$; 12: end for 13: Output: $(\widehat{\varphi },\widehat{\theta })$

Algorithm 2 Inference procedure

1: Input: Test data $(\mathcal{A},y)$, S, K, $\widehat{\varphi }=(\widehat{{\varphi }_1},{\widehat{\varphi }}_2)$;   2: for$s=1,2,\dots ,S$do   3:  Sample $z_s$ from $p_{{\varphi }_1}\left(z\right|y)$;   4:  Initialise ${\widehat{x}}_s$ with ${\mathcal{A}}^{*}\left(y\right)$ and a with zeros;   5:  for $k=1,2,\dots ,K$ do   6:   Update ${\widehat{x}}_s$ and a with $h_{{\varphi }_2}({\widehat{x}}_s,E(\mathcal{A},y,{\widehat{x}}_s),R\left({\widehat{x}}_s\right),a,z_s)$;   7:  end for   8: end for   9: Output: ${\left\{{\widehat{x}}_s\right\}}_{s=1}^S$; 10: Evaluate $({\widehat{\mathbb{E}}}_{p\left(x\right|y)}\left[x\right],{\widehat{\mathrm{Cov}}}_{p\left(x\right|y)}\left[x\right])$ by Equation (6)

Proposition 2.

Let $p_{\varphi }\left(x\right|y,z,\mathcal{A})=\mathcal{N}\left(x\right|x^K\left(z\right),\beta I)$ and the approximate posterior $p_{\varphi }\left(x\right|y)=\int p_{{\varphi }_2}\left(x\right|y,z,\mathcal{A})p_{{\varphi }_1}\left(z\right|y)\mathrm{d}z$ be a mixture of Gaussian distributions, with z being the mixture variable. Then, given samples ${\left\{z_i\right\}}_{i=1}^S$ of z from $p_{\varphi }\left(z\right|y)$, and the corresponding $x^K\left(z\right)$ denoted by ${\left\{x_i\right\}}_{i=1}^S$, the mean ${\mathbb{E}}_{p_{\varphi }\left(x\right|y)}\left[x\right]$ and the covariance ${\mathrm{Cov}}_{p_{\varphi }\left(x\right|y)}\left[x\right]$ of $p_{\varphi }\left(x\right|y)$ can be estimated by:

$$\begin{array}{c}{\widehat{\mathbb{E}}}_{p_{\varphi }\left(x\right|y)}\left[x\right]=\frac{1}{S}\sum _{i=1}^S{x}_i,\\ {\widehat{\mathrm{Cov}}}_{p_{\varphi }\left(x\right|y)}\left[x\right]=\beta I+\frac{1}{S}\sum _{i=1}^S{x}_i{x}_i^t-{\widehat{\mathbb{E}}}_{p_{\varphi }\left(x\right|y)}\left[x\right]{\widehat{\mathbb{E}}}_{p_{\varphi }\left(x\right|y)}{\left[x\right]}^t.\end{array}$$

(6)

Proof. 

For the mean ${\mathbb{E}}_{p_{\varphi }\left(x\right|y)}\left[x\right]$, by dentition, there holds the following:

$$\begin{array}{cc}\hfill {\mathbb{E}}_{p_{\varphi }\left(x\right|y)}\left[x\right]& ={\int }_{x}x{p}_{\varphi }\left(x\right|y)\mathrm{d}x={\int }_{x}x{\int }_z{p}_{{\varphi }_2}\left(x\right|y,z,\mathcal{A})p_{{\varphi }_1}\left(z\right|y)\mathrm{d}z\mathrm{d}x\hfill \\ \hfill & ={\int }_z{\int }_{x}x{p}_{{\varphi }_2}\left(x\right|y,z,\mathcal{A})\mathrm{d}x{p}_{{\varphi }_1}\left(z\right|y)\mathrm{d}z={\int }_z{x}^K\left(z\right)p_{{\varphi }_1}\left(z\right|y)\mathrm{d}z.\hfill \end{array}$$

Thus, $\frac{1}{S}{\sum }_{i=1}^S{x}_i$ is an unbiased estimator of ${\mathbb{E}}_{p_{\varphi }\left(x\right|y)}\left[x\right]$. Similarly, for the covariance ${\mathrm{Cov}}_{p_{\varphi }\left(x\right|y)}\left[x\right]$, by the standard bias variance decomposition,

$${\mathrm{Cov}}_{p_{\varphi }\left(x\right|y)}\left[x\right]=\int x{x}^T{p}_{\varphi }\left(x\right|y)\mathrm{d}x-{\mathbb{E}}_{p_{\varphi }\left(x\right|y)}\left[x\right]{\mathbb{E}}_{p_{\varphi }\left(x\right|y)}{\left[x\right]}^T.$$

Now, the first term on the right hand side, there holds the following:

$$\begin{array}{cc}\hfill \int x{x}^T{p}_{\varphi }\left(x\right|y)\mathrm{d}x& ={\int }_{x}x{x}^T{\int }_z{p}_{{\varphi }_2}\left(x\right|y,z,\mathcal{A})p_{{\varphi }_1}\left(z\right|y)\mathrm{d}z\mathrm{d}x\hfill \\ \hfill & ={\int }_z{\int }_{x}x{x}^T{p}_{{\varphi }_2}\left(x\right|y,z,\mathcal{A})\mathrm{d}x{p}_{{\varphi }_1}\left(z\right|y)\mathrm{d}z\hfill \\ \hfill & ={\int }_z({\mathrm{Cov}}_{p_{{\varphi }_2}\left(x\right|y,z,\mathcal{A})}\left[x\right]+{\mathbb{E}}_{p_{{\varphi }_2}\left(x\right|y,z,\mathcal{A})}\left[x\right]{\mathbb{E}}_{p_{{\varphi }_2}\left(x\right|y,z,\mathcal{A})}{\left[x\right]}^T)p_{{\varphi }_1}\left(z\right|y)\mathrm{d}z\hfill \\ \hfill & =\beta I+{\int }_z{\mathbb{E}}_{p_{{\varphi }_2}\left(x\right|y,z,\mathcal{A})}\left[x\right]{\mathbb{E}}_{p_{{\varphi }_2}\left(x\right|y,z,\mathcal{A})}{\left[x\right]}^T{p}_{{\varphi }_1}\left(z\right|y)\mathrm{d}z.\hfill \end{array}$$

Consequently, $\beta I+\frac{1}{S}{\sum }_{i=1}^S{x}_i{x}_i^t-{\widehat{\mathbb{E}}}_{p_{\varphi }\left(x\right|y)}\left[x\right]{\widehat{\mathbb{E}}}_{p_{\varphi }\left(x\right|y)}{\left[x\right]}^t$ is an unbiased estimator of the covariance ${\mathrm{Cov}}_{p_{\varphi }\left(x\right|y)}\left[x\right]$.    □

5. Numerical Experiments and Discussions

Now we showcase the proposed cVAE framework for learned reconstruction with quantified aleatoric uncertainty with numerical experiments on positron emission tomography (PET). PET is a pillar of modern diagnostic imaging, allowing for non-invasive, sensitive, and specific detection of functional changes in a number of diseases. Most conventional PET reconstruction algorithms rely on penalized maximum likelihood estimates, using a hand crafted prior (e.g., total variation and anatomical); see the work of [28] for an overview of classical reconstruction techniques. More recently, learning-based approaches have been proposed. While these techniques have been successful, they still lack the capability to provide uncertainty estimates; see the work of [29,30,31,37] for several recent studies on UQ in PET reconstruction, although none of them is based on deep learning.

For the experiments below, we employ a three-layer CNN as the recurrent unit $h_{{\varphi }_2}$ and fix $K=10$ iterations for each sampling step, cf. Figure 4 for a schematic illustration; VGG style encoders are then used for both $h_{{\varphi }_1}$ and $h_{\theta }$, cf. Figure 5. We train the network on a synthetic dataset consisting of elliptical phantoms and test it on the public medical imaging dataset BrainWeb [49] (available from https://brainweb.bic.mni.mcgill.ca/, last accessed on 10 January 2020). Throughout, the training pair $(x,y)\in {\mathbb{R}}^{128\times 128}\times {\mathbb{R}}^{30\times 183}$, and the forward map is the Radon transform, which is normalised. Different peak values of x are used to indicate the count level: $1\times {10}^4$ and $1\times {10}^2$ for respectively moderate and low count levels. The observation y is generated by corrupting the sinogram $\mathcal{A}x$ by Poisson noise entrywise, i.e., $y_i\sim \mathrm{Pois}\left({\left(\mathcal{A}x\right)}_i\right)$, where $\mathrm{Pois}(·)$ denotes the Poisson random variable. The hyper-parameter $\beta$ is tuned in a trial-and-error manner and fixed at $\beta =5\times {10}^{-3}$ below. The experiments are conducted on a desktop with two Nvidia GeForce 1080 Ti GPUs and Intel i7-7700K CPU 4.20GHz×8. It is trained for $T=1\times {10}^5$ batches, each of which contains 10 randomly generated $(x,y)$ pairs on-the-fly. The training almost converges after $2\times {10}^4$ batches, and it takes around 11 hours to go over all $T=1\times {10}^5$ batches. The summarizing statistics reported below are computed from 1000 samples for each observation y generated by the trained network. The implementation uses the following public deep learning frameworks: Tensorflow (https://www.tensorflow.org/, [50]), Tensorflow Probability (https://www.tensorflow.org/probability, 51]), DeepMind Sonnet (github.com/deepmind/sonnet), and ODL (github.com/odlgroup/odl), all last accessed on 10 January 2020; the source code will be made available at github.com/chenzxyz/cvae accessed on 10 January 2020.

First, we compare the proposed cVAE approach with conventional and deep learning-based methods on all 181 phantoms in the BrainWeb dataset. It is compared with the following three benchmark methods: maximum likelihood EM (MLEM) [52], total variation (TV) [53] with non-negativity constraint, and learned gradient descent (LGD) [54]. MLEM and TV are two established reconstruction methods in the PET community, and LGD is an unrolled iterative method inspired by classical variational regularisation and exploit DNNs for iterative refinement. For MLEM, we use the odl inbuilt solver mlem, and for TV, we use the primal dual hybrid gradient method (implemented by odl.solvers.phgd). The regularisation parameter $\alpha$ for total variation prior is fixed at $\alpha =2\times {10}^{-1}$ and $\alpha =2$ for the moderate and low count levels, respectively, which is determined in a trial-and-error manner. The comparative results are summarised in

Table 1. Comparisons between cVAE mean and benchmark methods on 181 BrainWeb phantoms at two count levels: $1\times {10}^4$ (MC) and $1\times {10}^2$ (LC). The numbers in the table denote the SSIM/PSNR values.

	MLEM	TV	LGD	cVAE
MC	0.74/23.20	0.85/28.76	0.92/29.07	0.91/28.01
LC	0.64/21.55	0.62/22.58	0.59/21.68	0.64/23.10

, shown with two of the most popular image quality metrics, i.e., SSIM and PSNR, averaged over all 181 phantoms in the BrainWeb dataset. The results clearly show that cVAE can deliver state-of-the-art point estimates in terms of PSNR and SSIM, especially in the practically very important low-count case. Compared with these deterministic benchmark methods, cVAE can additionally provide information on uncertainty.

Next, we compare cVAE with a probabilistic approach [19], which reports state-of-the-art performance for aleatoric uncertainty (see [55] for theoretical interpretations). It employs (non-Bayesian) neural network ensembles to estimate predictive uncertainty, where each network in the ensemble learns similar values close to the training data, and different ones in regions of the space far from the training data. For the comparison, we train a mixture with three multivariate Gaussians (GM3) without adversarial samples, where the training of each component of the mixture is to fit a mean network and a variance network of the outputs using a Gaussian log-likelihood [56]. To stabilize the training procedure, we first train the mean network, and then train the variance network. Alternatively, one can train the mean network as a warm up and then train the mean and variance networks simultaneously, but it usually leads to worse results, and thus we do not present the corresponding results. The comparative quantitative results are given in

Table 2. PSNR and SSIM values for the reconstructions by the trained cVAE and GM3 on ten phantoms, with peak values $1\times {10}^4$ (MC) and $1\times {10}^2$ (LC). The column index refers to the Python style index of the phantom in the BrainWeb dataset. The PSNR and SSIM values are shown as cVAE/GM3.

		10	20	30	50	70
PNSR	MC	27.66/28.05	27.14/27.48	27.25/27.43	27.25/27.50	25.65/26.77
	LC	22.60/21.86	22.09/21.35	22.30/21.09	22.14/20.69	20.87/19.32
SSIM	MC	0.89/0.89	0.89/0.89	0.91/0.91	0.88/0.88	0.88/0.89
	LC	0.65/0.62	0.65/0.61	0.69/0.62	0.54/0.50	0.46/0.45
		90	100	110	130	150
PSNR	MC	24.98/26.83	26.91/27.74	27.81/28.43	27.96/29.33	30.86/32.57
	LC	20.52/19.02	21.39/19.74	22.22/20.61	22.48/21.32	25.78/23.67
SSIM	MC	0.88/0.91	0.90/0.91	0.92/0.92	0.94/0.94	0.96/0.96
	LC	0.60/0.53	0.62/0.55	0.61/0.57	0.60/0.57	0.73/0.67

, which present the PSNR and SSIM results for ten phantoms from the Brainweb dataset in order to shed fine-grained insights into the performance of the methods over different phantoms. It is observed that in the low count case, cVAE can provide better point estimates in terms of both SSIM and PSNR, which concurs with Figure 6; however, in the moderate count case, GM3 can sometimes deliver better results. In terms of the variance map, the results of GM3 contain more structural patterns and resemble the error map more closely.

To shed more insights into the variance by cVAE and the benchmark GM3, we show the cross-section plots with $0.95$ Highest Posterior Density (HPD) in Figure 7 for both moderate and low count levels. According to Proposition 2, the estimated covariance by cVAE contains two distinct sources, i.e., the sample variance and the variance $\beta$ of the conditional Gaussians $p_{\varphi }\left(x\right|y,z,\mathcal{A})=\mathcal{N}\left(x\right|x^K\left(z\right),\beta I)$. The latter is uniform across the pixels and acts as a background. We show the HPDs of cVAE with full variance (unbiased variance estimated by ${\widehat{\mathrm{Cov}}}_{p_{\varphi }\left(x\right|y)}\left[x\right]$) and the variance without the $\beta$ factor (i.e., ${\widehat{\mathrm{Cov}}}_{p_{\varphi }\left(x\right|y)}\left[x\right]-\beta I$). It is observed that the latter contains more structures in the credible intervals. Furthermore, the overall shape and magnitude of the HPDs by cVAE with the full variance and GM3 are fairly close to each other. It is noteworthy that in the cold regions (i.e., zero count), cVAE can provide almost zero variance upon subtracting the background variance, and it is able to indicate the contrast of variance to highlight the pixels, whereas the variances by GM3 are also relatively high. The comparison between the cross-section plots for low and moderate count cases (i.e., high and low noise levels) on the same ground-truth phantom indicates that cVAE does provide higher uncertainty for a higher noise level, which is intuitively consistent with the underlying statistical background.

Lastly, we evaluate all the methods on the phantoms with an artificially added tumour by changing the pixel values to the peak values in order to examine their capability of recovering novel features that are not presented in the training data. We (randomly) choose two phantoms from the BrainWeb dataset (Python style index: 10 and 110). A small tumour with a radius of 2 and a large tumour with a radius of 5 are added into the 10th and 110th phantoms, respectively. The corresponding reconstructions are shown in Figure 8 and Figure 9, respectively. It is observed that the tumours can be clearly reconstructed by the cVAE means for both count levels, except for the small tumour at low count levels. In the latter case, none of the methods can reasonably reconstruct the tumour since the data are very noisy in terms of low signal strength. The results by cVAE, LGD, and GM3 are comparable, at least visually. The ability to reconstruct tumours further indicates that cVAE does not miss out on important features that are not present in the training data, indicating a certain degree of robustness of the cVAE framework, so long as the signal is sufficiently strong, since many deep learning-based methods tend to miss the tumour due to the bias induced by tumour-free training data [57].

6. Conclusions

In this work, we have developed a general and flexible probabilistic computational framework for the uncertainty quantification of inverse problems in a purely data-driven setting. The approach is based on the conditional variational autoencoder loss and employs the deep neural network as a recurrent unit to repeatedly refine the samples using the observation and forward map, seeded by a probabilistic encoder conditioned on the observation. The efficiency of the framework is underpinned by the encoding of observations in a low-dimensional latent space. The significant potentials of the framework have been demonstrated on a PET image reconstruction with both moderate and low count levels, and the approach shows competitive performance when compared with several deterministic and probabilistic benchmark methods, especially within the low count regime.

There are several avenues for further study. First, the framework is flexible and general, and it is of much interest to evaluate its potentials on other computationally expensive imaging modalities (e.g., MRI, CT, and PET-MRI), especially in the under sampling and low-dose regime, for which there is a great demand on uncertainty quantification due to the lack of sufficient information. Such studies will also provide important insights into the statistical features of the framework. Second, it is of much interest to analyse the theoretical properties of the cVAE loss as an upper bound of the expected KL divergence (e.g., approximation error and asymptotic convergence). This line of research has long been outstanding in approximate inference and often provides theoretical guarantees for the overall inference procedure and guidelines for constructing efficient approximations. Third and last, it is imperative to develop scalable benchmarks for the uncertainty quantification of high-dimensional inverse problems. Several deep learning-based uncertainty quantification techniques have been proposed in the machine learning literature, but mostly on different types of uncertainties or without explicitly elucidating the sources of uncertainties (see the work of [25] for a preliminary study in medical imaging). A careful calibration of the obtained uncertainties is essential towards practical deployment, as recently highlighted in the review presented in [8].