The implicit function theorem and gradients through solvers

Before this lesson: lesson 04 (the chain rule as a product of Jacobians), lesson 05 (JVP and VJP, why reverse mode brackets the product from the loss side) and lesson 08 (matrix-free products such as HVPs). From linear algebra: lesson 03 (inverse, "solve, do not invert"), lesson 13 (conditioning) and lesson 14 (CG and GMRES). Time: about 70 minutes, plus the notebook.

---

Why this matters #

• Deep equilibrium models (DEQs) find a fixed point $z^* = f(z^*, x; \theta)$ and train with an implicit gradient, so memory does not grow with the number of solver iterations (Bai, Kolter and Koltun 2019).
• Bilevel optimization. Hyperparameter tuning and meta-learning have an inner problem $w^*(\lambda) = \arg\min_w$ (training loss) and an outer loss on validation data. The hypergradient $dw^*/d\lambda$ comes from the optimality condition of the inner problem (Lorraine, Vicol and Duvenaud 2020; implicit MAML, Rajeswaran et al. 2019).
• Optimization layers. OptNet (Amos and Kolter 2017) differentiates the solution of a quadratic program through its optimality (KKT) conditions.
• ML for science. A discretized steady PDE $R(s, m) = 0$ links a state $s$ (pressure, a wavefield) to parameters $m$ (permeability, wave speed). The adjoint method gives the gradient of a misfit with respect to millions of parameters for one extra linear solve. It is the same theorem.

1. Predict first #

2. The picture #

One dimension: the slope of a curve you never solve #

The equation $x^2 - \theta = 0$ with $x > 0$ defines $x(\theta) = \sqrt\theta$. Now forget the formula $\sqrt\theta$. When $\theta$ moves by $d\theta$, the solution must move by $dx$ so that $F(x, \theta) = x^2 - \theta$ stays zero. To first order, $F$ changes by

Set $dF = 0$ and solve for the ratio:

At $\theta = 4$, $x = 2$, so $dx/d\theta = 1/4$. This equals $\frac{d}{d\theta}\sqrt\theta = 1/(2\sqrt\theta)$, but we used only $F$ and its partial derivatives at the solution.

Where it fails. At $\theta = 0$ the solution is $x = 0$, so $F_x = 0$. There the two branches $\pm\sqrt\theta$ meet, the tangent of the curve is vertical, and for $\theta < 0$ there is no real solution. Such a point is a fold. The condition $F_x \ne 0$ rules out exactly this situation.

Interactive figure: loading…

The linear case in $n$ dimensions #

Let $F(x, \theta) = Ax - B\theta$ with an invertible $A \in \R^{n\times n}$ and $B \in \R^{n\times p}$. Then $x^* = A^{-1}B\theta$ and $dx^*/d\theta = A^{-1}B$. Here $F_x = A$ and $F_\theta = -B$, so $dx^*/d\theta = -F_x^{-1}F_\theta$. Near a solution, a smooth nonlinear $F$ behaves like its linear approximation, so the same formula holds locally. That is the implicit function theorem.

Fixed points: a feedback loop #

Let $x = f(x, \theta)$, and write $f_x = \partial f/\partial x \in \R^{n\times n}$ and $f_\theta = \partial f/\partial\theta \in \R^{n\times p}$ at the fixed point. A change $d\theta$ moves $x$ directly by $f_\theta\, d\theta$. That move goes through $f$ again and moves $x$ by $f_x$ times it, and so on:

This Neumann series converges when the spectral radius $\rho(f_x)$ (the largest $|\text{eigenvalue}|$) is below 1. Each solver iteration adds one term. The implicit derivative is the whole sum.

Interactive figure: loading…

Argmins: a minimizer is a root #

If $x^*(\theta)$ minimizes a smooth function $h(x, \theta)$ over $x$, then $\nabla_x h(x^*, \theta) = 0$. So the argmin is a root of $F = \nabla_x h$, and $F_x$ is the Hessian $\nabla^2_{xx} h$. Everything above applies.

3. Definitions and main results #

Notation. $F: \R^n \times \R^p \to \R^n$ (as many equations as unknowns). $F_x = \partial F/\partial x \in \R^{n\times n}$ and $F_\theta = \partial F/\partial\theta \in \R^{n\times p}$ are the partial Jacobians. $Dx^*(\theta) \in \R^{n\times p}$ is the Jacobian of the solution map $\theta \mapsto x^*(\theta)$. A map is $C^1$ if it is differentiable with continuous derivatives.

