Chapter 11Differentiating through optimization

In this chapter, we study how to differentiate through optimization problems, and more generally through nonlinear systems of equations.

11.1 Implicit functions

Implicit functions are functions that do not enjoy an explicit decomposition into elementary functions, for which automatic differentiation, as studied in Chapter 8, can therefore not be directly applied. We describe in this chapter techniques to differentiate through such functions and how to integrate them into an autodiff framework.

Formally, we will denote an implicit function by $\wv^\star(\lambdav)$ , where $\wv^\star \colon \Lambda \to \cW$ . One question is then how to compute the Jacobian $\partial \wv^\star(\lambdav)$ . As a first application one can consider sensitivity analysis of a system. For example, $\wv^\star(\lambdav)$ could correspond to the equilibrium state of a physical system and in this case, $\partial \wv^\star(\lambdav)$ would tell us about the sensitivity of the system to some parameters $\lambdav \in \Lambda$ .

11.1.1 Optimization problems

Another example is a function implicitly defined as the solution (assumed unique) of an optimization problem $\wv^\star(\lambdav) \coloneqq \argmax_{\wv \in \cW} f(\wv, \lambdav),$ where $f \colon \cW \times \Lambda \to \RR$ and $\cW$ denotes a constraint set. Note that we use an $\argmax$ for convenience, but the same applies when using an $\argmin$ .

11.1.2 Nonlinear equations

More generally, $\wv^\star(\lambdav)$ can be defined as the root of some function $F \colon \cW \times \Lambda \to \cW$ , i.e., $\wv^\star(\lambdav)$ is implicitly defined as the value satisfying the (potentially nonlinear) system of equations $F(\w, \lambdav) = \zeros$ for all $\lambdav \in \Lambda$ .

11.1.3 Application to bilevel optimization

Besides sensitivity analysis, another example of application is bilevel optimization. Many times, we want to minimize a function defined as the composition of a fixed function and the solution of an optimization problem. Formally, let $f, g \colon \cW \times \Lambda \to \RR$ . We consider the composition $h(\lambdav)$ defined as $h(\lambdav) \coloneqq g(\wv^\star(\lambdav), \lambdav), \quad \text{where} \quad \wv^\star(\lambdav) \coloneqq \argmax_{\wv \in \cW} f(\wv, \lambdav). \tag{11.1}$ This includes for instance hyperparameter optimization, where $f$ is an inner log-likelihood objective, $g$ is an outer validation loss, $\wv \in \cW$ are model parameters and $\lambdav \in \Lambda$ are model hyperparameters, such as regularization strength, as illustrated in Fig. 11.1. To minimize $h(\lambdav)$ , one generally resorts to a gradient descent scheme w.r.t. $\lambdav$ , which requires computing $\nabla h(\lambdav)$ . Assuming that $\wv^\star(\lambdav)$ is differentiable at $\lambdav$ , by the chain rule, we obtain the Jacobian $\partial h(\lambdav) = \partial_1 g(\wv^\star(\lambdav), \lambdav) \partial \wv^\star(\lambdav) + \partial_2 g(\wv^\star(\lambdav), \lambdav).$ Using $\partial h(\lambdav)^\top = \nabla h(\lambdav)$ (see Remark 2.4), we obtain the gradient $\nabla h(\lambdav) = \partial \wv^\star(\lambdav)^\top \nabla_1 g(\wv^\star(\lambdav), \lambdav) + \nabla_2 g(\wv^\star(\lambdav), \lambdav).$ The only problematic term is $\partial \wv^\star(\lambdav)$ , as it requires argmax differentiation. Indeed, most of the time, there is no explicit formula for $\wv^\star(\lambdav)$ and it does not decompose into elementary functions.

Figure 11.1: Hyperparameter optimization in nonlinear regression can be cast as a bi-level optimization problem. Each line corresponds to the estimator obtained by fitting some training data (in blue circles) using a different hyperparameter $\lambda$ . Formally, denoting $f$ the training objective, the estimators are $w^\star(\lambda) \coloneqq \argmin_w f(w; \lambda)$ . The goal is to find the best hyperparameter that fits some validation data (here in cyan diamonds), that is, minimizing $h(\lambda) \coloneqq g(w^\star(\lambda), \lambda)$ , where $g$ is the validation objective. A too small $\lambda_1$ leads to overfitting the training objective and performs badly on validation objective. Conversely, a larger $\lambda_2$ underfits both training and validation objectives. The optimal parameter $\lambda^\star$ minimizes the validation objective and may be obtained by iterating gradient descent w.r.t. $\lambda$ . This requires gradients of $h(\lambda) = g(w^\star(\lambda), \lambda)$ w.r.t. $\lambda$ .

11.2 Envelope theorems

In the special case $g = f$ , the composition $h$ defined in Eq. (11.1) is simply given by $h(\lambdav) = f(\wv^\star(\lambdav), \lambdav) = \max_{\wv \in \cW} f(\wv, \lambdav).$ That is, we no longer need argmax differentiation, but only max differentiation, which, as we shall now see is much easier. The function $h$ is often called a value function (Fleming and Rishel, 2012). The reason for the name “envelope” is illustrated in Fig. 11.2. We emphasize that there is not one, but several envelope theorems, depending on the assumptions on $f$ .

Figure 11.2: The graph of $h(\lambdav)= \max_{\wv \in \cW} f(\wv, \lambdav)$ is the upper-envelope of the graphs of the functions $\lambdav \mapsto f(\wv, \lambdav)$ for all $\wv \in \cW$ .

11.2.1 Bertsekas’ theorem

The following theorem is due to Bertsekas (1997, Proposition B.25).

Theorem 11.1 (Bertsekas’ theorem)

