Visualizing OLS and regularization

This notebook has three parts:

1. Two parameters: separate contour plots of the OLS loss for full-rank, nearly rank-deficient, and rank-deficient designs.
2. Projection geometry: separate two-dimensional diagrams showing what the OLS projector does.
3. Highly overparameterized regression: a random design with a controllably sparse true vector, followed by interactive ridge and lasso estimates across the regularization parameter $\lambda$.

The numerical routines are provided. Students should change the clearly marked parameters and interpret the figures.

# Install Plotly in the current notebook kernel only if needed.
import sys
import subprocess
import importlib.util

required = ["numpy", "plotly"]
missing = [package for package in required if importlib.util.find_spec(package) is None]
if missing:
    subprocess.check_call([sys.executable, "-m", "pip", "install", *missing])
    print("Installed:", ", ".join(missing))
else:
    print("All required packages are already installed.")

import numpy as np
import plotly.graph_objects as go
from plotly.subplots import make_subplots

np.set_printoptions(precision=3, suppress=True)
rng = np.random.default_rng(8)

Supplied numerical tools

The unregularized estimator is the minimum-norm OLS solution $X^+y$. We use

$$ \widehat\theta_{\mathrm{ridge}} =\arg\min_\theta\left\{ \frac{1}{2n}\|y-X\theta\|_2^2+\frac\lambda2\|\theta\|_2^2 \right\}, $$

and

$$ \widehat\theta_{\mathrm{lasso}} =\arg\min_\theta\left\{ \frac{1}{2n}\|y-X\theta\|_2^2+\lambda\|\theta\|_1 \right\}. $$

def ols(X, y):
    return np.linalg.pinv(X) @ y


def ridge(X, y, lam):
    d = X.shape[1]
    return np.linalg.solve(X.T @ X / len(y) + lam * np.eye(d), X.T @ y / len(y))


def soft_threshold(value, threshold):
    return np.sign(value) * np.maximum(np.abs(value) - threshold, 0.0)


def lasso_cd(X, y, lam, initial=None, max_iter=5000, tol=1e-9):
    n, d = X.shape
    beta = np.zeros(d) if initial is None else initial.copy()
    norms = np.sum(X**2, axis=0) / n
    for _ in range(max_iter):
        old = beta.copy()
        for j in range(d):
            residual_without_j = y - X @ beta + X[:, j] * beta[j]
            correlation = X[:, j] @ residual_without_j / n
            beta[j] = soft_threshold(correlation, lam) / norms[j]
        if np.linalg.norm(beta - old) <= tol * (1 + np.linalg.norm(old)):
            break
    return beta


def mse(y, prediction):
    return np.mean((y - prediction) ** 2)


def loss_grid(X, y, axis):
    theta_1, theta_2 = np.meshgrid(axis, axis)
    prediction = (X[:, 0, None, None] * theta_1
                  + X[:, 1, None, None] * theta_2)
    loss = 0.5 * np.mean((y[:, None, None] - prediction) ** 2, axis=0)
    return theta_1, theta_2, loss

1. OLS loss landscapes with two parameters

Each design has two columns. The three figures are deliberately separate:

• full rank: one elliptical minimum;
• nearly rank deficient: a long, narrow valley;
• rank deficient: a flat line of minimizers.

The red point is the pseudoinverse solution. In the rank-deficient case, red dashed points show other minimizers with the same loss.

u = np.array([1.0, 0.0, 1.0]) / np.sqrt(2)
v = np.array([0.0, 1.0, 0.0])
y_2d = np.array([1.2, -0.4, 2.0])
parameter_axis = np.linspace(-4, 4, 160)

def correlated_design(correlation):
    second = correlation * u + np.sqrt(max(0.0, 1 - correlation**2)) * v
    return np.column_stack([u, second])

cases = [
    ("Full rank", 0.25),
    ("Full rank", 0.5),
    ("Full rank", 0.75),
    ("Nearly rank deficient", 0.9),
    ("Nearly rank deficient", 0.995),
    ("Rank deficient", 1.0),
]

