📊 statistics

Extrapolating from Regularised Solutions for Solving Ill-Conditioned Linear Systems in Machine Learning

This paper introduces **autonugget**, a JAX-compatible Python package that improves the accuracy and stability of solving ill-conditioned linear systems for machine learning prototyping by combining multiple Tikhonov-regularized solutions via Richardson extrapolation, thereby eliminating the need for difficult nugget selection and enabling fully differentiable end-to-end training.

Original authors: Disha Hegde, Jon Cockayne, Chris. J. Oates

Published 2026-06-30

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Disha Hegde, Jon Cockayne, Chris. J. Oates

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 puzzle where the pieces are numbers. In the world of machine learning, this often means solving a specific type of math problem called a "linear system." Usually, these puzzles are tricky, but sometimes they are ill-conditioned.

Think of an ill-conditioned system like a house of cards that is about to collapse. If you try to touch it (or solve it) with standard tools, the slightest wobble causes the whole thing to fall apart, giving you a completely wrong answer.

The Old Way: The "Nugget" Fix

To stop the house of cards from falling, mathematicians have traditionally used a trick called Tikhonov regularization. Imagine you are trying to balance a wobbly table. You slide a small, heavy coin (a "nugget") under one of the legs to make it stable.

In math terms, this "coin" is a tiny number added to the problem to make it stable enough to solve. However, this method has three big headaches:

The Goldilocks Problem: You don't know how big the coin needs to be. If it's too small, the table still wobbles. If it's too big, you've changed the shape of the table too much, and your answer is now biased (wrong in a different way).
The "Black Box" Problem: Modern machine learning often needs to learn by trial and error, adjusting its knobs automatically. But if you use a rule that says, "Add a coin if the table wobbles," the computer can't figure out how to adjust its knobs because the rule breaks the math chain needed for learning.
Wasted Effort: To find the right size for the coin, computers often try many different sizes, throw away the results, and only keep the last one. It's like tasting ten different soups to find the right amount of salt, but only serving the tenth spoonful.

The New Solution: "Autonugget"

The authors of this paper created a new tool called autonugget. They wanted a way to solve these wobbly puzzles that is automatic, stable, and doesn't waste effort.

Here is how it works, using a simple analogy:

The "Extrapolation" Trick
Instead of guessing the perfect coin size, autonugget tries a few different coin sizes that are definitely big enough to make the table stable.

Imagine you have a wobbly table. You place a small coin under the leg, and it wobbles a little. You place a medium coin, and it's steadier. You place a large coin, and it's very steady.
The computer measures how the table behaves with the small, medium, and large coins.
Then, it uses a mathematical technique called Richardson Extrapolation. Think of this as drawing a smooth curve through those three points and extending the line backward to see what would happen if the coin size was zero.

By looking at the trend of the stable solutions, the computer can mathematically "predict" the answer for the unstable, zero-coin scenario with high accuracy.

Why This Matters

The paper claims that autonugget solves the three headaches of the old method:

No Guessing: It automatically figures out the right balance without you needing to pick a number.
Machine Learning Friendly: Because the math is smooth and predictable, it plays nicely with modern "end-to-end" training, allowing computers to learn and adjust their parameters automatically.
No Waste: It uses the results of all the different coin sizes it tried to build a better final answer, rather than throwing the extra data away.

The Results

The authors tested their tool against standard methods (like the ones built into common software libraries) on various difficult math problems. They found that:

Accuracy: Autonugget gave much more accurate answers than standard methods, especially when the problems were very "wobbly" (ill-conditioned).
Speed: It was fast enough to be used for rapid prototyping (quickly testing new ideas).
Training: When used in a machine learning task (specifically optimizing a Gaussian Process, which is a way of modeling data), autonugget helped the computer find the correct solution quickly, whereas other methods got stuck or found the wrong answer.

In short, autonugget is a smart, automatic tool that stabilizes tricky math problems by testing a few safe versions and mathematically predicting the perfect answer, making it easier and faster to build and train machine learning models.

Technical Summary: Extrapolating from Regularised Solutions for Solving Ill-Conditioned Linear Systems in Machine Learning

Problem Statement

In modern machine learning, rapid prototyping often requires solving linear systems of the form $Ax = b$, where $A$ is a symmetric positive definite (SPD) matrix. While standard direct methods (e.g., LU factorization) are efficient for well-conditioned systems, they fail or become unstable when $A$ is ill-conditioned. The standard approach to mitigate this is Tikhonov regularisation, solving $(A + \sigma I)^{-1}b$ with a "nugget" parameter $\sigma > 0$ .

However, the paper identifies three critical limitations in the current practice of using Tikhonov regularisation for prototyping:

Nugget Selection Difficulty: The critical minimum value $\sigma^\star$ required for stable numerical computation is unknown and sensitive to the specific entries of $A$ and $b$ . Choosing $\sigma$ too small leads to numerical instability; choosing it too large introduces bias.
Incompatibility with End-to-End Training: Data-adaptive procedures (e.g., selecting $\sigma$ based on a condition number threshold) often involve conditional logic that breaks automatic differentiation (AD), preventing gradient-based optimization of parameters embedded within $A$ and $b$ .
Data Inefficiency: Adaptive procedures often perform multiple linear solves to find a suitable $\sigma$ but discard all intermediate results, returning only the final solution. This is argued to be computationally wasteful.

