Adjoint hypergradient with a measured linear-solve residual

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.

diagram failed to render:
\begin{tikzpicture}[>=stealth]
  \node[draw,align=center] (d) at (0,0) {$d=\nabla_w g$};
  \node[draw,align=center] (solve) at (3.4,0) {solve $H^T\widetilde v=d$};
  \node[draw,align=center] (cross) at (6.8,0) {$B^T\widetilde v$};
  \node[draw,align=center] (out) at (10.4,0) {$\widehat h=c-B^T\widetilde v$};
  \node[draw,align=center] (c) at (8.7,1.6) {$c=\partial_\theta g$};
  \node[draw,align=center] (resid) at (3.4,-1.7) {$\rho=d-H^T\widetilde v$};
  \draw[->] (d.east) -- (solve.west);
  \draw[->] (solve.east) -- node[above] {$\widetilde v$} (cross.west);
  \draw[->] (cross.east) -- node[above] {$-$} (out.west);
  \draw[->] (c.east) -- node[right] {$+$} (out.north);
  \draw[dashed,->] (d.south) to[bend right=18] (resid.west);
  \draw[dashed,->] (solve.south) -- (resid.north);
\end{tikzpicture}