Title: Unrolled denoising networks provably learn to perform optimal Bayesian inference
Venue: NeurIPS 2024
Accepted: yes
Keywords: algorithm unrolling

approximate message passing (AMP)

inverse problems

denoising
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

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 ageold 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.

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

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 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 b 𝑥 minimizing 𝔼 ∥ 𝑥 -b 𝑥 ∥ 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 b 𝑓 1 , b 𝑓 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 b 𝑓 ℓ .

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 𝑑 . The central reason behind this is that state evolution allows us to map the training dynamics of the network, which a priori lives in 𝐿 2 ( ℝ 𝑑 ) , to training dynamics over the space of functions 𝐿 2 ( ℝ ) , where the resulting learning problem amounts to that of one-dimensional score estimation . As a 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 a practical alternative that achieves the same performance but does not require explicit knowledge of the true signal prior.

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.

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.

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.

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 𝑥 is a deterministic sparse vector, and the goal is to converge to 𝑥 . 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-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, LZB22] from a different perspective than ours, was an upper bound on the Hessian of the empirical risk of an unrolled ISTA network in a neighborhood around random initialization.

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.

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

2.1 Compressed sensing

In compressed sensing, we are given noisy linear measurements 𝑦 ∈ ℝ 𝑚 obtained from an unknown signal 𝑥 ∈ ℝ 𝑑 via the observation process

where 𝐴 ∈ ℝ 𝑚 × 𝑑 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 b 𝑥 such that the mean squared error (MSE) 1 𝑑 𝔼 ∥ b 𝑥 -𝑥 ∥ 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

where 𝑓 ℓ : ℝ → ℝ 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

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

where 𝑿 ∼ 𝑝 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 ( GFOM s) [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].

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 𝑝 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 b 𝑓 ℓ : ℝ 2 → ℝ chosen from F . Set b 𝑥 0 = 0 𝑑 ∈ ℝ 𝑑 and b 𝑣 0 = 𝑦 𝑖 ∈ ℝ 𝑚 , for a given training input 𝑦 𝑖 . Then for each layer ℓ ∈ [ 0 , 𝐿 -1 ] , our network computes the forward pass

where b 𝜏 ℓ = ∥ b 𝑣 ℓ ∥ 2 / √ 𝑚 and 𝜕 1 denotes differentiation with respect to the first input parameter. The notation b 𝑓 ℓ ( · ; b 𝜏 ℓ ) denotes applying the scalar function b 𝑓 ℓ ( · , b 𝜏 ℓ ) entrywise. We emphasize that b 𝑓 ℓ is tied to 𝜕 1 b 𝑓 ℓ ; that is, we are taking the derivative of the MLP to compute the Onsager term. We refer to our unrolled architecture as an LDNet ( L earned D enoising Net work).

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 function L 𝐶𝑆 = 1 𝑁 ˝ 𝑁 𝑖 = 1 1 𝑑 ∥ b 𝑥 𝑖 𝐿 -𝑥 𝑖 ∥ 2 2 for compressed sensing or L 𝑀𝐸 = 1 𝑁 ˝ 𝑁 𝑖 = 1 1 𝑑 ∥ b 𝑥 𝑖 𝐿 b 𝑥 𝑖 ⊤ 𝐿 -𝑥 𝑖 𝑥 𝑖 ⊤ ∥ 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 , where we iteratively train the -th denoiser on the mean needed for our theory results.

layerwise training ℓ b 𝑓 ℓ 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 b 𝑓 ℓ 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

Algorithm 1: Layerwise Training

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.

3 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:

𝑝 is 𝑅 -sub-Gaussian with 𝔼 𝑿 ∼ 𝑝 [ 𝑿 ] = 0 .

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.

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 b 𝑓 by b 𝑥 ℓ and b 𝑣 ℓ . The Onsager term in AMP ensures that for any iteration ℓ , the distribution of 𝐴 ⊤ b 𝑣 ℓ + b 𝑥 ℓ 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:

where 𝐴 ⊤ 𝑗 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

Given a class of scalar functions F , we define the complexity of F 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.

Main result. Wecan 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

such that the following holds.

Let 𝐿 be any positive integer. Consider an LDNet of depth 𝐿 with MLP denoisers b 𝑓 ℓ 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 ( b 𝑥 𝐿 , b 𝑣 𝐿 ) and denoiser b 𝑓 𝐿 at the output layer of the LDNet (see (5) ) satisfy

where MSE AMP ( 𝐿 ) ≜ 𝔼 𝑥,𝐴 GLYPH<2> 1 𝑑 ∥ 𝑓 ∗ 𝐿 ( 𝐴 ⊤ 𝑣 𝐿 + 𝑥 𝐿 ; 𝜏 𝐿 ) -𝑥 ∥ 2 GLYPH<3> is 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 𝑑 .

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

