Methodology#

In this section we dive more deeply into the algorithms implemented in CausalML. To provide a basis for the discussion, we review some of the frameworks and definitions used in the literature.

We use the Neyman-Rubin potential outcomes framework and assume Y represents the outcome, W represents the treatment assignment, and X_i the observed covariates.

If you are deciding which of these methods to use on a given problem, start with Choosing an Estimator; this page documents how each method works.

Supported Algorithms#

CausalML currently supports the following methods:

Meta-Learner Algorithms#

Relevant classes: BaseSRegressor/BaseSClassifier, BaseTRegressor/BaseTClassifier, BaseXRegressor/BaseXClassifier, BaseRRegressor/BaseRClassifier and BaseDRRegressor/BaseDRClassifier in causalml.inference.meta, each taking any scikit-learn-compatible model as its base learner. Preconfigured variants (XGBTRegressor, LRSRegressor, …) are listed in the API Reference.

A meta-algorithm (or meta-learner) is a framework to estimate the Conditional Average Treatment Effect (CATE) using any machine learning estimators (called base learners) [18].

A meta-algorithm uses either a single base learner while having the treatment indicator as a feature (e.g. S-learner), or multiple base learners separately for each of the treatment and control groups (e.g. T-learner, X-learner and R-learner).

Confidence intervals of average treatment effect estimates are calculated based on the lower bound formula (7) from [16].

S-Learner#

S-learner estimates the treatment effect using a single machine learning model as follows:

Stage 1

Estimate the average outcomes \(\mu(x)\) with covariates \(X\) and an indicator variable for treatment \(W\):

\[\mu(x,w) = E[Y \mid X=x, W=w]\]

using a machine learning model.

Stage 2

Define the CATE estimate as:

\[\hat\tau(x) = \hat\mu(x, W=1) - \hat\mu(x, W=0)\]

Including the propensity score in the model can reduce bias from regularization induced confounding [38].

When the control and treatment groups are very different in covariates, a single linear model is not sufficient to encode the different relevant dimensions and smoothness of features for the control and treatment groups [1].

T-Learner#

T-learner [18] consists of two stages as follows:

Stage 1

Estimate the average outcomes \(\mu_0(x)\) and \(\mu_1(x)\):

\[\begin{split}\mu_0(x) = E[Y(0)|X=x] \\ \mu_1(x) = E[Y(1)|X=x]\end{split}\]

using machine learning models.

Stage 2

Define the CATE estimate as:

\[\hat\tau(x) = \hat\mu_1(x) - \hat\mu_0(x)\]

X-Learner#

X-learner [18] is an extension of T-learner, and consists of three stages as follows:

Stage 1

Estimate the average outcomes \(\mu_0(x)\) and \(\mu_1(x)\):

\[\begin{split}\mu_0(x) = E[Y(0)|X=x] \\ \mu_1(x) = E[Y(1)|X=x]\end{split}\]

using machine learning models.

Stage 2

Impute the user level treatment effects, \(D^1_i\) and \(D^0_j\) for user \(i\) in the treatment group based on \(\mu_0(x)\), and user \(j\) in the control groups based on \(\mu_1(x)\):

\[\begin{split}D^1_i = Y^1_i - \hat\mu_0(X^1_i) \\ D^0_i = \hat\mu_1(X^0_i) - Y^0_i\end{split}\]

then estimate \(\tau_1(x) = E[D^1|X=x]\), and \(\tau_0(x) = E[D^0|X=x]\) using machine learning models.

Stage 3

Define the CATE estimate by a weighted average of \(\tau_1(x)\) and \(\tau_0(x)\):

\[\tau(x) = g(x)\tau_0(x) + (1 - g(x))\tau_1(x)\]

where \(g \in [0, 1]\). We can use propensity scores for \(g(x)\).

R-Learner#

R-learner [24] uses the cross-validation out-of-fold estimates of outcomes \(\hat{m}^{(-i)}(x_i)\) and propensity scores \(\hat{e}^{(-i)}(x_i)\). It consists of two stages as follows:

Stage 1

Fit \(\hat{m}(x)\) and \(\hat{e}(x)\) with machine learning models using cross-validation.

Stage 2

Estimate treatment effects by minimising the R-loss, \(\hat{L}_n(\tau(x))\):

\[\hat{L}_n(\tau(x)) = \frac{1}{n} \sum^n_{i=1}\big(\big(Y_i - \hat{m}^{(-i)}(X_i)\big) - \big(W_i - \hat{e}^{(-i)}(X_i)\big)\tau(X_i)\big)^2\]

where \(\hat{e}^{(-i)}(X_i)\), etc. denote the out-of-fold held-out predictions made without using the \(i\)-th training sample.

Doubly Robust (DR) learner#

DR-learner [17] estimates the CATE via cross-fitting a doubly-robust score function in two stages as follows. We start by randomly split the data \(\{Y, X, W\}\) into 3 partitions \(\{Y^i, X^i, W^i\}, i=\{1,2,3\}\).

Stage 1

