Backpropagation by the Implicit Function Theorem

model derivation

Derivation. The full forward (jvp) and adjoint (vjp) equations for an algebraic factor. Cost: one $k\times k$ linear solve. The parameter cotangent turns out to be rank one.

Sources: original to this vault (design and analysis; no single paper).

Theory (CT-ML wiki): Lens

Setup

$r_{\Theta }(x_o,x_u)=\Theta v_d(x_o,x_u)=0$, with $x_o\in {\mathbb{R}}^p$ clamped, $x_u\in {\mathbb{R}}^q$ inferred, $r\in {\mathbb{R}}^k$, $\Theta \in {\mathbb{R}}^{k\times m}$. Assume the well-posed case $k=q$ and $J_u$ invertible at the solution (Inference as Root Finding; failure is Branches and the Discriminant).

Write, using The Veronese Parametrisation §“Jacobians, exactly”,

$$J_o=\Theta D{v}_d(x{)}_{[:,o]}\in {\mathbb{R}}^{k\times p},J_u=\Theta D{v}_d(x{)}_{[:,u]}\in {\mathbb{R}}^{k\times q},v:=v_d(x)\in {\mathbb{R}}^m$$

All three are exact; no automatic differentiation is involved anywhere.

Total differential

Differentiating the constraint $r_{\Theta }(x_o,x_u)=0$, which holds identically in a neighbourhood:

$J_o\mathrm{d}{x}_o+J_u\mathrm{d}{x}_u+(\mathrm{d}\Theta )v=0$

(the third term because $r=\Theta v$ is linear in $\Theta$, so its differential in $\Theta$ is just $(\mathrm{d}\Theta )v$). Hence the implicit function theorem:

$\boxed{\mathrm{d}{x}_u=-J_u^{-1}(J_o\mathrm{d}{x}_o+(\mathrm{d}\Theta )v)}$

giving the two sensitivities

$$\frac{\partial x_u}{\partial x_o}=-J_u^{-1}{J}_o\in {\mathbb{R}}^{q\times p},\frac{\partial x_u}{\partial {\Theta }_{j\ell }}=-J_u^{-1}{e}_j{v}_{\ell }$$

Forward mode (jvp)

Given tangents $({\dot{x}}_o,\dot{\Theta })$:

${\dot{x}}_u=-J_u^{-1}(J_o{\dot{x}}_o+\dot{\Theta }v)$

One solve. Cost $O(k^3)$ for the factorisation, $O(k^2)$ per additional tangent — so a factorisation of $J_u$ is reused across the whole batch of tangents.

Reverse mode (vjp) — the adjoint

This is what the put of a lens needs. Given a cotangent ${\bar{x}}_u\in {\mathbb{R}}^q$ arriving from downstream, we want ${\bar{x}}_o$ and $\bar{\Theta }$ satisfying $⟨{\bar{x}}_u,{\dot{x}}_u⟩=⟨{\bar{x}}_o,{\dot{x}}_o⟩+⟨\bar{\Theta },\dot{\Theta }{⟩}_F$ for all tangents.

Substitute:

$$⟨{\bar{x}}_u,{\dot{x}}_u⟩=-{\bar{x}}_u^{\top }{J}_u^{-1}(J_o{\dot{x}}_o+\dot{\Theta }v)$$

Define the adjoint variable

$\boxed{\lambda =-J_u^{-\top }{\bar{x}}_u\in {\mathbb{R}}^k}$

— one linear solve with the transpose. Then

$$⟨{\bar{x}}_u,{\dot{x}}_u⟩={\lambda }^{\top }{J}_o{\dot{x}}_o+{\lambda }^{\top }\dot{\Theta }v=⟨J_o^{\top }\lambda ,{\dot{x}}_o⟩+⟨\lambda v^{\top },\dot{\Theta }{⟩}_F$$

using ${\lambda }^{\top }\dot{\Theta }v=\mathrm{tr}(v{\lambda }^{\top }\dot{\Theta })=⟨\lambda v^{\top },\dot{\Theta }{⟩}_F$. Therefore

$\boxed{{\bar{x}}_o=J_o^{\top }\lambda ,\bar{\Theta }=\lambda v_d(x{)}^{\top }}$

Three things to notice

1. The parameter cotangent is rank one

$\bar{\Theta }=\lambda v^{\top }$ is an outer product. Every data point contributes exactly a rank-one update to the parameter, and a minibatch of size $B$ contributes rank $\le B$.

Consequences:

• The Grassmannian update of The Parameter is a Grassmannian is a rank-$B$ perturbation of a $k$-plane — a subspace-tracking problem, costing $O(Bkm)$ rather than $O(k{m}^2)$.
• Storage of the accumulated gradient can be $O(B(k+m))$ rather than $O(km)$.
• The Gauss–Newton matrix ${\sum }_i{J}_i^{\top }{J}_i$ inherits the same low-rank structure.

This is a real structural gift and it comes directly from the linearity in $\Theta$.

2. The cost is one $k\times k$ solve, shared

$J_u$ is factorised once during inference (Newton needs it anyway) and reused for the adjoint. So the backward pass is essentially free given the forward pass — one triangular solve. Compare an unrolled equilibrium solver, which pays $O(\text{iterations})$ in memory and time.

3. Where it breaks

Exactly on the discriminant, where $J_u$ is singular. The conditioning $\kappa (J_u)$ is a cheap, exact early-warning signal; see Branches and the Discriminant. Near $\Delta$ the honest options are to regularise the solve ($\lambda =-(J_u^{\top }{J}_u+\epsilon I{)}^{-1}{J}_u^{\top }{\bar{x}}_u$, i.e. a Levenberg–Marquardt/Tikhonov damping) and report that you did, or to refuse.

The overdetermined / minimising case

When $k>q$ there is generically no root and inference minimises $\sigma (r)=\frac{1}{2}\Vert r{\Vert }_M^2$ instead. The stationarity condition is $J_u^{\top }Mr=0$, and differentiating that gives the sensitivity

$\frac{\partial x_u}{\partial x_o}=-(J_u^{\top }M{J}_u+{\sum }_i(Mr{)}_i{\nabla }_{x_u}^2{r}_i{)}^{-1}{J}_u^{\top }M{J}_o$

The bracketed matrix is the full Hessian; dropping the second term gives the Gauss–Newton approximation $(J_u^{\top }M{J}_u{)}^{-1}$, which is exact when $r=0$ and cheap always. Note that the second-derivative term is also exactly computable here (the ${\nabla }^2{r}_i$ are again $\Theta$ times a fixed monomial Hessian), so unlike in neural settings the Gauss–Newton approximation is a choice, not a necessity.

Why this is not “differential algebra”

The README’s table lists the implicit backward pass as “implicit function theorem / differential algebra”. Everything above is the implicit function theorem — ordinary multivariable calculus. No differential algebra appears, and none is needed. Differential algebra does have genuine roles here, but they are different jobs:

job	tool	note
derivative of the solve	implicit function theorem	this note
is this polarity well-posed at all?	elimination / dimension theory	Composition is Elimination
composing two factors into one	Gröbner / border bases	Composition is Elimination
factors that are DAEs, $r(x,\dot{x})=0$	differential elimination, index reduction	Differential Algebra and DAE Factors

The last row is where differential algebra genuinely enters a backward pass, and even there it is index reduction of the adjoint DAE rather than differentiation as such.

Related: Inference as Root Finding, The Parameter is a Grassmannian, Lens, Differential Algebra and DAE Factors