Unrolled denoising networks provably learn to perform optimal Bayesian inference

Authors: aayush  karan
kulin  shah
sitan  chen
yonina c eldar
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 19:27:06

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.

Main paper: 
 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, a 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 a prior for the signal, and the measurements they observe naturally induce a posterior over what the underlying signal could have been.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 to the true prior distribution of the data. In practice, however, this prior is not known a 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 that can improve upon this practicality-performance tradeoff by learning from samples drawn from the data distribution [GL10, MLE21, LTG + 20, SWED23]. The idea is to craft a neural network architecture, each of whose layers is 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 it 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. Thus, at least empirically, algorithm unrolling seems to achieve the best of both worlds, marrying the domain-aware 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 non-algorithmic aspects like representational power and generalization bounds (see Section 1.1 for a detailed discussion).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 on 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 C.1).Approximate message passing and unrolling. Consider the standard Bayesian setup where we observe a noisy measurement 𝑦 of some signal 𝑥, and would like to output an estimate 𝑥 minimizing 𝔼∥𝑥 -𝑥 ∥ 2 . Information-theoretically, the Bayes-optimal estimator for this task is the posterior mean 𝔼[𝑥 | 𝑦], 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 a 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 a 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 𝜂 1 , 𝜂 2 , . . ., we consider generic denoisers given by neural networks 𝑓 1 , 𝑓 2 , . . . and unroll the iterations of AMP into layers of a 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 𝑓 ℓ .Unrolled AMP architectures and variants thereof were originally proposed and empirically investigated by Borgerding et al. [BS16, BSR17] 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 to 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 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-dimensional limit, achieves the same mean squared error as Bayes AMP.Our proof is based on a 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-independent even when the second moment 𝔼 𝑥∼𝑞 ∥𝑥 ∥ 2 scales with the dimension 𝑑. 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 a 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-conditioned and ill-conditioned settings. • We introduce rank-one matrix estimation as a new "model organism" for probing the properties of unrolled networks. To our knowledge, despite its prominence in the theoretical literature on 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 a 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.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 𝑥 is a deterministic sparse vector, and the goal is to converge to 𝑥. 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 (•; 𝜆), in addition to simple parametric families of denoisers like 5-wise linear functions and splines. For these parametric denoisers, on simple priors like Bernoulli-Gaussian 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-parametric method.In the product-prior setting in which we prove our results, this is indeed a 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 a 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 a 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 message-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. 
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 of 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.In compressed sensing, we are given noisy linear measurements 𝑦 ∈ ℝ 𝑚 obtained from an unknown signal 𝑥 ∈ ℝ 𝑑 via the observation processEQUATIONwhere 𝐴 ∈ ℝ 𝑚×𝑑 and 𝜍 ∈ ℝ 𝑚 is a random noise vector with i.i.d. entries drawn from the distribution N(0, 𝜎 2 ). Throughout, we assume that 𝑥 ∼ 𝑝 x for product prior 𝑝 x ≜ 𝑝 ⊗𝑑 , where 𝑝 is some distribution over ℝ. The compressed sensing problem aims to recover the unknown signal 𝑥 with estimate 𝑥 such that the mean squared error (MSE) 1 𝑑 𝔼∥ 𝑥 -𝑥 ∥ 2 is minimal. Throughout this work we focus on the proportional asymptotic setting where we implicitly work with a sequence of such compressed sensing problems indexed by dimension 𝑑, where 𝑑 and 𝑚 = 𝑚(𝑑) jointly tend to infinity and 𝑚(𝑑)/𝑑 → 𝛿 for absolute constant 𝛿 > 0.[DMM09] proposed the following approximate message passing (AMP) algorithm for estimating 𝑥 given 𝐴, 𝑦. The algorithm starts with 𝑥 0 = 0 and 𝑣 0 = 𝑦 and proceeds by 𝑥 ℓ+1 = 𝑓 ℓ (𝐴 ⊤ 𝑣 ℓ + 𝑥 ℓ )(2)EQUATIONwhere 𝑓 ℓ : ℝ → ℝ is a scalar denoiser applied entrywise, 𝑓 ′ ℓ -1 is also applied entrywise, and ⟨•⟩ denotes an entrywise average. The last term in Eq. ( 3) is commonly referred to as the Onsager term. Importantly, the AMP iterates asymptotically satisfy a distributional recursion called state evolution [BM11] . Suppose the entries of 𝐴 are given by 𝐴 𝑖 𝑗 ∼ N(0, 1/𝑚). Define the state evolution parameters (𝜏 ℓ ) via the scalar recursion𝜏 2 ℓ+1 = 𝜎 2 + 1 𝛿 𝔼[(𝑓 ℓ (𝑿 + 𝜏 ℓ 𝒁 ) -𝑿 ) 2 ] with 𝜏 0 = 𝜎 2 + 1 𝛿 𝔼[𝑿 2 ] ,where 𝑿 ∼ 𝑝 and 𝒁 ∼ N(0, 1). Then it is known that as 𝑑 → ∞, the empirical distribution over entries of 𝐴 ⊤ 𝑣 ℓ + 𝑥 ℓ converges in a certain sense to the one-dimensional distribution over 𝑿 + 𝜏 ℓ 𝒁 [BM11] . While updates of AMP can be run with any choice of differentiable 𝑓 ℓ , there is an asymptotically optimal choice that depends on the underlying prior 𝑝 x , and the resulting optimal algorithm is called Bayes AMP. In particular, let 𝜏 * 2 0 = 𝜎 2 + 1 𝛿 𝔼[𝑿 2 ], and define𝑓 * ℓ = 𝔼[𝑿 |𝑿 + 𝜏 * ℓ 𝒁 = •]andEQUATIONwhere 𝑿 ∼ 𝑝 and 𝒁 ∼ N(0, 1). Then setting 𝑓 ℓ = 𝑓 * ℓ for all ℓ in Eqs.(2) and (3) yields Bayes AMP.In the asymptotic limit, i.e. as 𝑚, 𝑑 → ∞, 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 a 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 𝑓 * ℓ . 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 𝜏 2 ℓ with 1 𝑚 ∥𝑣 ℓ ∥ 2 2 , which is typically enough to stabilize Bayes AMP. The fact that this is a valid estimate follows by state evolution, which ensures that in the infinite dimensional limit, the entries of 𝑣 ℓ are distributed according to N(0, 𝜏 2 ℓ ) [BM11] .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 𝑝 x . We consider replacing each scalar denoiser 𝑓 ℓ 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, a 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 {(𝑦 𝑖 , 𝑥 𝑖 )} 𝑁 𝑖=1 generated according to Eq. (1)with 𝑥 𝑖 ∼ 𝑝 x for all 𝑖. Let 𝐿 denote the number of layers in our unrolled network, and let F denote a family of MLPs with fixed architecture (i.e. fixed depth and width) constrained to a two-dimensional input and one-dimensional output. For each ℓ ∈ [0, 𝐿 -1], initialize an MLP𝑓 ℓ : ℝ 2 → ℝ chosen from F . Set 𝑥 0 = 0 𝑑 ∈ ℝ 𝑑 and 𝑣 0 = 𝑦 𝑖 ∈ ℝ 𝑚 , for a given training input 𝑦 𝑖 .Then for each layer ℓ ∈ [0, 𝐿 -1], our network computes the forward passEQUATIONwhere 𝜏 ℓ = ∥ 𝑣 ℓ ∥ 2 / √ 𝑚 and 𝜕 1 denotes differentiation with respect to the first input parameter. The notation 𝑓 ℓ ( • ; 𝜏 ℓ ) denotes applying the scalar function 𝑓 ℓ ( • , 𝜏 ℓ ) entrywise. We emphasize that 𝑓 ℓ is tied to 𝜕 1 𝑓 ℓ ; 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. Naïvely, one might consider training the 𝐿-layer network end-to-end on the mean squared errors of the network estimates -i.e., either with loss functionL 𝐶𝑆 = 1 𝑁 𝑁 𝑖=1 1 𝑑 ∥ 𝑥 𝑖 𝐿 -𝑥 𝑖 ∥ 2 2 for compressed sensing or L 𝑀𝐸 = 1 𝑁 𝑁 𝑖=1 1 𝑑 ∥ 𝑥 𝑖 𝐿 𝑥 𝑖⊤ 𝐿 -𝑥 𝑖 𝑥 𝑖⊤ ∥ 2𝐹 for rank-one matrix estimation.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 𝑓 ℓ on the mean squared error loss for the layer-ℓ estimate. If Ψ denotes an LDNet with 𝐿 layers, let Ψ[0 : ℓ] denote the subnetwork that consists only of the first ℓ + 1 layers of Ψ, with denoiser 𝑓 ℓ at layer ℓ for 0 ≤ ℓ ≤ 𝐿 -1. Then our training procedure follows Algorithm 1. Note we initialize the ℓ-th denoiser weights with the previous learned denoiser before training -while this is not relevant to our theoretical results in Section 3, empirically, we find that this initialization is necessary to avoid being trapped in 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.Input: Training data D, LDNet Ψ for ℓ = 0 to ℓ = 𝐿 -1 do if ℓ > 0 then Initialize 𝑓 ℓ ← 𝑓 ℓ -1 ;Freeze learnable weights in 𝑓 𝑘 for 𝑘 < ℓ;Train Ψ[0 : ℓ] on D;// Optional finetuning step Unfreeze learnable weights in 𝑓 𝑘 for 𝑘 < ℓ and train Ψ[0 : ℓ] on D; Output: Fully trained ΨThe proof of optimality of Bayes AMP among all implementations of AMP for the problems we consider, as 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 at layer ℓ 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 is validated in both our theory and experiments. 
provably learning bayes amp: 
 We now provide theoretical guarantees that our unrolled denoising network can learn Bayesoptimal denoisers when trained in a layerwise fashion. Consider any prior 𝑝 x = 𝑝 ⊗𝑑 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 𝜏 ≥ 0, let 𝑝 (•; 𝜏) denote the density of the convolution 𝑝 ★ N(0, 𝜏 2 ). We assume that:1. 𝑝 is 𝑅-sub-Gaussian with 𝔼 𝑿 ∼𝑝 [𝑿 ] = 0. 2. The score function 𝜕 1 𝑝 (•; 𝜏) is 𝐵-Lipschitz for all 𝜏 ≥ 𝜎 2 ,where 𝜎 2 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 sub-Gaussianity holds for any distribution with bounded support (see Section 2.5 in [Ver18] ) and the Lipschitzness of the score function is a 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, a 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 𝑝, and formally state our main result. The full proof is provided in Appendix B.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 𝑓 by 𝑥 ℓ and 𝑣 ℓ . The Onsager term in AMP ensures that for any iteration ℓ, the distribution of 𝐴 ⊤ 𝑣 ℓ + 𝑥 ℓ asymptotically behaves as if every coordinate is i.i.d. as 𝑑 → ∞ for any ℓ (Lemma 1). Therefore, it suffices to learn the denoiser for any fixed coordinate.2 Without loss of generality we consider the first coordinate and try to learn the denoiser function by minimizing the following objective:EQUATIONwhere 𝐴 ⊤ 𝑗 denotes the 𝑗th row of 𝐴 ⊤ and 𝑥 ( 𝑗 ) denotes the 𝑗th coordinate of the vector 𝑥. As the learning and generalization guarantee is identical for all ℓ, 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 𝑝, we will work with the following notion of the complexity of a class of scalar functions from [AZLL19] , which we will apply to the denoisers that arise in Bayes AMP: Definition 1 (Scalar function complexity [AZLL19] ). Let 𝐶 > 0 be some sufficiently large absolute constant (e.g. 10 4 ). Given any smooth function 𝜙 : ℝ → ℝ and a parameter 𝛼 > 0, we define complexity of 𝜙 at scale 𝛼 as follows. Suppose 𝜙 admits a power series expansion𝜙 (𝑧) = ∞ 𝑖=0 𝑐 𝑖 𝑧 𝑖 . Then ℭ 𝜀 (𝜙, 𝛼) ≜ ∞ ∑︁ 𝑖=0 1 + (log(1/𝜀)/𝑖) 𝑖/2 • (𝐶𝛼) 𝑖 |𝑐 𝑖 | and ℭ 𝑠 (𝜙, 𝛼) ≜ 𝐶 ∞ ∑︁ 𝑖=0 (𝑖 + 1) 1.75 𝛼 𝑖 |𝑐 𝑖 |.Given a class of scalar functions F , we define the complexity ofF at scale 𝛼 by ℭ 𝜀 (F, 𝛼) = sup 𝜙 ∈F ℭ 𝜀 (𝜙, 𝛼) (and similarly ℭ 𝑠 (F, 𝛼) = sup 𝜙 ∈F ℭ 𝑠 (𝜙, 𝛼)).Intuitively, ℭ 𝜀 (F, 𝛼) and ℭ 𝑠 (F, 𝛼) both captures how much functions in function class F can be approximated using a low-degree polynomial. For any function 𝜙 and any 𝛼, the above complexities are related by ℭ 𝑠 (𝜙, 𝛼) ≤ ℭ 𝜀 (𝜙, 𝛼) ≤ ℭ 𝑠 (𝜙, 𝑂 (𝛼)) × poly(1/𝜀) because (𝐶 √︁ log(1/𝜀)/ √ 𝑖) 𝑖 ≤ 𝑒 𝑂 (log 1/𝜀 ) = poly(1/𝜀) for all 𝑖. We provide more intuition on how ℭ 𝜀 (𝜙, 𝛼) and ℭ 𝑠 (𝜙, 𝛼) scales with 𝜀, 𝛼, under mild assumptions on 𝜙 in Section B.5.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 𝜀 2 ∈ (0, 1) and 𝜀 1 ∈ (0, 1/ℭ 𝑠 (F, 𝑅(log 1/𝜀 2 ) 3/2 )), there exists𝑀 0 = poly(ℭ 𝜀 1 (𝑓 * , 𝑅(log 1/𝜀 2 ) 3/2 ), 1/𝜀 1 ) and 𝑁 0 = poly(𝐿ℭ 𝑠 (𝑓 * , 𝑅(log 1/𝜀 2 ) 3/2 ), 1/𝜀 1 )such that the following holds.Let 𝐿 be any positive integer. Consider an LDNet of depth 𝐿 with MLP denoisers 𝑓 ℓ given by the MLP architecture in Eq. (8) with 𝑚 ≥ 𝑀 0 neurons. Suppose the network is trained by running gradient descent from random initialization with step size 𝜂 = Θ(1/(𝜀 1 𝑚)) on 𝑛 ≥ 𝑁 0 samples of the form (𝑦 𝑖 , 𝑥 𝑖 ), where each training example is generated by independently sampling Gaussian matrix 𝐴 with entries i.i.d. from N(0, 1/𝑚), sampling 𝑥 𝑖 ∼ 𝑝 x = 𝑝 ⊗𝑑 , and forming 𝑦 𝑖 = 𝐴𝑥 𝑖 + 𝜍 for 𝜍 ∼ N(0, 𝜎 2 • Id).After 𝑇 = Θ(ℭ 𝑠 (𝑓 * ℓ , 𝑅(log 1/𝜀 2 ) 3/2 ) 2 /𝜀 2 1) steps of gradient descent, with high probability the activations ( 𝑥 𝐿 , 𝑣 𝐿 ) and denoiser 𝑓 𝐿 at the output layer of the LDNet (see (5)) satisfy𝔼 𝑥,𝐴 1 𝑑 ∥ 𝑓 𝐿 (𝐴 ⊤ 𝑣 𝐿 + 𝑥 𝐿 ; 𝜏 𝐿 ) -𝑥 ∥ 2 ≲ MSE AMP (𝐿) + 𝑅 2 𝐵 2 𝛿𝜎 7 𝐿+1 (𝜀 1 + 𝜀 2 ) + 𝑜 𝑑 (1) , where MSE AMP (𝐿) ≜ 𝔼 𝑥,𝐴 1 𝑑 ∥𝑓 * 𝐿 (𝐴 ⊤ 𝑣 𝐿 + 𝑥 𝐿 ; 𝜏 𝐿 ) -𝑥 ∥ 2is the error achieved by running 𝐿 steps of Bayes AMP.Observe that the level of overparametrization in terms of number of samples 𝑛 and number of hidden neurons 𝑚 needed in Theorem 2 is dimension-free, unlike in typical NTK analyses. This happens because state evolution effectively allows us to convert the learning problem in 𝑑 dimensions to a learning problem in 1 dimension: we can effectively assume that the entries of 𝐴 ⊤ 𝑣 ℓ + 𝑥 ℓ converge in an appropriate sense to the distribution of 𝑿 + 𝜏 ℓ 𝒁 for 𝑿 ∼ 𝑝 and 𝒁 ∼ N (0, 1).This ensures that the learning objective effectively reduces to minimizing 𝔼[(𝑓 ℓ (𝑿 + 𝜏 ℓ 𝒁 ) -𝑿 ) 2 ] over a parametrized family of denoisers 𝑓 ℓ . The latter objective is often referred to as the score matching objective (in one dimension), which is minimized by the Bayes-optimal denoiser𝑓 * ℓ = 𝔼[𝑿 |𝑿 + 𝜏 ℓ 𝒁 = •] at each layer ℓ.A 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 (𝑿 + 𝜏 ℓ 𝒁, 𝑿 ).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-𝑠 polynomials are needed, then the runtime of the algorithm must scale as 𝑑 𝑂 (𝑠 ) . This would be prohibitively expensive if 𝑠 is increasing in the dimension 𝑑. Fortunately however, because we are able to reduce to one-dimensional training dynamics, we ultimately achieve much more favorable scaling in 𝑑. 
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 to be computationally optimal? b) In these settings, does our network learn the optimal denoisers? c) Are there settings where our methods offer a performance advantage over AMP?Implementation details. We set 𝑚 = 250, 𝑑 = 500 and fix a random Gaussian sensing matrix 𝐴 ∈ ℝ 250×500 . We consider two choices of prior for our experiments: Bernoulli-Gaussian and ℤ 2 (i.e. uniform over {1, -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 dataset {𝑦 𝑖 , 𝑥 𝑖 } 𝑁 𝑖=1 with 𝑁 = 2 15 samples by sampling from the prior and using Eq. (1). We train layerwise with finetuning as in Algorithm 1.For each prior, we also implemented Bayes AMP using the corresponding optimal denoiser. As an additional "semi-prior-aware" baseline, we replace the MLP denoisers in LDNet with "guided denoisers" that have the same functional form as the optimal denoisers but contain trainable parameters; see Appendix E for more details on the precise functional forms used. By convention, we report performance results for all methods by the normalized mean squared error (NMSE)∥ 𝑥 -𝑥 ∥ 2 2 /∥𝑥 ∥ 2 2 .Bernoulli-Gaussian prior. Here, each entry of 𝑥 is independently drawn from a standard normal distribution and set to 0 with probability 1 -𝜀; i.e., 𝑝 x = 𝑝 ⊗𝑑 where 𝑝 (𝑥) = 𝜀 N(0, 1; 𝑥) + (1 -𝜀) 𝛿 (𝑥), where 𝛿 denotes the Dirac delta at 𝑥 = 0. To match the setting considered in the prior work of Borgerding et al. [BSR17] , we set the masking probability to be 𝜀 = 0.1 and the measurement noise to be 𝜎 2 = 2 • 10 -5 .We plot the NMSE that our unrolled network and baselines achieve in decibels (dB); that is, 10 log 10 (NMSE), in Figure 1 . LDNet almost perfectly matches the NMSE of 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 is logarithmic, the difference in error achieved is negligible.ℤ 2 prior. Here each entry of 𝑥 is chosen from {-1, 1} with probability 1 2 ; i.e., 𝑝 x = 𝑝 ⊗𝑑 for 𝑝 (𝑥) = 1 2 𝛿 -1 (𝑥)+ 1 2 𝛿 +1 (𝑥). To examine a higher noise regime and to ensure Bayes AMP converged within a reasonable number of iterations, we set the measurement noise to be 𝜎 2 = 0.075. Figure 1 demonstrates that LDNet again recovers Bayes AMP at every iteration, even slightly outperforming by layer 15, achieving an NMSE of 0.4267 (an improvement of 1.28%).From Figure 2 we can observe qualitatively that the learned MLP denoisers recover the functional form for the optimal denoiser at each iteration, as our theory suggests. Interestingly, although the denoisers were trained relative to a fixed sensing matrix, they appear to learn the Bayes AMP denoiser that is measurement-independent, and in Appendix D.3 we show that the performance of these learned denoisers actually transfers to other randomly drawn sensing matrices 𝐴. LDNet on the Bernoulli-Gaussian and ℤ 2 priors relative to their optimal denoisers over a range of inputs in (-2, 2). The state evolution input 𝜏 ℓ to each denoiser is set to be its empirical estimate.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 [GL10] and LAMP [BSR17] reparameterize 𝐴 ⊤ in Eqs. (2) and (3) as a new learnable matrix 𝐵, 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 𝑑 → ∞ limit. In the context of our unrolling method, we posit that learning the matrix 𝐵 can be thought of as learning finite dimensional corrections to the Bayes AMP iterations. In Appendix D.1, we demonstrate that the lower the signal dimensionality, the larger the performance improvement of LDNet with learned matrix 𝐵 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 of the measurement matrix 𝐴, 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 to accommodate non-product priors. In D.2, we demonstrate LDNet to surpass Bayes AMP in a non-separable mixture-of-gaussians prior. 
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 overparametrization needed for our result to hold is independent of the dimension, unlike existing results in the NTK literature. One important consequence of these results is that a one-dimensional score function is learnable with gradient descent, for which only representational, as opposed to algorithm, results existed previously in the literature [MW23] .We supplemented our theory with extensive numerical experiments, confirming that LDNet can 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 optimallye.g. inference in low dimensions, non-Gaussian designs, and non-product priors -we demonstrate that LDNet outperforms Bayes AMP. We thus establish unrolling denoisers as a powerful, practical addition to the algorithmic toolkit for Bayesian inverse problems.One limitation is that our theoretical results are currently limited to the product prior setting. The non-product setting is difficult because even though state evolution is known here [BMN20] , proving an unrolled network converges to the right denoisers essentially amounts to proving that one can learn the score functions of a general data distribution. Additionally, it is subtle to define the right architecture, as the denoisers are no longer scalar, and a 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 in denoising error implies closeness in the state evolution parameter 𝜏 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 as 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 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 "non-asymptotic corrections" that these extra learnable parameters are imposing.Theory for unrolling ISTA. The existing theory for algorithm unrolling almost exclusively focuses on 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 a Bayesian setting: the signal 𝑥 is a deterministic vector assumed to have some level of sparsity, and the goal is to converge to 𝑥. 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 iteration-efficient than vanilla ISTA [MB16, CLWY18, XWG + 16, LC19, CLWY21].For example, [MB16] showed that if each layer performs a proximal splitting step with respect to a Gram matrix which admits a factorization with certain nice properties, then the iterates computed by each layer converge to 𝑥 at an accelerated rate. [CLWY18] showed that if the learned weights are such that the activations converge to 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 a setting of weights for which the iterates converge at 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.[LC19] 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 a modified Polyak-Lojasiewicz inequality around initialization and thus the network would 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 𝜂 st (•; 𝜆) with trainable parameter 𝜆 and trainable weight matrices playing the role of 𝐴 ⊤ in the AMP update (see Eqs. (2) and (3)). They found empirically that their unrolled network outperforms AMP with soft-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-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.. Among the various direct follow-ups to [BS16, BSR17], e.g. [ITW19, MJC21, MMB17], 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 a 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 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 is 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 a computational-statistical gap [BPW18] .On the experimental side, our focus was on synthetic setups instead of image recovery, with an emphasis on probing which aspects of the problem setup and which trainable parameters allow unrolled AMP to outperform Bayes AMP.The following result shows that we can characterize the behavior of the AMP iterates 𝑥 ℓ in the limit as 𝑑 → ∞.Lemma 1 (Asymptotic characterization of AMP iterates [BM11] ). Let 𝐴 ∈ ℝ 𝑚×𝑑 be a sensing matrix with i.i.d. entries 𝐴 𝑖 𝑗 ∼ N(0, 1/𝑚). Assume 𝑚/𝑑 → 𝛿 ∈ (0, ∞). Consider a sequence of vectors {𝑥 (𝑑), 𝜂 (𝑑)} indexed by dimension whose empirical distribution converges weakly to probability measures 𝑝 𝑥 and 𝑝 𝜂 on ℝ with bounded moments. Then, for any pseudo-Lipschitz function 𝜓 : ℝ 𝑡 → ℝ and all ℓ ≥ 0, almost surelylim 𝑑→∞ 1 𝑑 𝑑 ∑︁ 𝑖=1 𝜓 (𝑥 (𝑖 ) ℓ , 𝑥 (𝑖 ) ) = 𝔼[𝜓 (𝑓 ℓ (𝑿 + 𝜏 ℓ 𝒁 ), 𝑿 )] lim 𝑑→∞ 1 𝑑 𝑑 ∑︁ 𝑖=1 𝜓 (𝑥 (𝑖 ) -(𝐴 ⊤ 𝑖 𝑣 ℓ + 𝑥 (𝑖 ) ℓ ), 𝑥 (𝑖 ) ) = 𝔼[𝜓 (𝜏 ℓ 𝒁, 𝑿 )] .where 𝑿 ∼ 𝑝 𝑥 and 𝒁 ∼ N (0, 1). The state evolution parameters 𝜏 0 , 𝜏 1 , . . . are recursively defined as followsEQUATIONObserve that to minimize the variance 𝜏 2 ℓ at each iteration, the optimal choice of denoiser 𝑓 ℓ is𝑓 * ℓ (𝑥) = 𝔼[𝑿 | 𝑿 + 𝜏 ℓ 𝒁 = 𝑥].Here we prove a 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 a given layer of our unrolled AMP architecture as a one-hidden-layer ReLU network in the following form: letting 𝑤 [𝑡 ] 𝑗 denote the weight of the 𝑗th neuron after 𝑡 steps gradient descent, we considerEQUATIONWe initialize the entries of weights 𝑤 [0] 𝑗 and biases 𝑏 [0] 𝑗 to be i.i.d. from N (0, 1/𝑚) and entries of 𝑎 [0] 𝑗 to be i.i.d. from N (0, 𝜀 𝑎 ) for some fixed 𝜀 𝑎 ∈ (0, 1]. We only train the weights 𝑤 𝑗 of hidden layers to simplify the analysis and freeze the bias term 𝑏 𝑗 and output layer weights 𝑎 𝑗 at initialization. To update the weights at time 𝑡, we take one step of online gradient descent with respect to the loss in Eq. ( 6) with step size 𝜂 on a fresh sample (𝑥, 𝑦, 𝐴); 3 here we use the learned denoisers from the previous layers to compute 𝑥 ℓ and 𝑣 ℓ .We begin by proving the following claim that training a layer of the unrolled network on gradient descent from random initialization will converge to a denoiser that is competitive with the optimal denoiser under the objective in Eq. (6).Lemma 2 (Learning the denoiser within 𝐿 2 error). For every ℓ, assume every coordinate of 𝑣 ℓ and 𝑥 ℓ are sub-Gaussian random variables with constant 𝑅. Then, for every 𝜀 2 ∈ (0, 1) and 𝜀 1 ∈ (0, 1/ℭ 𝑠 (F, 𝑅(log 1/𝜀 2 ) 3/2 )), there exists 𝑀 0 = poly(ℭ 𝜀 1 (𝑓 * ℓ , 𝑅(log 1/𝜀 2 ) 3/2 ), 1/𝜀 1 ) and 𝑁 0 = poly(ℭ 𝑠 (𝑓 * ℓ , 𝑅(log 1/𝜀 2 ) 3/2 ), 1/𝜀 1 ) such that for every 𝑚 ≥ 𝑀 0 and 𝑛 ≥ 𝑁 0 , choosing learning rate 𝜂 = Θ(1/(𝜀 1 𝑚)) and running gradient descent from random initialization for 𝑇 = Θ(ℭ 𝜀 (𝜙, 𝑅(log 1/𝜀 2 ) 3/2 ) 2 𝑝 2 /𝜀 2 1 ), with high probability, 5 we can apply Lemma 3 to obtain a neural network that achieves 𝜀 1 error. 3 In our experiments, we keep the measurement matrix 𝐴 fixed and only sample fresh (𝑥, 𝑦) pairs, but in our theoretical result, we assume that gradient descent uses fresh samples (𝑥, 𝑦, 𝐴) 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 𝐴. 4 Here we 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. In Lemma B.4, we can write ∥𝑊 𝑡 +1 -𝑊 * ∥ 2 𝐹 = ∥𝑊 𝑡 -𝜂∇𝐿 𝐹 (Z,𝑊 𝑡 ) -𝑊 * ∥ 2 𝐹 and therefore, getting the equality of𝔼 𝑥,𝑦,𝐴 ( 𝑓 ℓ (𝐴 ⊤ 1 𝑣 ℓ + 𝑥 (1) ℓ ; 𝜃 𝑇 ) -𝑥 (1) ) 2 ≤ min 𝑔∈F 𝔼 (𝑔(𝐴 ⊤ 1 𝑣 ℓ + 𝑥 (1) ℓ ) -𝑥 (1) ) 2 + 𝜀 1 + 𝜀 2 .𝑣 ℓ + 𝑥 (𝑖 ) ℓ | ≲ 𝑅(log(1/𝜀 2 )) 3/2 ,2𝜂 ⟨𝑊 𝑡 -𝑊 * , ∇𝐿 𝐹 (Z,𝑊 𝑡 )⟩ = (∥𝑊 𝑡 -𝑊 * ∥ 2 𝐹 -∥𝑊 𝑡 +1 -𝑊 * ∥ 2 𝐹 )/2𝜂 + (𝜂/2)∥∇𝐿 𝐹 (Z,𝑊 𝑡 )∥ 2𝐹 and using this equality in Eq.(B.7) of [AZLL19] . The rest of the proof remains the same.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, to the optimal error for predicting a sample from the univariate prior 𝑝 given a Gaussian corruption. This will follow by state evolution. Lemma 4. Let 𝑓 ℓ (•, 𝜃 𝑡 ) be the learned neural network using gradient descent after 𝑡 timesteps such that the conditions of Lemma 2 satisfies. Then, with high probability, we havelim 𝑑→∞ 𝔼 𝑥,𝐴 1 𝑑 ∥ 𝑓 ℓ (𝐴 ⊤ 𝑣 ℓ + 𝑥 ℓ , 𝜃 𝑇 ) -𝑥 ∥ 2 ≤ 𝔼[( 𝑓 * ℓ (𝑿 + 𝜏 ℓ 𝒁 ) -𝑿 ) 2 ] + 𝜀 1 + 𝜀 2 .Additionally, the following statement holds with high probability:𝔼[( 𝑓 ℓ (𝑿 + 𝜏 ℓ 𝒁, 𝜃 𝑇 ) -𝑿 ) 2 ] ≤ 𝔼[( 𝑓 * ℓ (𝑿 + 𝜏 ℓ 𝒁 ) -𝑿 ) 2 ] + 𝜀 1 + 𝜀 2 .Proof. Observe that 𝐴 ⊤ 𝑗 𝑣 ℓ + 𝑥 ( 𝑗 ) ℓ follows the same distribution for all 𝑗. Therefore, using Lemma 2, we obtain that𝔼 𝑥,𝐴 ( 𝑓 ℓ (𝐴 ⊤ 𝑗 𝑣 ℓ + 𝑥 ( 𝑗 ) ℓ , 𝜃 𝑇 ) -𝑥 ( 𝑗 ) ) 2 ≤ min 𝑔 𝔼 (𝑔(𝐴 ⊤ 𝑗 𝑣 ℓ + 𝑥 ( 𝑗 ) ℓ ) -𝑥 ( 𝑗 ) ) 2 + 𝜀 1 + 𝜀 2 .As the distribution of 𝐴 ⊤ 𝑗 𝑣 ℓ + 𝑥 ( 𝑗 ) ℓ is the same for all 𝑗 ∈ [𝑑], we use the same learned denoiser for all the coordinates. Therefore, we can rewrite the above equation as1 𝑑 𝑇 ∑︁ 𝑖=1 𝔼 𝑥,𝐴 ∥ 𝑓 (𝐴 ⊤ 𝑣 ℓ + 𝑥 ℓ , 𝜃 𝑇 ) -𝑥 ∥ 2 ≤ min 𝑔 1 𝑑 𝔼 ∥𝑔(𝐴 ⊤ 𝑣 ℓ + 𝑥 ℓ ) -𝑥 ∥ 2 + 𝜀 1 + 𝜀 2 .Now, we want to use state evolution as 𝑑 → ∞. Note that as 𝑑 → ∞, using Lemma 1 with 𝜓 function as 𝜓 (𝑎, 𝑏) = (𝑎 -𝑏) 2 , we havelim 𝑑→∞ 1 𝑑 ∥𝑔(𝐴 ⊤ 𝑣 ℓ + 𝑥 ℓ ) -𝑥 ∥ 2 = 𝔼[(𝑔(𝑿 + 𝜏 ℓ 𝒁 ) -𝑿 ) 2 ]for any function 𝑔 ∈ F . As the quantity inside expectation converges to 𝔼[(𝑔(𝑿 + 𝜏 ℓ 𝒁 ) -𝑿 ) 2 ],using monotone convergence theorem, we have lim𝑑→∞ min 𝑔 𝔼[ 1 𝑑 ∥𝑔(𝐴 ⊤ 𝑣 ℓ + 𝑥 ℓ ) -𝑥 ∥ 2 ] = min 𝑔 lim 𝑑→∞ 𝔼[ 1 𝑑 ∥𝑔(𝐴 ⊤ 𝑣 ℓ + 𝑥 ℓ ) -𝑥 ∥ 2 ] = min 𝑔 𝔼[(𝑔(𝑿 + 𝜏 ℓ 𝒁 ) -𝑿 ) 2 ].Using this result and applying limits on both sides of Lemma 2, we obtain lim𝑑→∞ 𝔼 𝑥,𝐴 1 𝑑 ∥ 𝑓 ℓ (𝐴 ⊤ 𝑣 ℓ + 𝑥 ℓ , 𝜃 𝑇 ) -𝑥 ∥ 2 ≤ min 𝑔 𝔼[(𝑔(𝑿 + 𝜏 ℓ 𝒁 ) -𝑿 ) 2 ] + 𝜀 1 + 𝜀 2 .The minimizer of 𝔼[(𝑔(𝑿+ 𝜏 ℓ 𝒁 ) -𝑿 ) 2 ] is given by 𝑓 * ℓ (•) = 𝔼[𝑿 | 𝑿 + 𝜏 ℓ 𝒁 = •].Using this 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 to 𝔼[( 𝑓 ℓ (𝑿 + 𝜏 ℓ 𝒁, 𝜃 𝑇 ) -𝑿 ) 2 ] using the monotone convergence theorem. This gives the second result of the lemma statement.The right-hand side of the bound in the above Lemma corresponds to the minimum mean squared error achievable for denoising at noise scale 𝜏 ℓ , where 𝜏 ℓ is 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. (𝑿 + 𝜏 ℓ 𝒁 ) -𝑿 ) 2 ] ≲ 𝔼[(𝑓 * ℓ (𝑿 + 𝜏 ℓ 𝒁 ) -𝑿 ) 2 ] + 𝑉 2 𝐵 2 𝜎 6 | 𝜏 ℓ -𝜏 ℓ | where 𝑓 * ℓ (𝑥) = 𝔼[𝑿 |𝑿 + 𝜏 ℓ 𝒁 = 𝑥] and 𝑓 * ℓ (𝑥) = 𝔼[𝑿 |𝑿 + 𝜏 ℓ 𝒁 = 𝑥].Proof. Rewriting the error between 𝑓 * ℓ and 𝑓 * ℓ , we haveEQUATION)The term in Eq. ( 10) can be upper bounded by (𝐵( 𝜏 ℓ -𝜏 ℓ )) 2 / 𝜏 4 ℓ because 𝑓 ℓ is Lipschitz by assumption. Using Cauchy-Schwartz for the term in Eq. ( 11), the second term is upper bounded byEQUATIONThe squared term in Eq. ( 12) can be upper bounded by𝔼[( 𝑓 * ℓ (𝑿 + 𝜏 ℓ 𝒁 ) -𝑿 ) 2 ] = 𝔼[(( 𝑓 * ℓ (𝑿 + 𝜏 ℓ 𝒁 ) -𝑓 * ℓ (𝑿 + 𝜏 ℓ 𝒁 )) 2 )] + 2𝔼[( 𝑓 * ℓ (𝑿 + 𝜏 ℓ 𝒁 ) -𝑓 * ℓ (𝑿 + 𝜏 ℓ 𝒁 )) (𝑓 * ℓ (𝑿 + 𝜏 ℓ 𝒁 ) -𝑿 )] + 𝔼[(𝑓 * ℓ (𝑿 + 𝜏 ℓ 𝒁 ) -𝑿 ) 2 ] .We can upper bound the first term above using Lemma 6. Likewise, the second term can be bounded by Cauchy-Schwarz and Lemma 6. In this way, using that 𝑉 2 ≥ 𝜏 ℓ ≥ 𝜏 ℓ ≥ 𝜎 2 , we obtain that𝔼[( 𝑓 * ℓ (𝑿 + 𝜏 ℓ 𝒁 ) -𝑿 ) 2 ] ≲ 𝔼[(𝑓 * ℓ (𝑿 + 𝜏 ℓ 𝒁 ) -𝑿 ) 2 ] + 𝑉 2 𝐵 2 𝜎 6 ( 𝜏 ℓ -𝜏 ℓ ) 2 + 𝑉 2 𝐵 𝜎 3 | 𝜏 ℓ -𝜏 ℓ | .Combining this bound with Eq. ( 13), we obtain the bound.The proof above relies on the following "score perturbation lemma" showing that the optimal denoiser is Lipschitz with respect to the noise scale.Lemma 6 (Score perturbation lemma). Let 𝑝 (•; 𝜏) be the density function of 𝑿 + 𝜏𝒁 . Let the score function 𝜕 1 log 𝑝 (•; 𝜏) be 𝐵-Lipschitz continuous. Then, the 𝐿 2 error between the optimal denoisers at two noise scales 𝜏 ℓ and 𝜏 ℓ is given by𝔼 ( 𝑓 * ℓ (𝑿 + 𝜏 ℓ 𝒁 ) -𝑓 * ℓ (𝑿 + 𝜏 ℓ 𝒁 )) 2 ≤ 𝜏 2 ℓ 𝐵 2 ( 𝜏 ℓ -𝜏 ℓ ) 2 + 𝐵 2 ( 𝜏 ℓ -𝜏 ℓ ) 4 𝜏 4 ℓ 𝜏 2 ℓ 𝑉 2 .Proof. Denote the probability density function of 𝑿 + 𝜏 ℓ 𝒁 random variable at value 𝑥 as 𝑝 (𝑥; 𝜏 ℓ ).Assuming 𝜕 1 log 𝑝 (•; 𝜏 ℓ ) is 𝐵-Lipschitz function and using Lemma C.11 of [LLT22] , we have𝔼 𝑥∼𝑝 (•;𝜏 ℓ ) (𝜕 1 log 𝑝 (𝑥; 𝜏 ℓ ) -𝜕 1 log 𝑝 (𝑥; 𝜏 ℓ )) 2 ≲ 𝐵 2 ( 𝜏 ℓ -𝜏 ℓ ) 2 + 𝐵 2 ( 𝜏 ℓ -𝜏 ℓ ) 4 𝔼 (𝜕 1 log 𝑝 (𝑥; 𝜏 ℓ )) 2 . Using Tweedie's formula (𝜕 1 log 𝑝 (𝑥; 𝜏 ℓ ) = (𝔼[𝑿 |𝑿 + 𝜏 ℓ 𝒁 = 𝑥] -𝑥)/𝜏 2 ℓ ), we have 𝔼 𝑥∼𝑝 (•;𝜏 ℓ ) (𝜕 1 log 𝑝 (𝑥; 𝜏 ℓ ) -𝜕 1 log 𝑝 (𝑥; 𝜏 ℓ )) 2 = 𝔼 𝑥∼𝑝 (𝑥;𝜏 ) ( 𝜏 2 ℓ -𝜏 2 ℓ )(𝔼[𝑿 |𝑿 + 𝜏 ℓ 𝒁 = 𝑥] -𝑥) + 𝜏 2 ℓ (𝔼[𝑿 |𝑿 + 𝜏 ℓ 𝒁 = 𝑥] -𝔼[𝑿 |𝑿 + 𝜏 ℓ 𝒁 = 𝑥]) 𝜏 2 ℓ 𝜏 2 ℓ This implies that 𝔼 (𝔼[𝑿 |𝑿 + 𝜏 ℓ 𝒁 = 𝑥] -𝔼[𝑿 |𝑿 + 𝜏 ℓ 𝒁 = 𝑥]) 2 ≲ 𝜏 2 ℓ 𝔼 𝑥∼𝑝 (•;𝜏 ℓ ) (𝜕 1 log 𝑝 (𝑥; 𝜏 ℓ ) -𝜕 1 log 𝑝 (𝑥; 𝜏 ℓ )) 2 + ( 𝜏 2 ℓ -𝜏 2 ℓ ) 2 𝜏 4 ℓ 𝜏 4 ℓ 𝔼[(𝔼[𝑿 |𝑿 + 𝜏 ℓ 𝒁 = 𝑥] -𝑥) 2 ] ≲ 𝜏 2 ℓ 𝐵 2 ( 𝜏 ℓ -𝜏 ℓ ) 2 + 𝐵 2 ( 𝜏 ℓ -𝜏 ℓ ) 4 𝜏 4 ℓ 𝜏 2 ℓ 𝔼 (𝔼[𝑿 |𝑿 + 𝜏 ℓ 𝒁 = 𝑥] -𝑥) 2where the last inequality uses the fact that 𝜏 ℓ ≥ 𝜏 ℓ because 𝜏 ℓ is obtained using optimal denoiser. Additionally, we have 𝜏 0 = 𝜎 2 + (1/𝛿)𝔼[𝑿 2 ] ≤ 𝜏 2 + 𝑉 2 . Using law of total variance, we have var[𝑿 ] = 𝔼[var(𝑿 |𝑿 + 𝜏 ℓ 𝒁 )] + var(𝔼[𝑿 |𝑿 + 𝜏 ℓ 𝒁 ]). Therefore, we have 𝔼[𝑿 2 ] ≥ 𝔼[(𝔼[𝑿 |𝑿 + 𝜏 ℓ 𝒁 = 𝑥] -𝑥) 2] which implies the claimed bound.We are now ready to conclude the proof of our main result.Proof of Theorem 2. Recall that the state evolution recursion for 𝜏 ℓ is defined using the optimal denoiser 𝑓 * ℓ at layer ℓ in Lemma 1. Similarly, the state evolution recursion for 𝜏 ℓ is applied using learned denoiser 𝑓 ℓ ( • , 𝜃 𝑇 ). At some places, we will use 𝑓 ℓ ( • ) to denote 𝑓 ℓ ( • , 𝜃 𝑇 ) for brevity. By the state evolution recursion for 𝜏 ℓ , the error between 𝜏 ℓ and 𝜏 ℓ is given𝜏 2 ℓ -𝜏 2 ℓ = 1 𝛿 (𝔼[( 𝑓 ℓ -1 (𝑿 + 𝜏 ℓ -1 𝒁 ) -𝑿 ) 2 ] -𝔼[𝑓 * ℓ -1 (𝑿 + 𝜏 ℓ -1 𝒁 ) -𝑿 ) 2 ]) ≲ 1 𝛿 𝑉 2 𝐵 2 𝜎 6 | 𝜏 ℓ -1 -𝜏 ℓ -1 | 2 + 𝜀 1 + 𝜀 2, where the last inequality follows from Lemma 4 and Lemma 5. As 𝜏 ℓ + 𝜏 ℓ ≥ 2𝜎 for all ℓ and 𝜏 ℓ ≥ 𝜏 ℓ , we have| 𝜏 ℓ -𝜏 ℓ | ≲ 𝑉 2 𝐵 2 𝛿𝜎 7 | 𝜏 ℓ -1 -𝜏 ℓ -1 | + 𝜀 1 + 𝜀 2 𝛿𝜎 .Solving this recurrence, we have that the error in the state evaluation parameters after 𝐿 layers is upper bounded by| 𝜏 𝐿 -𝜏 𝐿 | ≲ 𝑉 2 𝐵 2 𝛿𝜎 7 𝐿 (𝜀 1 + 𝜀 2 ).Combining this bound with Lemma 4 and Lemma 5, we have lim𝑑→∞ 𝔼 𝑥,𝐴 1 𝑑 ∥ 𝑓 ℓ (𝐴 ⊤ 𝑣 ℓ + 𝑥 ℓ , 𝜃 𝑇 ) -𝑥 ∥ 2 ≲ 𝔼[(𝑓 * ℓ (𝑿 + 𝜏 ℓ 𝒁 ) -𝑿 ) 2 ] + 𝑉 2 𝐵 2 𝛿𝜎 7 𝐿+1 (𝜀 1 + 𝜀 2 ) .Applying state evolution (Lemma 1), we get that lim𝑑→∞ 𝔼 𝑥,𝐴 1 𝑑 ∥ 𝑓 ℓ (𝐴 ⊤ 𝑣 ℓ + 𝑥 ℓ , 𝜃 𝑡 ) -𝑥 ∥ 2 ≲ lim 𝑑→∞ 𝔼 𝑥,𝐴 1 𝑑 ∥𝑓 * ℓ (𝐴 ⊤ 𝑣 ℓ + 𝑥 ℓ , 𝜃 𝑡 ) -𝑥 ∥ 2 + 𝑉 2 𝐵 2 𝛿𝜎 7 𝐿+1 (𝜀 1 + 𝜀 2 )as claimed.In Definition 1, we primarily focus on the setting where 𝐶𝛼 ≥ 1. Here we provide some bounds on ℭ 𝜀 (𝜙, 𝛼) in this regime. This in turn provides the bound for ℭ 𝑠 (𝜙, 𝛼). First note that we always haveEQUATIONFirst consider the case where 𝜙 = 𝑖 𝑐 𝑖 𝑧 𝑖 is a degree-𝑘 polynomial. Eq. ( 14) then impliesEQUATION)This bound is useful when the degree is high, e.g. when 𝑘 = ω (𝐶𝛼 log(1/𝜀)). We also have the following naive bounds when the degree is small or intermediate. When 𝑘 ≥ log(1/𝜀),EQUATIONIn the context of our learning result, it suffices for the true denoisers 𝑓 * ℓ 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 7 (Jackson's theorem [Jac11] ). Let 𝑘 ∈ ℤ, 𝐵, 𝑅 ≥ 0. If a function 𝑓 : ℝ → ℝ is 𝐵-Lipschitz, then there exists a polynomial 𝑞 of degree-𝑘 for whichEQUATION)As the true denoisers are Lipschitz and the prior 𝑝 is assumed to be 𝑅-sub-Gaussian so that the effective support over which the true denoisers must be approximated has radius 𝑂 (𝑅), this implies that we can approximate them pointwise to error 𝛿 with polynomials of degree 𝑘 = 𝑂 (𝑅/𝛿). To apply any of the complexity bounds in Eqs. ( 15)-( 17) to these polynomials, it remains to bound the coefficient of largest magnitude. For this, we can apply a result like the following:EQUATIONIn 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 ℤ 2 prior: Example 1. When 𝑝 = ℤ 2 , the optimal denoisers used in Bayes AMP are of the form tanh •/𝜏 2 for various 𝜏 ≥ 𝜎, where 𝜎 2 is the variance of the measurement noise. This function is 𝐵 = 1/𝜎 2 -Lipschitz, and the effective support of 𝑝 ★ N(0, 𝜏 2 ) is of radius 𝑅 = 𝑂 (1). By standard results, tanh •/𝜏 2 can be 𝜂-approximated over [-𝑂 (1), 𝑂 (1)] with a polynomial of degree-𝑘 = 𝑂 (log(𝜎/𝜂)/𝜎 2 ), see e.g. Lemma 2 in [LSSS14] . Applying Lemma 8 to bound the coefficients of this polynomial by (𝜎/𝜂) 𝑂 (1/𝜎 2 ) and then applying the bound in Eq. (16), we conclude that the denoisers in the case of ℤ 2 prior are 𝜂-approximated by polynomials 𝑞 of complexityEQUATIONWhile compressed sensing involves a linear noisy transformation of its underlying signal, rankone matrix estimation offers a nonlinear counterpart. Here, the unknown signal 𝑥 ∈ ℝ 𝑑 is first transformed into a rank-one matrix 𝑥𝑥 ⊤ . The observed signal is given byEQUATIONwhere 𝐺 ∈ ℝ 𝑑 ×𝑑 is a symmetric matrix with entries 𝐺 𝑖 𝑗 ∼ N(0, 1 𝑑 ) for 𝑖 ≤ 𝑗 and 𝜆 > 0 denotes the signal-to-noise (SNR) ratio. Again, we assume 𝑥 ∼ 𝑝 x for product prior 𝑝 x ≜ 𝑝 ⊗𝑑 with 𝑝 some distribution over ℝ. Then, the rank-one matrix estimation problem attempts to recover an estimate 𝑥 for 𝑥 such that the Frobenius mean squared error (MSE) 1 𝑑 𝔼∥ 𝑥 𝑥 ⊤ -𝑥𝑥 ⊤ ∥ 2 𝐹 is minimal. The asymptotic setting corresponds to working with a sequence of such problems indexed by dimension 𝑑, where 𝑑 → ∞.The approximate message passing (AMP) algorithm for estimating 𝑥 given 𝑌 proposed in [RF12] takes the formEQUATIONwhere 𝑓 ℓ : ℝ → ℝ again indicates a scalar denoiser applied componentwise. There are various possible choices for initialization; here we consider 𝑥 0 = 1 ∈ ℝ 𝑑 and 𝑣 0 = 𝑌 𝑥 0 .As with compressed sensing, AMP iterates for rank-one matrix estimation satisfy a state evolution recursion. Define parameters 𝜇 ℓ and 𝜏 ℓ evolving according to the scalar equationsEQUATIONEQUATIONwhere 𝑿 ∼ 𝑝 and 𝒁 ∼ N(0, 1). Then for 𝑑 → ∞, the empirical distribution over entries of 𝑣 ℓ converges asymptotically to the one-dimensional distribution over 𝜇 ℓ 𝑿 + √ 𝜏 ℓ 𝒁 [RF12] . Bayes AMP thus corresponds to the choice of denoising functions optimizing the posterior mean on 𝑿 given 𝜇 ℓ 𝑿 + √ 𝜏 ℓ 𝒁 :EQUATIONwhere 𝜇 * ℓ and 𝜏 * ℓ are obtained by substituting 𝑓 * ℓ 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 𝑝 is the uniform distribution over {1, -1}, the Bayes AMP algorithm is information-theoretically optimal, i.e., it asymptotically matches the MSE 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] . 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 as with compressed sensing, choosing the state evolution parameters 𝜇 ℓ 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 𝜏 ℓ with 1 𝑑 ∥𝑥 ℓ ∥ 2 2 using Eq. ( 25) and 𝜇 ℓ with √︃ | 1 𝑑 ∥𝑣 ℓ ∥ 2 2 -1 𝑑 ∥𝑥 ℓ ∥ 2 2 | using the infinite dimensional distribution of the components of 𝑣 ℓ . The latter holds when the prior 𝑝 has zero mean and unit variance; in general, the estimate must be scaled down by an additional factor of𝔼 𝑿 ∼𝑝 [𝑿 2 ], which can be estimated from data. Unless expressed otherwise, we assume 𝔼 𝑥∼𝑝 [𝑥] = 0 and 𝔼 𝑥∼𝑝 [𝑥 2 ] = 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.Here, we are given training data {(𝑌 𝑖 , 𝑥 𝑖 )} 𝑁 𝑖=1 generated by Eq. ( 21) with all 𝑥 𝑖 ∼ 𝑝 x . Let F denote a family of MLPs with fixed architecture constrained to three-dimensional inputs and a one-dimensional output. Set 𝑥 0 = 1 ∈ ℝ 𝑑 and 𝑣 0 = 𝑌 𝑥 0 . Initializing MLP 𝑓 ℓ ∈ F for all ℓ ∈ [0, 𝐿 -1], we obtain an 𝐿-layer unrolled network that computesEQUATIONat every iteration, where 𝜇 ℓ = √︃1 𝑑 ∥𝑣 ℓ ∥ 2 2 -1 𝑑 ∥𝑥 ℓ ∥ 2 2 and 𝜏 𝑡 = 1 𝑑 ∥ 𝑥 ℓ ∥ 2 2 .Here 𝑓 ℓ ( • ; 𝜇 ℓ , 𝜏 ℓ ) denotes applying the scalar function 𝑓 ℓ ( • ; 𝜇 ℓ , 𝜏 ℓ ) entrywise.Implementation details. We set 𝑑 = 2 10 and fix 𝜆 = 1.5. We focus on the Gaussian and ℤ 2 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 {𝑌 𝑖 , 𝑥 𝑖 } 𝑁 𝑖=1 of size 𝑁 = 2 12 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 of the Bayes-optimal denoisers. LDNet denoisers. While later iterations are able to recover the Bayes-optimal denoisers, it is worth noting the high approximation error at early iterations as shown in Figure 4 . Early iterations 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 is tolerated. In Figure 5 , we set the parameters of 𝐵 = 𝐴 ⊤ to be learnable alongside the layerwise denoisers (see Algorithm 2). Holding 𝛿 = 1 2 fixed while scaling down 𝑚 has the effect of widening the gap between the NMSE performance of Bayes AMP versus LDNet with trainable 𝐵. Another regime to which existing theory for Bayes AMP largely breaks down is when the entries of the sensing matrix 𝐴 are non-Gaussian. While it was previously observed that learning 𝐵 can help for ill-conditioned 𝐴 [BSR17], we find that there are advantages even for well-conditioned (but non-Gaussian) sensing matrices. Figure 6 plots NMSE for two sensing matrices Ã ∈ ℝ 250×500 : one obtained by truncating a random orthogonal matrix 𝑄 ∈ ℝ 500×500 (condition number 1), and the other by truncating a Gram matrix 𝑋 ⊤ 𝑋 ∈ ℝ 500×500 with 𝑋 𝑖 𝑗 ∼ N(0, 1/𝑚) (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 Ã⊤ in Eqs. (2) and (3) with the "𝐵" matrix learned by LDNet. The baselines Mixture of Gaussians provide a quite general class of non-product priors. We consider 𝑝 = 1 2 N(𝜇 1 , Σ 1 ; 𝑥) + 1 2 N(𝜇 2 , Σ 2 ; 𝑥), where we choose 𝜇 1 and 𝜇 2 at random with each coordinate chosen from N(0, 1), ensuring that 𝜇 1 and 𝜇 2 do not define the same direction. For each of the covariance matrices, we begin by a choosing random vector with each coordinate drawn uniformly from [1, 2] and normalize so that vector has norm √ 𝑑. We take this to be the diagonalization (i.e. eigenvalues) of the covariance, and conjugate by a randomly drawn orthogonal matrix. As Figure 7 demonstrates, multi-dimensional LDNet significantly outperforms the Bayes AMP baseline on both priors. On a product ℤ 2 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 a 63.32% improvement. We see that the denoising functions are roughly transferable to 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 𝐴 rather than in expectation over random Gaussian 𝐴, 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 a fundamental "optimal" denoiser that is independent of 𝐴. To further verify this claim, we froze the learned MLP denoiser weights for the Bernoulli-Gaussian prior and replaced the 𝐴 matrix in Eqs. (5) with other randomly sampled Gaussian matrices 𝐴 ′ ∈ ℝ 250×500 . As shown in Figure 8 , this leads to minimal changes in the NMSE profile.For prior 𝑝 and 𝑿 ∼ 𝑝, 𝒁 ∼ N(0, 1), the Bayes optimal denoiser in Bayes AMP is given byEQUATIONfor compressed sensing andEQUATIONfor rank-one matrix estimation. For the priors examined in our experiments, we write out the setting-specific 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 beEQUATION)and parameterize the denoiser viaEQUATIONas done in [BS16, BSR17] .ℤ 2 prior. The compressed sensing optimal denoiser [DM14] can be written asEQUATIONwhich we parameterize asEQUATIONFor rank-one matrix estimation, the optimal denoiser can be similarly computed to beEQUATIONparametrized asEQUATIONGaussian prior. For rank-one matrix estimation, the optimal denoiser for a Gaussian prior isEQUATION)Mixture of Gaussians prior. The calculation of the Bayes-optimal denoiser for a mixture of Gaussians prior in rank-one matrix estimation is slightly more involved, so we provide some more details about the calculation. Given 𝑝 = (39) Thus, we have (41)EQUATION• 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. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 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 out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of 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 or 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 as well as 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 or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?Justification: Our research is theoretical 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 consideration due to laws or regulations in their jurisdiction). 
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:• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a 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. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
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 is theoretical in nature.Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was 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. 
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 to 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 as much detail as 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. 
institutional review board (irb) approvals or equivalent for research with human subjects: 
 Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?Answer: [Yes]Justification: Our research is 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.• 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.In our experiments, we learn the denoiser using all coordinates, but in our theoretical result, we focus on learning using only a single coordinate. The latter is less sample-efficient but more convenient for our proof. It should be possible to prove a guarantee for the later, but it 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.The reason we can do this is that this condition fails to hold only with probability at most 𝜀 2 , and whenever it fails to hold, we can output 0 and pay an additional 𝜀 2 •𝔼[𝑥 2 ]. Alternatively, we could also modify the learner network to implement the indicator of |𝐴 ⊤ 𝑖 𝑣 ℓ + 𝑥 ℓ 𝑖 | ≲ 𝑅(log(1/𝜀 2 )) 3/2 using an appropriate linear combination of ReLU activations without learnable parameters. 
 
 Appendix:  