Fit a propensity score model \(\hat{e}(x)\) with machine learning using \(\{X^1, W^1\}\), and fit outcome regression models \(\hat{m}_0(x)\) and \(\hat{m}_1(x)\) for treated and untreated users with machine learning using \(\{Y^2, X^2, W^2\}\).

Stage 2

Use machine learning to fit the CATE model, \(\hat{\tau}(X)\) from the pseudo-outcome

\[\phi = \frac{W-\hat{e}(X)}{\hat{e}(X)(1-\hat{e}(X))}\left(Y-\hat{m}_W(X)\right)+\hat{m}_1(X)-\hat{m}_0(X)\]

with \(\{Y^3, X^3, W^3\}\)

Stage 3

Repeat Stage 1 and Stage 2 again twice, rotating the three partitions so that each one trains the CATE model exactly once. First use \(\{Y^2, X^2, W^2\}\), \(\{Y^3, X^3, W^3\}\), and \(\{Y^1, X^1, W^1\}\) for the propensity score model, the outcome models, and the CATE model. Then use \(\{Y^3, X^3, W^3\}\), \(\{Y^1, X^1, W^1\}\), and \(\{Y^2, X^2, W^2\}\) for the propensity score model, the outcome models, and the CATE model. The final CATE model is the average of the 3 CATE models.

Doubly Robust Instrumental Variable (DRIV) learner#

We combine the idea from DR-learner [17] with the doubly robust score function for LATE described in [10] to estimate the conditional LATE. Towards that end, we start by randomly split the data \(\{Y, X, W, Z\}\) into 3 partitions \(\{Y^i, X^i, W^i, Z^i\}, i=\{1,2,3\}\).

Stage 1

Fit propensity score models \(\hat{e}_0(x)\) and \(\hat{e}_1(x)\) for assigned and unassigned users using \(\{X^1, W^1, Z^1\}\), and fit outcome regression models \(\hat{m}_0(x)\) and \(\hat{m}_1(x)\) for assigned and unassigned users with machine learning using \(\{Y^2, X^2, Z^2\}\). Assignment probability, \(p_Z\), can either be user provided or come from a simple model, since in most use cases assignment is random by design.

Stage 2

Use machine learning to fit the conditional LATE model, \(\hat{\tau}(X)\) by minimizing the following loss function

\[\begin{split}L(\hat{\tau}(X)) = \hat{E} &\left[\left(\hat{m}_1(X)-\hat{m}_0(X)+\frac{Z(Y-\hat{m}_1(X))}{p_Z}-\frac{(1-Z)(Y-\hat{m}_0(X))}{1-p_Z} \right.\right.\\ &\left.\left.\quad -\Big(\hat{e}_1(X)-\hat{e}_0(X)+\frac{Z(W-\hat{e}_1(X))}{p_Z}-\frac{(1-Z)(W-\hat{e}_0(X))}{1-p_Z}\Big) \hat{\tau}(X) \right)^2\right]\end{split}\]

with \(\{Y^3, X^3, W^3\}\)

Stage 3

Similar to the DR-Learner Repeat Stage 1 and Stage 2 again twice with different permutations of partitions for estimation. The final conditional LATE model is the average of the 3 conditional LATE models.

CATE Model Evaluation#

Cross-validation cannot be used to select among fitted CATE models because the ground-truth treatment effect is never observed for any unit – only one of the two potential outcomes is realized. The methods below are surrogate metrics that score fitted CATE models using only observed data, so that models such as the S/T/X/R/DR-learners above can be compared and selected among.

These metrics evaluate effect-magnitude accuracy – how close a model’s CATE estimate is to the truth – which is a distinct question from targeting quality – whether a model correctly ranks units by benefit – addressed by RATE.

DR (Doubly Robust) pseudo-outcome loss#

Following [17], we construct the same doubly-robust pseudo-outcome used in the DR-learner above,

\[\phi = \frac{W-\hat{e}(X)}{\hat{e}(X)(1-\hat{e}(X))}\left(Y-\hat{m}_W(X)\right)+\hat{m}_1(X)-\hat{m}_0(X)\]

with \(\hat{e}(x)\), \(\hat{m}_0(x)\), and \(\hat{m}_1(x)\) cross-fit so that \(\phi_i\) never depends on a model that saw unit \(i\) during training. Under either correct propensity or correct outcome-model specification, \(E[\phi \mid X]\) is an unbiased estimate of the true CATE, which lets \(\phi\) stand in for the unobserved ground truth when scoring a fitted CATE model \(\hat\tau(x)\):

\[\text{DR loss}(\hat\tau) = \frac{1}{n}\sum_i \left(\hat\tau(X_i) - \phi_i\right)^2\]

Lower is better. [23] found DR-based metrics dominate across 78 benchmark datasets for CATE model selection.

Plug-in T-learner loss#

A simpler, non-doubly-robust baseline: fit a cross-fit T-Learner, and use its held-out estimate \(\hat\mu_1(x) - \hat\mu_0(x)\) directly as a proxy for the true CATE:

\[\text{T-loss}(\hat\tau) = \frac{1}{n}\sum_i \left(\hat\tau(X_i) - (\hat\mu_1(X_i) - \hat\mu_0(X_i))\right)^2\]

This is biased under a misspecified outcome model, unlike the DR loss above. Despite that, [23] found it is never dominated across their benchmark datasets, making it a useful complement to DR-based scoring rather than a replacement.

RATE#

The Rank-Weighted Average Treatment Effect (RATE) of [34] scores a model as a targeting rule: not how close \(\hat\tau(x)\) is to the truth, but whether the units it ranks highest are the ones that benefit most.

Rank units by a prioritization score \(\hat\tau(x)\), treat the top \(q\) fraction, and measure the gain over treating everyone. That is the Targeting Operator Characteristic (TOC) curve:

\[\text{TOC}(q) = E\left[Y(1)-Y(0) \mid \hat\tau(X) \ge F^{-1}_{\hat\tau}(1-q)\right] - E\left[Y(1)-Y(0)\right]\]

where \(F^{-1}_{\hat\tau}(1-q)\) is the \((1-q)\) quantile of the score, so the conditioning set is exactly the top-\(q\) fraction. \(\text{TOC}(1) = 0\) by construction, since selecting everyone gives back the overall ATE, and a curve that is positive at small \(q\) means the rule has found units with above-average benefit. RATE reduces the curve to a single number with a weight function \(\alpha(q)\):

\[\text{RATE} = \int_0^1 \alpha(q)\,\text{TOC}(q)\,dq\]

Two weightings are implemented, and the choice between them is one of statistical power rather than correctness:

  • weighting="autoc" (the default) sets \(\alpha(q) = 1/q\). It weights the highest-priority units most heavily and is most powerful when the benefit is concentrated in a small subgroup.

  • weighting="qini" sets \(\alpha(q) = q\), which recovers the Qini coefficient. It is more powerful when treatment effects are diffuse across the population.

causalml.metrics.rate_score returns the scalar; causalml.metrics.get_toc and causalml.metrics.plot_toc return and draw the curve behind it. The integral is evaluated on the discrete quantile grid as a midpoint-weighted mean whose weights are normalized to sum to one, so the absolute scale can differ slightly from the continuous definition above while the ordering of models is unchanged.

Passing return_ci=True adds a standard error, a confidence interval and a p-value for \(H_0: \text{RATE} = 0\) – the null that the ranking is no better than random – from a half-sample bootstrap that draws \(n/2\) units without replacement, the resampling scheme covered by the functional central limit theorem in [34].

One assumption is easy to miss: without a ground-truth effect column, the subset ATE inside the TOC is a plain difference in mean outcomes between the treated and control units of each band, which is unbiased under randomization but not under confounding. For observational data, pass cross-fit AIPW pseudo-outcomes – the same \(\phi\) as in the DR loss above – as treatment_effect_col, so that the ranking is scored against a doubly-robust estimate rather than a raw difference in means.

Tree-Based Algorithms#

Relevant classes: the uplift trees UpliftTreeClassifier and UpliftRandomForestClassifier, which split on the divergence criteria below, and the causal trees CausalTreeRegressor and CausalRandomForestRegressor, regression trees that split on treatment-effect heterogeneity and estimate leaves honestly by default (see Honest estimation). All live in causalml.inference.tree.

Uplift Tree#

The Uplift Tree approach consists of a set of methods that use a tree-based algorithm where the splitting criterion is based on differences in uplift. [27] proposed three different ways to quantify the gain in divergence as the result of splitting [12]:

\[D_{gain} = D_{after_{split}} (P^T, P^C) - D_{before_{split}}(P^T, P^C)\]

where \(D\) measures the divergence and \(P^T\) and \(P^C\) refer to the probability distribution of the outcome of interest in the treatment and control groups, respectively. Three different ways to quantify the divergence, KL, ED and Chi, are implemented in the package.

Throughout this subsection the outcome is binary, so \(P^T\) and \(P^C\) are each described by a single number: let \(p\) be the sample mean of the outcome in the treatment group and \(q\) the sample mean in the control group, both computed within the node being scored. Each divergence below sums over the two outcome classes \(\{1, 0\}\) – which is where the second term in each expression comes from – and \(D_{gain}\) then combines the parent’s divergence with its children’s as above [12].

KL#

The Kullback-Leibler (KL) divergence is given by:

\[KL(P^T : P^C) = p\log\frac{p}{q} + (1-p)\log\frac{1-p}{1-q}\]

At \(p = 0\) and \(p = 1\) the expression is taken at its limit, \(-\log(1-q)\) and \(-\log q\) respectively.

ED#

The Euclidean Distance is given by:

\[ED(P^T : P^C) = (p - q)^2 + \big((1-p) - (1-q)\big)^2 = 2(p - q)^2\]

where the notation is the same as above.

Chi#

Finally, the \(\chi^2\)-divergence is given by:

\[\chi^2(P^T : P^C) = \frac{(p - q)^2}{q} + \frac{(p - q)^2}{1 - q}\]

