What's in a Prior? Learned Proximal Networks for Inverse Problems
Authors: Zhenghan Fang
Sam Buchanan
Jeremias Sulam
Venue: ICLR 2024 poster (2024)
Keywords: Inverse problems
Proximal operators
Plug-and-play
Explicit regularizer
Convergent PnP
Input convex neural networks
Primary area: representation learning for computer vision, audio, language, and other modalities

Abstract: Proximal operators are ubiquitous in inverse problems, commonly appearing as part of algorithmic strategies to regularize problems that are otherwise ill-posed. Modern deep learning models have been brought to bear for these tasks too, as in the framework of plug-and-play or deep unrolling, where they loosely resemble proximal operators. Yet, something essential is lost in employing these purely data-driven approaches: there is no guarantee that a general deep network represents the proximal operator of any function, nor is there any characterization of the function for which the network might provide some approximate proximal. This not only makes guaranteeing convergence of iterative schemes challenging but, more fundamentally, complicates the analysis of what has been learned by these networks about their training data. Herein we provide a framework to develop *learned proximal networks* (LPN), prove that they provide exact proximal operators for a data-driven nonconvex regularizer, and show how a new training strategy, dubbed *proximal matching*, provably promotes the recovery of the log-prior of the true data distribution. Such LPN provide general, unsupervised, expressive proximal operators that can be used for general inverse problems with convergence guarantees. We illustrate our results in a series of cases of increasing complexity, demonstrating that these models not only result in state-of-the-art performance, but provide a window into the resulting priors learned from data.

Published as a conference paper at ICLR 2024
WHAT'S IN A PRIOR?
LEARNED PROXIMAL NETWORKS FOR INVERSE PROBLEMS
Zhenghan Fang"
Sam Buchanan"
Mathematical Institute for Data Science Johns Hopkins University z zfagg238 jhu.edu
Toyota Technological Institute at Chicago sam@t ttiiiicc. edu
Jeremias Sulam
Mathematical Institute for Data Science Johns Hopkins University jsulaml @ jiuu jdu
ABSTRACT
Proximal operators are ubiquitous in inverse problems, commonly appearing as part of algorithmic strategies to regularize problems that are otherwise ill-posed. Modern deep learning models have been brought to bear for these tasks too, in the framework of plug-and-play or deep unrolling, where they loosely resem- ble proximal operators. Yet, something essential is lost in employing these purely data-driven approaches: there no guarantee that general deep network repre- sents the proximal operator of any function, nor is there any characterization of the function for which the network might provide some approximate proximal. This not only makes guaranteeing convergence of iterative schemes challenging but, more fundamentally, complicates the analysis of what has been learned by these networks about their training data. Herein we provide framework de- velop learned proximal networks (LPN), prove that they provide exact proximal operators for a data-driven nonconvex regularizer, and show how new train- ing strategy, dubbed proximal matching, provably promotes the recovery of the log-prior of the true data distribution. Such LPN provide general, unsupervised. expressive proximal operators that can be used for general inverse problems with convergence guarantees. We illustrate our results ing complexity, demonstrating that these models not only result in state-of-the- performance, but provide a window into the resulting priors learned from data.
series of cases of increas-
INTRODUCTION
Inverse problems concern the task of estimating underlying variables that have undergone degra- dation process, such as in denoising, deblurring, inpainting, or compressed sensing (Bertero et al., 2021; Ongie et al., 2020). Since these problems are naturally ill-posed, solutions to any of these problems involve, either implicitly or explicitly, the utilization of priors, or models, about what type of solutions are preferable (Engl et al., 1996; Benning & Burger, 2018; Arridge et al., 2019). Traditional methods model this prior directly, by constructing regularization functions that promote specific properties in the estimate, such as for it to be smooth (Tikhonov & Arsenin, 1977). piece- wise smooth (Rudin et al., 1992; Bredies al., 2010), for it to have sparse decomposition under a given basis or even a potentially overcomplete dictionary (Bruckstein et al., 2009; Sulam et al., 2014). On the other hand, from machine learning perspective, the complete restoration mapping can also be modeled by a regression function and by providing a large collection of input-outpun (or clean-corrupted) pairs of samples (McCann al., 2017: Ongic et al., 2020; Zhu et al., 2018).
"Equal contribution.

---

Published as a conference paper at ICLR 2024
An interesting third alternative has combined these two approaches by making the insightful obser- vation that many iterative solvers for inverse problems incorporate the application of the proximal operator for the regularizer. Such proximal step can be loosely interpreted and, as a result, off-the-shelf strong-pers poron-peroormnn denoising algorithms (as those given by modern deep learning methods) can be employed notable example where proximal operators are replaced with such denoisers (Venkatakrishnan et al., 2013; Zhang al., 2017b; Meinhardt et al., 2017: Zhang et al., 2021; Kamilov al., 2023b; Tachella et al., 2019), but these can be applied more broadly to solve inverse problems, as well (Romano et al., 2017: Romano & Elad, 2015). While this strategy works very well in practice, little is known about the approximation properties of these methods. For instance, do these denoising networks actually (i.e., provably) provide a proximal operator for some regularization function? Moreover, and from variational perspective, would this regularization function recover the correct regularizer, such as the (log) prior of the data distribution? ? Partial answers to some of these questions exist, but how to address all of them in single framework remains unclear (Hurault et al., 2022b; Lunz et al., 2018; Cohen et al., 2021a; Zou et al., 2023; Goujon et al., 2023) (see a thorough discussion of related works in Appendix A). More broadly, the ability to characterize a data-driven (potentially noncon- vex) regularizer that enables good restoration is paramount in applications that demand notions of robustness and interpretability, and this remains an open challenge.
denoising step
subroutine. The Plug-and-Play (PnP) framework is
In this work, we address these questions by proposing a new class of deep neural networks, termed learned proximal networks (LPN), that exactly implement the proximal operator of a gen- eral learned function. Such LPN implicitly learns a regularization function that can be charac- terized and evaluated, shedding light onto what has been learned from data. In turn, we present new training problem, which we dub proximal matching, that provably promotes the recovery of the correct regularization term (i.e., the log of the data distribution), which need not be con- vex. Moreover, the ability of LPNs to implement exact proximal operators allows for guaranteed convergence to critical points of the variational problem, which we derive for PnP reconstruction algorithms under additional assumptions on the trained LPN. We demonstrate through exper- iments on that our LPNs can recover the correct underlying data distribution, and further show that LPNs lead to state-of-the-art reconstruction performance on image deblurring, CT recon- struction and compressed sensing, while enabling precise characterization of the data-dependent prior learned by the model. Code for reproducing all experiments is made publicly available at https: /itpt/tttuuuu.u.o com/ Sulau am Grw.prroeeeaeee proximal www.lllneettsss...mmo
BACKGROUND
Consider an unknown signal in an Euclidean space', x E R", and a known measurement operator, A: F" + RT The goal of inverse problems is to recover x from measurements y A(x) + v RT where v is a noise or nuisance term. This problem is typically ill-posed: infinitely many solutions x may explain (i.e. approximate) the measurement y (Benning Burger, 2018). Hence, prior is needed to regularize the problem, which can generally take the form
1 min - Ily A(x)] R(x),
(2.1)
for a function R(x): R" + IR promoting solution that is likely under the prior distribution of x. We will make no assumptions on the convexity of R(x) in this work.
Proximal operators Originally proposed by Moreau (1965) as a generalization of projection oper- ators, proximal operators are central in optimizing the problem (2.1) by means of proximal gradient descent (PGD) (Beck, 2017), alternating direction method of multipliers (ADMM) (Boyd et al., 2011), or primal dual hybrid gradient (PDHG) (Chambolle & Pock, 2011). For a given functional R as above, its proximal operator prox, is defined by
pr)) == argmin Ily x1|2 + R(x).
(2.2)
When R is non-convex, the solution this problem may not be unique and the proximal mapping is set-valued. Following (Gribonval& & Nikolova, 2020), we define the proximal operator of a function
'The analyses in this paper he generalized directly more general Hilbert spaces.
2

---

Published as a conference paper at ICLR 2024
R as a selection of the set-valued mapping: f(y) is proximal operator of R if and only if f(y) argmin, 21lly - x||² + R(x) for each y E IR. A key result in (Gribonval & Nikolova, 2020) is that the continuous proximal of (potentially nonconvex) function can be fully characterized as the gradient of a convex function, as the following result formalizes.
Proposition 2.1. |Characterization of continuous proximal operators, (Gribonval & Nikolova, 2020, Corollary 1)] Let y C Th be non-empty and open and f ::y + IR" be a continuous function. Then, is a proximal operator of a function R TR. + RR U (+oo} if and only if there exists a convex differentiable function vo such that f(y) Vvl(y) for each y EY.
It is worth stressing the differences between R and v. While f is the proximal operator of R. i.e prox = f. f is also the gradient of a convex v. Vy = (see Figure 1). Furthermore, R may be non-convex, while iir must be convex. As can be expected, there exists a precise relation between R and di and we will elaborate further on this connection shortly. The characterization of proximals of convex functions is similar but additionally requiring f be non-expansive (Moreau, 1965). Hence. by relaxing the nonexpansivity, we obtain broader class of proximal operators. As we will show later, the ability to model proximal operators of non-convex functions will prove very useful in practice, as the log- priors of most real-world data are indeed non-convex.
R prux
Figure 1: Sketch of Prop. 2.1 for R(-) 1-11..
Plug-and-Play This paper closely relates to the Plug-and-Play (PnP) framework (Venkatakrishnan et al., 2013). PnP employs off-the-shelf denoising algorithms to solve general inverse problems within an iterative optimization solver, such as PGD (Beck, 2017; Hurault et al., 2022b), ADMM (Boyd et al., 2011; Venkatakrishnan et al., 2013). half quadratic splitting (HQS) (Geman & Yang, 1995; Zhang et al., 2021),, primal-dual hybrid gradient (PDHG) (Chambolle & Pock. 2011), and Douglas-Rachfowe splitting (DRS) (Douglas & Rachford, 1956; Lions & Mercier, 1979; Combettes & Pesquet, 2007; Hurault et al., 2022b). Inspired by the observation that py) resembles the maximum a posteriori (MAP) denoiser at with a log-prior Rt PnP replaces the explicit solution of this step with generic denoising algorithms, such as BM3D (Dabov et al., 2007: Venkatakrishnan et al., 2013) or CNN-based denoisers (Meinhardt et al., 2017: Zhang et al., 2017b; 2021; Kamilov et al., 2023b), bringing the benefits of advanced denoisers to general inverse problems. While useful in practice, such denoisers are not in general proximal operators. Indeed, modern denoisers need not be MAP estimators at all, but instead typically approximate a minimum mean squared error (MMSE) solution. Although deep learning denoisers have achieved impressive results when used with PnP, little is known about the implicit prior-if any--nncoden in these denoisers, thus diminishing the interpretability of the reconstruction results. Some convergence guarantees have been derived for PnP with MMSE denoisers (Xu et al., 2020), chiefly relying on the assumption that the denoiser is non-expansive (which can be hard to verify or enforce in practice). Furthermore, when interpreted as proximal operators, the prior in MMSE denoisers can be drastically different from the original (true data) prior Gribonval (201 1), raising concerns about correctness. There is broad family of works that relate to the ideas in this work, and we expand on them in Appendix A.
3 LEARNED PROXIMAL NETWORKS
First, we seek a way to parameterize a neural network such that its mapping is the proximal operator of some (potentially nonconvex) scalar-valued functional. Motivated by Proposition 2.1, we will seek network architectures that parameterize gradients of convex functions. A simple way to achieve this is by differentiating a neural network that implements a сonvex function: given a scalar-valued neural network, th : IR" + IR, whose output is convex with respect to its input, we can parameterize a LPN as fe Vuke which can be efficiently evaluated via back propagation. This makes LPN gradient field and a conservative vector field-an of an explicit convex function. Fortunately, this is not an entirely new problem. Amos al. (2017) proposed input convex neural networks (ICNN) that guarantee to parameterize convex functions by constraining the network weights to be non-negative and the nonlinear activations convex and non-decreasing". Consider a single-layer neural network characterized by the weights W E gmxn, bias b E RT and a scalar non-linearity g RR - R. Such a network, at y, given by z- = g (W + b). With this notation, we now move define our LPNs.
II this paper, the "log-prior" " a data distribution Px means its negative log-likelihond, - -log Px- *Otee ways to parameterize gradients of convex functions exist (Richter-Powell et al., 2021), but come with other constraints and limitations (see discussion in Appendix F.1).
3

---

Published as a conference paper at ICLR 2024
Proposition 3.1 (Learned Proximal Networks), Consider a scalar-valued (K + 11-layered neural network the R + R defined by vo(y) = w zK + b and the recursion
21 (HLY + b1).
Z (WWa H&y + b&), k [2
where 0 == (. (W+) (Hg. b&) ) } are learnable parameters, and g is convex, non- decreasing and C² scalar function, and WR & and w have non-negative entries. Let fo be the gradient map of the w.r.t its input, i.e. So = Vyvo. Then, there exists a fanction Ro R IR U (+oo)
such that fo(y) E prox Rg (y), V
The simple proof of this result follows by combining properties of ICNN from Amos et al. (2017) and the characterization of proximal operators from Gribonval & Nikolova (2020) (see Ap- pendix C.I). The C2 condition for the nonlinearity" g is imposed to ensure differentiability of the ICNN wy and the LPN So. which will become useful in proving convergence for PnP algorithms with LPN in Section 4. Although this rules out popular choices like Rectifying Linear Units (Re- LUs), there exist several alternatives satisfying these constraints. Following (Huang et al., 2021), adopt the seftplus function g(x) = z/t log(1 exp (3z)), f-smooth approximation of ReLU. Importantly, LPN can be highly expressive (representing any continuous proximal operator) under reasonable settings, given the universality of ICNN (Huang et al., 2021).
Networks defined by gradients of ICNN have been explored for inverse problems: Cohen et al. (2021a) used such networks to learn gradients of data-driven regularizers, thereby enforcing the learned regularizer to be convex. While this is useful for the analysis of the optimization problem, this cannot capture nonconvex log-priors that exist in most cases of interest. On the other hand, Hurault al. (2022b) proposed parameterizing proximal operators as f(y) = y Vg(y), where Vg is L-Lipschitz with L < 1. In practice, this is realized only approximately by regularizing its Lipschitz constant during training (see discussion in Appendix A). Separately, gradients of ICNNs are also important in data-driven optimal transport (Makkuva et al., 2020; Huang et al., 2021).
Recovering the prior from its proximal Once an LPN fe is obtained, we would like to recover its prox prx--rimmtivee,, Rg. This is important, as this function is precisely the regularizer in the variational objective, minx 2 lly - -Axxx)) 1 + (x).. Thus, being able to evaluate Rg at arbitrary points provides explicit information about the prior, enhancing interpretability of the learned regularizer. We start with the relation between f. Rg and vhe from Gribonval & Nikolova (2020) given by
Ro(foly)) foly)) 2 1l vy)
(3.1)
Given our parameterization for fe. all quantities are easily computable (via a forward pass of the LPN in Proposition 3.1). However, the above equation only allows to evaluate the regularizer Rg at points in the image of fo. Jo sy(y), and not at an arbitrary point x. Thus, we must invert fo, i.e. find y such that fe(y) or even surjective. Thus, as in Huang et al. (2021), we add a quadratic term to taa the (y; a) = (o(y)) + 2 l1, with 0, turning the strongly convex, and its gradient map, fo Vele, invertible and bijective. To compute this inverse, it suffices to minimize the strongly convex objective
This inverse is nontrivial, since in general an LPN may not be invertible
min oo(y; (x,y),
(3.2)
which has a unique global minimizer ý satisfying the first-order optimality condition fo()) = a) x: the inverse we seek. Hence, computing the inverse amounts to solving a con- vex optimization problem- addressed by a variety of solvers, e.g. conjugate gradients.
Another feasible approach to invert Se is to simply optimize miny IIJe(y) x1I2. using, e.g., first- order methods. This problem is nonconvex in general, however, and thus does not allow for global convergence guarantees. Yet, empirically find this approach work well on multiple datasets, yielding a solution ý with small mean squared error (y) xI| We summarize the procedures for estimating the regularizer from an LPN in Algorithm 2 and Appendix D.I.
*Proposition 3.1 also holds if the nonlincarities are different, which we omit for simplicity of presentation.
Note our use of prox-primitive to refer to the function R with respect the operator prox
4

---

Published as a conference paper at ICLR 2024
3.1 TRAINING LEARNED PROXIMAL NETWORKS VIA PROXIMAL MATCHING
To solve inverse problems correctly, it is crucial that LPNs capture the true proximal operator of the underlying data distribution. Given an unknown distribution Px, the goal of training LPN is to learn the proximal operator of its log, prox-log == s Unfortunately, paired ground-truth samples {x,, f (xx)} do not exist in common settings- the prior distributions of many types of real-world data are unknown, making supervised training infeasible. Instead, we seek to train an LPN using only i.i.d. samples from the unknown data distribution in an unsupervised way.
To this end, we introduce novel loss function that we call proximal matching. Based on the obser- vation that the proximal operator is the maximum a posteriori (MAP) denoiser for additive Gaussian noise, i.c. for samples y = x + av with x ~ Px, V ~ (O,0I), I), by minimizing a loss of the form
we train LPN perform denoising
E [d([foy):x)) x
(3.3)
where d is a suitable metric. Popular choices for d include the squared l2 distance (y) x x1 the l distance [y) - xI| or the Learned Perceptual Image Patch Similarity (LPIPS, (Zhang et al., 2018)), all of which have been used to train deep learning based denoisers (Zhang et al., 2017a; Yu et al., 2019; Tian et al., 2020). However, denoisers trained with these losses do not approximate the MAP denoiser, nor the proximal operator of the log-prion prox log The squared lg distance, for instance, leads to the minimum mean square error (MMSE) estimategiven by the mean of the posterior, E[x y). Similarly, the & distance leads the conditional marginal median of the posterior- and not its maximum. As a concrete example, Figure 2 illustrates the limitations of these metries for learning the proximal operator of the log-prior of a Laplacian distribution.
We thus propose a new loss function that promotes the recovery of the correct proximal, dubbed proximal matching loss:
60m(0:7) = E (1Jo()) - -x112)], m)(x)= =1 11- (TT72) exp
> 0. (3.4)
Crucially, fPM only depends on px (and Gaussian noise), allowing (approximate) proximal learn- ing given only finite i.i.d. samples. Intuitively, my can be interpreted as an approximation to the Dirac function controlled by y Hence, minimizing the proximal matching loss LPM amounts to maximizing the posterior probability Pxly (F)))), and therefore results in the MAP denoiser (and equivalently, the proximal of -ogioror). We now make this precise and show that minimizing LrM yields the proximal operator of the log-prior almost surely as 7 > O.
Theorem 3.2 (Learning via Proximal Matching). Consider a signal ~ Px. where is bounded and Px is a continuous density, S and a noisy observation y = x + ov, where v ~ ~N 1) and IR + R be defined as in (3.4). Consider the optimization problem
O. Let (1(x)
f
argmin lim Exy [m, (Ils(y) x||2)]
(3.5)
f measurable 7'20
Then, almost surely for almost all y). (y) = argmax, Px (c) prox_ o log Px (y).
We defer the proof to Appendix C.2 and instead make few remarks. First, while the result above was presented for the loss defined in (3.4) for simplicity. this holds in greater generality for loss functions satisfying specific technical conditions (see Appendix C.2). Second, an analogous result for discrete distributions can also be derived, and we include this companion result in Theorem B.1, Appendix I. Third, the Gaussian noise level a acts as scaling factor on the learned regularizer, as indicated by (y) prox_a" log p.. (y). Thus varying the noise level effectively varies the strength of the regularizer. Lastly, to bring this theoretical guarantee to practice, we progressively decrease until a small positive amount during training according to schedule function 7(-) for an empirical sample (instead of the expectation). and pretrain LPN with E loss before proximal matching. We include an algorithmic description of training via proximal matching in Appendix D.2, Algorithm 3. Connections between the proximal matching loss (3.4) and prior work on impulse denoising and modal regression are discussed in Appendix A.
Before moving on, we summarize the results of this section: the parameterization in Proposition 3.1 guarantees that LPN implement a proximal operator for some regularizer function; the optimization
Thht is, x admits continuous probability density p with respect the Lebesgue measure on T...
5

