Least Squares

Classes to approximate least-squares solution to regularised linear matrix equations.

Primarily used in Coreax to approximate regularised matrix inverses.

Given a matrix \(A\in\mathbb{R}^{n\times n}\) and some regularisation parameter \(\lambda\in\mathbb{R}_{\ge 0}\), the regularised inverse of \(A\) is \((A + \lambda I_n)^{-1}\) where \(I_n\) is the \(n\times n\) identity matrix.

Coreset algorithms which use the regularised inverse of large kernel matrices can become prohibitively expensive as the data size increases. To reduce this computational cost, this regularised inverse can be approximated by various methods, depending on the acceptable levels of error.

The RegularisedLeastSquaresSolver in this module provides the functionality required to approximate the regularised inverse of matrices. Furthermore, this class allows for “block-inversion” where given an invertible matrix \(B\), the block array

\[\begin{split}A = \begin{bmatrix}B & 0 & \dots & 0 \\ 0 & 0 & \dots & 0 \\ \vdots & \ddots & \dots & \vdots \\ 0 & 0 & \dots & 0\end{bmatrix},\end{split}\]

where only the top-left block contains non-zero elements, can be “inverted” to give

\[\begin{split}A^{-1} := \begin{bmatrix}B^{-1} & 0 & \dots & 0 \\ 0 & 0 & \dots & 0 \\ \vdots & \ddots & \dots & \vdots \\ 0 & 0 & \dots & 0\end{bmatrix}.\end{split}\]

This functionality allows for iterative coreset algorithms which require inverting growing arrays to have fully static array shapes and thus be JIT-compatible.

To compute these “inverses” in JAX, we require the target and identity array passed in solve() to be a matrix of zeros except for ones on the diagonal up to the dimension of the non-zero block.

class coreax.least_squares.RegularisedLeastSquaresSolver

Bases: Module

Base class for solving regularised linear matrix equations via least-squares.

Given an array \(A \in \mathbb{R}^{n \times n}\), a regularisation parameter \(\lambda \in \mathbb{R}_{\ge 0}\), and an array of targets \(B \in \mathbb{R}^{n \times m}\) the least-squares solution to the regularised linear equation \((A + \lambda I_n)B = X\) has solution \(X = (A + \lambda I_n)^{-1}B\) where \(I_n \in \mathbb{R}^{n\times n}\) is the identity matrix.

abstract solve(array, regularisation_parameter, target, identity)

Compute the least-squares solution to a regularised linear matrix equation.

Parameters:

• array (Shaped[Array, 'n n']) – \(n \times n\) array, corresponds to \(A\)

• regularisation_parameter (float) – Regularisation parameter for stable inversion of array; negative values will be converted to positive

• target (Shaped[Array, 'n m']) – \(n \times m\) array of targets, corresponds to \(B\)

• identity (Shaped[Array, 'n n']) – Identity matrix

Return type:

Shaped[Array, 'n n']

Returns:

Approximation of the regularised least-squares solution \(X\)

solve_stack(arrays, regularisation_parameter, targets, identity, in_axes=(0, None, 0, None))

Compute least-squares solutions to stack of regularised linear matrix equations.

Parameters:

• arrays (Shaped[Array, 'l n n']) – Horizontal stack of arrays with shape \(l \times n \times n\)

• regularisation_parameter (float) – Regularisation parameter for stable inversion of arrays; negative values will be converted to positive

• targets (Shaped[Array, 'l n m']) – Horizontal stack of targets with shape \(l \times n \times m\)

• identity (Shaped[Array, 'n n']) – Identity matrix

• in_axes (Union[int, None, Sequence[Any]]) – An integer, None, or sequence of values specifying which array axes of parameters to solve() to map over. See vmap() documentation for further information.

Return type:

Shaped[Array, 'l n n']

Returns:

Approximation of the regularised least-squares solutions

class coreax.least_squares.MinimalEuclideanNormSolver(rcond=None)

Bases: RegularisedLeastSquaresSolver

Find minimal-norm least-squares solution to the regularised linear matrix equation.

Computes the solution that approximately solves the regularised linear matrix equation. The equation may be under-, well-, or over-determined. If array is full rank, then the solution is exact, up to floating-point errors. Else, the solution minimises the Euclidean 2-norm. If there are multiple minimising solutions, the one with the smallest 2-norm is returned.