where the notation is again the same as above.

DDP#

Another Uplift Tree algorithm that is implemented is the delta-delta-p (\(\Delta\Delta P\)) approach by [9], where the sample splitting criterion is defined as follows:

\[\Delta\Delta P=|(P^T(y|a_0)-P^C(y|a_0) - (P^T(y|a_1)-P^C(y|a_1)))|\]

where \(a_0\) and \(a_1\) are the outcomes of a Split A, \(y\) is the selected class, and \(P^T\) and \(P^C\) are the response rates of treatment and control group, respectively. In other words, we first calculate the difference in the response rate in each branch (\(\Delta P_{left}\) and \(\Delta P_{right}\)), and subsequently, calculate their differences (\(\Delta\Delta P = |\Delta P_{left} - \Delta P_{right}|\)).

IDDP#

Build upon the \(\Delta\Delta P\) approach, the IDDP approach by [28] is implemented, where the sample splitting criterion is defined as follows:

\[IDDP = \frac{\Delta\Delta P^*}{I(\phi, \phi_l, \phi_r)}\]

where \(\Delta\Delta P^*\) is defined as \(\Delta\Delta P - |E[Y(1) - Y(0) \mid X \in \phi]|\) and \(I(\phi, \phi_l, \phi_r)\) is defined as:

\[\begin{split}I(\phi, \phi_l, \phi_r) = H(\frac{n_t(\phi)} {n(\phi)}, \frac{n_c(\phi)}{n(\phi)}) * 2 \frac{1+\Delta\Delta P^*}{3} + \frac{n_t(\phi)}{n(\phi)} H(\frac{n_t(\phi_l)}{n(\phi)}, \frac{n_t(\phi_r)}{n(\phi)}) \\ + \frac{n_c(\phi)}{n(\phi)} * H(\frac{n_c(\phi_l)}{n(\phi)}, \frac{n_c(\phi_r)}{n(\phi)}) + \frac{1}{2}\end{split}\]

where the entropy H is defined as \(H(p,q)=(-p*log_2(p)) + (-q*log_2(q))\) and where \(\phi\) is a subset of the feature space associated with the current decision node, and \(\phi_l\) and \(\phi_r\) are the left and right child nodes, respectively. \(n_t(\phi)\) is the number of treatment samples, \(n_c(\phi)\) the number of control samples, and \(n(\phi)\) the number of all samples in the current (parent) node.

IT#

Further, the package implements the Interaction Tree (IT) proposed by [32], where the sample splitting criterion maximizes the G statistic among all permissible splits:

\[G(s^*) = \max_s G(s)\]

where \(G(s)=t^2(s)\) and \(t(s)\) is defined as:

\[t(s) = \frac{(y^L_1 - y^L_0) - (y^R_1 - y^R_0)}{\sigma \sqrt{1/n_1 + 1/n_2 + 1/n_3 + 1/n_4}}\]

where \(\sigma=\sqrt{\sum_{i=1}^4 w_i s_i^2}\) is a pooled estimator of the constant standard deviation, and \(w_i=(n_i-1)/\sum_{j=1}^4(n_j-1)\). Further, \(y^L_1\), \(s^2_1\), and \(n_1\) are the sample mean, the sample variance, and the sample size for the treatment group in the left child node, respectively. Similar notation applies to the other quantities.

Note that this implementation deviates from the original implementation in that (1) the pruning techniques and (2) the validation method for determining the best tree size are different.

CIT#

Also, the package implements the Causal Inference Tree (CIT) by [31], where the sample splitting criterion calculates the likelihood ratio test statistic:

\[\begin{split}LRT(s) = -n_{\tau L}/2 * ln(n_{\tau L} SSE_{\tau L}) -n_{\tau R}/2 * ln(n_{\tau R} SSE_{\tau R}) + \\ n_{\tau L1} ln n_{\tau L1} + n_{\tau L0} ln n_{\tau L0} + n_{\tau R1} ln n_{\tau R1} + n_{\tau R0} ln n_{\tau R0}\end{split}\]

where \(n_{\tau}\), \(n_{\tau 0}\), and \(n_{\tau 1}\) are the total number of observations in node \(\tau\), the number of observations in node \(\tau\) that are assigned to the control group, and the number of observations in node \(\tau\) that are assigned to the treatment group, respectively. \(SSE_{\tau}\) is defined as:

\[SSE_{\tau} = \sum_{i \in \tau: t_i=1}(y_i - \hat{y_{t1}})^2 + \sum_{i \in \tau: t_i=0}(y_i - \hat{y_{t0}})^2\]

and \(\hat{y_{t0}}\) and \(\hat{y_{t1}}\) are the sample average responses of the control and treatment groups in node \(\tau\), respectively.

Note that this implementation deviates from the original implementation in that (1) the pruning techniques and (2) the validation method for determining the best tree size are different.

CTS#

The final Uplift Tree algorithm that is implemented is the Contextual Treatment Selection (CTS) approach by [35], where the sample splitting criterion is defined as follows:

\[\hat{\Delta}_{\mu}(s) = \hat{p}(\phi_l \mid \phi) \times \max_{t=0, ..., K}\hat{y}_t(\phi_l) + \hat{p}(\phi_r \mid \phi) \times \max_{t=0, ..., K}\hat{y}_t(\phi_r) - \max_{t=0, ..., K}\hat{y}_t(\phi)\]

where \(\phi_l\) and \(\phi_r\) refer to the feature subspaces in the left leaf and the right leaves respectively, \(\hat{p}(\phi_j \mid \phi)\) denotes the estimated conditional probability of a subject’s being in \(\phi_j\) given \(\phi\), and \(\hat{y}_t(\phi_j)\) is the conditional expected response under treatment \(t\).

Honest estimation#

A tree that estimates each leaf’s treatment effect on the same rows that chose its splits inherits the selection bias of the split search: the search favors partitions whose in-sample effects look large, so re-using those rows makes the leaf estimates systematically too extreme. Honest estimation [4] removes this bias by splitting the sample in two – the tree structure is grown on one half and each leaf’s per-group outcome means are re-estimated on the other – at the cost of fitting each half with less data.

CausalTreeRegressor and CausalRandomForestRegressor estimate honestly by default (honesty=True), matching the defaults of grf and EconML. UpliftTreeClassifier and UpliftRandomForestClassifier accept the same parameters with honesty=False by default. estimation_sample_size (default 0.5) sets the held-out fraction, and the split is stratified on treatment.

The trade is variance: each half sees only part of the data, so an honest tree is shallower and noisier per leaf. Pass honesty=False when halving the sample leaves too little to split on.

On CausalRandomForestRegressor, honesty composes with the bootstrap: each tree’s bootstrap counts weight both the structure fit and the leaf re-estimation. This is not the subsample-without-replacement scheme of generalized random forests [5].

Pruning#

Both tree families grow greedily and can overfit. Each offers a validation-based way to cut the fitted tree back; both are off by default because they change fitted trees.

UpliftTreeClassifier prunes against a held-out sample: prune(X, treatment, y) collapses an internal node whose split does not improve the treatment-effect estimate on that sample, controlled by min_gain and prune_rule, cascading upward until no further merge helps. Passing prune_fraction to fit runs the same procedure as part of the fit: that fraction of the rows is held out, stratified on treatment and outcome, the tree grows on the remainder, and pruning uses the holdout.

CausalTreeRegressor uses cost-complexity pruning, the scheme behind scikit-learn’s ccp_alpha: a penalty per leaf defines a sequence of nested subtrees, and a larger ccp_alpha keeps a smaller one. Setting ccp_alpha="cv" selects the subtree by cv_folds-fold cross-validation instead, scoring each candidate with the honest splitting objective on the held-out fold; it requires honesty=True and completes the algorithm of [4]. Prefer it on a single tree, where the gains grow with the noise level – on CausalRandomForestRegressor the same option measured no gain, since averaging across trees already removes the variance that pruning targets.

Neural Network Algorithms#

Relevant classes: DragonNet in causalml.inference.tf or causalml.inference.jax, and CEVAE in causalml.inference.torch or causalml.inference.jax.

Both methods below are optional backends. Install them with the tf, torch or jax extras, e.g. pip install causalml[tf].

CEVAE#

The Causal Effect Variational Autoencoder (CEVAE) [22] targets the case where the confounders are never observed directly. Instead of assuming that the covariates \(X\) satisfy unconfoundedness, it assumes \(X\) are noisy proxies for a latent confounder \(Z\) — for example, a zip code standing in for socio-economic status — and recovers the treatment effect by inferring \(Z\).

The generative model is

\[Z \sim p(Z), \quad X \sim p(X \mid Z), \quad W \sim p(W \mid Z), \quad Y \sim p(Y \mid W, Z)\]

where each distribution is parameterised by a neural network. The outcome model uses a disjoint pair of networks for \(p(Y \mid W=0, Z)\) and \(p(Y \mid W=1, Z)\), so a treatment arm with few observations does not have its outcome head diluted by the other arm.

Because \(Z\) is latent, the posterior \(p(Z \mid X, W, Y)\) is intractable. CEVAE approximates it with an inference network \(q(Z \mid X, W, Y)\) trained jointly with the generative model by maximising the evidence lower bound (ELBO). Two auxiliary networks, \(q(W \mid X)\) and \(q(Y \mid X, W)\), are trained alongside it, because at prediction time only \(X\) is available: the treatment and outcome that the inference network expects have to be supplied by the auxiliary heads before \(Z\) can be inferred for a new unit.

The CATE estimate marginalises the outcome heads over the inferred latent:

\[\hat\tau(x) = \mathbb{E}_{Z \sim q(Z \mid X=x)} \big[\mathbb{E}[Y \mid W=1, Z] - \mathbb{E}[Y \mid W=0, Z]\big]\]

CausalML wraps Pyro’s implementation in causalml.inference.torch.CEVAE and also ships a JAX/flax.nnx backend.

DragonNet#