4.1 Compressed sensing

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) ∥ b 𝑥 -𝑥 ∥ 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

Guided Denoisers

0.65

Guided Denoisers

10

Bayes AMP

Bayes AMP

0.60

NMSE (non-dB)

NMSE (dB)

-20

0.55

-30

0.50

-40

0.45

2

4

6

8

10

12

14

2

4

6

8

10

12

14

Layers

Layers

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, 10log 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%).

LDNet denoisers. 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 𝐴 .

1.5

MLP Denoiser 0

2.0

MLP Denoiser 3

2.0

MLP Denoiser 6

2.0

MLP Denoiser 14

Guided Denoiser 0

Guided Denoiser 3

Guided Denoiser 6

Guided Denoiser 14

1.0

Bayes Denoiser 0

1.5

1.5

Bayes Denoiser 6

1.5

Bayes Denoiser 14

1.0

1.0

1.0

0.5

0.5

0.5

0.0

0.0

0.0

0.0

0.5

0.5

0.5

0.5

1.0

1.0

-1.0

1.0

1.5

1.5

1.5

1.5

2.0

2.0

2.0

2.01.51.00.50.0

0.5

1.0

1.5

2.0

2.0-1.5 1.00.50.0

0.5

1.0

1.5

2.0

2.01.5 1.00.50.0

0.5

1.0

1.5

2.0

2.01.5 1.00.50.0

0.5

1.0

1.5

2.0

80

MLP Denoiser 0

1.00

MLP Denoiser 3

1.00

MLP Denoiser 7

1.00

MLP Denoiser 14

Guided Denoiser 0

Guided Denoiser 3

Guided Denoiser 7

Guided Denoiser 14

9'0

Bayes Denoiser 0

0.75

0.75

Bayes Denoiser 7

0.75

Bayes Denoiser 14

0.4

0.50

0.50

0.50

0.2

0.25

0.25

0.25

20.0

0.00

00°0

00°0

0.2

0.25

0.25

0.25

0.4

0.50

0.50

0S'0-

0.6

0.75

0.75

0.75

0.8

1.00

1.00

1.00

2.0 1.5 1.0 0.5 0.00.51.0

1.5

2.0

2.0 1.5 1.0 0.50.00.51.0

1.5

2.0

2.0 1.5 1.0 0.5 0.00.5

1.0

1.5

2.0

2.0-1.5 1.0 0.50.00.51.01.5

