Automatic Selection of the Nugget for Linear System Solves in Machine Learning

TMLR Paper6817 Authors

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 resolve 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`. A distinguishing feature of `autonugget` is the ability to combine multiple linear solves using Richardson extrapolation, improving in accuracy over approximations based on a single nugget.

Submission Type: Regular submission (no more than 12 pages of main content)

Assigned Action Editor: ~Mathurin_Massias1

Submission Number: 6817