DragonNet [29] starts from the sufficiency of the propensity score: if adjusting for \(X\) suffices to identify the treatment effect, then adjusting for \(e(X) = P(W=1 \mid X)\) alone also suffices. A representation of \(X\) therefore only has to retain what predicts treatment, and a representation trained purely to predict the outcome may discard exactly that.

The architecture makes this explicit with three heads on a shared trunk \(Z(x)\) of three fully-connected layers:

  • a propensity head \(\hat{e}(x)\), a single sigmoid unit on \(Z(x)\);

  • two outcome heads \(\hat{Q}(0, x)\) and \(\hat{Q}(1, x)\), each its own stack of layers.

Attaching the propensity head directly to \(Z(x)\) — with no hidden layer of its own — is what forces the shared trunk to keep the propensity information rather than discarding it in favour of outcome prediction.

The second contribution is targeted regularization, which adds a trainable scalar \(\epsilon\) and an extra loss term. Writing \(\hat{y} = w\hat{Q}(1,x) + (1-w)\hat{Q}(0,x)\) and

\[h(w, x) = \frac{w}{\hat{e}(x)} - \frac{1-w}{1 - \hat{e}(x)}\]

the perturbed prediction is \(\tilde{y} = \hat{y} + \epsilon \, h(w, x)\), and the training objective becomes

\[\mathcal{L} = \mathcal{L}_{\text{DragonNet}} + \beta \sum_i \big(y_i - \tilde{y}_i\big)^2\]

with \(\beta\) the ratio argument. This is a one-step correction in the spirit of TMLE: at the optimum the fitted model satisfies the estimating equation for the ATE, which is what gives the resulting estimator its asymptotic guarantees. It is enabled by default via targeted_reg=True.

The CATE estimate is the difference of the two outcome heads:

\[\hat\tau(x) = \hat{Q}(1, x) - \hat{Q}(0, x)\]

CausalML ships DragonNet as causalml.inference.tf.DragonNet (TensorFlow) and a JAX/flax.nnx backend.

Value optimization methods#

Relevant classes: CounterfactualUnitSelector and CounterfactualValueEstimator in causalml.optimize. PolicyLearner in the same module learns a treatment-assignment policy directly.

The package supports methods for assigning treatment groups when treatments are costly. To understand the problem, it is helpful to divide the population into four categories according to how a unit’s outcome responds to being treated:

  • Persuadables. Those who will have a favourable outcome if and only if they are treated.

  • Sure things. Those who will have a favourable outcome whether or not they are treated.

  • Lost causes. Those who will never have a favourable outcome whether or not they are treated.

  • Sleeping dogs. Those who will have a favourable outcome if and only if they are not treated.

These four response types are also referred to as compliers, always-takers, never-takers and defiers. Note that those names are standard in the instrumental variables literature for a different partition of the population – one defined by whether a unit takes the treatment in response to its assignment, which is the sense used in the LATE section below (see [2] for that one). The two partitions do not coincide, so the uplift names are used here.

Counterfactual Unit Selection#

[21] propose a method for selecting units for treatments using counterfactual logic. Suppose the following benefits for selecting units belonging to the different categories above:

  • Persuadables: \(\beta\)

  • Sure things: \(\gamma\)

  • Lost causes: \(\theta\)

  • Sleeping dogs: \(\delta\)

If \(X\) denotes the set of individual’s features, the unit selection problem can be formulated as follows:

\[\operatorname*{argmax}_X \beta P(\text{persuadable} \mid X) + \gamma P(\text{sure thing} \mid X) + \theta P(\text{lost cause} \mid X) + \delta P(\text{sleeping dog} \mid X)\]

The problem can be reformulated using counterfactual logic. Suppose \(W = w\) indicates that an individual is treated and \(W = w'\) indicates he or she is untreated. Similarly, let \(F = f\) denote a favourable outcome for the individual and \(F = f'\) an unfavourable outcome. Then the optimization problem becomes:

\[\operatorname*{argmax}_X \beta P(f_w, f'_{w'} \mid X) + \gamma P(f_w, f_{w'} \mid X) + \theta P(f'_w, f'_{w'} \mid X) + \delta P(f_{w'}, f'_{w} \mid X)\]

Note that the above simply follows from the definitions of the relevant user segments. [21] then use counterfactual logic ([26]) to solve the above optimization problem under certain conditions.

N.B. The current implementation in the package is highly experimental.

Counterfactual Value Estimator#

The counterfactual value estimation method implemented in the package predicts the outcome for a unit under different treatment conditions using a standard machine learning model. The expected value of assigning a unit into a particular treatment is then given by

\[\mathbb{E}[(v - cc_w)Y_w - ic_w]\]

where \(Y_w\) is the probability of a favourable event (such as conversion) under a given treatment \(w\), \(v\) is the value of the favourable event, \(cc_w\) is the cost of the treatment triggered in case of a favourable event, and \(ic_w\) is the cost associated with the treatment whether or not the outcome is favourable. This method builds upon the ideas discussed in [36].

Probabilities of causation#

