CATE Lasso: Conditional Average Treatment Effect Estimation with High-Dimensional Linear Regression

Masahiro Kato and Masaaki Imaizumi. 2023. Preprint (arXiv:2310.16819). 日本語版

This study estimates the treatment effect for each individual, the conditional average treatment effect (CATE), when the number of covariates exceeds the sample size. High-dimensional regression normally leans on sparsity, meaning that most coefficients are zero. Here no sparsity is imposed on the regression model for either treatment. The penalty falls instead on the difference between the coefficients of the two models, and the resulting Lasso estimator is shown to be consistent.

The problem

The CATE is the difference between the expected outcomes of two treatments, conditional on covariates, so it describes how far the effect varies from one individual to another. When it is estimated by linear regression and the covariate dimension is larger than the sample size, the least-squares solution loses properties such as consistency. The usual answer is a method built on sparsity, the Lasso being the obvious example. But its premise, that most of the coefficients explaining the outcome under a given treatment are zero, is not always reasonable.

The proposed method

The starting point is an assumption about structure. The coefficients of the two regression models split into a part whose value changes with the treatment and a part that is common to both. Since the CATE is the difference between the two models, the common part cancels and only the treatment-dependent part survives. So even when neither model is sparse on its own, the nonzero coefficients of the CATE are confined to the part that changes with the treatment. We call this implicit sparsity and build a Lasso whose ℓ1 penalty applies to the difference between the two coefficient vectors rather than to the coefficients themselves.

Main results and conditions

Under implicit sparsity the proposed estimator is consistent for the CATE. The interesting part is the contrast with the individual models: consistency cannot be established for the coefficients of either regression, yet it can be established for their difference. The condition is that few coefficients change with the treatment, and the case where each model is itself sparse falls inside that condition. Nuisance parameters drop out as well. Nothing has to be estimated for the propensity score or for the conditional expectation, which is an advantage over inverse probability weighting and doubly robust methods. In simulations the root mean squared error is smaller than that of T-learners built on least squares or on the ordinary Lasso. As more coefficients are allowed to change with the treatment, the performance of the proposed method moves toward that of least squares.

Relation to earlier work

The use of an ℓ1 penalty in high-dimensional linear regression begins with Tibshirani's Lasso (1996, Regression Shrinkage and Selection Via the Lasso). Methods for CATE estimation are usually compared inside the metalearner framework organized by Künzel and coauthors (2019, Metalearners for estimating heterogeneous treatment effects using machine learning), and the estimator proposed here belongs to the T-learner, which fits a separate regression for each treatment. The R-learner of Nie and Wager uses Robinson's decomposition (2021, Quasi-oracle estimation of heterogeneous treatment effects), and Kennedy's DR-learner regresses on a doubly robust pseudo-outcome (2023, Towards optimal doubly robust estimation of heterogeneous causal effects). Both need an estimate of the propensity score or of the conditional expectation, and with high-dimensional covariates that estimation itself calls for something like a sparsity assumption. Our method sidesteps that step and uses only the structure of the CATE.

Where it applies

The setting we have in mind is customer or medical data with many recorded variables, where the question is who responds to the treatment. It fits when you can argue that many factors explain the outcome while only a few change how the treatment works.

Paper and materials

BibTeX

@misc{cate-lasso,
  author       = {Masahiro Kato and Masaaki Imaizumi},
  title        = {CATE Lasso: Conditional Average Treatment Effect Estimation with High-Dimensional Linear Regression},
  year         = {2023},
  eprint       = {2310.16819},
  archivePrefix = {arXiv},
  url          = {https://arxiv.org/abs/2310.16819},
}