3.1 The implicit function theorem #

Theorem (implicit function theorem). Let $F$ be $C^1$ near $(x_0, \theta_0)$, with $F(x_0, \theta_0) = 0$ and $F_x(x_0, \theta_0)$ invertible. Then there are open sets $U \ni \theta_0$ and $V \ni x_0$ and a unique map $x^*: U \to V$ such that, for every $\theta \in U$, $x^*(\theta)$ is the only solution of $F(x, \theta) = 0$ in $V$. The map $x^*$ is $C^1$, $x^*(\theta_0) = x_0$, and

Proof of the derivative formula (given that $x^*$ exists and is $C^1$).

1. For every $\theta \in U$, $\phi(\theta) := F(x^*(\theta), \theta) = 0$. A function that is zero everywhere has derivative zero: $D\phi(\theta) = 0$.
2. $\phi$ is the composition $\theta \mapsto (x^*(\theta), \theta) \mapsto F$. The chain rule (lesson 04) gives $D\phi = F_x\, Dx^* + F_\theta$.
3. So $F_x\, Dx^* = -F_\theta$.
4. $F_x$ is invertible at $(x_0, \theta_0)$. Its determinant is continuous, so it stays nonzero on a smaller $U$. Multiply on the left by $F_x^{-1}$. $\blacksquare$

Shapes: $(n\times n)^{-1}(n\times p) = n\times p$. The formula uses only $F$ at the solution. It does not care how $x^*$ was found.

Why the solution exists and is unique (sketch)

1. Define $\Psi(x, \theta) = (F(x, \theta), \theta)$, a map from $\R^n\times\R^p$ to itself.
2. Its Jacobian is block triangular, $\begin{pmatrix} F_x & F_\theta \\ 0 & I \end{pmatrix}$, with determinant $\det F_x \ne 0$ at $(x_0, \theta_0)$.
3. By the inverse function theorem, $\Psi$ has a $C^1$ inverse near $\Psi(x_0, \theta_0) = (0, \theta_0)$.
4. The solution is the first block of $\Psi^{-1}(0, \theta)$. It is $C^1$ because $\Psi^{-1}$ is, and it is unique near $x_0$ because $\Psi$ is one-to-one there (Blondel and Roulet, §11.3.4). $\blacksquare$

Fixed points. For $x^* = f(x^*, \theta)$, set $F(x, \theta) = x - f(x, \theta)$. Then $F_x = I - f_x$ and $F_\theta = -f_\theta$, so

The condition is that $I - f_x$ is invertible, that is, 1 is not an eigenvalue of $f_x$. A contraction ($\|f_x\| < 1$ in some induced norm) is enough, but it is not necessary: a repelling fixed point also has a derivative (worked example (b)).

Argmins. For $x^*(\theta) = \arg\min_x h(x, \theta)$ with $h \in C^2$, set $F = \nabla_x h$:

where $\nabla^2_{x\theta} h \in \R^{n\times p}$ holds the mixed partials $\partial^2 h/\partial x_i \partial\theta_j$. It needs an invertible Hessian at $x^*$, for example $h$ strongly convex in $x$.

3.2 The VJP form: one adjoint solve #

Training needs the gradient of a scalar $L(\theta) = \ell(x^*(\theta), \theta)$, not the $n\times p$ Jacobian. Let $\bar x = \nabla_x \ell(x^*, \theta) \in \R^n$ be the upstream gradient, and $\partial_\theta \ell \in \R^p$ the direct dependence of $\ell$ on $\theta$ (often zero).

Proposition (adjoint form). Let $u \in \R^n$ solve $F_x^\top u = \bar x$. Then $\nabla_\theta L = \partial_\theta \ell - F_\theta^\top u$.

Proof.

1. Chain rule: $\nabla_\theta L = \partial_\theta \ell + (Dx^*)^\top \bar x$.
2. Insert the theorem: $(Dx^*)^\top \bar x = -F_\theta^\top F_x^{-\top} \bar x$, where $F_x^{-\top} = (F_x^\top)^{-1}$.
3. Bracket from the right, as reverse mode does (lesson 05): $u = F_x^{-\top}\bar x$, that is, $u$ solves $F_x^\top u = \bar x$.
4. Then $(Dx^*)^\top \bar x = -F_\theta^\top u$: one VJP of $\theta \mapsto F(x^*, \theta)$ with cotangent $u$. $\blacksquare$

