GAN-based Priors for Quantifying Uncertainty
Dhruv V. Patel Assad A. Oberai
stat.ML cs.CV cs.LG physics.comp-ph


[
    [
    January 1, 2004
===================




Bayesian inference is used extensively to quantify the uncertainty in an inferred field given the measurement of a related field when the two are linked by a mathematical model. Despite its many applications, Bayesian inference faces challenges when inferring fields that have discrete representations of large dimension, and/or have prior distributions that are difficult to characterize mathematically. In this work we demonstrate how the approximate distribution learned by a deep generative adversarial network (GAN) may be used as a prior in a Bayesian update to address both these challenges.  We demonstrate the efficacy of this approach on two distinct, and remarkably broad, classes of problems. The first class leads to supervised learning algorithms for  image classification with superior out of distribution detection and accuracy, and for image inpainting with built-in variance estimation. The second class leads to unsupervised learning algorithms for  image denoising and for  solving physics-driven inverse problems.






§ INTRODUCTION



Quantifying uncertainty in an inference problem amounts to making a prediction and quantifying the confidence in that prediction. In the context of an image recovery problem, this may be understood as follows. A typical computer vision algorithm uses a noisy version of an image and prior knowledge to produce the  inferred image which can be interpreted as the “best guess” of the original image. Quantifying uncertainty in this context involves generating an estimate of the level of confidence in the best guess, in addition to the guess itself. 





Bayesian inference provides a principled approach for quantifying uncertainty. As shown in the following section, it treats the inferred vector as a multivariate stochastic vector and leads to an expression for its distribution. This expression can be used to estimate the most likely solution (the maximum a-posteriori estimate, or the MAP), the mean, the variance, or any other population parameter of interest. Thus  providing a recipe for thoroughly quantifying the uncertainty in an inference problem. For the image recovery problems considered in this paper, Bayesian inference not only provides the best guess of the true image, but also a means to estimate measures of uncertainty such as the pixel-wise variance or the likelihood of an out-of-distribution input. 


The knowledge of uncertainty in a prediction can directly influence the downstream action that depends on the inference. Consider an image recovery problem where two distinct inputs lead to similar recovered images: those of a traffic sign with a high speed limit. However, for the first input the predicted variance is small, while for the second input it is large. Further, the set of likely images in the second set also includes images of a Stop Sign. Then the appropriate action for the two inputs, determined after solving the inference problem and quantifying uncertainty, is very different. For the first input, the appropriate action is one of continued motion, whereas for the second input it is to slow down. Similar examples can be drawn from other areas like medical imaging, high frequency trading and autonomous systems (<cit.>). 




The knowledge of uncertainty can also be useful in determining the optimal location of a sensor. Consider an image recovery problem, where the goal is to infer the signal, and associated uncertainty, using limited amount of measurement data. In this problem a user can leverage information about the spatial distribution of uncertainty to choose the location with maximum uncertainty as next measurement location.

This task falls within the fields of active learning and/or design of experiments (<cit.>) and is particularly useful in applications like satellite imaging, where each measurement requires significant time and/or resources. 




In Figure <ref>, we demonstrate how the proposed GAN-based Bayesian inference algorithm for quantifying uncertainty is useful in the scenarios described above. In the first example it is used to compute a measure that detects, with perfect accuracy, that the input to an image classification algorithm is not from the dataset that was used to train it - the so-called out of distribution (OOD) detection problem. In the second example it computes the pixel-wise variance in an image inpainting task, which is then used to determine the location of the subsequent window to be revealed in an optimal iterative strategy.  We return to these applications in greater detail in Section 4.








 §.§ BAYESIAN INFERENCE
 

Bayesian inference is a well-established technique for quantifying uncertainties in inference problems (<cit.>, <cit.>). It has found applications in diverse fields  such as geophysics,

climate modeling,

chemical kinetics,

heat conduction,

astrophysics,

materials modeling,

and the detection and diagnosis of disease.

The two critical ingredients of this technique are - an informative prior distribution representing the prior belief about the parameters to be inferred and an efficient method for sampling from the posterior distribution. In this manuscript we describe how deep generative adversarial networks (GANs) can be effectively used in these roles.


We consider the setting where we wish to infer a vector of parameters y∈ℝ^N from the measurement of a related vector x∈ℝ^P. We allow for two broad classes of problems. In one class (labeled Class 1 in Section 2) the map from y to x is known through a forward model x = f(y). These types of problems are often referred to as inverse problems. In the other class (labeled  Class 2 in Section 2) this map is not known and must be inferred from prior data. In the discussion that follows, we apply Bayesian inference to Class 1 problems and point out two main challenges. We note that the same challenges apply to problems in Class 2 as well. 

 A noisy measurement of x is denoted by x̂ = f(y) + η, where the vector η∈ℝ^P represents noise.  While the forward map f is typically well-posed, its inverse is not, and hence to infer y from the measurement x̂ requires techniques that account for this ill-posedness. Classical techniques based on regularization tackle this ill-posedness by using additional information about the sought solution field explicitly or implicitly (<cit.>). Bayesian inference offers a different approach to this problem by modeling the unknown solution as well as measurements as random variables. This addresses the ill-posedness of the inverse problem, and allows for the characterization of the uncertainty in the inferred solution.

The notion of a prior distribution plays a key role in Bayesian inference.  Through multiple observations of the field y, denoted by the set 𝒮 = {y^(1), ⋯, y^(S)}, we have some prior knowledge of y that can be utilized when inferring y from x̂. This is used to build, or intuit, a prior distribution for y, denoted by p^ prior_Y(y). Some typical examples include Gaussian process prior with specified co-variance kernels, Gaussian Markov random fields, 
Gaussian priors defined through differential operators, 
and hierarchical Gaussian priors.

These priors promote smoothness and/or structure in the inferred solution and can be expressed explicitly in an analytical form.

Another key component of Bayesian inference is a distribution that represents the likelihood of x given an instance of y, denoted by p^ l(x|y). This is often determined by the distribution of the error in the model.

Given this, 

the posterior distribution of y, determined using Bayes' rule after accounting for the observation x̂ is given by, 

    p^ post_Y(y|x)    =   1/ℤ p^ l(x|y) p^ prior_Y(y)

Here, ℤ



is the prior-predictive distribution of y.


The posterior distribution characterizes the uncertainty in y; however for vectors of large dimension characterizing this distribution explicitly is a challenging task. Consequently the expression above is used to perform tasks that are more manageable. These include determining estimates such as the maximum a-posteriori estimate (MAP), expanding the posterior distribution in terms of other distributions that are simpler to work with (<cit.>), or using techniques like Markov Chain Monte-Carlo (MCMC) to generate samples that are close to the samples generated by the true posterior distribution (<cit.>).

In summary, despite its numerous applications, Bayesian inference faces significant challenges. These include defining a reliable and informative prior distribution for x when the set 𝒮  is difficult to characterize analytically, and efficiently sampling from the posterior distribution when the dimension of x is large; a typical situation in many practical science and engineering applications.
















 §.§ OUR CONTRIBUTION

The main contribution of this paper are:


    
  * A novel method for performing Bayesian inference involving complex priors and high dimensional posterior. We utilize the distribution learned by a GAN as a surrogate for the prior distribution and reformulate the inference problem in the low-dimensional latent space of the GAN. Furthermore, we provide a theoretical analysis of the weak convergence of the posterior density learned by the proposed method to the true posterior density.
    
  * Application of this method to problems where the map from the inferred to the measured vector is known a-priori. This leads to novel unsupervised algorithms for physics-based inverse problems and image denoising problems with quantitative measures of uncertainty. 
    
  * Application of this method to problems where the map from the inferred to the measured vector is not known and is determined from data. This leads to novel algorithms for image classification and image inpainting with quantitative measures of uncertainty. 
    
  * Demonstration of the utility of quantifying uncertainty in the detection of out-of-distribution (OOD) samples and in active learning. 






 §.§ RELATED WORK


The main idea developed in this paper tackles the challenges described above by training a generative adversarial network (GAN) using the sample set 𝒮, and then using the distribution learned by the GAN as the prior distribution in Bayesian inference. 

Related work in this area can be organized by considering the two broad classes of problems this idea is applied to. 

In one class of problems, which are referred to as inverse problems, the map x = f(y), that is the map from the inferred field to the measurement is known. The use of sample-based priors for solving inverse problems has a rich history (<cit.>). As does the idea of reducing the dimension of the parameter space by mapping it to a lower-dimensional space (<cit.>). However, the use of learning-based deep generative models like GANs in these tasks is novel. Recently, several authors have considered the use of learning-based methods for solving inverse problems arising in different domains. These include the use of deep convolutional neural networks (CNNs) or recurrent neural networks (RNNs) to solve physics-driven inverse problems (<cit.>) and use of deep generative models like VAEs and GANs to solve inverse problems arising in computer vision (<cit.>). There is also a growing body of work dedicated to using GANs as a regularizer in solving inverse problems (<cit.> and in compressed sensing (<cit.>). However, these approaches differ from ours in that they solve the inverse problem as an optimization problem and do not rely on Bayesian inference; as a result, they add regularization in an ad-hoc manner and do not attempt to quantify the uncertainty in the inferred field. 
More recently, the approach described in (<cit.>) utilizes GANs in a Bayesian setting; however the GAN is trained to approximate the posterior distribution and not the prior, as in our case.


In another class of problems the forward map is not known; however in its lieu pair-wise instances of y and x are available. In this case the algorithms most closely related to our approach are the so-called hybrid methods, where an invertible  generator is trained to learn p(x) and is linked to another network that is trained to produce p(y|x) (<cit.>). This algorithm is then applied to image classification problems, where for a given input image x, p(x) is used to determine the likelihood of the input and p(y|x)  is used to infer the probability of the corresponding label. In contrast to this, we train a Wasserstein GAN to learn the joint density p(u), where u = [x,y] and then use this as the prior in a Bayesian update of the posterior density for a given input x. The sampling problem for the posterior is reduced to the latent space of the GAN, which is of smaller dimension, and a Markov Chain Monte-Carlo algorithm is trained to generate samples of  p(u|x) and fully characterize the posterior density. 

We note that deep learning based Bayesian networks, where the network weights are stochastic parameters that are determined using Bayesian inference, are another means of quantifying uncertainty (<cit.>), and have recently been applied to semantic image-segmentation and super-resolution (<cit.>). 



















§ PROBLEM FORMULATION



We consider the problem where we wish to infer the vector y from the noisy measurement of a related vector x.

We consider two broad classes of problems of this type. In one class we assume that the forward operator which maps y to x, that is x = f(y), is known. While in the other, we assume that f(y) is not known and any relation between x and y must be determined from data. 




 §.§ CLASS 1: THE MAP f(y) IS KNOWN

These problems are commonly referred to as inverse problems and  the map x = f(y) is called the forward map. In this class of problems, this forward map is usually well-defined and is assumed to be known either through physics-based principles or through other modeling paradigms.


A noisy measurement of x is denoted by x̂ = f(y) + η, where η∼ p_η is the noise vector. In addition, it is assumed that the sample set 𝒮 = {y^(1), ⋯, y^(S)}, which contains multiple  realizations of y drawn from the distribution P_Y, is also known. The goal is to use the prior information encoded in 𝒮, the noisy measurement x̂, and the forward map f to determine the distribution of the vector y.

The prior information for this class of problems is built from the distribution of y alone. Thus problems in this class fall into unsupervised learning category,  where training only requires instances of one type of data (the vector to be inferred). The need to have access to pairwise samples of x and y is circumvented by the knowledge of the forward operator. An example problem in this class is that of image denoising, where x represents a noisy image, the operator f is the identity, and y is the de-noised image. Another example, which is drawn from physics-based inverse problems, is where one wishes to infer the initial temperature field from a measurement of the temperature field at a later time. Here x represents the temperature field at a finite time T > 0, f is the forward in time heat conduction operator, and y is the temperature at the initial time. A large number of other physics-based inverse problems can also be cast in this form. 


Given x̂, using Bayes' rule we may write the posterior distribution of y as,

    p^ post_Y(y|x)    =   1/ℤ p^ l(x|y) p^ prior_Y(y) 
       =   1/ℤ p_η(x̂ - f(y)) p_Y(y).

where ℤ is the prior-predictive distribution of y and ensures that the posterior integrates to one. 


In order to efficiently sample the posterior density, we first train a GAN using the set 𝒮 whose elements are sampled from P_Y.  We let z∼ p_Z(z) characterize the latent vector space of the GAN, and let g(z) and d(y) denote its generator and discriminator, respectively. In Appendix <ref>, assuming that (a) the stationarity conditions for the adversarial loss function are satisfied and (b) that the set of basis functions obtained by taking the derivative of the discriminator with respect to its weights forms a complete basis in L^∞(Ω_y), in the limit of infinite capacity, we prove that  

    _y∼ p_Y [m(y)] = _z∼ p_Z [ m (g(z))],

for sufficiently smooth m(y). This equation states that once a GAN is trained using the sample set 𝒮, it may be used to evaluate any population parameter for p_Y by sampling from p_Z and then passing the samples through the generator. Since the dimension of the latent space is much smaller than that of y, this represents an efficient means of evaluating population parameters. 



In order to turn this into an expression for evaluating  a population parameter for the posterior density, we select m(y) = l(y) p_η(x̂ - f(y))/ℤ, substitute this expression on both side of (<ref>), and use (<ref>) to arrive at,

    _y∼ p^ post_Y[l(y)] 
           =   _z∼ p^ post_Z[l(g(z))],

where 

    p^ post_Z(z|x) ≡1/ℤ p_η(x̂ - f(g(z))) p_Z(z).

The distribution p^ post_Z is the analog of p^ post_Y in the latent vector space. The measurement x̂ updates the prior distribution for y to the posterior distribution. Similarly, it updates the prior distribution for z, p_Z, to the posterior distribution, p^ post_Z, defined above. 

Equation (<ref>) implies that sampling from the posterior distribution of y is equivalent to sampling from the posterior distribution for z and passing the sample through the generator g. That is,

    y∼ p^ post_Y (y| x̂) ⇒y = g(z), z∼ p^ post_Z(z|x̂).

Since the dimension of z is typically smaller than that of y, this represents an efficient approach to sampling from the posterior of y. 

The left hand side of (<ref>) is an expression for a population parameter of the posterior.

The right hand side of this equation describes how this parameter may be evaluated by sampling z (instead of y) from  p_Z^ post. In practise this is accomplished by generating an MCMC approximation, p^ mcmc_Z(z|x) ≈ p^ post_Z(z| x) using the definition in (<ref>), and thereafter sampling z from this distribution. This circumvents the calculation of the prior-predictive distribution of y (denoted by ℤ), which would be necessary when using (<ref>) directly. Then from (<ref>), any desired population parameter for posterior distribution may be approximated as 

    l (y)   ≡   _y∼ p^ post_Y[l(y)]  
       ≈   ∑_n = 1^N_ samp l(g(z))/N_ samp,   z∼  p^ mcmc_Z(z| x).

For all the numerical experiments in this paper we have used this approach to evaluate population parameters.





 §.§ CLASS 2: THE MAP f(y) IS NOT KNOWN

We now consider problems where the relation between x and y is not known and must be inferred from data. 
We denote by u = [x,y] the joint vector and recognize that the measurement has the form x̂ = 1_xu + η, where 1_x is the indicator function that extracts  components of x from u, and η is  the noise vector drawn from the distribution p_η. Further we assume that the sample set 𝒮= {u^(1), ⋯, u^(S)} contains multiple measurements of u drawn from the distribution P_U. The goal is to use the prior information encoded in 𝒮 and the new, noisy measurement x̂ to determine the distribution for the corresponding vector y, and perhaps also the de-noised version of x.

Since the prior information is built from the joint distribution P_U this class of problems is one of supervised learning where training requires pair-wise instances of x and y. An example problem in this class is that of image classification, where x represents an image and y represents the corresponding one-hot encoded label vector. Another example is that of image inpainting, where x represents the portion of an image that is revealed and y represents the portion that is occluded. 

The use of GANs as priors in this class of problems closely parallels the development for problems treated in the previous section. Therefore, rather than repeating the entire development below, we only highlight the important steps and salient differences. 

Using Bayes' rule we may write the posterior distribution of u = [x, y] as,

    p^ post_U(u|x)    =   1/ℤ p^ l(x|u) p^ prior_U(u) 
       =   1/ℤ p_η(x̂ - 1_x (u)) p_U(u),

where ℤ



is the prior-predictive distribution of u. In order to efficiently sample the posterior density we train a GAN using the set 𝒮 to generate a prior.  As before, we let z∼ p_Z(z) characterize the latent vector space of the GAN, and let g(z) denote its generator. 

Under the assumptions of the result derived in  Appendix <ref>, we have 

    _u∼ p_U [m(u)] = _z∼ p_Z [ m (g(z))],

for sufficiently smooth m(u). We choose m(u) = l(u) p_η(x̂ - 1_x (u))/ℤ, substitute it in (<ref>), and make use of (<ref>)  to arrive at,

    _u∼ p^ post_U[l(u)] 
           =   _z∼ p^ post_Z[l(g(z))],

where 

    p^ post_Z(z|x) ≡1/ℤ p_η(x̂ - 1_x (g(z))) p_Z(z).

The distribution p^ post_Z is the analog of p^ post_U in the latent vector space. We use (<ref>) to generate an MCMC approximation, p^ mcmc_Z(z|x) ≈ p^ post_Z(z| x) of the posterior. Thereafter, from (<ref>) we conclude that any population parameter for the posterior can be approximated as 

    l (u)   ≡   _u∼ p^ post_U[l(u)]  
       ≈   ∑_n = 1^N_ samp l(g(z))/N_ samp,   z∼  p^ mcmc_Z(z| x).

We note that this approach allows us to compute population parameters for the entire vector u, which includes the vector y, which is not observed, as well as the vector x, for which a noisy measurement, x̂, is available. While it is clear that parameters related to y are useful, in some instances it is also useful to estimate population parameters related to x. A case in point is the image classification problem considered in the following section. In this problem x̂ represents the input image and y represents its label. Here computing parameters associated with y provide information about the label for an image. In addition, computing x^ map is useful since large values of the quantity x^ map - x̂, which measures the distance between the mode of the posterior distribution and the input image, are strongly correlated with input images that lie outside of the range of the prior, thus enabling the detection of out of distribution (OOD) samples. 



  
Summary We have described a method for probing the posterior distribution in two broad classes of problems when the prior is defined by a GAN. The steps of our algorithm are: (1)  Train a GAN using the sample set 𝒮 to learn the prior distribution. (2) Reformulate the posterior distribution in the latent space of the GAN.(3) Run a Markov chain Monte Carlo algorithm to generate samples from this low-dimensional posterior distribution. (4) Use MCMC-generated samples to compute population parameters that quantify the uncertainty in the inference.


In the following section we apply  the above algorithm to a broad class of problems where we draw inferences from noisy measurements and quantify uncertainty in these inferences. Wherever possible, we compare our predictions with related methods and/or benchmark solutions and also highlight the role of uncertainty quantification in downstream tasks. In Appendix <ref>, we also derive  a computationally efficient approach for estimating the MAP for the posterior density of the latent vector under the assumptions of Gaussian noise and prior. 





§ EXPERIMENTS

In this section we apply our method to the two broad classes considered earlier. One where the forward map is known and another where it is inferred from data. In each case we apply our method to determine important population parameters that include y^ mean, var (y) and x̂ - x^ map. Thereafter, we use these to answer important questions like: Is the input to the inference problem consistent with the prior data it was trained on? Do we have confidence in the inference? How do we utilize this knowledge in order to design the next measurement. 




In all cases we use a Wasserstein GAN-GP (<cit.>) to learn the prior density (architecture described in Appendix <ref>). We also ensure that the target images are not chosen from the set used to train the GAN. We sample from the posterior using Hamiltonian Monte Carlo (<cit.>) and implement it using Tensorflow-probability library. We use initial step size of 1.0 for HMC and adapt it following (<cit.>) based on the target acceptance probability. We use 64k samples with burn-in period of 0.5. We select these parameters to ensure convergence of chains. Using the HMC sampler we compute the population parameters of interest.














 §.§ IMAGE CLASSIFICATION

This problem belongs to the second class, where the forward map is not known and must be learnt from data. The objective of this task is to infer the label y along with its uncertainty for a given input image x̂. This predictive uncertainty estimation is crucial in deep learning applications where high-stakes decisions are made based on the output of a model (<cit.>). 
It has been shown that in real-world scenarios, where a model might encounter inputs that are anomalous to its training data distribution, many models produce overconfident predictions (<cit.>) raising serious concerns about AI safety (<cit.>). In this situation, it is desirable that such out-of-distribution (OOD) data points are detected upfront before making any prediction. A useful probabilistic predictive model should therefore flag all OOD data points, maintain high levels of accuracy on in-distribution data points, and provide a measure of confidence in its predictions.

In order to achieve this goal, we compute three different quantities: x̂ - x^ map for OOD detection, y^ mean for prediction, and var (y) as a measure of confidence in the prediction. 

We consider the MNIST database of hand-written digits and use 55k images and the corresponding labels to train a  WGAN-GP. Thereafter, we use the MNIST test set to test the performance of our algorithm for in-distribution data, and NotMNIST test set for OOD data. Our approach of learning and inferring the joint distribution is closely related to hybrid models and hence we compare our performance against the most recent hybrid models in Table <ref>. 

We determine whether a given test image is OOD based on a rejection rule. If this condition is satisfied then following <cit.> we set the probability of each label to be equal. We then quantify the performance of the rejection rule by reporting the average entropy of the labels for all test samples from the OOD set and the false positive rate (FPR = # of in-distribution samples rejected as OOD/ # of in-distribution samples). Thereafter, for all in-distribution samples that are correctly identified, we report the accuracy of predicting the label, which is determined from y^ mean. We consider two rejection rules: x̂ - x^ map > c_1, and x̂ - x^ map +var(y)> c_2. 

The performance of x̂ - x^ map > c_1  rejection rule can be discerned in Figure <ref>, where we observe that it perfectly segregates the in-distribution and OOD samples. This is also apparent in Table <ref>, where it yields zero FPR and maximum entropy. Its accuracy for the in-distribution samples is also quite high. This accuracy can be further improved by using the combined rejection rule x̂ - x^ map +var(y)> c_2, since it rejects some incorrectly labeled in-distribution samples with high variance as OOD. However, this comes at the cost of a slightly higher FPR. The usefulness of var(y) as a predictor of accuracy is evident in Figure <ref>, where we observe that most correctly labeled samples have low variance (avg. value = 0.0026 ± 6e-3) when compared with their incorrectly labeled counterparts (avg. value = 0.025 ± 5e-3). 



In Table <ref> we compare the performance of our GAN-based approach with two hybrid invertible flow-based models <cit.>. The explicit nature of these models allows the joint density to be decomposed into generative and discriminative components, enabling a way to explore the generative-discriminative trade-off by introducing a weighted likelihood objective with a scaling parameter λ. Values of  λ<1 favor discriminative performance, while λ>1 favors generative performance. In this context our approach may be regarded as one  where the generative and discriminative components are equally weighted, and is therefore close to the choice λ =1. Given this, in Table <ref>, we have compared our approach with hybrid models where λ≈ 1. We note that with the x̂ - x^ map≤ c_1 rejection rule our model performs competitively with both the scaled (<cit.>) and the un-scaled hybrid models (<cit.>) for both in-distribution and OOD datasets. With the x̂ - x^ map +var(y)> c_2 rejection rule it outperforms both hybrid models giving maximum accuracy and entropy but with non-zero FPR.














 §.§ IMAGE INPAINTING
 
Image inpainting also belongs to second class of problems, where the forward map is unknown and has to be learnt from data. In this case the quantity to be inferred (y) is the occluded region of an image, and the measurement (x̂) is the noisy version of its visible portion. The goal is to recover the entire image (u = [x, y]).  

While there has been great interest in recent years in developing efficient deep learning-based image inpainting algorithms (<cit.>), most of it has focused on deterministic algorithms that lack the ability to quantify uncertainty in a prediction. In contrast, we use the algorithm in Section <ref> to perform probabilistic image inpainting.

We consider MNIST dataset and use 55k images to train a WGAN-GP. We generate measurements by selecting an image from the test set, occluding a significant region and then adding Gaussian noise. With this as input we use our algorithm to generate samples from the posterior distribution of the entire image (both occluded and retained regions). From these samples we evaluate the relevant population parameters, u^ map, u^ mean, and var(u). These results are shown in Figure <ref> and indicate that the map and mean images are close to the true image, even in the presence of significant occlusion and noise. The image of pixel-wise variance reveals that we are most uncertain along the boundaries of the digits and around the occlusion window. 









In Figure <ref>, we demonstrate how uncertainty  may be used in active learning/design of experiment, where the goal is to determine the optimal location for a measurement. We begin with an input where the entire image is occluded and in every subsequent step, we allow for a small 7×7 pixel window to be revealed. We select the window with the largest average pixel-wise variance. As the iterations progress, the MAP estimate converges to the true digit, and the variance decreases. In about 4 iterations we arrive at a very good guess for the digit. The performance of this approach is quantified in Figure <ref>, where we have plotted reconstruction error versus the number windows for this strategy, and a strategy where the subsequent window is selected randomly. The variance-driven strategy consistently performs better. 
We are not aware of any other methods for computing uncertainty in recovered images that have been applied to drive an active learning task in image inpainting. While methods based on dropout (<cit.>) or variational inference (<cit.>) could be extended to accomplish this, this has not been done thus far.





Results for the variance-based window selection strategy applied to the CelebA dataset are shown in Figure <ref>. We observe that the algorithm produces realistic images at each iteration; however, the initial variance is large indicating large uncertainty. As more windows are sampled using the active learning  strategy, the variance reduces and by the 7th iteration a good approximation of the true image is obtained, even though only a small, noisy portion is revealed. Additional results for this dataset are discussed in Appendix <ref>.  
















































 §.§ A PHYSICS-DRIVEN INFERENCE PROBLEM

We now consider a problem from Class 1, where the forward map is known. In particular,  we consider the problem of inverse heat conduction, where the goal is to infer the initial temperature distribution (at t=0) in a domain given a noisy measurement of temperature  at later time (t=1) and the thermal conductivity of the material. The forward map is the solution of the time-dependent heat conduction problem with uniform conductivity, κ = 0.64, in a square domain of length L = 2π with Dirichlet boundary conditions. This operator maps the initial temperature field (y) to the temperature field at later time (x). Its  discrete version is obtained by discretizing the time-dependent linear heat conduction equation using the central difference scheme in space and backward Euler scheme in time. Much like a blurring kernel, the forward operator smooths the initial temperature distribution, and the extent of smoothing increases with κ× t. 

We consider a family of initial temperatures where the background is zero, and the temperature on a rectangular sub-domain varies linearly from  2 units on the left edge to 4 units on the right edge. This distribution is parameterized  by four parameters, {ξ_i}_i=1^4, which are the coordinates of the lower left and upper right corners of the rectangular region. The sample set 𝒮 is created by sampling each parameter from a uniform distribution  and is used to train the GAN  prior. The posterior distribution is sampled using the HMC sampler. 

In the top two rows of Figure <ref>, we have plotted the true initial condition, the noise-free temperature at t=1, and the noisy temperature measurement (σ_x = 1) used as input in the GAN-based prior approach. The corresponding MAP, mean and pixel-wise variance estimated by the MCMC approximation are shown next. We observe that the MAP is very close to the true initial temperature distribution and the variance is concentrated along the edges of the rectangle where the uncertainty is the largest. In the following columns we have plotted the MAP estimate obtained assuming L_2 and H^1 Gaussian priors, which are often used to solve these types of problems, and are clearly much less accurate. 

For this problem the “true” posterior  can be reduced to the 4-dimensional space of parameters, and sampled by generating initial conditions corresponding to the values of these parameters. A simple Monte-Carlo approximation can be performed to compute the mean and the pixel-wise variance for the true posterior (last two columns of Figure <ref>). By comparing these with the mean and the pixel-wise variance (columns 5 & 6) estimated by the GAN-based prior, we conclude that GAN-based posterior has converged to the true posterior.  



In the bottom rows of Figure <ref>, we plot similar results for initial conditions generated from the MNIST database when the GAN prior was also trained on the MNIST database. The measurement is made at t=0.2. Once again we observe that mean and the map estimated by our approach is very close to the true initial condition, while the MAP solutions obtained from the L_2 and H^1 priors are inaccurate. The pixel-wise variance illustrates the uncertainty in determining the boundary of the digits. 










 §.§ IMAGE DENOISING

As another example of a problem where the forward map is known (Class 1 problem) we consider image denoising. Here the forward map is the identity, the measured data is the noisy image and the inferred field is its de-noised version. We consider the MNIST dataset and use 55k images to train the GAN. 
We add Gaussian noise with zero mean and specified variance (σ_x) to the test images and use these as  measurements to recover the distribution of likely images using the MCMC approach. In Figure <ref>, we have plotted the noisy input image, the MAP estimate, and the pixel-wise mean and variance. For low and medium noise levels (σ_x = 0.1, 1), we are able to recover the original image with good accuracy, the pixel-wise variance is small overall, and is largest around the boundary of the recovered digit. For the highest noise level (σ_x = 10), however, the image recovered by the MAP is incorrect in 2/3 cases, and would be misleading if viewed by itself. However, when viewed in conjunction with the estimated variance, which is large, it is clear that the confidence in the inference is small and the inferred image ought not be trusted for downstream tasks. The dependence of the average per-pixel variance in the recovered image on the variance of noise in the measured image is shown in Figure <ref>, and it increases with noise. Additional results for CelebA are discussed in Appendix C.2











§ CONCLUSIONS


The ability to quantify the uncertainty in an inference problem is useful in developing confidence in that inference, identifying measurements that are outliers, and in designing strategies to improve the confidence. 

In this paper we have described how this may be accomplished when solving a Bayesian inference problem by using GANs as priors. Since GANs can learn complex distributions of a wide variety of fields from their samples, this approach can be applied to a  range of problems in computer vision and physics-driven inference. This includes those where the operator that maps the inferred field to measurement is known (so-called inverse problems) and those where this map is not known and must be inferred from data. It derives its efficiency by mapping the posterior distribution to the latent space, whose dimension is often much smaller than that of the inferred field. We have presented applications of this approach to image classification, image inpainting, image denoising and physics-driven inverse problems. 










plainnat 







§ WEAK CONVERGENCE OF THE PRIOR DENSITY 


Let the generator of the Wasserstein GAN be given by g(z;θ), where z∈ℝ^M is the latent vector, and θ∈ℝ^N_θ is the vector of weights. The vector z is selected from the distribution p_Z(z). Note that g: ℝ^M →ℝ^N. 

Let the discriminator of the GAN be given by d(y;ϕ), where y∈ℝ^N, and ϕ∈ℝ^N_ϕ be the vector of weights. Note that d: ℝ^N →ℝ(0,1). 

Assume that the GAN is trained using a set of samples of y, drawn from p_Y(y). 

Then under the following assumptions:

    
  * The stationarity conditions for the adversarial loss function are satisfied.
    
  * In the limit of infinite capacity (N_ϕ→∞), the set of basis functions obtained by taking the derivative of the discriminator with respect to its weights forms a complete basis in L^∞(Ω_y)
 
For a sufficiently smooth m(y), we prove that 
    |_y∼ p_Y [m(y)]| =  | _z∼ p_Z [ m (g(z))]|.

That is we establish the weak convergence of the density obtained by using a GAN as a prior to the true prior density.









  
Proof. For the loss function, consider 

    L (θ,ϕ)    =   _y∼ p_Y  [ ρ( 1- d(y;ϕ))]  
           +  _z∼ p_Z [ρ(d(g(z; θ); ϕ))].

Here ρ is a monotone real-valued function which defines the GAN family being analyzed. For example, for the Wasserstein GAN, ρ(ξ) = ξ.

The optimal values of the weights are given by 

    θ^*, ϕ^* =  _θ (_ϕ (L (θ,ϕ))).



The necessary conditions for these optimal values are 

    ∂ L (θ^*,ϕ^*) /∂ϕ   =   0
    ∂ L (θ^*,ϕ^*) /∂θ   =   0.


Using the definition of the loss function (<ref>) in (<ref>), we have 

    _x∼ p_Y [ ρ'( 1- d(y;ϕ^*)) ∂ d/∂ϕ (y;ϕ^*) ] =         
        _z∼ p_Z [ρ'(d(g(z; θ^*); ϕ^*)) ∂ d/∂ϕ (g(z; θ^*);ϕ^*)].



Similarly, using (<ref>) in (<ref>), we have 

    _z∼ p_Z [ρ'(d(g(z; θ^*); ϕ^*)) ∂ d/∂x (g(z; θ^*);ϕ^*) ·
    ∂g/∂θ (z; θ^*) ] = 0.


For the Wasserstein GAN, ρ(ξ) = ξ and ρ'(ξ) = 1. As a result (<ref>) reduces to 

    _y∼ p_Y [ ∂ d/∂ϕ (y;ϕ^*) ]    =   _z∼ p_Z [ ∂ d/∂ϕ (g(z; θ^*);ϕ^*)],

and (<ref>) reduces to,

    _z∼ p_Z [ ∂ d/∂y (g(z; θ^*);ϕ^*) ·∂g/∂θ (z; θ^*) ]    =   0.


Let w_a(y) ≡∂ d/∂ϕ_a (y;ϕ^*), then for , a = 1, ⋯, N_ϕ.  (<ref>) implies 

    _y∼ p_Y [ w_a (y) ]    =   _z∼ p_Z [ w_a (g(z; θ^*))].

As N_ϕ→∞, this equation implies that 






the push forward of the measure in the latent space under the function g(z) weakly converges to the measure associated with distribution of y. We note with increasing number of weights in the discriminator, the relation above is required to hold for an increasing number of test functions, w_a. In addition, we have implicitly assumed that the generator is rich enough, that is it has enough weights/layers, such that this relation can actually be satisfied. To make this clear consider the extreme case of a generator with a single weight; in this case there is no way that (<ref>) will be satisfied for a large number N_ϕ. Thus in order for this relation to hold for a large N_ϕ, we must also provide the generator with a large N_θ. 

Let 𝒱≡ span{w_a(y), a = 1, ⋯, N_ϕ}. Then from (<ref>) for any v ∈𝒱, we have 

    _y∼ p_Y [v (y) ]    =   _z∼ p_Z [ v (g(z))].

In the equation above, and hereafter, we have suppressed the arguments θ^* and ϕ^* for ease of notation, implicitly assuming that the relations hold at the optimal values of weights. 
Now consider a sufficiently smooth function m(y) which defines the point estimate we wish to compute, and let m̅(y) be its L^∞ projection on to 𝒱. That is,

    m̅(y) = _v ∈𝒱 m(y) - v(y) _L^∞(Ω_y).

and let ϵ =  m(y) - m̅(y) _L^∞(Ω_y). 
Given the assumption that the functions w_a form a complete basis in L^∞(Ω_y), we note that as N_ϕ→∞, ϵ→ 0.

Now consider the difference, 

    |_y∼ p_Y [m(y)] - _z∼ p_Z [ m (g(z))]| 
          ≤ 
    |_y∼ p_Y [m(y)] - _y∼ p_Y [m̅(y)] 
               + _z∼ p_Z [ m̅ (g(z))]- _z∼ p_Z [ m (g(z))]| 
          ≤ |_y∼ p_Y [m(y)- m̅(y)]| 
               + | _z∼ p_Z [ m (g(z))- m̅ (g(z))]| 
          ≤_y∼ p_Y [ϵ] + | _z∼ p_Z [ϵ] 
           = 2 ϵ.

In the equation above, the first inequality is obtained by recognizing that m̅(y) ∈𝒱 and using (<ref>), the second inequality is a consequence of the triangle inequality and the third is due to the definition of ϵ. Now in the limit N_ϕ→∞, 𝒱 tends to a complete basis, therefore ϵ→ 0 and we have the desired result.




§ EXPRESSION FOR THE MAXIMUM A-POSTERIORI ESTIMATE 




The techniques described in Section 2.1 focus on sampling from the posterior distribution and computing approximations to population parameters. These techniques can be applied in conjunction with any distribution used to model noise and the latent space vector; that is, any choice of p_η (likelihood) and p_Z (prior). In this section we consider the special case when Gaussian models are used for noise and the latent vector. In this case, we can derive a simple optimization algorithm to determine the maximum a-posteriori estimate (MAP) for p^ post_Z(z|x). This point is denoted by z^ map in the latent vector space and represents the most likely value of the latent vector in the posterior distribution. It is likely that the operation of the generator on z^ map, that is g(z^ map), will yield a value that is close to y^ map, and may be considered as a likely solution to the inference problem. 

We consider the case when the components of the latent vector are iid with a normal distribution with zero mean and unit variance. This is often the case in many typical applications of GANs. Further, we assume that the components of noise vector are defined by a normal distribution with zero mean and a covariance matrix Σ. Using these assumptions in (<ref>), we have 

    0.9!
    p^ post_Z(z|x) ∝exp( -1/2( | Σ^-1/2 (x̂ - f(g(z)))|^2  + |z|^2 )^≡ r(z)).


The MAP estimate for this distribution is obtained by minimizing the negative of the argument of the exponential. That is 

    z^ map = _z r(z).

This minimization problem may be solved using any gradient-based optimization algorithm. The input to this algorithm is the gradient of the functional r with respect to z, which is given by 

    ∂ r/∂z = H^T Σ^-1 (f(g(z)) - x̂) + z,

where the matrix H is defined as 

    H≡∂f(g(z))/∂z = ∂f/∂y∂g/∂z.

Here ∂f/∂y is the derivative of the forward map f with respect to its input x, and ∂g/∂z is the derivative of the generator output with respect to the latent vector. In evaluating the gradient above we need to evaluate the operation of the matrices ∂f/∂y and ∂g/∂z on a vector, and not the matrices themselves. The operation of ∂g/∂z on a vector can be determined using a back-propagation algorithm within the GAN; while the operation of ∂f/∂y can be determined by making use of the adjoint of the linearization of the forward operator. 



Once z^ map is determined, one may evaluate g(z^ map) by using the GAN generator. This represents the value of the field we wish to infer at the most likely value value of latent vector. Note that this is not the same as the MAP estimate of p^ post_Y(y|x). 

We note that the result derived above applies to Class 1 problems, that is inverse problems. A similar result can also be derived for problems from Class 2, by using (<ref>) as a starting point and repeating the steps outlined above. In this case z^ map is the minimizer of 
'
    r ≡  | Σ^-1/2 (x̂ - R(g(z)))|^2  + |z|^2,

where R is the restriction operator. The gradient for this optimization problem is given by 

    ∂ r/∂z = H^T Σ^-1 (R(g(z)) - x̂) + z,

where the matrix H is defined as 

    H  = R∂g/∂z.











§ ADDITIONAL RESULTS

In this section we provide additional results for both MNIST and CelebA dataset for different tasks discussed in the main paper.


 §.§ MNIST

First we provide additional examples in Figure <ref> for variance-based adaptive measurement window selection procedure described in Section <ref>.




Figure <ref> shows additional results for the inpainting + denoising task, where an MNIST digit is occluded with masks of different sizes at different locations. Note that the variance is high where the occlusion mask is located indicating lower confidence in reconstructed image in that location. 


 
 


 §.§ CelebA
 
For the CelebA dataset, we trained  WGAN-GP model using more than 200k celebrity facial images and perform inference using remaining test set images. The input images were cropped to a 64 × 64 RGB image and were normalized between [-1, 1]. 


Once the GAN was trained, the HMC algorithm was used for posterior sampling and inference on a complimentary set of images (not used for training). In Figure <ref> we show some additional results for variance-based adaptive measurement window selection procedure for CelebA dataset. 



Next, in figure <ref> we show some additional results for image recovery task for CelebA dataset. Once again we note that the MAP estimate and the mean is close to the true image. On the other hand, the closest image from the training set (in an L_2 sense) is not as accurate. This points to the utility of using the GAN as an interpolant in the latent vector space.  









§ ARCHITECTURE AND TRAINING DETAILS 

We use the WGAN-GP model for learning prior density. The tuned value of hyper-parameters is shown in Table <ref>. 





We use the same generator and discriminator architecture for the MNIST and the physics-based inference problem; whereas for the CelebA dataset we use a slightly different architecture to accommodate different input image size. The layout of both these architecture is shown in Figure <ref> and <ref>. Some notes regarding nomenclature used in  Figure <ref> and <ref>.






    
  * Conv (H × W × C | s=n) indicates convolutional layer with filer size of HxW and number of filters=C with stride=n.
    
  * FC (x,y) indicates fully connected layer with x neurons in input layer and y neurons in output layer.
    
  * BN = Batch norm,  LN = Layer norm.
    
  * TrConv = Transposed Convolution.
    
  * LReLU = Leaky ReLU with α=0.2.
































































































































































