6.4 Many predictors: ridge and lasso

With many predictors, least squares fits the noise. Ridge and lasso add a price for large coefficients, which pulls them towards zero. A little bias buys a big drop in variance.

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.axes import Axes
from sklearn import datasets
from sklearn.base import BaseEstimator
from sklearn.cross_decomposition import PLSRegression
from sklearn.decomposition import PCA
from sklearn.linear_model import Lasso, LassoCV, LinearRegression, Ridge, RidgeCV, lasso_path
from sklearn.model_selection import KFold, cross_val_score
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import PolynomialFeatures, StandardScaler

from statly import style

style.use()
rng = np.random.default_rng(1)

diabetes = datasets.load_diabetes(as_frame=True, scaled=False)
X, y = diabetes.data, diabetes.target.to_numpy()  # 442 patients: 10 measurements, progression
expand = PolynomialFeatures(degree=2, include_bias=False).fit(X)
names = expand.get_feature_names_out(X.columns)
useful = names != "sex^2"  # sex is coded 1 or 2, so its square adds nothing new
X_many = expand.transform(X)[:, useful]  # the measurements, their squares and pairwise products
NAMES = names[useful]
Z = StandardScaler().fit_transform(X_many)  # every predictor on the same scale
FOLDS = KFold(10, shuffle=True, random_state=1)


def cv_error(model: BaseEstimator, features: np.ndarray) -> float:
    """10-fold cross-validated prediction error (RMSE)."""
    return float(
        -cross_val_score(
            model, features, y, cv=FOLDS, scoring="neg_root_mean_squared_error"
        ).mean()
    )

6.4.1 When there are too many predictors

The diabetes patients have 10 measurements. A curve or an interaction might matter, so add every square and every product of two measurements: 64 predictors for 442 patients. Least squares fits the training data better with every predictor it gets, but on new patients the 64-predictor model is worse than the 10-predictor one:

ridge_cv = make_pipeline(StandardScaler(), RidgeCV(alphas=np.logspace(-1, 4, 50)))
lasso_cv = make_pipeline(StandardScaler(), LassoCV(cv=FOLDS, max_iter=50_000))
ERRORS = [
    ("least squares, 10", cv_error(LinearRegression(), X.to_numpy())),
    ("least squares, 64", cv_error(LinearRegression(), X_many)),
    ("ridge, 64", cv_error(ridge_cv, X_many)),
    ("lasso, 64", cv_error(lasso_cv, X_many)),
]
lasso_cv.fit(X_many, y)  # once more on all patients, to read off the chosen penalty


def plot_bars(ax: Axes) -> None:
    """One bar per model: its cross-validated error."""
    names, values = zip(*ERRORS)
    colors = [style.GREY, style.CORAL, style.BLUE, style.GREEN]
    ax.barh(names, values, color=colors, alpha=0.7)
    for i, value in enumerate(values):
        ax.text(value + 0.3, i, f"{value:.1f}", va="center", fontsize=9)
    ax.invert_yaxis()
    ax.set(xlabel="cross-validated error (RMSE)", xlim=(50, max(values) + 3))


fig, ax = plt.subplots(figsize=(8, 2.6))
plot_bars(ax)
plt.show()

Figure 6.4.1: Cross-validated error of least squares with the 10 measurements and with all 64 predictors, and of ridge and lasso with all 64 (penalty chosen by cross-validation). The extra predictors hurt least squares; the penalised models use them without paying for it.

With 64 coefficients to set from 442 patients, least squares bends the fit to the noise: variance, in the language of Chapter 1.3. The solution is not to throw predictors away by hand, but to make the model cautious about every coefficient.

6.4.2 A price for large coefficients

Least squares minimises the sum of squared residuals. Ridge regression minimises the same sum plus a penalty: λ times the sum of the squared coefficients. Lasso uses λ times the sum of their absolute sizes instead:

$$\text{ridge: } \sum_i \Big(y_i - \beta_0 - \sum_j \beta_j x_{ij}\Big)^2 + \lambda \sum_j \beta_j^2, \qquad \text{lasso: } \sum_i \Big(y_i - \beta_0 - \sum_j \beta_j x_{ij}\Big)^2 + \lambda \sum_j |\beta_j|$$