Cost. Forming $Dx^*$ needs $p$ solves with $F_x$ (one per column of $F_\theta$) and stores $np$ numbers. The adjoint form needs one solve with $F_x^\top$ and one VJP, for any $p$.

Matrix-free solves. The products $v \mapsto F_x^\top v$ are VJPs of $F$ in $x$, so a Krylov solver finds $u$ without forming $F_x$ (linear-algebra lesson 14). Use conjugate gradient (CG) when $F_x$ is symmetric positive definite. This is the case for the argmin of a strongly convex $h$, where each product is an HVP of $h$ (lesson 08). Otherwise use GMRES or BiCGSTAB.

Fixed points. The adjoint system is $(I - f_x)^\top u = \bar x$, and $\nabla_\theta L = \partial_\theta \ell + f_\theta^\top u$. When $f$ is a contraction, the plain iteration

converges to $u$, with one VJP of $f$ per step. Its rate is the spectral radius of $f_x$, the same as the forward iteration, because $f_x^\top$ and $f_x$ have the same eigenvalues.

Interactive figure: loading…

3.3 Unrolled versus implicit #

Unrolling runs $x_{k+1} = f(x_k, \theta)$ for $K$ steps and backpropagates through all of them. If the iteration starts at the fixed point ($x_0 = x^*$, so every $x_k = x^*$), the chain rule gives $dx_K/d\theta = \sum_{j=0}^{K-1} f_x^j f_\theta$: the first $K$ terms of the Neumann series. The missing tail is $f_x^K (I - f_x)^{-1} f_\theta$, so the error shrinks like $\rho^K$, where $\rho$ is the spectral radius of $f_x$. From another start the error also decays geometrically (in worked example (b) it is proportional to $K\rho^K$).

	Unrolled (backprop through $K$ steps)	Implicit (theorem at $x^*$)
What it differentiates	the $K$-step algorithm, exactly	the solution map $\theta \mapsto x^*(\theta)$
Memory	all $K$ iterates and the intermediates of every step: grows like $K$	$x^*$ and one graph of $f$ at $x^*$: independent of $K$
Backward cost	$K$ VJPs of one step	one linear solve (some VJPs of $F$ in $x$) plus one VJP in $\theta$
Error against the true $Dx^*$	truncation (about $\rho^K$) plus the forward error	proportional to the forward error $\lVert\hat x - x^*\rVert$, plus the error of the adjoint solve
Needs	a differentiable solver	$F$ differentiable at $x^*$ and $F_x$ invertible; any forward solver

Accuracy of the implicit gradient. In practice you apply the formula at an approximate solution $\hat x$. Blondel et al. (2022, Theorem 1) show: if $F_x$ is well conditioned near $x^*$, $F_x$ and $F_\theta$ are Lipschitz there, and $F_\theta$ is bounded at $x^*$, then the Jacobian estimate at $\hat x$ has error at most $C\,\|\hat x - x^*\|$. So the gradient is about as accurate as the forward solve. The constant $C$ grows as the smallest singular value of $F_x$ shrinks, that is, as $F_x$ gets close to singular (linear-algebra lesson 13).

When unrolling is the right choice. When you want the gradient of what you actually computed: a few inner steps by design (MAML-style adaptation, learned optimizers), or a solution that is not unique or not smooth.

3.4 Ridge regression: a hyperparameter gradient #

The inner problem is $h(w, \lambda) = \frac12\|Xw - y\|^2 + \frac{\lambda}{2}\|w\|^2$ with $X \in \R^{N\times d}$, $y \in \R^N$ and $\lambda > 0$. Its optimality condition is

The partial Jacobians are $F_w = X^\top X + \lambda I =: A_\lambda$ (symmetric positive definite) and $F_\lambda = w$ (a $d\times1$ column). The theorem gives

For a validation loss $V(w^*(\lambda))$ with $\bar w = \nabla_w V(w^*)$, the adjoint form gives $dV/d\lambda = -w^{*\top} u$ with $A_\lambda u = \bar w$ (here $A_\lambda^\top = A_\lambda$). One solve, and no closed form of $w^*(\lambda)$ is needed, only its optimality condition.

Corollary (the norm shrinks). $\frac{d}{d\lambda}\frac12\|w^*\|^2 = w^{*\top}\frac{dw^*}{d\lambda} = -w^{*\top}A_\lambda^{-1}w^* < 0$ whenever $w^* \ne 0$, because $A_\lambda^{-1} \succ 0$ (linear-algebra lesson 07).

