Research map / Probabilistic and scientific modeling / Differential equations and PDEs
Gradients of Functions of Large Matrices
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Derives adjoint systems for the Lanczos and Arnoldi iterations and implements matrix-free reverse-mode gradients in JAX (matfree), showing competitiveness on GP model selection, PDE learning, and Laplace-approximated Bayesian neural networks.
Problem
Differentiating functions of large matrices (e.g. log-determinants, matrix exponentials) computed via matrix-free Lanczos/Arnoldi iterations, which previously lacked efficient, exact reverse-mode gradients for machine learning pipelines.
Approach
Treat the Lanczos and Arnoldi iterations as constrained dynamical systems and apply the adjoint method with Lagrange multipliers, yielding closed-form adjoint systems whose backward pass has the same linear time and memory complexity as the forward pass and uses only matrix-vector products.
Key results
- On five UCI GP datasets, the Lanczos-adjoint method matches GPyTorch RMSE (e.g. 0.09 vs 0.09 on elevators) and reaches lower training losses, but is roughly 20x slower per epoch.
- For PDE learning with a 250k-parameter MLP, Arnoldi adjoints achieve 7.7e-02 s/epoch versus 7.2e-02 for Dopri5 and 2.7e-01 for Tsit5, with comparable test loss (6.1e-03).
- For a Laplace-approximated VAN, Lanczos with 150 MVs gives marginal likelihood -154,475 versus -876,444 for the diagonal GGN approximation and AUROC 0.9 versus 0.5, though the diagonal has lower ECE (0.2 vs 0.5).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use Lanczos with adjoint gradients instead of GPyTorch gradients | exact Gaussian process model selection · elevators, protein, kegg_dir, kegg_undir | RMSE: 0.09 $\pm$ 0.002 vs 0.09 $\pm$ 0.003; 0.39 $\pm$ 0.005 vs 0.39 $\pm$ 0.005; 0.12 $\pm$ 0.002 vs 0.12 $\pm$ 0.005; 0.12 $ | No clear effect |
| use Lanczos with adjoint gradients instead of GPyTorch gradients | exact Gaussian process model selection · kin40k | RMSE: 0.10 $\pm$ 0.010 -> 0.12 $\pm$ 0.004 | Hurt |
| use Arnoldi adjoints instead of Diffrax Dopri5 | PDE coefficient-field learning · PDE coefficient-field learning | Runtime per epoch: 7.2e-02 $\pm$ 3.4e-05 vs 7.7e-02 $\pm$ 1.8e-05 | No clear effect |
| use Arnoldi adjoints instead of Diffrax Tsit5 | PDE coefficient-field learning · PDE coefficient-field learning | Runtime per epoch: 2.7e-01 $\pm$ 1.1e-05 -> 7.7e-02 $\pm$ 1.8e-05 | Helped |
| use Arnoldi adjoints instead of Diffrax solvers | PDE coefficient-field learning · PDE coefficient-field learning | Parameter RMSE: Arnoldi: 2.9e-04 $\pm$ 4.4e-05; Dopri5: 2.6e-04 $\pm$ 5.0e-05; Tsit5: 2.7e-04 $\pm$ 5.2e-05 | No clear effect |
| use Lanczos instead of a stochastic diagonal GGN approximation | Bayesian VAN calibration · Bayesian VAN calibration | Joint likelihood: test (log): -5,260.9 $\pm$ 290.2 -> -66.5 $\pm$ 11.2 | Helped |
| use Lanczos instead of a stochastic diagonal GGN approximation | Bayesian VAN calibration · Bayesian VAN calibration | Expected calibration error: 0.2 $\pm$ 0.003 -> 0.5 $\pm$ 0.01 | Hurt |
| use Lanczos instead of a stochastic diagonal GGN approximation | Bayesian VAN calibration · places365 out-of-distribution | AUROC (out-of-dist.): 0.5 $\pm$ 0.010 -> 0.9 $\pm$ 0.03 | Helped |
| increase Lanczos budget from 50 MVs to 150 MVs (ablation) | Bayesian VAN calibration · Bayesian VAN calibration | Marginal likelihood (log): -192,757 -> -154,475 | Helped |
Surprising findings
- Lanczos-based calibration has worse expected calibration error (0.5) than the diagonal approximation (0.2) despite better likelihood and AUROC.
- GPyTorch-style runtime is about 20x faster per epoch than the JAX adjoint implementation due to matrix-vector-product backends.
Limitations
- Assumes the Krylov basis Q fits into memory (full reorthogonalisation); partial reorthogonalisation gradients are left to future work.
- Only real-valued matrices are considered; forward-mode derivatives are not covered.
- Runtime in the GP case study is about 20x slower than GPyTorch because of matrix-vector-product backends.
Experimental setup
Datasets: UCI regression datasets (elevators, protein, kin40k, kegg_dir, kegg_undir) SuiteSparse bcsstk18 ImageNet places365
Benchmarks: UCI GP regression (RMSE, training loss, runtime) PDE coefficient learning (loss, parameter RMSE, runtime) Bayesian VAN calibration (marginal likelihood, ECE, AUROC)
Metrics: RMSE training loss runtime per epoch test loss parameter RMSE marginal likelihood expected calibration error AUROC
Baselines: GPyTorch Diffrax Dopri5 Diffrax Tsit5 diagonal GGN approximation
Builds on: adjoint method implicit differentiation stochastic trace estimation Lanczos iteration Arnoldi iteration
Closest papers by meaning
- Nearly Optimal Approximation of Matrix Functions by the Lanczos Method (NeurIPS 2024 · similarity 0.87)
- Linearization Turns Neural Operators into Function-Valued Gaussian Processes (ICML 2025 · similarity 0.87)
- Fast and Stable Gradient Approximation for Bilinear Forms of Hermitian Matrix Functions (NeurIPS 2026 · similarity 0.86)
- Standard Gaussian Process is All You Need for High-Dimensional Bayesian Optimization (ICLR 2025 · similarity 0.84)
- Bayes-optimal learning of an extensive-width neural network from quadratically many samples (NeurIPS 2024 · similarity 0.84)
- Scalable Derivative Gaussian Processes via Exact Gradient Reduction (NeurIPS 2026 · similarity 0.83)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.