Learning Proximal Operators to Discover Multiple Optima
Authors: Lingxiao Li
Noam Aigerman
Vladimir Kim
Jiajin Li
Kristjan Greenewald
Mikhail Yurochkin
Justin Solomon
Venue: ICLR 2023 poster (2023)

Abstract: Finding multiple solutions of non-convex optimization problems is a ubiquitous yet challenging task. Most past algorithms either apply single-solution optimization methods from multiple random initial guesses or search in the vicinity of found solutions using ad hoc heuristics. We present an end-to-end method to learn the proximal operator of a family of training problems so that multiple local minima can be quickly obtained from initial guesses by iterating the learned operator, emulating the proximal-point algorithm that has fast convergence. The learned proximal operator can be further generalized to recover multiple optima for unseen problems at test time, enabling applications such as object detection. The key ingredient in our formulation is a proximal regularization term, which elevates the convexity of our training loss: by applying recent theoretical results, we show that for weakly-convex objectives with Lipschitz gradients, training of the proximal operator converges globally with a practical degree of over-parameterization. We further present an exhaustive benchmark for multi-solution optimization to demonstrate the effectiveness of our method.

Published as a conference paper at ICLR 2023
LEARNING PROXIMAL OPERATORS TO DISCOVER MULTIPLE OPTIMA
Lingxiao Li MIT CSAIL lingxi seggxiaa@@mit.ed
Noam Aigerman
Vladimir G. Kim
Adobe Research gegemmm rmermant newddaabbe..mom
Adobe Research vok i www.mmmaae......ooommmm
Jiajin Li
Kristjan Greenewald
Stanford University jlajaii@ttaffrrdeeeu
IBM Research, MIT-IBM Watson Al Lab kristjan h www.eweewwlddeeeeemmmmmmm imm. com
Mikhail Yurochkin
Justin Solomon
IBM Research, MIT-IBM Watson AI Lab mi mikhail, yurochkine imm. com
MIT CSAIL jsolomon@mit.c edu
ABSTRACT
Finding multiple solutions of non-convex optimization problems is ubiquitous yet challenging task. Most past algorithms either apply single-solution optimiza- tion methods from multiple random initial guesses search in the vicinity of found solutions using ad hoc heuristics. We present an end-to-end method to learn the proximal operator of family of training problems that multiple local min- ima can be quickly obtained from initial guesses by iterating the learned operator, emulating the proximal-poin algorithm that has fast convergence. The learned proximal operator can be further generalized to recover multiple optima for un- seen problems at test time, enabling applications such as object detection. The key ingredient in our formulation is proximal regularization term, which elevates the convexity our training loss: by applying recent theoretical results, we show that for weakly-convex objectives with Lipschitz gradients, training of the proxi- mal operator converges globally with practical degree of wver-parameterization We further present exhaustive benchmark for multi-solution optimization demonstrate the effectiveness of our method.
INTRODUCTION
Searching for multiple optima of optimization problem is ubiquitous yet under-explored task. In applications like low-rank recovery (Ge et al., 2017), topology optimization (Papadopoulos et al., 2021), object detection (Lin et al., 2014), and symmetry detection (Shi et al., 2020), it is desirable to recover multiple near-optimal solutions, either because there are many optima due the fact that the optimization objective does not capture user preferences precisely. Even for single-solution non-convex optimization, typical methods look for multiple local optima from random initial guesses before picking the best local optimum. Additionally, it is often desirable to obtain solutions to a family of optimization problems with parameters not known in advance, for instance, the weight of regularization term, without having to restart from scratch.
global
Formally. we define multi-solutton optimization (MSO) problem to be the minimization mingex (x),, where E TT encodes parameters of the problem, X' the search space of the variable 2, and f- Rd + R. is the objective function depending on т. The goal of MSO is to identify multiple solutions for each 7 E T, i.c., the set {a* E X : fxx) = min fo (x)} which can contain more than one element or even infinitely many elements. In this work, we assume that X C Rd bounded and that d is small, and that T is, in a loose sense, a continuous space, such that the objective fr changes continuously further assume that each f- is differentiable almost everywhere. As finding all global minima in the general case is extremely challenging, realistically our goal is to find diverse set of local minima.
varies. To make gradient-based methods viable,

---

Published as a conference paper at ICLR 2023
As a concrete example, for object detection, T could parameterize the space of images and X' could be the 4-dimensional space of bounding boxes (ignoring class labels). Then, fr(r) could be the minimum distance between the bounding box E X' and any ground truth box for image r E T. Minimizing f-(x) would yield all object bounding boxes for image т. Object detection can then be cast as solving this MSO on training set of images and extrapolating
unseen images (Section 5.5). Object detection is a singular example of MSO where the ground truth annotation is widely available. In such cases, supervised learning can solve MSO by predicting fixed number of solutions together with confidence scores using set-based loss such as the Hausdorff distance. Unfortunately, such annotation is not available for most optimization problems in the wild where we only have access to the objective functions - this is the setting that our method aims to tackle.
Our work inspired by the proximal-point algorithm (PPA), which applies the proximal operator of the objective function to an initial point iteratively to refine it to local minimum. PPA is known to converge faster than gradient descent even when the proximal operator is approximated, both theoretically (Rockafellar, 1976; 2021) and empirically (e.g., Figure 2 of Hoheisel et al. (2020)). If the proximal operator of the objective function is available, then MSO can be solved efficiently by running PPA from variety of initial points. However, obtaining good approximation of the proxi- mal operator for generic functions is difficult, and typically we have to solve a separate optimization problem for each evaluation of the proximal operator (Davis & Grimmer, 2019).
In this work, we approximate the proximal operator using a neural network that is trained using a straightforward loss term including only the objective and a proximal term that penalizes deviation from the input point. Crucially, our training does not require accessing the ground truth proxi- mal operator. Additionally, neural parameterization allows us to learn the proximal operator for all (f)+er by treating T as an input to the network along with an application-specifi encoder. Once trained, the learned proximal operator allows us to effortlessly run PPA from any initial point to arrive at a nearby local minimum; from generative modeling point of view, the learned proximal operator implicitly encodes the solutions of an MSO problem as the pushforward of a prior distri- bution by iterated application of the operator. Such formulation bypasses the need to predict fixed number of solutions and can represent infinitely many solutions. The proximal term in our loss promotes the convexity of the formulation: applying recent results (Kawaguchi & Huang, 2019), we show that for weakly-convex objectives with Lipschitz gradients- in particular, objectives with bounded second derivatives wiith practical degrees of prrmmttramttiiiiinnnn, training converges globally and the ground truth proximal operator is recovered (Theorem 3.1 below). Such a global convergence result is not known for any previous
method (Chen et al., 2021).
Literature on MSO is scarce, so we huild a benchmark with a wide variety of applications including level set sampling, non-convex sparse recovery, max-cut, 3D symmetry detection, and object detec- tion in images. When evaluated on this benchmark, our learned proximal operator reliably produces high-quality results compared to reasonable alternatives, while converging in a few iterations.
2 RELATED WORKS
Learning to optimize. scarning-to-optimin (L20) methods utlilize past optimization experience to optimize future problems more effectively; see (Chen et al., 2021) for a survey. Model-free L20 uses recurrent neural networks to discover new optimizers suitable for similar problems (Andrychowicz et al., 2016; Li & Malik, 2016; Chen et al., 2017: Cao et al., 2019); while shown to be practical, these methods have almost no theoretical guarantee for the training to converge (Chen al., 2021). In comparison, we learn a
proximal operator so that at test time we do not need access to objective functions or their gradients, which can be costly to evaluate (c.g. symmetry detection in Section 5.4) or unavailable (c.g. object detection in Section 5.5). Model-based L20 substitutes components of a specialized optimization framework or schematically unrolls an opti- mization procedure with neural networks. Related to proximal methods, Gregor & LeCun (2010) emulate a few iterations of proximal gradient descent using neural networks for sparse recovery with an el regularizer, extended to non-convex regularizers by Yang et al. (2020); similar technique is applied to susceptibility-tent tensor imaging in Fang et al. (2022). Gilton et al. (2021) propose a deep equilibrium model with proximal gradient descent for inverse problems in imaging that circumvents expensive backpropagation of unrolling iterations. Meinhardt et al. (2017) use a fixed denoising neural network as surrogate proximal operator for inverse imaging problems. All these works use
2

---

