Operator Splitting with Hamilton-Jacobi-based Proximals

Operator splitting algorithms are a cornerstone of modern first-order optimization, decomposing complex problems into simpler subproblems solved via proximal operators. However, most functions lack closed-form proximal operators, which has long restricted these methods to a narrow set of problems. Hamilton-Jacobi-based proximal operator (HJ-Prox) is a recent derivative-free Monte Carlo technique based on Hamilton-Jacobi PDE theory, that approximates proximal operators numerically. In this work, we introduce a unified framework for operator splitting via HJ-Prox, which allows for deployment of operator splitting even when functions are not proximable. We prove that replacing exact proximal steps with HJ-Prox in algorithms such as proximal point, proximal gradient descent, Douglas–Rachford splitting, Davis–Yin splitting, and primal–dual hybrid gradient preserves convergence guarantees under mild assumptions. Numerical experiments demonstrate HJ-Prox is competitive and effective on a wide variety of statistical learning tasks.

Many machine learning and statistics problems are solved by breaking a difficult optimization task into smaller steps. These steps often rely on mathematical tools called proximal operators, which are easy to use for some common functions but difficult or unavailable to write down for others. When these operators are unavailable, researchers usually need to design problem specific special-purpose solvers. We show that we can replace these analytical proximal operator steps with a general numerical approximation based on Hamilton-Jacobi equations. Our method only needs to evaluate the function itself. We prove that this approximation can be used inside splitting optimization algorithms, including proximal gradient descent, Douglas-Rachford, Davis-Yin, and primal-dual hybrid gradient methods. This makes optimization tools easier to apply to a broader range of machine learning problems. In experiments on sparse regression, image denoising, trend filtering, multitask learning, and gene expression analysis, our approach closely matches specialized existing methods while requiring less problem-specific mathematical work.