for title, correlation in cases:
    X = correlated_design(correlation)
    theta_1, theta_2, loss = loss_grid(X, y_2d, parameter_axis)
    theta_hat = ols(X, y_2d)
    singular_values = np.linalg.svd(X, compute_uv=False)

    fig = go.Figure()
    fig.add_trace(go.Contour(
        x=parameter_axis, y=parameter_axis, z=loss,
        contours=dict(coloring="lines", showlabels=True),
        line=dict(width=2), colorscale="Viridis", colorbar_title="loss"
    ))
    fig.add_trace(go.Scatter(
        x=[theta_hat[0]], y=[theta_hat[1]], mode="markers",
        marker=dict(size=12, color="red"), name="pseudoinverse solution"
    ))
    if np.linalg.matrix_rank(X) == 1:
        shifts = np.linspace(-2.5, 2.5, 9)
        other = np.column_stack([theta_hat[0] + shifts, theta_hat[1] - shifts])
        fig.add_trace(go.Scatter(
            x=other[:, 0], y=other[:, 1], mode="markers",
            marker=dict(size=7, color="red", symbol="x"),
            name="other minimizers"
        ))
    fig.update_layout(
        title=(f"{title}: correlation={correlation}, rank={np.linalg.matrix_rank(X)}, "
               f"singular values={np.round(singular_values, 3)}"),
        xaxis_title="theta 1", yaxis_title="theta 2",
        width=720, height=580
    )
    fig.show()

Questions

1. Compare the spacing of the isolevel curves. In which direction does curvature disappear?
2. How do the singular values explain the change from ellipses to a flat valley?
3. In the rank-deficient case, verify that every red cross gives the same fitted vector (X\theta).
4. Why is the pseudoinverse point the closest minimizer to the origin?

2. Full-rank plane and rank-deficient line in 3D

Each observation is displayed as $(x_{i1},x_{i2},y_i)$, with the response on the vertical axis. For the full-rank design, choose $\theta^\star\in\mathbb R^2$ and generate

$$y_i=x_i^\top\theta^\star.$$

The graph of the model is a plane, and all noiseless observations lie exactly on it. For the rank-deficient design, feature vectors are restricted to one direction. The model becomes a line in 3D. We retain the same response values and use OLS to project them onto the responses attainable along that line.

geometry_rng = np.random.default_rng(16)
number_of_points = 14
theta_star_2d = np.array([1.4, -0.8])

# Full-rank features and noiseless responses y = X theta_star.
X_full = geometry_rng.uniform(-2.2, 2.2, size=(number_of_points, 2))
y_same = X_full @ theta_star_2d
theta_full = ols(X_full, y_same)
yhat_full = X_full @ theta_full

grid_values = np.linspace(-2.5, 2.5, 35)
grid_x1, grid_x2 = np.meshgrid(grid_values, grid_values)
grid_y = theta_star_2d[0] * grid_x1 + theta_star_2d[1] * grid_x2

fig = go.Figure()
fig.add_trace(go.Surface(x=grid_x1, y=grid_x2, z=grid_y,
                         opacity=0.55, colorscale="Blues", showscale=False,
                         name="y = x transpose theta star"))
fig.add_trace(go.Scatter3d(x=X_full[:, 0], y=X_full[:, 1], z=y_same,
                           mode="markers", marker=dict(size=6, color="crimson"),
                           name="observations"))
fig.update_layout(
    title=(f"Full rank: observations lie on y=xᵀθ★; "
           f"θ★={theta_star_2d}, OLS={np.round(theta_full, 3)}"),
    scene=dict(xaxis_title="x 1", yaxis_title="x 2", zaxis_title="y"),
    width=820, height=650)
fig.show()

# Collapse features onto one direction, while retaining the same y vector.
feature_direction = np.array([1.0, 0.65])
feature_direction /= np.linalg.norm(feature_direction)
coordinates = X_full @ feature_direction
X_rank1 = coordinates[:, None] * feature_direction[None, :]