Note

This solver does not give time savings and instead acts as a robust default option useful for comparing other solvers to.

Parameters:

rcond (Optional[float]) – Cut-off ratio for small singular values of array. For the purposes of rank determination, singular values are treated as zero if they are smaller than rcond times the largest singular value of array. The default value of None will use the machine precision multiplied by the largest dimension of the array. An alternate value of -1 will use machine precision.

rcond: Optional[float] = None

solve(array, regularisation_parameter, target, identity)

Compute the least-squares solution to a regularised linear matrix equation.

Parameters:

• array (Shaped[Array, 'n n']) – \(n \times n\) array, corresponds to \(A\)

• regularisation_parameter (float) – Regularisation parameter for stable inversion of array; negative values will be converted to positive

• target (Shaped[Array, 'n m']) – \(n \times m\) array of targets, corresponds to \(B\)

• identity (Shaped[Array, 'n n']) – Identity matrix

Return type:

Shaped[Array, 'n n']

Returns:

Approximation of the regularised least-squares solution \(X\)

class coreax.least_squares.RandomisedEigendecompositionSolver(random_key, oversampling_parameter=25, power_iterations=1, rcond=None)

Bases: RegularisedLeastSquaresSolver

Approximate solution to regularised linear equations via random eigendecomposition.

Solving regularised linear equations involving large arrays can become prohibitively expensive as size increases. To reduce this computational cost, the solution can be approximated by various methods. RandomisedEigendecompositionSolver is a class that does such an approximation using the randomised eigendecomposition of the input array.

Warning

Input arrays must be Hermitian for this method to have predictable behaviour. We do not check this.

Using Algorithm 4.4. and 5.3 from [halko2009randomness] we approximate the eigendecomposition of a Hermitian matrix. The parameters oversampling_parameter and power_iterations present a trade-off between speed and approximation quality. See [halko2009randomness] for discussion on choosing sensible parameters; the defaults chosen here are cautious.

Parameters:

• random_key (ArrayLike) – Key for random number generation

• oversampling_parameter (int) – Number of random columns to sample; the larger the oversampling_parameter, the more accurate, but slower the method will be

• power_iterations (int) – Number of power iterations to do; the larger the power_iterations, the more accurate, but slower the method will be

• rcond (Optional[float]) – Cut-off ratio for small singular values of the array. For the purposes of rank determination, singular values are treated as zero if they are smaller than rcond times the largest singular value of a. The default value of None will use the machine precision multiplied by the largest dimension of the array. An alternate value of -1 will use machine precision.

random_key: ArrayLike

oversampling_parameter: int = 25

power_iterations: int = 1

rcond: Optional[float] = None

randomised_eigendecomposition(array)

Approximate the eigendecomposition of Hermitian matrices.

Using Algorithm 4.4. and 5.3 from [halko2009randomness] we approximate the eigendecomposition of a matrix. The parameters oversampling_parameter and power_iterations present a trade-off between speed and approximation quality. See [halko2009randomness] for discussion on choosing sensible parameters, the defaults chosen here are cautious.

Given the matrix \(A \in \mathbb{R}^{n\times n}\) and \(r=\) oversampling_parameter, we return a diagonal array of eigenvalues \(\Lambda \in \mathbb{R}^{r \times r}\) and a rectangular array of eigenvectors \(U\in\mathbb{R}^{n\times r}\) such that we have \(A \approx U\Lambda U^T\).

Parameters:

array (Shaped[Array, 'n n']) – Array to be decomposed

Return type:

tuple[Shaped[Array, 'oversampling_parameter'], Shaped[Array, 'n oversampling_parameter']]

Returns:

Eigenvalues and eigenvectors that approximately decompose the array

solve(array, regularisation_parameter, target, identity)

Compute the least-squares solution to a regularised linear matrix equation.

Parameters:

• array (Shaped[Array, 'n n']) – \(n \times n\) array, corresponds to \(A\)

• regularisation_parameter (float) – Regularisation parameter for stable inversion of array; negative values will be converted to positive

• target (Shaped[Array, 'n m']) – \(n \times m\) array of targets, corresponds to \(B\)

• identity (Shaped[Array, 'n n']) – Identity matrix

Return type:

Shaped[Array, 'n n']

Returns:

Approximation of the regularised least-squares solution \(X\)