DML GATE Example

This notebook demonstrates how to estimate Group Average Treatment Effects (GATE) using Double Machine Learning. GATE allows us to understand treatment effect heterogeneity by estimating the average effect within specific groups defined by covariates.

Mathematical Formulation

Let $Y$ be the outcome, $D$ the binary treatment, and $X$ the covariates. We define the Conditional Average Treatment Effect (CATE) as:

\tau(X) = \mathbb{E}[Y(1) - Y(0) \mid X]

For a set of groups $G_1, \dots, G_K$ (where $G_k(X)$ is an indicator that unit falls into group $k$), the GATE for group $k$ is:

\theta_k = \mathbb{E}[\tau(X) \mid G_k(X)=1]

We will:

1. Generate synthetic data with heterogeneous treatment effects.
2. Perform Exploratory Data Analysis (EDA).
3. Estimate GATEs using gate_esimand with both automatic quantile groups and custom user-defined groups.

Generate data

We generate observational data with a nonlinear outcome model, nonlinear treatment assignment, and a heterogeneous (nonlinear) treatment effect $\tau(X)$.

EDA

Inference: Estimating GATE

We use gate_esimand to estimate Group Average Treatment Effects. This function works by leveraging the orthogonal signal from the DoubleML framework.

Methodology

1. Orthogonal Signal: We fit a DoubleML IRM model to obtain a score \psi(W;\hat{\eta}) such that:

   \mathbb{E}[\psi(W; \hat{\eta}) \mid X] \approx \tau(X)

   where $W = (Y, D, X)$ and \hat{\eta} are the estimated nuisance parameters.

2. Best Linear Predictor (BLP): We assume a linear model for the CATE using group indicators $G(X)$:

   $\tau(X) \approx \sum_k \theta_k G_k(X)$

   We estimate coefficients $\theta$ by solving the BLP optimization:

   \hat{\theta} = \arg\min_{\theta} \sum_{i=1}^N (\psi(W_i) - \theta^\top G(X_i))^2

When $G(X)$ consists of mutually exclusive group indicators, \hat{\theta}_k corresponds to the average of the orthogonal signal within group $k$, providing a consistent estimate of the GATE.

1. GATE by CATE Quantiles

If no groups are provided, gate_esimand automatically creates groups based on quantiles of the estimated CATE (Conditional Average Treatment Effect).

2. GATE by User-Defined Groups

We can also define custom groups based on covariates. For example, let's group users by their tenure:

• < 1 year
• 1-3 years
• > 3 years