Unrolled denoising networks provably learn to perform optimal Bayesian inference

We thank the reviewer for their thoughtful questions and remarks.

• On the smoothness and product distribution assumption: Please see the global response for the discussion about the assumptions on the prior.

• On the notion of the complexity of denoisers: We think that this confusion is because of a typo/underspecification in the statement of theorem 2 and theorem 3. The notion of complexity is applied to the underlying optimal denoiser $f_\ell^*$ (i.e. the function $\phi$). In theorem 2/theorem 3, $M_0$ and $N_0$ are polynomial in the sum of the complexities of $f_{\ell}^*$ for $\ell$ from $0$ to $L-1$. Moreover, $M_0$ and $N_0$ are polynomial in $\sum_{\ell \in [0, L-1]} C(f_{\ell}^*, R(\log \frac{1}{\epsilon})^{1.5})$ and $\frac{1}{\epsilon}$. We apologize for the confusion caused by the typo and will update it in the revised draft of the paper.

• On applying function classes and Lemma 3:

  • When we refer to the function class $\mathcal{F}$, we are indeed referring to target functions $F^*(x)$ with bounded coefficients. Moreover, to obtain the learning result for the $\ell^{th}$ layer denoiser, we apply Lemma 3 representing $f_{\ell}^*(x)$ as a target function $F^*(x)$.
  • Note that this is equivalent to setting $p=1$ and $a_1=1$, along with $\phi_1 = f_\ell^*$ (where $a_1^*=1$ and $w_{1, 1}^*=\langle 1, 0\rangle$, $w_{1, 2}^*=\langle 0, 1\rangle$; note we are treating "$x$" as the two dimensional input $\langle x, 1 \rangle$ into the target function).
  • Recall that $f_{\ell}^*$ is the optimal denoiser at noise level $\tau_\ell^*$ (i.e., $f_{\ell}^*(y) = E[ \mathbf{x} | \mathbf{x} + \tau_\ell^* \mathbf{z} = y ]$); therefore, it is an infinite order smooth function. To get an intuition for why this is the case, note that by Tweedie’s formula, the optimal denoiser is, up to a shift by $x$ and scaling, given by the derivative of $p(x; \tau_\ell^*)$, i.e. the log-density of the prior convolved with some Gaussian. Even if the prior were a mixture of Dirac deltas, the prior convolved with Gaussian is a mixture of Gaussians, and the log-density of this has derivatives at all orders. In particular, $p(x; \tau_\ell^*)$ is always positive and continuous (due to the convolution with a Gaussian, no matter the prior), so all higher order derivatives of $\frac{\nabla p(x; \tau_\ell^*)}{p(x; \tau_\ell^*)}$ exist by applying the quotient rule.
  • We thank the reviewer for bringing up these details and will expand on them in the revisions.

• On the right-hand-side error of the inequality: This is a typo. The right-hand side of theorem 2 should be $v^L$ and $x^L$ instead of $\hat{v}^L$ and $\hat{x}^L$ as mentioned in the formal version of the theorem in the appendix (theorem 3). This follows from the closeness of the state-evolution parameters ($\tau^L$ and $\hat{\tau}^L$) of iterates $x^L$ and $\hat{x}^L$. We thank the reviewer for pointing this out, and we will update it in the revised draft.

• On experiments with non-product signal priors: We have subsequently extended our experiments to non-product priors. We do see similar advantages to unrolling over Bayes-AMP show up in this setting; please see the global response and the supplemental second figure in the attached pdf for a discussion of the non-product prior experiments.

We thank the reviewer for their detailed comments.

• On a relatively niche topic within the NeurIPS community: While understanding the Bayes-AMP algorithm in itself is a relatively niche area for the NeurIPS community, our theoretical results show that the learnability of Bayes-AMP amounts to characterizing training dynamics for the score-matching objective. Given the prevalence of score matching within the literature on diffusion models, which is of very wide interest in the NeurIPS community, we believe that this line of work involving unrolling and AMP offers valuable insight into the dynamics and approximation capabilities of diffusion models. Some of this has been explored in work by [Mei24] and [MW24] (see our related work section in the Appendix).

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

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

• AMP on non-Gaussian measurements:

  • We agree with the reviewer that AMP has these limitations, but empirically we find that unrolling offers a powerful way to interpolate between the theoretically optimal guarantees of AMP in “nice” Gaussian settings and the practically useful guarantees of neural network-based methods for inference in real-world settings.
  • For example, we find in our “learnable B” experiments that by learning auxiliary matrix parameters, unrolling can surpass the performance of Bayes AMP (and baselines such as ISTA) on non-Gaussian designs despite working with an architecture inspired by AMP. In particular, we examine a random orthogonal design (fig. 4) and a random Gram matrix design (attached fig. in pdf in global response) and show our method can surpass Bayes-AMP. While our theory only explains why unrolling can compete with AMP, it opens up the potential for future work exploring why learning auxiliary matrix parameters can ameliorate the non-robustness of AMP to other designs.

• On the conjectured to be optimal among all poly-time algorithms: We thank the reviewer for bringing this up. We will in our revision refer to the precise GFOM optimality results for compressed sensing and low-rank matrix estimation to qualify the optimality of AMP.

• On the denoiser learned by unrolling can outperform the optimal denoiser:

  • With regards to learning the denoiser, as in Figure 1, the unrolled denoisers outperform Bayes AMP by a relatively small margin. The NMSE in Figure 1 is averaged over $2^{15}$ samples with fixed, randomly chosen design, and we observed this improvement consistently for other randomly sampled designs as well, so we do not believe the small improvement of MLP denoisers over vanilla AMP in Figure 1 is from random fluctuations. The broader point we want to emphasize is that Bayes AMP is only known to be “optimal” up to $o_d(1)$ margin in NMSE, and the fact that the MLP denoisers we learn improve over this in the finite-dimensional setting is consistent with this state of knowledge.
  • We also observed consistently that as the signal dimension decreased, we could outperform Bayes AMP by learning the “B” matrix (or measurement transpose) as detailed in Figure 3. Here, we believe the learned B experiments take advantage of finite-dimensional sub-optimalities of AMP as a GFOM, rather than just finite-dimensional sub-optimalities of the denoisers used by Bayes AMP.

• On the intuitive understanding of the fine-tuning:

  • We observed finetuning was not largely impactful for the compressed sensing experiments but was helpful for rank-one PCA. Note that PCA has different strengths of initialization for AMP of the signal estimate while compressed sensing has a “canonical” initialization ($x_0=0$). PCA thus induces a “burn-in” phase where the AMP estimates reflect almost no signal at early iterations (see Figure 6 of the submission).
  • We suspect that finetuning assists in maintaining the consistency of the learned denoisers across this burn-in phase where the correlation between the estimate and MMSE estimate is low, communicating the loss objectives of the later denoisers to earlier iterations.
  • We also observe that finetuning improves validation loss most during this “burn-in” phase, but notably not as much during later iterations where the estimate is close to the MMSE estimate.

• On the intuition of why sequential learning avoid local minima:

  • Sequential learning takes advantage of the following result: Suppose we have learned the Bayes-AMP denoisers in all previous layers. Then the denoiser that minimizes the MSE estimate of the next layer is the corresponding Bayes-AMP denoiser as well (see thm 3 in [MW22]). This simplifies the learning optimization immensely, as we only need the current layer’s denoiser to globally optimize the MSE loss function.
  • Another way to view this is that we know that our learned network should provide the minimum MSE estimate after each layer, but training end-to-end throws away this “intermediate layer” information and only uses the final layer’s optimality.

[MW22] Statistically Optimal First Order Algorithms: A Proof via Orthogonalization. Andrea Montanari, Yuchen Wu.

• Intuition behind initializing at the solution of the previous layer: As for initializing at the solution of the previous layer, the intuition is as follows. As the Bayes-AMP estimates converge, so do the denoisers, so we expect that the denoiser from the previous layer is close to the optimal denoiser, at least closer than a random initialization.

• On quantification of learning rate and number of steps: By sufficiently small learning rate, we mean that the learning rate is $\Theta(1/(\epsilon*m))$ and by a sufficiently large number of steps, we mean polynomial in the complexity of the denoiser and $1/\epsilon$ (See the statement of Theorem 3).

We sincerely thank the reviewer for their effort in reviewing our work and address their comments below.

• On the product distribution assumption: Please see the global response for the discussion about the assumptions on the prior. We have also subsequently expanded experiments to the non-product distribution regime, also included in the global response and the second supplementary figure in the attached pdf.

• On the flexibility of the architecture of the MLP denoiser: The MLP parameterizing the denoiser can in practice be set to any architecture (convolutional net, U-net, etc.), and doesn’t have to be constrained to the three-layer, 70-neuron wide MLP we used.

• On the understanding how architectural variations impact the results: We examined different depths and widths of our MLP parameterization and found that the behavior of unrolling was relatively robust to variations in the choice of the denoiser architecture. The adjustments we made were ultimately meant to control for the amount of error in learning the denoiser, and the architecture we eventually reported results for worked sufficiently well for our experiments while also being quite reasonable to train.

• On the practical likelihood of Assumption 1: Assumption 1 comprises two parts: 1) sub-Gaussianity and 2) Lipschitz continuity of the $\tau-$Gaussian smoothed score function (score function of $x + \tau z$). We believe both assumptions are mild, hold for a large distribution class, and are used in prior literature (e.g., the sub-Gaussianity of a distribution holds for any distribution with bounded support and the Lipschitz continuity of the score function holds for mixtures of spherical Gaussians with bounded means). Please also see the general response on the additional discussions about these assumptions. We will include a discussion about the generality of the assumption in the revised draft.

• The definition of variable $B$: The variable $B$ denotes the Lipschitz constant of the score function of $p(x; \tau)$ where $p(x; \tau)$ is a probability density of random variable $x + \tau z$ where $x$ and $z$ corresponds to the data distribution and standard Gaussian distribution respectively (Please see Assumption 1 for the complete definition of variable $B$).

Thank you for your detail responses I will update my score accordingly to the thoughtful respones the authors have provided.