Bilevel learning, hypergradients, and implicit layers

Let $p,d\ge 1$, let $\mathrm{\Theta }\subseteq {\mathbb{R}}^p$ and $W\subseteq {\mathbb{R}}^d$ be open, and let $f:\mathrm{\Theta }\times W\to \mathbb{R}$ be $C^2$ and $g:W\times \mathrm{\Theta }\to \mathbb{R}$ be $C^1$. Here $C^1$ means continuously differentiable on the open domain, and $C^2$ means that the derivative map is itself continuously differentiable. Gradients are column vectors, vector norms are Euclidean, and the norm of a matrix or linear map is the induced Euclidean operator norm.

For $(\theta ,w)\in \mathrm{\Theta }\times W$, define the inner stationarity residual $F(\theta ,w):={\mathrm{\nabla }}_{w}f(\theta ,w)\in {\mathbb{R}}^d$, its state Jacobian $H(\theta ,w):=D_{w}F(\theta ,w)\in {\mathbb{R}}^{d\times d}$, and its parameter Jacobian $B(\theta ,w):=D_{\theta }F(\theta ,w)\in {\mathbb{R}}^{d\times p}$. Thus $H$ is the inner Hessian and $B$ is the mixed derivative of the inner gradient. A point $({\theta }_0,w_0)$ is stationary when $F({\theta }_0,w_0)=0$. If a $C^1$ map $w^{\ast }:P\to Q$, for open neighborhoods $P\subseteq \mathrm{\Theta }$ of ${\theta }_0$ and $Q\subseteq W$ of $w_0$, satisfies $w^{\ast }({\theta }_0)=w_0$ and $F(\theta ,w^{\ast }(\theta ))=0$ for all $\theta \in P$, it is a local stationary branch. The associated outer objective is $J(\theta ):=g(w^{\ast }(\theta ),\theta )$, and its column gradient ${\mathrm{\nabla }}_{\theta }J$ is called the hypergradient of that branch. The definition of a local stationary branch imposes no global-minimality requirement.

For $\mu >0$, the phrase globally $\mu$-strongly convex in $w$ means that $f(\theta ,\cdot )$ is defined on all of ${\mathbb{R}}^d$ and, for every $u,v\in {\mathbb{R}}^d$, satisfies

$f(\theta ,v)\ge f(\theta ,u)+{\mathrm{\nabla }}_{w}f(\theta ,u{)}^T(v-u)+\frac{\mu }{2}\mathrm{\parallel }v-u{\mathrm{\parallel }}_2^2\mathrm{.}$

Continuously differentiable maps on open Euclidean domains

The C2 derivative-map convention for finite computational graphs

Input positive dimensions $p,d$, open sets $\mathrm{\Theta }\subseteq {\mathbb{R}}^p$ and $W\subseteq {\mathbb{R}}^d$, $C^2$ and $C^1$ scalar functions $f:\mathrm{\Theta }\times W\to \mathbb{R}$ and $g:W\times \mathrm{\Theta }\to \mathbb{R}$, a point $(\theta ,w)\in \mathrm{\Theta }\times W$, and finite real-arithmetic computational graphs whose outputs equal $f$ and $g$ on open neighborhoods of the evaluated inputs. The graph primitives for $f$ are $C^2$ and those for $g$ are $C^1$. Here $C^1$ means total differentiability with continuous derivative, $C^2$ means the derivative map is $C^1$, and all graph evaluations must remain on one unique strict branch. Define $F={\mathrm{\nabla }}_{w}f$, $H=D_{w}F$, $B=D_{\theta }F$, $c={\mathrm{\partial }}_{\theta }g$ with $w$ held fixed, and $d={\mathrm{\nabla }}_{w}g$; all gradients are column vectors. Also input a tolerance $\tau \ge 0$, a positive integer solve budget $N$, a number $\mu >0$, and a valid certificate that the evaluated inner Hessian is symmetric and satisfies $x^{T}Hx\ge \mu \mathrm{\parallel }x{\mathrm{\parallel }}_2^2$ for every $x\in {\mathbb{R}}^d$. Finally input a deterministic solver $S$ with the following contract: on a real symmetric positive-definite linear operator and a right-hand side, $S$ terminates within $N$ iterations and returns either a failure status or a vector $\tilde{v}$ and a report saying whether its own stopping test passed. The algorithm does not trust that report as an accuracy certificate; it recomputes the residual. All derivative and solver operations in this definition use exact real arithmetic.

