ridge_regression

Contents

ridge_regression#

class braintools.optim.ridge_regression(X, Y, alpha=1.0, *, fit_intercept=True)#

Fit a ridge (L2-penalised least-squares) regression in closed form.

Minimises

\[\|Y - X W - b\|_F^2 + \alpha \|W\|_F^2,\]

the objective of sklearn.linear_model.Ridge; the intercept \(b\) is not penalised. The solution is computed from the thin SVD \(X_c = U S V^\top\) of the centred design matrix as \(W = V \operatorname{diag}(s / (s^2 + \alpha)) U^\top Y_c\), which works for any shape (including more features than samples) and reduces to the minimum-norm least-squares solution at alpha=0.

Parameters:

  • X (array_like or brainunit.Quantity) – Design matrix, shape (n_samples, n_features).

  • Y (array_like or brainunit.Quantity) – Targets, shape (n_samples,) or (n_samples, n_targets).

  • alpha (float or brainunit.Quantity, default 1.0) – Non-negative L2 penalty. When X carries a unit, alpha must be a Quantity in that unit squared.

  • fit_intercept (bool) – Fit an unpenalised intercept by centring X and Y first.

Returns:

coef of shape (n_features,) or (n_features, n_targets) and intercept.

Return type:

RegressionResult

Raises:

  • ValueError – If the shapes are invalid or a concrete alpha is negative.

  • brainunit.UnitMismatchError – If alpha does not have the unit of X squared.

See also

linear_regression

the alpha=0 case.

lasso_regression, elastic_net_regression

RLS

online ridge regression with P0 = I / alpha.

Notes

Pure; supports jax.jit, jax.vmap (for example over a grid of alpha) and jax.grad. Gradients with respect to Y and alpha always exist; gradients with respect to X go through the SVD and need distinct non-zero singular values of the centred X. A negative alpha that is only known at run time (traced under jax.jit or jax.vmap) cannot raise, so it returns NaN coefficients and intercept instead.

Examples

>>> import jax
>>> import jax.numpy as jnp
>>> import braintools
>>> X = jax.random.normal(jax.random.PRNGKey(0), (50, 3))
>>> Y = X @ jnp.array([1., -2., 0.5]) + 3.
>>> res = braintools.optim.ridge_regression(X, Y, alpha=1e-6)
>>> print(jnp.round(res.coef, 3), round(float(res.intercept), 3))
[ 1.  -2.   0.5] 3.0
>>>
>>> # vmap over a grid of penalties
>>> alphas = jnp.array([0.1, 1., 10.])
>>> coefs = jax.vmap(lambda a: braintools.optim.ridge_regression(X, Y, a).coef)(alphas)
>>> coefs.shape
(3, 3)
Contents