2.0

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 [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.

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 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 optimally e.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.

Acknowledgments

AK and SC thank Demba Ba for illuminating discussions about unrolling at an early stage of this project. KS and SC thank Vasilis Kontonis for many fruitful conversations about provable score estimation. AK is 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-Zhu, 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 arXiv: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 , 35:36733-36747, 2022.

[BKM + 19] Jean Barbier, Florent Krzakala, Nicolas Macris, Léo Miolane, and Lenka Zdeborová. Optimal errors and phase transitions in high-dimensional generalized linear models. Proceedings of the National Academy of Sciences , 116(12):5451-5460, 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-separable functions. Information and Inference: A Journal of the IMA , 9(1):33-79, 2020.

[BPW18] Afonso S Bandeira, Amelia Perry, and Alexander S Wein. Notes on computationalto-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 a statistical learning perspective. In Compressed Sensing in Information Processing , pages 247-277. Springer, 2022.

[BS16] Mark Borgerding and Philip Schniter. Onsager-corrected deep learning for sparse linear inverse problems. In 2016 IEEE Global Conference on Signal and Information Processing ( GlobalSIP) , pages 227-231. IEEE, 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-4308, 2017.

[BT09] Amir Beck and Marc Teboulle. A fast iterative shrinkage-thresholding algorithm for linear inverse problems. SIAM journal on imaging sciences , 2(1):183-202, 2009.

[CKS24] Sitan Chen, Vasilis Kontonis, and Kulin Shah. Learning general gaussian mixtures with efficient score matching. arXiv preprint arXiv:2404.18893 , 2024.

[CLWY18] 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.

[CLWY21] Xiaohan Chen, Jialin Liu, Zhangyang Wang, and Wotao Yin. Hyperparameter tuning is all you need for lista. Advances in Neural Information Processing Systems , 34:1167811689, 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 (ISIT) , pages 185-189, 2016.

[DDDM04] Ingrid Daubechies, Michel Defrise, and Christine De Mol. An iterative thresholding algorithm for linear inverse problems with a sparsity constraint. Communications on Pure and Applied Mathematics: A Journal Issued by the Courant Institute of Mathematical Sciences , 57(11):1413-1457, 2004.

[DM14] Yash Deshpande and Andrea Montanari. Information-theoretically optimal sparse pca. In 2014 IEEE International Symposium on Information Theory , pages 2197-2201. IEEE, 2014.

[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-18919, 2009.

[GKL24] 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 on 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 on Theory of Computing , pages 348-357, 2024.

[ITW19] Daisuke Ito, Satoshi Takabe, and Tadashi Wadayama. Trainable ista for sparse signal recovery. IEEE Transactions on Signal Processing , 67(12):3113-3125, 2019.

[Jac11] Dunham Jackson. Über die Genauigkeit der Annäherung stetiger Funktionen durch ganze rationale Funktionen gegebenen Grades und trigonometrische Summen gegebener Ordnung . Dieterich'schen Universität-Buchdruckerei, 1911.

[LC19] Jialin Liu and Xiaohan Chen. Alista: Analytic weights are as good as learned weights in lista. In 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 , 35: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.

[LZB20] Chaoyue Liu, Libin Zhu, and Misha Belkin. On the linearity of large non-linear models: when and why the tangent kernel is constant. Advances in Neural Information Processing Systems , 33:15954-15964, 2020.

[LZB22] Chaoyue Liu,Libin Zhu,and Mikhail Belkin. Loss landscapes and optimization in overparameterized 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.00285 , 2016.

[Mei24] Song Mei. U-nets as belief propagation: Efficient classification, denoising, and diffusion in generative hierarchical models. arXiv preprint arXiv:2404.18444 , 2024.

[MJC21] Osman Musa, Peter Jung, and Giuseppe Caire. Plug-and-play learned gaussianmixture approximate message passing. In ICASSP 2021-2021 IEEE International Conference on 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 on 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 arXiv:2201.05101 , 2022.

[MW23] Song Mei and Yuchen Wu. Deep networks as denoising algorithms: Sample-efficient learning of diffusion models in high-dimensional graphical models. arXiv preprint arXiv:2309.11420 , 2023.

[NT10] Deanna Needell and Joel A. Tropp. Cosamp: iterative signal recovery from incomplete and inaccurate samples. Communications of the ACM , 53(12):93-100, 2010.

[RF12] Sundeep Rangan and Alyson K Fletcher. Iterative estimation of constrained rank-one matrices in noise. In 2012 IEEE International Symposium on Information Theory Proceedings , pages 1246-1250. 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 arXiv:2304.09802 , 2023.

[ 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-19649, 2023.

[ She12] Alexander A Sherstov. Making polynomials robust to noise. In Proceedings of the forty-fourth annual ACM symposium on Theory of computing , pages 747-758, 2012.

[ SPP + 23] Shaik Basheeruddin Shah, Pradyumna Pradhan, Wei Pu, Ramunaidu Randhi, Miguel RD Rodrigues, and Yonina C Eldar. Optimization guarantees of unfolded ista and admm networks with smooth soft-thresholding. arXiv preprint arXiv:2309.06195 , 2023.

[ ST99] Roel Snieder and Jeannot Trampert. Inverse problems in 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 J Thouless, Philip W Anderson, and Robert G Palmer. Solution of'solvable 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 in 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.

A Further related work

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.

Comparison to LDAMP [MMB17]. 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.

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 𝑥 ℓ 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 surely

where 𝑿 ∼ 𝑝 𝑥 and 𝒁 ∼ N ( 0 , 1 ) . The state evolution parameters 𝜏 0 , 𝜏 1 , . . . are recursively defined as follows

Observe that to minimize the variance 𝜏 2 ℓ at each iteration, the optimal choice of denoiser 𝑓 ℓ is 𝑓 ∗ ℓ ( 𝑥 ) = 𝔼 [ 𝑿 | 𝑿 + 𝜏 ℓ 𝒁 = 𝑥 ] .

B.2 Learning guarantee for a single layer of unrolling

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 consider

We 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 b 𝑥 ℓ and b 𝑣 ℓ .

Lemma 2 (Learning the denoiser within 𝐿 2 error) . For every ℓ , assume every coordinate of b 𝑣 ℓ and b 𝑥 ℓ 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

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

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,

This is a consequence of the following result on training neural networks in the NTK regime:

Lemma 3 (Theorem 1 of [AZLL19] 4 ) . Consider a target function 𝐹 ∗ : ℝ 𝑑 → ℝ of the following form

where each 𝜙 : ℝ → ℝ is infinite-order smooth and weights satisfy ∥ 𝑤 ∗ 𝑖, 1 ∥ , ∥ 𝑤 ∗ 𝑖, 2 ∥ ≤ 1 and | 𝑎 ∗ 𝑖 | ≤ 1 . Additionally, assume that ∥ 𝑥 ∥ ≤ 𝐵 . Then, for every 𝜀 ∈ ( 0 , 1 /( 𝑝 ℭ 𝜀 ( 𝜙, 𝐵 ))) , there exists 𝑀 0 = poly ( ℭ 𝜀 ( 𝜙, 𝐵 ) , 1 / 𝜀 ) and 𝑁 0 = poly ( ℭ 𝜀 ( 𝜙, 𝐵 ) , 1 / 𝜀 ) such that for every 𝑚 ≥ 𝑀 0 and 𝑛 ≥ 𝑁 0 , choosing learning rate 𝜂 = ˜ Θ ( 1 /( 𝜀𝑚 )) and running gradient descent from random initialization for 𝑇 = ˜ Θ ( ℭ 𝜀 ( 𝜙, 𝐵 ) 2 𝑝 2 / 𝜀 2 ) , with high probability,

Additionally, the absolute value of the neural network is bounded by | 𝑁 ( 𝑥 ; 𝜃 𝑡 )| ≲ ˜ Θ ( ℭ 𝜀 ( 𝜙, 𝐵 )) for all 𝑥 with ∥ 𝑥 ∥ ≤ 𝐵 for all 𝑡 .

Proof of Lemma 2. As the distribution over b 𝑣 ℓ is 𝑅 -sub-Gaussian, for a fixed 𝐴 , we have | 𝐴 ⊤ 𝑖 b 𝑣 ℓ | ≤ 𝑅 ∥ 𝐴 𝑖 ∥ √︁ log ( 1 / 𝜀 2 ) with probability at least 1 -𝜀 2 . Additionally, because of 𝐴 𝑖 𝑗 ∼ N ( 0 , Id / 𝑚 ) , we have ∥ 𝐴 𝑖 ∥ ≲ log ( 1 / 𝜀 2 ) with probability at least 1 -𝜀 2 . Combining both bounds, we have | 𝐴 ⊤ 𝑖 b 𝑣 ℓ | ≲ 𝑅 ( log ( 1 / 𝜀 2 )) 3 / 2 . Similarly, using sub-Gaussianity of b 𝑥 ℓ , we have | b 𝑥 ( 𝑖 ) ℓ | ≲ 𝑅 √︁ log ( 1 / 𝜀 2 ) . This gives us that | 𝐴 ⊤ 𝑖 b 𝑣 ℓ + b 𝑥 ( 𝑖 ) ℓ | ≲ 𝑅 ( log ( 1 / 𝜀 2 )) 3 / 2 with probability at least 1 -𝜀 2 .

By only considering samples satisfying | 𝐴 ⊤ 𝑖 b 𝑣 ℓ + b 𝑥 ( 𝑖 ) ℓ | ≲ 𝑅 ( log ( 1 / 𝜀 2 )) 3 / 2 , 5 we can apply Lemma 3 to obtain a neural network that achieves 𝜀 1 error.

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 b 𝑓 ℓ (· , 𝜃 𝑡 ) be the learned neural network using gradient descent after 𝑡 timesteps such that the conditions of Lemma 2 satisfies. Then, with high probability, we have

Additionally, the following statement holds with high probability:

Proof. Observe that 𝐴 ⊤ 𝑗 b 𝑣 ℓ + b 𝑥 ( 𝑗 ) ℓ follows the same distribution for all 𝑗 . Therefore, using Lemma 2, we obtain that

As the distribution of 𝐴 ⊤ 𝑗 b 𝑣 ℓ + b 𝑥 ( 𝑗 ) ℓ is the same for all 𝑗 ∈ [ 𝑑 ] , we use the same learned denoiser for all the coordinates. Therefore, we can rewrite the above equation as

Now, we want to use state evolution as 𝑑 →∞ . Note that as 𝑑 →∞ , using Lemma 1 with 𝜓 function as 𝜓 ( 𝑎, 𝑏 ) = ( 𝑎 -𝑏 ) 2 , we have

for any function 𝑔 ∈ F . As the quantity inside expectation converges to 𝔼 [( 𝑔 ( 𝑿 + b 𝜏 ℓ 𝒁 ) -𝑿 ) 2 ] , using monotone convergence theorem, we have

Using this result and applying limits on both sides of Lemma 2, we obtain

The minimizer of 𝔼 [( 𝑔 ( 𝑿 + b 𝜏 ℓ 𝒁 ) -𝑿 ) 2 ] is given by b 𝑓 ∗ ℓ (·) = 𝔼 [ 𝑿 | 𝑿 + b 𝜏 ℓ 𝒁 = ·] . 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 𝔼 [( b 𝑓 ℓ ( 𝑿 + b 𝜏 ℓ 𝒁 , 𝜃 𝑇 ) -𝑿 ) 2 ] using the monotone convergence theorem. This gives the second result of the lemma statement.

B.3 Stability of optimal denoisers

The right-hand side of the bound in the above Lemma corresponds to the minimum mean squared error achievable for denoising at noise scale b 𝜏 ℓ , where b 𝜏 ℓ 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 b 𝜏 ℓ 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.

Lemma 5. Let prior 𝑿 be such that the score function 𝜕 1 𝑝 (· ; 𝜏 ) is 𝐵 -Lipschitz continuous for all 𝜏 where 𝑝 (· ; 𝜏 ) denotes the probability density function of random variable 𝑿 + 𝜏 𝒁 . Additionally, assume that the variance of 𝑿 is bounded by 𝑉 . Then, we have

where b 𝑓 ∗ ℓ ( 𝑥 ) = 𝔼 [ 𝑿 | 𝑿 + b 𝜏 ℓ 𝒁 = 𝑥 ] and 𝑓 ∗ ℓ ( 𝑥 ) = 𝔼 [ 𝑿 | 𝑿 + 𝜏 ℓ 𝒁 = 𝑥 ] .

Proof. Rewriting the error between 𝑓 ∗ ℓ and b 𝑓 ∗ ℓ , we have

The term in Eq. (10) can be upper bounded by ( 𝐵 ( b 𝜏 ℓ -𝜏 ℓ )) 2 / b 𝜏 4 ℓ because b 𝑓 ℓ is Lipschitz by assumption. Using Cauchy-Schwartz for the term in Eq. (11), the second term is upper bounded by

The squared term in Eq. (12) can be upper bounded by

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 ≥ b 𝜏 ℓ ≥ 𝜏 ℓ ≥ 𝜎 2 , we obtain that

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 b 𝜏 ℓ is given by

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

Using Tweedie's formula ( 𝜕 1 log 𝑝 ( 𝑥 ; 𝜏 ℓ ) = ( 𝔼 [ 𝑿 | 𝑿 + 𝜏 ℓ 𝒁 = 𝑥 ] -𝑥 )/ 𝜏 2 ℓ ), we have

𝔼 𝑥 ∼ 𝑝 (· 𝜏 ℓ ) GLYPH<2> ( 𝜕 log 𝑝 ( 𝑥 ; 𝜏 ℓ ) -𝜕 log 𝑝 ( 𝑥 ; 𝜏 ℓ )) 2

This implies that

where the last inequality uses the fact that b 𝜏 ℓ ≥ 𝜏 ℓ 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.

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 𝜏 ℓ is defined using the optimal denoiser 𝑓 ∗ ℓ at layer ℓ in Lemma 1. Similarly, the state evolution recursion for b 𝜏 ℓ is applied using learned denoiser b 𝑓 ℓ ( · , 𝜃 𝑇 ) . At some places, we will use b 𝑓 ℓ ( · ) to denote b 𝑓 ℓ ( · , 𝜃 𝑇 ) for brevity. By the state evolution recursion for 𝜏 ℓ , the error between 𝜏 ℓ and b 𝜏 ℓ is given

where the last inequality follows from Lemma 4 and Lemma 5. As b 𝜏 ℓ + 𝜏 ℓ ≥ 2 𝜎 for all ℓ and b 𝜏 ℓ ≥ 𝜏 ℓ , we have

Solving this recurrence, we have that the error in the state evaluation parameters after 𝐿 layers is upper bounded by

Combining this bound with Lemma 4 and Lemma 5, we have

Applying state evolution (Lemma 1), we get that

as claimed.

B.5 Bounding the complexity of denoisers

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 have

First consider the case where 𝜙 = ˝ 𝑖 𝑐 𝑖 𝑧 𝑖 is a degree𝑘 polynomial. Eq. (14) then implies

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 / 𝜀 ) ,

When 𝑘 ≤ log ( 1 / 𝜀 ) ,

In 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 which

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:

Lemma 8 (Corollary of Lemma 4.1 from [She12]) . Given a polynomial 𝑞 ( 𝑧 ) = ˝ 𝑘 𝑖 = 0 𝑐 𝑖 𝑧 𝑖 for which sup 𝑧 ∈[ 0 , 1 ] | 𝑞 ( 𝑧 )| ≤ 𝑀 ,

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 ℤ 2 prior:

Example 1. When 𝑝 = ℤ 2 , the optimal denoisers used in Bayes AMP are of the form tanh GLYPH<0> ·/ 𝜏 2 GLYPH<1> 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 GLYPH<0> ·/ 𝜏 2 GLYPH<1> 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 complexity

C Rank-one matrix estimation

C.1 Preliminaries

While 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 by

where 𝐺 ∈ ℝ 𝑑 × 𝑑 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 b 𝑥 for 𝑥 such that the Frobenius mean squared error (MSE) 1 𝑑 𝔼 ∥ b 𝑥 b 𝑥 ⊤ -𝑥𝑥 ⊤ ∥ 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 form

where 𝑓 ℓ : ℝ → ℝ again indicates a scalar denoiser applied componentwise. There are various possible choices for initialization; here we consider 𝑥 0 = b 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 equations

where 𝑿 ∼ 𝑝 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 𝜇 ℓ 𝑿 + √ 𝜏 ℓ 𝒁 :

where 𝜇 ∗ ℓ 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.

C.2 Unrolled architecture

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 = b 1 ∈ ℝ 𝑑 and 𝑣 0 = 𝑌𝑥 0 . Initializing MLP b 𝑓 ℓ ∈ F for all ℓ ∈ [ 0 , 𝐿 -1 ] , we obtain an 𝐿 -layer unrolled network that computes

at every iteration, where b 𝜇 ℓ = √︃ GLYPH<12> GLYPH<12> 1 𝑑 ∥ 𝑣 ℓ ∥ 2 2 -1 𝑑 ∥ 𝑥 ℓ ∥ 2 2 GLYPH<12> GLYPH<12> and b 𝜏 𝑡 = 1 𝑑 ∥ b 𝑥 ℓ ∥ 2 2 . Here b 𝑓 ℓ ( · ; b 𝜇 ℓ , b 𝜏 ℓ ) denotes applying the scalar function b 𝑓 ℓ ( · ; b 𝜇 ℓ , b 𝜏 ℓ ) entrywise.

C.3 Results

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.

1.00

Guided Denoisers

Guided Denoisers

0.95

Bayes AMP

0.9

Bayes AMP

0.90

NMSE

0.8

0.85

0.80

0.7

0.75

0.6

0.70

0.5

0

5

10

15

20

25

30

0

5

10

15

20

25

30

Layers

Layers

Gaussian prior. For a Gaussian prior, each component 𝑥 𝑖 of 𝑥 is drawn 𝑥 𝑖 ∼ N ( 0 , 1 ) , so 𝑝 x = 𝑝 ⊗ 𝑑 for 𝑝 = N ( 0 , 1 ) . As expected, LDNet tracks Bayes AMP to convergence at an NMSE of 0.6931 with a slightly quicker convergence.

ℤ 2 prior. As with compressed sensing, each component of 𝑥 is drawn from {-1 , 1 } with probability 1 2 . Again, LDNet slightly outperforms Bayes AMP until convergence at an (informationtheoretically optimal [DM14]) NMSE of 0.5243 .

1.00

MLP Denoiser 0

1.5

0.75

Guided Denoiser 0

Bayes Denoiser 0

1.0

Bayes Denoiser 10

Guided Denoiser 10

1.0

Bayes Denoiser 20

Guided Denoiser 20

1.0

Bayes Denoiser 30

Guided Denoiser 30

05'0

0.25

0.5

0.5

0.5

0.00

0.0

0.0

0.0

0.25

05'0=

0.5

0.5

0.5

0.75

1.0

1.0

1.0

1.00

2.0

1.51.0

0.5

0.0

0.5

1.0

1.5

2.0

2.01.5 1.00.50.0

0.5

1.0

1.5

2.0

2.0

1.5

1.00.5

0.0

0.5

1.0

1.5

2.0

2.0 1.5 1.00.50.0

0.5

1.0

1.5

2.0

1.00

MLP Denoiser 0

1.0

1.00

1.00

Guided Denoiser 0

Guided Denoiser 10

Guided Denoiser 20

Guided Denoiser 30

0.75

Bayes Denoiser 0

Bayes Denoiser 10

0.75

Bayes Denoiser 20

0.75

Bayes Denoiser 30

0.50

0.5

0.50

0.50

0.25

0.25

0.25

20.00

0°0

00°0

0.00

0.25

0.25

0.25

+ 05'0-

0.5

0.50

0.50

0.75

0.75

0.75

1.00

1.0 +

1.00

1.00

2.0

1.5 1.00.50.0

0.5

1.0

1.5

2.0

2.0

1.51.00.50.00.51.0

1.5

2.0

2.0

0°0S0-0-S'-

0.5

1.0

1.5

2.0

2.0

1.5 1.0 0.5 0.00.5

1.01.5

2.0

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.

D Beyond Bayes AMP performance

D.1 Learning auxiliary parameters

m =200

m = 150

m = 100

-5

X

Learned B

-5

X

LearnedB

-5

x

Learned B

Bayes AMP

Bayes AMP

Bayes AMP

-10

10

-10

15

15

15

(aP)

-20

20

-20

NMSE

-25

25

25

-30

O-

-30

35

35

-35

-40

40

-40

45

45

45

2

4

6

8

10

12

14

2

4

6

8

10

12

14

2

4

6

8

10

12

14

Layers

Layers

Layers

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 𝐵 . Over five randomly drawn measurement matrices per dimension regime, we find that, on average, LDNet outperforms Bayes AMP by 7 . 2750% when 𝑚 = 200, 16 . 4364% when 𝑚 = 150, and 37 . 0605% when 𝑚 = 100. Here, percentage is measured by | Bayes -LDNet | | Bayes | × 100% in NMSE (dB).

TruncatedRandomOrthogonal

TruncatedRandomGram

2.5

x

Learned B

Bayes AMP

-5

5.0

ISTA

COSAMP

-10

7.5

15

(dB)

10.0

20

NMSE

12.5

-25

15.0

-30

LearnedB

17.5

-35

Adjusted BayesAMP

Adjusted ISTA

20.0

-40

COSAMP

2

4

6

8

10

12

14

2

4

6

8

10

12

14

Layers

Layers

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

considered all drastically underperform our unrolled network. In fact, the underperformance of 'adjusted' Bayes AMP demonstrates that the LDNet denoisers are strongly coupled with 𝐵 , 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, non-Gaussian 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 in the regime of product priors. But what happens when our underlying signal is drawn from a non-product distribution?

The modifications to AMP are minimal, as detailed in [BMN20], amounting to 𝑑 -dimensional denoisers 𝑓 ℓ : ℝ 𝑑 → ℝ 𝑑 and replacing the average derivative of the scalar denoiser in the Onsager term with the normalized divergence 1 𝑑 div 𝑓 ℓ in Eqs. (3) and (23). In this non-separable setting, AMP still satisfies a one-dimensional state evolution recursion [BMN20]. In fact, in the asymptotic limit, in some sense our theoretical guarantees carry over if for generic 𝑑 -dimensional priors , minimizing the score-matching objective via gradient descent (now in 𝑑 dimensions instead of 1 dimension) can learn the Bayes-optimal denoiser with gradient descent. This is a question of immense interest within the theory and practice of diffusion generative modeling and remains an important open direction in this area.

In practice, for compressed sensing, unrolled AMP has been shown to be performant on image datasets [MMB17], which serve as 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 handcrafted priors, where we can plot a baseline achieved by Bayes AMP.

LDNet for non-product priors. We work in the low-dimensional regime 𝑑 = 10, where 𝑥 ∈ ℝ 𝑑 . LDNet requires small modifications to the layer iterations defined by Eqs. (27) and (28). The family of MLPs F parametrizing the denoisers have three hidden layers with 1000 neurons and GELU activations, with input dimension 𝑑 + 2 and output dimension 𝑑 . We take b 𝑓 ℓ ∈ F for ℓ ∈ [ 0 , 𝐿 -1 ] , and we replace ⟨ 𝜕 1 b 𝑓 ℓ ⟩ in Eq. (28) with 1 𝑑 div b 𝑓 ℓ . To avoid backpropagating through the Jacobian during training, we omit the finetuning step in Algorithm 1.

Signal distributions. To analyze the performance of 𝑑 -dimensional LDNet, we consider two priors on 𝑥 : product ℤ 2 and a mixture of Gaussians . In both instances we work with a dataset of size 𝑁 = 2 12 generated by sampling a signal and using Eq. (21). The product ℤ 2 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 𝑑 -dimensional as opposed to one-dimensional.

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.

Mixture of Gaussians

0.95

3.5

Bayes AMP

Bayes AMP

0.90

3.0

0.85

2.5

NMSE

0.80

2.0

0.75

1.5

0.70

1.0

0.65

0.5

2

4

6

8

10

12

2

4

6

8

10

12

Layers

Layers

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.

D.3 Learned denoiser dependence on sensing matrix

-5

10

Original A

15

-20

(8P)

NMSE

25

-30

35

40

45

2

4

6

8

10

12

14

Layers

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.

E Explicit expressions for various Bayes-optimal denoisers

For prior 𝑝 and 𝑿 ∼ 𝑝 , 𝒁 ∼ N ( 0 , 1 ) , the Bayes optimal denoiser in Bayes AMP is given by

for compressed sensing and for 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 be

and parameterize the denoiser via

as done in [BS16, BSR17].

ℤ 2 prior. The compressed sensing optimal denoiser [DM14] can be written as

which we parameterize as

For rank-one matrix estimation, the optimal denoiser can be similarly computed to be

Gaussian prior. For rank-one matrix estimation, the optimal denoiser for a Gaussian prior is

parametrized as

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 𝑝 = 1 𝑘 ˝ 𝑘 𝑖 = 1 N ( 𝜇 𝑖 , Σ 𝑖 ; 𝑥 ) , where 𝜇 𝑖 ∈ ℝ 𝑑 and Σ 𝑖 ∈ ℝ 𝑑 × 𝑑 are invertible positive semidefinite symmetric covariance matrices, convolution with √ 𝜏 ℓ 𝒁 ∼

parametrized as

N ( 0 , 𝜏 ℓ 𝐼 𝑑 ) results in the mixture of Gaussians distribution ˜ 𝑝 = 1 𝑘 ˝ 𝑘 𝑖 = 1 N ( 𝜇 ℓ 𝜇 𝑖 , 𝜇 2 ℓ Σ 𝑖 + 𝜏 𝑡 𝐼 𝑑 ; 𝑥 ) . For any mixture ˝ 𝑘 𝑖 = 1 𝜆 𝑖 N ( 𝜇 𝑖 , 𝑄 𝑖 ; 𝑥 ) , the score is given by [CKS24]

Thus, we have

so by Tweedie's formula, the posterior mean on 𝑿 given 𝜇 ℓ 𝑿 + √ 𝜏 ℓ 𝒁 = 𝑦 is

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Justification: We provide complete proofs of our theoretical claim and experimental figures for the experiments.

Answer: [Yes]

Guidelines:

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

2. Limitations

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

Justification: The limitations are discussed in the last section of the paper.

Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes]

Guidelines:

The authors are encouraged to create a separate "Limitations" section in their paper.

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 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 well-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 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 a 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 reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

While the authors might fear that complete honesty about limitations might be used by reviewers as 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 to not penalize honesty concerning limitations.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

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?

Answer: [Yes]

Guidelines:

Justification: We provide complete set of assumption in Assumption ?? 1.

The answer NA means that the paper does not include theoretical results.

All assumptions should be clearly stated or referenced in the statement of any theorems.

All the theorems, formulas, and proofs in the paper should be numbered and cross-referenced.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Theorems and Lemmas that the proof relies upon should be properly referenced.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

4. Experimental Result Reproducibility

Answer: [Yes]

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 of whether the code and data are provided or not)?