Decision: 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.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: 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.

Furthermore, the chosen architecture could be considered impractical in practice due to its fixed width and depth.

However, these weaknesses also present opportunities for future research.
questions: 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.

* Could the authors comment on the practical likelihood of Assumption 1?

* 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: Limitations are clearly discussed
paper_summary: paper_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.

The 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.
strength_weakness: strength_weakness: 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.
questions: questions: - L80-81:
> "*and thus achieve mean-squared error which is conjectured to be optimal among all polynomial-time algorithms*"

As 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)

- 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?

- Can the authors comments what "*sufficiently small*" learning rate and "*sufficiently large*" number of steps exactly mean?

- 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}$?

- 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)

- 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: Limitations are discussed in Section 5.
paper_summary: paper_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.
strength_weakness: strength_weakness: - Some aspects regarding the readability and presentation could be improved, especially as they relate to the main theorems. I have noted my questions below. 
- 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?
questions: 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?
- 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.
- Have the authors conducted experiments in the case when the distribution is not a product prior? Do similar results extend to this case?
limitations: 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.

Comments: title: Response to authors
comment: I thank the authors for their rebuttal, which has addressed my questions. I maintain my score and recommendation towards acceptance.
Thank you for your detail responses I will update my score accordingly to the thoughtful respones the authors have provided.
Thank you to the authors for this detailed discussion on the sample complexity. This really helps in terms of better understanding how to apply these bounds and making them more transparent. It's a good idea to include them in the appendix as illustrative examples of the theory.

