Fixed-Point Automatic Differentiation of Forward--Backward Splitting Algorithms for Partly Smooth Functions

TL;DR

This paper develops a fixed-point automatic differentiation method for proximal splitting algorithms applied to partly smooth functions, enabling efficient sensitivity analysis and parameter learning in structured optimization problems.

Contribution

It introduces Fixed-Point Automatic Differentiation (FPAD), reducing memory overhead and improving convergence speed over traditional reverse-mode AD for differentiating fixed-point solutions.

Findings

AD of fixed-point iterations converges to the derivative of the solution mapping.

FPAD reduces memory usage compared to reverse-mode AD.

Numerical experiments demonstrate effective sensitivity analysis and parameter learning.

Abstract

A large class of non-smooth practical optimization problems can be written as minimization of a sum of smooth and partly smooth functions. We examine such structured problems which also depend on a parameter vector and study the problem of differentiating its solution mapping with respect to the parameter which has far reaching applications in sensitivity analysis and parameter learning problems. Under partial smoothness and other mild assumptions, we apply Implicit (ID) and Automatic Differentiation (AD) to the fixed-point iterations of proximal splitting algorithms. We show that AD of the sequence generated by these algorithms converges (linearly under further assumptions) to the derivative of the solution mapping. For a variant of automatic differentiation, which we call Fixed-Point Automatic Differentiation (FPAD), we remedy the memory overhead problem of the Reverse Mode AD and moreover provide faster convergence theoretically. We numerically illustrate the convergence and convergence rates of AD and FPAD on Lasso and Group Lasso problems and demonstrate the working of FPAD on prototypical image denoising problems by learning the regularization term.