$\newcommand{\diag}{\operatorname{diag}}$We have the problem: $$\min_{w\in\mathbb{R}^{d}}\left( \frac{1}{n}\sum_{i=1}^{n} \left( \langle w,x_{i}\rangle-y_{i} \right)^{2} +2\lambda||w||_1\right),$$ with the assumption that: $$\sum_{i=1}^nx_ix_i^T=\diag(\sigma_1^2,...,\sigma_d^2).$$

Is there a closed-form solution in this case?

I have that: $$(X^TX)^{-1}=\diag\left(\sigma_1^{-2},...,\sigma_d^{-2}\right),$$ and so I think the answer is: $$w\,^j=y\,^j\max\left\{0,1-\lambda \frac{n}{|y^j|}\right\},$$ for $y\,^j=\displaystyle\sum_{i=1}^n\frac{y_ix_i\,^j}{\sigma_i^2}$ , but I'm not sure.


I'm going to go through @cardinal's derivation of the closed form lasso solution when $X^T X = I$, found here, with minor modifications.

I'll be assuming that $\sigma^2_i > 0$ for all $i$. This is justified because if we have a $\sigma^2_i = 0$ then this tells us that the $i$th column of $X$ is all 0, and I think it's reasonable to exclude such a case. I'll let $X^T X = D$. Note that this also means that $X$ is full rank and the OLS solution $\hat \beta$ is uniquely defined.

I'm also going to modify your notation to better match that in the answer that I'm referencing. To that end, I shall be solving $$ \hat \beta_\lambda = \text{argmin}_{\beta \in \mathbb R^p } \frac 12 \vert \vert Y - X\beta\vert \vert^2_2 + \lambda \vert \vert \beta \vert \vert_1. $$

This is identical to your problem but I can add more details here if you'd like.

Following @cardinal's derivation, we have that we need to solve $$ \hat \beta_\lambda = \text{argmin } \frac 12 (Y^T Y - 2 Y^T X \beta + \beta^T X^T X \beta) + \lambda \vert \vert \beta \vert \vert_1 $$

$$ = \text{argmin } -Y^T X \beta + \frac 12 \beta^T D \beta + \lambda \vert \vert \beta \vert \vert_1. $$

Noting that the OLS solution is $\hat \beta = (X^T X)^{-1} X^T Y = D^{-1}X^T Y$, we have that $$ \hat \beta_\lambda = \text{argmin } -\hat \beta^T D \beta + \frac 12 \beta^T D \beta + \lambda \vert \vert \beta \vert \vert_1 $$

$$ = \text{argmin } \sum_{j=1}^p -\hat \beta_j \beta_j \sigma^2_j + \frac{\sigma^2_j}2 \beta_j^2 + \lambda | \beta_j|. $$

We are optimizing over each $\beta_j$ separately, so we can solve each term of this sum separately. This means we need to minimize $\mathcal L_j$ where $$ \mathcal L_j = -\hat \beta_j \beta_j \sigma^2_j + \frac{\sigma^2_j}2 \beta_j^2 + \lambda | \beta_j|. $$

Following a completely analagous argument to the linked answer, we find that $$ (\hat \beta_\lambda)_j = \mathrm{sgn}(\hat \beta_j) \left(|\hat \beta_j| - \frac{\lambda}{\sigma^2_j}\right)^+. $$

Furthermore, $\hat \beta = D^{-1} X^T Y \implies \hat \beta_j = \frac{X_j^T Y}{\sigma_j^2}$ so we have that $$ \left(|\hat \beta_j| - \frac{\lambda}{\sigma^2_j}\right)^+ = \frac 1 {\sigma^2_j} \left(|X_j^T Y| - \lambda\right)^+ $$

so it turns out that a predictor $X_j$ gets zeroed out exactly when it would if the design matrix was orthonormal, not just orthogonal. So we can see that in this case with $X^T X = D \neq I$, the variable selection is no different than if $X^T X = I$, but the actual coefficients $\hat \beta_\lambda$ are scaled according to the predictor variances.

As a final note, I will turn this solution into one that resembles yours which means we need to multiply $\hat \beta$ by something to get $\hat \beta_\lambda$. If $(\hat \beta_\lambda)_j \neq 0$ then we have that $$ (\hat \beta_\lambda)_j = \text{sgn}(\hat \beta_j) \left( \vert \hat \beta_j \vert - \frac{\lambda}{\sigma^2_j} \right) = \hat \beta_j - \text{sgn}(\hat \beta_j) \frac{\lambda}{\sigma^2_j} $$

$$ = \hat \beta_j \left( 1 - \frac{\lambda}{\sigma^2_j \vert \hat \beta_j \vert} \right) $$

since $\frac{a}{\vert a \vert} = \text{sgn}(a)$.

Noting that $(\hat \beta_\lambda)_j = 0$ exactly when $$ \vert \hat \beta_j \vert - \frac{\lambda}{\sigma^2_j} \leq 0 \iff \vert \hat \beta_j \vert \leq \frac{\lambda}{\sigma^2_j} \iff 1 \leq \frac{\lambda}{\sigma^2_j \vert \hat \beta_j \vert} \iff 1 - \frac{\lambda}{\sigma^2_j \vert \hat \beta_j \vert} \leq 0, $$

we see that we could alternatively express $\hat \beta_\lambda$ as $$ (\hat \beta_\lambda)_j = \hat \beta_j \left( 1 - \frac{\lambda}{\sigma^2_j \vert \hat \beta_j \vert} \right)^+. $$

So this is very close to what you had but not exactly the same.

I always like to check derivations like this against well-known libraries if possible, so here's an example in R:

## generating `x`
n = 1000
p = 5
sigma2s = 1:p
x = svd(matrix(rnorm(n * p), n, p))$u %*% diag(sqrt(sigma2s))

## check this
# t(x) %*% x

## generating `y`
betas = 1:p
y = x %*% betas + rnorm(nrow(x), 0, .5)

lambda = 2

## using a well-known library to fit lasso
penalized(y, x, lambda1 = lambda)@penalized

## using closed form solution
betahat = lm(y ~ x - 1)$coef
ifelse(betahat > 0, 1, -1) * sapply(abs(betahat) - lambda / sigma2s, function(v) max(c(0, v)))

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.