Linear regression

What it is and when you reach for it

Linear regression predicts a continuous target as a weighted sum of features: ŷ = w₁x₁ + … + w_dx_d + b. That restriction to a hyperplane is the whole point. It makes the model trainable in closed form, interpretable coefficient by coefficient, cheap enough to fit thousands of times in a hyperparameter sweep, and honest about what it cannot represent. Among its neighbors, it sits one step below logistic regression (same linear score, squashed through a sigmoid for classification), one step above predicting the mean (which is linear regression with no features), and at the base of everything with "linear layer" in its name: the final layer of nearly every deep network is a linear regression or classification head over learned features. You reach for it first when you need a baseline, when you need coefficients you can explain to a stakeholder or a reviewer, when you have far more rows than you have compute budget, and whenever the honest answer to "do we need a deep model here?" has not yet been established. A tuned gradient-boosted tree or a neural net that cannot beat a well-featurized linear model is telling you something about your data, and a linear baseline is the only way to hear it.

The math

Least squares from Gaussian maximum likelihood

Why squared error and not, say, absolute error? The clean answer comes from a noise model. Assume each observation is the linear signal plus independent Gaussian noise: y_i = x_i^Tw + ε_i with ε_i ~ N(0, σ²). The likelihood of the dataset is the product of Gaussian densities, and its logarithm is

log p(y | X, w) = −(n/2) log(2πσ²) − (1 / 2σ²) Σ_i (y_i − x_i^Tw)².

The first term does not involve w, and the second is a negative constant times the sum of squared residuals. Maximizing the likelihood over w is therefore exactly minimizing L(w) = ‖y − Xw‖²: least squares is not an arbitrary choice of loss, it is the maximum-likelihood estimator under Gaussian noise. The assumption is also the fine print. If the noise is heavy-tailed, a few outliers dominate the squared loss and drag the fit; under Laplace noise the same derivation yields absolute error and the median takes the place of the mean. When people say least squares is "sensitive to outliers", this is the precise sense: it is optimal for a noise distribution that makes outliers extremely improbable.

The normal equations

Stack the samples into X (n × d, with a column of ones appended so the intercept is just another weight) and expand the loss: L(w) = y^Ty − 2w^TX^Ty + w^TX^TXw. The gradient is ∇L = −2X^Ty + 2X^TXw, and setting it to zero gives the normal equations:

X^TX w = X^Ty.

Because X^TX is positive semidefinite, the loss is convex and any stationary point is a global minimum. The equations also have a geometry worth keeping in your head: they say X^T(y − Xw) = 0, that is, the residual is orthogonal to every column of X. The prediction Xw* is the orthogonal projection of y onto the column space of X, the closest point to y that the features can express.

        y
        |\
        | \   residual  r = y − Xw*
        |  \  (orthogonal to the plane: Xᵀr = 0)
        |   \
  ------+----*------------------------
       /    Xw*  = projection of y
      /     column space of X (all reachable predictions)
          

Why you do not invert X^TX: a worked conditioning example

The textbook next step is w = (X^TX)⁻¹ X^Ty, and production code should essentially never do it. The problem is conditioning. The condition number κ(A) = σ_max/σ_min measures how much a matrix amplifies relative error, and forming X^TX squares it: κ(X^TX) = κ(X)². Any algorithm that touches X^TX, whether it inverts it or solves it by Cholesky, pays the squared condition number; an algorithm that factorizes X directly pays only κ(X).

Here is the smallest example that shows the squaring destroying real information. Take two nearly collinear columns:

X = [[1, 1], [1, 1.0001]],  so   X^TX = [[2, 2.0001], [2.0001, 2.00020001]].

The singular values of X are σ₁ ≈ 2.00005 and σ₂ ≈ 5.0 × 10⁻⁵, so κ(X) ≈ 4 × 10⁴ and κ(X^TX) ≈ 1.6 × 10⁹. Now look at where the information actually lives. det(X^TX) = 2 × 2.00020001 − 2.0001² = 10⁻⁸: everything that distinguishes the two columns sits in the eighth decimal digit of the entry 2.00020001. Float32 carries about seven decimal digits, with a spacing of roughly 2.4 × 10⁻⁷ between representable numbers near 2.0, so merely storing X^TX in float32 rounds each entry by up to about 1.2 × 10⁻⁷, an error more than ten times larger than the determinant itself. The computed determinant of the stored matrix is rounding noise; it can even come out negative, at which point the "covariance" matrix is numerically indefinite and a Cholesky factorization simply fails. The rule-of-thumb bound says the relative error of a solve scales like κ times machine epsilon: through the normal equations that is κ²(X) · ε ≈ 1.6 × 10⁹ × 1.2 × 10⁻⁷ ≈ 190, meaning no correct digits at all, while a QR factorization of X pays κ(X) · ε ≈ 5 × 10⁻³, about two to three good digits. Same data, same precision, and the only difference is whether the algorithm ever formed X^TX. Squaring the condition number halves the number of digits you get to keep, and float32 does not have that many digits to spare.