Interactive figure: loading…

Envelope theorem: when no solve is needed

If the outer loss is the optimal value itself, $v(\theta) = \min_x h(x, \theta)$, then $\nabla v(\theta) = \partial_\theta h(x^*, \theta)$. The implicit term $(Dx^*)^\top \nabla_x h$ vanishes because $\nabla_x h(x^*, \theta) = 0$, so no solve is needed. This is Danskin's theorem (Blondel and Roulet §11.2). It is why the gradient of a log-partition function or of a dual objective is simple.

Answers to section 1. (1) Unrolled memory grows linearly in $K$: every iterate and the intermediates of every evaluation of $f$ are stored. Yes: the theorem gives the exact gradient from $z^*$ alone, with one adjoint solve $(I - f_z)^\top u = \nabla\ell(z^*)$ and one VJP $f_\theta^\top u$; memory does not depend on $K$. (2) Not every weight: in worked example (c), $w_2^* = 0$ at $\lambda = 1$ and $dw_2^*/d\lambda = +1/8$, so $|w_2^*|$ grows. But $\|w^*\|$ always shrinks (the corollary).

4. Worked example #

(a) A 2-D root. Two unknowns $x = (x_1, x_2)$ and one parameter $\theta$:

At $\theta = 2$, $x = (1, 1)$ is a root: $F = (1 + 1 - 2,\; 1 - 1) = (0, 0)$.

1. Partial Jacobians: $F_x = \begin{pmatrix} 1 & 2x_2 \\ x_2 & x_1 \end{pmatrix} = \begin{pmatrix} 1 & 2 \\ 1 & 1 \end{pmatrix}$, with $\det F_x = 1 - 2 = -1 \ne 0$, so the theorem applies. $F_\theta = (-1, 0)^\top$.
2. $F_x^{-1} = \frac{1}{-1}\begin{pmatrix} 1 & -2 \\ -1 & 1 \end{pmatrix} = \begin{pmatrix} -1 & 2 \\ 1 & -1 \end{pmatrix}$.
3. $Dx^* = -F_x^{-1}F_\theta = F_x^{-1}(1, 0)^\top = (-1, 1)^\top$.
4. Check without the theorem: $x_2 = 1/x_1$ gives $x_1 + 1/x_1^2 = \theta$. Then $d\theta/dx_1 = 1 - 2/x_1^3 = -1$ at $x_1 = 1$, so $dx_1/d\theta = -1$, and $dx_2/d\theta = -x_1^{-2}\, dx_1/d\theta = 1$. ✓
5. Prediction: at $\theta = 2.1$, $x \approx (1, 1) + 0.1\,(-1, 1) = (0.9, 1.1)$. The exact root is $(0.9209, 1.0858)$. The error, about $0.02$, is second order in $d\theta = 0.1$.
6. Where it fails: $\det F_x = x_1 - 2x_2^2$. On the curve $x_2 = 1/x_1$ it is zero when $x_1^3 = 2$, that is at $x_1 = 2^{1/3} \approx 1.26$ and $\theta = 3 \cdot 2^{-2/3} \approx 1.89$. There the branch through $(1, 1)$ meets a second branch (at $\theta = 2$ its root is $x_1 = (1 + \sqrt5)/2 \approx 1.618$, with $dx_1/d\theta \approx +1.89$). The theorem is local: each branch has its own derivative.

(b) A fixed point and its unrolled gradient. $f(x, \theta) = (x^2 + \theta)/4$ at $\theta = 3$. The fixed points solve $x = (x^2 + 3)/4$, that is $x^2 - 4x + 3 = 0$: $x^* = 1$ and $x^* = 3$.

1. At $x^* = 1$: $f_x = x/2 = 1/2$ and $f_\theta = 1/4$, so $dx^*/d\theta = (1/4)/(1 - 1/2) = 1/2$. Check: $x^*(\theta) = 2 - \sqrt{4 - \theta}$ has derivative $1/(2\sqrt{4 - \theta}) = 1/2$ at $\theta = 3$. ✓
2. At $x^* = 3$: $f_x = 3/2 > 1$, so the iteration moves away from this point. But $1 - f_x = -1/2 \ne 0$, and the theorem gives $dx^*/d\theta = (1/4)/(-1/2) = -1/2$ (check: $2 + \sqrt{4 - \theta}$ has derivative $-1/2$). The theorem needs invertibility, not a contraction.
3. At $\theta = 4$ the two fixed points merge at $x = 2$, where $f_x = 1$ and $1 - f_x = 0$: the theorem fails.
4. Unrolled, started at $x_0 = x^* = 1$: $dx_K/d\theta = \frac14\big(1 + \frac12 + \cdots + \frac{1}{2^{K-1}}\big) = \frac12(1 - 2^{-K})$. Started at $x_0 = 0$, forward-mode differentiation of the $K$-step map gives the second row:

$K$	1	2	5	10	20
start at $x^*$	0.25	0.375	0.484	0.49951	0.4999995
start at $0$	0.25	0.344	0.465	0.49797	0.4999961

The implicit answer $0.5$ costs one division. Unrolling needs about 20 steps, and 20 stored iterates, for 5–6 correct digits.

(c) A ridge hypergradient. Three training samples, two features, $\lambda = 1$: $X = \begin{pmatrix} 1 & 1 \\ 1 & 0 \\ 0 & 1 \end{pmatrix}$, $y = (1, 2, 0)$. One validation sample $x_{\text{val}} = (0, 1)$, $y_{\text{val}} = 1$, so $V = \frac12(x_{\text{val}}^\top w^* - y_{\text{val}})^2 = \frac12(w_2^* - 1)^2$.

1. $X^\top X = \begin{pmatrix} 2 & 1 \\ 1 & 2 \end{pmatrix}$ and $X^\top y = (1 + 2,\; 1 + 0) = (3, 1)$.
2. $A_1 = X^\top X + I = \begin{pmatrix} 3 & 1 \\ 1 & 3 \end{pmatrix}$, with $\det A_1 = 8$ and $A_1^{-1} = \frac18\begin{pmatrix} 3 & -1 \\ -1 & 3 \end{pmatrix}$.
3. Optimality: $A_1 w^* = (3, 1)$ gives $w^* = \frac18(9 - 1,\; -3 + 3) = (1, 0)$.
4. $dw^*/d\lambda = -A_1^{-1}w^* = -\frac18(3, -1) = (-3/8,\; 1/8)$.
5. Forward route: $\bar w = \nabla_w V = (0,\; w_2^* - 1) = (0, -1)$, so $dV/d\lambda = \bar w^\top dw^*/d\lambda = -1/8$.
6. Adjoint route: solve $A_1 u = \bar w$: $u = \frac18(1, -3)$. Then $dV/d\lambda = -w^{*\top}u = -1/8$. ✓
7. $dV/d\lambda < 0$, so increasing $\lambda$ lowers the validation loss.
8. The norm: $\frac{d}{d\lambda}\frac12\|w^*\|^2 = w^{*\top}(-3/8,\, 1/8) = -3/8 < 0$, while $w_2^*$ grows.
9. Check with the closed form $w^*(\lambda) = \big(3\lambda + 5,\; \lambda - 1\big)/\big((\lambda + 1)(\lambda + 3)\big)$: its derivative at $\lambda = 1$ is $(-3/8, 1/8)$. ✓

5. In code #

NumPy and SciPy check the numbers of section 4. The solver is a black box; the derivative comes from the equation at the solution.

import numpy as np
from scipy.optimize import fsolve

# (a) 2-D root: implicit Jacobian vs finite differences of a black-box solver
F = lambda x, t: [x[0] + x[1]**2 - t, x[0] * x[1] - 1]
root = lambda t: fsolve(F, [1.0, 1.0], args=(t,), xtol=1e-13)
x = root(2.0)
Fx, Ft = np.array([[1, 2 * x[1]], [x[1], x[0]]]), np.array([-1.0, 0.0])
h = 1e-5
print(-np.linalg.solve(Fx, Ft), (root(2 + h) - root(2 - h)) / (2 * h))   # [-1. 1.] [-1. 1.]
print(root(2.1))                                          # [0.92094247 1.08584415]

# (b) fixed point x = (x^2 + t)/4 at t = 3: unrolled derivative by forward mode, from x0 = 0
def unrolled(t, K, x=0.0, dx=0.0):
    for _ in range(K):
        x, dx = (x * x + t) / 4, (2 * x * dx + 1) / 4    # dx_{k+1} = f_x dx_k + f_t
    return dx
print([round(unrolled(3.0, K), 7) for K in (1, 2, 5, 10, 20)])
# [0.25, 0.34375, 0.4650059, 0.4979725, 0.4999961]   -> implicit: (1/4)/(1 - 1/2) = 0.5

