Classification Trees for Propensity Score Estimation

Why Estimate Propensity Scores With a Tree?

A propensity score model estimates treatment assignment from pretreatment covariates. A classification tree offers a flexible way to do this: instead of specifying one smooth equation for the relationship between covariates and treatment, it repeatedly divides the covariate space into regions. Within each final region, the estimated score is the observed fraction treated. This can capture threshold effects and interactions without requiring them to be written in advance.

That flexibility has a cost. A tree produces a stepwise, piecewise-constant estimate, and its leaves can be small. A leaf with only a few units may have a treatment proportion near zero or one simply because of a handful of assignments. The tree is therefore a model for the treatment mechanism, not a guarantee of adequate overlap or covariate balance.

Definition: Let \(A_i\) equal 1 if unit \(i\) receives treatment and 0 otherwise, and let \(X_i\) denote that unit’s pretreatment covariates. The propensity score is \(e(X_i)=P(A_i=1\mid X_i)\). A classification tree partitions the covariate space into terminal leaves. For a leaf \(L\), the basic tree estimate is \(\hat e_L=n_{1L}/n_L\), where \(n_{1L}\) is the number treated and \(n_L\) is the total number of units in that leaf.

The tree must use covariates measured before treatment; including post-treatment variables would make the estimated assignment model inappropriate for the usual causal purpose. If the score is to support a causal analysis, causal identification also requires the relevant conditions, such as conditional exchangeability and positivity, for the target population. Those conditions are not made true by using a tree. A well-fitted treatment classifier alone does not establish them.

How Recursive Partitioning Chooses Splits

A classification tree starts with all \(n\) units in a root node. At a node, it considers candidate splits: for example, whether age is below a threshold, or whether a categorical covariate belongs to a particular group. Each split sends units into child nodes. The algorithm selects a split that improves a chosen measure of treatment-class purity, then considers further splits within the resulting nodes. This repeated process is recursive partitioning.

One common impurity measure is the Gini impurity. If the treated proportion in a node is \(q\), the untreated proportion is \(1-q\), and the Gini impurity is \(2q(1-q)\). It is zero when all units have the same treatment status and largest when the node is evenly divided between treated and untreated units. For a candidate split producing child nodes \(j\), the weighted child impurity is the sum of each child’s impurity multiplied by its share of the parent’s units. The split’s improvement is the parent impurity minus this weighted value.

G(q)=1-\{q^2+(1-q)^2\}=2q(1-q),\qquad \Delta G=G_{\mathrm{parent}}-\sum_j\frac{n_j}{n_{\mathrm{parent}}}G(q_j).

This criterion has a simple justification. Suppose a class label is predicted by randomly drawing a treatment category according to the observed proportions in the node. The probability that this random prediction is correct is \(q^2+(1-q)^2\), so the probability of error is \(1-\{q^2+(1-q)^2\}=G(q)\). Thus, a split with positive \(\Delta G\) reduces the weighted expected classification error under this rule. In practice, trees can use other criteria, such as entropy; the selected criterion and stopping rules should be specified rather than assumed to be universal.

Worked Example: Selecting a Root Split

Consider 16 hypothetical units, with eight treated and eight untreated. The covariates are a binary indicator \(X_1\) and a binary indicator \(X_2\). The observed counts are:

Covariate group	Treated	Untreated	Total
\(X_1=1,\ X_2=1\)	4	0	4
\(X_1=1,\ X_2=0\)	2	2	4
\(X_1=0,\ X_2=1\)	0	4	4
\(X_1=0,\ X_2=0\)	2	2	4

At the root, \(q=8/16=0.5\), so \(G_{\mathrm{root}}=2(0.5)(0.5)=0.5\). Splitting on \(X_1\) gives an \(X_1=1\) child with six treated and two untreated, and an \(X_1=0\) child with two treated and six untreated. Each child has treated proportion either \(6/8=0.75\) or \(2/8=0.25\), so each has Gini impurity \(2(0.75)(0.25)=0.375\). The weighted child impurity is \((8/16)(0.375)+(8/16)(0.375)=0.375\), giving \(\Delta G=0.5-0.375=0.125\).

Splitting on \(X_2\) at the root gives four treated and four untreated units in each child. Each child’s impurity is \(2(0.5)(0.5)=0.5\), so the weighted impurity remains 0.5 and the improvement is zero. By the Gini criterion, \(X_1\) is the better root split.

Leaf Probabilities and Stopping Rules

After selecting a split, the tree repeats the search within each child. It stops when it reaches a specified constraint or when no eligible split improves the criterion sufficiently. Common controls include a minimum number of units in a leaf, a minimum improvement required for a split, a maximum tree depth, and pruning. Pruning removes splits that add too little useful structure according to a chosen evaluation rule. These settings determine how complex the tree can become.

Once the tree is grown or pruned, each unit receives the treated fraction in its terminal leaf. In the worked table, splitting the \(X_1=1\) child on \(X_2\) gives one leaf with four treated units out of four and another with two treated out of four. Their estimates are \(4/4=1\) and \(2/4=0.5\), respectively. Splitting the \(X_1=0\) child gives estimates \(0/4=0\) and \(2/4=0.5\). These are empirical proportions, not known assignment probabilities. In particular, a score of exactly zero or one from a pure leaf should be treated as a warning about the data and tree size, not as proof that treatment was impossible or certain for units with those covariates.

