Linear regression from scratch in R

One of the assignments in my certificate coursework this term was to fit a linear regression without using lm(), R’s built-in function for it, and then check your answer against lm(). It’s the sort of assignment that sounds like busywork and turns out to be the one you learn the most from. I did a version of it in Octave for Andrew Ng’s class back in 2011, but I mostly remember fiddling with the learning rate until the plot of the cost went down. This time I wanted to understand why the fiddling was necessary.

This is my writeup of it, partly so I remember it, and partly because a few people asked how the back to school thing was going. The answer is: hard, slow, and more fun than I expected. R is still strange to me. I’ll get to that.

The data

R ships with a small dataset called mtcars, from a 1974 issue of Motor Trend magazine: 32 cars, with their fuel economy, weight, horsepower, number of cylinders and so on. It’s the “hello world” of R. The question for this exercise is the obvious one: how much does a car’s weight predict its miles per gallon?

Heavier cars use more fuel, no surprise. The question is which line to draw through the points. A line is mpg = β₀ + β₁ × weight, so we need to pick two numbers, the intercept β₀ and the slope β₁.

What “best” means

For any line, each car has a residual: its actual mpg minus what the line predicts. Least squares picks the line that minimizes the sum of the squared residuals. Why squared and not, say, the absolute value? There are good statistical reasons (if the errors are normally distributed, least squares gives you the maximum-likelihood estimate), but the practical reason is that squares are smooth and easy to differentiate, so you can find the minimum with calculus. Absolute values have a kink at zero that makes the math much uglier.

It helps to write it with matrices. Put a column of 1s next to the weights to make a 32×2 matrix X (the 1s are what the intercept multiplies), put the mpg values in a vector y, and put β₀ and β₁ in a vector β. Then the predictions are Xβ, the residuals are y − Xβ, and the thing we’re minimizing is the squared length of that vector:

S(β) = (y − Xβ)ᵀ(y − Xβ)

Way one: the normal equations

At the minimum, the gradient of S with respect to β is zero. Expanding S and differentiating (this is the bit of the assignment where I was grateful for the math degree, even a rusty one) gives:

∂S/∂β = −2Xᵀ(y − Xβ) = 0

which rearranges to

XᵀX β = Xᵀy

These are called the normal equations. There’s a geometric way to see them that finally made them click for me. Xβ, for any β, is some combination of the two columns of X, so all the possible predictions live in a flat 2-dimensional plane inside 32-dimensional space. The y vector is generally not in that plane. The closest point in the plane to y is the one you get by dropping a perpendicular from y onto it, and “perpendicular” means the residual vector y − Xβ is at right angles to both columns of X. That’s exactly what Xᵀ(y − Xβ) = 0 says. “Normal” here means perpendicular. I’d used the name for twenty years without ever wondering where it came from.

For our data, XᵀX is a 2×2 matrix and Xᵀy is a 2-vector, small enough to see:

XᵀX = | 32        102.952  |      Xᵀy = |  642.9    |
      | 102.952   360.901  |            | 1909.753  |

The 32 is the number of cars, 102.952 is the sum of the weights, 360.901 is the sum of the squared weights, and Xᵀy is the sum of mpg and the sum of weight × mpg. Solving the 2×2 system by hand (the determinant is 949.72) gives β₀ = 37.285 and β₁ = −5.344. In R:

X <- cbind(1, mtcars$wt)
y <- mtcars$mpg

beta <- solve(t(X) %*% X, t(X) %*% y)
beta
#           [,1]
# [1,] 37.285126
# [2,] -5.344472

%*% is matrix multiplication, t() is transpose, and solve(A, b) solves Ax = b. It’s tempting to write solve(t(X) %*% X) %*% t(X) %*% y instead, which is the formula as it appears in textbooks, but that computes an actual matrix inverse, which is slower and numerically worse. Solving the system directly is the right habit. (lm() itself goes one step further and uses a QR decomposition of X, which avoids forming XᵀX at all, because squaring the matrix also squares its condition number. More on condition numbers in a minute.)

And checking against R’s own answer:

coef(lm(mpg ~ wt, data = mtcars))
# (Intercept)          wt
#   37.285126   -5.344472

