EnKSGD: A Class of Preconditioned Black Box Optimization and Inversion Algorithms

SIAM Journal on Scientific Computing, Volume 46, Issue 3, Page A2101-A2122, June 2024.
Abstract. In this paper, we introduce the ensemble Kalman–Stein gradient descent (EnKSGD) class of algorithms. The EnKSGD class of algorithms builds on the ensemble Kalman filter (EnKF) line of work, applying techniques from sequential data assimilation to unconstrained optimization and parameter estimation problems. An essential idea is to exploit the EnKF as a black box (i.e., derivative-free, zeroth order) optimization tool if iterated to convergence. In this paper, we return to the foundations of the EnKF as a sequential data assimilation technique, including its continuous-time and mean-field limits, with the goal of developing faster optimization algorithms suited to noisy black box optimization and inverse problems. The resulting EnKSGD class of algorithms can be designed to both maintain the desirable property of affine-invariance and employ the well-known backtracking line search. Furthermore, EnKSGD algorithms are designed to not necessitate the subspace restriction property and to avoid the variance collapse property of previous iterated EnKF approaches to optimization, as both these properties can be undesirable in an optimization context. EnKSGD also generalizes beyond the [math] loss and is thus applicable to a wider class of problems than the standard EnKF. Numerical experiments with empirical risk minimization type problems, including both linear and nonlinear least squares problems, as well as maximum likelihood estimation, demonstrate the faster empirical convergence of EnKSGD relative to alternative EnKF approaches to optimization. Reproducibility of computational results. This paper has been awarded the “SIAM Reproducibility Badge: Code and Data Available” as a recognition that the authors have followed reproducibility principles valued by SISC and the scientific computing community. Code and data that allow readers to reproduce the results in this paper are available at https://github.com/0x4249/EnKSGD and in the supplementary material (M156114_Supplementary_Materials.zip [106KB]).