Augmented Balancing Weights as Linear Regression

Shikhar Tyagi, Arvind Pandey, Bhupendra Singh, Vrijesh Tripathi

2026-08-08

Introduction

The AugBalWeight package implements the methodology established in Bruns-Smith, Dukes, Feller, and Ogburn (2026) (Journal of the Royal Statistical Society Series B, DOI: 10.1093/jrsssb/qkaf019).

The paper establishes novel numeric equivalences showing that combining outcome regression models with balancing weights (automatic debiased machine learning) is numerically equivalent to a single linear model with coefficients that are a weighted combination of estimated OLS coefficients and the base outcome model coefficients.

Quick Start with LaLonde (1986) Job Training Dataset

We demonstrate the estimation of the Average Treatment Effect on the Treated (ATT) using the canonical LaLonde (1986) job training dataset.

# Load canonical LaLonde dataset
data(lalonde_data)

# Specify covariates
covariates <- c("age", "educ", "black", "hisp", "married", "re74", "re75", "age2", "educ2", "re742")
X <- as.matrix(lalonde_data[, covariates])
Y <- lalonde_data$re78
Z <- lalonde_data$treat

# Estimate ATT using Ridge-augmented L2 balancing weights
fit_att <- aug_bal_att(
  Y = Y,
  Z = Z,
  X = X,
  type = "l2",
  outcome_model = "ridge",
  tuning_method = "cv_outcome"
)

# Print ATT estimate and confidence interval
print(fit_att)
#> 
#> Augmented Balancing Weights Estimation
#> --------------------------------------
#> Point Estimate : 1483.83679 
#> Std. Error     : 327.66037 
#> 95% CI         : [ 841.63426 , 2126.03931 ]
#> Weight Type    : l2 
#> Outcome Model  : ridge 
#> Hyperparameters: lambda =   0.1 | delta =   0.1

Summary and Coefficient Comparison

We can examine the summary table comparing OLS coefficients, base outcome model coefficients, and implied augmented coefficients \(\hat{\beta}_{\text{aug}}\).

summary(fit_att)
#> 
#> Summary of Augmented Balancing Weights Model
#> ============================================
#> Point Estimate : 1483.83679 
#> Std. Error     : 327.66037 
#> 95% CI         : [ 841.63426 , 2126.03931 ]
#> Penalty Params : lambda =   0.1 | delta =   0.1 
#> 
#> Regression Coefficients & Feature Shift Table:
#>             OLS  BaseModel  Augmented      ObsShift      ImpShift
#> X1   -134.06205  -82.41191 -119.56702      -2.52710      -2.55457
#> X2    647.37955  449.37634  585.78934      -0.03347      -0.03730
#> X3     60.19372  144.92185  131.80979      -0.00295       0.00246
#> X4  -1008.40079 -561.18427 -797.73324      -0.00379       0.00023
#> X5   -246.14568 -187.20791 -235.59689      -0.03408      -0.01937
#> X6      0.31414    0.31541    0.31388     555.98738     553.84072
#> X7      0.58169    0.58081    0.58099     -27.56096     -29.59524
#> X8      1.41088    0.73698    1.22327    -200.49724    -201.22000
#> X9    -24.87905  -15.99406  -22.11486      -0.64977      -0.73226
#> X10    -0.00001   -0.00001   -0.00001 3026435.14216 3012637.19426
#> 
#> Balancing Weights Summary:
#>    Min. 1st Qu.  Median    Mean 3rd Qu.    Max. 
#> -0.3201  0.7789  1.0097  1.0000  1.3099  1.6211

Balance Diagnostics

The package includes safe graphical routines for diagnosing covariate imbalance pre- and post-balancing.

# Plot covariate balance diagnostic
plot(fit_att, which = 1)


# Plot distribution of estimated balancing weights
plot(fit_att, which = 2)

Double Lasso (\(\ell_\infty\) Balancing)

For high-dimensional settings, double_lasso performs \(\ell_\infty\) balancing weights combined with a lasso outcome model, demonstrating the double selection phenomenon (\(I_{\text{aug}} = I_\lambda \cup I_\delta\)).

fit_lasso <- double_lasso(
  Y = Y[Z == 0],
  X_p = X[Z == 0, ],
  target_mean = colMeans(X[Z == 1, ]),
  lambda = 0.05,
  delta = 0.05
)

cat("Active outcome features :", fit_lasso$active_outcome, "\n")
#> Active outcome features : 1 2 3 4 5 6 7 8 9
cat("Active balance features :", fit_lasso$active_balance, "\n")
#> Active balance features : 1 6 7 8 9 10
cat("Active union features   :", fit_lasso$active_union, "\n")
#> Active union features   : 1 2 3 4 5 6 7 8 9 10

References