theta_rank1 = ols(X_rank1, y_same)
yhat_rank1 = X_rank1 @ theta_rank1

t_line = np.linspace(coordinates.min() - 0.5, coordinates.max() + 0.5, 150)
x_line = t_line[:, None] * feature_direction[None, :]
y_line = x_line @ theta_rank1

fig = go.Figure()
fig.add_trace(go.Scatter3d(x=x_line[:, 0], y=x_line[:, 1], z=y_line,
                           mode="lines", line=dict(width=8, color="royalblue"),
                           name="rank-one fitted model"))
fig.add_trace(go.Scatter3d(x=X_rank1[:, 0], y=X_rank1[:, 1], z=y_same,
                           mode="markers", marker=dict(size=6, color="crimson"),
                           name="same observed y values"))
fig.add_trace(go.Scatter3d(x=X_rank1[:, 0], y=X_rank1[:, 1], z=yhat_rank1,
                           mode="markers", marker=dict(size=5, color="green"),
                           name="projected fitted values"))

for i in range(number_of_points):
    fig.add_trace(go.Scatter3d(
        x=[X_rank1[i, 0], X_rank1[i, 0]],
        y=[X_rank1[i, 1], X_rank1[i, 1]],
        z=[y_same[i], yhat_rank1[i]],
        mode="lines", line=dict(color="black", width=2, dash="dash"),
        showlegend=(i == 0), name="residual"))

fig.update_layout(
    title=(f"Rank deficient: same y projected onto a fitted line; "
           f"rank(X)={np.linalg.matrix_rank(X_rank1)}"),
    scene=dict(xaxis_title="x 1", yaxis_title="x 2", zaxis_title="y"),
    width=820, height=650)
fig.show()

print("Full-rank maximum absolute residual:", np.max(np.abs(y_same - yhat_full)))
print("Rank-deficient normal-equation residual X.T @ (y-yhat):",
      X_rank1.T @ (y_same - yhat_rank1))
print("Ranks:", np.linalg.matrix_rank(X_full), np.linalg.matrix_rank(X_rank1))

Questions

1. Why do all red points lie exactly on the plane in the full-rank figure?
2. In the rank-deficient figure, why must every fitted green point lie on one line?
3. The black segments are vertical in the displayed coordinates. In which space is OLS performing an orthogonal projection?
4. Verify that $X^\top(y-\widehat y)=0$ although individual residuals are nonzero.
5. Can the rank-deficient model recover $\theta^\star$ uniquely? Can it still produce a unique fitted response vector?
6. Change feature_direction. How do the line and projection error change?

3. A highly overparameterized sparse model

We now use one high-dimensional problem throughout:

$$n<d,\qquad y=X\theta^\star+\varepsilon.$$

The true vector has a controllable number of nonzero entries. Change sparsity, rerun the remaining cells, and compare ridge with lasso.

# Main parameters to manipulate.
n_train = 35
n_test = 1500
d = 140
sparsity = 8       # Try 3, 8, 20, 60.
noise_sd = 0.7
data_seed = 24

data_rng = np.random.default_rng(data_seed)
X_train = data_rng.normal(size=(n_train, d))
X_test = data_rng.normal(size=(n_test, d))

# Standardize using training-column scales.
column_scales = X_train.std(axis=0)
X_train = X_train / column_scales
X_test = X_test / column_scales

support = data_rng.choice(d, size=sparsity, replace=False)
theta_true = np.zeros(d)
theta_true[support] = data_rng.choice([-1.0, 1.0], sparsity) * data_rng.uniform(1.0, 2.5, sparsity)

y_train_clean = X_train @ theta_true
y_train = y_train_clean + noise_sd * data_rng.normal(size=n_train)
y_test_clean = X_test @ theta_true
theta_ols = ols(X_train, y_train)

