Denoising Score-Matching for Uncertainty Quantiﬁcation in Inverse Problems
Abstract
Deep neural networks have proven extremely efﬁcient at solving a wide range of inverse problems, but most often the uncertainty on the solution they provide is hard to quantify.
In this work, we propose a generic Bayesian framework for solving inverse problems, in which we limit the use of deep neural networks to learning a prior distribution on the signals to recover.
We adopt recent denoising score matching techniques to learn this prior from data, and subsequently use it as part of an annealed Hamiltonian Monte-Carlo scheme to sample the full posterior of image inverse problems.
We apply this framework to Magnetic Resonance
Image (MRI) reconstruction and illustrate how this approach not only yields high quality reconstructions but can also be used to assess the uncertainty on particular features of a reconstructed image.
Introduction
Generative modeling has enjoyed remarkable success in recent years with models such as Generative
Adversarial Networks (GANs)
[Goodfellow et al., 2014] reaching extremely high quality results on complex high resolution images [Karras et al., 2020].
Yet, GANs are still prone to issues including unstable training and mode collapse, i e. a lack of diversity in generated images.
In addition,
GANs only provide a convenient way to sample from the learned distribution, they do not give access to the density function itself.
Another class of commonly used models, Variational AutoEncoders (VAEs) [Kingma and Welling, 2014] have also been able to reach high quality sample resolutions [Vahdat and Kautz, 2020], but can only provide a bound on the density function of the model.
In some practical applications of these generative models however, the most desirable aspect is not necessarily to be able to sample from the model, or directly have access to its density functions, but have access to gradients of the log density function, i.e. its score .
In particular, in this work we focus on the generic problem of performing Bayesian inference to solve an inverse problem (i.e. deconvolution, denoising, inpainting, etc), using a learned generative model as a Bayesian prior.
Such problems can typically be solved by gradient-based inference methods such as Variational Inference (VI), and Langevin or Hamiltonian Monte-Carlo (HMC), but all these methods have in common that they only rely on the score of the generative model .
Recent work [Lim et al., 2020, Song and Ermon, 2019] has investigated various approaches to directly train a deep neural network (coined, a score network ) to estimate the score of a density function instead of the density itself.
One of the most promising and scalable approach relies on denoising score-matching, where the score network is trained using a denoising loss on an augmented dataset where various amounts of Gaussian noise is added to the data.
Not only has this approach already been demonstrated to reach convincing results in score estimation, but Song and Ermon [2020] for instance demonstrated that these score networks can be used to sample high resolution natural images through Langevin dynamics.
In this work, we develop the ﬁrst application of deep denoising score matching to Magnetic Resonance
Image (MRI) reconstruction and its use for Uncertainty Quantiﬁcation (UQ) in the context of imaging inverse problems.
Related Works
Score Matching Score Matching was introduced by Hyvärinen [2005].
It was later revisited in Vincent [2011] and Alain and Bengio [2013] in the context of Deep Learning and Denoising
Auto-Encoders (DAE) in particular.
Recently, Song and Ermon [2019] applied these ﬁndings to images showing the potential of this approach combined with sampling techniques.
They reﬁned their work in Song and Ermon [2020] to take it to high dimension and provide some techniques for hyper-parameter setting.
Plug-and-Play Priors
The idea of using a denoising model as a prior for inverse problems solving was introduced by Venkatakrishnan et al.
[2013] who used non-learned denoisers to replace the proximity operator in the ADMM algorithm
[Boyd, 2011].
More recently, Meinhardt et al.
[2017] have made use of learned denoiser networks in Plug-and-Play inverse problems solving.
Annealed
HMC sampling can be seen as a fuzzy-version of these approaches with theoretical grounds.
Deep Uncertainty Quantiﬁcation in Inverse Problems
The idea of using generative models for inverse problems such as data imputation and denoising was for instance suggested in an early VAE work [Rezende et al., 2014].
In this approach, a generative network is ﬁrst trained on high quality data, and inference is then performed in the latent space of the model typically using VI.
The main advantage comes from the reduced dimensionality of the parameter space.
Recent examples of this approach include [Wu et al., 2018, Böhm et al., 2019], and an application to MRI can for instance be found in [Edupuganti et al., 2019].
An important aspect of this approach is that it relies on the known likelihood of observations.
A different approach was proposed in Adler and Öktem [2018] based on a conditional Wasserstein
Generative Adversarial Network (cWGAN).
In this formulation, a generative model is trained to sample high quality images conditioned on degraded observations.
At test time, independent samples from the posterior distribution induced by the GAN are obtaind by simply sampling different latent space variables.
Similar frameworks based on other conditional generative models have alse been proposed, such as [Denker et al., 2020] using conditional Normalizing Flows, or [Tonolini et al., 2019] using conditional VAEs.
While this conditional generative model approach allows for fast inference, thanks to ancestral sampling, it is worth noting that none of these models include an explicit data consistency step, nor do they make use of a test-time likelihood.
As a direct consequence, for instance, these models cannot be used on data with different noise levels as seen during training.
Deep Denoising Score Matching for Posterior Inference
Our main objective in this work is to perform probabilistic inference over a Bayesian model typically represented as
[Equation]
where y are some measurements and the posterior p ( x | y ) is the distribution of possible solutions x compatible with observations and p ( x ) a prior knowledge.
The problem-speciﬁc likelihood p ( y | x ) encodes the forward process of the model and accounts for observational noise, while the prior p ( x ) encapsulates any a priori information we have on the solution of the problem.
Our goal for UQ is to sample solutions x belonging to that posterior.
Multiple inference techniques can be leveraged for sampling from this posterior p ( x | y ) , but for high-dimensional problems modern techniques rely on gradient-based methods, including Variational Inference [Hoffman et al., 2013] and Langevin Diffusion, or HMC [Neal, 2011].
All of these techniques have in common that they only require having access to the score ∇ x log p ( x ) of the target distribution.
We illustrate in particular in
Appendix B how Metropolis Hastings-calibrated HMC can be conducted with only knowledge of the score.
Twoterms will enter in the score of the posterior distribution in Equation 1, the score of the likelihood, and the score of the prior.
In many problems, such as in the MRI problem presented later in this work, the likelihood score can be derived analytically, only the score of the prior remains unknown but can be learned from data by score matching.
Deep Denoising Score Matching As originally identiﬁed by Vincent [2011] and Alain and Bengio [2013], the score of a given target distribution P can be modeled using a DAE, i.e. by introducing an auto-encoding function r : R n × R (cid:55)→ R n trained to reconstruct under an (cid:96) 2 loss a true x ∼ P given a noisy version x (cid:48)
= x + n with n ∼ N (0 , σ 2 I ) .
An optimal denoiser r (cid:63) would then be achieved for:
[Equation]
where p σ 2 = p ∗ N (0 , σ 2 )
In other words, the optimal denoiser is closely related to the score we wish to learn and when the noise variance σ 2 tends to zero, should exactly match the score of the target density.
In practice to learn this score efﬁciently, we adopt the residual noise-conditional denoising score matching technique proposed by Lim et al.
[2020], and train a model to minimize the following Denoising Score Matching loss:
[Equation]
In this formulation, the network r θ ( x , σ ) is now directly modeling the score ∇ x log p σ 2 ( x ) of the
Gaussian-convolved target distribution.
Annealed Hamiltonian-Monte Carlo Sampling Given the noise-conditional neural scores learned with the procedure described above, it is now possible to use a variety of inference methods to access the Bayesian posterior.
In this work, we adopt an annealed HMC procedure which provides an efﬁcient way to obtain parallel batches of independent samples from the target posterior despite the high dimensionality of the problem.
This is a sampling procedure closely related to the Annealed Langevin
Diffusion proposed in [Song and Ermon, 2019, 2020], but beneﬁts from the faster Hamiltonian dynamics and Metropolis-Hastings calibration (see Appendix B).
To build our procedure, we consider a Gaussian-convolved version of our target density:
[Equation]
where σ 2 plays the role of the inverse temperature found in classical annealing.
The likelihood can be obtained analytically, and in the case of a Gaussian likelihood takes the following form:
[Equation]
) + cst , where σ 2 n is the noise variance in the measurements.
As for the prior term, the noise-conditional score network introduced above already models the score of log p σ 2 .
This distribution is gradually annealed to low temperatures and the chain progressively moves towards a point in the target distribution.
Application to Bayesian Inverse Problems
The MRI Problem Magnetic Resonance Imaging (MRI) is a non-invasive modality used to probe soft tissues.
Compressed sensing, introduced for MRI by Lustig et al.
[2007] is used to reduce its signiﬁcant acquisition time, and recently, deep learning approaches(see for example [Schlemper et al., 2018, Hammernik et al., 2018, Pezzotti et al., 2020]) have been shown to perform extremely well on the reconstruction problem.
The idealized reconstruction problem is usually formalized in the following way, for the single coil setting, with uniform acquisition:
[Equation]
where y is the acquired Fourier coefﬁcients, also called the k-space data, M Ω is a mask, F is the classical 2D Fourier transform, x is the real anatomical image, and n ∼ N (0 , σ 2 n ) is measurement noise (we set σ n = 0 .
1 in our experiments).
The data we consider is the single coil data with the
Proton-Density (PD) contrast from the fastMRI dataset introduced in [Zbontar et al., 2018].
For comparison with deep learning approaches, we use the state-of-the-art Primal-Dual net enhanced with a U-net for image correction (UPDNet) from [Ramzi et al., 2020].
The under-sampling was done retrospectively using an acceleration factor of 4 and a random mask as described for knee images in [Zbontar et al., 2018].
Results We use a simple residual U-net inspired by Ronneberger et al.
[2015] with ResNet building blocks from [He et al., 2016] as our score network.
In order to promote the regularity of the learned scores, we regularize the spectral norm of each convolutional layer, and we ﬁnd that setting the spectral norm to (cid:39) 2 yields the best results.
More details about the network architecture and training can be found in appendix A. We then sample from the posterior following our annealed HMC procedure down to relatively low temperatures, and apply one ﬁnal denoising step on the last sample from the chain, following the Expected Denoised Sample (EDS) scheme detailed in [Jolicoeur-Martineau et al., 2020].
Figure 1 compares samples from the MRI posterior to the ground truth (top left), zero-ﬁlled image (top second left), and UPDNet reconstruction (top center).
We see that although individual samples carry slightly less details than the neural network reconstruction, conﬁdence in any particular part of the image can be gauged by looking for variability or stability across multiple independent posterior samples.
We highlight in particular the red region, where posterior samples show signiﬁcant variability, indicating that this part of the image is poorly constrained by data.
In contrast, a direct neural network reconstruction (top center) does not match the ground truth in that region, and does not provide an estimate of uncertainty which may lead a physician to misinterpret the image.
Conclusions and Discussions
We have presented in this work the ﬁrst (to the best of our knowledge) instance of a framework for Bayesian inverse problems based on Deep Denoising Score Matching and applied to MRI reconstruction and Uncertainty Quantiﬁcation.
We illustrated the merits of this approach on an MRI example where the ability to sample from the full posterior highlights what features present in a reconstruction are not actually constrained by data.
We also showed that the Deep Denoising Score Matching approach can effortlessly be used for dimensions as high as 320 × 320 × 2 .
It is therefore much more scalable than GAN- or VAE-based models.
We note however that ﬁnding an optimal HMC annealing schedule and temperature-adaptive step-size proved difﬁcult.
In general we did not manage to lower the temperature below a certain level, prompting us to resort to a denoising step on the last samples from our HMC chains.
This is a direction for further research.
Acknowledgments and Disclosure of Funding
We acknowledge the ﬁnancial support of the Cross-Disciplinary Program on Numerical Simulation of
CEA for the project entitled SILICOSMIC.
We also acknowledge the French Institute of development and ressources in scientiﬁc computing (IDRIS) for their AI program allowing us to use the Jean Zay supercomputer’s GPU partitions.
A MRI experiments A.1 Network architecture
The network is a 3-scale U-net with residual blocks composed of 3 convolutions followed by a batch normalisation.
Each batch normalisation is followed by a ReLU non-linearity except the last one.
A projection is used for the input of residual blocks whose number of input channels is not the same as that of output channels.
Each scale uses 2 residual blocks for the downsampling path and 2 residual blocks for the upsampling path.
Downsampling is performed via average pooling and upsampling is performed via up-convolution (as designed in [Ronneberger et al., 2015]), in order to avoid checkerboard artefacts.
We use the following sequence of number of channels: [32, 64, 128].
In order to deal with MRI images, that are complex in nature, we concatenate the real and imaginary part of the image at the input of the network, forming effectively a 2-channel 320 ×image.
For input noise level conditioning, we concatenate a noise standard deviation map to the input and in the lower scale of the network.
Following the recommendation of Song and Ermon [2020], we also divide the output of the network by the absolute value of the noise power.
Finally, we use Spectral Norm regularisation from [Yoshida and Miyato, 2017] in order to make sure the score does not take inconsistent values in low-density regions.
The Spectral Norm regularisation indeed forces the spectral norm (maximum eigenvalue) of each layer to a certain value which in turn lowers the Lipschitz constant of the network, preventing it to blow in unseen regions.
The inﬂuence of the
Spectral Norm has been studied in the case of Plug-and-Play approaches by Ryu et al.
A.2 Network training
We use the Adam optimizer for network training, with a learning rate of 10 − 4 .
We add a white
Gaussian noise of variance σ 2 s to the images scaled by a factor of 10 6 .
σ s is drawn on-the-ﬂy from a
Gaussian distribution of variance s = 50 .
This means that at training time, the standard deviation of the noise can be negative, following the recommendation of Lim et al.
[2020] to go from extrapolation to interpolation.
A.3 HMC procedure
The HMC procedure starts from a zero-ﬁlled reconstruction of the image with added white Gaussian noise of variance σ init = 100 .
For the reduction of noise standard deviation at reduction step, we use a factor γ = 0 . 995 , following the recommendation of Song and Ermon [2020].
We take a step size α dependent of the sampling temperature at step i , α = (cid:15) (cid:16) [Table]
A.4 Zoom on Reconstruction
The zoom on this part of the reconstruction shown in Figure 2, illustrates a failure of the Neural network reconstruction.
It indeed produces a sharp artifact that might hinder the examiner’s conclusions.
The bayesian sampling allows us to see that this region is not necessarily very sharp.
B Score-Based Hamiltonian Monte Carlo
We used HMC for sampling from the posterior distribution.
The sampling procedure follows the leapfrog integrator presented bellow:
[Equation]
where α is the step size, m the momentum and M is a preconditioning matrix taking into account the space metric, but in our case the identity matrix.
The gradient of the log posterior is computed as ∇ x log p ( x t | y )
= ∇ x log p ( y | x t )
+ ∇ x log p ( x t ) , using the inverse Bayes theorem in equation 1.
It is worth noting that the discretized HMC with ﬁnal step size (cid:15) is not guaranteed to converge to the desired target distribution.
This is the motivation for the Metropolis-Hastings algorithm
[Metropolis et al., 1953, Hastings, 1970], which corrects for the discretization errors with a correction step.
Concretely, a new step is proposed using Equation 6, but it is accepted with probability of the form min { 1 , p ( x ∗ ) q ( x n | x ∗ ) /p
( x n ) q ( x ∗ | x n ) } where x n is the last accepted point in the chain, x ∗ is the proposed update, p is the target density, and q is the proposal corresponding to taking a step of
Equation 6.
In the context of a score-only approach, we do not have access to p and therefore cannot compute directly this acceptance probability.
We can nonetheless estimate from the scores themselves by using a path integral:
[Equation]
This integral can be evaluated numerically to arbitrary precision (at the cost of additional score evaluations), but we ﬁnd excellent agreement with the theoretical value on toy model densities like the two-moons distribution with a simple 4 points Simpson integration rule.
This illustrates that for gradient-based MCMC samplers, which also includes Metropolis Adjusted
Langevin Algorithm, these algorithms can fully be implemented with only knowledge of the score.