Getting ridge regression right

Picture ridge regression in your head. It probably looks like this:

β = (X^TX/n + λI)⁻¹X^TY/n

or

β = (Cov(X, X) + λI)⁻¹Cov(X, Y)

The accompanying proof will motivate this construction by claiming that we wish to penalize ∥β∥² which you might nod along the first time you read and think to yourself “this makes sense, big coefficients overfit, small coefficients underfit”.

Deep in my heart I feel discontent. Just because these variables need a larger coefficient to explain the same variance doesn’t make them any less worthy of β! By this formulation if one of the variables is simply a transformation of another, i.e. X₂ = 5X₁, then ridge regression will “prefer” X₂ over X₁ simply because its coefficient will be smaller.

Of course, textbook authors have foreseen this and in their infinite wisdom might include a disclaimer like “It is often good to normalize your variables when performing ridge regression”. Elements of Statistical Learning goes one step further to say “The ridge solutions are not equivariant under scaling of the inputs, and so one normally standardizes the inputs before solving” which is taking it a step further by adding standardization (i.e. de-meaning before normalizing) to the mix.

The normalization is mentioned as an afterthought to the regression, with a range between presenting it as a “trick” or as “probably important”. This is a common sin of statistical education which has birthed the machine learning intellectual disease of “I don’t know why but this seems to help so we do it”. It also annoyingly defers a crucial step in the modeling to the data preprocessing stage.

Instead, let’s penalize ∥β ⊙ σ_X∥². Other than scale invariance, we can attach a motivation: “penalize for expected prediction variance if all but one variable are missing”. We could alternatively say: “With probability p, variables may be independently missing, and we wish to minimize expected MSE”. Then, we can just write:

β = diag(σ_X⁻¹)(Corr(X) + λI)⁻¹diag(σ_X⁻¹)Cov(X, Y)

This is really just rewriting the normalization into the regression but I claim that it is much more straightforward to formulate ridge regression as adding λ to the diagonal of the correlation matrix, without having to add any further caveats.

In code:

def ridge_regression(X, y, lam):
    n = X.shape[0]
    m = X.shape[1]
    xtx = (X.T @ X) / n
    xty = (X.T @ y) / n
    xsd = np.diag(xtx) ** 0.5
    invxsd = np.reciprocal(xsd, where=xsd>0, out=np.zeros(m))
    xcorr = xtx * invxsd[:,np.newaxis] * invxsd[np.newaxis,:]
    invxcorr = np.linalg.inv(xcorr + lam * np.eye(m))
    beta = invxcorr @ (invxsd * xty) * invxsd
    return beta

Should you de-mean the X above? If you so wish. Should you add an intercept column to X? If you so wish. But you have no choice of applying the ridge correctly.

Dropout is ridge is dropout

Notice the motivation we gave above: “With probability p, variables may be independently missing, and we wish to minimize expected MSE”. If it sounds like dropout it’s because it is! Let’s formalize:

𝔼[MSE] = 𝔼[∥Y − β^T(M ⊙ X)∥₂]

where M ∈ {0, 1}^m × n is a random matrix of the same dimension as X indicating whether a particular entry is dropped or not. With dropout probability p the entries of M will be equal to 1 with probability (1 − p). We can rewrite β^T(M ⊙ X) = M^T(β ⊙ X) then expand:

𝔼[MSE] = Y^TY/n − 2(1 − p)(β ⊙ X)^TY/n + 𝔼[M^T(β ⊙ X)^T(β ⊙ X)M]/n

For the last term, the probability of entries in the diagonal terms is (1 − p) and on the off-diagonal terms is (1 − p)². So:

𝔼[M^T(β ⊙ X)^T(β ⊙ X)M] = (1 − p)²(β ⊙ X)^T(β ⊙ X) − p(1 − p)diag(β² ⊙ σ_X)

Putting it back together we just have:

𝔼[MSE] = 𝔼[∥Y − (1 − p)β^T(M ⊙ X)∥₂]−p(1 − p)diag(β² ⊙ σ_X)/n

Meaning that if we fit a regression with random dropout of variables per row, and multiply the β by (1 − p) – which is the standard dropout method, we will arrive at the same coefficients as ridge. The equivalency is when:

λ = p/(1 − p)

So in essence the ridge λ is equal to the dropout odds.

In code:

def dropout_regression(X, y, p):
    n = X.shape[0]
    m = X.shape[1]
    M = np.random.rand(n,m) > p
    Xm = X * M
    xtx = (Xm.T @ Xm) / n
    xty = (Xm.T @ y) / n
    xsd = np.diag(xtx) ** 0.5
    invxsd = np.reciprocal(xsd, where=xsd>0, out=np.zeros(m))
    xcorr = xtx * invxsd[:,np.newaxis] * invxsd[np.newaxis,:]
    invxcorr = np.linalg.inv(xcorr)
    beta = invxcorr @ (invxsd * xty) * invxsd * (1-p)
    return beta