Based on our discussion, I believe my comments about readability have been addressed. While the smooth product prior distribution still feels like a limitation, the experiments suggest that there are interesting avenues for future work. As such, I will raise my score.
We thank the reviewer again for clarifying the subtle details on how the complexity quantity scales with the denoiser.

In Definition 1 in the submission, the complexity $\mathcal{C}(\phi, \alpha)$ is given by the equation

$\mathcal{C} ( \phi, \alpha ) = \sum_{i=0}^\infty \Big( \frac{ \alpha C \sqrt{ \log (1 / \epsilon) } }{ \sqrt{i} } \Big)^i | c_i |$.

Since we know that the denoisers $\phi = f_\ell^*$ are infinite-order smooth, we can leverage Jackson’s Theorem to get a polynomial approximation of $f_\ell^*$ up to degree $k$, and this degree $k$ is what controls how $\mathcal{C}(\phi, \alpha)$ scales.

When $\alpha = R(\log 1/\epsilon)^{1.5}$, plugging into definition 1, the leading term takes the form $\frac{\sqrt{k} {R^{k}}}{k^{\frac{k}{2}}} \cdot (\log{1/\epsilon})^{2k} || c ||$ where $||  c ||$ denotes $L_2$ norm of $[c_0,  c_1, \ldots, c_k]$ vector.

