Computing OLS When the Design Matrix Is Too Large for Memory
In Ordinary Least Squares, you need the matrix product $X^T X$ (where $X$ is $n \times p$) to compute the coefficient estimate $\hat{\beta} = (X^T X)^{-1} X^T Y$. But suppose $n$ is so large that you cannot fit the full matrix $X$ in memory.
How would you compute $X^T X$ (and $X^T Y$) without storing the entire design matrix? Describe a specific, numerically accurate method -- not an approximate algorithm like SGD.
Hints
- Think about the dimensions: $X$ is $n \times p$ but $X^T X$ is only $p \times p$. You need $O(p^2)$ memory for the result, not $O(np)$ for the input.
- Write out the $(j,k)$ entry of $X^T X$: it is $\sum_{i=1}^{n} x_{ij} x_{ik}$. This is a sum over rows -- you can accumulate it one row at a time.
- Process data in chunks: for each block $X_k$, compute $X_k^T X_k$ and add it to a running accumulator. After one pass, you have the exact $X^T X$.
Worked Solution
How to Think About It: The key observation is a dimension mismatch that works in your favor. $X$ is $n \times p$ and might have billions of rows, but $X^T X$ is only $p \times p$. In most regression problems, $p$ is at most a few thousand features, so $X^T X$ fits comfortably in memory even when $X$ does not. The question is whether you can build $X^T X$ without loading all of $X$ at once. The answer is yes, because matrix multiplication can be decomposed into a sum over rows.
Key Insight: The matrix product $X^T X$ is a sum of rank-1 outer products -- one per row of $X$. You can accumulate this sum one row (or one chunk) at a time.
The Method:
- Initialize a $p \times p$ accumulator matrix $S = 0$ and a $p \times 1$ vector $t = 0$.
2. Stream through the data one row (or chunk of rows) at a time. For each row $x_i \in \mathbb{R}^p$ with response $y_i$: - Update: $S \leftarrow S + x_i x_i^T$ (outer product, $p \times p$) - Update: $t \leftarrow t + x_i y_i$ (scaled vector, $p \times 1$)
In chunk form, if you load a block $X_k$ of $m$ rows at a time: - $S \leftarrow S + X_k^T X_k$ - $t \leftarrow t + X_k^T Y_k$
- After one pass through all $n$ rows: $S = X^T X$ and $t = X^T Y$.
- Solve: $\hat{\beta} = S^{-1} t$ (or better, use Cholesky decomposition of $S$).
Why this works: Matrix multiplication is associative and can be decomposed:
$$X^T X = \sum_{i=1}^{n} x_i x_i^T$$
This is just the definition of matrix multiplication written as a sum over rows. Each outer product $x_i x_i^T$ is a $p \times p$ matrix, and you accumulate them into $S$. Memory usage is $O(p^2)$ for the accumulator plus $O(mp)$ for the current chunk -- independent of $n$.
Practical Considerations:
- Numerical stability: Summing many outer products can accumulate floating-point error. For better accuracy, use compensated summation (Kahan summation) or process data in double precision. If features have very different scales, center and scale them first.
- Chunk size: Larger chunks are faster (BLAS-optimized matrix multiply on blocks), but use more memory. A chunk of $m = 10{,}000$ rows with $p = 1{,}000$ features uses about 80 MB -- a good practical trade-off.
- Parallelization: Each chunk's contribution $X_k^T X_k$ is independent, so you can compute them in parallel (e.g., MapReduce) and sum the results. This is the basis of distributed OLS in systems like Spark.
- Why not SGD? SGD gives an approximate solution that depends on step sizes and number of passes. The chunked outer product method gives the exact OLS solution (up to floating-point precision) in a single pass. When you need the exact answer -- e.g., for inference, confidence intervals, or regulatory reporting -- this is the right approach.
Answer: Compute $X^T X = \sum_i x_i x_i^T$ by streaming through the data one row (or chunk) at a time, accumulating the $p \times p$ outer product sum. This requires $O(p^2)$ memory regardless of $n$, gives the exact OLS solution, and parallelizes trivially.
Intuition
This problem tests whether you understand that matrix multiplication is a sum and sums can be computed incrementally. The formula $X^T X = \sum_i x_i x_i^T$ is trivial mathematically but operationally powerful: it decouples the computation from the data size. You never need all of $X$ in memory at once -- you just need each row long enough to compute its outer product contribution.
This pattern is ubiquitous in large-scale computing. Any sufficient statistic that takes the form of a sum can be computed in a streaming fashion: sample means, covariance matrices, gradient sums. It is the theoretical foundation of MapReduce-style distributed computation. The practical lesson for quant work: when someone says "the data is too big," the first question is whether the computation can be decomposed into a sum over independent pieces. If so, you stream or parallelize -- you do not reach for approximate methods.