Let $h(\lambdav) \coloneqq \max_{\wv \in \cW} f(\wv, \lambdav)$ and $\wv^\star(\lambdav) \coloneqq \argmax_{\wv \in \cW} f(\wv, \lambdav),$ where $f \colon \cW \times \Lambda \to \R$ is concave in $\w$ and $\cW$ is a compact convex set. If the maximum $\wv^\star(\lambdav)$ is unique, then the function $h$ is differentiable with gradient $\nabla h(\lambdav) = \nabla_2 f(\wv^\star(\lambdav), \lambdav).$ If furthermore $f$ is convex in $\lambdav$ , then $h(\lambdav)$ is convex.

A variant of the theorem can be stated when $\cW = \RR^P$ .

Theorem 11.2 (Unbounded case)

Suppose $\cW = \RR^P$ . If $f$ is strongly concave in $\wv$ , then the function $h$ is differentiable with gradient $\nabla h(\lambdav) = \nabla_2 f(\wv^\star(\lambdav), \lambdav).$ If furthermore $f$ is convex in $\lambdav$ , then $h(\lambdav)$ is convex.

Informally, Theorem 11.1 and Theorem 11.2 mean that we can treat $\wv^\star(\lambdav)$ as if it were a constant of $\lambdav$ . In other words, we do not need to differentiate through $\wv^\star(\lambdav)$ , even though it depends on $\lambdav$ .

Remark 11.1 (Historical note)

Bertsekas’ theorem or its unbounded variant are often incorrectly cited as Danskin’s theorem in the machine learning literature, because Bertsekas himself called it that way (Bertsekas, 1997, Proposition B.25). However, because the assumptions differ, they should rather be thought as variants or corollaries of Danskin’s original theorem. See Section 11.2.3 for more details.

We now illustrate Theorem 11.2, to differentiate through a minimum, $h(\lambdav) = \min_{\wv \in \cW} f(\wv, \lambdav)$ , assuming $f(\wv, \lambdav)$ is strongly convex in $\wv$ .

Example 11.1 (Differentiating the minimum of a convex quadratic)

Let us define $h(\lambda) \coloneqq \min_{w \in \R} f(w, \lambda)$ , where $f(w, \lambda) \coloneqq \frac{\lambda}{2} w^2 + b w + c$ and $\lambda > 0$ . The derivative of $f$ w.r.t. $w$ is $\nabla_1 f(w, \lambda) = \lambda w + b$ . The second derivative of $f$ w.r.t. $w$ is $\nabla_1^2 f(w, \lambda) = \lambda$ . Therefore $f$ is strongly convex in $w$ . Let $w^\star(\lambda)$ be the minimum of $f$ in $w$ . Setting $\nabla_1 f(w, \lambda)$ to zero, we get $w^\star(\lambda) = -\frac{b}{\lambda}$ . The derivative of $f$ w.r.t. $\lambda$ is $\nabla_2 f(w, \lambda) = \frac{1}{2} w^2$ . From Theorem 11.2, we have $h'(\lambda) = \nabla_2 f(w^\star(\lambda), \lambda) = \frac{1}{2} (w^\star(\lambda))^2 = \frac{1}{2} \frac{b^2}{\lambda^2}$ . Let us check that the derivative is indeed correct. Plugging $w^\star(\lambda)$ back into $f(w, \lambda)$ , we get the analytical formula $h(\lambda) = -\frac{1}{2} \frac{b^2}{\lambda} + c$ . Using $(\frac{1}{\lambda})' = -\frac{1}{\lambda^2}$ , we indeed obtain the same expression for $h'(\lambda)$ .

Theorem 11.1 has a simple interpretation for functions that are affine in $\lambdav$ as shown below.

Example 11.2 (Convex conjugate)

Let $f(\wv, \lambdav) \coloneqq \langle \wv, \lambdav \rangle - \Omega(\wv)$ with $\Omega$ convex. We then have $h(\lambdav) = \max_{\wv \in \cW} ~ \langle \wv, \lambdav \rangle - \Omega(\wv) \eqqcolon \Omega^*(\lambdav)$ , where $\Omega^*$ denotes the convex conjugate of $\Omega$ . Since $f$ is convex and assuming $\cW$ is convex compact, the conditions of Theorem 11.1 are met. Since $\nabla_2 f(\wv, \lambdav) = \wv$ , we obtain $\nabla h(\lambdav) = \nabla \Omega^*(\lambdav) = \wv^\star(\lambdav)$ . In other words, in this special case, the gradient of the max is equal to the argmax. This is due to the fact that $f(\wv, \lambdav)$ is affine in $\lambdav$ .

Another application is saddle point optimization.

Example 11.3 (Saddle point problem)

Consider the saddle point problem $\min_{\lambdav \in \Lambda} \max_{\wv \in \cW} f(\wv, \lambdav)$ . If it is difficult to minimize w.r.t. $\lambdav$ but easy to maximize w.r.t. $\wv$ , we can rewrite the problem as $\min_{\lambdav \in \Lambda} h(\lambdav)$ , where $h(\lambdav) \coloneqq \max_{\wv \in \cW} f(\wv, \lambdav)$ , and use $\nabla h(\lambdav)$ to perform (projected) gradient descent w.r.t. $\lambdav$ .

11.2.2 Rockafellar’s theorem

A related theorem (Rockafellar and Wets, 2009, Theorem 10.31) can be proved under different assumptions on $f$ .

Theorem 11.3 (Rockafellar’s envelope theorem)

Let $f \colon \cW \times \Lambda \to \R$ and $\cW$ be a compact convex set. Let $h(\lambdav) \coloneqq \max_{\wv \in \cW} f(\wv, \lambdav)$ and $\wv^\star(\lambdav) \coloneqq \argmax_{\wv \in \cW} f(\wv, \lambdav).$ If $f$ is continuously differentiable in $\lambdav$ for all $\wv \in \cW$ , $\nabla_1 f$ is continuous and the maximum $\wv^\star(\lambdav)$ is unique, then the function $h$ is differentiable with gradient $\nabla h(\lambdav) = \nabla_2 f(\wv^\star(\lambdav), \lambdav).$

Compared to Bertsekas’ theorem, Rockafellar’s theorem does not require $f$ to be concave in $\w$ , but it requires stronger assumptions on the differentiability of $f$ .

