Greedy Learning to Optimize with Convergence Guarantees

Abstract

Learning to optimize (L2O) is an approach that leverages training data to accelerate the solution of optimization problems. Many approaches use unrolling to parametrize the update step and learn optimal parameters. This requires choosing an unroll length in advance and typically incurs memory costs that grow with the number of unrolled iterations. Moreover, learned optimizers often lack convergence guarantees. We propose a greedy L2O framework that learns iteration-dependent parameters sequentially, by minimizing the average objective value after one additional optimization step. This converts the finite-unroll training problem into a sequence of one-step learning problems, allowing training to proceed until a stopping criterion is met while maintaining constant device-memory requirements with respect to the number of training iterations. We apply this framework to a preconditioned variant of Polyak's heavy-ball method, using several parameterizations of the learned preconditioners, including a convolutional parameterization for imaging problems. For linear parameterizations, the greedy learning subproblem is convex whenever the training objectives are convex, and admits closed-form solutions for least-squares objectives. We prove convergence in the training set even when the preconditioners are not necessarily symmetric nor positive definite. Convergence on a class of unseen functions is also obtained under certain assumptions. We test our learned algorithms on two inverse problems, image deblurring and Computed Tomography, on which learned convolutional preconditioners demonstrate improved empirical performance over classical optimization algorithms such as Nesterov's Accelerated Gradient Method and the quasi-Newton method L-BFGS.