# (c) ridge hypergradient with one solve
X, y = np.array([[1.0, 1.0], [1.0, 0.0], [0.0, 1.0]]), np.array([1.0, 2.0, 0.0])
lam = 1.0
A = X.T @ X + lam * np.eye(2)
w = np.linalg.solve(A, X.T @ y)
wbar = np.array([0.0, w[1] - 1.0])                        # gradient of V = (w2 - 1)^2 / 2
u = np.linalg.solve(A, wbar)                              # adjoint solve (A is symmetric)
print(w, -np.linalg.solve(A, w), -w @ u)                  # [1. 0.] [-0.375 0.125] -0.125

In PyTorch you attach the implicit gradient to a solver with a custom backward. The forward pass runs the solver without recording a graph; the backward pass solves the adjoint equation at $x^*$:

import torch

class FixedPoint(torch.autograd.Function):
    """x* = f(x*, theta) for f(x, theta) = (x^2 + theta) / 4, with an implicit backward."""
    @staticmethod
    def forward(ctx, theta):
        x = torch.zeros_like(theta)
        for _ in range(200):                  # any solver; no graph is recorded here
            x = (x * x + theta) / 4
        ctx.save_for_backward(x)
        return x

    @staticmethod
    def backward(ctx, xbar):
        (x,) = ctx.saved_tensors
        f_x, f_theta = x / 2, 0.25            # partial derivatives of f at x*
        u = xbar / (1 - f_x)                  # adjoint solve (I - f_x)^T u = xbar
        return f_theta * u                    # f_theta^T u

theta = torch.tensor(3.0, dtype=torch.float64, requires_grad=True)
x = FixedPoint.apply(theta)
x.backward()
print(x.item(), theta.grad.item())            # 1.0 0.5

In JAX, jax.lax.custom_root and the library JAXopt (Blondel et al. 2022) do the same for roots and optimality conditions.

6. Where it shows up #

• Deep equilibrium models. A DEQ layer finds $z^* = f(z^*, x; \theta)$ with Broyden or Anderson iterations and trains with the implicit VJP, so memory does not grow with the number of iterations. Cheaper variants truncate the adjoint solve or set $u = \bar z$ ("Jacobian-free backprop"); their gradients are approximate by design (deep learning).
• Bilevel optimization. Lorraine, Vicol and Duvenaud (2020) tune millions of hyperparameters with a truncated Neumann series of HVPs for the inverse Hessian. Implicit MAML (Rajeswaran et al. 2019) solves the adjoint system with CG on HVPs.
• Optimization layers. OptNet differentiates the KKT conditions of a quadratic program. For a loss $s + r$ with smooth $s$ and non-smooth $r$ (an $\ell_1$ penalty), the proximal-gradient fixed point $x = \operatorname{prox}_{\eta r}(x - \eta\nabla s(x))$ serves as the equation (optimization).
• Matrix identities. $d(X^{-1}) = -X^{-1}\, dX\, X^{-1}$ (check yourself 2) and the VJP of a linear solve are the theorem applied to $XY - I = 0$ and to $Ax - b = 0$ (linear-algebra lesson 16).
• ML for science. For a discretized steady PDE $R(s, m) = 0$, the gradient of a misfit with respect to all parameters $m$ costs one adjoint solve with $(\partial R/\partial s)^\top$: for example one Helmholtz system per source and frequency in frequency-domain full-waveform inversion. An implicit time step (backward Euler for stiff terms) is a root problem, differentiated the same way. The time-continuous case is the adjoint method for ODEs, used by neural ODEs (Chen et al. 2018; ML for science).

7. Common mistakes #

