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

2026-06-29

AI Summary

This work addresses the challenge of selecting the regularization parameter (nugget) in ill-posed linear systems arising in machine learning, where existing adaptive methods lack compatibility with automatic differentiation and suffer from computational inefficiency. To overcome these limitations, we introduce autonugget, a lightweight Python package fully compatible with JAX’s automatic differentiation framework. Our approach uniquely integrates Richardson extrapolation with Tikhonov regularized solutions computed across multiple nugget values, thereby preserving end-to-end differentiability while avoiding the information loss inherent in single-solution strategies. Experimental results demonstrate that autonugget significantly enhances solution accuracy and training stability without compromising rapid prototyping capabilities.

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.

Problem

• ill-conditioned linear systems
• Tikhonov regularization
• nugget selection
• automatic differentiation
• rapid prototyping

Innovation

• ill-conditioned linear systems
• Tikhonov regularization
• Richardson extrapolation
• automatic differentiation
• rapid prototyping