📊 statistics

Affine Tracing: A New Paradigm for Probabilistic Linear Solvers

This paper unifies Bayesian probabilistic linear solvers and probabilistic iterative methods under the framework of non-stationary affine PIMs, introduces "affine tracing" to automatically generate such solvers with posterior covariance computation, and demonstrates its effectiveness through a probabilistic multigrid solver for Gaussian process approximation.

Original authors: Disha Hegde, Marvin Pförtner, Jon Cockayne

Published 2026-05-12

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Disha Hegde, Marvin Pförtner, Jon Cockayne

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to solve a giant, complex puzzle where the pieces are numbers, and the goal is to find a specific arrangement (the solution) that fits perfectly. In the world of computers, this is called solving a "linear system." Usually, computers just crunch the numbers until they get an answer. But what if the computer runs out of time or energy before it's 100% sure?

This is where Probabilistic Linear Solvers (PLSs) come in. Instead of just giving you a single answer, they say, "Here is my best guess, and here is a map of how uncertain I am about it." It's like a weather forecast: instead of saying "It will rain at 2 PM," it says, "There's a 70% chance of rain, but if I run the simulation for another hour, that might change."

The Old Debate: Two Schools of Thought

For a long time, scientists thought there were two completely different ways to build these "uncertainty-aware" solvers:

The "Bayesian" Approach: This method starts with a "prior belief" (a guess) and updates it step-by-step as it gathers new clues. It's like a detective updating their theory as they find new evidence.
The "Iterative" Approach: This method takes a standard, fast computer algorithm and tries to "lift" it into a probability world. It's like taking a race car and trying to add a weather forecast system to it.

The paper argues that this split is a fake argument. The authors show that the "Bayesian" detectives are actually just a special, slightly rigid version of the "Iterative" race cars. They are all driving the same car; they just have different paint jobs.

The Big Problem: The "Hand-Writing" Bottleneck

Here is the catch: Building these probabilistic race cars is incredibly hard. To make a standard algorithm (like a Multigrid solver, which is a super-fast way to solve puzzles) work with probability, you have to do massive amounts of complex math by hand. You have to write out new formulas for every single new algorithm you want to use.

It's like if every time you wanted to build a new type of bicycle, you had to invent the physics of gears and chains from scratch, rather than just using a standard toolkit. This meant that even though these probabilistic methods were great in theory, nobody could actually build them for complex problems because the math was too tedious and prone to errors.

The Solution: "Affine Tracing" (The Magic Tracer)

The authors introduce a new tool called Affine Tracing. Think of this as a "magic highlighter" or a "symbolic tracer."

Instead of a human mathematician sitting down to write out the complex probability formulas for a new solver, you simply feed the computer's standard code into this "tracer."

The tracer watches the code run.
It sees every math operation the code does.
It automatically builds a "computational graph" (a flowchart of the math).
It then automatically transforms this graph to calculate the uncertainty (the probability distribution) without a human ever having to write a single new formula.

It's like having a robot that watches you build a house and automatically writes the blueprint for the insurance policy, rather than you having to calculate the insurance risk yourself.

The "Equality Saturation" Trick

There was one tricky part: sometimes the math involves a "magic inverse" (a specific type of prior) that makes the numbers cancel out in beautiful ways, but it's hard to see those cancellations in a giant flowchart.

To fix this, the authors used a technique called Equality Saturation. Imagine you have a tangled ball of yarn representing a complex math equation. Equality saturation is like a machine that tries every possible way to untangle and rearrange that yarn to find the shortest, simplest path. It finds the most efficient way to cancel out the messy parts so the computer doesn't have to do unnecessary work.

What They Built and Tested

To prove this works, the authors didn't just talk about it; they built a Probabilistic Multigrid Solver.

Multigrid is a very sophisticated, high-speed solver used for big problems.
Using their "Affine Tracing" tool, they automatically generated the probabilistic version of this solver.
They tested it on a problem called Gaussian Process approximation (which is used for things like predicting weather patterns or smoothing out noisy data).

The Results:

Speed: Their automatically generated solver was faster than a human-written probabilistic solver because it didn't get bogged down in unnecessary calculations.
Accuracy: It converged (found the answer) much faster than existing probabilistic methods.
Real-World Test: They tested it on a massive dataset of global temperatures (ERA5). The new solver produced a much smoother and more accurate map of the temperature than the older methods, doing it in the same amount of time.

The Bottom Line

This paper breaks down the wall between two different mathematical worlds. It proves that if you have a standard, fast computer solver, you can automatically turn it into a "probabilistic" solver that tells you how sure it is, without needing a PhD in algebra to write the code.

They call this Affine Tracing. It's a "plug-and-play" system for uncertainty: you plug in your standard solver, and the system automatically gives you the uncertainty map, making it possible to use the most advanced, fast solvers in the world while still knowing how much you can trust the answer.

Technical Summary: Affine Tracing: A New Paradigm for Probabilistic Linear Solvers

Problem Statement

