Compressed Sensing with Invertible Generative Models and Dependent Noise
Abstract
We study image inverse problems with invertible generative priors, speciﬁcally normalizing ﬂow models.
Our formulation views the solution as the maximum a posteriori (MAP) estimate of the image given the measurements.
Our general formulation allows for any differentiable noise model with long-range dependencies as well as non-linear differentiable forward operators.
We establish theoretical recovery guarantees for denoising and compressed sensing under our framework.
We also empirically validate our method on various inverse problems including 1-bit compressed sensing and denoising with highly structured noise patterns
Introduction
Inverse problems seek to reconstruct an unknown signal from observations (or measurements ), which are produced by some process that transforms the original signal.
Because such processes are often lossy and noisy, inverse problems are typically formulated as reconstructing x from its measurements
[Equation]
where f is a known deterministic forward operator and δ is an additive noise which may have a complex structure itself.
An impressively wide range of applications can be posed under this formulation with an appropriate choice of f and δ , such as compressed sensing [1, 2], computed tomography [3], magnetic resonance imaging (MRI) [4], and phase retrieval [5, 6].
In general for a non-invertible forward operator f , there can be potentially inﬁnitely many signals that match given observations.
Thus the recovery algorithm must critically rely on a priori knowledge about the original signal to ﬁnd the most plausible solution among them.
Sparsity has classically been a very inﬂuential structural prior for various inverse problems [1, 2, 7].
Alternatively, recent approaches introduced deep generative models as a powerful signal prior, showing signiﬁcant gains in reconstruction quality compared to sparsity priors [8–10].
Our contributions
• We present a general formulation for obtaining maximum a posteriori (MAP) estimation reconstructions for dependent noise and general forward operators.
Notably, our method can leverage generative models for both the original image and the noise.
• We empirically show that our method achieves excellent reconstruction in the presence of structured non-Gaussian noise.
We also demonstrate the efﬁcacy of our method on inverse problems with non-linear forward operators.
• We provide the initial theoretical characterization of likelihood-based priors for two common linear inverse problems: denoising and compressed sensing.
Notably, unlike Bora et al.
[8], we make no structural assumptions on the generator and provide bounds that depend on the likelihood of the prior.
Background
Invertible Generative Models.
Invertible generative models, also known as normalizing ﬂows , are a class of likelihood-based generative models that approximate complex distributions by mapping a simple noise (e.g. standard Gaussian) through an invertible function G [11, 12].
Since G is invertible, change of variables formula allows us to efﬁciently compute the log-density of x as log p ( x ) = log p ( z ) + log | det J G − 1 ( x ) | , where J G − 1 is the Jacobian of G − 1 .
Thus likelihood computation is straightforward as long as the log-determinant term can be efﬁciently evaluated.
Importantly, the invertibility of G guarantees that it has an unrestricted range and can generate samples that are out-of-distribution, albeit at lower probability.
This is a key distinction from a GAN, whose generator has a restricted range and cannot even represent out-of-distribution samples.
Related Work.
While vast literature exists on various inverse problems, the idea of using a deep generative prior was introduced relatively recently by Bora et al.
In that work, the generator from a GAN [13] was shown to be an effective prior for compressed sensing.
Several recent studies have investigated different ways to extend this result.
For example, Dhar et al.
[14] considered allowing sparse deviations on the generator’s output and Van Veen et al.
[10] showed that the generator can be replaced by an untrained Deep Image Prior [15] for imaging tasks.
More recently, Wu et al.
[16] applied techniques from meta-learning to improve the reconstruction speed, and Ardizzone et al.
[17] showed that one can implicitly learn the inverse process by modelling the forward process with a ﬂow model, .
Asim et al.
[9] proposed to use a normalizing ﬂow prior and reported excellent reconstruction performance, especially on out-of-distribution images.
Our Method
Notations and Setup.
(cid:107) · (cid:107) denotes (cid:96) 2 norm, and (cid:12) denotes element-wise multiplication.
We assume a differentiable forward operator f and a noise distribution p ∆ (which itself can be a generative model).
Thus an observation is generated via y = f
( x ) + δ where x ∼ p and δ ∼ p ∆ , where p ( x ) is assumed to be a ﬂow model deﬁned by an invertible mapping G .
3.1 MAP Formulation
For a given observation y , our method tries to recover x as the MAP estimate of the conditional distribution p ( x | y ) where p ( x ) is given by a ﬂow model deﬁned by an invertible mapping G :
[Equation]
where L MAP ( x ; y )
= − log p ∆ ( y − f ( x )) − log p ( x ) .
Recalling that the generative procedure for the ﬂow model is z ∼ N ( 0 , I ) , x = G ( z ) , we arrive at the following loss function:
[Equation]
Model Smoothing.
Since our objective Equation (2) depends on the density p ( x ) given by the ﬂow model, our recovery of x depends heavily on the quality of density estimates from the model.
Unfortunately likelihood-based models exhibit counter-intuitive properties, such as assigning higher density on out-of-distribution examples.
Thus we introduce the smoothing parameter β ≥that controls the degree to which we rely on p ( x ) :
[Equation]
3.2 Prior Work
This paper generalizes the work of Bora et al.
[8] and Asim et al.
[9], so we describe the methods proposed in those papers in detail.
Importantly, we show that their approaches are special cases of our MAP formulation.
Note that both papers considered linear inverse problems, i.e. f ( x ) =
A x .
GAN Prior: [8] considers the loss L Bora ( z ; y )
= (cid:107) y − AG ( z ) (cid:107)
+ λ (cid:107) z (cid:107) 2 , which tries to project the input y onto the range of the generator G with (cid:96) 2 regularization on the latent variable.
We note that G here is a GAN and is not invertible as it maps a noise to a higher-dimensional vector.
Thus the image we wish to recover may not be in the range of G .
This inability to represent arbitrary images leads to suboptimal reconstructions, as conﬁrmed by Asim et al.
[9] and our experiments.
Flow Prior: [9] consider the loss L ( z ; y )
= (cid:107) y − AG ( z ) (cid:107) 2 + γ (cid:107)
z (cid:107) , which tries to simultaneously match the observation and regularize z to be close to zero.
This is a special case of our MAP loss for isotropic Gaussian noise δ ∼ N ( 0 , γ I ) with a volume-preserving ﬂow model p (i.e. the logdeterminant term is constant).
Continuing from eq.
(2), L G ( z ; y ) =
12 σ 2 (cid:107) y − AG ( z ) (cid:107)
+ (cid:107) z (cid:107) 2 + C, where C absorbs the constant log-determinant term in log p ( x )
= log p ( z ) + log (cid:107) det J G − 1 ( x ) (cid:107) .
This recovers the (cid:96) 2 -regularized version of L Asim .
Theoretical Analysis
Here we present a preliminary theoretical analysis for denoising and compressed sensing.
Unlike most prior work, we take a probabilistic approach and avoid speciﬁc structural assumptions on x , such as sparsity or being generated from a low-dimensional Gaussian prior.
Detailed proofs can be found in Appendix A.
Recovery Guarantees for Denoising.
Suppose we observe y = x ∗ + δ with Gaussian noise δ ∼ N ( 0 , σ 2 I ) with (cid:107) δ (cid:107)
=
r .
We perform MAP inference by minimizing the following loss with gradient descent: L MAP ( x ) = 12 σ 2 (cid:107)
y − x (cid:107) 2 + q ( x ) , where we write q ( x ) (cid:44) − log p ( x ) .
Theorem 4.1.
Let x ∗ be a local maximum of the model p ( x ) and y
= x ∗ + δ be the noisy observation.
Assume that q satisﬁes local strong convexity within the ball around x ∗ deﬁned as
B dr ( x ∗ ) (cid:44) (cid:8) x ∈ R d : (cid:107) x − x ∗ (cid:107) ≤ r (cid:9) , i.e. the Hessian of q satisﬁes H q ( x ) (cid:23) µI ∀ x ∈ B dr ( x ∗ ) for some µ > 0 .
Then gradient descent starting from y on the loss function L MAP ( x ) converges to ¯ x , a local minimizer of L MAP ( x ) , that satisﬁes (cid:107) ¯ x − x ∗ (cid:107) ≤ 1 µσ 2 +1 (cid:107) δ (cid:107) .
This theorem shows that a well-conditioned model with large µ leads to better denoising and conﬁrms that our MAP formulation encourages reconstructions with high density.
Recovery Guarantees for Compressed Sensing.
Here we show a worst-case bound for noisy compressed sensing where we observe y =
A x ∗ + δ with noise level (cid:15) = (cid:107) δ (cid:107) and A ∈ R m × d has entries drawn i.i.d.
from N (0 , 1 /m ) .
Note that the following constrained minimization problem is equivalent to the original unconstrained minimization of eq.
(3) for some β via Lagrange multipliers.
Theorem 4.2.
For a given x ∗ ∈ R d and ρ = log p ( x ∗ ) , deﬁne S ( ρ )
=
{ x | log p ( x ) ≥ ρ } to be the set of images with density higher than the ground truth image x ∗ .
Recall that the observation for x ∗ is y =
A x ∗ + δ .
When we perform the MAP inference by solving the following optimization problem,
[Equation]
we have
[Equation]
, where w ( · ) is the
Gaussian mean width, and the expectation is over the randomness of A .
We note that in high dimensions, Gaussian mean width scales with the set’s normalized surface area.
Since S ( ρ ) becomes smaller as ρ increases, this suggests the above error bound is tighter for ground truth images x ∗ with higher density.
Experiments
G r ound T r u t h
[Figure]
N o i s y
I npu
t
O u r s A s i m e t a l . B o r a e t a l . B M 3 D (a) Result of removing MNIST noise from CelebA-HQ faces.
No- tice that without any understanding of the complex noise structure, baseline methods fail to produce good reconstructions.
[Figure]
Our experiments are designed to evaluate our MAP formulation in two scenarios: (1) complex noise with dependencies and (2) non-linear forward operator.
A detailed description of the datasets, models and hyperparameters are provided in the appendix.
Baseline Methods.
We compare our approach to the methods of [8] and [9], as they are two recently proposed approaches that use deep generative prior on inverse problems.
When applicable, we also compare against BM3D [22] and LASSO [23] with Discrete Cosine Transform basis as appropriate.
We point out that the baselines methods are not designed to make use of the noise distribution, whereas our method does utilize it.
Thus, the experiments are not meant to be taken as direct comparisons, but rather as an empirical evidence that the MAP formulation beneﬁts from knowing the noise structure.
5.1 Results
Denoising MNIST Digits.
The measurement process is y = 0 .
5 · x + δ MNIST , where δ MNIST represents MNIST digits added at different locations and color channels.
Each digit itself comes from a ﬂow model trained on the MNIST dataset.
As shown in Figure 1a, our method successfully removes MNIST noise.
While the method of Bora et al.
[8] also removes MNIST digits because its outputs are forced to be in the range of the DCGAN used, the reconstructions are far from the ground truth – particularly on out-of-distribution samples that are not human faces.
Noisy Compressed Sensing.
Now we consider the measurement process y
=
A x + δ sine where
A ∈ R m × d is a random Gaussian measurement matrix and δ has positive mean with variance that follows a sinusoidal pattern (details provided in Appendix B).
Figure 1b shows that our method is able to make a better use of additional measurements.
Interestingly in Figure 2c, all three methods with deep generative prior produced plausible human faces.
However, the reconstructions from Asim et al.
[9] and Bora et al.
[8] signiﬁcantly differ from the ground truth images.
We posit that this is due to the implicit Gaussian noise assumption made by the two methods, again showing the beneﬁts of explicitly incorporating the knowledge of noise distribution.
[Figure]
(a) Result of 1-bit noisy compressed sensing at different measurement counts. Our method achieves the same reconstruction performance using to measurements
G r ound T r u t h
O u r s
A s i m e t a l . B o r a e t a l .
G r ound T r u t h
O u r s
A s i m e t a l .
B o r a e t a l .
L A SS O
[Figure]
measurements.
1-bit Compressed Sensing.
This task considers a combination of a non-linear forward operator as well as a non-Gaussian noise.
The measurement process is the quantized version of noisy compressed sensing, i.e. y = sign ( A x ) + δ sine .
Note that this the most extreme form of quantized compressed sensing, since y only contains the noisy sign { +1 , − 1 } of the measurements.
Because the gradient of sign function is zero everywhere, we use Straight-Through Estimator [24] for backpropagation.
See Figure 2a and Figure 2b for a comparison of our method to the baselines at varying numbers of measurements.
Conclusion
We propose a novel method that generalized [9] to solve inverse problems for general differentiable forward operators and structured noise.
The power of our approach stems from the ﬂexiblity of invertible generative models which can be combined in a modular way to solve MAP inverse problems in very general settings, as we demonstrate.
The central theoretical question that remains open is to analyze the optimization problem we formulated.
In this paper we empirically minimize this loss using gradient descent, but some theoretical guarantees would be desirable, possibly under assumptions, e.g. random weights following the framework of [25].
Acknowledgments and Disclosure of Funding
This research has been supported by NSF Grants CCF 1763702,1934932, AF 1901292, 2008710, 2019844 research gifts by Western Digital, WNCG IAP, computing resources from TACC and the
Archie Straiton Fellowship.
A Omitted Proofs A.1 Proof for Denoising
Proof of Theorem 4.1.
We ﬁrst show that gradient descent with sufﬁciently small learning rate will converge to ¯ x , the locally-optimal minimizer of L MAP ( x ) = 1 σ 2 (cid:107) x − y (cid:107) 2 + q ( x ) , where we write q ( x ) (cid:44) − log p ( x ) (we subsume the constant coefﬁcient 12 into 1 σ 2 without loss of generality).
Notice in the ball B dr ( x ∗ ) :
= (cid:8) x ∈ R d
| (cid:107) x − x ∗ (cid:107) ≤ r (cid:9) , L is (cid:0) µ + 1 σ 2 (cid:1) strongly-convex.
We next show there is a stationary point ¯ x ∈ B dr ( x ∗ ) of L ( x ) .
[Equation]
From strong convexity of q , (cid:104)∇ q (¯ x ) − ∇ q ( x ∗ ) , ¯ x − x ∗ (cid:105) ≥ µ (cid:107) ¯ x − x ∗ (cid:107) 2 .
Thus
[Equation]
Finally, by Cauchy-Schwartz inequality (cid:104)
y − x ∗ , ¯ x − x ∗ (cid:105) ≤ (cid:107)
y − x ∗ (cid:107) · (cid:107) ¯ x − x ∗ (cid:107) .
So we get (cid:107) ¯ x − x
[Equation] (cid:107) ≤
[Equation]
+
[Equation]
2 (cid:107) y − x ∗ (cid:107) ≤ (cid:107) δ (cid:107) ≤ r , i.e., ¯ x ∈ B dr ( x ∗ ) .
Notice L is (cid:0) µ + 1 σ 2 (cid:1) strongly-convex in B dr ( x ∗ ) , which contains the stationary point ¯ x .
Therefore ¯ x is a local minimizer of L ( x ) .
Also note that we implicitly require q to be twice differentiable, meaning in a compact set B dr ( x ∗ ) its smoothness is upper bounded by a constant M .
Thus gradient descent starting from y ∈ B dr ( x ∗ ) with learning rate smaller than 1 M will converge to ¯ x without leaving the (convex) set B dr ( x ∗ ) .
A.2 Proof for Compressed sensing
Deﬁnition A.1 (Gaussian mean width) .
The Gaussian mean width of a set K ⊂ R d is deﬁned as:
[Equation]
where M ( K ) = { x − y : x , y ∈ K } is the Minkowski sum of K and − K .
Intuitively, Gaussian mean width measures the complexity of the set K .
Theorem A.2 (Adapted from Theorem 6.1 by Vershynin [26]) .
Let K ⊂ R d be an arbitrary bounded set, and A ∈ R m × d be a random matrix with its entries sampled iid from Gaussian distribution
N (0 , 1) , observation y =
A x ∗ + δ ,
where 1 √ m (cid:107) δ (cid:107) 2 = (cid:15) , and x ∗ is unknown.
Choose ˆ x to be any vector satisfying ˆ x ∈ K and 1 √ m (cid:107)
A ˆ x − y (cid:107) 2 ≤ (cid:15).
Then
[Equation]
One consequence of this theorem is that m =
Θ( w ( K ) 2 ) measurements are sufﬁcient to guarantee small recovery error.
Remark:
The statement of Theorem A.2 is only different from the original Theorem 6.1 in [26] where we replaced the (cid:96) 1 norm assumption by a stronger (cid:96) 2 norm assumption, and therefore is a slightly weaker version of the original theorem.
As indicated in [26], this is still valid.
To see this, notice that for any ˆ δ ∈ R m , we have (cid:13)(cid:13)(cid:13)
ˆ δ (cid:13)(cid:13)(cid:13) 1 ≤ √ m (cid:13)(cid:13)(cid:13) ˆ δ (cid:13)(cid:13)(cid:13) 2 .
So the (cid:96) 2 bound 1 √ m (cid:13)(cid:13)(cid:13)
ˆ δ (cid:13)(cid:13)(cid:13) 2 ≤ (cid:15) gives
[Equation]
Theorem A.3 (Restatement of Theorem 4.2) .
For a given x ∗ and q = log p ( x ∗ ) , deﬁne S ( q ) =
{ x | log p ( x ) ≥ q } .
Recall that the observation for x ∗ is y =
A x ∗ + δ where A ∈ R m × d has entries drawn i.i.d.
from N (0 , 1 /m )
and the noise level is (cid:15) =
(cid:107) δ (cid:107) 2 .
When we perform the MAP inference by solving the following problem,
[Equation]
we have:
[Equation]
where the expectation is over the randomness of A .
Proof.
Since E (cid:107) ¯ x − x
[Equation]
(cid:107)
[Equation]
≤ E
[Equation] sup [Equation] , the proof directly follows from Theorem A.2 by choosing the set K to be S ( q ) .
We only need to verify that ¯ x is indeed inside S ( q ) .
Notice that both ¯ x and x ∗ are in the feasible set { x : (cid:107) A x − y (cid:107) 2 ≤ (cid:15) } .
Since ¯ x maximizes log p ( x ) within this set, clearly its density should be at least that of x ∗ , i.e. log p (¯ x ) ≥ log p ( x ∗ ) = q .
Thus ¯ x ∈ S ( q ) .
Note that in our setting the variance of A ’s entries is set to be 1 /m instead of 1 from Theorem A.2.
Since the Gaussian mean width w ( S ( q )) is invariant to the scaling of A , we correct the scaling on the noise term and obtain E (cid:104) sup x ∈ S ( q ) (cid:107) ¯ x − x ∗ (cid:107) 2 (cid:105) ≤ √ 8 π ( w ( S ( q )) √ m + (cid:15) ) .
Remark:
Here we show the recovery guarantee when we optimize over the constrained version as in eq.
Note that for a suitable constant β , the constrained version min x : (cid:107) A x − y (cid:107) 2 ≤ (cid:15) {− log p ( x ) } is equivalent to min x {− log p ( x ) + β (cid:107)
A x − y (cid:107) 22 } .
B Additional Experimental Results and Details
Here we include experimental results and details not included in the main text.
Across all the experiments, we individually tuned the hyperparameters for each method.
B.1 Experimental Details
Dataset For MNIST, we used the default split of 60,000 training images and 10,000 test images of [27].
For CelebA-HQ, we used the split of 27,000 training images and 3,000 test images as provided by [28].
During evaluation, the following Python script was used to select 1000 MNIST images and 100 CelebA-HQ images from their respective test sets: np.random
Note that CelebA-HQ images were further resized to 64 × 64 resolution.
Noise Distribution For the sinusoidal noise used in the experiments, the standard deviation of the k -th pixel/row is calculated as: σ k = 0 .
1 · (cid:0) exp (cid:0) sin(2 π · k 16 ) (cid:1) − 1 (cid:1) /
( e − 1) , clamped to be in range [0 . 001 , 1] .
For Figure 4b, we used vary the coefﬁcient 0 .
1 to values in { 0 . 05 , 0 . 1 , 0 . 2 , 0 . 3 , 0 . 4 } .
For the radial noise used in the additional experiment below, the standard deviation of each pixel with (cid:96) 2 distance is d from the center pixel (31, 31) is computed as: σ k = 0 . 1 · exp( − 0 .
005 · d 2 ) , clamped to be in range [0 . 001 , 1000] .
B.2 Additional Result: Removing R ADIAL Noise
Consider the measurement process y
= x + δ radial , where each pixel follows a Gaussian distribution, but with variance that decays exponentially in distance to the center point.
For a pixel whose (cid:96)distance to the center pixel is d , the standard deviation is computed as σ ( d ) =
exp (cid:0) − 0 .
005 · d 2 (cid:1) .
See
Figure 3 and Figure 4a for reconstructions as well as PSNR plot comparing the methods considered.
G r ound
T r u t h
[Figure]
N o i s y
I npu
t
O u r s
A s i m e t a l .
B o r a e t a l .
[Figure]
M
[Figure]
[Figure]
[Figure]
B.3 Additional Result: 1-bit Compressed Sensing
Figure 4b shows the performance of each method at different noise scales for a ﬁxed number of measurements.
We observe that our method performs consistently better at all noise levels.
C Model Architecture and Hyperparameters
For the RealNVP models we trained, we used multiscale architecture as was done in [29], with residual networks and regularized weight normalization on convolutional layers.
Following [28], we used 5-bit color depth for the CelebA-HQ model.
Hyperparameters and samples from the models can be found in Table 1 and Figure [Figure]
Hyperparameter
CelebA-HQ
MNIST
[Table]
5e
[Table]
Batch
[Table]
Image
[Table] [Table]
depth
[Table]
bits
[Table]
bits
Number of
[Table]
Number of scales
Residual blocks per scale
Learning rate halved every 60 epochs
[Table]
Weightnorm regularization 1e −5e −