Affine Tracing: A New Paradigm for Probabilistic Linear Solvers

2026-05-11

AI Summary

This work addresses the complexity and manual effort involved in implementing probabilistic linear solvers (PLS) by introducing an affine tracking framework that unifies the theoretical foundations of Bayesian PLS and probabilistic iterative methods. The framework leverages symbolic tracing to automatically transform standard affine iterative schemes into probabilistic solvers capable of computing posterior covariances, while also enabling algebraic optimizations. Key contributions include establishing that Bayesian PLS constitutes a special case of affine probabilistic iterative methods (PIM), proving the calibration property of practical affine PIMs, and presenting the first general-purpose automatic construction methodology. By integrating symbolic tracing, affine computational graphs, and equation saturation techniques, the authors successfully generate a probabilistic multigrid solver, demonstrating its efficacy and practicality in Gaussian process approximation tasks.

Abstract

Probabilistic linear solvers (PLSs) return probability distributions that quantify uncertainty due to limited computation in the solution of linear systems. The literature has traditionally distinguished between Bayesian PLSs, which condition a prior on information obtained from projections of the linear system, and probabilistic iterative methods (PIMs), which lift classical iterative solvers to probability space. In this work we show this dichotomy to be false: Bayesian PLSs are a special case of non-stationary affine PIMs. In addition, we prove that any realistic affine PIM is calibrated. These results motivate a focus on (non-stationary) affine PIMs, but their practical adoption has been limited by the significant manual effort required to implement them. To address this, we introduce affine tracing, an algorithmic framework that automatically constructs a PIM from a standard implementation of an affine iterative method by passing symbolic tracers through the computation to build an affine computational graph. We show how this graph can be transformed to compute posterior covariances, and how equality saturation can be used to perform algebraic simplifications required for computation under specific prior choices. We demonstrate the framework by automatically generating a probabilistic multigrid solver and evaluate its performance in the context of Gaussian process approximation.

Problem

• probabilistic linear solvers
• affine iterative methods
• implementation complexity
• probabilistic iterative methods
• Bayesian PLSs

Innovation

• Affine Tracing
• Probabilistic Linear Solvers
• Affine Iterative Methods
• Equality Saturation
• Bayesian Numerics