fig = go.Figure()
fig.add_trace(go.Scatter(x=np.arange(d), y=theta_true, mode="markers",
                         marker=dict(size=8, color=np.where(theta_true != 0, "crimson", "lightgray")),
                         name="true coefficients"))
fig.update_layout(title=f"Sparse truth: d={d}, n={n_train}, sparsity={sparsity}",
                  xaxis_title="coefficient index", yaxis_title="coefficient value",
                  width=1000, height=430)
fig.show()

print("rank(X) =", np.linalg.matrix_rank(X_train))
print("number of parameters minus observations =", d - n_train)
print("minimum-norm OLS training MSE =", mse(y_train, X_train @ theta_ols))
print("minimum-norm OLS clean test MSE =", mse(y_test_clean, X_test @ theta_ols))

Questions before regularizing

1. Why can OLS interpolate the noisy training responses when (d>n)?
2. Is the minimum-norm OLS estimate sparse?
3. Compare the true support with the largest OLS coefficients.
4. Increase sparsity. Which regularizer do you expect to become relatively more attractive?

4. Ridge estimates across (\lambda)

The slider displays the complete ridge estimate for each (lambda). The black crosses are the true coefficients and the gray points are unregularized minimum-norm OLS.

Small (lambda) stays near OLS; large (lambda) shrinks every coefficient toward zero.

ridge_lambdas = np.logspace(-4, 1.2, 24)
ridge_estimates = np.array([ridge(X_train, y_train, lam) for lam in ridge_lambdas])
ridge_train = np.array([mse(y_train, X_train @ beta) for beta in ridge_estimates])
ridge_test = np.array([mse(y_test_clean, X_test @ beta) for beta in ridge_estimates])
ridge_parameter_error = np.mean((ridge_estimates - theta_true[None, :])**2, axis=1)

ridge_frames = []
for index, lam in enumerate(ridge_lambdas):
    ridge_frames.append(go.Frame(
        name=str(index),
        data=[
            go.Scatter(x=np.arange(d), y=theta_true, mode="markers", name="truth",
                       marker=dict(color="black", size=8, symbol="x")),
            go.Scatter(x=np.arange(d), y=theta_ols, mode="markers", name="OLS min norm",
                       marker=dict(color="lightgray", size=5)),
            go.Scatter(x=np.arange(d), y=ridge_estimates[index], mode="markers", name="ridge",
                       marker=dict(color="royalblue", size=6))
        ],
        layout=go.Layout(title=(f"Ridge: lambda={lam:.4g}; train={ridge_train[index]:.3f}; "
                                f"test={ridge_test[index]:.3f}; parameter MSE={ridge_parameter_error[index]:.3f}"))
    ))

fig = go.Figure(data=ridge_frames[0].data, frames=ridge_frames)
fig.update_layout(
    xaxis_title="coefficient index", yaxis_title="coefficient value",
    width=1050, height=500,
    sliders=[dict(currentvalue={"prefix": "ridge lambda: "},
                  steps=[dict(method="animate", label=f"{lam:.3g}",
                              args=[[str(i)], {"mode": "immediate", "frame": {"duration": 0}}])
                         for i, lam in enumerate(ridge_lambdas)])]
)
fig.show()

fig = go.Figure()
fig.add_trace(go.Scatter(x=ridge_lambdas, y=ridge_train, mode="lines+markers", name="train MSE"))
fig.add_trace(go.Scatter(x=ridge_lambdas, y=ridge_test, mode="lines+markers", name="clean test MSE"))
fig.add_trace(go.Scatter(x=ridge_lambdas, y=ridge_parameter_error,
                         mode="lines+markers", name="parameter MSE"))
fig.update_layout(title="Ridge errors across lambda", xaxis_type="log",
                  xaxis_title="lambda", yaxis_title="error", width=800, height=480)
fig.show()

Ridge questions