11.2.3 Danskin’s theorem

Finally, we state Danskin’s original theorem (Danskin, 1966).

Theorem 11.4 (Danskin’s theorem)

Let $\cW$ be a compact topological space. Let $f \colon \cW \times \Lambda \to \R$ be continuous and continuously differentiable with respect to $\lambdav \in \Lambda$ . Let the value function and the set of maximizers be defined as $h(\lambdav) \coloneqq \max_{\wv \in \cW} f(\wv, \lambdav) \quad \text{and} \quad \cW^\star(\lambdav) \coloneqq \argmax_{\wv \in \cW} f(\wv, \lambdav).$ Then, the directional derivative of $h$ at $\lambdav \in \Lambda$ in direction $\dv \in \Lambda$ , denoted $h'(\lambdav; \dv)$ , exists and is given by $h'(\lambdav; \dv) = \max_{\wv \in \cW^\star(\lambdav)} \langle \nabla_2 f(\wv, \lambdav), \dv \rangle.$

We sum up envelope theorems and their assumptions in Table 11.1.

Table 11.1: Comparison of different envelope theorems.

Theorem	Unique $\wv^\star$ ?	Concave in $\w$ ?	Smooth in $\wv$ ?	Result
Danskin	No	No	No	Dir. Deriv.
Bertsekas	Yes	Yes	No	Gradient
Rockafellar	Yes	No	Yes	Gradient

11.3 Implicit function theorem

11.3.1 Univariate functions

The implicit function theorem (IFT) provides conditions under which an implicit relationship of the form $F(w, \lambda) = 0$ can be rewritten as a function $w = w^\star(\lambda)$ locally, and provides a way to compute its derivative w.r.t. $\lambda$ .

Theorem 11.5 (Implicit function theorem, univariate case)

Let
$F \colon \R \times \R \to \R$ . Assume $F(w, \lambda)$ is a continuously differentiable function in a neighborhood $\cU$ of $(w_0, \lambda_0)$ such that $F(w_0, \lambda_0) = 0$ and $\partial_1 F(w_0, \lambda_0) \neq 0$ . Then there exists a neighborhood $\cV \subseteq \cU$ of $(w_0, \lambda_0)$ in which there is a function $w^\star(\lambda)$ such that

$w^\star(\lambda_0) = w_0$ ,
$F(w^\star(\lambda), \lambda) = 0$ for all $\lambda$ in the neighborhood $\cV$ ,
$\partial w^\star(\lambda) = -\frac{\partial_2 F(w^\star(\lambda), \lambda)}{ \partial_1 F(w^\star(\lambda), \lambda)}$ .

We postpone the proof to the multivariate case and begin with a classical example of application of the theorem.

Figure 11.3: The circle equation $F(x, y) \coloneqq x^2 + y^2 - 1 = 0$ is not a function from $y \in [-1, 1]$ to $x \in [-1, 1]$ , as there are always two possible $x$ values, $x = \sqrt{1 - y^2}$ or $x = -\sqrt{1 - y^2}$ . However, locally around some point $(x_0, y_0)$ , e.g., such that $x_0 > 0$ and $y_0 > 0$ (upper-right quadrant), the function $x = x^\star(y) = \sqrt{1 - y^2}$ is well-defined. The implicit function theorem gives conditions for such a function to exist locally and provides a way to compute its derivative.

Example 11.4 (Equation of the unit circle)

We use $w \equiv x$ and $\lambda \equiv y$ for clarity. Let $F(x, y) \coloneqq x^2 + y^2 - 1$ . In general, we cannot rewrite the unit circle equation $F(x, y) = 0$ as a function from $y$ to $x$ , because for every $y \in [-1, 1]$ , there are always two possible $x$ values, namely, $x = \sqrt{1 - y^2}$ or $x = -\sqrt{1 - y^2}$ . However, locally around some point $(x_0, y_0)$ , e.g., such that $x_0 > 0$ and $y_0 > 0$ (upper-right quadrant), the function $x = x^\star(y) = \sqrt{1 - y^2}$ is well-defined. Using $\partial_1 F(x, y) = 2x$ and $\partial_2 F(x, y) = 2y$ , we get $\partial x^\star(y) = -\frac{\partial_2 F(x^\star(y), y)}{\partial_1 F(x^\star(y), y)} = - \frac{2y}{2 x^\star(y)} = -\frac{y}{\sqrt{1 - y^2}}$ in that neighborhood (the upper right quadrant in this case). This is indeed the same derivative expression as if we used the chain rule on $\sqrt{1 - y^2}$ and is well-defined on $y \in [0, 1)$ .

In the above simple example, we can easily derive an explicit function relating $y$ to $x$ in a given neighborhood, but this is not always the case. The IFT gives us conditions guaranteeing that such function exists and a way to differentiate it, but not a way to construct such a function. In fact, finding $w^\star(\lambda)$ such that $F(w^\star(\lambda), \lambda) = 0$ typically involves a root finding algorithm, an optimization algorithm, a nonlinear system solver, etc.

Example 11.5 (Polynomial)

Let $F(w, \lambda) = w^5 + w^3 + w - \lambda$ . According to the Abel-Ruffini theorem (Tignol, 2015), quintics (polynomials of degree $5$ ) do not enjoy roots in terms of radicals and one must resort to numerical root finding. In addition, odd-degree polynomials have real roots. Moreover, $\partial_1 F(w, \lambda) = 5w^4 + 3w^2 + 1$ is strictly positive. Therefore, by the intermediate value theorem, there must be only one root $w^\star(\lambda)$ such that $F(w^\star(\lambda), \lambda) = 0$ . This unique root can for example be found by bisection. Using the IFT, its derivative is found to be $\partial w^\star(\lambda) = (5w^\star(\lambda)^4 + 3 w^\star(\lambda)^2 + 1)^{-1}$ .