If any graph evaluation is outside its declared domain, return DOMAIN_ERROR. If any piecewise primitive is at a branch boundary or has no unique strict active branch, return NONDIFFERENTIABLE and no gradient. Otherwise, use reverse mode on $f$ and take the state block of its returned gradient as the inner residual $r_{\mathrm{i}\mathrm{n}}:={\mathrm{\nabla }}_{w}f(\theta ,w)$. Use reverse mode on $g$ and take its parameter and state blocks as $c:={\mathrm{\partial }}_{\theta }g(w,\theta )$ and $d:={\mathrm{\nabla }}_{w}g(w,\theta )$. In these expressions $c$ is the derivative with $w$ held fixed, and $d$ and $c$ are column vectors.

For any $u\in {\mathbb{R}}^d$, apply forward-over-reverse Hessian-vector differentiation to $f$ at $(\theta ,w)$ in direction $(0,u)\in {\mathbb{R}}^p\times {\mathbb{R}}^d$. Let ${\pi }_{\theta }$ and ${\pi }_w$ denote projection onto the parameter and state coordinates of the returned vector. Define $M_{H}u:={\pi }_w(D(\mathrm{\nabla }f)(\theta ,w)(0,u))$ and $M_{B}u:={\pi }_{\theta }(D(\mathrm{\nabla }f)(\theta ,w)(0,u))$. Hessian symmetry makes $M_{H}u=Hu=H^{T}u$ and $M_{B}u=B^{T}u$.

Call $S$ on $M_H$ and right-hand side $d$ with tolerance $\tau$ and budget $N$. If $S$ returns failure, return $(r_{\mathrm{i}\mathrm{n}},\text{SOLVER\_FAILURE})$ with no hypergradient. If it returns a vector $\tilde{v}$, compute the actual linear residual $\rho :=d-M_H\tilde{v}$ and the reported hypergradient estimate $\widehat{h}:=c-M_B\tilde{v}$. Return $(\widehat{h},r_{\mathrm{i}\mathrm{n}},\rho ,\text{CONVERGED})$ exactly when the solver reports its stopping test passed and $\mathrm{\parallel }\rho {\mathrm{\parallel }}_2\le \tau$; otherwise return the same numerical values with status TOLERANCE_OR_BUDGET_NOT_MET. The procedure makes no floating-point accuracy claim.

Stationarity residual, Hessian and mixed derivative, outer objective, and column-gradient convention

Finite C1 and C2 real-arithmetic graphs and strict-branch convention

Reverse-mode gradient procedure

Exact reverse-mode gradient on a fixed strict branch

Forward-over-reverse Hessian-vector procedure

Exact Hessian-vector products on a fixed strict branch

Coordinate definition by continuous iterated partial derivatives through order two

Coordinate partials of a totally differentiable map

Symmetry of the C2 Hessian

The solver is a supplied finite procedure, so its termination contract is an input requirement. The algorithm verifies the residual even when the solver reports success. Its estimate is still returned with a non-success status when the solver produces a vector whose residual misses the requested tolerance.

Under this algorithm's derivative-map convention, $Df$ is a $C^1$ map. Apply the total-derivative theorem to the coordinate representation of $Df$: each coordinate partial of $Df$ is the corresponding entry of $D(Df)$ and varies continuously. The coordinates of $Df$ are the first partials of $f$, so this supplies the continuous second coordinate partials required by the coordinate-$C^2$ definition and the Hessian-symmetry corollary. That symmetry applies to the full Hessian of $f$, hence its state block is $H=H^T$ and its parameter block in the state direction is $B^T$, as used above.

The solid arrows show the vector data used for the adjoint estimate: the state derivative of the outer loss is solved against the transposed inner Hessian, then the mixed derivative maps the adjoint vector back to parameter space, where it is subtracted from the direct parameter derivative. The dashed arrows show the independently recomputed linear residual. The state-space vectors have dimension $d$ and the parameter-space vectors have dimension $p$. This diagram summarizes the data flow; the derivative and error proofs do not rely on it.

