Unrolled denoising networks provably learn to perform optimal Bayesian inference
Authors: Aayush Karan
Kulin Shah
Sitan Chen
Yonina C. Eldar
Venue: NeurIPS 2024 poster (2024)
Keywords: algorithm unrolling
approximate message passing (AMP)
inverse problems
denoising
Primary area: learning_theory
TL;DR: We show that an algorithmically inspired architecture similar to reverse diffusion denoising can provably learn Bayes optimal inference with gradient descent, and we show the score-matching objective is learnable in one-dimension along the way.

Abstract: Much of Bayesian inference centers around the design of estimators for inverse problems which are optimal assuming the data comes from a known prior. But what do these optimality guarantees mean if the prior is unknown? In recent years, algorithm unrolling has emerged as deep learning's answer to this age-old question: design a neural network whose layers can in principle simulate iterations of inference algorithms and train on data generated by the unknown prior.  Despite its empirical success, however, it has remained unclear whether this method can provably recover the performance of its optimal, prior-aware counterparts.

In this work, we prove the first rigorous learning guarantees for neural networks based on unrolling approximate message passing (AMP). For compressed sensing, we prove that when trained on data drawn from a product prior, the layers of the network approximately converge to the same denoisers used in Bayes AMP. We also provide extensive numerical experiments for compressed sensing and rank-one matrix estimation demonstrating the advantages of our unrolled architecture \--- in addition to being able to obliviously adapt to general priors, it exhibits improvements over Bayes AMP in more general settings of low dimensions, non-Gaussian designs, and non-product priors.

Unrolled denoising networks provably learn to perform optimal Bayesian inference
Aayush Karan"
Kulin Shah
Harvard SEAS akaranl @g. harvardaddd
UT Austin kulinshah@utexas.edd
Sitan Chen
Yonina C. Eldar
Harvard SEAS
Weizmann Institute of Science
sitan@sessharvvrd.eedu
yonina
Abstract
Much of Bayesian inference centers around the design of estimators for inverse problems which are optimal assuming the data comes from a known prior. But what do these optimality guarantees mean if the prior unknown? In recent years, algorithm unrolling has emerged as deep learning's answer to this age- old question: design a neural network whose layers can in principle simulate iterations of inference algorithms and train on data generated by the unknown prior. Despite its empirical success, however, it has remained unclear whether this method can provably recover the performance of its optimal, prior-aware counterparts.
In this work, we prove the first rigorous learning guarantees for neural networks based on unrolling approximate message passing (AMP). For compressed sensing, we prove that when trained on data drawn from a product prior, the layers of the network approximately converge the same denoisers used in Bayes AMP. We also provide extensive numerical experiments for compressed sensing and rank-one matrix estimation demonstrating the advantages of our unrolled architecture in addition to being able to obliviously adapt to general priors, it exhibits improvements over Bayes AMP in more general settings low dimensions, non- -a--aassi designs, and non-product priors.
Introduction
Inverse problems within engineering and the sciences [BBDM21, ST99, Vog02] have inspired the development of a rich toolbox of algorithms for inferring unknown signals given noisy measurements. For instance, classic approach to solving sparse linear inverse problems is to solve the LASSO using an iterative algorithm like ISTA [DDDM04] or FISTA [BT09]. While these methods are easy to implement and remarkably performant, they are not designed to exploit distributional information about the underlying signal, which often comes from domain knowledge. In contrast, Bayesian methods like message passing and variational inference offer a natural framework for designing estimators that incorporate this kind of information: the algorithm designer crafts prior for the signal, and the measurements they observe naturally induce a posterior over what the underlying signal could have been.
Equal contribution
38th Conference on Neural Information Processing Systems (NeurIPS 2024).

---

