Q UANTIZED C OMPRESSED S ENSING WITH S CORE - B ASED G ENERATIVE M ODELS
Anonymous authors
Paper under double-blind review
A BSTRACT
We consider the general problem of recovering a high-dimensional signal from noisy quantized measurements.
Quantization, especially coarse quantization such as 1-bit sign measurements, leads to severe information loss and thus a good prior knowledge of the unknown signal is helpful for accurate recovery.
Motivated by the power of score-based generative models (SGM, also known as diffusion models) in capturing the rich structure of natural signals beyond simple sparsity, we propose an unsupervised data-driven approach called quantized compressed sensing with SGM (QCS-SGM), where the prior distribution is modeled by a pre-trained SGM.
To perform posterior sampling, an annealed pseudo-likelihood score called noise perturbed pseudo-likelihood score is introduced and combined with the prior score of SGM.
The proposed QCS-SGM applies to an arbitrary number of quantization bits.
Experiments on a variety of baseline datasets demonstrate that the proposed
QCS-SGM significantly outperforms existing state-of-the-art algorithms by a large margin for both in-distribution and out-of-distribution samples.
Moreover, as a posterior sampling method, QCS-SGM can be easily used to obtain confidence intervals or uncertainty estimates of the reconstructed results.
We have submitted our code in the supplementary material and will release it upon acceptance.
I NTRODUCTION
Many problems in science and engineering such as signal processing, computer vision, machine learning, and statistics can be cast as linear inverse problems:
[Equation]
where A ∈ R M × N is a known linear mixing matrix, n ∼ N ( n ; 0 , σ 2 I ) is an i.i.d.
additive Gaussian noise, and the goal is to recover the unknown signal x ∈ R N × 1 from the noisy linear measurements y ∈ R M × 1 .
Among various applications, compressed sensing (CS) provides a highly efficient paradigm that makes it possible to recover a high-dimensional signal from a far smaller number of M ≪ N measurements (Cand`es & Wakin, 2008).
The underlying wisdom of CS is to leverage the intrinsic structure of the unknown signal x to aid the recovery.
One of the most widely used structures is sparsity , i.e., most elements of x are zero under certain transform domains, e.g., wavelet and Fourier transformation (Cand`es & Wakin, 2008).
In other words, the standard CS exploits the fact that many natural signals in real-world applications are (approximately) sparse or compressible.
This direction has spurred a hugely active field of research during the last two decades, including efficient algorithm design (Tibshirani, 1996; Beck & Teboulle, 2009; Tropp & Wright, 2010; Kabashima, 2003; Donoho et al., 2009), theoretical analysis (Cand`es et al., 2006; Donoho, 2006; Kabashima et al., 2009; Bach et al., 2012), as well as all kinds of applications (Lustig et al., 2007; 2008), just to name a few.
Despite their remarkable success, traditional CS is still limited by the achievable rates since the sparsity assumptions, whether being naive sparsity or block sparsity (Duarte & Eldar, 2011), are still too simple to capture the complex and rich structure of natural signals.
For example, most natural signals are not strictly sparse even in the specified transform domains, and thus simply relying on sparsity alone for reconstruction could lead to inaccurate results.
Indeed, researchers have proposed to combine sparsity with additional structure assumptions, such as low-rank assumption (Fazel et al., 2008; Foygel & Mackey, 2014), total variation (Cand`es et al., 2006; Tang et al., 2009), to further improve the reconstruction performances.
Nevertheless, these hand-crafted priors can only apply, if not exactly, to a very limited range of signals and are difficult to generalize to other cases.
To address this problem, driven by the success of generative models (Goodfellow et al., 2014; Kingma & Welling, 2013; Rezende & Mohamed, 2015), there has been a surge of interest in developing CS methods with data-driven priors (Bora et al., 2017; Hand & Joshi, 2019; Asim et al., 2020; Pan et al., 2021).
The basic idea is that, instead of relying on hand-crafted sparsity, the prior structure of the unknown signal is learned through a generative model, such as VAE (Kingma & Welling, 2013) or
GAN (Goodfellow et al., 2014).
Notably, the past few years have witnessed a remarkable success of one new family of probabilistic generative models called diffusion models, in particular, score matching with Langevin dynamics (SMLD)
(Song & Ermon, 2019; 2020) and denoising diffusion probabilistic modeling (DDPM)
(Ho et al., 2020; Nichol & Dhariwal, 2021), which have proven extremely effective and even outperform the state-of-the-art (SOTA)
GAN (Goodfellow et al., 2014) and VAE (Kingma & Welling, 2013) in density estimation and generation of various natural sources.
As both DDPM and SMLD estimate, implicitly or explicitly, the score (i.e., the gradient of the log probability density w.r.t. to data), they are also referred to together as score-based generative models (SGM) (Song et al., 2020).
Several CS methods with SGM have been proposed very recently (Jalal et al., 2021a;b; Kawar et al., 2021; 2022; Chung et al., 2022) and perform quite well in recovering x with only a few linear measurements in (1).
Nevertheless, the linear model (1) ideally assumes that the measurements have infinite precision, which is not the case in realistic acquisition scenarios.
In practice, the obtained measurements have to be quantized to a finite number of Q bits before transmission and/or storage (Zymnis et al., 2009; Dai & Milenkovic, 2011).
Quantization leads to information loss which makes the recovery particularly challenging.
For moderate and high quantization resolutions, i.e., Q is large, the quantization impact is usually modeled as a mere additive Gaussian noise whose variance is determined by quantization distortion (Dai & Milenkovic, 2011; Jacques et al., 2010).
Subsequently, most CS algorithms originally designed for the linear model (1) can then be applied with some modifications.
However, such an approach is apparently suboptimal since the information about the quantizer is not utilized to a full extent (Dai & Milenkovic, 2011).
This is true especially in the case of coarse quantization, i.e., Q is small.
An extreme and important case of coarse quantization is 1-bit quantization, where Q = 1 and only the sign of the measurements are observed (Boufounos & Baraniuk, 2008).
Apart from extremely low cost in storage, 1-bit quantization is particularly appealing in hardware implementations and has also proven robust to both nonlinear distortions and dynamic range issues (Boufounos & Baraniuk, 2008).
Consequently, there have been extensive studies on quantized CS, particularly 1-bit CS, in the past decades and a variety of algorithms have been proposed, e.g., (Zymnis et al., 2009; Dai & Milenkovic, 2011; Plan & Vershynin, 2012; 2013; Jacques et al., 2013; Xu & Kabashima, 2013;
Xu et al., 2014; Awasthi et al., 2016; Meng et al., 2018; Jung et al., 2021; Liu et al., 2020; Liu & Liu, 2022).
However, most existing methods are based on standard CS methods, which, therefore, inevitably inherit their inability to capture rich structures of natural signals beyond sparsity.
While several recent works (Liu et al., 2020; Liu & Liu, 2022) studied 1-bit CS using generative priors, their main focuses are VAE and/or GAN, rather than SGM.
1. C ONTRIBUTIONS
• We propose a novel framework, termed as Quantized Compressed Sensing with Score-based Generative Models (QCS-SGM in short), to recover unknown signals from noisy quantized measurements.
QCS-SGM applies to an arbitrary number of quantization bits, including the extremely coarse 1-bit sign measurements.
To the best of our knowledge, this is the first time that SGM has been utilized for quantized CS (QCS).
• We consider one popular SGM model called NCSNv2 (Song & Ermon, 2020) and verify the effectiveness of the proposed QCS-SGM in QCS on various real-world datasets including MNIST, Cifar-10, CelebA 64 × 64 , and the high-resolution FFHQ 256 × 256 .
Using the pretrained SGM as a generative prior, for both in-distribution and out-of-distribution samples, the proposed QCS-SGM significantly outperforms existing SOTA algorithms by a large margin in QCS.
Perhaps surprisingly, even in the extreme case of 1-bit sign measurements, QCS-SGM can still faithfully recover the original images with much fewer measurements than the original signal dimension, as shown in Figure 1.
• We propose a noise-perturbed pseudo-likelihood score as a principled way to incorporate information from the measurements, which is one key to the success of QCS-SGM.
Interestingly, when degenerated to the linear case (1), if A is row-orthogonal, then QCS-SGM reduces to a form similar to Jalal et al.
(2021a).
While the annealing term (denoted as γ 2 t in (4) of Jalal et al.
(2021a) is added heuristically as one additional hyper-parameter, it appears naturally within our framework and admits an analytical solution.
More importantly, for general matrices A , our method generalizes and significantly outperforms Jalal et al.
(2021a).
It is expected that the idea of noise-perturbed pseudo-likelihood score can be generalized to other conditional generative models as a general rule.
• As one sampling method, QCS-SGM can easily yield multiple
samples with different random initializations, whereby one can obtain confidence intervals or uncertainty estimates.
1. R ELATED WORKS
The proposed QCS-SGM is one kind of data-driven method for CS which generally include two lines of research: unsupervised and supervised.
The supervised approach trains neural networks from end to end using pairs of original signals and observations (Jin et al., 2017; Zhang & Ghanem, 2018; Aggarwal et al., 2018; Wu et al., 2019;
Yao et al., 2019; Gilton et al., 2019; Yang et al., 2020; Antun et al., 2020).
Despite good performance on specific problems, the supervised approach is severely limited and requires re-training even with a slight change in the measurement model.
The unsupervised approach, including the proposed QCS-SGM, only learns a prior and thus avoids such problem-specific training since the measurement model is only known and used for inference (Bora et al., 2017; Hand & Joshi, 2019; Asim et al., 2020; Pan et al., 2021).
Among various others, the CSGM framework proposed in Jin et al.
(2017) is the most popular one which learns the prior with generative models.
Recent works Liu et al.
; Liu & Liu (2022) extended CSGM to non-linear observations including 1-bit CS, which are the two most related studies to current work.
However, the main focuses of Liu et al.
; Liu & Liu (2022) are limited to VAE and GAN (in particular DCGAN (Radford et al., 2015)).
As is widely known, GAN suffers from unstable training and less diversity due to the adversarial training nature while VAE uses a surrogate loss.
Another significant drawback of priors with VAE/GAN is that they can have large representation errors or biases due to architecture and training, which can be easily caused by inappropriate latent dimensionality and/or mode collapse (Asim et al., 2020).
With the advent of SGM such as SMLD (Song & Ermon, 2019; 2020) and
DDPM (Ho et al., 2020; Nichol & Dhariwal, 2021), there have been several recent studies (Jalal et al., 2021a;b; Kawar et al., 2021; 2022; Chung et al., 2022; Daras et al., 2022) in CS using SGM as a prior and outperform conventional methods. Interestingly, Daras et al. (2022) performs posterior sampling in the “latent space” of a pre-trained generative model.
Nevertheless, all the existing works on SGM for CS focus on linear measurements (1) without quantization, which therefore lead to performance degradation when directly applied to the quantized case (2). Moreover, even in the linear case (1), as illustrated in Corollary 1.2 and Remark 3, the likelihood score computed in Jalal et al. (2021a;b) is not accurate for general matrices A . By contrast, we provide a simple yet effective approximation of the likelihood score which significantly outperforms Jalal et al.
(2021a) for general matrices even in the linear case. Another related work (Qiu et al., 2020) studied 1-bit CS where the sparsity is implicitly enforced via mapping a low dimensional representation through a known ReLU generative network.
Wei et al.
(2019) also considers a non-linear recovery using a generative model assuming that the original signal is generated from a low-dimensional latent vector.
B ACKGROUND 2. I NVERSE P ROBLEMS FROM Q UANTIZED M EASUREMENTS
The inverse problem from noisy quantized measurements is posed as follows (Zymnis et al., 2009)
y
= Q ( Ax + n ) , (2) where the goal is to recover the unknown signal x ∈ R N × 1 from quantized measurements y ∈ R M × 1 , where A ∈ R M × N is a known linear mixing matrix, n ∼ N ( n ; 0 , σ 2 I ) is an i.i.d.
additive Gaussian noise with known variance, and Q ( · ) : R M × 1 → Q M × 1 is an element-wise 1 quantizer function which maps each element into a finite (or countable) set of codewords Q , i.e., y m =
Q ( z m + n m ) ∈ Q , or equivalently ( z m + n m ) ∈ Q − 1 ( y m ) , m = 1 , 2 , ..., M , where z m is the m -th element of z =
Ax .
Usually, the quantization codewords Q correspond to intervals (Zymnis et al., 2009), i.e., Q − 1 ( y m ) =
[ l y m , u y m ) where l y m , u y m denote prescribed lower and upper thresholds associated with the quantized measurement y m , i.e., l y m ≤ ( z m + n m ) <
u y m .
For example, for a uniform quantizer with Q quantization bits (resolution), the quantization codewords Q
= { q r } 2 Q r =1 consist of 2 Q elements, i.e.,
[Equation]
(3) where ∆ > 0 is the quantization interval, and the lower and upper thresholds associated q r are
[Equation]
In the extreme 1-bit case, i.e., Q = 1 , only the sign values are observed, i.e., y = sign( Ax + n ) ,
[Equation]
5) where the quantization codewords Q = {− 1 , +1 } and the associated thresholds are l − 1 =
−∞ , u − 1 = 0 and l +1 = 0 , u +1 = + ∞ , respectively.
2. S CORE - BASED G ENERATIVE M ODELS
For any continuously differentiable probability density function p ( x ) , if we have access to its score function, i.e., ∇ x log p ( x ) , we can iteratively sample from it using Langevin dynamics (Turq et al., 1977; Bussi & Parrinello, 2007;
Welling & Teh, 2011)
[Equation]
where z t ∼ N ( z t ; 0 , I ) , α t > 0 is the step size, and T is the total number of iterations.
It has been demonstrated that when α t is sufficiently small and T is sufficiently large, the distribution of x T will converge to p ( x ) (Roberts & Tweedie, 1996;
Welling & Teh, 2011).
In practice, the score function ∇ x log p ( x ) is unknown and can be estimated using a score network s θ ( x ) via score matching (Hyv¨arinen, 2006; Vincent, 2011).
However, the vanilla Langevin dynamics faces a variety of challenges such as slow convergence.
To address this challenge, inspired by simulated annealing (Kirkpatrick et al., 1983; Neal, 2001), Song & Ermon (2019) proposed an annealed version of
Langevin dynamics, which perturbs the data with Gaussian noise of different scales and jointly estimates the score functions of noise-perturbed data distributions.
Accordingly, during the inference, an annealed Langevin dynamics (ALD) is performed to leverage the information from all noise scales.
Specifically, assume that p β (˜ x | x ) =
N (˜ x ; x , β 2 I ) , and so we have p β (˜ x ) =
(cid:82) p data ( x ) p β (˜ x | x ) d x , where p data ( x ) is the data distribution.
Consider we have a sequence of noise scales { β t } Tt = satisfying β max = β 1 > β 2 > · · · > β T = β min > 0 .
The β min → 0 is small enough so that p β min (˜ x )
≈ p data ( x ) , and β max is large enough so that p β max (˜ x ) ≈ N ( x ; 0 , β 2max I ) .
The noise conditional score network (NCSN) s θ ( x , β ) proposed in Song & Ermon (2019)
aims to estimate the score function of each p β t (˜ x ) by optimizing the following weighted sum of score matching objective
[Equation]
After training the NCSN, for each noise scale, we can run K steps of Langevin MCMC to obtain a sample for each p β t (˜ x ) as
[Equation]
The sampling process is repeated for t = 1 , 2 , ..., T sequentially with x 01 ∼ N ( x ; 0 , β 2max I ) and x 0 t +1 = x
Kt when t < T .
As shown in Song & Ermon (2019), when K → ∞ and α t → 0 for all t , the final sample x KT will become an exact sample from p β min (˜ x )
≈ p data ( x ) under some regularity conditions.
Later, by a theoretical analysis of the learning and sampling process of NCSN, an improved version of NCSN, termed NCSNv2, was proposed in Song & Ermon (2020) which is more stable and can scale to various datasets with high resolutions.
Q UANTIZED CS WITH SGM
In the case of quantized measurements (2) with access to y , the goal is to sample from the posterior distribution p ( x | y ) rather than p ( x ) .
Subsequently, the Langevin dynamics in (6) becomes
[Equation]
where the conditional ( posterior ) score ∇ x log p ( x | y ) is required.
One direct solution is to specially train a score network work that matches the score ∇ x log p ( x | y ) .
However, this method is rather inflexible and needs to re-train the network when confronted with different measurements.
Instead, inspired by conditional diffusion models (Jalal et al., 2021a), we resort to computing the posterior score using the Bayesian rule from which ∇ x log p ( x | y ) is decomposed into two terms
[Equation]
which include the unconditional score ∇ x log p ( x ) (we call it prior score ), and the conditional score ∇ x log p ( y | x ) (we call it likelihood score ), respectively.
The prior score ∇ x log p ( x ) can be easily obtained using a trained score network such as NCSN or NCSNv2, so the remaining goal is to compute the likelihood score ∇ x log p ( y | x ) from the quantized measurements.
3. N OISE - PERTURBED PSEUDO - LIKELIHOOD SCORE
In the case without noise perturbing in x , one can calculate the likelihood score ∇ x log p ( y | x ) directly from the measurement process (2) and then substitute it into (10) to obtain the posterior score ∇ x log p ( x | y ) .
However, this does not work well and might even diverge (Jalal et al., 2021a).
The reason is that in NCSN/NCSNv2, as shown in (7), the estimated score is not the original prior score ∇ x log p data ( x ) but rather a noise-perturbed prior score ∇ ˜ x log p β t (˜ x ) for each noise scale β t .
Therefore, the likelihood score directly computed from (2) does not match the noise-perturbed prior score.
To address this problem, we propose to calculate the noise-perturbed pseudo-likelihood score ∇ ˜ x log p β t ( y | ˜ x ) by explicitly accounting for such the noise perturbation in NCSN/NCSNv2. Unfortunately, for NCSN/NCSNv2, the ∇ ˜ x log p β t ( y | ˜ x ) is intractable in the general case.
To this end, we propose a simple yet effective approximation by introducing two assumptions.
Assumption 1.
The prior p ( x ) is uninformative w.r.t.
p (˜ x | x ) , i.e., p ( x | ˜ x ) ∝ p (˜ x | x ) .
Remark 1.
This assumption is asymptotically accurate when the perturbed noise in ˜ x becomes negligible, i.e., β t → 0 when t is large.
An illustration of this is shown in Appendix A.
Assumption 2.
The matrix A is row-orthogonal, i.e., AA T is a diagonal matrix.
Remark 2.
The measurement matrices A are usually carefully designed in CS.
Interestingly, most popular CS matrices A satisfy this assumption, e.g., discrete Fourier transform (DFT) in MRI, discrete cosine transform (DCT), Hadamard matrices, as well as general row-orthogonal matrices
Vehkaper¨a et al. (2016).
For another class of widely used i.i.d.
Gaussian matrices, this assumption is approximated satisfied, as verified in Appendix B.
Notably, it has been theoretically demonstrated that row-orthogonal matrices are mostly preferred for CS due to both superior performance and easy implementation (Vehkaper¨a et al., 2016).
Therefore, this assumption does not impose a serious limitation in the case of CS, and we leave the study of general A as a future work.
Subsequently, we obtain our main result as shown in Theorem 1.
Theorem 1.
(noise-perturbed pseudo-likelihood score)
For each noise scale β t > 0 , if p β t (˜ x | x ) =
N (˜ x ; x , β 2 t I ) , under Assumption 1 and Assumption 2, the noise-perturbed pseudo-likelihood score ∇ ˜ x log p β t ( y | ˜ x ) for the quantized measurements y in (2) can be computed as
[Equation] [Equation]
where G ( β t , y , A , ˜ x ) =
[ g 1 , g 2 , ..., g M ] T ∈ R M × 1 with each element being
[Equation]
where a Tm ∈ R 1 × N denotes the m -th row vector of A and
[Equation]
Proof.
The proof is shown in Appendix C.
In the extreme 1-bit case, a more concise result can be obtained, as shown in Corollary 1.1:
Corollary 1.1.
(noise-perturbed pseudo-likelihood score in the 1-bit case)
For each noise scale β t > 0 , if p β t (˜ x | x ) =
N (˜ x ; x , β 2 t I ) , under Assumption 1 and Assumption 2, the noise-perturbed pseudo-likelihood score ∇ ˜ x log p β t ( y | ˜ x ) for the sign measurements y = sign( Ax + n ) in (5) can be computed as
[Equation]
) where G ( β t , y , A , ˜ x ) =
[ g 1 , g 2 , ..., g M ] T ∈ R M × 1 with each element being
[Equation]
where ˜ z m = a Tm ˜ x √ σ 2 + β [Equation] ∥ [Equation]
Tm [Equation] 22 and Φ( z ) = 1 √ 2 π (cid:82) z
−∞ e − t 2 2 dt is the cumulative distribution function (CDF) of the standard normal distribution.
Proof.
It can be easily proved following the proof in Theorem 1.
Interestingly, if no quantization is used, or equivalently Q → ∞ , the model in (2) reduces to the standard linear model (1).
Accordingly, results in Theorem 1 can be modified as in Corollary 1.2:
Corollary 1.2.
(noise-perturbed pseudo-likelihood score without quantization)
For each noise scale β t > 0 , if p β t (˜ x | x ) =
N (˜ x ; x , β 2 t I ) , under Assumption 1, the noise-perturbed pseudo-likelihood score ∇ ˜ x log p β t ( y | ˜ x ) for linear measurements
y =
Ax + n in (1) can be computed as
[Equation]
Moreover, under Assumption 2, the ∇ ˜ x log p β t ( y | ˜ x ) can be further simplified as
[Equation]
where G linear ( β t , y , A , ˜ x ) = [ g 1 , g 2 , ..., g M ] T ∈ R M × 1 with each element being
[Equation]
Proof.
The proof is shown in Appendix D.
Remark 3.
In contrast to the quantized case, here we obtain a closed-form approximation (16) for any kind of matrices A .
Interestingly, if Assumption 2 is further imposed, it reduces to (18), which is similar to Jalal et al.
(2021a) without quantization where the annealing term β 2 t (cid:13)(cid:13) a Tm (cid:13)(cid:13) 2 2 in (18) corresponds to γ 2 t in Jalal et al.
(2021a).
However, γ 2 t in Jalal et al.
(2021a) is added heuristically as an additional hyper-parameter while we derive it in closed-form under the principled perturbed pseudo-likelihood score.
More importantly, for general matrices A , our closed-form result (16) significantly outperforms Jalal et al.
(2021a), as demonstrated in
Appendix E. Therefore, we not only explain the necessity of an annealing term in Jalal et al.
(2021a)
but also extend and improve Jalal et al.
(2021a)
in the general case.
3. P OSTERIOR SAMPLING VIA ANNEALED L ANGEVIN DYNAMICS
By combining the prior score from SGM and noise-perturbed pseudo-likelihood score in Theorem 1, we obtain the resultant Algorithm 1, namely, Quantized CS with SGM (QCS-SGM in short).
Algorithm
Input: { β t } Tt =1 , ϵ, K, y , A , σ 2 , quantization codewords Q and thresholds { [ l q , u q ) | q ∈ Q}
Initialization: x 0 1 ∼ U (0 , 1) 1 for t = 1 to T do α t ← ϵβ 2 t /β 2 T for k = 1 to K do
Draw z kt ∼ N ( 0 , I )
Compute G ( β t , y , A , x k − 1 t ) as (12) (or (15) for 1-bit) x kt = x k − 1 t + α t (cid:2) s θ ( x k − 1 t
[Equation]
x
[Equation]
← x Kt
Output: ˆx = x KT
E XPERIMENTS
In this section, we empirically demonstrate the efficacy of the proposed QCS-SGM in various scenarios.
The widely-used i.i.d.
Gaussian matrix A , i.e., A ij ∼ N (0 , 1 /M ) are considered.
Datasets : Three popular datasets are considered: MNIST (LeCun & Cortes, 2010) , Cifar- (Krizhevsky & Hinton, 2009), and CelebA (Liu et al., 2015), and the high-resolution Flickr Faces
High Quality (FFHQ) (Karras et al., 2018). MNIST (LeCun & Cortes, 2010) are grayscale images of size 28 × 28 pixels so that the input dimension for MNIST is N = 28 × 28 = 784 per image.
Cifar- (Krizhevsky & Hinton, 2009) consists of natural RGB images of size 32 × 32 pixels, resulting in[-N = 
-]
N = 32 × 32 × 3 = 3072 inputs per image.
For CelebA dataset (Liu et al., 2015), we cropped each face image to a 64 × 64 RGB image, resulting in N = 64 × 64 × 3 = 12288 inputs per image. The
FFHQ are high-resolution RGB images of size 256 × 256 , so that N = 256 × 256 × 3 =per image. Note that to evaluate the out-of-distribution (OOD) performance, we also cropped FFHQ to size 64 × 64 , as OOD samples for CelebA. All images are normalized to range [0 , 1] .
QCS-SGM :
Regarding the SGM models, we adopt the popular NCSNv2 (Song & Ermon, 2020) in all cases.
Specifically, for MNIST, we train a NCSNv2 (Song & Ermon, 2020) model on the MNIST training dataset with a similar training set up as Cifar10 in Song & Ermon (2020), while we directly use the for Cifar-10, and CelebA, and FFHQ, which are available in this Link.
Therefore, the prior score can be estimated using these pre-trained NCSNv2 models.
After observing the quantized measurements as (2), we can infer the original x by posterior sampling via QCS-SGM in Algorithm 1.
It is important to note that, we select images x that are unseen by the pre-traineirec SGM models .
For details of the experimental setting, please refer to Appendix F. We have submitted the code in the supplementary material and will release it upon acceptance.
4. 1- BIT Q UANTIZATION
First, we perform experiments on the extreme case, i.e., 1-bit measurements.
Specifically, we consider images of the MNIST (LeCun & Cortes, 2010) and CelebA datasets (Liu et al., 2015) in the same setting as Liu & Liu (2022).
Apart from QCS-SGM, we also show results of LASSO (Tibshirani, 1996) (for CelebA, Lasso with a DCT basis (Ahmed et al., 1974) is used), CSGM (Bora et al., 2017), [Figure] BIPG (Liu et al., 2020), and OneShot (Liu & Liu, 2022).
For Lasso (Lasso-W), CSGM, BIPG, and OneShot, we follow the setting as the open-sourced code of Liu & Liu (2022).
The typical reconstructed images from 1-bit measurements with fixed M ≪ N are shown in Figure 2 for both MNIST and CelebA. Perhaps surprisingly, QCS-SGM can faithfully recover the images from 1-bit measurements with M ≪
N measurements while other methods fail or only recover quite vague images.
To quantitatively evaluate the effect of M , we compare different algorithms using three popular metrics, i.e., peak signal-to-noise ratio (PSNR) and structural similarity index (SSIM) for different M , and the results are shown in Figure 3.
It can be seen that, under both metrics, the proposed QCS-SGM significantly outperforms all the other algorithms by a large margin.
Please refer to Appendix H for more results.
Multiple samples & Uncertainty Estimates : It is worth pointing out that, as one kind of sampling method, QCS-SGM can yield multiple samples with different random initialization so that we can easily obtain confidence intervals or uncertainty estimates of the reconstructed results.
Please refer to Appendix G for more typical samples of QCS-SGM.
In contrast, OneShot (Liu & Liu, 2022) can easily get stuck in a local minimum with different random initializations (Liu & Liu, 2022).
4. M ULTI - BIT Q UANTIZATION
Then, we evaluate the efficacy of QCS-SGM in the case of multi-bit quantization, e.g., 2-bit, 3-bit.
The results on Cifar-10 and CelebA are shown in Figure 4.
The results of the linear case without quantization are also shown for comparison using (18) in Corollary 1.2.
As expected, with the increase of quantization resolution, the reconstruction performances get better.
One typical result for high-resolution FFHQ 256 × 256 images is shown in Figure 1.
Moreover, to evaluate the quantization effect under “fixed budget” setting, i.e., Q × B remain the same for different Q and M , we conduct experiments for CelebA in the fixed case of Q × M = 12288 .
Interestingly, as shown in Table 1, under ‘fixed budget”, while there is an apparent increase in PSNR for the multi-bit case, the perception metric SSIM remains about the same as 1-bit case.
For more results, please refer to Appendix H.
4. O UT - OF -D ISTRIBUTION P ERFORMANCE
In practice, the signals of interest may be out-of-distribution (OOD) because the available training dataset has bias and is unrepresentative of the true underlying distribution.
To assess the performance of QCS-SGM for OOD datasets, we choose the CelebA as the training set while evaluating results on
[Table] [Figure] OOD images randomly sampled from the FFHQ dataset, which have features of variation that are rare among CelebA images (e.g., skin tone, age, beards, and glasses) (Jalal et al., 2021b).
As shown in Figure 5, even on OOD images, the proposed QCS-SGM can still obtain significantly competitive results compared with other existing algorithms.
C ONCLUSION
In this paper, a novel framework called Quantized Compressed Sensing with Score-based Generative
Models (QCS-SGM) is proposed to reconstruct a high-dimensional signal from noisy quantized measurements.
Thanks to the power of SGM in capturing the rich structure of natural signals beyond simple sparsity, the proposed QCS-SGM significantly outperforms existing SOTA algorithms by a large margin under coarsely quantized measurements on a variety of experiments, for both indistribution and out-of-distribution datasets.
There are still some limitations in QCS-SGM.
For example, it is limited to the case where measurement matrix A (approximately) satisfies that AA T is a diagonal matrix.
The generalization of general matrices is an important future work.
Generative models are an active area of research with ongoing rapid improvements and we believe that SGM could open up new opportunities in quantized CS in the future, and hope that our work as a first step could inspire more studies on this fascinating topic. A V ERIFICATION OF A SSUMPTION Mathematically, we can write the exact p ( y | ˜ x ) as [Equation]
where from the Bayes’ rule,
[Equation]
For NCSN/NCSNv2, recall that the likelihood p (˜ x | x ) follows a Gaussian distribution N (˜ x ; x , β 2 I ) , where β is the variance of the perturbed Gaussian noise and, by definition, β → 0 in later steps of the reverse diffusion process in NCSN/NCSNv2. As a result, from (20), the likelihood p (˜ x | x ) will dominate the posterior p ( x | ˜ x ) as β → 0 and therefore p ( x | ˜ x ) ∝ p (˜ x | x ) , indicating that the prior p ( x ) becomes uninformative. Note that this is particularly the case when the entropy of p ( x ) is high, which is usually the case for generative models capable of generating diverse images. While this assumption is crude at the beginning of the reverse diffusion process when the noise variance β is large, it is asymptotically accurate as the process goes forward.
The effectiveness of this assumption is also empirically supported by various experiments in Section 4.
verifying the effectiveness of the
Assumption 1 for this toy example.
A toy example: In the following, we illustrate the assumption in a toy example where x reduces to a scalar random variable x and the associated prior p ( x ) follows a Gaussian distribution, i.e., p ( x ) =
N ( x ; 0 , σ 20 ) , where σ 2 is the prior variance.
The likelihood p (˜ x | x ) in this case is simply p (˜ x | x ) =
N (˜ x ; x, β 2 ) .
Then, from (20), after some algebra, it can be computed that the posterior distribution p ( x | ˜ x ) is
[Equation]
where
[Equation]
Under the Assumption 1, i.e., p ( x | ˜ x ) ∝ p (˜ x | x ) , we obtain an approximation of p ( x | ˜ x ) as follows
[Equation]
where
[Equation]
By comparing the exact result (22) and approximation result (24), it can be easily seen that for a fixed σ 20 > 0 , as β → 0 , we have m pseudo → m post and v pseudo → v post .
In Figure 6, similar to
NCSN/NCSNv2, we anneal β as β t = β max ( β min β max ) t − 1 T − 1 geometrically and compare m pseudo , v pseudo with m exact , v exact as t increase from 1 to T .
It can be seen in Figure 6 that the approximated values m pseudo , v pseudo , especially the variance v pseudo , approach to the exact values m exact , v exact very quickly, verifying the effectiveness of the Assumption 1 under this toy example.
B V ERIFICATION OF A SSUMPTION
Before verifying it in the i.i.d.
Gaussian case, we first briefly add a comment on the general case when AA T is not an exact diagonal matrix.
As demonstrated later in Appendix C, the Assumption 2 is introduced to ensure that the covariance matrix σ 2 I + β 2 t AA T of Gaussian noise is diagonal so that we can obtain a closed-form solution for the likelihood score in the quantized case, which is otherwise intractable.
If
Suppose that the elements of A follow i.i.d.
Gaussian, i.e., A ij ∼ N (0 , σ 2 ) .
Next, we investigate the elements of the matrix C
= AA T =
{ C ij } , i, j = 1 ...
M .
Regarding the diagonal elements C ii , by definition, it reads as
[Equation]
As C ii is the sum of square of N i.i.d.
Gaussian random variables A ij , j = 1 ...N , C ii follows a
Gamma distribution, i.e., Γ (cid:0) N 2 , 2 σ 2 (cid:1) .
The mean and variance of C ii can be computed as Nσ 2 and 2 Nσ 4 .
Regarding the off-diagonal elements C ij , i ̸
= j , by definition, it reads as
[Equation]
As A in and A jn are independent Gaussian for i ̸
= j , it can be computed that the mean and variance of C ij are 0 and Nσ 4 , respectively.
When σ 2 = 1 /M and M = αN , where α > 0 is the constant measurement ratio, the variance of C ij is 1 α 2 N → 0 as N → ∞ .
As a result, all the off-diagonal elements of AA T will tend to zero as N → ∞ .
From another perspective, in Figure 7, we compute the ratio between the average magnitude of off-diagonal elements C ij and the diagonal elements C ii when M = αN with α = 0 .
It can be seen that as N increases, the magnitude of the off-diagonal elements becomes negligible.
As a result, when A is i.i.d.
Gaussian, the matrix AA T can be well approximated as a diagonal matrix in the high-dimensional case.
C P ROOF OF T HEOREM
Proof.
Let us denote z
= Ax .
For each noise scale β t > 0 , under the Assumption 1, we obtain
[Equation]
(27) then we can equivalently write
[Equation]
(28) where w ∼ N ( 0 , I ) .
As a result, z =
Ax = A (˜ x + β t w ) =
A˜x + β t
Aw .
Then, from (2), we obtain y
[Equation]
Q ( A ˜
[Equation]
˜ [Equation] ) , (29) where ˜ n = n + β t
Aw .
Since n ∼ N ( 0 , σ 2 I ) and w ∼ N ( 0 , I ) and are independent to each other, it can be concluded that ˜ n is also Gaussian with mean zero and covariance σ 2 I + β 2 t AA T , i.e., ˜ n ∼ N ( ˜ n ; 0 , σ 2 I +
Subsequently, under the Assumption 2, i.e., AA T is a diagonal matrix, each element ˜ n m ∼ N (˜ n m ; 0 , σ 2 + β 2 t (cid:13)(cid:13) a Tm (cid:13)(cid:13) 2 2 ) of ˜ n will be independent to each other and thus from (38) we can obtain a closed-form solution for the likelihood distribution (we will refer it as a pseudo-likelihood due to the assumptions used) p ( y | ˆ z =
A ˜ x ) as follows
[Equation]
where, from the definition of quantizer Q
[Equation]
where Φ( z ) = 1 √ 2 π (cid:82) z
−∞ e − t 2 2 dt is the cumulative distribution function of the standard normal distribution and
[Equation]
As a result, it can be calculated that the noise-perturbed pseudo-likelihood score ∇ ˜ x log p β t ( y | ˜ x ) for the quantized measurements y in (2) can be computed as
[Equation]
where G ( β t , y , A , ˜ x ) =
[ g 1 , g 2 , ..., g M ] T ∈ R M × 1 with each element being
[Equation]
which completes the proof.
[Figure]
D P ROOF OF C OROLLARY 1.
Proof.
Similarly in the proof of Theorem 1, let us denote z
= Ax .
For each noise scale β t > 0 , under the Assumption 1, we can equivalently write
[Equation]
where w ∼ N ( 0 , I ) .
As a result, z =
Ax = A (˜ x + β t w ) =
A˜x + β t
Aw .
Then, in the case of the linear model, from (1), we obtain
[Equation]
where ˜ n = n + β t
Aw .
Since n ∼ N ( 0 , σ 2 I ) and w ∼ N ( 0 , I ) and are independent to each other, it can be concluded that ˜ n is also Gaussian with mean zero and covariance σ 2 I + β 2 t AA T , i.e., ˜ n ∼ N (˜ n ; 0 , σ 2 I + β 2 t AA T ) .
Therefore, a closed-form solution for the likelihood distribution p ( y | ˆ z =
A ˜ x ) can be obtained as follows
[Equation]
As a result, we can readily obtain a closed-form solution for the noise-perturbed pseudo-likelihood score ∇ ˜ x log p β t ( y | ˜ x ) as follows
[Equation]
Furthermore, if AA T is a diagonal matrix, the inverse of matrix (cid:0) σ 2 I + β 2 t AA T (cid:1) become trivial since it is a diagonal matrix with the m -th diagonal element being σ 2 + β 2 t (cid:13)(cid:13) a Tm (cid:13)(cid:13) 2 2 .
After some simple algebra, we can obtain the equivalent representation in (18), which completes the proof.
E C OMPARISON WITH ALD IN J ALAL ET AL . (2021 A ) IN THE LINEAR CASE
As shown in Corollary 1.2, in the special case without quantization, our results in Theorem 1 can be reduced to a form similar to the ALD in Jalal et al.
(2021a).
However, there are several different important differences.
First, our results are derived in a principled way as noise-perturbed pseudolikelihood score and admit closed-form solutions, while the results in Jalal et al.
(2021a) are obtained [Figure] heuristically by adding an additional hyper-parameter (and thus needs fine-tuning).
Second, the results in Jalal et al.
(2021a) are similar to an approximate version (18) of ours, which holds only when AA T is a diagonal matrix.
For general matrices A , we have a closed-form approximation (16).
We compare our results with ALD (Jalal et al., 2021a) and the results are shown in Figureand Figure 9.
It can be seen that when at low condition number of A when AA T is approximately a diagonal matrix, our with diagonal approximation (18) performs similarly as (16), both of which are similar to ALD (Jalal et al., 2021a), as expected.
However, when the condition number of A is large so that AA T is far from a diagonal matrix, our method with the (16) significantly outperforms ALD in Jalal et al.
(2021a).
F D ETAILED E XPERIMENTAL S ETTINGS
In training NCSNv2 for MNIST, we used a similar training setup as Song & Ermon (2020) for Cifar as follows.
Training: batch-size: 128 n-epochs: 500000 n-iters: 300001 snapshot-freq: 50000 snapshot-sampling: true anneal-power: 2 log-all-sigmas: false.
Please refer to Song & Ermon (2020) and associated open-sourced code for details of training. For Cifar-10, CelebA, and FFHQ, we directly use the pre-trained models available in this Link.
When performing posterior sampling using the QCS-SGM in 1, for simplicity, we set a constant value ϵ = 0 .
for all quantized measurements (e.g., 1-bit, 2-bit, 3-bit) for MNIST, Cifar10 and [-ϵ = 0 . 
-]CelebA.
For the high-resolution FFHQ 256 × 256 , we set ϵ = 0 . 00005 for 1-bit and ϵ = 0 . for 2-bit and 3-bit case, respectively. For all linear measurements for MNIST, Cifar10, and CelebA, we set ϵ = 0 . 00002 .
It is believed that some improvement can be achieved with further fine-tuning of ϵ for different scenarios.
For MNIST and Cifar-10, we set β 1 = 50 , β T = 0 .
01 , T = 232 ; for CelebA, we set β 1 = 90 , β T = 0 . 01 , T = 500 ; for FFHQ, we set β 1 = 348 , β T = 0 .
01 , T = 2311 which are the same as Song & Ermon (2020).
The number of steps K in QCS-SGM for each noise scale is set to be K = 5 in all experiments.
For more details, please refer to the submitted code.
G M ULTIPLE S AMPLES AND U NCERTAINTY E STIMATES
As one kind of posterior sampling method, QCS-SGM can yield multiple samples with different random initialization so that we can easily obtain confidence intervals or uncertainty estimates of the reconstructed results.
For example, typical samples as well as mean and std are shown in Figurefor MNIST and CelebA in the 1-bit case. [Figure]
H A DDITIONAL R ESULTS
Some additional results are shown in this section.
Figure 11 and Figure 12 show results with relatively large value of M in the same setting as Figureand Figure 4, respectively. [Figure] [Figure]
Figure 13 shows the quantitative results of QCS-SGM for different quantization bits.
Figure 14 and Figure 15 show the reconstructed images of QCS-SGM for Cifar10 and CelebA in the fixed budget case of Q × M = 3072 and Q × M = 12288 , respectively.
[Figure] [Figure] [Table] I C OMPARISON WITH N EUMANN N ETWORKS G ILTON ET AL . (2019)
In the case of no quantization we compare our method with one popular method called Neumann
Networks (Gilton et al., 2019).
The results for Cifar10 are shown in Table 2.
It can be seen that our method outperforms Neumann Networks by a large margin.
We do not show quantization comparisons for quantized CS since Neumann Networks does not support it.
[Figure] [Figure] J C OMPARISON IN THE EXACTLY SAME SETTING AS L IU & L IU (2022) Note that there is a slight difference in the modeling of measurement matrix A and noise n between ours and that in Liu & Liu (2022).
In Liu & Liu (2022), it is assumed that A is an i.i.d.
Gaussian matrix where each element A ij follows A ij ∼ N (0 , 1) and that n is an i.i.d.
Gaussian with variance σ 2 , i.e., n i ∼ N (0 , σ 2 ) .
However, in practice, the measurement matrix A is usually normalized so that the norm of each column equals 1.
As a result, in our setting, we assume that each element of i.i.d.
Gaussian matrix follows A ij ∼ N (0 , 1 /M ) .
Mathematically, there is a one-to-one correspondence between the two settings, but the simulation setting is different due to the different measurement size M .
As a result, for an exact comparison with results in Liu & Liu (2022), we also conducted experiments assuming exactly the same setting as Liu & Liu (2022), i.e., A ij ∼ N (0 , 1) , and n i ∼ N (0 , σ 2 ) .
The results are shown in Figures 20, 21, 22, 23.
[Figure]
[Figure]
[Figure]