Worked Example: How Leaf Size Affects a Score

Suppose a hypothetical terminal leaf contains four units: three treated and one untreated. Its estimated propensity score is \(\hat e_L=3/4=0.75\). If one of those assignments had instead been untreated, while the covariate values and leaf membership stayed fixed, the estimate would be \(2/4=0.50\). A change in just one of four assignments changes the estimated score by \(0.25\).

For comparison, a leaf with 40 units and 30 treated has the same estimated score, \(30/40=0.75\). Changing one assignment from treated to untreated changes it to \(29/40=0.725\), a difference of \(0.025\). These calculations do not account for the additional instability from re-selecting the tree splits, but they show why a small leaf gives a coarse estimate. Its possible proportions are restricted to increments of \(1/n_L\).

A rough binomial standard error for a fixed leaf proportion is \(\sqrt{\hat e_L(1-\hat e_L)/n_L}\). For the four-unit leaf it is \(\sqrt{0.75(0.25)/4}=0.2165\); for the 40-unit leaf it is \(\sqrt{0.75(0.25)/40}=0.0685\). These are illustrative within-leaf calculations, not complete uncertainty estimates for a tree chosen from the same data.

Diagnostics and Failure Modes

Inspect the fitted tree as a treatment model and as an input to the intended causal analysis. Start with terminal-leaf sizes and treated and untreated counts. Very small leaves, empty treatment groups, or estimated scores at zero and one indicate that the tree may be too finely partitioned for stable probability estimation. Also examine whether modest changes to the sample or tuning choices substantially alter the selected splits or leaf estimates.

Next, assess the estimated scores in context. Check the distribution of scores among treated and untreated units and look for regions without plausible counterparts. Then evaluate covariate balance using the planned analysis design; compare covariates, including important nonlinear features or interactions where appropriate. A tree can classify treatment well yet leave important covariate differences. Conversely, optimizing classification accuracy is not the same as optimizing balance.

A practical safeguard is to limit tree complexity using minimum leaf sizes and pruning, then evaluate the resulting model rather than relying on the training fit alone. Cross-validation can help compare tuning choices for predictive performance, but a choice that predicts treatment well is not automatically the choice that gives the best balance for a causal estimand. The criterion used to choose complexity should match the purpose, and balance must still be checked.

Worked Example: Interpreting a Pure Leaf

A hypothetical fitted tree produces a terminal leaf containing five units, all treated. The basic estimate is \(\hat e_L=5/5=1\). This means that all five observed units in this leaf were treated; it does not establish that the population treatment probability for every unit with these covariates is exactly one.

If a sixth, untreated unit with the same covariate pattern were observed and assigned to this leaf, the estimate would become \(5/6\approx0.8333\). The apparent certainty was therefore sensitive to very little data. The leaf also has no untreated observations to serve as direct within-leaf counterparts for treated units. Before using this score, one should consider whether the leaf is too small, whether a less complex tree is warranted, and whether the target analysis has adequate overlap. Neither smoothing a probability nor changing a split can create information about a missing treatment group.

How Trees Compare With Other Score Models

A logistic propensity model specifies a linear predictor for the log odds of treatment unless terms such as transformations and interactions are added. A generalized additive model allows flexible smooth functions of covariates while retaining a structured additive form. A tree instead makes sequential divisions and assigns a common score within each final leaf. It can represent certain interactions and threshold patterns naturally, but its estimates jump at split boundaries and may depend strongly on the observed sample and tuning choices.

These approaches are alternatives for estimating the assignment mechanism, not substitutes for checking the design. Logistic and additive models can be more stable when their structure is appropriate, while trees can offer an interpretable partition and adapt to non-smooth patterns. No model should be selected solely because it has high classification accuracy. The key questions are whether the score estimates are credible in the covariate regions used for analysis and whether the resulting design adequately balances the relevant pretreatment covariates.

Key takeaway: A classification tree estimates propensity scores as treatment proportions in recursively formed terminal leaves. The split criterion controls how leaves are formed, while stopping and pruning rules control their size and complexity. Small or pure leaves signal potentially unstable scores; inspect leaf counts, overlap, and covariate balance before using them.

Check Your Understanding

Answer each question, then compare your reasoning with the solutions below.

A node has 12 treated and 8 untreated units. Calculate its Gini impurity.
A split creates two equal-sized children with treatment proportions 0.8 and 0.4. Calculate the weighted child impurity and split improvement if the parent treatment proportion is 0.6.
What score does a leaf with 7 treated and 13 untreated units assign, and what is the leaf size?
Why is a tree leaf containing only treated units a diagnostic concern even if its estimated score is easy to calculate?

Solutions. (1) The treated proportion is \(12/20=0.6\), so \(G=2(0.6)(0.4)=0.48\). (2) The child impurities are \(2(0.8)(0.2)=0.32\) and \(2(0.4)(0.6)=0.48\). Their weighted value is \((0.32+0.48)/2=0.40\). The parent impurity is \(2(0.6)(0.4)=0.48\), so the improvement is \(0.48-0.40=0.08\). (3) The leaf score is \(7/(7+13)=0.35\), and its size is 20 units. (4) Its estimated score is one, but the estimate is based on a finite sample and there are no untreated observations in that leaf. It may be unstable and indicates a lack of within-leaf representation of both treatment groups; it does not prove the population probability is exactly one.