As you pointed out, the key tradeoff is the approximation error in Jackson’s theorem, which modulates the degree of the polynomial expansion, which dictates how this leading term scales.

Jackson’s theorem states that the error of approximation $f_\ell^*$ with a degree $k$ polynomial is at most a constant times $|\nabla f_\ell^*|/k$. Since we further assume the score function is $B$-Lipschitz, this is at most $B/k$. Thus, if we want an approximation error $\delta$ in the polynomial approximation, we require $k \geq \frac{B}{\delta}$.

Putting this all together, if you admit error $\delta$ in the polynomial approximation of $f_\ell^*$ and error $\epsilon$ in SGD, then the complexity scales polynomially with leading term

$\left(\frac{\delta}{B}\right)^{\frac{B}{2\delta}} \cdot (R(\log{1/\epsilon})^{2})^{B/\delta} || c ||$.

**Example**: For $Z_2$ prior, we know that the denoiser function is given by $\tanh(y / \tau_\ell^2)$. In this case, the denoiser function is Lipschitz function with $\frac{1}{\tau_\ell^2} \leq \frac{1}{\sigma^2}$ constant. Additionally, it is sub-Gaussian with constant sub-Gaussianity constant. Moreover, $|| c || \leq O((R \log 1/\epsilon)^{10B/\delta})$ (see Lemma 23 in [1], the polynomial in Lemma 23 is bounded with constant at most 2 for $\delta < 1$ because it is $\delta$-approximation of $\tanh$ (a bounded function)). Using this bound on $|| c ||$ and $B=1/\sigma^2, R=O(1)$, the complexity of $\delta-$approximation of the denoiser is