The safe path is a factorization of X itself. QR writes X = QR with Q orthonormal and R upper triangular; then ‖y − Xw‖ = ‖Q^Ty − Rw‖ and the solution comes from back-substituting the triangular system Rw = Q^Ty. Orthogonal transformations do not change lengths, so nothing gets amplified, and Householder QR is backward stable. This is exactly what lstsq routines do: LAPACK's gelsy driver is QR with column pivoting, gelsd is an SVD-based solve that also handles rank deficiency gracefully, and both are reachable from torch.linalg.lstsq, while JAX's jnp.linalg.lstsq goes through the SVD. In practice the rule is short: to solve a least-squares problem, call lstsq; if you must solve the normal equations (for example because ridge already regularized them), use a Cholesky solve, and reserve explicit matrix inversion for the rare case where you need the inverse itself, such as coefficient covariance matrices in statistics.

Gradient descent: the scalable alternative

A QR factorization costs O(nd²), which is wonderful until n or d stops fitting. When the data streams from disk, when d is in the millions (text n-grams, hashed features), or when the "linear regression" is the last layer of a network being trained anyway, you trade the exact solve for iterations of gradient descent. With L(w) = (1/2n)‖Xw − y‖² the gradient is ∇L = X^T(Xw − y)/n, one matrix-vector product per step, and the update is w ← w − η ∇L. Convexity guarantees convergence for any step size below 2/L where L = σ_max²/n is the largest eigenvalue of the Hessian X^TX/n, and the convergence speed is governed by the same villain as before: the error contracts by roughly a factor (κ − 1)/(κ + 1) per step, where κ is the condition number of the Hessian. Badly scaled features make κ huge and gradient descent crawl, which is why the implementation below standardizes columns first; after standardization every feature has unit variance, the Hessian's eigenvalues are all O(1), and a single global learning rate works. Minibatch SGD is the same story with noisier gradients, and it is the version that scales to datasets that never fit in memory.

Ridge regression

When columns are nearly collinear, as in the worked example, the least-squares solution exists but is wildly sensitive: tiny noise in y swings the coefficients by enormous amounts along the ill-conditioned direction. Ridge regression adds a penalty, minimizing ‖y − Xw‖² + λ‖w‖², with solution (X^TX + λI) w = X^Ty. The effect on conditioning is direct: every eigenvalue of X^TX gains λ, so the condition number drops from σ_max²/σ_min² to (σ_max² + λ)/(σ_min² + λ), and directions the data barely determines get shrunk toward zero instead of exploding. Probabilistically, ridge is the MAP estimate under a Gaussian prior w ~ N(0, τ²I), with λ = σ²/τ²; the penalty is literally a prior belief that coefficients are small. Two practical rules follow from the math. The penalty is not scale-invariant, so standardize features before ridging or λ punishes coefficients for the units their features happen to be measured in. And never penalize the intercept: shrinking b toward zero would bias every prediction toward zero for no reason, which is why implementations either center the data first or exclude the intercept column from the penalty, as the code below does. Numerically, you still avoid the explicit inverse: ridge is exactly ordinary least squares on an augmented system, X stacked on top of √λ·I and y padded with zeros, so one lstsq call solves it stably.

Implementation, twice

First the closed-form path. Both versions append a ones column so the intercept is an ordinary weight, and both hand the problem to the library's lstsq, which factorizes X rather than forming X^TX. Ridge reuses the same routine through the augmented system, keeping the intercept column out of the penalty block. The JAX version turns on float64 explicitly, because JAX defaults to float32 and the conditioning section above is exactly the argument for why a least-squares solver wants the extra digits.

import torch

def add_intercept(X):
    """Append a ones column: the bias becomes just another weight."""
    ones = torch.ones(X.shape[0], 1, dtype=X.dtype, device=X.device)
    return torch.cat([X, ones], dim=1)

def fit_lstsq(X, y):
    """Least squares via torch.linalg.lstsq.

    lstsq factorizes X itself (QR with pivoting on CPU), never
    forming X^T X, so the solve pays cond(X) rather than cond(X)^2.
    Returns w of shape (d + 1,); the last entry is the intercept.
    """
    Xb = add_intercept(X)
    return torch.linalg.lstsq(Xb, y.unsqueeze(-1)).solution.squeeze(-1)