While an implicit function is differentiable at a point if the assumptions of the IFT hold in a neighborhood of that point, the reciprocal is not true: failure of the IFT assumptions does not necessarily mean that the implicit function is not differentiable, as we now illustrate.

Example 11.6 (IFT conditions are not necessary for differentiability)

Consider $F(w, \lambda) = (w - \lambda)^2$ . We clearly have that $F(w^\star(\lambda), \lambda) = 0$ if we define $w^\star(\lambda) = \lambda$ , the identity function. It is clearly differentiable for all $\lambda$ , yet the assumptions of the IFT fail, since we have $\partial_1 F(w, \lambda) = 2(w-\lambda)$ and therefore $\partial_1 F(0, 0) = 0$ .

11.3.2 Multivariate functions

We now present the IFT in the general multivariate setting. Informally, if $F(\w^\star(\lambdav), \lambdav) = \zeros$ , then by the chain rule, we have $\partial_1 F(\w^\star(\lambdav), \lambdav) \partial \w^\star(\lambdav) + \partial_2 F(\w^\star(\lambdav), \lambdav) = \zeros,$ meaning that the Jacobian $\partial \w^\star(\lambdav)$ , assuming that it exists, satisfies $-\partial_1 F(\w^\star(\lambdav), \lambdav) \partial \w^\star(\lambdav) = \partial_2 F(\w^\star(\lambdav), \lambdav).$ The IFT gives us conditions for the existence of $\partial \w^\star(\lambdav)$ .

Theorem 11.6 (Implicit function theorem, multivariate case)

Let us define $F \colon \cW \times \Lambda \to \cW$ . Assume $F(\wv, \lambdav)$ is a continuously differentiable function in a neighborhood of $(\wv_0, \lambdav_0)$ such that $F(\wv_0, \lambdav_0) = \zeros$ and $\partial_1 F(\wv_0, \lambdav_0)$ is invertible, i.e., its determinant is nonzero. Then there exists a neighborhood of $\lambdav_0$ in which there is a function $\wv^\star(\lambdav)$ such that

$\wv^\star(\lambdav_0) = \wv_0$ ,
$F(\wv^\star(\lambdav), \lambdav) = \zeros$ for all $\lambdav$ in the neighborhood,
$-\partial_1 F(\wv^\star(\lambdav), \lambdav) \partial \wv^\star(\lambdav) = \partial_2 F(\wv^\star(\lambdav),\lambdav)$

$\iff \partial \wv^\star(\lambdav) = -\partial_1 F(\wv^\star(\lambdav), \lambdav)^{-1} \partial_2 F(\wv^\star(\lambdav), \lambdav).$

We begin with a simple unconstrained optimization algorithm.

Example 11.7 (Unconstrained optimization)

Assume we want to differentiate through $\wv^\star(\lambdav) = \argmin_{\wv \in \R^P} f(\wv, \lambdav)$ , where $f$ is strictly convex in $\wv$ , which ensures that the solution is unique. From the stationary conditions, if we define $F(\wv, \lambdav) \coloneqq \nabla_1 f(\wv, \lambdav)$ , then $\wv^\star(\lambdav)$ is uniquely characterized as the root of $F$ in the first argument, i.e., $F(\wv^\star(\lambdav), \lambdav) = \zeros$ . We have $\partial_1 F(\wv, \lambdav) = \nabla^2_1 f(\wv, \lambdav)$ , the Hessian of $f$ in $\wv$ , and $\partial_2 F(\wv, \lambdav) = \partial_2 \nabla_1 f(\wv, \lambdav)$ , the cross derivatives of $f$ in $\wv$ and $\lambdav$ . Therefore, assuming that the Hessian is well-defined and invertible at $(\wv^\star(\lambdav), \lambdav)$ , we can use the IFT to differentiate through $\wv^\star(\lambdav)$ and obtain $\partial \wv^\star(\lambdav) = -(\nabla^2_1 f(\wv^\star(\lambdav), \lambdav))^{-1} \partial_2 \nabla_1 f(\wv^\star(\lambdav), \lambdav)$ .

Next, we generalize the previous example, by allowing constraints in the optimization problem.

Example 11.8 (Constrained optimization)

Now, assume we want to differentiate through $\wv^\star(\lambdav) = \argmin_{\wv \in \cC} f(\wv, \lambdav)$ , where $f$ is strictly convex in $\wv$ and $\cC \subseteq \cW$ is a convex set. A solution is characterized by the fixed point equation $\wv^\star(\lambdav) = P_\cC(\wv^\star(\lambdav) - \eta \nabla_1 f(\wv^\star(\lambdav), \lambdav))$ , for any $\eta > 0$ , where $P(\yv) \coloneqq \argmin_{\xv \in \cC} \|\xv - \yv\|_2^2$ is the Euclidean projection of $\yv$ onto $\cC$ . Therefore, $\wv^\star(\lambdav)$ is the root of $F(\wv, \lambdav) = \wv - P_\cC(\wv - \eta \nabla_1 f(\wv, \lambdav))$ (see Chapter 16). We can differentiate through $\wv^\star(\lambdav)$ using the IFT, assuming that the conditions of the theorem apply. Note that $\partial_1 F(\wv, \lambdav)$ requires the expression of the Jacobian $\partial P_\cC(\yv)$ . Fortunately, $P_\cC(\yv)$ and its Jacobian are easy to compute for many sets $\cC$ (Blondel et al., 2021).

11.3.3 JVP and VJP of implicit functions

To integrate an implicit function $\w^\star(\lambdav)$ in an autodiff framework, we need to be able to compute its JVP or VJP. This is the purpose of the next proposition.

Proposition 11.1 (JVP and VJP of implicit functions)

