Causal ML: Why naive model selection fails

Simulation notebook

Goals:

• Illustrate why naive implementations of model selection using Lasso are problematic

Data generating process

We consider a DGP with correlated variables:

• \(p\geq10\) covariates drawn from a multivariate normal distribution: \(X \sim N(0,\Sigma)\), where we set in \(\Sigma\) the diagonal entries (variance) to one and the off-diagonal entries (covariance) to 0.5.

• The outcome model is \(Y = \underbrace{\beta_0 + \beta_1 X_1 + ... + \beta_{10} X_{10}}_{\text{CEF }f(X)} + e\), where \(X_j\) is the \(j\)-th column of \(X\) and \(e \sim N(0,2)\)

• We consider the following parameters \(\beta_0 = 0\), \(\beta_1 = 1\), \(\beta_2 = 0.9\), …, \(\beta_9 = 0.2\), \(\beta_{10} = 0.1\) such that the first 10 variables have a decreasing impact on the outcome and any further variables are just irrelevant noise with no true predictive power

We set \(p=99\) and work with sample size \(N=100\).

# Load the packages required for later
library(tidyverse)
library(mvtnorm)
library(glmnet)

set.seed(1234) # For replicability

# Define the relevant parameters
n = 100
var = 1
cov = 0.5
p = 99

beta = c(0,seq(1,0.1,-0.1),rep(0,p-10))
sig = matrix(cov,p,p)
diag(sig) = var

A causal question and different advice

You are interested in the causal effect of variable \(X_1\) on outcome \(Y\). You know that you should control for additional variables, but not really which specification to choose. You organize a meeting with some colleagues and ask for advice. You receive three different answers:

• Professor: Looks in a coffee cup and tells you that it is obvious that you should control for \(X_2-X_{10}\) using OLS and leaves.

The remaining PhD students just finished a course in “Big Data, Machine Learning and AI” and are very excited about variable selection via Lasso because this is exactly what you are asking for:

• PhD student 1: Tells you that you should just throw everything in a Lasso regression and report the resulting coefficient of \(\beta_1\) as your causal effect.

• PhD student 2: Is skeptical because the causal effect is shrunken towards zero and proposes to leave the coefficient of \(\beta_1\) unpenalized and to only penalize the other variables.

A simulation study to figure out what works

We run a simulation study that draws 1000 random samples based on our DGP and stores the estimates resulting from the three proposed procedures (runs several minutes).

repl = 1000

beta1_ols = beta1_lasso = beta1_lasso_unpen = rep(NA,repl)

for (i in 1:repl) {
  x = cbind(rep(1,n),rmvnorm(n,sigma=sig))
  y = x %*% beta + rnorm(n,0,2)
  
  # OLS
  beta1_ols[i] = lm(y~x[,2:11])$coefficients[2]
  
  # Plain Lasso (specifying lambda.min.ratio speeds up)
  cv_lasso = cv.glmnet(x[,-1],y,lambda.min.ratio=0.001) 
  lasso_temp = glmnet(x[,-1],y,lambda=cv_lasso$lambda.min)
  beta1_lasso[i] = lasso_temp$beta[1]
  
  # Lasso with unpenalized X_1
  cv_lasso = cv.glmnet(x[,-1],y,penalty.factor=c(0,rep(1,p-1)),lambda.min.ratio=0.001)
  lasso_temp = glmnet(x[,-1],y,lambda=cv_lasso$lambda.min,penalty.factor=c(0,rep(1,p-1)))
  beta1_lasso_unpen[i] = lasso_temp$beta[1]
}

Let’s check the resulting estimator distributions of the proposed strategies:

Professor OLS

Does it work?

Yes! The distribution of the estimator has a mean (dashed line) close to the true value of one (solid line) \(\Rightarrow\) unbiased:

df = data.frame(x=beta1_ols)
ggplot(df,aes(x=x)) + geom_histogram(bins=30,aes(y =..density..)) + 
              geom_vline(xintercept = c(beta[2],mean(df$x)),linetype=c('solid','dashed'))

Why? (i) The model was chosen without looking at the data \(\Rightarrow\) no problems with post-selection issues. (ii) The correct model was chosen (however the Prof. managed to do this). (iii) OLS is used, which is the best unbiased estimator for such a setting.

Student 1’s plain Lasso

Does it work?

No! The estimator is clearly downward biased.

df = data.frame(x=beta1_lasso)
ggplot(df,aes(x=x)) + geom_histogram(bins=30,aes(y =..density..)) + 
  geom_vline(xintercept = c(beta[2],mean(df$x)),linetype=c('solid','dashed'))

Why? Student 2 is right. The penalization shrinks the parameter of interest towards zero. In several replications the variable of interest is not even selected, which explains the spike at zero:

paste("Share of replications where variable of interest not selected:",round( sum(beta1_lasso == 0)) / repl *100,"%")

[1] "Share of replications where variable of interest not selected: 2.8 %"

Student 2’s Lasso with unpenalized parameter of interest

Does it work?

No! The estimator is upward biased:

df = data.frame(x=beta1_lasso_unpen)
ggplot(df,aes(x=x)) + geom_histogram(bins=30,aes(y =..density..)) + 
  geom_vline(xintercept = c(beta[2],mean(df$x)),linetype=c('solid','dashed'))

Why? Lasso is not aware of the causal problem at hand. Thus, it could use the omitted variable bias (OVB) for prediction purposes. Correcting OVB requires to build up regression coefficients of confounding variables, but this is costly because of the penalization. Leveraging the OVB in the unpenalized coefficient is “for free”. On the other hand, this must not be the case. With a sufficiently low penalty term, it could work. But plain Lasso does not care about an unbiased parameter \(\beta_1\).

To understand this better, let’s zoom into the coefficient path of the last replication:

cv_lasso_unpen = cv.glmnet(x[,-1],y,penalty.factor=c(0,rep(1,p-1)),lambda.min.ratio=0.002)
plot(cv_lasso_unpen)

lasso_unpen = glmnet(x[,-1],y,penalty.factor=c(0,rep(1,p-1)),lambda.min.ratio=0.002)
plot(lasso_unpen,xvar = "lambda",label=T)

The empty model consists now of the unpenalized variable and it leverages the full omitted variable bias for prediction. The OVB is reduced as the other variables receive at least parts of their coefficients. In principle, there would be a sweet spot on the trajectory of \(\beta_1\), but cross-validation seems to systematically favor models where the OVB in \(\beta_1\) is used for prediction.

Such things happen if the ML method is not “taught” that we are interest in a causal effect in the correct way.

Take-away

• If somebody tells us the correct model without looking at the data, life is good.

• In the likely case that not, we need to be careful how we leverage the ML machinery for causal analysis.

• There are a lot of interesting things to be learned about how to teach ML that we are estimating causal effects.

Suggestions to play with the toy model

Feel free to play around with the code. This is useful to sharpen and challenge your understanding of the methods. Think about the consequences of a modifications before you run it and check whether the results are in line with your expectation. Some suggestions:

• Modify DGP (correlation of covariates, betas, noise term, …)

• Increase the number of observations

• Implement other ad hoc ideas that one could come up with