(the intercept is not penalised). Large coefficients now cost something, so the fit only uses them where they earn their keep. λ sets the price: 0 gives least squares; a huge λ sets all coefficients to zero.

The two penalties behave differently, and a picture with two coefficients shows why:

LEAST_SQUARES = np.array([2.0, 0.8])
SHAPE = np.array([[1.0, 0.45], [0.45, 0.6]])  # how the squared error grows around its minimum
b1, b2 = np.meshgrid(np.linspace(-1.6, 3.2, 300), np.linspace(-1.6, 3.2, 300))
offsets = np.stack([b1 - LEAST_SQUARES[0], b2 - LEAST_SQUARES[1]], axis=-1)
error = np.einsum("...i,ij,...j->...", offsets, SHAPE, offsets)
angles = np.linspace(0, 2 * np.pi, 2000)


def boundary_point(points: np.ndarray) -> np.ndarray:
    """The point on the region's edge with the smallest squared error."""
    d = points - LEAST_SQUARES
    return points[np.argmin(np.einsum("ni,ij,nj->n", d, SHAPE, d))]


def plot_region(ax: Axes, points: np.ndarray, title: str) -> None:
    """Error contours, the penalty region, and where they meet."""
    ax.contour(
        b1,
        b2,
        error,
        levels=[0.1, 0.4, 0.9, 1.6, 2.5, 3.6],
        colors=style.INK,
        linewidths=0.8,
        alpha=0.6,
    )
    ax.fill(points[:, 0], points[:, 1], color=style.BLUE, alpha=0.2)
    ax.plot(*LEAST_SQUARES, "o", ms=6, color=style.INK)
    ax.plot(*boundary_point(points), "o", ms=8, color=style.CORAL)
    ax.axhline(0, color=style.INK, lw=0.8)
    ax.axvline(0, color=style.INK, lw=0.8)
    ax.set(
        title=title,
        xlabel="coefficient 1",
        ylabel="coefficient 2",
        aspect="equal",
        xlim=(-1.6, 3.2),
        ylim=(-1.6, 3.2),
    )


circle = np.c_[np.cos(angles), np.sin(angles)]
diamond = circle / np.abs(circle).sum(axis=1, keepdims=True)
fig, axes = style.subplots(2, height=3.6, width=8.4)
plot_region(axes[0], circle, "Ridge: a circle")
plot_region(axes[1], diamond, "Lasso: a diamond")
fig.tight_layout()
plt.show()

Figure 6.4.2: Two coefficients. The ellipses are contours of the squared error, smallest at the least squares estimate (dot). The penalty limits the coefficients to the shaded region: a circle for ridge, a diamond for lasso. The best fit is where the ellipses first touch the region (red). The diamond’s corners lie on the axes, so lasso often lands on one, setting a coefficient to exactly zero.

Ridge shrinks every coefficient a little but keeps them all. Lasso shrinks too, and sets some to exactly zero: it selects predictors. That makes lasso models easier to read.

Because the penalty adds up coefficient sizes, the predictors must be on the same scale first, or a predictor measured in grams would be punished more than one in kilograms. Standardise before fitting.

6.4.3 The paths

As λ grows, the coefficients travel towards zero. Following each one gives its path:

alphas_ridge = np.logspace(-1, 5, 80)
ridge_paths = np.array([Ridge(alpha=a).fit(Z, y).coef_ for a in alphas_ridge])
alphas_lasso, lasso_paths, _ = lasso_path(
    Z, y - y.mean(), alphas=np.logspace(-1.5, 1.6, 80), max_iter=50_000
)
chosen = lasso_cv[-1].alpha_


def plot_ridge(ax: Axes) -> None:
    """Ridge coefficient paths."""
    ax.plot(alphas_ridge, ridge_paths, color=style.BLUE, lw=0.8, alpha=0.6)
    ax.set_xscale("log")
    ax.invert_xaxis()
    ax.set(title="Ridge", xlabel="penalty λ (smaller to the right)", ylabel="coefficient")