def fit_ridge(X, y, lam):
    """Ridge as plain least squares on an augmented system.

    min ||Xw - y||^2 + lam * ||w||^2  is exactly lstsq on
    [X; sqrt(lam) * I] with y padded by zeros. The intercept
    column stays out of the penalty block, so it is not shrunk.
    """
    n, d = X.shape
    Xb = add_intercept(X)
    pen = torch.zeros(d, d + 1, dtype=X.dtype, device=X.device)
    pen[:, :d] = lam ** 0.5 * torch.eye(d, dtype=X.dtype, device=X.device)
    A = torch.cat([Xb, pen])
    b = torch.cat([y, torch.zeros(d, dtype=X.dtype, device=X.device)])
    return torch.linalg.lstsq(A, b.unsqueeze(-1)).solution.squeeze(-1)

def predict(X, w):
    return add_intercept(X) @ w

import jax
jax.config.update("jax_enable_x64", True)  # least squares wants float64
import jax.numpy as jnp

def add_intercept(X):
    """Append a ones column: the bias becomes just another weight."""
    return jnp.concatenate([X, jnp.ones((X.shape[0], 1), X.dtype)], axis=1)

def fit_lstsq(X, y):
    """Least squares via jnp.linalg.lstsq (SVD-based).

    The SVD route factorizes X itself, never forming X^T X, so the
    solve pays cond(X) rather than cond(X)^2, and it handles
    rank-deficient X gracefully. Last entry of w is the intercept.
    """
    Xb = add_intercept(X)
    w, *_ = jnp.linalg.lstsq(Xb, y, rcond=None)
    return w

def fit_ridge(X, y, lam):
    """Ridge as plain least squares on an augmented system.

    min ||Xw - y||^2 + lam * ||w||^2  is exactly lstsq on
    [X; sqrt(lam) * I] with y padded by zeros. The intercept
    column stays out of the penalty block, so it is not shrunk.
    """
    n, d = X.shape
    Xb = add_intercept(X)
    pen = jnp.concatenate(
        [lam ** 0.5 * jnp.eye(d, dtype=X.dtype),
         jnp.zeros((d, 1), X.dtype)], axis=1)
    A = jnp.concatenate([Xb, pen])
    b = jnp.concatenate([y, jnp.zeros(d, X.dtype)])
    w, *_ = jnp.linalg.lstsq(A, b, rcond=None)
    return w

def predict(X, w):
    return add_intercept(X) @ w

Now the gradient-descent version, the shape of the computation that survives when the closed form does not scale. Both implementations standardize columns so one global learning rate works, pick that rate from the spectral norm of the standardized matrix (the safe step is anything below 2/L; the code uses 1/L), and fold the standardization back into raw-feature coefficients at the end so the returned weights are directly comparable to the closed form. The PyTorch tab leans on autograd, since in practice this loop is what an optimizer does inside a larger model; the JAX tab expresses the same loop as a lax.scan over jitted steps, which compiles the whole training loop into one XLA program.

import torch

def fit_gd(X, y, steps=500, lam=0.0):
    """Full-batch gradient descent on (1/2n)||Xw + b - y||^2.

    Standardizing columns first is what makes a single global
    learning rate work: afterwards the Hessian eigenvalues are
    all O(1), so cond is small and convergence is fast.
    """
    n, d = X.shape
    mu, sigma = X.mean(0), X.std(0).clamp_min(1e-12)
    Xs = (X - mu) / sigma

    # largest Hessian eigenvalue L = sigma_max(Xs)^2 / n; step 1/L
    L = torch.linalg.matrix_norm(Xs, ord=2) ** 2 / n
    lr = 1.0 / L

    w = torch.zeros(d, dtype=X.dtype, requires_grad=True)
    b = torch.zeros((), dtype=X.dtype, requires_grad=True)
    for _ in range(steps):
        resid = Xs @ w + b - y
        loss = (resid ** 2).mean() / 2 + lam * (w ** 2).sum() / 2
        loss.backward()
        with torch.no_grad():
            w -= lr * w.grad
            b -= lr * b.grad
        w.grad = None
        b.grad = None

    # fold standardization back into raw-feature coefficients
    with torch.no_grad():
        w_raw = w / sigma
        b_raw = b - (w_raw * mu).sum()
    return w_raw, b_raw

import jax
import jax.numpy as jnp

