Simplifying the Matrix Form of the Solution to Ridge Regression

Question

I'm trying to understand how to obtain the solution to an objective function by solving for the parameter vector $\theta$ in ridge regression. I found an example here from Naomi which takes an example from The Elements of Statistical Learning by Hastie et al. and seems understandable, but I'm having trouble understanding the steps between obtaining matrix form, and solving for $\theta$ to get the solution:

It looks like in the book they also just say "the ridge regression solutions are easily seen to be..." and handwave the steps for solving theta.

When I try to solve, I thought that maybe you could expand $$(X\theta-Y)^2 = (X)^2(θ)^2-2XθY+Y^2$$ so you can start moving theta to one side, but this seems wrong and I don't really understand where the identity matrix $I$ comes from in the result, so maybe I'm just approaching this wrong altogether? Am I missing some fundamental knowledge here?

Answer 1

Your initial thoughts about expanding are on the right track, but you're going to need to deploy some matrix calculus. Your function $$J(\theta) = (\mathbf{X}'\theta - Y')^{T}(\mathbf{X}'\theta - Y') + \lambda \theta^{T}\theta,$$ for a fixed scalar $\lambda$, is a function of $\theta$, which is a vector. So you need the higher-dimensional generalizations of the derivative in order to deal with differentiating scalars, vectors, and matrices with respect to scalars, vectors, and matrices. It's often useful, especially at first, to explicitly keep track of the types of each element in whatever expression you're differentiating as well as the type of element you are trying to find.

Here, $J(\theta)$ is a scalar. I'm going to assign quantities to the dimensions of the other quantities to make things explicit: $\theta$ is a $p$ dimensional vector, $\mathbf{X}'$ is an $n \times p$ matrix, and $Y'$ is an $n$ dimensional vector. All the vectors here will be column vectors, meaning that $\theta$ can be written to look like a $p \times 1$ dimensional matrix and $Y'$ as an $n \times 1$ matrix.

Now if you follow the link above, you'll see a discussion of some ambiguity in the notation of matrix calculus, as there are two alternate conventions in how the derivatives are taken. But based on the form of the solution given for $\theta$, it looks like the form of the relevant rule you want is: $$ \frac{\partial ||u||_{2}^{2}}{\partial u} = \frac{\partial u^{T}u}{\partial u} = 2u \frac{\partial u}{\partial x}$$ where $u(x)$ is a vector-valued function of the vector $x$. This is basically the chain rule for derivatives applied to a quadratic, but for vector-valued functions. In your function $J(\theta)$, there are two quadratics of this form being summed together. The second term is easier to deal with: $$\frac{\partial \lambda\theta^{T}\theta}{\partial \theta} = 2\lambda\theta $$ One thing to observe here is that even though the inner product $\theta^{T}\theta = ||\theta||_{2}^{2}$ is a scalar, the derivative of this scalar with respect to the vector $\theta$ is a vector.

The above summand was easier to deal with because there $u(\theta) = \theta$, so we have the trivial 'inner' partial derivative $\frac{\partial u(\theta)}{\partial \theta} =1$. For the first summand, we have $u(\theta) = \mathbf{X}'\theta - Y'$. It's still just a quadratic, we just now have a non-trivial inner partial derivative in our chain rule. Differentiating, we have