Methodology: The autonugget Framework

The authors propose autonugget, a Python package designed to provide stable, automatic, and differentiable solutions to ill-conditioned linear systems. The core methodology relies on Richardson extrapolation to combine multiple regularised solutions.

1. Extrapolation Estimator

The method assumes that for $\sigma \ge \sigma^\star$ , the numerical approximation $\hat{x}_\sigma$ of the regularised solution $x_\sigma = (A + \sigma I)^{-1}b$ is reliable. The algorithm:

Selects a set of nuggets $\Sigma = \{h\sigma_i\}_{i=1}^n$ based on a reference set $\Sigma_{ref}$ and a scaling factor $h$ .
Computes the solutions $\hat{x}_{h\sigma_i}$ for each nugget.
Fits a polynomial $f_n$ to the data points $(\sigma_i, \hat{x}_{\sigma_i})$ coordinate-wise.
Extrapolates this polynomial to $\sigma = 0$ to estimate the true solution $x$ .

2. Balancing Error and Stability

The paper provides a theoretical analysis in finite precision arithmetic to determine the optimal scaling factor $h$ . The total error is decomposed into:

Extrapolation Error: Decreases as $h \to 0$ (smaller nuggets) and as the number of points $n$ increases.
Numerical Error: Increases as $h \to 0$ because the condition number $\kappa(A + h\sigma_{min}I)$ grows, amplifying machine precision errors.

The authors derive a bound (Equation 8) to select $h$ such that these two error sources are balanced. For near-singular matrices, the optimal $h$ scales as $\epsilon_{MP}^{1/(n+1)}$ , where $\epsilon_{MP}$ is machine precision. This implies that using more data points ( $n$ ) allows for the use of larger nuggets while maintaining stability, thereby conferring additional numerical stability.

3. Automatic Differentiation

A distinguishing feature of autonugget is its compatibility with JAX. To compute gradients $\nabla_\theta x_\theta$ where $A$ and $b$ depend on parameters $\theta$ , the package implements a custom forward differentiation rule based on the identity:
$\nabla_\theta x_\theta = -A^{-1}_\theta (\nabla_\theta A_\theta) A^{-1}_\theta b_\theta + A^{-1}_\theta (\nabla_\theta b_\theta)$
This requires three linear solves. Crucially, autonugget automatically selects appropriate nuggets for each of these three systems independently, ensuring that the differentiation process remains stable even when the intermediate systems are ill-conditioned.

Key Contributions

autonugget Package: A lightweight, JAX-compatible Python tool for solving ill-conditioned linear systems without manual tuning of regularisation parameters.
Extrapolation Strategy: A method that utilizes multiple linear solves not just for parameter selection, but to construct a higher-accuracy estimate of the unregularised solution ( $\sigma \to 0$ ) via Richardson extrapolation.
Theoretical Error Bounds: Derivation of error bounds that explicitly balance extrapolation accuracy against numerical instability, providing a principled method for selecting the nugget scaling factor $h$ .
Differentiable Solver: A custom implementation that enables end-to-end training of machine learning models involving ill-conditioned linear systems, overcoming the AD incompatibility of standard adaptive solvers.

Empirical Results

The paper evaluates autonugget against several baselines (standard LU, fixed-nugget LU, condition-number-adaptive LU, Conjugate Gradient, SVD, and Truncated SVD) across a spectrum of Gram matrices with varying condition numbers.

Solver Accuracy: autonugget consistently outperforms standard LU and SVD methods, particularly for severely ill-conditioned systems (large kernel length-scales). It also outperforms adaptive baselines (LU-adapt), demonstrating that the extrapolation step provides accuracy gains beyond simple nugget selection.
Derivative Accuracy: In gradient computations, autonugget provides significantly more accurate derivatives than baselines. Standard LU fails due to instability, while fixed-nugget methods introduce bias. autonugget maintains stability and accuracy even when the system is severely ill-conditioned.
Gaussian Process Regression: In a hyperparameter optimization task (inferring kernel length-scales), autonugget enables the optimizer to converge to the true length-scale. In contrast, baselines either diverged (LU, SVD), converged to incorrect values due to bias (LU-fix), or produced non-smooth loss landscapes (LSTSQ, TSVD).

Significance and Claims

The paper positions autonugget as a practical tool for the "rapid prototyping" phase of machine learning development. The authors argue that while sophisticated numerical methods exist for specific problem structures, there is a gap in general-purpose, differentiable routines for ill-conditioned systems.

The significance of the work lies in:

Demystifying Nugget Selection: Removing the need for manual or heuristic tuning of regularisation parameters during the prototyping phase.
Enabling End-to-End Training: Allowing gradient-based optimization in scenarios where standard solvers would fail or break the computational graph.
Efficiency: Demonstrating that the "wasteful" multiple solves required by adaptive methods can be repurposed to improve the final solution's accuracy via extrapolation.

The authors acknowledge limitations, noting that the current implementation is restricted to symmetric positive definite matrices and incurs a higher computational cost (proportional to the number of extrapolation points $n$ ) than a single direct solve. They suggest future work could extend the framework to non-SPD matrices, optimize low-level implementations, or incorporate probabilistic error bars.