Let $p,d\ge 1$, let $\mathrm{\Theta }\subseteq {\mathbb{R}}^p$ and $W\subseteq {\mathbb{R}}^d$ be open, and let $f:\mathrm{\Theta }\times W\to \mathbb{R}$ be $C^2$ and $g:W\times \mathrm{\Theta }\to \mathbb{R}$ be $C^1$. Here $C^1$ means total differentiability with continuous derivative, and $C^2$ means the derivative map is $C^1$. Define $F(\theta ,w)={\mathrm{\nabla }}_{w}f(\theta ,w)$, $H(\theta ,w)=D_{w}F(\theta ,w)$, and $B(\theta ,w)=D_{\theta }F(\theta ,w)$. Suppose $F({\theta }_0,w_0)=0$ and $H({\theta }_0,w_0)$ is invertible. Then there are open neighborhoods $P\subseteq \mathrm{\Theta }$ and $Q\subseteq W$ of ${\theta }_0$ and $w_0$ and a unique $C^1$ local stationary branch $w^{\ast }:P\to Q$ through $({\theta }_0,w_0)$, meaning $w^{\ast }({\theta }_0)=w_0$ and $F(\theta ,w^{\ast }(\theta ))=0$ for every $\theta \in P$. After shrinking $P$ if necessary, $H(\theta ,w^{\ast }(\theta ))$ is invertible and

$D{w}^{\ast }(\theta )=-H(\theta ,w^{\ast }(\theta ){)}^{-1}B(\theta ,w^{\ast }(\theta ))\mathrm{.}$

Write $H^{-T}:=(H^{-1}{)}^T$. For $J(\theta )=g(w^{\ast }(\theta ),\theta )$, where ${\mathrm{\partial }}_{\theta }g$ holds $w$ fixed and gradients are column vectors, the hypergradient is

${\mathrm{\nabla }}_{\theta }J(\theta )={\mathrm{\partial }}_{\theta }g(w^{\ast }(\theta ),\theta )-B(\theta ,w^{\ast }(\theta ){)}^{T}H(\theta ,w^{\ast }(\theta ){)}^{-T}{\mathrm{\nabla }}_{w}g(w^{\ast }(\theta ),\theta )\mathrm{.}$

The local branch need not be a global inner optimizer. In the separate global regime where, for every $\theta \in P$, $f(\theta ,\cdot )$ is defined on all of ${\mathbb{R}}^d$ and, for some $\mu >0$, satisfies $f(\theta ,v)\ge f(\theta ,u)+{\mathrm{\nabla }}_{w}f(\theta ,u{)}^T(v-u)+\frac{\mu }{2}\mathrm{\parallel }v-u{\mathrm{\parallel }}_2^2$ for every $u,v\in {\mathbb{R}}^d$, each stationary point is the unique global minimizer.

Bilevel stationary-problem notation and strong-convexity convention

Local branch existence, uniqueness, and derivative formula

Finite-dimensional chain rule

Let $p,d\ge 1$, let $\mathrm{\Theta }\subseteq {\mathbb{R}}^p$ and $W\subseteq {\mathbb{R}}^d$ be open, and let $f:\mathrm{\Theta }\times W\to \mathbb{R}$ be $C^2$ and $g:W\times \mathrm{\Theta }\to \mathbb{R}$ be $C^1$, where $C^1$ means total differentiability with continuous derivative and $C^2$ means the derivative map is $C^1$. Let $w^{\ast }:P\to W$ be a $C^1$ map on an open neighborhood $P$ of ${\theta }_0$, with $w^{\ast }({\theta }_0)=w_0$ and $F(\theta ,w^{\ast }(\theta ))=0$ for $F(\theta ,w):={\mathrm{\nabla }}_{w}f(\theta ,w)$, and define $J(\theta ):=g(w^{\ast }(\theta ),\theta )$, $H(\theta ,w):=D_{w}F(\theta ,w)$, and $B(\theta ,w):=D_{\theta }F(\theta ,w)$. Assume $H({\theta }_0,w_0)$ is invertible. Write $c(w):={\mathrm{\partial }}_{\theta }g(w,{\theta }_0)$ and $d(w):={\mathrm{\nabla }}_{w}g(w,{\theta }_0)$, and use the induced Euclidean operator norm for matrices. If the adjoint algorithm at $w_0$ returns a vector $\tilde{v}$, define $\rho :=d(w_0)-H({\theta }_0,w_0{)}^T\tilde{v}$ and $\widehat{h}:=c(w_0)-B({\theta }_0,w_0{)}^T\tilde{v}$. Suppose the evaluated Hessian is symmetric and has a valid lower-bound certificate $x^{T}H({\theta }_0,w_0)x\ge \mu \mathrm{\parallel }x{\mathrm{\parallel }}_2^2$ for every $x$, for some $\mu >0$. Then