---

Published as a conference paper at ICLR 2024
problem in (3.2) then provides way to evaluate this regularizer function arbitrary points; and lastly, Theorem 3.2 shows that if we want the LPN to recover the correct proximal (of the log-prior of data distribution), then proximal matching is the correct learning strategy for these networks.
4 SOLVING INVERSE PROBLEMS WITH LPN
Once an LPN is trained, it can be used to solve inverse problems within the PnP framework (Venkatakrishnan et al., 2013) by substituting any occurrence of the proximal step prox a with the leamed proximal network fa. As with any PnP method, our LPN can be flexibly plugged into a wide range of iterative algorithms, such as PGD, ADMM, or HQS. Chiefly, and in contrast to pre- vious Pnl approaches, our LPN-PnP approach provides the guarantee that the employed denoiser is indeed a proximal operator. As we will now show, this enables convergence guarantees absent any additional assumptions on the learned network. We provide an instance of solving inverse problems using LPN with PnP-ADMM in Algorithm I, and another example with PnP-PGD in Algorithm 4.
Convergence Guarantees in Plug-and- Play Frameworks Because LPNs are by construction proximal operators, PnP schemes with plug-in LPNs correspond to iterative algorithms for minimizing the variational objective (2.1), with the implicitly-define regularizer Rs associ- ated to the LPN. As result, convergence guarantees for PnP schemes with LPNs follow readily from convergence analyses of the corresponding optimization proce- dure, under suitably general assumptions. We state and discuss such a guarantee for using an LPN with PnP-ADMM (Algo-
A Algorithm 1 Solving inverse problem with LPN and PnP-ADMM
Input: Trained LPN fe. operator A, measurement y. initial Xg, number of iterations K, penalty param-
eter p
1: Ug 0, 20 <- X0
2: for k = to K - do
3:
Xk+1 argmin f A(x)1 + 11
Ui -x1
4:
u++11 f un + X++1 Z Zk+1 + fo (ux+1 xx+1)
6: end for
Output: XK
rithm 1) in Theorem 4ur proof appeals the nonconvex ADMM analysis of Themelis & Patrinos (2020).
Theorem 4.1 (Convergence guarantee for running PnP-ADMM with LPNs). Consider the sequence of iterates (X&; Ug Zz.. k E {0,1....} defined by Algorithm run with linear measurement operator A and an LPN fe with softplus activations, trained with O < a < 1. Assume further that the penalty parameter p satisfies p ||AT All Then the sequence of iterates (Xk, Ug, 2a. converges to a limit point (x*, u z²) which is a fixed point of the PnP-ADMM iteration (Algorithm 1).
We defer the proof of Theorem 4.1 to Appendix C.4.2. There, we moreover show that X2 converges to critical point of the regularized reconstruction cost (2.1) with regularization function R XR where Rg is the implicitly-defiie parameters of the PnP algorithm (A p for PnP-ADMM)). In addition,
regularizer associated to fe (i.e. fo = prox n. ) and the regular-
ization strength À depends we emphasize that Theorem 4.1 requires the bare minimum of assumptions on the trained LPN: it holds for any LPNs by construction, under assumptions that are all actionable and achievable in practice (on network weights, activation, and strongly convex parameter). This should be contrasted to PnP schemes that utilize black-box denoiser convergence guarantees in this setting require re- strictive assumptions on the denoiser, such as contractivity (Ryu et al., 2019). (firm) nonexpansivity (Sun al., 2019; 2021: Cohen et al., 2021a;b; Tan et al., 2023), or other Lipschitz constraints (Hu- rault et al., 2022a;b), which are difficult to verify or enforce in practice without sacrificing denoising performance. Alternatively, other PnP schemes sacrifice expressivity for principled approach by enforcing that the denoiser takes restrictive form, such being (Gaussian) MMSE denoiser (Xu et al., 2020), a linear denoiser (Hauptmann et al., 2023), or the proximal operator of an implicit convex function (Sreehari et al., 2016; Teadoro et al., 2018).
The analysis of LPNs we use to prove Theorem 4.1 is general enough to be extended straight- forwardly to other PnP optimization schemes. Under similarly-minimal level of assumptions to Theorem 1, we give in Theorem B.2 (Appendix B.2) a convergence analysis for PnP-PGD (Algo- rithm 4), which tends to perform slightly worse than PnP-ADMM in practice.

---

Published as a conference paper at ICLR 2024
Sas
524
2
3
3
3
3
3
9
A--233
5
Sym
2--265
22-221
6
6
6
6
4 A-OL
4
A
)
P-30
2.0
2-02
0003
4-0
2-233
A-10
Figure 3: Left: log-p Rg learned by LPN on MNIST (computed over 100 test images), evaluated at images corrupted by (a) additive Gaussian noise, and (b) convex combination of two images (1 A)x + Ax'. Right: the prior evaluated at individual examples.
5 EXPERIMENTS
Learned proximal
Convex potential
Log-prior R
3.5
b
3.0
We evaluate LPN on datasets of increasing com- plexity, from analytical one-dimensional example of Laplacian distribu- tion to image datasets of increasing dimensions: MNIST (28 x 28) (LeCun,
prone
Referonce
2.5
2,0
1.5
-1
1.0
0.5
-2
0.0
Figure 2: The proximal fe. convex potential thg, and log-prior Rs learned by LPN via the squared Eg loss, é loss, and proximal matching loss LrM for a Laplacian distribution (ground truth in gray).
1998), CelebA (128 x 128)
(Liu et al., 2018), and Mayo-CT (512 x 512)
(McCollough, 2016). We demonstrate how the ability of LPN to learn an exact proximal for the correct prior reflects on natural values for the obtained log- likeeioooso... Importantly, we showcase the performance of LPN for real-world inverse problems on CelebA and Mayo-CT, for deblurring, sparse-view tomographic reconstruction, and compressed sensing, comparing it with other state-of-the-art unsupervised approaches for (unsupervised) image restoration. See full experimental details in Appendix E.
5.1 WHAT IS YOUR PRIOR?
Learning soft-thresholding from Laplacian distribution We first experiment with a distribu- tion whose log-prior has a known proximal operator, the 1-D Laplacian distribution p(x; (,, 6) exp ( | Letting 0, 6 1 for simplicity, the negative log likelihood (NLL) is the & 28
1
norm, ((x) |a|| log( and its proximal can be written is the soft-threshodding function prox ( (z) = sign(x) max [2| 1, 0). We train a LPN on i.i.d. samples from the Laplacian and Gaussian noise, as (3.3), and compare different loss functions, including the proximal matching {0.5, 0.3, 0.1} in LPM (see (3.4)).
loss LpM. for which we consider different
As seen in Figure when using either the & or l loss, the learned prox differs from the correct soft- thresholding function. Indeed. verifying our analysis in Section 3.1, these yield the posterior mean and median, respectively, rather than the posterior mode. With the matching loss LPAs (? = 0.1 in (3.4)), the learned proximal matches much more closely the ground-truth prox. The third panel in Figure 2 further depicts the learned log-prior Rg associated with each LPN fe. computed using the algorithm in Section 3. Note that Rg does not match the ground-truth log-prior | for & and & losses, but converges to the correct prior with LPM (see more results for different y in Appendix G.1). Note that we normalize the offset of learned priors by setting the minimum value to 0 for visualization: the learned priooroor& Rg has an arbitrary offset (since we only estimate the log-prior). In other words, LPN is only able learn the relative density of the distribution due to the intrinsic scaling symmetry of the proximal operator.
Learning a prior for MNIST Next, we train an LPN on MNIST, attempting to learn a general restoration method for hand-written digits digits-and through it, a prior of the data. For images, implement the LPN with convolution layers; see Appendix E.2 for more details. Once the model is leareed, we evaluate the obtained prior on series of inputs with different types and degrees of
7

---

Published as a conference paper at ICLR 2024
Elared
PYP-M2D
RFCICM
PIP-S SMALE
Our
Grandd
Blundd and leary
PYFERED K.S.
RFCECM TESCE
Ground Tush
INGLE AID
(a) Ohtur
1.0, naviar 0.02.
(b) aMer 1.0, Ouoise 0.04.
Figure 5: Visual results for deblurring on CelebA using Plug-and-Play with different denoisers (BM3D, DrCNN, the gradient step (GS) Prox-DRUNet, and our LPN), for different Gaussian blur kernel standard deviation Jblur and noise standard deviation noosee PSNR and SSIM are presented above each prediction.
Table 1: Deblurring CelebA, over 20 samples. Bold (underline) for the best (second best) score.
METHOD
Oueer
Cunise
Oiese PSNR(T)
2, Dasice
Oaler PSNR(1)
desise
PSNR(1)
SSIM(1)
SSIM(T)
PSNR(()
SSIM
SSIM(")
Blurred and Noisy
27.0 1,6
80 03
24.9
1.0
63 +
24.0 1.7
69 + 04
22.8 :: 1.3
.54 +: 04
PuP-BM3D /Neakatakrishman 2013)
29.5
05
PnP-DnCNN DrCNN (Zhang al., 2017a)
32.3 2.6
90 03
30.9
2.1
+
29.5
2.0
04
28.3 +
.79 05
PuP-OS (Hurault 2022bj
33,0 3,0
92 03
31.4
2.4
.89 + .3
36.1
=
2,5
.8 444
29.3
+ 2.3
.84 .05
Ours
2,9
92
.03
31.3
2.3
30.
.84
perturbations in order to gauge how such modifications to the data are reflected by the learned prior. Figure 3a visualizes the change prior Rg after adding increasing levels of Gaussian noise. As expected, as the noise level increases, the values reported by the log-prior also increases, reflecting that noisier images are less likely according to the data distribution of real images.
The lower likelihood upon perturbations of the samples is general. We depict examples with image blur in Appendix G.2, and also present a study that depicts the non-convexity of the log-prior in Figure 3b: we eval- uate the learned prior at the convex combination of two samples, Ax + (1 - A)x' of two testing images x and x', with À E [0, As de- picted in Figure 3b, as À goes from to the learned prior first in- creases and then decreases, exhibit- ing nonconvex shape. This is nat- ural, since the convex combination of two images no longer resembles a natural image, demonstrating that the true prior should indeed be noncon- vex. As see, LPN can correctly learn this qualitative property in the prior, while existing approaches ing convex priors, either hand-crafted (Tikhonov & Arsenin, 1977; Rudin et al., 1992: Mallat, 1999; Beck & Teboulle, 2009; Elad & Aharon, 2006; Chambolle & Pock, 2011) or data-driven (Mukherjee al., 2021: Cohen et al., 2021a), a), are suboptimal by not faithfully capturing such nonconvexity. All these results collectively show that LPN can learn a good approximation the prior of images from data samples, and the learned prior either recovers the correct log-prior when it is known (in the Lapla- cian example), or provides a prior that coincides with human preference of natural, realistic images. With this at hand, we now move to address more challenging inverse problems.
FDP
AR
Ours
Ground Truth
0.233
0.878
33.16
Sparse-view tomographic reconstruction.
Sparsity (Wovelet)
Ours
Ground Truth
28 cB, 0.722
31.44 0870
38 82 d8, 0.388
(b) Compressed sensing (compression rate 1/16).
Figure 4: Results on the Mayo-CT dataset (details in text).
5.2 SOLVING INVERSE PROBLEMS WITH LPN
CelebA We now showcase the capability of LPN for solving realistic inverse problems. We begin by training an LPN on the CelebA dataset, and employ the PnP- methodology for deblur- ring. We compare with state-of-the-art PnP approaches: PnP-BM3D (Venkatakrishhan al., 2013),
8

---

Published as a conference paper at ICLR 2024
which uses the BM3D denoiser (Dabov et al., 2007), PnP-DnCNN, which uses DnCNN as the denoiser (Zhang et al., 2017a) and PnP-GS using the gradient step proximal denoiser called Prox- DRUNet (Hurault et al., 2022b). Both DrCNN and Prox-DRUNet have been trained on CelebA. As shown in Table 1, LPN achieves state-of-the-art result across multiple blur degrees, noise levels and metries considered. As visualized in Figure 5, LPN significantly improves the quality of the blurred image, demonstrating the effectiveness of the learned prior for solving inverse problems. Compared to the state-of-the-art methods, LPN can produce sharp images with comparable vi- sual quality, while allowing for the evaluation of the obtained prior-which is impossible with any of the other methods.
Table 2:
Numerical results for inverse problems Mayo-CT, computed over 128 test images.
METHOD
PSNR (T)
SSIM (1)
Tomographic reconstruction
Mayo-CT We train LPN on the public Mayo- CT dataset (McCollough, 2016) of Computed Tomography (CT) images, and evaluate for two inverse tasks: sparse-view CT reconstruc- tion and compressed sensing. For sparse-view CT reconstruction. we compare with filtered back-prope nack-projection (FBP) (Willemink & Noël, 2019), the adversarial regularizer (AR) method of (Lunz et al., 2018) with an explicit regu- larizer, and its improved and subsequent ver- sion using unrolling (UAR) (Mukherjee et al., 2021). UAR is trained solve the inverse problem for a specific measurement operator (i.... task-specific). while both AR and LPN are generic regularizers that are applicable to any measurement model (i.e., task-agnostie). In other words, the comparison with UAR is not completely fair, but we still include it here for a broader comparison.
FBP
21.29
.203
Operator-agnestic AR (Lunz et al., 2018)
33.48
.890
Ours
34.14
.891
Operator-specifec
UAR (Mukherjee et al., 2021)
34.76
.897
Compressed sensing (compression rate 1/16)
Sparsity (Wavelet) AR (Lunz al., 2018) Ours
26.54
.666
29.71
.712
38.03
.919
Compressed sensing (compression rate 1/4)
Sparsity (Wavelet) AR (Lunz et al., 2018) Ours
36.80
.921
37.94
.920
44.05
.973
Following Lunz et al. (2018), we simulate CT sinograms using a parallel-beam geometry with 200 angles and 400 detectors, with an undersampling rate of 200 400 experimental details. As visualized in Figure 4a, compared to the baseline FBP, LPN can signif- icantly reduce noise in the reconstruction. Compared to AR, LPN result slightly sharper, with higher PNSR. The numerical results in Table show that our method significantly improves over the baseline FBP, outperforms the unsupervised counterpart AR, and performs just slightly worse than the supervised approach UAR-without even having had access to the used forward operator. Figure 4b and Table show compressed sensing results with compression rates of & and To
30%. See Appendix E.4 for
5122
1 LPN significantly outperforms the baseline and AR, demonstrating better generalizability to different for- ward operators and inverse problems.
6 CONCLUSION
The learned proximal networks presented in this paper are guaranteed to parameterize proximal operators. We showed how the prox-primitive, regularizer function of the resulting proximal (pa- rameterized by an LPN) can be recovered, allowing explicit characterization of the prior learned from data. Furthermore, via proximal matching, LPN can approximately learn the correct prox (i.e. that of the log-prior) of an unknown distribution from only i.i.d. samples. When used to solve gen- eral inverse problems, LPN achieves state-of-the-art results while providing more interpretability by explicit characterization of the (nonconvex) prior, with convergence guarantees. The ability to not only provide unsupervised models for general inverse problems but, chiefly, to characterize the priors learned from data open exciting new research questions of uncertainty quantification (Angelopoulos et al., 2022; Teneggi et al., 2023; Sun Bouman, 2021), sampling (Kadkhodaie & Simoncelli, 2021; Kawar et al., 2021; Chung et al., 2022; Kawar et al., 2022; Feng al., 2023), equivariant learning (Chen et al., 2023a; 2021; 2022a), learning without ground-truth (Tachella et al., 2023; 2022; Gao et al., 2023), and robustness (Jalal al., 2021a; Darestani et al., 2021), I), all of which constitute matter of ongoing work.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
This research has been supported by NIH Grant P41EB031771, as well as by the Toffler Charitable Trust and by the Distinguished Graduate Student Fellows program of the KAVLI Neuroscience Discovery Institute.
REFERENCES
Amir Adler, Yacov Hel-Or, and Michael Elad. A shrinkage learning approach for single image super- 2010: 11th European
resolution with overcomplete representations. In Computer Conference on Computer Vision, Heraklion, Crete, Greece, September 5-11, 2010, Proceedings, Part II I1, pp. 622-635. Springer, 2010.
Jonas Adler and Ozan Öktem. Learned primal-dual reconstruction. IEEE transactions on medical imaging, 37(6):1322-133 2018.
Hemant K Aggarwal, Merry Mani, and Mathews Jacob. Modl: Model-based deep learning archi- tecture for inverse problems. IEEE transactions on medical imaging, 38(2):394-405 2018.
Brandon Amos, Lei Xu, and Zico Kolter, Input convex neural networks. In Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pp- 146-155. PMLR, 2017. URL https: /ttp////ooooiccccecccooocooccccc ngs ml wtt.ssml....r.r.. press/v7C/ amos 1 7b. html.
Anastasios N Angelopoulos, Amit Pal Kohli, Stephen Bates, Michael Jordan, Jitendra Malik, Thayer Alshaabi, Srigokul Upadhyayula, and Yaniv Romano. Image-to-image regression with distribution-Fee uncertainty quantification and applications in imaging. International Confer- ence on Machine Learning, pp. 717-730. PMLR, 2022.
Simon Arridge, Peter Maass, Ozan Öktem, and Carola-Bibiane Schönlieb. Solving inverse problems using data-driven models. Acta Numerica, 28:1-174, 2019.
Hédy Attouch, Jérôme Bolte, Patrick Redont, and Antoine Soubeyran. Proximal alternating mini- mization and projection methods for nonconvex problems: An approach based on the Kurdyka- EOjasiewicz inequality. Mathematics of Operations Research, 35(2):4338443 May 2010. ISSN 0364-765X. doi: 10.1287/moor. 100.0449. URL http: //t::///xxxx// doi wto/10.rg11.... 1287/moor. 1100. 0449.
Hédy Attouch, Jérôme Bolte, and Benar Fux Svaiter. Convergence of descent methods for semi- algebraic and tame problems: proximal algorithms, newwadd-bckweer splitting, and regularized Gauss-Seide methods. Mathematical Programming. A Publication of the Mathematical Pro- gramming Society, 137(1):91-129, February 2013. ISSN 0025-5610, 1436-4646. doi: 10.1007/ s10107-011 1-0484-9. URL https::///// /ttp:://oiiioooo org 0.. 0110://///1100001010000000000000
Thilo Balke, Fernando Davis Rivera, Cristina Garcia-Cardona, Soumendu Majee, Michael Thomp- son McCann, Luke Pfister, and Brendt Egon Wohlberg. Scientific computational imaging code (scico). Journal of Open Source Software, 7(LA-UR-22-2855), 2022.
Amir Beck. First-order methods in optimization. SIAM, 2017.
Amir Beck and Mare Teboulle. fast iterative shrinkage-threeshollnng algorithm for linear inverse problems. SIAM journal imaging sciences, 2(1):183-202, 2009.
Martin Benning and Martin Burger. Modern regularization methods for inverse problems. Acta numerica, 27:1-111 2018
Mario Bertero, Patrizia Boccacci, and Christine De Mol. Introduction to inverse problems in imag- ing. CRC press, 2021.
Dimitri P Bertsekas. Nonlinear Programning. Athena Scientific, 2016. ISBN 9781886529052. URL https::// /ma rket awdroid.
10

---

Published as a conference paper at ICLR 2024
Radu Ioan Bot, Emoö Robert Csetnek, and Szilárd Csaba László. An inertial forward backward algorithm for the minimization of the sum of two nonconvex functions. EURO Journal on Computational Optimization, 4(1):3-25, February 2016. ISSN 2192-4414. doi: 10.1007/ $13675-01150045-8 URL pettpp://i //d/do....
000000071111155550000000553 5-00055--.
Stephen Boyd, Neal Parikh, Eric Chu, Borja Peleato, Jonathan Eckstein, et al. Distributed optimiza- tion and statistical learning via the alternating direction method of multipliers. Foundations and Trends@ in Machine learning, 3(1):1-122, 2011.
Stephen Boyd and Lieven Vandenberghe. Convex optimization. Cambridge university press, 2004.
Kristian Bredies, Karl Kunisch, and Thomas Pock. Total generalized variation. SIAM Journal Imaging Sciences, 33)):492-526, 2010.
Alfred M Bruckstein, David L Donoho, and Michael Elad. From sparse solutions of systems of equations to sparse modeling of signals and images. SIAM review, 51(1):44-81, 2009.
Antonin Chambolle and Thomas Pock. A first-order primal-dual algorithm for convex problems with applications to imaging. Journal of mathematical imaging and vision, 40:1 20-155,, 2011.
Dongdong Chen, Julián Tachella, and Mike E Davies. Equivariant imaging: Learning beyond the range space. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp- 4379-4388, 2021.
Dongdong Chen, Julián Tachella, and Mike E Davies. Robust equivariant imaging: a fully unsuper- vised framework for learning to image from noisy and partial measurements. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 5647-5656, 2022a.
Dongdong Chen, Mike Davies, Matthias Ehrhardt, Carola-Bibiane Schönlieb, Ferdia Sherry, and Julián Tachella. Imaging with equivariant deep learning: From unrolled network design to fully unsupervised learning. IEEE Signal Processing Magazine, 40(1): 147, 2023a.
Hongrui Chen, Holden Lee, and Jianfeng Lu. Improved analysis of score-based generative modeling: User-friendly bounds under minimal smoothness assumptions. In International Conference on PMLR, 2023b.
Machine Learning, pp.
Tianlong Chen, Xiaohan Chen, Wuyang Chen, Zhangyang Wang, Howard Heaton, Jialin Liu, and Wotao Yin. Learning to optimize: A primer and a benchmark. The Journal of Machine Learning Research, 23(1):8522-863 2022b.
Hyungjin Chung, Jeongsol Kim, Michael T Mccann, Marc L. Klasky, and Jong Chul Ye. Diffusion posterior sampling for general noisy inverse problems. arXiv preprint arXiv:2209. 2022.
Regev Cohen, Yochai Blau, Daniel Freedman, and Ehud Rivlin. It has potential: Gradient-driven de-
noisers for convergent solutions to inverse problems. Advances Neural Information Processing Systems, 34:1522181611 2021a.
Regev Cohen, Michael Elad, and Peyman Milanfar. Regularization by denoising via Fixed-Poin projection (RED-PRO). SIAM journal on imaging sciences, 14(3):174-1406, January 2021b. doi: 10.1137/20M1337168 URL http::://bb..orr/1...111111 1
Patrick L Combettes and Tean-Christople Pesquet. A douglas-rachford splitting approach to nons- mooth convex variational signal recovery. IEEE Journal of Selected Topics in Signal Processing, 1(4):564-574, 2007.
Kostadin Dabov, Alessandro Foi, Vladimir Katkovnik, and Karen Egiazarian. Image denoising by sparse 3-d cenesfoooommmmin collaborative filtering. IEEE Transactions image processing, 16 (8):2080-209)) 2090, 2007.
Mohammad Zalbagi Darestani, Akshay S Chaudhari, and Reinhard Heckel. Measuring robustness in deep learning based compressive sensing. In International Conference on Machine Learning, pp. 2434.2444. PMLR, 2021.
11

---

Published as a conference paper at ICLR 2024
Mauricio Delbracio and Peyman Milanfar. Inversion by direct iteration: An alternative to denois- ing diffusion for image restoration. March 2023. URL /attp://arxiv...o www.a.tww.........t..o....ootttt org/ abs 2303.
11435.
Albrecht Dold, Lectures on Algebraic Topology. Springer Science & Business Media, Decem- ber 2012. ISBN 9783642678219, URL attps://play
detailss i2 d- d=P- xrCQAAQBAJ.
Jim Douglas and Henry H Rachford. On the numerical solution of heat conduction problems in two and three space variables. Transactions of the American mathematical Society, 82(2):421-453 1956.
Michael Elad and Michal Aharon. Image denoising via sparse and redundant representations over learned dictionaries. IEEE Transactions Image processing, 15(12):3736-3745 2006.
Heinz Werner Engl, Martin Hanke, and Andreas Neubauer. Regularization of inverse problems, volume 375. Springer Science & Business Media, 1996.
Zhenghan Fang, Kuo-Wei Lai, Peter van Zijl, Xu Li, and Jeremias Sulam. Deepsti: Towards tensor reconstruction using fewer orientations in susceptibility tensor imaging. Medical image analysis, 87:102829, 2023.
Berthy Feng. Jamie Smith, Michael Rubinstein, Huiwen Chang, Katherine L Bouman, and William T Freeman, Score-based diffusion models as principled priors for inverse imaging, arXiv preprint arXiv:234..11551, 2023.
Yunlong Feng, Jun Fan, and Johan A K Suykens. A statistical learning approach to modal regression. Journal of machine learning research: JMLR, 21(2):1-35, 2020. ISSN 1532-4435, 1533-7928. URL https:tp///.mccc..tm / fmlr. org/ papers /atper///v111111111 .html.
Angela F Gao, Oscar Leong. He Sun, and Katherine Bouman. Image reconstruction without explicit priors. In ICASSP 2023-2023 JEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 1-5. IEEE, 2023.
Donald Geman and Chengda Yang. Nonlinear image recovery with half-quadratic regularization. IEEE transactions on Image Processing, 4(7):932 946. 1995.
Davis Gilton, Greg Ongie, and Rebecca Willett. Neumann networks for linear inverse problems in imaging. IEEE Transactions on Computational Imaging, 6:328 343, 2019.
Davis Gilton, Gregory Ongie, and Rebecca Willett. Deep equilibrium architectures for inverse problems imaging. IEEE Transactions on Computational Imaging, 7:1123-1133, 2021.
Tilmann Gneiting. Making and evaluating point forecasts. Journal of the American Statistical
Association, 106(494):746-66, 2011. ISSN 0162-1459. URL http://///// /www www.irsttooooooooo stable/4141640
Alexis Goujon, Sebastian Neumayer, and Michael Unser. Learning weakly convex regularizers for convergent Image- Recontrrrtttnn algorithms. August 2023. URL abs/2308 10542.
Karol Gregor and Yann LeCun. Learning fast approximations of sparse coding. In Proceedings of the 27th international conference on international conference on machine learning, pp. 399-406, 2010.
Rémi Gribonval. Should penalized least squares regression be interpreted as maximum posteriori estimation? IEEE transactions signal processing: a publication of the IEEE Signal Processing Society, 59(5):2405-244 2410, May 2011. ISSN 1053-587X, 1941-0476. doi: 10. 109/TSP..01
2107908. URL http:
/tgg/10.00 109//00/555 2011 .21 2107900..
Rémi Gribonval and Mila Nikolova. A characterization of proximity operators. Journal of Mathe- matical Imaging and Vision, 676777) 2020.
12

---

Published as a conference paper at ICLR 2024
Andreas Hauptmann, Subhadip Mukherjee, Carola-Bibiane Schönlieb, and Ferdia Sherry. Conver- gent regularization in inverse problems and linear plug-and-play denoisers. July 2023. URL
/ttt:://axxiv. org/ abs/ /207 2207.09441..
Heinrich. The mode functional not elicitable. Biometrika, 101(1):245-25) 251, 2014. ISSN 0006- 3444. URL http: : www. aetor.orro stable/abeeee//eee/eee 43305608.
Chin-Wei Huang, Ricky T Q Chen, Christos Tsirigotis, and Aaron Courville. Convex potential flows: Universal probability distributions with optimal transport and convex optimization. In International Conference Learning Representations, 2021. URL net/ (ddd=eePPHHH sPxJ.
Samuel Hurault, Arthur Leclaire, and Nicolas Papadakis. Gradient step denoiser for convergent plug-and-play In International Conference on Learning Representationss, 2022a.
Samuel Hurault, Arthur Leclaire, and Nicolas Papadakis. Proximal denoiser for convergent plug- and-play optimization with nonconvex regularization. In International Conference Machine Learning, pp. 9503-9505. PMLR, 2022b.
Ajil Jalal, Marius Arvinte, Giannis Daras, Eric Price, Alexandros G Dimakis, and Jon Tamir. Robust compressed sensing mri with deep generative priors. Advances Neural Information Processing Systems, 34:14938- 14944, 2021a.
Ajil Jalal, Sushrut Karmalkar, Alexandros G Dimakis, and Eric Price. Instance-optimal compressed sensing via posterior sampling. arxiv preprint arXiv:2106. /438, 2021b.
Ziwei Ji and Matus Telgarsky. Directional convergence and alignment in deep learning. In H Larochelle, Ranzato, R Hadsell, M Balcan, and H Lin (eds.), Advances in Neu- ral Information Processing Systems, volume 33, pp. 17176-17186. Curran Associates, Inc., neurips cc/ http//ppeefeeepppppeeeeeeeeeeeee
2020. URL fi ffle:/cffebbtts11s0001111aoddddd
Paperaapppeeee-eee pdf.
Zahra Kadkhodaie and Eero Simoncelli. Stochastic solutions for linear inverse problems using the prior implicit in denoiser. Adv. Neural Inf. Process. Syst., 34:13242-13254, 2021.
Ulughek S Kamilov, Charles A Bouman, Gregery T Buzzard, and Brendt Wohlberg. Plug-and- Play methods for integrating physical and learned models in computational imaging: Theory, algorithms, and applications. IEEE Signal Processing Magazine, 40(1):85-97, January 2023a. ISSN 1558-0792, doi: 10.1 D0.1000//PP200221199993 URL http: /ttp//xx////o//// doi 1100// MSP 2022 3199595.
Ulugbek Kamilov, Charles A Bouman, Gregery T Buzzard, and Brendt Wohlberg. Plug-and- play methods for integrating physical and learned models in computational imaging: Theory, algorithms, and applications. IEEE Signal Processing Magazine, 40(1):85 97, 2023b.
Bahjat Kawar, Gregory Vaksman, and Michael Elad. stochastically. May 2021.
SNIPS: Solving noisy inverse problems
Bahjat Kawar, Michael Elad, Stefano Ermon, and Jiaming Song. Denoising diffusion restoration models. January 2022.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 1rX.v:1412.6980, 2014.
Erich Kobler, Teresa Klatzer, Kerstin Hammernik, and Thomas Pock. Variational networks: Connecting variational methods and deep learning. In Pattern Recognition, Lecture Notes in Computer Science, pp. 281-293. Springer, Cham, September 2017. ISBN 9783319667089, 9783319667096. doi: com/ httpccaapttt///ee 00700000000000 ----9-3-----10111 3-31 66709-6_26 2_..
23. URL httph:t///// www.iiinn.kiing..... springer.
Erich Kobler, Alexander Effland, Karl Kunisch, and Thomas Pock. Total deep variation for linear inverse problems. Proceedings of the IEEE/CVF Conference on computer vision and pattern recognition, pp. 7549-7558, 2020.
13

---

Published as a conference paper at ICLR 2024
Kuo-Wei Wei Lai, Manisha Aggarwal, Peter van Zijl. Xu Li, and Jeremias Sulam. Learned proximal networks for quantitative susceptibility mapping. In Medical Image Computing and Computer Assisted Intervention-MICCAl 2020: 23rd International Conference, Lima, Peru, October 4-8, 2020, Proceedings, Part II 23, pp. 125-135. Springer, 2020.
Kuo-
Yann LeCun. The mnist database of handwritten digits. http://yann. lecun, com/exdb/mmist// 1998.
Jaakko Lehtinen, Jacob Munkherg, Jon Hasselgren, Samuli Laine, Tero Karras, Miika Aittala, and Learning image restoration without clean data. In Jennifer Dy and Andreas Krause (eds.). Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 2965 29744. PMLR, 10-15 Jul 2018. URL https: httpp//oceingg...e// mlr press/ /tt///8011eeee Lete/eehtiie innn 8a html.
Timo Aila. Noise2Noise:
Guoyin Li and Ting Kei Pong. Douglas-Rachfort splitting for nonconvex optimization with application to nonconvex feasibility problems. Mathematical Programming. Publication of the Mathematical Programming Society, 159())371-401, September 2016. ISSN 0025- 5610, 1436-4646. doi: 10. 10.1007711110--055-096----- URL https:////// http://i...or..o././//
110107-015-0961-
Housen Li, Johannes Schwab, Stephan Antholzer, and Markus Haltmeier. Nett: Solving inverse problems with deep neural networks. Inverse Problems, 36(6):065005, 2020.
Pierre-Louis Lions and Bertrand Mercier. Splitting algorithms for the sum of two nonlinear opera- tors. SIAM Journal on Numerical Analysis, 16(6):964-979) 1979.
Jialin Liu, Xiaohan Chen, Zhangyang Wang, and Wotao Yin. ALISTA: Analytic weights are as good as learned weights in LISTA. In International Conference on Learning Representations, 2019.
URL https: http//ppeeevviieeee net worum?
Jaming Liu, Xiaojian Xu, Weijie Gan, Ulugbek Kamilov, et al. Online deep equilibrium learning for regularization by denoising. Advances in Neural Information Processing Systems, 35:25363- 25376, 2022.
Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Large-scale celebfaces attributes (celeba) dataset. Retrieved August, 15(201 1, 2018.
Ta Lê Loi. Lecture 1: O-minimal structures. In The Japanese-Australian Workshop on Real and Complex Singularities: JARCS III, volume 43, pp. 19-31, Australiaan National University, Mathe- matical Sciences Institute, January 2010. URL htt.a//pp.//pppppeeeeaaaaaadorr.
Stanislaw Lojasiewicz. Une propriété topologique des sous-ensembles analytiques réels. Les équations aux dérivées partielles, 117:87-89, 1963.
Sebastian Lunz, Ozan Öktem, and
Schönlieb. Adversarial regularizers in inverse problems. Advances in neural information processing systems, 31, 2018.
Ashok Makkuva, Amirhossein Taghvaei, Sewoong Oh, and Jason Lec. Optimal transport mapping via input convex neural networks. In International Conference on Machine Learning, pp. 6672- 6681. PMLR, 2020.
Stéphane Mallat. A wavelet tour of signal processing. Elsevier, 1999.
Morteza Mardani, Qingyun Sun, David Donoho, Vardan Papyan, Hatef Monajemi, Shreyas Vasanawala, and John Pauly. Neural proximal gradient descent for compressive imaging. Ad- vances in Neural Information Processing Systems, 31, 2018.
Michael T McCann, Kyong Hwan Jin, and Michael Unser. Convolutional neural networks for inverse problems in imaging: A review. JEEE Signal Processing Magazine, 34(6):85-95, 2017.
McCollough. Tu-fg-207a-04: overview of the low dose et grand challenge. Medical physics, 43 (695-335) 2016.
14

---

Published as a conference paper at ICLR 2024
Tim Meinhardt, Michael Moller, Caner Hazirbas, and Daniel Cremers. Learning proximal operators: Using denoising networks for regularizing inverse imaging problems. In Proceedings of IEEE International Conference Computer Vision, pp. 1781-1790, 2017.
Vishal Monga, Yuelong Li, and Yonina Eldar. Algorithm unrolling: Interpretable, efficient deep learning for signal and image processing. IEEE Signal Processing Magazine, 38(2):18-44 March 2021. ISSN 1538-0792. doi: 10.1 109//M/P2200000010000 URL http: /xx .dxx.o. doi .rrg/ 10. 1109 /M00//MPPPPM 2020 3016905.
Jean-Jacques Moreau. Proximité et dualité dans espace hilbertien. mathématique de France, 93:273-299, 1965.
Bulletin de la Société
Subhadip Mukherjee, Sören Dittmer, Zakhar Shumaylov, Sebastian Lunz, Ozan Öktem, and Carola-Bibiana Schönlieb. Learned convex regularizers for inverse problems. arXiv preprint SBIi:2008.028833 2020.
Subhadip Mukherjee, Marcello Carioni, Ozan Öktem, and Carola-Bibiane Schönlieb. End-to-end reconstruction meets data-driven regularization for inverse problems. Advances in Neural Infor- mation Processing Systems, 34:21413-21425, 2021.
2021.
Gregory Ongie, Ajil Jalal, Christopher A Metzler, Richard G Baraniuk, Alexandros G Dimakis, and Rebecca Willett. Deep learning techniques for inverse problems in imaging. IEEE Journal on Selected Areas Information Theory, 1(1):39-56, 2020.
Jack Richter-Powell, Jonathan Jon Lorraine, and Brandon Amos. Input convex gradient networks. arXiv preprint arXiv:2111.12187, 2021.
Tyrell Rockafellar and Roger J-B Wets. Variational Analysis. Grundlehren der mathematis- chen Wissenschaften. Springer- Verlag Berlin Heidelberg, edition, 1998. ISBN 9783642024313, 9783540627722.
doi:
URL https: http://wwwwww.. www.ngepprneerr.. com/
us/book /tt/b/oo///035540777777770
Yaniv Romano and Michael Elad. Boosting of image denoising algorithms. SIAM Journal Imaging Sciences, 182):1187- 2015.
Yaniv Romano, Michael Flad, and Peyman Milanfar. denoising (red). SIAM Journal on Imaging Sciences,
The little engine that could: Regularization by
2017.
Leonid Rudin, Stanley Osher, and Emad Fatemi. Nonlinear total variation based noise removal algorithms. Physica D: nonlinear phenomena, 60(1-4):259-263 1992.
Walter Rudin. Principles of mathematical analysis. McGraw-Hir New York, 3 edition, 1976. ISBN 9780070542358. URL httph:////// / openlibrarye org.orgggoooo.oo http//ookk//tt/ 0L51 00t///0551999999 95991M. opds.
Emnsst Ryu, Jialin Liu, Sicheng Wang, Xiaohan Chen, Zhangyang Wang, and Wotao Yin. Plug- and-Play methods provably converge with properly trained denoisers. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 5546-5557. PMLR, 2019. URL http: http://pooceeenn....o/ mmr...mmm.......m............mm. wwwpppesssvvv.... / ryul 9a. html.
Vineet R Shenoy, Tim K Marks, Hassan Mansour, and Suhas Lohit. Unrolled ippg: Video heart rate estimation via unrolling proximal gradient descent. In 2023 IEEE International Conference Image Processing (ICIP), pp. 2715-2719 IEEE, 2023.
Suhas Sreehari, S V Venkatakrishnan, Brendt Wohlberg, Gregery T Buzzard, Lawrence F Drummy, Jeffrey Simmons, and Charles Bouman. Plug-and-Play priors for bright field electron mography and sparse interpolation. IEEE Transactions on Computational Imaging, 2(4):408- 423. December 2016. ISSN 2333-9403. //dx. doi. www.orrg///. 110111099990 TT.//TCI.. 201 2599778.
doi: 10.1109/01220112599778 URL http:
Elias M Stein and Rami Shakarchi. Real Analysis: Measure Theory, Integration, and Hilbert Spaces. Princeton University Press, April 2005. ISBN 9780691113869. URL https:
hiplay.goog. wtw/storseeeeeee httm://ttr//ooos//////////o
15

