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.
49,280 characters · 22 sections · 77 citation commands
Direct Bias-Correction Term Estimation for Average Treatment Effect Estimation
We consider the problem of estimating the average treatment effect (ATE) in causal inference Imbens2015causalinference. Methods for estimating ATEs are typically designed to eliminate bias arising from treatment assignment and the estimation of nuisance parameters, aiming for (asymptotic) unbiasedness and efficiency.
We begin by formulating the problem. There are two treatments, denoted by $1$ and $0$.\footnote{In some cases, only treatment $1$ is referred to as the treatment, while treatment $0$ is referred to as the control. For simplicity, we refer to them as treatment $1$ and treatment $0$ throughout this study.} For each treatment $d \in \{1, 0\}$, let $Y(d) \in {\mathbb{R}}$ denote the potential outcome under treatment $d$. The treatment assignment indicator is denoted by $D \in \{1, 0\}$, and the observed outcome is given by $Y = \mathbbm{1}[D = 1] Y(1) + \mathbbm{1}[D = 0] Y(0)$, meaning that we observe $Y(d)$ only if the unit is actually assigned to treatment $d$. Each unit is characterized by $K$-dimensional covariates $X \in \mathcal{X} \subset {\mathbb{R}}^K$, where $\mathcal{X}$ denotes the covariate space. For $n$ units indexed by $1, 2, \dots, n$, let $\mathcal{D} \coloneqq \{(X_i, D_i, Y_i)\}_{i=1}^{n}$ denote the observed data, where each $(X_i, D_i, Y_i)$ is an i.i.d. copy of $(X, D, Y)$ generated from an underlying distribution $P_0$. Our goal is to estimate the ATE, defined as \[ \tau_0 \coloneqq {\mathbb{E}}\big[Y(1) - Y(0)\big], \] where the expectation is taken over the distribution $P_0$. Note that we can also apply our method for the ATE for the treated group (ATT). For the details about ATT estimation, see Appendix (ref).
Let $e_0(X) = P_0(D = 1 \mid X)$ denote the probability of assigning treatment $1$ given covariates $X$, which is known as the propensity score. Throughout this study, we impose the following conditions, commonly referred to as the unconfoundedness and common support assumptions.
When $e_0(x)$ is not constant, a distributional shift arises between the observed outcomes in the treatment and control groups, denoted by $\mathcal{G}_1$ and $\mathcal{G}_0$, respectively, where $\mathcal{G}_d \coloneqq \{i \in \{1, 2, \dots, n\} \colon D_i = d\}$. This shift induces bias in the sample mean, $ \frac{1}{|\mathcal{G}_d|} \sum_{i \in \mathcal{G}_d} Y_i = \frac{1}{|\mathcal{G}_d|} \sum_{i \in \mathcal{G}_d} Y_i(d)$, which deviates from ${\mathbb{E}}\left[Y(d)\right]$ and thus prevents the sample mean difference, $ \frac{1}{|\mathcal{G}_1|} \sum_{i \in \mathcal{G}_1} Y_i - \frac{1}{|\mathcal{G}_0|} \sum_{i \in \mathcal{G}_0} Y_i $, from being an unbiased estimator of the ATE.
To address this issue, several debiased estimators have been proposed under standard regularity conditions. In this section, we introduce two representative estimators, the inverse probability weighting (IPW) estimator and the augmented IPW (AIPW) estimator, as follows:
Bias-correction term. In both estimators, the term \[ h_0(D, X) \coloneqq h(D, X) \coloneqq \frac{\mathbbm{1}[D = 1]}{e_0(X)} - \frac{\mathbbm{1}[D = 0]}{1 - e_0(X)} \] is crucial. This term, referred to as the bias-correction term, is central to ATE estimation Schuler2024introductionmodern. A common approach is to estimate $e_0$ using logistic regression and then plug the resulting estimate $\widehat{e}^{{\mathrm{L}}}_n$ into $h$. Note that the bias-correction term is also referred to as the Riesz representer Chernozhukov2021automaticdebiased or the clever covariates vanderLaan2006targetedmaximum. We use the term bias-correction term because the Riesz representer is closely connected to the automatic debiased machine learning literature, and the clever covariates is closely connected to the targeted maximum likelihood estimation (TMLE) literature.
For example, in a typical one-step bias correction, we first construct an ATE estimator as $\widehat{\tau}^{\mathrm{DM}}_n \coloneqq \frac{1}{n}\sum^n_{i=1}\left(\widehat{\mu}_n(1, X) - \widehat{\mu}_n(0, X)\right)$, where $\widehat{\mu}_n$ is an estimator of $\mu_0$. This estimator is known as the direct method (DM) or naive plug-in estimator. To obtain an efficient estimator, we add the bias-correction term $\frac{1}{n}\sum^n_{i=1}h_0(D_i, X_i)(Y_i - \widehat{\mu}_n(D_i, X_i))$ to the first-stage DM estimator $\widehat{\tau}^{\mathrm{DM}}_n$, yielding the AIPW estimator.
In this study, we propose a method to estimate the bias-correction term, also called the Riesz representer or the clever covariates. For example, we can estimate the bias-correction term by estimating the propensity score $e_0$ using the maximum likelihood estimation. However, our interest is not in propensity score estimation but in bias-correction term estimation. As the well-known Vapnik principle states, we should avoid such an intermediate problem and ideally aim to estimate the target objective in a more direct manner Vapnik1998statisticallearning. Following this principle, this study considers estimating $h_0(D, X)$ by directly minimizing the estimation error for the true $h_0(D, X)$.
The technical challenge is that the target objective $h_0$ is unknown. To address this issue, we employ techniques developed in the direct density-ratio estimation (DRE) literature Sugiyama2012densityratio. In direct DRE, the goal is to minimize the empirical risk between the true density ratio and its model, even though the true density ratio is unknown. It is known that empirical risk minimization is feasible even without knowledge of the true propensity score. Since the inverse propensity score can be viewed as a density ratio, we can extend these existing methods to our setting. For causal inference researchers who are unfamiliar with DRE, we review the DRE literature in Appendix (ref).
Our motivation is also closely aligned with studies on Riesz regression Chernozhukov2021automaticdebiased and covariate balancing weights Imai2013estimatingheterogeneity,Deville1992calibrationestimators, which also aim to estimate the bias-correction term in a direct manner. Studies in covariate balancing focus on the balancing property of propensity score estimator and estimate them using the property. Chernozhukov2021automaticdebiased proposes Riesz regression which represents the bias-correction term as the Riesz representer. Although the derivation process is different, we derive the objective function that is the same as Chernozhukov2021automaticdebiased by using the DRE techniques. Further, we generalize our objective by using the Bregman divergence as well as DRE in Sugiyama2011densityratio. From this generalization, we further connect our approach to the covariate balancing by showing the equivalence between our objective and empirical balancing through the duality arguments discussed in Zhao2019covariatebalancing and BrunsSmith2025augmentedbalancing.
This study has the following four contributions: (i) a general framework for directly estimating the bias-correction term (also called the Riesz representer or clever covariates) via Bregman divergence minimization; (ii) our proposed framework includes Riesz regression in Chernozhukov2021automaticdebiased and the tailored loss in Zhao2019covariatebalancing as special cases; (iii) under our framework, we show that there are appropriate choices of bias-correction term models and Bregman divergences under which covariate balancing is automatically realized as the dual of the Bregman divergence minimization problem (automatic covariate balancing); (iv) we provide a theoretical analysis of the estimator.
Our first contribution is the proposal of a framework for direct bias-correction term estimation via Bregman divergence minimization. We estimate the bias-correction term by directly minimizing the estimation error of the true bias-correction function $h_0$, measured by the Bregman divergence, $\mathrm{BR}^\dagger_g\big(h_0\mid h\big) \coloneqq {\mathbb{E}}\Big[g\big(h_0(D, X)\big) - g\big(h(D, X)\big) - \partial g\big(h(d, X)\big) \big(h_0(D, X) - h(D, X)\big)\Big]$, where $g$ is a differentiable and strictly convex function. By changing $g$, we can measure the error using various metrics, such as the squared loss or KL divergence loss. Since the Bregman divergence involves the unknown function $h_0$, direct optimization is infeasible. To address this issue, we propose minimizing an alternative objective function, defined as $\mathrm{BR}_g\big(h\big) \coloneqq {\mathbb{E}}\Big[ - g\big(h(D, X)\big) + \partial g\big(h(D, X)\big) h(D, X) - \partial g\big(h(1, X)\big) - \partial g\big(h(0, X)\big)\Big]$. Minimizing the original Bregman divergence $\mathrm{BR}^\dagger_g\big(h\big)$ is equivalent to minimizing $\mathrm{BR}_g\big(h\big)$, which does not depend on the unknown function. That is, we establish the equivalence: $h^* \coloneqq \operatorname*{arg\,min}_{h \in{\mathcal{H}}} \mathrm{BR}^\dagger_g\big(h_0\mid h\big) = \operatorname*{arg\,min}_{h \in{\mathcal{H}}} \mathrm{BR}_g\big(h\big)$. The resulting objective function can then be approximated using an empirical risk function.
Our second contribution is the unification of existing literature. Our proposed Bregman divergence minimization objective includes Riesz regression in Chernozhukov2021automaticdebiased (when using the squared loss) and the tailored loss in Zhao2019covariatebalancing (when using the KL divergence loss). Furthermore, our framework also integrates covariate balancing methods Imai2013covariatebalancing,Hainmueller2012entropybalancing,Zubizarreta2015stableweights,Chan2015globallyefficient,Wong2017kernelbased. If we use linear models to approximate the bias-correction term and train the model with the squared loss (Riesz regression), the dual problem coincides with the optimization problem in stable balancing weights. If we model the bias-correction term via the propensity score with logistic models and train the model with the KL divergence loss (tailored loss), the dual problem becomes the same as the optimization problem in entropy balancing weights. Kato2025directdebiased, a subsequent work of this study, refers to this property as automatic covariate balancing. See Table (ref) in Section (ref) and Figure (ref) in Appendix.
Our third main contribution is the theoretical analysis of the estimator obtained via direct bias-correction term estimation. Since we estimate $r_0$ using empirical risk minimization, we establish bounds on the estimation error using empirical process theory. Furthermore, we present examples of ATE estimators that incorporate the bias-correction term estimated using our framework and conduct simulation studies. Using standard ATE estimation techniques, we demonstrate that our method yields a $\sqrt{n}$-consistent ATE estimator.
As a side product of our contributions, we find that we can import various existing results from the DRE literature. Since Riesz regression is essentially the same as LSIF, various results about convergence rate analysis and optimization methods have already been established. For example, Kanamori2012statisticalanalysis shows the convergence rate when using a reproducing kernel hilbert space (RKHS) for the density ratio, or equivalently the bias-correction term. Kato2021nonnegativebregman shows the rate when using neural networks, which has been further refined in Zheng2022anerror. Rhodes2020telescopingdensiyratio and Kato2021nonnegativebregman point out the overfitting problem characteristic of DRE estimation and propose techniques to avoid the problem. Lin2023estimationbased finds that nearest neighbor matching can be interpreted as density ratio estimation, and it can also be interpreted as a special case of LSIF or Riesz regression (See Appendix (ref)). These findings not only help deepen our understanding of Riesz regression, but also prevent unnecessary reinvention. For example, the covariate adaption method proposed in Chernozhukov2025automaticdebiased uses Riesz regression, but it is essentially the same as covariate adaption with a density ratio estimated via LSIF Kanamori2009aleastsquares, except for the regression adjustment. While Chernozhukov2022riesznet proposes neural networks and random forests for Riesz regression, the techniques for estimating the density ratio have also been proposed in the DRE literature Kanamori2012statisticalanalysis,Abe2019,Rhodes2020telescopingdensiyratio,Kato2021nonnegativebregman.
In this study, we consider estimating $h_0$ by minimizing the empirical risk associated with the Bregman divergence between $h_0$ and its estimator $h\colon \{1, 0\} \times \mathcal{X} \to {\mathbb{R}}$.
Let $g\colon {\mathbb{R}} \to {\mathbb{R}}$ be a differentiable and strictly convex function. Given $d\in\{1, 0\}$, we define the Bregman divergence between $h_0$ and $h$ as $\mathrm{br}^\dagger_g\big(h_0(d, x)\mid h(d, x)\big) \coloneqq g\big(h_0(d, x)\big) - g\big(h(d, x)\big) - \partial g\big(h(d, x)\big) \big(h_0(d, x) - h(d, x)\big)$, where $\partial g$ denotes the derivative of $g$. Then, we define the average Bregman divergence as $\mathrm{BR}^\dagger_g\big(h_0\mid h\big) \coloneqq {\mathbb{E}}\Big[g\big(h_0(D, X)\big) - g\big(h(D, X)\big) - \partial g\big(h(d, X)\big) \big(h_0(D, X) - h(D, X)\big)\Big]$. Then, we estimate $h_0$ by $h^* = \operatorname*{arg\,min}_{h\in {\mathcal{H}}} \mathrm{BR}^\dagger_g\big(h_0\mid h\big)$. By dropping the term that is irrelevant to learning, we have \[h^* = \operatorname*{arg\,min}_{h\in {\mathcal{H}}} \mathrm{BR}_g\big(h\big),\] \[\text{where}\quad \mathrm{BR}_g\big(h\big) \coloneqq {\mathbb{E}}\Big[ - g\big(h(D, X)\big) + \partial g\big(h(D, X)\big) h(D, X) - \partial g\big(h(1, X)\big) + \partial g\big(h(0, X)\big)\Big].\] This can be shown as follows:
Here, we dropped terms irrelevant to the optimization and used ${\mathbb{E}}[\mathbbm{1}[D = 1]h_0(1, X)\mid X] = {\mathbb{E}}[e_0(X)h_0(1, X)\mid X] = 1$ and ${\mathbb{E}}[\mathbbm{1}[D = 0]h_0(0, X)\mid X] = -1$.
Thus, surprisingly, we demonstrate that the least squares estimate for the unknown true bias-correction term $h_0$ can be defined by an objective function that does not explicitly include $h_0$ itself. As discussed in the following subsection, this objective function can be easily approximated using observations.
Then, we estimate the bias-correction term $h_0$ by minimizing an empirical Bregman divergence as \[ \widehat{h}_n \coloneqq \operatorname*{arg\,min}_{h \in {\mathcal{H}}}\widehat{\mathrm{BR}}_g\big(h\big) + \lambda J(h), \] where $J(h)$ is some regularization function and \[\widehat{\mathrm{BR}}_g(h) \coloneqq \frac{1}{n}\sum^n_{i=1}\Big( - g(h(D_i, X_i)) + \partial g(h(D_i, X_i)) h(D_i, X_i) - \partial g(h(1, X_i) + \partial g(h(0, X_i))\Big).\]
By changing $g$, we can obtain various loss functions for estimating the bias-correction term, as shown in the subsequent subsections. In particular, if we use the squared loss in the Bregman divergence, we obtain Riesz regression in Chernozhukov2021automaticdebiased, which is originally called Least-Squares Importance Fitting (LSIF) in the DRE literature Kanamori2009aleastsquares. Note that kernel mean matching by Gretton2009covariateshift is also the same as, or a variant of, LSIF. If we use the KL divergence, we obtain the tailored loss in Zhao2019covariatebalancing, which is originally called KLIEP in the DRE literature Sugiyama2008directimportance. Furthermore, as we discuss in Section (ref), if we use linear models for $h_0$ and train them with the squared loss, the covariate balancing property is automatically obtained, as shown in BrunsSmith2025augmentedbalancing. If we model $h_0$ using the propensity score $e_0$ approximated via logistic models and train it with the tailored loss, the covariate balancing property is automatically obtained, as shown in Zhao2019covariatebalancing. We demonstrate the correspondence of the existing methods in Table (ref). Also see Figure (ref) in Appendix for the relationship among bias-correction term estimation via Bregman divergence minimization, density ratio estimation, and covariate balancing, summarized in Kato2025directdebiased and Kato2025unifiedtheory.
Our least squares method for direct bias-correction term estimation can be obtained by using a squared loss $g^{\text{SL}}(h) = (h-1)^2$. By substituting this function into the Bregman divergence, we formulate the estimation problem as $h^* \coloneqq \operatorname*{arg\,min}_{h \in {\mathcal{H}}}\mathrm{BR}_{g^{\text{SL}}}\big(h\big)$, where \[\mathrm{BR}_{g^{\text{SL}}}\big(h\big) = {\mathbb{E}}\Big[- 2\big(h(1, X) - h(0, X)\big) + h(D, X)^2\Big].\] Then, we estimate the bias-correction term as $\widehat{h}_n \coloneqq \operatorname*{arg\,min}_{h \in {\mathcal{H}}}\widehat{\mathrm{BR}}_{g^{\mathrm{SL}}}\big(h\big) + \lambda J(h)$, where $ \widehat{\mathrm{BR}}_{g^{\mathrm{SL}}}(h) = \frac{1}{n}\sum^n_{i=1}\big(- 2\big(h(1, X_i) - h(0, X_i)\big) + h(D_i, X_i)^2\big)^2}$. This objective function is the same as the one used in Chernozhukov2021automaticdebiased. This type of estimation method is referred to as LSIF in density-ratio estimation Kanamori2009aleastsquares.
Consider $g^{\mathrm{KL}}(h) = |h|\log |h| - |h|$, which is a convex function. By substituting this function into the Bregman divergence, we formulate the estimation problem as $h^* \coloneqq \operatorname*{arg\,min}_{h \in {\mathcal{H}}}\mathrm{BR}_{g^{\mathrm{KL}}}\big(h\big)$, where \[\mathrm{BR}_{g^{\mathrm{KL}}}\big(h\big) \coloneqq {\mathbb{E}}\Big[|h(D_i, X_i)| - \log(|h(1, X)|) - \log(|h(0, X)|)\Big].\]
Then, we estimate the bias-correction term as $\widehat{h}_n \coloneqq \operatorname*{arg\,min}_{h \in {\mathcal{H}}}\widehat{\mathrm{BR}}_{g^{\mathrm{KL}}}\big(h\big) + \lambda J(h)$, where $ \widehat{\mathrm{BR}}_{g^{\mathrm{KL}}}(h) = \frac{1}{n}\sum^n_{i=1}\big(|h(D_i, X_i)| -\log \big(|h(1, X_i)|\big) - \log \big(|h(0, X_i)|\big)\big)\big)}$. This estimation method corresponds to unnormalized Kullback–Leibler (UKL) minimization in DRE Nguyen2010estimatingdivergence, which generalizes the KL importance estimation procedure (KLIEP). Also see Appendix (ref).
Next, as a variant of the KL divergence loss, we propose the tailored loss. Let us redefine a model ${\mathcal{H}}$ as a set of functions $h(1, \cdot)\colon {\mathcal{X}} \to (1, \infty)$ and $h(0, \cdot)\colon {\mathcal{X}} \to (-1, -\infty)$; that is, we restrict the space of $h$. This restriction is justified from the form of $h_0$ and the common support assumption. Let us consider $g^{\text{TL}}(h) = (|h|-1)\log\left(|h| - 1\right) - |h|$. By substituting this function, we obtain \[\mathrm{BR}_{g^{\text{TL}}}\big(h\big) \coloneqq {\mathbb{E}}\Big[\log\left(|h(D, X)| - 1\right) + |h(D, X)| - \log\left(|h(1, X)| - 1\right) - \log\left(|h(0, X)| - 1\right)\Big].\] Note that it holds that $\mathrm{BR}_{g^{\text{TL}}}\big(h\big) \coloneqq {\mathbb{E}}\Big[- \mathbbm{1}[D = 0]\log\left(|h(1, X)| - 1\right) - \mathbbm{1}[D = 1]\log\left(|h(0, X)| - 1\right) + \mathbbm{1}[D = 1]h(1, X) - \mathbbm{1}[D = 0]h(0, X)\Big]$. Then, we estimate the bias-correction term as $ \widehat{h}_n \coloneqq \operatorname*{arg\,min}_{h \in {\mathcal{H}}}\widehat{\mathrm{BR}}_{g^{\text{TL}}}\big(h\big) $, where the empirical Bregman divergence becomes $ \widehat{\mathrm{BR}}_{g^{\text{TL}}}(h) = \frac{1}{n}\sum^n_{i=1}\big(\mathbbm{1}[D_i = 0]\log \big(|h(1, X_i)| - 1\big) + \mathbbm{1}[D_i = 1]\log \big(|h(0, X_i)| - 1\big) + \mathbbm{1}[D_i = 1]|h(1, X_i)| - \mathbbm{1}[D_i = 0]|h(0, X_i)|\big)X_i)|} $.
Under specific choices of Riesz regression models and Bregman divergence, we can automatically enforce the covariate balancing property. The key tool is the duality relationship between the Bregman divergence minimization problem and the covariate balancing optimization problem. This result is shown in Kato2025directdebiased, and we introduce the result for reference.
Consider a linear model \[h_\beta(D, X) = \Phi(D, X)^\top \beta,\] where $\Phi \colon \{1, 0\} \times {\mathcal{X}} \to {\mathbb{R}}^p$ is a basis function. For this model, using the squared loss (Riesz regression) automatically achieves covariate balancing, as discussed in BrunsSmith2025augmentedbalancing.
Specifically, under linear models, by duality, this MSE minimization problem is equivalent to solving
where $\bm{0}_p$ is the $p$-dimensional zero vector. This optimization problem matches that used to obtain stable weights Zubizarreta2015stableweights.
It enforces the covariate balancing condition $\sum^n_{i= 1}\widehat{w}_i \Phi(D_i, X_i) - \left(\sum^n_{i=1}\Big(\Phi(1, X_i) - \Phi(0, X_i)\Big)\right) = \bm{0}_p$, where $\widehat{w}_i = \Phi(D_i, X_i)^\top \widehat{\beta}$.
Another advantage of using linear models is that we can write the entire ATE estimation with a single linear model, as shown by BrunsSmith2025augmentedbalancing.
We can model the Riesz representer by modeling the propensity score as \[h_\beta(D, X) = \mathbbm{1}[D = 1]r_\beta(1, X) - \mathbbm{1}[D = 0]r_\beta(0, X),\] where $r_\beta(1, X) = \frac{1}{e_{\beta}(X)}$, $r_\beta(0, X) = \frac{1}{1 - e_{\beta}(X)}$, $e_{\beta}(X) \coloneqq \frac{1}{1 + \exp\big(-\beta^\top \Phi(X)\big)}$, and $\Phi \colon {\mathcal{X}} \to {\mathbb{R}}^p$ is a basis function. Note that we do not include $D$, unlike the basis function used in linear models. For this model, if we use the KL-divergence–flavored convex function defined in Section (ref), which corresponds to the tailored loss in Zhao2019covariatebalancing, we automatically achieve covariate balancing.
Define $\widehat{\beta} \coloneqq \operatorname*{arg\,min}_{\beta}\frac{1}{n}\sum^n_{i=1}\sum_{d\in\{1, 0\}}\Big(\mathbbm{1}[D_i = d]\left( - \log \left(\frac{1}{r_\beta(d, X_i) - 1}\right) + r_\beta(d, X_i)\right)(d, X_i)\Big)}$, and denote $r_{\widehat{\beta}}$ by $\widehat{r}$. Under logistic models, by duality, the KL divergence-flavored loss is equivalent to solving
This optimization problem matches that used in entropy balancing Hainmueller2012entropybalancing. Note that this objective function is derived from $\widehat{\mathrm{BR}}_{g^{\text{TL}}}(h)$ when we use the logistic model specified in this section.
As a result, we obtain $\sum^n_{i= 1}\Big(\mathbbm{1}[D_i = 1]\widehat{w}_i \Phi(X_i) - \mathbbm{1}[D_i = 0]\widehat{w}_i \Phi(X_i)\Big) = \bm{0}_p$, where $\widehat{w}_i = \widehat{r}(X_i).$
This model has the advantage that we can use a basis function $\Phi(X)$ independent of $D$. Moreover, it naturally achieves covariate balance in the sense that the covariate distributions match between the treated and control groups. Additionally, it allows us to automatically impose nonnegativity on $h(1, X)$ and $- h(0, X)$, which may be violated in linear models. Note that $ h_0(1, X) = \frac{1}{e(X)}$ and $ h_0(1, X) = \frac{1}{1 - e(X)}$.
We first discuss the advantages of using logistic models over linear models. One benefit of using logistic models is that we can simplify the basis function by making it independent of $D$. Furthermore, we can express covariate balancing in a clearer form as $\sum^n_{i= 1}\Big(\mathbbm{1}[D_i = 1]\widehat{w}_i \Phi(X_i) - \mathbbm{1}[D_i = 0]\widehat{w}_i \Phi(X_i)\Big) = \bm{0}_p$, while under linear models, $\sum^n_{i= 1}\widehat{w}_i \Phi(D_i, X_i) - \left(\sum^n_{i=1}\Big(\Phi(1, X_i) - \Phi(0, X_i)\Big)\right) = \bm{0}_p$ is attained, but it is somewhat harder to interpret. Moreover, using logistic models incorporates more information about the form of the bias-correction term, which includes the inverse propensity function. Logistic models also naturally impose restrictions such that $h(1, X)\in (1, \infty)$ and $h(0, X) \in (-\infty, -1)$ under the common support assumption.
In contrast, if we use linear models, we can express the entire ATE estimator with a single linear model, as shown in BrunsSmith2025augmentedbalancing. Furthermore, we can obtain the estimator of the bias-correction term as a closed-form solution. In addition, as discussed in Kato2025nearestneighbor, a subsequent work of this study. nearest neighbor matching is also an instance of linear models trained via Riesz regression (squared loss). We introduce the result in Appendix (ref) for reference.
Ultimately, there is no clear dominance between the use of linear and logistic models. Moreover, we can also use more complex models, such as random forests and neural networks. The choice of model should be made based on the data and application, and once the model is selected, we can determine appropriate specifications that ensure covariate balancing automatically.
This section provides an estimation error analysis for $h_0$ estimated by the direct bias-correction term estimation method. We can use various models for ${\mathcal{H}}$, including RKHS and neural networks.
We define a model of the bias-correction term $h_0$ by $h(D, X) = \zeta^{-1} \circ f(D, X) = \zeta^{-1} (f(D, X))$, where $\zeta$ is a continuously differentiable and globally Lipschitz link function, and $f$ is some basic model. For example, if we use linear model for the bias-correction term $h_0$, we can write $h(D, X) = \Phi(D, X)^\top \beta$, where $\zeta$ is the identity function, $f(D, X) = \Phi(D, X)^\top \beta$, $\Phi$ is some basis function and $\beta$ is the corresponding parameter. If we use logistic model for the bias-correction term $h_0$, we can use logistic link for $\zeta$, and $f(D, X) = \Phi(X)^\top \beta$.
First, we investigate the case with RKHS regression. Let ${\mathcal{F}}^{\mathrm{RKHS}}$ be a class of RKHS functions, and define $\widehat{f}^{\mathrm{RKHS}}_n \coloneqq \operatorname*{arg\,min}_{f \in {\mathcal{F}}^{\mathrm{RKHS}} } \widehat{{\mathcal{L}}}_n(\zeta^{-1}\circ f) + \lambda \|f\|^2_{{\mathcal{F}}}$, where $\|\cdot \|^2_{{\mathcal{F}}}$ is the RKHS norm. Then, we define an estimator as $h^{\mathrm{RKHS}} \coloneqq \zeta^{-1} \circ \widehat{f}^{\mathrm{RKHS}}_n$ We analyze the estimation error by employing the results in Kanamori2012statisticalanalysis, which study RKHS-based LSIF in DRE. We define the following localized class of RKHS functions as a technical device: ${\mathcal{F}}^{\text{RKHS}}_M \coloneqq \big\{f \in {\mathcal{F}}^{\text{RKHS}}\colon I(f) \leq M\big\}$ for some norm $I(f)$ of $f$. We also define ${\mathcal{H}}^{\text{RKHS}} \coloneqq \left\{\zeta^{-1}\circ f \colon f\in {\mathcal{F}}^{\mathrm{RKHS}}\right\}$. We then make the following assumption using this localized class.
For the details of the definition of the bracketing entropy, see Appendix (ref) and Definition 2.2 in VandeGeer2000empiricalprocesses.
Under these preparations, we establish an estimation error bound.
The proof is provided in Appendix (ref), following the approach of Kanamori2012statisticalanalysis. The parameter $\gamma$ is determined by the function class to which $f_0$ belongs.
Second, we provide an estimation error analysis when we use neural networks for ${\mathcal{H}}$. Our analysis is mostly based on Kato2021nonnegativebregman and Zheng2022anerror. We define Feedforward neural networks (FNNs) as follows:
For the model ${\mathcal{F}}^{\mathrm{FNN}}$, we define $\widehat{f}^{\mathrm{FNN}}_n \coloneqq \operatorname*{arg\,min}_{f \in {\mathcal{F}}^{\mathrm{FNN}} } \widehat{{\mathcal{L}}}_n(\zeta^{-1}\circ f)$. Then, we define an estimator as $\widehat{h}^{\mathrm{FNN}}_n \coloneqq \zeta^{-1}\circ \widehat{f}^{\mathrm{FNN}}_n$.
For the estimator, we can prove an estimation error bound. Let us make the following assumption.
Let $\mathrm{Pdim}({\mathcal{F}}^{\mathrm{FNN}})$ be the pseudo-dimension of ${\mathcal{F}}^{\mathrm{FNN}}$. For the definition, see Anthony1999neuralnetwork and Definition 3 in Zheng2022anerror. Then, we prove the following estimation error bound:
The proof is provided in Appendix (ref), following the approach of Zheng2022anerror. This result directly implies the minimax optimality of the proposed method when $f_0$ belongs to a H\"older class.
This section introduces the AIPW estimator with nuisance parameters estimated using our proposed direct bias-correction term estimation. We prove that under certain conditions, the proposed estimator is asymptotically normal. Note that this result is well known in the literature except for the use of nuisance parameters estimated via our direct bias-correction term estimation. The purpose of this section is not to provide novel methodological or theoretical results but to present an application of our proposed method.
We analyze the AIPW estimator with an estimated propensity score. Recall that the AIPW estimator is defined as $\widetilde{\tau}^{\mathrm{AIPW}}_n = \frac{1}{n}\sum^n_{i=1}\Big(\widehat{h}_n(D_i, X_i)\left(Y_i - \widehat{\mu}_n(D_i, X_i)\right) + \widehat{\mu}_n(1, X_i) - \widehat{\mu}_n(0, X_i)\Big)$, which is also called the DR estimator.
We first make the following assumption.
For example, the Donsker condition holds when the bracketing entropy of ${\mathcal{H}}$ is finite. In contrast, it is violated in high-dimensional regression or series regression settings where the model complexity diverges as $n \to \infty$. For neural networks, the assumption holds if both the number of layers and the width are finite. However, if these quantities grow with the sample size, the assumption is no longer valid.
Even if the Donsker condition does not hold, we can still establish asymptotic normality by employing sample splitting Klaassen1987consistentestimation. There are various ways to implement sample splitting, and one of the most well-known is cross-fitting, used in double machine learning Chernozhukov2018doubledebiased. In DML, the dataset is split into several folds, and the nuisance parameters are estimated using only a subset of the folds. This ensures that in $\widehat{h}_n(D_i, X_i)\left(Y_i - \widehat{\mu}_n(D_i, X_i)\right) + \widehat{\mu}_n(1, X_i) - \widehat{\mu}_n(0, X_i)$, the observations $(X_i, D_i, Y_i)$ are not used to construct $\widehat{\mu}_n$ and $\widehat{r}_n$. For more details, see Chernozhukov2018doubledebiased.
Under these assumptions, we show the asymptotic normality of $\widetilde{\tau}^{\mathrm{AIPW}}_n$. We omit the proof. For details, see Schuler2024introductionmodern, for example.
Here, $V^*$ matches the efficiency bound given as the variance of the efficient influence function Vaart1998asymptoticstatistics. Thus, this estimator is efficient.
If we follow the standard DRE approach, we may formulate the problem as the direct estimation of $r_0(1, X)$. For example, when using LSIF, the risk is given by ${\mathbb{E}}\big[ - 2r(1, X)\big] + {\mathbb{E}}\big[\mathbbm{1}[D = 1]r(1, X)^2\big]$, which corresponds to a part of our risk: ${\mathbb{E}}\Big[ - 2r(1, X) - 2r(0, X) + \mathbbm{1}[D = 1]r(1, X)^2 + \mathbbm{1}[D = 0]r(0, X)^2\Big]$. Thus, our proposed method is closely connected to LSIF. However, the standard DRE approach does not address whether it is suitable for bias-correction term estimation. In fact, we can estimate $r_0$ by minimizing the LSIF risk, but our proposed method adopts a different risk: the sum of ${\mathbb{E}}\big[ - 2r(1, X)\big] + {\mathbb{E}}\big[\mathbbm{1}[D = 1]r(1, X)^2\big]$ and ${\mathbb{E}}\big[ - 2r(0, X)\big] + {\mathbb{E}}\big[\mathbbm{1}[D = 0]r(0, X)^2\big]$, which is directly related to the bias-correction term.
We assess the performance of our method through simulation studies, evaluating ATE estimation error. We denote our direct bias-correction term estimation methods as DBC (LS) when using the squared loss, and DBC (TL) when using the tailored loss. We compare our approach with ATE estimators using propensity score estimated by maximum likelihood estimation (MLE), CBPS Imai2013covariatebalancing, and RieszNet Chernozhukov2022riesznet. Because our DBC (LS) is equivalent to Resz regression, we include RieszNet primarily as a numerical check of equivalence, noting architectural differences. In this section, for simplicity, we do not apply cross-fitting. We also conduct experiments in Appendices (ref) and (ref) using synthetic and semi-synthetic data, respectively, in which we apply cross-fitting.
We consider two different dimensions for \(X\), setting \(K=3\) and \(K=10\), and two different outcome models. This results in a total of four experimental settings. In all cases, the true ATE is fixed at \(\tau_0 = 5.0\). To generate synthetic data, we first sample covariates \(X_i\) from a multivariate normal distribution \(\mathcal{N}(0, I_K)\), where \(I_K\) denotes the \(K \times K\) identity matrix. The propensity score is then defined as $e_0(X_i) = \frac{1}{1 + \exp\bigl(-h(X_i)\bigr)}$, where $ h(X_i) = \sum_{j=1}^3 \alpha_j X_{i,j} + \sum_{j=1}^3 \beta_j X_{i,j}^2 + \gamma_{1} X_{i,1} X_{i,2} + \gamma_{2} X_{i,2} X_{i,3} + \gamma_{3} X_{i,1} X_{i,3} $. The coefficients \(\alpha_j\), \(\beta_j\), and \(\gamma_j\) are independently drawn from \(\mathcal{N}(0,0.5)\). Given these propensity scores, the treatment assignment \(D\) is sampled accordingly. The outcome is then generated under two models, referred to as Model 1 and Model 2. In Model 1, we specify $Y_i = \left(X^\top_i \beta\right)^2 + 1.1 + \tau_0 D_i + \varepsilon_i$, where \(\varepsilon_i \sim \mathcal{N}(0,1)\) and \(\tau_0 = 5.0\). In Model 2, the outcome is generated as $Y_i = X^\top_i \beta + \left(X^\top_i \beta\right)^2 + 3\sin(X_{i,1}) + 1.1 + \tau_0 D_i + \varepsilon_i$.
We model $h_0$ by modeling $e_0$. To model $e_0$, we use a three-layer neural network with an Exponential Linear Unit (ELU) activation function for each hidden layer ($100$ nodes per layer). The final output layer applies a sigmoid function to ensure that the estimated propensity scores remain in \((0,1)\). We use this model for our method, logistic regression, and CBPS. For RieszNet, we adopt the DragonNet architecture proposed in Shi2019adaptingneural, following Chernozhukov2022automaticdebiased. For each method, including ours, we compute both the IPW and AIPW estimators using the estimated scores. Additionally, we include the direct method (DM) estimator with neural networks for comparison. In each case, the expected conditional outcomes are estimated using a three-layer neural network (100 nodes per hidden layer, with ELU activation). As a baseline, we also consider the DM estimator with linear models.
The sample size is fixed at \(n = 3000\). As noted earlier, we evaluate two values of \(K\) (\(K=3\) and \(K=10\)) and two outcome-model specifications (Model 1 and Model 2), resulting in four experimental configurations. Each setting is repeated 500 times. We report the MSEs and biases of the resulting ATE estimates in Table (ref) for \(n = 3000\). Overall, the results indicate that our direct bias-correction approach achieves competitive or superior estimation accuracy compared with logistic regression and CBPS, highlighting the benefits of explicitly estimating the bias-correction term in the ATE context. RieszNet tends to outperform our method, but we consider this to be partly due to differences in the regression models. While RieszNet employs DragonNet, we use a simpler implementation. We do not employ such models, as model complexity is not our primary focus. Nevertheless, we emphasize that our method outperforms most existing approaches while exhibiting comparable performance to RieszNet.
This study proposed direct bias-correction term estimation in ATE estimation. Instead of focusing on estimating the propensity score itself, our approach directly minimizes the estimation error of the bias-correction term, leveraging empirical risk minimization techniques. We demonstrated that this direct approach enhances estimation accuracy by avoiding the intermediate step of propensity score estimation. Additionally, our method was analyzed through the lens of Bregman divergence minimization, providing a generalized framework.
\onecolumn