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

Abstract: Rapid prototyping of algorithms is a critical step in modern machine learning. Most algorithms exploit linear algebra, creating a need for lightweight numerical routines which -- while potentially sub-optimal for the task at hand -- can be rapidly implemented. For the numerical solution of ill-conditioned linear systems of equations, the standard solution for prototyping is Tikhonov-regularised inversion using a nugget. However, selection of the size of nugget is often difficult, and the use of data-adaptive procedures precludes automatic differentiation, introducing instabilities into end-to-end training. Further, while data-adaptive procedures perform multiple linear solves to select the size of nugget, only the result of one such solve is returned, which we argue is wasteful. This paper aims to circumvent the above difficulties, presenting autonugget; a Python package for automatic and stable numerical solution of linear systems suitable for rapid prototyping, and fully compatible with automatic differentiation using JAX. autonugget combines multiple linear solves using Richardson extrapolation to determine the solution of the ill-conditioned system, improving in accuracy over approximations based on a single nugget.

• The paper introduces autonugget, which employs Richardson extrapolation from multiple regularised solutions to accurately resolve ill-conditioned linear systems.
• It proposes a data-adaptive nugget selection method that balances numerical instability with extrapolation error to achieve unbiased results.
• The approach’s AD compatibility significantly improves gradient-based optimization in kernel methods and GP hyperparameter tuning.

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

Problem Formulation and Motivation

The paper addresses the numerical solution of ill-conditioned linear systems $Ax = b$, where $A$ is symmetric positive definite but potentially has a high condition number, as commonly encountered in kernel methods (e.g., Gaussian processes, Bayesian optimisation) and probabilistic numerics. The standard approach in rapid prototyping pipelines is to leverage Tikhonov (nugget) regularisation, seeking $x_\sigma = (A + \sigma I)^{-1} b$ for small $\sigma > 0$ to mitigate numerical instability.

Despite its prevalence, this approach has several deficiencies:

• Nugget selection is nontrivial: The smallest stable nugget $\sigma_\star$ is instance-dependent and typically selected via condition number heuristics, making it incompatible with AD frameworks (e.g., JAX) due to introduced conditional control flow.
• Suboptimal data efficiency: Adaptive nugget search routines perform multiple solves but discard all but the final candidate solution.
• Bias-variance conflation: Nugget regularisation for numerical stability (where $\sigma \to 0$ is desired) becomes confounded with regularisation for statistical generalisation ($\sigma > 0$ as a hyperparameter).

The authors propose autonugget, an automatic and stable solver, which combines data-adaptive nugget selection with Richardson extrapolation, yielding improved accuracy and full AD compatibility for end-to-end machine learning training.

Methodology: Richardson Extrapolation on Nuggeted Solutions

A critical empirical insight is that as the nugget $\sigma \to 0$, the direct solver loses numerical fidelity below a certain problem- and architecture-dependent threshold $\sigma_\star$. Above this threshold, approximate solutions $\hat{x}_\sigma$ are numerically precise approximations of $A$0, but with increasing bias relative to the true unregularised solution $A$1.

The key idea exploited in autonugget is that one can compute several solutions $A$2, at successively smaller nuggets $A$3, and extrapolate these reliable values back to $A$4 via polynomial (Richardson) extrapolation. This removes bias without reintroducing instability:

Figure 1: Illustration of autonugget—extrapolation from multiple stable nuggeted solutions enables a more accurate approximation than any individual regularised solve.

The methodology has several technical components:

• A polynomial model is fit to each coordinate of $A$5 as a function of $A$6, using a Vandermonde system for the chosen nodes $A$7 (typically, geometric grids are recommended).
• The solution is extrapolated to $A$8 to approximate $A$9.
• The nugget values are selected by approximately balancing the error due to numerical instability (which increases as $x_\sigma = (A + \sigma I)^{-1} b$0 decreases) with the extrapolation error (which decreases as $x_\sigma = (A + \sigma I)^{-1} b$1 decreases), yielding a trade-off equation whose solution quantifies the minimum safe nugget.
• Custom AD rules are implemented to enable differentiation through the entire extrapolation process, an essential property for modern ML pipelines.

Theoretical Analysis

The analysis is separated into exact arithmetic and finite precision settings.