Same to six decimal places. So each extra 1,000 lb of weight costs about 5.3 mpg, and weight alone explains about 75% of the variation in mpg across these cars (R² = 0.753). A car weighing nothing would get 37 mpg, which is a good reminder not to read too much into an intercept that’s far outside your data.

Way two: gradient descent

The normal equations are exact, and for two parameters they’re instant. Gradient descent is what you use when you can’t do that, when there are millions of parameters, or the model isn’t linear and there’s no closed form at all. Every neural network is trained with some version of it, which is why the course has us do it on a problem where we already know the answer.

The idea is to start somewhere, compute the gradient (which direction is uphill), take a small step the other way, and repeat. Using a cost of S/(2n), so the numbers don’t depend on how many rows you have, the gradient is Xᵀ(Xβ − y)/n, and the loop is:

gradient_descent <- function(X, y, alpha, iters) {
  beta <- c(0, 0)
  n <- nrow(X)
  for (i in 1:iters) {
    grad <- t(X) %*% (X %*% beta - y) / n
    beta <- beta - alpha * grad
  }
  beta
}

α is the learning rate, the size of the step. This is where I ran into the same thing I did in 2011. With α = 0.01, it takes about 13,800 iterations to get both numbers within 0.001 of the right answer. With α = 0.1 it takes about 1,400. With α = 0.15, about 900. And with α = 0.2, it doesn’t converge at all. The numbers swing back and forth and grow on every step, and by iteration 50 the slope is in the hundreds of millions.

That cliff between 0.15 and 0.2 is not an accident, and the explanation is the part of this assignment I’m proudest of figuring out. For a quadratic cost like this one, gradient descent is stable only if α is less than 2/λ, where λ is the largest eigenvalue of XᵀX/n. For our data the eigenvalues of XᵀX/n are about 12.2 and 0.076. So the limit is 2/12.2 ≈ 0.164. Anything below it converges, anything above it blows up.

The small eigenvalue explains why it’s slow even when it does converge. The cost surface is a long, narrow valley. The ratio of the two eigenvalues, about 160, is how much steeper the valley’s sides are than its floor. The step size has to be small enough not to overshoot across the steep direction, which means it’s tiny in the shallow direction, and you crawl along the floor of the valley for hundreds of steps. That ratio is the condition number I mentioned above.

The cause is mundane. The weights are all around 3, never near 0, so the intercept column (all 1s) and the weight column point in almost the same direction, and the two parameters get tangled up: changing the intercept and changing the slope have nearly the same effect on the predictions. The fix is the one Andrew Ng’s class taught as “feature scaling” without, as far as I remember, ever saying why. Subtract the mean from the weights and divide by the standard deviation:

z <- (mtcars$wt - mean(mtcars$wt)) / sd(mtcars$wt)
Xs <- cbind(1, z)

Now the column has mean 0, so it’s perpendicular to the column of 1s, both eigenvalues are close to 1, and the valley is round. Gradient descent with α = 0.1 gets within 0.001 in about 95 iterations, and with α = 0.5 in about 15. You then have to convert the coefficients back to the original units (the slope divides by the standard deviation, 0.978, and the intercept absorbs the mean), and you get the same 37.285 and −5.344 as before.

So the 2011 version of me, nudging α up and down until the plot behaved, was searching for a number below 2/λ without knowing it existed. That stings a little, and it’s exactly why I signed up for this.

About R

I promised I’d come back to R. Coming from C#, the biggest adjustment has been that everything is a vector. mtcars$wt - mean(mtcars$wt) subtracts a number from all 32 weights at once, with no loop, and writing a loop is usually a sign you’re doing it wrong (my gradient descent function has one only because the iterations really are sequential). Indexing starts at 1. Assignment is <-, although = mostly works too, and every R person I’ve asked has a strong opinion about which one you should use. And the error messages are, let’s say, an acquired taste.

But I’ll admit the lm(mpg ~ wt, data = mtcars) formula syntax has won me over. Writing a model as a formula, then calling summary() on the result and getting the coefficients, standard errors, t-values, p-values and R² in one table, is a really nice piece of design. Doing that in C# would be a weekend. The next stretch of the program is more regression, with more than one predictor, and then the standard errors in that table start to matter. I’m looking forward to finding out what I’ve been getting wrong about them.