Extracted main text — title through conclusion, appendix excluded. This is what our citation measures are computed over, published so the extraction can be checked by eye.
101,663 characters · 29 sections · 124 citation commands
Augmented balancing weights as linear regression
\pagenumbering{arabic}
Combining outcome modeling and weighting, as in augmented inverse propensity score weighting (AIPW) and other doubly robust (DR) or double machine learning (DML) estimators, is a core strategy for estimating causal effects using observational data. A growing body of literature finds weights by solving a “balancing weights” optimization problem to estimate weights directly, rather than by first estimating the propensity score and then inverting. DR versions of these estimators are referred to by a number of terms, including augmented balancing weights athey2018approximate, hirshberg2021augmented, automatic debiased machine learning chernozhukov2022automatic, and generalized regression estimators deville1992calibration; see ben2021balancing for a review. Moreover, this strategy has been applied to a wide range of linear estimands via the Riesz representation theorem hirshberg2021augmented, chernozhukov2022riesz. In this paper, we consider augmented balancing weights in which the estimators for both the outcome model and the balancing weights are based on penalized linear regressions in some possibly infinite basis; in addition to all high-dimensional linear models, this broad class includes popular nonparametric models such as kernel regression and certain forms of random forests and neural networks.
We first show that, somewhat surprisingly, augmenting any regularized linear outcome regression (the “base learner") with linear balancing weights is numerically equivalent to a single linear outcome regression applied to the target covariate profile. The resulting coefficients are an affine (and often convex) combination of the base learner model coefficients and unregularized OLS coefficients; the hyperparameter for the balancing weights estimator directly controls the regularization path defining the affine combination. In the extreme case where the weighting hyperparameter is set to zero --- which we show can easily occur in practice --- the entire procedure is equivalent to estimating a single, unregularized OLS regression.
We specialize these results to ridge and lasso regularization ($\ell_2$ and $\ell_\infty$ balancing, respectively) and show that augmenting an outcome regression estimator with balancing weights generally corresponds to a form of undersmoothing. Most notably, we show that an augmented balancing weight estimator that use (kernel) ridge regression for both outcome and weighting models --- which we refer to as “double ridge” --- collapses to a single, undersmoothed (kernel) ridge regression estimator.
We leverage these results to prove novel statistical results for double ridge estimators and to make progress towards practical hyperparameter tuning, which remains an open problem in this area. We first make explicit the connection between asymptotic results for double kernel ridge estimators singh2020kernel and prior results on optimal undersmoothing for a single kernel ridge outcome model mou2023kernel, showing that the latter is also semiparametrically efficient. This generalizes the argument in robins2007comment that “OLS is doubly robust” to a much broader class of penalized parametric and non-parametric regression estimators. As a complementary analysis, we next adapt existing finite sample error analysis results for single ridge regression dobriban2018high to derive the finite-sample-exact bias and variance of double ridge estimators. Using these expressions, we can compute oracle hyperparameters for any given data-generating process.
Finally, we illustrate our results with several numerical examples. We first explore hyperparameter tuning for double ridge regression in an extensive simulation study on 36 data-generating processes, and compare three practical methods to the optimal hyperparameter computed using our finite sample analysis. Surprisingly, asymptotic theory and our simulation results suggest equating the hyperparameters for the outcome and weighting models. We caution against the naive application of hyperparameter tuning based solely on cross-validating the weighting model, forms of which have been suggested previously. This approach can lead to setting the weighting hyperparameter to exactly zero --- and therefore recovering standard OLS --- even in scenarios where OLS is far from optimal. We emphasize this point by applying our results to the canonical lalonde1986evaluating study, highlighting that researchers can inadvertently recover OLS in practice.
Broadly, our results provide important insights into the nexus of causal inference and machine learning. First, these results open the black box on the growing number of methods based on augmented balancing weights and AutoDML --- methods that can sometimes be difficult to taxonomize or understand. We show that, under linearity, these estimators all share an underlying and very simple structure. Our results further highlight that estimation choices for augmented balancing weights can lead to potentially unexpected behavior. At a high level, as causal inference moves towards incorporating machine learning and automation, our work highlights how the traditional lines between weighting and regression-based approaches are becoming increasingly blurred.
Second, our results connect two approaches to “automate” semiparametric causal inference. AutoDML and related methods exploit the fact that we can estimate a Riesz representer without a closed form expression for a wide class of functionals. The estimated Riesz representer then augments a base learner by bias correcting a plug-in estimator of the functional. Older approaches, such as undersmoothing goldstein1992optimal,newey1998undersmoothing, twicing kernels newey2004twicing, and sieve estimation newey1994asymptotic,shen1997methods, avoid estimation of the Riesz representer, tuning the base learner regression fit such that an additional bias correction is not required. Achieving this optimal tuning in practice has long been a hurdle for the implementation of these methods. Subject to certain conditions, both approaches can yield estimators that are asymptotically efficient. We show that if all required tuning parameters are defined in terms of an $\ell_2$-norm constraint, then these approaches can be numerically identical even in finite samples. We use these equivalences to make progress toward practical hyperparameter selection and find promising directions for new theoretical analysis.
In Section (ref) we introduce the problem setup, identification assumptions, and common estimation methods; we also review balancing weights and previous results linking balancing weights to outcome regression models. In Section (ref) we present our new numerical results, and in Sections (ref) and (ref) we cache out the implications for $\ell_2$ and $\ell_\infty$ balancing weights specifically. Building on our numerical results, Section (ref) explores both asymptotic and finite sample statistical results for kernel ridge regression. Section (ref) illustrates our results with a simulation study and application to canonical data sets. Section (ref) offers some other directions for future research. The appendix includes extensive additional technical discussion and extensions.
\paragraph{Balancing weights and AutoDML.} With deep roots in survey calibration methods and the generalized regression estimator deville1992calibration,lumley2011connections,gao2022soft, a large and growing causal inference literature uses balancing weights estimation in place of traditional inverse propensity score weighting (IPW). ben2021balancing provide a recent review; we discuss specific examples at length in Section (ref) below. This approach typically balances features of the covariate distributions in the different treatment groups, with the aim of minimising the maximal design-conditional mean squared error of the treatment effect estimator. Of particular interest here are augmented balancing weights estimators that combine balancing weights with outcome regression; see, for example, athey2018approximate, hirshberg2021augmented, ben2021augmented.
A parallel literature in econometrics instead focuses on so-called automatic estimation of the Riesz representer, of which IPW are a special case, where “automatic" refers to the fact that we can estimate the Riesz representer without obtaining a closed form expression. Estimating the Riesz representer directly, under the assumption that it is linear in some basis, dates back at least to robins2008higher; see also robins2007comment. The corresponding augmented estimation framework has more recently come to be known as Automatic Debiased Machine Learning, or AutoDML; see, among others, chernozhukov2022locally, chernozhukov2022riesznet, chernozhukov2022automatic, and chernozhukov2022riesz. This approach has also been applied in a range of settings, including to corrupted data agarwal2021causal, to dynamic treatment regimes chernozhukov2022dynamic, and to address noncompliance singh2022late. As we discuss in Appendix (ref), the AutoDML approach nearly always employs cross-fitting and is typically motivated by asymptotic properties rather than achieving minimax design-conditional mean squared error.
\paragraph{Numerical equivalences for balancing weights.} Many seminal papers highlight connections between weighting approaches, such as balancing weights and IPW, and outcome modeling; see bruns2022outcome for discussion. Most relevant are a series of papers that show numerical equivalences between linear regression and (exact) balancing weights, especially robins2007comment,kline2011oaxaca,chattopadhyay2021implied, and between kernel ridge regression and forms of kernel weighting kallus2020generalized, hirshberg2019minimax. We discuss these equivalences at length in (ref). Finally, as we discuss in (ref), there are close connections between balancing weights and Empirical Likelihood hellerstein1999el, newey2004el.
The core results in our paper are numeric equivalences for existing estimation procedures, and as such these results hold absent any causal assumptions or statistical model. Nonetheless, a primary motivation for this work is the task of estimating unobserved counterfactual means in causal inference, as well as estimating the broad class of linear functionals described in chernozhukov2018learning. We briefly review the corresponding setup, emphasizing that this is purely for interpretation.
Let $X, Y, Z$ be random variables defined on $\mathcal{X} , \mathbb{R}, \mathcal{Z}$ with joint probability distribution $p$. To begin, consider the example of a binary treatment, $\mathcal{Z} = \{0,1\}$ and covariates $X$. Define potential or counterfactual outcomes $Y(1)$ and $Y(0)$ under assignment to treatment and control, respectively. Under SUTVA rubin1980randomization, we observe outcomes $Y = Z Y(1) + (1 - Z)Y(0)$. To estimate the average treatment effect, $\mathbb{E}[ Y(1) - Y(0) ]$, we first estimate the means of the partially observed potential outcomes. We initially focus on estimating $\mathbb{E}[ Y(1) ]$; a symmetric argument holds for $\mathbb{E}[Y(0)]$.
Let $m(x,z) \coloneqq \mathbb{E}[Y \mid X = x, Z = z]$ be the outcome model, $e(x) \coloneqq \mathbb{P}[Z = 1 \mid X = x]$ be the propensity score, and $\alpha(x,z) = z / e(x)$ be the inverse propensity score weights (IPW). Under the additional assumptions of conditional ignorability, $Y(1) \protect\mathpalette{\protect\independenT}{\perp} Z \mid X$, and overlap, $\mathbb{E}[\alpha(X,Z)^2] < \infty$, we have that $\mathbb{E}[ Y(1) ]$ is identified by $ \mathbb{E}[ m(X,1) ]$, a linear functional of the observed data distribution.
There are three broad strategies for estimating $\mathbb{E}[ Y(1) ]$. First, the identifying functional above suggests estimating the outcome model, $m(x, 1)$ among those units with $Z=1$, and plugging this into the regression functional, $\mathbb{E}[ m(X,1) ]$. Second, the equality $\mathbb{E}[ m(X,1) ] = \mathbb{E}[Z/e(X) Y] = \mathbb{E}[\alpha(X,Z) Y]$ suggests estimating the inverse propensity score weights, $\alpha(x,z) = z/e(x)$, and plugging these into the weighting functional. Finally, we can combine these two via the doubly robust functional robins1994estimation: $$ \mathbb{E}[m(X,1) + \alpha(X,Z)(Y - m(X,1)]. $$ This functional has the attractive property of being equal to $\mathbb{E}[ m(X,1) ]$ even if either one of $\alpha$ or $m$ is replaced with an arbitrary function of $X$ and $Z$, hence the term “doubly robust.” Doubly robust estimators have been studied extensively in semiparametric theory; note that $m(X,1) + \alpha(X,Z) (Y - m(X,Z))-\psi(m)$ coincides with the efficient influence function for $\psi(m)$ under a nonparametric model (see kennedy2022semiparametric for a review of the relevant theory).See chernozhukov2018double, kennedy2022semiparametric for recent overviews of the active literature in causal inference and machine learning focused on estimating versions of this functional.
Our results apply well beyond the example above. In particular, they apply to any functional of the form
where $\mathcal{Z}$ is an arbitrary set; $Z$ a random variable with support $\mathcal{Z}$; and $h$ is a real-valued, mean-squared continuous linear functional of $m$ chernozhukov2018learning,hirshberg2021augmented, chernozhukov2022automatic. Following chernozhukov2022automatic, chernozhukov2022riesz, we can generalize the weighting functional to this general class of estimands via the Riesz representer, which is a function $\alpha(X,Z) \in L_2(p)$ such that, for all square-integrable functions $f \in L_2(p)$:
As in the counterfactual mean example, we can identify the more general target functional in (ref) via the outcome regression functional in (ref), via the Riesz representer functional in (ref) with $f=m$, or via the doubly robust functional
Estimators of this DR functional are augmented in the sense that they augment the “plug-in,” “outcome regression,” or “base learner” estimator of $\mathbb{E}[h(X, Z, m)]$ with appropriately weighted residuals; or, equivalently, that augment the weighting estimator with an appropriate outcome regression. This is the class of estimators to which our results apply. Doubly robust estimators have been studied extensively in semiparametric theory. In particular, $h(X, Z, m) + \alpha(X,Z) (Y - m(X,Z))-\psi(m)$ coincides with the efficient influence function for $\psi(m)$ under a nonparametric model (see kennedy2022semiparametric for a review of the relevant theory). In future work we will explore whether we can extend our results to a different class of functionals that admit DR functional forms, first introduced by robins2008higher, and to the superset of such functionals characterized by rotnitzky2021characterization.
The core idea behind balancing weights is to estimate the Riesz representer directly --- rather than via an analytic functional form (e.g., by estimating the propensity score and inverting it). As a result, balancing weights do not require a known analytic form for the Riesz representer chernozhukov2022riesz, are often much more stable zubizarreta2015stable, and offer improved control of finite sample covariate imbalance zhao2019covariate. We briefly describe two primary motivations for this approach.
First, a central property of the Riesz representer is that the corresponding weights, $w(X, Z) = \alpha(X,Z)$, are the unique weights that satisfy the population balance property property in (ref) for all square-integrable functions $f \in L_2(p)$. For our target estimand $\psi(m)$ we only need to satisfy the condition in (ref) for the special case of $f=m$. If we are willing to assume that $m$ lies in a model class $\mathcal{F} \subset L_2(p)$, then it suffices to balance functions in that class. This is achieved by minimizing the imbalance over $\mathcal{F}$:
As we discuss next, balancing weights minimize a (penalized) sample analog of Equation (ref).
Alternatively, chernozhukov2022automatic consider finding weights $f$ that minimize the mean-squared error for $\alpha(X,Z)$:
Automatic estimation of the Riesz representer, also known as Riesz regression chernozhukov2024riesz_regression, minimizes a sample analog of Equation (ref). When $\mathcal{F}$ is convex, then up to choice of hyperparameters (see (ref) below), the solutions to Equations (ref) and (ref) are equivalent.
In this paper, we consider the special case in which the outcome models are linear in some basis expansion of $X$ and $Z$. This is an extremely broad class that encompasses linear and polynomial models of arbitrary functions of $X$ and $Z$ and with dimension possibly larger than the sample size, as well as non-parametric models such as reproducing kernel Hilbert spaces gretton2012kernel, the Highly-Adaptive Lasso benkeser2016hal, the neural tangent kernel space of infinite-width neural networks jacot2018neural, and “honest” random forests agarwal2022randomforests. However, this class excludes models for $m$ that are fundamentally non-linear in their parameters, like general neural networks or generalized linear models with a non-linear link function. We extend our results to arbitrary nonlinear balancing weights in Appendix (ref).
Under linearity, the imbalance over all $f \in \mathcal{F}$ has a simple closed form. Because our results concern numeric equivalences, we will focus on the finite sample version of the linear balancing weights problem. Let $\mathcal{F} = \{ f(x,z) = \theta^\top\phi(x,z) : \Vert \theta \Vert \leq 1 \}$ where $\Vert \cdot \Vert$ can be any norm on $\mathbb{R}^d$. The general setup constrains $\Vert \theta \Vert \leq r$; we set $r = 1$ without loss of generality, which simplifies exposition below. Let $\Vert \cdot \Vert_*$ be the dual norm of $\Vert \cdot \Vert$; that is, $ \Vert v \Vert_* \coloneqq \sup_{\Vert u \Vert \leq 1} u^\top v.$ Many common vector norms have familiar, closed-form, dual norms, e.g., the dual norm of the $\ell_2$-norm is the $\ell_2$-norm; and the dual norm of the $\ell_1$-norm is the $\ell_\infty$-norm. Let $X_p, Y_p, Z_p$ be $n$ i.i.d. samples from the distribution $p$ of the observed data. Define the feature map $\phi : \mathcal{X} \times \mathcal{Z} \rightarrow \mathbb{R}^d$ and let $\phi_j : \mathcal{X} \times \mathcal{Z} \rightarrow \mathbb{R}$ denote the mapping for the $j$th feature. Define $\Phi_p \coloneqq \phi(X_p, Z_p)$ and let $\Phi_q \coloneqq h(X_p, Z_p, \phi)$ denote the target features. We will write $\hat{\mathbb{E}}$ for sample averages; define $\mkern1.5mu\overline{\mkern-1.5mu\Phi\mkern-1.5mu}\mkern 1.5mu_p \coloneqq \hat{\mathbb{E}}[ \Phi_p ]$ and $\mkern1.5mu\overline{\mkern-1.5mu\Phi\mkern-1.5mu}\mkern 1.5mu_q \coloneqq \hat{\mathbb{E}}[ \Phi_q ]$. For exposition, we assume that $d < n$ and that $\Phi_p$ has rank $d$. We emphasize that this is not necessary for our results --- one can replace $\mathbb{R}^d$ with an infinite-dimensional Hilbert space $\mathcal{H}$ and relax the rank restriction. See (ref) for a formal presentation of the high-dimensional ($d>n$) setting.
In what follows we write $w$ for the $1 \times n$ vector $w(\Phi_p)$, to highlight the fact that we will estimate $w$ directly rather than as an explicit function of $X$ or $\Phi_p$. Using the derivation above, we can directly calculate the finite sample imbalance as: $$ \widehat{\text{Imbalance}_{\mathcal{F}}}(w) = \Vert \tfrac{1}{n} w \Phi_p - \bar{\Phi}_q \Vert_*.$$
Now we can write the penalized sample analog of balancing weights optimization problem in (ref) equivalently as either:
Furthermore, we can write the equivalent problem in (ref) as:
where we use the terminology “Riesz regression” from chernozhukov2024riesz_regression. For any parameter $\delta_2 > 0$ and corresponding constrained problem solution $\hat{w}$, there exists a parameter $\delta_3 > 0$ such that $\hat{w} = \delta_3 \Phi_p \hat{\theta}$, where $\hat\theta$ is the solution to the Riesz regression form. As a result, for any norm $\Vert \cdot \Vert$, the penalized and constrained forms will always produce weights that are linear in $\Phi_p$ ben2021balancing. Therefore, since the problems are equivalent, we typically use a generic $\delta$ to denote the regularization parameter, and will specify the particular form only if necessary. In Appendix (ref) we illustrate several concrete examples for this problem and in Appendix (ref) we consider alternative dispersion parameters and discuss popular forms of balancing that constrain the weights to be non-negative.
Our first main result demonstrates that any linear balancing weights estimator is equivalent to applying OLS to the re-weighted features. Our second result provides a novel analysis of augmented balancing weights, demonstrating that augmenting any linear balancing weights estimator with a linear outcome regression estimator is equivalent to a plug-in estimator of a new linear model with coefficients that are a weighted combination of estimated OLS coefficients and the coefficients of the original linear outcome model.
Our first result is that estimating $\psi(m)$ with any linear balancing weights is equivalent to fitting OLS for the regression of $Y_p$ on $\Phi_p$ and then applying those coefficients to the re-weighted target feature profile. The key idea for this result begins with the simple unregularized regression prediction for $\psi(m)$, $\mkern1.5mu\overline{\mkern-1.5mu\Phi\mkern-1.5mu}\mkern 1.5mu_q \hat{\beta}_{\text{ols}}$.
Note that here we have written the OLS coefficients using the pseudo-inverse $\dag$. For clarity in the main text, we focus on the full rank setting, where $(\Phi_p^\top \Phi_p)^\dag = (\Phi_p^\top \Phi_p)^{-1}$; we provide a proof for the general setting in (ref). In (ref), we extend Proposition (ref) to non-linear balancing weights, including those wiht a non-negativity constraint.
We can interpret this result via a contrast with standard regularization. Regularized regression models navigate a bias-variance trade-off by regularizing estimated coefficients $\hat\beta_{\text{reg}}$ relative to $\hat{\beta}_{\text{ols}}$, leading to $\mkern1.5mu\overline{\mkern-1.5mu\Phi\mkern-1.5mu}\mkern 1.5mu_q \hat{\beta}_{\text{reg}}$. The balancing weights approach instead keeps $\hat{\beta}_{\text{ols}}$ fixed and regularizes the target feature distribution by penalizing the implied feature shift, $\widehat{\Delta}^\delta = \hat{\Phi}^\delta_q - \mkern1.5mu\overline{\mkern-1.5mu\Phi\mkern-1.5mu}\mkern 1.5mu_p$.
We emphasize that this is a new and quite general result. As we discuss in (ref), it has been shown previously that for exact balancing weights, $ \hat{\mathbb{E}}[\hat{w}_\text{exact} Y_p ] = \mkern1.5mu\overline{\mkern-1.5mu\Phi\mkern-1.5mu}\mkern 1.5mu_q \hat{\beta}_\text{ols}$. However, (ref) holds for any weights of the form $w = \theta \Phi_p^\top$ with arbitrary $\theta \in \mathbb{R}^d$. In Sections (ref) and (ref), we consider the particular form of $\hat{\Phi}_q^\delta$ for $\ell_2$ and $\ell_\infty$ balancing, respectively.
We can immediately extend this to augmented balancing weights, which regularize both the coefficients and the feature shift. Let $\hat\beta^{\lambda}_\text{reg}$ be the coefficients of any regularized linear model for the relationship between $Y_p$ and $\Phi_p$, where the superscript $\lambda$ indicates dependence on a hyperparameter (e.g., estimated by regularized least squares). We consider augmenting $\hat{\mathbb{E}}\left[\hat{w}^\delta \circ Y_p\right]$ with $\hat\beta^{\lambda}_\text{reg}$ using the doubly robust functional representation in Equation (ref). The augmented estimator is:
Many recently proposed estimators have this form; see e.g., athey2018approximate, ben2021balancing. If the weighting model and outcome model have different bases, our result applies to a shared basis by either combining the dictionaries as in chernozhukov2022automatic or by applying an appropriate projection as in hirshberg2021augmented.
We apply Proposition (ref) to the first term of the right-hand side of (ref) to yield the following result. As this result is purely numerical, it applies to arbitrary vectors $\hat\beta^{\lambda}_\text{reg} \in \mathbb{R}^d$, but substantively we think of $\hat\beta^{\lambda}_\text{reg}$ as the estimated coefficients from an outcome model.
This is our central numerical result for augmented balancing weights: when both the outcome and weighting models are linear, the augmented estimator is equivalent to a linear model applied to the target features $\Phi_q$, with coefficients that are element-wise affine combinations of the base learner coefficients, $\hat{\beta}_\text{reg}^\lambda$, and the coefficients $\hat{\beta}_{\text{ols}}$ from an OLS regression of $Y_p$ on $\Phi_p$. (The coefficients are additionally convex combinations of $\hat{\beta}_\text{reg}^\lambda$ and $\hat{\beta}_{\text{ols}}$ when the covariance matrix is diagonal.) In Sections (ref) and (ref) below, we analyze some of the properties of the augmented estimator for $\ell_2$ and $\ell_\infty$ balancing weights problems respectively.
The regularization parameter for the balancing weights problem, $\delta$, parameterizes the path between $\hat{\beta}_\text{reg}^\lambda$ and $\hat{\beta}_\text{ols}$. To see this, consider the cases where $\delta \to 0$ and $\delta \to \infty$. As $\delta \to 0$ the balancing weights problem prioritizes minimizing balance over controlling variance, and $\widehat{\Delta}^\delta_j \to \Delta_j$ for all $j$. (Recall that we assume $\mkern1.5mu\overline{\mkern-1.5mu\Phi\mkern-1.5mu}\mkern 1.5mu_{p,j} = 0$ for all $j$. Thus, $\Delta_j = \mkern1.5mu\overline{\mkern-1.5mu\Phi\mkern-1.5mu}\mkern 1.5mu_{q,j}$ and $\widehat{\Delta}_j^\delta = \hat{\Phi}_{q,j}^\delta$. So $\widehat{\Delta}_{j}^\delta \rightarrow \Delta_j$ is equivalent to $\hat{\Phi}_{q}^\delta \rightarrow \mkern1.5mu\overline{\mkern-1.5mu\Phi\mkern-1.5mu}\mkern 1.5mu_{q,j}$.) In this case, $a_j^\delta = \widehat{\Delta}_j^\delta / \Delta_j \to 1$, and the weights fully “de-bias” the original outcome model by recovering unregularized regression, $\hat{\beta}_\text{aug} \to \hat{\beta}_\text{ols}$. In Section (ref), we will see that when chosen by cross-validation, $\delta$ sometimes equals exactly $0$ in applied problems; thus even when $\hat{\beta}_\text{reg}^\lambda$ is a sophisticated regularized estimator, the final augmented point estimate can nonetheless be numerically equivalent to the simple OLS plug-in estimate. Conversely, as $\delta \to \infty$, the balancing weights problem prioritizes controlling variance, leading to uniform weights and $\widehat{\Delta}_j \to 0$. In this case, $a_j^\delta = \widehat{\Delta}_j^\delta / \Delta_j \rightarrow 0$, the weighting model does very little, and $\hat{\beta}_\text{aug} \rightarrow \hat{\beta}_\text{reg}^\lambda$.
It is also instructive to consider two other extremes: unregularized outcome model and unregularized balancing weights. First, consider the special case of fitting an unregularized linear regression outcome model, i.e., $\hat{\beta}_\text{reg}^\lambda = \hat{\beta}_{\text{ols}}$. Then Proposition (ref) reproduces the result, originally due to robins2007comment, that “OLS is doubly robust" kline2011oaxaca. This is because $\hat{\beta}_{\text{aug}} = \hat{\beta}_{\text{ols}}$ for arbitrary linear weights $\hat{\theta}^\delta \in \mathbb{R}^d$. Thus, OLS augmented by any choice of linear balancing weights collapses to OLS alone. Equivalently, we can view OLS alone as an augmented estimator that combines an OLS base learner with linear balancing weights.
A similar result holds for unregularized balancing weights, i.e., exact balancing weights. Let $\hat{w}_\text{exact}$ be the solution to a balancing weights problem in (ref) with hyperparameter $\delta = 0$, and let $\hat{\beta}_\text{reg}^\lambda \in \mathbb{R}^d$ be arbitrary coefficients. Then from the balance condition, $\hat{\Phi}_q = \mkern1.5mu\overline{\mkern-1.5mu\Phi\mkern-1.5mu}\mkern 1.5mu_q$, $a_j^\delta = 1$ for all $j$, and we have that $\hat{\beta}_{\text{aug}} = \hat{\beta}_{\text{ols}}$. Thus, the augmented exact balancing weights estimator also collapses to the OLS regression estimator. Equivalently, the augmented exact balancing weights estimator collapses to the unaugmented exact balancing weights estimator. zhao2017entropy use a very similar result to argue that entropy balancing, a form of exact balancing weights, is doubly robust.
Finally, before we turn to new results for $\ell_2$ and $\ell_\infty$ balancing, we briefly comment on several points that are discussed in more detail in the Appendix.
In this section, we study $\ell_2$ balancing weights estimators, which are commonly used in the context of kernel balancing gretton2012kernel, hirshberg2019minimax, kallus2020generalized, ben2021multilevel and for panel data methods abadie2010synthetic, ben2021augmented. We first show that the regularization path $a_j^\delta$ from (ref) follows typical ridge regression shrinkage, with a smooth decay. Moreover, augmenting with $\ell_2$ balancing weights is equivalent to boosting with ridge regression, and always overfits relative to the unaugmented outcome model alone. We then show that when the outcome model used to augment $\ell_2$ balancing weights is also a ridge regression (which we refer to as “double ridge”), the augmented estimator is itself equivalent to a single, generalized ridge regression, albeit undersmoothed relative to the base learner. These results extend immediately to the RKHS setting of “double kernel ridge” estimation, combining kernel balancing weights and kernel ridge regression. In Section (ref), we show the implications of these numeric results for undersmoothing in the statistical sense.
While the following results hold for arbitrary covariance matrices, in the main text we simplify the presentation by assuming that $\Phi_p^\top\Phi_p$ is diagonal; that is, $ (\Phi_p^\top\Phi_p) = \text{diag}(\sigma^2_1, \sigma^2_2, ..., \sigma^2_d)$, with $\sigma^2_j > 0$. We show that this is without loss of generality for $\ell_2$ balancing in (ref).
Following Remark (ref) above, $\ell_2$ balancing weights, including kernel balancing weights, have a closed form that is always linear in $\mkern1.5mu\overline{\mkern-1.5mu\Phi\mkern-1.5mu}\mkern 1.5mu_q$. Our next result applies this closed form to Proposition (ref) to derive the regularization path that results from augmenting an arbitrary linear outcome model with $\ell_2$ balancing weights. Although this is an immediate consequence of Proposition (ref), the resulting form of the augmented estimator has unique structure that warrants a new result.
In this case, the $a_j^\delta$ are exactly equal to the standard regularization path of ridge regression. To see this, recall that ridge regression with penalty $\delta$ shrinks the $\hat{\beta}_\text{ols}$ coefficients as follows:
This is identical to the expression in (ref) but with $\hat{\beta}_\text{reg}^\lambda$ set to $0$: Ridge regression shrinks $\hat{\beta}_\text{ols}$ towards $0$ with regularization path $a_j^\delta$, while $\ell_2$ augmenting shrinks $\hat{\beta}_\text{ols}$ towards $\hat{\beta}_\text{reg}^\lambda$ with the same regularization path.
As an illustration, the right panel of Figure (ref) shows $\hat{\beta}_{\ell_2}$ (on the y-axis) for ten covariates, with $\delta$ increasing from $0$ (on the x-axis). The dots on the left pick out $\hat{\beta}_{\text{ols}}$; when $\delta = 0$, then $a_j^0 = 1$ and $\hat{\beta}_{\ell_2} = \hat{\beta}_{\text{ols}}$. The limit on the right shows $\hat{\beta}_{\text{reg}}^\lambda$. The smooth regularization path is characteristic of ridge regression shrinkage.
We can also view $\hat{\beta}_{\ell_2}$ as the output of a single iteration of a ridge boosting procedure, fit using $Y_p$ and $\Phi_p$ alone. See buhlmann2003boosting and park2009l_2 for detailed discussion; newey2004twicing makes a similar connection in the context of twicing kernels.
So for a fixed $\delta$, the augmented $\ell_2$ balancing estimator is equivalent to estimating a new outcome model coefficient estimator $\hat{\beta}_{\ell_2}$ that overfits relative to $\hat{\beta}_\text{reg}^\lambda$ (in the sense of having smaller in-sample training error), and then applying that model to $\Phi_q$.
Surprisingly --- and in contrast to the general result in (ref) --- the augmented coefficients $\hat{\beta}_{\ell_2}$ are the same for every target covariate profile $\Phi_q$. To see this, note that (ref) shows that $\ell_2$ balancing weights are always linear in $\mkern1.5mu\overline{\mkern-1.5mu\Phi\mkern-1.5mu}\mkern 1.5mu_q$. Therefore, the corresponding regularization path $a_j^\delta$ does not depend on the target profile $\Phi_q$; it depends only on $\delta$ and the source distribution variances $\sigma_j^2$. This property is closely related to universal adaptability in the computer science literature on multi-group fairness kim2022universal. The particular $\Phi_q$ may nonetheless impact the choice of $\delta$ in hyperparameter selection, e.g., via cross-validating imbalance, which in turn influences the degree of overfitting; we do find this to be the case theoretically in (ref).
(ref) holds for arbitrary linear outcome model coefficient estimators $\hat{\beta}_\text{reg}^\lambda \in \mathbb{R}^d$; we now state the corresponding result for a “double ridge” estimator, where the base learner outcome model is itself fit via ridge regression. The key takeaway is that the implied augmented coefficients are undersmoothed relative to the base learner ridge coefficients.
For this section, we will consider the following generalized ridge regression, sometimes known as “adaptive” ridge regression grandvalet1998least. Let $\Lambda \in \mathbb{R}^{d \times d}$ be a diagonal matrix with $j$th diagonal entry $\lambda_j \geq 0$. Then the generalized ridge coefficients are:
Standard ridge regression is the special case where the $\lambda_j$ all take the same value and so $\Lambda = \lambda I$. As above, the generalized ridge coefficients can be rewritten as shrinking the OLS coefficients:
We now demonstrate that the augmented $\ell_2$ balancing weights estimator with base learner $\hat{\beta}_\text{ridge}^\Lambda$ is equivalent to a plug-in estimator using generalized ridge with smaller hyperparameters, $\hat{\beta}_\text{ridge}^\Gamma$, where $\Gamma$ is a diagonal matrix with $j$th diagonal entry $\gamma_j \in [0,\lambda_j]$.
The same result holds for kernel ridge regression; see (ref).
In this setting, augmenting with balancing weights is equivalent to undersmoothing the original outcome model fit. In particular, we can use the expansion in (ref) to see the undersmoothing in $\hat{\beta}^\Gamma_\text{ridge}$ explicitly: \[ \frac{\sigma^2_j}{\sigma^2_j + \gamma_j} = \underbrace{\left(\frac{\sigma^2_j}{\sigma^2_j + \lambda_j} \right)}_{\text{outcome model}} \underbrace{\left(\frac{\sigma_j^2 + \lambda_j + \delta}{\sigma_j^2 + \delta}\right)}_{\text{augmentation}}, \] where the first term is the shrinkage from the original generalized ridge model alone, and the second term is due to augmenting with $\ell_2$ balancing weights. Importantly, the second term is in $[1, \frac{\sigma^2_j + \lambda_j}{\sigma^2_j}]$ and therefore partially reverses the shrinkage of the original estimate. In Section (ref), we connect this to undersmoothing in the statistical sense.
In this section, we study $\ell_\infty$ balancing weights estimators, which are widely used in the balancing weights literature zubizarreta2015stable, athey2018approximate and in the AutoDML literature chernozhukov2022automatic. In the main text, we consider the special case where the covariance matrix $\Phi_p^\top\Phi_p$ is diagonal; that is, $ (\Phi_p^\top\Phi_p) = \text{diag}(\sigma^2_1, \sigma^2_2, ..., \sigma^2_d)$, with $\sigma^2_j > 0$. Unlike with $\ell_2$ balancing, this is no longer without loss of generality. We discuss this general case in (ref).
For diagonal covariance, we first show that $\ell_\infty$ balancing has a closed form: it is equivalent to applying a soft-thresholding operator to the feature shift from $\mkern1.5mu\overline{\mkern-1.5mu\Phi\mkern-1.5mu}\mkern 1.5mu_p$ to $\mkern1.5mu\overline{\mkern-1.5mu\Phi\mkern-1.5mu}\mkern 1.5mu_q$. We then write the resulting augmented estimator as applying coefficients $\hat{\beta}_{\ell_\infty}$ to $\Phi_q$ and show that $\hat{\beta}_{\ell_\infty}$ is a sparse, element-wise convex combination of the base learner coefficients and OLS coefficients. When the outcome model is also fit via the lasso, we use the resulting representation to demonstrate a familiar “double selection” phenomenon belloni2014inference, where $\hat{\beta}_{\ell_\infty}$ inherits the non-zero coefficients of both the base learner and the weighting model. This is a form of undersmoothing in the $\ell_0$ “norm,” in the sense that $\hat{\beta}_{\ell_\infty}$ always has at least as many non-zero coefficients as the base learner, $\hat{\beta}_{\text{reg}}$.
We first define the soft-thresholding operator and show that the $\ell_\infty$ balancing problem has a closed form solution.
For intuition, compare the (un-augmented) $\ell_\infty$ balancing weights estimator to the lasso-based coefficient estimates hastie2009elements:
where we simplify $\hat{\Phi}_q^\delta$ here to emphasize the connections between the methods. Whereas lasso performs soft-thresholding on the OLS coefficients (regularizing the outcome regression), $\ell_\infty$ balancing performs soft-thresholding on the implied feature shift to the target features.
We can then plug the closed-form solution for the weights into (ref).
The augmented coefficients $\hat{\beta}_{\ell_\infty}$ are an element-wise convex combination of $\hat{\beta}_\text{reg}^\lambda$ and $\hat{\beta}_\text{ols}$. For features where the mean feature shift $\Delta_j$ is small (relative to $\delta$), $\hat{\beta}_{\ell_\infty}$ is equivalent to the base learner coefficient $\hat{\beta}_\text{reg}^\lambda$. The remaining coefficients are interpolated linearly toward the $\hat{\beta}_\text{ols}$ coefficients.
(ref) summarizes these results and their implications for the augmented estimator. As with (ref), we generate simple simulated data with $d = 10$. In the left panel, we plot the coefficients from lasso regression of $Y_p$ on $\Phi_p$ as a function of the lasso regularization parameter. The regularization path begins with the black dots, which represent the OLS coefficients. Each lasso coefficient (represented by a colored line) then shrinks linearly to exactly zero, due to the soft-thresholding operator. The middle panel plots the reweighted covariates using $\ell_\infty$ balancing weights between $\Phi_p$ and $\Phi_q$ solved in the constrained form. The black dots represent $\mkern1.5mu\overline{\mkern-1.5mu\Phi\mkern-1.5mu}\mkern 1.5mu_q$, corresponding to exact balance. Then as the weight regularization parameter increases, the reweighted covariates shrink linearly to exactly zero, just as in lasso. The right panel plots coefficients for the augmented estimator that combines a baseline outcome model fit $\hat{\beta}_\text{reg}^\lambda$ with $\ell_\infty$ balancing weights. The lines correspond to $\hat{\beta}_{\ell_\infty}$ as defined in (ref). The regularization path begins at the black dots, where $\hat{\beta}_{\ell_\infty} = \hat{\beta}_\text{ols}$, and eventually converges to $\hat{\beta}_\text{reg}^\lambda$, showing the usual soft-thresholding behavior. The order at which the coefficients go to zero reflects the size of $\mkern1.5mu\overline{\mkern-1.5mu\Phi\mkern-1.5mu}\mkern 1.5mu_q$, because the regularization path depends on the weight coefficients from the middle panel. Thus, the augmented estimator shrinks $\hat{\beta}_\text{ols}$ toward $\hat{\beta}_\text{reg}^\lambda$ but via a soft-thresholding operator applied to the feature shift, $\Delta_j$.
In the case where $\hat{\beta}_\text{reg}^\lambda$ is itself fit via lasso, as studied in chernozhukov2022automatic, then we recover a familiar double selection phenomenon belloni2014inference.
The lasso coefficients have a sparsity pattern generated by soft-thresholding the OLS coefficients. The augmented estimator then shrinks from OLS toward $\hat{\beta}_\text{reg}^\lambda$ by soft-thresholding the implied feature shift to the target features. As a result, wherever the lasso coefficients are non-zero or the weight coefficients are non-zero, the final augmented coefficients are also non-zero. The “included coefficients” for the final estimator are then the union of the coefficients included in either individual model. Therefore, augmenting a lasso outcome model with $\ell_\infty$ balancing also exhibits a form of undersmoothing in the $\ell_0$ “norm”, $\| \hat{\beta}_{\ell_\infty}\|_0$, in the sense that there are always at least as many non-zero coefficients as for the unaugmented lasso outcome model. However, this will not correspond to undersmoothing the base learner in the traditional sense, because in general there will not exist a lasso hyperparameter $\lambda$ that will produce sparsity pattern $I_\text{aug}$.
As noted by, for example, tang2023ultra, the double selection estimator may suffer from imprecision due to adjustment for covariates that are associated with treatment but not outcome. One could in principle remove covariates that are only predictive of the treatment, but this can jeopardize statistical inference. We refer to moosavi2023costs for a discussion on navigating this trade-off.
The results above are numerical: they hold without any statistical or causal assumptions. However, the connection between augmented estimators and outcome models also presents statistical insights that we discuss here. In particular, we leverage the numerical result that double (kernel) ridge regression --- which uses ridge regression for fitting both the outcome and weighting models --- is equivalent to a single, undersmoothed outcome ridge regression plug-in estimator.
First, we consider an asymptotic analysis in (ref): we use this equivalence to make explicit the connection between asymptotic results for augmented balancing weights with kernel ridge regression and prior results on optimal undersmoothing of a kernel ridge plug-in estimator. As a result, optimally undersmoothed kernel ridge regression inherits whatever guarantees can be proven for augmented ridge regression. An implication is that we can generalize the insight from robins2007comment that OLS is doubly robust to a wider class of non-parametric estimators. This equivalence also suggests an appropriate hyperparameter scheme when the outcome regression is an element of an RKHS.
Second, we consider a finite sample analysis in (ref): we use this equivalence to derive the finite-sample design-conditional mean squared error of augmented kernel ridge regression. We then use this expression to characterize finite-sample-optimal hyperparameter tuning. We turn to hyperparameter tuning in practice in the next section.
We now use our results in (ref) to make explicit the connection between two otherwise distinct sets of asymptotic results. First, wong2018kernel and singh2021debiased argue that double kernel ridge regression can deliver $\sqrt{n}$-consistent estimation of functionals in certain scenarios. wong2018kernel also proposes an optimally undersmoothed $\ell_2$ balancing weights estimator. Separately, hirshberg2019minimax and mou2023kernel propose optimally undersmoothed (single) kernel ridge outcome regression. Since, as we have shown in (ref) (see also Remark 2), these three procedures are equivalent, we can connect these results and show that plug-in estimators based on optimally undersmoothed kernel ridge regression or $\ell_2$ balancing weights can be $\sqrt{n}$-consistent. Moveover, results on RKHSs suggest a simple heuristic for hyperparameter choice. We give the high-level argument here and defer additional technical details to Appendix (ref).
To move from numerical results to statistical results, we must place some constraints on the data generating process. Assume that we observe $n$ iid samples of $(x_i,y_i,z_i)$ from $p$. Define $K \in \mathbb{R}^{n \times n}$ to be the kernel matrix with $i,j$-th entry $K_{ij} = k( (x_i, z_i), (x_j, z_j))$. Let $\sigma_j^2$ denote the eigenvalues of $K$. We assume that $\sigma_j^2 = \sigma^2>0$ is constant for all $j$; we can relax this at the cost of additional complexity. The “single” kernel ridge regression outcome regression estimator with parameter $\lambda$ has coefficient estimates: \[ \hat{\beta}_\text{ridge}^\lambda = (K + \lambda I)^{-1} y. \] Applying (ref), the augmented “double kernel ridge” estimator with hyperparameter $\delta$ is equivalent to a plug-in estimate for a new kernel ridge model: \[ \hat{\beta}_\text{aug} = (K + \gamma I)^{-1} y, \qquad \mbox{with} \;\; \gamma = \frac{\lambda \delta}{\sigma^2 + \lambda + \delta}. \] For statistical guarantees, we must typically allow the hyperparameters to change with $n$; let $\gamma_n$, $\lambda_n$ and $\delta_n$ then denote sequences of hyperparameters. In the doubly robust framework, one can choose $\lambda_n$ and $\delta_n$ in a way that is MSE-optimal for prediction purposes whilst ensuring that the bias of the augmented estimator is small. For two functions of $n$, $f_n$ and $g_n$, let $f_n \asymp g_n$ denote that $f_n = O(g_n)$ and $g_n = O(f_n)$. Then due to special properties of RKHS geometry, it follows that $\delta_n$ can be of the same order as $\lambda_n$, that is $\delta_n \asymp \lambda_n$ singh2020kernel. In the next section, we consider setting $\delta = \lambda$ for hyperparameter tuning in practice; our (ref) then implies that $\gamma_n\asymp \lambda_n^2$. We note more generally that (ref) implies that $\gamma_n\asymp \lambda_n \delta_n$.
There are two important cases to consider. When the RKHS is finite dimensional, the choice $\lambda_n = \delta_n = n^{-1/2}$ is optimal for controlling the prediction error for both the outcome and weighting models caponnetto2007optimal,singh2020kernel. The augmented estimator is then equivalent to a single ridge regression with hyperparameter $\gamma_n \asymp n^{-1}$, which matches the rate of hirshberg2019minimax,mou2023kernel. Hence, this approach will always undersmooth relative to the MSE-optimal hyperparameter for a single ridge regression.
When the RKHS is infinite-dimensional, we find that the undersmoothed hyperparameter implied by the augmented procedure can take on a range of asymptotic rates, both faster and slower than $n^{-1}$, depending on effective dimension and smoothness; we give concrete examples in the Appendix. This somewhat contrasts with the results in hirshberg2019minimax,mou2023kernel. In this sense, (ref) generalizes the standard undersmoothing arguments, which typically change the regularization schedule from $n^{-1/2}$ to $n^{-1}$.
We now use our numerical equivalences to write out the exact finite-sample mean squared error of the augmented kernel ridge estimator: by re-writing the augmented balancing weights estimator as a single outcome model, we can immediately leverage existing results from dobriban2018high.
Following their setup, we define the diagonal matrix $\hat{\Sigma} \coloneqq \tfrac{1}{n} \Phi_p^T\Phi_p$; if $\hat{\Sigma}$ is not diagonal, we can apply the rotation in (ref). We consider ridge regression with rescaled hyperparameter $\lambda$ and solution $(\hat{\Sigma} + \lambda I)^{-1} \Phi_p Y_p/n$; this is equivalent to standard ridge regression above with hyperparameter $n \lambda$, and also accommodates kernel ridge regression with appropriate choice of $\Phi_p$. Assume that $Y_p = \Phi_p \beta_0 + \epsilon$ with $\beta_0 \in \mathbb{R}^d$, and where $\epsilon \in \mathbb{R}^n$ are iid with mean zero and variance $\sigma^2$. Then the exact, design-conditional, squared bias and variance of the ridge regression prediction applied to a new iid sample $(\Phi_\text{new} , Y_\text{new}) \sim p$ are:
Applying (ref), we can simlarly derive the squared bias and variance of an augmented ridge estimator for our linear functional estimand; we denote these quantities $B_q^2$ and $V_q$ respectively. We express the bias and variance in terms of the two hyperparameters, $\lambda$ and $\delta$:
In the next section, we compare --- numerically and via simulation --- existing hyperparameter selection schemes to the optimal trade-off between $B_q^2$ and $V_q$. However, first we note that the analysis above opens up exciting new avenues for both theoretical and methodological work. One could theoretically analyze the mean squared error to understand how the optimal $\delta$ scales with the problem parameters; for example, by using proportionate asymptotics from random matrix theory as in the high-dimensional ridge regression literature hastie2022surprises. Second, our analysis here suggests a novel, more complex hyperparameter selection scheme based directly on the finite sample analysis. We leave this to future work.
This section illustrates our results in practice. We first explore hyperparameter tuning for double ridge regression, comparing practical methods to the optimal hyperparameter computed using our results from (ref). Following our asymptotic results in (ref), we recommend equating the weighting and outcome model hyperparameters in practice. We then apply both double ridge and lasso-augmented $\ell_\infty$-balancing to two versions of the canonical lalonde1986evaluating application. An important theme throughout is that some approaches for hyperparameter selection frequently lead to $\delta = 0$, which collapses the augmented estimate to OLS alone --- even in settings where this is far from optimal. Overall, we take this as a warning that existing hyperparameter tuning schemes can be potentially misleading when applied naively.
We begin with practical hyperparameter tuning for the special case of double ridge, building on the MSE expression in Section (ref). There is an active literature on selecting hyperparameters for augmented balancing weights estimators and double machine learning estimators more broadly kallus2020generalized, wang2020minimal, ben2021balancing, bach2024hyperparameter. We contribute to this literature by comparing practical hyperparameter tuning schemes with an oracle hyperparameter tuning scheme based on (ref).
Reflecting empirical practice, we focus here on choosing hyperparameters sequentially: we first select the outcome model hyperparameter $\lambda$ (e.g. by cross-validation) and then select the weighting model hyperparameter $\delta$. Ultimately, we find strong performance for both CV imbalance and CV outcome hyperparameters, as defined below. We especially recommend the latter as a reasonable starting point in practice. In additional to theoretical support from our asymptotic analysis, the outcome model hyperparameter scheme does not require any additional algorithm or code after having fit the initial outcome model.
\paragraph{Oracle hyperparameter.}
To compute oracle hyperparameters, we first compute the prediction-MSE-optimal $\lambda$ using the standard ridge regression MSE expression, and then we use (ref) to compute the corresponding optimal $\delta$ for the linear functional estimand:
While there is not a closed form for $\delta^*$, we can nonetheless directly compute this optimal hyperparameter and characterize its behavior under a range of scenarios. We draw several conclusions about optimal $\delta^*$ for a wide range of DGPs of the form $Y_p = \Phi_p \beta_0 + \epsilon$. First, $\delta^*$ is generally increasing in the noise, $\sigma^2$: larger $\sigma^2$ typically implies larger $\delta^*$. Second, $\delta^*$ generally depends on the target mean, $\mathbb{E}[\Phi_q]$; that is, two DGPs that are identical except for $\mathbb{E}[\Phi_q]$ can have different values of $\delta^*$. The optimal hyperparameter, however, does not depend on the magnitude of the shift in the target mean: replacing $\mathbb{E}[\Phi_q]$ with $c \mathbb{E}[\Phi_q]$ for $c \neq 0$, scales both the bias and variance by $c^2$, leaving $\delta^*$ unchanged.
\paragraph{Practical hyperparameter.} We compare the oracle hyperparameter with three implementable practical proposals. In all cases, we first pick $\lambda$ by cross-validating the mean squared error of a ridge outcome model.
Before presenting simulation results, we provide a preliminary analytic discussion, comparing these practical schemes to the behavior of the oracle $\delta^*$. For the first two proposals: just like the oracle, both depend on the target mean $\mathbb{E}[\Phi_q]$ and are invariant to re-scaling. However, these two approaches are mechanically independent of the outcomes $Y_p$, unlike the oracle $\delta^*$ which, in general, depends on the variance of the outcomes. On the other hand, the last proposal depends on the outcomes $Y_p$ but is mechanically independent of $\mathbb{E}[\Phi_q]$.
This suggests that any one of these tuning parameter approaches cannot perform well across all DGPs. In future work, if we pursue a theoretical analysis of the oracle hyperparameter, e.g. in a proportionate asymptotics framework, we may be able predict when either the outcomes or the covariate shift is more important. In this work we begin by demonstrating that no one tuning scheme does uniformly best in simulations.
To assess the behavior of these hyperparameter tuning schemes, we conduct an extensive simulation study using 36 distinct data-generating processes, 30 synthetic and 6 semi-synthetic; see (ref) for a detailed discussion. For each DGP, we directly compute the oracle hyperparameter using the results in Section (ref). We then compute values from the three practical hyperparameter tuning methods discussed above. The mean squared error that we consider is design-conditional, and so we draw samples of the covariates for each DGP only once.
Table (ref) presents a summary of the MSE for the three methods across the 36 DGPs. Overall, we find that the CV outcome approach of choosing $\delta = \lambda$ and the CV imbalance approach both perform well in practice: these two achieve the lowest MSE in 35 of the 36 DGPs, with CV imbalance performing slightly better on average. By contrast, selecting $\delta$ via CV for the Riesz loss has numerical stability problems that compromises performance. The performance for the outcome and balance approaches, on the other hand, seem to degrade gracefully and rarely perform catastrophically. Taken together, these preliminary findings suggest researchers should begin with these two tuning methods as defaults.
\paragraph{Recovering the OLS point estimate.} As we discuss above (see, e.g., (ref)), when $\delta = 0$ the point estimate for the augmented balancing weights estimator is numerically identical to the OLS point estimate. Thus, when a hyperparameter tuning procedure chooses $\delta = 0$ in practice, researchers are simply estimating the equivalent of OLS --- even if they are unaware they are doing so. This is especially problematic in settings where OLS is far from optimal kobak2020optimal, hastie2022surprises. In our synthetic and semi-synthetic DGPs, $\delta = 0$ is never optimal, and is usually associated with a very large error driven by extreme variance --- see for example, (ref) in the Appendix. Thus the fact that hyperparameter tuning procedures can return $\delta = 0$ in these DGPs represents a pathological case.
In our simulation study, we find that, when cross validating the Riesz loss, over half of all draws returned $\delta = 0$. By contrast, none of the other methods returned $\delta = 0$ in the synthetic DGPs, though, as we discuss below, we do observe exact zeros for $\delta$ occasionally when cross-validating imbalance in the standard LaLonde dataset. This further highlights the numeric instability of hyperparameter tuning via CV for the Riesz loss, at least in the settings we consider here. We further suggest that in these cases, practitioners assess the sensitivity of the $\delta = 0$ results to the particular tuning procedure used or to the random choice of cross-validation splits.
We now illustrate our equivalence and hyperparameter tuning results on real-world datasets. Following chernozhukov2022automatic, we focus on the canonical lalonde1986evaluating data set evaluating a job training program in the National Supported Work (NSW) Demonstration. The primary outcome of interest is annual earnings in 1978 dollars.
For these illustrations, we estimate the Average Treatment Effect on the Treated (ATT), $\mathbb{E}[Y(1) - Y(0) \mid Z = 1]$. We recover the missing conditional mean $\mathbb{E}[Y(0) \mid Z = 1]$ using the setup from Example (ref) in Appendix (ref), where the source and target populations are the control and treated units respectively. Thus $\Phi_p$ and $\Phi_q$ correspond to the feature expansion $\phi(X)$ applied to the covariates in the control group and treated group respectively. We consider two different features expansions of the original covariates: (1) a “short” set of 11 covariates used in dehejia1999causal;\footnote{These are: age, years of education, Black indicator, Hispanic indicator, married indicator, 1974 earnings, 1975 earnings, age squared, years of education squared, 1974 earnings squared, and 1975 earnings squared.} and (2) an expanded, “long” set of 171 interacted features used in farrell2015robust.
Our goal is to explicate how augmented estimators under different hyperparameter tuning schemes undersmooth in practice in both low and high-dimensional settings. In some cases, the augmented estimator collapses to exactly OLS as we document above. (ref) contains extensive additional analyses, including dataset summaries, additional results from the Infant Health Development Program (IHDP), and sensitivity of these numerical results to cross-fitting.
Following chernozhukov2022automatic, we first consider the expanded set of 171 features for lalonde1986evaluating used in farrell2015robust. (ref) shows estimates for ridge-augmented $\ell_2$ balancing (top row) and lasso-augmented $\ell_\infty$ balancing (bottom row). We explicitly characterize these results in terms of undersmoothing in (ref). The left two panels of each row show the cross-validation curves for the outcome regression and balancing weights, respectively. The right panels show the point estimate as a function of the weighting hyperparamter $\delta$, holding the outcome model hyperparameter $\lambda$ fixed; the black triangle represents the OLS plug-in point estimate. For context, the corresponding experimental estimate is \$1,794 dehejia1999causal. The green and red dotted lines correspond to hyperparameters chosen by cross-validating balance and the Riesz loss, respectively. For the double ridge estimate, the purple line corresponds to $\delta = \hat{\lambda}$, the outcome hyperparameter selected via cross validation.
(ref) highlights that both the imbalance and the point estimate are highly nonlinear close to zero. Thus, even small departures from OLS (at $\delta = 0$) lead to large changes in the point estimate --- in (ref) we give some suggestive evidence that the variance blows up relative to the bias in this range. We can also assess the sensitivity of the point estimate to the hyperparameter selection scheme. In this case, choosing $\delta$ via CV balance leads to meaningfully larger choices than via other methods.
Finally, the selected $\delta$ is always strictly greater than zero for this high-dimensional dataset. However, we find this is sensitive to small perturbations in the problem parameters. For example, when we perturb $\mathbb{E}[\Phi_q]$ by adding a small value to all the even elements, then the cross-validated $\ell_2$ Riesz loss chooses $\delta = 0$ in 38% of draws of the cross-validation splits. As suggested by (ref) and our simulation results, this is likely to result in extremely large mean squared error.
Finally, we apply double ridge to the “short” version of the lalonde1986evaluating data set with 11 features. (ref) shows the cross-validation curves for the outcome and weighting models, as well as the point estimate as a function of the balance hyperparameter, with the OLS estimate given by the black triangle. As above, the green, red, and purple dotted lines correspond to hyperparameters chosen by cross-validating balance, cross-validating the Riesz loss, and choosing $\delta=\lambda$ respectively.
Unlike for the “long” dataset in Figure (ref), (ref) does not display quite as stark nonlinearity around zero. Importantly, however, setting $\delta$ by cross-validating imbalance or the Riesz loss yields $\delta = 0$ (up to numerical imprecision), which reduces the augmented estimator to exactly the estimate from a simple OLS regression --- even though the base learner ridge outcome model is heavily regularized. By contrast, our preferred hyperparameter tuning scheme of choosing $\delta = \lambda$ results in an estimate that is roughly \$400 dollars smaller than the OLS estimate.
We have shown that augmenting a plug-in regression estimator with linear balancing weights results in a new plug-in estimator with coefficients that are shrunk towards --- in some cases all the way to --- the estimates from OLS fit on the same observations. We generalize this equivalence for different choices of outcome and weighting regressions. In the asymptotic setting, we draw the explicit connection between augmented estimators and undersmoothing for the special case of kernel ridge regression. Then we derive the design-conditional finite sample MSE for the double ridge estimator, and use it to solve numerically for oracle hyperparameters. We compare the oracle hyperparameters with three practical tuning schemes and then illustrate our results on the canonical LaLonde data set. In the Appendix also explore many extensions, including to nonlinear weights and to high-dimensional features.
There are many promising avenues for future research. The fundamental connection between doubly robust estimation and undersmoothing opens up several theory directions. While we focus on the special case of kernel ridge regression in Section (ref), we anticipate that these connections will hold more broadly. Similarly, while our focus in this paper has been on interpreting balancing weights as a form of linear regression, the converse is also valid: we could instead focus on how many outcome regression-based plug-in estimators are, in fact, a form of balancing weights; see lin2022regression for connections between outcome modeling and density ratio estimation.
We also anticipate that the MSE we derive in (ref) is a starting place for future theoretical analysis that can inform practice. We demonstrate in our simulation study that existing hyperparameter selection methods cannot perform uniformly well over all DGPs. We expect that analyzing the optimal hyperparameters, for example in a proportionate asymptotics regime, can either help devise new tuning schemes or inform which tuning method will work best on the dataset at hand.
We conjecture that these results may provide new insights into the estimation of causal effects in the proximal causal inference framework tchetgen2020introduction. This framework uses proxy variables to identify causal effects in the presence of unmeasured confounding. Estimation has been complicated by the fact that, in the absence of strong parametic assumptions, estimators of proximal causal effects are solutions to ill-posed Fredholm integral equations. ghassami2022minimax and kallus2021causal recently proposed tractable nonparametric estimators in this setting. They use an “adversarial” version of double kernel ridge regression --- allowing the weighting and outcome models to have different bases --- to estimate the solution to the required Fredholm integral equations. Our results apply immediately to standard augmented estimators with different bases for the outcome and weighting models, either via a union basis chernozhukov2022automatic or by applying an appropriate projection as in hirshberg2021augmented, and extending these results to proximal causal effect estimators might help in constructing new proximal balancing weights, matching, or regression estimators with attractive asymptotic properties.
Finally, many common panel data estimators are forms of augmented balancing weight estimation abadie2010synthetic, ben2021augmented, arkhangelsky2021synthetic. We plan to use the numeric results here, especially the results for simplex-constrained weights in Appendix (ref), to better understand connections between methods and to inform inference.