Gradient Boosting for Propensity Score Estimation

Why Boost Treatment-Assignment Models?

A propensity score model estimates how treatment assignment varies with measured pretreatment covariates. Flexible models can represent nonlinearities and interactions without requiring each one to be specified in advance. Gradient boosting builds such a model in stages: each new shallow tree is fitted to treatment-assignment errors left by the current model, then added as a controlled correction. The result is a single additive prediction function, not a collection of independently fitted trees.

The aim is not simply to classify treatment correctly. For propensity score use, the predictions should be useful estimates of treatment probabilities across the covariate patterns relevant to the analysis. A highly accurate classifier can still produce scores with poor overlap or fail to balance important covariates.

Definition: Let \(A_i\in\{0,1\}\) be the treatment indicator and \(X_i\) the vector of measured pretreatment covariates for unit \(i\). The propensity score is \(e(X_i)=P(A_i=1\mid X_i)\). A boosted propensity model constructs a function \(f_M(X_i)\) by adding \(M\) tree-based updates and reports \(\hat e(X_i)=1/(1+\exp[-f_M(X_i)])\).

The score is intended to estimate treatment assignment conditional on the measured covariates, not the outcome. Fit the treatment model using \(A\) and pretreatment \(X\), without using post-treatment variables or outcomes to choose treatment predictions. Causal interpretation additionally requires the relevant causal assumptions, including consistency, conditional exchangeability given the measured covariates, and positivity. Boosting does not establish those assumptions or repair unmeasured confounding.

The Boosting Update

The model works on the log-odds scale. Given a current function \(f_{m-1}\), the fitted probability for unit \(i\) is \(p_i^{(m-1)}=1/(1+\exp[-f_{m-1}(X_i)])\). The binary logistic loss, summed over the sample, is

L(f)=\sum_{i=1}^{n}\left[-A_i\log p_i-(1-A_i)\log(1-p_i)\right].

Differentiating an observation's loss with respect to its current log-odds gives \(p_i-A_i\). Thus the negative gradient, the direction in which a small update can reduce loss, is the residual \(r_i=A_i-p_i\). A positive residual means the current model assigns too little treatment probability to that unit; a negative residual means it assigns too much.

At iteration \(m\), fit a shallow regression tree \(h_m(X)\) to these residuals, usually by minimizing squared error. The tree divides covariate space into leaves and gives each unit in a leaf the same fitted correction. Add a fraction of that correction to the current log-odds:

f_m(X_i)=f_{m-1}(X_i)+\nu h_m(X_i),\qquad \hat e_m(X_i)=\frac{1}{1+\exp[-f_m(X_i)]},

where \(0<\nu\leq 1\) is the learning rate. In the simple gradient-residual version used here, a leaf's fitted value is the mean residual among the training observations in that leaf. Some software uses second-order or line-searched leaf updates instead; the exact update convention should be checked when interpreting software output.

Initialization: A common starting model is the constant log-odds of the observed treatment prevalence, \(f_0=\log\{\bar A/(1-\bar A)\}\), where \(\bar A=n^{-1}\sum_i A_i\). This gives every unit the initial fitted probability \(\bar A\). If the sample contains only treated units or only controls, this log-odds is not finite, and this initialization cannot be used as stated.

The learning rate and the number and depth of trees jointly determine model complexity. A smaller \(\nu\) takes more cautious steps; it commonly requires more iterations but can reduce the effect of any one tree. A deeper tree can form more leaves and capture more complex interactions, but may fit idiosyncrasies in the training sample. Shallow trees, often stumps or trees with a small number of splits, constrain each update so the overall function becomes flexible gradually.

Why the Residual-Fitting Step Can Reduce Loss

The gradient connection explains why residuals are used instead of repeatedly fitting treatment labels from scratch. Suppose the current fitted probabilities are strictly between zero and one, and a tree assigns a constant value \(c_\ell\) to each leaf \(\ell\). The derivative of the loss when the model is updated to \(f+t h\), evaluated at \(t=0\), is

\left.\frac{d}{dt}L(f+t h)\right|_{t=0} =\sum_i(p_i-A_i)h(X_i) =-\sum_{\ell}c_\ell\sum_{i\in\ell}r_i.

If each leaf value is the mean residual in that leaf, then \(c_\ell=\bar r_\ell\), so \(\sum_{i\in\ell}r_i=n_\ell\bar r_\ell\). Therefore,