def fit_gd(X, y, steps=500, lam=0.0):
    """Full-batch gradient descent on (1/2n)||Xw + b - y||^2.

    Standardizing columns first is what makes a single global
    learning rate work: afterwards the Hessian eigenvalues are
    all O(1), so cond is small and convergence is fast.
    lax.scan compiles the whole loop into one XLA program.
    """
    n, d = X.shape
    mu = X.mean(0)
    sigma = jnp.maximum(X.std(0), 1e-12)
    Xs = (X - mu) / sigma

    # largest Hessian eigenvalue L = sigma_max(Xs)^2 / n; step 1/L
    L = jnp.linalg.norm(Xs, ord=2) ** 2 / n
    lr = 1.0 / L

    def loss(params):
        w, b = params
        resid = Xs @ w + b - y
        return (resid ** 2).mean() / 2 + lam * (w ** 2).sum() / 2

    def step(params, _):
        g_w, g_b = jax.grad(loss)(params)
        w, b = params
        return (w - lr * g_w, b - lr * g_b), None

    init = (jnp.zeros(d, X.dtype), jnp.zeros((), X.dtype))
    (w, b), _ = jax.lax.scan(step, init, None, length=steps)

    # fold standardization back into raw-feature coefficients
    w_raw = w / sigma
    return w_raw, b - w_raw @ mu

Note the regularization bookkeeping: the gradient-descent loss uses a mean over samples while the closed-form ridge uses a sum, so the same amount of shrinkage requires lam_gd = lam_lstsq/n. This factor-of-n mismatch is one of the most common reasons two ridge implementations "disagree".

Using it on a real shape of problem

A realistic tabular shape: fifty thousand rows, sixty-four features, moderate noise. The snippet fits both paths and checks that they agree with each other and with the ground truth. With noise standard deviation 0.5 and n = 50,000, the standard error of each coefficient is about 0.5/√n ≈ 0.002, so recovered weights should sit within a few thousandths of the truth; exact digits are machine- and seed-dependent.

import torch
torch.manual_seed(0)

n, d = 50_000, 64
X = torch.randn(n, d, dtype=torch.float64)
w_true = torch.randn(d, dtype=torch.float64)
y = X @ w_true + 3.0 + 0.5 * torch.randn(n, dtype=torch.float64)

w = fit_lstsq(X, y)
print((w[:-1] - w_true).abs().max())   # ~1e-2 or below: recovery
print(w[-1])                           # ~3.0: the intercept

w_gd, b_gd = fit_gd(X, y, steps=500)
print((w_gd - w[:-1]).abs().max())     # tiny: both paths agree
print(((predict(X, w) - y) ** 2).mean())  # ~0.25 = noise variance

import jax
jax.config.update("jax_enable_x64", True)
import jax.numpy as jnp

key = jax.random.PRNGKey(0)
kx, kw, ke = jax.random.split(key, 3)

n, d = 50_000, 64
X = jax.random.normal(kx, (n, d), jnp.float64)
w_true = jax.random.normal(kw, (d,), jnp.float64)
y = X @ w_true + 3.0 + 0.5 * jax.random.normal(ke, (n,), jnp.float64)

w = fit_lstsq(X, y)
print(jnp.abs(w[:-1] - w_true).max())  # ~1e-2 or below: recovery
print(w[-1])                           # ~3.0: the intercept

w_gd, b_gd = fit_gd(X, y, steps=500)
print(jnp.abs(w_gd - w[:-1]).max())    # tiny: both paths agree
print(((predict(X, w) - y) ** 2).mean())  # ~0.25 = noise variance

What to expect: the training mean squared error should flatten out near the true noise variance (0.25 here), and pushing it lower than that is by definition fitting noise. On this well-conditioned random design, gradient descent matches the closed form in a few hundred steps; if you delete the standardization and rescale one column by 1000, watch the same loop take thousands of steps or diverge, which is the condition-number story made visible.

Applications

Linear regression earns its keep first as the universal baseline. Any tabular prediction task, churn, revenue, delivery time, energy use, starts with a linear fit because it is the cheapest strong model and because everything more complex must justify itself against it. In forecasting, linear models with engineered features remain the production workhorse: short-term electricity load forecasting regresses demand on temperature, lagged load, and calendar dummies; retail demand and capacity planning do the same with promotions and seasonality features. Quantitative finance is built on it: CAPM and the Fama-French factor models are literally ordinary least squares of returns on factor returns, and "beta" is a regression coefficient. Experimentation platforms use regression adjustment to shrink A/B test variance, and covariate-adjustment schemes like CUPED are regression in light disguise. In the sciences it is the default instrument-calibration and effect-estimation tool, and econometrics is to a first approximation the study of when a regression coefficient can be read causally.