$\left( \delta \sigma^2 \right)^{\frac{1}{2\sigma^2 \delta}} \cdot (C \log{1/\epsilon})^{ \frac{12}{\sigma^2\delta} }$

for some large constant $C$.


We will include this explanation in the appendix of our revision for clarity. Thank you again for engaging closely with our work!


[1] Goel et al. Reliably Learning the ReLU in Polynomial Time.
Thank you to the authors for their detailed response and answering my questions. I appreciate the additional experiments that were conducted on examples with non-product distributions. I was hoping to further ask about the sample complexity:
- **Notion of complexity:** Thank you for clarifying that the function $\phi$ indeed should be the optimal denoiser $f_{\ell}^*$. I was hoping to make sure I understand this point further. If the notion of complexity then depends on $f_{\ell}^*$, how should I think about how the quantity $C(f_{\ell}^*, R(\log1/\epsilon)^{1.5})$ scales depending on the distribution $p$ chosen? In particular, you mention in another bullet point that $f_{\ell}^*$ is an infinite-order smooth function under the Theorem's assumptions. In this case, what would the complexity quantity scale like? There is an example in the paper stating that polynomials of degree $\ell$ have complexity $\mathrm{poly}(\alpha^{\ell},\log(1/\epsilon)^{\ell})$ and that for cases where the score is sufficiently smooth, one could apply Jackson's theorem to approximate them. What will be the tradeoff in terms of approximation in sample complexity (e.g., the approximation error + complexity when choosing an $n$-degree polynomial)? Are there examples of distributions $p$ for which we can explicitly calculate or easily bound $C(f_{\ell}^*, R(\log1/\epsilon)^{1.5})$?