---

Published as a conference paper at ICLR 2024
Jeremias Sulam, Boaz Ophir, and Michael Elad. Image denoising through multi- learnt dictio- naries. In 2014 IEEE International Conference on Image Processing (ICIP), pp. 808 812. IEEE, 2014.
Jeremias Sulam, Aviad Aberdam, Amir Beck, and Michael Elad. On multi-layer basis pursuit, efficient algorithms and convolutional neural networks. IEEE transactions on pattern analysis and machine intelligence, 12(8):166-1198) 2019.
Jeremias Sulam, Ramchandran Muthukumar, and Raman Arora. Adversarial robustness of super- vised sparse coding. Advances in neural information processing systems, 33:21 10-2121, 2020.
He Sun and Katherine L Bouman, Deep probabilistic imaging: Uncertainty quantification and multi- modal solution characterization for computational imaging. In Proceedings of the AAAI Confer- ence on Artificial Intelligence, volume 35, pp. 2628-2637, 2021.
Yu Sun, Brendt Wohlberg, and Ulugbek S Kamilov. An online Plug-and-Play algorithm for reg- ularized image reconstruction. IEEE Transactions on Computational Imaging, 5(3):395-408 September 2019. ISSN 2333-9403. doi: 10. (0.1109/TC122119229935 org/10. 110111109/11111 /CII. 2019 2893568.
URL http: do.....
Yu Sun, Zihui Wu, Xiaojian Xu, Brendt Wohlberg, and Ulugbek S Kamilov. Scalable Plug-and- Play ADMM with convergence guarantees. IEEE Transactions on Computational Imaging, 7: 849-863, 2021. ISSN 2333-9403. doi: (..109//CC2222133999990 URL bttp:://ak doi. exg/10.1r
DB0.I..222..090005..0000 3094062.
Julián Tachella, Yoann Altmann, Nicolas Mellado, Aongus McCarthy, Rachael Tobin, Gerald S Buller, Jean-Yves Tourneret, and Stephen McLaughlin. Real-time 3d reconstruction from single- photon lidar data using haug-and-plan point cloud denoisers. Nature communications, 10(1):4984 2019.
Julián Tachella, Dongdong Chen, and Mike Davies. Unsupervised learning from incomplete mea- surements for inverse problems. Advances in Neural Information Processing Systems, 35:4983- 4995, 2022.
Julián Tachella, Dongdong Chen, and Mike Davies. Sensing theorems for unsupervised learning in linear inverse problems. Journal of Machine Learning Research, 24(39):1-45 2023.
Hong Ye Tan, Subhadip Mukherjee, Junqi Tang, and Carola-Bibiane Schönlieb. Provably conver- gent Plug-and-Play Quasi-Newton methods. March 2023. URL http::////xxxii..rrrgg 0503.07273
Jacopo Teneggi. Matthew Tivnan, Web Stayman, and Jeremias Sulam. How to trust your diffusion model: A convex optimization approach to conformal risk control. International Conference on Machine Learning. pp. 33940-33960. PMLR, 2023.
Afonso M Teodoro, Jose M Bioucas-Dias. and Mario A T Figueiredo. A convergent image fusion algorithm using Scene- Adapted Baussian-MixtureeBaaa
denoising. IEEE transactions on image processing: publication of the IEEE Signal Processing Society, September 2018. ISSN 1057- 7149, 1941-0042, doi: 10.1 1..1199911.2201888677777 URL http: // //dx. dx. doi orgg wttg/1..100//....1.............. 1109/
TIP 2018 8 2 2869788888 69777.
Andreas Themelis and Panagiotis Patrinos. Douglas Douhlas-Rachford splitting and ADMM for nonconvex Rachford
optimization: Tight convergence results. SIAM journal on optimization: a publication of the Society for Industrial and Applied Mahematics, 3001)))111,,3311111111111111 January 2020. ISSN 1052-6234. doi: 10.1137/ 11/1138/11111 163993. URL /ttptttt////.i..og///..../ 1.t.////....../...... 8M1 63993.
Chunwei Tian, Lunke Fei, Wenxian Zheng, Yong Xu, Wangmeng Zuo, and Chia-Was Wen Lin. Deep learning on image denoising: An overview. Neural Networks, 131:251-275, 2020.
Andrey N Tikhonov and Vasiliy Y Arsenin. Solutions of ill-posed problems. vh winston & sons, 1977.
Baharch Tolooshams, Satish Mulleti, Demba Ba, and Yonina C Eldar. Unrolled compressed blind- deconvolution. IEEE Transactions Signal Processing, 2023.
16

