lasso_regression

Contents

lasso_regression#

class braintools.optim.lasso_regression(X, Y, alpha=1.0, *, fit_intercept=True, max_iter=10000, tol=None)#

Fit a lasso (L1-penalised) linear regression with FISTA.

Minimises the objective of sklearn.linear_model.Lasso,

\[\frac{1}{2 n} \|Y - X W - b\|_F^2 + \alpha \|W\|_1,\]

with \(n\) the number of samples, by FISTA [1] with gradient-based adaptive restart [2] in a lax.while_loop. The intercept is not penalised, and each target column is an independent problem.

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 L1 penalty. When X or Y carry units, alpha must be a Quantity in the unit of X times the unit of Y.

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

  • max_iter (int) – Maximum number of FISTA iterations (static).

  • tol (float | None) – Relative tolerance on the first-order optimality residual (the stopping rule is described in elastic_net_regression()). Default 1e-8 for float64 data and 1e-7 (close to the rounding floor) for float32 data. tol=0 runs all max_iter iterations unless an exact minimiser is reached.

Returns:

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

Return type:

RegressionResult

Raises:

  • ValueError – If the shapes are invalid, a concrete alpha is negative, or max_iter/tol are invalid.

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

See also

elastic_net_regression

mixed L1 and L2 penalties.

ridge_regression

closed-form L2 penalty.

Notes

Pure; supports jax.jit and jax.vmap. The iteration count is data dependent, so reverse-mode differentiation (jax.grad) is not supported. For large alpha (above max|X_c^T Y_c| / n) every coefficient is exactly zero.

The loop stops when, for every target column, the largest entry of the proximal gradient mapping \(L (Z_k - W_{k+1})\) (a first-order optimality residual, zero exactly at the minimiser) is at most tol * max|X_c^T Y_c| / n; see elastic_net_regression(). 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.

References

Examples

>>> import jax
>>> import jax.numpy as jnp
>>> import braintools
>>> X = jax.random.normal(jax.random.PRNGKey(0), (100, 5))
>>> Y = X @ jnp.array([3., 0., 0., -2., 0.]) + 1.
>>> res = braintools.optim.lasso_regression(X, Y, alpha=0.1)
>>> print(jnp.round(res.coef, 2), round(float(res.intercept), 2))
[ 2.86 -0.    0.   -1.9   0.  ] 0.98
Contents