Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Matrix-Free Least Squares Solvers: Values, Gradients, and What to Do With Them

Hrittik Roy, Søren Hauberg, Nicholas Krämer

TL;DR

This work recasts least-squares as a differentiable, matrix-free operator by deriving reverse-mode gradients through adjoint methods and implementing an efficient $\texttt{LstSq}$ solver in JAX using $\text{LSMR}$. It then leverages the null-space method to solve constrained neural-network optimization problems by reframing constraints as a sequence of least-squares problems, enabling easy integration with Optax. The authors demonstrate practical benefits: scalable sparsity and equivariance constraints on models with millions of parameters, improved conservativeness for score-based models, and faster, more reliable Gaussian-process calibration via differentiable LS. Overall, the paper provides a pragmatic toolkit that makes differentiable linear algebra accessible for a broad range of ML applications with explicit, data-efficient gradients through LS solvers.

Abstract

This paper argues that the method of least squares has significant unfulfilled potential in modern machine learning, far beyond merely being a tool for fitting linear models. To release its potential, we derive custom gradients that transform the solver into a differentiable operator, like a neural network layer, enabling many diverse applications. Empirically, we demonstrate: (i) scalability by enforcing weight sparsity on a 50 million parameter model; (ii) imposing conservativeness constraints in score-based generative models; and (iii) hyperparameter tuning of Gaussian processes based on predictive performance. By doing this, our work represents the next iteration in developing differentiable linear-algebra tools and making them widely accessible to machine learning practitioners.

Matrix-Free Least Squares Solvers: Values, Gradients, and What to Do With Them

TL;DR

This work recasts least-squares as a differentiable, matrix-free operator by deriving reverse-mode gradients through adjoint methods and implementing an efficient solver in JAX using . It then leverages the null-space method to solve constrained neural-network optimization problems by reframing constraints as a sequence of least-squares problems, enabling easy integration with Optax. The authors demonstrate practical benefits: scalable sparsity and equivariance constraints on models with millions of parameters, improved conservativeness for score-based models, and faster, more reliable Gaussian-process calibration via differentiable LS. Overall, the paper provides a pragmatic toolkit that makes differentiable linear algebra accessible for a broad range of ML applications with explicit, data-efficient gradients through LS solvers.

Abstract

This paper argues that the method of least squares has significant unfulfilled potential in modern machine learning, far beyond merely being a tool for fitting linear models. To release its potential, we derive custom gradients that transform the solver into a differentiable operator, like a neural network layer, enabling many diverse applications. Empirically, we demonstrate: (i) scalability by enforcing weight sparsity on a 50 million parameter model; (ii) imposing conservativeness constraints in score-based generative models; and (iii) hyperparameter tuning of Gaussian processes based on predictive performance. By doing this, our work represents the next iteration in developing differentiable linear-algebra tools and making them widely accessible to machine learning practitioners.
Paper Structure (26 sections, 3 theorems, 47 equations, 5 figures, 7 tables)

This paper contains 26 sections, 3 theorems, 47 equations, 5 figures, 7 tables.

Table of Contents

  1. Introduction
  2. Least squares: values, gradients, and what to do with them
  3. Values: matrix-free least-squares in JAX
  4. Gradients: Adjoint of the least-squares problem
  5. What to do with them: constrained optimization of neural networks
  6. Experiments
  7. Efficiency of custom gradients
  8. Gaussian process calibration via differentiable least-squares
  9. Constrained optimization
  10. Limitations and conclusion
  11. Least squares redux
  12. Background on the method of adjoints
  13. Proof of \ref{['theorem:lsq_gradients']}
  14. Tall case
  15. Wide case
  16. ...and 11 more sections

Key Result

Theorem 1

Let $\mathbf{A}(\theta)$ be a full-rank matrix, dependent on parameters $\theta$, and accessed through matrix-vector products $(\theta, \mathbf{v}) \mapsto \mathbf{A}(\theta)\mathbf{v}$. Let $\mathbf{b}$ be a known vector, and let $\lambda \in \mathbb{R}$ be a known regularization weight. Let $\mu$

Figures (5)

  • Figure 1: Combine the null-space projection with a standard Optax optimizer using https://github.com/h-roy/nuox.
  • Figure 2: SGD & penalty method fail the constraint, unlike the null-space method.
  • Figure 3: Automatic differentiation versus a custom vector-Jacobian product (VJP). Our custom VJP is five to ten times faster than unrolling the solver's loop.
  • Figure 4: Calibration with negative log-marginal-likelihood (LML, baseline) versus evaluating the fit of the predictive mean (PRED, ours). PRED is over ten times faster (left), consistently achieves lower test losses (root mean square error on test set, $p$-value is $3.11\%$), which correlates with a better mean-data fit (right).
  • Figure 5: Bidiagonalization vs. CG.

Theorems & Definitions (7)

  • Theorem 1: Gradients of LstSq
  • proof
  • Example 2: Bidiagonalization vs. CG
  • Theorem 3
  • proof
  • Theorem 4
  • proof