• Differentiating an unconverged solution. The formula assumes $F(\hat x, \theta) = 0$. Check the forward residual first; the gradient error is about proportional to $\|\hat x - x^*\|$.
• Forgetting the transpose. The adjoint solve uses $F_x^\top$. It equals $F_x$ only when $F_x$ is symmetric, as for an argmin (a Hessian). For a DEQ, $I - f_x$ is not symmetric.
• Sign errors in the fixed-point form. With $F = x - f$: $Dx^* = +(I - f_x)^{-1} f_\theta$ and $\nabla_\theta L = \partial_\theta \ell + f_\theta^\top u$. If $\ell$ depends on $\theta$ directly, do not drop $\partial_\theta \ell$.
• Letting autograd record the solver. Then you pay unrolled memory anyway and get the unrolled gradient. Run the solver without a graph (torch.no_grad()), evaluate $f$ once at $x^*$ with a graph for the VJPs, and attach the implicit gradient in a custom backward.
• Using the theorem where it does not hold. $F_x$ singular or nearly singular (folds and bifurcations), non-unique minimizers (ridge with $\lambda = 0$ and rank-deficient $X$), or $F$ not differentiable at $x^*$ (a ReLU kink, a constraint that becomes active).
• The plain adjoint iteration without a contraction. $u \leftarrow \bar x + f_x^\top u$ does not converge when the spectral radius of $f_x$ is $\ge 1$. Use GMRES instead.

8. Check yourself #

1. The circle $x^2 + \theta^2 = 1$ defines $x(\theta)$ near $(x, \theta) = (0.6, 0.8)$. Find $dx/d\theta$ with the theorem. Where on the circle does it fail for $x$ as a function of $\theta$?
   Answer

   $F_x = 2x = 1.2$ and $F_\theta = 2\theta = 1.6$, so $dx/d\theta = -1.6/1.2 = -4/3$ (check: $x = \sqrt{1 - \theta^2}$ gives $-\theta/x = -0.8/0.6$). It fails at $x = 0$, $\theta = \pm1$: there $F_x = 0$, the two branches $\pm\sqrt{1 - \theta^2}$ meet and the tangent is vertical.
2. Derive $d(X^{-1})$ from the theorem.
   Answer

   $Y = X^{-1}$ solves $F(Y, X) = XY - I = 0$. Differentiate: $dX\,Y + X\,dY = 0$, so $dY = -X^{-1}\,dX\,X^{-1}$. The partial derivative in $Y$ is the linear map $dY \mapsto X\,dY$, which is invertible exactly when $X$ is.
3. The fixed point $x = 0.8x + \theta$. Give $dx^*/d\theta$, the unrolled derivative after $K$ steps from $x_0 = 0$, and the smallest $K$ with relative error below 1%.
   Answer

   Implicit: $1/(1 - 0.8) = 5$. Unrolled: $x_K = \theta(1 + 0.8 + \cdots + 0.8^{K-1})$, so $dx_K/d\theta = 5(1 - 0.8^K)$ and the relative error is $0.8^K$. $0.8^{20} \approx 0.0115$ and $0.8^{21} \approx 0.0092$, so $K = 21$.
4. $x^* \in \R^{1000}$ and $\theta \in \R^{10^6}$. Compare forming $Dx^*$ with the adjoint route for $\nabla_\theta \ell(x^*)$.
   Answer

   $Dx^* = -F_x^{-1}F_\theta$ has $10^9$ entries (8 GB in float64) and needs $10^6$ solves with $F_x$ (one LU factorization, then $10^6$ pairs of triangular solves). The adjoint route needs one solve $F_x^\top u = \bar x$ and one VJP $-F_\theta^\top u$, and stores nothing of size $n\times p$.
5. Show that $\|w^*(\lambda)\|$ of ridge regression decreases in $\lambda$, and say why one coordinate can still grow.
   Answer

   $\frac{d}{d\lambda}\frac12\|w^*\|^2 = -w^{*\top}A_\lambda^{-1}w^* < 0$ for $w^* \ne 0$, since $A_\lambda^{-1} \succ 0$. The direction $dw^*/d\lambda = -A_\lambda^{-1}w^*$ has a negative inner product with $w^*$, but the off-diagonal entries of $A_\lambda^{-1}$ mix the coordinates, so a single coordinate can grow (worked example (c): $dw_2^*/d\lambda = +1/8$).
6. The forward iteration $x \leftarrow f(x, \theta)$ converges with local rate 0.9 (the spectral radius of $f_x$ at $x^*$). Does the adjoint iteration $u \leftarrow \bar x + f_x^\top u$ converge, and how many steps does each extra correct digit cost?
   Answer

   Yes. Its iteration matrix $f_x^\top$ has the same eigenvalues as $f_x$, so the error shrinks by about 0.9 per step: $\ln 10/\ln(1/0.9) \approx 22$ steps per digit. A Krylov solver such as GMRES can need far fewer.

Recall cardsAnswer in your head first, then reveal. Add the cards to your review deck to see them again just before you would forget.

1. State the derivative formula of the implicit function theorem for $F(x^*(\theta), \theta) = 0$.

   Show answer

