Tuning Lasso Propensity Scores for Causal Balance

Why the Prediction-Optimal Penalty May Not Be Causally Useful

A lasso propensity model uses a penalty to trade fit against complexity. Cross-validation can choose the penalty that predicts treatment assignment well in held-out observations. But the causal aim of a propensity score is different: to support a comparison of treated and untreated units with similar measured pre-treatment covariates. A penalty with the smallest prediction loss is therefore not automatically the penalty that produces the most useful adjusted comparison.

This tutorial concerns choosing among lasso penalties using both prediction and causal diagnostics. The covariate set is assumed to have been defined in advance; deciding which variables belong in it is a separate task. The treatment is \(A\), where \(A=1\) denotes treatment and \(A=0\) denotes control, and \(X\) is the vector of measured pre-treatment covariates. The propensity score is \(e(X)=P(A=1\mid X)\). A fitted lasso model gives an estimated score \(\hat e_i\) for each unit.

The usual assumptions that give propensity-score adjustment a causal interpretation remain necessary: consistency of the observed outcome with the potential outcome under the received treatment; conditional exchangeability, \(Y^a\perp A\mid X\) for each treatment level \(a\); and positivity, \(0<P(A=1\mid X)<1\) in the target population. Covariates used for adjustment must precede treatment. Good measured balance is a diagnostic, not proof of exchangeability, correct measurement, or absence of unmeasured confounding.

What Cross-Validation Optimizes

For a logistic lasso, the fitted propensity model estimates treatment probabilities from \(X\), with a penalty that shrinks coefficients and can set some to zero. Larger \(\lambda\) means stronger shrinkage. In \(K\)-fold cross-validation, the model is fit on \(K-1\) folds and its treatment predictions are evaluated on the held-out fold; the process is repeated so every fold is held out. A common criterion is average held-out log loss:

Here \(\hat e_i^{(-k(i))}\) is the prediction for unit \(i\) from a model not trained on that unit's fold. Smaller loss indicates better out-of-fold prediction according to this scoring rule. It does not directly measure whether treated and control covariate distributions will be comparable after weighting or matching.

The distinction can be seen from the weighting identity. If \(e(X)\) is the true propensity score and positivity holds, weighting treated units by \(1/e(X)\) and controls by \(1/\{1-e(X)\}\) gives the same covariate distribution in expectation, up to a common normalizing constant. For any integrable function \(g(X)\),

The identity uses \(E(A\mid X)=e(X)\). It explains why a correct score can balance covariates in the population under inverse-probability weighting. In a finite sample, a penalized estimate need not equal the true score, and predictive accuracy does not guarantee small empirical imbalance. Conversely, a modest sacrifice in prediction loss can sometimes improve the covariate balance and overlap relevant to the planned causal comparison.

Measuring Balance and Checking Overlap

For a continuous covariate \(X_j\), a signed standardized mean difference (SMD) compares its adjusted treated and control means, scaled by a fixed reference standard deviation. One common reference is the unweighted pooled within-group standard deviation before adjustment:

The signed difference shows which group has the larger adjusted mean; \(|\mathrm{SMD}_j|\) summarizes the magnitude. For a binary covariate, use the adjusted proportions in the numerator and a fixed unweighted pooled standard deviation such as \(\sqrt{\{p_{1j}(1-p_{1j})+p_{0j}(1-p_{0j})\}/2}\) in the denominator. Use the same reference denominator across candidate penalties, so differences in SMDs reflect changes in adjusted group differences rather than changes in scale.

The adjusted means and proportions depend on the target estimand and adjustment method. For an average treatment effect (ATE) weighting comparison, for example, use weights \(A_i/\hat e_i\) for treated units and \((1-A_i)/(1-\hat e_i)\) for controls, then calculate each group’s weighted mean. Other estimands require weights matched to that estimand. Balance should be checked for all relevant measured covariates, not just variables retained with nonzero lasso coefficients.

Overlap is a separate requirement. Inspect the estimated score distributions by treatment group and the resulting weights. Scores near zero or one can create large inverse-probability weights and indicate that some covariate profiles have few plausible comparators. No single score-range rule is universal: judge the diagnostics in relation to the target population and the intended adjustment. A score-based procedure cannot create information where treatment groups do not overlap.

A Practical Penalty-Selection Workflow

A defensible choice uses a prespecified set of candidate penalties and evaluates prediction, balance, and overlap together. Balance thresholds such as an absolute SMD below \(0.10\) are conventions, not guarantees of no bias; choose and justify criteria in the context of the analysis. The penalty should not be selected by searching for the most favorable treatment-effect estimate.

Limitations, Diagnostics, and Common Mistakes

Penalty tuning cannot repair a poorly chosen feature set, unmeasured confounding, or lack of positivity. A model may show good average balance while leaving important subgroup differences, nonlinear relationships, or tail behavior inadequately represented. This tutorial’s SMD comparison is useful but does not replace substantive review of the covariates and score distributions. Later diagnostics can examine high-dimensional balance patterns in more depth.

Cross-validation folds should preserve the treatment structure sufficiently for stable fitting and evaluation; very small or highly unbalanced samples can make fold-level estimates noisy. Because a balance-guided choice is itself data-adaptive, it should be prespecified where possible and described as a selection procedure. Repeatedly changing the candidate set until a favorable balance table appears can make the reported diagnostic misleading.

• Choosing only the smallest prediction loss. Prediction loss measures how well treatment is predicted, not whether the planned adjusted groups are balanced. Compare balance and overlap before choosing.
• Using an absolute difference to report direction. An absolute SMD has no sign. Use the signed treated-minus-control difference divided by the reference standard deviation, then use its absolute value to describe magnitude.
• Treating an SMD threshold as proof. A small measured SMD does not prove exchangeability or remove bias from unmeasured variables. It is a diagnostic for measured covariates.
• Ignoring poor overlap because balance looks good. Weighting can make a summary balance statistic appear favorable while relying on a few extreme weights. Inspect score and weight distributions as well.
• Choosing by the causal estimate. Selecting \(\lambda\) because it produces a preferred effect estimate uses outcome results to tune the design and can distort inference. Base the choice on prespecified prediction and design diagnostics.