Nonlinear Splitting for Gradient-Based Unconstrained and Adjoint Optimization

Abstract

High dimensional and/or nonconvex optimization remains a challenging and important problem across a wide range of fields, such as machine learning, data assimilation, and partial differential equation (PDE) constrained optimization. Here we consider gradient-based methods for solving unconstrained and constrained optimization problems, and introduce the concept of nonlinear splitting to improve accuracy and efficiency. For unconstrained optimization, we consider splittings of the gradient to depend on two arguments, leading to semi-implicit gradient optimization algorithms. In the context of adjoint-based constrained optimization, we propose a splitting of the constraint $F(\mathbf{x},θ)$, effectively expanding the space on which we can evaluate the ``gradient''. In both cases, the formalism further allows natural coupling of nonlinearly split optimization methods with acceleration techniques, such as Nesterov or Anderson acceleration. The framework is demonstrated to outperform existing methods in terms of accuracy and/or runtime on a handful of diverse optimization problems. This includes low-dimensional analytic nonconvex functions, high-dimensional nonlinear least squares in quantum tomography, and PDE-constrained optimization of kinetic equations, where the total number of high-dimensional kinetic solves is reduced by a factor of three compared with standard adjoint optimization.