Method details and design decisions

Daijiro Kabata

This vignette documents the method implemented by psave() at the level of formulas, records the small number of places where the package deliberately deviates from the paper’s reference code (with justification), and situates psAve among related software. Throughout, “the paper” is Kabata, Stuart & Shintani (2024).

Notation

For subjects $i = 1, \dots, n$: covariates $X_i$ (a $p$-column numeric matrix after full dummy expansion of factors), binary treatment $A_i \in \{0, 1\}$, and outcome $Y_i$. The propensity score is $e(X_i) = \Pr(A_i = 1 \mid X_i)$. The prognostic score is $g(0, X_i)$, the predicted outcome under the untreated condition, estimated from untreated units only (Hansen 2008).

psave() fits $M$ candidate propensity score models $\hat e_1, \dots, \hat e_M$ (on all $n$ units) and $K$ candidate prognostic models $\hat g_1, \dots, \hat g_K$ (on the $n_0$ untreated units, predicted for all $n$), then selects convex mixing weights $\gamma$ (prognostic) and $\lambda$ (propensity) on a simplex grid.

The simplex grid and tie-breaking

Both $\gamma$ and $\lambda$ range over the probability simplex discretized with increment step (default 0.05). The exported function simplex_grid(M, step) enumerates the grid in integer arithmetic: with $S = 1/\text{step}$ steps, it lists every integer composition $(c_1, \dots, c_M)$ with $\sum_m c_m = S$ and $c_m \ge 0$, returned as $c/S$. The number of grid points is exactly $\binom{S + M - 1}{M - 1}$:

library(psAve)

nrow(simplex_grid(2, step = 0.05))   # choose(21, 1) = 21
#> [1] 21
nrow(simplex_grid(3, step = 0.05))   # choose(22, 2) = 231
#> [1] 231
nrow(simplex_grid(4, step = 0.05))   # choose(23, 3) = 1771
#> [1] 1771

The enumeration order is a documented part of the method because it defines the tie-breaking rule. The grid is generated by recursive descent, with $c_1$ running from $S$ down to 0, then $c_2$ descending over the remainder, and so on:

head(simplex_grid(3, step = 0.25))
#>      [,1] [,2] [,3]
#> [1,] 1.00 0.00 0.00
#> [2,] 0.75 0.25 0.00
#> [3,] 0.75 0.00 0.25
#> [4,] 0.50 0.50 0.00
#> [5,] 0.50 0.25 0.25
#> [6,] 0.50 0.00 0.50
tail(simplex_grid(3, step = 0.25), 3)
#>       [,1] [,2] [,3]
#> [13,]    0 0.50 0.50
#> [14,]    0 0.25 0.75
#> [15,]    0 0.00 1.00

The first row puts all weight on the first candidate; the last row puts all weight on the last. Ties in any grid search are resolved by taking the first row attaining the minimum, within a 1e-9 relative tolerance of the minimum, so ties favor learners listed earlier in ps.methods / prog.methods. This makes “first minimum” a reproducible formula rather than an accident of grid construction. The tolerance is deliberate: the criterion values are computed with floating-point matrix algebra whose lowest-order bits can differ across BLAS implementations, so an exact bitwise which.min() would not be reproducible across machines, whereas the tolerant first-minimum rule is. (The reference code used expand.grid() order with which.min(); the package’s order differs but is equally deterministic and is fixed by simplex_grid() by construction.)

The grid size grows quickly with $M$: with step = 0.05, $M = 6$ already gives 53,130 points. psave() warns when the grid exceeds $10^5$ rows and suggests a coarser step; in that case the full criterion path is also not stored.

Step 1: Candidate models

Candidate propensity score models are fit on all $n$ units and predict $\Pr(A = 1 \mid X)$ in-sample. The built-in engines are pinned to fixed hyperparameters (chosen to mirror the SuperLearner wrapper defaults used in the paper, so the direct and SL.* routes agree); they can be overridden per learner via control =, and the resolved values are stored in fit$info$learners. Any "SL.*" label is passed through to SuperLearner verbatim, which is the exact-replication path for the paper. Alternatively, a user-supplied matrix of candidate scores can be given via ps.matrix.

Each candidate column is clipped to clip = c(0.01, 0.99) before averaging (as in the paper). No re-clipping is applied after averaging: a convex combination of values inside the clipping interval cannot leave it.

Candidate prognostic models use the same learner menu in regression mode (or probability mode for family = binomial()), fit only on untreated units and predicted for all $n$ units.

Step 2: $\gamma$ — the model-averaged prognostic score

$\gamma$ minimizes the unweighted mean squared prediction error among untreated units:

$$\hat\gamma = \arg\min_{\gamma \in \mathcal{S}_K} \; \frac{1}{n_0} \sum_{i : A_i = 0} \Bigl( Y_i - \sum_{k=1}^{K} \gamma_k \, \hat g_k(X_i) \Bigr)^2,$$

where $\mathcal{S}_K$ is the simplex grid. The model-averaged prognostic score is $\bar g(X) = \sum_k \hat\gamma_k \hat g_k(X)$. For family = binomial() the same formula is the Brier score; note the paper’s simulations validated continuous outcomes (see Limitations). With a single prognostic candidate, $\gamma = 1$ and the grid search is skipped.

Step 3: $\lambda$ — the model-averaged propensity score

For each grid point $\lambda$, the averaged score is $\bar e_\lambda(X_i) = \sum_m \lambda_m \hat e_m(X_i)$, and the implied inverse-probability weights are

$$\text{ATT:}\quad W_i = \begin{cases} 1 & A_i = 1 \\ \dfrac{\bar e_\lambda(X_i)}{1 - \bar e_\lambda(X_i)} & A_i = 0 \end{cases} \qquad\qquad \text{ATE:}\quad W_i = \begin{cases} \dfrac{1}{\bar e_\lambda(X_i)} & A_i = 1 \\ \dfrac{1}{1 - \bar e_\lambda(X_i)} & A_i = 0. \end{cases}$$

$\lambda$ is selected to minimize one of four criteria, evaluated at every grid point.

criterion = "logloss"

The negative Bernoulli log-likelihood of treatment (prediction accuracy, the Xie et al. 2019 lineage):

$$\mathrm{LL}(\lambda) = -\frac{1}{n} \sum_{i=1}^{n} \Bigl[ A_i \log \bar e_\lambda(X_i) + (1 - A_i) \log\bigl(1 - \bar e_\lambda(X_i)\bigr) \Bigr],$$

finite by clipping. This is the only criterion that does not require an outcome.

criterion = "smd"

The mean over covariate columns $j$ of the weighted absolute standardized mean difference:

$$\mathrm{wASMD}_j(\lambda) = \frac{\Bigl| \dfrac{\sum_i A_i W_i X_{ij}}{\sum_i A_i W_i} - \dfrac{\sum_i (1 - A_i) W_i X_{ij}}{\sum_i (1 - A_i) W_i} \Bigr|}{s_j}, \qquad s_j = \mathrm{sd}(X_j \mid A = 1),$$

where $s_j$ is the plain (unweighted, $n - 1$ denominator) sample standard deviation in the treated group — for both the ATT and the ATE, per the paper’s supplement (not the pooled SD). Under the ATT, $W_i = 1$ for treated units, so the first term is the raw treated mean.

criterion = "ks"

The mean over covariates of the weighted Kolmogorov–Smirnov statistic, using the proper weighted empirical CDF in each arm:

$$F^w_a(x) = \frac{\sum_{i : A_i = a} W_i \, \mathbf{1}(X_{ij} \le x)}{\sum_{i : A_i = a} W_i}, \qquad \mathrm{wKS}_j(\lambda) = \max_{x} \bigl| F^w_1(x) - F^w_0(x) \bigr|$$

over the observed values of $X_j$. For a binary covariate this reduces to the absolute difference in weighted proportions.

criterion = "prog" (the default; the paper’s “Prog (Ave)”)

The wASMD formula above applied to a single column: the model-averaged prognostic score $\bar g$ (when prog.target = "average", the headline recommendation) or a single candidate $\hat g_k$ (when prog.target names a learner — the paper’s “Prog ($g_k$)” variants). The denominator is $\mathrm{sd}(\bar g \mid A = 1)$, again unweighted and for both estimands.

Vertex mode

average = FALSE runs the identical machinery on the $M$ vertices of the simplex only, i.e., it selects the single best candidate propensity score by the chosen criterion (the “best single learner” variants computed alongside Table 1 of the paper’s supplement).

Criterion vs. display conventions for the SMD denominator

One subtlety deserves explicit documentation. The paper (and its reference implementation) standardizes every covariate — including binary ones — by the plain sample SD in the treated group. cobalt, by contrast, standardizes binary covariates by $\sqrt{p_1 (1 - p_1)}$ (the population formula) when it standardizes them at all, and its bal.tab() display leaves binary covariates as raw differences in proportions by default.

psAve resolves this as follows:

• The selection criteria are paper-faithful. The smd and prog criteria use uniform sample-SD standardization for all columns (implemented by passing bin.vars = FALSE for every column to cobalt::col_w_smd()), so the selected $\lambda$ and the reported criterion.value match the published method exactly. For the ks criterion the two conventions coincide for binary variables, so auto-detection is harmless there.
• The display follows cobalt. The balance field of the fitted object, summary(), plot(type = "balance"), and bal.tab() use cobalt’s native conventions, so the numbers you see agree with what cobalt::bal.tab() reports for the same matched or weighted analysis elsewhere in your workflow.

For criterion = "prog" the denominator is a single positive constant across the whole $\lambda$ grid, so the argmin — and therefore the selected score — is invariant to this choice; only the reported criterion value depends on it.

Five documented fixes relative to the paper’s reference code

The published algorithm is implemented exactly; the published reference code contained implementation artifacts that this package fixes. These are documented so that anyone comparing psAve output against the original scripts understands where and why small discrepancies arise.

1. Integer simplex grid. The reference code built the grid with expand.grid() over $\{0, 0.05, \dots, 1\}^M$ and kept rows with rowSums(gr) == 1 — an exact floating-point equality test. With $M = 4$ and step 0.05, this silently dropped 187 of the 1,771 valid simplex points (about 10.6%), including potentially optimal mixtures. simplex_grid() enumerates integer compositions, so every valid grid point is present by construction.
2. Proper weighted-eCDF KS. The reference Fks computed ks.test() on the covariate values multiplied by the weights within each arm, which is not the paper’s weighted-eCDF definition (multiplying values by weights changes the distribution’s support, not its weighting). The package uses the paper’s definition via cobalt::col_w_ks().
3. binomial() family for binary responses. The reference code fit binary models inside SuperLearner with gaussian(link = "logit"). The package uses binomial() throughout for treatment models (and for binary-outcome prognostic models).
4. No train/test-inconsistent scale(). The reference code applied scale() separately to fitting and prediction sets, standardizing them by different means and SDs. The package passes raw covariates to all engines (the pinned engines are either scale-invariant or fit and predict on identical rows).
5. Strict complete-case handling. The reference Fasmd applied na.omit() to a covariate vector while indexing full-length treatment and weight vectors, silently misaligning rows in the presence of missing data. psave() refuses incomplete data outright: any NA in the treatment, the covariates, or the outcome (when used) is an error, never a silent drop.

A sixth, minor note: candidate scores are clipped before averaging and never re-clipped afterward (see Step 1); the reference code’s clipping constants 0.01/0.99 are exposed as the clip argument.

Relation to other software

No other R package implements propensity score model averaging over a mixing-weight simplex (Xie et al. 2019 published no software). The closest relatives, and how psAve differs:

• WeightIt::method_super with SL.method = "method.balance" (the Balance Super Learner; Pirracchio & Carone 2018). Also builds a convex combination of candidate propensity models with weights chosen for balance rather than prediction. Differences: its criterion is covariate balance (a user-chosen cobalt-computed statistic), whereas psAve’s default criterion is prognostic-score balance, which targets the covariate directions that matter for outcome bias; the Balance Super Learner optimizes over the continuous simplex via SuperLearner’s machinery with cross-validated predictions, whereas psAve uses the paper’s fixed grid with in-sample candidate predictions and a documented tie-break; and method_super lives inside a weighting workflow, whereas psAve returns a bare score vector equally usable for matching (MatchIt::matchit(distance = )) and weighting.
• twang (McCaffrey, Ridgeway & Morral 2004) tunes a single model class — the number of boosting iterations of a GBM — against covariate balance. It selects within one learner; psAve averages across heterogeneous learners.
• PSweight provides a broad estimand/weighting framework (including overlap weights) with a single user-specified propensity model per analysis; it does model use, not model selection or averaging.
• cobalt is not a competitor but the substrate: the balance statistics that define psAve’s criteria are computed by cobalt::col_w_smd() and cobalt::col_w_ks(), and bal.tab() works on psave objects directly.

Limitations

Honesty about scope, matching the paper’s evidence base:

• Binary treatment only. The method is defined, and was evaluated, for a two-arm comparison. Multi-category and continuous treatments are out of scope.
• ATT and ATE only. The supplement’s criterion formulas are validated for these two estimands; other estimands (ATO, ATM, …) are refused with an error rather than extrapolated.
• Continuous outcomes in the validating simulations. The paper’s simulations used continuous $Y$. family = binomial() is permitted — the $\gamma$ criterion is then the Brier score, an unchanged formula — but its finite-sample performance was not studied in the paper.
• In-sample candidate predictions. As in the paper’s worked example, candidate models predict on their own training data (no cross-fitting in v1). Flexible learners can therefore produce optimistically extreme candidate scores; clipping bounds the damage, the diagnostics table shows every candidate under every criterion, and plot(type = "distribution") makes overfit candidates visible. Balance-targeted selection (unlike log-loss selection) is not rewarded for overfit treatment prediction, which is part of the method’s rationale.
• Design-based standard errors downstream do not reflect score estimation. See the weighting vignette; bootstrap the full pipeline if that matters for your application.

References

Hansen, B. B. (2008). The prognostic analogue of the propensity score. Biometrika, 95(2), 481–488. doi:10.1093/biomet/asn004

Kabata, D., Stuart, E. A., & Shintani, A. (2024). Prognostic score-based model averaging approach for propensity score estimation. BMC Medical Research Methodology, 24, 228. doi:10.1186/s12874-024-02350-y

McCaffrey, D. F., Ridgeway, G., & Morral, A. R. (2004). Propensity score estimation with boosted regression for evaluating causal effects in observational studies. Psychological Methods, 9(4), 403–425. doi:10.1037/1082-989X.9.4.403

Pirracchio, R., & Carone, M. (2018). The Balance Super Learner: A robust adaptation of the Super Learner to improve estimation of the average treatment effect in the treated based on propensity score matching. Statistical Methods in Medical Research, 27(8), 2504–2518. doi:10.1177/0962280216682055

Stuart, E. A., Lee, B. K., & Leacy, F. P. (2013). Prognostic score-based balance measures can be a useful diagnostic for propensity score methods in comparative effectiveness research. Journal of Clinical Epidemiology, 66(8 Suppl), S84–S90. doi:10.1016/j.jclinepi.2013.01.013

Xie, Y., Zhu, Y., Cotton, C. A., & Wu, P. (2019). A model averaging approach for estimating propensity scores by optimizing balance. Statistical Methods in Medical Research, 28(1), 84–101. doi:10.1177/0962280217715487