What's in a Prior? Learned Proximal Networks for Inverse Problems

We sincerely thank the reviewer for their careful review of our work. We have addressed all of your questions and comments, and please don't hesitate to let us know if there are any pending matters.

1. The end the proof from Theorem 4.1 is I think not true...

Thank you for pointing this out. We have corrected the proof. The argument that the algorithm converges to a fixed-point still holds, due to the continuity of LPN. Specifically, since both $f_\theta$ and $\nabla h$ are continuous, the mapping from $x_k$ to $x_{k+1}$, $x_{k+1} = f_\theta(x_k - \eta \nabla h(x_k))$ is continuous. Then, given that the iterative algorithm converges, and that the mapping from $x_k$ to $x_{k+1}$ is continuous, the algorithm must converge to a fixed point. Please see the last paragraph of the proof of Theorem 4.1 and Lemma C.2 for the updated proof.

2. The targeted objective function h + eta^-1 phi contains the step size ...

Thank you for pointing this out. Indeed, this 1/eta weighing of the regularizer might seem strange, but we would like to argue that in our setting it is reasonable and may even be useful.

First, we would like to comment on the cause of $1/\eta$ weighing. For an objective $f(x) + g(x)$, the proximal gradient descent (PGD) algorithm applies ${prox}_{\eta g}(x-\eta \nabla f(x))$

at each iteration. Note that the proximal operator $prox_{\eta g}$ needs to be scaled appropriately by $\eta$. If we follow this strictly, the objective will not change with eta. However, in our case, the proximal operator is fixed at $prox_{\phi}$, and we did not scale it accordingly by $\eta$, which requires evaluation of $prox_{\eta \phi}$. Thus our calculation has a subtle difference from classic PGD, creating the $1/\eta$ weighing in the objective.

Second, this weighing may be useful, as it provides a way to change the regularization weight without retraining LPN. Indeed, eta can be treated as a knob to control the regularization weight: smaller eta means higher weight on the regularization. However, as the reviewer points out, eta being too large would cause the algorithm to diverge (note that in practice $\eta > 1/L$ does not necessarily mean the algorithm will diverge, i.e. the actual upper bound may be larger than $1/L$), so this tuning is not full-spectrum. We do not see how this necessarily impacts the performance, as there exist other ways to tune the regularization weight, e.g., via Gaussian noise level at training.

3. The algorithm analyzed for convergence is just the classical Proximal Gradient Descent (PGD) (or Forward-Backward) ...

We used the result from Bot et al. (2016) because it includes coercivity of the objective as a hypothesis, whereas the main result in Attouch et al. (2013), i.e. Theorem 5.1, writes the result instead in the form “if the iterate sequence is bounded, then it converges”. Although it is straightforward to argue that coercivity of the objective implies iterate boundedness, we assumed that for a ML venue, it would be easier for readers to parse the proof if we appealed to a result with hypotheses that match our setting. In the main body of the submission, we made a note of the precedence of Attouch et al. (2013) via a citation in the “Convergence Guarantees in Plug-and-Play Frameworks” paragraph at the top of page 6 in the main body. If you want us to change the proof to appeal to Attouch et al. (2013) and argue boundedness of the iterate sequence using coercivity instead of appealing to Bot et al. (2016), we can do so.

4. "Definable" is an important notion that should be defined. Moreover, the stability properties of Definable functions ...

We have added a new section to the appendices, Section C.4.1, which contains what we believe is a far more complete and readable overview of the relevant aspects of the KL property that we need to prove our convergence results. The submission originally referenced Attouch et al. 2010 for these facts, which seems to us to be a standard and sufficient reference here (allowing one to avoid tedious verifications), but the revision has more precise references to proofs and spells out all details in the verifications. On this note, we would like to note that you are asking for what seems to us to be a very high standard of rigor for a machine learning conference submission, and we hope you will acknowledge our contribution in meeting it. Please let us know if you have further feedback in this connection.

Thank you for the detailed answer. The majority of my questions have been clarified. In particular, the clarification on o-minimal structures improves significantly the proof.

I have however still some doubts on the following points.

• First the authors experiment with the ADMM algorithms when they showed convergence of the PGD algorithm. This is for me a major limitation and I do not understand the point of proving convergence if not experimenting with the same algorithm afterwards.
• Second the authors compare with the literature on the Celeba dataset. Their network is trained on Celeba when the compared methods seems not to be trained on Celeba. This looks like unfair comparison.

I decided to raise my score to 5.

We thank the reviewer for their questions and comments.

1. It is better to give some details on the training of the learned proximal networks. For example, how to learn the non-negative weights?

Thank you for the question. We have provided extensive training details, such as learning rate, Gaussian noise level, schedule of proximal matching $\gamma$, and batch size in the Appendix E. Sorry for missing the details about how to enforce non-negative weights. During training, we apply weight clipping, i.e., projecting negative weights to 0, at each iteration to ensure the weights are nonnegative. Additionally, in the experiments on CelebA and Mayo-CT, we enforce the weights to be nonnegative at initialization by initializing the weights according to a Gaussian distribution and then taking the exponential. We observed that this initialization helped the training converge faster. We have added these details to Appendix E. In addition, we will also release our code publicly for reproductivity.

In Appendix E.1, Details of Laplacian Experiment:

• To enforce nonnegative weights, weight clipping is applied, projecting the negative weights to zero at each training iteration.

In Appendix E.3, Details of CelebA Experiment:

• We observed that initializing the respective weights to be nonnegative, by initializing them according to a Gaussian distribution and then taking the exponential, helped the training converge faster. Therefore, we applied such initialization in the experiments on CelebA and Mayo-CT. The same weight clipping as in Appendix E.1 is applied to ensure the weights stay nonnegative throughout training.

2. How to set the \alpha in \phi_\theta in the training?

This is a user-defined parameter. We discussed its impact and trade-off in Section 3, the “Recovering the prior from its proximal” paragraph, and footnote 4. The values we used for each experiment are included in Appendix E.

3. How to set \eta in Theorem 4.1 in the training?

The $\eta$ is the step size of gradient descent while using a trained LPN for solving inverse problems in the PnP framework. It is not a parameter during training. Sorry about the confusion. As noted in Theorem 4.1, $\eta$ should be between 0 and $1 / \|A^TA\|$ to guarantee convergence of the PnP iterates with LPN. In practice, $\eta$ is a hyperparameter that can be tuned by the user.

4. The methods compared in this paper are relatively old. It is better to compare the recent methods. Then we could find out the superiority of the proposed method.

We thank the reviewer for this suggestion. Indeed, we have now added further comparison with a newer, highly relevant and state-of-the-art method (which was also suggested by another reviewer) :

• Hurault, Samuel, Arthur Leclaire, and Nicolas Papadakis. "Proximal denoiser for convergent plug-and-play optimization with nonconvex regularization." ICML 2022.

Note, however, that no other method exists that approximates the solution of a MAP problem based on learned networks that provide exact proximal operators. Our LPN is the first to provide such a guarantee. We refer the reviewers to Figure 5 and Table 1 in the revised manuscript for the new results. The results from LPN are comparable and on-par with those from GS networks, while our method provides the added advantage of providing an exact proximal operator – which GS networks might not (we explain in detail why this is the case in our response to Reviewer 4, comment 11).

Thank you for your questions and comments.

1. A straightforward solution to obtain the proximal map of log p(x) is to initially train an energy-based model, E, and then compute the proximal map of og E. The authors should either compare this approach with their proposed method or provide clarification on why this is not considered viable or advisable. For instance, one might anticipate encountering similar challenges as those faced when training energy models when training Learned Proximal Networks (LPNs).

Thanks for bringing up this very relevant comment. There are two main reasons why the above suggested by the reviewer is not straightforward: First, training an energy-based generative model E is a statistically complex task, particularly for distributions supported in high-dimensional spaces. While it is true that good approximate models can be obtained, computing the proximal operator will then boil down to the explicit optimization of a non-convex problem, precluding us from any precise guarantees. Our approach circumvents both of these by computing the exact proximal directly (even if the prior is non-convex), i.e., by learning to compute the maximum of a posterior, which is a significantly easier task than computing the complete posterior and then finding its maximum (which is the approach entailed by the energy model).

2. In a similar vein, a natural point of comparison for the proposed approach would involve using other state-of-the-art unsupervised methods that rely on generative models as priors. While the paper does make a comparison with the less recent adversarial regularizer, it would be valuable to assess its performance against other more recent unsupervised methods.

Thank you for the suggestion. In our revised manuscript, we have added a new and more recent comparison method (Hurault et al. 2022), based on the Gradient Step (GS) denoiser.

• Hurault, Samuel, Arthur Leclaire, and Nicolas Papadakis. "Proximal denoiser for convergent plug-and-play optimization with nonconvex regularization." ICML 2022.

We chose this method because it is closely related to ours, for the following reasons. First, this method also concerns the topic of solving inverse problems via a variational setup and seeking the Maximum a Posteriori estimate with a data-driven prior. Second, this method also attempts to train a denoiser as a proximal mapping and study convergence guarantees when using the trained denoiser within the PnP framework. Lastly, this method was also suggested by another reviewer as a state-of-the-art method that is comparable to our proposed methodology. Therefore we believe it is a very relevant comparison.

We refer the reviewers to Figure 5 and Table 1 in the revised manuscript for the new results. Our results are comparable and on-par with those from the GS networks, while our method provides the added advantage of providing an exact proximal operator – which GS networks might not (we explain in detail why this is the case in our response to Reviewer 4, comment 11).