Justification: We provide detailed information about the setup of our experiments.

The answer NA means that the paper does not include experiments.

Guidelines:

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

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

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or 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. In 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 to reproducing or verifying the results.

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, as described in supplemental material?

Justification: We currently don't provide the code for the open access but we are planning to do it soon.

Answer: [No]

Guidelines:

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

The answer NA means that paper does not include experiments requiring code.

While we encourage the release of code and data, we understand that this might not be possible, so '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 authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

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://nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

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.

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Justification: We provide detailed information about the training and test details.

Answer: [Yes]

Guidelines:

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The answer NA means that the paper does not include experiments.

The full details can be provided either with the code, in appendix, or as supplemental material.

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

7. Experiment Statistical Significance

Justification: The goal of our analysis is to validate theoretical findings on the synthetic data.

Answer: [No]

Guidelines:

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.

The answer NA means that the paper does not include experiments.

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 assumptions made should be given (e.g., Normally distributed errors).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

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

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.

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.

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?

8. Experiments Compute Resources

Justification: The estimated amount of training time for our final experiments is around 100 CPU hours.

Answer: [Yes]

Guidelines:

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The answer NA means that the paper does not include experiments.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

9. Code Of Ethics

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

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.

Answer: [NA]

Guidelines:

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The answer NA means that the authors have not reviewed the NeurIPS 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).

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

