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

Sheheryar Mehmood and Peter Ochs

University of Tübingen, Tübingen, Germany

A large class of non-smooth practical optimization problems can be written as minimization of a sum of smooth and partly smooth functions. We consider 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 optmization problems. We show that under partial smoothness and other mild assumptions, Automatic Differentiation (AD) of the sequence generated by proximal splitting algorithms converges 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 practical image denoising problem by learning the regularization term.

§ INTRODUCTION

Given a parameter $\u\in\spPrm$, we consider a composite optimization problem of the form

\begin{equation} \label{prob:min} \tag{$\mathcal P$}
\min_{\x\in\spVar} F(\x, \u), \quad F\coloneqq f+g \,,
\end{equation}

where $\map{f}{\spVar\times\spPrm}{\R}$ is $C^2$-smooth and $\map{g}{\spVar\times\spPrm}{\eR}$ is possibly non-smooth with a simple proximal mapping when the second argument is fixed. Let $\smap\Argmap\spPrm\spVar$ be such that $\Argmap(\u)\coloneqq \Argmin_{\x\in\spVar} F(\x, \u)$ gives us the set of solutions of (<ref>) depending on $\u$. Often, we are interested in computing the variation of $$, which has many far-reaching applications, for example, in Machine Learning, as we discuss in Section~\ref{ssec:intro:apps}. Clearly, (generalized) differentiation governs this information.

In general, it can be quite challenging to ensure the differentiability of $$ on some open set $U⊂$. However, in special situations, suitable approaches exist. While there are analysis tools for generalized differentiation, from a practical and computational perspective the most frequently used approach is to identify points or neighbourhoods where classical derivatives are well-defined and efficiently computable. In fact for locally Lipschitz continuous functions, thanks to Rademacher's Theorem, differentiability is asserted almost everywhere which is the foundation for constructing a commonly used generalized derivative, the Clarke subdifferential [37].

In this paper, we consider computational approaches for computing the classical derivatives under special conditions within a generally non-smooth setup and thereby broadening the understanding of these situations by providing strong theoretical guarantees. Exemplarily when $g$ in \eqref{prob:min} is 0 and $_^2 f (^*, )̆$ is positive definite for $⊂()̆$, we can employ the Implicit Function Theorem on the necessary optimality condition $_ f (^*, )̆ = 0$ to establish the local differentiability of the solution mapping around $$̆. A promising alternative to the so-called Implicit Differentiation (ID) approach, which has favourable computational aspects with fruitful potential to generalization to more complicated problems, is to iteratively construct a sequence of derivatives with increasing accuracy to the derivative of the solution mapping. In particular, we first run an algorithm for say $K\in\N$ iterations to approach the critical point of interest $\xK\approx\xmin$ with the desired accuracy (called forward pass). After that we compute the derivative $D_{\u} \xK$ by performing Automatic Differentiation (AD) on the generated computational graph of the forward pass. Good convergence, stability or robustness properties of the algorithm can often be transferred to the derivative sequence. Nevertheless, AD, which essentially relies on the classic chain rule for differentiation along the composition of iteration mapping in the forward pass, must be used cleverly to circumvent computational bottleneck such as memory overhead. Both convergence and memory issues have been addressed previously [46, 35, 48, 68, 64, 74, 18] in standard setups.

However, when $g$ is non-zero, which is mostly the case in practical applications in Machine Learning when low complexity (e.g., sparse or low-rank) solutions are sought through the formulation of a regularization [80, 85, 44, 73]. Using ID may be possible theoretically but mostly infeasible computationally. Similarly, using AD on forward–backward splitting algorithms, which are perfectly suited for problems of the form (<ref>), can be problematic theoretically and practically. Such methods involve a proximal mapping which in itself is an optimization problem and is not differentiable in general. For example, consider Proximal Gradient Descent (PGD) [60] which is an iterative algorithm expressed as $\xkp \coloneqq \pgd (\xk, \u)$ where $\map{\pgd}{\spVar\times\spPrm}{\spVar}$ is defined as

\begin{equation} \label{eq:prox:grad}
\pgd (\w, \u) \coloneqq \argmin_{\x\in\spVar} g(\x, \u) + \scal{\grad_{\x} f(\x, \u)}{\x - \w} + \frac{1}{2\alpha} \norm{\x - \w}^2 \,.
\end{equation}

Computing the derivative of $\xkp$ with respect to a parameter vector $\u$ in terms of the derivative of $$ requires differentiating $$, which clearly depends on the non-smoothness properties of $g$.

Nevertheless, non-smoothness in practical problems appears in a structured manner. In particular, low-complexity regularization is governed by the low-dimensional activity of constraint sets associated to the optimal solution $$ which, under some assumptions, are stable under minor perturbations. This activity often reveals some hidden smoothness, as it usually defines a set that is a smooth manifold $$ containing $$, i.e., the problem restricted to $$ is smooth. Fortunately, under some assumptions, this activity is identified by the algorithm in \eqref{eq:prox:grad} in a finite number of iterations [59, 58]. That is, the sequence generated by the algorithm eventually lies in $$. This is exactly the starting point for the research in this paper on various theoretical and computational aspects of AD for such algorithms. The key to enjoy this property is provided by the choice of partly-smooth functions [56]. As an example, we consider the Lasso problem [80]
\[
\min_{\x} \frac{1}{2} \norm[2]{A\x - \b}^2 + \lambda \norm[1]{\x} \,.
\]
Solving this problem yields a sparse solution for large enough $λ$, a simple but illustrative example of a low-complexity regularizer. The activity in this case is the support of $$ given by $^* { i : x^*_i ≠0 } $. It remains fixed if we change $λ$ slightly and consider the variation in $(λ)$ and is identified by forward--backward splitting algorithms [59, 58] such as PGD in a finite number of iterations. This simple idea is lifted to a broad theoretical and computational grounding with the framework of partly smooth functions. This powerful framework gives a more general Implicit Function Theorem for \eqref{prob:min} that is applicable in some non-smooth scenarios, i.e., when $g$ and hence $F$ is partly smooth (see Theorem~\ref{thm:IFT}). This not only facilitates ID but also paves way for AD since the update step of PGD is also an optimization problem of the form \eqref{eq:prox:grad} with partly smooth objective (see Section~\ref{sec:AD} for more detail).

\textbf{Contributions:} The use of AD on Proximal Gradient Descent (PGD) and Accelerated Proximal Gradient (APG) method (also known as FISTA [13]) requires theoretical justification as to whether the derivative of the iterates will converge to the true derivative of the solution mapping or not. It also begs for a memory-efficient alternative of AD, which we propose under the name of Fixed-Point Automatic Differentiation (FPAD) and demonstrate its efficiency in Sections~\ref{sssec:ADFPI:FP} and \ref{sec:FPAD}.  FPAD sequences share the same convergence rates as that of the iterates of the optimization algorithm and are most flexible to explain the differentiablity of the solution mapping of \eqref{prob:min}. In particular, our main contributions are the following:
\begin{enumerate}
\item We show convergence of the sequences obtained by AD and FPAD applied to APG to the derivative of the solution mapping of \eqref{prob:min} when $F$ belongs to a large subclass of convex partly-smooth functions which encompasses various practical problems [81].
\item We also establish convergence rates for FPAD and show that the rates of the iterates of PGD and APG are simply mirrored in the corresponding FPAD sequences.
\item We establish an equivalence of ID applied to the fixed-point equation of PGD and FPAD applied to PGD.
\item We provide experimental demonstration of convergence and convergence rates of FPAD and AD applied to PGD and APG for different practical applications, including a bilevel formulation for learning the regularizer of an image denoising problem [77] beyond the smooth setting.
\end{enumerate}

\subsection{Applications} \label{ssec:intro:apps}

Consider an optimization problem where the objective of one optimization problem (the upper or outer problem) depends on the solution of another problem (the lower or inner problem). Solving such a problem, also called the bilevel problem [40], by gradient based methods requires computing the derivative of the solution mapping of the inner problem. Bilevel problems are ubiquitous in Machine Learning. We will discuss two of the major applications here. For other applications like Meta-Learning and the so-called Learning-to-Optimize or L2O, we request the reader to go through some recent surveys [33, 52].

\textbf{Hyperparameter Learning:} When training a neural network or any machine learning model, we always encounter preselecting the values of some parameters before staring the training process. These so-called hyperparameters play an important role in the output of the model and therefore have to be selected optimally. They remain fixed during the training process, i.e., solving the lower-level problem. They are learned by optimizing a different loss function (since it depends on different dataset called the validation set) which depends on the trained parameters from the lower problem. They can be learned through grid search [15] or Bayesian Optimization [66, 27]. More recently the trend has shifted to gradient-based approach to compensate for the growing number of hyperparameters [14, 41, 54, 39, 62, 72]. This approach involves differentiating the upper level problem with respect to the hyperparamters by differentiating the solution mapping of the lower-level problem through Implicit or Automatic Differentiation and has been put to test to solve many practical problems like image denoising [41, 54], image segmentation [70] and data cleaning [45].

\textbf{Implicit Models:} In recent years, there has been a tendency towards embedding a function defined by an implicit equation inside the prediction function of the machine learning model. Training such models is equivalent to solving a bilevel problem with the implicit equation being the lower problem and the training of the model being the upper problem. Examples of implicit models include Deep Equilibrium Models [7, 8, 84], Neural ODE's [32] and Optimization Layers [4, 65, 2, 16].
When the implicit equation in such a model represents some inherent structure of the practical application, the learned model may provide better interpretability and generalization.

% ************************
% >>>>> RELATED WORK <<<<<
% ************************
\subsection{Related Work}

We discuss related work whose methodology applies to certain classes of non-smooth problems. In [39], the authors considered algorithms where the update step is Lipschitz continuous and hence weakly differentiable. This holds for proximal splitting algorithms for solving problems of type \eqref{prob:min} such as PGD in \eqref{eq:prox:grad}. This allowed them to use AD on such algorithms thanks to the chain rule. Unfortunately, convergence of the sequence of derivatives that is generated by AD cannot be shown. In particular, it is unclear what exactly would be the correct generalized derivative concept of the solution mapping. Moreover, weak derivatives are not defined pointwise which makes their use in optimization problematic. Ochs et al. [69, 70] replaced the non-smooth problem with a finite number of steps of a smooth algorithm or its corresponding fixed-point equation. The smooth algorithm was obtained by carefully selecting the kernel generating distance function for Bregman Proximal Gradient [26, 10]. They were then able to solve a bilevel problem through AD and ID. The key advantage is the structure preserving approximation by a classical derivative, which can be employed to optimization. However, no convergence guarantees for the derivative sequence is shown.

Bolte et al. [20] used the conservative calculus developed in [21, 22] to rigorously study the implicit differentiation performed by using standard autograd packages like TensorFlow [1], PyTorch [71], and JAX [25] etc. on possibly non-smooth implicit equations. Thereby, they provide a mathematical framework for the types of ``derivatives'' that are used in many practical applications. The same tools in [21, 22] were applied to study the application of autograd packages to differentiate the output of a non-smooth iterative algorithm in [23].

