Affine Tracing: A New Paradigm for Probabilistic Linear Solvers

TL;DR

This paper introduces affine tracing, a framework that automates the creation of probabilistic linear solvers from standard affine iterative methods, enabling uncertainty quantification in linear system solutions.

Contribution

It demonstrates how to automatically generate probabilistic iterative methods using symbolic tracing and algebraic simplification, unifying Bayesian and affine probabilistic methods.

Findings

Affine tracing can automatically produce probabilistic multigrid solvers.

The generated solvers provide uncertainty quantification in Gaussian process approximation.

The framework reduces manual effort in implementing probabilistic linear solvers.

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.