Let $\wv^\star \colon \Lambda \to \cW$ be a function implicitly defined as the solution of $F(\wv^\star(\lambdav), \lambdav) = \zeros$ , for some function $F \colon \cW \times \Lambda \to \cW$ . Define $\begin{align*} A & \coloneqq -\partial_1 F(\w^\star(\lambdav), \lambdav) \\ B & \coloneqq \partial_2 F(\w^\star(\lambdav),\lambdav). \end{align*}$ Assume the assumptions of the IFT hold. The JVP $\t \coloneqq \partial \wv^\star(\lambdav) \v$ in the input direction $\v \in \Lambda$ is obtained by solving the linear system $A \t = B \v.$ The VJP $\partial \wv^\star(\lambdav)^* \u$ in the output direction $\u \in \cW$ is obtained by solving the linear system $A^* \r = \u.$ Using the solution $\r$ , we get $\partial \wv^\star(\lambdav)^* \u = \partial \wv^\star(\lambdav)^* A^* \r = B^*\r.$

Note that in the above linear systems, we can access $A$ and $B$ as linear maps, the JVPs of $F$ . Their adjoints, $A^*$ and $B^*$ , correspond to the VJPs of $F$ . To solve these systems, we can therefore use matrix-free solvers as detailed in Section 9.4. For example, when $A$ is symmetric positive semi-definite, we can use the conjugate gradient method (Hestenes et al., 1952). When $A$ is not symmetric positive definite, we can use GMRES (Saad and Schultz, 1986) or BiCGSTAB (Vorst and Vorst, 1992).

11.3.4 Proof of the implicit function theorem

We prove the theorem using the inverse function theorem presented in Theorem 11.7. Define $f(\lambdav, \w) = (\lambdav, F(\w, \lambdav))$ which goes from $\RR^Q \times \RR^P$ onto $\RR^Q \times \RR^P$ . The Jacobian of $f$ is $\jac f(\lambdav, \w) = \begin{pmatrix} \idm & 0 \\ \jac_2 F(\w, \lambdav) & \jac_1 F(\w, \lambdav) \end{pmatrix}.$ So at $\w_0, \lambdav_0$ , we have $\det(\jac f(\lambdav_0, \w_0)) = \det(\idm)\det( \jac_1 F(\w_0, \lambdav_0)) \neq 0$ since we assumed $\partial_1 F(\wv_0, \lambdav_0)$ invertible. By the inverse function theorem, the function $f$ is then invertible in a neighborhood $N$ of $f(\lambdav_0, \w_0) = (\lambdav_0, \zeros)$ . In particular, it is invertible in $N \cap \{(\lambdav, \zeros), \lambdav \in \RR^Q\}$ . The solution of the implicit equation in a neighborhood of $\lambdav_0$ is then $( \lambdav, \w^*(\lambdav)) = f^{-1}(\lambdav, \zeros)$ . By the inverse function theorem, $f^{-1}$ is continuously differentiable inverse and so is $\w^*(\lambdav)$ . The derivative $\partial \w^*(\lambdav)$ is obtained from the differential of the inverse as $\begin{pmatrix} \sim & \sim \\ \partial \w^*(\lambdav) & \sim \end{pmatrix} = \partial f^{-1}(\lambdav, \zeros),$ and by the inverse function theorem, we have $\partial f^{-1}(\lambdav, \zeros) = (\partial f(\lambdav, \w^*(\lambdav)))^{-1}.$ Using the block matrix inversion formula, we therefore obtain $\begin{pmatrix} \A & \B \\ \C & \D \end{pmatrix}^{-1} = \begin{pmatrix} \sim & \sim \\ -(\D - \C\A^{-1}\B)^{-1}\C\A^{-1} & \sim \end{pmatrix},$ we get the claimed expression. Though we expressed the proof in terms of Jacobians and matrices, the result naturally holds for the corresponding linear operators, JVPs, VJPs, and their inverses.

11.4 Adjoint state method

11.4.1 Differentiating nonlinear equations

We describe in this section the adjoint state method (a.k.a. adjoint method, method of adjoints, adjoint sensitivity method). The method can be used to compute the gradient of the composition of an explicit function and an implicit function, defined through an equality constraint (e.g., a nonlinear equation). The method dates back to Céa (1986).

Suppose a variable $\s \in \cS$ (which corresponds to a state in optimal control) is implicitly defined given some parameters $\w \in \cW$ through the (potentially nonlinear) equation $c(\s, \w) = \zeros$ , where $c \colon \cS \times \cW \to \cS$ . Assuming $\s$ is uniquely determined for all $\w \in \cW$ , this defines an implicit function $\s^\star(\w)$ from $\cW$ to $\cS$ such that $c(\s^\star(\w), \w) = \zeros$ . Given an objective function $L \colon \cS \times \cW \to \RR$ , the goal of the adjoint state method is then to compute the gradient of $L(\w) \coloneqq L(\s^\star(\w), \w).$ However, this is not trivial as $\s^\star(\w)$ is an implicit function. For instance, this can be used to convert the equality-constrained problem $\begin{align*} \min_{\w \in \cW} L(\s, \w) \quad \text{s.t.} \quad c(\s, \w) = \zeros. \end{align*}$ into the unconstrained problem $\begin{align*} \min_{\w \in \cW} L(\s^\star(\w), \w). \end{align*}$ Access to $\nabla L(\w)$ allows us to solve this problem by gradient descent.

Proposition 11.2 (Adjoint state method)

Let $c:\cS \times \cW \rightarrow \cS$ be a mapping defining constraints of the form $c(\s, \w)$ . Assume that for each $\w \in \cW$ , there exists a unique $\s^\star(\w)$ satisfying $c(\s^\star(\w), \w)= \zeros$ and that $\s^\star(\w)$ is differentiable. The gradient of $L(\w) \coloneqq L(\s^\star(\w), \w),$ for some differentiable function $L:\cS \times \cW \rightarrow \RR$ , is given by $\nabla L(\w) = \nabla_2 L(\s^\star(\w), \w) + \partial_2 c(\s^\star(\w), \w)^* \r^\star(\w),$ where $\r^\star(\w)$ is the solution of the linear system $\partial_1 c(\s^\star(\w), \w)^* \r = -\nabla_1 L(\s^\star(\w), \w).$