def plot_lasso(ax: Axes) -> None:
    """Lasso coefficient paths, with the chosen penalty and the names of the survivors."""
    ax.plot(alphas_lasso, lasso_paths.T, color=style.GREEN, lw=0.8, alpha=0.6)
    ax.axvline(chosen, color=style.INK, ls="--", lw=1)
    ax.set_xscale("log")
    ax.invert_xaxis()
    ax.set(title="Lasso", xlabel="penalty λ (smaller to the right)")


fig, axes = style.subplots(2, height=3.4, width=10, sharey=True)
plot_ridge(axes[0])
plot_lasso(axes[1])
fig.tight_layout()
plt.show()

Figure 6.4.3: Coefficients of the 64 standardised predictors as the penalty grows (to the left). Left: ridge shrinks them all smoothly. Right: lasso drops them one by one to exactly zero; the dashed line is the penalty chosen by cross-validation.

6.4.4 Choosing the penalty

λ is a knob that the data cannot set by least squares, since λ = 0 always fits the training data best. Cross-validation sets it: try a range of λ, and keep the one with the lowest error on the held-out folds (or the simplest within one standard error, Chapter 6.3).

fold_rmse = np.sqrt(lasso_cv[-1].mse_path_)  # one row per λ, one column per fold
path_errors = fold_rmse.mean(axis=1)
path_se = fold_rmse.std(axis=1, ddof=1) / np.sqrt(FOLDS.n_splits)
alphas = lasso_cv[-1].alphas_
kept = [(Lasso(alpha=a, max_iter=50_000).fit(Z, y).coef_ != 0).sum() for a in alphas[::10]]


def plot_lambda(ax: Axes) -> None:
    """The cross-validation curve over λ, with the count of predictors kept."""
    ax.fill_between(
        alphas, path_errors - path_se, path_errors + path_se, color=style.GREEN, alpha=0.2, lw=0
    )
    ax.plot(alphas, path_errors, color=style.GREEN)
    ax.axvline(chosen, color=style.INK, ls="--", lw=1)
    ax.set_xscale("log")
    ax.invert_xaxis()
    top = ax.secondary_xaxis("top")
    top.set_xticks(alphas[::10], [str(k) for k in kept])
    top.set_xlabel("predictors kept", fontsize=9)
    ax.set(xlabel="penalty λ (smaller to the right)", ylabel="cross-validated error (RMSE)")


fig, ax = plt.subplots(figsize=(8, 3.4))
plot_lambda(ax)
plt.show()

Figure 6.4.4: Cross-validated error of the lasso across penalties, with ± one standard error, and the number of predictors left in the model (top). Too little penalty overfits, too much leaves too few predictors.

At the chosen penalty, the lasso keeps 17 of the 64 predictors. The strongest are BMI times blood pressure, sex, age times sex and s5 (a blood fat); 47 squares and products were dropped as noise. Here the extra predictors carry little real signal, so the penalised models only match the plain 10-measurement model. But they get there on their own, without anyone choosing predictors by hand, and that is the point: with hundreds of candidates, choosing by hand is hopeless.

Ridge or lasso? When many predictors each matter a little, ridge often predicts slightly better. When a few matter and the rest are noise, lasso does, and it gives a readable model. The elastic net mixes the two penalties.

6.4.5 Choosing subsets

Before penalties, the standard was subset selection: try models with different sets of predictors and keep the best. With the 10 original diabetes measurements there are $2^{10} = 1{,}024$ possible models, few enough to fit them all: best subset selection. For each size, keep the subset with the smallest residual sum of squares; then choose the size.

The training RSS always prefers the biggest model, so the size must be chosen by something that charges for predictors. With d predictors, n cases and $\hat\sigma^2$ the noise variance estimated from the full model:

$$C_p = \frac{1}{n}\big(\text{RSS} + 2d\hat\sigma^2\big), \qquad \text{BIC} = \frac{1}{n}\big(\text{RSS} + \log(n)\, d\hat\sigma^2\big), \qquad R^2_{\text{adj}} = 1 - \frac{\text{RSS}/(n - d - 1)}{\text{TSS}/(n - 1)}$$