• In exact arithmetic, the extrapolation error decays as $x_\sigma = (A + \sigma I)^{-1} b$2, where $x_\sigma = (A + \sigma I)^{-1} b$3 is the base nugget and $x_\sigma = (A + \sigma I)^{-1} b$4 is the number of extrapolation nodes. More nodes allow higher-order convergence to the unbiased solution.
• In floating-point arithmetic, polynomial extrapolation amplifies noise, so the number and location of nodes must be selected to balance floating-point error with bias.
• The authors develop explicit error decompositions and provide practical guidance for automatically identifying $x_\sigma = (A + \sigma I)^{-1} b$5 based on stochastic estimates of the smallest eigenvalue of $x_\sigma = (A + \sigma I)^{-1} b$6 and the observed conditioning.

Figure 2: Identification of the critical nugget value $x_\sigma = (A + \sigma I)^{-1} b$7; extrapolation is performed only with $x_\sigma = (A + \sigma I)^{-1} b$8 to ensure numerical stability.

Implementation: The autonugget Package

autonugget is implemented as a drop-in replacement for NumPy/SciPy solvers, supporting both direct usage and gradient-based optimisation under JAX. The reference grid for extrapolation, polynomial order, and nugget selection strategy are configurable.

A notable aspect is the integration of forward-mode AD custom rules, which ensure stability and correctness during the computation of derivatives of the linear system solution w.r.t. parameters, a core requirement for differentiable programming in machine learning.

Empirical Evaluation

Comprehensive experiments compare autonugget against standard routines (LU, fixed-nugget, adaptive-nugget, conjugate gradient, SVD, least-squares, truncated SVD) across a variety of matrix ensembles and kernel-generated problems.

Accuracy: Solution and Derivatives

Autonugget consistently achieves lower error than all baselines, especially for ill-conditioned systems (i.e., large kernel length-scale $x_\sigma = (A + \sigma I)^{-1} b$9):

Figure 3: Relative solver errors—autonugget outperforms standard direct and iterative solvers, with more pronounced gains for ill-conditioned matrices.

For the more challenging task of differentiating through the solver (e.g., in hyperparameter optimisation), autonugget exhibits strong stability and accuracy, uniquely retaining validity where SVD- and pseudo-inverse-based approaches fail due to loss of differentiability or instability:

Figure 4: Relative errors in automatic differentiation solutions—autonugget achieves superior or comparable accuracy for matrix derivatives, even for large $\sigma > 0$0 and extreme ill-conditioning.

Machine Learning Application: GP Hyperparameter Learning

In a realistic use-case, autonugget is used for gradient-based optimisation of GP kernel hyperparameters via leave-one-out cross-validation. Standard solvers, including fixed-nugget or SVD approaches, induce biased or unstable gradients, leading to poor convergence or pathologically biased parameter estimates. By contrast, autonugget enables smooth, unbiased optimisation dynamics:

Figure 5: Gradient-based cross-validation for GP regression—autonugget supports stable, accurate optimisation of highly ill-conditioned, small-sample kernel problems.

Auxiliary Results

Absolute error distributions and performance under alternative kernel types, as well as wall-clock timings, are systematically reported. While the extrapolation overhead implies higher computational cost (one extra solve per degree of extrapolation), autonugget remains practical for prototyping. In extreme scaling regimes, lower-cost variants (e.g., autonugget-cond) are suggested as alternatives.

Figure 6: Absolute error heatmaps—autonugget variants maintain minimal errors over a wide spectrum of conditioning and sizes.

Figure 7: Generalisation to non-kernel random matrices—autonugget matches direct solvers except where ill-conditioning otherwise defeats standard methods.

Implications and Future Directions

The integration of nugget extrapolation with AD compatibility satisfies an important and persistent need in the machine learning toolkit for robust, tuning-free linear solvers applicable across model classes (e.g., kernel methods, GP inference, differentiable ODE/PDE solvers).

Theoretically, the results formalise the bias-variance trade-off in numerical regularisation and specify how Richardson extrapolation can restore consistency without sacrificing computational efficiency. This cleanly decouples statistical regularisation from numerical stabilisation, enabling more principled uncertainty quantification and optimisation. The extension to large-scale or non-SPD systems, probabilistic error quantification for the extrapolated solution, and partial replacement of polynomial models with statistical regression are immediate areas for future work.

On the practical side, further software optimisations (e.g., batched or SVD-accelerated extrapolation) may reduce runtime burden, enabling application to even larger or more complex systems.

Conclusion

This work provides a rigorous and practical framework for solving ill-conditioned linear systems in machine learning pipelines, leveraging multiple stable regularised solutions and polynomial extrapolation to eliminate bias while ensuring stability. The method is empirically validated for both solver and derivative accuracy, is fully AD-compatible, and is deployed as an accessible Python package. Future advances in extrapolation theory, solver implementation, and probabilistic error accounting are expected to further enhance practical impact (2606.30328).