elastic_net_regression

elastic_net_regression#

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

Fit an elastic-net (L1 + L2 penalised) linear regression with FISTA.

Minimises the objective of sklearn.linear_model.ElasticNet,

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

with \(\rho\) = l1_ratio and \(n\) the number of samples, by FISTA [1] with gradient-based adaptive restart [2] in a lax.while_loop. The intercept is not penalised. Each target column is an independent problem (as in scikit-learn, not a multi-task penalty).

Parameters:

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

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

  • alpha (float, default 1.0) – Non-negative overall penalty strength.

  • l1_ratio (float, default 0.5) – Mixing parameter \(\rho \in [0, 1]\): 1 is the lasso, 0 is ridge (with the 1/(2n) data-term scaling).

  • 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 (see Notes). 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, a concrete l1_ratio is outside [0, 1], or max_iter/tol are invalid.

  • brainunit.UnitMismatchError – If X or Y carry physical units: the L1 and L2 terms then have different units, so the objective is not defined. Pass mantissas.

See also

lasso_regression

the l1_ratio=1 case.

ridge_regression

closed-form L2 penalty.

Notes

Pure; supports jax.jit and jax.vmap (a batched while_loop runs until every member has converged). The iteration count is data dependent, so reverse-mode differentiation (jax.grad) is not supported; forward mode (jax.jvp) works but differentiates the iterations, not the solution map.

Stopping rule. Let \(L = \sigma_{\max}(X_c)^2 / n + \alpha (1 - \rho)\) be the Lipschitz constant of the smooth part, \(Z_k\) the extrapolated point and \(W_{k+1}\) the proximal-gradient step from it. The gradient mapping \(G_k = L (Z_k - W_{k+1})\) lies in \(\nabla f(Z_k) + \alpha \rho\, \partial \|W_{k+1}\|_1\) and vanishes exactly at a minimiser, so it measures the violation of the optimality (KKT) conditions in the units of the gradient. The loop stops once, for every target column \(j\),

\[\max_i |G_k|_{ij} \le \mathrm{tol} \cdot \max_i |X_c^\top Y_c|_{ij} / n,\]

i.e. the residual is small relative to the gradient at \(W = 0\). The test is invariant to rescaling X and Y, and unlike a test on the step \(|W_{k+1} - Z_k|\) it does not stop early when the curvature is small (more features than samples, ill-conditioned X). The coefficient error is roughly the residual divided by the smallest curvature on the active set, so for very ill-conditioned problems lower tol (and raise max_iter).

A penalty that is only known at run time (traced under jax.jit or jax.vmap) cannot raise: a negative alpha (or an l1_ratio outside [0, 1] that makes a penalty negative) returns NaN coefficients and intercept, without iterating.

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.])
>>> res = braintools.optim.elastic_net_regression(X, Y, alpha=0.1, l1_ratio=0.7)
>>> print(jnp.round(res.coef, 2))
[ 2.79 -0.    0.   -1.88  0.  ]