Christof [36] considered a class of bilevel-problems where the non-smooth lower problem is replaced with an elliptic variational inequality \cite[\DRSecs{2.1}{2A}]{DR09}. He showed that the solution mapping for this type of inequality is Fr{\'e}chet differentiable and can be obtained by solving an analogous elliptic variational equality.

In [18], Bertrand et al. consider bilevel optimization problems with non-smooth, lasso-like lower-level problems. They compute the derivative of the solution mapping by implicit differentiation of the fixed-point equation for Proximal Gradient Descent and Proximal Coordinate Descent. They also show linear convergence of forward mode automatic differentiation of the two algorithms to the derivative obtained by the corresponding fixed-point equations. However, as opposed to our work, their framework is quite restrictive in the sense that it is only applicable to the cases where the non-smooth component in \eqref{prob:min} is expressed as
\[
\R^N\times\R\ni(\x, u) \mapsto g(\x, u) = \sum_{i=1}^{N} g_i (x_i, u) \,,
\]
with $g_i×$ and therefore does not include many interesting practical regularizers like $ l_∞$ norm, $l_2,1$ group norm and nuclear norm [81].

Riis [74] studied the bilevel problems with partly smooth lower problems of the form \eqref{prob:min}. In particular, he showed the convergence of AD applied to APG. He showed that when a non-degeneracy assumption \eqref{eq:ND} holds, the derivative sequence converges to the derivative of the minimizer. He also showed piecewise smoothness \cite[\DRSecs{2.4}{2D}]{DR09} of the solution mapping and the update step for APG when \eqref{eq:ND} is violated. However, the memory-overhead of Reverse Mode AD is not addressed. Moreover assumptions on the objective $F$ such as strongly convex is restrictive and excludes many practical problems like the following matrix-decomposition optimization problem
\[
\min_{X, Y \in \R^{M\times N}} \frac{1}{2} \norm[2]{X+Y-A}^2 + \gamma \norm[1]{X} + \lambda \norm[\mathrm{nuc}]{Y} \,,
\]
where $[2]·$, $[1]·$ and $[nuc]·$ are the Frobenius norm, $l_1$ norm and nuclear norm defined on matrices in $^M×N$ respectively. The above problem is used in background subtraction tasks, e.g., in surveillance camera setup. Assuming that the regularization parameters $γ>0$ and $λ>0$ are carefully chosen, given a matrix $A∈^M×N$, for example, an image sequence (each of its rows is a vectorized image frame), solving the above problem yields the sparse background sequence $X^*$ and a slowly-changing (hence low-rank) foreground sequence $Y^*$. The above problem is not strongly convex but can still satisfy a weaker condition like Assumption~\ref{ass:RPD}.

Chambolle and Pock [30] apply a Piggyback-style technique [49] to learn the linear operator $K$ in the following saddle point problem.
\begin{equation} \label{prob:sad}
\min_{\x\in\spE} \max_{\y\in\spF} \scal{K\x}{\y} + g(\x) - f^* (\y) \,,
\end{equation}
where $g$ and $f^*$ are assumed to be convex and $C^2$-smooth. Given a linear operator $K$ and the solution $(^*, ^*)$ of \eqref{prob:sad}, they show that the gradient of the loss function $L(, )$, given by $L(K)l(^*, ^*)$, can be computed through the adjoint states $(X^*, Y^*)$ which are obtained by solving yet another saddle-point problem
\begin{equation} \label{prob:sad:adj}
\min_{X\in\spE} \max_{Y\in\spF} \scal{KX}{Y} + \frac{1}{2} \scal{\grad^2 g(\x^*)X}{X} - \frac{1}{2} \scal{\grad^2 f^*(\x^*)Y}{Y} + \scal{\grad l (\x^*, \y^*)}{(X, Y)} \,.
\end{equation}
While $(^*, ^*)$ appear in the update step of the primal-dual algorithm applied to \eqref{prob:sad:adj}, the authors show that at every iteration $k∈$, they can be replaced by $(, )$ which are generated by applying the primal-dual algorithm to \eqref{prob:sad}. More recently, the convergence guarantees for this technique are given with even weaker assumptions of $g$ and $f^*$ being convex and $C^1$ [19].

\section{Notation and Preliminaries} \label{sec:NP}

In the rest of the paper, we use $_0 ∪0$ and define $[K]0,…,K$ for any $K∈$. We denote the set of extended real numbers by $∪{-∞, +∞} $ and assume that $$ and $$ are Euclidean spaces of dimensions $N$ and $P$ respectively, each equipped with an inner product $··$ and an induced norm $·$. The corresponding dual spaces are written as $^*$ and $^*$. The set of all linear maps from $$ to $$ is denoted by $(, )$. The elements of $,  (and  )$ and $(, )$ are denoted by lower-case normal (e.g., $t$), lower-case bold (e.g., $$) and upper-case normal (e.g., $A$) letters respectively. For $∈^N$ we write $(x_1 , …, x_N ) = (x_i)_i=1,…,N$ and for $A∈^N×P$ we write $A =(a_i,j)_i=1,…,N,j=1,…,P$. In the following subsections, we provide a gentle introduction to the fundamentals of our contribution to make the results widely accessible.

\subsection{Riemannian Geometry} \label{ssec:NP:RG}

We recap here a few definitions and results from Riemannian Geometry. For further study, we refer the reader to [55, 31].

\begin{DEF}[Manifold]
We say that $\manif\subset\spVar$ is a $C^2$-smooth $m$-dimensional submanifold of $\spVar$ if for every point $\x\in\manif$ there exists an open set $X\subset\spVar$ and a $C^2$-smooth map $\map{\Phi}{X}{\R^{N-m}}$ such that $\x\in X$, the derivative $D \Phi (\x)$ is surjective, and $U\cap\manif = \Phi^{-1} (0) = \set{\y \in U : \Phi (y) = 0}$.
\end{DEF}

In the rest of the paper we will refer a $C^2$-smooth $m$-dimensional submanifold by simply manifold.

\begin{DEF}[Tangent and Normal Spaces]
Let $\manif\subset\spVar$ be a manifold and $\x\in\manif$. We say that $\v\in\spVar$ is a tangent vector of $\manif$ at $\x$ if there exists $\epsilon>0$ and a $C^1$-smooth curve $\map{\gamma}{(-\epsilon, \epsilon)}{\manif}$ on $\manif$ with $\gamma (0) = \x$ and $\dot\gamma (0) = \v$. The set of all tangent vectors of $\manif$ at $\x$ constitute $\spTan\x$, the tangent space of $\manif$ at $\x$. We define the normal space of $\manif$ at $\x$ by $\spNor\x \coloneqq (\spTan\x)^\perp$, the orthogonal complement of $\spTan\x$.
\end{DEF}

We define the functions $ \map{\projTan, \projNor}{\manif}{\spLin (\spVar, \spVar)} $ that provides for any $\x\in\manif$, the projection onto the tangent space $\projTan (\x) \coloneqq \proj{\spTan{\x}}$ and the projection onto the normal space $\projNor (\x) \coloneqq \proj{\spNor{\x}}$.

\begin{DEF}[Riemannian Gradient]
Let $\manif\subset\spVar$ be a manifold, $\map{f}{\manif}{\R}$ be a function and $\x\in\manif$. We say that $f$ is $C^2$-smooth at $\x$ if there exists a neighborhood $X\subset\spVar$ of $\x$ and a $C^2$-smooth function $\map{\tilde{f}}{X}{\R}$ such that $\tilde{f}$ agrees with $f$ on $\manif \cap X$. In this case, we call $\tilde{f}$ a smooth extension of $f$ around $\x$. We call $\grad_{\manif} f(\x) \in\spTan\x$, the Riemannian gradient of $f$ at $\x$ if for all $\v\in\spTan\x$
\[
\scal{\grad_{\manif} f(\x)}{\v} = (f \circ \gamma )^\prime (0) \,,
\]
where $\map{\gamma}{(-\epsilon, \epsilon)}{\manif}$ is any $C^1$-curve with $\gamma (0) = \x$ and $\dot\gamma (0) = \v$.
\end{DEF}

The Riemannian Gradient of $f$ can alternatively be expressed in terms of the gradient of the smooth extension $\tilde{f}$ of $f$ by
\begin{equation} \label{eq:Riem:grad}
\grad_{\manif} f (\x) = \projTan (\x) \grad \tilde{f} (\x) \,.
\end{equation}
Note that this gradient does not depend on the choice of both the curve $\gamma$ and the smooth extension $\tilde{f}$.

\begin{DEF}[Riemannian Hessian]
Let $\manif\subset\spVar$ be a manifold, $\map{f}{\manif}{\R}$ be a function and $\x\in\manif$. We call $\map{\grad_{\manif}^2 f(\x)}{\spTan\x}{\spTan\x}$, the Riemannian Hessian of $f$ at $\x$ if for all $\v\in\spTan\x$
\[
\scal{\grad_{\manif}^2 f(\x) \v}{\v} = (f \circ \gamma )^{\prime\prime} (0) \,,
\]
where $\map{\gamma}{(-\epsilon, \epsilon)}{\manif}$ is any $C^1$-curve with $\gamma (0) = \x$ and $\dot\gamma (0) = \v$.
\end{DEF}

We can similarly express the Riemannian Hessian $\grad_{\manif}^2 f(\x) \in \spLin (\spTan\x, \spTan\x)$ by using the smooth extension $\tilde{f}$, i.e.,
\begin{equation} \label{eq:Riem:hess}
\grad_{\manif}^2 f (\x) = \projTan (\x) \grad^2 \tilde{f} (\x) + \wein{\cdot}{\projNor (\x) \grad \tilde{f} (\x)} \,,
\end{equation}
which however requires the mapping $\wein{\cdot}{\w} \in \spLin (\spTan\x, \spTan\x)$ for $\w\in\spNor\x$ is called the Weingarten map and is defined by
\[
\v \mapsto \wein{\v}{\w} \coloneqq - \projTan (\x) \mathrm{d} W [\v] \,,
\]
where $W$ is a local extension of $\w$ to a normal vector field on $\manif$. $\wein{\cdot}{\w}$ is independent of the choice of normal field $W$ \cite[Proposition~II.2.1]{Cha06} and caters for the change of the tangent space as we move away from $\x$. It vanishes when the manifold is affine, i.e., $\manif = \x + \spTan\x $ reducing the Hessian expression to $\grad_{\manif}^2 f (\x) = \projTan (\x) \grad^2 \tilde{f} (\x)$. To extend the domain of $\grad_{\manif}^2 f(\x)$ to $\spVar$, we apply $\projTan (\x)$ before applying $\grad_{\manif}^2 f(\x) \in \spLin (\spTan\x, \spTan\x)$.

\subsection{Partial Smoothness} \label{ssec:PS}
We are now ready to define partial smoothness of a function as introduced in [56], which defines large class of functions including the loss and regularization functions that are typically used in Machine Learning (see \cite[Section~3]{VDP+17} for an extensive list of examples). The definition lifts key properties for handling constraint sets (for example, active set approaches) to the world of non-smooth functions in more generality.
\begin{DEF}[Partial Smoothness] \label{def:part:C2}
Let $\map{f}{\spVar}{\eR}$ be proper and lower semi-continuous and $\manif\subset\spVar$ be a set. We say that $f$ is partly smooth at a point $\x\in\manif$ relative to $\manif$ if the following conditions hold:
\begin{enumerate}
\item (Regularity:) $f$ is regular at $\x$ and $\partial f (\x) \neq \emptyset$.
\item (Smoothness:) $\manif$ is a $C^2$-smooth manifold and $\rstDom f \manif$ is $C^2$ around $\x$.
\item (Sharpness:) $\spNor\x = \spPar \partial f(\x)$.
\item (Continuity:) $\partial f$ is continuous at $\x$ relative to $\manif$.
\end{enumerate}
We call $f$ partly smooth relative to $\manif$ if $f$ is partly smooth at every $\x\in\manif$ relative to $\manif$.
\end{DEF}
\begin{REM}
The first condition in Definition~\ref{def:part:C2} is automatically satisfied when $f$ is convex on $\ri (\dom f)$ and $\x \in \ri (\dom f)$. However when $f$ is non-convex, this assumption needs to be verified additionally. For more on regularity, we request the reader to look into \cite[Chapter~8]{RW98}.
\end{REM}

\subsection{Matrix Analysis} \label{ssec:NP:MA}
In this section we provide some results which will be useful in the later sections. We start with showing the convergence of a generalized linear fixed-point iteration sequence.
\begin{LEM} \label{lem:LGFPI:conv}
Let $( B_k )_{k\in\N}$ and $ ( \bk)_{k\in\N}$ be sequences in $\spLin(\spVar, \spVar)$ and $\spVar$ with limits $B$ and $\b$ respectively. When $\rho (B) < 1$, the sequence $( \xk )_{k\in\N}$, with $\xz\in \spVar$, generated by
\[
\xkp \coloneqq B_k \xk + \bk \,,
\]
converges to $\x^*\coloneqq(I - B)^{-1} \b$. The convergence is linear when $( B_k)_{k\in\N_0}$ and $(\bk)_{k\in\N}$ converge linearly.
\end{LEM}
\begin{proof}
Following \cite[Proposition~2.7]{Rii20}, we set $\zk\coloneqq\xk - \x^*$ with $\x^* \coloneqq (I-B)^{-1}\b$, $\yk\coloneqq (B_k - B)\x^* + \bk - \b$ and $\wk\coloneqq (B_k - B)\zk $ which leads to
\begin{equation*}
\begin{aligned}
\zkp &= (B_k \xk + \bk) - (B\x^* + \b) \\
&= B \zk + \wk + \yk \,.
\end{aligned}
\end{equation*}
Since $\wk = o(\zk)$, for any $\epsilon\in (0, 1-\rho)$ we can find $K\in\N$ such that $\norm\wk / \norm\zk \leq \epsilon/2 $ for all $k\geq K$ where we abbreviate $\rho\coloneqq\rho (B)$. Also, there exists a norm $\norm[\epsilon/2]{\cdot}$ on $\spLin (\spVar, \spVar)$ such that $\norm[\epsilon/2]{B}\leq \rho+\epsilon/2$. Therefore, for the vector norm $\map{\norm{\cdot}}{\spVar}{\R}$ which is consistent with the norm $\map{\norm[\epsilon/2]{\cdot}}{\spLin(\spVar, \spVar)}{\R}$, we obtain
\begin{equation*}
\begin{aligned}
\norm{\zkp} &\leq (\rho + \epsilon/2) \norm\zk + \norm\wk + \norm\yk \\
&\leq (\rho + \epsilon) \norm\zk + \norm\yk \,.
\end{aligned}
\end{equation*}
Recursive Expansion of the above expression yields
\begin{equation*}
\norm{\zkp} \leq (\rho + \epsilon)^{k-K+1} \norm{\zK} + \sum_{i=K}^{k} (\rho + \epsilon)^{k-i} \norm{\y^{(i)}} \,.
\end{equation*}
Using the fact that $\yk\to 0$, for any $\delta>0$, we can make $K$ large enough to obtain $\norm{\y^{(i)}} \leq (1-\rho-\epsilon) \delta/2$ for all $i\geq K$. This allows for the following reduction
\begin{equation*}
\begin{aligned}
\norm{\zkp} &\leq (\rho + \epsilon)^{k-K+1} \norm{\zK} + \sum_{i=K}^{k} (\rho + \epsilon)^{k-i} (1-\rho-\epsilon) \frac{\delta}{2} \\
&\leq (\rho + \epsilon)^{k-K+1} \norm{\zK} + (1-\rho-\epsilon) \frac{\delta}{2} \sum_{i=0}^{\infty} (\rho + \epsilon)^i \\
&= (\rho + \epsilon)^{k-K+1} \norm{\zK} + \frac{\delta}{2} \,.
\end{aligned}
\end{equation*}
Also, the quantity $(\rho + \epsilon)^{k-K+1} \norm{\zK}$ eventually becomes smaller than $\delta/2$ as $k\geq K$ grows. In particular, there exists $N\geq K$ such that $(\rho + \epsilon)^{k-K+1} \norm{\zK} < \delta/2$ and $\norm{\zkp} < \delta$ for all $k\geq N$. This concludes the proof for convergence. The proof of linear convergence of $\xk$ under the linear convergence of $B_k$ and $\bk$ (and hence $\yk$) follows from \cite[Lemma~6.18]{Ber21}.
\end{proof}
Lemma~\ref{lem:LGFPI:conv} allows us to prove the following more general statement. We use Painlev{\'e}-Kuratowski notion of set convergence \cite[Section~4B]{RW98} to define the convergence of $( V_k)_{k\in\N_0}$ to $V$. However, the convergence of the projections mappings $( \proj{V_k})_{k\in\N_0}$ to $\proj V$ can also be alternatively used for this purpose since the two definitions are equivalent for subspaces \cite[4.9~Proposition]{RW98}.
\begin{COR} \label{cor:LGFPI:conv}
Let $( V_k)_{k\in\N_0} \subset \spVar$ be a sequence of subspaces that converge to $V\subset\spVar$. Let $( B_k )_{k\in\N_0}$ and $( \bk )_{k\in\N_0}$ be sequences in $\spLin(\spVar, \spVar)$ and $\spVar$ with limits $B$ and $\b$ respectively such that $\im B_k \subset V_{k+1}$ and $\bk\in V_{k+1}$ for all $k\geq0$. When $\rho (B\circ\proj V) < 1$, the sequence $( \xk )_{k\in\N}$, with $\xz\in \spVar$, generated by
\[
\xkp \coloneqq B_k \xk + \bk \,,
\]
converges to $(I - B\circ\proj V)^{-1} \b$ which lies in $V$. The convergence is linear when $( V_k)_{k\in\N_0}$, $( B_k)_{k\in\N_0}$ and $( \bk)_{k\in\N}$ converge linearly.
\end{COR}
\begin{proof}
Since $\xk\in V_{k}$ for all $k\in\N$ we can rewrite the iteration mapping as
\[
\xkp \coloneqq (B_k\circ\proj{V_k}) \xk + \bk \,.
\]
First note that $\xmin\coloneqq(I - B\circ\proj V)^{-1} \b$ is well defined and unique. The (linear) convergence of $\xk$ to $\xmin$ then follows from Lemma~\ref{lem:LGFPI:conv} because $B_k\circ\proj{V_k}\to B\circ\proj V$, $\rho (B\circ\proj V) < 1$ and the projection onto a subspace is a linear operator. To check if $\xmin = (B\circ\proj V) \xmin + \b$ lies in $V$, we observe that $\b\in V$ and $\im (B) \subset V$, because for any $\z\in\spVar$, $B_k\z\in V_{k+1}$ which implies $B\z = \lim_{k\to\infty} B_k\z \in V$ by assumption of converging sequence of subspaces.
\end{proof}

\subsection{Automatic Differentiation} \label{ssec:NP:AD}
Automatic or Algorithmic differentiation, as the name suggests is a way of computing the derivatives of a function, given as a computer program, automatically. The key is the classic chain rule and a representation of a function as composition of a finite number of elementary functions like polynomials, logarithms, exponential and trigonometric functions etc. From classical Calculus, the derivatives of these elementary functions are known and therefore are combined via chain rule to obtain the derivative of the given function. There are two main modes of AD, namely the forward and the reverse mode which we briefly recall here. For more on AD and its applications, the reader is requested to look into [50, 12] and the references therein.

Let $\mathcal X$, $\mathcal Y$, and $\mathcal Z$ be Euclidean spaces. Consider a function $\map{h}{\mathcal X}{\mathcal Z}$ given by
\[
h \coloneqq (h_3 h_2) \circ h_1 \,,\quad h_1 \textnormal{ composed with the product } h_3 h_2 \,
\]
where $\map{h_1}{\mathcal X}{\mathcal Y}$, $\map{h_2}{\mathcal Y}{\mathcal Z}$, and $\map{h_3}{\mathcal Y}{\R}$ are functions with known derivatives. We compute the derivative of $h$ at some $\x \in \mathcal X$ via the two modes of AD in Sections~\ref{sssec:AD:FM} and \ref{sssec:AD:RM}. We first use some intermediate variables to break down our example in the following manner
\begin{equation}
\begin{aligned}
\y &= h_1 (\x) \,, \\
\w &= h_2 (\y) \,, \; a = h_3 (\y) \,, \\
\z &= a\w \,.
\end{aligned}
\end{equation}
\begin{figure}[h]
\centering
\tikzset{font={\fontsize{10pt}{12}\selectfont}}
\begin{tikzpicture}[->, semithick, node distance=3.5cm, >=stealth', auto]
\tikzstyle{place} = [circle, thick, draw=black!75, fill=black!10, minimum size=6mm]
\begin{scope}

\node [place] (x) [label=above:$\x$] {};
\node [place] (y) [right of=x, label=above:$\y$] {};
\node [place] (w) [right of=y, yshift=+1.6cm, label=above:$\w$] {};
\node [place] (a) [right of=y, yshift=-1.6cm, label=above:$a$] {};
\node [place] (z) [right of=w, yshift=-1.6cm, label=left:$\times$, label=above:$\z$] {};

\path[every node/.style={sloped,anchor=south,auto=false}]
(x) edge node [above, pos=0.5] {$Dh_1 (\x)$} (y)
(y) edge node [above, pos=0.5] {$Dh_2 (\y)$} (w)
(y) edge node [above, pos=0.5] {$Dh_3 (\y)$} (a)
(w) edge node [above, pos=0.5] {$a$} (z)
(a) edge node [above, pos=0.5] {$\w$} (z);

\end{scope}
\end{tikzpicture}
\caption{Computational Graph of $ h(\bm{x}) $.}
\label{fig:CG}
\end{figure}
This allows us to construct the computational graph for $h$ as shown in Figure~\ref{fig:CG}. A node represents a variable while a directed edge shows the dependence of one variable on another (represented by the nodes it connects) through some elementary function or operation. In the figure, the nodes are labelled by the variables while the edges are denoted by the derivative (Jacobian) of the function they represent.
\subsubsection{Forward Mode AD} \label{sssec:AD:FM}
The forward mode derivative $\zd$ of $\z$ is represented by placing a dot on it and is computed by starting with some $\xd\in\mathcal X$ and computing the directional derivative of $\z$ with respect to $\x$ along $\xd$, i.e.,
\begin{equation}
\begin{aligned}
\yd &= Dh_1 (\x) \xd \,, \\
\dot\w &= h_2 (\y) \yd\,, \; \dot a = Dh_3 (\y) \yd \,, \\
\zd &= a\dot\w + \dot a\w \,.
\end{aligned}
\end{equation}
\begin{figure}[h]
\centering
\tikzset{font={\fontsize{10pt}{12}\selectfont}}
\begin{tikzpicture}[->, semithick, node distance=3.5cm, >=stealth', auto]
\tikzstyle{place} = [circle, thick, draw=black!75, fill=black!10, minimum size=6mm]
\begin{scope}

\node [place] (x) [label=above:$\x$] {};
\node [place] (y) [right of=x, label=above:$\y$] {};
\node [place] (w) [right of=y, yshift=+1.6cm, label=above:$\w$] {};
\node [place] (a) [right of=y, yshift=-1.6cm, label=above:$a$] {};
\node [place] (z) [right of=w, yshift=-1.6cm, label=left:$\times$, label=above:$\z$] {};
\node [dashed, place] (t) [left of=x, label=above:$t$] {};

\path[every node/.style={sloped,anchor=south,auto=false}]
(t) edge [dashed] node [above, pos=0.5] {$\xd=D\xi(t)$} (x)
(x) edge node [above, pos=0.5] {$Dh_1 (\x)$} (y)
(y) edge node [above, pos=0.5] {$Dh_2 (\y)$} (w)
(y) edge node [above, pos=0.5] {$Dh_3 (\y)$} (a)
(w) edge node [above, pos=0.5] {$a$} (z)
(a) edge node [above, pos=0.5] {$\w$} (z);

\end{scope}
\end{tikzpicture}
\caption{Depiction of Forward Mode AD of $ h $.}
\label{fig:FAD}
\end{figure}
We can also motivate the forward mode AD in the following manner. Assume that $\x$ is obtained by applying a differentiable function $\map{\xi}{\R}{\mathcal X}$ on some scalar $t\in\R$, i.e., $\x=\xi(t)$. The forward mode derivative $\xd$ is then interpreted as the derivative of $\xi$ at $t$, i.e., $\xd = D\xi(t)$ (see Figure~\ref{fig:FAD}). In other words $\xd$ is the derivative of the variable $\x$ with respect to $t$, at $t$, i.e., $\xd=\dydx{\x}{t}(t)$. In fact, all the forward mode derivatives $\yd$, $\dot\w$, $\dot a$ and $\zd$ are interpreted as the derivatives $\dydx{\y}{t}(t)$, $\dydx{\w}{t}(t)$, $\dydx{a}{t}(t)$ and $\dydx{\z}{t}(t)$. The intermediate variables $\y$, $\w$ and $a$ need not be stored since they can be computed alongside the derivatives $\yd$, $\dot\w$, $\dot a$ and $\zd$. However, even with no memory overhead, forward mode AD becomes impractical when the dimension of $\mathcal X$ is significantly larger than that of $\mathcal Z$, since, for example, computing the derivative $D h(\x)$ requires $\dim (\mathcal X)$ runs of the forward mode, once for each directional derivative.

\subsubsection{Reverse Mode AD} \label{sssec:AD:RM}
The reverse mode derivatives are represented by bars on the variables. In this mode, we start at the output and move towards the input. That is, given $\bar\z\in\mathcal Z^*$, we compute $\bar\x\in\mathcal X^*$ by
\begin{equation} \label{eq:ADeg:RM}
\begin{aligned}
\bar a &= \bar\z \w\,, \; \bar\w = a \bar\z \,, \\
\bar\y &= \bar a Dh_3 (\y) + \bar\w Dh_2 (\y) \,, \\
\bar\x &= \bar\y Dh_1 (\x) \,.
\end{aligned}
\end{equation}
We can motivate the reverse mode analogously to the forward mode. Assume that $\z$ is fed into a scalar-valued function $\map{\zeta}{\mathcal Z}{\R}$ to obtain $s = \zeta(\z)$. The reverse mode derivative is then treated as the derivative of $\zeta$ at $\z$, i.e., $\bar\z=D\zeta(\z)=\dydx{s}{\z}(\z)$ (Figure~\ref{fig:RAD}). Similarly, $\bar\x$, $\bar a$, $\bar\w$ and $\bar\y$ are interpreted as the derivatives $\dydx{s}{\x}(\x)$, $\dydx{s}{a}(a)$, $\dydx{s}{\w}(\w)$ and $\dydx{s}{\y}(\y)$ respectively. This scheme is referred to as back-propagation in the Machine Learning community [78] and is useful when the dimension of the output space is very small, e.g., a (scalar-valued) loss function. Its computation cost is independent of the size of the input thus making it the preferred mode in Deep Learning. However, since we are moving backward when computing the derivatives, we cannot compute the intermediate variables $\y$, $\w$ and $a$ in parallel anymore. We must store them beforehand to efficiently use this mode which causes a memory overhead, as can be seen by the dependence of each step in \eqref{eq:ADeg:RM}.
\begin{figure}[h]
\centering
\tikzset{font={\fontsize{10pt}{12}\selectfont}}
\begin{tikzpicture}[->, semithick, node distance=3.5cm, >=stealth', auto]
\tikzstyle{place} = [circle, thick, draw=black!75, fill=black!10, minimum size=6mm]
\begin{scope}

\node [place] (x) [label=above:$\x$] {};
\node [place] (y) [right of=x, label=above:$\y$] {};
\node [place] (w) [right of=y, yshift=+1.6cm, label=above:$\w$] {};
\node [place] (a) [right of=y, yshift=-1.6cm, label=above:$a$] {};
\node [place] (z) [right of=w, yshift=-1.6cm, label=left:$\times$, label=above:$\z$] {};
\node [dashed, place] (s) [right of=z, label=above:$s$] {};

\path[every node/.style={sloped,anchor=south,auto=false}]
(z) edge [dashed] node [above, pos=0.5] {$\bar\z=D\zeta(\z)$} (s)
(x) edge node [above, pos=0.5] {$Dh_1 (\x)$} (y)
(y) edge node [above, pos=0.5] {$Dh_2 (\y)$} (w)
(y) edge node [above, pos=0.5] {$Dh_3 (\y)$} (a)
(w) edge node [above, pos=0.5] {$a$} (z)
(a) edge node [above, pos=0.5] {$\w$} (z);

\end{scope}
\end{tikzpicture}
\caption{Depiction of Reverse Mode AD of $ h $.}
\label{fig:RAD}
\end{figure}

\subsection{AD of Fixed-Point Iterations} \label{ssec:NP:ADFPI}

Given $\map{g}{\spVar\times\spPrm}{\spVar}$, $\xz\in\spVar$, and $\u\in\spPrm$, we consider the following fixed-point iterations for $k\geq0$
\begin{equation} \label{itr:IM} \tag{IM}
\xkp \coloneqq g(\xk, \u) \,.
\end{equation}
A fixed-point $\xmin$ of $g(\cdot, \u)$, i.e., which satisfies $\xmin = g(\xmin, \u)$ for some given $\u$ is guaranteed to exist and is obtainable by \eqref{itr:IM} under certain standard assumptions \cite[Theorem~4.1.3]{AH05}. Our goal is to compute the derivative of $\xmin$ as a function depending on $\u$. One way to do so is through Implicit Differentiation. When $\rho (D_{\x} g (\xmin, \u)) < 1$, we apply Implicit Function Theorem \cite[Theorem~1B.1]{DR09} on the fixed-point equation $ \xmin = g(\xmin, \u) $ to obtain
\begin{equation} \label{eq:FPI:IFT}
D_{\u} \xmin = (I - D_{\x} g (\xmin, \u))^{-1} D_{\u} g (\xmin, \u) \,.
\end{equation}
Obviously this requires inverting the linear operator $I - D_{\x} g (\xmin, \u)$ which become impractical when $\dim \spVar$ become very large. A more favourable strategy to approximate this derivative would be by running \eqref{itr:IM} for a finite number of iterations, say $K\in\N$, to obtain $\xK$ and computing its derivative with respect to $\u$ by unrolling the iterative mapping as depicted in Figure~\ref{fig:IM} and applying AD.
\begin{figure}[h]
\centering
\tikzset{font={\fontsize{8pt}{12}\selectfont}}
\begin{tikzpicture}[->, semithick, node distance=4cm, >=stealth', auto]
\tikzstyle{place} = [circle, thick, draw=black!75, fill=black!10, minimum size=6mm]
\begin{scope}

\node [place] (xz) [label=above:$\xz$] {};
\node [place] (x1) [right of=xz, label=above:$\x^{(1)}$] {};
\node [place] (xKm) [right of=x1, label=above:$\x^{(K-1)}$] {};
\node [place] (xK) [right of=xKm, label=above:$\xK$] {};
\node [place] (u) [below of=xz, label=above left:$\u$] {};

\path[every node/.style={sloped,anchor=south,auto=false}]
(xz) edge node [above, pos=0.5] {$D_{\x} g (\xz, \u)$} (x1)
(xKm) edge node [above, pos=0.5] {$D_{\x} g (\x^{(K-1)}, \u)$} (xK)
(x1) edge [dotted] node {} (xKm)
(u) edge [bend right] node [above, pos=0.6] {$D_{\u} g (\xz, \u)$} (x1)
(u) edge [bend right] node [above, pos=0.6] {$D_{\u} g (\x^{(K-2)}, \u)$} (xKm)
(u) edge [bend right] node [above, pos=0.6] {$D_{\u} g (\x^{(K-1)}, \u)$} (xK);

\end{scope}
\end{tikzpicture}
\caption{Depiction of Unrolling of \eqref{itr:IM}.}
\label{fig:IM}
\end{figure}
When $g$ is continuously differentiable, the computational graph in Figure~\ref{fig:IM} allows us to apply the two modes of AD from Sections~\ref{ssec:AD:FAD} and \ref{ssec:AD:RAD} to \eqref{itr:IM}. Thus, for the forward mode, given $\xzd \coloneqq 0\in\spVar$ and $\ud\in\spPrm$, we perform the following iterations for $k = 0, \ldots, K-1$
\begin{equation} \label{itr:IM:FAD} \tag{IM-F}
\xkpd \coloneqq D_{\x} g (\xk, \u) \xkd + D_{\u} g (\xk, \u) \ud \,,
\end{equation}
to obtain the forward derivative $\xKd$. We similarly apply reverse mode AD to \eqref{itr:IM} by starting with $\xKb\coloneqq \xminb\in\spVar^*$, and $\uzb\coloneqq 0\in\spPrm^*$ and computing for $n = 0,\ldots,K-1$,
\begin{equation} \label{itr:IM:RAD} \tag{IM-R}
\begin{aligned}
\xkb &\coloneqq \xkpb D_{\x} g (\xk, \u) \\
\unpb &\coloneqq \unb + \xkpb D_{\u} g (\xk, \u) \,,
\end{aligned}
\end{equation}
where $k \coloneqq K - n - 1$. Notice that in \eqref{itr:IM:RAD}, we introduce a new index $n$ which increases as we move from right (output) to left (input) along the graph in Figure~\ref{fig:IM} to distinguish it from $k$. The subscript $K$ in the reverse mode derivatives here is due to the fact that we are computing the derivative of $\xK$. The output of reverse mode AD applied to \eqref{itr:IM} is $\ubK \coloneqq \uKb$. AD provides the correct estimate for $D_{\u} \xmin$ under certain regularity assumptions which are independent of the choice of $\xz$ [46].

As stated earlier, reverse mode is preferable when the input of the function has a significantly larger dimensionality than its output. This happens for instance when $\xK$ is simply fed into a real-valued function, e.g., a loss or penalty function [7] and the gradient of this function with respect to $\u$ is being sought. Due to the memory overhead however, reverse mode AD also becomes impractical since $K$ can get very large depending on the convergence speed of \eqref{itr:IM}. This problem has been studied extensively and several techniques have been proposed to address it, for instance, Checkpointing [82, 38] and Truncating [83]. Nevertheless, in the context of fixed-point iteration a better alternative exists [46, 35, 6, 9, 79] which we call Fixed-Point Automatic Differentiation (FPAD). Surprisingly though, FPAD is hardly applied in Machine Learning context as it even improves the convergence of the derivative sequence [64].

\subsubsection{Fixed-Point Automatic Differentiation (FPAD)} \label{sssec:ADFPI:FP}
Since the convergence of $\xKd$ and \eqref{itr:IM:RAD} to $D_{\u}\xmin\ud$ and $\xminb D_{\u}\xmin$ respectively does not depend on the starting point $\xz$, we can choose it anywhere. One way of interpreting FPAD is by setting $\xz$ to $\xmin$ (or $\x$ such that $\norm{\x - \xmin}$ is small). We note that all the nodes in Figure~\ref{fig:IM} will contain (approximate) copies of $\xmin$. Therefore, we can simply replace $\xk$ with $\xmin$ (or $\x$) at every iteration in \eqref{itr:IM:FAD} to obtain its FPAD variant. That is, we start with $\xzh = 0\in\spVar$, and $\ud\in\spPrm$ and perform
\begin{equation} \label{itr:IM:FPFAD} \tag{IM-\^F}
\xkph \coloneqq D_{\x} g (\x, \u) \xkh + D_{\u} g (\x, \u) \ud \,,
\end{equation}
for $k = 0, \ldots, K-1$, which just changes the points where the derivatives are evaluated, to the fixed-point $\xmin$ (or its estimate $\x$). Similarly for reverse mode, we set $\xKt\coloneqq \xminb\in\spVar^*$, and $\uzt\coloneqq 0\in\spPrm^*$ and get the following modified iterations of \eqref{itr:IM:RAD} for $n = 0,\ldots, K-1$
\begin{equation} \label{itr:IM:FPRAD} \tag{IM-\~R}
\begin{aligned}
\xkt &\coloneqq \xkpt D_{\x} g (\x, \u) \\
\unpt &\coloneqq \unt + \xkpt D_{\u} g (\x, \u) \,,
\end{aligned}
\end{equation}
where $k \coloneqq K - n - 1$. Notice that, we distinguish FPAD variables from AD variables by using hat (instead of dot) for forward mode and tilde (instead of bar) for reverse mode. Also, we omit the subscript $K$ for reverse mode FPAD derivatives for the reason which will become obvious once we look at a different interpretation of FPAD.

\textbf{Neumann Series Interpretation:} The above motivation for FPAD was taken from Section~4 of [46]. Another way \cite[Section~4.3]{ORBP16} to motivate this technique is by comparing \eqref{itr:IM:FPFAD} and  \eqref{itr:IM:FPRAD} with the Neumann Series expansion of $I - D_{\x} g (\xmin, \u)$ in \eqref{eq:FPI:IFT}. That is, one can write $D_{\u} \xmin$ alternatively as
\begin{equation*}
D_{\u} \xmin = (I - D_{\x} g (\xmin, \u))^{-1} D_{\u} g (\xmin, \u) = \sum_{i=0}^{\infty} D_{\x} g (\xmin, \u)^i D_{\u} g (\xmin, \u) \,.
\end{equation*}
For $r\in\N_0$, the partial sums $X_r \coloneqq \sum_{i=0}^{r-1} D_{\x} g (\x, \u)^i D_{\u} g (\x, \u)$ can be expressed recursively in two different ways, i.e.,
\begin{equation} \label{eq:DFPI:FAD}
\begin{aligned}
X_{r+1} &= \sum_{i=0}^{r} D_{\x} g (\x, \u)^i D_{\u} g (\x, \u) \\
&= D_{\u} g (\x, \u) + \sum_{i=1}^{r} D_{\x} g (\x, \u)^i D_{\u} g (\x, \u) \\
&= D_{\u} g (\x, \u) + D_{\x} g (\x, \u) \sum_{i=0}^{r-1} D_{\x} g (\x, \u)^i D_{\u} g (\x, \u) \\
&= D_{\x} g (\x, \u) X_r + D_{\u} g (\x, \u)
\end{aligned}
\end{equation}
\begin{equation} \label{eq:DFPI:RAD1}
\begin{aligned}
X_{r+1} &= \sum_{i=0}^{r} D_{\x} g (\x, \u)^i D_{\u} g (\x, \u) \\
&= D_{\x} g (\x, \u)^{r} D_{\u} g (\x, \u) + \sum_{i=0}^{r-1} D_{\x} g (\x, \u)^i D_{\u} g (\x, \u) \\
&= D_{\x} g (\x, \u)^{r} D_{\u} g (\x, \u) + X_r \,,
\end{aligned}
\end{equation}
where $\x$ is an estimate of $\xmin$, as in \eqref{itr:IM:FPFAD} or \eqref{itr:IM:FPRAD}. We further simplify \eqref{eq:DFPI:RAD1} by
\begin{equation} \label{eq:DFPI:RAD2}
\begin{aligned}
Y_{r+1} &= Y_r D_{\x} g (\x, \u) \\
X_{r+1} &= Y_r D_{\u} g (\x, \u) \,.
\end{aligned}
\end{equation}
with $Y_0 \coloneqq I $. Comparing \eqref{itr:IM:FPFAD} with the last equality in \eqref{eq:DFPI:FAD} and \eqref{itr:IM:FPRAD} with \eqref{eq:DFPI:RAD2} we conclude that $\xkh = X_k \ud$ and $\unt = \xminb X_n$.
\begin{REM} \label{rem:FPRAD}
The index for $\xkt$ in \eqref{itr:IM:FPRAD} was borrowed from $\xkb$ in \eqref{itr:IM:RAD}. However, a more natural indexing will lead to a following alternate form of \eqref{itr:IM:FPRAD}
\begin{equation*}
\begin{aligned}
\tilde\y^{(n+1)} &= \tilde\y^{(n)} D_{\x} g (\x, \u) \\
\unpt &= \unt + \tilde\y^{(n)} D_{\u} g (\x, \u) \,.
\end{aligned}
\end{equation*}
This suggests that $K$ has no apparent role to play and \eqref{itr:IM:FPRAD} can be performed indefinitely. Also, due to $\rho (D_{\x} g (\x, \u)) < 1$, $\tilde\x^{(K-n-1)} \to 0$ as $n \to \infty$, we get a natural stopping criterion $\norm{\tilde\x^{(K-n-1)}} < \epsilon$ for reverse mode FPAD.
\end{REM}

% ***************************
% >>>>> PROBLEM SETTING <<<<<
% ***************************
\section{Problem Setting} \label{sec:PS}

We aim for a theoretically sound and efficient AD strategy for computing the derivative of the solution of \eqref{prob:min} with respect to the parameter $\u$. Let $\manif\subset\spVar$ be a $C^2$-smooth manifold and $\setPrm\subset\spPrm$ be an open set. We consider the composite optimization problem \eqref{prob:min} where $\map{f, g}{\spVar\times\spPrm}{\eR}$ satisfy the following assumption. Examples of problems which satisfy this assumption are given in [81]. A few examples are also provided in Section~\ref{sec:Exp} for numerical demonstration of our results.
\begin{ASS}[Convex Partly Smooth Objective] \label{ass:CPSO}
$\map{f}{\spVar\times\spPrm}{\R}$ is $C^2$-smooth, $\map{g}{\spVar\times\spPrm}{\eR}$ is partly smooth relative to $\manif\times\setPrm$ and for every $\u\in\setPrm$, $f(\cdot, \u)$ and $g(\cdot, \u)$ are convex and $f(\cdot, \u)$ has an $L$-Lipschitz continuous gradient.
\end{ASS}
\begin{REM}
The Lipschitz constant $L$ can be made dependent on $\u$ without changing the following analysis. We assume it to be independent of $\u$ for brevity.
\end{REM}

The next assumption that we require is reminiscent of conditions for the Implicit Function Theorem \cite[Theorem~1B.1]{DR09}. Let $\xmin \in \manif$ be a minimizer of $F(\cdot, \u^*)$ for $\u^*\in\setPrm$. A key requirement for using the implicit function theorem for obtaining the derivative of the solution mapping at $\u^*$ is the positive definiteness of $\grad_{\manif}^2 F (\xmin, \u^*)$. However, this alone will not suffice for proving the convergence of the derivative iterates of (Accelerated) Proximal-Gradient Descent on $\spTan\xmin$ (see Theorem~\ref{thm:APG:FAD}). We additionally need positive definiteness of either $\grad_{\x}^2 f (\xmin, \u)$ or $\grad_{\manif}^2 g (\xmin, \u)$. For this purpose, we impose the following assumption on $f$ and $g$ at some $(\xmin, \u^*) \in \manif\times\setPrm$.
\begin{ASS}[Restrictive Positive Definiteness] \label{ass:RPD}
\begin{enumerate}
\item \label{itm:RPD-i}The Hessian $\grad_\manif^2 F (\xmin, \u^*)$ is positive definite on $\spTan\xmin $, i.e.,
\begin{equation} \tag{RPD-i} \label{eq:RPD-i}
\grad_\manif^2 F (\xmin, \u^*) \succ 0 \,.
\end{equation}
\item \label{itm:RPD-ii}Moreover, $\projTan (\xmin) \grad_{\x}^2 f (\xmin, \u^*)$ or $ \grad_\manif^2 g(\xmin, \u^*) $ is positive definite on $\spTan{\xmin}$. More precisely, it holds that
\begin{equation} \tag{RPD-ii} \label{eq:RPD-ii}
\rstDom{\projTan (\xmin) \grad_{\x}^2 f (\xmin, \u^*)}{\spTan\xmin} \succ 0 \quad \text{or} \quad \grad_\manif^2 g (\xmin, \u^*) \succ 0 \,.
\end{equation}
\end{enumerate}
\end{ASS}
Finally the solution of \eqref{prob:min} must change in a stable manner as the parameter $\u$ changes. This is guaranteed (see Theorem~\ref{thm:IFT}) by the following assumption on $f$ and $g$ at some $\u^*\in\setPrm$ and $\xmin \in \argmin_{\x\in\spVar} F(\x, \u^*)$, which is a standard condition when working with partly-smooth functions.
\begin{ASS}[Non-degeneracy] \label{ass:ND}
The non-degeneracy condition is satisfied, i.e.,
\begin{equation} \tag{ND} \label{eq:ND}
0 \in \ri \partial_{\x} F (\xmin, \u^*) \,.
\end{equation}
\end{ASS}
\begin{REM}
\begin{enumerate}
\item The requirement in \eqref{eq:ND} is stronger than the Fermat's rule
\[0 \in \partial_{\x} F (\xmin, \u^*) \,,\]
and therefore automatically implies optimality of $\xmin$.
\item Unlike Assumption~\ref{ass:CPSO}, Assumptions~\ref{ass:RPD} and \ref{ass:ND} are required to be satisfied by $f$ and $g$ at a specific point in $\manif\times\setPrm$.
\item Lewis [56] calls $\xmin\in\manif$ a strong critical point of $F(\cdot, \u^*)$ relative to $\manif$ when Assumptions~\ref{ass:RPD}\ref{itm:RPD-i} and \ref{ass:ND} are satisfied by $f$ and $g$ at $(\xmin, \u^*)$.
\end{enumerate}
\end{REM}
Under these assumptions, we can differentiate the solution mapping of \eqref{prob:min} on an open subset of $\setPrm$. The following result is derived from [56, 81, 74]
\begin{THM}[Differentiation of Solution Mapping] \label{thm:IFT}
Let $f$ and $g$ satisfy Assumption~\ref{ass:CPSO} and $(\xmin, \u^*)\in\manif\times\setPrm$ be such that Assumptions~\ref{ass:RPD} and \ref{ass:ND} are satisfied. Then there exists an open neighbourhood $U\subset\setPrm$ of $\u^*$ and a continuously differential mapping $\map{\argmap}{U}{\manif}$ such that for all $\u\in U$,
\begin{enumerate}
\item $\argmap(\u)$ is the unique minimizer of $F(\cdot, \u)$, \label{itm:IFT:i}
\item \eqref{eq:ND} and \eqref{eq:RPD-i} are satisfied at $(\u, \argmap (\u))$, and \label{itm:IFT:ii}
\item the derivative of $\argmap$ is given by
\begin{equation} \label{eq:IFT}
D \argmap (\u) = -\grad_{\manif}^2 F (\argmap (\u), \u)^{\dagger} D_{\u} \grad_{\manif} F (\argmap (\u), \u) \,,
\end{equation}
where $\grad_{\manif}^2 F (\argmap (\u), \u)^{\dagger}$ denotes the pseudoinverse of $\grad_{\manif}^2 F (\argmap (\u), \u)$. \label{itm:IFT:iii}
\end{enumerate}
\end{THM}
\begin{proof}
Assumptions~\ref{ass:RPD} and \ref{ass:ND} ensure that $\xmin$ is a strong critical point of $F(\cdot, \u^*)$ relative to $\manif$ \cite[Definition~5.6]{Lew02}. The transversal embedding condition is also satisfied for $\manif\times\setPrm$ \cite[Assumption~5.1]{Lew02}. Thus \ref{itm:IFT:i} and \ref{itm:IFT:ii} follow from \cite[Theorem~5.7]{Lew02}. The expression for the derivative of the solution mapping follows from \cite[Theorem~1]{VDP+17} (see also \cite[Lemma~7.26]{Rii20}).
\end{proof}
\begin{REM}
Assumption~\ref{ass:ND} is necessary for the differentiability of the solution mapping in general. However one can still study the sensitivity analysis of the solution mapping when \eqref{eq:ND} is violated (see [43] or \cite[Section~7.4.2]{Rii20}), which however requires other assumptions and is not pursued in this paper.
\end{REM}
Theorem~\ref{thm:IFT} elegantly extends Implicit Differentiation to the setting when the objective is partly smooth and provides a theoretically justified way to computing the derivative of the solution mapping. However, it is not practical in the sense that the Riemannian Gradient and Hessian in general cannot be obtained through autograd packages and computing them by-hand can become cumbersome. One remedy is to use automatic differentiation on PGD or APG simply because the derivative of the update step of the two algorithms for various practical applications can be computed by standard autograd packages [1, 71, 25]. But as already highlighted in Section~\ref{sec:intro}, this raises two problems: (i) Does the derivative of the sequence of iterates converge to the derivative of the solution mapping? (ii) If yes, then is there a way to circumvent the memory-overhead problem of the reverse mode AD? We answer the first question in Section~\ref{sec:AD} and use the tools from Section~\ref{sssec:ADFPI:FP} (FPAD) to address the second question in Section~\ref{sec:FPAD}. In Section~\ref{sec:FPID}, we show that FPAD applied to PGD is similar to applying ID to the fixed-point equation of PGD. First we provide a brief overview of APG. We recall the finite identification and local-linear convergence of APG when the objective in \eqref{prob:min} is partly-smooth [58].

\subsection{Accelerated Proximal Gradient (APG)}
Introduced in [13] as Fast Iterative Shrinkage/Thresholding Algorithm or FISTA, the algorithm is the application of the celebrated Nesterov Acceleration [67] to Proximal Gradient Descent [60]. For convex problems of the form \eqref{prob:min}, APG in general exhibits convergence like $\O (1/k^2)$ in objective values as compared to PGD which converges like $\O(1/k)$ [29]. In fact, APG is the optimal algorithm for the convex problems of type \eqref{prob:min} when we only have access to the first order information of the smooth component $f$ of the objective. Convergence of the iterates of APG has been established [28] for a specific range of parameters.
\begin{Algorithm}[Accelerated Proximal Gradient (APG)] \ \label{alg:APG}
\begin{itemize}
\item \key{Initialization:} $\xz=\xzm\in \spVar$, $\u\in\spPrm$, $0 < \sslow\leq\ssup < 2/L$.
\item \key{Parameter:} $(\alpha_k)_{k\in\N} \in [\sslow, \ssup]$ and $(\beta_k)_{k\in\N} \in [0, 1]$.
\item \key{Update $k\geq 0$:}
\begin{equation} \label{itr:APG} \tag{APG}
\begin{aligned}
\yk &\coloneqq (1+\beta_k ) \xk - \beta_k \xkm \\
\wk &\coloneqq \yk - \alpha_k \grad_{\x} f(\yk, \u) \\
\xkp &\coloneqq \bwd[k] (\wk, \u) \,.
\end{aligned}
\end{equation}
\end{itemize}
\end{Algorithm}
Algorithm~\ref{alg:APG} shows the update steps of APG. In \eqref{itr:APG}, $\map{\bwd}{\spVar\times\spPrm}{\spVar}$ is defined as
\begin{equation} \label{eq:prox}
\bwd (\w, \u) \coloneqq \argmin_{\x\in\spVar} \alpha g(\x, \u) + \frac{1}{2} \norm{\x - \w}^2 \,.
\end{equation}
The iterations in \eqref{itr:APG} can be written more compactly as $\xkp \coloneqq \pgd[k] (\xk + \beta_k (1+\beta_k) \xkm)$ with $ \map{\pgd}{\spVar\times\spPrm}{\spVar} $ given by
\begin{equation*}
\pgd (\x, \u) \coloneqq \bwd (\x - \alpha \grad_{\x} f (\x, \u), \u) \,.
\end{equation*}
A good choice of $\beta_k$ which also guarantees the convergence of the iterates is $(k-1)/(k+q)$ with $q>2$ [28]. Note that by setting $\beta_k$ to $0$, \algref{alg:APG} reduces to PGD.

To be able to differentiate $\pgd$ by using Theorem~\ref{thm:IFT}, the first thing that we need is to ensure that the iterates $\xk$ lie on the manifold. This problem was addressed for PGD in [59] and later for APG in [58] for partly-smooth objective. It was shown that the iterates $\xk$ generated by PGD and APG eventually lie on the manifold $\manif$. In particular, there exists $K\in\N$ such that for all $k\geq K$, $\xk\in\manif$. Furthermore, the iterates converge to $\xmin$ at a linear rate. The following lemma summarizes the results in [58].

\begin{LEM}[Activity Identification and Linear Convergence of \algref{alg:APG}] \label{lem:APG}
Let $f$ and $g$ satisfy Assumption~\ref{ass:CPSO} and $(\xmin, \u)\in\manif\times\setPrm$ be such that Assumptions~\ref{ass:RPD} and \ref{ass:ND} are satisfied. Let $\alpha_k \in [\sslow, \ssup]$ and $\beta_k \in [0, 1]$ converge to $\alpha_*$ and $\beta_*$ respectively such that $-1/(1 + 2\beta_*) < \underline{\lambda}$ where $\underline{\lambda}$ is the smallest eigenvalue of $D_{\x} \pgd[*] (\xmin, \u)$. If the sequence $(\xk)_{k\in\N}$ generated by \algref{alg:APG} converges to $\xmin$ then $\xk$ eventually lies on $\manif$ and converges locally linearly to $\xmin$.
\end{LEM}

\begin{REM}
The authors in [58] consider a more general inertial Forward Backward or iFB algorithm (see Algorithm~1 in [58]), which encompasses different variants of FISTA. We stick to the more well-known version, i.e., \algref{alg:APG}. However, it should be noted that our results lie on the analysis in [58] and with little effort, can be extended to iFB algorithm.
\end{REM}

\section{Automatic Differentiation (AD) of APG} \label{sec:AD}
In many practical applications, analytic solutions to the proximal mapping are known. They are expressed as compositions of elementary functions which means that $\bwd$ and $\pgd$ defined in \eqref{eq:prox} and \eqref{eq:prox:grad} respectively can be differentiated by using standard autograd packages. Therefore AD appears to be a more natural choice for estimating the derivative of solution mapping of \eqref{prob:min} as opposed to \eqref{eq:IFT}. In this section we apply AD on \algref{alg:APG} and show convergence of the generated derivative sequence.

\subsection{Forward Mode AD of APG} \label{ssec:AD:FAD}
We first perform forward mode AD on APG.
\begin{Algorithm}[Forward Mode AD of \algref{alg:APG}] \ \label{alg:APG:FAD}
\begin{itemize}
\item \key{Initialization:} $\xzd=\xzmd=0\in\spVar$, $\u, \ud\in\spPrm$, $0 < \sslow\leq\ssup < 2/L$; $(\xk)_{k\in\N}$, $(\yk)_{k\in\N}$ and $(\wk)_{k\in\N}$ from \eqref{itr:APG}.
\item \key{Parameter:} $(\alpha_k)_{k\in\N} \in [\sslow, \ssup]$ and $(\beta_k)_{k\in\N} \in [0, 1]$.
\item \key{Update $k\geq 0$:}
\begin{equation} \label{itr:APG:FAD} \tag{APG-F}
\begin{aligned}
\ykd &\coloneqq (1+\beta_k ) \xkd - \beta_k \xkmd \\
\wkd &\coloneqq (I - \alpha_k \grad_{\x}^2 f(\yk, \u)) \ykd - \alpha_k \grad_{\x\u} f(\yk, \u)\ud \\
\xkpd &\coloneqq D \bwd[k] (\wk, \u) (\wkd, \ud) \,.
\end{aligned}
\end{equation}
\end{itemize}
\end{Algorithm}
The sequence $(\xkd)_{k\in\N}$ generated by \algref{alg:APG:FAD} is well defined when $\bwd$ is differentiable for which we refer to the following Lemma. The proximal mapping in \eqref{itr:APG} by definition is a solution mapping of an optimization problem given in \eqref{eq:prox}. Therefore, for computing it derivative, we resort to a result which is similar to Theorem~\ref{thm:IFT}.

\begin{LEM} \label{lem:prox:diff}
Let $f$ and $g$ satisfy Assumption~\ref{ass:CPSO} and $(\xmin, \u^*)\in\manif\times\setPrm$ be such that Assumptions~\ref{ass:RPD} and \ref{ass:ND} are satisfied. For any $\alpha\in[\sslow, \ssup]$, the mapping $\bwd$ defined in \eqref{eq:prox} is differentiable at any $(\w, \u)$ sufficiently close to $(\xmin - \alpha \grad_{\x} f (\xmin, \u^*), \u^*)$ with derivative given by
\begin{equation} \label{eq:prox:derv}
\begin{aligned}
D_{\x} \bwd (\w, \u) &= \big ( \projTan (\x) + \wein{\cdot}{\projNor (\x) (\x - \w)} + \alpha \grad_{\manif}^2 g(\x, \u) \big ) ^\dagger \\
D_{\u} \bwd (\w, \u) &= - \alpha D_{\x} \bwd (\w, \u) D_{\u} \grad_\manif g(\x, \u) \,,
\end{aligned}
\end{equation}
for $\x = \bwd (\w, \u)$.
\end{LEM}
\begin{proof}
We prove this result by verifying the assumptions similar to Assumptions~\ref{ass:CPSO}, \ref{ass:RPD} and \ref{ass:ND} for \eqref{eq:prox} at $(\w^*, \u^*)$ for $\w^*\coloneqq \xmin-\alpha \grad_{\x} f (\xmin, \u^*)$ and then using Theorem~\ref{thm:IFT}. It is straightforward to check the validity of Assumption~\ref{ass:CPSO}. Assumption~\ref{ass:RPD} is satisfied from the strong convexity of \eqref{eq:prox}. Finally from $\xmin = \bwd (\w^*, \u^*)$ and \eqref{eq:ND} we note that $\w^* \in \xmin + \ri \alpha \partial_{\x} g (\xmin, \u^*)$, i.e., Assumption~\ref{ass:ND} is also satisfied.
\end{proof}
\begin{REM} \label{rem:prox:ND}
The expressions in \eqref{eq:prox:derv} only make sense when $\bwd (\w, \u) \in \manif$ and the non-degeneracy condition holds for \eqref{eq:prox} at the given inputs $(\w, \u)$. The first condition is satisfied by all the iterates of \eqref{itr:APG} after a finite number of iterations thanks to Lemma~\ref{lem:APG}. As for non-degeneracy, in [74], $\bwd$ is shown to be piecewise smooth under further assumptions when non-degeneracy condition is violated. Hence, earlier in the iterations of \eqref{itr:APG} when $\xk\notin\manif$ and $\wk$ is not yet close enough to $\xmin - \alpha \grad_{\x} f(\xmin, \u)$, we have that $\xk$ is only piecewise $C^1$-differentiable with respect to $\u$. Even for large enough $k$ when $\bwd$ is differentiable, we cannot guarantee the differentiablity of $\xk$ due to the non-differentiablity of the earlier iterates.
\end{REM}
\textbf{Warm-Start:} Following Remark~\ref{rem:prox:ND}, we assume that the non-degeneracy condition is satisfied by \eqref{eq:prox} for every $(\wk, \u)$ in \eqref{itr:APG:FAD} so that the sequence $( \xkd )_{k \in\N}$ is well-defined. This is possible, for example, if we initially perform \eqref{itr:APG} alone for a large number of iterations, say $N\in\N$, to obtain $\x^{(N)}$. We then reset $\x^{(N)}$ to $\xz$ and perform \eqref{itr:APG} and \eqref{itr:APG:FAD} in a normal way from there on. This allows us to provide convergence guarantees for \algref{alg:APG:FAD} under the given assumptions. We consider this as a sever drawback which we remedy by our Fixed-Point AD approach, since actually the mapping $\xz \mapsto \x^{(N)}$ depends on $\u$ and this cannot be simply ignored for differentiation by warm-starting.

\begin{THM} \label{thm:APG:FAD}
Let $f$ and $g$ satisfy Assumption~\ref{ass:CPSO} and $(\xmin, \u)\in\manif\times\setPrm$ be such that Assumptions~\ref{ass:RPD} and \ref{ass:ND} are satisfied. Let $\alpha_k \in [\sslow, \ssup]$ and $\beta_k \in [0, 1]$ converge to $\alpha_*$ and $\beta_*$ respectively such that $-1/(1 + 2\beta_*) < \underline{\lambda}$ where $\underline{\lambda}$ is the smallest eigenvalue of $D_{\x} \pgd[*] (\xmin, \u)$. For any $\xz$ sufficiently close to $\xmin$ with $(\xk)_{k\in\N}$ converging to $\xmin$, the sequence $( \xkd )_{k \in\N}$ generated by \algref{alg:APG:FAD} converges to $D \argmap(\u) \ud$.
\end{THM}

\begin{proof}
Lemma~\ref{lem:APG} already ensures that the sequence $\xk$ eventually lies in $\manif$ under the given assumptions. Therefore we assume that $\xk\in\manif$ for all $k\in\N$ and define linear operators $Q_k, R_k, S_k, M_k$ and vector $\bk$ for $ k\geq0 $ by:
\begin{equation} \label{eq:linops:k}
\begin{aligned}
Q_k &\coloneqq \projTan (\xkp) + \wein{\cdot}{\projNor (\xkp)(\xkp - \wk)} + \alpha_k \grad_{\manif}^2 g(\xkp, \u) \\
R_k &\coloneqq Q_k^\dagger (I - \alpha_k \grad_{\x}^2 f(\yk, \u)) \\
S_k &\coloneqq -\alpha_k Q_k^{\dagger} D_{\u} (\grad_{\x} f (\yk, \u) + \grad_{\manif} g (\xkp, \u)) \\
M_k &\coloneqq \begin{bmatrix} (1+\beta_k) R_k & -\beta_k R_k \\ I & 0 \end{bmatrix} \\
T_k &\coloneqq \begin{bmatrix} S_k & 0 \\ 0 & S_k \end{bmatrix} \,,
\end{aligned}
\end{equation}
and expand \eqref{itr:APG:FAD} to obtain
\begin{align*}
\xkpd &= Q_k^\dagger \wkd - \alpha_k Q_k^\dagger D_{\u} \grad_\manif g(\x, \u) \ud \\
&= Q_k^\dagger (I - \alpha_k \grad_{\x}^2 f(\yk, \u)) \ykd - \alpha_k Q_k^\dagger \grad_{\x\u} f(\yk, \u)\ud - \alpha_k Q_k^\dagger D_{\u} \grad_\manif g(\x, \u) \ud \\
&= R_k \ykd + S_k \ud \\
&= (1 + \beta_k) R_k \xkd - \beta_k R_k \xkmd + S_k \ud \,.
\end{align*}
Thus, by setting $\zkd \coloneqq (\xkd, \xkmd)$ and $\zzd=0\in \spVar\times\spVar$, we can write the forward mode automatic differentiation of \eqref{itr:APG} in a compact manner as
\begin{equation} \label{eq:glfpi:FAD}
\zkpd \coloneqq M_k \zkd + T_k (\ud, 0) \,.
\end{equation}
To show convergence we will verify, one-by-one, all the conditions in Corollary~\ref{cor:LGFPI:conv}. We start by noticing that $T_k (\s, 0) \in V_k \coloneqq \spTan\xkp \times \spTan\xk $ and $\im M_k = V_k$. Also, from the continuity of the respective functions, as $ \xk \to \xmin, \alpha_k \to \alpha_* $ and $\beta_k \to \beta_*$, we obtain
\begin{equation}
\begin{aligned}
\yk &\to \xmin \\
\wk &\to \w^* \coloneqq \xmin - \alpha \grad_{\x} f (\xmin, \u) \\
Q_k &\to Q_* \coloneqq \projTan (\xmin) + \wein{\cdot}{\projNor (\xmin) (\alpha_*\grad_{\x} f (\xmin, \u))} + \alpha_* \grad_{\manif}^2 g(\xmin, \u) \\
R_k &\to R_* \coloneqq Q_*^\dagger (I - \alpha_* \grad_{\x}^2 f(\xmin, \u)) \\
S_k &\to S_* \coloneqq -\alpha_* Q_*^{\dagger} D_{\u} (\grad_{\x} f (\xmin, \u) + \grad_{\manif} g (\xmin, \u)) =-\alpha_* Q_*^{\dagger} D_{\u} \grad_{\manif} F (\xmin, \u) \\
M_k &\to M_*\coloneqq \begin{bmatrix} (1+\beta) R_* & -\beta R_* \\ I & 0 \end{bmatrix} \\
T_k &\to T_* \coloneqq \begin{bmatrix} S_* & 0 \\ 0 & S_* \end{bmatrix} \\
V_k &\to V \coloneqq \spTan\xmin \times \spTan\xmin \,.
\end{aligned}
\end{equation}
Finally we need to show that $\rho (M_*\circ\proj V) < 1$. This will follow from \cite[Corollary~4.9]{LFP17} once we establish that $\rho (R_* \projTan (\xmin)) < 1$ since $R_* = D_{\x} \pgd[*](\xmin, \u)$. We therefore rewrite
\begin{equation*}
\begin{aligned}
\rho (R_* \projTan (\xmin)) &= \rho (Q_*^{\dagger} (I - \alpha_* \grad_{\x}^2 f(\xmin, \u)) \projTan (\xmin)) \\
&= \rho (Q_*^{\dagger} (\projTan (\xmin) - \projTan (\xmin) \alpha_* \grad_{\x}^2 f(\xmin, \u) \projTan (\xmin))) \\
&\leq \rho (Q_*^{\dagger}) \rho (\projTan (\xmin) - \projTan (\xmin) \alpha_* \grad_{\x}^2 f(\xmin, \u) \projTan (\xmin)) \,.
\end{aligned}
\end{equation*}
Let $\mu_g$ and $\mu_f$ be the smallest eigenvalues of the self-adjoint linear operators $\grad_{\manif}^2 g (\xmin, \u)$ and $\projTan (\xmin) \grad_{\x}^2 f (\xmin, \u) \projTan (\xmin)$ respectively. From Assumption~\ref{ass:RPD}, either $\mu_f$ or $\mu_g$ is positive. By noting that $ \mu_f I_{\spTan\xmin} \preceq \rstDom{\projTan (\xmin) \grad_{\x}^2 f (\xmin, \u)}{\spTan\xmin} \preceq L I_{\spTan\xmin}$ and $0<\alpha_*<2/L$ we find that $\rho (\projTan (\xmin) - \projTan (\xmin) \alpha_* \grad_{\x}^2 f(\xmin, \u) \projTan (\xmin)) \leq \max(\abs{1-\alpha_* \mu_f}, \abs{1-\alpha_* L})$ which is less than $1$ when $\mu_f > 0$. On the other hand $Q_* = \grad_{\manif}^2 (\alpha_* g(\cdot, \u) + 1/2 \norm{\cdot - \w^*})(\xmin) \succeq (1+\alpha_* \mu_g)I_{\spTan\xmin}$ due to the $1+\alpha_*\mu_g$-strong convexity of \eqref{eq:prox} from which we conclude that $\rho (Q^{\dagger}_*) \leq 1/(1+\alpha_*\mu_g)$ which is less than $1$ when $\mu_g > 0$.

It remains to verify that the limit of $\eqref{eq:glfpi:FAD}$ is the same as $(D \argmap (\u) \ud, D \argmap (\u) \ud)$. For that, we evaluate
\begin{equation*}
\begin{aligned}
(I - M_*\circ\proj V) (D \argmap (\u) \ud, D \argmap (\u) \ud) &= \begin{bmatrix} I - R_* - \beta R_* & \beta R_* \\ -I & I \end{bmatrix} \begin{bmatrix} D \argmap (\u) \ud \\ D \argmap (\u) \ud \end{bmatrix} \\
&= \begin{bmatrix} (I - R_*) D \argmap (\u) \ud \\ 0 \end{bmatrix} \,.
\end{aligned}
\end{equation*}
This simplifies our problem to showing that $ (I - R_*) D \argmap (\u) = S_* $. But since
\begin{equation*}
\begin{aligned}
I - R_* &= I - Q_*^\dagger \big( I - \alpha_* \grad_{\x}^2 f(\xmin, \u) \big) \\
&= I - Q_*^\dagger \big( \projTan (\xmin) - \alpha_* \projTan (\xmin) \grad_{\x}^2 f(\xmin, \u) \big) \\
&= Q_*^\dagger \big( Q_* - \projTan (\xmin) + \alpha_* \projTan (\xmin) \grad_{\x}^2 f(\xmin, \u) \big) \\
&= Q_*^\dagger \big( \wein{\cdot}{\projNor (\xmin) (\alpha_*\grad_{\x} f (\xmin, \u))} + \alpha_* \grad_{\manif}^2 g(\xmin, \u) + \alpha_* \projTan (\xmin) \grad_{\x}^2 f(\xmin, \u) \big) \\
&= \alpha_* Q_*^\dagger \grad_{\manif}^2 F(\xmin, \u) \,,
\end{aligned}
\end{equation*}
we have
\begin{equation*}
\begin{aligned}
(I - R_*)^{-1} S_* &= \frac{1}{\alpha_*} \grad_{\manif}^2 F(\xmin, \u)^\dagger Q_* \big( -\alpha_* Q_*^{\dagger} D_{\u} \grad_{\manif} F (\xmin, \u) \big) \\
&= -\grad_{\manif}^2 F(\xmin, \u)^\dagger D_{\u} \grad_{\manif} F (\xmin, \u) \big) \,.
\end{aligned}
\end{equation*}
\end{proof}

\begin{REM} \label{rem:APG:FAD}
\begin{enumerate}
\item In [74], there are slight errors that we fix here. In \cite[Lemma~7.31]{Rii20} and afterwards in the proofs, the Weingarten term of \eqref{eq:prox:derv} is missing. Also in \cite[Proposition~2.7]{Rii20}, for linear convergence of $(d^k)_{k\in\N}$, the sequences $(A_k)_{k\in\N}$ and $(b^k)_{k\in\N}$ ought to converge linearly.
\item We refer the reader to Section~4 in [58] for a detailed discussion on the spectral properties of $M_*$ and how they can affect the convergence and convergence rate of APG (resp. forward mode AD of APG) as compared to PGD (resp. forward mode AD of PGD).
\end{enumerate}
\end{REM}

\subsection{Reverse Mode AD of APG} \label{ssec:AD:RAD}
As mentioned in Sections~\ref{ssec:NP:AD} and \ref{ssec:NP:ADFPI}, the forward mode AD is not practical when the size of the input of the function is much bigger than the size of its output. In this case we prefer to use the reverse mode AD.
\begin{Algorithm}[Reverse Mode AD of \algref{alg:APG}] \ \label{alg:APG:RAD}
\begin{itemize}
\item \key{Initialization:} $\u\in\spPrm$, $\xKb=\xminb\in\spVar^*$, $ \yKb=0\in\spVar^*$ and $\uzb=0\in\spPrm^*$, $0 < \sslow\leq\ssup < 2/L$; $(\xk)_{k\in [K-1]}$, $(\yk)_{k\in [K-1]}$ and $(\wk)_{k\in [K-1]}$ from \algref{alg:APG}.
\item \key{Parameter:} $(\alpha_k)_{k\in\N} \in [\sslow, \ssup]$ and $(\beta_k)_{k\in\N} \in [0, 1]$.
\item \key{Update $(n=0,\ldots,K-1$, $k \coloneqq K - n - 1 )$:}
\begin{equation} \label{itr:APG:RAD} \tag{APG-R}
\begin{aligned}
\wkb &\coloneqq \xkpb D_{\x} \bwd[k] (\wk, \u) \\
\ykb &\coloneqq \wkb (I - \alpha_k \grad_{\x}^2 f(\yk, \u)) \\
\unpb &\coloneqq \unb + \xkpb D_{\u} \bwd[k] (\wk, \u) - \alpha_k \wkb \grad_{\x\u} f(\yk, \u) \\
\xkb &\coloneqq (1+\beta_k ) \ykb - \beta_k \ykpb \,.
\end{aligned}
\end{equation}
% \item \key{Output:} $\ubK \coloneqq \uKb$.
\end{itemize}
\end{Algorithm}
The convergence guarantees for reverse mode are similar to that of Theorem~\ref{thm:APG:FAD}.
\begin{THM} \label{thm:APG:RAD}
Let $f$ and $g$ satisfy Assumption~\ref{ass:CPSO} and $(\xmin, \u)\in\manif\times\setPrm$ be such that Assumptions~\ref{ass:RPD} and \ref{ass:ND} are satisfied. Let $\alpha_k \in [\sslow, \ssup]$ and $\beta_k \in [0, 1]$ converge to $\alpha_*$ and $\beta_*$ respectively such that $-1/(1 + 2\beta_*) < \underline{\lambda}$ where $\underline{\lambda}$ is the smallest eigenvalue of $D_{\x} \pgd[*] (\xmin, \u)$. For any $\xz$ sufficiently close to $\xmin$ with $(\xk)_{k\in\N}$ converging to $\xmin$, the sequence $( \uKb )_{K \in\N}$ generated by \algref{alg:APG:RAD} converges to $\xminb D \argmap(\u)$.
\end{THM}
\begin{proof}
We already showed the convergence of \eqref{itr:APG:FAD} and because the two modes essentially aim to compute the same quantity in the end, the proof follows.
\end{proof}

\section{Fixed-Point Automatic Differeniation (FPAD) of APG} \label{sec:FPAD}
The ease of using AD on \algref{alg:APG} combined with the result shown in Theorem~\ref{thm:APG:FAD} provides a powerful tool for estimating the derivative of the solution mapping of \eqref{prob:min}. However, as already noted in Sections~\ref{ssec:NP:AD} and \ref{ssec:NP:ADFPI}, the reverse mode AD has a memory overhead and therefore can become impractical when the optimization algorithm converges slowly and is required to run for a longer period of time making $K$ very large. Fortunately, in our context the curse of memory can be easily lifted by using Fixed-Point Automatic Differentiation (see Section~\ref{sssec:ADFPI:FP}) on \algref{alg:APG}. In particular, we make the following contributions. (i) We perform both forward and reverse mode FPAD on \algref{alg:APG} and show the convergence results for the corresponding derivative sequences separately. (ii) We note that FPAD not only solves the memory issue for reverse mode, it is also possible to show linear convergence for the FPAD derivative sequences under Assumptions~\ref{ass:CPSO}, \ref{ass:RPD} and \ref{ass:ND} as opposed to AD sequences which require additional assumptions for linear convergence. (iii) We remedy the differentiability issue when truncating the sequence $ (\xk)_{k\in\N} $ as discussed in Remark~\ref{rem:prox:ND}. (iv) Lastly, we would like to emphasize that we show the convergence results for FPAD for \textit{approximate} fixed-points and not just for \textit{exact} fixed-points. This is useful in practice because we may get very close to, while still never attain, the solution of \eqref{prob:min} even by running \eqref{itr:APG} for a considerably large number of iterations. In our analysis, the corresponding result \textit{at} the fixed-point can be easily recovered.

\subsection{Forward Mode FPAD of APG} \label{ssec:FPAD:FM}
\algref{alg:APG:FPFAD} performs forward mode FPAD of \eqref{itr:APG}. We make a slight change in the update step when moving from \eqref{itr:APG:FAD} to \eqref{itr:APG:FPFAD}. Instead of using $-\grad_{\x} f (\x, \u)$ as a direction for computing $\w$ (see the update for $\wk$ in \eqref{itr:APG}), we instead use $\bnu\coloneqq \proj{\partial_{\x} g (\x, \u)} (-\grad_{\x} f (\x, \u))$ to compute $\w \coloneqq \x + \alpha \bnu $. Notice that $\bnu$ is well defined thanks to the projection theorem \cite[Theorem~3.16]{BC11}.
\begin{Algorithm}[Forward Mode FPAD of \algref{alg:APG}] \ \label{alg:APG:FPFAD}
\begin{itemize}
\item \key{Initialization:} $\xzh=\xzmh=0\in\spVar$, $\u, \ud\in\spPrm$, $\x\in\spVar$ with small $\norm{\x-\xmin}$, $0 < \sslow\leq\ssup < 2/L$.
\item \key{Parameter:} $\alpha \in [\sslow, \ssup]$ and $\beta \in [0, 1]$.
\item \key{Define:} $\w \coloneqq \x + \alpha \proj{\partial_{\x} g (\x, \u)} (-\grad_{\x} f (\x, \u))$.
\item \key{Update $k\geq 0$:}
\begin{equation} \label{itr:APG:FPFAD} \tag{APG-\^F}
\begin{aligned}
\ykh &\coloneqq (1+\beta) \xkh - \beta \xkmh \\
\wkh &\coloneqq \ykh - \alpha \grad_{\x}^2 f(\x, \u)\ykh - \alpha \grad_{\x\u} f(\x, \u)\ud \\
\xkph &\coloneqq D\bwd (\w, \u) (\wkh, \ud) \,,
\end{aligned}
\end{equation}
\end{itemize}
\end{Algorithm}

\textbf{Reason for Changing the Direction:} The reason for choosing such a peculiar direction $ \bnu = \proj{\partial_{\x} g (\x, \u)} (-\grad_{\x} f (\x, \u)) $ in \eqref{itr:APG:FPFAD} becomes apparent when we demonstrate what we achieve by doing so. Clearly for an exact fixed-point, we have $\bnu = \grad_{\x} f(\xmin, \u)$. Observe that $\bwd (\w, \u) = \x$ since $\w \in (I + \alpha \partial_{\x} g(\cdot, \u)) (\x) $ and $\bwd (\w, \u) = (I + \alpha \partial_{\x} g(\cdot, \u))^{-1} (\w) $. Thus the expression for $D\bwd (\w, \u)$ in \eqref{eq:prox:derv} suggests that the projections $\projTan$ and $\projNor$ are computed at $\x$ and the sequence $\xkh$ eventually lie in $\spTan{\x}$. This is crucial since we will use it to approach the expression $\grad_{\manif}^2 F (\x, \u)$ (see \eqref{eq:Riem:hess}). One may wonder, any subgradient of $g(\cdot, \u)$ at $\x$ could have served the purpose. The reason for specifically choosing $\bnu$ is that $\bnu$ is (by definition) the closest subgradient of $g(\cdot, \u)$ to $-\grad_{\x} f(\x, \u)$. This is important because we require $\bnu + \grad_{\x} f(\x, \u)$ as small as possible (see Theorem~\ref{thm:APG:FPFAD}, in particular, \eqref{eq:IFT:approx}). In the following Lemma, we show that $\bnu$ is close enough to $-\grad_{\x} f (\xmin, \u)$ and $\bwd$ is differentiable at $\w$ whenever $\x$ is close enough to $\xmin$ relative to $\manif$.
\begin{LEM} \label{lem:sg:seq}
Let $f$ and $g$ satisfy Assumption~\ref{ass:CPSO} and $(\xmin, \u)\in\manif\times\setPrm$ be such that Assumptions~\ref{ass:RPD} and \ref{ass:ND} are satisfied. Then for any sequence $( \xk)_{k\in\N}$ in $\manif$ with limit $\xmin$, the sequence $( \bnuk)_{k\in\N}$ defined by
\[
\bnuk \coloneqq \proj{\partial_{\x} g (\xk, \u)} (-\grad_{\x} f (\xk, \u)) \,,
\]
converges to $-\grad_{\x} f (\xmin, \u)$. Moreover, for any $\alpha\in[\sslow, \ssup]$, there exists $K\in\N$ such that the mapping $\bwd$ is differentiable at $\xk + \alpha\bnuk$ for all $k\geq K$.
\end{LEM}

\begin{proof}
Once we show the first part of the lemma, the differentiability of $\bwd$ will follow from Lemma~\ref{lem:prox:diff}. Assume that $\xk\neq \xmin$ because otherwise the proof is trivial. From subdifferential continuity of $g(\cdot, \u)$ relative to $\manif$, we note that $\partial_{\x} g (\xk, \u)$ converges to $\partial_{\x} g (\xmin, \u)$. Since $-\grad_{\x} f (\xmin, \u) \in \partial_{\x} g (\xmin, \u)$, from the definition of convergence of $\partial_{\x} g (\xk, \u)$ \cite[Chapter~4]{RW98}, there exists a sequence $(\bmuk )_{k\in\N}$ with limit $-\grad_{\x} f (\xmin, \u)$ such that for every $k\in\N, \bmuk\in\partial_{\x} g (\xk, \u)$. By the definition of projection, we note that for every $k\in\N$
\begin{equation*}
\begin{aligned}
\dist (-\grad_{\x} f (\xk, \u), \partial_{\x} g (\xk, \u)) &\coloneqq \inf_{\y\in\partial_{\x} g (\xk, \u)} \norm{\y + \grad_{\x} f (\xk, \u)} \\
&= \norm{\bnuk + \grad_{\x} f (\xk, \u)} \\
&\leq \norm{\bmuk + \grad_{\x} f (\xk, \u)} \,.
\end{aligned}
\end{equation*}
But since $\grad_{\x} f (\xk, \u)$ converges to $\grad_{\x} f (\xmin, \u)$ and $\norm{\bmuk + \grad_{\x} f (\xk, \u)}$ converges to $0$, the result follows.
\end{proof}
\begin{REM}
In practice, when $\x$ is close enough to $\xmin$, one can avoid computing the projection onto $\partial_{\x} g(\x, \u)$. That is, the original direction $-\grad_{\x} f(\x, \u)$ in the gradient step can be used instead in \eqref{itr:APG:FPFAD} due to Lemma~\ref{lem:sg:seq}.
\end{REM}

\noindent We are now ready to prove the convergence of \algref{alg:APG:FPFAD}.

\begin{THM} \label{thm:APG:FPFAD}
Let $f$ and $g$ satisfy Assumption~\ref{ass:CPSO} and $(\xmin, \u)\in\manif\times\setPrm$ be such that Assumptions~\ref{ass:RPD} and \ref{ass:ND} are satisfied. Let $ \alpha \in [\sslow, \ssup]$ and $\beta \in [0, 1]$ be such that $-1/(1 + 2\beta) < \underline{\lambda}$ where $\underline{\lambda}$ is the smallest eigenvalue of $D_{\x} \pgd (\xmin, \u)$. For any $\x\in\manif$ sufficiently close to $\xmin$, the sequence $( \xkh )_{k \in\N}$ generated by \algref{alg:APG:FPFAD} converges to $\phi(\x, \u) \ud$ where
\begin{equation} \label{eq:IFT:approx}
\phi (\x, \u) = - \big ( \grad_{\manif}^2 F(\x, \u) - \wein{\cdot}{\alpha \projNor (\x) (\bnu + \grad_{\x} f (\x, \u))} \big )^\dagger D_{\u} \grad_{\manif} F (\x, \u) \,,
\end{equation}
which for $\x=\xmin$, recovers $D\argmap (\u)$ in \eqref{eq:IFT}. Moreover, the convergence of $\xkh$ is linear and its convergence rate approaches that of \algref{alg:APG} as $\x$ approaches $\xmin$.
\end{THM}

\begin{proof}
First we argue that $\grad_{\manif}^2 F(\x, \u) - \wein{\cdot}{\alpha \projNor (\x) (\bnu + \grad_{\x} f (\x, \u))}$ is bijective on $\spTan\x$. When $\x$ is arbitrarily close to $\xmin$, $\grad_{\manif}^2 F(\x, \u)$ is invertible due to Assumption~\ref{ass:RPD} and $\norm{\bnu + \grad_{\x} f (\x, \u)} $ is sufficiently small due to Lemma~\ref{lem:sg:seq}. Therefore $\wein{\cdot}{\alpha \projNor (\x) (\bnu + \grad_{\x} f (\x, \u))}$ which is also linear in the second argument, has a sufficiently small spectral norm and its effect on the eigenvalues of $\grad_{\manif}^2 F(\x, \u)$ becomes arbitrarily small.

Now, as in Theorem~\ref{thm:APG:FAD}, we similarly define the following linear operators
\begin{equation} \label{eq:linops}
\begin{aligned}
Q &\coloneqq \projTan (\x) + \wein{\cdot}{\projNor (\x)(-\alpha\bnu)} + \alpha \grad_{\manif}^2 g(\x, \u) \\
R &\coloneqq Q^\dagger (I - \alpha \grad_{\x}^2 f(\x, \u)) \\
S &\coloneqq -\alpha Q^{\dagger} D_{\u} \grad_{\manif} F (\x, \u) \\
M &\coloneqq \begin{bmatrix} (1+\beta) R & -\beta R \\ I & 0 \end{bmatrix} \\
T &\coloneqq \begin{bmatrix} S & 0 \\ 0 & S \end{bmatrix} \,,
\end{aligned}
\end{equation}
and $\zzh \coloneqq (\xzh, \xzmh) $, we can write \eqref{itr:APG:FPFAD} more compactly as:
\begin{equation} \label{eq:FPRAD:FPI}
\zkph \coloneqq M\zkh + T (\ud, 0) \,.
\end{equation}
Note that even though we do not show it explicitly, all the operators in \eqref{eq:linops} depend on $\x, \u$ and $\alpha$. $M$ additionally depends on $\beta$. For small $\norm{\x - \xmin}$, we can show that $\rho (M) < 1$ analogously to the proof of Theorem~\ref{thm:APG:FAD}, which implies that the sequence $\zkph$ converges to $(I-M)^{-1} T (\ud, 0)$. Since $(I-M) (\phi(\x, \u) \ud, \phi(\x, \u) \ud) = ((I-R) \phi(\x, \u) \ud, 0)$, our goal is to show that $(I-R) \phi(\x, \u) = S$. We therefore simplify
\begin{equation*}
\begin{aligned}
I - R &= I - Q^\dagger \big( I - \alpha \grad_{\x}^2 f(\x, \u) \big) \\
&= Q^\dagger \big( Q - \projTan (\x) + \alpha \projTan (\x) \grad_{\x}^2 f(\x, \u) \big) \\
&= Q^\dagger \big( \wein{\cdot}{\projNor (\x) (\alpha \bnu} + \alpha \grad_{\manif}^2 g(\x, \u) + \alpha \projTan (\x) \grad_{\x}^2 f(\x, \u) \big) \\
&= \alpha Q^\dagger \big ( \grad_{\manif}^2 F(\x, \u) - \wein{\cdot}{\alpha \projNor (\x) (\bnu + \grad_{\x} f (\x, \u))} \big ) \,,
\end{aligned}
\end{equation*}
to obtain
\begin{equation} \label{eq:FPFAD:conc}
(I - R)^{-1} S = - \big ( \grad_{\manif}^2 F(\x, \u) - \wein{\cdot}{\alpha \projNor (\x) (\bnu + \grad_{\x} f (\x, \u))} \big )^\dagger D_{\u} \grad_{\manif} F (\x, \u) \,,
\end{equation}
which is by definition $\phi (\x, \u)$.

When $\x = \xmin$, the term $\wein{\cdot}{\alpha \projNor (\x) (\bnu + \grad_{\x} f (\x, \u))}$ in \eqref{eq:IFT:approx} vanishes (see Lemma~\ref{lem:sg:seq}) leaving us with $D \argmap (\u)$ given in \eqref{eq:IFT}.

The convergence of $\xkh$ is linear with rate $\rho (M)$ because \eqref{eq:FPRAD:FPI} is a linear fixed-point iteration. Since the rate of convergence of \eqref{itr:APG} is $\rho (M_*)$ \cite[Theorem~4.13]{LFP17} and $M$ approaches $M_*$ as $\x$ approaches $\xmin$, the claim follows.
\end{proof}

\begin{REM} \label{rem:APG:FPFAD}
\begin{enumerate}
\item \label{itm:affine:manif}When $\manif$ is an affine manifold, the Weingarten term in \eqref{eq:IFT:approx} vanishes and the expression for $\phi$ simplifies to
\begin{equation} \label{eq:IFT:approx:affine}
\phi (\x, \u) = - \grad_{\manif}^2 F(\x, \u) ^\dagger D_{\u} \grad_{\manif} F (\x, \u) \,.
\end{equation}
Thus in this case, \algref{alg:APG:FPFAD} converges to a quantity resembling $D\argmap (\u)$ in \eqref{eq:IFT} with $\xmin$ replaced by $\x$.
\item The additional assumption on the sequence $\wk$ in \eqref{itr:APG} in Remark~\ref{rem:prox:ND} to ensure the differentiability of the proximal mapping in \eqref{itr:APG:FAD} is not required here. The iterations in \eqref{itr:APG:FPFAD} are already well-defined under Assumption~\ref{ass:ND} when $\x$ is close enough to $\xmin$.
\end{enumerate}
\end{REM}

\subsection{Reverse Mode FPAD of APG} \label{ssec:FPAD:RM}
We perform reverse mode FPAD on \algref{alg:APG} by modifying \eqref{itr:APG:RAD} similarly by using $\bnu\coloneqq \proj{\partial_{\x} g (\x, \u)} (-\grad_{\x} f (\x, \u))$ as the direction to compute $\w \coloneqq \x + \alpha \bnu$.
\begin{Algorithm}[Reverse Mode FPAD of \algref{alg:APG}] \ \label{alg:APG:FPRAD}
\begin{itemize}
\item \key{Initialization:} $\u\in\spPrm$, $\xKt=\xminb\in\spVar^*$, $ \yKt=0\in\spVar^*$ and $\uzt=0\in\spPrm^*$, $0 < \sslow\leq\ssup < 2/L$.
\item \key{Parameter:} $\alpha \in [\sslow, \ssup]$ and $\beta \in [0, 1]$.
\item \key{Define:} $\w \coloneqq \x + \alpha \proj{\partial_{\x} g (\x, \u)} (-\grad_{\x} f (\x, \u))$.
\item \key{Update $(n\geq 0$, $k\coloneqq K-n-1)$:}
\begin{equation} \label{itr:APG:FPRAD} \tag{APG-\~R}
\begin{aligned}
\wkt &\coloneqq \xkpt D_{\x} \bwd (\w, \u) \\
\ykt &\coloneqq \wkt (I - \alpha \grad_{\x}^2 f(\x, \u)) \\
\unpt &\coloneqq \unt + \xkpt D_{\u} \bwd (\w, \u) - \alpha \wkt \grad_{\x\u} f(\x, \u) \\
\xkt &\coloneqq (1+\beta ) \ykt - \beta \ykpt \,,
\end{aligned}
\end{equation}
\end{itemize}
\end{Algorithm}
The following result shows the convergence of \algref{alg:APG:FPRAD}.
\begin{THM}
Let $f$ and $g$ satisfy Assumption~\ref{ass:CPSO} and $(\xmin, \u)\in\manif\times\setPrm$ be such that Assumptions~\ref{ass:RPD} and \ref{ass:ND} are satisfied. Let $ \alpha \in [\sslow, \ssup]$ and $\beta \in [0, 1]$ be such that $-1/(1 + 2\beta) < \underline{\lambda}$ where $\underline{\lambda}$ is the smallest eigenvalue of $D_{\x} \pgd (\xmin, \u)$. For any $\x\in\manif$ sufficiently close to $\xmin$, the sequence $( \unt )_{n \in\N}$ generated by \algref{alg:APG:FPRAD} converges to $\xminb \phi(\x, \u)$ where $\phi$ is defined in \eqref{eq:IFT:approx}. Moreover, the convergence of $\unt$ is linear and its convergence rate approaches that of \algref{alg:APG} as $\x$ approaches $\xmin$.
\end{THM}

\begin{proof}
We will make use of the linear mappings defined in \eqref{eq:linops}. Since $D_{\x} \bwd (\w, \u) = Q^\dagger$, we write
\[
\ykt = \xkpt Q^\dagger (I - \alpha \grad_{\x}^2 f(\x, \u)) = \xkpt R \,,
\]
which gives us
\[
\xkt = (1+\beta) \xkpt R - \beta \xkppt R = (\xkpt, \xkppt) \begin{bmatrix} (1+\beta) R \\ -\beta R \end{bmatrix} \,.
\]
By defining $\zkt \coloneqq (\xkt, \xkpt) $, we get the update $\zkt = \zkpt M^*$ which can be expanded to $\zkt = \zKt (M^*)^{K-k} = \zKt (M^*)^{n+1}$ where we set $\xKpt = 0$. We now write
\begin{equation*}
\begin{aligned}
\unpt &= \unt - \alpha\xkpt Q^\dagger D_{\u} \grad_\manif g (\x, \u) - \alpha \xkpt Q^\dagger \grad_{\x\u} f(\x, \u) \\
&= \unt - \alpha\xkpt Q^\dagger D_{\u} (\grad_\manif g (\x, \u) + \grad_{\x} f(\x, \u)) \\
&= \unt - \alpha\xkpt Q^\dagger D_{\u} \grad_\manif F (\x, \u) \\
&= \unt + \xkpt S \,.
\end{aligned}
\end{equation*}
We similarly define $\vnt \coloneqq (\unt, \unpt) $ with $ \uzmt \coloneqq 0$ and obtain $ \vnt = \vnmt + \zkt T $. By recursively expanding this expression, we arrive at
\begin{equation*}
\begin{aligned}
\vnt &= \vnmt + \zKt (M^*)^n T \\
&= \vzt + \zKt T + \ldots + \zKt (M^*)^n T \\
&= \zKt (I + \ldots + (M^*)^n) T \,. \\
\end{aligned}
\end{equation*}
When $\rho (M) < 1$, we note that $\vnt \to \zKt (I - M^*)^{-1} T $. In order to compute the inverse of
\[
I - M^* = \begin{bmatrix} I - R - \beta R & -I \\ \beta R & I \end{bmatrix} \,,
\]
we use \cite[Theorem~1]{LS02} to compute the Schur complement of the bottom right block which equals $I - R$. The inverse therefore is given by
\[
(I - M^*)^{-1} = \begin{bmatrix} (I - R)^{-1} & (I - R)^{-1} \\ \beta R (I - R)^{-1} & I - \beta R (I - R)^{-1} \end{bmatrix} \,.
\]
Since $\zKt = (\xminb, 0)$ we get
\begin{equation*}
\begin{aligned}
\zKt (I - M^*)^{-1} T &= (\xminb (I - R)^{-1}, \xminb (I - R)^{-1}) T \\
&= (\xminb (I - R)^{-1} S, \xminb (I - R)^{-1} S) \,,
\end{aligned}
\end{equation*}
which is equal to $ (\xminb \phi (\x, \u), \xminb \phi (\x, \u)) $ from \eqref{eq:FPFAD:conc}. The argument for linear convergence is similar to what we showed in the proof of Theorem~\ref{thm:APG:FPFAD}.
\end{proof}
\begin{REM}
Again, when $\manif$ is affine \algref{alg:APG:FPRAD} converges to an approximation of $D\argmap(\u)$ in \eqref{eq:IFT} at $\x$ (see Remark~\ref{rem:APG:FPFAD}\ref{itm:affine:manif}).
\end{REM}

\section{Implicit Differentiation of Fixed-Point Equation} \label{sec:FPID}
In our setting, a promising counterpart to (FP)AD of \algref{alg:APG} is Implicit Differentiation applied to the fixed-point equation for PGD, i.e.,
\begin{equation} \label{eq:PGD:FPE}
\xmin = \pgd (\xmin, \u) \,.
\end{equation}
for some $(\xmin, \u)\in\manif\times\setPrm$ with $\xmin=\argmin_{\x\in\spVar} F(\x, \u)$ where $\pgd$ is defined in \eqref{eq:prox:grad}. Implicit Differentiation is more common in Machine Learning Community [4, 2, 7] for differentiating through a fixed-point equation. In this section, we show that for computing the derivative of the solution mapping, both FPAD of \algref{alg:APG} and ID of \eqref{eq:PGD:FPE} solve the same system of Linear Equation. Following our results from the previous section, we show the equivalence of the two methods when they are evaluated at any point in some neighbourhood of the fixed-point.

Differentiating \eqref{eq:PGD:FPE} with respect to $\u$ yields
\begin{equation} \label{eq:FPE:ID}
D_{\u} \xmin = D_{\x} \bwd (\w^*, \u) \big ( (I - \alpha \grad_{\x}^2 f (\xmin, \u) ) D_{\u} \xmin - \alpha \grad_{\x\u} f (\xmin, \u) \big ) + D_{\u} \bwd (\w^*, \u) \,,
\end{equation}
where $\w^* = \xmin - \grad_{\x} f (\xmin, \u)$. When $I - D_{\x} \pgd (\xmin, \u)$ is invertible, we obtain an expression for $D_{\u} \xmin$ which can be shown to be equal to $\argmap (\u)$ in \eqref{eq:IFT}. However, like FPAD, we consider a more general setting where we want to evaluate $D_{\u} \xmin$ from \eqref{eq:FPE:ID} by using an estimate $\x\in\manif$ of $\xmin$, i.e.,
\begin{equation} \label{eq:FPE:ID:approx}
X = D_{\x} \bwd (\w, \u) \big ( (I - \alpha \grad_{\x}^2 f (\x, \u) ) X - \alpha \grad_{\x\u} f (\x, \u) \big ) D_{\u} \x + D_{\u} \bwd (\w, \u) \,,
\end{equation}
where $\w = \x + \alpha \bnu$ with $\bnu = \proj{\partial_{\x} g(\x, \u)} (-\grad_{\x} f(\x, \u))$ and we replace $D_{\u} \xmin$ with $X$. The following result shows that $\phi (\x, \u)$ defined in \eqref{eq:IFT:approx} solves \eqref{eq:FPE:ID:approx}.

\begin{THM} \label{thm:PGD:ID}
Let $f$ and $g$ satisfy Assumption~\ref{ass:CPSO} and $(\xmin, \u)\in\manif\times\setPrm$ be such that Assumptions~\ref{ass:RPD} and \ref{ass:ND} are satisfied. For any $\x\in\manif$ sufficiently close to $\xmin$ and $\alpha \in [\sslow, \ssup]$, solving \eqref{eq:FPE:ID:approx} with respect to $X$ yields the solution $\phi (\x, \u)$ defined in \eqref{eq:FPE:ID:approx}.
\end{THM}
\begin{proof}
Using definitions of $R$ and $S$ in \eqref{eq:linops} and noting that $D_{\x} \bwd (\w, \u) = Q^{\dagger}$, we get
\begin{equation*}
\begin{aligned}
X &= D_{\x} \bwd (\w, \u) \big ( (I - \alpha \grad_{\x}^2 f (\x, \u) ) X - \alpha \grad_{\x\u} f (\x, \u) \big ) + D_{\u} \bwd (\w, \u), \u) \\
&= Q^\dagger (I - \alpha \grad_{\x}^2 f (\x, \u) ) X - \alpha Q^\dagger D_{\u} \grad_{\x} f (\x, \u) - \alpha Q^\dagger D_{\u} \grad_{\manif} g (\x, \u) \\
&= R X - \alpha Q^\dagger D_{\u} \grad_{\manif} F (\x, \u) \\
&= R X + S \,.
\end{aligned}
\end{equation*}
Therefore $(I - R)^{-1} S$ solves \eqref{eq:FPE:ID:approx} which we know from \eqref{eq:FPFAD:conc} is the same as $\phi (\x, \u)$.
\end{proof}
\begin{REM}
\begin{enumerate}
\item Similarly, when $\manif$ is an affine manifold, ID of \eqref{eq:PGD:FPE} computes a simpler quantity (given in \eqref{eq:IFT:approx:affine}) which resembles $D\argmap (\u)$ in \eqref{eq:IFT}.
\item Theorems~\ref{thm:APG:FAD}, \ref{thm:APG:FPFAD} and \ref{thm:PGD:ID} assert that Fixed-Point Automatic Differentiation of \algref{alg:APG} and Implicit Differentiation of \eqref{eq:FPE:ID} are equivalent in the sense that they aim at solving a system of linear equation to compute $\phi (\x, \u)$.
\end{enumerate}
\end{REM}

% ***********************
% >>>>> EXPERIMENTS <<<<<
% ***********************
\section{Experiments} \label{sec:Exp}

We now put our results to test by applying them on some practical applications. We first show, in Section~\ref{ssec:Exp:CRD}, the experimental validation of convergence and convergence rates of AD and FPAD applied to \eqref{itr:APG} proved in Sections~\ref{sec:AD} and \ref{sec:FPAD} on Lasso [80] and Group Lasso [85]. We verify with our experiments that the convergence rate of FPAD of \eqref{itr:APG} is at least as good as that of AD of \eqref{itr:APG}. In Section~\ref{ssec:Exp:ID}, we solve a bilevel optimization problem where we learn the parameters of the regularizer for image denoising. The parameter learning in this case is the upper-level problem and solving it requires to compute the derivative of the solution mapping of the lower-level problem, i.e., the denoising problem. We perform all the experiments in PyTorch [71].

\subsection{Convergence Rate Demonstration} \label{ssec:Exp:CRD}
Given a matrix $A\in\R^{M\times N}$ and a vector $\b\in\R^M$, the lasso problem [80] aims to find a sparse vector $\x^*\in\R^N$ which roughly solves $A\x = \b$, formulated as the following optimization problem
\begin{equation} \label{prob:lasso} \tag{L}
\min_{\x\in\R^N} F_\mathrm{l}(\x, A, \b, \lambda), \quad F_\mathrm{l}(\x, A, \b, \lambda) \coloneqq \frac{1}{2} \norm[2]{A\x - \b}^2 + \lambda \norm[1]{\x} \,,
\end{equation}
for some $\lambda>0$. Here $\norm[2]{\cdot}$ and $\norm[1]{\cdot}$ correspond to $l_2$ and $l_1$ norms defined on $\R^N$. The second part of the objective in \eqref{prob:lasso} controls the sparsity of the solution $\xmin$ of the problem where larger $\lambda$ makes the solution sparser. In a similar setting, given multiple vectors $\b_i \in \R^M$ for $i = 1,\ldots,L$, we are often interested in finding sparse solutions $\xmin_i \in \R^N$ which approximately satisfy $A\x_i = \b_i$ for all $i$ in such a way that all the solutions share the same sparsity pattern. This amounts to solving the following problem
\begin{equation} \label{prob:group:lasso} \tag{GL}
\min_{X\in\R^{N\times L}} F_{\mathrm{gl}}(X, A, B, \lambda), \quad F_{\mathrm{gl}}(X, A, B, \lambda) \coloneqq \frac{1}{2} \norm[2]{AX - B}^2 + \lambda \norm[2,1]{X} \,.
\end{equation}
where we stack the vectors $\b_i$ as column vectors in the matrix $B\in\R^{M\times L}$. Here $\norm[2]{\cdot}$ is the Frobenius norm and $\norm[2,1]{\cdot}$ is the $l_{2,1}$ group norm defined on $\R^{N\times L}$ by
\[
\norm[2,1]{X} \coloneqq \sum_{i} \sqrt{\sum_{j} x_{i,j}^2} \,.
\]
The problem of finding the group sparsity pattern is called group lasso [85] in the literature. In our experiments, we set $u\coloneqq (A, \b, \lambda)$ for \eqref{prob:lasso} and $u\coloneqq (A, B, \lambda)$ for \eqref{prob:group:lasso}.

\begin{figure}[h]
\centering
\begin{subfigure}[b]{0.98\textwidth}
\centering
\includegraphics[width=0.98\textwidth]{figs/legends.pdf}
\label{sfig:cr:legs}
\end{subfigure}
\begin{subfigure}[b]{0.45\textwidth}
\centering
\includegraphics[width=0.98\textwidth]{figs/Lasso.pdf}
\label{sfig:cr:lasso}
\end{subfigure}
\begin{subfigure}[b]{0.45\textwidth}
\centering
\includegraphics[width=0.98\textwidth]{figs/GLasso.pdf}
\label{sfig:cr:glasso}
\end{subfigure}
\caption{Error Plots for various sequences obtained from \eqref{prob:lasso} (left) and \eqref{prob:group:lasso} (right). The solid lines represent PGD and its derivative sequences while the dotted lines represent all the APG sequences. The lines without any markers correspond to the optimization algorithms. The lines with circular markers represent forward mode sequences while those with square markers represent Reverse Mode sequences. Markers for FPAD sequences are larger than their AD counterparts. Our FPAD sequences are more stable and show faster convergence to the true derivative.}
\label{fig:conv:rates}
\end{figure}

Since the $l_1$ norm and the $l_{2,1}$ group norm are partly smooth [81], the two problems presented above fit the description of \eqref{prob:min} and satisfy Assumptions~\ref{ass:CPSO} and \ref{ass:RPD} for all $\u$ when $A$ is full-rank and $M\geq N$. We verify Assumption~\ref{ass:ND} experimentally by first solving the problem for given $\u$ and check if \eqref{eq:ND} is satisfied or not. As an example, we explain this procedure for \eqref{prob:lasso}. Let $\xmin$ be the solution of \eqref{prob:lasso} for some given $\u$. The vector $\v \coloneqq -\grad_{\x} f(\x, \u) / \lambda = A^T (\b - A\xmin) / \lambda$ must be a subgradient of $\norm[1]{\cdot}$ at $\xmin$. Therefore for all $i$, $x^*_i \neq 0$ simply implies $v_i = 0$ and for \eqref{eq:ND} to hold at $(\xmin, \u)$, $v_i$ must lie in the open interval $(-1, 1)$ whenever $x^*_i = 0$.

For both problems, we perform each algorithm from \algref{alg:APG} to \algref{alg:APG:FPRAD} twice, once with $\beta_k=0$ (PGD) and once with $\beta_k>0$ (APG) to generate $20$ sequences in total; $10$ for each problem (see Figure~\ref{fig:conv:rates}). For simplicity, we do not construct the full Jacobian of the iterates of PGD and APG for AD and FPAD. We instead just perform the AD and FPAD algorithms for some $\ud\in\R^N$ and $\xminb\in\R^{1\times N}$ (a row vector). For computing the error sequences, for instance $\norm[2]{\xk-\xmin}$ for \eqref{itr:APG} and $\norm[2]{\xkd - \xmind}$, we run \eqref{itr:APG} for a long time to compute a very good estimate of $\xmin$ and use it to compute $\xmind=D\argmap (\u)\ud$ and $\ub=\xminb D\argmap (\u)$ where $D\argmap (\u)$ is defined in \eqref{eq:IFT}. Again, as an example, we outline our approach for computing $D\argmap (\u)$ for \eqref{prob:lasso}. The Riemannian Hessian of $\norm[1]{\cdot}$ vanishes everywhere while its Riemannian Gradient is $(\sign (x_1), \ldots, \sign (x_N))$ \cite[Example~5.2.1]{Lia16} where $\map{\sign}{\R}{\set{-1, 0, 1}}$ is defined by
\[
\sign(t) \coloneqq \begin{cases}
-1, & \textnormal{if} \ t < 0 \,; \\
\;0, & \textnormal{if} \ t = 0 \,; \\
+1, & \textnormal{if} \ t > 0 \,.
\end{cases}
\]
The affine manifold $\manif$ and $\spTan\xmin$ are both given by $\set{\x\in\R^N:\supp{\x}\subset\supp{\xmin}}$, where $\supp{\x}$ denotes the support of a vector $\x\in\R^N$ defined as $\supp{\x} \coloneqq \set{i\in\N : x_i \neq 0}$.
In other words, for all $\x$,
\begin{equation} \label{eq:lasso:manif}
\x\in\manif=\spTan\xmin \quad\iff\quad (x^*_i = 0 \implies x_i = 0) \,.
\end{equation}
This allows us to compute the Riemannian Gradient and Hessian of $F_{\mathrm{l}}$ as
\begin{equation*}
\begin{aligned}
\grad_{\manif} F_{\mathrm{l}} (\x, \u) &= \proj{\spTan\xmin} (A^T(A\x - \b)) + \lambda \grad_{\manif} \norm[1]{\cdot} (\x) \\
\grad_{\manif}^2 F_{\mathrm{l}} (\x, \u) &= \proj{\spTan\xmin} \circ A^TA \circ \proj{\spTan\xmin} \,.
\end{aligned}
\end{equation*}
Thus from \eqref{eq:lasso:manif}, for all $i, j \notin \supp{\xmin}$, $(\grad_{\manif} F_{\mathrm{l}} (\x, \u))_i = 0$ and $(\grad_{\manif}^2 F_{\mathrm{l}} (\x, \u))_{i,j} = 0$ and we can solve a reduced linear system of $\abs{\supp{\xmin}}$ equations. The derivative $D_{\u} \grad_{\manif} F(\x, \u)$ can be computed by using autograd package.

For \eqref{prob:lasso}, we set $M=1000$, and $N=250$ and for \eqref{prob:group:lasso}, we set $M=1000$, $N=100$, and $L=40$. We generate $A$ by sampling each $a_{i, j} $ from a uniform distribution with parameters $0$ and $1$ and $\b$ and $B$ by sampling each $b_i$ and $b_{i,j}$ from a normal distribution with mean $0$ and standard deviation $1$. We run each algorithm on \eqref{prob:lasso} for $K=5000$ iterations and on \eqref{prob:group:lasso} for $K=1000$ iterations. The computed forward mode sequences are $(\xkd)_{k\in[K]}$ and $(\xkh)_{k\in[K]}$ while the reverse mode sequences are $(\unb)_{n\in[K]}$ and $(\unt)_{n\in[K]}$. For a fair comparison, we use the same initializations for all algorithms. For performing FPAD on an optimization algorithm, we use its own output as approximate fixed-point. For example, we run PGD on \eqref{prob:lasso} for $K=5000$ iterations and use this $\xK$ as $\x$ in \eqref{itr:APG:FPFAD} and \eqref{itr:APG:FPRAD} for performing Forward and Reverse Mode FPAD on PGD. We use a constant step size for all iterations $\alpha_k = 1/\norm[\mathrm{op}]{A^TA}$, where $\norm[\mathrm{op}]{\cdot}$ is the operator norm and for APG, we set $\beta_k \coloneqq (k-1) / (k+q)$ with $q=5$ and for FPAD on APG, we set $\beta\coloneqq\beta_K$. For computing $\xmin$, we run \eqref{itr:APG} for a large number of iterations until it satisfies $\norm[2]{\xmin - \pgd (\xmin, \u)} < 10^{-12}$. We use Conjugate Gradient Method [51] to compute $\xmind$ and $\ub$ for the two problems.

In Figure~\ref{fig:conv:rates}, we plot all the error sequences. The left figure shows the sequences computed for \eqref{prob:lasso} while the right one shows those obtained from \eqref{prob:group:lasso}. In both plots, we note that APG algorithm (dashed blue lines) converges faster than PGD (solid blue lines). This acceleration effect is translated to the FPAD sequences; the sequences obtained by applying FPAD on APG (dashed red lines with circle for forward mode and square markers for Reverse Mode) converge more quickly as compared to those generated by applying FPAD on PGD (see corresponding solid red lines). For both problems, PGD requires more time to converge and therefore does not provide satisfying results for any derivative sequence (all the solid lines move slowly towards $0$). For \eqref{prob:group:lasso}, the forward and reverse mode AD of APG (dashed green lines) yield good results but after $K$ iterations, we observe that FPAD gives a better estimate (dashed red lines). For \eqref{prob:lasso}, the behaviour of forward mode AD on APG (dashed green line with circular marker) in the beginning and that of reverse mode AD on APG (dashed green line with square marker) becomes bad in the end is quite erratic. One possibility for this situation might be that \eqref{itr:APG} may not behave well in the initial iterations. Note that the derivative is a local quantity and therefore convergence of the generated sequence of derivatives may be hampered in initial iterates that are far away from this local neighbourhood.

% Conclusion
In the above experiment we saw that FPAD iterates behave as well as the original iterates and the convergence for the two types of sequences is very similar. We noted that the FPAD sequences also benefit from inertia just like the optimization algorithm they are being applied to. We further found out that AD applied in a naive way may not provide a good result in the end and the algorithm must be initially run for some time before starting to construct the computational graph. These findings together with the enormous memory-gain boost our FPAD as a practical approach for computing derivative in such a setting.

\subsection{Learning Regularizers for Image Denoising} \label{ssec:Exp:ID}

We now demonstrate the effectiveness of FPAD by using it to solve a simple yet practical bilevel optimization problem. Let $\spImg\coloneqq\R^{N_x\times N_y\times N_c}$ and $\spFeat\coloneqq\R^{N_x\times N_y\times N_c\times N_f}$, we consider the following optimization problem
\begin{equation} \label{prob:ID}
\min_{\recovered\in\spImg} \frac{1}{2}\norm[2]{\recovered - \noisy}^2 + \lambda \norm[\mathrm{r}]{\imgLinOp \recovered} \,,
\end{equation}
where $\lambda > 0$ and $\imgLinOp\in\spLin (\spImg, \spFeat$). When $N_c\in\set{1,3}$, $N_f=2$, $\imgLinOp$ is a convolution operator computing the forward $x$ and $y$-derivative of its input image with Neumann boundary condition and $\norm[\mathrm{r}]{\cdot}$ is an appropriate norm like the $l_1$ norm or the $l_{2,1}$ group norm defined by
\[
\norm[2,1]{\p} \coloneqq \sum_{i,j} \sqrt{\sum_{k,r} p_{i,j,k,r}^2} \,,
\]
the above problem corresponds to total variation image denoising [77] which aims to produce a denoised image $\recovered^*$ from a noisy image $\noisy$. Here the first term in the objective, also known as the data or the fidelity term makes sure that $\recovered^*$ is not too far away from $\x$. The second term, called the regularization term, ensures the sparsity of the derivative of the image which relies on the assumption that a natural image is ``sort of'' piecewise constant and is controlled by regularization parameter $\lambda$. In [29], the authors list various variants of this model, for instance, image deblurring, inpainting and zooming and list appropriate algorithms to solve the problem, which our experiment easily generalizes to. The best way to solve \eqref{prob:ID} is by applying APG to its Fenchel-Rockafellar dual
\[
\min_{\edges\in\spFeat} \frac{1}{2}\norm[2]{\noisy - \imgLinOp^*\edges}^2 \quad \st \quad \norm[\mathrm{r},*]{\edges} \leq \lambda \,,
\]
and using the dual solution $\edges^*$ to compute its primal solution $\recovered^* = \noisy - \imgLinOp^*\edges^*$. Here $\norm[\mathrm{r},*]{\cdot}$ is the dual norm of $\norm[\mathrm{r}]{\cdot}$.

\begin{figure}[!t]
\def\evolwidth{0.115}
\captionsetup[subfigure]{font=tiny,aboveskip=1pt,labelformat=empty}
\centering

\subfig{\evolwidth}{figs/denoisinig/train/}{Ground}{1}{}
\subfig{\evolwidth}{figs/denoisinig/train/}{Noisy}{1}{\ ($18.58$)}
\subfigevol{\evolwidth}{train}{1}{0}{$18.65$}
\subfigevol{\evolwidth}{train}{1}{1}{$19.52$}
\subfigevol{\evolwidth}{train}{1}{2}{$19.33$}
\subfigevol{\evolwidth}{train}{1}{10}{$19.76$}
\subfigevol{\evolwidth}{train}{1}{20}{$19.85$}
\subfigevol{\evolwidth}{train}{1}{30}{$19.82$}
\hfill%

\vspace{1pt}
\subfig{\evolwidth}{figs/denoisinig/train/}{Ground}{2}{}
\subfig{\evolwidth}{figs/denoisinig/train/}{Noisy}{2}{\ ($19.05$)}
\subfigevol{\evolwidth}{train}{2}{0}{$20.25$}
\subfigevol{\evolwidth}{train}{2}{1}{$22.34$}
\subfigevol{\evolwidth}{train}{2}{2}{$22.25$}
\subfigevol{\evolwidth}{train}{2}{10}{$22.29$}
\subfigevol{\evolwidth}{train}{2}{20}{$22.35$}
\subfigevol{\evolwidth}{train}{2}{30}{$22.52$}
\hfill%

\vspace{1pt}
\subfig{\evolwidth}{figs/denoisinig/train/}{Ground}{3}{}
\subfig{\evolwidth}{figs/denoisinig/train/}{Noisy}{3}{\ ($18.24$)}
\subfigevol{\evolwidth}{train}{3}{0}{$20.44$}
\subfigevol{\evolwidth}{train}{3}{1}{$23.40$}
\subfigevol{\evolwidth}{train}{3}{2}{$23.26$}
\subfigevol{\evolwidth}{train}{3}{10}{$22.83$}
\subfigevol{\evolwidth}{train}{3}{20}{$23.34$}
\subfigevol{\evolwidth}{train}{3}{30}{$23.33$}
\hfill%

\vspace{1pt}
\subfig{\evolwidth}{figs/denoisinig/train/}{Ground}{4}{}
\subfig{\evolwidth}{figs/denoisinig/train/}{Noisy}{4}{\ ($18.06$)}
\subfigevol{\evolwidth}{train}{4}{0}{$19.06$}
\subfigevol{\evolwidth}{train}{4}{1}{$21.05$}
\subfigevol{\evolwidth}{train}{4}{2}{$20.82$}
\subfigevol{\evolwidth}{train}{4}{10}{$20.81$}
\subfigevol{\evolwidth}{train}{4}{20}{$21.09$}
\subfigevol{\evolwidth}{train}{4}{30}{$21.14$}
\hfill%

\vspace{1pt}
\subfig{\evolwidth}{figs/denoisinig/train/}{Ground}{5}{}
\subfig{\evolwidth}{figs/denoisinig/train/}{Noisy}{5}{\ ($18.07$)}
\subfigevol{\evolwidth}{train}{5}{0}{$19.45$}
\subfigevol{\evolwidth}{train}{5}{1}{$23.15$}
\subfigevol{\evolwidth}{train}{5}{2}{$23.02$}
\subfigevol{\evolwidth}{train}{5}{10}{$22.59$}
\subfigevol{\evolwidth}{train}{5}{20}{$22.81$}
\subfigevol{\evolwidth}{train}{5}{30}{$22.86$}
\hfill%

\vspace{1pt}
\subfig{\evolwidth}{figs/denoisinig/valid/}{Ground}{1}{}
\subfig{\evolwidth}{figs/denoisinig/valid/}{Noisy}{1}{\ ($18.95$)}
\subfigevol{\evolwidth}{valid}{1}{0}{$20.26$}
\subfigevol{\evolwidth}{valid}{1}{1}{$21.36$}
\subfigevol{\evolwidth}{valid}{1}{2}{$21.22$}
\subfigevol{\evolwidth}{valid}{1}{10}{$21.47$}
\subfigevol{\evolwidth}{valid}{1}{20}{$21.58$}
\subfigevol{\evolwidth}{valid}{1}{30}{$21.63$}
\hfill%

\vspace{1pt}
\subfig{\evolwidth}{figs/denoisinig/valid/}{Ground}{2}{}
\subfig{\evolwidth}{figs/denoisinig/valid/}{Noisy}{2}{\ ($17.99$)}
\subfigevol{\evolwidth}{valid}{2}{0}{$19.35$}
\subfigevol{\evolwidth}{valid}{2}{1}{$21.34$}
\subfigevol{\evolwidth}{valid}{2}{2}{$21.19$}
\subfigevol{\evolwidth}{valid}{2}{10}{$20.96$}
\subfigevol{\evolwidth}{valid}{2}{20}{$21.18$}
\subfigevol{\evolwidth}{valid}{2}{30}{$21.16$}
\hfill%

\caption{Performance of our model depicted over the epochs on train (first five rows) and test images (last two rows). The first two columns correspond to ground truth $\ground[i]$ and the noisy images $\noisy[i]$ respectively while the remaining six columns correspond to the recovered image $\recovered[i]\iter K (\Theta_l)$ for $l\in\set{0,1,2,10,20,30}$. Below each image in brackets, we show PSNR value of each image relative to the ground truth.}
\label{fig:ID:evol}
\end{figure}

Instead of simply using $\imgLinOp$ as the discrete derivative operator, one may instead try to learn it. This problem has also been well-studied and numerous strategies have been proposed (see e.g., [76, 34, 54, 53]). We take inspiration from [34] and define $\imgLinOp$ by $\imgLinOp\u \coloneqq (\lrndFil{1} \u, \ldots, \lrndFil{N_f} \u)$ where each $\lrndFil{r}$ is a linear combination of $N_b$ basis filters $\basFil{s} \in \spLin (\spImg, \spImg)$, i.e., $\lrndFil{r} = \sum_s \theta_{r, s} \basFil{s}$ for some $\Theta\in\R^{N_f\times N_b}$. Since $\imgLinOp$ is linear in $\Theta$, we can ignore $\lambda$ by simply absorbing it into $\Theta$. An example of $\basFil{s}$ is the $2$D DCT basis [3]. Our goal is to learn the weight matrix $\Theta$. In particular, given a training dataset $\trainset\coloneqq\set{(\noisy[i], \ground[i]) : i=1,\ldots,M_{\textnormal{tr}}}$, we solve the following bilevel problem
\begin{equation} \tag{BL} \label{prob:BL}
\begin{aligned}
\min_{\Theta} \quad & \sum_{i=1}^{M_{\textnormal{tr}}} J_i (\Theta) \,, \quad J_i (\Theta) \coloneqq \frac{1}{2}\norm[2]{\recovered[i]^* (\Theta) - \ground[i]}^2 \\
\st \quad & \recovered[i]^* (\Theta) = \argmin_{\recovered\in\spImg} \frac{1}{2}\norm[2]{\recovered - \noisy[i]}^2 + \norm[\mathrm{r}]{\imgLinOp (\Theta) \recovered} \,,\quad i = 1,\ldots,M_{\textnormal{tr}} \,.
\end{aligned}
\end{equation}
The best way to solve \eqref{prob:BL} is by using stochastic algorithms [24], e.g., Stochastic Gradient method. The idea is to randomly sample $(\noisy[i], \ground[i])$ from $\trainset$ to compute $\grad J_i (\Theta) = D_{\Theta}\recovered[i]^* (\Theta) (\recovered[i]^* (\Theta) - \ground[i])$ and use the update $\Theta\leftarrow\Theta - \tau \grad J_i (\Theta)$. In practice, stochastic algorithms consist of epochs where each epoch involves randomly shuffling the training dataset and drawing the samples $(\noisy[i], \ground[i])$ sequentially from the shuffled set until it is exhausted [47]. We will use \algref{alg:APG:FPRAD} to compute $\grad J_i (\Theta)$. The following algorithm lists the steps for solving \eqref{prob:BL}.
\begin{Algorithm} \ \label{alg:BL}
\begin{itemize}
\item \key{Initialization:} $\Theta_0 \in\R^{N_f \times N_b}$.
\item \key{Parameter:} $K\in\N$, $(\tau_l)_{l\in\N_0}$.
\item \key{Define:} $l\coloneqq0$, $i\coloneqq1$ and $\dataset\coloneqq\trainset$.
\item \key{Update:}
\begin{itemize}
\item Draw $(\noisy[i], \ground[i])$ uniformly from $\dataset$ and set $\dataset\leftarrow\dataset\backslash\set{(\noisy[i], \ground[i])}$.
\item Run \algref{alg:APG} with $\edges[i]\iter 0\coloneqq \imgLinOp\noisy[i]$ on
\begin{equation} \label{prob:ID:FRDual}
\min_{\edges\in\spFeat} \frac{1}{2}\norm[2]{\noisy[i] - \imgLinOp (\Theta)^*\edges}^2 \quad \st \quad \norm[\mathrm{r},*]{\edges} \leq 1 \,,
\end{equation}
for $K$ iterations to obtain $\edges[i]\iter K$.
\item Compute $\recovered[i] (\Theta)\iter K = \noisy[i] - \imgLinOp^* (\Theta) \edges[i]\iter K(\Theta)$ and $J_i\iter K (\Theta) = \frac{1}{2} \norm[2]{\recovered[i]\iter K (\Theta) - \ground[i]}^2$.
\item Compute $\grad J_i\iter K (\Theta)$ by some standard autograd package. Run \algref{alg:APG:FPRAD} for $K$ iterations for Back-propagation through \eqref{prob:ID:FRDual}.
\item Update $\Theta_{l+1} \coloneqq \Theta_{l} - \tau_l \grad J_i\iter K (\Theta_{l})$.
\item Increment $i \leftarrow i + 1$.
\item When $i>M_{\textnormal{tr}}$; increment $l \leftarrow l + 1$, reset $i\coloneqq 1$ and replenish $\dataset\coloneqq\trainset$.
\end{itemize}
\end{itemize}
\end{Algorithm}

\begin{figure}[!t]
\captionsetup[subfigure]{font=tiny,aboveskip=1pt,labelformat=empty}
\centering

\subfig{0.32}{figs/denoisinig/test/}{Ground}{1}{}
\subfig{0.32}{figs/denoisinig/test/}{Noisy}{1}{\ ($18.75$)}
\subfig{0.32}{figs/denoisinig/test/}{Recovered}{1}{\ ($24.79$)}
\hfill%

\vspace{1ex}
\subfig{0.32}{figs/denoisinig/test/}{Ground}{2}{}
\subfig{0.32}{figs/denoisinig/test/}{Noisy}{2}{Noisy ($18.92$)}
\subfig{0.32}{figs/denoisinig/test/}{Recovered}{2}{\ ($25.4$)}
\hfill%

\caption{Performance of our model after training for $30$ epochs on two unseen (larger) examples, with each row corresponding to one example. The columns respectively show the ground truth $\ground$, noisy $\noisy$ and recovered images $\recovered\iter K (\Theta_{30})$.}
\label{fig:ID:test}
\end{figure}

For our experimental setup, we use a training dataset of $5$ images and a test dataset of $2$ images, each being a colored image patch normalized to $[0, 1]$ with size $50\times50\times3$, taken from [63]. To each clean image, we add a Gaussian noise with standard deviation $40/255$ to generate the noisy images. We run \algref{alg:BL} for $30$ epochs and define the learning rate sequence $\tau_l \coloneqq 10^{-4} / (\floor{l/4}+1)$, i.e., we start with a learning rate of $10^{-4}$ and decay it linearly after every $4$ epochs. We use the built-in SGD optimizer of PyTorch for this purpose and set the momentum to $0.75$. We learn $N_f = 24$ filters by using $5\times5$ DCT bases without the constant basis vector, giving us $N_b = 24$ basis filters and a weight matrix $\Theta$ of size $24\times24$. For initialization, we draw each entry of $\Theta$ from a uniform distribution with parameters $0$ and $0.01$. We run \algref{alg:APG} and \algref{alg:APG:FPRAD} for $K=500$ iterations each. In the two algorithms, we set $\beta_k\coloneqq (k-q)/(k+1)$ and $\beta\coloneqq (K-q)/(K+1)$ with $q=5$ and use a constant step size $\alpha_k\coloneqq 1/L^2$ where $L$ is an upper bound on the operator norm of $\imgLinOp$. This upper bound can be computed by noting that the operator norm of a convolution operator with kernel matrix $C$ is bounded from above by $\norm[1]{C}$. In our experiments, we use PyTorch to implement the forward pass which involves running \algref{alg:APG} to obtain $\edges[i]\iter K$ for any given $\Theta$ followed by computing the recovered image $\recovered[i]\iter K$ and the batch loss $J_i\iter K$. \algref{alg:APG} is implemented as a function with custom gradient where we use \algref{alg:APG:FPRAD} to implement the backward pass.

In Figure~\ref{fig:ID:evol}, we show the evolution of our model over epochs by showing its performance on train image (first five rows) and test images (last two rows). Before training, the model outputs an image (third column) which is closer to the noisy image (second column). This happens because, the entries of $\Theta$ are very small and the regularization has close to no impact on the output. However, after only the first epoch (fourth column), the model starts to improve and is already close to the ground truth image (first column). We can also see how only with such a small dataset, the model learns to perform really well on unseen images like two images of same size as that of the training images (bottom two rows in Figure~\ref{fig:ID:evol}) and two larger images of size $481\times381\times3$ (the two rows in
Figure~\ref{fig:ID:test}). Figure~\ref{fig:ID:filters} displays the $N_f=24$ learned filters, i.e., the $24$ kernels each corresponding to some $\lrndFil{i}$ for $i=1,\ldots,N_f$.

\begin{figure}
\captionsetup[subfigure]{font=tiny,aboveskip=1pt,labelformat=empty}
\def\filtwidth{0.07}

\centering
\foreach \filt in {1,...,12}{%
\subfigfilt{\filtwidth}{\filt}
\hfill%

\vspace{1ex}
\centering
\foreach \filt in {13,...,24}{%
\subfigfilt{\filtwidth}{\filt}
\hfill%

\caption{$N_f=24$ learned filters after $30$ epochs.}
\label{fig:ID:filters}
\end{figure}

The above example clearly shows the effectiveness of FPAD because if we had used AD instead, the memory overhead would have been of the order of $K N_x N_y N_c = 3.75\times10^{6}$ which equates to $28.6$ MiB as opposed to $N_x N_y N_c = 7500$ or $58.6$ KiB of FPAD.

% **********************
% >>>>> CONCLUSION <<<<<
% **********************
\section{Conclusion} \label{sec:Conc}
In this paper, we studied the problem of differentiating the solution mapping of a structured parametric optimization problem. We showed that when the objective of this problem is partly smooth and some other minor assumptions are satisfied, the classical derivative of the solution mapping can be estimated by automatic differentiation of proximal splitting algorithms like Proximal Gradient Descent. We showed that the memory overhead of the reverse mode AD can be overcome by using our Fixed-Point Automatic Differentiation technique. We also showed that in terms of convergence speed, FPAD outperforms AD both in theory and in practice. We also showed the working of FPAD by solving a bilevel optimization problem to learn the regularizers for image denoising.

\section*{Acknowledgments}
Sheheryar Mehmood and Peter Ochs are supported by the German Research Foundation
(DFG Grant OC 150/4-1).

% *******************
% >>>>> APPENDIX <<<<<
% *******************

\appendix
% \section{Proofs}

% ************************
% >>>>> bibliography <<<<<
% ************************
\bibliographystyle{ieee}
\begin{thebibliography}{10}\itemsep=-1pt

[1]
M.~Abadi, P.~Barham, J.~Chen, Z.~Chen, A.~Davis, J.~Dean, M.~Devin,
S.~Ghemawat, G.~Irving, M.~Isard, M.~Kudlur, J.~Levenberg, R.~Monga,
S.~Moore, D.~G. Murray, B.~Steiner, P.~Tucker, V.~Vasudevan, P.~Warden,
M.~Wicke, Y.~Yu, and X.~Zheng.
\newblock Tensorflow: A system for large-scale machine learning.
\newblock In {\em 12th $\{$USENIX$\}$ Symposium on Operating Systems Design and
Implementation ($\{$OSDI$\}$ 16)}, pages 265--283, 2016.

[2]
A.~Agrawal, B.~Amos, S.~Barratt, S.~Boyd, S.~Diamond, and J.~Z. Kolter.
\newblock Differentiable convex optimization layers.
\newblock In {\em Advances in Neural Information Processing Systems 32}, pages
9562--9574. Curran Associates, Inc., 2019.

[3]
N.~Ahmed, T.~Natarajan, and K.~R. Rao.
\newblock Discrete cosine transform.
\newblock {\em IEEE transactions on Computers}, 100(1):90--93, 1974.

[4]
B.~Amos and J.~Z. Kolter.
\newblock {O}pt{N}et: Differentiable optimization as a layer in neural
\newblock In D.~Precup and Y.~W. Teh, editors, {\em Proceedings of the 34th
International Conference on Machine Learning}, volume~70 of {\em Proceedings
of Machine Learning Research}, pages 136--145, International Convention
Centre, Sydney, Australia, 06--11 Aug 2017. PMLR.

[5]
K.~Atkinson and W.~Han.
\newblock {\em Theoretical numerical analysis}, volume~39.
\newblock Springer, 2005.

[6]
\newblock Post-convergence automatic differentiation of iterative schemes.
\newblock {\em Nuclear Science and Engineering}, 125(1):12--18, 01 1997.

[7]
S.~Bai, J.~Z. Kolter, and V.~Koltun.
\newblock Deep equilibrium models.
\newblock In H.~Wallach, H.~Larochelle, A.~Beygelzimer, F.~d\textquotesingle
Alch\'{e}-Buc, E.~Fox, and R.~Garnett, editors, {\em Advances in Neural
Information Processing Systems}, volume~32. Curran Associates, Inc., 2019.

[8]
S.~Bai, V.~Koltun, and J.~Z. Kolter.
\newblock Multiscale deep equilibrium models.
\newblock In H.~Larochelle, M.~Ranzato, R.~Hadsell, M.~F. Balcan, and H.~Lin,
editors, {\em Advances in Neural Information Processing Systems}, volume~33,
pages 5238--5250. Curran Associates, Inc., 2020.

[9]
M.~C. Bartholomew{--}Biggs.
\newblock Using forward accumulation for automatic differentiation of
implicitly-defined functions.
\newblock {\em Computational Optimization and Applications}, 9(1):65--84, 1998.

[10]
H.~H. Bauschke, J.~Bolte, and M.~Teboulle.
\newblock A descent lemma beyond lipschitz gradient continuity: First-order
methods revisited and applications.
\newblock {\em Mathematics of Operations Research}, 42(2):330--348, 05 2017.

[11]
H.~H. Bauschke and P.~L. Combettes.
\newblock {\em Convex Analysis and Monotone Operator Theory in Hilbert Spaces}.
\newblock Springer Publishing Company, Incorporated, 1st edition, 2011.

[12]
A.~G. Baydin, B.~A. Pearlmutter, A.~A. Radul, and J.~M. Siskind.
\newblock Automatic differentiation in machine learning: a survey.
\newblock {\em Journal of Marchine Learning Research}, 18:1--43, 2018.

[13]
A.~Beck and M.~Teboulle.
\newblock A fast iterative shrinkage-thresholding algorithm for linear inverse
\newblock {\em SIAM Journal on Imaging Sciences}, 2(1):183--202, March 2009.

[14]
\newblock Gradient-based optimization of hyperparameters.
\newblock {\em Neural computation}, 12:1889--900, 09 2000.

[15]
J.~Bergstra and Y.~Bengio.
\newblock Random search for hyper-parameter optimization.
\newblock {\em Journal of Machine Learning Research}, 13:281--305, 2012.

[16]
Q.~Berthet, M.~Blondel, O.~Teboul, M.~Cuturi, J.-P. Vert, and F.~Bach.
\newblock Learning with differentiable pertubed optimizers.
\newblock In H.~Larochelle, M.~Ranzato, R.~Hadsell, M.~Balcan, and H.~Lin,
editors, {\em Advances in Neural Information Processing Systems}, volume~33,
pages 9508--9519. Curran Associates, Inc., 2020.

[17]
\newblock {\em Hyperparameter selection for high dimensional sparse learning:
application to neuroimaging}.
\newblock PhD thesis, Universit{\'e} Paris-Saclay, 2021.

[18]
Q.~Bertrand, Q.~Klopfenstein, M.~Massias, M.~Blondel, S.~Vaiter, A.~Gramfort,
and J.~Salmon.
\newblock Implicit differentiation for fast hyperparameter selection in
non-smooth convex learning.
\newblock {\em arXiv preprint arXiv:2105.01637}, 2021.

[19]
L.~Bogensperger, A.~Chambolle, and T.~Pock.
\newblock {Convergence of a Piggyback-style method for the differentiation of
solutions of standard saddle-point problems}.
\newblock working paper or preprint, 1 2022.

[20]
J.~Bolte, T.~Le, E.~Pauwels, and T.~Silveti-Falls.
\newblock Nonsmooth implicit differentiation for machine-learning and
\newblock In M.~Ranzato, A.~Beygelzimer, Y.~Dauphin, P.~Liang, and J.~W.
Vaughan, editors, {\em Advances in Neural Information Processing Systems},
volume~34, pages 13537--13549. Curran Associates, Inc., 2021.

[21]
J.~Bolte and E.~Pauwels.
\newblock A mathematical model for automatic differentiation in machine
\newblock In H.~Larochelle, M.~Ranzato, R.~Hadsell, M.~Balcan, and H.~Lin,
editors, {\em Advances in Neural Information Processing Systems}, volume~33,
pages 10809--10819. Curran Associates, Inc., 2020.

[22]
J.~Bolte and E.~Pauwels.
\newblock Conservative set valued fields, automatic differentiation, stochastic
gradient methods and deep learning.
\newblock {\em Mathematical Programming}, 188(1):19--51, 2021.

[23]
J.~Bolte, E.~Pauwels, and S.~Vaiter.
\newblock Automatic differentiation of nonsmooth iterative algorithms.
\newblock {\em arXiv preprint arXiv:2206.00457}, 2022.

[24]
L.~Bottou, F.~E. Curtis, and J.~Nocedal.
\newblock Optimization methods for large-scale machine learning.
\newblock {\em Siam Review}, 60(2):223--311, 2018.

[25]
J.~Bradbury, R.~Frostig, P.~Hawkins, M.~J. Johnson, C.~Leary, D.~Maclaurin,
G.~Necula, A.~Paszke, J.~Vander{P}las, S.~Wanderman-{M}ilne, and Q.~Zhang.
\newblock {JAX}: composable transformations of {P}ython+{N}um{P}y programs,

[26]
L.~M. Bregman.
\newblock The relaxation method of finding the common point of convex sets and
its application to the solution of problems in convex programming.
\newblock {\em USSR computational mathematics and mathematical physics},
7(3):200--217, 1967.

[27]
E.~Brochu, V.~M. Cora, and N.~De~Freitas.
\newblock A tutorial on bayesian optimization of expensive cost functions, with
application to active user modeling and hierarchical reinforcement learning.
\newblock {\em arXiv preprint arXiv:1012.2599}, 2010.

[28]
A.~Chambolle and C.~Dossal.
\newblock On the convergence of the iterates of the “fast iterative
shrinkage/thresholding algorithm”.
\newblock {\em Journal of optimization theory and applications},
166(3):968--982, 2015.

[29]
A.~Chambolle and T.~Pock.
\newblock An introduction to continuous optimization for imaging.
\newblock {\em Acta Numerica}, 25:161--319, 2016.

[30]
A.~Chambolle and T.~Pock.
\newblock Learning consistent discretizations of the total variation.
\newblock {\em SIAM Journal on Imaging Sciences}, 14(2):778--813, 2021.

[31]
\newblock {\em Riemannian Geometry: A Modern Introduction}, volume~98.
\newblock Cambridge University Press, 2006.

[32]
R.~T.~Q. Chen, Y.~Rubanova, J.~Bettencourt, and D.~K. Duvenaud.
\newblock Neural ordinary differential equations.
\newblock In S.~Bengio, H.~Wallach, H.~Larochelle, K.~Grauman, N.~Cesa-Bianchi,
and R.~Garnett, editors, {\em Advances in Neural Information Processing
Systems 31}, pages 6571--6583. Curran Associates, Inc., 2018.

[33]
T.~Chen, X.~Chen, W.~Chen, H.~Heaton, J.~Liu, Z.~Wang, and W.~Yin.
\newblock Learning to optimize: A primer and a benchmark.
\newblock {\em arXiv preprint arXiv:2103.12828}, 2021.

[34]
Y.~Chen, T.~Pock, R.~Ranftl, and H.~Bischof.
\newblock Revisiting loss-specific training of filter-based mrfs for image
\newblock In {\em German Conference on Pattern Recognition}, pages 271--281.
Springer, 2013.

[35]
\newblock Reverse accumulation and attractive fixed points.
\newblock {\em Optimization Methods and Software}, 3(4):311--326, 1994.

[36]
\newblock Gradient-based solution algorithms for a class of bilevel
optimization and optimal control problems with a nonsmooth lower level.
\newblock {\em SIAM Journal on Optimization}, 30(1):290--318, 2020.

[37]
\newblock {\em Optimization and Nonsmooth Analysis}.
\newblock Wiley New York, 1983.

[38]
B.~Dauvergne and L.~Hasco{\"e}t.
\newblock The data-flow equations of checkpointing in reverse automatic
\newblock In {\em International Conference on Computational Science}, pages
566--573. Springer, 2006.

[39]
C.-A. Deledalle, S.~Vaiter, J.~Fadili, and G.~Peyr{\'e}.
\newblock Stein unbiased gradient estimator of the risk (sugar) for multiple
parameter selection.
\newblock {\em SIAM Journal on Imaging Sciences}, 7(4):2448--2487, 2014.

[40]
S.~Dempe, V.~Kalashnikov, G.~A. Prez-Valds, and N.~Kalashnykova.
\newblock {\em Bilevel Programming Problems: Theory, Algorithms and
Applications to Energy Networks}.
\newblock Springer Publishing Company, Incorporated, 2015.

[41]
\newblock Generic methods for optimization based modeling.
\newblock In {\em Proceedings of the Fifteenth International Conference on
Artificial Intelligence and Statistics}, volume~22 of {\em Proceedings of
Machine Learning Research}, pages 318--326, La Palma, Canary Islands, 21--23
Apr 2012. PMLR.

[42]
A.~L. Dontchev and R.~T. Rockafellar.
\newblock {\em Implicit Functions and Solution Mappings}, volume 543.
\newblock Springer, 2009.

[43]
J.~Fadili, J.~Malick, and G.~Peyr{\'e}.
\newblock Sensitivity analysis for mirror-stratifiable convex functions.
\newblock {\em SIAM Journal on Optimization}, 28(4):2975--3000, 2018.

[44]
\newblock {\em Matrix rank minimization with applications}.
\newblock PhD thesis, PhD thesis, Stanford University, 2002.

[45]
L.~Franceschi, M.~Donini, P.~Frasconi, and M.~Pontil.
\newblock Forward and reverse gradient-based hyperparameter optimization.
\newblock In {\em International Conference on Machine Learning}, pages
1165--1173. PMLR, 2017.

[46]
J.~C. Gilbert.
\newblock Automatic differentiation and iterative processes.
\newblock {\em Optimization Methods and Software}, 1(1):13--21, 1992.

[47]
I.~Goodfellow, Y.~Bengio, and A.~Courville.
\newblock {\em Deep learning}.
\newblock MIT press, 2016.

[48]
A.~Griewank, C.~Bischof, G.~Corliss, A.~Carle, and K.~Williamson.
\newblock Derivative convergence for iterative equation solvers.
\newblock {\em Optimization methods and software}, 2(3-4):321--355, 1993.

[49]
A.~Griewank and C.~Faure.
\newblock Piggyback differentiation and optimization.
\newblock In L.~T. Biegler, M.~Heinkenschloss, O.~Ghattas, and B.~van
Bloemen~Waanders, editors, {\em Large-Scale PDE-Constrained Optimization},
pages 148--164, Berlin, Heidelberg, 2003. Springer Berlin Heidelberg.

[50]
A.~Griewank and A.~Walther.
\newblock {\em Evaluating Derivatives: Principles and Techniques of Algorithmic
\newblock SIAM, Philadelphia, 2nd edition, 2008.

[51]