Gradient Boosting for Treatment Mechanism Models

Why Boost a Treatment Mechanism?

Inverse probability weighting depends on the probability of receiving the treatment actually observed, conditional on the covariates that guide treatment. A linear logistic model can miss nonlinearities and interactions in those covariates. Gradient boosting offers a flexible alternative: it builds a predictor through successive, usually shallow, decision trees, with each tree improving the current fit under a chosen loss function.

Flexibility is not itself a guarantee of good weights. A boosted classifier can predict treatment labels accurately while producing probabilities that are too close to zero or one, or while failing to represent the treatment mechanism in relevant parts of the covariate distribution. The modeling target is therefore not merely correct classification. It is a useful estimate of the conditional treatment probabilities, together with adequate overlap for the weighting comparison.

For a causal interpretation, flexible estimation does not replace the causal assumptions. In particular, treatment must be conditionally exchangeable given the included pretreatment covariates, treatment must be consistently defined, and positivity must hold: \(0<P(A=1\mid X)<1\) for covariate values in the target population. The model should not include variables measured after the treatment decision as if they were pretreatment confounders. These conditions concern the causal design and target population, not the choice of a boosting algorithm.

Logistic Gradient Boosting for Binary Treatment

Encode treatment as \(A_i=1\) for treated and \(A_i=0\) for untreated. A common model represents the log odds of treatment as a function \(F(X_i)\), with predicted probability \(p_i=\operatorname{expit}\{F(X_i)\}=1/\{1+\exp[-F(X_i)]\}\). Logistic boosting minimizes the Bernoulli negative log-likelihood, also called binomial deviance up to a constant or scale convention.

A standard additive boosting fit begins with a constant log-odds value, often the logit of the sample treatment prevalence, and then adds tree predictions in stages. At stage \(m\), the tree is fit to information about the current loss gradient. A learning rate \(\nu\), usually between zero and one, shrinks each tree's contribution; the number of stages and tree complexity control how much structure the fit can learn.

To see the derivatives, note that \(\log p_i=F_i-\log(1+\exp F_i)\) and \(\log(1-p_i)=-\log(1+\exp F_i)\). Substitution gives \(\ell_{\mathrm{loss}}=\log(1+\exp F_i)-A_iF_i\), so differentiation yields \(p_i-A_i\); differentiating again gives \(p_i(1-p_i)\). A tree can be fitted to the negative gradients \(A_i-p_i\). In a Newton-style implementation, a terminal leaf \(R\) may receive the approximate update

This leaf update follows by taking a second-order Taylor approximation to the total loss in that leaf. The approximate change for a common increment \(\gamma\) is \(-G_R\gamma+\tfrac12 H_R\gamma^2\), where \(G_R=\sum_{i\in R}(A_i-p_i)\) and \(H_R=\sum_{i\in R}p_i(1-p_i)\). Differentiating this quadratic and setting its derivative to zero gives \(\gamma=G_R/H_R\), provided \(H_R>0\). Implementations may instead use gradient-only updates or regularize the denominator, so the exact leaf calculation depends on the specified boosting procedure.

From Predicted Probabilities to Weights

After choosing a fitted boosting model, predict probabilities for each person's observed pretreatment covariates. For binary treatment, the fitted probability of the treatment actually received is \(\hat p_{A_i}(X_i)=\hat e(X_i)\) if \(A_i=1\), and \(1-\hat e(X_i)\) if \(A_i=0\). The unstabilized weight is the reciprocal of that probability. For a stabilized weight, multiply by the marginal probability of the person's observed treatment, estimated in a simple baseline setting by the sample treatment proportion for treated people and its complement for untreated people.

where \(\bar A=n^{-1}\sum_{i=1}^n A_i\). These formulas describe a single binary treatment decision. For multiple treatment categories, the denominator is the fitted probability of the observed category; a multinomial classifier uses a probability for each category. The choice of numerator is tied to the target population and estimand, not a device for repairing a misspecified denominator.

Why Correct Weights Balance Covariates

The weighting identity explains the role of the denominator. Let \(e_a(X)=P(A=a\mid X)\), and suppose \(e_a(X)>0\) on the target population's covariate support. For any integrable function \(g(X)\), the expectation of \(I(A=a)g(X)/e_a(X)\) equals the marginal expectation of \(g(X)\). This statement concerns the true probabilities; fitted weights approximate it only to the extent that the fitted probabilities are adequate.

For the proof, condition on \(X\). Because \(g(X)\) is fixed given \(X\) and \(E[I(A=a)\mid X]=e_a(X)\), the conditional expectation of the weighted term is \(g(X)e_a(X)/e_a(X)=g(X)\). Taking expectations over \(X\) gives the identity. Positivity makes the ratio defined; integrability ensures the expectations exist. In a finite sample, estimated weights need not balance every covariate exactly, so balance diagnostics remain necessary. \(\square\)

Choosing Complexity and Checking Failure Modes

Tree depth, learning rate, and number of boosting rounds jointly control flexibility. A shallow tree captures limited interactions; deeper trees can represent more complex patterns but may fit noise. A smaller learning rate usually requires more rounds to reach a similar fit. Select complexity using a prespecified validation strategy and inspect fitted probabilities and weights in the analysis sample. Very low training loss is not evidence by itself that the weighting model is suitable.

Assess overlap by examining the distribution of \(\hat e(X)\) in each treatment group and the distribution of the observed-treatment probabilities. Also inspect weight summaries, the largest weights, and covariate balance before and after weighting. For example, an effective sample size summary is

A small effective sample size relative to \(n\) indicates that a few observations carry much of the weighted analysis. It is a warning, not a universal pass/fail threshold. Poor overlap is a property of the treatment assignment and covariate support; a more flexible classifier cannot create information where treatment probabilities are near zero. Truncating weights may reduce variance but changes the weighting procedure and can introduce bias; it should be reported and justified rather than treated as a substitute for adequate positivity.

Gradient boosting differs from a conventional logistic model chiefly in how it represents the covariate function: successive trees can capture nonlinearities and interactions without specifying them in advance. Random forests also use trees, but aggregate many tree predictions rather than building a stagewise sequence directed by the current loss gradient. Neither approach guarantees calibrated probabilities, covariate balance, or causal validity. A later ensemble strategy may combine candidate probability models, but it does not remove the need to check overlap, weights, and balance.

Common Pitfalls

• Using predicted classes instead of probabilities. A label such as “treated” or “untreated” cannot be inverted to form a propensity weight. Use the predicted probability for the observed treatment.
• Using the probability of the wrong category. For an untreated person, the binary denominator is \(1-\hat e(X)\), not \(\hat e(X)\).
• Equating classification accuracy with causal adequacy. Accuracy at a threshold does not reveal whether probabilities are suitable for weighting or whether weighted covariates are balanced.
• Ignoring extreme weights. A very small fitted probability for an observed treatment produces a large weight and can make the estimate unstable.
• Including post-treatment covariates in the treatment model. Predictors must reflect information available before the treatment decision for the intended causal adjustment.