Uncertainty-Driven Adaptive Sampling via GANs
Abstract
We propose an adaptive sampling method for the linear model, driven by the uncertainty estimation with a generative adversarial network (GAN) model.
Speciﬁcally, given a forward observation model that provides partial measurements y about an unknown parameter x , we show how to build a GAN model to estimate its posterior p ( x | y ) .
We then leverage our approximate posterior to perform sequential adaptive sampling by actively selecting the measurement with the current maximal uncertainty.
We empirically demonstrate that our posterior estimate contracts rapidly towards the correct mode, while outperforming the state-of-the-art approaches even for other criteria for which they are speciﬁcally trained, such as PSNR or SSIM
Introduction
This work focuses on the following inverse problem, where we seek to recover a signal x ∈ C P from partial observations y ∈ C N obtained by subsampling a unitary transform matrix A ∈ C P × P :
[Equation]
where η ∈ C N is a noise vector;
ω ⊆
[ P ] := { 1 , . .
. , P } is an index set of allowable sampling locations with cardinality N ; and P ω is a diagonal matrix
such that ( P ω ) ii = 1 if i ∈ ω, 0 otherwise.
The compressed sensing (CS) problem above is classical [1, 2], where the literature revolves around using statistical or geometric priors to circumvent its ill-posed nature since N (cid:28) P .
In particular, the
Bayesian perspective on CS imposes prior distributions p ( x ) on a random vector x , aiming then at constructing a tractable posterior distribution p ( x | y ) which captures the probability for a sample x to have generated the observed data y [3, 4].
The geometric perspective in contrast leverages convex models based on the (cid:96) 1 -norm, and then seeks to obtain point estimates [1, 2].
Contemporary approaches to posterior modeling in the CS setting can be roughly split into two camps, with one computing point estimates of the mean and the variance, assuming an underlying Gaussian distribution, such as [5], and the other approximating the posterior with a generative adversarial network (GAN) or a variational autoencoder, such as [6, 7, 8].
These data-driven, deep learning-based methods have also signiﬁcantly improved practical reconstruction performance in CS problems [9, 10, 11, 12] as compared to the Bayesian or (cid:96) 1 -norm based estimates.
By data-driven, we mean that these approaches explicitly build underlying prior or posterior models using training data, as opposed to explicitly relying on carefully-constructed mathematical prior and geometric models, such as Laplace priors, the (cid:96) 1 -norm, or Besov smoothness.
The successes of these approaches have led the broader medical imaging community to re-think the problem of optimizing the undersampling patterns also in a data-driven, adaptive (closed-loop) fashion [5, 13, 14, 15], instead of considering ﬁxed (open-loop), random undersampling [1, 2, 16].
While there are earlier CS works that also use a data-driven perspective, they still perform the sampling in an open-loop fashion [17, 18, 19].
There are also closed-loop, but not data-driven sequential CS approaches, [3, 20, 21], which rely on mathematical models and do not use no training data.
In this work, we argue that the practical beneﬁts of the generative methods have not been fully exploited in the CS setting.
To this end, we propose a closed-loop, data-driven, Generative Adaptive [Figure]
! <latexit sha1_base64="5cWOcuqDHtRYcwH44kbj8B5Cdh0=">AAAB7nicbVDLSgNBEOyNrxhfUY9eBoPgKeyKoN6CXjxGMA9IljA7mU2GzGOZmRXCko/w4kERr36PN//G2WQPmljQUFR1090VJZwZ6/vfXmltfWNzq7xd2dnd2z+oHh61jUo1oS2iuNLdCBvKmaQtyyyn3URTLCJOO9HkLvc7T1QbpuSjnSY0FHgkWcwItk7q9JWgI1wZVGt+3Z8DrZKgIDUo0BxUv/pDRVJBpSUcG9ML/MSGGdaWEU5nlX5qaILJBI9oz1GJBTVhNj93hs6cMkSx0q6kRXP190SGhTFTEblOge3YLHu5+J/XS218HWZMJqmlkiwWxSlHVqH8dzRkmhLLp45gopm7FZEx1phYl1AeQrD88ippX9SDy/rNw2WtcVvEUYYTOIVzCOAKGnAPTWgBgQk8wyu8eYn34r17H4vWklfMHMMfeJ8/yoyPOw==</latexit>
Sampling (GAS) method based on the recently proposed approaches of [6, 7].
In contrast, the GAN approaches in CS so far has been only used to obtain point estimates for recovery [22, 23].
By leveraging the ﬂexible posterior modeling of GANs [6, 8] with their ability to exploit conditionals [24, 7], we develop a GAN model where the generator can simultaneously provide ( i ) a good image reconstruction, and ( ii ) uncertainty quantiﬁcation in the image domain, thus naturally leading to an end-to-end , sequential adaptive sampling approach.
Notation.
Throughout the text, we refer to random vectors as boldface straight letter x , y and vector as boldface italic, x , y .
We use capital boldface letters for matrices.
Methods 2.1 Reconstruction via conditional sampling
Given subsampled measurements via Equation (1), we consider two approaches to estimate p ( x | y ) .
In the ﬁrst approach, we obtain the posterior mean and its variance, as a summary estimate of the posterior distribution.
This is typically achieved via the problem of signal reconstruction:
[Equation]
where y ω speciﬁes the dependency of y on ω ; f θ refers to a deep neural network with parameters θ ; (ˆ x θ , ˆ s 2 θ ) denote the point estimates of the conditional mean and variance of the model, respectively.
Such models are typically trained using supervised learning by minimizing a Gaussian negative log-likelihood with heteroscedastic variance [25, 5].
In the second approach, we can have the ability draw samples of the posterior distribution:
[Equation]
where z is a random vector drawn from a simple distribution; G θ ( y ω , z ) is a deep generative model mapping the input to a sample from p ( x | y ω ) .
Such models typically rely on conditional generative adversarial networks [26, 6, 7] or variational autoencoders [27, 24].
2.2 Sampling optimization
The ideal sampling optimization algorithm would tailor the mask to each instance of x ∼ p ( x ) solving [Equation]
[Equation]
which is impossible since this requires using the unknown ground truth signal
x at testing time.
Two main approaches have been explored to address this problem.
Fixed (open-loop) sampling.
A majority of data-driven mask design approaches use ﬁxed masks [28, 29, 30, 31].
Indeed, the subsampling mask is built using training data and kept ﬁxed at inference time.
Formally, the problem of choosing the subsampling pattern corresponds to ﬁnding a subset ω that satisﬁes
[Equation]
where we are constrained with a maximal sampling budget N and want to ﬁnd a mask that minimizes a given loss function (cid:96) .
We estimate the expected loss using the empirical estimate with the training samples.
We let ˆ x θ ( y ω
=
P ω
A x ) denote an estimate of the mean, obtained either directly through
Equation (2) or by averaging on z for Equation (3), i.e., ˆ x θ ( y ω ) =
E z
[ˆ x θ ( y ω , z )] =
E z
[ G θ ( y ω , z )] .
Adaptive (closed-loop) sampling.
In contrast with Equation 5, adaptive sampling optimizes a heuristic (cid:96) available at test time.
For a ﬁxed, unknown data sample x , we use information of the previously obtained measurements y ω t − 1 to determine what should be acquired at time t .
The heuristic can be human crafted or learned from data, e.g. a score assigned by a neural network[5].
By conditioning on the previous samples and using the heuristic (cid:96) , we solve
[Equation]
where the variables v t typically refer to individual pixels to be sequentially observed from the reference image x ∼ p ( x ) .
Here we adopt a greedy approach, without trying to plan on a longer horizon.
This has been explored, for instance, in [32, 13, 14].
Such approaches rely on performing additional computation or training an additional network in a supervised fashion in order to guide the planning.
The work we present in the next section is based on a greedy approach, but could easily be combined with a long-term planning approach in future work.
2.3 GAS: Generative adaptive sampling
Departing from the idea of explicitly training a policy to guide adaptive sampling [32, 13, 5], GAS simply minimizes the variance of an approximate posterior.
This turns out to yield a strong sampling policy despite the networks involved being trained soley for conditional generation .
We start by modeling an approximate posterior p θ ( x | y ω ) on any mask ω by a conditional GAN [26, 7] or its Wasserstein counterpart [33, 6].
Given the trained model G θ ∗ ( y ω , · ) that gives samples from the approximate posterior p ( x | y ω ) , adaptive sampling follows naturally.
Starting at t = 0 from unconditional sampling, i.e. sampling from the prior, we iteratively sample the locations v t with the largest variance in the transformed domain A :
[Equation]
(7) where Var [ A x | y ω t ] denotes the point-wise variance on the vector
A x .
It is estimated at each stage by computing the empirical variance using samples from G θ ∗ ( y ω t , · ) .
The chosen element is then added in the sampling mask as ω t +1 = { ω t , v t } .
Experiment
We test the applicability of GAS on MNIST, where we explore subsampling in pixel ( A = I )
We use the same architecture and optimizer for both the Wasserstein GAN (WGAN) [34] and the Neural Conditioner (NC) [7].
The architectures used are a Unet [35] with residual blocks (ResUNet), used in [7], and a cascade of residual networks (c-ResNet) used in [5].
Both architecture are use dataconsistency layers [10], that integrate the observations into the ﬁnal image, preventing the networks from discarding them.
We generate undersampled images by dynamically sampling masks at random or by sampling from a Gaussian centered in the middle of the image (variable-density sampling), using various sampling rates between 0 .
5% and 20% sampling rate.
Throughout the experiments, we used 5 samples from the posterior to compute empirical posterior means and variances.
In addition to these baselines, we also trained reconstruction networks to further expand the comparison.
We refer the model as Recon.
Further details of the implementation are provided in the appendix, where.
we also included additional experiments on CIFAR10.
Our results are summarized in Table 1 below, and we compare the following sampling models.
• VDS: Variable-Density sampling.
The model is evaluated on the variable-density masks used during training, a common baseline in CS works [36].
• LBC: Learning-based compressive sensing [19].
A ﬁxed sampling mask is trained and applied on all testing data.
We use their stochastic version for scalability reasons [37].
• GAS: Our generative, uncertainty-driven adaptive compressed sensing approach.
In the sampling experiments, the baselines are evaluated on various sampling rates up to 20% sampling rates.
The LBC and adaptive methods are evaluated on the whole range from 0 to 20% .
The models are test on several sampling rates and sampling masks that they never saw during training.
We report the PSNR, SSIM [38] as well as the downstream classiﬁcation accuracy, using a small residual network trained on the original MNIST digits, and evaluated on a batch of samples from the conditional generator.
The ﬁnal classiﬁcation result is ensembled from this batch of predictions.
As we consider results at different sampling rates, we summarize the overall performance by computing the area under curve (AUC) of the different metrics.
We then compute the average and standard deviations of these test curves.
The curves, illustrating the performance of the models at different sampling rates, are provided in the appendix.
We see on Tables 1 that both GAN models strongly outperform the LBC baseline in the image domain.
In the Fourier domain, the performance gap is smaller, due to the more homogeneous distribution of the different spectra.
The WGAN c-ResNet adaptive is still matches the LBC performance [Table] but the
LBC is much more competitive in this setting and tends to outperform GAS.
This is likely due to the fact that the LBC mask is trained to minimize PSNR [Table] whereas our model uses the uncertainty in an unsupervised fashion.
We would also like to emphasize that the model learns to perform unconditional sampling despite never seeing an empty mask during training, as shown in the Appendix
[Table]
Outlook
We presented a novel approach, rooted in Bayesian modeling, for adaptive sampling.
We showed that, in line with recent advances [7], GANs performs well at the task of learning a large number of different conditional distributions.
Having been trained on a limited set of sampling rates and masks, the model is able to convincingly generalize to unseen masks and sampling rates, and is even able to perform unconditional sampling without having been explicitly trained on this task.
In addition, the samples from the conditional are able to provide meaningful uncertainty information, and was successfully leveraged as a greedy, adaptive sampling policy.
This success is due to the unique ability of generative models to capture complex, multimodal distributions and cheaply sample from them at inference time.
This could make them prime candidates for a more uncertainty-aware deep learning approach to inverse problems.
Our future work will demonstrate the applicability of this method to problems where sampling does not occur in image domain, such as MRI, with Fourier domain sampling. Acknowledgments
This work has received funding from the European Research Council (ERC) under the European
Union’s Horizon 2020 research and innovation program (grant agreement n ◦ 725594 - time-data)m from Hasler Foundation Program: Cyber Human Systems (project number 16066) and from the
Ofﬁce of Naval Research (ONR) (grant n ◦ N62909-17-1-2111).
A Implementation details
We carried out our experiments on the MNIST dataset.
The training set contains 60000 handwritten digits of size 28 × 28 .
We held out 10000 digits for validation purposes.
The test set consists of 10000 digits.
We zero-padded all images to 32 × 32 for convenience.
We ran all experiments on a
Tesla K40 GPU, using batch size 64 , and initialized the networks using the default Pytorch values.
A.1 Experiment setting
We experiment on two main generator architectures, a UNet[1] with residual blocks, called ResUNet for short, as used in [2].
The architecture used is detailed in Table 2 below.
We also used a cascade of residual networks, called c-ResNet [3], where one block of the architecture is given in Table 3.
The network was composed of 3 such blocks.
Our choice was motivated by the fact that encoder-decoder architectures are required by the high dimensional input of this type of conditional GAN.
Regarding the discriminator, we used a ResNet-like architecture, following [2, 4].
The architecture is reported in
Table 4.
The cWGAN models were trained with gradient penalty [4] and the generators are followed with a sigmoid to keep the range in [0 , 1] , and followed by a data-consistency layer [5, 6], which has the role of replacing the generated data with the observed ones at the last layer.
All models were trained using Adam [7] with learning rate 10 − 4 , β = (0 . 5 , 0 . 999) and weight decay 10 − 4 .
All models were trained for 300 epochs.
The images were undersampled using masks that were either drawn at random or drawn with a probability decaying away from the center.
The settings used are given in greater detail in the Table 5 below.
The Neural Conditioner (NC) [2] generator and discriminator were updated every other round.
In the WGAN case, we optimized the generator once, followed by four updates of the discriminator, as done in [8].
Similarly to [2], we found gradient regularization of the discriminator to be essential.
However, we obtained very poor training results using the authors proposed gradient penalty, and observed that the gradient penalty of [4] led to a stable training.
We used the gradient penalty with a weight of 10 , following [8] in the WGAN case, and a value of 1 for the NC, but did not observe signiﬁcant differences between these two weights.
We note in passing that contrarily to [2], we did not observe any particular effect of spectral normalization of the encoder part of the generator. Reconstruction networks.
Using the same architectures as the GAN generators (ResUNet and c-ResNet), we also trained two reconstruction networks by minimizing the (cid:96) 2 distance between ground truth images and reconstructed images.
Formally, given a training set { ( x i , y ω i ,i ) } ni =1 , the network f θ is trained by minimizing 1 n (cid:80) ni =1 (cid:107) x i − f θ ( y ω i ,i ) (cid:107) 22 . We used Adam [7] with learning rate 2 · 10 − 4 , β = (0 . 0 , 0 . 9) and weight decay = 10 − 4 . We trained the models for 200 epochs using a batch size 128 .
The masks used are the same as the GANs, described in Table 5 below.
A.2 Sampling experiments
After training, we carried out several rounds of evaluation for different types of sampling as reported in Table 1.
We hereafter detail the different setting we used.
As mentioned in the main body of the paper, we chose to estimate the posterior mean and variances using 5 samples from the posterior.
We found that 2 samples were sufﬁcient for the adaptive sampling with GAS to yield good results.
We also observed that more samples 20 to 100 made little improvement in the sampling quality or image quality, except that it suppressed the noise of the empirical posterior mean and variance a bit more efﬁciently.
For the baseline, variable-density sampling, we constructed both random and Gaussian masks using the settings of Table 5, but including also sampling rates that were never seen during training, including rates above 20% . [Table] [Table]
[Table] Regarding the LBC approach [10], there are two parameters to be adjusted, namely the number of training samples to be used at each stage of the greedy mask optimization [Table] and the size of candidate locations to be used.
We set these to 64 and 256 respectively.
The criterion used for the optimization was to select at each round the sample that maximally decreased the mean squared error in the empirical posterior mean.
The mask was then evaluated at different stages on the testing set and the results reported.
Our adaptive sampling procedure, GAS, was also averaged on 5 samples from the generator, and the optimization was carried out by selecting at each round the point with the largest empirical variance in the transformed domain. [Table]
B Additional results B.1 Additional MNIST results
We provide additional results that further illustrate the performance of our models and of the different baselines considered.
In Figure we show additional image domain results, using the NC model [2] using the ResUNet architecture.
We also present some quantitative results on Figure 4 are respectively summarized by AUC curves in Table 1. We then provide some images obtained with our reconstruction algorithms on Figure [Figure] [Figure]
B.2 CIFAR10 experiment
Finally, we provide results on CIFAR10, for which the pixelwise subsampling task is much more challenging.
We chose to use grayscale versions of the images.
Experiment setting.
For CIFAR10, we modiﬁed the training procedures slightly.
For the WGAN we use ExtraAdam [11], TTUR [12] with learning rates of lr G = 10 − 4 and lr D = 3 · 10 − 4 , one sided gradient penalty [13] and betas of (0 , 0 . 9) .
Number of epochs was set to 200 with the learning rates being halved every 50 steps.
Batchsize was set to 128 and we used 8 GPUs in total, yielding an effective batch size of 1024 .
For the NC we retain batch size 128 and the Adam settings, simply moving to 8 GPUs and increasing training time to 200 epochs.
Results.
Figure 6 and ﬁg. 7 showcase GAS results using a WGAN and NC respectively.
We evaluate
VDS and GAS on a subset of the test set of CIFAR10 composed of 1000 images and report the results on Table 6.
The results are detailed in the plots of Figure 8.
[Table]
Observe that while the model saw only masks going from 0 .
5% to 20% during training, it generalizes well up to 30% .
[Figure]