A cause is said to be necessary for an outcome if the outcome would not have occurred in the absence of the cause. A cause is said to be sufficient for an outcome if the outcome would have occurred in the presence of the cause. A cause is said to be necessary and sufficient if both of the above two conditions hold. [33] show that we can calculate bounds for the probability that a cause is of each of the above three types.

To understand how the bounds for the probabilities of causation are calculated, we need special notation to represent counterfactual quantities. Let \(y_t\) represent the proposition “\(y\) would occur if the treatment group was set to ‘treatment’”, \(y^{\prime}_c\) represent the proposition “\(y\) would not occur if the treatment group was set to ‘control’”, and similarly for the remaining two combinations of the (by assumption) binary outcome and treatment variables.

Then the probability that the treatment is sufficient for \(y\) to occur can be defined as

\[PS = P(y_t \mid c, y^{\prime})\]

This is the probability that the \(y\) would occur if the treatment was set to \(t\) when in fact the treatment was set to control and the outcome did not occur.

The probability that the treatment is necessary for \(y\) to occur can be defined as

\[PN = P(y^{\prime}_c \mid t, y)\]

This is the probability that \(y\) would not occur if the treatment was set to control, while in actuality both \(y\) occurs and the treatment takes place.

Finally, the probability that the treatment is both necessary and sufficient is defined as

\[PNS = P(y_t, y^{\prime}_c)\]

and states that \(y\) would occur if the treatment took place; and \(y\) would not occur if the treatment did not take place. PNS is related with PN and PS as follows:

\[PNS = P(t, y)PN + P(c, y^{\prime})PS\]

In bounding the above three quantities, we utilize observational data in addition to experimental data. The observational data is characterized in terms of the joint probabilities:

\[P_{TY} = {P(t, y), P(c, y), P(t, y^{\prime}), P(c, y^{\prime})}\]

Given this, [33] use the program developed in [8] to obtain sharp bounds of the above three quantities. The main idea in this program is to turn the bounding task into a linear programming problem (for a modern implementation of their approach see here).

Using the linear programming approach and given certain constraints together with observational data, [33] find that the sharp lower bound for PNS is given by

\[\max\{0, P(y_t) - P(y_c), P(y) - P(y_c), P(y_t) - P(y)\}\]

and the sharp upper bound is given by

\[\min\{P(y_t), P(y^{\prime}_c), P(t, y) + P(c, y^{\prime}), P(y_t) - P(y_c) + P(t, y^{\prime}) + P(c, y)\}\]

They use a similar routine to find the bounds for PS and PN. The get_pns_bounds() function calculates the bounds for each of the three probabilities of causation using the results in [33].

Selected traditional methods#

Relevant classes: NearestNeighborMatch and MatchOptimizer in causalml.match, the propensity models in causalml.propensity, IVRegressor in causalml.inference.iv for 2SLS, and TMLELearner in causalml.inference.meta.

The package supports selected traditional causal inference methods. These are usually used to conduct causal inference with observational (non-experimental) data. In these types of studies, the observed difference between the treatment and the control is in general not equal to the difference between “potential outcomes” \(\mathbb{E}[Y(1) - Y(0)]\). Thus, the methods below try to deal with this problem in different ways.

Matching#

The general idea in matching is to find treated and non-treated units that are as similar as possible in terms of their relevant characteristics. As such, matching methods can be seen as part of the family of causal inference approaches that try to mimic randomized controlled trials.

While there are a number of different ways to match treated and non-treated units, the most common method is to use the propensity score:

\[e_i(X_i) = P(W_i = 1 \mid X_i)\]

Treated and non-treated units are then matched in terms of \(e(X)\) using some criterion of distance, such as \(k:1\) nearest neighbours. Because matching is usually between the treated population and the control, this method estimates the average treatment effect on the treated (ATT):

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

See [30] for a discussion of the strengths and weaknesses of the different matching methods.

Inverse probability of treatment weighting#

The inverse probability of treatment weighting (IPTW) approach uses the propensity score \(e\) to weigh the treated and non-treated populations by the inverse of the probability of the actual treatment \(W\). For a binary treatment \(W \in \{1, 0\}\):

\[\frac{W}{e} + \frac{1 - W}{1 - e}\]

In this way, the IPTW approach can be seen as creating an artificial population in which the treated and non-treated units are similar in terms of their observed features \(X\).

One of the possible benefits of IPTW compared to matching is that less data may be discarded due to lack of overlap between treated and non-treated units. A known problem with the approach is that extreme propensity scores can generate highly variable estimators. Different methods have been proposed for trimming and normalizing the IPT weights ([15]). An overview of the IPTW approach can be found in [7].

2-Stage Least Squares (2SLS)#