---

Published as a conference paper at ICLR 2024
Lou Van den Dries. Tame Topology and O-minimal Structures. Cambridge University Press, May 1998. ISBN 9780521598385.
doi: 10.1017/CB0998051522299 URL https: //ppa. com store http:///c....cccco/ooo/////aaaaa
Lou van den Dries and Chris Miller. On the real exponential field with restricted analytic functions.
Israel Journal of Mathematics, 85(1):19-5,, February 1994. ISSN 0021-2172, 1565-8511. doi:
10.1007/BF02758635. URL httph:t////cc. org. org/ /t/g///..0 100.1000//0225585555555
Lou van den Dries and Chris Miller
Geometric categories and o-minimal struc-
tures. Duke Mathematical Journal, 84(2)497-540, August 1996.
ISSN 0012-
7094, 1547-7398. doi: 1012255880122-79444---4466---
URL https://
50012-7094-96-08111 full.
Singanallur V Venkatakrishnan, Charles Bouman, and Brendt Wohlberg. Plug-and-play priors for model based reconstruction. In 2013 IEEE Global Conference on Signal and Information Processing, pp. 945-948. IEEE, 2013.
Martin Willemink and Peter B Noël. The evolution of image reconstruction for ct-from filtered back projection to artificial intelligence. European radiology, 29:2185-2195, 2019.
Xiaojian Xu, Yu Sun, Jiaming Liu, Brendt Wohlberg, and Ulugbek S Kamilov. Provable convergence of Plug-and aan--ndd-laa Play priors with MMSE denoisers. IEEE Signal Processing Letters, 27:1280-1284, 2020. ISSN 1558-2361. doi: 10. 10.10099..P..200.33066999 URL /ttt:://dx.cr org/10. 1109 1109/LSP 2020 3006390.
Ming Yan and Wotao Yin. Self equivalence of the alternating direction method of multipliers. In Roland Glowinski, Stanley J Osher, and Wotao Yin (eds.), Splitting Methods in Communication, Imaging, Science, and Engineering, pp. 165-194. Springer International Publishing, Cham, 2016. ISBN 9783319415895, doi: 10. 10007978---3199411555 1
URL https: //ooi. org/
Songhyun Yu, Bumjun Park, and Jechang Jeong. Deep iterative down-up for image denoising. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition workshops, pp. 0-0, 2019.
Kai Zhang, Wangmeng Zuo, Yunjin Chen, Deyu Meng, and Lei Zhang. Beyond gaussian denoiser: Residual learning of deep cnn for image denoising. IEEE transactions on image processing, 26 11:3142-3155 2017a.
Kai Zhang, Wangmeng Zuo, Shuhang Gu, and Lei Zhang. Learning deep enn denoiser prior for image restoration. In Proceedings of the IEEE conference computer vision and pattern recog- nition, pp. 3929-3938, 3938, 2017b.
Kai Zhang, Luc Van Gool, and Radu Timofte. Deep unfolding network for image super-resolution. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 3217-3226, 2020.
Kai Zhang, Yawei Li, Wangmeng Zuo, Lei Zhang, Luc Van Gool, and Radu Timofte. Plug-and-play image restoration with deep denoiser prior. IEEE Transactions on Pattern Analysis and Machine Intelligence, 44(10):630--6376 2021.
Richard Zhang, Phillip Isola, Alexci A Efros, Eli Shechtman, and Oliver Wang. The unreasonable effectiveness deep features as a perceptual metric. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 586-595, 2018.
Bo Zhu, Jeremiah Liu, Stephen F Cauley, Bruce R Rosen, and Matthew S Rosen. Image recon- struction by aamain-transfom manifold learning. Nature, 555(7697):487-49 2018.
Zihao Zou, Jiaming Liu, Brendt Wohlberg, and Ulugbek S Kamilov. Deep equilibrium learning of explicit regularizers for imaging inverse problems. arXiv preprint arXiv:2303.0533 2023.
17

---

Published as a conference paper at ICLR 2024
A RELATED WORKS
Deep Unrolling In addition to Plug-and-Play, deep unrolling is another approach using deep neu- ral networks to replace proximal operators for solving inverse problems. Similar PnP, the deep unrolling model is parameterized by an unrolled iterative algorithm, with certain (proximal) steps replaced by deep neural nets. In contrast to PnP,, the unrolling model trained in an end-to-end fashion by paired data of ground truth and corresponding measurements from specific forward op- erators. Truncated deep unrolling methods unfold the algorithm for a fixed number of steps (Gregor & LeeCun, 2010; Adler et al., 2010; Liu et al., 2019; Aggarwal et al., 201 8; Adler & Öktem, 2018; Zhang et al., 2020; Monga et al., 2021: Gilton et al., 2019; Tolooshams et al., 2023; Kobler et al., 2017: Chen et al., 2022b; Mardani et al., 2018; Sulam et al., 2019), while infinite-ster step models have been recently developed based on deep equilibrium learning (Gilton et al., 2021; Liu et al., 2022; Zou et al., 2023). In future work, LPN improve the performance and interpretability of deep un- rolling methods in e.g., medical applications (Lai et al., 2020; Fang et al., 2023; Shenoy et al.. 2023) or in cases that demand the analysis of robustness (Sulam et al., 2020). The end-to-end supervision in unrolling can also help increase the performance of LPN-based methods for inverse problems in general.
Explicit Regularizer A series of works have been dedicated to designing explicit data-driven regularizer for inverse problems, such as RED (Romano et al., 2017), AR (Lunz et al., 2018), ACR (Mukherjee et al., 2020), UAR (Mukherjee et al., 2021) and others (Li et al., 2020; Kobler et al., 2020; Cohen et al., 2021a; Zou et al., 2023; Goujon et al., 2023). Our work contributes a new angle to this field, by learning a proximal operator for the log-prior and then recovering the regularizer from the learned proximal.
Gradient Denoiser Gradient step (GS) denoisers (Cohen et al., 2021a; Hurault et al., 2022a;b) are a cluster of recent approaches that parameterize a denoiser as a gradient descent step using the gradient map of a neural network. Although these works share similarities to our LPN, there are a few key differences.
1. Parameterization. In GS denoisers, the denoiser is defined as gradient descent step: f Id Vg. where Id represents the identity operator, and g is a scalar-valued function that is either directly parameterized by neural network (Cohen et al., 2021a), or implicitly defined by a network N R + EK as g(x) = x N(x) [ (Hurault et al., 2022a;b). Cohen et al. (2021a) also experiment with denoiser architecture analogous to our LPN architecture. but find its denoising performance to be inferior to the GS denoiser (we will discuss this further in the final bullet below). In order to have accompanying convergence guarantees when used in PnP schemes, these GS parameterizations demand special struc- tures on the learned denoiser-in in particular, Lipschitz constraints Vg-which can be challenging to enforce in practice.
2. Proximal operator guarantee. The GS denoisers in Cohen et al. (2021a); Hurault et al. (2022a) are not priori guaranteed to be proximal operators. Hurault et al. (2022b) pro- posed to constrain the GS denoiser to be a proximal operator by limiting the Lipschitz constant of Vg. also exploiting the characterization of Gribonval & Nikolova (2020). How- ever, as a result, their denoiser necessarily has a bounded Lipschitz constant. Furthermore, in practice, such a constraint is not strictly enforeed, but instead realized by adding a regu- larization term on the spectral norm of the network during training. Such a regularization only penalizes large Lipschitz constants, but does not guarantee that the Lipschitz constant will be lower than the required threshold. Additionally, the regularization is only com- puted at training data points, thus either not regularizing the network's s behavior globally resulting in loose upper-bounds for it. In other words, such proximal GS denoiser is only "encouraged" to resemble proximal, but it is not guaranteed. On the other hand, our LPN provides the guarantee that the learned network will always parameterize proximal operator.
3. Training. All GS denoiser methods used the conventional lz loss for training. We propose the proximal matching loss and show that it is essential for the network leam the correct proximal operator of the log-prior of data distribution. Indeed, we attribute the inferior performance of the ICNN-based architecture that Cohen et al. (2021a) experiment with,
18

---

Published as a conference paper at ICLR 2024
which is analogous to LPN, to the fact that their experiments train this architecture on MMSE- based denoising, where "regression to the mean" on multimodal and nonlin- car natural image data hinders performance (see, e.g., Delbracio & Milanfar (2023) in this connection). The key insight that powers our successful application of LPNs in experi- ments is the proximal matching training framework, which allows us to make full use of the constrained capacity of the LPN in representing highly expressive proximal operators (corresponding (nearly) maximum a-posteriori estimators for data distributions).
Comparisons to Difflusion Models Recently, score-based diffusion models have proven very effi- cient for unconditional and conditional image generation. There are several key differences between our work and diffusion models. First, conditional diffusion models do not minimize a variational problem as we do in this paper (as in (2.1)), but instead provide samples from the posterior distribu- tion. Moreover, the diffusion models rely on inverting diffusion process which requires MMSE denoiser, and-just as in the case of regular denoisers- do not approximate any MAP estimate, whereas we concerned with networks that compute MAP estimate for learned prior. In terms of strict ad advantages, one should again note that our approach solves (provides MAP estimate) for a denoising problem with a single forward -aass, whereas sampling with diffusion models requires a large sequence of forward passes of denoising network. Lastly, but also importantly, our method provides an exact proximal operator for a learned prior distribution. Diffusion models have no such guarantee: all these results provide samples from an approximate posterior distribution, which relies on the approximation qualities of the MMSE denoiser that do not exist for general cases (Chen et al., 2023b).
Proximal Matching Loss and Mode-Seeking Regression Objectives In the literature on both deep learning-basee denoising and statistical methodology. prior works have explored training schemes that promote learning the mode of distribution (or, in our denoising setting, the con- ditional mode/MAP estimate of the prior). On the methodological side, it is noted that training with respect single objective function cannot lead the optimal denoiser being the mode uniformly over aufficiently-exresssvvy classes of denoisers and priors, a concept formalized as inelicitability of the mode functional (Gneiting, 2011; Heinrich, 2014). In contrast, our nonparametric result on the proximal matching loss, Theorem 3.2, characterizes the minimizer of limit of a sequence of losses. This is both outside the framework of the preceding references, and distinct from what occurrs in practice, where we attempt to minimize the proximal matching loss with a sufficiently small param- e > O. We expect this latter setting to coincide with correct learning of the mode/MAP estimate of the prior in practical settings of interest, when y is much smaller than the 'characteristic scale' of the prior. Prior work has also considered modal regression in an abstract statistical learning setting (Feng et al., 2020), where in contrast our proximal matching- objective, an approach based on kernel density estimation was advanced.
In the literature on learning deep denoisers, we note that previous work (Lehtinen al., 2018) used an annealed version of an "Eo loss" for mode approximation. with motivation similar to that of proximal matching. Their loss takes different form, (,(f(y) x/: + <)7, where is a small constant and y E [0, 2] is the annealing parameter. Their loss is designed for learning from corrupted targets with random impulse noise, and does not recover the mode of the posterior (as the case of proximal matching), but rather the zero-crossing of the Hilbert transform of the probability density function.
B ADDITIONAL THEOREMS
B. LEARNING VIA PROXIMAL MATCHING (DISCRETE CASE)
Theorem B.1 (Learning via Proximal Matching (Discrete Case)). Consider signal ~ P(x), with P(x) a discrete distribution, and a noisy observation y = x - de, where E ~ No I) and
a > O. Let m, : R + IR be defined by m,(x)) = 1 111-ey eexp
7, Consider the optimization
This definition of m. differs slightly from the one in (3.4), but they are equivalent in terms of minimization objective as they only differ by a scaling constant.
19