For least squares, $C_p$ and the AIC of Chapter 2.3 choose the same model.

X10 = X.to_numpy()
y_all = y
NAMES10 = list(X.columns)
full_rss = ((y_all - LinearRegression().fit(X10, y_all).predict(X10)) ** 2).sum()
sigma2 = full_rss / (len(y_all) - X10.shape[1] - 1)
tss = ((y_all - y_all.mean()) ** 2).sum()
results = []
for mask in range(1, 2**10):
    cols = [j for j in range(10) if mask >> j & 1]
    rss = (
        (y_all - LinearRegression().fit(X10[:, cols], y_all).predict(X10[:, cols])) ** 2
    ).sum()
    results.append((len(cols), rss, cols))
sizes = np.arange(1, 11)
best = {}
for d in sizes:
    of_size = [r for r in results if r[0] == d]
    best[d] = of_size[int(np.argmin([r[1] for r in of_size]))]


def plot_all_models(ax: Axes) -> None:
    """RSS of every model by size, with the best per size."""
    ax.scatter(
        [r[0] + rng.uniform(-0.15, 0.15) for r in results],
        [r[1] / 1e6 for r in results],
        s=3,
        color=style.GREY,
        alpha=0.4,
    )
    ax.plot(sizes, [best[d][1] / 1e6 for d in sizes], "o-", color=style.CORAL, ms=5)
    ax.set(
        title="1,024 models",
        xlabel="number of predictors",
        ylabel="RSS (millions)",
        xticks=sizes,
    )


def plot_criteria(ax: Axes) -> None:
    """Cp, BIC and adjusted R² for the best model of each size, rescaled to 0 (worst) to 1 (best)."""
    n = len(y_all)
    rss = np.array([best[d][1] for d in sizes])
    criteria = {
        "Cp (AIC)": -(rss + 2 * sizes * sigma2) / n,
        "BIC": -(rss + np.log(n) * sizes * sigma2) / n,
        "adjusted R²": 1 - (rss / (n - sizes - 1)) / (tss / (n - 1)),
    }
    for (name, values), color in zip(criteria.items(), (style.BLUE, style.GREEN, style.GOLD)):
        scaled = (values - values.min()) / (values.max() - values.min())
        ax.plot(sizes, scaled, "-", color=color, label=name)
        top = int(np.argmax(values))
        ax.plot(sizes[top], 1, "o", ms=9, color=color)
    ax.legend(fontsize=8, loc="lower right")
    ax.set(
        title="Choosing the size",
        xlabel="number of predictors",
        ylabel="criterion (best = 1)",
        xticks=sizes,
    )


fig, axes = style.subplots(2, height=3.2, width=9.5)
plot_all_models(axes[0])
plot_criteria(axes[1])
fig.tight_layout()
plt.show()

Figure 6.4.5: Best subset selection on the 10 diabetes measurements. Left: the residual sum of squares of all 1,024 models (grey) by size, and the best of each size (red). Right: Cp, BIC and adjusted R² of the best model of each size, each rescaled so its best value is at the top; the dots mark the size each would choose. Adjusted R² picks 8 predictors, Cp 6 and BIC 5: BIC charges the most per predictor, adjusted R² the least.

With 30 predictors there would be a billion subsets. Forward stepwise selection starts with no predictors and adds the most helpful one at a time; backward starts with all and removes the least helpful. They fit only about p²/2 models. All subset methods make hard in-or-out decisions that jump around from sample to sample, and the p-values of the final model are too optimistic, because the same data chose and tested it. The lasso makes the same kind of choice smoothly.

6.4.6 Fewer, combined predictors: PCR and PLS

A third way to tame many predictors is to replace them by a few combinations. Principal component regression (PCR) first finds the principal components of the predictors (Chapter 7.1), the directions in which they vary most, and regresses the outcome on the first M of them. Partial least squares (PLS) chooses its directions to vary a lot and to be related to the outcome. M is chosen by cross-validation, like λ.

components = np.arange(1, 31)
pcr_errors = [
    cv_error(make_pipeline(StandardScaler(), PCA(m), LinearRegression()), X_many)
    for m in components
]
pls_errors = [
    cv_error(make_pipeline(StandardScaler(), PLSRegression(m, scale=False)), X_many)
    for m in components
]