As shown in the proof below, $\r^\star(\w)$ corresponds to a Lagrange multiplier. The linear system can be solved using matrix-free solvers.

11.4.2 Relation with envelope theorems

Because $\s$ is uniquely determined for any $\w \in \cW$ by $c(\s, \w) = \zeros$ , we can alternatively rewrite $L(\w)$ as the trivial minimization or maximization, $\begin{align*} L(\w) &= \min_{\s \in \cS} L(\s, \w) \quad \text{s.t.} \quad c(\s, \w) = \zeros\\ &= \max_{\s \in \cS} L(\s, \w) \quad \text{s.t.} \quad c(\s, \w) = \zeros. \end{align*}$ Therefore, the adjoint state method can be seen as an envelope theorem for computing $\nabla L(\w)$ , for the case when $\w$ is involved in both the objective function and in the equality constraint.

11.4.3 Proof using the method of Lagrange multipliers

Classically, the adjoint state method is derived using the method of Lagrange multipliers. Let us introduce the Lagrangian associated with $L$ and $c$ , $\cL(\s, \w, \r) \coloneqq L(\s, \w) + \langle \r, c(\s, \w) \rangle,$ where $\r \in \cS$ is the Lagrange multiplier associated with the equality constraint $c(\s, \w) = \zeros$ . In the optimal control literature, $\r$ is often called the adjoint variable or adjoint state. The gradients of the Lagrangian are $\begin{align*} \nabla_\s \cL(\s, \w, \r) &= \nabla_1 L(\s, \w) + \partial_1 c(\s, \w)^* \r \\ \nabla_\w \cL(\s, \w, \r) &= \nabla_2 L(\s, \w) + \partial_2 c(\s, \w)^* \r \\ \nabla_\r \cL(\s, \w, \r) &= c(\s, \w), \end{align*}$ where $\partial_i c(\s, \w)^*$ are the adjoint operators. Setting $\nabla_\r \cL(\s, \w, \r)$ to zero gives the constraint $c(\s, \w) = \zeros$ . Setting $\nabla_\s \cL(\s, \w, \r)$ to zero gives the so-called adjoint state equation $\partial_1 c(\s, \w)^* \r = -\nabla_1 L(\s, \w).$ Solving this linear system w.r.t. $\r$ at $\s = \s^\star(\w)$ gives the adjoint variable $\r^\star(\w)$ . We then get $\begin{align*} \nabla L(\w) &= \nabla_2 \cL(\s^\star(\w), \w, \r^\star(\w)) \\ &= \nabla_2 L(\s^\star(\w), \w) + \partial_2 c(\s^\star(\w), \w)^* \r^\star(\w), \end{align*}$ which concludes the proof.

11.4.4 Proof using the implicit function theorem

A more direct proof is possible thanks to the implicit function theorem (Section 11.3). Using the chain rule, we get $\nabla L(\w) = \nabla_2 L(\s^\star(\w), \w) + \partial \s^\star(\w)^* \nabla_1 L(\s^\star(\w), \w),$ where $\partial \s^\star(\w)^*$ is the VJP of $\s^\star$ , a linear map from $\cS$ to $\cW$ .

Computationally, the main difficulty is to apply $\partial \s^\star(\w)^*$ to the vector $\u = \nabla_1 L(\s^\star(\w), \w) \in \cS$ . Using the implicit function theorem (Section 11.3) on the implicit function $c(\s^\star(\w), \w) = \zeros$ , and Proposition 11.1, we get the linear system $A^* \r = \u$ , where $A^* \coloneqq \partial_1 c(\s^\star(\w), \w)^*$ is a linear map from $\cS$ to $\cS$ . After solving for $\r$ , we get $\partial \s^\star(\w)^* \u = B^* \r$ , where $B^* \coloneqq \partial_2 c(\s^\star(\w), \w)^*$ is a linear map from $\cS$ to $\cW$ . Putting everything together, we get $\nabla L(\w) = \nabla_2 L(\s^\star(\w), \w) + \partial_2 c(\s^\star(\w), \w)^* \r.$

11.4.5 Reverse mode as adjoint method with backsubstitution

In this section, we revisit reverse-mode autodiff from the perspective of the adjoint state method. For clarity, we focus our exposition on feedforward networks with input $\x \in \cX$ and network weights $\w = (\w_1, \ldots, \w_k) \in \cW_1 \times \ldots \times \cW_K$ , defined by $\begin{align*} \s_0 &\coloneqq \x \in \cX \\ \s_1 &\coloneqq f_1(\s_0, \w_1) \in \cS_1 \\ & \hspace{6pt} \vdots \\ \s_K &\coloneqq f_K(\s_{K-1}, \w_K) \in \cS_K \\ f(\w) &\coloneqq \s_K. \tag{11.2}\end{align*}$ Here we focus on gradients with respect to the parameters $\w$ , hence the notation $f(\w)$ . We can use the adjoint state method to recover reverse-mode autodiff, and prove its correctness in the process. While we focus for simplicity on feedforward networks, our exposition can be generalized to computation graphs.

Feedforward networks as the solution of a nonlinear equation

While we defined the set of intermediate computations $\s = (\s_1, \ldots, \s_K) \in \cS_1 \times \ldots \times \cS_K$ as a sequence of operations, they can also be defined as the unique solution of the nonlinear equation $c(\s, \w) = \zeros$ , where $c(\s, \w) \coloneqq \begin{pmatrix} \s_1 - f_1(\x, \w_1) \\ \s_2 - f_2(\s_1, \w_2) \\ \vdots \\ \s_K - f_K(\s_{K-1}, \w_K) \end{pmatrix}.$ This defines an implicit function $\s^\star(\w) = (\s_1^\star(\w), \ldots, \s_K^\star(\w))$ , the solution of this nonlinear system, which is given by the variables $\s_1, \ldots, \s_K$ defined in Eq. (11.2). The output of the feedforward network is then $f(\w) = \s_K^\star(\w)$ .