Published as a conference paper at ICLR 2023
schematics of proximal methods to design neural network that is then trained with strong supervi- sion. In contrast, we learn the proximal operator directly, requiring only access to the objectives; we do not need ground truth for inverse problems during training.
Existing L20 methods are not designed to recover multiple solutions: without a proximal term like in (2), the learned operator can degenerate even with multiple starts (Appendix D.3).
Finding multiple solutions. Many heuristic methods have been proposed to discover multiple so- lutions including niching (Brits et al., 2007; Li, 2009), parallel multi-starts (Larson & Wild, 2018), and deflation (Papadopoulos et al., 2021). However, all these methods do not generalize to similar but unseen problems.
Predicting multiple solutions at test time is universal in deep learning tasks like multi-label classi- fication (Tsoumakas & Katakis, 2007) and detection (Liu et al., 2020). The typical solution is to ask the network to predict a fixed number of candidates along with confidence scores to indicate how likely each candidate is a solution (Ren al., 2015; Li al., 2019; Carion al., 2020). Then the solutions will be chosen from the candidates using heuristics such as non-man non-maximum suppression (Neubeck & Van Gool, 2006). Models that output fixed number of solutions without taking into account the unordered set structure can suffer from "discontinuity" issues: a small change in set space requires large change the neural network outputs (Zhang al., 2019). Furthermore, this approach cannot handle the case when the solution set is continuous.
Wasserstein gradient flow. Our formulation (2) corresponds to one step of JKO discretization of the Wasserstein gradient flow where the energy functional is the the linear functional dual to the MSO objective function (Jordan et al., 1998; Benamou et al., 2016). See the details in Appendix E. Compared to recent works neural Wasserstein gradient flows (Mokrov al., 2021; Hwang al., 2021; Bunne et al., 2022), where a separate network parameterizes the pushforward map for every JKO step, functional's linearity makes the pushforward map identical for each step, allowing end-to-end training using single neural network. We additionally let the network input a parameter T, in effect learning continuous family of JKO-discretized gradient flows.
METHOD
3.1 PRELIMINARIES
Given the objective fr Rd + R of an MSO problem parameterized by т, the corresponding proximal operator (Moreau, 1962; Rockafellar, 1976; Parikh & Boyd, 2014) is defined, for fixed X E Rs0, as
A prox(#;T) == arg min f. (v) (v) + 2 -
(I)
The weight À in the proximal term 1/211y- ? controls how close (r;x(2;; T) is x: increasing
À will reduce [Iprox(;;;)) - 2-122- For the arg min in (1) to be unique, a sufficient condition is that fr is É-weakly convex with { A, that f-(v) 2 lly - -=|2 is strongly convex. The class of weakly convex functions is deceivingly broad: for instance, any twice differentiable function with bounded second derivatives (e.g. any C2 function on a compact set) is weakly convex. When the function is convex, prox(2;7) T) is precisely one step of the backward Euler discretization of integrating the vector field fr with time step 1/A (see Section 4.1.1 of Parikh & Boyd (2014)).
The proximal-poine algorithm (PPA) for finding a local minimum of fr iterates
= prox 110) E N
with initial point d0 (Rockafellar, 1976). In practice prox(;;T) often can only be approximated, resulting in ineract PPA. When the objective function is locally indistinguishable from convex function and x is sufficiently close to the set of local minima, then with reasonable stopping crite- rion, inexact PPA converges linearly to a local minimum of the objective: the smaller À is, the faster the convergence rate becomes (Theorem 2.1 2-2.3.3 of Rockafellar (2021)).
'The usual convention is to use the reciprocal of X in front of the proximal term. We use different
convention to associate À with the convexity of (1).
3

---

Published as a conference paper at ICLR 2023
3.2 LEARNING PROXIMAL OPERATORS
The fast convergence rate of PPA makes it a strong candidate for MSO: to obtain a diverse set of solutions for any т T. we only need to run few iterations of PPA from random initial points. The proximal term penalizes big jumps and prevents points from collapsing to single solution. However, rumning a subroutine to approximate prox(2; 7) for every pair (z, T) can be costly.
To overcome this issue, we learn the operator prox(:;-) given access to (fTTEET. A naive way to learn prox(::-) to first solve (1) to produce ground truth for a large number of (z, r) pairs independently using gradient-
methods and then learn the operator using mean-squared error loss. However, this approach is costly as the space X' x T can be large. Moreover, this procedure requires stopping criterion for the minimization in (I), which is hard to design a priori.
Instead, we formulate the following end-to-end optimization over the space of functions:
EasH T~U ( r)) + ( r) - x112]
min 1: X
(2)
+X-
where x is sampled from 14, a distribution on X, and r is sampled from v, a distribution on T. To get (2) from (I), we essentially substitute y with the output (T) and integrate over the product probability distribution μ @ v.
To solve (2), parameterize X X - X using a neural network with additive and mul- tiplicative residual connections (Appendix B). Intuitively, the implicit regularization of neural net- works aligns well with the regularity of prox(::-): for a fixed r the proximal operator prox(:; 7) is 1-Lipschitz in local regions where f- is convex, while as the parameter T varies fr changes con- tinuously so (2; r) should not change too much. To make (2) computationally practical during training, we realize v as a training dataset. For the choice of je, we employ an importance sampling technique from Wang & Solomon (2019) as opposed to using (ni the uniform distribution over X, so that the learned operator can refine near-optimal points (Appendix C). To train D, we sample a
of (z, to evaluate the expectation and optimize using Adam (Kingma & Ba, 2014). For problems where the space T is structured (e.g. images or point clouds), we first embed into a Euclidean feature space through an encoder before passing it to t. Such encoder is trained together with operator network 4. This allows to use efficient domain-specific encoder (e.g. convolutional networks) to facilitate generalization to unseen T.
To extract multiple solutions at test time for a problem with parameter 7, we sample batch of x's from (X) and apply the learned ((,, T) to the batch of samples a few times. Each application of & approximates single step of PPA. From a distributional perspective, for k E N20, we can view thethe operator & applied & times (#*)* (U"), concentrates on the set of local minima approximates as k increases. An advantage of representation is that it can represent arbitrary number of solutions even when the set of minima is continuous (Figure 2). This procedure differs from those in existing L20 methods (Chen al., 2021): test time, we do not need access {St (Jr)reT or their gradients, which can be costly to evaluate or unavailable; instead we only need T (e.g. in the case of object detection, T is an image).
generative model so that the pushforward distribution,
3.3 CONVERGENCE OF TRAINING
We have turned the problem of finding multiple solutions for each fr in the space X' into the problem of finding a single solution for (2) in the space of functions. If the f-'s 's are é-weakly convex with g < À and 1r, v have full support, then the arg min in (1) is unique for every pair (x, r) and hence the functional solution of (2) is the unique proxmal operator (r)
If in addition the gradients of the objectives are Lipschitz, using recent learning theory results (Kawaguchi & Huang, 2019) we can show that with practical degrees of over-parameterization, gradient descent on neural network parameters of & converges globally during training. Suppose our training dataset is S = [(T, R) C X x Define the training loss, discretized version of (2) using S, to be, for g : X x T + X.
A fr. (o(2,, Ti)) + (lo(x; - 10l12
L(g) =
(3)
4

---

Published as a conference paper at ICLR 2023
Theorem 3.1 (informal). Suppose for any T, the objective fr is differentiable, g-weakly convex, neural network with (2(n) total
and frr is <-Lipschitz with S X. Then for any parameters' and common activation units, when the initial weights are drawn from a Gaussian distribution, with high probability, gradient descent its weights using fixed learning rate will eventually reach the minimum loss ming:XxT L(g). The number of iterations needed to achieve > O training error is O((2 (),,), and when this occurs, if { < X, then the mean- aaaan--uaredd error of the learned proximal operator compared to the true one is (2)(())E)) on training data.
We state and prove Theorem 3.1 formally in Appendix A. Even though the optimization over net- work weights is non-convex, training can still result in a globally minimal loss and the true proximal operator can be recovered. In Appendix D.2, we empirically verify that when the objective is the l norm, the trained operator converges to the true proximal operator, the shrinkage operator. In Appendix D.3, we study the effect of in relation to the weakly-convex constant & for the 21) cosine problem and compare an L.20 particle-swarm method (Cao et 2019).
We note a few gaps between Theorem 3.1 and our implementation. batching instead of gradient descent. Second, instead of feed-forward networks, we use network architecture with residual connections (Figure B.1), which works better empirically. Under these conditions, global convergence results can still be obtained, e.g., via (Allen-Zhh et al.. 2019, Theo- rems 6 and 8), but with large polynomial bounds in n. H for the network parameters. Another gap is caused by the restriction of the function class of the objectives. In several applications in Section 5, the objective functions are not weakly convex or have Lipschitz gradients, or we deliberately choose small for faster PPA convergence; we empirically demonstrate that our method remains effective.
First, we use SGD with mini-
PERFORMANCE MEASURES
Metrics. Designing single-valued met- ric for MSO is challenging since one needs to consider the diversity of the solutions well each solution's 's level of optimality. For an MSO problem with parameter and ob- jective fr. the output of an MSO algorithm can be represented as a (possibly infinite) set of solutions {IaJo C X with objective values ua = (#a) Suppose we have ac- cess to ground truth solutions {ys ha C X with vy := (21)). Pick a threshold E R and denote A, := {xa Uo & (), B. := {yg vig S 1). Let W be a random variable that is uniformly distributed on X. Define a random variable
1
Figure 1: Interpretation of Dj. In this example, the witness W is drawn uniformly from the union of four squares. If Ar (resp. B,) is the set of red (resp. blue) points, then P(D, 1/4, since D; is only non zon-zero when W is in the right- most square. This aligns well with the intuition that 3/4 of the red points match with the blue ones. In com- parison, the Hausdorff distance between A, and B, approximately I, which is the same as the Hausdorff distance between the orange point and B,, despite the fact most of red points are close to the blue ones.
0) 3/4 and P(D,
0.5)
D = A (W) - TB. (TA, (W))l2 + [|TBB (W) -A, (TB, (W))
(4)
- We call W a witness of D., as it witnesses how different A and B, are near W. To summarize the law of Dr, define the witnessed divergence and wifnessed precision at 8 > 0 as
where (x)))= = arg min,es ||== 2||2-
WD; = E[D,] and WP = P(D, < <))
(5)
Witnesses help handle unbalanced clusters that can appear in the solution sets These metrics are agnostic to duplicates, unlike the chamfer distance optimal transport metrics. Compared alter- natives like the Hausdorff distance, WD, remains low if a small portion of A,, B, are mismatched. We illustrate these metrics in Figure One can interpret WD, as weighted chamfer distance whose weight is proportional to the volume of the & eo-Vooooo cell at each point in either set.
gradient descent until convergence on randomly sampled particles in X' for every E T. We use this method to generate approximated ground truth solutions to compute the metrics in (5) when the ground truth is not available. This method is not directly comparable to ours since it cannot generalize to unseen
Particle Descent: Ground Truth Generation. A naïve method for MSO is
We use Ω notation in the standard way, i.e., ((()) + 3k Nzo such that f E R(n log* n).
5

---

Published as a conference paper at ICLR 2023
T's at test time. Remarkably, for highly non-convex objectives, particle descent produce worse solutions than the ones obtained using the learned proximal operator (Figure D.7).
Learning Gradient Descent Operators. As there is ceadily-available application-agnen base- line for MSO, we propose the following method that learns iterations the gradient descent opera- tor. Fix Q N21 and step size 1j O. We optimize an operator Vi via
min V: x - 7~p
Ezval
(x, 7) - (2(x;;) (x):
(6)
where 42 (2; T) is the result of Q steps of gradient descent on fr starting at i.e., (](z;)) = I, and V (2; r) = (2:T)
f (2; Each iteration of minimizing (6) requires Q evaluations V.r. which can be costly (e.g., for symmetry detection in Section 5.4). We use importance sampling similar to Appendix C. An ODE interpretation is that V performs Q iterations of forward Euler on the gradient field fr. whereas the learned proximal operator performs a single iteration of backward Euler. We choose Q 10 for all experiments except for symmetry detection (Section 5.4) where we choose Q = I because otherwise the training will take > 200 hours. As will see in Figure D.6, aside from slower training, this approach struggles with non-smooth objectives due to the fixed step size 17, while the learned proximal operator has no such issues.
APPLICATIONS
We consider five applications to benchmark our MSO method, chosen to highlight the ubiquity of MSO in diverse settings. We abbreviate POL for proximal operator learning (proposed method), GOL for gradient operator learning (Section 4), and PD for particle descent (Section 4). Further details about each application can be found in Appendix D. The source code for all experiments can be found at https: // /gggttttbbb...////g////g com weiinniiall9///ee.
5.1 SAMPLING FROM LEVEL SETS
Formulation. Level sets provide concise and resolution-free implicit shape repre- sentation (Museth et al., 2002; Park et al., 2019; Sitzmann et al., 2020). Yet they are less intuitive to work with, even for straightforward tasks on discretized do- mains (meshes, point clouds) like visual- izing or integration on the domain. We present an MSO formulation to sample from level sets, enabling the adaptation of downstream tasks level sets.
Figure 2:
Visualization of the solutions for the conic section problem. Red, green, and blue indicate the so- lutions by PD, GOL, and POL respectively. See Fig- ure D.3 for more examples.
Given a family of functions {gr X
R*TTET, for each t suppose we want to sample from the O-level set (0). We formulate an MSO
problem with objective /=(x)= == ||9- (<r) [ whose global optima are precisely g,+ We do not need assumptions level set topology or that the implicit function represents distance field, unlike most existing methods (Park et al., 2019; Deng et al., 2020; Chen et al., 2020).
Benchmark. We consider sampling from conic sections. We keep this experiment simple so to visualize the solutions easily. Let X' = [-5, 5]2 and T = [|1, For T = (A, B, C, D, E, F) E T, define 9r to be 9- (2) == Az Bay + Cy² + Dx + Ey +. Since f- = ( is defiedd on compact X, it satisfies the conditions of Theorem 3.1 for a large X, but a large À corresponds
to small PPA step size. Empirically, small À for POL gave decent results compared to GOL: Figure
illustrates that POL consistently produces sharper level sets for both hyperbolas (B² - 4AC 0)
and ellipses (B2- 4AC < 0). Figure D.4 shows that POL yields significantly higher WP than GOL for small 8, implying that details are well recovered. Figure D.5 verifies that iterating the trained operator of POL converges much faster than that of GOL. It is straightforward to extend this setting to sample from more complicated implicit shapes parameterized by T.
5.2 SPARSE RECOVERY
Formulation. In signal processing, the sparse recovery problem aims to recover a signal x² E X C Rd from noisy measurement y R. distributed according y Ax e, where A E Rmxd

---

Published as a conference paper at ICLR 2023
m < d, and is measurement noise (Beck & Teboulle, 2009). In applications like imaging and speech recognition, the signals are sparse, with few non-zero entries (Marques et al., 2018). Hence, the goal of sparse recovery is to recover sparse x" given A and y.
A common way to encourage sparsity is to solve least-squares plus an e norm on the signal:
min xEx
A 2112 + a||x>>p
for a, > and | == (x? + c)p/ for a small to prevent instability. We consider the non-convex case where 0 < p < 1. Compared to convex alternatives like in LASSO (p = 1), non-convex t norms require milder conditions under which the global optima of (7) are the desired sparse x² (Chartrand Staneva, 2008; Chen & Gu, 2014).
To apply our MSO framework, we define T = (a.p) E T and S+ to be the objective (7) with corresponding o,p. Compared to existing methods for non-convex sparse recovery (Lai et al., 2013), our method recover multiple solutions from the non-convex landscape for family of a's and p's without having to restart. The user can adjust parameters a, p to quickly generate candidate solutions before choosing a solution based on their preference.
T = [0, 1] x [0.2,0.5] 0.5). We consider highly non-convex S norms with p 00.20..5, to test method's limits. We choose d 8 and sparse signal x²' uniformly in X' with half of the coordinates set to 0. We then sample entries in A i..dd from N(0, 1) and generate y Ax + where (N 0.1). Although ||=:|p is not weakly convex, POL achieves decent results (Figure D.6). Notably, POL often reaches a better objective than PD (Figure D.7) while retaining diversity, even though POL uses a much bigger step size (1/A = 0.1 compared to PD's 10-5) and needs to leam different operator for entire family of T. In Figure D.8, we additionally compare POL with proximal gradient descent (Tibshirani et al.. 2010) for p- 1/2 where the corresponding thresholding formula has closed-form (Cao al., 2013). Remarkably, we have observed superior performance of POL against such strong baseline.
Benchmark. Let X' = |- ,
4, and sample the
5.3 RANK-2 RELAXATION OF MAX-CUT
Formulation. MSO can be applied to solve combinatorial problems that admit smooth non-convex relaxations. Here, we consider the classical problem of finding the maximum cut of an undirected
graph G = (V, E), where V = {1..... ,n), E C V V, with edge weights {wir} } C R so that wiy 0 if (i.j) & E. The goal to find {ar] +11,+11 to maximize Σi Wsy (1 x12;).
Burer et al. (2002) propose solving minser Ei, Wiy cos(ff: -,)) a rank-2 non-convex relaxation of the max-cut problem. This objective inherits weak convexity from cosine, so it satisfies the conditions of Theorem 1. In practice, instead of using angles as the variables which are ambiguous up to 2n, we represent each variable as point on the unit circle 5 so we choose X' = (8*)² and T be the space of all edge weights with n vertices. For T {Tih} E TT corresponding to graph with edge weights ) we define, for x E X.
fx(x) = Σ Tijal xz.
(8)
After minimizing fr, we can find cuts using Goemans & Williamson procedure (1995). In- stead of using heuristics find optima near solution (Burer al., 2002), our method can help the user effortlessly explore the set of near-optimal solutions without hand-designed heuristics.
Benchmark. We apply our formulation to Kg. the complete graph with 8 vertices. Hence X (S*)* C R6, We choose T = [0, 1 as there are 28 edges in Kg. We mix two types of random graphs with 8 vertices in training and testing: Senōss-ééyiieee graphs with p = 0.5 and Kş with uniform edge weights in [0, 1]. Figure 3 shows that POL can generate diverse set of max cuts. Quantitatively, compared to GOL, POL achieves better witnessed metrics (Figure D.10).
5.4 SYMMETRY DETECTION OF 3D SHAPES
Formulation. Geometric symmetries are omnipresent in natural and man-made objects. Knowing symmetries can benefit downstream tasks in geometry and vision (Mitra et al., 2013; Shi et al., 2020; Zhou et al., 2021). We consider the problem of finding all reflection symmetries of 3D surface.
7

---

Published as a conference paper at ICLR 2023
Figure 3: 18 different max cuts (max cut value 10) graph generated by our method. Red and blue vertices indicate the two vertex set separated by the cut. Vertex 0 is set to blue to remove the duplicates obtained by swapping the colors. See Figure D.12 for more results.
Let T be shape representation (e.g. point cloud, multi-view scan), and let M, C R$ denote the corresponding triangular mesh that is available for the training set. As reflections are determined by the reflectional plane, set X' 52 x R20. where unit normal n E 5 C R and intercept d E Rzo (we assume d 2 0 to remove the ambiguity of (-n,-0) representing the same plane). Let R₂ R R denote the corresponding reflection. Perfect symmetries of M, satisfy Ra (M,) = Mr. Let 8T R - R be the (unsigned) distance field of M, given by s, (p) mingeM, ||p 4l12- Inspired by Podolak et al. (2006), we define the MSO objective to be
(n,, d) X' denotes the plane with
f+(x) (s-
where a batch of p is sampled uniformly from M, when evaluating the expectation. Although fr stochastic, since zero with probability one. Compared to existing methods that either require ground truth symmetries obtained by human annotators (Shi et al., 2020) detect only small number of symmetries (Gao et al., 2020), our method applied to (9) finds arbitrary numbers of symmetries including continuous and can generalize unseen shapes, without needing ground truth symmetries supervision.
point-to-non meso distances compute 8y, perfect symmetries will make (9)
Benchmark. We detect reflection symmetries for mechanical parts in the MCB dataset (Kim et al., 2020). We choose T be the space of 3D point clouds representing mechanical parts. From the mesh of each shape, we sample 2048 points with their normals uniformly and use DGCNN (Wang al., 2019) encode the oriented point clouds. Figure show our method's 's results on selection of models in the test dataset; for per-iteration PPA results of our method, see Figure D.13. Figure D.1 D.11 shows that POL. achieves much higher witnessed precision compared to GOL.
Figure 4: Symmetry detection results. Each reflection is represented as a colored line segment rep- resenting the normal of the reflection plane with endpoint the plane. Pink indicates better objective values, while blue indicates worse. Our method capable of detecting complicated dis- crete symmetries well continuous families of cylindrical reflectional symmetries.
5.5 ОВJЕCТ DETECTION IMAGES
Formulation. Identifying objects an image is a central problem in vision on which recent works have made significant progress (Ren et al., 2015; Carion al., 2020; Liu al., 2021). We consider a simplified task where we drop the class labels and predict only bounding boxes. Let b (x,, y, w. h) X [0, 1]a denote box with (normalized) center coordinates (x, width w, and height h. We choose T to be the space of images. Suppose an image has K, ground truth object bounding boxes {b] ] K- We define the MSO objective to be f-(r)= = min ||] - =
its minimizers are exactly {b] J.K., Although the objective may seem trivial, its gradients reveal the Ll_ Voronoi diagram formed by be's when training the proximal operator. Different from exist- ing approaches, we encode the distribution of bounding boxes conditioned on each image in the learned proximal operator without needing to predict confidence scores or a fixed number of boxes. A similar idea based on diffusion is recently proposed by Chen al. (2022).
8

---

Published as a conference paper at ICLR 2023
Benchmark. We apply the above MSO formulation the COC02017 dataset (Lin et al., 2014). As r is an image, we fine-tune ResNet-50 (He et al., 2016) to encode 7 into a vector & that can be consumed by the operator network (Figure B.1).
Table 1: Object detection results. &D (resp. WP..) is the witnessed divergence (resp. precision) in (5) with oe (i.e. keeping all solutions), averaged over 10 trials (standard deviation 10 3)) Precision and recall are computed with Hungarian matching as no confidence score is available for the usual greedy matching (see Appendix D.8). FRCNN(.S) (Ren et al., 2015) means keeping predictions with confidence 2 8% for Faster R-CNN.
METHOD
2D
WP
PRECISION
RECALL
FRCNN(.80)
0.140
0.624
0.778
0.650
FRCNN(.95)
0.162
0.589
0.887
0,515
0.161
0.481
0.139
0.577
GOL
0.251
0.243
0,508
0.282
POL (OURS)
0.149
0.590
0.817
0,442
In addition to GOL, we design a baseline method FN that uses the same ResNet-50 backbone and predicts fixed number of boxes using the chamfer distance the training loss. Table compares the proposed methods with alternatives and the highly- on the test dataset. Since we do not output confidence scores, the metrics are computed solely based on the set of predicted boxes. Our method achieves significantly better results than FN and GOL. Compared to the Faster R-CNN, we achieve slightly worse results with 40.7% fewer network parameters. While Faster R-CNN contains highly-specializel modules such as the regional proposal network, in our method we simply feed the image feature vector output by ResNet-50 to a general- purpose operator network. Incorporating specialized architectures like region proposal networks into our proximal operator learning framework for object detection is an exciting future direction. We visualize the effect of PPA using the learned proximal operator in Figure 5. Further qualitative results (Figure 2)) and details can be found in Appendix D.8.
Faster R-CNN (Ren et al., 2015)
Figure 5: First 4 iterations of PPA using the learned proximal operator 20 randomly initialized boxes (leftmost column). Only a few iterations are needed for the boxes to form distinctive clusters.
CONCLUSION
Our work provides a straightforward and effective method to learn the proximal operator of MSO problems with varying parameters. Iterating the learned operator on randomly initialized points ef- ficiently yields multiple optima to the MSO problems. Beyond promising results on our benchmark tasks, see many exciting future directions that will further improve our pipeline.
A current limitation is that at test time the optimal number of iterations to apply the learned operator not known ahead of time (see end of Appendix D.1). One way to overcome this limitation would be to train another network that estimates when to stop. This measurement can be the objective itself if the optimum value is known a priori (e.g., sampling from level sets) or the gradient norm if objectives are smooth. One other future direction is to learn proximal operator that adapts to multiple X's. This way, the user can easily experiment with different X's and to enable PPA with growing step sizes for super-linear convergence (Rockafellar, 1976; 2021). Another direction is to study how much we can relax the assumption that X' is a low-dimensional Euclidean space. Our method could remain effective when X' is low-dimension submanifold of a high-dimensional Euclidean space. The challenges would be to constrain the proximal operator to a submanifold and to design proximal term that is more suitable than the ambient L norm.
9

---

Published as a conference paper at ICLR 2023
Reproducibility statement. The complete source code for all experiments be found at https:// btttps://githie /ithhub........ www.///nggiaalii4///1
Detailed instructions are given in README md. We have further included tutorial how to extend the framework custom problems-see "Extending to custom problems" section where we include a toy physics problem of finding all configurations of an elastic spring. For all our experiments, the important details are provided in the main text, while the remaining details needed to reproduce results exactly are included the appendix.
Acknowledgements We thank Chenyang Yuan for suggesting the rank-2 relaxation of max-cut problems. The MIT Geometric Data Processing group acknowledges the generous support of Army Research Office grants W91 INF201016 and W911NF2) 102933, 10293, of Air Force Office of Scientific Research award FA9550-19-1-03/ of National Science Foundation grants IIS-1838071 and CHS- 1955697, from the CSAIL Systems that Learn program, from the MIT-IBM Watson Al Laboratory, from the Toyota-CSAll Joint Research Center, from gift from Adobe Systems, and from Google Research Scholar award.
REFERENCES
Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. convergence theory for deep learning via over- parameterization. In International Conference Machine Learning, pp. 252. PMLR, 2019.
Luigi Ambrosio, Nicola Gigli, and Giuseppe Savaré. Gradient flows: in metric spaces and in the space of probability measures. Springer Science & Business Media, 2005.
Marcin Andrychowicz, Misha Denil, Sergio Gomez, Matthew Hoffman, David Pfau, Tom Schaul, Brendan Shillingford, and Nando De Freitas. Learning to learn by gradient descent by gradient descent. In Advances in neural information processing systems, pp. 3981-3989, 2016.
Amir Beck and Mare Teboulle. A fast iterative shrinkage-threshollinng algorithm for linear inverse problems. SIAM journal on imaging sciences, 2(1): 2009.
David Benamou, Guillaume Carlier, Quentin Mérigot, and Edouard Oudet. Discretization of functionals involving the monge-ampére operator. Numerische mathematik, 134(3):611-63) 2016.
Brits, Andries Petrus Engelbrecht, and Frans van den Bergh. Locating multiple optima using particle swarm optimization. Applied Mathematics and Computation, 189(2):1859-1883 2007.
Charlotte Bunne, Laetitia Papaxanthos, Andreas Krause, and Marco Cuturi. Proximal optimal trans- port modeling of population dynamics. In International Conference on Artificial Intelligence and Statistics, pp. 6511 1-228. PMIR, 2022.
Samuel Burer, Renato DC Monteiro, and Yin Zhang. Rank-two relaxation heuristics for max-cut and other binary quadratic programs. SIAM Journal on Optimization, 12(2):503-522 2002.
Wenfei Cao, Jian Sun, and Zongben Xu. Fast image deconvolution using closed-form thresholding formulas of lq (q= 12. 23) regularization. Journal of visual communication and image represen- tation, 24(1):31 11,11,1,260 2013.
Yue Cao, Tianlong Chen, Zhangyang Wang. and Yang Shen. Learning to optimize in swarms. Advances in neural information processing systems, 32, 2019.
Nicolas Carion, Francisco Massa, Gabriel Synnaeve, Nicolas Usunier, Alexander Kirillov, and Sergey Zagoruyko. End-to-end object detection with transformers. In European Conference on Computer Vision, pp. 213-229. Springer, 2020.
Rick Chartrand and Valentina Staneva. Restricted isometry properties and nonconvex compressive sensing. Inverse Problems, 24(3):035020, 2008.
Laming Chen and Yuantao Gu. The convergence guarantees of a non-convex approach for sparse recovery. IEFE Transactions Signal Processing, 62(15):3754-376 2014.
10

---

Published as a conference paper at ICLR 2023
Shoufa Chen, Peize Sun, Yibing Song, and Ping Luo. Diffusiondet: Diffusion model for object detection. arXiv preprint 1.X0:2221..09788, 2022.
Tianlong Chen, Xiaohan Chen, Wuyang Chen, Howard Heaton, Jialin Liu, Zhangyang Wang, and Wotao Yin. Learning to optimize: primer and a benchmark. arXiv preprint urXiv:2103. 12828, 2021.
Yutian Chen, Matthew W Hoffman, Sergio Gómez Colmenarejo, Misha Denil, Timothy P Lilli- crap, Matt Botvinick, and Nando Freitas. Learning to learn without gradient descent by gradient descent. International Conference on Machine Learning, pp. 748-756. PMLR, 2017.
Zhiqin Chen, Andrea Tagliasacchi, and Hao Zhang. Bsp-net: Generating compact meshes via binary space partitioning. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 45-54, 2020.
Damek Davis and Benjamin Grimmer. Proximally guided stochastic subgradient method for nons- mooth, nonconvex problems. SIAM Journal on Optimization, 29(3):1908-1999 2019.
Boyang Deng, Kyle Genova, Soroosh Yazdani, Sofien Bouaziz, Geoffrey Hinton, and Andrea Tagliasacchi. Cvxnet: Learnable convex decomposition. In Proceedings of the IEEE/CVF Con- ference Computer Vision and Pallern Recognition, pp. 31-44, 2020.
Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using real nvp. arXiv preprint arXiv:1605.000 08803, 2016.
Zhenghan Fang, Kuo-Wei Lai, Peter van Zijl, Xu Li, and Jeremias Sulam. Deepsti: Towards tensor reconstruction using fewer orientations in susceptibility tensor imaging. arXiv preprint arXiv:2209.04504 2022.
Lin Gao, Ling-Xiao Zhang, Hsien- Yu Meng, Yi-Hui Ren, Yu-Kun Lai, and Leif Kobbelt. Prs-net: Planar reflective symmetry detection net for 3d models. IEEE Transactions Visualization and Computer Graphics, 27(6):3007-3088, 2020.
Rong Ge, Chi Jin, and Yi Zheng. No spurious local minima in nonconvex low rank problems: A unified geometric analysis. In International Conference on Machine Learning, pp. 1233-1242. PMLR, 2017.
Davis Gilton, Gregory Ongie, and Rebecca Willett. Deep equilibrium architectures for inverse problems in imaging. IEEE Transactions on Computational Imaging, 7:1123-1133, 2021.
Michel X Goemans and David P Williamson. Improved approximation algorithms for maximum cut and satisfiability problems using semidefinite programming. Journal of the ACM (JACM), 42(6); 1115-1145, 1995.
Karol Gregor and Yann LeCun. Learning fast approximations of sparse coding. In Proceedings of the 27th International Conference on Machine Learning, pp.
2010.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Tim Hohcisel, Maxime Laborde, and Adam Oberman. A regularization interpretation the proxi- mal point method for weakly convex functions. Journal of Dynamics & Games, 7(1):79, 2020.
Hyung Hwang, Cheolhyeong Kim, Min Sue Park, and Hwijae Son. The deep minimizing move- ment scheme. arXiv preprint arXiv:2 14851, 2021.
Richard Jordan, David Kinderlehrer, and Felix Otto. The variational formulation of the fokker- planck equation. SIAM journal on mathematical analysis, 29(1):1-17, 1998.
Kenji Kawaguchi and Jiaoyang Huang. Gradient descent finds global minima for generalizable deep neural networks practical sizes. In 2019 57th Annual Allerton Conference on Communication, Control, and Computing (Allerton), pp. 92-99. IEEE, 2019.

---

Published as a conference paper at ICLR 2023
Sangpil Kim, Hyung- gun Chi, Xiao Hu, Qixing Huang, and Karthik Ramani. A large-scale anno- tated mechanical components benchmark for classification and retrieval tasks with deep neural networks. In Proceedings of 16th European Conference on Computer Vision (ECCV), 2020.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv: 2014.
Ming-Jun Lai, Yangyang Xu, and Wotao Yin. Improved iteratively reweighted least squares for unconstrained smoothed \ell minimization. SIAM Journal Numerical Analysis, 51(2):927- 957,0013 2013.
Jeffrey Larson and Stefan M Wild. Asynchronously parallel optimization solver for finding multiple minima. Mathematical Programming Computation, 10(3):303-333 300-2333,333 2018.
Ke Li and Jitendra Malik. Learning optimize. arXiv preprint arXiv: 2016.
Lingxiao Li, Minhyuk Sung, Anastasia Dubrovina, Li Yi, and Leonidas J Guibas. Supervised fitting of geometric primitives to 3d point clouds. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp.
2019.
Xiaodong Li. Niching without niching parameters: particle swarm optimization using ring topol- ogy. IEEE Transactions on Evolutionary Computation, 14(1):150-169, 2009.
Tsung-Yi Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C Lawrence Zitnick. Microsoft coco: Common objects in context. In European conference on computer vision, pp. 740-755. Springer, 2014.
Li Liu, Wanli Ouyang, Xiaogang Wang, Paul Fieguth, Jie Chen, Xinwang Liu, and Matti Pietikäinen. Deep learning for generic object detection: A survey. International journal of computer vision, 128(2):261-318, 2020.
Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. arXiv preprint arXiv:2103. 14030, 2021.
Elaine Crespo Marques, Nilson Maciel, Lirida Naviner, Hao Cai, and Jun Yang. review of sparse recovery algorithms. IEEE access, 7:1300-1322, 2018.
Tim Meinhardt, Michael Moller, Caner Hazirbas, and Daniel Cremers. Learning proximal operators: Using denoising networks for regularizing inverse imaging problems. In Proceedings the IEEE International Conference Computer Vision, pp. 1781-1790, 2017.
Niloy J Mitra, Mark Pauly, Michael Wand, and Duygu Ceylan. Symmetry in 3d geometry: Extrac- tion and applications. In Computer Graphics Forum, volume 32, pp. 1-23. Wiley Online Library, 2013.
Petr Mokrov, Alexander Korotin, Lingxiao Li, Aude Genevay, Justin Solomon, and Evgeny Burnaev. Large-scale wasserstein gradient flows. arXiv preprint arXiv:2 2021.
Jean Jacques Moreau. Fonctions convexes duales et points proximaux dans un espace hilbertien. Comples rendus hebdomadaires des séances de l'Académie des sciences, 255:2897-2899, 1962.
Ken Museth, David E Breen, Ross T Whitaker, and Alan H Barr. Level set surface editing operators. In Proceedings of the 29th annual conference on Computer graphics and interactive techniques, pp. 330-338, 2002.
Alexander Neubeck and Luc Van Gool. Efficient non-maximum suppression. In 18th International Conference on Pattern Recognition (I6) volume 3, pp. 850-855. IEEE, 2006.
Ioannis PA Papadopoulos, Patrick E Farrell, and Thomas Surowiee. Computing multiple solutions of topology optimization problems. SIAM Jourmal on Scientific Computing, 43(3):A 1555-A1522, 2021.
Neal Parikh and Stephen Boyd. Proximal algorithms. Foundations and Trends in optimization, 1(3): 127-239, 2014.
12

---

Published as a conference paper at ICLR 2023
Jeong Joon Park, Peter Florence, Julian Straub, Richard Newcombe, and Steven Lovegrove. Deepsdf: Learning continuous signed distance functions for shape representation. In Proceed- ings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 165-174, 2019.
Joshua Podolak, Philip Shilane, Aleksey Golovinskiy, Szymon Rusinkiewicz, and Thomas Funkhouser. A planar-reflective symmetry transform for 3d shapes. ACM Trans. Graph. 25 (3):549-559, jul 2006. ISSN 0730-0301. doi: 10.1145/1419111.144992 URL https: //doi. otg/10..g1.... 1145/111111
1141923.
Shaoqing Ren, Kaiming He, Ross Girshick, and Jian Sun. Faster r-cnn: Towards real-time object detection with region proposal networks. Advances in neural information processing systems, 28: 91-99, 2015.
R Tyrrell Rockafellar. Monotone operators and the proximal point algorithm. SIAM journal on control and optimization, 14(5):877-898, 1976.
Tyrrell Rockafellar. Advances in convergence and scope of the proximal point algorithm. J. Nonlinear and Convex Analysis, 2021.
Yifei Shi, Junwen Huang. Hongjia Zhang, Xin Xu, Szymon Rusinkiewicz, and Kai Xu. Symme- trynet: learning to predict reflectional and rotational symmetries of 3d shapes from single-view rgb-d images. ACM Transactions Graphics (TOG), 39(6):1-14, 2020.
Vincent Sitzmann, Julien NP Martel, Alexander W Bergman, David B Lindell, and Gordon Wetzstein. Implicit neural representations with periodic activation functions. arXiv preprint arXiv:2006.0966 2020.
Ryan Tibshirani et Proximal gradient descent and acceleration. Lecture Notes, 2010.
Grigorios Tsoumakas and loannis Katakis. Multi-label classification: An overview. International Journal of Data Warehousing and Mining (IJDWM), 3(3):1-13, 2007.
Yue Wang and Justin M. Solomon. Prne:: Self-supervised learning for partial-to- to--trtial registration. In 33rd Conference on Neural Information Processing Systems, 2019.
Yue Wang, Yongbin Sun, Ziwei Liu, Sanjay E Sarma, Michael M Bronstein, and Justin M Solomon. Dynamic graph enn for learning on point clouds. Acm Transactions Graphics (tog), 38(5): 1-12, 2019.
Chengzhu Yang, Yuantao Gu, Badong Chen, Hongbing Ma, and Hing Cheung So. Learning proxi- mal operator methods for nonconvex sparse recovery with theoretical guarantee. IEEE Transac- fions Signal Processing, 08:5244-5259, 2020.
Yan Zhang, Jonathon Hare, and Adam Prugel-Bennett. Deep prediction networks. Advances in Neural Information Processing Systems, 32:3212-3222, 2019.
Yichao Zhou, Shichen Liu, and Yi Ma. Nerd: Neural 3d reflection symmetry detector. In Pro- ceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 15940- 15949, 2021.
13

---

Published as a conference paper at ICLR 2023
A CONVERGENCE OF TRAINING
We formally state and prove Theorem 3.1 via the following Proposition A. and Proposition A.2. Proposition A.I. Suppose
1. C Rd for some N21;
2. for any T E T, the objective Jr differentiable, E-weakly convex, and Vr is (-Lipschitz i.e.
[ (22)lll S cl|=1 - x2ll2.
with & A.
3. the activation function (x) used is proper, real analytic, monotonically increasing and 1-Lipschitz, egg. sigmoid, hyperbolic tangent.
For any 8 0, 2 2, n N21, assume is an H H-layer feed-forwadd neural network with hidden layer sizes m.,..... HH satisfying
my mH- D(H2 log Hn? (5))
m 2 (2log (H /))) my 2 92()..
Let D denote the total number of weights in $. Then D (1(m). Moreover, there exists a learning rate E R such that for any dataset $ {(x1, 2= of size with the training loss I. defined (over random Gaussian initial weights 8 of 1), there exists = (x,(( + S)/e) such that L(D(,,,;')) $ L + c, where [10$11 stays bounded, L == mingexxT-xx L(g) is the global minimum of the functional L, JkeN is the sequence generated by gradient descent gh-1 = 8 - 11 2 VoL(Ø(., (),)),,,;;)) and or depends only on and the initialization 0
as (3), for any O, with probability least
Proof of Proposition A.I. The theorem is an application of Theorem 1 in Kawaguchi & Huang (2019) with the following modifications.
For E [n]. define (x)) == fx,(x) (a:) + I1 n. 1 To check Assumption 1 of Kawaguchi & Huang (2019), observe
Vzl(x)= V.. (z) + M(x xx));
V2(()= (z) + Ala-
Hence the assumption that fr, is g-weakly convex implies that
(²) Ma (²) +(a)
Hence is convex. The assumption that Vr. is Ç-Lipschitz implies, for any 21, & x TT,
[[(((z))) (12))))
| fr, (z1) (z2) x1)11
< V (x1) (22( + + 22112
< (A + O] x2l12-
Hence Vl, (A + C)-Liscchiz
An input vector to the neural network + is the concatenation (z, 7) R Kawaguchi & Huang (2019) assume that the input data points are normalized have unit length. This is not an issue, as we can scale down (I,,, Tτ) uniformly to be contained in unit ball, then pad 7, one extra coordinate to make [z,, r1)l1 = 1 for all E [n]. similar to the argument given in the footnotes before Assumption 2.1 of Allen
al. (2019).
Lastly, we mention explicitly lower bounds for the layer sizes that are used in the proof of Theorem 1 of Kawaguchi & Huang (2019) (see the paragraph below Lemma 3), instead of stating a single bound on the total number of weights in the statement of Theorem 1. This is because Theorem only states that there exists a network of size Q(n) for which training converges, whereas every network satisfying the layer-wise bounds will have the same convergence guarantee.
[
14

---

Published as a conference paper at ICLR 2023
Next show that once the training loss is away from the global minimum, can guarantee that the approximation error on the training data in the mean-squared sense is small: i.c, the learned operator (
: 8) is close to the true proximal operator (1).
Proposition A.2. Suppose for any T E T, the objective fo is differentiable and E-weakly convex with X, where À is the proximal regularization weight of the training loss L(g) defined in (3). Let 0 be the weight of the network t such that L [ :; S L* + where L == mingex T-x L(g) is the global minimum of the functional T.. Let prox(:;-) be the true proximal operator defined in (1). Then the mean- mean-squares squared error on the training data is bounded by
1 (2,, 03:0) - prox(2;; 2 x
(10)
Proof. Clearly L2 = L(prox(::))) }), ie., the minimum of L is achieved with the true proxi- mal operator. Define ha X - R by h₂(2) == (x) + ž11 xiller so that we can write L(g) 1 Ei=// (g( T:)). By the assumption on weak convexity, each h; is (A - E)-strongly -
convex. This implies for any v, y E X,
hx(x) 2 h4( + Vha(v) (x y) A- x -v1
(11)
The minimum of h; is achieved at prox Ti) by the definition of prox. Differentiability and con- vexity imply (proxx(( Tx)) 0. Hence setting y prox(æ;; 7.) in (11) implies, for any x E X,
À { || (xo 7
(x) h, (prox(x;; 2
2
Now by the definition of (3),
2 L(2(;;0)) = 0)) (prox((-))
(.(;, Tz; 0)) h,, (ps)
i=1
2
c (,;;,, pro ( m))l
2
Rearranging terms we obtain the desired result.
B NETWORK ARCHITECTURES
The network architecture we use to parameterize the operators for both POL. and GOL is identical and is shown in Figure B.I. The encoder of T will be chosen depending the application. For our conic section (5.1), sparse recovery (5.2), and max-cut (5.3) benchmarks, the encoder is just the identity map. For symmetry detection (5.4), is a point cloud and we use DGCNN (Wang al., 2019). For object detection (5.5), we use ResNet-50 (He et al., 2016). Inspired by Dinh et al. (2016), include both additive and multiplicative coupling in the residual blocks. At the same time, since we do not need bijectivity of the operator (and proximal operators should not be) nor access to the determinant of the Jacobian, we do not restrict ourselves to a map with triangular structure in Dinh et al. (2016). We use 3 residual blocks for all applications, except for symmetry detection where we use blocks which give slightly improved performance.
Our architecture is economical: the model size (excluding the application-specific encoder) is under 2MB for all applications we consider. This also makes iterating the operators fast at test time. Note that the application- -specifcc encoder only needs to be run once at each test time 7 as the encoded vector z can reused (Figure B)
C IMPORTANCE SAMPLING VIA UNFOLDING PPA
Directly optimizing (2) (6) using mini-batching may not yield operator that can refine near- optimal solution, if μ is taken to be (n', the uniform measure on X' (more precisely, the d- dimensional Lebesgue measure restricted to X' and normalized to probability distribution). Instead,
15

---

Published as a conference paper at ICLR 2023
Operator Network
Residual Block
Application
Encoder
Residual Block
Residual Block
Residual Block
$ 22, T)
Tnput
Network
Output
(shared by residunl blocks)
Fully-connected
(2
connecccd
Intermediatc
(not shared)
Figure B.I: Network architecture for the operators used in POL and GOL. We ReLU as the activation after all intermediate linear layers, except for predicting the scaling in the residual block, where we use HardTanh to ensure s E [-2, 2]. For the shared 2-layer fully-connected network, the hidden layer sizes are 256, 128. For the 3-layer fully-connected network in each residual block, the hidden layer sizes are 128, 128, 128.
we would like to sample from a distribution that puts more probability density on near-optimal solutions. We achieve this goal as follows, inspired by Wang & Solomon (2019). Let & denote the network with weights after training iterations. For k E 220 denote μ == (#* ) (unif(X()). For रम Σ v=tt Then, for training iteration / + 1, we optimize
fixed K E N21, we set the objective (2) or (6) with the constructed μ. Note this modification does not introduce any bias for POL (similarly for GOL), in the sense that the optimal solution (2) is still the true proximal operator since has full support, yet it puts more density in near-optimal regions as increases. In practice, we choose K or K 10. For the choice of other hyper- pareeetarreters, see Appendix D.1.
D DETAILED RESULTS
D.1 HYPER-PARAMETERS
Unless mentioned otherwise, the following hyper- parameearsessrss are used.
In each training iteration of POL and GOL, we sample 32 problem parameters from the training dataset of T, and 256 of w's from unif(,X') when computing (2) or (6) using the importance sampling trick Appendix C. The learning rate of the operator is kept 10-1 11 for both POL. and GOL, and by default we train the operator network for 2 x 10S iterations. This is sufficient for the loss to converge for both POL and GOL in most cases. Since GOL requires multiple evaluations of the gradient of the objective, it typically trains two more times slower than POL.. For the proximal weight À of POL, we choose based on the scale of the objective and the dimension of X'; see Table D.I. All training done on single NVIDIA RTX 3090 GPU.
For the step size 7 in GOL, we start with 1/2 (so same step size as POL in the forward/backwarr Fuler sense) and then slowly increase it (so fewer iterations needed for convergence) without
16

---

Published as a conference paper at ICLR 2023
Table D.I: Choices of À for all applications considered. X' is the search space of solutions, and is the dimension of the Euclidean space where we embed X' (so it might be greater than the intrinsic dimension of X).
Application
X
d
conic section (5.1)
-5,
5
2
0.1
sparse recovery (5.2)
-2.
8
10.0
max-cut (5.3)
(5)))
16
10.0
symmetry detection (5.4)
S R
1.0
object detection (5.5)
[0,1]
4
1.0
degrading the metrics. When evaluating (6), we set Q = 10 in all experiments except for symmetry detection, where choose a step size small enough so as to not miss significant minima and a sufficient number of iterations for the loss (i.e. the objectives) to fully converge.
Q 1 because otherwise the training will take 200 hours. For PD,
For evaluation, the number of iterations to apply the trained operators is chosen to be enough so that the objective converges. This number will be chosen separately for each application and method. By default, 1024 solutions are extracted from each method, and 1024 witnesses are sampled to compute WD; and WP averaged over test dataset and over 10 trials with standard deviation provided (in most cases the standard deviation is two orders of magnitude smaller than the metries). We filter solutions that do not lie in X.
limitation for both POL and GOL that when the solution set is continuous, too many applications of the learned operator can cause the solutions to collapse. We suspect this is because even with the importance sampling trick (Appendix C), during training the operators may never see enough input that are near-optimal to learn the correct refinement needed to recover the continuous solution set. future direction is to have another network to predict a confidence score for each x E so that at test time the user knows when to stop iterating the operator, e.g., when the objective value and its gradient are small enough; see the discussion Section 6.
D.2 CONVERGENCE TO THE PROXIMAL OPERATOR
To empirically verify Proposition A.2, that our method can faithfully approximate the true proximal operators of the objectives, we conduct the following simple experiments. We consider the function (x) = ||=|||, for x E X = [] and treat T as a singleton. proximal operator prox(2) = arg miny |Ivll + |1v x112 is known in closed form as the shrinkage operation, defined coordinate- wise as:
Ft 1/2
r; 2 1/2
prox(x);
0
| 1/2
(12)
04 + 1/2
I; S -1/2.
For each dimension of d = 2,4, 16,32, we train an operator network (Figure B.1) using (2) as
the loss with learning rate 10-3 Figure D.1 shows the mean--quuare-errror [xxx)) (>))) scaled by 1/a and averaged over 1024 samples the training iterations, where Φ- is the shrinkage opera- tion (12). We see that the trained operator indeed converges to &- as predicted by Proposition A.2, and the convergence speed faster in smaller dimensions.
D.3 EFFECT OF THE PROXIMAL TERM
In this section, study the necessity of the proximal term 4/2| 7) 2=1 in (2). Without such a term, the learned operator can degenerate. For example, consider (1) in (Chen et al., 2021), which minimizes ming Σin wx,(2,, with Zt+1 = notation). Suppose x² one global optimum of but is not the only one. Then (fx,...) == clearly minimizes the objective, yet the update steps will always set a, = x= regardless of the initial positions.
((z,, (z(x)),... for all (with adapted
To further illustrate the effect of different choices of X, consider the 2D cosine function f(x) Σ 10cos(2 for X [-5,5] and singleton T. This function is E-weakly convex
17

---

Published as a conference paper at ICLR 2023
0.000
0.010
0.00
2000
2000
0000
8000
10000
training iteration
Figure D.1: Convergence to the true proximal operator of f(x) ||2|11-
A 400
=
=
-2
-2
Figure D.2: Left: the result of POL after 10 iterations with X = 0, 1, 10, 400 for the 2D cosine function which has weakly convex constant 40 400. Right: particle swarms recovered by Cao et al. (2019) for after 10 iterations from 256 independent runs. The population size of the swarm is (default value in their source code).
with ş = 2 < 400 and has global minima forming grid (all local minima are global minima). On the left of Figure D.2, we see that when X = 400 > Ç-in which case the condition of Theo- rem 3.1 is met-POL recovers all optima. In comparison, for 10, the outer ring solutions is missing, and with À 0, most optima missing in the grid.
To demonstrate how existing L20 methods can fail to recover multiple solutions, we conduct the same experiment the L20 particle-swarm method by Cao et al. (2019), which recovers swarm of particles that are close to optima. We use the default parameters in the provided source code except changing the objective to the 2D cosine function and the standard deviation of the initial random particles to 1. As the method by Cao et al. (2019) could produce particles outside X' = [-5,5] we add an additional term 0.01 ||=]] to the objective f(x); without such a term the particle swarm simply collapses
single point far away from the origin. The results are shown on the right of Figure D.2. We see that even with 256 independent random starts and with population size 4, this method fails recover most optima, in particular non-positive quadrants.
D.4 SAMPLING FROM CONIC SECTIONS
Setup. For this problem, the training dataset contains 220 samples of T E T. while the test dataset has size 256. In our implementation and similarly in other benchmarks do not store the dataset on disk, but instead generate them on the fly with fixed randomness. The r's are sampled uniformly in T. PD is run for 5 x 10* steps with learning rate 1.0. For step sizes, choose À 0.1 for POL and = 1.0 for GOL. We found that the training of GOL explodes when 17 > 1.0. Mcanwhile, POL is able to take bigger (I/a 10.0) steps while staying stable during training (but might fail to recover
18

---

Published as a conference paper at ICLR 2023
solutions due to large step size). To obtain solutions, we use iterations for POL, while for GOL use 100 iterations since it converges slower (and more iterations won't 't improve the results).
Results. We visualize for the conic section problem in Figure D.3 for 16 randomly chosen T. In Figure D.4 we plot of ó vs. WP? (5) to quantitatively verify how good POL and GOL are at recovering the level sets, where we treat the results by PD as the ground truth. Both visually and quantitatively, see that POL outperforms GOL. Figure D.5 compares the convergence speed when applying the learned iterative operators at test time: clearly POL converges much faster.
O
Figure D.3: Visualization of the solutions for the conic section problem. Red indicates the solutions by PD which treat as ground truth. Green and blue indicate the solutions by GOL (Section and POL (proposed method) respectively.
GOUD 00.70 POLG 001)
GOL EX2 POULy 159)
6000 108)
POLL (98)
0.1
23
24
Rlge 150
Ang 600)
Figure D.4: The plot of 8 vs. WP for the conic section problem 0.05, 0.1, 0.2). The vertical dashed line indicates WD,, 50 equally spaced values are used to draw the plot. Here Pgt indicates the percentage of PD solutions that have objectives S and p similarly indicates the percentage of solutions for each method with objectives below t. We sample 1024 witnesses to compute WP averaged over 256 test problem instances. The plot is averaged over 10 trials of witness sampling (the fill-in region's 's width indicates the standard deviation). Here the standard deviations are all less than 10-3 so the fill-in regions are too small to be visible.
152
PO
PD
234
C0L
220
GOL
20
GOL
Figure D.5: Convergence speed comparison at time for the conic section problem. For POL and GOL, the --axis is the number of iterations used. For PD, the --axis is the number of gradient descent steps, multiplied by 100. The horizontal axis shows the number iterations, and the vertical axis shows the value of f (2)) averaged over all current solutions (fill-in region's 's width indicates standard deviation). The three plots shown correspond to the problem instances the first three columns in Figure D.3. Once the operator has been trained, POL converges in less than 5 steps, while GOL converges slower (GOL is already trained with the largest step size without causing training to explode).
D.5 NON-CONVEX SPARSE RECOVERY
Setup. For this problem, the training dataset contains 1024 samples of (a, p). while the test dataset has 128 samples. The r's are sampled uniformly in T [0, 1] x [0.2, 0.5). We extract 4096 solutions from each method after training. For PD, 5 x 10 steps of gradient step with
19

---

Published as a conference paper at ICLR 2023
learning rate 10-5. We found that due to the highly nonconvex landscape of the problem, bigger learning rates will cause PD to miss significant local minima. For step sizes, we choose À 10 for POL (so this corresponds to step size 0.1 for backward Euler) and 1) 0.1 for GOL. To obtain solutions, POL. requires less than 20 iterations to converge, while for GOL. over 100 iterations are needed.
Results. We show the histogram of the solutions' objective values for PD, GOL, and POL in Fig- D.7 for 4 problem instances. Figure D.6 visualizes the solutions for problem instances pro- jected onto the last two coordinates. GOL fails badly in all instances. Remarkably, despite the non-convexity of the problem and the much larger step size (0.1 compared to 10-5), POL yields solutions on par or better than PD when is small. For instance, for the second and third columns Figure D.6 (corresponding optimal solutions that POL. (in blue) captures. As such the results of PD can be suboptimal, so we do compute witness metrics here.
second and third columns in Figure D.7), PD (in red) misses near-
Figure D.6: Visualization of the solutions' objective values for the conic section problem. Red indicates the solutions by PD which we treat as ground truth. Green and blue indicate the solutions by GOL and POL (proposed method) respectively.
(0.830 0.465)
(0,202, 0.211)
(1.759.0:278)
(0.2090..05))
FD
3031
200
1.00
122
Figure D.7: Histograms of the objectives for non non-convex sparse recovery on 4 problem instances. (a. p) the top, corresponding sparsity-induest function of the form a a||=>|pp
We denote
Comparison with proximal gradient descent. For p 1/2, 2/3. thresholding formulas exist for the ( norm (Cao et al., 2013). That is, the proximal operator of ||>l|"" has a closed-form. This allows us to apply proximal gradient descent (Tibshirani et al., 2010) to solve (7). When (i.c. when the problem reduces LASSO), this reduces to the popular iterative soft-threshodding algorithm (ISTA) which converges significantly faster than gradient descent.
We compare the convergence speed of POL to that of proximal gradient descent (denoted PGD) for the p = 1/2 case. We also include PD for reference. The generation of data (i.c. A and in (7)) is the same as before. For POL. we use the same setup with À 10 as before (corresponding to step size 0.1) except we restrict to 1/2 during training and we train only for 1000 steps (note the test-time unseen during training). For PGD we use 0.04 as the step size because using 0.05 for the step size would lead to divergence of PGD the objective would go to infinity. For PD we use 0.05 as the step size. We run all three methods for 200 steps (for POL this corresponds to 200 steps of PPA after training) and visualize the convergence and histograms of the objective for each method in Figure D.8. We see that POL converges faster than PGD even when PGD highly specialized to the p 1/2 case (where the thresholding formula has closed form).
20

---

Published as a conference paper at ICLR 2023
0.627
PD
1
PGD
11000
PGD
POL
POL
2-
0
50 1010 150
2001
[.(x]
Figure D.8: Convergence (left figure) and histograms (right figure) of the objective for the 1/2 non-convex sparse recovery problem with a = 0.627. For the convergence figure on the left, the horizontal axis shows the number of iterations, and the vertical axis shows the value of fa(),), averaged over all current solutions (fill-in region's 's width indicates standard deviation). PGD denotes the proximal gradient descent method (Tibshhraaii et al., 2010) with the closed-form thresholding formula by Cao et al. (2013).
Other coursity-inducing regularizers. Our method can be applied to sparse recovery problems with other sparsity-inducing regularizers in straightforward manner. Consider minimax concave penalty (MCP) from Yang et al. (2020) defined component- wise as:
MCP(#;
: :
|
T
1 |
T
1
47
2r
which is r-weakly convex. We repeat the same setup in Section 5.2 but with objectives
[+(x) = Ar y
MCP(5:;),
(13)
fort (0.5, 2). Note PGD viable to solve (13) because the proximal operator of MCP has closed- form. We run PGD for 2 x 10t iterations with step size 10-4 to make sure converges fully. We show the histogram of the solutions' objective values for PD, GOL, POI., and PGD in Figure D.9. Our results are consistent with those in Figure D.7: POL is on par with PD and significantly outperforms GOL, POL also performs better than PGD which only applicable because the regularizer MCP has a dlosed-form proximal operator.
1758
1900
PD
PO
FD
PB
COL
1000
FACT
F-14
Figure D.9: Histograms of the objectives for sparse recovery with minimax concave penalty (MCP) on 4 problem instances.
D.6 RANK-2 RELAXATION OF MAX-CU -UX
Setup. An additional feature of (8) that the variables are constrained to (S C R Hence for POL and GOL we always project the output of the operator network to the constrained set (normal- izing to unit length before computing the loss or before iterating), while for PD we apply projection after each gradient step.
We generate a training dataset of 20 graphs and a test dataset of 1024 graphs using the procedure described in Section 5.3: half of the graphs will be Erdōs-Rény: graphs with p 0.5 and the remaining half being Kg with edge weights drawn from [0. 1] uniformly. For PD, we use learning rate 10-4 For step sizes of POL and GOL, we choose À 10.0 and 10.0.
We choose to directly feed the edge weight vector 7 E 28 to the operator network (Appendix B). We find this simple encoding works better than alternatives such graph convolutional networks.
21

---

Published as a conference paper at ICLR 2023
This is likely because 2 X' (8))) requires order information from the encoded graph pooling operation can be detrimental for the operator network architecture, Designing an equivariant operator network that is capable of effectively consuming larger graphs is an interesting direction for future work.
Results. If a cut happens to be local minimum of the relaxation, then it is maximum cut (Theo- rem 3.4 of Burer et al. (2002)). However, finding all the local minima of the relaxation is not enough to find all max cuts as max cuts can also appear as saddle points (see the discussion after Theorem 3.4 of Burer et al. (2002)). Hence solving the MSO (8) is not enough to identify all the max cuts. Nevertheless, we can still compare POL and GOL against PD based solely on the relaxed MSO prob- lem corresponding the objective (8). In Figure D.10, we plot ò vs. WP, (5) to verify the quality of
10
GOUF (224)
POL 081)
3.450
GOLLe
6300
1.00
POU (1.0]
POL(p 0000)
POLL- 1289
0.0
65
170
13
20
07
100
05
13
1.
20
Agg 0.50)
410g
flg - 1.00
Figure D.10: The plot of 6 vs. WP for the max-cut problem. See the caption of Figure D.4 for the meaning of the symbols. As different edge weights lead to different minima values, we choose the threshold t in relative manner: the actual threshold used will be t times the best objective value found by PD.
the solutions obtained by POL and GOL compared to PD. We see that POL more faithfully recovers the solutions generated by PD with consistently higher witnessed precision.
Empirically, we found the proposed POL can identify a diverse family of cuts. We visualize the multiple cuts obtained by POL for a number of graphs in Figure D.12. Although some cuts are not maximal, they are likely due to the relaxation - not all fractional solutions correspond to a cut and not because of the proposed method. As evident in Figure 1.10, they are still very close to the local minima of (8) generated by PD.
D.7 SYMMETRY DETECTION OF 3D SHAPES
Setup. Since the variables in (9) is constrained to X' = 52 x R20. we always project the output of the operator network to the constrained set: for (n, d) E X. normalize to have unit length and take absolute value of d. The same projection is applied after each gradient step in PD.
To generate training and test datasets, et al., 2020) but filter out meshes with more than 5000 triangles and keep up to 100 meshes per category make the categories more balanced. During each training iteration, a fresh batch of point clouds are sampled (these are r's) from the meshes in the current batch. For step sizes, we choose == 1.0 for POL and 10.0 for GOL. The training of POL and GOL takes about 30 hours. For PD, we run gradient descent for 500 iterations for each model, which is sufficient for convergence.
the original train/test split of the MCB dataset (Kim
We use the official implementation DGCNN by Wang et al. (2019) as the encoder with the modi- fication that we change the input channels to 6-dimension to consume oriented point clouds and we turn off the dropout layers which do not improve performance.
The objective (9) involves which requires woint-to-mesh projection. We implemented custom CUDA functions to speed up the projection. Even so, it remains the bottleneck of training. Since GOL requires multiple evaluations, it is extremely slow and can take more than a week. As such, we set Q = in (6). Both POL. and GOL. are trained for 10 iterations with batch size 8. At test time iterating the operator networks does not need to evaluate the objective nor the s, 's; moreover, only point clouds are needed.
We show the witness metrics in Figure D.11; quantitatively, POL exhibits far higher wit- nessed precision values than GOL. We show a visualization of iterations of PPA with the learned proximal operator in Figure D.13. In particular, discrete reflectional symmetries as well as a continuous family of reflectional symmetries for cylin- drical objects.
Results.
method is capable of detecting complicated
22

---

Published as a conference paper at ICLR 2023
0.75
0.00
WP
0.55
COUG 0.2)
COUG 01.11
00.|--
0.64]
00
POLP wil)
POLL (.29)
POUL
(150]
POL.- 100
0.0 600
6/05
0650
0.00
6/2
0.50
0.075
0.000
3.00
0.005
0.03
0.015
0.00
0.00
0.205
3.500
C/TS 0.00
Ssg. -
A0g 6.0)
46g 2280
Figure D.11: The plot of ā vs. WP for symmetry detection on the test dataset of Kim et al. (2020) (t = 10 10 1002 1002.10-1), See the caption of Figure D.4 for the meaning of the various notations. We do not show WD: as the vertical bars here because GOL's WD, much higher than POL's and is out of the range for the horizontal axis. We sample 1024 witnesses to compute WP?. averaged over 10 trials of witness sampling (the fill-in region's 's width indicates the standard deviation).
D.8 OBJECT DETECTION IN IMAGES
Setup. We use the training and validation split of COC02017 (Lin et al., 2014) as the training and test dataset, keeping only images with at most 10 ground truth bounding boxes. For training, common augmentation techniques such as random resize/crop, horizontal flip, and random RGB shift, to generate a 400 x 400 patch from each training batch image, with batch size 32. For evaluation, we crop a 400 x 400 image patch from each test image. For step sizes, we choose À 1.0 for POL and = 1.0 for GOL. We train both POL and GOL for 10t steps. This takes about 100 hours. To extract solutions, we use 100 iterations for POL (for most images it only needs 5 iterations to converge) and 1000 iterations for GOL (the convergence is very slow so we run it for a large number of iterations).
We fine-tune Py Torch's pretrained ResNet-50 (He et al., 2016) with the following modifications. We first delete the last fully-connected layer. Then we add an additional linear layer to turn the 2048 channels into 256. We then add sinusoidal positional encodings to pixels in the feature image output by ResNet-50 followed by a fully-connected layers with hidden layer sizes 256 256, 256. Finally average pooling is used to obtain single feature vector for the image.
For Faster R-CNN (FRCNN), we use the pretrained model from PyTorch with ResNet-50 backbone and regional proposal network. It should be noted that FRCNN is designed for different task that includes prediction of class labels, and thus is trained with more supervision (object class labels) than our method and it uses additional loss terms for class labels.
For the alternative method FN that predicts a fixed number of boxes, we attach a fully-connected layer of hidden sizes [256, 256, 80] with ReLU activation to consume the pooled feature vector from ResNet-50. The output vector of dimension 80 is then reshaped to 20 x 4, representing the box parameters of 20 boxes. We use chamfer distance between the set of predicted boxes and the set of ground truth boxes as the training loss.
Results. In Table we compute witness metrics and traditional metries including precision and recall. As our method does not output confidence scores, we cannot use common evaluation metrics such average precision. To calculate precision and recall, which normally would require an order given by the confidence scores, we instead build a bipartite graph between the predicted boxes and the ground truth, adding an edge if the Intersection over Union (IoU) between two boxes greater than 0.5. Then we consider predictions that appear in the Hungarian max matching true positives, and the unmatched ones false positives. Then precision defined as the number of true positives over the total number of predictions, while recall is defined as the number of true positives over the total number of ground truth boxes. When computing metrics for POL and GOL, we run mean-shift algorithm with RBF bandwidth 0.01 to find the centers of clusters and them as the predictions. As shown in Figure 5, the clusters formed by POL are usually extremely sharp after a few steps, and any reasonable bandwidth will result in the same clusters.
In Figure D.14, we show the detection results by our method for a large number of test images chosen at random.
23

---

Published as a conference paper at ICLR 2023
Figure 1. Visualization of the cuts obtained by applying Goemans-Williamoon Williamson procedure randomly selected solutions of POL. The graphs are chosen among the ones that have at least solutions uniformly at random without human intervention. Each row contains the multiple solutions the same graph with binary weights. Two colors indicate the two vertex sets separated by the cut. Dashed lines indicate edges in the cut. For each graph annotate the number of cuts top. The Oth vertex is always in blue to remove the obvious duplicates obtained by swapping the two colors on each vertex.
24

---

Published as a conference paper at ICLR 2023
Figure D.13: Visualization of PPA with learned proximal operator on selected models from the test dataset of Kim et al. (2020). Iterations 0, 1, 2, 5, 10, 15, 20 are shown, where the Oth iteration contains the initial samples from un). Pink indicates lower objective value (9), while light blue indicates higher.
25

---

Published as a conference paper at ICLR 2023
SET
Figure D.14: Randomly selected object detection results by POL COC017 validation split. Each test image 400 x 400 patch cropped from the original image with a random scaling by a number between 0.5 and In each image display all 1024 bounding boxes without clustering, most of which are perfectly overlapping. For some images there is bounding box for the whole image (check if the image has orange border). There is class label associated with each box.
26

---

Published as a conference paper at ICLR 2023
E CONNECTION TO WASSERSTEIN GRADIENT FLOWS
In this section we show that we can view (2) as solving the JKO discretization of Wasserstein gra- dient flows every time step, under the assumption that the measures along the JKO discretization are absolutely continuous.
If F()) linear functional of the form F(u) = f fdu on P(X), the space of probability distri- butions in X' with X' compact, then the JKO discretization of the gradient flow of F at step t + with step size 1/A is
( fdp 2wzeenn). MEP2(X)
>
M4+1 = arg min
arg
where W2 is the Wasserstein-2 distance and we assume μe is absolutely continuous. Let
F(u): fdw+ 2
W2 (P..
Let us also define another functional such that for Borel map T :x - X.
G(T) = / ((()))dm()))+ + 2 [[T(x)- x1 =l
First given μ E P(X), since X' is compact (so in particular all probability distributions have finite second moments), by Brenier's theorem (Ambrosio et al., 2005, Theorem 6.2.4), there exists Borel map T (the Monge map) such that 14 μ and ( ()) such ja and we have G(T) = F(n), and thus (E F(J) 2 miny G(7).
[[T(z) ==līdu,(z) Hence for
Next given a Borel T, μ T#Me. By Brenier's theorem, T be Monge map corresponding to W2(me, 14) so that μ = Tape and f [[))(()) =l[)dr((>) = W2(H4. S [TT(2) (x). This shows that F(u) S G(T) and hence min F(A) S miny G(T). Thus min F(u) miny G(T).
iie has full support, then the best T == arg miny G(T) is obtained pointwise and becomes the proximal operator of f (cf. (2)). In particular, 7" does not depend on Ptc-
27


Reviews:
[{"review_id": "u2PmnogB0_", "reviewer": "Reviewer_T3qd", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper proposes a method to learn proximal mappings which have multiple beneficial properties as optimisers, alternative to gradient descent and variants.", "summary_of_the_review": "The paper reads well and is motivated well. Proximal operators tend to perform much better and more stable than gradient updates, therefore learning prox operators is a viable idea for improving \"learning to optimise\" tasks.\n\n", "strength_and_weaknesses": "Strength: The formulation is neat and clean -- the paper is well written. On top of that, theoretical results are provided which is rare in learning to optimise literature.\n\nWeakness: Would be great if more experiments regarding approximation of true proximal operators were done.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "s4w0NWuy74e", "reviewer": "Reviewer_4Grh", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "The authors consider the general problem of multi-solution optimization (MSO),\nwhere one is interested in obtaining a sampling (perhaps all) of the minimizers\nof each of a family of (nonconvex) functions $\\{ f_{\\tau} : \\tau \\in \\mathcal{T} \\}$\ndefined on a low-dimensional domain $\\mathcal{X}$ (here, $\\mathcal{T}$\nrepresents parameters -- an example given is bounding box detection, where\n$\\mathcal{T}$ is a class of digital images and $\\mathcal{X}$ is the\nfour-dimensional euclidean space of bounding box parameters). The authors\npropose a method for this problem that consists of learning a proximal operator\n(parameterized by a neural network) for the functions $f_{\\tau}$\n(assuming access to ground truth) over a training set of parameters $\\tau$ and\ninputs $x$ (the loss/training setup is similar to what one would formulate for\nmulti-task learning); at test time (given a new $\\tau$), a set of putative\nsolutions is generated by repeatedly iterating the learned operator from\nvarious random initializations (in an attempt to simulate the proximal point\nalgorithm). The authors propose an evaluation metric for these families of\nputative solutions in general MSO problems, given ground truth (it is described\nas a weighted chamfer distance between the putative set and the ground truth\nset, with weights proportional to a certain Voronoi tiling of the space\n$\\mathcal{X}$ -- hence more robust to outliers than standard metrics). They\ncompare their approach to a benchmark of gradient descent from repeated random\ninitialization and an approach that attempts to learn an operator that\napproximates a finite number of gradient descent iterations (trained similarly\nto theirs) on several synthetic and practical MSO tasks, showing that their\napproach performs well.\n\n", "summary_of_the_review": "The method is clearly presented, and the paper seems to do a useful service\nof organizing lots of thought and baseline approaches to general MSO problems.\nThe approach taken here to learning a proximal operator and then applying it as\nin the proximal point iteration has been studied elsewhere, but this study in\nthe MSO context seems new and useful. The theory presented here seems mostly\n\"plug and play\" with regards to existing results, so it could be changed to a\nMSO-relevant context by applying a different result from the literature.\n\n", "strength_and_weaknesses": "## Strengths\n\nThe authors present a method applicable to general MSO problems that strikes a\ngood balanced between computational/optimization principles (beginning with the\nproximal point iteration) and learning-based approaches (learning a proximal\noperator on datasets). This suggests the ability to provide theoretical\nguarantees for the approach.\n\nThe experiments are extensive, and strike a good balance between synthetic\ntasks (which allow specific behaviors of the method to be understood) and\nimportant applications (symmetry detection and object detection on big\ndatasets, comparing to reasonable benchmarks). Ablations and comparisons to\nother general MSO baselines are given in the main text and the extensive\nexperimental appendices. It seems reasonable to imagine that with further\ntuning, the approach could be even more competitive relative to strong\nbaselines.\n\n\n## Weaknesses\n\nReferences to a large body of literature on learning proximal operators for\ninverse problems (esp. in imaging) seem to be omitted. See e.g. [1-3] below --\nand note that in particular the paradigm in many of the most recent works of\nthis line is different from that of LISTA (which is referenced and discussed),\nwhere the optimization problem is used to derive an architecture. These\ncomparisons would be helpful in situating the present work in this broader\ncontext (there may be intuitions developed here that can be applied in those\nsettings, and vice versa).\n\n[1] http://arxiv.org/abs/2102.07944\n\n[2] http://dx.doi.org/10.1109/JSAIT.2020.2991563\n\n[3] https://arxiv.org/abs/2209.04504\n\nThe assumption of a Lipschitz gradient in Theorem 3.1 does not seem to make\nsense to me in the context of multi-solution optimization. For instance,\nconsidering the example of bounding box detection given by the authors in the\nintroduction where $f_{\\tau}(x)$ is the minimum distance of the bounding box\n$x$ to a ground truth box in the image $\\tau$, at any time where there are\nmultiple boxes in the image $\\tau$, there will be boxes $x'$ that are at an\nequal distance from two or more ground truth boxes, implying that $x'$ is a\npoint of discontinuity for $\\nabla f_{\\tau}$ (and therefore the gradient is not\nLipschitz). More generally, it seems like this issue should arise any time\n$f_{\\tau}$ is a metric and there are multiple minima (this feels like a result\nfrom geometry, involving regularity of the distance function at the cut locus).\nThe authors may want to apply a global convergence theorem that does not\nrequire smooth activation functions and Lipschitz gradients, as Kawaguchi and\nHuang's does -- for example, Du et al., or Allen-Zhu et al, which are\nformulated for networks with ReLUs (hence non-Lipschitz gradients are an\ninherent part of the problem). Some discussion of this assumption in section\n3.3 seems warranted (I did not see it discussed anywhere -- the authors discuss\nGD vs. SGD, the network architecture, weak convexity).\n\nFor the experiments on symmetry detection, it would be\ninteresting to see how the approach compares to strong baselines that are\nspecifically tuned for these tasks (rather than just POL vs. GOL). For object\ndetection, the comparisons to Faster R-CNN are very interesting and seem\npromising, but would comparisons to more recent networks like YOLO (etc.) be\npossible also?\n\n## Questions / Minor Points\n\nI am curious whether other baseline approaches were attempted than just the one\nproposed at the bottom of page 5 -- it seems to me like this proposed approach\nmight be more challenging to optimize than the proximal loss (2), which is like\napproximating a single map which is iterated at test time, given that here we\nneed to approximate $Q>1$ iterations of the gradient map in general (which\nseems like it might depend a lot on the architecture of the network $\\Psi$\nbeing suitable). Is it possible to solve this problem for $Q=1$, then iterate\nthe learned mapping to simulate many steps of gradient descent (possibly with\nsome stability-promoting regularization?), or the like? \nI did not understand why the \"particle descent\" approach on the previous page\nis \"not comparable\" -- can't one just run gradient descent from many random\ninitializations given a \"test\" $\\tau$?\n\nFor the sparse recovery experiment, it could be that learning in this setting\nhas an advantage over GD due to the low-dimensionality ($m=4$, $d=8$) of the\nexperimental setup.\n\nWhy are DGCNN and ResNet used in the symmetry detection and object detection\nexperiments as encoders for the raw data? Is there a restriction that not only\n$x$ but also $\\tau$ should be low-dimensional here?\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "5IpnM0ot34", "reviewer": "Reviewer_HUdf", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "Authors propose a new estimator, inspired from proximity operator theory.\nThe proposed estimator is supposed to be able to recover multiple local minima.\nAuthors proposed multiple experiments on toy problems\n", "summary_of_the_review": "In conclusion I did not fully understand the idea and the contribution of the paper. The theoretical and practical impact of the paper seems limited.\n", "strength_and_weaknesses": "\nMajor concerns:\nI do not know if it is because this paper is too far from my research area, but I really had troubles to understand the idea and the originality of the paper.\n\n- I found the paper very verbose \"the proximal term in our loss promotes convexity of the formulation\". I am not sure to understand how the paper is related to Wasserstein graident flow. Is it only because they also parametrize le fixed-point oerator by a neural network?\n- At the end of section 3.2 authors talk about computational gains, I did not understand were do these gains come from. Would it be possible to write the training algorithm, with the cost of each step in order to understand better? In the same vein, it cold be nice to make the difference with usual L20 method clearer, from myself I did not understand.\n\nExperiments:\nI cannot comment on all the experiments, but it feels like a succession of toy examples, without proving that the method is actually useful.\nI can comment on the sparse recovery benchmark: why authors used an $\\ell$-p norm, which is not alpha semi-convex, despite their exist alpha-semi-convex sparsity promoting norm such as the Minimax Concave Penalty, or SCAD, see for instance [1]. Morover this benchmark is performed in dimesion $d=8$, and $m=4$, which is too small to be considered as a serious experiment. For problem of this scale, it seems that one can directly solve the original $\\ell_0$ constrained problem [2]\n\n[1] Breheny, P. and Huang, J., 2011. Coordinate descent algorithms for nonconvex penalized regression, with applications to biological feature selection. The annals of applied statistics, 5(1), p.232.\n\n[2] Marmin, A., Castella, M. and Pesquet, J.C., 2019, May. How to globally solve non-convex optimization problems involving an approximate ℓ 0 penalization. In ICASSP 2019-2019 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP) (pp. 5601-5605). IEEE.\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster