Super Learner Ensembles for Treatment Probabilities

Why Stack Treatment-Probability Models?

Inverse-probability weighting depends on estimating the probability of the treatment actually received, conditional on the covariate and treatment history available at that decision. A single model can miss important structure: a linear predictor may be too restrictive, while a highly adaptive learner may be unstable in sparse regions. Super Learner stacking offers a principled way to combine a prespecified collection of candidate algorithms, using their out-of-fold predictions to choose combination weights.

The ensemble is a prediction procedure, not a guarantee of good overlap or causal identification. Its weights are chosen to predict treatment assignments under a specified loss. The causal interpretation of the resulting inverse-probability weights still requires the relevant identification assumptions, including consistency, sequential exchangeability, and positivity.

For a binary treatment, a candidate predicts \(p_i^{(m)}=P(A_i=1\mid H_i)\); its probability for the observed treatment is \(q_i^{(m)}=(p_i^{(m)})^{A_i}(1-p_i^{(m)})^{1-A_i}\). For a treatment with \(J\) categories, write \(p_{ij}^{(m)}=P(A_i=j\mid H_i)\), where each candidate's probabilities sum to one across categories. The observed-treatment probability is \(q_i^{(m)}=p_{i,A_i}^{(m)}\).

Cross-Validated Stacking

Partition the \(n\) observations into \(K\) folds. For each fold \(k\), fit every candidate algorithm using observations outside that fold, then predict treatment probabilities for observations in fold \(k\). Thus, \(p_{ij}^{(m,-k)}\) is candidate \(m\)'s prediction for category \(j\) on an observation not used to fit that fold-specific model. For binary treatment, the prediction for the observed treatment is \(q_i^{(m)}\) as defined above, using the corresponding fold-specific prediction.

Choose ensemble weights \(\alpha=(\alpha_1,\ldots,\alpha_M)\) from the probability simplex \(\mathcal{S}_M=\{\alpha_m\geq0:\sum_{m=1}^M\alpha_m=1\}\). The stacked prediction for category \(j\) is \(\hat p_{ij}=\sum_m\alpha_m p_{ij}^{(m,-k(i))}\). The usual selection criterion is the mean cross-validated negative log likelihood, also called multiclass log loss:

The selected weights minimize this criterion over \(\mathcal{S}_M\). After selection, the candidate algorithms are usually refitted on the full training data, and their full-data predictions are combined using the selected weights. For a new history \(h\), this gives \(\hat g(a\mid h)=\sum_m\hat\alpha_m\hat g_m(a\mid h)\). Because each component is a probability distribution and the weights are nonnegative and sum to one, the ensemble is also a probability distribution.

Proof. For each observation \(i\), let \(c_{im}=p_{i,A_i}^{(m,-k(i))}\), so \(c_{im}\geq\epsilon\). The term in the objective is \(-\log(\sum_m\alpha_m c_{im})\). The argument of the logarithm is affine in \(\alpha\), is at least \(\epsilon\) on the simplex, and is at most one because it is a mixture of probabilities. The function \(-\log(x)\) is continuous and convex for \(x>0\). Its composition with a positive affine function is therefore continuous and convex in \(\alpha\); summing these terms preserves both properties. The simplex is compact and nonempty, so a continuous function on it attains its minimum. Convexity means any local minimum is global, although the minimizing weights need not be unique. \(\square\)

In practice, candidate predictions can be zero or numerically extremely small. The stated lower-bound condition is a convenient sufficient condition for the result; clipping probabilities to a prespecified small positive floor can enforce it. Clipping changes predictions and should be reported, not treated as an invisible numerical fix.

Why Use Log Loss?

Log loss evaluates the probability assigned to the treatment that actually occurred. It penalizes confident errors sharply: assigning probability \(0.01\) to the observed treatment incurs loss \(-\log(0.01)\), much larger than assigning probability \(0.4\). Under a correctly specified conditional probability target, its population optimum is the true treatment mechanism.

Proof. Subtract the true-probability loss from the proposed loss and collect terms: \(\sum_{j:g_j>0}g_j\log(g_j/q_j)\). To show this is nonnegative, use \(\log x\leq x-1\) for \(x>0\), with equality only at \(x=1\). Applying it to \(x=q_j/g_j\) gives \[ \sum_{j:g_j>0}g_j\log\left(\frac{g_j}{q_j}\right) \geq \sum_{j:g_j>0}(g_j-q_j) = 1-\sum_{j:g_j>0}q_j \geq 0. \] Equality requires \(q_j=g_j\) on the support of \(g\), with no probability assigned outside that support. This proves the claim. \(\square\)

This result explains why log loss is a natural stacking criterion for probabilities, rather than a loss on hard treatment classifications. It does not imply that minimizing empirical cross-validated loss identifies the true mechanism in a finite sample, nor that low prediction loss ensures adequate positivity.

From Ensemble Probabilities to Weights

For a single discrete treatment decision, the unstabilized weight for observation \(i\) is the reciprocal of the fitted probability of the treatment actually received. In a longitudinal setting, the denominator contribution at time \(t\) is \(\hat g_t(A_{it}\mid H_{it})\), and the unstabilized weight through time \(T\) is the product of the reciprocals of these contributions. A stabilized weight replaces each denominator factor with a specified numerator probability, often based on a reduced history:

The stacking procedure supplies the denominator predictions when it is used to estimate the treatment mechanism; the numerator is a separate modeling or marginal-probability choice. For a single decision, if the observed treatment is \(A_i=a\), calculate \(\hat g(a\mid H_i)\) first, then take its reciprocal (or divide the selected numerator probability by it). For multivalued treatment, use the probability for the observed category, not the probability of some other category.

Diagnostics, Alternatives, and Limitations

Inspect both the prediction procedure and the weights it creates. Useful checks include the distribution of predicted probabilities for observed treatments, the minimum and lower quantiles of those probabilities, the distribution of resulting weights, and effective sample size. For weights \(W_i\), a common effective-sample-size summary is \((\sum_i W_i)^2/\sum_i W_i^2\). A small value relative to \(n\) indicates that a few observations dominate weighted summaries. For longitudinal weights, inspect cumulative products as well as time-specific factors, since moderate factors can compound.

Compare cross-validated log loss across candidate learners and the ensemble, and check that fold construction respects the independent sampling units. Repeated observations from one person should not be split across training and validation folds as if independent. A large gap between training and out-of-fold performance suggests overfitting. Also check for near-zero probabilities and unstable weights; low average log loss can coexist with poor predictions in a clinically important subgroup. Weight truncation may reduce variance but changes the estimator and can introduce bias.

Stacking differs from selecting the single candidate with the smallest cross-validated loss: it optimizes over mixtures, which can outperform every candidate, as the first worked example shows. It also differs from an unweighted average, which fixes weights in advance rather than learning them from out-of-fold predictions. The advantage depends on the candidate library and the validation design. A misspecified or narrow library limits what the ensemble can achieve; a very large library can make weight selection noisy. The selected convex weights do not prove that the true treatment mechanism is in the library, and they do not correct structural positivity violations.