In machine learning, the final layer $\s^\star_K(\w)$ is typically fed into a loss $\ell$ , to define $L(\w) \coloneqq \ell(\s_K^\star(\w); \y).$ Note that an alternative is to write $L(\w)$ as $L(\w) = \min_{\s \in \cS} \ell(\s; \y) \quad \text{s.t.} \quad c(\s, \w) = \zeros.$

More generally, if we just want to compute the VJP of $\s_K^\star(\w)$ in some direction $\u_K \in \cS_K$ , we can define the scalar-valued function $L(\w) \coloneqq \ell(\s_K^\star(\w); \u_K) \coloneqq \langle \s_K^\star(\w), \u_K \rangle$ so that $\partial f(\w)^* \u_K = \nabla L(\w).$ Let us define $\u \in \cS_1 \times \dots \times \cS_{K-1} \times \cS_K$ as $\u \coloneqq (\zeros, \ldots, \zeros, \nabla_1 \ell(f(\w); \y))$ (gradient of the loss $\ell$ case) or $\u \coloneqq (\zeros, \ldots, \zeros, \u_K)$ (VJP of $f$ in the direction $\u_K$ case). Using the adjoint state method, we know that the gradient of this objective is obtained as $\nabla L(\w) = \partial_2 c(\s(\w), \w)^* \r^\star(\w),$ for $\r^\star(\w)$ the solution of the linear system $\partial_1 c(\s(\w), \w)^* \r = -\u.$

Solving the linear system using backsubtitution

The JVP of the constraint function $c$ at $\s^\star(\w)$ , materialized as a matrix, takes the form of a block lower-triangular matrix $\jac_1 c(\s^\star(\w), \w) = \begin{pmatrix} \idm & 0 & \dots & \dots & 0 \\ - A_1 & \idm & \ddots & & \vdots\\ 0 & -A_2 & \idm & \ddots & \vdots \\ \vdots & \ddots & \ddots & \ddots & 0 \\ 0 & \dots & 0 & -A_K & \idm \end{pmatrix},$ where $A_k \coloneqq \jac_1 f_k(\s_{k-1}, \w_k)$ . Crucially the triangular structure of the JVP stems from the fact that each intermediate activation only depends on the past intermediate activations. Therefore, the constraints, corresponding to the lines of the Jacobian, cannot introduce non-zero values beyond its diagonal. The VJP takes the form of a block upper-triangular matrix $\jac_1 c(\s^\star(\w), \w)^* = \begin{pmatrix} \idm & - A_1^* & 0 & \dots & 0 \\ 0 & \idm & -A_2^* & \ddots & \vdots\\ \vdots & \ddots & \idm & \ddots & 0 \\ \vdots & & \ddots & \ddots & -A_K^* \\ 0 & \dots & \dots & 0 & \idm \end{pmatrix}.$ Solving an upper triangular system like $\partial_1 c(\s(\w), \w)^*\r = \u$ can then be done efficiently by backsubstitution. Starting from the last adjoint state $\r_K = \u$ , we can compute each adjoint state $\r_k$ from that computed at $k+1$ . Namely, for $k \in (K-1,\dots,1)$ , we have $\begin{align*} \r_k - A_{k+1}^*\r_{k+1} = \zeros \iff \r_k = \partial f_{k+1}(\s_k, \w_{k+1})^*\r_{k+1}. \end{align*}$ The VJPs with respect to the parameters are then obtained by $\partial_2 c(\s(\w), \w)^* \r = \begin{pmatrix} \partial_2 f_1(\x, \w_1)^* \r_1 \\ \partial_2 f_2(\s_1(\w), \w_2)^* \r_2 \\ \vdots \\ \partial_2 f_K(\s_{K-1}(\w), \w_K)^* \r_K \end{pmatrix},$ recovering reverse-mode autodiff.

The Lagrangian perspective of backpropagation for networks with separate parameters $\w = (\w_1, \dots, \w_K)$ is well-known; see for instance LeCun (1988) or Recht (2016). The Lagrangian perspective of backpropagation through time (Werbos, 1990) for networks with shared parameter $\w$ is discussed for instance by Franceschi et al. (2017). Our exposition uses the adjoint state method, which can itself be proved either using the method of Lagrange multipliers (Section 11.4.3) or by the implicit function theorem (Section 11.4.4), combined with backsubtitution for solving the upper-triangular linear system. Past works often minimize over $\w$ but we do not require this, as gradients are not necessarily used for optimization. Our exposition also supports computing the VJP of any vector-valued function $f$ , while existing works derive the gradient of a scalar-valued loss function.

11.5 Inverse function theorem

11.5.1 Differentiating inverse functions

In some cases (see for instance Section 12.4.4), it is useful to compute the Jacobian of an inverse function $f^{-1}$ . The inverse function theorem below allows us to relate the Jacobian of $f^{-1}$ with the Jacobian of $f$ .

Theorem 11.7 (Inverse function theorem)

Assume $f \colon \cW \to \cW$ is continuously differentiable with invertible Jacobian $\partial f(\wv_0)$ at $\wv_0$ . Then $f$ is bijective from a neighborhood of $\wv_0$ to a neighborhood of $f(\w_0)$ . Moreover, the inverse $f^{-1}$ is continuously differentiable near $\omegav_0 = f(\wv_0)$ and the Jacobian of the inverse $\partial f^{-1}(\omegav)$ is $\partial f(\wv) \partial f^{-1}(\omegav)= I \Leftrightarrow \partial f^{-1}(\omegav) = (\partial f(\wv))^{-1},$ with $\wv = f^{-1}(\omegav)$ .

11.5.2 Link with the implicit function theorem