One of the basic requirements for identifying the treatment effect of \(W\) on \(Y\) is that \(W\) is orthogonal to the potential outcome of \(Y\), conditional on the covariates \(X\). This may be violated if both \(W\) and \(Y\) are affected by an unobserved variable, the error term after removing the true effect of \(W\) from \(Y\), that is not in \(X\). In this case, the instrumental variables approach attempts to estimate the effect of \(W\) on \(Y\) with the help of a third variable \(Z\) that is correlated with \(W\) but is uncorrelated with the error term. In other words, the instrument \(Z\) is only related with \(Y\) through the directed path that goes through \(W\). If these conditions are satisfied, in the case without covariates, the effect of \(W\) on \(Y\) can be estimated using the sample analog of:

\[\frac{Cov(Y_i, Z_i)}{Cov(W_i, Z_i)}\]

The most common method for instrumental variables estimation is the two-stage least squares (2SLS). In this approach, the cause variable \(W\) is first regressed on the instrument \(Z\). Then, in the second stage, the outcome of interest \(Y\) is regressed on the predicted value from the first-stage model. Intuitively, the effect of \(W\) on \(Y\) is estimated by using only the proportion of variation in \(W\) due to variation in \(Z\). Specifically, assume that we have the linear model

\[Y = W \alpha + X \beta + u = \Xi \gamma + u\]

Here for convenience we let \(\Xi=[W, X]\) and \(\gamma=[\alpha', \beta']'\). Assume that we have instrumental variables \(Z\) whose number of columns is at least the number of columns of \(W\), let \(\Omega=[Z, X]\), 2SLS estimator is as follows

\[\hat{\gamma}_{2SLS} = \left[\Xi'\Omega (\Omega'\Omega)^{-1} \Omega' \Xi\right]^{-1}\left[\Xi'\Omega(\Omega'\Omega)^{-1}\Omega'Y\right].\]

See [3] for a detailed discussion of the method.

LATE#

In many situations the treatment \(W\) may depend on subject’s own choice and cannot be administered directly in an experimental setting. However one can randomly assign users into treatment/control groups so that users in the treatment group can be nudged to take the treatment. This is the case of noncompliance, where users may fail to comply with their assignment status, \(Z\), as to whether to take treatment or not. In general there are 3 types of users in this situation, classified by whether they take the treatment in response to their assignment:

  • Compliers Those who will take the treatment if and only if they are assigned to the treatment group.

  • Always-Takers Those who will take the treatment regardless of which group they are assigned to.

  • Never-Takers Those who will not take the treatment regardless of which group they are assigned to.

However one assumes that there is no Defier for identification purposes, i.e. those who will only take the treatment if they are assigned to the control group.

In this case one can measure the treatment effect of Compliers,

\[\hat{\tau}_{Complier}=\frac{E[Y|Z=1]-E[Y|Z=0]}{E[W|Z=1]-E[W|Z=0]}\]

This is Local Average Treatment Effect (LATE). The estimator is also equivalent to 2SLS if we take the assignment status, \(Z\), as an instrument.

Targeted maximum likelihood estimation (TMLE) for ATE#

Targeted maximum likelihood estimation (TMLE) [19] provides a doubly robust semiparametric method that “targets” directly on the average treatment effect with the aid from machine learning algorithms. Compared to other methods including outcome regression and inverse probability of treatment weighting, TMLE usually gives better performance especially when dealing with skewed treatment and outliers.

Given binary treatment \(W\), covariates \(X\), and outcome \(Y\), the TMLE for ATE is performed in the following steps

Step 1

Use cross fit to estimate the propensity score \(\hat{e}(x)\), the predicted outcome for treated \(\hat{m}_1(x)\), and predicted outcome for control \(\hat{m}_0(x)\) with machine learning.

Step 2

Scale \(Y\) into \(\tilde{Y}=\frac{Y-\min Y}{\max Y - \min Y}\) so that \(\tilde{Y} \in [0,1]\). Use the same scale function to transform \(\hat{m}_i(x)\) into \(\tilde{m}_i(x)\), \(i=0,1\). Clip the scaled functions so that their values stay in the unit interval.

Step 3

Let \(Q=\log(\tilde{m}_W(X)/(1-\tilde{m}_W(X)))\). Maximize the following pseudo log-likelihood function

\[\begin{split}\max_{h_0, h_1} -\frac{1}{N} \sum_i & \left[ \tilde{Y}_i \log \left(1+\exp\left(-Q_i-h_0 \frac{1-W}{1-\hat{e}(X_i)}-h_1 \frac{W}{\hat{e}(X_i)}\right) \right) \right. \\ &\quad\left.+(1-\tilde{Y}_i)\log\left(1+\exp\left(Q_i+h_0\frac{1-W}{1-\hat{e}(X_i)}+h_1\frac{W}{\hat{e}(X_i)}\right)\right)\right]\end{split}\]

Step 4

Let

\[\begin{split}\tilde{Q}_0 &= \frac{1}{1+\exp\left(-Q-h_0 \frac{1}{1-\hat{e}(X)}\right)},\\ \tilde{Q}_1 &= \frac{1}{1+\exp\left(-Q-h_1 \frac{1}{\hat{e}(X)}\right)}.\end{split}\]

The ATE estimate is the sample average of the differences of \(\tilde{Q}_1\) and \(\tilde{Q}_0\) after rescale to the original range.