Such an approach comes at a cost. On one hand, this method often comes with strong optimality guarantees in the well-specified setting where the algorithm designer has access the true prior distribution of the data. In practice, however, this prior is not known priori and hence must be inferred, and any mismatch between the inferred prior and the true distribution will adversely affect the performance of the estimator in ways that are poorly understood compared to what is known in the well-specified setting [BCPS21, BHMS22, MS22].
In recent years, algorithm unrolling has emerged as a scalable solution for developing estimators tradeoff by learning from samples drawn from the data distribution [GL10, MLE21, LTG* 20, SWED23]. The idea is to craft neural network architecture, each of whose layers expressive enough to implement one step of an existing, hand-crafted iterative algorithm (e.g. ISTA). Then, instead of explicitly setting the weights of the network so that implements that algorithm, one trains the network on examples of the inference problem at hand using stochastic gradient descent. Remarkably, the algorithm that the network converges to tends to perform at least as well as (and often better than) the hand-crafted algorithm being unrolled, e.g. in the number of layers and iterations necessary to achieve a certain level of error.
that can improve upon this pocticcaliyy-perfrmmaa
Thus, at least empirically, algorithm unrolling seems to achieve the best of both worlds, marrying the domain-awart power of classical iterative methods with the remarkable learning capabilities of neural networks. From a theoretical perspective however, our understanding of its performance is rather limited, as existing guarantees are centered around algorithiii aspects like representational power and generalization bounds (see Section 1.1 for detailed discussion).
In particular, the following fundamental learning question remains open:
Can an unrolled network trained with stochastic gradient descent provably obtain an estimator competitive with the best prior-aware algorithms?
In this work, we give the first rigorous learning guarantees for this question, focusing the well-studied setting of compressed sensing (see Section 2.1 for definitions). In addition, we provide the first empirical evidence in the affirmative for the problem of rank-one matrix estimation (Section 1.
Approximate message passing and unrolling. Consider the standard Bayesian setup where we observe noisy measurement y of some signal x, and would like to output an estimate minimizing E||x - x|
the Bayes-optimal estimator for this task is the posterior mean E [x y]. but in many settings of interest this estimator may not be computable by a polynomial-time algorithm. Among computationally efficient estimators, for a wide variety of such inference tasks it is conjectured [MW22b, CMW20, MW22a, BKM*19] that a certain family of iterative algorithms called approximate message passing (AMP) is optimal. We give self-contained exposition of this method in Section 2. Roughly speaking, one can think of AMP as a more advanced version of ISTA where the denoiser at each step can be tuned depending on the prior, and additionally there is a crucial momentum term inspired by correction from statistical physics [TAP77]. AMP with the optimal tuning of the denoiser is called Bayes AMP.
Motivated by the appealing theoretical properties of AMP, in this work we investigate the training dynamics of neural networks given by unrolling this algorithm. In place of the prior-dependent
denoisers 7/1, 17.... ........... .. we consider generic denoisers given by neural networks fi, f.... and unroll the iterations of AMP into layers of neural network (see Section 2.2 for details). For the theoretical results in this work, we focus on the setting where the only trainable parameters in
the network are the ones parametrizing the denoisers fr.
Unrolled AMP architectures and variants thereof were originally proposed and empirically investigated by Borgerding et al. [BS16, BSR17J and follow-ups [MMB17, MJC21, ITW19]. These works found that unrolled AMP can significantly outperform unrolled ISTA as well as a version of AMP with soft threshholding denoisers in terms of convergence; i.e., the number of layers needed to achieve a certain MSE.
Despite these compelling experimental results, to our knowledge, there is still little understanding as whether these networks can actually recover the performance of Bayes AMP. The main theoretical result of this work is to give the first proof that unrolled AMP networks trained with
2

---

gradient descent converge to the same denoisers as Bayes AMP and thus achieve mean squared error which is conjectured to be optimal among all polynomial-time algorithms for compressed sensing:
Theorem 1 (Informal, see Theorem 2). For compressed sensing with Gaussian sensing matrix, if the prior on the signal is a product distribution with smooth, Sub-Gaussian marginals, then an unrolled network based on AMP which is trained with gradient descent on polynomially many samples will converge in polynomially many iterations to an estimator which, in the infinite-dimen dimensional limit, achieves the same mean squared error Bayes AMP.
Our proof is based novel synthesis of state evolution, the fundamental distributional recursion driving analyses of AMP, together with neural tangent kernel (NTK) analysis of training dynamics for overparametrized networks. Crucially, unlike in typical applications of NTK analysis, the level of overparametrization needed in our network is dimension-indee indeeeddedd even when the second moment Exx-/|<[[| scales with the dimension d. The central reason behind this is that state evolution allows us to map the training dynamics of the network, which a priorí lives in (2(R"), to training dynamics over the space of functions (R), where the resulting learning problem amounts to that of one-dimensional score estimation. As result, our learning guarantee only requires overparametrization scaling inverse polynomially in the target error.
Experiments. We complement these theoretical results with extensive numerical experiments. Our main empirical contributions are as follows:
We demonstrate that our theoretically motivated unrolled network learns the same optimal denoisers as Bayes AMP, providing practical alternative that achieves the same performance but does not require explicit knowledge of the true signal prior.
We observe that introducing auxiliary trainable parameters along with learnable denoisers further improves performance over AMP in low-dimensional settings (where the asymptotic optimality of Bayes AMP does not apply) and when the sensing matrix is non-Gaussian, both in well- well-conditioned and ill-oonitiooee settings.
We introduce rank-one matrix estimation new "model organism" for probing the properties of unrolled networks. To our knowledge, despite its prominence in the theoretical literature AMP, rank-one matrix estimation has not yet been studied in the context of algorithm unrolling.
The general approach of unrolling with learned denoisers is lesser utilized in the algorithm unrolling literature, which instead largely emphasizes learning auxiliary parameters around domain-specific entities - e.g. measurement matrices or sparse coding dictionaries - while fixing denoisers typically to soft thresholding function. Our results indicate that learned denoisers can in fact capture distributional priors and are composable with these domain-specific learned parameters, providing a valuable addition to the algorithmic toolkit for practitioners of both AMP and unrolling.
1.1 Related work
We provide an extensive review of prior work in the appendix. Here we discuss the works most directly related to ours.
Theory for unrolling ISTA. The existing theory for algorithm unrolling almost exclusively focuses on unrolled ISTA (often called LISTA) and compressed sensing. Unlike the present work, they do not consider a Bayesian setting: the signal x is a deterministic sparse vector, and the goal is to converge to x. Instead of proving learning guarantees, most of them are representational in nature, arguing that under certain settings of the weights in LISTA, the estimator computed by the network can be more iteration iteration-efficient efficient than vanilla ISTA [MB16, CLWY18, XWG+16, LC19, CLWY21]. The works of [SARE23, SBR23, BRS22] proved generalization bounds for unrolled networks; these are statistical rather than computational in nature. Finally, recent work of [SPP*23] studied optimization aspects of LISTA, and their main theorem, motivated by the NTK literature [LZB20, L.Z222] from different perspective than ours, was an upper bound on the Hessian of the empirical risk of an unrolled ISTA network in neighborhood around random initialization.
3

---

Unrolled AMP. Borgerding et al. [BS16, BSR17] were the first to propose unrolling AMP for compressed sensing. In contrast to the architecture we consider, they primarily considered fixed soft-thresholding denoisers (st(;; 2), in addition to simple parametric families of denoisers like 5-wise linear functions and splines. For these parametric denoisers, on simple priors like Bernoulli-Gaussin they found that the learned network could approach the performance of the "oracle" estimator that knows the support of the underlying signal - see the Appendix for further discussion.
Other learning-based approaches. Here we further motivate the setting we consider by contrasting with other possible approaches to learning optimal inference algorithms from data. Perhaps the most obvious would be to simply try to directly learn an approximation to the density function for the prior, e.g. via kernel density estimation or some other non-paraeetric method. In the product-prior setting in which we prove our results, this is indeed viable approach in theory. But in practice, unlike algorithm unrolling, this will not scale gracefully to general high-dimensional distributions [GMOV19].
A more scalable approach might training a diffusion model on the data distribution [HD05]. The learned score network could then be used to approximately implement Bayes AMP. Our method is roughly special case of this: whereas little is known about provable score estimation in general [SCK23, GKL24, CKS24], our theoretical results demonstrate that layerwise training of unrolled networks is a viable, provably correct way to implicitly estimate the score of the data distribution. Furthermore, unrolling accommodates additional trainable parameters to improve robustness to real-world deviations from the stylized models studied in theory.
Finally, we mention the recent theoretical work of [IS24], which shows that semidefinite programs can simulate AMP. While this is not a learning result, it has similar motivation of reproducing the performance guarantees of AMP using a more robust suite of algorithmic tools.
Other theory for unrolling. [MW23] established sample complexity bounds for learning graphical models via diffusion models by unrolling the variational inference algorithms used for score estimation into a ResNet and bounding the number of parameters needed for the network to express these algorithms. Similarly, [Mei24] showed that the popular U-Net architecture can simulate wessage-passing algorithms. These works can be interpreted as giving representational guarantees for algorithm unrolling, whereas in contrast, the focus of our work is on proving learning guarantees.
2 Preliminaries on Bayes AMP and unrolling
Here we give an overview of the Bayes AMP algorithm in the compressed sensing setting. We refer the reader to Appendix C for a full treatment the rank-one matrix estimation setting. For convenience, when it is clear from context, we use Bayes AMP to refer to the general algorithm used in either setting.
2.1 Compressed sensing
In compressed sensing, we are given noisy linear measurements E R obtained from an unknown signal x Rd via the observation process
y = Ax +
(1)
where A E mxd and ş E Rt is random noise vector with i.i.d. entries drawn from the distribution N(O, o Throughout, we assume that x ~ px for product prior px A pdd where p is some distribution over R. The compressed sensing problem aims to recover the unknown signal x with estimate X such that the mean squared error (MSE) - is minimal. Throughout this work we focus on the proportional asymptotic setting where we implicitly work with sequence of such compressed sensing problems indexed by dimension d, where d and m = m(d) d) jointly tend to infinity and m(d)/d 8 for absolute constant 8 > 0.
[DMM09] proposed the following approximate message passing (AMP) algorithm for estimating given A, y. The algorithm starts with x 0 and un = y and proceeds by
X++1 f((A 02 + xx)
(2)
1
By = - Ax + 5 Jo-- (fr (A 0z-1 + xx(1))).
(3)
4

---

where ft: : R. - R is a scalar denoiser applied entrywise, ft is also applied entrywise, and (-) denotes an entrywise average. The last term in Eq. (3) commonly referred as the Onsager term. Importantly, the AMP iterates asymptotically satisfy distributional recursion called state evolution [BM11]. Suppose the entries of A are given by Aij N(O, 1 1/m). Define the state evolution parameters (I+) via the scalar recursion
1 F/11 = o + -I + 2 - =x) with = o² F[
where ~ P and Z ~ N)., Then is known that as - co, the empirical distribution over entries of A 0y + X converges in certain sense to the one-dimensional distribution over X + TZ [BM11]. While updates of AMP can be run with any choice of differentiable fo there is an asymptotically optimal choice that depends on the underlying prior px, and the resulting optimal algorithm is called Bayes AMP. In particular, let =0= 0² EE(22,, and define
fi = E|X + = and
[+1 0² El - -)
(4)
where X p and Z N(O, 1). Then setting fo fi² for all l in Eqs. (2) and (3) yields Bayes AMP.
In the asymptoric limit, i.e. as m, d - co, Bayes AMP has strong theoretical properties. In the setting above, it is conjectured to obtain the optimal MSE over all polynomial-time algorithms [BKM* 19] and has been proven to be optimal over quite general class of algorithms known as general first-order methods (GFOMs) [CMW20, MW22b].
In practice, however, Bayes AMP is subtly nontrivial to implement. For starters, one must know to construct f Furthermore, using the exact recursion in Eq. (4) can often lead the algorithm to diverge in finite dimensions. One instead estimates the state evolution parameters from the previous iterates, i.e. replacing r² with 1 11221122 which is typically enough to stabilize Bayes AMP. The fact that this is valid estimate follows by state evolution, which ensures that in the infinite dimensional limit, the entries of ü, are distributed according to N(O, r/) [BM11].
2.2 Unrolling Bayes AMP
The aforementioned challenges in realizing the conjectured optimality of Bayes AMP in practice motivate the need for a robust method that does not require knowledge of the prior distribution px. We consider replacing cach scalar denoiser fr in Eq. (2) or (22) with a multilayer perceptron (MLP) that learns the "right" denoiser function to use at each iteration of AMP. As we will see, prudent training approach is enough to provably ensure that our unrolled network learns the optimal denoiser at each layer, effectively recovering Bayes AMP even without explicit knowledge of the prior.
Architecture. Suppose we are given training data ((y', x') generated according to Eq. (1) with x² ~ Px for all i. Let denote the number of layers in our unrolled network, and let F denote family of MLPs with fixed architecture (i.e. fixed depth and width) constrained two-dimensional input and one-dimensional output. For each & [0, I. - 1], initialize an MLP
fr R - & chosen from F. Set Xu O. E Rd and ōy y E R" for given training input y' Then for each layer & E [0, L - 1], our network computes the forward pass
1 (o f +
Eo+1 = fi(A Uz + Xy; Tr) and
6p41
= AR +
(5)
where テ = 1/elle/ vm and o denotes differentiation with respect to the first input parameter.
The notation f(( ; denotes applying the scalar function fr( Tr) entrywise. We emphasize that fo is tied to à fif that is, we are taking the derivative of the MLP to compute the Onsager term. We refer to our unrolled architecture as an LDNet (Learned Denoising Network).
Training. Naively, one might consider training the I-laaee network end-to-end on the mean squared errors the network estimates - i.c., either with loss function Les = 1 V - N for compressed sensing or LME ΣN | 112 for rank-one matrix estimation.
5

---

However, as we observed empirically (echoed by findings in [MMB17]), such an approach gets trapped in suboptimal local assignments of denoising functions.
Instead, we employ layerwise training, where we iteratively train the (-th denoiser fo on the mean squared error loss for the layer-f estimate. If Y denotes an LDNet with layers, let V[O f] denote the subnetwork that consists only of the first 1 layers of Y, with denoiser fo at layer for 0 s f -- 1. Then our training procedure follows Algorithm 1. Note we initialize the e-th denoiser weights with the previous learned denoiser before training while this is not relevant to theoretical results in Section 3, empirically, to avoid being trapped suboptimal regions of parameters. Likewise, we include an optional finetuning step that further reduces approximation error in the learned denoisers but is not needed for our theory results.
find that this initialization is necessary
Algorithm 1: Layerwise Training
Input: Training data D, LDNet for == 0 to f L do
f 0 then
Initialize fr - fr-1
Freeze learnable weights in fk for f; Train P f] on D;
// Optional finetuning step
Unfreeze learnable weights in fa for k and train V[O: f] on D;
Output: Fully trained y
The proof of optimality of Bayes AMP among all implementations of AMP for the problems we consider, given in [CMW20, MW22b], strongly motivates our training method: assuming we have learned optimal denoisers up to layer ( 1, one can show that the minimum mean squared error layer f is achieved by the denoiser used in Bayes AMP. This gives a heuristic sense for how layerwise training facilitates learning optimal denoisers, and this intuition validated in both our theory and experiments.
Provably learning Bayes AMP
We now provide theoretical guarantees that our unrolled denoising network can learn Bayes- optimal denoisers when trained in a layerwise fashion. Consider any prior Px pdd for which satisfies the following assumption. The product prior setting is quite standard and widely studied within the theory literature on AMP (e.g. [BM11, DAM16)).
Assumption 1. Given 2 0, let p(:; denote the density of the convolution * N(O, We assume that:
p R-sub-Gaussian with Ex-p[X] 0.
The score function T) is B-Lipschitz for all o², where o² is the variance of the entries of in Eq. (1).
Both assumptions are relatively mild and hold for a large class of distributions. For example, the uub-Gaussianity holds for any distribution with bounded support (see Section 2.5 [Ver18]) and the Lipschitzness of the score function is consequence of regularizing properties of heat flow (see& e.g. Lemma 4 in [MS23].
Our main guarantee (see Theorem 2 below) is that under Assumption 1, suitable unrolled architecture trained with SGD on examples of compressed sensing tasks can compete with Bayes AMP. In Section 3.1, we define the training objective and architecture, describe how our bounds will depend on the underlying prior p, and formally state our main result. The full proof provided in Appendix B.

---

3.1 Proof preliminaries and theorem statement
To prove our learning guarantee, we start by proving that the error in using the learned denoiser in one unrolled layer is small. Denote the sequence generated by the learned denoiser f by X and 8y. The Onsager term AMP ensures that for any iteration f, the distribution of A' De + &t asymptotically behaves as if every coordinate is i.i.d. as d - co for any l (Lemma 1). Therefore, suffices to learn the denoiser for any fixed coordinate." Without loss of generality we consider the first coordinate and try learn the denoiser function by minimizing the following objective:
min E[ + +x )].
(6)
where A] denotes the jth row of A and x(1) denotes the jth coordinate of the vector As the learning and generalization guarantee is identical for all e, we will occasionally drop & from the subscript when the context is clear.
Denoiser complexity. To quantify the complexity of learning Bayes AMP in terms of the underlying prior p, we will work with the following notion of the complexity of class of scalar functioos from [AZLL19], which we will apply to the denoisers that arise in Bayes AMP:
Definition 1 (Scalar function complexity [AZLL19])). Let C > 0 be some sufficiently large absolute constant (e.g. 10*). Given any smooth function & R R and a parameter
define
complexity of at scale a as follows. Suppose & admits power series expansion (zz) Σi= o, Then
a)) A
Σ [1+ (log (log(1/))/))33) (Ca)"led and (, a) & C 1=0
>
(i+ (i+ 1.75 a lc/c.l
1=0
Given a class of scalar functions F, we define the complexity of F at scale a by (F, suP&EF (g a) (and similarly (F, a) suppe C,( ot).
Intuitively, (F(F.a) and C,(F, a) both captures how much functions in function class F can be approximated using low-degree polynomial. For any function and any a, the above complexities are related by Cg(g, a) x (g) s (E,(0,(()) x poly(1/e) because (CvIog(1/e)/Vl Vi)' Olog 1/e) poly(1/E) for all i. We provide more intuition on how (g) a) and (, a) scales with e, a, under mild assumptions on & in Section B.5.
Main result. We can now formally state the main theoretical guarantee of this work, namely that layerwise training of LDNet results in performance matching that of Bayes AMP for compressed sensing:
Theorem 2. Suppose the prior distribution satisfies Assumption 1. Then, for every E2 E (0, 1) and E E 10, (F, 1 there exists
Mo = po (f"," 1/e2) 3/2)) 1/E1) and No = Pf R(log 1 1/ 3122 1/t))
such that the following holds.
Let be any positive integer. Consider an LDNet of depth with MLP denoisers fe given by the MLP architecture in Eq. (8) with m 2 Mo neurons. Suppose the network is trained by running gradient descent from random initialization with step size n (011/(mm)) the form x"), where each training example is generated by independently sampling Gaussian matrix A with entries ii.d. from N(O, 1/m), sampling x² Px pd, and forming y Axf g for s N(0, o² Id).
2 No samples of
II our experiments, learn the denoiser using all coordinates, but in our theoretical result, focus on learning using only single coordinate. The latter is less sample-efficien but more convenient for our proof. It should be possible to prove a guarantee for the later, but is more cumbersome as we need to prove that the samples obtained by different coordinates are close to being i.i.d. for some notion of closeness, and then generalize the result of [AZLL19] to allow for such samples.
7

---

After T = , R(log 1/22)3/2)31E2) steps of gradient descent, with high probability the activations (H 51.) and denoiser fi. at the output layer of the LDNet (see (5)) satisfy
Ex.Al I + XLi TL
xx|²] X MSEAMP +
B2
1+1 (E1 + (1),
80
where MSEAMP (L) A Ex, I u1 x.) - -x1?²] is the error achieved by running steps of Bayes AMP.
Observe that the level of overparametrization in terms of number of samples n and number of hidden neurons needed in Theorem is dimension-free, unlike in typical NTK analyses. This happens because state evolution effectively allows us to convert the learning problem dimensions
learning problem in dimension: we can effectively assume that the entries of A A 0y+ + Xy converge in an appropriate sense to the distribution of X + 7 for X ~ and (0, 1).
- =) This ensures that the learning objective effectively reduces to minimizing E[ r,2) - parametrized family of denoisers fr. The latter objective is often referred to the score matching objective (in one dimension), which minimized by the Bayes-optimal denoiser f = =| X + 7 =] at each layer f. key component in the proof of Theorem 2 is thus to show that gradient descent can learn this optimal denoiser given one-dimensional training data of the form (X+ + zzZ, X).
As we will show, the runtime for gradient descent is largely dictated by the extent to which these denoisers can be polynomially approximated.. A priori, one might expect that if degree-s polynomials are needed, then the runtime of the algorithm must scale as dOss) This would be prohibitively expensive if is increasing in the dimension d. Fortunately however, because we are able to reduce to miedimmnsional training dynamics, ultimately achieve much more favorable scaling in
4 Experiments
We now empirically demonstrate the performance of our proposed architecture and training scheme for unrolling Bayes AMP in a variety of statistical settings. Throughout these experiments, we are motivated by the following questions: a) Can our method empirically match the performance of Bayes AMP in settings where the latter is conjectured be computationally optimal? b) In these settings, does our network learn the optimal denoisers? c) Are there settings where our methods offer performance advantage over AMP?
4.1 Compressed sensing
Implementation details. We
250, d 500 and fix random Gaussian sensing matrix A E 250000000 We consider two choices of prior for our experiments: Bernoulli- Carnoulli-Gausian and Z2 (i.e. uniform over (1)). For our unrolled architecture, the family F of learned MLP denoisers was restricted to three hidden layers, each with 70 neurons and GELU activations. This particular architectural choice was the most convenient for our experiments, but our experimental findings are not particularly sensitive to this. We randomly generated a train and validation [=l with N = 215 samples by sampling from the prior and using Eq. (1). We train layerwise with finetuning in Algorithm 1.
dataset (y', x
For each prior, we also implemented Bayes AMP using the corresponding optimal denoiser. As an additional "semi- "ssemi-prior-eee arior-aware" baseline, replace the MLP denoisers in LDNet with "guided denoisers" that have the same functional form as the optimal denoisers but contain trainable parameters; Appendix for more details the precise functional forms used. By convention, we report performance results for all methods by the normalized mean squared error (NMSE) || - x |
Bernoulli-Gaussian prior. Here, each entry x is independently drawn from a standard normal distribution and set to with probability 1 - E; i.e., Px p where p(x) e N(O,, 1; (1 E) 8(x), where 8 denotes the Dirac delta 0. To match the setting considered in the

---

Bemnull-Gaussian
Gurted Denosess Bayee
LOWet
0.65
LOCaE Guikel Donokers Sayes
-10
0.80
NV00 (530
-20
NMSE
-30
-40
Figure 1: LDNet for Compressed Sensing. On the left, plot the NMSE (in dB) obtained by LDNet and Bayes AMP baselines on the Bernoulli- Gaussian prior. the right, we plot NMSE (not in dB) achieved on the 22 prior LDNet (along with the guided denoisers) achieves virtually identical performance to the conjectured computationally optimal Bayes AMP.
prior work of Borgerding et al. [BSR17], set the masking probability to be 0.1 and the measurement noise to be o² 2 10-5
We plot the NMSE that our unrolled network and baselines achieve in decibels (dB); that is, 10 (NMSE), in Figure 1. LDNet almost perfectly matches the NMSE Bayes AMP at each layer/iteration. Over 15 layers, our network converges to an NMSE of -44.9313 dB, as compared to Bayes AMP converging to -45.3280 dB. As the scale logarithmic, the difference in error achieved is negligible.
Ze prior. Here each entry of x chosen from {-1,1} with probability 1; i.e., Px = ped for p(x) ( To examine higher noise regime and to ensure Bayes AMP converged within reasonable number of iterations, we set the measurement noise be o² 0.075. Figure demonstrates that LDNet again recovers Bayes AMP at every iteration, even slightly outperforming by layer 15, achieving an NMSE of 0. 0.4266 (an improvement of 1.28%)
LDNet denoisers. From Figure we can observe qualitatively that the learned MLP denoisers recover the functional form for the optimal denoiser each iteration, Interestingly, although the denoisers were trained relative to fixed sensing matrix, they appear to learn the Bayes AMP denoiser that is ceeasurement-indeeeenee show that the performance of these learned denoisers actually transfers to other randomly drawn sensing matrices A.
theory suggests.
and Appendix D.3 we
sater
RSF REFIRI
-33
114
-10
+1.0
+15
-12
PJ
-025
-350
-30
10.5
4.5
Figure 2: Learned Denoisers for Compressed Sensing. We plot layerwise denoising functions learned by and Z, priors relative to their optimal denoisers over range of inputs in (-2, 2). The state evolution input t each denoiser is to be its empirical estimate.
LDNet on the
9

---

4.2 Beyond Bayes AMP performance
Much of the algorithm unrolling literature focuses on learning auxiliary parameters while using fixed denoisers, as opposed to learning the denoisers themselves. In particular, unrolled methods like LISTA [GL11] and LAMP [BSR17] reparameterize A Eqs. (2) and (3) as a new learnable matrix B, which results in faster convergence than classical, learning-free counterparts like ISTA and AMP.
This does not necessarily contradict Bayes AMP's conjectured optimality for compressed sensing, which only applies in the d - limit. In the context of our unrolling method, posit that learning the matrix B can be thought of as learning finite dimensional corrections to the Bayes AMP iterations. Appendix D.1, we demonstrate that the lower the signal dimensionality, the larger the performance improvement of LDNet with learned matrix B over Bayes AMP.
Finite dimensionality is not the only deviation in design from where Bayes AMP is (conjectured) optimal. In fact we can consider non-Gaussian designs the measurement matrix A, and we show in D.1 that LDNet outperforms Bayes AMP in both well-conditioned and ill-conditioned settings. Furthermore, we can relax the assumption that the signal is drawn for a product prior and extend LDNet accommodate non-product priors. In D.2, we demonstrate LDNet surpass Bayes AMP in a non-separable cexture-of-ggussiinn prior.
5 Outlook
In this work we gave the first proof that unrolled denoising networks can compete with optimal prior-aware algorithms simply via gradient-based training on data. Our proof used a novel synthesis of state evolution with NTK theory, and notably, the level of overparametrizatio needed for our result hold is independent of the dimension, unlike existing results in the NTK literature. One important consequence of these results is that one-dimensional score function learnable with gradient descent, for which only representational, as opposed to algorithm, results existed previously in the literature [MW23].
We supplemented recover Bayes AMP performance and Bayes-optimal denoisers without knowledge of the signal prior. Moreover, for various settings where Bayes AMP is not conjectured to perform optimally e.g. inference in low dimensions, non-Gaussian designs, and non non-prrduct priors - we demonstrate that I.Nett outperforms Bayes AMP. We thus establish unrolling denoisers powerful, practical addition to the algorithmic toolkit for Bayesian inverse problems.
theory with extensive numerical experiments, confirming that LDNet can
One limitation is that our theoretical results are currently limited to the product prior setting. The non-pro non-product setting is difficult because even though state evolution is known here [BMN20], proving unrolled network converges to the right denoisers essentially amounts to proving that one can learn the score functions of general data distribution. Additionally, it is subtle to define the right architecture, the denoisers are longer scalar, and generic feedforward architecture would be difficult to prove rigorous guarantees for (and to scale in practice).
In addition, our theoretical results do not immediately extend to the rank-one matrix estimation setting. While closeness denoising error implies closeness in the state evolution parameter r for compressed sensing, this is not immediate for rank-one matrix estimation, where multiple choices for parameters μ and lead to the same denoising error. This is reflected in Figure 4, where the learned denoisers at early iterations achieve the same MSE the Bayes optimal denoisers, but the functional forms are completely different. We leave the extension of our compressed sensing results to rank-one matrix estimation as an open question.
Finally, while our experiments suggest that including auxiliary trainable parameters like the "B matrix" offers significant performance performacntrro advantages once one departs from the asymptotic, Gaussian design setting in which Bayes AMP is believed to be optimal, these are not yet supported by theory. It is an intriguing open question whether one can use some of the insights from the aforementioned representational results for ISTA to rigorously characterize the thom-asymptotic corrections" that these extra learnable parameters are imposing.
10

---

Acknowledgments
AK and SC thank Demba Ba for illuminating discussions about unrolling at an early stage this project. KS and SC thank Vasilis Kontonis for many fruitful conversations about provable score estimation. AK supported by the Paul and Daisy Soros Fellowship for New Americans. KS is supported by the NSF AI Institute for Foundations of Machine Learning (IFML).
References
[AZLL19] Zeyuan Allen-Zhut Yuanzhi Li, and Yingyu Liang. Learning and generalization in overparameterized neural networks, going beyond two layers, Advances in neural information processing systems, 32, 2019.
[BBDM21] Mario Bertero, Patrizia Boccacci, and Christine De Mol. Introduction to inverse problems in imaging. CRC press, 2021.
[BCPS21] Jean Barbier, Wei-Kuo Chen, Dmitry Panchenko, and Manuel Sáenz. Performance of bayesian linear regression in a model with mismatch. arXiv preprint mrxiv:2107.06936, 2021.
[BHMS22] Jean Barbier, TianQi Hou, Marco Mondelli, and Manuel Sáenz. The price of ignorance: how much does it cost to forget noise structure in low-rank matrix estimation? Advances in Neural Information Processing Systems, 15:36733-36744 2022.
[BKM*19] Jean Barbier, Florent Krzakala, Nicolas Macris, Léo Miolane, and Lenka Zdeborová. Optimal errors and phase transitions high-dimensional generalized linear models. Proceedings of the National Academy of Sciences, 116(12):545-54460 2019.
[BM11] Mohsen Bayati and Andrea Montanari. The dynamics of message passing on dense graphs, with applications to compressed sensing. IEEE Transactions on Information Theory, 57(2):764-785, 2011.
[BMN20] Raphael Berthier, Andrea Montanari, and Phan-Minh Nguyen. State evolution for approximate message passing with non-scparable functions. Information and Inference: A Journal of the IMA, 911::33-7, 2020.
[BPW18] Afonso S Bandcira, Amelia Perry, and Alexander S Wein. Notes on computational- to-statistical gaps: predictions using statistical physics. Portugaliae mathematica, 75(2):159-186, 2018.
[BRS22] Arash Behboodi, Holger Rauhut, and Ekkehard Schnoor. Compressive sensing and neural networks from statistical learning perspective. In Compressed Sensing in Information Processing, pages 247-277. Springer, 2022.
[BS16] Mark Borgerding and Philip Schniter. Insager-correcte deep learning for sparse linear inverse problems. In 2016 IEEE Global Conference on Signal and Information Processing (GlobalSIP), pages 227-231. IEEF, 2016.
[BSR17] Mark Borgerding, Philip Schniter, and Sundeep Rangan. Amp-inspired deep networks for sparse linear inverse problems. IEEE Transactions on Signal Processing, 65(16))4293-433) 4308, 2017.
[BT09] Amir Beck and Marc Teboulle. A fast iterative shrinkagg-thressooddngg algorithm for linear inverse problems. SIAM journal on imaging sciences, 2(1): 2009.
[CKS24] Sitan Chen, Vasilis Kontonis, and Kulin Shah. Learning general gaussian mixtures with efficient score matching. arXiv preprint arxiv:2404. 18893, 2024.
[CIWY18] Xiaohan Chen, Jialin Liu, Zhangyang Wang, and Wotao Yin. Theoretical linear convergence of unfolded ista and its practical weights and thresholds. Advances in Neural Information Processing Systems, 31, 2018.
11

---

[CLWY21] Xiaohan Chen, Jialin Liu, Zhangyang Wang, and Wotao Yin. Hyperparameter tuning is all you need for lista. Advances Neural Information Processing Systems, 34:11678- 11689, 2021.
[CMW20] Michael Celentano, Andrea Montanari, and Yuchen Wu. The estimation error of general first order methods. In Conference on Learning Theory, pages 1078-1141. PMLR, 2020.
[DAM16] Yash Deshpande, Emmanuel Abbe, and Andrea Montanari. Asymptotic mutual information for the binary stochastic block model. In 2016 IEEE International Symposium on Information Theory (ISII), pages 185-189, 2016.
[DDDM04] Ingrid Daubechies, Michel Defrise, and Christine De Mol. An iterative thresholding algorithm for linear inverse problems with sparsity constraint. Communications on Pure and Applied Mathematics: Journal Issued by the Courant Institute of Mathematical Sciences, 57(11):1413-1457 2004.
[DM14] Yash Deshpande and Andrea Montanari. pea. 201 IEEE International Symposium on Information Theory, pages 2197-2201. IEEE, 2014.
optimal sparse
[DMM09] David L Donoho, Arian Maleki, and Andrea Montanari. Message-passing algorithms for compressed sensing. Proceedings of the National Academy of Sciences, 106(45):18914-1819 2009.
[GKL224] Khashayar Gatmiry, Jonathan Kelner, and Holden Lee. Learning mixtures of gaussians using diffusion models. arXiv preprint arXiv:2404. 18869, 2024.
[GL10] Karol Gregor and Yann LeCun. Learning fast approximations of sparse coding. In Proceedings of the 27th international conference on international conference on machine learning, pages 399-406 2010.
[GMOV19] Weihao Gao, Ashok V Makkuva, Sewoong Oh, and Pramod Viswanath. Learning one-hidden-layer neural networks under general input distributions. In The 22nd International Conference Artificial Intelligence and Statistics, pages 1950-1959. PMLR, 2019.
[HD05] Aapo Hyvärinen and Peter Dayan. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research, 6(4), 2005.
[IS24] Misha Ivkov and Tselil Schramm. Semidefinite programs simulate approximate message passing robustly. In Proceedings of the 56th Annual ACM Symposium Theory of Computing, pages 348-357, 2024.
[TTW19] Daisuke Ito, Satoshi Takabe, and Tadashi Wadayama. Trainable ista for sparse signal recovery. IEEE Transactions on Signal Processing, 67(12):3113-31 3125, 2019.
[Jac1]] Dunham Jackson. Über die Genauigkeit der Annäherung stetiger Funktionen durch ganze rationale Funktionen gegebenen Grades und trigonometrische Summen gegebener Ordnung. Dieterich's scchnnnhssnnnnn Uhiversität-Bucchruukk
1911.
[LC19] Jialin Liu and Xiaohan Chen. Alista: Analytic weights are as good as learned weights in lista. International Conference on Learning Representations (ICLR), 2019.
[LLT22] Holden Lee, Jianfeng Lu, and Yixin Tan. Convergence for score-based generative modeling with polynomial complexity. Advances in Neural Information Processing Systems, 05:22870-22882 2022.
[LSSS14] Roi Livni, Shai Shalev-Shwartz and Ohad Shamir. On the computational efficiency of training neural networks. Advances in neural information processing systems, 27, 2014.
[LTG+20] Yuelong Li, Mohammad Tofighi, Junyi Geng, Vishal Monga, and Yonina C Eldar. Efficient and interpretable deep blind image deblurring via algorithm unrolling. IEEE Transactions on Computational Imaging, 6:666-681, 2020.
12

---

[LZB20] Chaoyue Liu, Libin Zhu, and Misha Belkin. On the linearity of large non-linear models: when and why the tangent kernel constant. Advances in Neural Information Processing Systems, 33:15954-15964, 2020.
[1.ZB22] Chaoyue Liu, Libin Zhu, and Mikhail Belkin. Loss landscapes and optimization in over- parameterized non-linear systems and neural networks. Applied and Computational Harmonic Analysis, 59:85-116, 2022.
[MB16] Thomas Moreau and Joan Bruna. Understanding trainable sparse. coding via matrix factorization. arXiv preprint arXiv:1609.0028: 2016.
[Mei24] Song Mei. U-nets as belief propagation: Efficient classification, denoising, and diffusion in generative hierarchical models. arXiv preprint arX1v:2404.18444, 2024.
[MJC21] Osman Musa, Peter Jung, and Giuseppe Caire. Plug-and-ple learned gaussian- mixture approximate message passing. In ICASSP 2021-2021 IEEE International Conference Acoustics, Speech and Signal Processing (ICASSP), pages 4855-4859 IEEE, 2021.
[MLE21] Vishal Monga, Yuelong Li, and Yonina C Eldar. Algorithm unrolling: Interpretable, efficient deep learning for signal and image processing. IEEE Signal Processing Magazine, 38(2):18 44, 2021.
[MMB16] Christopher A Metzler, Arian Maleki, and Richard G Baraniuk. From denoising to compressed sensing. IEEE Transactions Information Theory, 62(9):5117-5144, 2016.
[MMB17] Chris Metzler, Ali Mousavi, and Richard Baraniuk. Learned d-amp: Principled neural network based compressive image recovery. Advances in neural information processing systems, 30, 2017.
[MS22] Sumit Mukherjee and Subhabrata Sen. Variational inference in high-dimensional linear regression. Journal of Machine Learning Research, 23(304):1-56, 2022.
[MS23] Dan Mikulincer and Yair Shenfeld. On the lipschitz properties of transportation along heat flows. In Geometric Aspects of Functional Analysis: Israel Seminar (GAFA) 2020-2022, pages 269-290. Springer, 2023.
[MV21] Andrea Montanari and Ramji Venkataramanan. Estimation of low-rank matrices via approximate message passing. The Annals of Statistics, 49(1):321-345, 2021.
[MW22a] Andrea Montanari and Alexander S Wein. Equivalence of approximate message passing and low-degree polynomials in rank-one matrix estimation. arXiv preprint arXiv:2212.06996, 2022.
[MW22b] Andrea Montanari and Yuchen Wu. Statistically optimal first order algorithms: A proof via orthogonalization. arXiv preprint arXhv:2201.05101 2022.
[MW23] Song Mei and Yuchen Wu. Deep networks as denoising algorithms: Sample-efficient learning of diffusion models nigh-dimensional graphical arXiv:2309.1 1110009.11140 11400, 2023.
graphiiaa models. arXiv preprint
[NT10] Deanna Needell and Joel Tropp. Cosamp: iterative signal recovery from incomplete and inaccurate samples. Communications of the ACM, 53(12)993-100, 2010.
[RF12] Sundeep Rangan and Alyson Fletcher. Iterative estimation of constrained rank-one matrices in noise. In 2012 IEEE International Symposium on Information Theory Proceedings, pages 1246-12 12500 IEEE, 2012.
[SARE23] Avner Shultzman, Eyar Azar, Miguel RD Rodrigues, and Yonina C Eldar. Generalization and estimation error bounds for model-based neural networks. arXiv preprint
2023.
13

---

[SBR23] Ekkehard Schnoor, Arash Behboodi, and Holger Rauhut. Generalization error bounds for iterative recovery algorithms unfolded as neural networks. Information and Inference: A Journal of the IMA, 12(3):2267-2299, 2023.
[SCK23] Kulin Shah, Sitan Chen, and Adam Klivans. Learning mixtures of gaussians using the ddpm objective. Advances in Neural Information Processing Systems, 36:19636-1969, 2023.
[She12] Alexander A Sherstov. Making polynomials robust noise. In Proceedings of the forty-fourth annual ACM symposium on Theory of computing, pages 747-758, 2012.
[S) Shaik Basheeruddin Shah, Pradyumna Pradhan, Wei Pu, Ramunaidu Randhi, Miguel RD Rodrigues, and Yonina Eldar. Optimization guarantees of unfolded ista and admm networks with smooth soft- nof-thhresholdig. arXiv preprint arXiv:2309.06119, 2023.
[ST99] Roel Snieder and Jeannot Trampert. Inverse problems geophysics. Springer, 1999.
[SWED23] Nir Shlezinger, Jay Whang, Yonina C Eldar, and Alexandros G Dimakis, Model-based deep learning. Proceedings of the IEEE, 2023.
[TAP77] David Thouless, Philip W Anderson, and Robert G Palmer. Solution solvalee model of a spin glass'. Philosophical Magazine, 35(3):593-601, 1977.
[Ver18] Roman Vershynin. High-Dimensional Probability: An Introduction with Applications in Data Science. Cambridge Series Statistical and Probabilistic Mathematics. Cambridge University Press, 2018.
[Vog02] Curtis R Vogel. Computational methods for inverse problems. SIAM, 2002.
XWG+16] Bo Xin, Yizhou Wang, Wen Gao, David Wipf, and Baoyuan Wang. Maximal sparsity with deep networks? Advances in Neural Information Processing Systems, 29, 2016.
14

---

A Further related work
Theory for unrolling ISTA. The existing theory for algorithm unrolling almost exclusively focuses unrolled ISTA (often called LISTA) and compressed sensing. The focus of these works is rather different from ours. For starters, they do not work in Bayesian setting: the signal x is deterministic vector assumed to have some level of sparsity, and the goal is to converge x. As mentioned above, these results do not prove convergence guarantees for the learning algorithm. Instead, most of them argue that under certain settings of the weights in the LISTA architecture (and variants), the resulting estimator computed by the network can be more deraatin-efffcieent than vanilla ISTA [MB16, CLWY18, XWG 16, IC19, CLWY21].
For example, [MB16] showed that if each layer performs a proximal splitting step with respect to a Gram matrix which admits factorization with certain nice properties, then the itcrates computed by each layer converge to x at an accelerated rate. [CLWY18] showed that if the learned weights are such that the activations converge the ground truth vector, then they have to have a certain structure; under a specialized LISTA architecture that imposes this structure, they prove there is setting of weights for which the iterates converge a linear rate. [XWG*16] showed a similar result for iterative hard thresholding and argued that the learned parameters can potentially reduce the RIP constant of the sensing matrix and thus speed up convergence. [1C199] identified a certain tied parametrization of the weights for LISTA that can also achieve linear convergence with fewer trainable parameters (see also the follow-up [CLWY21]).
Apart from these works, the works of [SARE23, SBR23, BRS22] proved generalization bounds for unrolled networks. These guarantees pertain to questions of sample complexity for empirical risk minimization, instead of computational complexity of learning these networks via gradient descent, and are thus orthogonal to the thrust of our work.
Finally, recent work of [SPP*23] studied optimization aspects of LISTA, and their main theorem was a bound on the Hessian of the empirical risk of an unrolled ISTA network in a neighborhood around random initialization. Under the unproven condition that the associated NTK at initialization is sufficiently well-conditioned, this would imply that the empirical risk satisfies inequality around initialization and thus the network would
a modified converge exponentially quickly to the empirical risk minimizer. While we also draw upon tools from the NTK theory, our focus is not just on optimizing the empirical loss, but on proving that the learned network generalizes to achieve mean squared error competitive with the theoretically optimal prior-aware algorithm, AMP. In addition, our results our end-to-end and apply to unrolled AMP instead of unrolled LISTA.
Learned AMP. Borgerding et al. [BS16, BSR17] were the first to propose unrolling AMP for compressed sensing. In contrast to the architecture we consider, they primarily considered fixed soft-thresholding denoisers ( A) with trainable parameter À and trainable weight matrices playing the role of AT in the AMP update (see Eqs. (2) and (3)). They found empirically that their unrolled network outperforms AMP with soft- coft-thresholding denoisers. They also considered unrolling vector AMP, a more powerful version of AMP, and showed that even when the sensing matrix is ill-conditioned, the network essentially matches the performance of vector AMP. They also considered some simple parametric families of denoisers like 5 5-wise linear functions and splines and found that they could approach the performance of the "oracle" estimator that knows the support of the underlying signal.
Comparison to LDAMP [MMB17]. Among the various direct follow-ups to [BS16, BSR17], e.g. [ITW19, MJC21, MMB17J, the most relevant to ours is the follow-up work of [MMB17] extended this to denoisers given by convolutional neural networks of nontrivial depth (roughly 20 layers). They experimentally demonstrated that these networks performed quite well on compressive image recovery tasks. Using proof technique of [MMB16] and under a certain monotonicity assumption on the score functions of the data distribution, they showed that Bayes AMP is Bayes optimal (see Lemma 1 therein).
While the fact that they employ a generic architectures for the denoiser step in AMP and demonstrate the effectiveness of the resulting unrolled architecture is closely related to the spirit of the present work, there are important differences. We note that their theoretical result does not not explain how unrolled AMP, when trained on data with gradient descent, learns to
15

---

compete with Bayes AMP, only that in certain situations, optimally tuning the denoisers in AMP can achieve Bayes optimality. Furthermore, the monotonicity assumption they make restrictive (e.g. it does not even apply to the two-point prior given by the uniform distribution over (+1}). Furthermore, in general Bayes AMP need not be Bayes optimal, i.e. in situations where there is computational-statiiicca gap [BPW18].
On the experimental side, our focus was on synthetic setups instead of image recovery, with emphasis on probing which aspects of the problem setup and which trainable parameters allow unrolled AMP outperform Bayes AMP.
B Proof of main theorem
B.1 State evolution and learner function
The following result shows that we can characterize the behavior of the AMP iterates Xy in the limit as d co.
Lemma 1 (Asymptotic characterizatton of AMP iterates [BM11]). Let A R maxd be sensing matrix with i..dd entries Aij ~ N(0, 1/m). Assume m/d - 8 E (0, oo). Consider a sequence of vectors ((d))n(d)) indexed by dimension whose empirical distribution converges weakly to probability measures Px and Py on R with bounded moments. Then, for any pseudo-Lipschitz function & I and all 2 O, almost surely
d
lim Σ
()) = +, Mxx
d
lim
(x (0) (A] xx1)) (1)) =[M(T,
doo
where X ~ Px and Z NNO, 1). The state evolution parameters To, T.... are recursively defined as follows
o = 0² Exxp,,
8
2,1 = 0² + Exxp.2-NNOO -X
8
Observe that to minimize the variance r? at each iteration, the optimal choice of denoiser fr is fr (x) E[X rzZ ==Z=x] x
B.2 Learning guarantee for single layer of unrolling
Here we prove learning result for one layer of unrolled AMP, which we later extend to give an end-to-end learning result for training the full unrolled network.
Learner function. We parametrize the scalar denoiser in given layer of our unrolled AMP architecture as a mee-hidden-layee ReLU network in the following form: letting w denote the weight of the jth neuron after steps gradient descent, we consider
f(x; 0,)
Σ a ReLU( x b,),
(8)
We initialize the entries of weights w/) and biases //01 to be i.i.d. from NN00, 1/m) and entries of a [0) to be i.i.d. from (0 Ea) for some fixed &a (0, 1]. We only train the weights w, of hidden layers to simplify the analysis and freeze the bias term by and output layer weights aj at initialization. To update the weights at time take step of online gradient descent
16

---

with respect to the loss in Eq. (6) with step size on a fresh sample (x, A); here we use the learned denoisers from the previous layers to compute x and Üy.
We begin by proving the following claim that training a layer of the unrolled network on gradient descent from random initialization will converge optimal denoiser under the objective in Eq. (6).
denoiser that is competitive with the
Lemma (Learning the denoiser within L2 error). For every t, assume every coordinate of गe and X are sub- Sub-Gaussian random variables with constant R. Then, for every E2 (0, 1) and E (0, (F, (R,,oog (/222)()))) there exists
Mo = poly((,, (fi', R(log 1 312)) 1/E1) and No = poly R(log 1/ 3/22)312),11 1/ai)
such that for every m 2 Mo and n 2 No, choosing learning rate r = 6(1/(e)m) and rumning gradient descent from random initialisation for T = @ R(log 1/22) /222)332222222/ E with high probability,
Ex,y, ( ;,) - -x s
[[ - () + + ge
min
This is a consequence of the following result on training neural networks in the NTK regime:
Lemma 3 (Theorem 1 of [AZLL19]*). Consider a target function F R2 - R of the following form
F*(x) Σ a (w)
where each + R + R is infinite-order smooth and weights satisfy || |W S 1 and |a;| s 1. Additionally, assume that ||x|| s B. Then, for every E E (0, ( B))), there exists Mo Poly((,, B), 1/e) and No = poly(C,,($ B), 1/e) such that for every m 2 Mo and n 2 No. choosing learning rate n ((1/(em) and running gradient descent from random initialization for T = (6 2) p² / )e,,, with high probability,
E[(N(x;0) 0r) - y s inf El(((x) - y) t gCF
Additionally, the absolute value of the newral network is bounded by IN 0,1 S (G,,(,) for all x with ||x|| s B for all t.
Proof of Lemma 2. As the distribution over 0, is R-sub Suuu-GGasssin, for a fixed A, we have |A s R||A,| Vlog( 1/c2) with probability at least 1 - Additionally, because of A. (N ld/ m), we have || & log(1/E2) with probability at least 22. Combining both bounds, we have a $ ((log(1/e2))) 3/2 Similarly, using sub-Gaussianity of &,, we have (xx( S R (log(1//2)). This gives us that [A, o + x(01 & R(log(1/ 1/22)) 3/2 with probability at least E2-
By only considering samples satisfying [AJ 5y x S ((loo(1//2)))1 we can apply Lemma to obtain neural network that achieves error.
'In our experiments, we keep the measurement matrix A fixed and only sample fresh (x, y) pairs, but in our theoretical result, we assume that gradient descent uses fresh samples (x. y. to avoid technical difficulties regarding dependencies between the errors at different layers of unrolled architecture. We expect that with more work, one can extend the proof to fixed A.
*Here state the result in terms of gradient descent instead of stochastic gradient descent as in [AZLL19] However, the same proof of [AZLL19] goes through upon slightly modifying Lemma B.4 therein. Lemma B.4, we can write [| W' [|W,- V Z, W.) w'1l and therefore, getting the equality of (w; W VL Z, W)) (I]W;- W 1 ||W /1/))/2 (1/2))/VL((Z Will? and using this equality in (B.7) of [AZLL19]. The rest of the proof remains the same.
5The reason we can do this that this condition fails hold only with probability at most #2, and whenever it fails to hold, can output O and pay additional [[[22] Alternatively, could also modify the learner network to implement the indicator of [A, 0 + Ryil S R(log(1/e2)) using an appropriate linear combination of ReLU activations without learnable parameters.
17

---

Next, we relate the error in Lemma 2, which is for predicting the first coordinate of the signal given the noisy estimate provided by the previous layer of the unrolled network, the optimal error for predicting a sample from the univariate prior given a Gaussian corruption. This will follow by state evolution.
Lemma 4. Let fe(:,0) 0₂) be the learned neural network using gradient descent after timesteps such that the conditions of Lemma 2 satisfies. Then, with high probability, we have
lim Ex,A I +x,00,00))))))) x||²| s EIGE (X - + d=000
Additionally, the following statement holds with high probability:
+ x) EX
Proof. Observe that A + z()) follows the same distribution for all j. Therefore, using Lemma 2,
we obtain that
Ex.A Ex.A Ss.Al(f(AAJo x 07)
01) - -x(0)²] min E[ - + E2.
As the distribution of A7 + x(1) is the same for all j [d], the same learned denoiser for all the coordinates. Therefore, we can rewrite the above equation as
Σ ExA EIF(A xx x||2| s min d
EEllo((A + Rr) x1| +
Now, we want to use state evolution as d co. Note that d co, using Lemma with & function as (aa (a b) we have
lim 2 (lg(A oe + Rx) - x|12 E[(g(X x) d-000
for any function E F. As the quantity inside expectation converges to E[(g(X+ x)
using monotone convergence theorem, we have
lim min El 7 [l9(A xx) - x11 d
min lim El [1g 0. + x1|
9 d
min El 7,Z) x)
9
Using this result and applying limits on both sides Lemma 2, we obtain
lim Es.A | IIf(AA θ. + Rr. 07)
x|| min E[(g(X + ( +77]+1++8
The minimizer of (l + - fact, we obtain the first result of the lemma statement. Similar to the proof of the right side of the inequality, the quantity on the left side converges E [[ + 7, 07) - x)3] using the monotone convergence theorem. This gives the second result of the lemma statement.
X)²1
given by (())) E[X X = .1. Using this
B.3 Stability of optimal denoisers
The right-hand side of the bound in the above Lemma corresponds to the minimum mean squared eтror achievable for denoising at noise scale 7, where F the state evolution parameter corresponding to running the learned AMP iterations up to that layer of the unrolled network. To show that the learned network can compete with Bayes AMP, we need to relate ? to the corresponding state evolution parameter given by Bayes AMP. For this, we need the following stability result showing that the minimum mean-squared error is stable with respect to perturbations of the noise scale.
18

---

Lemma 5. Let prior X be such that the score function p);(;; is B-Lipschitz continuous for all where p(= r) denotes the probability density function of random variable X+ rZ. Additionally, assume that the variance of X is bounded by V. Then, we have
V2 x)
E[( [ + + x) &
Tr
Trl
σ6
where (x) = E X 7zZ and ff (x)
Proof.
Rewriting the error between fü and fi; we have
- = E(( (X+ 1.)) f; 42) + (X x22) -1))
(9)
=I( (X
fx -fx+ +x++))))
(10)
2X - (x (X +2 - X)] (11)
+ +EI (z) - -x)]].
(12)
The term in Eq. (10) can be upper bounded by (B(T; TT) 2 because f Lipschitz by
assumption. Using Cauchy-Schwartz for the term in Eq. (11), the second term is upper bounded by
2E 2E[(f (x+ -x
2E[ (X (x+ fr (X+ (X+ (X +x 1)
< 2BV - /
(13)
The squared term in Eq. (12) can be upper bounded by
E[G T - =) E[((F ((
+ 2E[(f
(2) 122) X)]
+ +E[(f (X -
We can upper bound the first term above using Lemma 6. Likewise, the second term can be bounded by Cauchy-Schward and Lemma 6. In this way, using that 2 N T I z σ² we obtain that
V2B² σ6
V2B
(X (x++x2) (X+
S EIIF -X²1 +2 x)
( )2
Combining this bound with Eq. (13), we obtain the bound.
The proof above relies the following "score perturbation lemma" showing that the optimal denoiser is Lipschitz with respect to the noise scale.
Lemma (Score perturbation lemma). Let p(:; 7) be the density function of + rZ. Let the score function d log p(:; r) be B-Lipschitz continuous. Then, the L2 error between the optimal denoisers at two noise scales ti and Ty is given by
B² T() V2
E | -
s PiB² -2)
Proof. Denote the probability density function of X T random variable value x as p(x; rt). Assuming o log p(: T²) is B-Lipschitz function and using Lemma C.11 of [LLT22], we have
Ex-p(sre (2) log (xx; T2) - a p(x; +) - B2 - T²) +B² (Tr - T2) (E) p(; =))
Using Tweedie's formula (0, log r2) [ + TZ x] x) //), we have
Ex-p(sre (o) po; rτ) à aog [])
-2 +
x!)
x)){² 2
=Ex-(xxz))
19

---

This implies that
[[(E[X]X+ x -
(7) - 27) PT
S ( (0, 10110g 12) - 1ogpp(x;;)))
E[(E[X|X + 122 = x] x)
$ ( -() B2 +
[| + 72 = x
where the last inequality uses the fact that F, t, because is obtained using optimal denoiser. Additionally, we have To = r² + (1/8)E[X*) Σ 2 V2, Using law of total variance, we have var X] = E[var(XX (X + TzZ)] + var(E[X|X TyZl). Therefore, we have E[X²] 2 E[(E[X|X TzZ = x] - x) which implies the claimed bound.
B.4 Putting everything together
We are now ready to conclude the proof of our main result.
Proof of Theorem 2. Recall that the state evolution recursion for t, is defined using the optimal denoiser ft at layer in Lemma 1. Similarly, the state evolution recursion for 7, is applied using learned denoiser fi( ., (,)) At some places, we will use fr( .) to denote fr( ,,) for brevity. By the state evolution recursion for Tr, the error between Ty and ř is given
(E((F- (X 7-1 x) - El,f,, 12)
8 σ6 1 t 11 +1 + 22)
1
1 12
where the last inequality follows from Lemma and Lemma 5. As Tr + f 20 for all and 2 Tú, we have
V2B²
81 +2
1-11+ -
807
80
Solving this recurrence, we have that the error in the state evaluation parameters after L layers upper bounded by
B2
12 - tzl S
(at + E2).
807
Combining this bound with Lemma 4 and Lemma 5, we have
(222211 1+1 (E) (2).
lim Ex, d-000
( II xr, x|| > E[ T, x)²)
807
Applying state evolution (Lemma 1), we get that
Exal 1 Geal=IficAA oy + &r. 0,) - x||
lim d=
aall=llf < lim Exal Exal llfe (ATor + 01) -x|| + V2 L+1
(21+
d
807
as claimed.
B.5 Bounding the complexity of denoisers
In Definition 1, we primarily focus on the setting where Ca 2 Here we provide some bounds on (g æ)) in this regime. This in turn provides the bound for (, a). First note that we always have
(og(( s cxp(log(1/))Ca//)> (1/z)
(14)
First consider the case where $ Σ, c,z is degree-k polynomial. Eq. (14) then implies
C,(6,a) k(1/o) max lci].
(15)
20

---

This bound is useful when the degree is high, e.g. when k = á @(C (1)//))) We also have the following naive bounds when the degree small or intermediate. When k 2 log( 1/e),
(,(6, S (log(1/e)Ca) Oloog 1/e)
max (slog(1/e)
|cil + (Ca)* (>log(1/e)
max Icil.
(16)
When k z log (1/)),
C a) Σ (log(1/e)Ca)0mh max
(17)
In the context of our learning result, it suffices for the true denoisers fr to be well-approximated by functions of low complexity, e.g. low-degree polynomials.
In our setting of Lipschitz denoisers from Assumption 1, we can get good approximation by low-degree polynomials via the following standard result:
Lemma (Jackson's theorem [Jac11]). Let k Z, B, If a function f RR. B-Lipschitz, then there exists polynomial of degree-k for which
R is
sup If(z) q(2)) & BR/k. lalsR
(18)
As the true denoisers are Lipschitz and the prior p is assumed to be Ssuu-Gaussian so that the effective support over which the true denoisers must be approximated has radius O(R,, this implies that we can approximate them pointwise to error 8 with polynomials of degree k (R/8). To apply any of the complexity bounds in Eqs. (15)- (17) to these polynomials, remains to bound the coefficient of largest magnitude. For this, we can apply a result like the following:
Lemma 8 (Corollary of Lemma 4.1 from [She12]). Given polynomial q(z) Σn ozt for which suP2c(o.1] [q(2)] M,
Σ |cil lall 4
(19)
In situations where the denoiser has additional smoothness properties, one can obtain even better polynomial approximations. To illustrate this, we provide an example in the special case of the Z2 prior:
Example 1. When p Zz the optimal denoisers used in Bayes AMP are of the form tanh(-/r2 for
various r 2 a, where x² the variance the measurement noise. This function is B 1/ps-iipschitz
and the effective support of N(O, r²) is of radius R = O(1). By standard results, tanh(-/ can be n-approximated over [-0 0(1) with a polynomial of degree-k
O(lo( /) o²),
see e.g. Lemma 2 in [LSSS14]] 4]. Applying Lemma 8 to bound the coefficients of this polynomial by (a [n) and then applying the bound in Eq. (16), we conclude that the denoisers in the case of Zz prior are j-approximated by polynomials q of complexity
(g(g,a) & (l/o)) ((lool)))((((gg)))))( ((og(11/)Co)) O(log//)) + (Ca) ((ogoto(oo)))))
(20)
C Rank-one matrix estimation
C.1 Preliminaries
While compressed sensing involves linear noisy transformation of its underlying signal, rank- one matrix estimation offers nonlinear counterpart. Here, the unknown signal R is first transformed into a rank-one matrix xx The observed signal is given by
A Y= G,
(21)
where G e dxd is symmetric matrix with entries Gij N(O, 11) for i S and À > 0 denotes the signal-to-noise (SNR) ratio. Again, we assume x Px for product prior Px A p pd with p some distribution over I.. Then, the rank-one matrix estimation problem attempts to recover an
21

---

estimate X for such that the Frobenius mean squared error (MSE) 1 _ - -xx" I| is minimal. The asymptotic setting corresponds working with a sequence such problems indexed by dimension d, where d - co.
The approximate message passing (AMP) algorithm for estimating given Y proposed in [RF12]
takes the form
Xx+1 = f(((e) Ur = Yx=xxxxxxx - -x--
(22)
(23)
where fi IR R again indicates scalar denoiser applied componentwise. There are various possible choices for initialization; here we consider x 1 E æd and un Yxp-
As with compressed sensing, AMP iterates for rank-one matrix estimation satisfy state evolution recursion. Define parameters er and 5, evolving according to the scalar equations
Mr+1 XE|Xfi(xXX + VT,Z)l
(24)
Tr+1 = =[f((xx
(25)
where X p and Z N(O, 1). Then for d - 00, the empirical distribution over entries of dz converges asymptotically to the one-dimensional distribution over μr + vtiZ [RF12]. Bayes AMP thus corresponds to the choice of denoising functions optimizing the posterior mean on X given J + vT,Z:
ft = E[X 14, v.
(26)
where iú and r are obtained by substituting fi as the denoiser in Eqs. (24) and (25).
As with compressed sensing, AMP has powerful theoretical guarantees for rank-one matrix estimation. In some cases, e.g. when p is the uniform distribution over {1. -1,,,,, the Bayes AMP algorithm is information-theorrtiiil achieved by the Bayes optimal estimator [DM14]. In general, Bayes AMP is conjectured to achieve asymptotically optimal MSE over all polynomial-time algorithms for this task [MV21] while being provably optimal over all GFOMs [MW22b].
optimal, i.e., it asymptotically matches the MSE
Despite these guarantees, Bayes AMP for rank-one matrix estimation suffers from the same implementation bottlenecks regarding knowledge of the true prior of the underlying signal. And with compressed sensing, choosing the state evolution parameters Jl and according to the true recursion can cause Bayes AMP to diverge, so in practice one estimates these parameters using the previous iterates. In particular, practitioners typically replace ty with 1 +|xx using
Eq. (25) and He with / llvell 3/ ||x+ using the infinite dimensional distribution of the components of oz. The latter holds when the prior p has zero mean and unit variance; in general, the estimate must be scaled down by an additional factor of -x [X*], [X2], which can be estimated from data. Unless expressed otherwise, we assume Ex-p [x] = 0 and [x = 1 in this setting throughout. These conditions are without loss of generality for compressed sensing, to which our theoretical results pertain, but not without loss of generality for rank-one matrix estimation.
C.2 Unrolled architecture
Here, we are given training data {(Y',x')) generated by Eq. (21) with all x² Px- Let F denote family of MLPs with fixed architecture constrained to three-dimensione inputs and onediimensional output. Set x2 = e Rd and an Yxg- Initializing MLP fo F for all f E [0, L 1], we obtain an L-layer unrolled network that computes
7e. = fo(ūr; De, 7r)
(27)
6r1 Yx (0 7,),
(28)
1| at every iteration, where f4, == 1 applying the scalar function fi( pe. 7r) entrywise.
and ア, -1 Here F.. : 7p) denotes
22

---

C.3 Results
Implementation details. We set d 20 and fix A 1.5. We focus on the Gaussian and Z priors for this setting. The family of MLPs F is constrained to have three hidden layers, each with 20 neurons and GELU activations, and we train over a dataset (Y*,'*) of size N = 21 obtained by sampling from the prior and using Eq. (21).
As with compressed sensing, we train with finetuning and also consider a baseline with the MLP denoisers in LDNet replaced with guided denoisers that learn parameters attached to the analytic forms the Bayes-optima denoisers.
Caussian
22
LUPA
LDNet
Guided Denolsers
Guiced Densitens
Sayea AMP
Dayss AMP
0.80
0.35
Figure 3: LDNet for Rank-One Matrix Estimation. On the left, we plot the NMSE obtained by LDNet and Bayes AMP on the Gaussian prior, while the right plots are Z2. LDNet matches Bayes AMP with slightly quicker convergence.
Gaussian prior. For a Gaussian prior, each component x; of x is drawn x; N(O, 1), so Px pd for p = N(O, 1). As expected, L.DNet tracks Bayes AMP to convergence at an NMSE of 0.6931 with slightly quicker convergence.
Z, prior. As with compressed sensing, each component of is drawn from (-1,1} with probability 1. Again, LDNet slightly outperforms Bayes AMP until convergence at an (information- theoretically optimal [DM14]) NMSE of 0.5243.
Seat beroarn
Seat Oepose
Deeti tero 20
Cprd cenoper
-32
43
220 425
226
423
Dendor Biyal Evacis
110
Beat
-0.0
4.5
430 425
120
Figure 4: Learned Denoisers for Rank-One Matrix Estimation. We plot layerwise denoising functions learned by LDNet the Gaussian and Z2 priors relative to their optimal denoisers range of inputs in (-2, 2). The state evolution inputs Pt, ti each denoiser to be their empirical estimates.
LDNet denoisers. While later iterations are able to recover the Bayes-optimal denoisers, worth noting the high approximation error at early iterations shown in Figure 4. Early iterations
23

---

correspond to when the AMP estimates stagnate at an NMSE of roughly 1.0, corresponding to a random, uninformed signal that can accommodate high denoiser error. Around iteration 10 is an inflection point when both LDNet and AMP transition from the uninformed regime to convergence at an informed NMSE (Figure 3), where low approximation error tolerated.
D Beyond Bayes AMP performance
D.1 Learning auxiliary parameters
m = 200
m 150
m = 100
Leamed
Ieamed Sayas
eemmed
AMP
10
-10
-10
-15
-15
-1
-3
30
-30
-30
-30
-45
-45
-45
2
14
10
14
Layers
Layers
Layers
Figure 5: Learned B with Decreasing Dimension. We hold 8 1/2 fixed while scaling m from 200 down to 100. Plots show NMSE (dB) performance of unrolling denoisers and learning vs. Bayes AMP for randomly drawn measurement matrices. There is an increasing gap in performance as m decreases.
In Figure 5, we set the parameters of B = A to be learnable alongside the layerwise denoisers (see Algorithm 2). Holding 8 = 1/ fixed while scaling down m has the effect of widening the gap between the NMSE performance of Bayes AMP versus LDNet with trainable B. Over five randomly drawn measurement matrices per dimension regime, we find that, on average, LDNet outperforms Bayes AMP by 7.2750% when m 200, 43643% when m 150, and 37.0605% when m = 100. Here, percentage is measured by [Dayes-
X 100% in NMSE (dB).
Bayes
Truncated Random Orthogonal
Truncated Rancom
-25
AVP
-50
1514
COSAMP
-10-
-75
100
-20
-22
-25
Lesmned Aquated Dayes AMP
-35
COSANI
30
34
1.
12
Layars
Layors
Figure 6: Non-Gaussian Measurements. On the left, we plot LDNet with learnable compared to several baselines for a random truncated orthogonal measurement matrix, and on the right, for random truncated Gram matrix. LDNet outperforms the other baselines in NMSE as well as convergence.
Another regime to which existing theory for Bayes AMP largely breaks down is when the entries of the sensing matrix are non-Gaussian. While was previously observed that learning can help for ill-conditioned A [BSR17], we find that there are advantages even for well-conditioned (but non-Gaussian) sensing matrices. Figure 6 plots NMSE for two sensing matrices Á e R 250x500, one obtained by truncating a random orthogonal matrix Q E 500x500 (condition number 1), and the other by truncating Gram matrix X X E RR 500x500 with Xıj N(0, 1/m) (condition number 1091). Also displayed are the iterative baselines of Bayes AMP, ISTA, and COSAMP [NT10]. For the truncated random Gram matrix, Bayes AMP and ISTA actually diverge, so we plot "adjusted" baselines replacing A" in Eqs. (2) and (3) with the "B" matrix learned by LDNet. The baselines
24

---

Algorithm 2: Learning B
Input: Training data D, LDNet V, Measurement matrix Initialize B = AT; I do
2 for if E > 0 then
Initialize f
Freeze Icarnable weights in and for k Train 4/0 t] on D;
6
7 Unfreeze learnable weights in and Jt for f;
Train 4/0 f on DD
Output: Fully trained Y, Learned B
considered all drastically underperform our unrolled network. In fact, the underperformance of "adjusted" Bayes AMP demonstrates that the LDNet denoisers are strongly coupled with B, suggesting that learning denoisers is beneficially composable with the traditional algorithm unrolling approach.
All told, we see how adding auxiliary learnable parameters can mitigate scenarios (e.g. finite dimensionality, oon-Gaasiian or ill-conditioned sensing matrices) where Bayes AMP is suboptimal or not known to be optimal.
D.2 Non-product priors
Thus far, our theoretical and experimental results have remained the regime of product priors. But what happens when our underlying signal is drawn from non-product distribution?
The modifications AMP are minimal, detailed in [BMN20], amounting d-dimensional denoisers ft Rd Rd and replacing the average derivative of the scalar denoiser in the Onsager term with the normalized divergence div in Eqs. (3) and (23). In this non-separable setting, AMP still satisfies one-dimen mne-immnniiona state evolution recursion [BMN20]. In fact, in the asymptotic limit, in some sense our theoretical guarantees carry over if for generic d-dimensional priors, minimizing the score mottehn objective via gradient descent (now in d dimensions instead of dimension) can learn the Bayes-optimal denoiser with gradient descent. This is question of immense interest within the theory and practice of diffusion generative modeling and remains an important open direction this area.
In practice, for compressed sensing, unrolled AMP has been shown to be performant on image datasets [MMB17], which serve prime examples of real-world, non-product signal priors. For our purposes, we focus on rank-one matrix estimation which, even in the product setting, remains unexplored in the unrolling literature. Additionally, we work with handerafted priors, where can plot a baseline achieved by Bayes AMP.
LDNet for non-pro non-product priors. We work in the low-dimensional regime d = 10, where x E Rd. LDNet requires small modifications to the layer iterations defined by Eqs. (27) and (28). The family of MLPs FF parametrizing the denoisers have three hidden layers with 1000 neurons and GELU activations, with input dimension d 2 and output dimension d. We take fo E F for
E [0, L 1], and we replace (ar in Eq. (28) with divvf.. To avoid backpropagating through the Jacobian during training, omit the finetuning step in Algorithm 1.
Signal distributions. To analyze the performance of d-dimensional LDNet, we consider two priors on product Z, and mixture of Gaussians. In both instances we work with dataset of size N = 212 generated by sampling a signal and using Eq. (21). The product Z₂ prior serves as a test example to see whether the multi-dimensional learnable denoiser provides additional performance gain over Bayes AMP when treating the product distribution as d-dimensional as opposed to oneddimensional.
25

---

Mixture of Gaussians provide a quite general class of non non-product priors. We consider p N(11, x) + 2/2 N(22, X2;x), where we choose ji and Hl at random with each coordinate chosen from N(0, 1), ensuring that iI and μl do not define the same direction. For cach of the covariance matrices, we begin by choosing random vector with each coordinate drawn uniformly from [1, 2] and normalize so that vector has norm Vd. We take this to be the diagonalization (i.e. eigenvalues) of the covariance, and conjugate by randomly drawn orthogonal matrix.
0.70
0.65
32
10
Figure 7: Multi-Dimensional LDNet for Rank-One Matrix Estimation. On the left, plot the NMSE obtained by LDNet and Bayes AMP on Z2, while the right plots are on the mixture of Gaussians. LDNet outperforms Bayes AMP by significant margins.
As Figure 7 demonstrates, multi-dimension LDNet significantly outperforms the Bayes AMP baseline on both priors. On a product Z, prior, LDNet achieves an NMSE of 0.6485 compared to Bayes AMP's 0.6864, marking a 5.52% improvement while also reaching convergence much faster. On the mixture of Gaussians prior, LDNet achieves NMSE 0.6881 compared to Bayes AMP's 1.8757, marking 63.32% improvement.
D.3 Learned denoiser dependence on sensing matrix
Random Al
Random A2
Random A3
Origina A
-55
MNCE
-25
-45
Figure 8: Transfer Experiments. Above we plot the NMSE (in dB) over 15 iterations for different choices of measurement matrices coupled with learned MI. denoisers, including the training-time sensing matrix. We see that the denoising functions are roughly transferable several random Gaussian measurement settings, suggesting the learning process is not coupled to the fixed sensing matrix seen during training.
Recall that our unrolled denoising network differs from the theoretical setting in two key ways: a), the network is trained assuming a fixed sensing matrix A rather than in expectation over random Gaussian A, and b), state evolution parameters are estimated from previous iterates to account for finite dimension corrections.
Indeed, despite these differences, the plots in Figure 2 suggest our network appears to learn fundamental "optimal" denoiser that is independent of To further verify this claim, we froze Gaassiinn prior and replaced the A matrix in
the learned MLP denoiser weights for the Eqs. (5) with other randomly sampled Gaussian matrices A' 250x500 As shown Figure 8, this leads to minimal changes in the NMSE profile.
26

---

Explicit expressions for various Bayes-optimal denoisers
For prior p and X p, N(O, 1), the Bayes optimal denoiser in Bayes AMP is given by
fi =IXIX + 1z y]
(29)
for compressed sensing and
= v
(30)
for rank-one one matrix estimation. For the priors examined in our experiments, we write out the -eeeiin- optimal denoiser along with the parameterized guided denoiser form (if relevant), where the parameters are learnable during training.
Bernoulli-Gaussian prior.
One can compute the optimal denoiser to be
y
f7(y)
(31)
(1 x2) (1 1-r
N(0.
N(0,ff+1:y)
and parameterize the denoiser via
Fely; 01, 02)
y
(32)
0
exp
02
10.7+r//0.)
y²
as done in [BS16, BSR17]
Z2 prior. The compressed sensing optimal denoiser [DM14] be written as
fi (y) = tanh (
(33)
7 2,
which parameterize
fr tanh y B
(34)
For rank-one matrix estimation, the optimal denoiser be similarly computed to be
f((y) tanh y He Tt
(35)
parametrized as
F(y: B) tanh (y. Bpr
(36)
Gaussian prior.
For rank-one matrix estimation, the optimal denoiser for Gaussian prior is
Hr
(y) y 1r +
(37)
parametrized as
F((y: y mr +
He
(38)
Mixture of Gaussians prior: The calculation of the Bayes-optimal denoiser for mixture of Gaussians prior in rank-one matrix estimation is slightly more involved, so we provide some more details about the calculation. Given Σ N(ps Σ; x), where μi Rd and Σ; e R are invertible positive semidefinite symmetric covariance matrices, convolution with vtZ
27

---

N(O, vala) results in the mixture of Gaussians distribution p # N(Holi 12 + x,)a; For any mixture Σ A, Qú x), the score is given by [CKS24]
2,N(Hi, 21;x)
2x ( pi).
(39)
(,,,
Thus, we have
k Σ
N(pelt 22 + rxla;x)
Vlog
(1721 + r,la) (xx - HeHl),
(40)
Σ, N(upp p xıddix)
F1
so by Tweedie's formula, the posterior mean on X given X + VTi is
11 +, (17 + tyla) - - -m(H1i)
y -
(41)
Σ, HC rylaiy)
i=1
28

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yss]
Justification: We provide complete proofs of our theoretical claim and experimental figures for the experiments.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations.
No or NA answer to this question will be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine include aspirational goals motivation long it is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
The limitations are discussed in the last section of the paper.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model sell-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g,, if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
29

---

Answer: [Yes]
Justification: We provide complete set of assumption in Assumption ?? 1.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross-referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the complemented by formal proofs provided in appendix or supplemental material.
of the paper should be
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided not)?
Answer: [Yes]
Justification:
We provide detailed information about the setup of our experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible important, regardless of whether the code and data are provided or not.
If the contribution is a dataset and/o or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to hosted model (e-g., in the case of large language model), releasing of model checkpoint, or other means that are appropriate to the research performed.
novel architecture, describing the architecture
either make it possible for others to replicate the model
While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., large language model), then there should either be way to access this model for reproducing the results way reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility.
the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path reproducing or verifying the results.
30

---

5. Open access to data and code
Question:
Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, described in supplemental material?
Answer: [No]
Justification: We currently don't provide the code for the open access but we are planning do it soon.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nig) cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: pettss//nipss cc/public/guie httliiccguies///eeeee Muuides/CooeSubmisse wwt.sssooooiocyyll for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Detain
Question:
Does the paper specify all the training and details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We provide detailed information about the training and test details.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper a level of detail that necessary to appreciate the results and make sense them.
The full details can be provided either with the code, in appendix, or supplemental material.
7. Experiment Statistical Significance
Question:
Does the paper report error bars suitably and correctly defined other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: The goal of our analysis to validate theoretical findings on the synthetic data.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
31

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula,
call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK to report 1-sigma error bars, but one should state The authors should preferably report 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are our of range (e-g-
negative error rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time execution) needed to reproduce the experiments?
Answer: [Yes]
Justification: The estimated amount of training time for our final experiments is around 100 CPU hours.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https: htaps:///nuriipp httppp/ic//.iiissiielllleeeeeeee
Answer: [NA]
Justification: Our research thcoretical and does not have any direct ethical impact.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is a special considcration due to laws regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: Our research is theoretical and does not have any direct societal impact.
Guidelines:
32

---

The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly specifc groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e-g,, gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of MI.).
11. Safeguards
Question:
Does the paper describe safeguards that have been put in place for responsible release of data or models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification: Our research theoretical and does not have any safety threats.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (c.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [NA]
Justification: Most of experiments are validating theoretical foundings and the synthetic data.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
33

---

For scraped data from a particular source (e.g,, website), the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
this information not available online, the authors are encouraged reach out to the asser's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
Our research theoretical in nature.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the panaset/code/mode as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question:
For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [Yes]
Justification:
Our research is theoretical in nature.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then much detail possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country institution) were obtained?
Question:
Answer: [Yes]
Justification: Our research theoretical and does not involve any human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
34

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
35


Reviews:
[{"review_id": "g0aWcVLW74", "reviewer": "Reviewer_F56a", "summary": "This paper studies the theoretical capabilities of unrolled denoising networks in the context of compressed sensing and rank-one PCA problems, with a focus on experimental validation. The authors present the first proof of learning guarantees for neural networks based on unrolling Approximate Message Passing (AMP). They demonstrate that, when trained with gradient-based methods, these networks can achieve performance on par with optimal prior-aware algorithms. Specifically, for compressed sensing under the assumption that the data is sampled from a product prior, the parameters of the network converge to the same denoisers used in Bayes AMP. The technical aspect of this proof employs a combination of state evolution and neural tangent kernel (NTK) analysis. Numerically, the authors validate their claims and show that the unrolled architecture can handle more general priors.", "questions": "* Have the authors numerically simulated the effects of different network architectures on the robustness of their estimates? This inquiry focuses on understanding how architectural variations impact the results, rather than solely assessing performance.\n\n* Could the authors comment on the practical likelihood of Assumption 1?\n\n* Could the authors please provide context for the variable B in Informal Theorem 2? Its meaning is clear from Formal Theorem 3, but not from the informal statement.", "limitations": "Limitations are clearly discussed", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The paper presents significant theoretical contributions to an important subfield of machine learning. It is well-motivated, clearly written, and aligns with the existing literature. This is a significant theoretical advancement that bridges the gap between optimal prior-aware algorithms and data-driven training methods.\n\nThe authors introduce a novel proof technique that combines state evolution and Neural Tangent Kernel (NTK) analysis. Unlike prior work, this proposed analysis is robust and independent of dimensionality.\n\nThe author provides sufficient numerical evidence to validate their theoretical contributions.", "weaknesses": "The paper's theoretical results are confined to settings of compressed sensing with data derived from a product prior, which is an unrealistic assumption compared to data from actual science and engineering applications.\n\nFurthermore, the chosen architecture could be considered impractical in practice due to its fixed width and depth.\n\nHowever, these weaknesses also present opportunities for future research.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for their effort in reviewing our work and address their comments below.\n\n- **On the product distribution assumption**: Please see the global response for the discussion about the assumptions on the prior. We have also subsequently expanded experiments to the non-product distribution regime, also included in the global response and the second supplementary figure in the attached pdf.\n\n- **On the flexibility of the architecture of the MLP denoiser**: The MLP parameterizing the denoiser can in practice be set to any architecture (convolutional net, U-net, etc.), and doesn’t have to be constrained to the three-layer, 70-neuron wide MLP we used. \n\n- **On the understanding how architectural variations impact the results**: We examined different depths and widths of our MLP parameterization and found that the behavior of unrolling was relatively robust to variations in the choice of the denoiser architecture. The adjustments we made were ultimately meant to control for the amount of error in learning the denoiser, and the architecture we eventually reported results for worked sufficiently well for our experiments while also being quite reasonable to train.\n\n\n- **On the practical likelihood of Assumption 1**: Assumption 1 comprises two parts: 1) sub-Gaussianity and 2) Lipschitz continuity of the $\\tau-$Gaussian smoothed score function (score function of $x + \\tau z$). We believe both assumptions are mild, hold for a large distribution class, and are used in prior literature (e.g., the sub-Gaussianity of a distribution holds for any distribution with bounded support and the Lipschitz continuity of the score function holds for mixtures of spherical Gaussians with bounded means). Please also see the general response on the additional discussions about these assumptions. We will include a discussion about the generality of the assumption in the revised draft.\n\n\n- **The definition of variable $B$**: The variable $B$ denotes the Lipschitz constant of the score function of $p(x; \\tau)$ where $p(x; \\tau)$ is a probability density of random variable $x + \\tau z$ where $x$ and $z$ corresponds to the data distribution and standard Gaussian distribution respectively (Please see Assumption 1 for the complete definition of variable $B$).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "XILIiiEDHT", "reviewer": "Reviewer_DSkx", "summary": "This work investigates unrolling approximate message passing (AMP) for solving compressive sensing under Gaussian design and separable prior distribution. The unrolling scheme consists of parametrizing the AMP denoiser at each iteration by the layer of a neural network, which is sequentially trained using observations from the signal distribution by minimizing the empirical mean-squared error.\n\nThe main result is to show that with a polynomially wide two-layer neural network trained under one-pass SGD with with polynomially many observations, unrolled AMP asymptotically achieves the same MSE as AMP implemented with the optimal denoiser.", "questions": "- L80-81:\n> \"*and thus achieve mean-squared error which is conjectured to be optimal among all polynomial-time algorithms*\"\n\nAs the authors probably know, there are a few counter examples to this conjecture, including for problems which are closely related to Compressive Sensing (e.g. noiseless phase retrieval, see []). I suggest the authors being more precise, by either refeering to the known optimality results within first order methods [CMW20], MW22b] or by adding a footnote saying the conjecture is believed to hold in a broad sense, since all known exceptions rely on rather fine tuned algorithms (like LLL in the noiseless phase retrieval case)\n\n- The observation in L303-304 that the denoiser learned by unrolling can outperform the optimal denoiser is surprising. I understand that optimality of AMP is an asymptotic statement - so this difference should get smaller with the dimension - but is this something you consistently observed throughout in the experiments? Is this improvement always within a $O(d^{-1/2})$ interval? If yes, how can it be distinguished from finite-size fluctuations?\n\n- Can the authors comments what \"*sufficiently small*\" learning rate and \"*sufficiently large*\" number of steps exactly mean?\n\n- The tricks for training discussed in Appendix B are intriguing. Do the authors have an intuitive understanding of why fine-tuning is necessary for successfully learning the a good denoiser for Compressive sensing with Gauss-Bernouilli prior but not rank-one matrix factorization and Compressive sensing with $\\mathbb{Z}_{2}$?\n\n- At a similar note to the above, what is the intuition of why sequential learning avoid bad local minima? Similarly, why initializing at the solution of the previous layer avoids local minima? How these local minima look like and how they compare with the optimal denoisers? This should be easy to visualise since the denoiser is a one dimensional function (as in Figure 2)\n\n- It would help readability to add in all figure captions the details on the plots, e.g. the dimensions $m, n, d, L$.", "limitations": "Limitations are discussed in Section 5.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "The paper is well-written and easy to follow. It is also self-contained - the small introduction to AMP in Section 2.1 is a nice addition since most NeurIPS readers will not be familiar with AMP. The results are interesting, since they suggest that the optimal AMP denoiser can be learned \"on the fly\" with unrolling.", "weaknesses": "The topic is relatively niche within the NeurIPS community. Also, AMP is designed to Gaussian measurements, and it is well-known that as an algorithm it is not very robust to other designs, therefore the results have a limited practical scope. Unrolling requires samples from the signal distribution.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their detailed comments.\n\n- **On a relatively niche topic within the NeurIPS community**: While understanding the Bayes-AMP algorithm in itself is a relatively niche area for the NeurIPS community, our theoretical results show that the learnability of Bayes-AMP amounts to characterizing training dynamics for the score-matching objective. Given the prevalence of score matching within the literature on diffusion models, which is of very wide interest in the NeurIPS community, we believe that this line of work involving unrolling and AMP offers valuable insight into the dynamics and approximation capabilities of diffusion models. Some of this has been explored in work by [Mei24] and [MW24] (see our related work section in the Appendix). \n\n[Mei24] Song Mei. U-nets as belief propagation: Efficient classification, denoising, and diffusion in generative hierarchical models.\n\n[MW23] Song Mei and Yuchen Wu. Deep networks as denoising algorithms: Sample-efficient learning of diffusion models in high-dimensional graphical models.\n\n- **AMP on non-Gaussian measurements**: \n  - We agree with the reviewer that AMP has these limitations, but empirically we find that *unrolling offers a powerful way to interpolate between the theoretically optimal guarantees of AMP in “nice” Gaussian settings and the practically useful guarantees of neural network-based methods for inference in real-world settings.* \n  - For example, we find in our “learnable B” experiments that by learning auxiliary matrix parameters, unrolling can surpass the performance of Bayes AMP (and baselines such as ISTA) on non-Gaussian designs despite working with an architecture inspired by AMP. In particular, we examine a random orthogonal design (fig. 4) and a random Gram matrix design (attached fig. in pdf in global response) and show our method can surpass Bayes-AMP. While our theory only explains why unrolling can compete with AMP, it opens up the potential for future work exploring why learning auxiliary matrix parameters can ameliorate the non-robustness of AMP to other designs.\n\n- **On the conjectured to be optimal among all poly-time algorithms**: We thank the reviewer for bringing this up. We will in our revision refer to the precise GFOM optimality results for compressed sensing and low-rank matrix estimation to qualify the optimality of AMP.  \n\n- **On the denoiser learned by unrolling can outperform the optimal denoiser**: \n  - With regards to learning the denoiser, as in Figure 1, the unrolled denoisers outperform Bayes AMP by a relatively small margin. The NMSE in Figure 1 is averaged over $2^{15}$ samples with fixed, randomly chosen design, and we observed this improvement consistently for other randomly sampled designs as well, so we do not believe the small improvement of MLP denoisers over vanilla AMP in Figure 1 is from random fluctuations. The broader point we want to emphasize is that Bayes AMP is only known to be “optimal” up to $o_d(1)$ margin in NMSE, and the fact that the MLP denoisers we learn improve over this in the finite-dimensional setting is consistent with this state of knowledge.\n  - We also observed consistently that as the signal dimension decreased, we could outperform Bayes AMP by learning the “B” matrix (or measurement transpose) as detailed in Figure 3. Here, we believe the learned B experiments take advantage of finite-dimensional sub-optimalities of AMP as a GFOM, rather than just finite-dimensional sub-optimalities of the denoisers used by Bayes AMP.\n\n- **On the intuitive understanding of the fine-tuning**:\n  - We observed finetuning was not largely impactful for the compressed sensing experiments but was helpful for rank-one PCA. Note that PCA has different strengths of initialization for AMP of the signal estimate while compressed sensing has a “canonical” initialization ($x_0=0$). PCA thus induces a “burn-in” phase where the AMP estimates reflect almost no signal at early iterations (see Figure 6 of the submission).\n  - We suspect that finetuning assists in maintaining the consistency of the learned denoisers across this burn-in phase where the correlation between the estimate and MMSE estimate is low, communicating the loss objectives of the later denoisers to earlier iterations.\n  - We also observe that finetuning improves validation loss most during this “burn-in” phase, but notably not as much during later iterations where the estimate is close to the MMSE estimate. \n\n- **On the intuition of why sequential learning avoid local minima**:\n  - Sequential learning takes advantage of the following result: Suppose we have learned the Bayes-AMP denoisers in all previous layers. Then the denoiser that minimizes the MSE estimate of the next layer is the corresponding Bayes-AMP denoiser as well (see thm 3 in [MW22]). This simplifies the learning optimization immensely, as we only need the current layer’s denoiser to globally optimize the MSE loss function.\n  - Another way to view this is that we know that our learned network should provide the minimum MSE estimate after each layer, but training end-to-end throws away this “intermediate layer” information and only uses the final layer’s optimality. \n\n[MW22] Statistically Optimal First Order Algorithms: A Proof via Orthogonalization. Andrea Montanari, Yuchen Wu.\n\n- **Intuition behind initializing at the solution of the previous layer**: As for initializing at the solution of the previous layer, the intuition is as follows. As the Bayes-AMP estimates converge, so do the denoisers, so we expect that the denoiser from the previous layer is close to the optimal denoiser, at least closer than a random initialization. \n\n- **On quantification of learning rate and number of steps**: By sufficiently small learning rate, we mean that the learning rate is $\\Theta(1/(\\epsilon*m))$ and by a sufficiently large number of steps, we mean polynomial in the complexity of the denoiser and $1/\\epsilon$ (See the statement of Theorem 3).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "wSk1wgSuZn", "reviewer": "Reviewer_mdyi", "summary": "This work seeks to understand why neural network-based approaches for inverse problems can outperform methods that incorporate hand-crafted priors. In the Bayesian setting, it is known that when the true prior is available, the optimal estimator (in a mean square sense) is the conditional mean. However, in practice, the true prior is never known, and learning-based methods, such as those based on unrolling, must implicitly estimate such a prior from samples of the unknown prior. While such approaches have shown to perform well empirically, it is unclear from a mathematical perspective if they are learning the true prior. This work addresses this by analyzing what estimators do unrolled AMP-based networks converge to. It is shown that when the underlying prior comes from a product distribution with subGaussian coordinates, then the network trained with SGD will converge to the optimal denoiser, if the prior were known. This convergence occurs in the high-dimensional asymptotic regime. The authors support their theoretical results with empirics showing that the learned denoisers converge to the optimal denoisers under the setting of their theorem, along with extensions to other problems, such as the nonlinear rank-1 matrix estimation problem.", "questions": "- I am confused about how the notion of complexity is applied in the Theorem 2. In particular, what is the function $\\phi$ in the statement of Theorem 2? Is it the underlying “optimal” denoiser $f_{\\ell}^*$ or is it the one-hidden layer ReLU network? It is also not clear in the statements of Lemma 3 and Theorem 3 in the appendix what the function classes $\\mathcal{F}$ are. Are these functions of the form $F^*$ with bounded coefficients $(a_i, w_{i,1},w_{i,2})$? Moreover, to apply Lemma 3, the $\\phi_i$’s need to be infinite-order smooth. How does this arise in the assumptions of the Theorem?\n- Also, why is the error on the right-hand side of the inequality in terms of $\\hat{v}_L$ and $\\hat{x}_L$ instead of $v_L$ and $x_L$, the outputs of the optimal Bayes AMP? It’s not clear why it should be in terms of the outputs of the learned AMP, since the claim is that the learned AMP’s ell2 error is close to the optimal Bayes AMP ell2 error.\n- Have the authors conducted experiments in the case when the distribution is not a product prior? Do similar results extend to this case?", "limitations": "The authors discuss the limitations of their theory in the Conclusion section, noting that the assumptions on the prior distribution and linearity/Gaussianity of the inverse problem are in particular, limited settings. There are no broader societal impacts.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The work is tackling an interesting and timely question, which is why do neural-network based methods outperform algorithms based on hand-crafted priors and can they reach the performance of “optimal” methods. \n- Toy empirical results help in supporting the theory that the training scheme and network architectures converge to the optimal Bayes denoisers.", "weaknesses": "- Some aspects regarding the readability and presentation could be improved, especially as they relate to the main theorems. I have noted my questions below. \n- The assumptions on the prior distribution being a smooth product distribution feels somewhat limiting. In particular, the authors note that their overparameterization result is dimension-free, but I wonder whether this is a result of the fact that with a smooth product distribution, the dynamics are characterized by the 1-d state evolution. Do the authors have a sense of whether the result would still be dimension-free if the product distribution assumption was relaxed?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their thoughtful questions and remarks.\n\n- **On the smoothness and product distribution assumption**: Please see the global response for the discussion about the assumptions on the prior.\n\n- **On the notion of the complexity of denoisers**:  We think that this confusion is because of a typo/underspecification in the statement of theorem 2 and theorem 3. The notion of complexity is applied to the underlying optimal denoiser $f_\\ell^*$ (i.e. the function $\\phi$). In theorem 2/theorem 3, $M_0$ and $N_0$ are polynomial in the sum of the complexities of $f_{\\ell}^*$ for $\\ell$ from $0$ to $L-1$. Moreover, $M_0$ and $N_0$ are polynomial in $\\sum_{\\ell \\in [0, L-1]} C(f_{\\ell}^*, R(\\log \\frac{1}{\\epsilon})^{1.5})$ and $\\frac{1}{\\epsilon}$. We apologize for the confusion caused by the typo and will update it in the revised draft of the paper.\n\n- **On applying function classes and Lemma 3**: \n   - When we refer to the function class $\\mathcal{F}$, we are indeed referring to target functions $F^*(x)$ with bounded coefficients. Moreover, to obtain the learning result for the $\\ell^{th}$ layer denoiser, we apply Lemma 3 representing $f_{\\ell}^*(x)$ as a target function $F^*(x)$. \n   - Note that this is equivalent to setting $p=1$ and $a_1=1$, along with $\\phi_1 = f_\\ell^*$ (where $a_1^*=1$ and $w_{1, 1}^*=\\langle 1, 0\\rangle $, $w_{1, 2}^*=\\langle 0, 1\\rangle$; note we are treating \"$x$\" as the two dimensional input $\\langle x, 1 \\rangle$ into the target function). \n   - Recall that $f_{\\ell}^*$ is the optimal denoiser at noise level $\\tau_\\ell^*$ (i.e., $f_{\\ell}^*(y) = E[ \\mathbf{x} | \\mathbf{x} + \\tau_\\ell^* \\mathbf{z} = y ]$); therefore, it is an infinite order smooth function. To get an intuition for why this is the case, note that by Tweedie’s formula, the optimal denoiser is, up to a shift by $x$ and scaling, given by the derivative of $p(x; \\tau_\\ell^*)$, i.e. the log-density of the prior convolved with some Gaussian. Even if the prior were a mixture of Dirac deltas, the prior convolved with Gaussian is a mixture of Gaussians, and the log-density of this has derivatives at all orders. In particular, $p(x; \\tau_\\ell^*)$ is always positive and continuous (due to the convolution with a Gaussian, no matter the prior), so all higher order derivatives of $\\frac{\\nabla p(x; \\tau_\\ell^*)}{p(x; \\tau_\\ell^*)}$ exist by applying the quotient rule.\n   - We thank the reviewer for bringing up these details and will expand on them in the revisions.\n\n- **On the right-hand-side error of the inequality**: This is a typo. The right-hand side of theorem 2 should be $v^L$ and $x^L$ instead of $\\hat{v}^L$ and $\\hat{x}^L$ as mentioned in the formal version of the theorem in the appendix (theorem 3). This follows from the closeness of the state-evolution parameters ($\\tau^L$ and $\\hat{\\tau}^L$) of iterates $x^L$ and $\\hat{x}^L$. We thank the reviewer for pointing this out, and we will update it in the revised draft.\n\n- **On experiments with non-product signal priors**: We have subsequently extended our experiments to non-product priors. We do see similar advantages to unrolling over Bayes-AMP show up in this setting; please see the global response and the supplemental second figure in the attached pdf for a discussion of the non-product prior experiments.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank all the reviewers for their careful reviews and constructive feedback. We address some general comments raised by multiple reviewers here and address the rest of the comments in the individual responses.  

- **On the product prior assumption**: We restricted our results to product priors for two main reasons: 
   - The product prior setting is quite standard and widely studied within the theory literature on AMP (e.g., see [BM11, DAM15]).
   - The state evolution for AMP, which establishes that the input to the denoisers behaves like samples from the data distribution smoothed with Gaussian noise, still holds for non-product distributions (see [BMN17]). So to show that unrolled AMP provably learns to perform Bayesian inference, one needs to show that unrolled AMP learns denoisers at different noise scales. However, this is a major open problem in the theory of diffusion models (understanding when “score estimation” can be performed in polynomial time) and therefore remains well outside the scope of this work. However, with product distributions, corresponding to scalar/one-dimensional score estimation, we are able to show such a result. That being said, it would be interesting to study the specific class of non-product priors given by **generative priors** coming from pushforwards along random neural networks, as studied by Aubin et al. [ALB20].


[BM11] Mohsen Bayati and Andrea Montanari. The dynamics of message passing on dense graphs, with applications to compressed sensing.

[DAM15] Yash Deshpande, Emmanuel Abbe, Andrea Montanari Asymptotic Mutual Information for the Two-Groups Stochastic Block Model.

[BMN17] Raphaël Berthier, Andrea Montanari, and Phan-Minh Nguyen. State Evolution for Approximate Message Passing with Non-Separable Functions.

[ALB20]  Benjamin Aubin, Bruno Loureiro, Antoine Baker. Exact asymptotics for phase retrieval and compressed sensing with random generative priors.

- **On the smoothness assumption**: 
   - We want to reiterate that our theoretical result only requires smoothness of the data distribution after the Gaussian convolution operation. More specifically, we require the Lipschitz constant of the score function of random variable $x + \tau z$, where $x$ is drawn from the prior and $z$ follows the standard Gaussian distribution. Note that this is a mild assumption and has been used in prior literature [CCL+23, CLL23].	

[CCL+23] Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions. Sitan Chen, Sinho Chewi, Jerry Li, Yuanzhi Li, Adil Salim, Anru R. Zhang. 2023.

[CLL23] Improved Analysis of Score-based Generative Modeling: User-Friendly Bounds under Minimal Smoothness Assumptions. Hongrui Chen, Holden Lee, Jianfeng Lu 2023.


- **Extension to experiments with a non-product prior**: 
   - We have subsequently extended our empirical results to non-product priors. Here, the learned denoiser is now a $d$-to-$d$ dimensional function as opposed to a scalar function, and the Onsager term is the averaged divergence of the denoiser. Our MLP architecture parametrizes this $d$-to-$d$ input-output structure as well. Our non-product setting is a mixture of Gaussians prior which in general can accommodate practical distributions through kernel density estimation. We outperform the Bayes-AMP baseline in low-dimensional settings; see Figure 2 in the attached pdf.

Meta-review: In recent years, algorithm unrolling has emerged as a deep learning approach to deal with  the origin of priors in bayesian inference: it involves designing a neural network whose layers can, in theory, simulate iterations of these "optimal" inference algorithms, and then train the network on data generated from an unknown prior.  Despite its empirical successes, the theoretical reasons for the effectiveness of this approach and the nature of the estimators these networks converge to have remained unclear. This work claims to provides the first rigorous learning guarantees for neural networks based on unrolling approximate message passing (AMP). In the context of compressed sensing, it is proven that when trained on data generated from a product prior, the network’s layers approximately converge to the same denoisers used in Bayes AMP. The authors also carry out numerical experiments to further demonstrate the advantages of this unrolled architecture: it can adapt to general priors without explicit knowledge, handle more general measurement scenarios (including non-sub-Gaussian measurements), and shows non-asymptotic improvements over the mean squared error achieved by Bayes AMP.

The reviewers all thought the paper was interesting. They raised some technical concerns that were mostly addressed during the response and revision. Given the rating and the discussions this is a clear accept.
Decision: Accept (poster)