# Matrix-Free Krylov Backpropagation Through Solver Layers

- ID: 2395
- Canonical URL: https://synthcore.org/idea/2395/matrix-free-krylov-backpropagation-through-solver-layers
- API JSON: https://synthcore.org/api/idea/2395.json
- API Markdown: https://synthcore.org/api/idea/2395.md
- Verification status: mech_ok_no_baseline
- Source: [arXiv:2608.16280](https://arxiv.org/abs/2608.16280)
- Category: training
- Solves: memory, speedup, stability
- ML areas: optimizer, training-dynamics, mlp, ssm, transformer
- Math tags: optimization, linear-algebra, numerical-analysis, dynamical-systems
- Ratings: usefulness 8/10; difficulty 6/10; novelty 5/10

## Idea description

Turn an iterative optimization or equilibrium computation inside a neural network into a differentiable layer whose backward pass solves the implicit adjoint system with conjugate gradients or GMRES using only automatic-differentiation matrix-vector products. This avoids storing unrolled iterations and avoids explicit Hessian or Jacobian construction, enabling longer solver horizons and lower-memory implicit architectures.

## Mathematical statement

Let $z^*(\theta)$ be the converged output of a solver layer defined by the residual equation $r(z,\theta)=0$, where $z$ may be an equilibrium state, optimized latent code, or finite-horizon control trajectory. For an outer loss $L(z^*(\theta),\theta)$, define $A=\partial r/\partial z$ and $b=\partial L/\partial z$. The implicit adjoint $\lambda$ satisfies $A^T\lambda=b$, and the parameter gradient is $\mathrm{d}L/\mathrm{d}\theta=\partial_\theta L-\lambda^T\partial_\theta r$. The matrix-free mechanism evaluates $A^Tv$ through reverse-mode autodiff and solves the adjoint system in a Krylov subspace without constructing $A$. If $A$ is symmetric positive definite with condition number $\kappa$, conjugate-gradient error decays at the standard geometric rate.

## Key formulas

- $$r(z^*(\theta),\theta)=0,\qquad A^T\lambda=b,\qquad A=\frac{\partial r}{\partial z},\quad b=\frac{\partial L}{\partial z}.$$
- $$\frac{\mathrm dL}{\mathrm d\theta}=\frac{\partial L}{\partial\theta}-\lambda^T\frac{\partial r}{\partial\theta},\qquad \lambda=\operatorname{KrylovSolve}\left(v\mapsto A^Tv,\ b,\ \varepsilon\right).$$
- $$\|e_j\|_2=\|b-A^T\lambda_j\|_2\leq\varepsilon\|b\|_2,$$
- $$\|\lambda_j-\lambda^*\|_A\leq2\left(\frac{\sqrt\kappa-1}{\sqrt\kappa+1}\right)^j\|\lambda_0-\lambda^*\|_A\quad\text{for symmetric positive-definite }A.$$

## Implementation notes

1. Integration point: implement a PyTorch or JAX solver layer that approximately minimizes an inner objective, computes a deep-equilibrium fixed point, or performs $K$ proximal optimization steps. During the forward pass, retain only the final state $z^*$ and enough information to recompute the residual $r(z,\theta)$; do not retain the full unrolled graph. 2. Pseudocode: solve forward until $\|r(z,\theta)\|$ is below tolerance; detach $z$; form $b=\nabla_zL$; define $\operatorname{matvec}(v)=\operatorname{vjp}_z[r](v)=A^Tv$ using autodiff; run GMRES for nonsymmetric residuals or conjugate gradients when the residual Jacobian is symmetric positive definite; stop at $\|b-A^T\lambda\|\leq\varepsilon\|b\|$; obtain parameter gradients as $\nabla_\theta L-\operatorname{vjp}_\theta[r](\lambda)$. 3. Computed versus estimated: matrix-vector products are computed by autodiff up to floating-point error; $A$ and its spectrum are never materialized. Estimate conditioning from GMRES or CG residuals and use damping $A\leftarrow A+\delta I$ if residuals stagnate. 4. First cheap experiment: use a 20-dimensional learned equilibrium MLP or a 10-step proximal sparse-coding layer on MNIST, comparing full unrolling, explicit-Jacobian backpropagation, and matrix-free GMRES. Sweep forward iterations and backward tolerances $\varepsilon\in\{10^{-2},10^{-3},10^{-4}\}$. The prediction is geometric Krylov residual decay for a well-conditioned linearized residual, with SPD iteration count scaling approximately as $O(\sqrt\kappa\log(1/\varepsilon))$; gradients should differ from full backprop by relative error on the order of $\varepsilon$, while peak memory remains approximately independent of the number of forward iterations.

## Verification

- Status: mech_ok_no_baseline
- Mechanism evidence: yes
- Mechanism confirmed: yes
- Practical verdict: inconclusive
- Verdict: Built a matrix-free PyTorch CG implicit-adjoint prototype using autodiff VJPs without materializing the Jacobian. The mechanism manifested: the SPD energy error stayed below the predicted geometric bound for all 10 checked iterations, observed iterations increased with condition number (1, 17, 30, 45, 59 for κ=1,4,16,64,256), and gradient error tracked tolerance (0.0073, 5.9e-5, and 4.5e-12). The toy implicit training run reached loss 8.302 versus 8.322 for 18-step unrolling, but this is only a near-tie and not evidence of a meaningful practical speed or memory win.

### Mechanism check

- Verdict: Built a matrix-free PyTorch CG implicit-adjoint prototype using autodiff VJPs without materializing the Jacobian. The mechanism manifested: the SPD energy error stayed below the predicted geometric bound for all 10 checked iterations, observed iterations increased with condition number (1, 17, 30, 45, 59 for κ=1,4,16,64,256), and gradient error tracked tolerance (0.0073, 5.9e-5, and 4.5e-12). The toy implicit training run reached loss 8.302 versus 8.322 for 18-step unrolling, but this is only a near-tie and not evidence of a meaningful practical speed or memory win.
- Confidence: 9/10
- Limitations: Only a linear SPD toy residual was tested; nonsymmetric GMRES, nonlinear equilibrium MLPs, MNIST-scale data, explicit-Jacobian baselines, wall-clock speed, and peak-memory scaling were not measured. The condition-number tests used diagonal matrices, and the training comparison was too small to establish a practical performance win.

### Practical benchmark

- Paired seeds: 8
- Baseline mean: 0
- Idea mean: 0
- p-value: 0.16725
- Paired wins: 5/8
- Benchmark verdict: no significant win

## Artifacts

- [bench_output.json](https://synthcore.org/code/811/bench_output.json)
- [krylov_bench.py](https://synthcore.org/code/811/krylov_bench.py)
- [krylov_experiment.py](https://synthcore.org/code/811/krylov_experiment.py)
- [report.md](https://synthcore.org/code/811/report.md)
- [report_bench_2026-09-01T111856.md](https://synthcore.org/code/811/report_bench_2026-09-01T111856.md)
- [Download all files as ZIP](https://synthcore.org/download/811)

## Disclaimer

AI-generated research hypothesis, automatically tested. Not peer-reviewed.