2. What condition does the implicit function theorem need (besides $F \in C^1$ and $F(x_0, \theta_0) = 0$)?

   Show answer

3. Derivative of a fixed point $x^* = f(x^*, \theta)$?

   Show answer

4. Compute $dx^*/d\theta$ for the fixed point of $f(x, \theta) = 0.9x + \theta$.

   Show answer

5. How do you get $\nabla_\theta L$ for $L = \ell(x^*(\theta))$ without forming $Dx^*$?

   Show answer

6. Why does the implicit gradient need memory independent of the number of solver iterations $K$?

   Show answer

7. What is $dw^*/d\lambda$ for ridge regression?

   Show answer

8. When is the implicit gradient inaccurate?

   Show answer

9. Why can unrolling be the right choice?

   Show answer

Add these cards to my reviewOpen the review →

9. Interview angle #

• "How does a DEQ backpropagate with constant memory?" Implicit function theorem at $z^*$: solve $(I - f_z)^\top u = \bar z$ with VJPs (fixed-point iteration or a Krylov method), then $\nabla_\theta \ell = f_\theta^\top u$. Only $z^*$ is stored.
• "How do you get the gradient of the validation loss with respect to weight decay without unrolling?" Differentiate the optimality condition: $dw^*/d\lambda = -H^{-1}w^*$ with $H$ the training Hessian (for ridge, $H = X^\top X + \lambda I$), and get $H^{-1}\bar w$ by CG on HVPs or a truncated Neumann series.
• "What is the VJP of $x = A^{-1}b$?" Solve $A^\top u = \bar x$; then $\bar b = u$ and $\bar A = -u\,x^\top$.
• "When would you unroll instead?" For the gradient of a few-step or early-stopped computation, or when the solution is not unique or not smooth.
• "What happens at a fold?" $F_x$ is singular, the derivative is infinite (vertical tangent) and two branches meet; near the fold the implicit gradient is ill-conditioned.

10. Practice #

The notebook for this lesson runs in Colab. Five exercises, each with a check cell and a folded solution:

1. newton_root and implicit_jacobian for the 2-D system of worked example (a); check with finite differences.
2. implicit_vjp: the adjoint solve with scipy.sparse.linalg.gmres and CG on a LinearOperator, for a system with 200 unknowns.
3. fixed_point and fixed_point_vjp for a small DEQ-like layer; compare with the unrolled gradient for $K = 5, 20, 100$.
4. ridge_hypergradient on the diabetes data set (scikit-learn), then gradient descent on $\log\lambda$.
5. The gradient of a softmax-regression solution with respect to its L2 weight, from the optimality condition.

11. Watch and read #

Implicit differentiation, what's going on here? | Chapter 6, Essence of calculus3Blue1Brown · 16 min

The one-dimensional picture of section 2: move along the curve so that F stays zero.

Efficient and Modular Implicit Differentiation (Machine Learning Research Paper Explained)Yannic Kilcher · 33 min

Roots, fixed points and optimality conditions as one framework, and the accuracy result of section 3.3.

Equilibrium approaches to deep learning: One (implicit) layer is all you needZico Kolter · Fields Institute · 74 min

Deep equilibrium models, fixed points and implicit backpropagation with constant memory.

Adjoint Equation of a Linear System of Equations - by implicit derivativeMachine Learning & Simulation · 29 min

The adjoint solve of section 3.2 for $Ax = b$, step by step.

• Kolter, Duvenaud and Johnson, Deep Implicit Layers (NeurIPS 2020 tutorial), chapter 2: naive differentiation through solvers, the implicit function theorem, and fixed-point JVPs and VJPs.
• Blondel et al. (2022), Efficient and Modular Implicit Differentiation: §2 (roots, fixed points, optimality conditions; one linear solve), Theorem 1 (error at most $C\|\hat x - x^*\|$).
• Blondel and Roulet, The Elements of Differentiable Programming (free), chapter 11: bilevel optimization, envelope theorems, the implicit function theorem and its proof, adjoint state method.
• Bai, Kolter and Koltun (2019), Deep Equilibrium Models.
• Lorraine, Vicol and Duvenaud (2020), Optimizing Millions of Hyperparameters by Implicit Differentiation.
• Rajeswaran et al. (2019), Meta-Learning with Implicit Gradients.
• Amos and Kolter (2017), OptNet: Differentiable Optimization as a Layer in Neural Networks.