LASSO/Ridge regression with adjustment for a covariate

Question

I'm trying to address the following analysis problem in high-dimensional biological data. The setup is bulk gene expression data where multiple cell types (tumor and immune cells) can contribute to detected gene expression signal in the mixture:

I would like to find the genes (predictor variables) whose expression levels are the most predictive for the changes in the immune signature (response variable). However, my goal is to untangle a few moving parts by accounting for the "levels of immune cells" in the mixture to find predictor genes that are more likely to be related to the tumor cells.

In other words, I would like to adjust the response variable for a covariate (level of immune cells in the mixture -- which is likely to show some confounding effects with the predictor variables I care about). A linear model version of what I'm trying to achieve would be probably something like this:

immune_signature = b0 + b1*immune_cell_level + b2*gene1 + b3*gene2 + ...

I would like the algorithm to apply shrinkage on coefficients b2 ... bN, but leave b1 alone.

What is the best way of doing this in LASSO/Ridge?

• Can we force the algorithm not to shrink the coefficient of a desired covariate (I'm using R)?

• Or is it customary to perform two separate analyses 1) One with the response variable I'm interested in, and 2) One with covariate I'd like to adjust for as the response variable, and compare the coefficients between two methods?

Thanks!

Answer 1

You can apply the "double machine learning" of Chernozhukov et al. (2017) outlined in this paper. To fit your problem into their framework it is useful to write it in the following notation

$$ Y_i = D_i\theta + X_i'\beta + \varepsilon_i$$

The variable $D_i$ is your primary variable of interest, $X_i$ are additional controls (including an intercept), and $\varepsilon_i$ is an error term. All three variables are assumed to be i.i.d.. The idea of double machine learning is running auxiliary regressions on different subsets of the variables. This allows you to run an algorithm like Lasso on only a subset of the covariates.

1. (Sample Splitting) Randomly split the sample, let's call it Sample A and B.
2. (Auxiliary Regressions) In sample A, do two lasso regressions:

$\qquad$ (a) $Y_i$ on $X_i$. The estimated coefficients are $\hat{\gamma}^Y$.

$\qquad$ (b) $D_i$ on $X_i$. The estimated coefficients are $\hat{\gamma}^D$.

3. (Main Regressions) In sample B, construct the residuals $\hat{U}_i^Y = Y_i - X_i'\hat{\gamma}^Y$ and $\hat{U}^{D} = D_i - X_i'\hat{\gamma}^D$. Then run the following regression: $$ \hat{U}_i^Y = \hat{U}_i^D \theta + \tilde{\varepsilon}_i $$ You can estimate robust standard errors in step (3) to construct confidence intervals for $\hat{\theta}$.

Note: (i) The coefficients $(\hat{\gamma}^Y,\hat{\gamma}^D)$ are different from $\beta$ and should not be confused. Similarly, $\varepsilon_i \ne \tilde{\varepsilon}_i$. With this approach you can only estimate $\theta$. (ii) This procedure works for a wide variety of high-dimensional procedures. You can use Ridge instead of Lasso in the second step, or neural networks, random forest, etc. (iii) In the published paper, they also outline an alternative procedure where you repeat the process, now changing the role of samples $A$ and $B$. You can then construct $\hat{\theta}$ by obtaining the average of the coefficients obtained from either labelling strategy. This can increase efficiency by utilizing your whole sample.

Why does it work?

(i) On one hand, Steps (2) and (3) can be viewed as analogs of the Frisch-Waugh-Lovell Theorem, which shows that we can always rewrite a multivariate linear regression in terms of auxiliary regressions with two sets of covariates $(D_i,X_i)$. On the other hand, Step (3) has a "Neyman-Orthogonality" property. In layman's terms this means that Step (3) is less sensitive to whether you estimated $(\hat{\gamma}^Y,\hat{\gamma}^D)$ accurately.

(ii) The additional sample splitting step is necessary to further control for the first stage uncertainty in estimating $(\hat{\gamma}^Y,\hat{\gamma}^D)$. The formal proof is non-trivial but it is based on the following intuitive idea: The estimation error from Step 2 is uncorrelated with the error in Step 3 because they are constructed from different samples.