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.
70,621 characters · 30 sections · 65 citation commands
Semi-Supervised Treatment Effect Estimation with Unlabeled Covariates for Prediction-Powered Causal Inference
{\flushleft{{\bf Keywords:} causal inference; prediction-powered inference; double machine learning; Riesz regression; semiparametric efficiency; semi-supervised learning}}
A core interest in causal inference is estimating treatment effects, including the average treatment effect Imbens2015causalinference. In the standard setup, we estimate such treatment effects from triples of covariates, a treatment indicator, and outcomes. As in other statistical analyses, accuracy depends not only on the statistical method but also on the amount and type of data available. While randomized controlled trials are the gold standard, they are often infeasible. Therefore, in many practical scenarios, we use observational data to perform causal inference. However, observational data are also not necessarily easy to collect. In particular, treatment variables and the corresponding outcomes are often costly, whereas covariates are usually easy to gather.
Under this practical scenario, we consider estimating ATEs more accurately using auxiliary unlabeled covariates, even when treatment variables and outcomes are missing. We also discuss the average treatment effect on the treated, because it is often the causal target in observational studies where treatment participation itself determines the relevant population. This setting corresponds to semi-supervised learning in machine learning, where we utilize both labeled and unlabeled data Chapelle2006semisupervisedlearning, which is also referred to as prediction-powered inference Angelopoulos2021gentleintroduction,Ilker2024predictionpowered.
In many applications, such unlabeled covariates are easy to gather. For example, in the United States, we may aim to estimate the ATE for the effect of a new scholarship. Although we may know the covariates for an enormous number of students, we can assign treatment, scholarship, to only a limited number of them. In such cases, the unlabeled covariates contain information about the population over which the treatment effect is averaged, even though they do not contain treatment indicators or outcomes.
We find that, under appropriate conditions, using unlabeled data allows us to construct an ATE estimator whose asymptotic variance, or equivalently, asymptotic MSE, is smaller than that of an estimator that ignores unlabeled data, as shown by Hahn1998ontherole. To support this finding, we develop an asymptotic efficiency bound, a lower bound on the asymptotic variance, when using unlabeled covariates, propose ATE estimators, and show that the resulting asymptotic variances match the efficiency bound. In the methodological and theoretical arguments, we consider two practical scenarios, called the one-sample and two-sample scenarios. In the one-sample scenario, we interpret the unlabeled covariates as part of a dataset with missing variables, outcomes and the treatment indicator. In the two-sample scenario, we assume that labeled and unlabeled data are two independent datasets. The distinction is important because the two scenarios lead to different tangent spaces, different efficient influence functions, and different ways of using the unlabeled covariates.
Our efficient ATE estimators are developed based on the efficient influence function implied by the efficiency bound. We then extend the same logic to ATT, where the target distribution depends on the propensity score and therefore requires an additional treatment-law correction. This object is also called a Neyman orthogonal score in the debiased machine learning literature Chernozhukov2018doubledebiased. The Neyman orthogonal scores include nuisance parameters, regression functions and a Riesz representer, which must be estimated before obtaining the ATE estimators. For the Riesz representer estimation, we employ generalized Riesz regression in Kato2025directdebiased,Kato2025directbias, which generalizes the Riesz regression in Chernozhukov2021automaticdebiased. The unified treatment in Kato2026aunified further clarifies how Bregman divergence, loss-link choices, automatic regressor balancing, and automatic Neyman orthogonalization are connected. We extend this generalized Riesz regression perspective to the semi-supervised setting.
We list our contributions as follows:
\paragraph{Scope of the contribution.} The role of generalized Riesz regression in this paper is different from its role in the general framework of Kato2026aunified. That framework studies how to fit Riesz representers under Bregman divergences and how loss-link choices induce automatic regressor balancing and automatic Neyman orthogonalization. In contrast, the present paper fixes a semi-supervised observation scheme and derives the efficient scores, efficiency bounds, and attainable estimators under that scheme. The new content is therefore the interaction between unlabeled covariates, stratum-specific efficiency theory, and ATE or ATT targets. The generalized Riesz regression objectives are used as implementable nuisance-learning devices for the representers implied by these efficient scores.
\paragraph{Related work.} The related topics of this study include debiased machine learning, efficiency under the two-sample case (stratified sampling scheme), treatment effect estimation with missing values, density-ratio estimation, and semi-supervised learning.
In treatment effect estimation, we typically aim to attain the $\sqrt{n}$-rate with the smallest asymptotic variance, or equivalently, asymptotic MSE. We provide an efficiency bound, which is a lower bound on the asymptotic variance among regular estimators. As discussed in Uehara2020offpolicy, when there are two independent datasets, we cannot apply the usual efficiency bounds developed for a single dataset. To derive efficiency bounds in such settings, existing studies employ the efficiency theory under the stratified sampling scheme Wooldridge2001asymptoticproperties. Using this scheme, efficiency bounds have been proposed for various settings, including multiple log data, active learning, learning from positive and unlabeled data, and external-validity problems. This study also employs this technique to develop efficiency bounds.
The efficiency bounds are derived from the efficient influence functions. Certain efficient influence functions take forms that allow the removal of bias caused by the estimation errors of the nuisance parameters. Debiased machine learning is a framework for estimating treatment effects by utilizing such properties Chernozhukov2018doubledebiased. We refer to efficient influence functions with these properties as Neyman orthogonal scores. Chernozhukov2022automaticdebiased reframes this framework by characterizing Neyman orthogonal scores using the Riesz representer. Chernozhukov2021automaticdebiased proposes Riesz regression, an end-to-end method for estimating the Riesz representer. Kato2025directbias and Kato2025directdebiased propose generalized Riesz regression by regarding the Riesz representer estimation problem as Bregman divergence minimization Sugiyama2011densityratio. Kato2026aunified provides a broader formulation that relates Bregman-Riesz fitting to automatic regressor balancing and automatic Neyman orthogonalization. The present paper uses this machinery for a different purpose. We keep the sampling scheme explicit and derive the efficiency bounds induced by semi-supervised covariates, rather than starting from a fixed single-sample Riesz functional. This is why the two-sample result requires separate labeled and unlabeled influence functions, and why the ATT result requires a debiased denominator.
This study generalizes treatment effect estimation under covariate shift Uehara2020offpolicy,Kato2024activeadaptive and in the positive-unlabeled (PU) learning setup Kato2025puate. PU learning is a classical problem, originally studied in Imbens1996efficientestimation, and recently reframed by duPlessis2015convexformulation as a modern statistical machine learning framework. Our sampling scheme arguments are significantly inspired by the works in this literature.
This study is also related to semi-supervised regression Azriel2022semisupervised,Kawakita2013semisupervisedlearning and treatment effect estimation with missing values Heckman1974shadowprices,Robins1994estimationregression,Kennedy2020efficientnonparametric. These studies clarify how auxiliary covariates or missingness mechanisms can improve estimation. Our focus differs because the target is the semiparametric efficiency bound for causal effects when the auxiliary observations contain only covariates.
A further distinction is that the unlabeled observations in this paper affect the target covariate distribution rather than providing additional outcomes or surrogate outcomes. This distinction is important for efficiency. The unlabeled sample can improve the estimation of the covariate averaging component, but it cannot directly reduce the conditional outcome-noise component.
In this section, we formulate our problem setting. We define potential outcomes and observations separately by following the Neyman–Rubin causal model Neyman1923surapplications,Rubin1974estimatingcausal. Then, we define the evaluation covariate density and the two sampling scenarios. .
There is a binary treatment $d \in\{1, 0\}$. Let us define the corresponding potential outcome by $Y(d)$. Let $X\in \mathcal{X}\subset \mathbb{R}^k$ be a $k$-dimensional covariate, where $\mathcal{X}$ is the space. For each $d\in\{1, 0\}$, assume that the conditional distribution of $Y(d)$ given $X$ has its density, and let $r_{Y(d), 0}(y(d)\mid X)$ be the probability density function.
This study focuses on the estimation of the average treatment effect, which is the expected value of $Y(1) - Y(0)$. We take the expectation over a distribution whose covariate probability density is given by \[\kappa_0(x).\] We call it the evaluation covariate density. This density function can differ from $p_0(x)$. We make the assumptions for $\kappa_0(x)$ in the following sections.
Under a given covariate density $\kappa_0(x)$, the ATE is defined as follows:
where ${\mathbb{E}}_{\kappa_0}[\cdot]$ denotes the expectation taken over the distribution whose covariate density is $\kappa_0(x)$.
In addition to ATE, we consider the average treatment effect on the treated. For a given evaluation covariate density $\kappa_0$, define \[ \rho_{0, \kappa} \coloneqq {\mathbb{E}}_{\kappa_0}\left[e_0(1\mid X)\right], \] where $e_0(1\mid X)=P(D=1\mid X)$. The ATT target under the evaluation density $\kappa_0$ is \[ \tau^{\text{ATT}}_{0, \kappa} \coloneqq \frac{{\mathbb{E}}_{\kappa_0}\left[e_0(1\mid X)\tau_0(X)\right]}{\rho_{0, \kappa}}. \] When $\kappa_0=p_0$, this target is the usual ATT in the one-sample population. When $\kappa_0=\kappa_{0, \beta}$ in the two-sample scenario, it is the ATT for the evaluation population determined by the mixture density. This parameter is not obtained by replacing the ATE density with the treated covariate density alone, because the treated covariate density itself depends on the unknown propensity score. This feature is the source of the treatment-law correction in Section (ref).
This section defines the sample, that is, observations of $X$, $D$, and $Y$. To define observations rigorously, we need to consider the censoring setting carefully. To discuss data augmentation within the theory of semiparametric efficiency, we introduce two DGPs. The first DGP is the one-sample scenario, where there is only one dataset, and in this dataset, treatment indicators and outcomes are observed only for a subset of units. We also refer to this setting as the censoring setting. The second DGP is the two-sample scenario, where there are two independent datasets; one of the datasets contains data with covariates, treatment indicator, and outcomes, while the other only contains covariates. We also refer to this setting as the case-control setting or the stratified sampling scheme. We define these two DGPs below.
\paragraph{One-sample scenario.} In the one-sample scenario, we observe a single dataset ${\mathcal{D}}$, defined as follows:
where $O_i \in \{1, 0\}$ is an observation indicator, $\widetilde{D}_i \in \{1, 0, \text{NA}\}$, and $\widetilde{Y}_i$ is the observable treatment indicator and outcome, defined as
$D_i \in \{1, 0\}$ is a treatment indicator, and $Y_i$ is the outcome defined as \[Y_i \coloneqq \mathbbm{1}[D_i = 1] Y_i(1) + \mathbbm{1}[D_i = 0] Y_i(0).\] Here, $\text{NA}$ denotes a missing value. Equivalently, we can write $\widetilde{Y}_i$ as \[\widetilde{Y}_i = \mathbbm{1}\big[O_i = 1, \widetilde{D}_i = 1\big] Y_i(1) + \mathbbm{1}\big[O_i = 1, \widetilde{D}_i = 0\big] Y_i(0) + \mathbbm{1}\big[O_i = 0\big] \text{NA}.\] In this setting, we assume $p_0(x) = \kappa_0(x)$.
Note that $\widetilde{D}$ and $\widetilde{Y}$ are observable, while $Y_i$ and $D_i$ are not observable when $O_i = 0$.
\paragraph{Two-sample scenario} In the two-sample scenario, we observe two stratified datasets, ${\mathcal{D}}_{\text{L}}$ and ${\mathcal{D}}_{\text{U}}$:
where $m$ and $l$ are the sample sizes of each dataset, and $Y_j$ is the observed outcome defined as \[Y_j = \mathbbm{1}[D_j = 1]Y_j(1) + \mathbbm{1}[D_j = 0]Y_j(0),\] and $D_j \in \{1, 0\}$ is a treatment indicator.
\paragraph{Difference between the two settings} We show an illustration that demonstrates the difference between the one-sample and two-sample scenarios in Figure (ref). In both settings, we can identify and estimate the ATE in the standard way if we ignore the unlabeled auxiliary covariates. That is, in the one-sample scenario, we can estimate the ATE only by using ${\mathcal{D}}$, while in the two-sample scenario, we can estimate the ATE only by using ${\mathcal{D}}_{\text{L}}$. However, the unlabeled covariates change the information available about the evaluation covariate distribution. We demonstrate that this additional information reduces the asymptotic variance of efficient estimators.
A summary of the differences is provided below:
Throughout this study, let $P(R)$ denote the distribution of a random variable $R$. For simplicity, we assume that the distribution $P(R)$ of a continuous random variable $R$ has a probability density, whose notation depends on the random variable. For a probability density or mass function $p$, we denote the expectation over $p$ by ${\mathbb{E}}_p[\cdot]$. If the dependence is clear from the context, we omit $p$ and simply denote it as ${\mathbb{E}}[\cdot]$. Similarly, let $\text{Var}\left(\cdot\right)$ be the variance operator. Let us denote the true mean and variance of $Y(d)$ conditioned on $X = x \in\mathcal{X}$ by $\mu_0(d, x) = \mathbb{E}[Y(d)\mid X=x]$ and $\sigma^2_0(d, x) = \text{Var}(Y(d) \mid X=x)$, respectively.
We make the following regularity assumption.
First, we consider the one-sample setting for the DGP, which is also referred to as the censoring setting. We redefine the DGP with its notations and assumptions in Section (ref). Then, for this DGP, we develop an efficiency bound in Section (ref). We propose our estimator in Section (ref) and show consistency in Section (ref) and asymptotic normality in Section (ref).
This section introduces and summarizes the notations and assumptions, while recapping the DGP of the one-sample scenario. As defined in Section (ref), the DGP of this scenario is
Let $\pi_0(o \mid X) = p(O = o\mid X)$ be the probability of observation $O = o$, $e_0(d\mid X) = P(D = d\mid X, O = 1)$ be the propensity score defined in the observed samples, and let $g_0(a\mid X) = P(O = 1, D = a\mid X) = e_0(a\mid X)\pi_0(1 \mid X)$ be the joint probability of $O = 1$ and $D = d$. Under this notation, the probability density $p_0(x, o, \widetilde{d}, \widetilde{y})$ is written as
For simplicity, we assume that the evaluation density $\kappa_0(x)$ is the marginal density of the covariates.
We also make the following assumptions.
Note that this assumption also implies the existence of a universal constant $0 < \epsilon' < 1/2$ such that $\epsilon' < \pi_0(1 \mid X)$.
First, we derive the efficiency bound for regular estimators, which provides a lower bound on asymptotic variances. The efficiency bound is characterized via the efficient influence function Vaart1998asymptoticstatistics. In this scenario, the influence function has the usual augmented inverse probability structure, with $g_0(d\mid X)$ replacing the propensity score because both treatment assignment and label observation must occur. The result is stated below, and the proof is provided in Appendix (ref).
Recall that $g_0(d\mid X)=\pi_0(1\mid X)e_0(d\mid X)$. The following proposition is Theorem 25.20 in Vaart1998asymptoticstatistics, which connects the efficient influence function to the efficiency bound.
Using this proposition, we can derive the following efficiency bound from Lemma (ref).
Here, note that the efficient influence function depends on the unknown $\mu_0, g_0$, which are referred to as nuisance parameters. Since the efficient influence function satisfies the equation ${\mathbb{E}}\left[\psi^{\text{OS}}(X, O, Y; \mu_0, g_0, \tau_0)\right] = 0$, if the nuisance parameters are known and the exact expectation is computed, we can obtain $\tau_0$ by solving for $\tau_0$ that satisfies this equation. Thus, the efficient influence function provides a direct estimating equation for an efficient estimator. Furthermore, the accuracy of the estimation of the nuisance parameters affects the estimation of $\tau_0$, the parameter of interest.
Based on the efficient influence function, we propose an ATE estimator defined as \[\widehat{\tau}^{\text{OS}\mathchar`-\text{eff}}_n \coloneqq \frac{1}{n}\sum^n_{i=1}S^{\text{OS}}\left(X_i, O_i, \widetilde{D}_i, \widetilde{Y}_i; \widehat{\mu}_{n, i}, \widehat{g}_{n, i}\right),\] where $\widehat{\mu}_{n, i}$ and $\widehat{g}_{n, i}$ are estimators of $\mu_0$ and $g_0$. Note that the estimators can depend on $i$. This estimator is an extension of the augmented inverse probability weighting estimator, also called a doubly robust estimator Bang2005doublyrobust. We say that an estimator is efficient if its asymptotic variance aligns with $V^{\text{OS}}$.
For estimating the regression function $\mu_0$, we can employ methods for conditional ATE estimation Wager2018estimationinference,Curth2021nonparametricestimation,Kennedy2024minimaxrates, as well as standard regression methods using parametric or nonparametric models Tsybakov2008introductionnonparametric,SchmidtHieber2020nonparametric. We can also use targeted maximum likelihood estimation to refine this estimation vanderLaan2011targetedlearning.
For estimating $g_0$, we can use logistic regression or other advanced methods, such as the covariate balancing propensity score Imai2011estimationheterogeneous,Hainmueller2012entropybalancing and Riesz regression Chernozhukov2021automaticdebiased. As Zhao2019covariatebalancing, BrunsSmith2025augmentedbalancing, and Kato2025directbias show, Riesz regression and covariate balancing methods are in a dual relationship, and Riesz regression can be interpreted as a special case of density ratio estimation Kato2025unifiedtheory. For details, see Section (ref) and Appendix (ref).
We explain how to construct estimators for $g_0$. In this study, we employ generalized Riesz regression, also referred to as Bregman-Riesz regression Kato2025directbias,Kato2025directdebiased,Kato2026aunified. In the efficient estimation of causal parameters, Neyman orthogonal scores play an important role and typically correspond to the efficient score. Specifically, asymptotically efficient estimators must be asymptotically linear with respect to the Neyman orthogonal scores. When the parameter of interest is linear in the regression functions, the Neyman orthogonal score can be decomposed into the Riesz representer and regression functions. In our semi-supervised setting, the Riesz representer is the component that determines how labeled residuals and unlabeled covariate averages are combined. In our framework, the Neyman orthogonal score is given by
where we replace $g_0$ with $\alpha_0$ in the original definition of $\psi^{\text{OS}}$, and $\alpha_0\left(O_i, \widetilde{D}_i, X_i\right) \coloneqq \frac{\mathbbm{1}[O_i = 1, \widetilde{D}_i = 1]}{g_0(1\mid X_i)} - \frac{\mathbbm{1}[O_i = 1, \widetilde{D}_i = 0]}{g_0(0\mid X_i)}$ is the Riesz representer. Riesz regression, as proposed by Chernozhukov2021automaticdebiased, is a method for estimating the Riesz representer in an end-to-end manner. Kato2025directbias shows that Riesz regression is a specific instance of density ratio estimation and can be generalized via Bregman divergence minimization Sugiyama2011densityratio. Kato2025directdebiased further reformulates and extends this approach as direct debiased machine learning (DDML) via generalized Riesz regression. The efficiency results below use high-level product-rate assumptions for the learned nuisance functions. Generalized Riesz regression is one way to construct the representer estimator under these assumptions, while detailed rates for RKHS and neural network classes can be imported from the general theory of Bregman-Riesz fitting.
\paragraph{Generalized Riesz regression.} Generalized Riesz regression estimates $\alpha_0$ by minimizing the Bregman divergence between the true Riesz representer $\alpha_0$ and its model $\alpha$. That is, the estimation error of $\alpha_0$ is measured using the Bregman divergence. The recent unified framework of Kato2026aunified emphasizes that the choice of Bregman divergence and link function determines both the geometry of the fitted representer and the associated balancing interpretation. For a twice differentiable convex function $f$ with bounded derivative, the population objective for Riesz representer estimation is written as
The empirical counterpart $\widehat{\text{BD}}_{f}(\alpha)$ replaces expectations with sample averages. Here, we used $\tau_0 = \mathbb{E}\left[\mathbb{E}\left[\widetilde{Y}\mid O = 1, \widetilde{D} = 1, X\right] - \mathbb{E}\left[\widetilde{Y}\mid O = 1, \widetilde{D} = 0, X\right]\right] = 1, \widetilde{D} = 0, X\right]}$. Minimizing this objective over a hypothesis class ${\mathcal{A}}$ yields an estimator of $\alpha_0$, that is, \[ \widehat{\alpha}^{\text{GRR}} \coloneqq \operatorname*{arg\,min}_{\alpha \in {\mathcal{A}}} \widehat{\text{BD}}_{f}(\alpha), \] where GRR denotes generalized Riesz regression.
The use of generalized Riesz regression allows us to naturally incorporate unlabeled covariates into the estimation of the Riesz representer. This is because, in (ref), we can approximate ${\mathbb{E}}\left[\Big(\partial f\left(\alpha\left(1, 1, X\right)\right)X\right)\Big) - \partial f\left(\alpha\left(1, 0, X\right)\right)X\right)\right]}}$ using unlabeled covariates, whereas ${\mathbb{E}}\left[\partial f\left(\alpha\left(O, \widetilde{D}, X\right)\right)X\right)\right]\alpha\left(O, \widetilde{D}, X\right)-f\left(\alpha\left(O, \widetilde{D}, X\right)\right)X\right)}}$ requires labeled data. That is,
where the second term can be evaluated using both labeled and unlabeled data. This is the point at which the semi-supervised structure enters the Riesz representer estimation problem. Note that unlabeled covariates can be utilized even when $g_0$ is estimated via maximum likelihood. However, the generalized Riesz regression approach is arguably more appropriate in an end-to-end formulation because it directly targets the representer that appears in the Neyman orthogonal score. Also see Kawakita2013semisupervisedlearning.
Let ${\mathcal{A}}$ denote the model class for $\alpha_0$. If we set $f(\alpha) = (\alpha - 1)^2$, then
This population objective corresponds to Riesz regression as in Chernozhukov2021automaticdebiased. We refer to this objective as least squares Riesz (LS-Riesz) regression.
Now, redefine ${\mathcal{A}}$ as the set of $\alpha$ such that $\alpha(1, 1, \cdot) > 1$ and $\alpha(1, 0, \cdot) < -1$, a condition that should hold under the common support assumption. For $f(\alpha) = (|\alpha| - 1)\log(|\alpha| - 1) + |\alpha|$ ($\alpha \in {\mathcal{A}}$), the corresponding Bregman divergence is
We refer to this objective as Kullback-Leibler Riesz (KL-Riesz) regression, since the choice of $f$ yields the KL divergence.
By replacing the expectations with the sample mean and minimizing the empirical objective for $\alpha$, we can estimate $\alpha_0$.
\paragraph{Interpretation.} As Kato2025directdebiased discusses, LS-Riesz regression corresponds to the stable balancing weights proposed in Zubizarreta2015stableweights, and KL-Riesz corresponds to the entropy balancing weights in Hainmueller2012entropybalancing. These correspondences were originally shown in the covariate balancing literature, such as in Zhao2019covariatebalancing and BrunsSmith2025augmentedbalancing. They can be derived from duality relationships. Kato2026aunified further interprets these relationships as automatic regressor balancing under suitable loss-link pairs.
Note that the duality depends on the model class used for $\alpha_0$, namely ${\mathcal{A}}$. For the duality between LS-Riesz and stable balancing weights, linear models must be used for ${\mathcal{A}}$, whereas for the duality between KL-Riesz and entropy balancing weights, logistic models for $\alpha_0$ are required. Therefore, the choice of ${\mathcal{A}}$ is not a purely computational choice; it determines which balancing equations are targeted by the Riesz representer estimator.
First, we prove the consistency result, that is, $\widehat{\tau}^{\text{OS}\mathchar`-\text{eff}}_n \xrightarrow{{\mathrm{p}}} \tau_0$ holds as $n\to \infty$. We can obtain this result relatively easily compared to asymptotic normality. We make the following assumption, which holds for most estimators of the nuisance parameters.
Then, the following consistency result holds. This result is given as a special case of Theorem (ref); therefore, we omit the proof.
This consistency structure is referred to as double robustness.
Next, we establish the asymptotic normality of our estimator. Unlike consistency, this requires stronger assumptions on the nuisance estimators, especially for the propensity score.
To prove asymptotic normality or $\sqrt{n}$-consistency, we must control the complexity of the nuisance parameter estimators. One simple approach is to assume the Donsker condition; however, it is well known that this condition often fails in high-dimensional regression settings. In such cases, asymptotic normality can still be attained using sample splitting, a common technique in this field Klaassen1987consistentestimation, which has recently been refined by Chernozhukov2018doubledebiased as cross-fitting.
\paragraph{Cross-fitting.} We estimate $\mu_0$ and $g_0$ using cross-fitting. Cross-fitting is a variant of sample splitting Chernozhukov2018doubledebiased. We randomly partition ${\mathcal{D}}$ into $L > 0$ folds (subsamples), and for each fold $b \in {\mathcal{L}} \coloneqq \{1,2,\dots, L\}$, the nuisance parameters are estimated using all other folds. Let the estimators for fold $b \in {\mathcal{L}}$ be denoted by $\widehat{\mu}^{(b)}_n$ and $\widehat{g}^{(b)}_n$. Let ${\mathcal{I}}^{(b)}$ be the index set of samples belonging to fold $b$. This construction separates nuisance estimation from score evaluation and avoids imposing a Donsker condition on the nuisance classes.
Various estimation methods may be used, including neural networks and Lasso, as long as they satisfy the convergence rate conditions in Assumption (ref). The pseudocode is shown in Algorithm (ref).
\paragraph{Asymptotic normality.} We present results for the case with cross-fitting, but similar results hold under the Donsker condition.
We make the following assumptions:
We define the estimator as \[ \widehat{\tau}^{\text{OS}\mathchar`-\text{eff}}_n \coloneqq \frac{1}{n}\sum_{b \in {\mathcal{L}}}\sum_{i \in {\mathcal{I}}^{(b)}}S^{\text{OS}}\left(X_i, O_i, \widetilde{D}_i, \widetilde{Y}_i; \widehat{\mu}^{(b)}_n, \widehat{g}^{(b)}_n\right) \] and show that the asymptotic normality holds as follows:
The proof is provided in Appendix (ref). The asymptotic variance of $\widehat{\tau}^{\text{OS}\mathchar`-\text{eff}}_n$ matches the efficiency bound. Therefore, Theorem (ref) also implies that $\widehat{\tau}^{\text{OS}\mathchar`-\text{eff}}_n$ is asymptotically efficient.
For inference, let $\widehat{\psi}^{\text{OS}}_i$ denote the influence function in Lemma (ref) evaluated at the cross-fitted nuisance estimators and at $\widehat{\tau}^{\text{OS}\mathchar`-\text{eff}}_n$. We estimate the asymptotic variance by \[ \widehat{V}^{\text{OS}} \coloneqq \frac{1}{n}\sum^n_{i=1}\Big(\widehat{\psi}^{\text{OS}}_i - \frac{1}{n}\sum^n_{i'=1}\widehat{\psi}^{\text{OS}}_{i'}\Big)^2. \] The resulting Wald interval uses $\widehat{V}^{\text{OS}}/n$ as the variance of $\widehat{\tau}^{\text{OS}\mathchar`-\text{eff}}_n$.
We now discuss alternative ATE estimators.
Next, we consider the two-sample scenario for the DGP, which is also referred to as the case-control setting and stratified sampling scheme. We reintroduce the notation and assumptions required for our analysis in Section (ref). Section (ref) presents the efficiency bound, and Section (ref) provides an ATE estimator under this setting. We establish consistency in Section (ref) and asymptotic normality in Section (ref). Finally, we compare the one-sample and two-sample scenarios in Section (ref).
As introduced in Section (ref), the DGP for the two-sample scenario is defined as
Let $e_0(d\mid X) = P(D = d\mid X)$ denote the propensity score. Then, the joint density $p_0(x, d, y)$ can be written as
For the evaluation density, we make the following assumption.
The following support condition is needed because all outcome and treatment information is contained in the labeled sample.
If $\beta < 1$, Assumption (ref) implies $q_0 \ll p_0$. Without this condition, there may be target covariate regions for which outcomes are never observed in the labeled sample.
We also impose the following assumptions.
Define \[ \omega_{0, \beta}(x) \coloneqq \frac{\kappa_{0, \beta}(x)}{p_0(x)}, \] and \[ v_{0, \beta}(d, x) \coloneqq \frac{e_0(d\mid x)}{\omega_{0, \beta}(x)} = \frac{p_0(d, x)}{\kappa_{0, \beta}(x)}. \] Thus, $1/v_{0, \beta}(d, x)$ is the efficient residual weight for the labeled sample. We also define \[ \tau_{p,0} \coloneqq {\mathbb{E}}_{p_0}\left[\tau_0(X)\right], \tau_{q,0} \coloneqq {\mathbb{E}}_{q_0}\left[\tau_0(Z)\right], \tau_0 = \beta \tau_{p,0} + (1 - \beta)\tau_{q,0}. \]
Following Uehara2020offpolicy, we derive the efficiency bound using the efficiency arguments under the two-sample scenario. In this scheme, there are two efficient influence functions, one for the labeled stratum and one for the unlabeled stratum. The proof is provided in Appendix (ref).
The centering in Lemma (ref) is stratum specific. The labeled stratum is centered by $\tau_{p,0}$ and the unlabeled stratum is centered by $\tau_{q,0}$. This is necessary because each stratum has its own sampling distribution.
As in the one-sample scenario, the efficient influence functions directly yield the following efficiency bound.
The parameter $\beta$ is fixed throughout the theorem. If $p_0 = q_0$, all values of $\beta$ define the same evaluation density. Only in such cases can $\beta$ be selected to improve precision without changing the estimand.
Based on the efficient influence functions, we define the estimator as \[ \widehat{\tau}^{\text{TS}\mathchar`-\text{eff}}_n \coloneqq \frac{1}{m}\sum^m_{j=1} S^{\text{TS}}_{(X, D, Y)}\big(X_j, D_j, Y_j; \widehat{\mu}, \widehat{v}_{\beta}\big) + \beta\frac{1}{m}\sum^m_{j=1} S^{\text{TS}}_{(X)}\big(X_j; \widehat{\mu}\big) + \left(1 - \beta\right)\frac{1}{l}\sum^l_{k=1} S^{\text{TS}}_{(X)}\big(Z_k; \widehat{\mu}\big), \] where $S^{\text{TS}}_{(X)}(x; \mu) \coloneqq \mu(1, x) - \mu(0, x)$. Here, $\widehat{\mu}$ and $\widehat{v}_{\beta}$ denote estimators of $\mu_0$ and $v_{0, \beta}$. Unlike in the one-sample scenario, we do not use the observation indicator $O$, since it is deterministically known whether a unit belongs to the labeled or unlabeled stratum. This distinction leads to theoretical differences from the one-sample scenario.
The two-sample efficient score also yields a semi-supervised generalized Riesz regression objective. The Riesz representer for the labeled residual part is \[ \alpha_{0, \beta}(D, X) \coloneqq \frac{\mathbbm{1}\big[D = 1\big]}{v_{0, \beta}(1, X)} - \frac{\mathbbm{1}\big[D = 0\big]}{v_{0, \beta}(0, X)}. \] It represents the linear functional $h \mapsto {\mathbb{E}}_{\kappa_{0, \beta}}\left[h(1, X) - h(0, X)\right]$ through expectations under the labeled distribution. For a twice differentiable convex function $f$, define
The empirical counterpart is
Then, we estimate $\alpha_{0, \beta}$ by \[ \widehat{\alpha}^{\text{TS}\mathchar`-\text{GRR}} \in \operatorname*{arg\,min}_{\alpha\in{\mathcal{A}}}\left\{\widehat{\text{BD}}^{\text{TS}}_{f, \beta}(\alpha) + \lambda J(\alpha)\right\}. \] This objective makes explicit how the unlabeled covariates enter generalized Riesz regression: the labeled sample controls the residual representer geometry, while both labeled and unlabeled covariates define the target linear functional. This connection follows the generalized Riesz regression perspective, where Bregman divergence and loss-link choices determine both representer fitting and automatic balancing behavior Kato2026aunified.
We impose the following assumption.
Then, the following consistency result holds.
Next, we establish the asymptotic normality of the estimator.
We now establish asymptotic normality in the following theorem, with the proof provided in Appendix (ref). In this result, we consider the asymptotic regime where the sample sizes $m$ and $l$ approach infinity while maintaining a fixed ratio $m:l = \alpha:(1-\alpha)$.
Thus, the proposed estimator is efficient with respect to the efficiency bound derived in Theorem (ref).
For inference, let $\widehat{\psi}^{\text{TS}}_{\text{L},j}$ and $\widehat{\psi}^{\text{TS}}_{\text{U},k}$ denote the two stratum influence functions in Lemma (ref) evaluated at the cross-fitted nuisance estimators and at $\widehat{\tau}^{\text{TS}\mathchar`-\text{eff}}_n$. We estimate the scaled variance by
The resulting Wald interval uses $\widehat{V}^{\text{TS}}(\beta)/N$ as the variance of $\widehat{\tau}^{\text{TS}\mathchar`-\text{eff}}_n$.
The difference between the one-sample and two-sample scenarios appears in the formulation of ATE estimators, the setup of Riesz regression, and the corresponding efficiency arguments. In the one-sample scenario, the observation indicator is random within a single superpopulation, and the efficient influence function uses $g_0(d\mid X) = P(O = 1, D = d\mid X)$. In the two-sample scenario, the sample membership is fixed by design, and the efficient influence function has separate labeled and unlabeled stratum components. This distinction is the reason why the two-sample bound in Theorem (ref) uses stratum-specific centering.
In many applications, we often have access to many more unlabeled data points than fully labeled ones, as unlabeled covariates are less costly to collect. The main results above use the fixed-ratio regime $m/N\to\alpha$. This section records the corresponding sequential implication when the unlabeled sample is asymptotically much larger than the labeled sample.
In the two-sample scenario, let $l/m\to\infty$ and normalize by $\sqrt{m}$. Then, the unlabeled stratum average is negligible, but the labeled stratum still contains both the residual component and the $p_0$-covariate averaging component when $\beta>0$.
The one-sample analogue requires a triangular-array formulation if the probability of observing labels changes with the sample size. To avoid conflating this regime with the fixed-DGP efficiency theory in Section (ref), we state the qualitative implication only. If an external covariate sample makes the empirical distribution of $X$ negligible relative to the labeled sample size, then the covariate averaging component in the one-sample efficiency bound is estimated with negligible additional noise, while the conditional outcome-noise component remains governed by the labeled observations. A fully formal triangular-array statement can be added separately if this regime is the focus of the application.
By using auxiliary unlabeled covariates, we can reduce the asymptotic variance of ATE estimators. As Hahn1998ontherole shows, for a labeled dataset $\left\{\big(X_i, D_i, Y_i\big)\right\}^{n^\dagger}_{i=1}$, the efficiency bound of ATE estimators $\widehat{\tau}$ is given as $V^\dagger \coloneqq {\mathbb{E}}\left[\frac{\sigma^2_0(1, X)}{P(D = 1\mid X)} + \frac{\sigma^2_0(0, X)}{P(D = 0\mid X)}\right] + {\mathbb{E}}\left[\big(\tau_0(X) - \tau_0\big)^2\right]$, and an efficient ATE estimator satisfies \[ \sqrt{n^\dagger}\left(\widehat{\tau} - \tau_0\right)\xrightarrow{{\mathrm{d}}} {\mathcal{N}}\left(0, V^\dagger\right). \] The corrected two-sample bound makes the efficiency gain transparent in the same-population case $p_0 = q_0$. If $m/N\to\alpha$ and $\beta = \alpha$, then the semi-supervised bound is \[ V^{\text{TS}}(\alpha) = \frac{1}{\alpha}{\mathbb{E}}_{p_0}\left[\frac{\sigma^2_0(1, X)}{e_0(1\mid X)} + \frac{\sigma^2_0(0, X)}{e_0(0\mid X)}\right] + {\mathbb{E}}_{p_0}\left[\Big(\tau_0(X) - \tau_0\Big)^2\right]. \] By contrast, using only the $m$ labeled observations and scaling by $\sqrt{N}$ gives \[ V^{\text{sup}} = \frac{1}{\alpha}{\mathbb{E}}_{p_0}\left[\frac{\sigma^2_0(1, X)}{e_0(1\mid X)} + \frac{\sigma^2_0(0, X)}{e_0(0\mid X)} + \Big(\tau_0(X) - \tau_0\Big)^2\right]. \] Therefore, \[ V^{\text{sup}} - V^{\text{TS}}(\alpha) = \left(\frac{1}{\alpha} - 1\right){\mathbb{E}}_{p_0}\left[\Big(\tau_0(X) - \tau_0\Big)^2\right]. \] This identity is the main variance-reduction message. Unlabeled covariates reduce the covariate averaging component but do not reduce the residual outcome-noise component.
The same message holds for ATT. In the same-population case $p_0=q_0$ with $\beta=\alpha$, let $\rho_0={\mathbb{E}}_{p_0}\left[e_0(X)\right]$ and $\Delta_0(X)=\tau_0(X)-\tau^{\text{ATT}}_0$. The labeled-only ATT bound under $\sqrt{N}$ scaling contains $\frac{1}{\alpha\rho_0^2}{\mathbb{E}}_{p_0}\left[e_0(X)^2\Delta_0(X)^2\right]$ as the treated-covariate averaging component. The semi-supervised ATT bound contains the same component without the factor $1/\alpha$. Therefore, \[ V^{\text{ATT}\mathchar`-\text{sup}} - V^{\text{ATT}\mathchar`-\text{TS}}(\alpha) = \left(\frac{1}{\alpha}-1\right)\frac{{\mathbb{E}}_{p_0}\left[e_0(X)^2\Delta_0(X)^2\right]}{\rho_0^2}. \] Thus, for ATT, unlabeled covariates reduce the treated-covariate averaging component weighted by the squared propensity score.
Uehara2020offpolicy investigates ATE estimation, equivalently, off-policy evaluation, under covariate shift. Our formulation includes their ATE estimation approach as a special case. When $\beta = 0$ in the two-sample scenario, the evaluation density is determined by the unlabeled covariate distribution $q_0$. In this case, $\kappa_{0,0}=q_0$ and the residual weights are density-ratio weighted through $q_0(X)/p_0(X)$. The fixed-ratio efficiency bound contains the labeled residual component and the unlabeled target-covariate averaging component. If $l/m\to\infty$, the latter vanishes under $\sqrt{m}$ normalization, which yields the usual covariate-shift form.
We add a small simulation to verify the efficiency-gain identity and to check the behavior under covariate shift. The numerical illustration uses oracle nuisance quantities so that the experiment isolates the variance-decomposition mechanism. The accompanying code is written in the style of generalized Riesz regression, with the residual weights treated as Riesz representers, but it does not import the genriesz package. The purpose of this section is therefore diagnostic rather than a full benchmark of nuisance-learning algorithms. The experiment is included to verify the main variance decomposition before adding additional implementation-specific variation from first-step learning. A full implementation can replace the oracle representers with estimators obtained from the empirical Bregman objectives in Sections (ref) and (ref), following the loss-link workflow of generalized Riesz regression Kato2026aunified.
In the same-population experiment, the semi-supervised estimators have smaller $N$-scaled variance than the supervised estimators for both ATE and ATT. In the covariate-shift experiment, the supervised estimators are biased because they average over the source covariate distribution, while the semi-supervised estimators target the evaluation distribution using the unlabeled covariates.
The results rely on several restrictions that are useful to state explicitly. First, all causal targets are identified under unconfoundedness, and the theory does not address unobserved confounding. Second, the two-sample scenario requires support of the evaluation density inside the labeled covariate distribution, because outcomes are observed only under $p_0$. Third, the mixture parameter $\beta$ is part of the estimand. Except in special cases such as $p_0=q_0$, changing $\beta$ changes the target population rather than merely improving precision. Fourth, the main asymptotic theory uses fixed stratum proportions. The many-unlabeled discussion records the limiting implication of that theory, while a one-sample regime with a label probability converging to zero requires a separate triangular-array formulation. Finally, the numerical illustration isolates the efficiency mechanism using oracle nuisances. It is intended to verify the variance decomposition, while full empirical evaluation of generalized Riesz regression requires replacing the oracle representers with the empirical Bregman objectives.
These limitations do not change the main message. Unlabeled covariates are useful when the causal target averages conditional effects over a covariate law that can be learned more accurately from the auxiliary sample. They do not create outcome information in regions without labeled support, and they do not remove the need for accurate residual correction. This distinction explains both the support condition and the efficiency-gain identities.
This study investigates semiparametric efficient estimation of ATE and ATT when auxiliary unlabeled covariates are accessible. We consider both one-sample and two-sample scenarios, and derive semiparametric efficiency bounds for each. Based on the corresponding efficient influence functions, we construct asymptotically efficient estimators via Neyman orthogonal scores. Our approach leverages generalized Riesz regression for estimating nuisance parameters, allowing flexible incorporation of unlabeled covariates. The main implication is that unlabeled covariates reduce the variance associated with averaging the conditional treatment effect over the evaluation covariate density, while the conditional outcome-noise component remains governed by the labeled data. The two-sample analysis also clarifies that the evaluation-density mixture parameter is part of the estimand, and that the efficient influence functions must be centered within each stratum. The ATT extension shows that the same principle continues to hold for debiased ratio targets, but the efficient score must also correct the treatment law and the treatment mass. The proposed framework performs prediction-powered causal inference and extends existing methods for treatment effect estimation under covariate shift, missing labels, and semi-supervised settings.
\onecolumn