$\mathrm{\parallel }\widehat{h}-{\mathrm{\nabla }}_{\theta }J({\theta }_0){\mathrm{\parallel }}_2\le \frac{\mathrm{\parallel }B({\theta }_0,w^{\ast }({\theta }_0))\mathrm{\parallel }\mathrm{\parallel }\rho {\mathrm{\parallel }}_2}{\mu }\mathrm{.}$

There is also a separate approximate-state bound under stronger hypotheses. Put $R(w,\theta ):={\mathrm{\nabla }}_{w}f(\theta ,w)$ and $A(w):=D_{w}R(w,{\theta }_0)=H({\theta }_0,w)$. Suppose $R$ is $C^2$ on an open neighborhood of $(w_0,{\theta }_0)$. Let $\tilde{w}$ lie in a $C^1$ local inverse chart: for $r:=R(\tilde{w},{\theta }_0)$, there is a local inverse $\mathrm{\Psi }$ of $w\mapsto R(w,{\theta }_0)$ with $\mathrm{\Psi }(0)=w_0$, defined on a convex open residual neighborhood containing $0$ and $r$, with $\tilde{w}=\mathrm{\Psi }(r)$ and $\mathrm{\parallel }D\mathrm{\Psi }(y)\mathrm{\parallel }\le \kappa$ along the segment from $0$ to $r$. Suppose the segment from $w_0$ to $\tilde{w}$ lies in a convex state neighborhood where $\mathrm{\parallel }A(w{)}^{-1}\mathrm{\parallel }\le M$, $\mathrm{\parallel }B({\theta }_0,w)\mathrm{\parallel }\le B_0$, and $A(w),B({\theta }_0,w)$ are Lipschitz with constants $L_A,L_B$. Suppose also that $c,d$ are Lipschitz there with constants $L_c,L_d$ and $\mathrm{\parallel }d(w_0){\mathrm{\parallel }}_2\le D_{\ast }$. If the adjoint algorithm at $\tilde{w}$ returns $\tilde{v}$, set $\rho :=d(\tilde{w})-A(\tilde{w}{)}^T\tilde{v}$ and $\widehat{h}:=c(\tilde{w})-B({\theta }_0,\tilde{w}{)}^T\tilde{v}$. Assume the certificate $A(\tilde{w})=A(\tilde{w}{)}^T$ and $x^{T}A(\tilde{w})x\ge \mu \mathrm{\parallel }x{\mathrm{\parallel }}_2^2$ for all $x$, with the same $\mu >0$. Then

$\mathrm{\parallel }\widehat{h}-{\mathrm{\nabla }}_{\theta }J({\theta }_0){\mathrm{\parallel }}_2\le \kappa (L_c+D_{\ast }(M^2{L}_A{B}_0+M{L}_B)+M{B}_0{L}_d)\mathrm{\parallel }r{\mathrm{\parallel }}_2+\frac{B_0}{\mu }\mathrm{\parallel }\rho {\mathrm{\parallel }}_2\mathrm{.}$

The first term bounds the implicit formula evaluated at an approximate state; it does not identify that formula with the derivative of a stopped iteration. The second term bounds the separate linear-solve error.

Definitions of F, H, B, J, and the hypergradient

The adjoint algorithm and its measured residual

Full-chart state and sensitivity perturbation bound

Derivative of the residual-chart path

Uniform derivative bound along a path controls its endpoint change

Exact local-branch hypergradient

Euclidean Cauchy–Schwarz inequality

