From 3e9fc6edd397a586ee69f6d443fc6a7370f8a262 Mon Sep 17 00:00:00 2001 From: AliFozooni Date: Sat, 15 Aug 2026 11:02:50 -0700 Subject: [PATCH] Correct continuous-treatment target transformation --- .../20-Plug-and-Play-Estimators.ipynb | 111 ++++++++++++------ 1 file changed, 73 insertions(+), 38 deletions(-) diff --git a/causal-inference-for-the-brave-and-true/20-Plug-and-Play-Estimators.ipynb b/causal-inference-for-the-brave-and-true/20-Plug-and-Play-Estimators.ipynb index 7129851..0de2ca7 100644 --- a/causal-inference-for-the-brave-and-true/20-Plug-and-Play-Estimators.ipynb +++ b/causal-inference-for-the-brave-and-true/20-Plug-and-Play-Estimators.ipynb @@ -99,7 +99,7 @@ " \n", "This seems very odd, because you are saying that the effect of the email can be a negative number, but bear with me. If we do a little bit of math, we can see that, on average or in expectation, this transformed target will be the treatment effect. This is nothing short of amazing. What I'm saying is that by applying this somewhat wacky transformation, I get to estimate something that I can't even observe. \n", " \n", - "To understand that, we need a bit of math. Because of random assignment, we have that $T \\perp Y(0), Y(1)$, which is our old unconfoundedness friend. That implies that $E[T, Y(t)]=E[T]*E[Y(t)]$, which is a consequence of independence.\n", + "To understand that, we need a bit of math. Because of random assignment, we have that $T \\perp Y(0), Y(1)$, which is our old unconfoundedness friend. That implies that $E[T Y(t)]=E[T]*E[Y(t)]$, which is a consequence of independence.\n", "\n", "\n", "Also, we know that\n", @@ -117,14 +117,14 @@ "&= 2E[Y(1)_i| X_i=x] * E[ T_i | X_i=x] - 2E[Y(0)_i| X_i=x]*E[(1-T_i)|X_i=x] \\\\\n", "&= 2E[Y(1)_i| X_i=x] * 0.5 - 2E[Y(0)_i| X_i=x]*0.5 \\\\ \n", "&= E[Y(1)_i| X_i=x] - E[Y(0)_i| X_i=x] \\\\\n", - "&= \\tau(x)_i\n", + "&= \\tau(x)\n", "\\end{align}\n", "$\n", " \n", - "So, this apparently crazy idea ended up being an unbiased estimate of the individual treatment effect $\\tau(x)_i$. Now, we can replace our infeasible optimization criterion with \n", + "So, this apparently crazy idea ended up being a conditionally unbiased pseudo-outcome for the conditional average treatment effect $\\tau(x)$. Now, we can replace our infeasible optimization criterion with \n", " \n", "$\n", - "E[(Y^*_i - \\hat{\\tau}(x)_i)^2]\n", + "E[(Y^*_i - \\hat{\\tau}(x))^2]\n", "$\n", " \n", "In simpler terms, all we have to do is use any regression machine learning model to predict $Y^*_i$ and this model will output treatment effect predictions. \n", @@ -147,7 +147,7 @@ "&= \\dfrac{1}{e(X_i)} E[Y(1)_i * T_i|X_i=x] - \\dfrac{1}{1-e(X_i)} E[Y(0)_i * (1-T_i)| X_i=x]\\\\\n", "&= \\dfrac{1}{e(X_i)} E[Y(1)_i|X_i=x] * E[T_i|X_i=x] - \\dfrac{1}{1-e(X_i)} E[Y(0)_i|X_i=x] * E[(1-T_i)| X_i=x]\\\\\n", "&= E[Y(1)_i|X_i=x] - E[Y(0)_i|X_i=x]\\\\\n", - "&= \\tau(x)_i\n", + "&= \\tau(x)\n", "\\end{align}\n", "$\n", " \n", @@ -713,52 +713,81 @@ "metadata": {}, "source": [ "$\n", - "Y^*_i = (Y_i- \\bar{Y})\\dfrac{(T_i - \\bar{T})}{\\sigma^2_T}\n", + "Y^*_i = (Y_i- \\bar{Y})\\dfrac{T_i - \\bar{T}}{\\sigma^2_T}\n", "$\n", " \n", - "In plain English, we would transform the original target by subtracting the mean from it, then we would multiply it by the treatment, from which we've also subtracted the mean from. Finally, we would divide it by the treatment variance. Alas, we have a target transformation for the continuous case. \n", + "In plain English, we transform the original target by subtracting its mean, multiply it by the demeaned treatment, and divide by the treatment variance. The bars and $\\sigma_T^2$ are sample versions of the population quantities used in the argument below. Alas, we have a target transformation for the continuous case.\n", "\n", "![img](./data/img/plug-and-play-estimators/genious.jpeg)\n", "\n", - "The question now is: does it work? As a matter of fact it does and we can go over a similar proof for why it works, just like we did in the binary case. First, lets call\n", + "The question now is: does it work? It does under a conditionally linear response and a common randomized treatment distribution. Suppose that, for units with $X_i=x$,\n", " \n", "$\n", - "V_i = \\dfrac{(T_i - \\bar{T})}{\\sigma^2_T}\n", + "Y_i = \\alpha(x) + \\tau(x)T_i + \\varepsilon_i, \\qquad E[\\varepsilon_i \\mid T_i,X_i]=0.\n", "$\n", " \n", - "notice that $E[V_i|X_i=x]=0$ because under random assignment $E[T_i|X_i=x]=\\bar{T}$. In other words, for every region of X, $E[T_i]=\\bar{T}$. Also $E[T_i V_i | X_i=x]=1$ because $E[T_i(T_i - \\bar{T})|X_i=x] = E[(T_i - \\bar{T})^2|X_i=x]$, which is the treatment variance. Finally, under conditional independence (which we get for free in the random treatment assignment case), $E[T_i e_i | X_i=x] = E[T_i | X_i=x] E[e_i | X_i=x]$.\n", + "Let $m=E[T_i]$, $v=Var(T_i)$ and\n", " \n", - "To show that this target transformation works, we need to remember that we are estimating the parameter for a local linear model\n", + "$\n", + "V_i = \\dfrac{T_i-m}{v}.\n", + "$\n", + " \n", + "If treatment is randomized independently of $X$, with the same distribution for every $x$, then $E[V_i \\mid X_i=x]=0$ and $E[T_iV_i \\mid X_i=x]=1$. Mean independence of the error also gives $E[\\varepsilon_iV_i \\mid X_i=x]=0$. Therefore,\n", + " \n", + "$\n", + "\\begin{align}\n", + "E[Y^*_i \\mid X_i=x] &= E[(Y_i-E[Y_i])V_i \\mid X_i=x] \\\\\n", + "&= (\\alpha(x)-E[Y_i])E[V_i \\mid X_i=x] \\\\\n", + "&\\quad + \\tau(x)E[T_iV_i \\mid X_i=x] + E[\\varepsilon_iV_i \\mid X_i=x] \\\\\n", + "&= \\tau(x).\n", + "\\end{align}\n", + "$\n", + " \n", + "Randomization is sufficient, but it is not the only setting where this idea applies. Suppose the variables in $X$ are enough to make the treatment as good as random, and define three conditional nuisance functions\n", + " \n", + "$\n", + "m(x)=E[T_i \\mid X_i=x], \\qquad \\mu(x)=E[Y_i \\mid X_i=x], \\qquad v(x)=Var(T_i \\mid X_i=x).\n", + "$\n", + " \n", + "Under the same conditionally linear response, residualizing gives the Robinson decomposition\n", + " \n", + "$\n", + "Y_i-\\mu(X_i)=\\tau(X_i)(T_i-m(X_i))+\\varepsilon_i.\n", + "$\n", + " \n", + "Provided $v(x)>0$, the correctly normalized transformed target is\n", " \n", "$\n", - "Y_i = \\alpha + \\beta T_i + e_i | X_i=x\n", + "Y^*_i = \\dfrac{(Y_i-\\mu(X_i))(T_i-m(X_i))}{v(X_i)}.\n", "$\n", " \n", - "In our example, those would be the linear models for the hot and cold days. Here, we are interested in the $\\beta$ parameter, which is our conditional elasticity or CATE. With all that, we can prove that\n", + "Indeed, conditioning on $X_i=x$ makes the first term below equal to $\\tau(x)$ and the second equal to zero:\n", " \n", "$\n", "\\begin{align}\n", - "E[Y^*_i|X_i=X] &= E[(Y_i-\\bar{Y})V_i | X_i=x] \\\\\n", - "&= E[(\\alpha + \\beta T_i + e_i - \\bar{Y})V_i | X_i=x] \\\\\n", - "&= \\alpha E[V_i | X_i=x] + \\beta E[T_i V_i | X_i=x] + E[e_i V_i | X_i=x] \\\\\n", - "&= \\beta + E[e_i V_i | X_i=x] \\\\\n", - "&= \\beta = \\tau(x)\n", + "E[Y^*_i \\mid X_i=x] &= \\tau(x)\\dfrac{E[(T_i-m(x))^2 \\mid X_i=x]}{v(x)} \\\\\n", + "&\\quad + \\dfrac{E[\\varepsilon_i(T_i-m(x)) \\mid X_i=x]}{v(x)} \\\\\n", + "&= \\tau(x).\n", "\\end{align}\n", "$\n", " \n", - "Bear in mind that this only works when the treatment is randomized. For non randomized treatment, we have to replace $\\bar{T}$ by $M(X_i)$, where $M$ is a model that estimates $E[T_i|X_i=x]$. \n", + "It is important that the denominator is the **conditional variance** $v(X_i)$, not the individual squared residual $(T_i-m(X_i))^2$. The latter is just one noisy realization and can be arbitrarily close to zero. Dividing by it would reduce the pseudo-outcome to a potentially unstable ratio $(Y_i-\\mu(X_i))/(T_i-m(X_i))$, whose unweighted conditional average does not generally equal $\\tau(x)$. In practice, the nuisance functions should be estimated, preferably out of fold when flexible models are used.\n", + " \n", + "Another route is to avoid dividing altogether and estimate $\\tau$ by minimizing the residual-on-residual loss\n", " \n", "$\n", - "Y^*_i = (Y_i- \\bar{Y})\\dfrac{(T_i - M(T_i))}{(T_i - M(T_i))^2}\n", + "\\sum_i \\left[(Y_i-\\hat{\\mu}(X_i))-\\hat{\\tau}(X_i)(T_i-\\hat{m}(X_i))\\right]^2.\n", "$\n", " \n", - "This will make sure that the term $\\alpha E[V_i | X_i=x]$ in the third line vanishes to zero and that the term $E[T_i V_i | X_i=x]$ goes to 1. Notice that you don't actually need $E[T_i V_i | X_i=x]$ to go to 1 if you just want to order units in terms of treatment effect. In other words, if you just want to know in which days demand is more sensitive to price increases but you don't need to know by how much, it doesn't matter if the $\\beta$ estimates are scaled up or down. If that is the case, you can omit the denominator.\n", + "This is the R-loss that we will revisit in Chapter 22. Where $T_i-\\hat m(X_i)\\neq 0$, it is equivalent to regressing the ratio $(Y_i-\\hat\\mu(X_i))/(T_i-\\hat m(X_i))$ with weights $(T_i-\\hat m(X_i))^2$; dropping those weights changes the objective. For a linear-in-parameters least-squares stage, multiplying every design row and target by the square root of its weight reproduces weighted least squares. That trick is not generally equivalent for a nonlinear learner, because scaling the features changes the learner's inputs and hypothesis class. Use the learner's native `sample_weight` support or optimize the R-loss directly instead.\n", + " \n", + "Finally, if we use the unnormalized product\n", " \n", "$\n", - "Y^*_i = (Y_i- \\bar{Y})(T_i - M(T_i))\n", + "\\widetilde{Y}_i=(Y_i-\\mu(X_i))(T_i-m(X_i)),\n", "$\n", " \n", - "If all that math seems tiresome, don't worry. The code is actually very simple. Once again, we transform our training target with the formulas seen above. Here, we have random treatment assignments, so we don't need to build a model that predicts prices. I'm also omitting the denominator, because here I only care about ordering the treatment effect." + "then $E[\\widetilde{Y}_i \\mid X_i=x]=v(x)\\tau(x)$. A constant $v(x)$ guarantees that omitting the denominator preserves the CATE ordering. If $v(x)$ varies across unit profiles, multiplying by it can change the ordering unless additional structure happens to preserve it. Treatment variance is constant in this example because price is randomized from the same distribution for every $X$. The following code uses sample means and omits that common variance because we only care about ordering the treatment effect." ] }, { @@ -891,7 +920,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "This time, the CATE's interpretation is non intuitive. Since we've removed the denominator from the target transformation, this CATE we are seeing is scaled by $Var(X)$. However, this prediction should still order the treatment effect pretty well. To see that, we can use the cumulative gain curve, just like we did before." + "This time, the prediction is not in CATE units. Since we've removed the denominator from the target transformation, the score is scaled by $Var(T)$, not $Var(X)$. Here that treatment variance is a positive constant because price was randomized from the same distribution for every unit profile, so the score should still order the treatment effect pretty well. If $Var(T_i \\mid X_i=x)$ changed with $x$, however, this scaling could change the ranking too. To evaluate the ranking in this example, we can use the cumulative gain curve, just like we did before." ] }, { @@ -938,38 +967,44 @@ " \n", "### Non Linear Treatment Effects\n", " \n", - "Having talked about the continuous case, there is still an elephant in the room we need to adress. We've assumed a linearity on the treatment effect. However, that is very rarely a reasonable assumption. Usually, treatment effects saturate in one form or another. In our example, it's reasonable to think that demand will go down faster at the first units of price increase, but then it will fall more slowly.\n", + "Having talked about the continuous case, there is still an elephant in the room we need to address. We've assumed that the response is linear in the treatment. However, that is very rarely a reasonable assumption. Usually, treatment effects saturate in one form or another. In our example, it's reasonable to think that demand will go down faster at the first units of price increase, but then it will fall more slowly.\n", " \n", "![img](./data/img/plug-and-play-estimators/non-linear-case.png)\n", " \n", - "The problem here is that **elasticity or treatment effect changes with the treatment itself**. In our example, the treatment effect is more intense at the beginning of the curve and smaller as prices get higher. Again,suppose you have two types of days: hot days (yellow) and cold days (blue) and we want to distinguish between the two with a causal model. The thing is that causal models should predict elasticity, but in the nonlinear case, the elasticity for hot and cold days could be the same, if we look at different price points in the curve (right image).\n", + "The problem here is that the **marginal treatment effect changes with the treatment itself**. In our example, the treatment effect is more intense at the beginning of the curve and smaller as prices get higher. Again, suppose you have two types of days: hot days (yellow) and cold days (blue), and we want to distinguish between the two with a causal model. A causal model should predict the slope of the response curve, but in the nonlinear case hot and cold days could have the same slope at different price points (right image).\n", " \n", - "There is no easy way out of this problem and I confess I'm still investigating what works best. For now, the thing that I do is try to think about the functional form of the treatment effect and somehow linearize it. For example, demand usually has the following functional form, where higher $\\alpha$s means that demand falls faster with each price increase\n", + "There is no easy way out of this problem and I confess I'm still investigating what works best. For now, the thing that I do is try to think about the functional form of the treatment effect and somehow linearize it. For example, one common model is the isoelastic demand curve, where $A_i>0$ controls the demand level and a higher $\\alpha$ means that demand falls more for a given percentage increase in price:\n", " \n", "$\n", - "D_i = \\dfrac{1}{P_i^{\\alpha}}\n", + "D_i = \\dfrac{A_i}{P_i^{\\alpha}}.\n", "$\n", "\n", - "So, if I apply the log transformation to both the demand $Y$ and prices $T$, I should get something that is linear.\n", + "If I apply the log transformation to both demand and price, I get something that is linear. In this log-log model, $-\\alpha$ is the price elasticity of demand.\n", "\n", "$\n", "\\begin{align}\n", - "log(D)_i &= log\\bigg(\\dfrac{1}{P_i^{\\alpha}}\\bigg) \\\\\n", - "&= log(1) - log(P_i^{\\alpha}) \\\\\n", - "&= log(1) - log(P_i^{\\alpha}) \\\\\n", - "&= - \\alpha * log(P_i) \\\\\n", + "\\log(D_i) &= \\log\\left(\\dfrac{A_i}{P_i^{\\alpha}}\\right) \\\\\n", + "&= \\log(A_i) - \\alpha\\log(P_i). \\\\\n", "\\end{align}\n", "$\n", " \n", - "Linearization is not so easy to do, as it involves some thinking. But you can also try stuff out and see what works best. Often, things like logs and square roots help.\n", + "This functional form is useful for estimating responsiveness, but it does not by itself produce an interior revenue-maximizing price when marginal cost is zero. Revenue is\n", + " \n", + "$\n", + "R_i(P_i)=P_iD_i(P_i)=A_iP_i^{1-\\alpha}.\n", + "$\n", + " \n", + "It is increasing when $\\alpha<1$, constant when $\\alpha=1$, and decreasing when $\\alpha>1$. Choosing an unconstrained optimal price therefore requires additional structure, such as positive marginal costs and a profit objective (which produces an interior markup when $\\alpha>1$), realistic price bounds, or a more flexible demand curve with non-constant elasticity. The isoelastic curve here should be read as an illustration of linearization, not as a complete pricing model.\n", + " \n", + "Linearization is not so easy to do, as it involves some thinking. But you can also try things out and see what works best. Often, transformations like logs and square roots help.\n", " \n", "## Key Ideas\n", "\n", - "We are now moving in the direction of estimating conditional average treatment effects using machine learning models. The biggest challenge when doing so is adapting a predictive model to one that estimates causal effects. Another way of thinking about it is that predictive models focus on estimating the outcome Y as a function of features X and possibly treatment T $Y = M(X, T) $ while causal models need to estimate the partial derivative of this output function on the treatment $ \\partial Y = \\partial M(X, T) $. This is far from trivial, because while we do observe the outcome Y, we can't observe $\\partial Y$, at least not on an individual level. As a consequence, we need to be creative when designing an objective function for our models. \n", + "We are now moving in the direction of estimating conditional average treatment effects using machine learning models. The biggest challenge when doing so is adapting a predictive model to one that estimates causal effects. Another way of thinking about it is that predictive models focus on estimating the outcome $Y$ as a function of features $X$ and possibly treatment $T$, $Y=M(X,T)$, while causal models need to estimate how that response changes with treatment, $\\partial M(X,T)/\\partial T$. This is far from trivial, because while we do observe the outcome $Y$, we can't observe an individual treatment effect or response derivative directly. As a consequence, we need to be creative when designing an objective function for our models. \n", " \n", - "Here, we saw a very simple technique of target transformation. The idea is to combine the original target Y with the treatment T to form a transformed target which is, in expectation, equal to the CATE. With that new target, we can plug any predictive ML model to estimate it and then the model predictions will be CATE estimates. As a side note, besides target transformation, this method also goes by the name of **F-Learner**. \n", + "Here, we saw a very simple technique of target transformation. The idea is to combine the original target $Y$ with the treatment $T$ to form a transformed target whose conditional expectation, under the stated identification and response assumptions, is equal to the CATE. With that new target, we can plug in a predictive ML model and use its predictions as CATE estimates. As a side note, besides target transformation, the binary-treatment method also goes by the name of **F-Learner**. \n", " \n", - "With all that simplicity, there is also a price to pay. The transformed target is a very noisy estimate of the individual treatment effect and that noise will be transferred to the model estimates in the form of variance. This makes target transformation better suited for big data applications, where variance is less of a problem due to sheer sample sizes. Another downside of the target transformation method is that it is only defined for binary or categorical treatments. We did our best to come up with a continuous version of the approach and even ended up with something that seemed to work, but up until now, there is no solid theoretical framework to back it up. \n", + "With all that simplicity, there is also a price to pay. The transformed target is a very noisy estimate of the individual treatment effect and that noise will be transferred to the model estimates in the form of variance. This makes target transformation better suited for big data applications, where variance is less of a problem due to sheer sample sizes. The original F-Learner transformation is for binary or categorical treatments. The continuous transformation in this chapter has a clean justification for a response that is conditionally linear in treatment: normalize the residual product by $Var(T\\mid X)$, or use the corresponding R-loss. It does not, however, recover an arbitrary nonlinear dose-response curve. Constant conditional treatment variance guarantees that dropping the variance does not change the ranking; when the variance changes with $X$, the ranking need not be preserved. \n", " \n", "Finally, we've ended with a discussion on non linear treatment effects and the challenges that come with it. Namely, when the treatment effect changes with the treatment itself, we might mistakenly think units have the same treatment response curve because they have the same responsiveness to the treatment, but actually they are just receiving different treatment amounts. \n", "\n", @@ -978,7 +1013,7 @@ " \n", "The things I've written here are mostly stuff from my head. I've learned them through experience. This means that they have **not** passed the academic scrutiny that good science often goes through. Instead, notice how I'm talking about things that work in practice, but I don't spend too much time explaining why that is the case. It's a sort of science from the streets, if you will. However, I am putting this up for public scrutiny, so, by all means, if you find something preposterous, open an issue and I'll address it to the best of my efforts.\n", " \n", - "Most of this chapter draws from Susan Atheys' and Guido W. Imbens' paper, *Machine Learning Methods for Estimating Heterogeneous Causal Effects*. Some material about target transformation can also be found on Pierre Gutierrez' and Jean-Yves Gérardy's paper, *Causal Inference and Uplift Modeling: A review of the literature*. Note that these papers only cover the binary treatment case. Another review of causal models for CATE estimation that references the F-Learner is *Meta-learners for Estimating Heterogeneous Treatment Effects using Machine Learning*, by Künzel et al, 2019. \n", + "Most of this chapter draws from Susan Atheys' and Guido W. Imbens' paper, *Machine Learning Methods for Estimating Heterogeneous Causal Effects*. Some material about target transformation can also be found on Pierre Gutierrez' and Jean-Yves Gérardy's paper, *Causal Inference and Uplift Modeling: A review of the literature*. Note that these papers only cover the binary treatment case. Another review of causal models for CATE estimation that references the F-Learner is *Meta-learners for Estimating Heterogeneous Treatment Effects using Machine Learning*, by Künzel et al, 2019. The residual-on-residual derivation builds on Robinson's 1988 paper, *Root-N-Consistent Semiparametric Regression*. The R-loss is developed by Nie and Wager in [*Quasi-Oracle Estimation of Heterogeneous Treatment Effects*](https://arxiv.org/abs/1712.04912) and is covered in more detail in Chapter 22.\n", "\n", " \n", "## Contribute\n",