1. Which (lambda) minimizes training error, clean test error, and parameter error? Are they identical?
2. Does ridge recover exact zeros?
3. Why can moderate shrinkage improve test error even though it worsens training error?
4. What happens to the estimates as (lambda\to\infty)?

5. Lasso estimates across (\lambda)

Lasso is computed from strong to weak regularization with warm starts. Its nonsmooth penalty can set coefficients exactly to zero. The title reports the number of active coefficients.

lasso_lambdas = np.logspace(-3, 0.5, 24)
lasso_descending = []
current = np.zeros(d)
for lam in lasso_lambdas[::-1]:
    current = lasso_cd(X_train, y_train, lam, current, max_iter=4500)
    lasso_descending.append(current.copy())
lasso_estimates = np.array(lasso_descending[::-1])

lasso_train = np.array([mse(y_train, X_train @ beta) for beta in lasso_estimates])
lasso_test = np.array([mse(y_test_clean, X_test @ beta) for beta in lasso_estimates])
lasso_parameter_error = np.mean((lasso_estimates - theta_true[None, :])**2, axis=1)
lasso_active = np.sum(np.abs(lasso_estimates) > 1e-6, axis=1)

lasso_frames = []
for index, lam in enumerate(lasso_lambdas):
    lasso_frames.append(go.Frame(
        name=str(index),
        data=[
            go.Scatter(x=np.arange(d), y=theta_true, mode="markers", name="truth",
                       marker=dict(color="black", size=8, symbol="x")),
            go.Scatter(x=np.arange(d), y=theta_ols, mode="markers", name="OLS min norm",
                       marker=dict(color="lightgray", size=5)),
            go.Scatter(x=np.arange(d), y=lasso_estimates[index], mode="markers", name="lasso",
                       marker=dict(color="darkorange", size=6))
        ],
        layout=go.Layout(title=(f"Lasso: lambda={lam:.4g}; active={lasso_active[index]}; "
                                f"train={lasso_train[index]:.3f}; test={lasso_test[index]:.3f}; "
                                f"parameter MSE={lasso_parameter_error[index]:.3f}"))
    ))

fig = go.Figure(data=lasso_frames[0].data, frames=lasso_frames)
fig.update_layout(
    xaxis_title="coefficient index", yaxis_title="coefficient value",
    width=1050, height=500,
    sliders=[dict(currentvalue={"prefix": "lasso lambda: "},
                  steps=[dict(method="animate", label=f"{lam:.3g}",
                              args=[[str(i)], {"mode": "immediate", "frame": {"duration": 0}}])
                         for i, lam in enumerate(lasso_lambdas)])]
)
fig.show()

fig = make_subplots(rows=1, cols=2, subplot_titles=("Errors", "Sparsity"))
fig.add_trace(go.Scatter(x=lasso_lambdas, y=lasso_train, mode="lines+markers", name="train MSE"), row=1, col=1)
fig.add_trace(go.Scatter(x=lasso_lambdas, y=lasso_test, mode="lines+markers", name="clean test MSE"), row=1, col=1)
fig.add_trace(go.Scatter(x=lasso_lambdas, y=lasso_parameter_error,
                         mode="lines+markers", name="parameter MSE"), row=1, col=1)
fig.add_trace(go.Scatter(x=lasso_lambdas, y=lasso_active,
                         mode="lines+markers", name="active coefficients"), row=1, col=2)
fig.update_xaxes(type="log", title_text="lambda")
fig.update_layout(title="Lasso errors and sparsity across lambda", width=1000, height=480)
fig.show()

Lasso and comparison questions

1. At which (lambda) is the number of active coefficients closest to the true sparsity?
2. Does that same (lambda) minimize coefficient error or test error?
3. Compare the ridge and lasso sliders. Which method better reveals the true support here?
4. Increase sparsity to 60. Does lasso retain its advantage?
5. Increase the noise or introduce correlated columns. How stable is support recovery?
6. Explain why L1 regularization is naturally matched to a sparse truth, whereas L2 regularization is often better for a dense collection of small coefficients.