State meta-learner assumptions and correct X-learner weighting

causal-inference-for-the-brave-and-true/21-Meta-Learners.ipynb

@@ -48,7 +48,17 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Here, we will use the same data we had before, regathing investment advertisement emails. Again, the goal here is to figure out who will respond better to the email. There is a little twist, though. This time, we will use non-random data to train the models and random data to validate them. Dealing with-non random data is a much harder task, because the meta learners will need to debias the data **AND** estimate the CATE."
+    "Here, we will use the same data we had before, regarding investment advertisement emails. Again, the goal here is to figure out who will respond better to the email. There is a little twist, though. This time, we will use non-random data to train the models and random data to validate them. Dealing with non-random data is a much harder task because the meta-learners must adjust for observed confounding **AND** estimate the CATE.\n",
+    "\n",
+    "It is important to be precise about what makes that adjustment causal. Meta-learners do not remove confounding merely because their component models are flexible. For the S-, T-, and X-learners below to identify $\\tau(x)$ from observational training data, we need the usual assumptions:\n",
+    "\n",
+    "1. **Consistency and no interference:** the observed outcome is the potential outcome under the treatment actually received, $Y_i=Y_i(T_i)$, and one unit's treatment does not change another unit's outcome.\n",
+    "2. **Conditional exchangeability:** all variables needed to control confounding are measured before treatment and included in $X$, so $(Y_i(0),Y_i(1)) \\perp T_i \\mid X_i$.\n",
+    "3. **Overlap:** both treatment levels are possible for the covariate profiles where we estimate effects, $0<e(x)=P(T_i=1\\mid X_i=x)<1$.\n",
+    "\n",
+    "Under these assumptions, the observed conditional outcome regressions identify the potential-outcome regressions: $E[Y_i\\mid T_i=t,X_i=x]=E[Y_i(t)\\mid X_i=x]$. Without exchangeability, a meta-learner can reproduce hidden confounding very accurately; without overlap, it must extrapolate a missing treatment arm. The raw treated and control groups do not have to have identical covariate distributions, but severe imbalance can reveal weak overlap and heavy dependence on model extrapolation. Balance and overlap are useful diagnostics, not properties automatically guaranteed by fitting a meta-learner. The outcome and propensity models must also be estimated well enough for the method being used.\n",
+    "\n",
+    "The randomized test set makes the gain-curve evaluation less vulnerable to confounding, but it does **not** repair confounding in the observational training set. Representativeness is a separate transport question. To deploy $\\hat\\tau(x)$ in a new population, that population must have covariate support covered by the training data and the treatment-response relationship must remain relevant there. Its covariate distribution also matters if we aggregate CATEs into a target-population ATE, $E_{\\text{target}}[\\tau(X)]$."
    ]
   },
   {
@@ -450,41 +460,48 @@
     "\\hat{M}_1(X) \\approx E[Y| T=1, X]\n",
     "$\n",
     " \n",
-    "Now, things start to take a turn. For the second stage, we impute the treatment effect for the control and for the treated using the models above\n",
+    "Now, things start to take a turn. For the second stage, we impute the treatment effect for control and treated units using the models above. Let $\\hat D_i^0$ denote an effect imputed for a control unit and $\\hat D_i^1$ an effect imputed for a treated unit:\n",
     " \n",
     "$\n",
-    "\\hat{\\tau}(X, T=0) = \\hat{M}_1(X, T=0) - Y_{T=0}\n",
+    "\\hat D_i^0 = \\hat M_1(X_i)-Y_i, \\qquad T_i=0,\n",
     "$\n",
     " \n",
     "$\n",
-    "\\hat{\\tau}(X, T=1) = Y_{T=1} - \\hat{M}_0(X, T=1)\n",
+    "\\hat D_i^1 = Y_i-\\hat M_0(X_i), \\qquad T_i=1.\n",
     "$\n",
     " \n",
-    "Then, we fit two more models to predict those effects\n",
+    "Then, we fit two more models to predict those effects. The subscript indicates the observed treatment arm used to fit each second-stage model:\n",
     " \n",
     "$\n",
-    "\\hat{M}_{\\tau 0}(X) \\approx E[\\hat{\\tau}(X)|T=0]\n",
+    "\\hat M_{\\tau 0}(x) \\approx E[\\hat D_i^0 \\mid X_i=x,T_i=0],\n",
     "$\n",
     " \n",
     "$\n",
-    "\\hat{M}_{\\tau 1}(X) \\approx E[\\hat{\\tau}(X)|T=1]\n",
+    "\\hat M_{\\tau 1}(x) \\approx E[\\hat D_i^1 \\mid X_i=x,T_i=1].\n",
     "$\n",
     " \n",
-    "If we apply this on the image we've shown before, $\\hat{\\tau}(X, T=0)$, the imputed treatment effect on the untreated, would be the red crosses and the red dashed line would be $\\hat{M}_{\\tau 0}(X)$. Notice that this model is wrong. Because $\\hat{\\tau}(X, T=0)$ was made using the regularized, simple model, estimated on the treated, $\\hat{M}_1$.  The treatment effect it imputes is non linear, since it **doesn't** capture the non-linearity in the Y variable.\n",
+    "If we apply this to the image shown before, $\\hat D^0$, the imputed treatment effect for controls, would be the red crosses and the red dashed line would be $\\hat M_{\\tau 0}(x)$. Notice that this model is wrong in the example. Its imputed target uses the regularized, simple outcome model $\\hat M_1$, estimated on the small treated sample. Because $\\hat M_1$ doesn't capture the non-linearity in $Y$, the treatment effect it imputes is wrong.\n",
     " \n",
-    "In contrast, the blue dots are the imputed treatment effect for the treated, $\\hat{\\tau}(X, T=1)$. These effects are estimated using the correct model, $M_0$, trained in the untreated, large, sample. As a result, since its imputed treatment effects are correct, we are able to train a correct second stage model $\\hat{M}_{\\tau 1}(X)$, shown by the blue line. \n",
+    "In contrast, the blue dots are the imputed treatment effects for treated units, $\\hat D^1$. These effects use the better estimated outcome model $\\hat M_0$, trained on the large control sample. As a result, we are able to train a better second-stage model $\\hat M_{\\tau 1}(x)$, shown by the blue line. \n",
     " \n",
     "![img](data/img/meta-learners/second-stage-x.png)\n",
     " \n",
-    "So we have one model that is wrong because we've impute the treatment effects wrongly and another model that is correct because we've imputed those values correctly. Now, we need a way to combine the two in a way that gives more weight to the correct model. Here is where the propensity score model comes to play. Let $\\hat{e}(x)$ be the propensity score model, we can combine the two second stage models as follows:\n",
+    "So we have one model whose imputed effects are poor and another whose imputed effects are better. Now, we need a way to combine them. This is where the propensity score model comes into play. Let $e(x)=P(T_i=1\\mid X_i=x)$ be the propensity score and $\\hat e(x)$ its estimate. A common X-learner choice combines the two second-stage models as follows:\n",
     " \n",
     "$\n",
-    "\\hat{\\tau(x)} = \\hat{M}_{\\tau 0}(X)\\hat{e}(x) +  \\hat{M}_{\\tau 1}(X)(1-\\hat{e}(x))\n",
+    "\\hat\\tau(x) = \\hat e(x)\\hat M_{\\tau 0}(x) + (1-\\hat e(x))\\hat M_{\\tau 1}(x).\n",
     "$\n",
     " \n",
-    "Since there are very few treated units, $\\hat{e}(x)$ is very small. This will give a very small weight to the wrong model $\\hat{M}_{\\tau 0}(X)$. \n",
+    "Since there are very few treated units in this example, $\\hat e(x)$ is small. This gives little weight to $\\hat M_{\\tau 0}(x)$, whose imputed target depends on the poorly estimated treated-outcome model $\\hat M_1$.\n",
     "\n",
-    "In contrast, $1-\\hat{e}(x)$ is close to one, so we will give a high weight to the correct model $\\hat{M}_{\\tau 1}(X)$. More generally, weighted average using the propensity score will make sure we give more weight to the CATE model that was estimated where the assigned treatment was more likely. In other words, we will favor the model that was trained using more data. The following image shows the estimated CATE given by the X-learner and the T-learner.\n",
+    "In contrast, $1-\\hat e(x)$ is close to one, so we give high weight to $\\hat M_{\\tau 1}(x)$, whose imputed target uses the well-estimated control-outcome model $\\hat M_0$.\n",
+    "\n",
+    "The general direction is worth spelling out because it can look reversed if we focus only on which arm fits the second-stage model:\n",
+    "\n",
+    "- When $\\hat e(x)$ is high, treated observations are relatively abundant near $x$, so $\\hat M_1(x)$ has more local support. The control-arm effect model $\\hat M_{\\tau 0}(x)$ uses $\\hat M_1(x)$ for its missing treated outcome, so it receives the high weight $\\hat e(x)$.\n",
+    "- When $\\hat e(x)$ is low, control observations are relatively abundant near $x$, so $\\hat M_0(x)$ has more local support. The treated-arm effect model $\\hat M_{\\tau 1}(x)$ uses $\\hat M_0(x)$ for its missing control outcome, so it receives the high weight $1-\\hat e(x)$.\n",
+    "\n",
+    "Thus, propensity weighting favors the effect estimate whose **counterfactual outcome component** is trained on the locally more common arm. It does not simply favor the second-stage model fitted on the more common arm. The propensity score is a practical weighting choice, not a guarantee of optimal weighting, and it cannot repair unmeasured confounding or a lack of overlap. The following image shows the estimated CATE given by the X-learner and the T-learner.\n",
     " \n",
     "![img](data/img/meta-learners/t-vs-x-learner.png)\n",
     " \n",
@@ -643,7 +660,9 @@
     " \n",
     "## Key Ideas\n",
     " \n",
-    "Again, the simplest thing we can do is using a single or S-learner with the treatment as a feature. This tends to work well when the treatment is not a weak predictor of the outcome. But if that's not the case, the S-learner tends to be biased towards zero or even drop the treatment entirely. Adding a bit more complexity, we can force the learner to pick up the treatment by using a T-learner. Here, we fit one Machine Learning model per treatment level. This works fine when there are enough samples for all treatment levels, but it can fail when one treatment level has a small sample size, forcing a model to be heavily regularized. To fix that, we can add another level of complexity using an X-learner, where we have two fitting stages and we use a propensity score model to correct potential mistakes from models estimated with very few data points.\n",
+    "Again, the simplest thing we can do is using a single or S-learner with the treatment as a feature. This tends to work well when the treatment is not a weak predictor of the outcome. But if that's not the case, the S-learner tends to be biased towards zero or even drop the treatment entirely. Adding a bit more complexity, we can force the learner to pick up the treatment by using a T-learner. Here, we fit one Machine Learning model per treatment level. This works fine when there are enough samples for all treatment levels, but it can fail when one treatment level has a small sample size, forcing a model to be heavily regularized. To improve on that, the X-learner adds a second fitting stage and uses a propensity-weighted blend to favor the arm-specific effect estimate whose counterfactual outcome model has more local support.\n",
+    " \n",
+    "When these learners are trained on observational data, their predictions have a causal interpretation only under consistency, conditional exchangeability, and overlap. Predictive flexibility and randomized validation are not substitutes for those identification assumptions.\n",
     " \n",
     "One big problem of these learners (except the S-learner) is that they assume a binary or categorical treatment. There is one additional learner that we haven't seen yet, which is more general: the R-learner. But don't worry, we will have an entire chapter dedicated to it. \n",
     " \n",