---

Published as a conference paper at ICLR 2024
problem
== argmin lim Exyy (m, (llf(y) - xx12)] I measurable 7'50
Then, almost surely (i.e., almost all y), (x) argmax, P(x - yly).
The proof is deferred to Appendix C.3.
B.2 CONVERGENCE OF PNP-PGD USING LPN
Theorem B.2 (Convergence guarantee for running PnP-PGD with LPNs). Consider the sequence of iterates Xk. k E 1..1....J. defined by Algorithm run with a linear measurement operator A and an LPN fe with softplus activations, trained with a < 1. Assume that the step size satisfies 0 < 1j < |/|| |l. Then, the iterales X5. converge to a fixed point x" of Algorithm 4: that is, there exists x" E Rn such that limg X₁ x", annd
fo (x" V)(xxx)) =x'
(B.1)
The proof is deferred to Appendix C.4.1.
C
PROOFS
In this section, we include the proofs for the results presented this paper.
C.I PROOF OF PROPOSITION 3.1
Proof. By Amos al. (2017, Proposition 1), the is convex. Since the activation g is differentiable, the is also differentiable. Hence, fo Viha is the gradient of a convex function. Thus, by Proposi- tion 2.1, fo is proximal operator of a function.
C.2 PROOF OF THEOREM 3.2
Proof. First, note by linearity of the expectation that for any measurable f, one has
lim Exy (m, (1lf(y) x||2)] lim Ex, [1)0)2(f(y)) x)] 7000
(C.1)
750
where 992/2 denotes the density of an isotropic Gaussian random variable with mean and variance 2/2 Because p(x) is continuous density with respect to the Lebesgue measure dx, by Gaussian conditioning, we have that the conditional distribution of x given y admits a density Px with respect to dx as well. Taking conditional expectations, we have
lim Exy [ x)] lim By Exly (p - -x]]
(C.2)
7 ~0
From here, we can state the intuition for the remaining portion of the proof. Intuitively, because the Gaussian density 902/2 concentrates more and more at zero as Y >> O, and meanwhile is nevertheless a probability density for every 7 > 0, the inner expectation over x y leads to simply replacing the integrand with its value at x f(y); the integrand is of course the conditional density of given y, and from here it is straightforward to argue that this leads the optimal f to be (almost surely) the conditional maximum posteriori (MAP) estimate, under our regularity assumptions on p(x).
To make this intuitive argument rigorous, we need to translate our regularity assumptions on p(x) into regularity of Pxly. interchange the 7 limit (C.2) with the expectation over y. and instantiate rigorous analogue of the heuristic "concentration" argument. First, we have by Bayes' rule and Gaussian conditioning
(xly Poa x)p(x)
(Yor p)()),
'oor readers familiar with signal processing Schwartz's theory of distributions, this could be alternately stated "the small-variance limit of the Gaussian density behaves like Dirac delta distribution".
20

---