The inverse function theorem can be used to prove the implicit function theorem; see proof of Theorem 11.6. Conversely, recall that, in order to use the implicit function theorem, we need to choose a root objective $F \colon \cW \times \Lambda \to \cW$ . If we set $\cW=\Lambda=\RR^Q$ and $F(\wv, \omegav) = f(\wv) - \omegav$ , with $f \colon \R^Q \to \R^Q$ , then we have that the root $\wv^\star(\omegav)$ satisfying $F(\wv^\star(\omegav), \omegav) = \zeros$ is exactly $\wv^\star(\omegav) = f^{-1}(\omegav)$ . Moreover, $\partial_1 F(\wv, \omegav) = \partial f(\wv)$ and $\partial_2 F(\wv, \omegav) = -I$ . By applying the implicit function theorem with this $F$ , we indeed recover the inverse function theorem.

11.5.3 Proof of inverse function theorem

We first give a proof of the formula assuming that $f^{-1}$ is well-defined and continuously differentiable in a neighborhood of $f(\w_0)$ . In that case, we have for any $\omegav$ in a neighborhood of $f(\w_0)$ , $f \circ f^{-1}(\omegav) = \omegav.$ Differentiating both sides w.r.t. $\omegav$ , we get $\partial f(f^{-1}(\omegav)) \partial f^{-1}(\omegav) = \idm,$ where $\idm$ is the identity function in $\RR^Q$ . In particular, for $\w = f^{-1}(\omegav)$ we recover the formula presented in the statement.

Now, it remains to show that invertibility of the JVP ensures that the function is invertible in a neighborhood of $f(\w_0)$ and that the inverse is continuously differentiable. For that, denote $\l=\partial f(\w_0)$ such that $\l^{-1}$ is well-defined by definition. $f$ is invertible with continuously differentiable inverse, if and only if $\l^{-1}(f(\w)) - f(\w_0)$ is invertible with continuously differentiable inverse. So without loss of generality, we consider $\partial f(\w_0) = \idm$ , $f(\w_0) = \zeros$ , $\w_0=\zeros$ .

As $f$ is continuously differentiable, there exists a neighborhood $\cN \coloneqq \{\w: \|\w-\w_0\|_2\leq \delta\}$ on which we have $\|\partial f(\w) - \idm\|_2\leq 1/2$ . In this neighborhood, the function $g(\w) = f(\w) - \w$ is contractive by the mean value theorem with contraction factor $1/2$ . For any $\omegav$ such that $\|\omegav-f(\w_0)\|_2 \leq \delta/2$ , the sequence $\w_{k+1} = \w_k - f(\w_k) - \omegav'$ remains in $\mathcal{N}$ and converges (since it is a Cauchy sequence by the contraction of $g$ ) to a unique fixed point $\w_\infty$ satisfying $\w_\infty = \w_\infty - f(\w_\infty) - \omegav \iff f(\w_\infty) = \omegav$ . This shows the existence of the inverse in the neighborhood $\cM \coloneqq \{\omegav: \|\omegav-\omegav_0\|_2\leq \delta/2\}$ of $\omegav_0 = f(\w_0)$ onto $\cN$ .

We tackle now the differentiability (hence the continuity) of $f^{-1}$ . For any $\omegav$ in the neighborhood of $\omegav_0$ with inverse $\w \coloneqq f^{-1}(\omegav) \in N$ , the JVP of $f$ at $\w$ satisfies by assumption $\|\partial f(\w) - \idm\|_2 \leq 1/2$ . Hence, $\a = \partial f(\w) -\idm$ defines a convergent series $\b = \sum_{k=0}^{+\infty} \a^k$ and one verifies easily that $\b = \partial f(\w)^{-1}$ , that is $\partial f(\w)$ is invertible and $\|\partial f(\w)^{-1}\|\leq 2$ . To compute the JVP of the inverse, we consider then $\partial f(\w)^{-1}$ as the candidate JVP and examine $\frac{\|f^{-1}(\omegav + \etav) - f(\omegav) - (\partial f(\w))^{-1}\etav\|_2} {\|\etav\|_2}.$ Denote then $\v$ such that $f^{-1}(\omegav + \etav) = \w + \v$ . As $g(\w) = f(\w) - \w$ is $1/2$ -contractive in $\cN$ , we have $\|\v -\etav\|_2 = \|g(\w + \v) - g(\w)\|_2 \leq 1/2\|\v\|_2$ . So $\|\v\|_2 \geq \|\etav\|/2$ . We then get $\begin{align*} & \frac{\|f^{-1}(\omegav + \etav) - f(\omegav) - (\partial f(\w))^{-1}\etav\|_2} {\|\etav\|_2} \\ & = \frac{\|\v - (\partial f(\w))^{-1}(f(\w + \v) - f(\w))\|_2} {\|\etav\|_2} \\ & \leq 4 \frac{\|f(\w+\v) - f(\w) - \partial f(\w)\v\|_2}{\|\v\|_2} \end{align*}$ As $\|\etav\|_2 \rightarrow 0$ , we have $\|\v\|_2 \rightarrow 0$ and so $\|f(\w+\v) - f(\w) - \partial f(\w)\v\|_2/\|\v\|_2 \rightarrow 0$ . Hence, $f^{-1}$ is differentiable with JVP $\partial f^{-1}(\omegav) = (\partial f(\w))^{-1} = (\partial f(f^{-1}(\omegav)))^{-1}$ . This shows that $f^{-1}$ is continuous and so $\partial f^{-1}(\omegav)$ is continuous as a composition of continuous functions.

11.6 Summary

Implicit functions are functions that cannot be decomposed into elementary operations and for which autodiff can therefore not be directly applied. Examples are optimization problems and nonlinear equations.
Envelope theorems can be used for differentiating through the min or max value (not solution) of a function.
More generally, the implicit function theorem allows us to differentiate through implicit functions. It gives conditions for the existence of derivatives and how to obtain them.
The adjoint state method can be used to obtain the gradient of the composition of an explicit function and of an implicit function, specified by equality constraints. It can be used to prove the correctness of reverse-mode autodiff.
The inverse function theorem can be used to differentiate function inverses.
In a sense, the implicit function theorem can be thought of as the mother theorem, as it can be used to prove envelope theorems, the adjoint state method and the inverse function theorem.