fig, ax = plt.subplots(figsize=(8, 3.2))
ax.plot(
    components, pcr_errors, "o-", ms=3, color=style.BLUE, label="principal component regression"
)
ax.plot(components, pls_errors, "o-", ms=3, color=style.CORAL, label="partial least squares")
ax.axhline(
    min(ERRORS[2][1], ERRORS[3][1]),
    color=style.INK,
    ls="--",
    lw=1,
    label="ridge and lasso (Section 6.4.1)",
)
ax.legend(fontsize=8)
ax.set(xlabel="components kept", ylabel="cross-validated error (RMSE)", ylim=(53, 62))
plt.show()

Figure 6.4.6: Cross-validated error of principal component regression (blue) and partial least squares (red) on the 64 diabetes predictors, by the number of components kept. PLS reaches its best with very few components because it looks at the outcome; PCR needs more, because the directions of largest spread in the predictors are not necessarily the ones that predict.

Both are close relatives of ridge regression, which shrinks the low-variance directions of the predictors most; PCR simply drops them. None of the three selects original predictors, so the result is harder to read than a lasso.

6.4.7 When there are more predictors than cases

Genetics, text and images routinely give more measurements than cases: p > n. Then least squares breaks completely. With as many predictors as cases, it can fit any outcome perfectly, even pure noise:

N_CASES = 50
noise_y = rng.normal(size=N_CASES)
noise_X = rng.normal(size=(N_CASES, 49))
new_y = rng.normal(size=500)
new_X = rng.normal(size=(500, 49))
counts = np.arange(1, 50, 2)
train_r2, test_ls, test_lasso = [], [], []
for p in counts:
    ls = LinearRegression().fit(noise_X[:, :p], noise_y)
    train_r2.append(ls.score(noise_X[:, :p], noise_y))
    test_ls.append(np.sqrt(np.mean((new_y - ls.predict(new_X[:, :p])) ** 2)))
    lasso = LassoCV(cv=5).fit(noise_X[:, :p], noise_y)
    test_lasso.append(np.sqrt(np.mean((new_y - lasso.predict(new_X[:, :p])) ** 2)))

fig, axes = style.subplots(2, height=3.1, width=9.5)
axes[0].plot(counts, train_r2, "o-", ms=3, color=style.CORAL)
axes[0].set(title="Training fit", xlabel="noise predictors", ylabel="training R²")
axes[1].plot(counts, test_ls, "o-", ms=3, color=style.CORAL, label="least squares")
axes[1].plot(counts, test_lasso, "o-", ms=3, color=style.BLUE, label="lasso")
axes[1].set_yscale("log")
axes[1].legend(fontsize=8)
axes[1].set(title="New cases", xlabel="noise predictors", ylabel="test error (log scale)")
fig.tight_layout()
plt.show()

Figure 6.4.7: Simulated: 50 cases, an outcome that is pure noise, and an increasing number of noise predictors. Left: least squares training R² climbs to 1 as the predictors approach 50, and its training error to 0. Right: its error on new cases explodes near p = n. The lasso, choosing λ by cross-validation, keeps the test error near that of predicting the average.

In high dimensions, training R², p-values and standard errors from least squares are meaningless, and every association needs checking on new data. Penalised methods, cross-validation and a sceptical eye are the tools. Even then, with thousands of predictors, many will look useful by chance (Chapter 3.10).

6.4.8 In one picture

• With many predictors, least squares overfits: low training error, high test error.
• Ridge penalises squared coefficients: it shrinks all of them, keeps all of them.
• Lasso penalises absolute coefficients: it shrinks and sets some to zero.
• Standardise predictors first; choose λ by cross-validation.
• Best subset and stepwise selection choose predictors in or out; size by Cp/AIC, BIC or adjusted R², or cross-validation.
• PCR and PLS regress on a few combinations of the predictors.
• With more predictors than cases, least squares fits noise perfectly; use penalties and new data.
• A little bias for a lot less variance is usually a good trade.

Next: curves, steps and splines, for relationships that are not straight.