The application closest to modern deep learning research is the linear probe. To measure what a frozen network has learned, you extract activations at some layer and fit a linear model from them to a property of interest; the probe's accuracy measures how linearly accessible that information is. Alain and Bengio's linear classifier probes made this a standard diagnostic, the CLIP paper reported linear-probe accuracy across dozens of datasets as its headline representation-quality metric, and interpretability work on language models probes for syntax, position, and world state with small linear maps (Hewitt and Manning's structural probe is a single linear transformation). The linear-probe workflow is exactly the ridge regression on this page, applied to a design matrix of frozen activations, and knowing the closed form means a probe sweep over layers and penalties costs seconds, not GPU-hours. The same trick appears as the standard evaluation for self-supervised vision models (SimCLR, DINO and their descendants report linear evaluation on ImageNet features) and in fine-tuning practice, where training only the final linear head is the cheapest adaptation baseline.

Against the real libraries

The reference implementations above are a few dozen lines; here is what the production libraries add and when you should switch. scikit-learn's LinearRegression is the same computation, routed through LAPACK's SVD-based gelsd driver via SciPy, with the ergonomics that actually matter in practice: sparse input support, multi-target fitting, sample weights, a positive=True option that switches to non-negative least squares, and seamless composition with Pipeline, StandardScaler, and cross-validation. Its Ridge adds a menu of solvers (Cholesky, SVD, LSQR, sparse conjugate gradient, SAG/SAGA for huge n) chosen automatically by problem shape, and RidgeCV exploits the SVD to evaluate leave-one-out cross-validation across a whole grid of penalties at essentially the cost of one fit, something a from-scratch loop cannot touch. Note the objective convention: sklearn's alpha multiplies ‖w‖² against a sum-of-squares data term, so it matches the lam of the closed-form ridge above directly, and equals n times the lam of the gradient-descent version.

statsmodels answers a different question. Where sklearn optimizes prediction, statsmodels.api.OLS delivers inference: standard errors, t-statistics and p-values per coefficient, confidence intervals, R² and adjusted R², AIC/BIC, heteroskedasticity-robust covariance options (HC0 through HC3), and a summary table that flags a large design-matrix condition number, the same diagnosis this page worked by hand. If anyone will ask "is this coefficient significantly different from zero?" or "how uncertain is this effect?", fit with statsmodels; if the model only needs to predict, sklearn is leaner. The from-scratch version is genuinely enough when you need speed and control inside a research loop: fitting linear probes on GPU-resident activations, solving thousands of small regressions inside a bigger algorithm, or differentiating through the fit itself, which the JAX version supports for free.

Verification is one fixed-seed script: generate float64 data, then check np.allclose(fit_lstsq(X, y)[:-1], LinearRegression().fit(X, y).coef_, atol=1e-8) and the same for the intercept, and for ridge compare fit_ridge(X, y, alpha) against Ridge(alpha=alpha).fit(X, y). Agreement to eight decimals in float64 is the expected outcome, not a lucky one; anything worse means an objective mismatch (the factor of n, a penalized intercept, unstandardized penalty) rather than "numerical noise", and hunting the mismatch is the best conditioning exercise there is.

Traps and misconceptions

"Just use inv(X.T @ X)." The formula is correct mathematics and poor numerics: forming X^TX squares the condition number and the explicit inverse adds instability and cost on top. The worked example above loses every digit in float32 this way. Solve with lstsq or a factorization; compute an inverse only when the inverse itself is the deliverable.

"Ridge is scale-free." The penalty λ‖w‖² compares coefficients across features as if their units were comparable. Measure a feature in grams instead of kilograms and its coefficient shrinks a thousandfold less. Standardize first, and keep the intercept out of the penalty entirely, or the fit is biased toward predicting zero.

"High R² means a good model." R² only says the features explain in-sample variance. It rises monotonically as you add features, junk included, says nothing about out-of-sample performance or causality, and can be high while residuals show blatant structure. Judge on held-out error and residual plots, and treat coefficient stories as causal only with a design that supports it.

"The coefficients are unstable, so the model is broken." Under collinearity the individual coefficients can swing wildly between resamples while predictions barely move, because only the sum of the collinear terms is pinned down. That is a statement about identifiability, not brokenness. If you need stable, interpretable coefficients, ridge (or dropping redundant features) is the cure; if you only need predictions, the instability may not matter at all.

"Gradient descent didn't converge, so the problem is non-convex." The least-squares loss is convex; a diverging or crawling loop means the step size exceeds 2/L or the features are badly scaled, making the condition number enormous. Standardize the columns and set the step from the spectral norm as the code does, and the same loop converges quickly. Blaming convexity here trains the wrong instinct for debugging the deep-learning loops where the same symptoms have the same cure.