Questions:

• Equation (3.2) should it be \psi_\thet(x,\alpha) ?

Right. Thank you for pointing this out and we have corrected it.

• Page 5, "demillustrates" should be corrected to "illustrates."

Fixed, thanks!

• Page 5, the terminology "proximal operator of target distribution" is frequently used but remains undefined. This terminology appears to be uncommon in the field and should be clarified or explained for the readers' benefit.

Thank you for the suggestion. It means the proximal operator of the log-prior of the data distribution we aim to learn, i.e. prox_{-log p_x}. We have rephrased it in the paper:

Page 5: "Figure 2 illustrates the limitations of these distance metrics for learning the true proximal operator of the log-prior of the underlying data distribution (a Laplacian, in this example)."

Moreover, to avoid any further confusion, we have explicitly defined “log-prior” the first time it appears in the paper: "In this paper, the “log-prior” of a data distribution px means its negative log-likelihood, −log px."

• It would be insightful to determine whether the log-likelihood computed using LPN matches the log-likelihood computed using other generative models.

We agree with this comment. However, we believe it is out of the scope of this work, which focuses on solving inverse problems with data-driven priors. Note that in this work, we did not attempt to use LPN as a generative model (e.g., using LPN to generate human faces from scratch), or compare it with any generative models. However, we do believe LPN has the potential to be adapted into generative models and we are very interested in exploring this connection in future work.

Thank you for your questions and comments.

1. The main contribution of the work is to propose a way to learn a MAP denoiser through a proximal loss under equation 3.4. Optimizing for this loss entails that the prior distribution is assumed to be a mixture of Gaussians (or Diracs when gamma tends to zero) around the training samples. Why is this a good prior? It seems to me that is a too simplistic prior and in the limit of gamma -> 0 a discontinuous prior.

Not quite: the proximal matching loss (Eq. 3.4) does not require the prior distribution to be a mixture of Gaussians or Diracs. There is almost no assumption on the prior to be learned, except the most basic regularity conditions, e.g., for the prior to be continuous in Thm 3.1 (and discrete in Thm B.1). Note that in the paper, immediately after Eq. 3.4, we do mention that the loss \ell_\gamma can be thought of as an \gamma-approximation to a Dirac, but this is just for intuition, and no conditions are imposed on the actual priors.

2. Although most of the paper was quite clear I found it unclear whether the final implementation enforced convexity on the prior or not. Under section 2 and section 3, a case was made for convexity: a class of learned proximal network was suggested to ensure convexity, and later under (3.1) another regularization was introduced to ensure strong convexity for recovering the prior. I think these sections do not flow well and the thread of logic is easily lost by the reader. The clarity can be improved. Also on a related note, if convexity is assumed, why is it a proper assumption? It most probably would be a too simplistic of an assumption for image priors.

We believe we see the confusion of the reviewer: We do not enforce convexity on the prior at any point. The review is correct in that convex priors are not well suited for natural images, but note that our method can learn nonconvex priors, as has been demonstrated in the MNIST example (Fig. 3b), and we view this as a main advantage of our method. We will further clarify on this point - thank you. In our presentation, while the function $\psi$ is convex, note that it is not the prior - the prior is denoted by $\phi$. Convexity on $\psi$ is enforced to ensure the LPN, which is defined as the gradient of $\psi$, is an exact proximal operator, as required by Proposition 2, which uses the fact that proximal operators of nonconvex priors are equivalent to gradients of convex functions. Please see Figure 1 for a clear illustration of this relation.

3. The details of the algorithm can be more elaborated in the main text. Specially, since the algorithm is very similar to the basic score-based diffusion generative algorithms for solving inverse problems, the differences between the two should be pointed out.

We thank the reviewer for this question. While our methodology is used to solve inverse problems, PnP in general, and our LPN in particular, are significantly different from score-based diffusion models. First: conditional diffusion models do not minimize a variational problem like we do in this paper (as in Eq. 2.1), but instead provide samples from the posterior distribution. Second: score-based sampling arises as the result of inverting a diffusion process which requires access to an MMSE denoiser, whereas we are only concerned with networks that compute a MAP estimate for a learned prior. Third, the reviewer should note that our LPN provides a (MAP) solution in a single call (or forward pass) for the denoising task, whereas sampling from the posterior as with diffusion models requires very expensive computation. For all these reasons, comparing with diffusion based models is out of the scope of this work. However, based on your suggestion, we have included a detailed discussion of these points in our revised version of the paper (please see Appendix A, “Comparison to Diffusion Models” paragraph). If the reviewer has specific works in mind that we should reference and comment on, we would be happy to do so. In terms of algorithm details, we have included the pseudo-code in Algorithms 1 and 4, and details of hyperparameters in Section E. We will also release our code publicly.