10. Broader Impacts

Justification: Our research is theoretical and does not have any direct societal impact.

Answer: [NA]

Guidelines:

The answer NA means that there is no societal impact of the work performed.

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 impact specific groups), privacy considerations, and security considerations.

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.

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, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

If 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 ML).

The authors should consider possible harms that could arise when the technology is being used as 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.

11. Safeguards

Answer: [NA]

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Justification: Our research is theoretical and does not have any safety threats.

The answer NA means that the paper poses no such risks.

Guidelines:

Released models that have a 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.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Justification: Most of our experiments are validating theoretical foundings and on the synthetic data.

Answer: [NA]

Guidelines:

The authors should cite the original paper that produced the code package or dataset.

The answer NA means that the paper does not use existing assets.

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.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

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

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

13. New Assets

Justification: Our research is theoretical in nature.

Answer: [NA]

Guidelines:

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 answer NA means that the paper does not release new assets.

The paper should discuss whether and how consent was obtained from people whose asset is used.

14. Crowdsourcing and Research with Human Subjects

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

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

Justification: Our research is theoretical in nature.

Answer: [Yes]

Guidelines:

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.

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

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

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?

Justification: Our research is theoretical and does not involve any human subjects.

Answer: [Yes]

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.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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?\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\n- 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$.\n\nLIMITATIONS:\nLimitations are discussed in Section 5.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\n* 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.\n\nLIMITATIONS:\nLimitations are clearly discussed", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-16"}