Implement Least Squares From Scratch With a QR Decomposition

Regression · Medium · Free problem

You are asked to fit an ordinary least squares regression $$Y = X\beta + \varepsilon, \qquad X \in \mathbb{R}^{n \times p},\; Y \in \mathbb{R}^{n},\; n > p,$$ in a language with no regression routine. Only basic array arithmetic is available.

(a) Derive the estimator $\hat\beta$ that minimizes $\|Y - X\beta\|^{2}$ and state the linear system it satisfies.

(b) Explain how to compute $\hat\beta$ numerically without forming $(X^{T}X)^{-1}$ explicitly, using a QR decomposition of $X$, and why that is preferable to solving the normal equations directly.

(c) Write the code.

Hints

  1. Expand $S(\beta) = (Y - X\beta)^{T}(Y - X\beta)$ and differentiate with respect to $\beta$; setting the gradient to zero gives the normal equations $X^{T}X\beta = X^{T}Y$.
  2. If $X = QR$ with $Q^{T}Q = I_p$ and $R$ upper triangular and invertible, then $X^{T}X = R^{T}R$ and $X^{T}Y = R^{T}Q^{T}Y$; cancel $R^{T}$.
  3. So $\hat\beta$ solves the triangular system $R\hat\beta = Q^{T}Y$ by back substitution. Build $Q, R$ with (modified) Gram-Schmidt or Householder reflections; the condition number of $R$ equals that of $X$, whereas $X^{T}X$ squares it.

Worked Solution

How to Think About It: Least squares is a calculus problem (minimize a quadratic) whose answer is a linear system, and a numerical-linear-algebra problem (solve that system stably). The normal equations are correct but numerically fragile; a QR factorization solves the same problem while working with $X$ itself rather than $X^{T}X$.

Approach: Derive the normal equations by differentiating the residual sum of squares, substitute $X = QR$ to obtain a triangular system, then implement Gram-Schmidt QR plus back substitution.

Formal Solution:

Part (a): The estimator

*Step 1 -- Objective.* $$S(\beta) = (Y - X\beta)^{T}(Y - X\beta) = Y^{T}Y - 2\beta^{T}X^{T}Y + \beta^{T}X^{T}X\beta.$$

*Step 2 -- First-order condition.* $$\nabla_\beta S = -2X^{T}Y + 2X^{T}X\beta = 0 \quad \Longrightarrow \quad X^{T}X\,\hat\beta = X^{T}Y \quad (\text{the normal equations}).$$ The Hessian $2X^{T}X$ is positive definite when $X$ has full column rank, so this is the unique minimizer, $\hat\beta = (X^{T}X)^{-1}X^{T}Y$. Geometrically the residual $Y - X\hat\beta$ is orthogonal to every column of $X$.

Part (b): Solving via QR

*Step 3 -- Substitute the factorization.* Write $X = QR$ with $Q \in \mathbb{R}^{n \times p}$ having orthonormal columns ($Q^{T}Q = I_p$) and $R \in \mathbb{R}^{p \times p}$ upper triangular, invertible when $X$ has full column rank. Then $$X^{T}X = R^{T}Q^{T}QR = R^{T}R, \qquad X^{T}Y = R^{T}Q^{T}Y,$$ so the normal equations become $R^{T}R\hat\beta = R^{T}Q^{T}Y$. Cancelling the invertible $R^{T}$: $$R\,\hat\beta = Q^{T}Y,$$ a $p \times p$ upper-triangular system solved by back substitution in $O(p^2)$ operations (after the $O(np^2)$ factorization).

*Step 4 -- Why not invert $X^{T}X$?* Forming $X^{T}X$ squares the condition number: $\kappa(X^{T}X) = \kappa(X)^2$. With nearly collinear regressors ($\kappa(X) \sim 10^{6}$, common for factor exposures) the normal equations lose roughly twice as many digits as the QR route, and can even report a singular matrix in floating point. QR also never needs an explicit inverse. For rank-deficient $X$, use a pivoted QR or the SVD instead.

*Step 5 -- Building $Q$ and $R$.* Modified Gram-Schmidt: for each column $j$, subtract its projections onto the already-orthonormalized columns $q_1, \ldots, q_{j-1}$ (recording the coefficients in $R_{ij}$), then normalize (recording $R_{jj}$). Householder reflections are more stable still and are what LAPACK uses; Gram-Schmidt is simpler to write and adequate for well-conditioned problems.

Part (c): Code

```python import numpy as np

def qr_mgs(X): # Modified Gram-Schmidt: X (n x p) = Q (n x p, orthonormal cols) @ R (p x p, upper) X = np.array(X, dtype=float) n, p = X.shape Q = np.zeros((n, p)); R = np.zeros((p, p)) for j in range(p): v = X[:, j].copy() for i in range(j): R[i, j] = Q[:, i] @ v v = v - R[i, j] * Q[:, i] R[j, j] = np.sqrt(v @ v) if R[j, j] == 0: raise ValueError("X is rank deficient") Q[:, j] = v / R[j, j] return Q, R

def back_substitute(R, b): # Solve R beta = b for upper-triangular R p = len(b); beta = np.zeros(p) for i in range(p - 1, -1, -1): beta[i] = (b[i] - R[i, i + 1:] @ beta[i + 1:]) / R[i, i] return beta

def ols(X, Y, intercept=True): X = np.asarray(X, dtype=float) if intercept: X = np.column_stack([np.ones(len(X)), X]) Q, R = qr_mgs(X) beta = back_substitute(R, Q.T @ np.asarray(Y, dtype=float)) resid = Y - X @ beta return beta, resid ```

Complexity: $O(np^2)$ for the factorization, $O(np)$ for $Q^{T}Y$, $O(p^2)$ for back substitution. (Optionally, standard errors follow from $(X^{T}X)^{-1} = R^{-1}R^{-T}$, again without inverting $X^{T}X$ directly.)

Answer: (a) $\hat\beta$ solves the normal equations $X^{T}X\hat\beta = X^{T}Y$, i.e. $\hat\beta = (X^{T}X)^{-1}X^{T}Y$. (b) Factor $X = QR$; then $\hat\beta$ solves the triangular system $R\hat\beta = Q^{T}Y$ by back substitution, which avoids forming $X^{T}X$ and squaring the condition number. (c) Modified Gram-Schmidt (or Householder) QR plus back substitution, as in the code above.

Intuition

OLS is a projection: $\hat\beta$ makes $X\hat\beta$ the orthogonal projection of $Y$ onto the column space of $X$, and a QR factorization is exactly an orthonormal basis for that column space. Solving $R\beta = Q^{T}Y$ avoids squaring the condition number that forming $X^{T}X$ would cause, which matters for the nearly collinear factor matrices common in return regressions; this is what production libraries (LAPACK's least-squares drivers, R's lm) do under the hood.

Open the full interactive solver →