$$\frac{||\mathbf{X}'\theta - Y'||_{2}^{2}}{\partial \theta} = \frac{(\mathbf{X}'\theta - Y')^{T}(\mathbf{X}'\theta - Y')}{\partial \theta} = 2(\mathbf{X}')^{T}(\mathbf{X}'\theta - Y') $$

The $(\mathbf{X}')^{T}$ factor in front is the 'inner' partial derivative, $\frac{\partial u(\theta)}{\partial \theta}$ here. Technically, we've actually applied another variation of the chain rule to the $\mathbf{X}'\theta$ term, since $\frac{\partial \mathbf{X}'\theta}{\partial \theta} = \frac{\partial \theta}{\partial \theta}(\mathbf{X}')^{T} = (\mathbf{X}')^{T}$, so we've really used 'the' chain rule twice on this term.

Combining these results, what we have so far is $$\nabla_{\theta}J(\theta) = 2\bigg((\mathbf{X}')^{T}(\mathbf{X}'\theta - Y') + \lambda \theta\bigg) \overset{\color{red}{set}}{=} 0. $$ Since we are trying to minimize the function $J(\theta)$, the usual intuition from calculus applies and we set this derivative equal to $0$. That nuisance factor of $2$ doesn't alter the nature of our eventual result, which is why there was originally a factor of $\frac{1}{2}$ in your function $J(\theta)$. I'm just going to drop it from now on to focus on the main argument.

Doing some algebra and rearranging things, we have

$$ (\mathbf{X}')^{T}(\mathbf{X}'\theta - Y') + \lambda \theta\\ = (\mathbf{X}')^{T}\mathbf{X}'\theta - (\mathbf{X}')^{T}Y' + \lambda\theta\\ = \bigg((\mathbf{X}')^{T}\mathbf{X}' + \lambda \mathbf{I}\bigg)\theta - (\mathbf{X}')^{T}Y' = 0, $$ where the last equality follows because multiplying the $p$ dimensional vector $\theta$ by the $p \times p$ identity matrix $\mathbf{I}$ just produces the $p$ dimensional vector $\theta$ again: $\mathbf{I}\theta = \theta$.

From here, it's just some straightforward arithmetic with matrices to solve for $\theta$:

$$ \bigg((\mathbf{X}')^{T}\mathbf{X}' + \lambda \mathbf{I}\bigg)\theta - (\mathbf{X}')^{T}Y' = 0 \\ \iff \bigg((\mathbf{X}')^{T}\mathbf{X}' + \lambda \mathbf{I}\bigg)\theta = (\mathbf{X}')^{T}Y' \\ \iff \theta = \bigg((\mathbf{X}')^{T}\mathbf{X}' + \lambda \mathbf{I}\bigg)^{-1}(\mathbf{X}')^{T}Y' $$

Something to note that is of interest here is that even if $(\mathbf{X}')^{T}\mathbf{X}'$ is not invertible, so that there is no unique least squares solution to the standard, unpenalized version of the problem where $\lambda = 0$, we instead have that $\bigg((\mathbf{X}')^{T}\mathbf{X}' + \lambda \mathbf{I}\bigg)$ is always invertible for $\lambda > 0$ — this is because of a result in linear algebra relating invertibility to (non-zero) eigenvalues, and how $\lambda \mathbf{I}$ affects the eigenvalues of $(\mathbf{X}')^{T}\mathbf{X}'$. The positivity of $\lambda$ ensures that we don't turn a non-zero eigenvalue of $(\mathbf{X}')^{T}\mathbf{X}'$ into $0$ through subtraction, even though a negative $\lambda$ would make a different eigenvalue of $(\mathbf{X}')^{T}\mathbf{X}'$that was originally $0$ nonzero.

Answer 2

Your idea of expanding and solving is on the right track, but you're missing a couple things.

First, these are vectors and matrices so the expansion of the loss is $$ (X\theta - Y)^T(X\theta - Y) = \theta^TX^TX\theta -2 \theta^TX^TY + Y^TY. $$

Next, we don't solve this for $\theta$ because this is our loss function so that'd be like trying to find a $\theta$ that leads to some particular loss. Instead, we want to find the value of $\theta$, which we'll call $\hat\theta$, that minimizes this loss. We can do that by taking the derivative with respect to $\theta$ and solving that for zero. This loss turns out to be convex so we know that that root is the global minimum of the loss.

Thus if $J$ is the loss as a function of $\theta$, and we have a penalty on the norm of $\theta$ (i.e. $\theta^T\theta$), we have $$ J(\theta) = (X\theta - Y)^T(X\theta - Y) +\lambda \theta^T\theta\\ =\theta^TX^TX\theta - 2\theta^TX^TY + Y^TY +\lambda \theta^T\theta \\ \implies \frac{\partial J}{\partial \theta} = 2X^TX\theta - 2X^TY + 2\lambda \theta \\ = 2\left[(X^TX + \lambda I)\theta - X^TY\right] $$ and the $I$ appears because I'm factoring $\theta$ out of $\lambda \theta$. Therefore $$ \frac{\partial J}{\partial \theta} \stackrel{\text{set}}= \mathbf 0\\ \implies (X^TX + \lambda I)\hat\theta = X^TY \\ \implies \hat\theta = (X^TX + \lambda I)^{-1}X^TY $$ where I'm using the fact that $X^TX + \lambda I$ is always invertible.