△ Mechanism confirmed, baseline not beaten 2026

Matrix-Free Krylov Backpropagation Through Solver Layers

Implementation & benchmark of arXiv:2608.16280 — PANDA: A Matrix-Free Differentiable NMPC Solver via Proximal Averaged Quasi-Newton with Adaptive Linesearch Algorithm

Usefulness8/10

Difficulty6/10

Novelty5/10

Training Memory Speedup Stability Optimizer Training dynamics Mlp Ssm Transformer Optimization Linear algebra Numerical analysis Dynamical systems

Source paper: PANDA: A Matrix-Free Differentiable NMPC Solver via Proximal Averaged Quasi-Newton with Adaptive Linesearch Algorithm arXiv:2608.16280 ⓘ · analyzed Sep 1, 2026

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

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.

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.$$

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.

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

Mechanism confirmed, baseline not beaten

Stage 1 · Toy mechanism gate: Passed ✓

Stage 2 · Mechanism transferred to benchmark: Confirmed ✓

Stage 2 · Practical benchmark result: Inconclusive

Methodology: 8 paired seeds, permutation test (p<0.05) against a learning-rate-tuned baseline. How verification works

Stage 1 — Mechanism check agent confidence 9/10

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.

Agent confidence: 9/10

Baseline: 18-step Richardson unrolling: final loss 8.3217648; initial loss 15.0569419
Idea: Matrix-free CG implicit adjoint: final loss 8.3020707; initial loss 15.0569419; gradient relative errors at tolerances 1e-2, 1e-4, 1e-8 were 7.27e-3, 5.93e-5, 4.55e-12; VJP absolute error 0

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.

How to run: python3 krylov_experiment.py

Stage 2 — Benchmark vs. tuned baseline (arXiv:2608.16280) Sep 2, 2026

Ahead of the baseline (−4.3%) but not statistically significant (p=0.1673, wins 5 / 8 paired seeds).

Benchmark: Tabular regression (Friedman)
Model: mlp_tiny
Paired seeds: 8
Baseline mean: 5.5859
Idea mean: 5.3442
Effect (Δ): -0.2417 (−4.3%; negative = idea better)
Wins: 5 / 8 paired seeds
p-value: 0.1673 (permutation test, 20 000 shuffles)
Smallest detectable effect: ±7.7%
Mechanism: Confirmed ✓
Practical effect: Inconclusive
Baseline tuning: swept over 3 configs

Limitations:

Only the tabular track was tested; vision, sequence, and dynamics transfer were not tested. The solver residual used a fixed 16-dimensional SPD operator rather than a learned nonlinear equilibrium Jacobian, and peak memory and wall-clock scaling versus substantially longer unroll horizons were not benchmarked. The benchmark used 8 seeds and 18 epochs, so the modest MSE advantage lacks statistical power under the required test.

How to run: python3 krylov_bench.py > bench_output.json

Verdict computed by deterministic test code from paired-seed statistics — not by the language model.

Stage 2 — Benchmark Sep 1, 2026 latest Failed ✗

Agent confidence: 9/10

Implemented a matched solver-layer comparison on the built-in Friedman#1 tabular track using the same MLP, solver horizon, Adam settings, and shared learning-rate sweep. Matrix-free CG implicit backprop reduced mean test MSE from 5.5859 to 5.3442 (paired delta -0.2417; 5/8 wins), but the permutation p-value was 0.16725, so this is not a significant benchmark win. The trained-model mechanism signature confirmed the predicted SPD CG energy-error bound at all tested iterations, but mechanism confirmation does not override the required task-metric significance criterion.

Baseline: Friedman#1 tabular, mlp_tiny, 18 epochs, 12 solver steps, lr=0.01 selected by 3-value baseline sweep; 8-seed mean test MSE 5.585879, std 0.798322
Idea: Same task/model/solver budget, matrix-free autodiff CG adjoint, lr=0.01 selected from the identical learning-rate union; 8-seed mean test MSE 5.344174, std 0.667608
Idea wins: 5/8 paired seeds
Mechanism signature: For trained seed-0 model with estimated condition number 16, observed relative energy errors at CG iterations 1,2,4,8 were 0.4988, 0.2889, 0.0970, 0.00916 versus bounds 1.2, 0.72, 0.2592, 0.03359; final relative residual 5.53e-05; confirmed=true
Paired delta: -0.241705 test MSE (idea - baseline)
Permutation p value: 0.16725

Benchmark result