Optimizing Tuning Parameters for Balance

smd.mean, smd.max, smd.rms

The mean, maximum, and root mean square of the absolute standardized mean differences computed for the covariates using col_w_smd(). The other allowable arguments include estimand (ATE, ATC, or ATT) to select the estimand, focal to identify the focal treatment group when the ATT is the estimand and the treatment has more than two categories, and pairwise to select whether mean differences should be computed between each pair of treatment groups or between each treatment group and the target group identified by estimand (default TRUE). Can be used with binary and multi-category treatments.

ks.mean, ks.max, ks.rms

The mean, maximum, or root-mean-squared Kolmogorov-Smirnov statistic, computed using col_w_ks(). The other allowable arguments include estimand (ATE, ATC, or ATT) to select the estimand, focal to identify the focal treatment group when the ATT is the estimand and the treatment has more than two categories, and pairwise to select whether statistics should be computed between each pair of treatment groups or between each treatment group and the target group identified by estimand (default TRUE). Can be used with binary and multi-category treatments.

ovl.mean, ovl.max, ovl.rms

The mean, maximum, or root-mean-squared overlapping coefficient compliment, computed using col_w_ovl(). The other allowable arguments include estimand (ATE, ATC, or ATT) to select the estimand, focal to identify the focal treatment group when the ATT is the estimand and the treatment has more than two categories, and pairwise to select whether statistics should be computed between each pair of treatment groups or between each treatment group and the target group identified by estimand (default TRUE). Can be used with binary and multi-category treatments.

mahalanobis

The Mahalanobis distance between the treatment group means, which is computed as \[ \sqrt{(\mathbf{\bar{x}}_1 - \mathbf{\bar{x}}_0) \Sigma^{-1} (\mathbf{\bar{x}}_1 - \mathbf{\bar{x}}_0)} \] where \(\mathbf{\bar{x}}_1\) and \(\mathbf{\bar{x}}_0\) are the vectors of covariate means in the two treatment groups, \(\Sigma^-1\) is the (generalized) inverse of the covariance matrix of the covariates (Franklin et al. 2014). This is similar to "smd.rms" but the covariates are standardized to remove correlations between and de-emphasize redundant covariates. The other allowable arguments include estimand (ATE, ATC, or ATT) to select the estimand, which determines how the covariance matrix is calculated, and focal to identify the focal treatment group when the ATT is the estimand. Can only be used with binary treatments.

energy.dist

The total energy distance between each treatment group and the target sample, which is a scalar measure of the similarity between two multivariate distributions. See Huling and Mak (2022) for details. The other allowable arguments include estimand (ATE, ATC, or ATT) to select the estimand, focal to identify the focal treatment group when the ATT is the estimand and the treatment has more than two categories, and improved to select whether the “improved” energy distance should be used, which emphasizes difference between treatment groups in addition to difference between each treatment group and the target sample (default TRUE). Can be used with binary and multi-category treatments.

kernel.dist

The kernel distance between treatment groups, which is a scalar measure of the similarity between two multivariate distributions. See Zhu, Savage, and Ghosh (2018) for details. Can only be used with binary treatments.

l1.med

The median L1 statistic computed across a random selection of possible coarsening of the data. See Iacus, King, and Porro (2011) for details. The other allowable arguments include l1.min.bin (default 2) and l1.max.bin default (12) to select the minimum and maximum number of bins with which to bin continuous variables and l1.n (default 101) to select the number of binnings used to select the binning at the median. covs should be supplied without splitting factors into dummies to ensure the binning works correctly. Can be used with binary and multi-category treatments.

r2, r2.2, r2.3

The post-weighting \(R^2\) of a model for the treatment given the covariates. Franklin et al. (2014) describe a similar but less generalizable metric, the “post-matching c-statistic”. The other allowable arguments include poly to add polynomial terms of the supplied order to the model and int (default FALSE) to add two-way interactions between covariates into the model. Using "r2.2" is a shortcut to requesting squares, and using "r2.3" is a shortcut to requesting cubes. Can be used with binary and continuous treatments. For binary treatments, the McKelvey and Zavoina \(R^2\) from a logistic regression is used; for continuous treatments, the \(R^2\) from a linear regression is used.

p.mean, p.max, p.rms

The mean, maximum, or root-mean-squared absolute Pearson correlation between the treatment and covariates, computed using col_w_corr(). Can only be used with continuous treatments.

s.mean, s.max, s.rms

The mean, maximum, or root-mean-squared absolute Spearman correlation between the treatment and covariates, computed using col_w_corr(). Can only be used with continuous treatments.

distance.cov

The distance covariance between the scaled covariates and treatment, which is a scalar measure of the independence of two possibly multivariate distributions. See Huling, Greifer, and Chen (2023) for details. Can only be used with continuous treatments.

Choosing a balance statistic

Given all these options, how should one choose? There has been some research into which yields the best results (Franklin et al. 2014; Griffin et al. 2017; Stuart, Lee, and Leacy 2013; Belitser et al. 2011; Parast et al. 2017), but the actual performance of each depends on the unique features of the data and system under study. For example, in the unlikely case that the true outcome model is linear in the covariates, using the "smd" or "mahalanobis" statistics will work well for binary and multi-category treatments. In more realistic cases, though, every measure has its advantages and disadvantages.

For binary and multi-category treatments, only "energy.dist", "kernel.dist", and "L1.med" reflect balance on all features of the joint covariate distribution, whereas the others summarize across balance statistics computed for each covariate ignoring the others. Similarly, for continuous treatments, only "distance.cov" reflects balance on all features of the joint covariate distribution. Given these advantages, "energy.dist" and "distance.cov" are my preferences. That said, other measures are better studied, possibly more intuitive, and more familiar to a broad audience.

Tuning GBM for balance

GBM has many tuning parameters that can be optimized, but the key parameter is the number of trees to use to calculate the predictions. WeightIt and twang both implement the methods described in McCaffrey, Ridgeway, and Morral (2004) for selecting the number of trees using a balance criterion. Here, we will do so manually both to understand the internals of these functions and illustrate the uses of bal.compute() and bal.init().

data("lalonde")

# Initialize balance
covs <- subset(lalonde, select = -c(treat, re78))
ks.init <- bal.init(covs,
                    treat = lalonde$treat,
                    stat = "ks.max",
                    estimand = "ATT")

# Fit a GBM model using `WeightIt` and `twang` defaults
fit <- gbm::gbm(treat ~ age + educ + married + race +
                    nodegree + re74 + re75,
                data = lalonde,
                distribution = "bernoulli",
                n.trees = 4000, interaction.depth = 3,
                shrinkage = .01, bag.fraction = 1)

trees_to_test <- seq(1, 4000)

p.mat <- predict(fit, type = "response",
                 n.trees = trees_to_test)

stats <- apply(p.mat, 2, function(p) {
    # Compute ATT weights
    w <- ifelse(lalonde$treat == 1, 1, p/(1-p))
    
    bal.compute(ks.init, weights = w)
})

stats[which.min(stats)]
#>       1408 
#> 0.09563649

From these results, we see that using 1408 trees gives us the lowest maximum KS statistic of 0.0956. Out of interest, we can plot the relationship between the number of trees and our balance statistic to see what it looks like:

library("ggplot2")
ggplot() +
    geom_line(aes(x = trees_to_test, y = stats)) +
    theme_bw() +
    labs(y = "ks.max", x = "n.trees")

Let’s compare this to the output of WeightIt:

library("WeightIt")
w.out <- weightit(treat ~ age + educ + married + race +
                    nodegree + re74 + re75,
                data = lalonde, estimand = "ATT",
                method = "gbm", n.trees = 4000,
                stop.method = "ks.max")

# Display the best tree:
w.out$info$best.tree
#> [1] 1408

# ks.max from weightit()
bal.compute(ks.init, weights = get.w(w.out))
#> [1] 0.09563649

We can see that weightit() also selects 1408 trees as the optimum and the resulting maximum KS statistic computed using the returned weights is equal to the one we computed manually. Using twang also produces the same results.