Let $p,m\ge 0$, let $\mathrm{\Theta }\subseteq {\mathbb{R}}^p$ be open with ${\theta }_0\in \mathrm{\Theta }$, let $D\subseteq {\mathbb{R}}^m$ be an open convex neighborhood of $z_{\ast }$, and let $G:U\to {\mathbb{R}}^m$ be $C^2$ on an open neighborhood $U$ of $D\times \{{\theta }_0\}$. Here $C^1$ means total differentiability with continuous derivative, $C^2$ means the derivative map is $C^1$, and vector and operator norms are Euclidean and induced Euclidean norms. Write $G_z:=D_{z}G$ and $G_{\theta }:=D_{\theta }G$ for the state and parameter derivatives. Suppose $z_{\ast }=G(z_{\ast },{\theta }_0)$, $G(D,{\theta }_0)\subseteq D$, $0\le q<1$, $\mathrm{\parallel }{G}_z(z,{\theta }_0)\mathrm{\parallel }\le q$ on $D$, and $G_z,G_{\theta }$ are Lipschitz in $z$ with constants $L_z,L_{\theta }$. Let $z_0:\mathrm{\Theta }\to {\mathbb{R}}^m$ be $C^1$, with all finite iterates at ${\theta }_0$ lying in $D$, and define $z_{k+1}(\theta )=G(z_k(\theta ),\theta )$. Let $\varphi$ be the exact local fixed-point branch, $s_{\ast }:=D\varphi ({\theta }_0)$, $e_K:=D{z}_K({\theta }_0)$, ${\delta }_0:=\mathrm{\parallel }{z}_0({\theta }_0)-z_{\ast }{\mathrm{\parallel }}_2$, $d_0:=\mathrm{\parallel }D{z}_0({\theta }_0)-s_{\ast }\mathrm{\parallel }$, and $C:=L_z\mathrm{\parallel }{s}_{\ast }\mathrm{\parallel }+L_{\theta }$. For every integer $K\ge 1$,

$\mathrm{\parallel }{e}_K-s_{\ast }\mathrm{\parallel }\le q^K{d}_0+CK{q}^{K-1}{\delta }_0,$

where $q^0=1$ when $q=0$ and $K=1$. For the fixed-point residual $r_K:=z_K({\theta }_0)-G(z_K({\theta }_0),{\theta }_0)$, the primal-state estimate is $(1-q)\mathrm{\parallel }{z}_K({\theta }_0)-z_{\ast }{\mathrm{\parallel }}_2\le \mathrm{\parallel }{r}_K{\mathrm{\parallel }}_2$. This primal residual estimate alone does not imply that the derivative of the finite $K$-step graph equals the exact fixed-point sensitivity.

For a separate finite shared-parameter recurrence, let $p_0,q_0\ge 1$, let ${\mathrm{\Theta }}_0\subseteq {\mathbb{R}}^{p_0}$ and $H\subseteq {\mathbb{R}}^{q_0}$ be open, let $T\ge 0$, let $h_0:{\mathrm{\Theta }}_0\to H$ and ${\mathrm{\Phi }}_t:{\mathrm{\Theta }}_0\times H\to H$ for $1\le t\le T$ and $\mathrm{\ell }:{\mathrm{\Theta }}_0\times H\to \mathbb{R}$ be $C^1$ on open domains containing the trajectory at $\theta \in {\mathrm{\Theta }}_0$, and suppose every piecewise primitive is evaluated in one unique strict open $C^1$ branch. Set $h_t={\mathrm{\Phi }}_t(\theta ,h_{t-1})$ and $J_T(\theta )=\mathrm{\ell }(\theta ,h_T)$; derivative matrices use Euclidean coordinates and transposes are Euclidean. Backpropagation differentiates this finite scalar graph: with $A_t=D_{\theta }{\mathrm{\Phi }}_t(\theta ,h_{t-1})$, $B_t=D_h{\mathrm{\Phi }}_t(\theta ,h_{t-1})$, ${\lambda }_T=(D_h\mathrm{\ell }(\theta ,h_T){)}^T$, and ${\lambda }_{t-1}=B_t^T{\lambda }_t$, its derivative is

${\mathrm{\nabla }}_{\theta }{J}_T(\theta )=(D_{\theta }\mathrm{\ell }(\theta ,h_T){)}^T+{\sum }_{t=1}^T{A}_t^T{\lambda }_t+(D{h}_0(\theta ){)}^T{\lambda }_0\mathrm{.}$

This is the derivative of the finite training computation. The exact implicit hypergradient is the derivative of a local stationary or fixed-point branch. In particular, a small training residual is not by itself an error bound for the derivative of a stopped iteration; contraction and sensitivity hypotheses such as those stated above are needed.

Finite fixed-point sensitivity and residual bounds

Gradient of a finite shared-parameter recurrence

Adjoint solve and measured residual

Exact local stationary-branch hypergradient