# Learning Proximal Operators to Discover Multiple Optima

arXiv: https://arxiv.org/abs/2201.11945 · HTML: https://mlsanity.com/paper/2201.11945/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 40 of its 95 references that are in the map.
- First posted on arXiv 28 Jan 2022.

## Builds on (most cited first, 30 of 40)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Microsoft COCO: Common Objects in Context](https://mlsanity.com/paper/1405.0312/) (2014)
- [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://mlsanity.com/paper/2103.14030/) (2021)
- [Faster R-CNN: Towards Real-Time Object Detection with Region Proposal Networks](https://mlsanity.com/paper/1506.01497/) (2015)
- [End-to-End Object Detection with Transformers](https://mlsanity.com/paper/2005.12872/) (2020)
- [DeepSDF: Learning Continuous Signed Distance Functions for Shape Representation](https://mlsanity.com/paper/1901.05103/) (2019)
- [Implicit Neural Representations with Periodic Activation Functions](https://mlsanity.com/paper/2006.09661/) (2020)
- [Dynamic Graph CNN for Learning on Point Clouds](https://mlsanity.com/paper/1801.07829/) (2018)
- [Density estimation using Real NVP](https://mlsanity.com/paper/1605.08803/) (2016)
- [A Convergence Theory for Deep Learning via Over-Parameterization](https://mlsanity.com/paper/1811.03962/) (2018)
- [Learning to learn by gradient descent by gradient descent](https://mlsanity.com/paper/1606.04474/) (2016)
- [Loss Surfaces, Mode Connectivity, and Fast Ensembling of DNNs](https://mlsanity.com/paper/1802.10026/) (2018)
- [Deep Ensembles: A Loss Landscape Perspective](https://mlsanity.com/paper/1912.02757/) (2019)
- [Deep Learning for Generic Object Detection: A Survey](https://mlsanity.com/paper/1809.02165/) (2018)
- [Snapshot Ensembles: Train 1, get M for free](https://mlsanity.com/paper/1704.00109/) (2017)
- [Pareto Multi-Task Learning](https://mlsanity.com/paper/1912.12854/) (2019)
- [Learning to Optimize: A Primer and A Benchmark](https://mlsanity.com/paper/2103.12828/) (2021)
- [BSP-Net: Generating Compact Meshes via Binary Space Partitioning](https://mlsanity.com/paper/1911.06971/) (2019)
- [Learning with a Wasserstein Loss](https://mlsanity.com/paper/1506.05439/) (2015)
- [Learning to Optimize](https://mlsanity.com/paper/1606.01885/) (2016)
- [CvxNet: Learnable Convex Decomposition](https://mlsanity.com/paper/1909.05736/) (2019)
- [PRNet: Self-Supervised Learning for Partial-to-Partial Registration](https://mlsanity.com/paper/1910.12240/) (2019)
- [No Spurious Local Minima in Nonconvex Low Rank Problems: A Unified Geometric Analysis](https://mlsanity.com/paper/1704.00708/) (2017)
- [Deep Equilibrium Architectures for Inverse Problems in Imaging](https://mlsanity.com/paper/2102.07944/) (2021)
- [Learning the Pareto Front with Hypernetworks](https://mlsanity.com/paper/2010.04104/) (2020)
- [Learning Proximal Operators: Using Denoising Networks for Regularizing Inverse Imaging Problems](https://mlsanity.com/paper/1704.03488/) (2017)
- [Proximal Optimal Transport Modeling of Population Dynamics](https://mlsanity.com/paper/2106.06345/) (2021)
- [Supervised Fitting of Geometric Primitives to 3D Point Clouds](https://mlsanity.com/paper/1811.08988/) (2018)
- [Large-Scale Wasserstein Gradient Flows](https://mlsanity.com/paper/2106.00736/) (2021)

## Built on by (most cited first, 4 of 4)

- [Towards Constituting Mathematical Structures for Learning to Optimize](https://mlsanity.com/paper/2305.18577/) (2023)
- [Beyond Scores: Proximal Diffusion Models](https://mlsanity.com/paper/2507.08956/) (2025)
- [Learn to Evolve: Self-supervised Neural JKO Operator for Wasserstein Gradient Flow](https://mlsanity.com/paper/2601.05583/) (2026)
- [Bifurcation Models: Learning Set-Valued Solution Maps with Weight-Tied Dynamics](https://mlsanity.com/paper/2605.07277/) (2026)

## Abstract (arXiv)

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.