Probabilistic linear solvers (PLSs) are designed to quantify the uncertainty incurred by limited computational resources when solving linear systems $Ax^\star = b$ . Historically, the literature has maintained a strict dichotomy between two approaches:

Bayesian PLSs: These place a prior distribution on the solution and condition it on information obtained from projections of the linear system (e.g., BayesCG). While theoretically grounded, these methods often suffer from poor calibration when using search directions that depend non-linearly on the solution (such as Lanczos directions), leading to incorrect uncertainty quantification.
Probabilistic Iterative Methods (PIMs): These lift classical iterative solvers into probability space by pushing forward an initial distribution through affine iteration maps. While these methods offer strong calibration properties, their practical adoption has been hindered by the significant manual effort required to derive and implement the necessary algebraic manipulations, particularly for complex, non-stationary methods or when utilizing inverse priors.

The core challenge addressed is the trade-off between convergence speed and calibration, compounded by the "implementation barrier" where deriving probabilistic versions of sophisticated iterative methods (like multigrid) requires bespoke, error-prone algebraic derivations for every new solver.

Methodology: Affine Tracing

The authors propose Affine Tracing, an algorithmic framework that automates the construction of PIMs from standard implementations of affine iterative methods. The methodology consists of three main components:

Symbolic Tracing: Instead of performing pen-and-paper derivations, the framework passes symbolic tracers through a standard implementation of an iterative method (e.g., Gauss-Seidel or Multigrid). By operator overloading, the tracer records all affine operations ( $x \mapsto Gx + f$ ) into a computational graph.
Graph Transformation for Covariance: To compute the posterior covariance $\Sigma_i$ , the framework transforms the constructed computational graph. It replaces all shift operations (addition of vectors) with identity nodes, effectively isolating the linear operator $G_{i:1}$ that maps the initial state to the current iterate. The covariance is then computed by passing the prior covariance through this transformed graph (or its transpose).
Equality Saturation for Inverse Priors: A critical requirement for applications like computation-aware Gaussian processes (CAGPs) is the use of an inverse prior ( $\Sigma_0 = A^{-1}$ ), which necessitates computing a "downdate" term $D_i$ such that $\Sigma_i = A^{-1} - D_i$ . The authors employ equality saturation (using the egglog library) to algebraically simplify the computational graph. This process identifies equivalent representations of the graph and applies rewrite rules (e.g., $A^{-1}A = I$ ) to cancel terms and isolate the downdate structure without explicitly computing $A^{-1}$ .

Key Theoretical Contributions

The paper establishes two fundamental theoretical results that unify the existing literature:

Unification of Bayesian and Iterative Methods: The authors prove that Bayesian PLSs with Gaussian iterates are a special case of non-stationary affine PIMs. Specifically, if the iteration maps are projection methods, the Bayesian update is equivalent to a PIM. This implies that the "Bayesian" label is often a misnomer for algorithms like BayesCG, which rely on non-affine dependencies.
Calibration of Affine PIMs: The paper proves that any realistic affine iterative method with the fixed-point property is strongly calibrated. This result significantly weakens the assumptions of previous work (which required stationarity and diagonalizability), showing that calibration is an emergent property of the fixed-point nature of affine solvers rather than a specific Bayesian construction.

Experimental Results

The framework was evaluated through the automatic generation of a probabilistic multigrid solver, a method that would otherwise require extensive manual derivation.

Implementation Efficiency: In experiments comparing the affine traced version of a 2-grid multigrid solver against a hand-implemented version, the traced version demonstrated superior performance. The time gap widened with increasing system dimension, attributed to the avoidance of explicit cross-covariance tracking required by manual implementations.
Synthetic CAGP Performance: In Gaussian process approximation tasks using the inverse prior, the probabilistic multigrid solvers (2-grid and 3-grid) generated via affine tracing achieved faster convergence in terms of Root Mean Square Error (RMSE) and Negative Log Likelihood (NLL) compared to state-of-the-art baselines like BayesCG and probabilistic Gauss-Seidel. The faster convergence of multigrid iterations compensated for the higher per-iteration cost.
Real-World Regression (ERA5): Applied to a large-scale temperature regression problem (ERA5 dataset), the CAGP-MG-2 solver produced posterior means that were smoother and captured finer structural details of the temperature field compared to other PLSs, while achieving rapid RMSE reduction.

Significance and Claims

The paper claims to resolve the calibration-convergence trade-off by shifting the focus to (non-stationary) affine PIMs. Its primary significance lies in democratizing the implementation of sophisticated probabilistic solvers. By automating the lifting process and the algebraic simplification required for inverse priors, the framework allows practitioners to experiment with advanced iterative methods (like multigrid) without needing to perform bespoke numerical analysis or algebraic manipulations.

The authors note that while the approach is powerful, it currently relies on the affine assumption, which excludes Krylov methods (the workhorse for sparse systems). They acknowledge that the equality saturation approach can be brittle for very large graphs and that the numerical stability of the extracted algorithms requires further investigation. However, the work establishes a new paradigm where probabilistic linear algebra can be constructed algorithmically from existing deterministic codebases.