Published as a conference paper at ICLR 2024
where * denotes convolution of densities; the denominator is the density of y, and it satisfies foo p > 0 since ς > 0. In particular, this implies that pxly is continuous function of (x, y), because p(x) continuous by assumption. We can then write, by the definition of convolution,
Exly (9)))2(f())= x)] = Pp-222 (fy)),
so following (C.2), we have
lim E (p/p/2([(y) x)] lim By [spy2/2 Pay (f(y)) N00
(C.3)
7'50
We are going to argue that the limit can be moved inside the expectation in (C.3) momentarily; for the moment, we consider the quantity that results after moving the limit inside the expectation. To treat this term, we apply a standard approximation to the identity argument to evaluate the limit of the preceding expression. (Stein & Shakarchi, 2005, Ch. 3, Example 3) implies that the densities (2y2/2 constitute approximation to the identity Y + O, and because Px is continuous, we can then apply (Stein Shakarchi, 2005, Ch. 3, Theorem 2.1) to obtain that
lim 9/2/2 * Px Pxy(f(y)) = Pxly (f
? 7'00
In particular, after justifying the interchange of limit and expectation in (C.3), we will have shown, by following our manipulations from (C.1), that
lim Exy (my (Ils()) x(|2)] = By (Pxy(((())) 7'00
(C.4)
We will proceed to conclude the proof from this expression, and justify the limit-expectation inter- change at the end of the proof. The problem at hand is equivalent to the problem
argmax Ey [Pxly (f()))
I measurable
Writing the expectation as an integral, we have by Bayes' rule as above
By [Pxiy (f(y))] = / Porly - f(x)))(((x))dd
Let us define an auxiliary function g g: RR. x R R by yxx,)) = (y x))(x). Then
Ey [Px (f())) / ,(f(y),))dy,
and moreover, for every y, o0.,y) y) is continuous and compactly supported, by continuity and bound- edness of the Gaussian density and the assumption that p(x) is continuous and the random variable p(x) is bounded. We have for any measurable
,(f(y),y)
S max y). xER"
(C.5)
Our aim is thus to argue that there is choice of measurable f such that the preceding bound can be made tight; this will imply that any measurable f maximizing the objective By [Pxyyy [J ([))) satisfies g(f(y),y) = maxxeR" g(x,y) y) almost surely, equivalently that f(y) E argmaxxeR> g(x. y) almost surely. The claim will then follow, because argmax g(x, y) argmax Pxly (x).
To this end, define h(y) maxxeR- g(x,y) Then by the Weierstrass theorem, is finite-valued, and for every y there exists some e E R" such that h(y) y). Because is continuous, it then follows from Rockafellar & Wets (1998, Theorem 1.17(c)) that h is continuous. Moreover, because is continuous and for every y. g.-.)) is compactly supported, g is in particular level- bounded in x locally uniformly in y in the sense of Rockafellar Wets (1998, Definition 1.16), and it follows that the set-valued mapping y + argmaxx x 9 y) R = I is compact-valued, by the Weierstrass theorem, and outer semicontinuous relative to R", by Rockafellar & Wets (1998, Example 5.22). Applying Rockafellar & Wets (1998, Exercise 14.9, Corollary 14.6), we conclude that the set- set-valued mapping y + argmaxx g(x, y) is measurable, and that in particular there exists measurable function f R² + R²" such that f* (y) E argmax, 9(x, y) for every y R". Thus, there is a measurable f atttinnnn the bound in (C.5), and the claim follows after we can justify the preceding interchange of limit and expectation.
21

---

Published as a conference paper at ICLR 2024
To justify the interchange of limit and expectation, will apply the dominated convergence theo- rem, which requires us to show an integrable (with respect to the density of y) upper bound for the function y /x1yle11/2(f(t)) x)]. For this, calculate
Exly [P<3/2(f(y) - x)] (Poz * (y)
Paa - x)p(()) -x]xx
R"
sup Poa (y - -x)p(x) / ()2/2(f(x)) - x)dx
(Poz * (y)
sup P (y - x)p(x²)
(Po= * (x)
by Hölder's inequality and the fact that 922/22 is probability density. Because the random variable x ~ p(x) is assumed bounded, the density p(x) has compact support, and the density p(x) is assumed continuous, there exists R. > 0 such that if ||x||2 > R then p(x) and M > 0 such that p(x) & M. We then have
sup Poa ly - x))((x) S M sup 6aa (y 1x x
This means that the supremum
attain nonzero value only points where ||*||
R. On
the other hand, for every y with ||l| 2 2R, whenever ||x|| S R the triangle inequality implies lly x|2 > || la ||||| 2 |y||-- Because the Gaussian density Paa is radial function,
conclude that |2yl2 2 2R, one has
sup Poa (y x)p(x)= S MMra(y/2) CM)4
where C > 0 depends only on At the same time, we always have
M
sup a (y - x))(x)= S
(2TT2)
Consequently, have the composite upper bound
M
llyl2 < 2R
sup Poaly x))(x)= < 2M (MMpso=(y)
(2
Ilyl|2 2 2R,
x
and by our work above
M 12
||yl| 2R
Exly [P/4/2(f(y) - xx] s (goz * p))
x
2M (.Maaaa(y)) (y)
||y|| > 2R.
Because 4p is the density y,, this upper bound is sufficient to apply the dominated convergence theorem to obtain
+ppis
lim Exy [x)))2((())) x)] By lim Exiy (p x)] 7'50
Combining this assertion with the argument surrounding (C.4), we conclude the proof.
Remark (Other loss choices). Theorem 3.2 also holds for any m,, such that m. is uniformly (in 7) bounded above, for each 7 0 uniquely minimized at 0, and supaeR (x)x) (x(()))2)) is an approximation to the identity as y > 0 (see (Stein & Shakarchi, 2005, Ch. 3, $2)).
C.3 PROOF THEOREM B.
Proof. For brevity, denote argmax, P(x e y) by MAP[x y), i.e., the maximum a posteri- ori estimate of x given y.
First, show that MAP(x | unique for almost all y.
22

---

Published as a conference paper at ICLR 2024
Consider y such that MAP x y] is not unique. There exists + j. such that
P(x: y) = P(x; y)
- ply x P(x:) ply px)))
ly x + o P(xi) = -- 1 xz)l + o² log P(x;)
ly, x; xj | 1| || 112 o² log P(x,) o² log P(x;).
i.e. y lies in a hyperplane defined by X x, (note that x, + xj). Denote the hyperplane by
Has =
(y, x; x
= 2
| xx, 1| o² log P + x² log P(xj)
2
Consider
u = Uits Hi.s
We have that Vy with non-unique MAP(x |
y).
Ei j.y Hi.s
- ycu.
Note that 4/ has zero measure countable union of zero-measure sets, hence the measure of all y with non-unique MAP(x yy] is zero. Hence, for almost all y. MAP(x y] is unique.
Next, show that for almost all y.
(y) = argmin Exly [exx]]
Note that
lim Ex,y [m, (Ilf(y) x11
7 70
-Ex.y
lim my (llf(y) - x112)
(>00
-Exyy
[11 [11
=Ex.y [xx)))x]]
Above, the first equality uses the monotone convergence theorem. Use the law of iterated expecta- tions,
Exy [1 = Ey [1 We will use this expression study the global minimizers of the objective. By conditioning. Exly [1 [[vja]] 2 min [xiy(_exxx,
and so
Ey Exly 11 sxx min ( 2
Because ply) > 0, it follows that every global minimizer of the objective f satisfies
Exly [1 = min Exyy [xxx]] a.s.
Hence, for almost all y,
(")) E argmin Exly [c
Finally, show that argmin, Exly [eapx MAP(x y]. The claim then follows from preceding work showing that MAP x y) is almost surely unique. Consider
Exly [ cxx]
Σ P(x; l)
Σ P(x l
Σ P(xi 1y) Σ P(xi 1x) xi=c P(x c(x=cly).
23

---

Published as a conference paper at ICLR 2024
Hence,
argmin Exly|" = argmax P(x = ccl)
= MAPP[x yl.
C.4 PROOFS OF PNP OPTIMIZATION RESULTS
In this section, restate and provide proofs of Theorem B.2 and Theorem 4.1. We prove Theo- rem B.2 under slightly more general assumptions, and state the conclusions both Theorems 4.1 and B.2 with more precision. The restated results are given below, as Theorem C.1 and Theorem C.4.
Before proceeding to proofs, let us briefly describe the common high-level 'recipe" underlying each plug-and-pie algorithm's proof. The recipe separates into two distinct steps:
1. Leverage general, black-box convergence analyses from the optimization literature. A plug-and-play algorithm is derived from a 'baseline' optimization algorithm; we there- fore appeal to convergence analyses from the literature of the relevant baseline algorithm. Because the regularization function associated to LPN is implicitly defined by the LPN architecture and need not be convex, it is necessary to appeal to general, 'black-box' con- vergence analyses which do not leverage special properties of the regularization function. We make use of convergence results nonconvex proximal gradient descent of Bot et al. (2016), and on nonconvex ADMM of Themelis & Patrinos (2020). To make the presen- tation self-contained, we reproduce key results from these works in context. The principal technical activity is therefore to translate the iterate sequence generated by the relevant PnP algorithm into a form that allows these convergence analyses be applied to it. Echoes of the same approach appear in prior work on convergent plug-and-pe and-play, for example work of Hurault et al. (2022b).
2. Establish general regularity properties of the regularization function associated to LPNs. To appeal to the aforementioned convergence analyses, it is necessary to ascertain a minimum level of regularity of the regularization function associated to an LPN, order to establish that it possesses the mandyka--.ojasiewic. (KI.) property (and, say, coerciv- ity). We give a self-contained overview of the KI. property and how we establish it in Appendix C.4.4 for clarity of presentation. We provide in Appendix C.4.3 technical lem- mas that establish that LPNs of the architecture specified in Proposition 3.1 satisfy these properties, regardless of the exact values of their parameters. These results are essentially consequences of differentiability and surjectivity of the LPN when < a < 1 is used as the strong convexity weight, and they enable to assert convergence guarantees for LPNs without any extra assumplions about the trained network.
We anticipate that this recipe will applicable to virtually any PnP scheme which there exists a convergence analysis under the KL property of the corresponding baseline optimization algorithm. Because our technical work in Appendix C.4.3 establishes the KL property and cocrcivity for the regularization function associated to LPNs with the architecture of Proposition 3.1, obtaining convergence analysis for such a PnP scheme with LPNs of this architecture only requires the first step of the above recipe. We expect approach in Appendix C.4.3 to extend straightforwardly to LPNs with novel architectures-for instance, different computational graphs or weight-sharing schemes-as long as the nonlinear activation functions do not grow rapidly (see the proofs for more precise statements).
C.4.1 PROOF OF THEOREM B.2 (PNP-PGD)
Theorem C.1 (Convergence guarantee for running PnP-PGD with LPNs). Consider the sequence of iterates Xs, (0,1....), }. surement operator A and an LPN fa with softplus activations, trained with 0 < a < 1. Assume
defined by Algorithum 4 run with a continuously differentiable mea-
The form these results are stated in makes them most convenient for purposes of our presentation, although the result need is originally due to (Attouch al., 2013).
24

---

TTii is very mild assumption. For example, when A is linear, the gradient of the data fidelity term V denotes the operator norm a linear operator and Published as a conference paper at ICLR 2024
further that the data fidelity term h(x) = 21lly - A(x)11 is definable in the o-minimal structure of Proposition C.II, Property 2 and has T.-Iipschitz gradient") and that the step size satisfies 0 < < L/L. Then, the iterates X2 converge to fixed point x' of Algorithm 4: that is, there exist x E I such that
fo (x" - -n/h((')) =x';
(C.6)
and lim Xx x'. Furthermore, x" is a critical point of h + RR, where Rg is the regular- ization function ossociated to the LPN So (i.e., So prox no
Before proceeding to the proof, we state a few settings and results from Bot et al. (2016) that are useful for proving Theorem C.1, for better readability.
Problem 1 ((Bot et al., 2016, Problem 1)). Let f R - (-20, +oo] be a proper, lower semicon- tinuous function which is bounded below and let h 12 + R be a Fréchet differentiable function
with Lipschitz continuous gradient, iee. there exists Lvh 2 such that ||Vh(x) Vh(x')ll S
Lval[x- x' for all x, x E RR. Consider the optimization problem
(P)
inf [f(x) + h(x)).
XER"
Algorithm C.1 ((Bot al., 2016, Algorithm 1)), Choose Xo, X1 R a, ã > 0, Z 0 and the sequences (0n)2211 Iulfiinin
fulfilling
vassa Vn
and
0 s Bn S 1121.
Consider the iterative scheme
(Vn 2 1) Xn+1 E argmin
(D,(U, (rgmin(Dp(U,x.)) +ap (U, Vh(x,)) B,, (nnx111-- +an (C.7)
Here, F R. + R is o-strongly convex, Fréchet differentiable and V is F continu- ous, with o; Lvp 0; Dr is the Bregman distance to F.
Theorem C.2 ((Bot et al., 2016, Theorem 13)). In the setting of Problem 1, choose a. 3 satisfying
o > LLa, + + a
(C.8)
a
Assume that f +++ h is coercive and that
H : RR x R + (-o0, tool, H(x = (f + (x)) + ||x x112, v(x,x') E Rm x R
2a
is KL " Let (X,)nEN be sequence generated by Algorithm C.I. Then the following statements are true:
1. ΣmEN |x++1
Xnll < +00
2. there exists x E crit(f + h) such that lim Xn
Now, we prove Theorem C.1.
This mild technical assumption is satisfied by an extremely broad array of nonlinear operators A: for example, any A which is polynomial in the input (in particular, linear A) is definable, and compositions inverses of definable functions are definable, so that definability of implies definability of h. See an extensive overview these ideas in Appendix C.4.4
has Lipschitz constant no larger than ||A" All. where A is the adjoint A.
In this work. the set of critical points of function is detined by crit(f) == {x 0 E @f(x)), where 0 the limiting (Mordukhovich) Fréchet subdi fferential mapping of (see definition in (Bot et al., 2016, Section 2)).
In this work, function being KI. means it satisfies the Appendix C.4.4 Definition C.I.
property (Lojasiewicz,, 1963),
25

---

Published as a conference paper at ICLR 2024
Proof of Theorem C.I. By Lemma C.7, there is coercive function Rg RD RU{+oc} such that fg = proX The idea the proof is to apply Theorem C.2 to our setting; this requires us to check that Algorithm 4 maps onto Algorithm C.I, and that our iiplicitly-definees objective function and parameter choices satisfy the requirements of this theorem. To this end, note that the application of fo in Algorithm can be written
Xk+1 fo (xx
1 Vh(xx))
argmin
|x - -xx (1V(x)V(( Ro(x')
2
x eRRR"
1 ||x - xx:l2² + xx, (x:xxx)) + Rx(x')
argmin x'ER"
2
xx xx11 n(x', Vh(xx)) n. 1 R
argmin x'ER"
||x
2
1)
showing that Algorithm corresponds to Algorithm C.I with the Bregman distance Dx = a d||x - y1 (and correspondingly Fxx) = |
parameter B B,, = 0, the step size on = x which satisfies = LVF = 1), the momentum 77, and == Ro In the framework of Bot et al. (2016), Algorithm 4 minimizes the implicitty-defined objective h + 1 Rg. Moreover, one checks that our choice of constant step size 0 < 17 1/L verifies the necessary condition (C.8), and because/ 2 O, coercivity of RG implies that h + 1 ( is coercive. Using Lemma C.13, we obtain that Rg is definable, and by assumption, h is also definable, that by Proposition C.I1. Properties 2 and 5, it follows that the objective /t 41+-7 +1- Rg is definable. Thus h- Rg is definable, continuously differentiable (by Lemma C.7), and proper (as a sum of real-valued functions. again by Lemma C.7), and therefore has the KL property, by Proposition C.11, Property 1. We can therefore apply Theorem C.2 to conclude convergence to a critical point of h+ Rg. Finally, by Lemma C.3 and the continuity of fo and Vh, we conclude convergence to fixed point, x- fo(x v)((()), which is identical to (C.6).
Lemma C.3 (Convergence Implies Fixed Point Convergence). Suppose F RR - Rn is contin- uous map that defines iterative process, Xt+1 F(xx). Assume Xx converges, i.e., 3 x" such that lim X2+ = x". Then, x" is a fixed point of F, ie.. x* = F(x*).
Proof.
x lim Xx lim X
lim F(xx) =
lim X F((x). k0000
1+000
k->00
k0+00
The fourth equality follows from continuity of F.
C.4.2 PROOF OF THEOREM 4.1 (PNP-ADMM)
We present in this section a proof of convergence for PnP-ADMM schemes which incorporate an LPN for the regularizer (Algorithm 1), following the recipe we have described in Appendix C.4. These guarantees are analogous to those we have proved in Theorem C. for the PnP-PGD scheme Algorithm 4 with LPNs. For simplicity, we will assume in this section (in contrast to the more general setting of Theorem C.I, and in agreement with the result stated in Theorem 4.1) that the measurement operator in the underlying inverse problem (2.1) is linear and acts in the standard basis, and accordingly we identify it with its matrix representation A. This means (among other things) that the data fidelity term x + lll - Ax| 2 convex.
Before proceeding to the proof, we note that Algorithm adopts an update ordering which is non- standard in the signal processing literature (c.f. (Venkatakrishhann et al., 2013; Kamilov et al., 2023a)) for technical reasons. We employ this update order due to its prevalence in the optimization liter- ature, notably in the analysis of Themelis & Patrinos (2020), and we emphasize that all our experiments are done following Algorithm 1. Although both the typical PnP- update order and the update order in Algorithm correspond to an ADMM algorithm for the same objective function, the analysis of these two iterative optimization procedures is different, and seems quire different technical assumptions (c.f. (Themelis Patrinos, 2020; Yan & Yin, 2016)). Prior work convergent PnP seems have also run into this barrier, suggesting it is not artifact of
26

---

Published as a conference paper at ICLR 2024
analysis: for example, Hurault et al. (2022b) study PnP variant of Douglas-Rachford Rachford splitting rather than ADMM, which is roughly analogous to the reversed-order of updates in Algorithm by reduction of Themelis & Patrinos (2020), and Sun et al. (2021) prove convergence of sequence of associated residuals in the standard-order PnP-ADMM rather than of the sequence of iterates itself.
Our proof will be based the work of Themelis & Patrinos (2020), which provides guarantees for ADMM in the nonconvex setting. We restate some of their results for convenience after stating our convergence result, then proceed the proof.
Theorem C.4 (Convergence guarantee for rumning PnP-ADME with LPNs). Consider the sequence of iterates (Xx: Ug, 3x), k E (0,1.... defined by Algorithm run with linear meusurement operator A and a LPN So with softplus activations, trained with a 1. Assume further that the penalty parameter p. satisfies
|| 4 All Then the sequence of ieeatees (Xk: Ug 2x) converges to a limit point (x", u', , z²) which satisfies the KKT conditions (of the augmented problem):
x" zz, u AA yy)
(C.9)
(Ro(2"),
where Rg is the regularization function associated to the LPN fo (i.e., fo prOXRsh, which is continuously differentiable. In particular, the primal limit critical point of the regularized reconstruction cost x + 2 b Ax + pRo(x), and the full limit iterate (x", uu, 2²) is fixed point of the PnP-ADMM iteration (Algorithm 1).
We restate convergence results of Themelis & Patrinos (2020) in lesser generality, given the addi- tional regularity properties present in our setting of interest.
Problem 2. Let h1 R I and h2 minimization problem
RR R be continuously differentiable. We consider the
min h (x) h2(x). xEX"
(C.10)
Algorithm C.2 (Themelis & Patrinos (2020, Equs. (1.2), (ADMM), (1.3))). Perform variable split- fing (C.10) to obtain equivalent problem
min
bu (x) + l2(22) S.L. x - 2 =0 0.
(C.11)
xER"
Fix > 0. Form the augmented Lagrangian (C.11) at level p. that is, the function
(x(x, 2. hx() h2(z)) + (y, 2) + P ||x z 2
(C.12)
2
and consider the following iteration."" 14;
x E argmin Lpl ,2, y),
y y p(xt z),
(C.13)
z E argmin Cp(xt, ,,
This iteration induces a set-valued map To R. x R" = R x R. x R". Given an initializa- tion (yo, a sequence of ADMM iterates (x,, Yk, kk,%% 2x) defined inductively by (Xi. Yk. 2x) E T, 74-1), k E 14.
Themelis & Patrinos (2020) provide the following convergence guarantee for Algorithm C.2 relative to the objective (C.10), under weak assumptions on h1 and h2:
"This corresponds setting the "relaxation parameter" À in Themelis & Patrinos (2020, Eqn. (ADMM))
The variable y defined here follows the notation of Themelis & Patrinos (2020), and in particular should not be confused with the measurements in the inverse problems framework. We hope the reader will forgive this conflict of notation.
27

---

Published as a conference paper at ICLR 2024
Theorem C.5 (Themelis & Patrinos (2020, Theorem 5.6, Theorem 4.1, Theorem 5. Sup- pose the objective h1 + h2 is coercive;" 17 that h1 continuously differentiable, its gradiens Vh, is L-Lipschitz, and there exists proper and lower semicontinuous. Moreover, suppose the penalty parameter is chosen so that p max 0), L}. and that the augmented Lagrangian (C.12) is Kl. function (see Ap- pendix C.4.4, Definition C. Then the sequence (Xk, Yk, kE)kE converges to limit point (x", yy, ,2") which satisfies the KKT conditions
R such that h 911 11 is convex; and h2 is
y Vh1(x")
y 0h2(z")
O.
Here, dhy is the (limiting) subdifferentaal mapping of h2z More concisely, the limit satisfies O (x1(x") + 0h2(x*), and in particular x" is a critical point for the objective (C.10).
It is standard to argue that Algorithm corresponds to Algorithm C.2 up to a simple reparameteri- zation (i.c., relabeling of variables), given that the LPN fe in Algorithm is a proximal operator.
Lemma C.6. An ADMM sequence (Xk, Vk, zk) generated by the update (C.13) is linearly isomor- phic to a sequence generated by the update rule
x+ E prox
1/ (z - u)
u = u (xt - 2)
+
(C.14)
zt E prox 1/2 (xt + +ut),
z in the sense that if (x4, u/g, Z/) is the corresponding sequence of iterates generated by this update rule with initialization Za 2n and un J,y, then we have Xx. = X/y, Z₁ zzy and u, y for every ks E No-
Proof. Notice that we can write in (C.12) by completing the square
2
2
(,(x, hx(x) + h2(z) + P y p
p
(x 2)
(C.15)
2
'In obtaining the result stated here from Themelis & Patrinos (2020, Theorem 5.6), we simplify the "image function" expressions (Themelis & Patrinos, 2020, Definition 5.1) using the simple constraint structure of the ADMM problem (C.11): in particular, in checking (Themelis & Patrinos. 2020, Assumption II), we have that (s) infxEx= (x)
s} les( (-s). In particular 101 h1 and (o 'u Id, so that (Themelis Patrinos, 2020, Assumption IL.A4) is implied by (Themelis & Patrinos, 2020, Assumption ILAA)). This also allows us to translate the convergence guarantees of Themelis & Patrinos (2020, Theorems 5.6, 5.8) from applying to the sequence of iterate images under the constraint maps to the sequence of iterates themselves.
8} fu (s) and similarly (22(s) infæeæ= (1v(2)
This implics that the objective function of the equivalent penalized version (C.11) level bounded and admits a solution (recall that the latter a consequence of the Weierstrass theorem, e.g. (Bertsekas, 2016, Proposition A.8(2)))
18 Such a a always exists, and satisfies |a| S L. Roughly speaking, smaller a value can chosen, the better-this is possible when h2 is 'more convex'.
19 Althouggh Themelis & Patrinos (2020) state their global convergence result, Theorem 5.8, only in the semialgebraic setting, inspection their arguments (notably (Themelis Patrinos, 2020, p. 163 top), and connecting discussion in (Li & Pong. 2016, Theorem 2, Remark 2(ii)), together with the equivalence between
DRS and ADMM in (Themelis & Patrinos, 2020, Theorem 5.5)) reveals that it is only necessary that the augmented Lagrangian C, is KL function.
20This is the same notion of subdifferential introduced in Appendix C.4.1 in order to state the results of Boj et al. (2016). We use the fact that the limiting subdifferential coincides (up to converting vanuue set-valued map into function) with the gradient for C function (Rockafellar & Wets, 1998, Theorem 9. 9.18, Corollary 9.19).
28

---

Published as a conference paper at ICLR 2024
in order to simplify the minimization operations in (C.13). Indeed, the iteration (C.13) then becomes equivalent to, with (C.15), the iteration
x E prox
z y
y = y + p(xt - z)
(C.16)
z + E proxx 1/2 (xt + 1 pxrt))
Introducing now the "scaled dual variable"
y,,, we have the equivalent update
x E prox h, (z - u)
u (xt - z)
(C.17)
xt prox f (xt + +u))
The equivalence claims in the statement of the lemma follow from this chain of reasoning.
L
With this preparation completed,
ready to give the proof of Theorem C.4.
Proof of Theorem C.4. Below, to avoid a notational conflict with the dual variables in Algo- rithm C.2, will write Ymeas for the measurements that define the data fidelity term in the inverse problem cost.
We have by assumption and Proposition 3.1 and Lemma C.7 that there exists a coercive C function
Rg : 1 + IR such that fe = proX Given the initialization uo O in Algorithm 1, applying Lemma C.6 implies that Algorithm 1 corresponds to a sequence of ADMM iterates (xx, Yk. Zz generated via Algorithm C.2 with the initialization yo O, the objectives hi (x) lly menss Ax and (2( pfg (z) and the correspondence yı pug. In addition, we observe that h; is smooth, nonnegative. convex, and satisfies /
| A where || denotes the operator norm. In turn, we know that z + h2(z)) + E is (strongly) convex: Lemma C.10 implies that fo is Lipschitz, and (Gribonval & Nikolova, 2020, Proposition 2(1)) implies that L-Lipschitz proximal operator's associated prox- prox-primitive is (1 - (1-1L)--wakij convex, from which it follows that the sum z r h2(z) 8 | is strongly convex. As result of these facts, every minimization operation in Algorithm C.2 has unique minimizer, and can interchange set inclusion operations with equalities when describing the ADMM sequence corresponding to Algorithm without any concern in the sequel.
Now, these instantiations of h1 and h2 verify the elementary hypotheses of Theorem C.5:
Nonnegativity h and coercivity of h2 imply that h1 hg coercive;
2. h is smooth and we take L |ATA| and 0, and therefore the hypothesis that p > ||A'||| verifies the conditions on the penalty parameter;
3. By Lemma C.7, h2 is real-valued and C² as above.
Finally, to check the KL property of the augmented Lagrangian Lp defined by (C.12), will follow Appendix C.4.4 and verify that Cp is definable in an o-minimal structure, then apply Propo- sition C.II, Property I, since Cp is C as a sum of C functions (both h and fg are so). To this end, note by Corollary C.12 that Cp is definable in the o-minimal structure asserted by Property 2 if both bn and hy are definable in that o-minimal structure. Proposition C.11, Property implies that h is definable, since is degree two polynomial. Then Lemma C. implies that he is definable in the same o-minimal structure, and thus follows from the preceding reasoning that C, has the KI. property.
29

---

Published as a conference paper at ICLR 2024
We therefore apply Theorem C.5 to obtain that the sequence of iterates of Algorithm 1 converges, and its limit point (x", u², z²) satisfies the KKT conditions (C.9)
u , "Ax" Ymeas),
u (R(=))) x - O,
where we have used fact that the limiting subdifferential coincides (up converting single-valued set-valued map into functiion) with the gradient for a C function (Rockafellar & Wets, 1998, Theorem 9.18, Corollary 9.19). In particular, simplifying gives
AA - -mmass) = - - Rx
(C.18)
which equivalent to the claimed critical point property. To see that this is also equivalent to being fixed point of Algorithm 1, it is convenient to the expression (C.17) for the ADMM update expression, which appeared in the proof of Lemma C.6. The KKT conditions (C.9) imply that Vxi(x*) Vh2(x*), and since both 1/1 and hh2 are differentiable and yield a strongly convex function when summed with a quadratic 211 112, we can express the action of their proximal operators
prox p (Id + Wh.) v -1
1,2.
(C.19)
From (C.17), we get that x+ (Id (ld+ -1 (z" -u)) The KKT conditions (C.9) imply that u' Vh (x") and z² = x", so that z uu = (Id (/:)))))(x')) and therefore indeed
x* Proceeding, we then have that u u² =z" + xx, which gives that u u', since x* = zz. Finally, we obtain from the previous two steps and the definition of h2 that
z = (Id VV)) + u²),
(C.20)
and the final KKT condition (C.9) implies that u (Zo Hence
= (Id+VRg) (Id +V RR) (z)) = z",
(C.21)
and we indeed conclude that (x", ,,, z²) is a fixed point of Algorithm 1.
C.4.3 REGULARITY THE REGULARIZATION FUNCTION LPNS
As discussed in the "recipe" of Appendix C.4, in this section we prove basic regularity properties of the regularization function associated to any LPN with the architecture of Proposition 3.1.
Lemma C.7 (Regularity Properties of LPNs). Suppose fo is LPN constructed following the
recipe in Proposition 3.1, with softplus activations (z) = (1/3) log(1 + exp(3x)). where 3 > an arbitary constant, and with strong convexity weight 0 < 1. Let foly) Vubo(y) ay be the defining equation of the LPN. Then there is real-valued function Rs : R - such that f9 pr Moreover, we have the following regularity properties:
1. Rg is coercive, i.e., we have Rx(x) + +o0 as ||x||2 + +
2. - is surjective and invertible, with inverse mapping fo R R which is continuous.
3. Rg is continuously differentiable. In particular, it holds
Rø(x) (1 (x), Violfo (x))) + a(1 a)
(C.22)
115o (x)11 - (x)).
Remark. Lemma C.7 does not, strictly speaking, require the softplus activation: the proof shows that any Lipschitz activation function with enough differentiability and slow growth at infinity, such as another smoothed verison of the ReLU activation, the GeLU, or the Swish activation, would also work.
30

---

Published as a conference paper at ICLR 2024
Proof of Lemma C.7. The main technical challenge will be to establish coercivity of Rg, which always exists as necessary. by Propositions 2.1 and 3.1. We will therefore pursue this estimate as the main line of the proof, establishing the remaining assertions in the result statement along the way.
By Proposition 3.1, there exists Rg such that fo = proxn, Now, using (Gribonval & Nikolova, 2020, Theorem 4(a)), for every y E R,,
Ro(So(y)) - (y, fo(y)) 3 (y) - (voly) + ()y11
Using the definition of fo and minor algebra, we rewrite this as Ro(foly)) (y, Veloy) ay) 1 (V + aayl (vo(y) +
= -1- (y (y) + a(1 a 1yl V (y)
(C.23)
At this point, we observe that by Lemma C.8, the map So R + R" is invertible and surjective, with a continuous inverse mapping. This establishes the second assertion that we have claimed. In addition, taking inverses in (C.23) implies (C.22) and as a consequence the fact that Rg is real- valued, and the fact that it is continuously differentiable on I" is then an immediate consequence of (Gribonval & Nikolova, 2020, Corollary 6(b)). To conclude, it only remains to show that Rg is coercive, which we will accomplish by lower bounding the RHS of (C.23). By Lemma C.9, vh is L--ipsshitz for constant L > 0. Thus, we have for every y (by the triangle inequality)
[bb S L + K
for a (finite) constant E R. depending only on 0. Now, the Cauchy-Schwar inequality implies from the previous two statements (and || cbll22 S L by the Lipschitz property of vha)
ax1 a)
Ro(fely)) 2 -(1 a)1 V/ (y) +
y [ - L - K,
2
a(1 a)
1
- a alllylaa +
y
L K.
We rewrite this estimate with some algebra as
a<11 a)
I
Ro(fo(y)) 2 lyyl|
a)) - a1
|yyle
2 - K.
2
Next, we notice that when 0 < a < 1, the coefficient axl - a) > 0; hence there is a constant
M > 0 depending only on and L such that for every y with |y|| 2 M, one has
a a:) |yl12 L(1 - - L>
a(1 a)
4 lyl12.
In turn, iterating this exact argument implies that there is another constant M' > 0 (depending only on a, L, and K)) such that whenever |yyl 2 2 M', one has
Rs Ro(fe(y) a(1 a) 1ly
8
We can therefore rewrite the previous inequality
Rx > a
a(1-a) 115 (x)
(C.24)
8
for every x such that (x)]l2 M'. To conclude, we will show that whenever ||x||2 - +o0, also have | +oc, which together with (C.24) will imply coercivity of Rg. To this end, write || ||tip for the Lipschitz seminorm:
llf(y) - f(y'))2
Il = sup
Ily - -y'll
yfy'
and note that I < / Ltpp +a. By Lemma C.10, Velu is Lov, Lipschitz continuous, thus fo is (Lve, + a)-iipchhitz continuous,
[)o So(y')ll (Ivo, +allly y'1l2.
31

---

Published as a conference paper at ICLR 2024
Thus, taking inverses, we have
1
fo Lo |xx|12
and it then follows from the triangle inequality that whenever x such that ||*||2 2 2(L + a) life (0) we have in fact
1
llfe ( (xx]12 2 2(L +a) xx12..
Combining this estimate with (C.24), we obtain that for every x such that ||x||e 2 2(L + a)llfa and ||x||2 2 2M + ax), it holds
a(1
(x
32(L + +a) 1x1
Taking limits in this last bound yields coercivity of Rø, and hence the claim.
Lemma C.8 (Invertibility of fe and Continuity of fo')) Suppose fe an LPN constructed
the recipe in Proposition 3.1, with softplus activations (xx) (1/3) log(1 + exp(3x)), where an arbitrary constant, and with strong convexity weight 0 < < 1. Then so R2 Tk is invertible and surjective, and fg : Rd is Co.
Proof. The proof uses the invertibility construction that describe informally in Section 3. By construction, we have fo = V + a Id, where Id denotes the identity operator on I" (i.e., Id(x) for every R.).
For fixed x E IR", consider the strongly convex minimization problem miny vb((y) + 1 y). By first-order optimality condition, the minimizers are exactly {y V + ay x). Furthermore, since the problem strongly convex, has a unique minimizer for each x E Ik" (Boyd & Vandenberghe, 2004). Therefore, for each x R", there exists unique y such that x (y)(y) + ay = fe
The argument above establishes that R R is injective and surjective; hence there exists an inverse fo R" + R". To conclude the proof, we will argue that fø is continuous. To this end, we use the characterization of continuity which states that function g : E= + R" is continuous if and only if for every open set U C. I2", we have that g (U) is open, where (U) {x E RR" g(x) E U} (e.g., (Rudin, 1976, Theorem 4.8)). To show that fg is continuous, is therefore equivalent to show that for every open set U C R"" one has that (U) is open. But this follows from invariance of domain, standard result in algebraic topology (c.g., (Dold, 2012, Proposition 7.4)), since fo is injective and continuous. We have thus shown that so is invertible, and that its inverse is continuous, as claimed.
Lemma C.9 (Lipschitzness of vha). Suppose fo an LPN constructed the recipe in Propo- sition 3.1, with softplus activations (xx) = (1/3) log(1 + exp(Bx)), where 3 > 0 is an arbitrary constant, and let the denote the convex potential function for the LPN. Then the is Lgg Lipschitz continuous for a constant Ligs > 0, i.e., (v) (y())) & Los lly y' 12, for all y,y' E RR..
Proof. Note that the derivative o' of the softplus activation satisfies (x) = 1/(1 exp(-Br)) which is no larger than 1, since exp(x) > 0 for æ E IR. If F map between Euclidean spaces we will write DF for its differential (a map from the domain of F to the space of linear operators from the domain of F the range of F). Hence the activation function g in Proposition 3.1 is 1-Lipschitz with respect to the Eg norm, since the induced (by elementwise application) map R + R defined by g(y) = [o(x)...... (xr)))" satisfies
T (
Dgly)
a' (In)
32

---

Published as a conference paper at ICLR 2024
which is bounded in operator norm by supx 1o (x)] 1. First, notice that
(1vo(y) vo('')l2 | (w'(ze (Zx (y) zK 1l22
- |||ll||/<()) 2x(22)12
by Cauchy-Schwarz Meanwhile, we have similarly
2z (y) - z1 (y') I |H - y2ll,
where || denotes the operator norm of a matrix, and for integer < l K +1
= + H Illyy y'1l2- By a straightforward induction, it follows that the is L-Lipschitz for constant 0 (depending only on 0).
[zx(y) - z (y')ll < || [Wlll|z--100
Lemma C.10 (Lipschitzness of Viby and LPNs fo). Suppose fo is an LPN constructed follwnntte the recipe in Proposition 3.1, with softplus activations o (xx) (1/3) log(1 + cxp(Bx)), where B 0 is an arbitrary constant, and with strong convexity weight 0 < a < 1. Let foly) Vihaly) + ay be the defining equation of the LPN. Then Vijo is -ipp-hip continuous, for constant Lowe O. In particular, Se is (ax -Lipschitt continuous.
Proof. The claimed expression for Lipschitzness of fe follows from the claimed expression for Lipschitzness of Vrhe. by differentiating and using the triangle inequality. We recall basic notions of Lipschitz continuity of differentiable mappings at the beginning of the proof of Lemma C.9, which will make use of below. We will upper bound Vtolltap by deriving an explicit expression for the gradient. By the defining formulas in Proposition 3.1, we have
voly) = w zK (y) + b.
The chain rule gives
Vvaa() = DzK Dzk(y)² w, where T denotes the adjoint of linear operator, so for any y, we have
[|Voo(y)- Vobo(y')l22
(Dzk (y) Dzk(y'))² w
< (DzK (y) D2x(y')) ||
[DzK(l) DZK (y") I1 ///2
I [)zz(() - DZK (yzk(y")) || l/w||2,
where the first inequality uses Cauwhy-Schwarz, the third line uses that the operator norm of linear operator is equal to that of its adjoint, and the third line uses that the operator norm is upper-bounded by the Frobenius norm. This shows that we obtain a Lipschitz property &g for Velo by obtaining one for the differential Dzæ of the LPN's last-layer features. To this end, we can use the chain rule to compute for any integer 1 <
ks K + 1 and any E R.
Dz of (Wazs- n (y) H&y ba) (y) + Hgō),
where g' is the derivative of the softplus activation function g, applied elementwise, and denotes elementwise multiplication, and similarly
Dz (5) g' (Hy H,)++b)) H
Now notice that for any vectors and y and any matrix A such that the sizes compatible, have v O (Ay)- = diag(v)Ay: Hence we can rewrite the above recursion in matrix form as
(y diag (g' (Wazz-1 + H&y bx)) (WDDz-i(y) + Hz];
D.())
and similarly
Dz, (y) = diag (g' (H₂y + H,
Dy(y)
33

---

Published as a conference paper at ICLR 2024
We will proceed with inductive argument. First, by the submultiplicative property of the Frobe- nius norm and the triangle inequality for the Frobenius norm, note that we have 1 < k < K +1
[ (y) S [Dx(y) Dl
DE(y') DD
[x(y) - (yEE
(
[[Dd())WDDzx-1(x)) - DN(y') D
Dx(y) - DIE
+ |IDx(y)Wallr (iz)) D2a-1
+ -1 (r))>////////()(()))))))))) - Da(r')Walle
$ (1 + (Alll)) [Dx Da(v'IlE
+ |[|xx(y)]=/|WWa D/ (y) - D
Now, we have shown above, g(xx) (1 + expl Bxx) -1 ~ for every R. This implies
[Dx S V
where me the output dimension of k-th layer. Moreover, we calculate with the chain rule
Be- Aa
gx(x) +1+ Amg2
and by L'Hôpital's rule, we have that lim (1+z)² = O, so that by continuity, g" is bounded for
R. It follows that g' is Lipschitz. Notice now that
[D) (yy) = Ilg' (Waze-1
H&Y + b&) - g (We W (yy) + H&y + ba)
I l1 I|
1z=- + Hallelly - y' y2l)2.)
where in the second line we used the fact that the derivative of an elementwise function is a diagonal matrix together with the triangle inequality and Cauchy- -chwwrr. However, we have already argued previously by induction that th is Lipschitz, and in particular each of its feature maps 25. is Lipschitz. We conclude that D is I. Lipschitz, and the Lipschitz constant depends only on O. This means that there are constants LE, Li depending only on and 0 such that
[y) (zx(y')llF S Lally - y'll2 + LA Dz) - Dzu (y))llF.
Meanwhile, following the same arguments as above, but in slightly simplified setting, we obtain
(zz Dz (y))lF=
||D (Di(y) H D₁ (y
< [ D (y) - Dy(r')l llg Lpp
H l
which demonstrates that Dz is also Lipschitz, with the I Lipschitz constant depending only on 0. By induction, we therefore conclude that there Tr > 0 such that
(y))) Via 12 < Lo lly y' 112
with Lve, depending only on 0 and ni.-
C.4.4 NONDEGENERACY OF THE PRIOR
For our convergence results for PnP with LPNs, we make use of general convergence analyses for nonsmooth nonconvex optimization from the literature. To invoke these results, we need to show that the regularizer Rd associated to the LPN fo satisfies a certain nondegeneracy property called the Kurdyka-E.jaccieee Lojaciewizz (KL) inequality. We provide a self-contained proof this fact in this sec- tion. We will instantiate various concepts from the literature in our specific setting (i.e., sacrificing generality for clarity), and include appropriate references to the literature for more technical aspects that are not essential to our setting.
34

---

Published as a conference paper at ICLR 2024
Definition C.1 (KL property; (Attouch al., 2010, Definition 3.1), (Boj al., 2016, Definition I)). A C' function f R" - IR is said to have the KL property (or to a KL function) at a point R" if there exists 0 1 +o0, neighborhood U с R" of Σ, and a continuous concave function :p: [0, v) + R such that
1. (0)) 0;
2. yp 2 0:
3. i is C' on (0,
4. for all (0,n), ('(s) >0, 0:
+ 7) the Kurdyka--ojaciewiiz inequality holds: 5. for all Euf (()) [(x) < (
|V/(2)|2 2 ('((x)- f()))
(C.25)
can be shown that any f E (R") has the KL property at every point T that not a critical point of f (Attouch et al., 2010, Remark 3.2(b).. At critical points E of such function f. the KL inequality a kind of nondegeneracy condition on f. This can be intuited from the (common) example of functions f witnessing the KL property via the function s(s) cs*/2, where > constant: in this case, (C.25) gives lower bound on the squared magnitude of the gradient near to a critical point in terms of the values of the function f. Every Morse function satisfies such an instantiation of the KL property: see (Attouch et al., 2010, beginning of $4). In general, the KL property need not require excessive regularity of f; we only state as such because of the special structure in our setting (c.f. Lemma C.7).
The KL property is useful in spite of the unwieldy Definition C.1 because broad classes of functions can be shown have the KL property, and these classes of functions satisfy convenient calculus that allows one to construct new KL functions from simpler primitives (somewhat analogous to the situation with convex functions in convex analysis). The class of such functions work are functions definable in an o-minimal structure. We refer to (Attouch et al., 2010, Definition 4.1) for the precise definition of this concept from model theory-for our purposes, it will suffice to understand that an o-minimal structure is collection of subsets of IR.", for each n E N, which satisfy certain algebraic properties and are called definable sets, or definable for short, and that definable function is one whose graph is contained in this collection (i.e., its graph is a definable set)al-and instead state several useful properties of functions definable an o-minimal structure that will suffice to complete our proof. We note that Attouch al. (2010, $4) give excellent readable optimization-mann motivated iontiizationmmoiivateedd overview of these properties in greater generality and depth, as do Ji & Telgarsky (2020), the latter moreover in deep learning context.
will in this
Proposition C.11. The following properties of o-minimal structures and functions definable in an o-minimal structure hold.
1. (Attouch et al., 2010, Theorem 4.1) If B - R. is continuously differentiable and definable in an o-minimal structure, then has the KL property (Definition C.I) at every x E R
2. (Wilkie) There is an o-minimal structure which the following functions are definable: the exponential function + exp(x) (for x E R., and polynomials of arbitrary degree in real variables Z1...., Fn (van den Dries Miller, 1994), (Van den Dries, 1998, Corollary 2.11).
3. (van den Dries Miller, 1996, $B, B.7(3)) f : - Oa and definable in an o-minimal structure, then its gradient x Vx(x) is definable.
Here and below, the use of the article "an" is intentional there may be multiple o-minimal structures in which function is definable, and in some applications it is important to distinguish between them. For our purposes in this work, this distinction is immaterial: will be content to simply work in 'maximal' o-minimal structure that contains all functions will need.
35

---

Published as a conference paper at ICLR 2024
4. ((van den Dries & Miller, 1996, $B. B.4)) A function f R -> T with (ft...., f is definable in an o-minimal structure if and only if each fo is definable (in the same struc- ture).
5. (Loi, 2010, Theorem 2.3(i The composition of functions definable in an o-minimal struc- ture is definable.
6. (Loi, 2010, Theorem 2.3(ii) If f.:: : IR" + IR is definable in an o-minimal structure, then its image imi(f f) is definable; if moreover f is invertible on its image, then its inverse f-1 (defined on if))) ) is definable."
These properties represent a rather minimal set that are sufficient for our purposs.. To illustrate how to use them to deduce slightly more accessible (and useful) corollaries, we provide the following result on definable functions that will be used repeatedly in our proofs.
Corollary C.12. If functions fi : R + R and f2 : R + R are definable in the o-minimal structure asserted by Proposition C.11, Property 2. then fi + f2 is definable in the same structure.
Proof. By Proposition C.1 Property 2, the function g = 21 + 22, for F1.2y R. definable since it is a polynomial. By Proposition C.11, Property 4, (2) is definable. So, by Proposition C.11, Property 5, f1 + f2 is definable since it is the composition of (2) and g.
Using these properties, we prove that the regularizer associated to an LPN with <a 1 is KL function. For concision, we call a function f "definable" if it is definable in the o-minimal structure whose existence is asserted by Proposition C.11, Property 2.
Lemma C.13 (Definability and KI. property of the prior Ro). Suppose So is an LPN constructed following the recipe in Proposition 3. with softplus activations a(x) = (1/33 log(1 ( + exp(Br)), where B3 > O is an arbitrary constant, and with strong convexity weight 0 < a < Then there is function Rg IK" + R such that fe = prox and Rg is definable and has the KL property (Definition C.I) at every point of IR". Moreover, this Rg can be chosen to simultaneously satisfy the conclusions of Lemma C.7.
Proof. We appeal to Lemma C.7, given that 0 < a 1, to obtain that a function Rg R" - R such that fe = proXR, R, exists, and moreover that this function Rg is C and satisfies the relation (C.22). The idea of the proof is use (C.22) show that Rg is definable in an o-minimal structure, then appeal to Proposition C.11, I, Property to obtain that Ra has the KL property at every point of IR.. By Corollary C.12, to show that Rg is definable, suffices show that four functions appearing in the representation (C.22) are definable: x + (fg (x), V (fø (x))), x + ll 12, x + II, (x) and x + th (f (x)). We can reduce further. By Proposition C.11, Property 2, the squared norm x + ||x|] is definable, and Proposition C.11, Property 5, namely that finite compositions of definable functions are definable, it suffices to show that the arguments of the norms appearing amongst these four functions are definable. Similarly, by Proposition 111 Properties 2, 4, and 5. for the inner product term appearing amongst these four functions, suffices to show that the two individual arguments of the inner product are definable. It then follows by one additional application of Proposition C. Property 5 that we need only argue that the following three functions are definable: x (x). + fg (x), and + V (x). Finally, by Proposition C.I Property 6 and Lemma C.7, which asserts that fo is invertible on R" with range equal to R", suffices merely to show that x + fo is definable to obtain that its inverse is definable.
To complete the proof, we argue in sequence below that each of these three functions are definable.
ICNN the The ICNN the is defined inductively in Proposition 3.1, and by our hypotheses, the activation function g(x) is an elementwise application of the softplus activation with parameter O: with minor abuse of notation, this function is
g(x) (1/ log(1 exp(3x)).
(C.26)
The proof of this property follows readily from the cited result and the preceding properties by noticing that if one defines map F: RD x R."" + RD R by F(z, y) = (f() f())), then gr(f-1) F-1 (gr((),, which expresses the graph of f-1 the inverse image of definable set by definable function.
36

---

Published as a conference paper at ICLR 2024
This (scalar) activation function is definable. To this, by Proposition C.11, Property 2, have that x + exp & is definable, and by Proposition C.11. Property 6, we have that I + log x is definable. It then follows from Proposition C.11, Properties 2 & 5 that is definable, composition of definable functions. We then conclude from Proposition C.11, Property 4 that the elementwise activation + g(x) is definable. Now, by the definition of 4/0 in Proposition 3.1, we have that 2 is definable as composition of definable functions (Proposition C.11, Properties 2 and 5), and arguing inductively, we have by the same reasoning that zu is definable for each k ....., K. Because the = w zς + b is an affine function of Zk, one additional application of Proposition C.11, Properties and establishes that thy is definable.
a finite
ICNN gradient Vele. We conclude this immediately from Proposition C.II, Property 3 and the definability of the since the C2 (because the composed affine maps and C² activations).
LPN fo. Recall that fx(x) (x)o axx Because we have shown above that Vel is definable, using Proposition C.11, Properties and 5 once more, we conclude that fo is definable.
Concluding. By our preceding reductions, we have shown that Rg definable. We conclude from Proposition C.11, Property that Rg has the KL property at every point of IR".
Remark. Note that byproduct the proof of Lemma C.13 is that all constituent functions of the LPN are also definable in o-minimal structure. This fact may be of interest for future work extending Lemma C.13 to priors associated with novel LPN architectures.
D ALGORITHMS
D.1 ALGORITIM FOR LOG-PRIOR EVALUATION
Algorithm 2 Log-prior evaluation for LPN
Input: Learned proximal network fo(+), wh(() that satisfies fo Vila. query point x
1: Find such that fe fy x, by solving miny dy(y; a) - (x,y) y) or miny [Ifo(y) x||
2: 4 (y, x) ||| - (y)
Output: R
The learned log-prior (i.e., regularizer function)
D.2 ALGORITHM FOR LPN TRAINING
Algorithm 3 Training the LPN with proximal matching loss
Input: Training dataset D. initial LPN parameter 6, loss schedule 7(-), noise standard deviation d, number of iterations K, network optimizer Optm(-.))
1: k 0
2: repeat
3:
Sample x ~ D. IO,
4:
y +xx+oc
5:
Crm + my (1.fo()) - x x122
6:
1 0+0 Optm(0, VoLPM)
B Update network parameters
7:
k+ +1
8: until k = K
Output: 0
> Trained LPN
37

---

Published as a conference paper at ICLR 2024
D.3 ALGORITHM FOR USING LPN WITH PNP-PGD TO SOLVE INVERSE PROBLEMS
Algorithm Solving inverse problems with LPN and PnP-PGD
Input: Trained LPN fo. measurement operator A, measurement y, data fidelity function h(x)
lly A(x) 12 initial estimation Xg, step size 1), number of iterations K
1: for k == to K do
2: X1++1 - fo (xk- v)((xx))
3: end for
Output: X
E EXPERIMENTAL DETAILS
E. DETAILS OF LAPLACIAN EXPERIMENT
The LPN architecture contains four linear layers and 50 hidden neurons at each layer, with 3 = 10 in softplus activation. The LPN is trained by Gaussian noise with 1, Adam optimizer (Kingma & Ba, 2014) and batch size of 2000. For either & or & loss, the model is trained for total of 20k iterations, including 10k iterations with learning rate Ir = le - 3, and another 10k with /r le 4. For the proximal matching loss, we initialize the model from the & checkpoint and train according to the schedule in Table 3. To enforce nonnegative weights of LPN, weight clipping is applied during training. projecting the negative weights to zero at each iteration.
Number of iterations
7 in LPM
Learning rate
2k:
0.5
1e
2k
0.5
le
4k
0.4
le
4k
0.3
1e
4k
0.2
1e
4k
0.1
1e
4k:
1e
Table 3: The schedule for proximal matching training of LPN in the Laplacian experiment.
E.2 DETAILS OF MNIST EXPERIMENT
The LPN architecture is implemented with four convolution layers and 64 hidden neurons at each layer, with a = 0.01 and softplus 3 10. The model trained on the MNIST training set contain- ing 50k images, with Gaussian noise with standard deviation
0.1 and batch size of 200. The LPN first trained by L loss for 20k iterations; and then by the proximal matching loss for 20k iterations. with initialized at 0..44 28- 17.92 and halved every 5k iterations. The learned prior evaluated on 100 MNIST test images. Conjugate gradient used to solve the convex inversion problem: miny v
(x, y) in prior evaluation.
E.3 DETAILS OF CELERA EXPERIMENT
We center-crop CelebA images from 178 x 218 to 128 X 128, and normalized the intensities to [0, 1). Since CelebA images are larger and more complex than MNIST, we use deeper and wider network. The LPN architecture includes 7 convolution layers and 256 hidden neurons per layer, with a = lE - and 3 = 100. For LPN training, we train two separate models with two levels 0.05 and 0.1. When applied for deblurring, the best model is selected for each blurring degree (abtur) and measurement noise level (anoise). We pretrain the network with l loss for 20% iterations with lr le 3. Then, train the LPN with proximal matching loss LPM for 20% iterations using Ir = le - 4, with the schedule of y similar to MNIST: initialized 0.64 /128 x 128 x 3 142, and multiplied by 0.5 every 5k iterations. A batch size of 64
of training noise:
38

---

Published as a conference paper at ICLR 2024
is used during training. 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, and Mayo-CT. The same weight clipping as in Appendix E.1 is applied to ensure the weights stay nonnegative throughout training. The training time of LPN on the CelehA dataset is about 6 hours on a NVIDIA RTX A5000 GPU.
applied such initialization in the experiments on CelebA
PnP algorithm and comparison methods We use PnP-ADMM to perform deblurring on CelebA for BM3D, DrCNN, and our LPN (see Algorithm 1). We implement the PnP-ADMM algorithm using the SCICO package (Balke et al., 2022). We implement DaCNN (Zhang et al., 2017a) using their public code 23 We implement the GS denoiser, Prox-DRUNet (Hurault al., 2022b), using their public code and follow their paper to use the Douglas-Rachfore when solving inverse problem, which performs the best with Prox-DRUNet based on their paper. Both DnCNN and Prox-DRUNet are trained on CelebA.
Rachford splitting (DRS) algorithm
E.4 DETAILS OF MAYO-CT EXPERIMENT
We use the public dataset from Clinic for the low-dose CT grand challenge (Mayo-CT) (Mc- Collough, 2016), which contains abdominal CT scans from 10 patients and a total of 2378 images of size 512 x 512. Following (Lunz et al., 2018), we use 128 images for testing and leave the rest for training. The LPN architecture contains 7 convolution layers with 256 hidden neurons per layer, with a = 1e 6 and 3 = 100. During training. we randomly crop the images to patches of size 128 128. At test time, LPN is applied to the whole image by sliding windows of the patch size with stride size of 64. The training procedure of LPN is the same as in CelebA, except that y in proximal matching loss is initialized to 0.64 /128 x 128 82. As in the CelebA experiment, we use LPN with PnP-ADMM for solving inverse problems.
Sparse-view CT Following Lunz et al. (2018), we simulate CT sinograms using a parallel-beam geometry with 200 angles and 400 detectors. The angles are uniformly spaced between -90" and 90%. White Gaussian noise with standard deviation = 2.0 is added to the sinogram data to simulate noise measurement. We implement AR in Py PyTroch based on its public TensorFlow code25; for UAR, we use the publicly available code and model weights 26
Compressed sensing For compressed sensing, we implement the random Gaussian sampling ma- trix following Jalal et al. (2021b), and add noise of 0.001 the measurements. The Wavelet- based sparse recovery method for compressed sensing minimizes the object 1 ly--A4 Ax + where the sensing matrix and W is suitable Wavelet transform. We select the "db4" Wavelet and X = 0.01. To solve the minimization problem in Wavelet-based approach, we use proximal gradient descent with step size of 0.5, stopping criterion ||x xxeli 1e 4, and maximum number of iterations = 1000.
F DISCUSSIONS
F.1 OTHER WAYS TO PARAMETERIZE GRADIENTS OF CONVEX FUNCTIONS VIA NEURAL NETWORKS
Input convex gradient networks (ICGN) (Richter-Powell al., 2021) provide another way to param- eterize gradients of convex functions. The model performs line integral over Positive Deni-Definite (PSD) Hessian matrices, where the Hessians are implicitly parameterized by the Gram product of Jacobians of neural networks, hence guaranteed to be PSD. However, this approach only permits single-layer networks order to satisfy crucial PDE condition in its formulation (Richter-Powel Powell et al., 2021), significantly limiting the representation capacity. Furthermore, the evaluation of the
httpps: : http:/:://gttuubboo/ooec httpg/tth....cccc//aaaaaaaa
Regularr aallariserr
httt:::///ittttee http://thhh...ccoccmcaiiiiillplp //nroo11 wtwprroooonggmeeeeeeeeeeeee ne_mmets__aaeeeeeeeeeeeegeeee_se _ag__ee_t_ddaa________e________e _ata_driver_____
regulari annularizaiion zation.
39

---

Published as a conference paper at ICLR 2024
convex function is less straightforward than ICNN, which is essential step in prior evaluation for LPN (see Section 3). We therefore adopt the differentiation- parameterization in this work and leave the exploration of other possibilities to future research.
G ADDITIONAL RESULTS
G.1 LEARNING SOFT-TIIRESHOLDING FROM LAPLACIAN DISTRIBUTION
Ese. 0.5
Con. 0.3
Eoe. 0.1
Learteed prodiaal ls
LPN prox. -1
-2
Corvex potential bl
LPN Reference
LPN
Log-srtor Ro
-2.5
0.0
2.5
-2.5 0.0
2.5
-2.5
0.0
2.5
-2.
0.0
2.5
-25
0.8
2.5
Figure The proximal operator fo. convex potential tho, and log-prior Rg learned by LPN via different losses: the square & loss, f loss, and the proposed proximal matching loss LPM with dif- ferent 7 E (0.5, 0.3, 0.1). The ground-truth data distribution is the Laplacian p(x) = exp(-Ix)), with log-prior - log p(x) = |æ| log()). With proximal matching loss, the learned proximal fa and log-prior Re progressively approach their ground-truth, proxj. and respectively, y shrinks from 0.5 to 0.1.
G.2 LEARNING PRIOR FOR MNIST IMAGE BLUR
9.19.12
Rg
Rg
R, 10.13
9
9
9
0.0
6.5
2.0
0.0
0.5
a- 1.0
2.0
4.0
Figure 7: The log-prior Rg learned by LPN on MNIST, evaluated images blurred by Gaussian kernel with increasing standard deviation a. Left: the prior over 100 test images. Right: the prior at individual examples.
40

---

Published as a conference paper at ICLR 2024
Besides perturbing the images by Gaussian noise and convex combination in Section 5.1, also evaluate the prior LPN at blurry images, with results shown in Figure Again, the prior increases as the image becomes blurrier, coinciding with the distribution of the hand-written digits in MNIST.
G.3 SOLVING INVERSE PROBLEMS USING LPN WITH PNP-PGD
Besides PnP-ADMM, we also test LPN's performance for solving inverse problems using PnP-PGD (proximal gradient descent). Table shows the numerical results for deblurring CelebA images: PGD is slightly less performant than ADMM in terms of PSNR.
Table 4: Numerical results for CelebA deblurring using LPN with PnP-PGD and PmP-ADMM, averaged over 20 test images.
Oblur = L, Onoise .02
Obhu 1, noise .04
METHOD
PSNR(()
SSIM(1)
PSNR
SSIM
LPN with PnP-PGD
32.7 2.9
.92 = .03
31.2 + 2.5
.89 11: .04 1
LPN with PaP-ADMM
33.0 2.9
=
31.3 1 2.3
41


Reviews:
[{"review_id": "Kjws32K3br", "reviewer": "Reviewer_fi4Y", "summary": "Plug-and-play methods are a framework used for solving inverse problems in image processing, and rely on proximal operators. It's been shown that the proximal operator for the linear inverse problem case is equivalent to a MAP denoiser. In the literature PnP methods almost always rely on MMSE denoiser due to difficulty of learning and designing MAP estimators. In this work, the authors suggest a way to learn a MAP denoiser which they then use in an ADMM algorithm to solve inverse problems. They show the algorithm converges to fixed points of the prior for a particular design of neural networks which guarantees strongly convex mapping. Finally they test the algorithm empirically on toy Laplacian distribution as well as MNIST, CelebA and Mayo-CT datasets.", "questions": "Please see the questions raised under the weaknesses section.", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 1, "contribution": 2, "strengths": "This work addresses a prevalent problem at the core of PnP methods, which is the use of MMSE denoisers instead of MAP denoisers, despite the fact that the convergence results for PnP holds for MAP denoisers. They offer a novel way to learn a MAP denoiser and provide convergence results.", "weaknesses": "- 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 \\to 0$ a discontinuous prior. \n- 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. \n- 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. \n- I find it surprising and contradictory that the PSNR values are higher for this methods compared to other methods that use MMSE denoisers. It is expected that for solving inverse problems using a MAP denoiser result in sharper images with lower PSNR (=higher MSE) while using MMSE denoiser result in higher PSNR (=lower MSE) and more blurry results. This is obvious since, as also stated by the authors, the MMSE denoiser pushed the image towards the mean of the posterior while MAP denoiser pushes the image towards the mode of the posterior. Obviously the mean of the posterior is equivalent to the MSE minimizer, so it should result in lower MSE. \n- Additionally, I found it confusing that the qualitative results for CelebA dataset seem pretty much the same across different methods, and not very good in general. The results look pretty blurry for the proposed method which is surprising given that the method relies on a MAP denoiser. \n- Finally, in recent years, there has been some score based diffusion models algorithms which use priors in denoisers to solve linear inverse problems, which result in significantly better performance. Considering the disadvantage in performance, what are the advantages of this methods? It seems necessary to include a comparison to those models.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "NtU1HXboLs", "reviewer": "Reviewer_8FVJ", "summary": "Inverse problems $y = A(x) + v$ are commonly formulated using regularized least squares:\n$$\\min_x \\frac{1}{2} \\|\\| y - A(x) \\|\\| + \\phi(x)$$\n\nwhere $\\phi$ is an appropriate regularizer. Solving this minimization problem often involves the use of proximal gradient methods. This process necessitates the selection of a regularizer, $\\phi$, and the knowledge of its proximal map. A natural approach, leading to Maximum A Posteriori (MAP) estimation, is to choose $\\phi(x) = -\\log p(x)$.\n\nThis paper presents a method for learning the proximal map of $\\phi(x) = -\\log p(x)$ from data, using only samples from $p(x)$. The authors then apply this proposed approach, in conjunction with proximal gradient methods, to various standard inverse problem tasks, demonstrating favorable performance.", "questions": "- Equation (3.2) should it be \\psi_\\thet(x,\\alpha) ? \n\n- Page 5, \"demillustrates\" should be corrected to \"illustrates.\"\n\n- 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.\n\n- It would be insightful to determine whether the log-likelihood computed using LPN matches the log-likelihood computed using other generative models.\n\n- \"Hidden\" in the appendix is that the LPN is first trained on the l1-loss, then it is trained with the proximal matching loss. This should be stated clearly in the main body of the paper. Why this is needed? How bad is LPN trained just with the proximal matching loss?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The use of proximal and descent methods in conjunction with deep neural networks to address inverse problems is a dynamic and exciting field of study. Numerous papers have explored these approaches, attempting to approximate Maximum A Posteriori (MAP) estimation in various ways or training regularizers in a supervised manner (which is not always feasible and not even correct). To the best of my knowledge, this paper stands out as the first to present a principled approach for training proximal maps of the log probability and effectively approximating MAP estimation.\n\nAdditionally, the numerical results presented in this paper are well-executed and show promise.", "weaknesses": "- 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).\n\n- 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.\n\n--- \nAfter rebuttal: The authors have addressed the concerns raised in the reviews and I have increased my score.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2tPJ3w9dRt", "reviewer": "Reviewer_hkxu", "summary": "The paper presents a framework for constructing learned proximal networks (LPN) that offer precise proximal operators for a data-driven regularizer. It also demonstrates a novel training strategy called proximal matching, which ensures that the resulting regularizer accurately captures the log-prior of the actual data distribution. Experiments test the effectiveness.", "questions": "My main concerns are the experimental details.\n1.\tIt is better to give some details on the training of the learned proximal networks. For example, how to learn the non-negative weights?\n2.\tHow to set the \\alpha in \\phi_\\theta in the training?\n3.\tHow to set \\eta in Theorem 4.1 in the training?\n4.\tThe 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.\n5.\tThe used datasets are very small. Deep learning can only be effective when there is a substantial amount of data.\n6.\tIt is best to verify the effectiveness of the proposed method on common real tasks, e.g., real denoising or deblurring.\n7.\tCompared with other methods, how is the training efficiency? It is better to verify it by experiment.\n8.\tAt the end of the Abstract, \"demonstrating\" should be \"Demonstrating\".", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1.\tThe paper forms a class of neural networks to guarantee to parameterize proximal operators. The idea is novel.\n2.\tSome theoretical results are provided.\n3.\tExperiments show the effectiveness.", "weaknesses": "1.\tThe paper is not well organized, making it hard to follow.\n2.\tThe experiments are a little weak. See the questions below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "eLHLWWhTNI", "reviewer": "Reviewer_XDuj", "summary": "In this work, the authors propose a way to learn a denoising map, called LPN, which is exactly the proximal operator of a scalar function which should approximate the true log image prior. They make use of the characterization of nonconvex proximity operators as gradient of convex functions from (Gribonva & Nikolova, 2020) and parametrize their denoiser as the gradient of (an ICNN.+ a quadratic term). They also suggest a specific loss, called \"proximal matching loss\" such that, when the denoiser is trained to denoise Gaussian noise with this loss, it should approximate the prox of the true log image prior. They also propose a convergence proof of the PnP-PGD algorithm with plugged LPN denoiser and experiment on different datasets and inverse problems.", "questions": "See above.", "limitations": "", "rating": 5, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- The paper is generally well presented and well written. \n- The idea to parameterize a denoiser as the gradient of an ICNNs is new in the PnP literature. \n- The most interesting contribution of the paper is for me the proximal matching loss. \n- The verification of the nonconvexity of the learned prior is also interesting.", "weaknesses": "Here are several potential issues that I spotted while reviewing the paper.\n\nMajor weaknesses : \n- The end the proof from Theorem 4.1 is I think not true, and the presented result not valid. Indeed, the definition of subdifferential subdiferential for convergence of nonconvex PGD is not the usual subdifferential but the limiting subdifferential. With this notion, the first equivalence from C.48 is not true when $\\phi_\\theta$ is nonconvex, and only the implication holds. Therefore, the algorithm can not converge to a fixed-point as presented, but only to a critical point of $h + \\eta^{-1}\\phi_\\theta$.\n- The targeted objective function $h + \\eta^{-1}\\phi_\\theta$ contains the stepsize (which is then a regularization parameter). This is uncommon in optimization and not desirable, as a change in the stepsize affects the objective function and the obtained result. Moreover, the stepsize (and thus the regularization parameter) must be bounded by $1/L$, is this limiting in practice and does it impact the performance ?\n- The algorithm analyzed for convergence is just the classical Proximal Gradient Descent (PGD) (or Forward-Backward). \nWhy do you use the convergence analysis from (Bot et. al, 2016) which is specific for an accelerated version of this scheme. You could use directly the (most commonly used) convergence result from (Attouch et. al, 2013) of PGD in the nonconvex setting.\n- \"Definable\" is an important notion that should be defined. Moreover, the stability properties of Definable functions ( by sum, composition, inverse, derivative) are extensively used without referencing the proof of these results. Same comment for the fact that the exponential is definable. \n- You experiment with PnP-ADMM when the presented theoretical result is with PGD. If this is correct, this is a major issue. Indeed, the theoretical convergence results of ADMM in the convex setting are not the same as the ones for PGD. \n\nMinor weaknesses : \n- Lemma C.2 :  I think that $f_\\theta$ is only invertible on $Im(f_\\theta)$ and $\\phi_\\theta$ is only differentiable on $Im(f_\\theta)$. If this is really true on the whole space, can you explain why ? \n- Section 1 : \"for almost any inverse problem, a proximal step for the regularization function is always present\" I do not understand this sentence and what you mean by \"present\".\n- Section 1 : In (Romano et. al, 2017), there is no prox involved. \n- Section 2 : Contrary to what you seem to explain, PnP is really not limited to ADMM !! \n- Section 3 : Parameterizing a network as the gradient of an ICNN as already been proposed in the literature, in particular in the Optimal Transport community. These works should be cited as well. \n- Section 3 : You explain that the chose parameterization is more general and universal than (Hurault et. al, 2022). I am very doubtful about this statement, given the fact that the ICNN parameterization is very (and I think way more) constraining. A way to support your affirmation would be to compare the performance with both denoisers for denoising and PnP restoration. \n- In most PnP algorithm, the Gaussian noise parameter $\\sigma$ on which the denoiser is trained acts as a regularization parameter. Here, you do not mention this parameter, and you do not explain how it is chosen. \n\nThese weaknesses explain my low score, and I am ready to raise my score if the authors answer these limits.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)