Research map / Probabilistic and scientific modeling / Uncertainty and Gaussian processes
Affine Tracing: A New Paradigm for Probabilistic Linear Solvers
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.
Shows Bayesian probabilistic linear solvers are special cases of affine probabilistic iterative methods, proves affine PIMs are strongly calibrated, and introduces affine tracing to automatically build PIMs (including probabilistic multigrid) from standard iterative solver code with equality-saturation simplification for inverse priors.
Problem
Probabilistic linear solvers are split into Bayesian and non-Bayesian (PIM) camps; implementing calibrated, fast PIMs such as multigrid requires bespoke, error-prone derivations of posterior covariances and inverse-prior downdates.
Approach
Treat Bayesian PLSs as non-stationary affine PIMs, prove strong calibration for any affine fixed-point method, and automate lifting by tracing an affine iterative implementation into a computational graph, then using equality saturation (egglog) to cancel A^{-1} terms for the inverse-prior downdate.
Key results
- The affine-traced two-grid multigrid (AffineTracer) is consistently faster in wall-time than the hand-implemented ProbMG, with the gap growing with system dimension.
- On the synthetic CAGP problem, both multigrid CAGP variants (CAGP-MG-2, CAGP-MG-3) converge faster in RMSE and NLL against wall-time than CAGP-CG and CAGP-GS, despite higher per-iteration cost.
- On ERA5 regression with 36450 training points, CAGP-MG-2 has RMSE dropping rapidly and yields smoother posterior means than CAGP-CG and CAGP-GS, though it appears to stagnate after a few iterations.
Surprising findings
- Bayesian PLSs with Gaussian iterates are exactly non-stationary affine PIMs, so the Bayesian/non-Bayesian dichotomy is false.
- The downdate structure (posterior covariance equal to inverse prior minus a downdate) emerges for any fixed-point affine iterative method.
Limitations
- Equality saturation approach found brittle; whole-graph simplification is too large to perform.
- Restricted to affine methods, excluding Krylov methods such as CG-based PIMs; non-affine operations raise errors.
- Numerical stability of the e-graph-optimal algorithm is unexplored.
Experimental setup
Datasets: ERA5 2 meter temperature (single time point, 1 Jan 2026) synthetic GP sample on 100x100 grid
Benchmarks: ERA5 regression synthetic CAGP regression
Metrics: wall-time RMSE NLL
Baselines: BayesCG probabilistic Gauss-Seidel CAGP-CG CAGP-GS ProbMG (hand-implemented)
Builds on: probabilistic iterative methods BayesCG computation-aware Gaussian processes equality saturation egg
Closest papers by meaning
- Stable and Scalable Probabilistic Numerical Solvers for Stiff and High-Dimensional ODEs (NeurIPS 2026 · similarity 0.85)
- Learning to Relax: Setting Solver Parameters Across a Sequence of Linear System Instances (ICLR 2024 · similarity 0.85)
- Turbocharging Gaussian Process Inference with Approximate Sketch-and-Project (NeurIPS 2025 · similarity 0.85)
- Improving Linear System Solvers for Hyperparameter Optimisation in Iterative Gaussian Processes (NeurIPS 2024 · similarity 0.84)
- Parallel Affine Transformation Tuning of Markov Chain Monte Carlo (ICML 2024 · similarity 0.84)
- Practical Bayesian Algorithm Execution via Posterior Sampling (NeurIPS 2024 · similarity 0.84)
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.