\left.\frac{d}{dt}L(f+t h)\right|_{t=0} =-\sum_{\ell}n_\ell\bar r_\ell^{\,2}\leq 0.

Unless all leaf mean residuals are zero, this derivative is negative: a sufficiently small positive step in the fitted tree's direction locally reduces logistic loss. This is a local justification, not a guarantee that an arbitrarily large step or many unvalidated iterations will improve generalization.

Worked Calculations

Worked Example: The First Tree Correction

Consider eight hypothetical units. A binary pretreatment covariate \(X\) has values Low and High. The treatment indicators are \(1,1,1,0\) for Low and \(0,0,0,1\) for High. The treatment prevalence is \(\bar A=4/8=0.5\), so the initial log-odds are \(f_0=\log(0.5/0.5)=0\) and each initial fitted probability is \(p_i^{(0)}=0.5\).

The residual \(r_i=A_i-p_i^{(0)}\) is \(+0.5\) for a treated unit and \(-0.5\) for a control. In the Low group, three residuals are \(+0.5\) and one is \(-0.5\), giving a leaf mean of \((0.5+0.5+0.5-0.5)/4=0.25\). In the High group, one residual is \(+0.5\) and three are \(-0.5\), giving \((0.5-0.5-0.5-0.5)/4=-0.25\). A stump splitting on \(X\) therefore predicts \(0.25\) for Low and \(-0.25\) for High.

With learning rate \(\nu=0.8\), the updated log-odds are \(0+0.8(0.25)=0.20\) for Low and \(0+0.8(-0.25)=-0.20\) for High. The corresponding scores are \(1/(1+\exp[-0.20])\approx0.5498\) and \(1/(1+\exp[0.20])\approx0.4502\). The tree has moved the groups' predictions in the directions suggested by their residuals, without assigning probabilities of zero or one.

The initial logistic loss is \(8[-\log(0.5)]=5.545\). After this update, the Low group's loss is \(3[-\log(0.5498)]+[-\log(1-0.5498)]\), and the High group's loss is \(1[-\log(0.4502)]+3[-\log(1-0.4502)]\). Their sum is approximately \(5.185\), lower than the initial loss. This is an in-sample treatment-prediction calculation, not evidence by itself that causal assumptions hold or covariates are adequately balanced.

Worked Example: A Second Sequential Update

Continue the same hypothetical data and use \(\nu=0.8\) again. After the first update, each Low unit has fitted probability about \(0.5498\). The three treated residuals are \(1-0.5498=0.4502\), and the control residual is \(0-0.5498=-0.5498\). Their mean is \((3(0.4502)-0.5498)/4\approx0.2002\).

For High, the fitted probability is about \(0.4502\). Its one treated residual is \(1-0.4502=0.5498\), and its three control residuals are each \(-0.4502\). The mean is \((0.5498-3(0.4502))/4\approx-0.2002\). A tree with the same split therefore fits corrections of approximately \(+0.2002\) for Low and \(-0.2002\) for High.

The next log-odds are \(0.20+0.8(0.2002)\approx0.3601\) for Low and \(-0.20+0.8(-0.2002)\approx-0.3601\) for High. Applying the logistic transformation gives scores of about \(0.5891\) and \(0.4109\). The second tree is fitted to the remaining residuals, which are smaller in magnitude than at initialization. This illustrates the sequential nature of boosting: later corrections depend on the current model.

Worked Example: Tree Depth and Memorization

Suppose a hypothetical training sample has four units with distinct values of a continuous pretreatment covariate, and treatment indicators \(1,0,1,0\). The prevalence is \(0.5\), so the initial score for every unit is \(0.5\) and each residual is \(+0.5\) or \(-0.5\). A sufficiently deep tree can create a separate leaf for each training unit. Its fitted leaf values are then the individual residuals.

With \(\nu=1\), the first update assigns log-odds \(+0.5\) to each treated unit and \(-0.5\) to each control. Their fitted scores are \(1/(1+\exp[-0.5])\approx0.6225\) and \(1/(1+\exp[0.5])\approx0.3775\). If later trees keep isolating those same observations, they can continue pushing training scores toward their observed labels, even if the covariate pattern does not predict treatment well for new units. A shallow tree would be unable to make a separate correction for every observation in this way.

The calculation illustrates a risk, not a rule that every deep tree overfits. Whether a depth is too great depends on sample size, covariate structure, learning rate, iteration count, and the stability of performance on held-out data.

Choosing Complexity and Checking the Result

In practice, the fitted model is determined by at least three choices: tree depth (or number of leaves), learning rate, and the number of boosting iterations. Tune these choices using the cross-validation principles introduced in the earlier material on cross-validation, rather than selecting them by training accuracy alone. Treatment-prediction loss is a useful predictive criterion, but it does not replace checks relevant to propensity score use.

After fitting, examine the distribution of estimated scores by treatment group, the amount of overlap, and whether important measured covariates are balanced under the intended analysis. Also compare training and validation performance and check whether predictions are unusually close to zero or one. A model that separates the observed treatment labels very well may yield extreme estimated propensities and poor overlap. Conversely, modest predictive accuracy does not by itself mean that the model is unsuitable: the goal is not to maximize classification accuracy.

A practical diagnostic sequence is:

Inspect score distributions.
Compare fitted scores for treated and control units and identify regions with little or no overlap. Very extreme scores can signal weak empirical support for comparisons.

Check covariate balance.
Assess whether relevant pretreatment covariates are balanced using the planned weighting or matching procedure. A favorable treatment-prediction loss does not guarantee balance.

Look for overfitting.
Compare training and validation loss, and assess whether increasing depth or iteration count creates unstable or extreme scores without useful validation improvement.

Reconsider the specification.
If diagnostics are poor, consider a smaller learning rate, shallower trees, fewer iterations, or a different model. Do not treat a different tuning choice as a remedy for absent overlap or unmeasured confounding.

Gradient boosting differs from a random forest in how it combines trees. A random forest fits trees on perturbed samples or covariate sets and aggregates their predictions; boosting fits each successive tree to correct the current model's errors. A single classification tree creates a simpler partition, while logistic regression and generalized additive models impose different functional forms and may be easier to interpret. Boosting can capture complex nonlinearities and interactions with less manual specification, but its predictions depend on tuning and can be less transparent. None of these alternatives guarantees causal validity or adequate overlap.

Common Mistakes and Limits

Using outcomes to fit treatment scores. The treatment-assignment model should use pretreatment covariates, not the outcome or post-treatment information. Including such information can compromise the intended design.
Equating classification with propensity estimation. A high percentage of correctly classified treatment labels is not the same as calibrated probabilities, overlap, or covariate balance.
Assuming more trees always help. Additional iterations can improve training fit while worsening generalization or producing scores too extreme for stable comparisons.
Treating the learning rate as the only complexity control. A small learning rate does not prevent a very deep tree from making overly specific corrections; depth and iteration count matter too.
Claiming causal effects from a fitted score alone. Boosting estimates assignment probabilities from observed covariates. It cannot ensure that all confounders were measured or that treated and control units have adequate common support.

Key takeaway: Gradient boosting estimates propensity scores by repeatedly fitting shallow trees to the residuals \(A_i-\hat e_i\) and adding learning-rate-scaled corrections on the log-odds scale. Its flexibility is useful only when the fitted scores are also examined for overlap, balance, stability, and overfitting.

Check Your Understanding

Answer each question, then compare with the concise solutions below.

If the current fitted probability is \(0.7\) for a control unit, what is its residual, and what direction does that residual suggest for a tree correction?
A sample has treatment prevalence \(0.25\). What constant initial fitted probability and log-odds does the prevalence initialization give?
Why can a deeper tree overfit treatment assignments even when the learning rate is small?
Does lower training logistic loss guarantee adequate covariate balance or causal validity? Explain.
State one key difference between how gradient boosting and a random forest build their trees.

Solutions. (1) The residual is \(A-p=0-0.7=-0.7\), indicating that the current probability is too high for this observation; a fitted correction tending downward is appropriate. (2) The initial probability is \(0.25\), and the log-odds are \(\log(0.25/0.75)=\log(1/3)\approx-1.099\). (3) A deeper tree can form more specific leaves and fit idiosyncratic patterns; a small learning rate scales each update but does not prevent repeated, overly specific updates. (4) No. Training loss measures fit to observed treatment assignments; balance must be checked for the planned analysis, and causal validity also depends on assumptions such as exchangeability and positivity. (5) Boosting fits each next tree to errors or gradients of the current model; a random forest aggregates trees fitted separately on perturbed samples or covariate sets.