Keywords— Causal Inference; Dimension Reduction; Factor Models; High-Dimensional Data; Auto-encoder

1 Introduction↩︎

Causal mediation analysis seeks to explain how a treatment affects an outcome by decomposing the total effect into pathways operating through intermediate variables, or mediators [1]–[4]. In the potential-outcomes framework [5], this amounts to studying counterfactual means whose contrasts yield the natural indirect effect, natural direct effect, and total effect. Classical methods typically assume that the relevant covariates and mediators are directly observed and low-dimensional. Modern biomedical studies often violate both assumptions: the causal mechanisms of interest are latent, while the available measurements are noisy, high-dimensional proxies. This setting arises across genomics and multi-omics applications. High-dimensional molecular measurements, such as metabolomic or multi-omics profiles, are often indirect views of lower-dimensional biological states [6], [7]. For example, oxidative stress has been proposed as a pathway linking depression to Alzheimer’s disease [8], [9], but reactive oxygen species are short-lived and difficult to measure directly [10]. Researchers therefore rely on downstream molecular signatures, including DNA damage and DNA methylation changes [11]–[13]. Such measurements may be nonlinear functions of latent biological processes [14], [15]; moreover, upstream covariate factors, treatments, and mediator factors can interact in ways that complicate causal identification and estimation [16]. These features motivate a mediation framework that can handle latent structure, high dimensionality, nonlinear measurement maps, and treatment–covariate–mediator interactions simultaneously.

Existing high-dimensional mediation methods address only parts of this problem. Sparse and variable-selection approaches can screen many candidate mediators [17]–[20], but they usually treat the observed variables as the mediators and do not model latent structure. Dimension-reduction methods based on PCA or linear factor models provide low-dimensional summaries [21], [22], but linear measurement maps can be restrictive in complex biological systems [23]. Autoencoders and related deep models can learn nonlinear representations [24], [25], yet applying them separately to covariates and mediators ignores the structural dependence between the two latent systems. Deep mediation estimators such as [26] improve robustness to nuisance model misspecification, but do not explicitly target nonlinear high-dimensional latent-factor settings. Variational autoencoder approaches for indirectly observed mediators [27] model latent mediators, but do not jointly accommodate latent nonlinear covariate structure, structural dependence between covariate and mediator representations, and interaction effects; see Appendix 7 for further discussion. Overall, existing methods address individual aspects of the problem but do not simultaneously address all the aforementioned challenges. The key missing ingredient is a framework that learns nonlinear covariate and mediator representations jointly, aligns them with the latent mediation structure, and then uses these representations in a statistically valid causal estimator.

To bridge this gap, we propose MediEncoder, a nonlinear representation-learning framework for high-dimensional causal mediation analysis. The method jointly learns low-dimensional representations of covariates and mediators through a structured encoder–decoder architecture. A cross-factor network links the treatment and covariate representation to the mediator representation, encouraging the learned features to respect the causal structure of the latent mediation problem. Exact recovery of the original latent factors is not required: it suffices to learn representations that preserve the information needed for the mediation functional, and our theory formalizes invariance under left-invertible transformations of the latent factors.

Our contributions are fourfold. (i) First, we introduce a nonlinear latent-factor framework for mediation analysis with high-dimensional covariates and mediators, avoiding linearity and sparsity requirements. (ii) Second, we develop MediEncoder, a coupled representation-learning algorithm that uses a cross-factor network to capture structural dependence between covariate and mediator representations. (iii) Third, we develop a modular asymptotic theory for the resulting cross-fitted efficient influence-function estimator. The theory formalizes that exact factor recovery is unnecessary: it suffices for the learned representation to approximate a left-invertible transformation of the latent factors, and the resulting inference remains valid under high-level factor-recovery and nuisance-rate conditions. (iv) Fourth, we demonstrate the empirical performance of MediEncoder in extensive simulations and in a biomedical application using data from the Alzheimer’s Disease Neuroimaging Initiative (ADNI).

2 Model Description↩︎

We consider a binary treatment (or exposure) \(A\in\{0,1\}\) and a real-valued outcome \(Y\in\mathbb{R}\). Let \(\boldsymbol{f}_X\in\mathbb{R}^{\bar p}\) denote a vector of pre-treatment latent covariate factors and let \(\boldsymbol{f}_M\in\mathbb{R}^{\bar q}\) denote a vector of latent mediator factors. The covariate factors may confound the treatment–mediator, treatment–outcome, and mediator–outcome relationships, whereas the mediator factors are the mechanisms through which part of the effect of \(A\) on \(Y\) may be transmitted. We use boldface notation for random latent vectors and write \(f_X\) and \(f_M\) for their generic values.

A compact structural representation of the latent mediation model is \[\label{eq:factor} \begin{align} \boldsymbol{f}_M &=(1-A)\{\mu_0^*(\boldsymbol{f}_X)+\boldsymbol{u}_{XM}\} +A\{\mu_1^*(\boldsymbol{f}_X)+\boldsymbol{u}'_{XM}\},\\ Y &=(1-A)\{\mu_0(\boldsymbol{f}_X,\boldsymbol{f}_M)+\epsilon'\} +A\{\mu_1(\boldsymbol{f}_X,\boldsymbol{f}_M)+\epsilon\}, \end{align}\tag{1}\] where \(\mu_0^*,\mu_1^*:\mathbb{R}^{\bar p}\to\mathbb{R}^{\bar q}\) and \(\mu_0,\mu_1:\mathbb{R}^{\bar p+\bar q}\to\mathbb{R}\) are unknown structural functions and \(\epsilon\), \(\epsilon'\), \(\boldsymbol{u}_{XM}\in \mathbb{R}^{\bar q}\), and \(\boldsymbol{u}_{XM}'\in \mathbb{R}^{\bar q}\) are centered noise terms.

In the applications motivating this work, \(\boldsymbol{f}_X\) and \(\boldsymbol{f}_M\) are not directly observed. Instead, we observe high-dimensional measurements \(X\in\mathbb{R}^p\) and \(M\in\mathbb{R}^q\), with \(p\gg \bar p\) and \(q\gg \bar q\), that serve as noisy proxies for the latent factors; see Figure 1. We assume the measurement model \[\label{eq:factors} X=\phi_X(\boldsymbol{f}_X)+\boldsymbol{u}_X, \qquad M=\phi_M(\boldsymbol{f}_M)+\boldsymbol{u}_M,\tag{2}\] where \(\phi_X:\mathbb{R}^{\bar p}\to\mathbb{R}^p\) and \(\phi_M:\mathbb{R}^{\bar q}\to\mathbb{R}^q\) are unknown loading functions, and \(\boldsymbol{u}_X\) and \(\boldsymbol{u}_M\) are centered idiosyncratic errors. The classical linear factor model is a special case of 2 , obtained by taking \(\phi_X(f_X)=\Gamma_X f_X\) and \(\phi_M(f_M)=\Gamma_M f_M\) for loading matrices \(\Gamma_X\in\mathbb{R}^{p\times \bar p}\) and \(\Gamma_M\in\mathbb{R}^{q\times \bar q}\).

We next define the causal estimands at the latent-factor level. Let \(\boldsymbol{f}_M^{(a)}\) denote the mediator factors that would be observed if the treatment were set to \(a\), and let \(Y^{(a,f_M)}\) denote the potential outcome if the treatment were set to \(a\) and the mediator factors were set to the value \(f_M\). We use the shorthand \(Y^{(a)}=Y^{(a,\boldsymbol{f}_M^{(a)})}\). For treatment level \(1\) relative to \(0\), the total effect decomposes as \[\textstyle \label{eq:NDENIE} \underbrace{\mathbb{E}\{Y^{(1)}-Y^{(0)}\}}_{\text{Total effect}} = \underbrace{\mathbb{E}\! \left[ Y^{(1,\boldsymbol{f}_M^{(1)})}-Y^{(1,\boldsymbol{f}_M^{(0)})}\right]}_{\text{Natural indirect effect}}+ \underbrace{\mathbb{E}\! \left[ Y^{(1,\boldsymbol{f}_M^{(0)})}-Y^{(0,\boldsymbol{f}_M^{(0)})}\right]}_{\text{Natural direct effect}} .\tag{3}\] Identification of parameters of the form \(\theta(a,a')=\mathbb{E}\{Y^{(a,\boldsymbol{f}_M^{(a')})}\}\) for any \(a,a'\in\{0,1\}\) yields identifying the natural indirect effect, natural direct effect, and total effect through 3 . Specifically, \(\mathrm{NIE}=\theta(1,1)-\theta(1,0)\), \(\mathrm{NDE}=\theta(1,0)-\theta(0,0)\), and \(\mathrm{TE}=\theta(1,1)-\theta(0,0)\). Hence, in what follows, we focus on the cross-world mean \[\theta_0:=\theta(1,0) =\mathbb{E}\!\bigl[Y^{(1,\boldsymbol{f}_M^{(0)})}\bigr].\]

To identify \(\theta_0\), we impose the following standard causal assumptions.

The first part rules out unmeasured confounding of the treatment–outcome and treatment–mediator relationships after conditioning on the latent covariate factors. The second part is the usual cross-world condition for natural direct and indirect effects; it rules out unmeasured mediator–outcome confounding not captured by \(\boldsymbol{f}_X\) and excludes treatment-induced mediator–outcome confounders.

Under Assumptions 1–3, the mediation functional is identified by the mediation formula [2], [3]. Let \(\mu_1(f_X,f_M):=\mathbb{E}\{Y\mid A=1,\boldsymbol{f}_X=f_X,\boldsymbol{f}_M=f_M\}.\) Then \[\label{theta:functional} \begin{align} \theta_0 &=\mathbb{E}\!\left[ \mathbb{E}\!\left\{ \mathbb{E}\left(Y\mid A=1,\boldsymbol{f}_X,\boldsymbol{f}_M\right) \mid A=0,\boldsymbol{f}_X \right\} \right]. \end{align}\tag{4}\]

A direct plug-in estimator based on 4 can be sensitive to nuisance-model misspecification. Following the semiparametric theory of [28], we instead use the efficient influence function-based estimator of \(\theta_0\). Define the nuisance functions \[\textstyle \label{eq:nuisance} \begin{align} \pi_1(f_X) &:=\Pr(A=1\mid \boldsymbol{f}_X=f_X), & \mu_1(f_X,f_M) &:=\mathbb{E}\{Y\mid A=1,\boldsymbol{f}_X=f_X,\boldsymbol{f}_M=f_M\}, \\ \pi_2(f_X,f_M) &:=\frac{p(f_M\mid A=0, f_X)}{p(f_M\mid A=1, f_X)}, & \mu_{10}(f_X) &:=\mathbb{E}\{\mu_1(\boldsymbol{f}_X,\boldsymbol{f}_M) \mid A=0,\boldsymbol{f}_X=f_X\}. \end{align}\tag{5}\] With these definitions, \(\theta_0=\mathbb{E}\{\mu_{10}(\boldsymbol{f}_X)\}\), and the efficient influence function is \[\label{eq:EIF} \textstyle \begin{align} \psi_0(O) &=\mu_{10}(\boldsymbol{f}_X)-\theta_0 +\frac{A}{\pi_1(\boldsymbol{f}_X)}\pi_2(\boldsymbol{f}_X,\boldsymbol{f}_M) \{Y-\mu_1(\boldsymbol{f}_X,\boldsymbol{f}_M)\}\\ &\quad+ \frac{1-A}{1-\pi_1(\boldsymbol{f}_X)} \{\mu_1(\boldsymbol{f}_X,\boldsymbol{f}_M)-\mu_{10}(\boldsymbol{f}_X)\} \,, \end{align}\tag{6}\] where \(O = (Y, A, \boldsymbol{f}_X, \boldsymbol{f}_M)\). In the oracle setting where \((\boldsymbol{f}_X,\boldsymbol{f}_M)\) are observed, replacing the nuisance functions in 6 by estimates and solving the corresponding estimating equation gives \[\label{eq:IF} \textstyle \begin{align} \hat{\theta}^{\mathrm{IF}}_0 =\mathbb{P}_n\bigg[&\hat{\mu}_{10}(\boldsymbol{f}_X) +\frac{A}{\hat{\pi}_1(\boldsymbol{f}_X)}\hat{\pi}_2(\boldsymbol{f}_X,\boldsymbol{f}_M) \{Y-\hat{\mu}_1(\boldsymbol{f}_X,\boldsymbol{f}_M)\}\\ &\quad+ \frac{1-A}{1-\hat{\pi}_1(\boldsymbol{f}_X)} \{\hat{\mu}_1(\boldsymbol{f}_X,\boldsymbol{f}_M)-\hat{\mu}_{10}(\boldsymbol{f}_X)\} \bigg], \end{align}\tag{7}\] where \(\mathbb{P}_n\) denotes the empirical average. However, in reality, we only observe \((X, M)\), (ultra)-high-dimensional proxies of \((\boldsymbol{f}_X, \boldsymbol{f}_M)\), and the latent factors in this oracle estimator are replaced by learned representations, as elaborated in the next section.

3 Proposed Method: MediEncoder↩︎

We now describe MediEncoder, a cross-fitted procedure for estimating the mediation functional \(\theta_0=\theta(1,0)=\mathbb{E}\{Y^{(1,\boldsymbol{f}_M^{(0)})}\}\). The method has two main components. First, it learns low-dimensional nonlinear representations of the high-dimensional covariates and mediators through a coupled autoencoder. Second, it uses these learned representations to estimate the nuisance functions in the efficient-influence function estimator from Section 2.

3.0.0.1 Coupled representation learning.

Let \(z_{XM}=(g_X,g_M,g_{XM},h_X,h_M)\in\mathcal{Z}_{\mathrm{NN}},\) where \(g_X:\mathbb{R}^p\to\mathbb{R}^{\tilde{p}}\) and \(h_X:\mathbb{R}^{\tilde{p}}\to\mathbb{R}^p\) are the covariate encoder and decoder, \(g_M:\mathbb{R}^q\to\mathbb{R}^{\tilde{q}}\) and \(h_M:\mathbb{R}^{\tilde{q}}\to\mathbb{R}^q\) are the mediator encoder and decoder, and \(g_{XM}:\{0,1\}\times\mathbb{R}^{\tilde{p}}\to\mathbb{R}^{\tilde{q}}\) maps the treatment and covariate representation to the mediator-representation space. For a training index set \(\mathcal{I}\), define the empirical loss \[\textstyle \label{eq:medi95loss} \begin{align} \mathcal{L}_{\lambda}(z_{XM};\mathcal{I}) =\frac{1}{|\mathcal{I}|}\sum_{i\in\mathcal{I}} \Big[& \lambda_1\|X_i-h_X(g_X(X_i))\|_2^2 +\lambda_2\|M_i-h_M(g_M(M_i))\|_2^2 \\ &+\lambda_3\|g_M(M_i)-g_{XM}(A_i,g_X(X_i))\|_2^2 \Big], \end{align}\tag{8}\] where \(\lambda=(\lambda_1,\lambda_2,\lambda_3)\) consists of nonnegative tuning parameters. The first two terms are reconstruction losses for \(X\) and \(M\). The third term couples the two representations by encouraging the mediator representation to be predictable from the treatment and the covariate representation, matching the structural dependence of the latent mediator factors on \((A,\boldsymbol{f}_X)\). Unlike a standard multi-view autoencoder, the coupling is directional and treatment-aware: the mediator representation is encouraged to be predictable from \((A,g_X(X))\), obeying the latent structural equation for \(\boldsymbol{f}_M\). After representation learning, only the encoders \(\hat{g}_X\) and \(\hat{g}_M\) are used in the influence function-based estimator; the decoders and \(g_{XM}\) act as regularizers during training.

The working dimensions \((\tilde{p},\tilde{q})\) need not equal the true latent dimensions \((\bar p,\bar q)\). The theoretical results in Section 4 require the learned representation to approximate a sufficiently informative, left-invertible transformation of the latent factors. Thus, exact recovery of the original factors is not required, but the chosen representation dimensions must be large enough to retain the information relevant for the mediation functional.

3.0.0.2 Cross-fitted estimation.

Algorithm 3 gives the full procedure for estimating \(\theta_0\). The data are split into four disjoint subsets. In each fold, these subsets play four distinct roles: representation training, tuning, nuisance estimation, and final evaluation. This separation ensures that the summands used to evaluate the influence function-based estimator are not used to train either the representation or the nuisance functions. The particular rotations below are one convenient choice; any rotations with the same disjointness property can be used.

The estimator in [eq:fold95estimator95method] is the empirical analogue of the efficient influence function representation in Section 2, with the latent factors replaced by fold-specific learned representations. The auxiliary propensity score \(\rho(f_X,f_M)=\Pr(A=1\mid \boldsymbol{f}_X=f_X,\boldsymbol{f}_M=f_M)\) is used only to estimate the mediator density ratio \(\pi_2\). Indeed, \[\begin{align}\label{eq:pi295ratio} \frac{p(f_M\mid A=0,f_X)}{p(f_M\mid A=1,f_X)} = \frac{\Pr(A=0\mid f_X,f_M)}{\Pr(A=1\mid f_X,f_M)} \frac{\Pr(A=1\mid f_X)}{\Pr(A=0\mid f_X)}, \end{align}\tag{9}\] which yields 11 after substituting \(\rho\) and \(\pi_1\).

The same construction estimates any mediation functional \(\theta(a,a')=\mathbb{E}\{Y^{(a,\boldsymbol{f}_M^{(a')})}\}\). Consequently, \[\widehat{\mathrm{NIE}}=\hat{\theta}^{\mathrm{IF}}(1,1)-\hat{\theta}^{\mathrm{IF}}(1,0),\qquad \widehat{\mathrm{NDE}}=\hat{\theta}^{\mathrm{IF}}(1,0)-\hat{\theta}^{\mathrm{IF}}(0,0),\qquad \widehat{\mathrm{TE}}=\hat{\theta}^{\mathrm{IF}}(1,1)-\hat{\theta}^{\mathrm{IF}}(0,0).\] Section 4 establishes conditions under which \(\hat{\theta}_0^{\mathrm{IF}}\) is \(\sqrt n\)-consistent and asymptotically normal.

4 Theoretical Analysis↩︎

In this section, we present theoretical results establishing the \(\sqrt n\)-consistency and asymptotic normality of the proposed estimator \(\hat{\theta}^{\rm IF}_0\). As is clear from Algorithm 3, the estimation procedure relies on two key ingredients: (i) constructing \((\hat{\boldsymbol{f}}_X,\hat{\boldsymbol{f}}_M)\), a surrogate of the latent factors \((\boldsymbol{f}_X,\boldsymbol{f}_M)\), from the high-dimensional observations \((X,M)\); and (ii) estimating the nuisance functions that enter the influence function representation by regressing either \(Y\) or \(A\) on these learned representations. Moreover, as discussed in the Section 1, we do not assume that the true factor dimensions are known. Instead, we only require an upper bound on them, and our theoretical framework is flexible enough to accommodate this setting. We impose the following regularity condition on the data-generating process, in addition to the standard causal identification assumptions in Assumptions 1–3.

Assumption 4 is mild. It requires the outcome noise and the relevant conditional mean contrast to have finite conditional second moments, which is a basic regularity requirement for establishing a central limit theorem. We further have the following requirement on the estimated factor surrogates.

In Appendix 14 (Proposition [prop:factor95recovery95sufficient]), we give a sufficient condition showing that Equation ?? follows from three interpretable ingredients: (i) the existence of oracle encoders that recover a left-invertible transformation of the latent factors up to denoising error; (ii) a local stability, or quadratic-growth, condition for the population MediEncoder objective around the oracle representation; and (iii) a small excess-risk bound for the fitted MediEncoder. This result shows that the factor condition is not tied to a particular algorithm, but can be verified whenever the representation learning objective is stable, and the learned encoder has small population excess risk. In linear factor models with pervasive factors, diversified projection provides one such example [29], [30], fan2023factor?. For nonlinear factor models, recent autoencoder theory [31] provides related recovery guarantees up to functional transformations. Translating these guarantees into ?? requires additional stability conditions on the learned bottleneck map; we discuss this connection in Appendix 14.

We now present our first main result which shows that if the estimated factor surrogates and the nuisance estimators converge at appropriate rates, then the final estimator \(\hat{\theta}^{\rm IF}\) is \(\sqrt n\)-consistent and asymptotically normal.

Theorem 1 only requires weak rate requirements on the learned representation and nuisance functions irrespective of what algorithm is being used. We now verify these rates under a compositional smoothness assumption for DNN-based estimators as used in Algorithm 3. However, the nuisances are estimated using learned representation \(\widehat{\boldsymbol{f}}\), but their population targets are functions of \(\boldsymbol{f}\) or equivalently \(\nu(\boldsymbol{f})\) as \(\nu\) is left-invertible; hence, the factor learning error \(\delta_{n, f}\) contributes to the nuisance parameter estimation. We start with a definition of the compositional function class, namely Hierarchical Compositional Model (HCM):

HCM includes ordinary Hölder classes as well as additive, single-index, and other low-dimensional compositional models. For instance, if \(\textstyle{h(x)=\sum_{j=1}^d h_j(x_j)}\) or \(h(x)=g(x^\top\beta_0)\) with \(\beta\)-smooth univariate components or link functions, then the effective index is \(\gamma^\star=\beta\) up to smooth linear and aggregation maps, whereas a generic \(d\)-variate Hölder function has \(\gamma^\star=\beta/d\). Thus \(\gamma^\star\) can be much larger than the ambient smoothness index, and DNN-based estimators adaptively achieve faster rates than those dictated by the ambient dimension. To that end, we require the following condition.

Theorem 2 implies that, if \(\gamma_j^\star>1/2\) and \(\delta_{n,f}=o(n^{-1/4})\), then all nuisance functions are estimable at rates faster than \(n^{-1/4}\). Consequently, Equation ?? in Theorem 1 holds, and \(\textstyle{\widehat\theta^{\rm IF}}\) is \(\textstyle{\sqrt{n}}\)-CAN. The rate for \(\mu_{10}\) additionally depends on the rate for \(\mu_1\), because \(\mu_{10}\) is estimated by regressing the estimated responses \(\widehat\mu_1(\widehat{\boldsymbol{f}}_X,\widehat{\boldsymbol{f}}_M)\) on \(\widehat{\boldsymbol{f}}_X\). Our theory extends the results of [32] to nonlinear factor estimation and nuisance estimation under logistic-link losses.

5 Experiments↩︎

5.1 Simulation and Ablation Study↩︎

We evaluate MediEncoder in nonlinear high-dimensional settings generated from the latent mediation model in Section 2. The target is \(\theta_0=\theta(1,0)\), estimated by the cross-fitted efficient influence function-based estimator in 7 . We compare with the baseline representation methods described in Appendix 8. Across \(B=500\) Monte Carlo replications, we report the empirical standard deviation (SD), root mean squared error (RMSE), average length of the nominal 95% confidence interval, and empirical coverage. The full data-generating process is given in Appendix 10.

5pt

Simulation settings. For each replication, we generate \(n\in\{100,300,800,1200,2000,3000\}\) observations and consider \(p+q\in\{1000,3000,5000,10000\}\). The latent dimensions are fixed at \(\bar p=\bar q=5\), while the learned representation dimensions are set to \(\tilde{p}=\tilde{q}=7\). The mediator and outcome noise standard deviations are \(\sigma_M=\sigma_Y=1\); the covariate noise level is set to \(\sigma_X=2\) for the first two dimension settings and \(\sigma_X=1.5\) for the larger settings. Observed covariates and mediators are generated from nonlinear additive Haar-wavelet loading functions with independent Gaussian idiosyncratic noise.

All representation and nuisance models are multilayer perceptrons with ReLU activations. The encoders \(g_X\) and \(g_M\) use two hidden layers of widths \(300\) and \(200\), and the decoders \(h_X\) and \(h_M\) use the reverse architecture. The cross-factor network \(g_{XM}\) uses one hidden layer of width \(50\), giving it lower capacity than the encoders and decoders. For nuisance estimation, we fit \(\pi_1\), \(\rho\), \(\mu_1\), and \(\mu_{10}\) using fully connected networks; the density ratio \(\pi_2\) is then computed from \(\hat{\pi}_1\) and \(\hat{\rho}\) using Bayes’ rule as in Algorithm 3. Networks are trained with Adam using initial learning rate \(10^{-3}\), weight decay \(10^{-3}\), a learning-rate decay factor of \(0.5\) every \(30\) epochs, and at most \(300\) epochs. We use mean squared error for regression losses and binary cross-entropy for classification losses, with early stopping based on validation loss.

Results. Table 1 reports a representative setting with \(p+q=3000\); additional results are given in Appendix 17. In this setting, MediEncoder has the smallest SD, RMSE, and average confidence-interval length for all displayed sample sizes, while maintaining coverage close to the nominal 95% level. The improvement becomes clearer as \(n\) increases: for example, at \(n=3000\), MediEncoder reduces RMSE from \(0.349\) for IMAVAE and \(0.343\) for VAE to \(0.290\). The appendix results show the same trend for moderate and large sample sizes, although very small samples can favor simpler representation learners because deep representations are harder to train stably. Figure 4 shows that the scaled errors are approximately centered and have nearly linear Q–Q behavior, consistent with the asymptotic normal approximation developed in Section 4.

Ablation study on \(\lambda_3\). To assess the role of the alignment loss, we set \(\lambda_3=0\) and compare this ablated version with the fully tuned MediEncoder. The ablation results in Appendix 17 show that removing the alignment term generally increases SD, RMSE, and confidence-interval length, with the largest gap in smaller samples. This supports the use of the cross-factor alignment component for stabilizing the learned mediator representation and improving estimation efficiency, while the advantage naturally becomes less pronounced as the sample size grows.

5.2 Real Data Application: Alzheimer’s Disease Neuroimaging Initiative (ADNI)↩︎

We apply MediEncoder to data from the Alzheimer’s Disease Neuroimaging Initiative (ADNI) [33] to study whether the effect of depressive symptoms on cognitive decline is mediated by DNA methylation. The treatment \(A\) is the binarized Geriatric Depression Scale (GDS), with \(A=1\) indicating \(\mathrm{GDS}>5\), a commonly used threshold for elevated depressive symptoms [34]. The outcome \(Y\) is the Alzheimer’s Disease Assessment Scale–Cognitive Subscale (ADAS-Cog), where higher scores indicate worse cognitive function [35], [36]. The covariates \(X\) include demographic and clinical variables, and the mediators \(M\) are high-dimensional DNA methylation features.

After preprocessing (Appendix 16), the analysis includes \(n=649\) subjects, with \(\dim(X)=171\) and \(\dim(M)=3{,}206\). Table 2 reports the estimated natural direct effect (NDE), natural indirect effect (NIE), and total effect (TE). The total effect estimate is positive, suggesting that elevated depressive symptoms are associated, under the causal assumptions in Section 2, with worse cognitive outcomes. The point estimates indicate that most of the effect operates through the direct pathway, while the estimated DNA-methylation-mediated indirect effect is positive but smaller. The bootstrap interval for the total effect excludes zero, whereas the intervals for the direct and indirect effects include zero, so the mediation finding should be interpreted cautiously.

6 Conclusion↩︎

We introduced MediEncoder, a coupled representation-learning framework for high-dimensional causal mediation analysis with nonlinear latent covariate and mediator structures. By combining a structured autoencoder with a cross-fitted efficient influence function-based estimator, the proposed method targets natural direct and indirect effects while accommodating complex high-dimensional measurements. Across simulations, MediEncoder achieved lower RMSE and more stable inference than competing methods, and the ablation study showed that the alignment term improves estimation accuracy and efficiency. In the ADNI application, we found a positive total effect of depressive symptoms on cognitive decline; the point estimates suggested that most of this effect operated through the direct pathway, while the estimated DNA-methylation-mediated indirect effect was smaller and more uncertain.

Limitations and future work. The proposed method relies on standard causal assumptions for natural mediation analysis, including no unmeasured mediator–outcome confounding. It is also currently developed for binary treatments, requires tuning representation dimensions and network hyperparameters, and the learned mediator representations may be difficult to interpret biologically. These limitations motivate future work on sensitivity analysis for unmeasured mediator–outcome confounding, extensions to continuous or multi-valued treatments, and methods for improving interpretability of the learned mediator representations.

7 Related Work↩︎

High-Dimensional Mediation Analysis. Recent work extends causal mediation analysis to high-dimensional regimes where the number of mediators or covariates may exceed the sample size. Early approaches focus primarily on high-dimensional mediators, employing sparsity assumptions and variable selection techniques such as Lasso to estimate mediation effects [20], [37]–[40]. Subsequent work incorporates interaction effects among treatment, mediators, and covariates, allowing for more flexible modeling of complex mechanisms [41]. More recent developments consider settings where both mediators and covariates are high-dimensional and may interact, proposing estimation and inference procedures that account for their joint structure [19]. Despite these advances, most existing methods rely on linear modeling assumptions and treat observed variables directly, which may be restrictive when the underlying mechanisms are nonlinear and only indirectly observed through noisy measurements.

Factor Models. Factor models provide a natural framework for representing high-dimensional data via a lower-dimensional set of latent variables. Classical linear factor models have been widely used in statistics and econometrics for dimension reduction [29], [30], [42]–[44]. These ideas have also been incorporated into causal inference and mediation analysis to handle high-dimensional covariates or mediators by projecting them onto latent spaces [21]. However, existing approaches predominantly rely on linear factor structures, assuming observed variables are linear combinations of latent factors. In many scientific applications, particularly in biological and genomic settings, the mapping from latent processes to observed measurements is often nonlinear. Recent work has explored nonlinear extensions of factor models using neural networks to allow more flexible representations [22], [31], [45], though their integration with mediation analysis remains limited. In parallel, factor-based methods have also been developed for high-dimensional average treatment effect (ATE) estimation under doubly robust frameworks. [22] proposed a factor-informed AIPW estimator combining factor models and deep learning for high-dimensional causal inference. Their setting can be viewed as an ATE analogue of our problem, but still primarily relies on linear factor structures without considering nonlinear factor generation mechanisms.

Autoencoders and More on Mediation Analysis. Deep representation learning methods, especially autoencoders, provide a flexible approach for learning nonlinear low-dimensional representations of high-dimensional data [24], [25]. They have been widely applied in large-scale settings for nonlinear dimension reduction and feature learning [46], [47]. Variational autoencoders further introduce probabilistic latent-variable models that enable structured and generative representations [48]. Related work on multi-view representation learning, including multi-view autoencoders and deep canonical correlation methods, aims to learn shared representations across multiple data modalities [49], [50]. [26] develop a deep learning-based mediation estimator that is robust to model misspecification. However, their approach learns representations independently and does not explicitly model for high-dimensionality and latent factor structure. In particular, it does not consider a shared representation capturing the dependence between covariates and mediators, which is important in high-dimensional biological systems driven by associated latent factors. [27] propose a variational autoencoder-based mediation framework for indirectly observed mediators; however, their method does not accommodate latent and nonlinear structures in covariates, does not model the structural dependence between the latent representations of covariates and mediators, and does not allow treatment–covariate or treatment–mediator interactions. In addition, their estimation procedure may be sensitive to model misspecification.

8 Overview of Baseline Methods↩︎

We consider several baseline methods for estimating low-dimensional latent representations (factors) from high-dimensional covariates \(X\) and mediators \(M\). These representations are subsequently used in downstream mediation effect estimation.

8.0.0.1 Diversified Projection.

Diversified projection method [30], fan2023factor? constructs a linear surrogate of latent factors using a projection matrix \(W \in \mathbb{R}^{p \times \tilde{r}}\) that satisfies boundedness and significance conditions. The estimated factors are obtained as \[\begin{align} \hat{\boldsymbol{f}}= p^{-1} W^\top x, \end{align}\] which can be viewed as an affine approximation of the true latent factor \(f\): \[\begin{align} \hat{\boldsymbol{f}} = H {\boldsymbol{f}} + \xi, \quad \text{with } H = p^{-1} W^\top B. \end{align}\] Under mild conditions, the signal term \(Hf\) dominates the noise \(\xi\), allowing \(\tilde{f}\) to serve as a consistent proxy for \(f\). This approach avoids explicit estimation of the factor loading matrix and does not require precise knowledge of the factor dimension.

8.0.0.2 Autoencoder.

Autoencoders (AE) learn nonlinear low-dimensional representations by jointly training an encoder \(g(\cdot)\) and a decoder \(h(\cdot)\). Given input \(x\), the latent representation is \(z = g(x)\) and the reconstruction is \(\hat{x} = h(z)\). The model is trained by minimizing the reconstruction loss: \[\begin{align} \min_{g,h} \; \sum_{i=1}^n \| x_i - h(g(x_i)) \|_2^2. \end{align}\] In our setting, autoencoders are applied separately to \(X\) and \(M\) to obtain latent representations, which may fail to capture shared structure across modalities.

8.0.0.3 Variational Autoencoder.

Variational autoencoders (VAE) extend autoencoders by introducing a probabilistic latent-variable model. The encoder defines a variational posterior \(q_\phi(z|x)\), while the decoder specifies a likelihood \(p_\theta(x|z)\). The model is trained by maximizing the evidence lower bound (ELBO), equivalently minimizing: \[\begin{align} \mathcal{L}_{\text{VAE}} = \mathbb{E}_{q_\phi(z|x)}\big[\|x - \hat{x}\|_2^2\big] + \mathrm{KL}\big(q_\phi(z|x) \,\|\, p(z)\big), \end{align}\] where \(p(z)\) is a prior distribution. Similar to AE, VAEs are typically applied independently to \(X\) and \(M\), which may not capture their joint dependence structure.

8.0.0.4 IMAVAE.

IMAVAE [27] extends the VAE framework to mediation analysis with indirectly observed mediators by incorporating treatment and covariates as auxiliary variables. Let \(u = t\) or \(u = (w,t)\) denote the auxiliary input. The model jointly learns an encoder \(q_\phi(z|x,u)\), decoder \(p_\theta(x|z)\), and predictor \(g_\gamma(z,u)\) for the outcome. The objective function is: \[\begin{align} \min_{\theta,\phi,\gamma} \; \Big\{ \alpha \, \mathcal{L}_{\theta,\phi}(\hat{x}, x) - \beta \, \mathcal{L}_{\theta,\phi}(x,u) + \mathcal{L}_{\phi,S,\lambda,\gamma}(\hat{y}, y) \Big\}, \end{align}\] where \(\mathcal{L}_{\theta,\phi}(x,u)\) corresponds to the ELBO term, \(\mathcal{L}_{\theta,\phi}(\hat{x}, x)\) is the reconstruction loss, and \(\mathcal{L}_{\phi,S,\lambda,\gamma}(\hat{y}, y)\) is the prediction loss for the outcome. This framework enables learning latent mediator representations informed by treatment and covariates, but remains tied to a specific generative model.

9 Details of Algorithm 3↩︎

9.1 Nuisance Function Estimation↩︎

• Estimate nuisance functions on \(\mathcal{I}_{\mathrm{nu}}^{(k)}\): \[\label{eq:nuisance95opts95method} \begin{align} \hat{\pi}_1^{(k)} &\in \arg\min_{g\in\mathcal{F}_{\pi}} \sum_{i\in\mathcal{I}_{\mathrm{nu}}^{(k)}} \Big[-A_i\log g(\widehat{\boldsymbol{f}}_{X,i}^{(k)}) -(1-A_i)\log\{1-g(\widehat{\boldsymbol{f}}_{X,i}^{(k)})\}\Big],\\ \hat{\mu}_1^{(k)} &\in \arg\min_{g\in\mathcal{F}_{\mu}} \sum_{\substack{i\in\mathcal{I}_{\mathrm{nu}}^{(k)}\\ A_i=1}} \Big[Y_i-g(\widehat{\boldsymbol{f}}_{X,i}^{(k)},\widehat{\boldsymbol{f}}_{M,i}^{(k)})\Big]^2,\\ \hat{\rho}^{(k)} &\in \arg\min_{g\in\mathcal{F}_{\rho}} \sum_{i\in\mathcal{I}_{\mathrm{nu}}^{(k)}} \Big[-A_i\log g(\widehat{\boldsymbol{f}}_{X,i}^{(k)},\widehat{\boldsymbol{f}}_{M,i}^{(k)}) -(1-A_i)\log\{1-g(\widehat{\boldsymbol{f}}_{X,i}^{(k)},\widehat{\boldsymbol{f}}_{M,i}^{(k)})\}\Big],\\ \hat{\mu}_{10}^{(k)} &\in \arg\min_{g\in\mathcal{F}_{\mu_{10}}} \sum_{\substack{i\in\mathcal{I}_{\mathrm{nu}}^{(k)}\\ A_i=0}} \Big[\hat{\mu}_1^{(k)}(\widehat{\boldsymbol{f}}_{X,i}^{(k)},\widehat{\boldsymbol{f}}_{M,i}^{(k)}) -g(\widehat{\boldsymbol{f}}_{X,i}^{(k)})\Big]^2. \end{align}\tag{10}\]

• Estimate the density ratio \(\pi_2(f_X,f_M)=p(f_M\mid A=0,f_X)/p(f_M\mid A=1,f_X)\) by Bayes’ rule: \[\label{eq:pi295bayes95method} \hat{\pi}_2^{(k)}(u,v) = \frac{1-\hat{\rho}^{(k)}(u,v)}{\hat{\rho}^{(k)}(u,v)} \frac{\hat{\pi}_1^{(k)}(u)}{1-\hat{\pi}_1^{(k)}(u)}.\tag{11}\]

9.2 Tuning Algorithm↩︎

The auxiliary regression \(\widehat m_1^{(\lambda)}\) is used only to select the representation learning tuning parameter \(\lambda\) and is not used in the final influence function-based estimator. After \(\hat{\lambda}\) is selected, Algorithm 3 refits the coupled autoencoder on the fold-specific representation-training set and estimates the nuisance functions on a separate nuisance-estimation set.

10 More on Simulation Study↩︎

10.1 Simulation Settings↩︎

The data generating process is as follows:

1. We generate latent covariates \(\boldsymbol{f}_X \in \mathbb{R}^{\bar{p}}\) from \(\boldsymbol{f}_X \sim \mathcal{U}(-1, 1)^{\bar{p}}\). The treatment variable \(A\) is generated according to a logistic model based on \(\boldsymbol{f}_X\): \[\Pr(A = 1 \mid \boldsymbol{f}_X) = \frac{\exp(\boldsymbol{f}_X^\top \alpha)}{1 + \exp(\boldsymbol{f}_X^\top \alpha)}, \qquad A \sim \mathrm{Bernoulli}\!\left( \frac{\exp(\boldsymbol{f}_X^\top \alpha)}{1 + \exp(\boldsymbol{f}_X^\top \alpha)} \right),\] where \(\alpha \in \mathbb{R}^{\bar{p}}\) is drawn once from \(\mathcal{U}(0,2)^{\bar{p}}\).

2. The latent mediators \(\boldsymbol{f}_M \in \mathbb{R}^{\bar{q}}\) are generated through a nonlinear transformation of \(\boldsymbol{f}_X\), the treatment indicator \(A\), and additive noise: \[\boldsymbol{f}_M = (1 - A)\big( \boldsymbol{\delta}_0 \boldsymbol{f}_X^{\circ 2} + \mathbf{u}_{XM} \big) + A\big( \boldsymbol{\delta}_1 \boldsymbol{f}_X^{\circ 2} + \mathbf{u}_{XM}' \big), \quad \mathbf{u}_{XM},\mathbf{u}'_{XM} \sim \mathcal{N}(0, \Sigma_U),\] where \(\boldsymbol{f}_X^{\circ 2}\) denotes the element-wise square of \(\boldsymbol{f}_X\), \(\boldsymbol{\delta}_0, \boldsymbol{\delta}_1 \in \mathbb{R}^{\bar{q} \times \bar{p}}\) have entries drawn once from \(\mathcal{U}(0.5, 1.5)\), and \(\Sigma_U = R_U \Lambda_U R_U^\top\), where \(R_U \in \mathbb{R}^{\bar q \times \bar q}\) is an orthonormal matrix and \(\Lambda_U = \mathrm{diag}(\lambda_1,\ldots,\lambda_{\bar q})\) with \(\lambda_j \sim \mathcal{U}(1, 2)\).

3. The outcome variable \(Y\) is generated via a nonlinear equation: \[\begin{align} Y &= (1 - A) \Big( \sin(5 \boldsymbol{f}_X)^\top \beta_0 + \log\!\big(1 + \exp(\boldsymbol{f}_X^\top \kappa_0)\big) + (\boldsymbol{f}_M \circ \boldsymbol{f}_X^{\circ 2})^\top \gamma_0 + \epsilon_Y' \Big) \\ &\quad + A \Big( \sin(5 \boldsymbol{f}_X)^\top \beta_1 + \log\!\big(1 + \exp(\boldsymbol{f}_X^\top \kappa_1)\big) + (\boldsymbol{f}_M \circ \boldsymbol{f}_X^{\circ 2})^\top \gamma_1 + \epsilon_Y \Big), \end{align}\] where \(\beta_0, \beta_1, \kappa_0, \kappa_1 \in \mathbb{R}^{\bar{p}}\), \(\gamma_0, \gamma_1 \in \mathbb{R}^{\bar{q}}\), each drawn once from \(\mathcal{U}(0.5, 1.5)\), and \(\epsilon_Y, \epsilon_Y' \sim \mathcal{N}(0, \sigma_Y^2)\).

4. We generate the observed high-dimensional covariates \(X \in \mathbb{R}^{p}\) and mediators \(M \in \mathbb{R}^{q}\) through nonlinear factor loadings with additive noise. Specifically, for each observed coordinate \(j = 1,\dots,p\) and \(k = 1,\dots,q\), we generate \[X_j = \phi_{X,j}(\boldsymbol{f}_X) + \boldsymbol{u}_{X,j}, \qquad M_k = \phi_{M,k}(\boldsymbol{f}_M) + \boldsymbol{u}_{M,k},\] where \(\boldsymbol{f}_X \in \mathbb{R}^{\bar p}\) and \(\boldsymbol{f}_M \in \mathbb{R}^{\bar q}\) denote the latent factors.

   Nonlinear loading families. We consider nonlinear factor loadings for \(g_{X,j}(\cdot)\) and \(g_{M,k}(\cdot)\) in wavelet (see 10.1.1 for more details) based on Haar atoms, and we use an additive structure over latent coordinates: \[g_{X,j}(\boldsymbol{f}_X) = \sum_{i=1}^{\bar p} h_{X,ji}(\boldsymbol{f}_{X_i}), \qquad g_{M,k}(\boldsymbol{f}_M) = \sum_{i=1}^{\bar q} h_{M,ki}(\boldsymbol{f}_{M_i}),\] where the univariate functions \(h_{X,jm}(\cdot)\) and \(h_{M,km}(\cdot)\) are instantiated using the wavelet bases.

5. The idiosyncratic components satisfy \[\mathbf{u}_X \sim \mathcal{N}(0, \sigma_X^2 I_p), \qquad \mathbf{u}_M \sim \mathcal{N}(0, \sigma_M^2 I_q),\] independently of the latent factors.

10.1.1 Wavelet Loading↩︎

In the simulation experiment 5, the loading functions \(g_{X,j}(\cdot)\) and \(g_{M,k}(\cdot)\) are constructed using nonlinear basis expansions. Specifically, we consider nonlinear loadings based on wavelet which follows an additive structure over latent coordinates.

Haar mother wavelet. We adopt the standard Haar mother wavelet \(\psi : \mathbb{R} \to \mathbb{R}\), defined by \[\psi(t) = \begin{cases} 1, & 0 \le t < 1, \\ -1, & 1 \le t < 2, \\ 0, & \text{otherwise}. \end{cases}\] This function satisfies \[\int_{\mathbb{R}} \psi(t)\, dt = 0, \qquad \|\psi\|_{L^2} = \sqrt{2}.\]

Wavelet atoms. For each wavelet atom \(\phi_\ell\), we independently sample an integer-valued scale parameter \(r_\ell \sim \mathrm{Unif}\{r_{\min},\dots,r_{\max}\}.\) To ensure that the wavelet atoms overlap with the latent support, we sample a location \(t_\ell \sim \mathrm{Unif}[x_{\min}, x_{\max}],\) where \([x_{\min}, x_{\max}]\) denotes the range of the latent variables. The corresponding translation parameter is then defined as \(s_\ell = \big\lfloor 2^{r_\ell} t_\ell \big\rfloor.\) Given the sampled pair \((r_\ell, s_\ell)\), the wavelet atom is \(\psi_{r_\ell,s_\ell}(t) = 2^{r_\ell/2}\, \psi\!\left(2^{r_\ell} t - s_\ell\right).\)

Additive nonlinear loading. For each latent coordinate \(i\), we select a collection of \(L\) wavelet atoms \(\{\phi_{\ell}(\cdot)\}_{\ell=1}^{L} \subset \{\psi_{r,s}(\cdot)\}.\) The loading functions are defined as \[\begin{align} g_{X,j}(f_X) = \sum_{i=1}^{\bar p} \sum_{\ell=1}^{L} \lambda_{X,ji\ell}^{\mathrm{wavelet}} \, \phi_{\ell}(f_{X_i}), \qquad g_{M,k}(f_M) = \sum_{i=1}^{\bar q} \sum_{\ell=1}^{L} \lambda_{M,ki\ell}^{\mathrm{wavelet}} \, \phi_{\ell}(f_{M_i}), \end{align}\] where \(\lambda_{X,jm\ell}^{\mathrm{wavelet}} \sim \mathcal{N} \Big(0, \frac{1}{L}\Big), \lambda_{M,km\ell}^{\mathrm{wavelet}} \sim \mathcal{N} \Big(0, \frac{1}{L}\Big),\) drawn once at initialization.

11 Proof of Theorem 1↩︎

For notational simplicity, we present the analysis for a single fold (e.g., \(k=1\), so \(|\mathcal{I}_4|\) is used for evaluation of \(\hat{\theta}_0^{\mathrm{IF}}\) and write \(n = |\mathcal{I}_4|\) for simplicity). The arguments for the remaining folds follow analogously by permuting the roles of \(\mathcal{I}_1,\ldots,\mathcal{I}_4\) (see Algorithm 3).

Decomposition of \(\hat{\theta}_0^{\mathrm{IF}}-\theta_0\). Our goal in this step is to decompose \(\hat{\theta}_0^{\mathrm{IF}}-\theta_0\) as the leading influence-function term and the remainder terms. Although the Algorithm 3 estimates the auxiliary propensity \(\rho(\boldsymbol{f}_X,\boldsymbol{f}_M) := \mathbb{P}(A=1\mid \boldsymbol{f}_X,\boldsymbol{f}_M)\) rather than estimating the mediator density ratio directly, we keep the notation \(\pi_2\) for readability in this part of the proof. Throughout the proof, \(\pi_2\) is understood as the Bayes-rule representation \(\pi_2(\boldsymbol{f}_X,\boldsymbol{f}_M) = \frac{1-\rho(\boldsymbol{f}_X,\boldsymbol{f}_M)}{\rho(\boldsymbol{f}_X,\boldsymbol{f}_M)} \cdot \frac{\pi_1(\boldsymbol{f}_X)}{1-\pi_1(\boldsymbol{f}_X)} .\)

We first claim that the estimator admits the decomposition \[\begin{align} \hat{\theta}_0^{\mathrm{IF}}-\theta_0 = \frac{1}{n}\sum_{i=1}^n \Psi_i + R_0+R_1+R_2+R_3+R_4, \end{align}\] where the oracle influence-function score \[\begin{align} \Psi_i &:= \mu_{10}(\boldsymbol{f}_{X,i}) + \frac{A_i}{\pi_1(\boldsymbol{f}_{X,i})} \pi_2(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}) \bigl( Y_i-\mu_1(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}) \bigr)\\ &\qquad + \frac{1-A_i}{1-\pi_1(\boldsymbol{f}_{X,i})} \bigl( \mu_1(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}) - \mu_{10}(\boldsymbol{f}_{X,i}) \bigr) -\theta_0. \end{align}\] The remainder terms collect the differences between the estimated nuisance quantities and their population counterparts. Specifically, \[\begin{align} R_0 &= \frac{1}{n}\sum_{i=1}^n A_i \pi_2(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}) \left( \frac{1}{\hat{\pi}_1(\hat{\boldsymbol{f}}_{X,i})} - \frac{1}{\pi_1(\boldsymbol{f}_{X,i})} \right) \bigl( Y_i-\mu_1(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}) \bigr), \\ R_1 &= \frac{1}{n}\sum_{i=1}^n A_i \frac{ \hat{\pi}_2(\hat{\boldsymbol{f}}_{X,i},\hat{\boldsymbol{f}}_{M,i}) - \pi_2(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}) }{ \hat{\pi}_1(\hat{\boldsymbol{f}}_{X,i}) } \bigl( Y_i-\mu_1(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}) \bigr), \\ R_2 &= \frac{1}{n}\sum_{i=1}^n \left[ \frac{1-A_i}{1-\hat{\pi}_1(\hat{\boldsymbol{f}}_{X,i})} - A_i \frac{\hat{\pi}_2(\hat{\boldsymbol{f}}_{X,i},\hat{\boldsymbol{f}}_{M,i})}{\hat{\pi}_1(\hat{\boldsymbol{f}}_{X,i})} \right] \bigl( \hat{\mu}_1(\hat{\boldsymbol{f}}_{X,i},\hat{\boldsymbol{f}}_{M,i}) - \mu_1(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}) \bigr), \\ R_3 &= \frac{1}{n}\sum_{i=1}^n (1-A_i) \bigl( \mu_1(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}) - \mu_{10}(\boldsymbol{f}_{X,i}) \bigr) \left( \frac{1}{1-\hat{\pi}_1(\hat{\boldsymbol{f}}_{X,i})} - \frac{1}{1-\pi_1(\boldsymbol{f}_{X,i})} \right), \\ R_4 &= \frac{1}{n}\sum_{i=1}^n \frac{A_i-\hat{\pi}_1(\hat{\boldsymbol{f}}_{X,i})}{1-\hat{\pi}_1(\hat{\boldsymbol{f}}_{X,i})} \bigl( \hat{\mu}_{10}(\hat{\boldsymbol{f}}_{X,i}) - \mu_{10}(\boldsymbol{f}_{X,i}) \bigr). \end{align}\] For readibility, we suppress the arguments of all functions throughout the following derivation. For example, we write \(\hat{\pi}_1 = \hat{\pi}_1(\hat{\boldsymbol{f}}_{X,i}), \pi_1 = \pi_1(\boldsymbol{f}_{X,i}), \hat{\pi}_2 = \hat{\pi}_2(\hat{\boldsymbol{f}}_{X,i},\hat{\boldsymbol{f}}_{M,i}), \pi_2 = \pi_2(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}),\) and similarly for \(\hat{\mu}_1\), \(\mu_1\) and \(\hat{\mu}_{10}\), \(\mu_{10}\).

We now prove this decomposition step by step. By definition of \(\hat{\theta}_0^{\mathrm{IF}}\), we have \[\begin{align} \label{eq:diff} \hat{\theta}_0^{\mathrm{IF}}-\theta_0 &= \frac{1}{n}\sum_{i=1}^n \Biggl[ \hat{\mu}_{10} + \frac{A_i}{\hat{\pi}_1} \hat{\pi}_2 \bigl( Y_i-\hat{\mu}_1 \bigr) + \frac{1-A_i}{1-\hat{\pi}_1} \bigl( \hat{\mu}_1 - \hat{\mu}_{10} \bigr) \Biggr] -\theta_0 . \end{align}\tag{12}\] Introduce \[\begin{align} Z_{1,i} &:= \frac{A_i}{\hat{\pi}_1} \hat{\pi}_2 \bigl( Y_i-\hat{\mu}_1 \bigr) - \frac{A_i}{\pi_1} \pi_2 \bigl( Y_i-\mu_1 \bigr),\\ Z_{0,i} &:= \frac{1-A_i}{1-\hat{\pi}_1} \bigl( \hat{\mu}_1-\hat{\mu}_{10} \bigr) - \frac{1-A_i}{1-\pi_1} \bigl( \mu_1-\mu_{10} \bigr). \end{align}\]

Then 12 can be rewritten as \[\label{eq:diff2} \begin{align} \hat{\theta}_0^{\mathrm{IF}}-\theta_0 &= \frac{1}{n}\sum_{i=1}^n \Psi_i + \frac{1}{n}\sum_{i=1}^n \bigl( \hat{\mu}_{10} - \mu_{10} \bigr) + \frac{1}{n}\sum_{i=1}^n Z_{1,i} + \frac{1}{n}\sum_{i=1}^n Z_{0,i}, \end{align}\tag{13}\] where 13 is obtained by adding and subtracting \(\Psi_i\) and regrouping terms. We next expand \(Z_{1,i}\) in 13 . By direct algebra, \[\begin{align} Z_{1,i} &= \frac{A_i}{\hat{\pi}_1} \hat{\pi}_2 \bigl( Y_i-\hat{\mu}_1 \bigr) - \frac{A_i}{\pi_1} \pi_2 \bigl( Y_i-\mu_1 \bigr) \\ &= A_i \Biggl[ \frac{\hat{\pi}_2}{\hat{\pi}_1} \bigl( Y_i-\hat{\mu}_1 \bigr) - \frac{\pi_2}{\pi_1} \bigl( Y_i-\mu_1 \bigr) \Biggr] \\ &= A_i \Biggl[ \frac{\hat{\pi}_2}{\hat{\pi}_1} \bigl( Y_i-\hat{\mu}_1 \bigr) - \frac{\hat{\pi}_2}{\hat{\pi}_1} \bigl( Y_i-\mu_1 \bigr) + \frac{\hat{\pi}_2}{\hat{\pi}_1} \bigl( Y_i-\mu_1 \bigr) - \frac{\pi_2}{\pi_1} \bigl( Y_i-\mu_1 \bigr) \Biggr] \\ &= A_i \Biggl[ \frac{\hat{\pi}_2}{\hat{\pi}_1} \bigl( \mu_1-\hat{\mu}_1 \bigr) + \left( \frac{\hat{\pi}_2}{\hat{\pi}_1} - \frac{\pi_2}{\pi_1} \right) \bigl( Y_i-\mu_1 \bigr) \Biggr] \\ &= A_i \frac{\hat{\pi}_2}{\hat{\pi}_1} \bigl( \mu_1-\hat{\mu}_1 \bigr) + A_i \left[ \frac{\hat{\pi}_2-\pi_2}{\hat{\pi}_1} + \pi_2 \left( \frac{1}{\hat{\pi}_1} - \frac{1}{\pi_1} \right) \right] \bigl( Y_i-\mu_1 \bigr) \\ &= A_i \frac{\hat{\pi}_2}{\hat{\pi}_1} \bigl( \mu_1-\hat{\mu}_1 \bigr) + A_i \frac{\hat{\pi}_2-\pi_2}{\hat{\pi}_1} \bigl( Y_i-\mu_1 \bigr) + A_i \pi_2 \left( \frac{1}{\hat{\pi}_1} - \frac{1}{\pi_1} \right) \bigl( Y_i-\mu_1 \bigr). \end{align}\] We next expand \(Z_{0,i}\) in 13 . Again, by direct algebra, \[\begin{align} Z_{0,i} &= \hat{\mu}_{10} - \mu_{10} + \frac{1-A_i}{1-\hat{\pi}_1} \bigl( \hat{\mu}_1-\hat{\mu}_{10} \bigr) - \frac{1-A_i}{1-\pi_1} \bigl( \mu_1-\mu_{10} \bigr) \\ &= \hat{\mu}_{10} - \mu_{10} + (1-A_i) \Biggl[ \frac{\hat{\mu}_1-\hat{\mu}_{10}}{1-\hat{\pi}_1} - \frac{\mu_1-\mu_{10}}{1-\pi_1} \Biggr] \\ &= \hat{\mu}_{10} - \mu_{10} + (1-A_i) \Biggl[ \frac{\hat{\mu}_1-\hat{\mu}_{10}}{1-\hat{\pi}_1} - \frac{\mu_1-\mu_{10}}{1-\hat{\pi}_1} + \frac{\mu_1-\mu_{10}}{1-\hat{\pi}_1} - \frac{\mu_1-\mu_{10}}{1-\pi_1} \Biggr] \\ &= \hat{\mu}_{10} - \mu_{10} + (1-A_i) \Biggl[ \frac{ \hat{\mu}_1-\hat{\mu}_{10}-\mu_1+\mu_{10} }{ 1-\hat{\pi}_1 } + (\mu_1-\mu_{10}) \left( \frac{1}{1-\hat{\pi}_1} - \frac{1}{1-\pi_1} \right) \Biggr] \\ &= \hat{\mu}_{10} - \mu_{10} + (1-A_i) \Biggl[ \frac{\hat{\mu}_1-\mu_1}{1-\hat{\pi}_1} - \frac{\hat{\mu}_{10}-\mu_{10}}{1-\hat{\pi}_1} + (\mu_1-\mu_{10}) \left( \frac{1}{1-\hat{\pi}_1} - \frac{1}{1-\pi_1} \right) \Biggr] \\ &= \hat{\mu}_{10} - \mu_{10} - \frac{1-A_i}{1-\hat{\pi}_1} \bigl( \hat{\mu}_{10}-\mu_{10} \bigr) + \frac{1-A_i}{1-\hat{\pi}_1} \bigl( \hat{\mu}_1-\mu_1 \bigr) \\ &\qquad + (1-A_i) (\mu_1-\mu_{10}) \left( \frac{1}{1-\hat{\pi}_1} - \frac{1}{1-\pi_1} \right) \\ &= \left[ 1-\frac{1-A_i}{1-\hat{\pi}_1} \right] (\hat{\mu}_{10}-\mu_{10}) + \frac{1-A_i}{1-\hat{\pi}_1} (\hat{\mu}_1-\mu_1) \\ &\qquad + (1-A_i) (\mu_1-\mu_{10}) \left( \frac{1}{1-\hat{\pi}_1} - \frac{1}{1-\pi_1} \right) \\ &= \frac{A_i-\hat{\pi}_1}{1-\hat{\pi}_1} (\hat{\mu}_{10}-\mu_{10}) + \frac{1-A_i}{1-\hat{\pi}_1} (\hat{\mu}_1-\mu_1) \\ &\qquad + (1-A_i) (\mu_1-\mu_{10}) \left( \frac{1}{1-\hat{\pi}_1} - \frac{1}{1-\pi_1} \right). \end{align}\] Substituting the expansions of \(Z_{1,i}\) and \(Z_{0,i}\) into 13 , we obtain \[\begin{align} \hat{\theta}_0^{\mathrm{IF}}-\theta_0 &= \frac{1}{n}\sum_{i=1}^n \Psi_i \\ &\quad + \frac{1}{n}\sum_{i=1}^n A_i \frac{\hat{\pi}_2}{\hat{\pi}_1} (\mu_1-\hat{\mu}_1) + \frac{1}{n}\sum_{i=1}^n A_i \frac{\hat{\pi}_2-\pi_2}{\hat{\pi}_1} (Y_i-\mu_1) + \frac{1}{n}\sum_{i=1}^n A_i \pi_2 \left( \frac{1}{\hat{\pi}_1}-\frac{1}{\pi_1} \right) (Y_i-\mu_1) \\ &\quad + \frac{1}{n}\sum_{i=1}^n \frac{A_i-\hat{\pi}_1}{1-\hat{\pi}_1} (\hat{\mu}_{10}-\mu_{10}) + \frac{1}{n}\sum_{i=1}^n \frac{1-A_i}{1-\hat{\pi}_1} (\hat{\mu}_1-\mu_1) \\ &\quad + \frac{1}{n}\sum_{i=1}^n (1-A_i) (\mu_1-\mu_{10}) \left( \frac{1}{1-\hat{\pi}_1} - \frac{1}{1-\pi_1} \right) \\ &= \frac{1}{n}\sum_{i=1}^n \Psi_i \\ &\quad + \frac{1}{n}\sum_{i=1}^n A_i \pi_2 \left( \frac{1}{\hat{\pi}_1}-\frac{1}{\pi_1} \right) (Y_i-\mu_1) \\ &\quad + \frac{1}{n}\sum_{i=1}^n A_i \frac{\hat{\pi}_2-\pi_2}{\hat{\pi}_1} (Y_i-\mu_1) \\ &\quad + \frac{1}{n}\sum_{i=1}^n \left[ \frac{1-A_i}{1-\hat{\pi}_1} - A_i \frac{\hat{\pi}_2}{\hat{\pi}_1} \right] (\hat{\mu}_1-\mu_1) \\ &\quad + \frac{1}{n}\sum_{i=1}^n (1-A_i) (\mu_1-\mu_{10}) \left( \frac{1}{1-\hat{\pi}_1} - \frac{1}{1-\pi_1} \right) \\ &\quad + \frac{1}{n}\sum_{i=1}^n \frac{A_i-\hat{\pi}_1}{1-\hat{\pi}_1} (\hat{\mu}_{10}-\mu_{10}). \end{align}\] Comparing with the definitions of \(R_0,R_1,R_2,R_3,R_4\), we obtain \[\begin{align} \label{eq:decomp} \hat{\theta}_0^{\mathrm{IF}}-\theta_0 = \frac{1}{n}\sum_{i=1}^n \Psi_i + R_0+R_1+R_2+R_3+R_4. \end{align}\tag{14}\] Reduction to the remainder terms. To establish the asymptotic linearity of \(\hat{\theta}_0^{\mathrm{IF}}\), it suffices to show that \(\sqrt n\,R_j = o_p(1), ~~ j=0,\ldots,4.\) Multiplying the decomposition in Equation 14 by \(\sqrt n\) yields \[\begin{align} \sqrt n\bigl(\hat{\theta}_0^{\mathrm{IF}}-\theta_0\bigr) &= \frac{1}{\sqrt n}\sum_{i=1}^n \Psi_i + \sqrt n\,R_0+\sqrt n\,R_1+\sqrt n\,R_2+\sqrt n\,R_3+\sqrt n\,R_4 \\ &=: \frac{1}{\sqrt n}\sum_{i=1}^n \Psi_i+T_1+T_2+T_3+T_4+T_5. \end{align}\] Thus, it remains to show that \(T_j=o_p(1)\) for \(j=1,\ldots,5\), in which case \[\sqrt n\bigl(\hat{\theta}_0^{\mathrm{IF}}-\theta_0\bigr) = \frac{1}{\sqrt n}\sum_{i=1}^n \Psi_i + o_p(1).\] We now write the remainder terms explicitly: \[\begin{align} T_1 &= \frac{1}{\sqrt n}\sum_{i=1}^n A_i \pi_2(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}) \left( \frac{1}{\hat{\pi}_1(\hat{\boldsymbol{f}}_{X,i})} - \frac{1}{\pi_1(\boldsymbol{f}_{X,i})} \right) \bigl( Y_i-\mu_1(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}) \bigr), \\ T_2 &= \frac{1}{\sqrt n}\sum_{i=1}^n A_i \frac{ \hat{\pi}_2(\hat{\boldsymbol{f}}_{X,i},\hat{\boldsymbol{f}}_{M,i}) - \pi_2(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}) }{ \hat{\pi}_1(\hat{\boldsymbol{f}}_{X,i}) } \bigl( Y_i-\mu_1(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}) \bigr), \\ T_3 &= \frac{1}{\sqrt n}\sum_{i=1}^n \left[ \frac{1-A_i}{1-\hat{\pi}_1(\hat{\boldsymbol{f}}_{X,i})} - A_i \frac{\hat{\pi}_2(\hat{\boldsymbol{f}}_{X,i},\hat{\boldsymbol{f}}_{M,i})}{\hat{\pi}_1(\hat{\boldsymbol{f}}_{X,i})} \right] \bigl( \hat{\mu}_1(\hat{\boldsymbol{f}}_{X,i},\hat{\boldsymbol{f}}_{M,i}) - \mu_1(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}) \bigr), \\ T_4 &= \frac{1}{\sqrt n}\sum_{i=1}^n (1-A_i) \bigl( \mu_1(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}) - \mu_{10}(\boldsymbol{f}_{X,i}) \bigr) \left( \frac{1}{1-\hat{\pi}_1(\hat{\boldsymbol{f}}_{X,i})} - \frac{1}{1-\pi_1(\boldsymbol{f}_{X,i})} \right), \\ T_5 &= \frac{1}{\sqrt n}\sum_{i=1}^n \frac{A_i-\hat{\pi}_1(\hat{\boldsymbol{f}}_{X,i})}{1-\hat{\pi}_1(\hat{\boldsymbol{f}}_{X,i})} \bigl( \hat{\mu}_{10}(\hat{\boldsymbol{f}}_{X,i}) - \mu_{10}(\boldsymbol{f}_{X,i}) \bigr). \end{align}\] Before proceeding with the analysis, we define the following sigma-algebras: \[\begin{align} \mathcal{G}_n^{T_1} &:= \sigma\!\Big( \{\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}, \boldsymbol{u}_{X,i},\boldsymbol{u}_{M,i}: i \in \mathcal{I}_4\}, \hat{\pi}_1 \Big), \\ \mathcal{G}_n^{T_2} &:= \sigma\!\Big( \{\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}, \boldsymbol{u}_{X,i},\boldsymbol{u}_{M,i} : i \in \mathcal{I}_4\}, \hat{\pi}_1,\hat{\rho} \Big), \\ \mathcal{G}_n^{T_3} &:= \sigma\!\Big( \{\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}, \boldsymbol{u}_{X,i},\boldsymbol{u}_{M,i}: i \in \mathcal{I}_4\}, \hat{\pi}_1,\hat{\rho},\hat{\mu}_1 \Big), \\ \mathcal{G}_n^{T_4} = \mathcal{G}_n^{T_5} &:= \sigma\!\Big( \{\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}, \boldsymbol{u}_{X,i},\boldsymbol{u}_{M,i}: i \in \mathcal{I}_4\}, \hat{\pi}_1,\hat{\mu}_{10} \Big)\,. \end{align}\] Analysis of \(T_1\). We consider a single fold \(\mathcal{I}_4\) for evaluation and \(|\mathcal{I}_4|= n\), and all summations below are taken over \(i\in \mathcal{I}_4\). We first study the conditional expectation of \(T_1\): \[\begin{align} \mathbb{E}[T_1 \mid \mathcal{G}_n^{T_1}] &= \mathbb{E}\!\left[ \frac{1}{\sqrt n}\sum_{i\in \mathcal{I}_4} A_i \pi_2(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}) \left( \frac{1}{\hat{\pi}_1(\hat{\boldsymbol{f}}_{X,i})} - \frac{1}{\pi_1(\boldsymbol{f}_{X,i})} \right) \bigl( Y_i-\mu_1(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}) \bigr) \;\middle|\; \mathcal{G}_n^{T_1} \right] \\ &= \frac{1}{\sqrt n}\sum_{i\in \mathcal{I}_4} \pi_2(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}) \left( \frac{1}{\hat{\pi}_1(\hat{\boldsymbol{f}}_{X,i})} - \frac{1}{\pi_1(\boldsymbol{f}_{X,i})} \right) \mathbb{E}\!\left[ A_i \bigl( Y_i-\mu_1(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}) \bigr) \;\middle|\; \mathcal{G}_n^{T_1} \right] \\ &= \frac{1}{\sqrt n}\sum_{i\in \mathcal{I}_4} \pi_2(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}) \left( \frac{1}{\hat{\pi}_1(\hat{\boldsymbol{f}}_{X,i})} - \frac{1}{\pi_1(\boldsymbol{f}_{X,i})} \right) \\ &\qquad\qquad\cdot \mathbb{E}\!\left[ A_i \mathbb{E}\!\left[ Y_i-\mu_1(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}) \;\middle|\; A_i=1,\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i} \right] \;\middle|\; \mathcal{G}_n^{T_1} \right] \\ &=0. \end{align}\] Therefore, \(\mathbb{E}[T_1 \mid \mathcal{G}_n^{T_1}] =0\). We next study the conditional variance of \(T_1\). Using conditional independence across observations given \(\mathcal{G}_n^{T_1}\), we have: \[\begin{align} \mathrm{Var}(T_1 \mid \mathcal{G}_n^{T_1}) &= \frac{1}{n}\sum_{i\in\mathcal{I}_4} \mathrm{Var}\!\Bigg( A_i \pi_2(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}) \left( \frac{1}{\hat{\pi}_1(\hat{\boldsymbol{f}}_{X,i})} - \frac{1}{\pi_1(\boldsymbol{f}_{X,i})} \right) \bigl( Y_i-\mu_1(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}) \bigr) \;\Bigg|\; \mathcal{G}_n^{T_1} \Bigg) \\ &\le \frac{1}{n}\sum_{i\in\mathcal{I}_4} \pi_2(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i})^2 \left( \frac{1}{\hat{\pi}_1(\hat{\boldsymbol{f}}_{X,i})} - \frac{1}{\pi_1(\boldsymbol{f}_{X,i})} \right)^2 \\ &\qquad\qquad\cdot \mathbb{E}\!\left[ A_i \bigl( Y_i-\mu_1(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}) \bigr)^2 \;\middle|\; \mathcal{G}_n^{T_1} \right] \\ &\le C \frac{1}{n}\sum_{i\in\mathcal{I}_4} \left( \frac{1}{\hat{\pi}_1(\hat{\boldsymbol{f}}_{X,i})} - \frac{1}{\pi_1(\boldsymbol{f}_{X,i})} \right)^2 \\ &= C \frac{1}{n}\sum_{i\in\mathcal{I}_4} \left| \frac{ \pi_1(\boldsymbol{f}_{X,i}) - \hat{\pi}_1(\hat{\boldsymbol{f}}_{X,i}) }{ \hat{\pi}_1(\hat{\boldsymbol{f}}_{X,i}) \pi_1(\boldsymbol{f}_{X,i}) } \right|^2 \\ &\le C \frac{1}{n}\sum_{i\in\mathcal{I}_4} \bigl( \hat{\pi}_1(\hat{\boldsymbol{f}}_{X,i}) - \pi_1(\boldsymbol{f}_{X,i}) \bigr)^2 \\ &= C \left\| \hat{\pi}_1(\hat{\boldsymbol{f}}_X) - \pi_1(\boldsymbol{f}_X) \right\|_{L_2(P_n)}^2 \\ &= C \left\| \hat{\pi}_1(\hat{\boldsymbol{f}}_X) - \pi_1\circ v_X^{-1}(\hat{\boldsymbol{f}}_X) + \pi_1\circ v_X^{-1}(\hat{\boldsymbol{f}}_X) - \pi_1(\boldsymbol{f}_X) \right\|_{L_2(P_n)}^2 \\ &\le C \Big( \big\| \hat{\pi}_1(\hat{\boldsymbol{f}}_X) - \pi_1\circ v_X^{-1}(\hat{\boldsymbol{f}}_X) \big\|_{L_2(P_n)} + \big\| \pi_1\circ v_X^{-1}(\hat{\boldsymbol{f}}_X) - \pi_1(\boldsymbol{f}_X) \big\|_{L_2(P_n)} \Big)^2 \\ &\le C \big\| \hat{\pi}_1(\hat{\boldsymbol{f}}_X) - \pi_1\circ v_X^{-1}(\hat{\boldsymbol{f}}_X) \big\|_{L_2(P_n)}^2 + C \big\| \pi_1\circ v_X^{-1}(\hat{\boldsymbol{f}}_X) - \pi_1(\boldsymbol{f}_X) \big\|_{L_2(P_n)}^2 . \end{align}\] The last inequality follows from Young’s inequality, equivalently \((a+b)^2\le 2a^2+2b^2\). The second inequality above uses \(A_i\le 1\), boundedness of \(\pi_2\), \(\pi_1\) and \(\hat{\pi}_1\), and the bounded conditional second moment in Assumption 4. It remains to bound the first empirical \(L_2( P_n)\) term. We first add and subtract its population counterpart: \[\begin{align} & \big\| \hat{\pi}_1(\hat{\boldsymbol{f}}_X) - \pi_1\circ v_X^{-1}(\hat{\boldsymbol{f}}_X) \big\|_{L_2( P_n)}^2 \\ &\qquad = \big\| \hat{\pi}_1 - \pi_1\circ v_X^{-1} \big\|_{L_2( P)}^2 + \left\{ \big\| \hat{\pi}_1(\hat{\boldsymbol{f}}_X) - \pi_1\circ v_X^{-1}(\hat{\boldsymbol{f}}_X) \big\|_{L_2( P_n)}^2 - \big\| \hat{\pi}_1 - \pi_1\circ v_X^{-1} \big\|_{L_2( P)}^2 \right\}\\ &\qquad = O_p(\delta_{n,\pi_1}^2+n^{-1/2})\,, \end{align}\] where the first term follows from the population \(L_2\) rate in Theorem 1, and the second term is obtained via Lemma 1 applied to the empirical average over the evaluation fold. Combining the above bounds, we obtain \(\mathrm{Var}(T_1\mid \mathcal{G}_n^{T_1}) = O_p(\delta_{n,\pi_1}^2+n^{-1/2}+\delta_{n,f}^2) = o_p(1).\) Together with \(\mathbb{E}[T_1\mid \mathcal{G}_n^{T_1}]=0\), Chebyshev’s inequality gives, for any \(\varepsilon>0\), \[\begin{align} \mathbb{P}(|T_1|>\varepsilon\mid \mathcal{G}_n^{T_1}) \le \frac{\mathrm{Var}(T_1\mid \mathcal{G}_n^{T_1})}{\varepsilon^2} = o_p(1). \end{align}\] Since \(0\le \mathbb{P}(|T_1|>\varepsilon\mid \mathcal{G}_n^{T_1})\le 1\), the dominated convergence argument gives \(\mathbb{P}(|T_1|>\varepsilon) = \mathbb{E}\!\left[ \mathbb{P}(|T_1|>\varepsilon\mid \mathcal{G}_n^{T_1}) \right] \to 0.\) Hence \(T_1=o_p(1)\).

Analysis of \(T_2\). Recall that the density ratio \(\hat{\pi}_2\) is represented as \(\hat{\pi}_2(\hat{\boldsymbol{f}}_X,\hat{\boldsymbol{f}}_M) = \frac{1-\hat{\rho}(\hat{\boldsymbol{f}}_X,\hat{\boldsymbol{f}}_M)}{\hat{\rho}(\hat{\boldsymbol{f}}_X,\hat{\boldsymbol{f}}_M)} \cdot \frac{\hat{\pi}_1(\hat{\boldsymbol{f}}_X)}{1-\hat{\pi}_1(\hat{\boldsymbol{f}}_X)} .\) We first study the conditional expectation of \(T_2\) given \(\mathcal{G}_n^{T_2}\): \[\begin{align} \mathbb{E}[T_2 \mid \mathcal{G}_n^{T_2}] &= \mathbb{E}\!\left[ \frac{1}{\sqrt n}\sum_{i\in \mathcal{I}_4} A_i \frac{ \hat{\pi}_2(\hat{\boldsymbol{f}}_{X,i},\hat{\boldsymbol{f}}_{M,i}) - \pi_2(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}) }{ \hat{\pi}_1(\hat{\boldsymbol{f}}_{X,i}) } \bigl( Y_i-\mu_1(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}) \bigr) \;\middle|\; \mathcal{G}_n^{T_2} \right] \\ &= \frac{1}{\sqrt n}\sum_{i\in \mathcal{I}_4} \frac{ \hat{\pi}_2(\hat{\boldsymbol{f}}_{X,i},\hat{\boldsymbol{f}}_{M,i}) - \pi_2(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}) }{ \hat{\pi}_1(\hat{\boldsymbol{f}}_{X,i}) } \mathbb{E}\!\left[ A_i \bigl( Y_i-\mu_1(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}) \bigr) \;\middle|\; \mathcal{G}_n^{T_2} \right] \\ &= \frac{1}{\sqrt n}\sum_{i\in \mathcal{I}_4} \frac{ \hat{\pi}_2(\hat{\boldsymbol{f}}_{X,i},\hat{\boldsymbol{f}}_{M,i}) - \pi_2(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}) }{ \hat{\pi}_1(\hat{\boldsymbol{f}}_{X,i}) }\\ &~~~~~~~~~~~~~~~~\cdot\mathbb{E}\!\left[ A_i\, \mathbb{E}\!\left[ Y_i-\mu_1(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}) \;\middle|\; A_i=1,\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i} \right] \;\middle|\; \mathcal{G}_n^{T_2} \right] \\ &=0. \end{align}\] We next study the conditional variance of \(T_2\): \[\begin{align} \mathrm{Var}(T_2 \mid \mathcal{G}_n^{T_2}) &= \frac{1}{n}\sum_{i\in \mathcal{I}_4} \mathrm{Var}\!\left( A_i \frac{ \hat{\pi}_2(\hat{\boldsymbol{f}}_{X,i},\hat{\boldsymbol{f}}_{M,i}) - \pi_2(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}) }{ \hat{\pi}_1(\hat{\boldsymbol{f}}_{X,i}) } \bigl( Y_i-\mu_1(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}) \bigr) \;\middle|\; \mathcal{G}_n^{T_2} \right) \\ &\le \frac{1}{n}\sum_{i\in \mathcal{I}_4} \left( \frac{ \hat{\pi}_2(\hat{\boldsymbol{f}}_{X,i},\hat{\boldsymbol{f}}_{M,i}) - \pi_2(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}) }{ \hat{\pi}_1(\hat{\boldsymbol{f}}_{X,i}) } \right)^2 \mathbb{E}\!\left[ A_i \bigl( Y_i-\mu_1(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}) \bigr)^2 \;\middle|\; \mathcal{G}_n^{T_2} \right] \\ &\le C \frac{1}{n}\sum_{i\in \mathcal{I}_4} \bigl( \hat{\pi}_2(\hat{\boldsymbol{f}}_{X,i},\hat{\boldsymbol{f}}_{M,i}) - \pi_2(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}) \bigr)^2 . \end{align}\] The second inequality uses \(A_i\le 1\), positivity of \(\hat{\pi}_1\), and the bounded conditional second moment in Assumption 4. Now, using the ratio representation of \(\pi_2\) together with the uniform boundedness of \((\pi_1, \hat{\pi}_1, \rho, \hat{\rho})\), we obtain \[\begin{align} &\left| \hat{\pi}_2(\hat{\boldsymbol{f}}_{X,i},\hat{\boldsymbol{f}}_{M,i}) - \pi_2(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}) \right| \\ &\le C \Big( \left| \hat{\rho}(\hat{\boldsymbol{f}}_{X,i},\hat{\boldsymbol{f}}_{M,i}) - \rho\circ v^{-1}(\hat{\boldsymbol{f}}_{X,i},\hat{\boldsymbol{f}}_{M,i}) \right| + \left| \hat{\pi}_1(\hat{\boldsymbol{f}}_{X,i}) - \pi_1\circ v_X^{-1}(\hat{\boldsymbol{f}}_{X,i}) \right| \\ &\qquad\qquad+ \left| \rho\circ v^{-1}(\hat{\boldsymbol{f}}_{X,i},\hat{\boldsymbol{f}}_{M,i}) - \rho(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}) \right| + \left| \pi_1\circ v_X^{-1}(\hat{\boldsymbol{f}}_{X,i}) - \pi_1(\boldsymbol{f}_{X,i}) \right| \Big), \end{align}\] where we use that the maps \(x\mapsto (1-x)/x\) and \(x\mapsto x/(1-x)\) are Lipschitz. Therefore, we obtain \[\begin{align} &\frac{1}{n}\sum_{i\in \mathcal{I}_4} \bigl( \hat{\pi}_2(\hat{\boldsymbol{f}}_{X,i},\hat{\boldsymbol{f}}_{M,i}) - \pi_2(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}) \bigr)^2 \\ &\le C \left\| \hat{\rho}(\hat{\boldsymbol{f}}_X,\hat{\boldsymbol{f}}_M) - \rho\circ v^{-1}(\hat{\boldsymbol{f}}_X,\hat{\boldsymbol{f}}_M) \right\|_{L_2( P_n)}^2 + C \left\| \hat{\pi}_1(\hat{\boldsymbol{f}}_X) - \pi_1\circ v_X^{-1}(\hat{\boldsymbol{f}}_X) \right\|_{L_2( P_n)}^2 \\ &\quad+ C \left\| \rho\circ v^{-1}(\hat{\boldsymbol{f}}_X,\hat{\boldsymbol{f}}_M) - \rho(\boldsymbol{f}_X,\boldsymbol{f}_M) \right\|_{L_2( P_n)}^2 + C \left\| \pi_1\circ v_X^{-1}(\hat{\boldsymbol{f}}_X) - \pi_1(\boldsymbol{f}_X) \right\|_{L_2( P_n)}^2 \\ &= O_p(\delta_{n,\rho}^2+n^{-1/2}) + O_p(\delta_{n,\pi_1}^2+n^{-1/2}) + O_p(\delta_{n,f}^2) + O_p(\delta_{n,f}^2) \\ &= O_p(\delta_{n,\rho}^2+\delta_{n,\pi_1}^2+\delta_{n,f}^2+n^{-1/2}). \end{align}\] where the inequality follows \((a+b+c+d)^2\le C(a^2+b^2+c^2+d^2)\). The first two rates follow by adding and subtracting their population \(L_2( P)\) counterparts, using the nuisance rates in Theorem 1, and applying Lemma 1. The last two rates follow from the Lipschitz continuity of \(\rho\) and \(\pi_1\) together with the factor-recovery rate in Theorem 1. Combining the above bounds, we obtain \(\mathrm{Var}(T_2\mid \mathcal{G}_n^{T_2}) = O_p(\delta_{n,\rho}^2+\delta_{n,\pi_1}^2+\delta_{n,f}^2+n^{-1/2})= o_p(1).\) Together with \(\mathbb{E}[T_2\mid \mathcal{G}_n^{T_2}]=0\), Chebyshev’s inequality gives, for any \(\varepsilon>0\), \[\begin{align} \mathbb{P}(|T_2|>\varepsilon \mid \mathcal{G}_n^{T_2}) \le \frac{\mathrm{Var}(T_2\mid \mathcal{G}_n^{T_2})}{\varepsilon^2} = o_p(1). \end{align}\] Hence \(T_2=o_p(1)\) by the similar argument as in proof of \(T_1\).

Analysis of \(T_3\). Consider \(T_3\), which can be expressed as \[\begin{align} T_3 &= \frac{1}{\sqrt n}\sum_{i\in \mathcal{I}_4} \frac{ (1-A_i) - A_i \frac{1-\hat{\rho}(\hat{\boldsymbol{f}}_{X,i},\hat{\boldsymbol{f}}_{M,i})}{\hat{\rho}(\hat{\boldsymbol{f}}_{X,i},\hat{\boldsymbol{f}}_{M,i})} }{1-\hat{\pi}_1(\hat{\boldsymbol{f}}_{X,i})} \bigl( \hat{\mu}_1(\hat{\boldsymbol{f}}_{X,i},\hat{\boldsymbol{f}}_{M,i}) - \mu_1(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}) \bigr) \\ &= \frac{1}{\sqrt n}\sum_{i\in \mathcal{I}_4} \frac{ (1-A_i) - A_i \frac{1-\rho(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i})}{\rho(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i})} }{1-\pi_1(\boldsymbol{f}_{X,i})} \bigl( \hat{\mu}_1(\hat{\boldsymbol{f}}_{X,i},\hat{\boldsymbol{f}}_{M,i}) - \mu_1(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}) \bigr) \\ &\quad+ \frac{1}{\sqrt n}\sum_{i\in \mathcal{I}_4} \left[ \frac{ (1-A_i) - A_i \frac{1-\hat{\rho}(\hat{\boldsymbol{f}}_{X,i},\hat{\boldsymbol{f}}_{M,i})}{\hat{\rho}(\hat{\boldsymbol{f}}_{X,i},\hat{\boldsymbol{f}}_{M,i})} }{1-\hat{\pi}_1(\hat{\boldsymbol{f}}_{X,i})} - \frac{ (1-A_i) - A_i \frac{1-\rho(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i})}{\rho(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i})} }{1-\pi_1(\boldsymbol{f}_{X,i})} \right] \\ &\qquad\qquad\cdot \bigl( \hat{\mu}_1(\hat{\boldsymbol{f}}_{X,i},\hat{\boldsymbol{f}}_{M,i}) - \mu_1(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}) \bigr) \\ &=: T_{31}+T_{32}. \end{align}\] We first study \(T_{31}\). Then \[\begin{align} \mathbb{E}[T_{31}\mid \mathcal{G}_n^{T_3}] &= \frac{1}{\sqrt n}\sum_{i\in \mathcal{I}_4} \frac{ \hat{\mu}_1(\hat{\boldsymbol{f}}_{X,i},\hat{\boldsymbol{f}}_{M,i}) - \mu_1(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}) }{1-\pi_1(\boldsymbol{f}_{X,i})} \\ &\qquad\qquad\cdot \mathbb{E}\!\left[ (1-A_i) - A_i \frac{1-\rho(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i})}{\rho(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i})} \;\middle|\; \mathcal{G}_n^{T_3} \right] \\ &= \frac{1}{\sqrt n}\sum_{i\in \mathcal{I}_4} \frac{ \hat{\mu}_1(\hat{\boldsymbol{f}}_{X,i},\hat{\boldsymbol{f}}_{M,i}) - \mu_1(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}) }{1-\pi_1(\boldsymbol{f}_{X,i})} \\ &\qquad\qquad\cdot \mathbb{E}\!\left[ \mathbb{E}\!\left[ (1-A_i) - A_i \frac{1-\rho(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i})}{\rho(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i})} \;\middle|\; \boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i} \right] \;\middle|\; \mathcal{G}_n^{T_3} \right] \\ &= \frac{1}{\sqrt n}\sum_{i\in \mathcal{I}_4} \frac{ \hat{\mu}_1(\hat{\boldsymbol{f}}_{X,i},\hat{\boldsymbol{f}}_{M,i}) - \mu_1(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}) }{1-\pi_1(\boldsymbol{f}_{X,i})} \cdot 0 \\ &=0, \end{align}\] where \[\begin{align} \mathbb{E}\!\left[ (1-A_i) - A_i \frac{1-\rho(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i})}{\rho(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i})} \;\middle|\; \boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i} \right]&= 1-\rho(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}) - \rho(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}) \frac{1-\rho(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i})}{\rho(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i})}\\ &=0. \end{align}\] Therefore, \(\mathbb{E}[T_{31}\mid \mathcal{G}_n^{T_3}]=0.\) Next, Given \(\mathcal{G}_n^{T_3}\), \[\begin{align} \mathrm{Var}(T_{31}\mid \mathcal{G}_n^{T_3}) &\le C\frac{1}{n}\sum_{i\in \mathcal{I}_4} \bigl( \hat{\mu}_1(\hat{\boldsymbol{f}}_{X,i},\hat{\boldsymbol{f}}_{M,i}) - \mu_1(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}) \bigr)^2 \\ &= C \left\| \hat{\mu}_1(\hat{\boldsymbol{f}}_X,\hat{\boldsymbol{f}}_M) - \mu_1(\boldsymbol{f}_X,\boldsymbol{f}_M) \right\|_{L_2( P_n)}^2 \\ &\le C \left\| \hat{\mu}_1(\hat{\boldsymbol{f}}_X,\hat{\boldsymbol{f}}_M) - \mu_1\circ v^{-1}(\hat{\boldsymbol{f}}_X,\hat{\boldsymbol{f}}_M) \right\|_{L_2( P_n)}^2 + C \left\| \mu_1\circ v^{-1}(\hat{\boldsymbol{f}}_X,\hat{\boldsymbol{f}}_M) - \mu_1(\boldsymbol{f}_X,\boldsymbol{f}_M) \right\|_{L_2( P_n)}^2 \\ &= O_p(\delta_{n,\mu_1}^2+n^{-1/2}) + O_p(\delta_{n,f}^2) \\ &= O_p(\delta_{n,\mu_1}^2+\delta_{n,f}^2+n^{-1/2}) = o_p(1). \end{align}\] Together with \(\mathbb{E}[T_{31}\mid \mathcal{G}_n^{T_3}]=0\), Chebyshev’s inequality and the same dominated-convergence argument as in the proof of \(T_1\) imply \(T_{31}=o_p(1)\).

We next study \(T_{32}\). Recall that \[\begin{align} T_{32} = \frac{1}{\sqrt n}\sum_{i\in\mathcal{I}_4} Z_{1,i}Z_{2,i}, \end{align}\] where \[\begin{align} Z_{1,i} &:= \frac{ (1-A_i) - A_i \frac{1-\hat{\rho}(\hat{\boldsymbol{f}}_{X,i},\hat{\boldsymbol{f}}_{M,i})}{\hat{\rho}(\hat{\boldsymbol{f}}_{X,i},\hat{\boldsymbol{f}}_{M,i})} }{1-\hat{\pi}_1(\hat{\boldsymbol{f}}_{X,i})} - \frac{ (1-A_i) - A_i \frac{1-\rho(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i})}{\rho(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i})} }{1-\pi_1(\boldsymbol{f}_{X,i})}, \\ Z_{2,i} &:= \hat{\mu}_1(\hat{\boldsymbol{f}}_{X,i},\hat{\boldsymbol{f}}_{M,i}) - \mu_1(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}). \end{align}\] By Cauchy–Schwarz, \[\begin{align} |T_{32}| &\le \frac{1}{\sqrt n} \left( \sum_{i\in\mathcal{I}_4}Z_{1,i}^2 \right)^{1/2} \left( \sum_{i\in\mathcal{I}_4}Z_{2,i}^2 \right)^{1/2} = \sqrt n \left( \frac{1}{n}\sum_{i\in\mathcal{I}_4}Z_{1,i}^2 \right)^{1/2} \left( \frac{1}{n}\sum_{i\in\mathcal{I}_4}Z_{2,i}^2 \right)^{1/2}. \end{align}\] By positivity and Lipschitz continuity of the maps \(x\mapsto (1-x)/x\) and \(x\mapsto 1/(1-x)\), \[\begin{align} |Z_{1,i}| &\le C \left| \hat{\rho}(\hat{\boldsymbol{f}}_{X,i},\hat{\boldsymbol{f}}_{M,i}) - \rho\circ v^{-1}(\hat{\boldsymbol{f}}_{X,i},\hat{\boldsymbol{f}}_{M,i}) \right| + C \left| \rho\circ v^{-1}(\hat{\boldsymbol{f}}_{X,i},\hat{\boldsymbol{f}}_{M,i}) - \rho(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}) \right| \\ &\quad+ C \left| \hat{\pi}_1(\hat{\boldsymbol{f}}_{X,i}) - \pi_1\circ v_X^{-1}(\hat{\boldsymbol{f}}_{X,i}) \right| + C \left| \pi_1\circ v_X^{-1}(\hat{\boldsymbol{f}}_{X,i}) - \pi_1(\boldsymbol{f}_{X,i}) \right|. \end{align}\] Therefore, using \((a+b+c+d)^2\le C(a^2+b^2+c^2+d^2)\), Lemma 1, and Assumptions in Theorem 1, we have: \[\begin{align} \left( \frac{1}{n}\sum_{i\in\mathcal{I}_4}Z_{1,i}^2 \right)^{1/2} &\le C \left\| \hat{\rho}(\hat{\boldsymbol{f}}_X,\hat{\boldsymbol{f}}_M) - \rho\circ v^{-1}(\hat{\boldsymbol{f}}_X,\hat{\boldsymbol{f}}_M) \right\|_{L_2( P_n)} \\ &\quad+ C \left\| \rho\circ v^{-1}(\hat{\boldsymbol{f}}_X,\hat{\boldsymbol{f}}_M) - \rho(\boldsymbol{f}_X,\boldsymbol{f}_M) \right\|_{L_2( P_n)} \\ &\quad+ C \left\| \hat{\pi}_1(\hat{\boldsymbol{f}}_X) - \pi_1\circ v_X^{-1}(\hat{\boldsymbol{f}}_X) \right\|_{L_2( P_n)} \\ &\quad+ C \left\| \pi_1\circ v_X^{-1}(\hat{\boldsymbol{f}}_X) - \pi_1(\boldsymbol{f}_X) \right\|_{L_2( P_n)} \\ &= O_p(\delta_{n,\rho}+n^{-1/2}) + O_p(\delta_{n,f}) + O_p(\delta_{n,\pi_1}+n^{-1/2}) + O_p(\delta_{n,f}) \\ &= O_p(\delta_{n,\rho}+\delta_{n,\pi_1}+\delta_{n,f}+n^{-1/2})\,. \end{align}\] Similarly, \[\begin{align} \left( \frac{1}{n}\sum_{i\in\mathcal{I}_4}Z_{2,i}^2 \right)^{1/2} &\le C \left\| \hat{\mu}_1(\hat{\boldsymbol{f}}_X,\hat{\boldsymbol{f}}_M) - \mu_1\circ v^{-1}(\hat{\boldsymbol{f}}_X,\hat{\boldsymbol{f}}_M) \right\|_{L_2( P_n)} \\ &\quad+ C \left\| \mu_1\circ v^{-1}(\hat{\boldsymbol{f}}_X,\hat{\boldsymbol{f}}_M) - \mu_1(\boldsymbol{f}_X,\boldsymbol{f}_M) \right\|_{L_2( P_n)} \\ &= O_p(\delta_{n,\mu_1}+n^{-1/2}) + O_p(\delta_{n,f}) \\ &= O_p(\delta_{n,\mu_1}+\delta_{n,f}+n^{-1/2}), \end{align}\] again by Lemma 1. Combining the above bounds gives \[\begin{align} |T_{32}| &= O_p\!\left[ \sqrt n (\delta_{n,\rho}+\delta_{n,\pi_1}+\delta_{n,f}+n^{-1/2}) (\delta_{n,\mu_1}+\delta_{n,f}+n^{-1/2}) \right] \\ &= o_p(1), \end{align}\] Combining the results for \(T_{31}\) and \(T_{32}\), we conclude \(T_{3}=o_p(1).\)

Analysis of \(T_4\). We first study the conditional expectation of \(T_4\) given \(\mathcal{G}_n^{T_4}\): \[\begin{align} \mathbb{E}[T_4 \mid \mathcal{G}_n^{T_4}] &= \frac{1}{\sqrt n}\sum_{i\in\mathcal{I}_4} \left( \frac{1}{1-\hat{\pi}_1(\hat{\boldsymbol{f}}_{X,i})} - \frac{1}{1-\pi_1(\boldsymbol{f}_{X,i})} \right) \\ &\qquad\qquad\cdot \mathbb{E}\!\left[ (1-A_i) \bigl( \mu_1(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}) - \mu_{10}(\boldsymbol{f}_{X,i}) \bigr) \;\middle|\; \mathcal{G}_n^{T_4} \right] \\ &= \frac{1}{\sqrt n}\sum_{i\in\mathcal{I}_4} \left( \frac{1}{1-\hat{\pi}_1(\hat{\boldsymbol{f}}_{X,i})} - \frac{1}{1-\pi_1(\boldsymbol{f}_{X,i})} \right) \\ &\qquad\qquad\cdot \mathbb{E}\!\left[ \mathbb{E}\!\left[ (1-A_i) \bigl( \mu_1(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}) - \mu_{10}(\boldsymbol{f}_{X,i}) \bigr) \;\middle|\; A_i,\boldsymbol{f}_{X,i} \right] \;\middle|\; \mathcal{G}_n^{T_4} \right] \\ &=0. \end{align}\] The last equality follows from the definition \(\mu_{10}(\boldsymbol{f}_X) = \mathbb{E}[ \mu_1(\boldsymbol{f}_X,\boldsymbol{f}_M) \mid A=0,\boldsymbol{f}_X ]\), which implies \[\begin{align} &\mathbb{E}\!\left[ (1-A_i) \bigl( \mu_1(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}) - \mu_{10}(\boldsymbol{f}_{X,i}) \bigr) \;\middle|\; A_i,\boldsymbol{f}_{X,i} \right] \\ &\qquad= (1-A_i) \mathbb{E}\!\left[ \mu_1(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}) - \mu_{10}(\boldsymbol{f}_{X,i}) \;\middle|\; A_i,\boldsymbol{f}_{X,i} \right] =0. \end{align}\] We next study the conditional variance: \[\begin{align} \mathrm{Var}(T_4 \mid \mathcal{G}_n^{T_4}) &\le \frac{1}{n}\sum_{i\in\mathcal{I}_4} \mathbb{E}\!\left[ (1-A_i) \bigl( \mu_1(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}) - \mu_{10}(\boldsymbol{f}_{X,i}) \bigr)^2 \left( \frac{1}{1-\hat{\pi}_1(\hat{\boldsymbol{f}}_{X,i})} - \frac{1}{1-\pi_1(\boldsymbol{f}_{X,i})} \right)^2 \;\middle|\; \mathcal{G}_n^{T_4} \right] \\ &\le C\frac{1}{n}\sum_{i\in\mathcal{I}_4} \left( \frac{1}{1-\hat{\pi}_1(\hat{\boldsymbol{f}}_{X,i})} - \frac{1}{1-\pi_1(\boldsymbol{f}_{X,i})} \right)^2 \\ &\le C \frac{1}{n}\sum_{i\in\mathcal{I}_4} \bigl( \hat{\pi}_1(\hat{\boldsymbol{f}}_{X,i}) - \pi_1(\boldsymbol{f}_{X,i}) \bigr)^2 \\ &= C \left\| \hat{\pi}_1(\hat{\boldsymbol{f}}_X) - \pi_1(\boldsymbol{f}_X) \right\|_{L_2( P_n)}^2 \\ &\le C \left\| \hat{\pi}_1(\hat{\boldsymbol{f}}_X) - \pi_1\circ v_X^{-1}(\hat{\boldsymbol{f}}_X) \right\|_{L_2( P_n)}^2 + C \left\| \pi_1\circ v_X^{-1}(\hat{\boldsymbol{f}}_X) - \pi_1(\boldsymbol{f}_X) \right\|_{L_2( P_n)}^2 \\ &= O_p(\delta_{n,\pi_1}^2+n^{-1/2}) + O_p(\delta_{n,f}^2) \\ &= O_p(\delta_{n,\pi_1}^2+\delta_{n,f}^2+n^{-1/2}) = o_p(1). \end{align}\] The second inequality uses \(1-A_i\le 1\), boundedness of \(\pi_1\) and \(\hat{\pi}_1\), and the bounded conditional second moment in Assumption 4. Together with \(\mathbb{E}[T_4\mid \mathcal{G}_n^{T_4}]=0\), Chebyshev’s inequality and the same bounded-convergence-in-probability argument as in the proof of \(T_1\) imply \(T_4=o_p(1)\).

Analysis of \(T_5\). Consider \(T_5\), which can be expressed as \[\begin{align} T_5 &= \frac{1}{\sqrt n}\sum_{i\in \mathcal{I}_4} \frac{A_i-\hat{\pi}_1(\hat{\boldsymbol{f}}_{X,i})}{1-\hat{\pi}_1(\hat{\boldsymbol{f}}_{X,i})} \bigl( \hat{\mu}_{10}(\hat{\boldsymbol{f}}_{X,i}) - \mu_{10}(\boldsymbol{f}_{X,i}) \bigr) \\ &= \frac{1}{\sqrt n}\sum_{i\in \mathcal{I}_4} \frac{A_i-\pi_1(\boldsymbol{f}_{X,i})}{1-\pi_1(\boldsymbol{f}_{X,i})} \bigl( \hat{\mu}_{10}(\hat{\boldsymbol{f}}_{X,i}) - \mu_{10}(\boldsymbol{f}_{X,i}) \bigr) \\ &\quad+ \frac{1}{\sqrt n}\sum_{i\in \mathcal{I}_4} \left[ \frac{A_i-\hat{\pi}_1(\hat{\boldsymbol{f}}_{X,i})}{1-\hat{\pi}_1(\hat{\boldsymbol{f}}_{X,i})} - \frac{A_i-\pi_1(\boldsymbol{f}_{X,i})}{1-\pi_1(\boldsymbol{f}_{X,i})} \right] \\ &\qquad\qquad\cdot \bigl( \hat{\mu}_{10}(\hat{\boldsymbol{f}}_{X,i}) - \mu_{10}(\boldsymbol{f}_{X,i}) \bigr) \\ &=: T_{51}+T_{52}. \end{align}\]

We first study \(T_{51}\). Then \[\begin{align} \mathbb{E}[T_{51}\mid \mathcal{G}_n^{T_5}] &= \frac{1}{\sqrt n}\sum_{i\in \mathcal{I}_4} \frac{ \hat{\mu}_{10}(\hat{\boldsymbol{f}}_{X,i}) - \mu_{10}(\boldsymbol{f}_{X,i}) }{1-\pi_1(\boldsymbol{f}_{X,i})} \mathbb{E}\!\left[ A_i-\pi_1(\boldsymbol{f}_{X,i}) \;\middle|\; \mathcal{G}_n^{T_5} \right] \\ &= \frac{1}{\sqrt n}\sum_{i\in \mathcal{I}_4} \frac{ \hat{\mu}_{10}(\hat{\boldsymbol{f}}_{X,i}) - \mu_{10}(\boldsymbol{f}_{X,i}) }{1-\pi_1(\boldsymbol{f}_{X,i})} \mathbb{E}\!\left[ \mathbb{E}\!\left[ A_i-\pi_1(\boldsymbol{f}_{X,i}) \;\middle|\; \boldsymbol{f}_{X,i} \right] \;\middle|\; \mathcal{G}_n^{T_5} \right] \\ &=0. \end{align}\] We next study the conditional variance: \[\begin{align} \mathrm{Var}(T_{51}\mid \mathcal{G}_n^{T_5}) &\le C\frac{1}{n}\sum_{i\in \mathcal{I}_4} \bigl( \hat{\mu}_{10}(\hat{\boldsymbol{f}}_{X,i}) - \mu_{10}(\boldsymbol{f}_{X,i}) \bigr)^2 \\ &= C \left\| \hat{\mu}_{10}(\hat{\boldsymbol{f}}_X) - \mu_{10}(\boldsymbol{f}_X) \right\|_{L_2( P_n)}^2 \\ &\le C \left\| \hat{\mu}_{10}(\hat{\boldsymbol{f}}_X) - \mu_{10}\circ v_X^{-1}(\hat{\boldsymbol{f}}_X) \right\|_{L_2( P_n)}^2 + C \left\| \mu_{10}\circ v_X^{-1}(\hat{\boldsymbol{f}}_X) - \mu_{10}(\boldsymbol{f}_X) \right\|_{L_2( P_n)}^2 \\ &= O_p(\delta_{n,\mu_{10}}^2+n^{-1/2}) + O_p(\delta_{n,f}^2) \\ &= O_p(\delta_{n,\mu_{10}}^2+\delta_{n,f}^2+n^{-1/2}) = o_p(1). \end{align}\] Together with \(\mathbb{E}[T_{51}\mid \mathcal{G}_n^{T_5}]=0\), Chebyshev’s inequality and the same bounded-convergence-in-probability argument as in the proof of \(T_1\) imply \(T_{51}=o_p(1)\).

We next study \(T_{52}\). Let \(Z_{1,i} := \frac{A_i-\hat{\pi}_1(\hat{\boldsymbol{f}}_{X,i})}{1-\hat{\pi}_1(\hat{\boldsymbol{f}}_{X,i})} - \frac{A_i-\pi_1(\boldsymbol{f}_{X,i})}{1-\pi_1(\boldsymbol{f}_{X,i})},\) and \(Z_{2,i} := \hat{\mu}_{10}(\hat{\boldsymbol{f}}_{X,i}) - \mu_{10}(\boldsymbol{f}_{X,i}).\) Then, we write \[\begin{align} T_{52} = \frac{1}{\sqrt n}\sum_{i\in\mathcal{I}_4} Z_{1,i}Z_{2,i}. \end{align}\] By Cauchy–Schwarz, \[\begin{align} |T_{52}| &\le \frac{1}{\sqrt n} \left( \sum_{i\in\mathcal{I}_4} Z_{1,i}^2 \right)^{1/2} \left( \sum_{i\in\mathcal{I}_4} Z_{2,i}^2 \right)^{1/2} = \sqrt n \left( \frac{1}{n}\sum_{i\in\mathcal{I}_4} Z_{1,i}^2 \right)^{1/2} \left( \frac{1}{n}\sum_{i\in\mathcal{I}_4} Z_{2,i}^2 \right)^{1/2}. \end{align}\] By positivity and Lipschitz continuity of the map \(x\mapsto (A_i-x)/(1-x)\) on the relevant range, \[\begin{align} |Z_{1,i}| &\le C \left| \hat{\pi}_1(\hat{\boldsymbol{f}}_{X,i}) - \pi_1(\boldsymbol{f}_{X,i}) \right| \\ &\le C \left| \hat{\pi}_1(\hat{\boldsymbol{f}}_{X,i}) - \pi_1\circ v_X^{-1}(\hat{\boldsymbol{f}}_{X,i}) \right| + C \left| \pi_1\circ v_X^{-1}(\hat{\boldsymbol{f}}_{X,i}) - \pi_1(\boldsymbol{f}_{X,i}) \right|. \end{align}\] Therefore, \[\begin{align} \left( \frac{1}{n}\sum_{i\in\mathcal{I}_4} Z_{1,i}^2 \right)^{1/2} &\le C \left\| \hat{\pi}_1(\hat{\boldsymbol{f}}_X) - \pi_1\circ v_X^{-1}(\hat{\boldsymbol{f}}_X) \right\|_{L_2( P_n)} + C \left\| \pi_1\circ v_X^{-1}(\hat{\boldsymbol{f}}_X) - \pi_1(\boldsymbol{f}_X) \right\|_{L_2( P_n)} \\ &= O_p(\delta_{n,\pi_1}+n^{-1/2}) + O_p(\delta_{n,f}) \\ &= O_p(\delta_{n,\pi_1}+\delta_{n,f}+n^{-1/2}). \end{align}\]

Similarly, \[\begin{align} |Z_{2,i}| &\le \left| \hat{\mu}_{10}(\hat{\boldsymbol{f}}_{X,i}) - \mu_{10}\circ v_X^{-1}(\hat{\boldsymbol{f}}_{X,i}) \right| + \left| \mu_{10}\circ v_X^{-1}(\hat{\boldsymbol{f}}_{X,i}) - \mu_{10}(\boldsymbol{f}_{X,i}) \right|, \end{align}\] and hence \[\begin{align} \left( \frac{1}{n}\sum_{i\in\mathcal{I}_4} Z_{2,i}^2 \right)^{1/2} &\le C \left\| \hat{\mu}_{10}(\hat{\boldsymbol{f}}_X) - \mu_{10}\circ v_X^{-1}(\hat{\boldsymbol{f}}_X) \right\|_{L_2(\mathbb{P}_n)} + C \left\| \mu_{10}\circ v_X^{-1}(\hat{\boldsymbol{f}}_X) - \mu_{10}(\boldsymbol{f}_X) \right\|_{L_2(\mathbb{P}_n)} \\ &= O_p(\delta_{n,\mu_{10}}+n^{-1/2}) + O_p(\delta_{n,f}) \\ &= O_p(\delta_{n,\mu_{10}}+\delta_{n,f}+n^{-1/2}). \end{align}\] Combining the above bounds gives \[\begin{align} |T_{52}| &= O_p\!\left[ \sqrt n (\delta_{n,\pi_1}+\delta_{n,f}+n^{-1/2}) (\delta_{n,\mu_{10}}+\delta_{n,f}+n^{-1/2}) \right] = o_p(1). \end{align}\] we obtain \(T_{52}=o_p(1)\). Together with \(T_{51}=o_p(1)\), this implies \(T_5=o_p(1).\)

Asymptotics. Combining the bounds for \(T_1,\ldots,T_5\), we obtain \[\begin{align} \sqrt n\bigl(\hat{\theta}_0^{\mathrm{IF}}-\theta_0\bigr) = \frac{1}{\sqrt n}\sum_{i=1}^n \Psi_i + o_p(1). \end{align}\] Since \(\mathbb{E}[\Psi_i]=0\) and \(\mathrm{Var}(\Psi_i)=\sigma_{\mathrm{eff}}^2<\infty\) by the central limit theorem, \[\begin{align} \frac{1}{\sqrt n}\sum_{i=1}^n \Psi_i \overset{d}{\longrightarrow} \mathcal{N}(0,\sigma_{\mathrm{eff}}^2). \end{align}\] Applying Slutsky’s theorem yields \[\begin{align} \sqrt n\bigl(\hat{\theta}_0^{\mathrm{IF}}-\theta_0\bigr) \overset{d}{\longrightarrow} \mathcal{N}(0,\sigma_{\mathrm{eff}}^2). \end{align}\]

12 Proof of Lemma 1↩︎

From definition we have: \[\left|\|f\|_n - \|f\|_2\right| = \frac{\left|\|f\|^2_n - \|f\|^2_2\right|}{\|f\|_n + \|f\|_2} \le \frac{\left|\|f\|^2_n - \|f\|^2_2\right|}{\|f\|_2}\] As a consequence, we have: \[\begin{align} \bbP\left(\left|\|f\|_n - \|f\|_2\right| > t\right) & \le \bbP\left(\left|\|f\|^2_n - \|f\|^2_2\right| > t\|f\|_2\right) \\ & = \bbP\left(\left|\frac{1}{n} \sum_{i = 1}^n (f^2(X_i) - \bbE[f^2(X)])\right| > t\|f\|_2\right) \\ & \le \frac{\var(f^2(X))}{nt^2 \|f\|_2^2} \\ & \le \frac{\bbE(f^4(X))}{nt^2 \|f\|_2^2} \\ & \le \frac{B^2\bbE(f^2(X))}{nt^2 \|f\|_2^2} = \frac{B^2}{nt^2} \,. \end{align}\] The last inequality follows from the fact that \(\|f\|_\infty \le B\). This completes the proof.

13 Proof of Theorem 2↩︎

The proof mostly follows using a similar argument as in the proof of Theorem 1 of fan2023factor?. We highlight the key parts here. As the proof for \(\hat{\pi}_1\) and \(\hat{\rho}\) are analogous, we skip the latter for brevity and present the proof of the other three nuisances here.

13.1 Rate for \(\widehat\mu_1\) under the transformed factor representation↩︎

Recall our notation \(\boldsymbol{f}= (\boldsymbol{f}_X, \boldsymbol{f}_M)\) denotes the unobserved factors, and \(\hat{\boldsymbol{f}}= (\hat{\boldsymbol{f}}_X, \hat{\boldsymbol{f}}_M)\) denotes the estimated factors, i.e. \(\hat{\boldsymbol{f}}_X = \hat{g}_X(X)\) and \(\hat{\boldsymbol{f}}_M = \hat{g}_M(M)\). By the regression model, we have: \[Y_i = \mu_1(\boldsymbol{f}_i) + \eps_i\] where \(\bbE[\eps_i \mid \boldsymbol{f}_i, \bu_{X, i}, \bu_{M, i}, A] = 0\) and \(\eps_i\) is sub-gaussian random variable. As \(\hat{\boldsymbol{f}}\) is a function of \((\boldsymbol{f}, \bu_X, \bu_X)\), and the encoder is estimated from a separate dataset, it is immediate that \(\bbE[\eps \mid \hat{\boldsymbol{f}}, A] = 0\). We use the notation \(\eta = \mu_1 \circ \nu^{-1}\). As per Theorem 2, \(\eta\) is bounded, Lipschitz, and moreover \(\eta \in \cH(d_\eta, \ell_\eta, \cP)\) with DAS \(\gamma^*_\eta\). Let \(\cI_1 = \{i \in \cI_{\rm nuis}: A_i = 1\}\). As per our assumption, \[\|\hat{\boldsymbol{f}}- \nu(\boldsymbol{f})\|_{L_2(P_1)} = O_p(\delta_{n,f}),\] Let \(\mathcal{G}_{\eta,n} := \cG(L_\eta,d_\eta,1,N_\eta,M_\eta,B_\eta)\) be a collection of truncated ReLU network class used to estimate \(\eta_1\), where \(d_\eta = \widetilde{p}+\widetilde{q}\) be the input dimension, \(L_\eta\) is the depth, \(N_\eta\) is the width, \(M_\eta\) is the truncation level, and \(B_\eta\) is the weight bound. We take \(M_\eta\) larger than \(\|\eta\|_\infty\). The estimator \(\hat{\eta}_1\) is defined as: \[\widehat\eta \in \arg\min_{g\in\mathcal{G}_{\eta,n}} \frac{1}{n_1} \sum_{i\in \cI_1} \left\{ Y_i-g(\hat{\boldsymbol{f}}_i) \right\}^2.\] By Theorem 4 of [32], since \(\eta \in \cH(d_\eta, \ell_\eta, \cP)\) and is supported on a bounded domain, there exist constants \(C_1,C_2,C_3,C_4>0\), depending only on the HCM parameters, such that for any integer \(N_\mu\ge 2\), there exists a ReLU network \[\eta^\dagger \in \mathcal{G}(C_1,d_\eta,1,C_2N_\eta,\infty,C_3N_\eta^{C_4})\] satisfying \[\|\eta^\dagger-\eta\|_{\infty} \le C N_\eta^{-2\gamma_\eta^\star}. \] After truncation at level \(M_\eta\), the approximation error does not increase, provided \(M_\eta>\|\eta\|_\infty\). Therefore, by enlarging constants, we may regard \(\eta^\dagger\in \mathcal{G}_{\eta,n}\) and retain the same approximation bound. Consequently, \[\label{eq:nn95approx95mu951} \|\eta^\dagger-\eta\|_{\infty}^2 \le C N_\eta^{-4\gamma_\eta^\star}.\tag{15}\] However, our estimator is evaluated at the learned representation \(\hat{\boldsymbol{f}}\), whereas the target is \(\mu_1(\boldsymbol{f}) = \eta(\nu(\boldsymbol{f}))\). Thus, the relevant oracle error is \[\|\eta^\dagger(\hat{\boldsymbol{f}})-\eta( \nu(\boldsymbol{f}))\|_{L_2(P_1)}^2.\] Using \((a+b)^2\le 2a^2+2b^2\), we have \[\begin{align} \mathbb{E}_1 \left[ \left\{ \eta^\dagger(\hat{\boldsymbol{f}})-\eta(\nu(\boldsymbol{f})) \right\}^2 \right] \le 2\mathbb{E}_1 \left[ \left\{ \eta^\dagger(\hat{\boldsymbol{f}})-\eta(\hat{\boldsymbol{f}}) \right\}^2 \right] + 2\mathbb{E}_1 \left[ \left\{ \eta(\hat{\boldsymbol{f}})-\eta(\nu(\boldsymbol{f})) \right\}^2 \right]. \end{align}\] The first term is controlled by the uniform approximation bound: \[\mathbb{E}_1 \left[ \left\{ \eta^\dagger(\hat{\boldsymbol{f}})-\eta(\hat{\boldsymbol{f}}) \right\}^2 \right] \le \|\eta^\dagger-\eta\|_\infty^2 \le C N_\eta^{-4\gamma_\eta^\star}.\] For the second term, since \(\eta\) is Lipschitz, \[\mathbb{E}_1 \left[ \left\{ \eta(\hat{\boldsymbol{f}})-\eta(\nu(\boldsymbol{f})) \right\}^2 \right] \le L_\eta^2 \mathbb{E}_1\|\hat{\boldsymbol{f}}- \nu(\boldsymbol{f})\|_2^2 = O_p(\delta_{n,f}^2).\] Therefore, \[\begin{align} \label{eq:approx95bound95dnn951} \|\eta^\dagger(\hat{\boldsymbol{f}})-\eta(\nu(\boldsymbol{f}))\|_{L_2(P_1)}^2 = O_p\left( N_\eta^{-4\gamma_\eta^\star} + \delta_{n,f}^2 \right). \end{align}\tag{16}\] Next, we bound the stochastic error. For any \(g\in\mathcal{G}_{\eta,n}\), define \(r_g(i) := g(\hat{\boldsymbol{f}}_i)-\eta(\nu(\boldsymbol{f}_i)).\) In particular, \[r_{\widehat \eta}(i) = \widehat \eta(\hat{\boldsymbol{f}}_i)-\eta(\nu(\boldsymbol{f}_i)), \qquad r_{\eta^\dagger}(i) = \eta^\dagger(\hat{\boldsymbol{f}}_i)-\eta(\nu(\boldsymbol{f}_i)).\] Since \(Y_i=\eta(\nu(\boldsymbol{f}_i))+\varepsilon_i,\) the empirical loss can be written as \[\frac{1}{n_1}\sum_{i\in\cI_1} \{Y_i-g(\hat{\boldsymbol{f}}_i)\}^2 = \frac{1}{n_1}\sum_{i\in\cI_1} \{\varepsilon_i-r_g(i)\}^2.\] By the empirical-risk minimizing property of \(\widehat\eta\), \[\frac{1}{n_1}\sum_{i\in\cI_1} \{\varepsilon_i-r_{\widehat\eta_1}(i)\}^2 \le \frac{1}{n_1}\sum_{i\in\cI_1} \{\varepsilon_i-r_{\eta^\dagger}(i)\}^2.\] Expanding both sides and canceling the common term \(n^{-1}\sum_i\varepsilon_i^2\), we obtain the basic inequality \[\begin{align} \label{eq:bound951} \|r_{\widehat\eta_1}\|_{n,1}^2 & \le \|r_{\eta^\dagger}\|_{n,1}^2 + \frac{2}{n_1}\sum_{i\in\cI_1} \varepsilon_i\left\{\widehat\eta_1(\hat{\boldsymbol{f}}_i)-\eta^\dagger(\hat{\boldsymbol{f}}_i)\right\}, \notag \\ & \le \|r_{\eta^\dagger}\|_{n,1}^2 +2\sup_{g\in\mathcal{G}_{\eta,n}} \left|\frac{1}{n_1}\sum_{i\in \cI_1} \varepsilon_i \left\{g(\hat{\boldsymbol{f}}_i)-\eta^\dagger(\hat{\boldsymbol{f}}_i)\right\}\right|, \end{align}\tag{17}\] where \[\|h\|_{n,1}^2 := \frac{1}{n_1}\sum_{i\in \cI_1}h_i^2.\] Thus, it remains to control the supremum of the empirical process. By Lemma 4 of [32], the suprema can be bounded in terms of the pseudo-dimension of the network class. Let \(V_n = \mathrm{Pdim}(\cG_{\eta, n})\) denotes the pseudo-dimension of \(\cG_{\eta, n}\). Then with probability at least \(1-e^{-t}\), \[\begin{align} \left| \frac{1}{n_1}\sum_{i\in \cI_1} \varepsilon_i \left\{ g(\hat{\boldsymbol{f}}_i)-\eta^\dagger(\hat{\boldsymbol{f}}_i) \right\} \right| \le C \left( \|g(\hat{\boldsymbol{f}})-\eta^\dagger(\hat{\boldsymbol{f}})\|_{n,1} + \sqrt{\frac{V_n \log{n_1}}{n_1}} \right) \sqrt{ \frac{V_n \log{n_1}}{n_1}+\frac{t}{n} } \end{align}\] simultaneously for all \(g\in\mathcal{G}_{\eta,n}\). Applying this inequality to \(g=\widehat\eta\) gives \[\begin{align} \left| \frac{1}{n_1}\sum_{i\in \cI_1} \varepsilon_i \left\{ \widehat\eta_1(\hat{\boldsymbol{f}}_i)-\eta^\dagger(\hat{\boldsymbol{f}}_i) \right\} \right| & \le C \left( \|\widehat\eta_1(\hat{\boldsymbol{f}})-\eta^\dagger(\hat{\boldsymbol{f}})\|_{n,1} + \sqrt{\frac{V_n \log{n_1}}{n_1}} \right) \sqrt{ \frac{V_n \log{n_1}}{n_1}+\frac{t}{n} } \\ & \le C \left( \|r_{\widehat\eta_1}\|_{n,1} + \|r_{\eta^\dagger}\|_{n,1} + \sqrt{\frac{V_n \log{n_1}}{n_1}} \right) \sqrt{ \frac{V_n \log{n_1}}{n_1}+\frac{t}{n} } \end{align}\] where the last line follows from the fact that: \[\|\widehat\eta_1(\hat{\boldsymbol{f}})-\eta^\dagger(\hat{\boldsymbol{f}})\|_{n,1} \le \|r_{\widehat\eta_1}\|_{n,1} + \|r_{\eta^\dagger}\|_{n,1} \,.\] This, along with Equation 17 yields: \[\begin{align} \|r_{\widehat\eta_1}\|^2_{n,1} \le \|r_{\eta^\dagger}\|^2_{n,1} + C \left( \|r_{\widehat\eta_1}\|_{n,1} + \|r_{\eta^\dagger}\|_{n,1} + \sqrt{\frac{V_n \log{n_1}}{n_1}} \right) \sqrt{ \frac{V_n \log{n_1}}{n_1}+\frac{t}{n} } \,. \end{align}\] We then apply Young’s inequality and change sides to obtain: \[\|r_{\widehat\eta_1}\|_{n,1}^2 \le C \left(\|r_{\eta^\dagger}\|^2_{n,1} + \frac{V_n \log{n_1}}{n_1}+\frac{t}{n}\right).\] Conditional on the representation-learning fold, \(r_{\eta^\dagger}\) is a fixed bounded function of the nuisance-fold observation. Therefore, by Bernstein’s inequality applied to the fixed oracle \(r_{\eta^\dagger}\), \[\|r_{\eta^\dagger}\|_{n,1}^2 \le C \|r_{\eta^\dagger}\|_{L_2(P_1)}^2 + C\frac{t}{n}\] with probability at least \(1-e^{-t}\). This, along with the previous equation, yields with probability \(\ge 1 - 2e^{-t}\): \[\label{eq:bound95295dnn} \|r_{\widehat\eta_1}\|_{n,1}^2 \lesssim \left(\|r_{\eta^\dagger}\|^2_{L_2(P_1)} + \frac{V_n \log{n_1}}{n_1}+\frac{t}{n}\right)\tag{18}\] It remains to bound \(V_\mu=\operatorname{Pdim}(\mathcal{G}_{\mu,n})\). By Theorem 7 of [51], we have: \[V_n \le C W_\eta L_\eta \log{W_\eta} \le C_1 N_\eta^2 L^2_\eta \log{(N^2_\eta L_\eta)} \,.\] Choosing \(L_\eta \asymp 1\) and \(N_\eta \asymp n^{1/(4\gamma_\eta^* + 2)}\), we conclude from Equation 16 and 18 : \[\|r_{\widehat\eta_1}\|_{n,1}^2 = O_p\left( n^{-\frac{2\gamma_\eta^*}{2\gamma_\eta^* + 1}}(\log{n})^a + \delta_{n, f}^2 \right) \,.\] Finally, applying the empirical-to-population norm comparison, as in Lemma 3 and the last step of the proof of Theorem 1 of [32], gives the same bound for the population norm: \[\|r_{\widehat\eta_1}\|_{L_2(P_1)}^2 = O_p\left( n^{-\frac{2\gamma_\eta^*}{2\gamma_\eta^* + 1}}(\log{n})^a + \delta_{n, f}^2 \right) \,.\]

13.2 Rate for \(\widehat \pi_1\) under the transformed factor representation↩︎

We now prove the rate for the treatment propensity nuisance \(\pi_1(\boldsymbol{f}_X):=\bbP(A=1\mid \boldsymbol{f}_X).\) Define as before, the relevant nuisance parameter, \[\eta := \pi_1\circ \nu^{-1}, \qquad \eta(\nu(\boldsymbol{f}_X))=\pi_1(\boldsymbol{f}_X).\] Let \(\cI_\pi=\cI_{\rm nuis}\) be the nuisance-estimation fold and write \(n=|\cI_\pi|\) for simplicity. Assume \[\|\hat{\boldsymbol{f}}_X-\nu(\boldsymbol{f}_X)\|_{L_2(P)} = O_p(\delta_{n,f}).\] From our assumption in Theorem 1, the propensity score is uniformly bounded away from \(0\) and \(1\), i.e., for some \(\kappa\in(0,1/2)\), \[\kappa \le \eta(\nu(\boldsymbol{f}_X)) \le 1-\kappa\] almost surely. We restrict the neural-network estimators to take values in \([\kappa/2,1-\kappa/2]\) by clipping. This does not affect the rate. Suppose \(\eta\) admits a bounded Lipschitz HCM extension to a compact set containing both \(\nu(\boldsymbol{f}_X)\) and \(\hat{\boldsymbol{f}}_X\) with probability tending to one. Let \(\eta\in \cH(d_\eta,\ell_\eta,\cP)\) with DAS \(\gamma_\eta^*\), where \(d_\eta=\widetilde{p}\) is the input dimension. Let \[\mathcal{G}_{\eta,n} := \cG(L_\eta,d_\eta,1,N_\eta,M_\eta,B_\eta)\] be the clipped ReLU network class. As per Line 7 of Algorithm 3 \[\widehat\eta \in \arg\min_{g\in\mathcal{G}_{\eta,n}} \frac{1}{n}\sum_{i\in\cI_\pi} \ell(A_i,g(\hat{\boldsymbol{f}}_{X,i})) := \arg\min_{g\in\mathcal{G}_{\eta,n}} \;\bbP_n(\ell(A, g(\hat{\boldsymbol{f}}_X)) \,\] where \[\ell(a,u) := -a\log u-(1-a)\log(1-u), \qquad u\in[\kappa/2,1-\kappa/2].\] As in the previous subsection, we can use the approximation theorem (e.g., Theorem 4 of [32]), there exists \[\eta^\dagger \in \cG(C_1,d_\eta,1,C_2N_\eta,\infty,C_3N_\eta^{C_4})\] such that \[\|\eta^\dagger-\eta\|_\infty \le C N_\eta^{-2\gamma_\eta^\star}.\] After clipping, we may take \(\eta^\dagger\in\mathcal{G}_{\eta,n}\) and retain \[\label{eq:pi195nn95approx95shorter} \|\eta^\dagger-\eta\|_\infty^2 \le C N_\eta^{-4\gamma_\eta^\star}.\tag{19}\] As in the proof for \(\widehat\mu_1\), the relevant oracle error is \[\|\eta^\dagger(\hat{\boldsymbol{f}}_X)-\eta(\nu(\boldsymbol{f}_X))\|_{L_2(P)}^2.\] Using the same decomposition as in the previous subsection, and the Lipschitzness of \(\eta\), \[\begin{align} \|\eta^\dagger(\hat{\boldsymbol{f}}_X)-\eta(\nu(\boldsymbol{f}_X))\|_{L_2(P)}^2 \lesssim \|\eta^\dagger-\eta\|_\infty^2 + \|\hat{\boldsymbol{f}}_X-\nu(\boldsymbol{f}_X)\|_{L_2(P)}^2 \,. \end{align}\] Therefore, \[\label{eq:pi195oracle95shorter} \|\eta^\dagger(\hat{\boldsymbol{f}}_X)-\eta(\nu(\boldsymbol{f}_X))\|_{L_2(P)}^2 = O_p\left( N_\eta^{-4\gamma_\eta^\star} + \delta_{n,f}^2 \right).\tag{20}\] By empirical risk minimization, \[\bbP_n\ell(A,\widehat\eta(\hat{\boldsymbol{f}}_X)) \le \bbP_n\ell(A,\eta^\dagger(\hat{\boldsymbol{f}}_X)).\] Adding and subtracting population risks gives \[\begin{align} \label{eq:pi195basic95shorter} P\{\ell(A,\widehat\eta(\hat{\boldsymbol{f}}_X)) - \ell(A,\eta(\nu(\boldsymbol{f}_X)))\} &\le P\{\ell(A,\eta^\dagger(\hat{\boldsymbol{f}}_X)) - \ell(A,\eta(\nu(\boldsymbol{f}_X)))\} \notag\\ &\qquad\quad+ (P-\bbP_n) \{\ell(A,\widehat\eta(\hat{\boldsymbol{f}}_X)) - \ell(A,\eta^\dagger(\hat{\boldsymbol{f}}_X))\}. \end{align}\tag{21}\] Since \(\hat{\boldsymbol{f}}_X\) is a function of \((\boldsymbol{f}_X,\bu_X)\) and \(A\) is indepndent of \(\bu_X\), \[\bbP(A=1\mid \nu(\boldsymbol{f}_X),\hat{\boldsymbol{f}}_X) = \pi_1(\boldsymbol{f}_X) = \eta(\nu(\boldsymbol{f}_X)).\] Hence, for any \(g\in\mathcal{G}_{\eta,n}\), \[\begin{align} P\{\ell(A,g(\hat{\boldsymbol{f}}_X)) - \ell(A,\eta(\nu(\boldsymbol{f}_X)))\} = \bbE\left[\mathrm{KL}\{\mathrm{Bern}(\eta(\nu(\boldsymbol{f}_X))) \mid \mid \mathrm{Bern}(g(\hat{\boldsymbol{f}}_X)) \} \right]. \end{align}\] By overlap and clipping, this KL divergence is equivalent to squared error: \[c_\kappa \{g(\hat{\boldsymbol{f}}_X)-\eta(\nu(\boldsymbol{f}_X))\}^2 \le \mathrm{KL}\{\cdot,\cdot\} \le C_\kappa \{g(\hat{\boldsymbol{f}}_X)-\eta(\nu(\boldsymbol{f}_X))\}^2 .\] Therefore, applying the lower bound to \(\widehat\eta\) and the upper bound to \(\eta^\dagger\), 21 yields \[\begin{align} \label{eq:pi195kl95shorter} \|\widehat\eta(\hat{\boldsymbol{f}}_X)-\eta(\nu(\boldsymbol{f}_X))\|_{L_2(P)}^2 &\lesssim \|\eta^\dagger(\hat{\boldsymbol{f}}_X)-\eta(\nu(\boldsymbol{f}_X))\|_{L_2(P)}^2 \notag\\ &\quad+ (P-\bbP_n) \{\ell(A,\widehat\eta(\hat{\boldsymbol{f}}_X)) - \ell(A,\eta^\dagger(\hat{\boldsymbol{f}}_X))\}. \end{align}\tag{22}\] The logistic loss is Lipschitz on \([\kappa/2,1-\kappa/2]\) since \[\left| \partial_u\ell(a,u) \right| \le \frac{2}{\kappa}.\] Thus, by the contraction inequality and the pseudo-dimension bound used in Lemma 4 of [32], if \[V_n:=\operatorname{Pdim}(\mathcal{G}_{\eta,n}),\] then with probability at least \(1-e^{-t}\), \[\begin{align} &(P-\bbP_n) \{\ell(A,\widehat\eta(\hat{\boldsymbol{f}}_X)) - \ell(A,\eta^\dagger(\hat{\boldsymbol{f}}_X))\} \\ &\qquad\le \zeta \|\widehat\eta(\hat{\boldsymbol{f}}_X)-\eta(\nu(\boldsymbol{f}_X))\|_{L_2(P)}^2 + C_\zeta \|\eta^\dagger(\hat{\boldsymbol{f}}_X)-\eta(\nu(\boldsymbol{f}_X))\|_{L_2(P)}^2+ C_\zeta \left( \frac{V_n\log n}{n} + \frac{t}{n} \right). \end{align}\] Choosing \(\zeta\) small enough and combining this with 22 , we obtain \[\begin{align} \|\widehat\eta(\hat{\boldsymbol{f}}_X)-\eta(\nu(\boldsymbol{f}_X))\|_{L_2(P)}^2 \lesssim \|\eta^\dagger(\hat{\boldsymbol{f}}_X)-\eta(\nu(\boldsymbol{f}_X))\|_{L_2(P)}^2 + \frac{V_n\log n}{n} + \frac{t}{n}. \end{align}\] Using 20 , \[\label{eq:pi195pre95balance95shorter} \|\widehat\eta(\hat{\boldsymbol{f}}_X)-\eta(\nu(\boldsymbol{f}_X))\|_{L_2(P)}^2 = O_p\left( N_\eta^{-4\gamma_\eta^\star} + \delta_{n,f}^2 + \frac{V_n\log n}{n} \right).\tag{23}\]

Finally, as in the proof for \(\widehat\mu_1\), \[V_n \lesssim \left( L_\eta^2N_\eta^2+d_\eta L_\eta N_\eta \right) \log(L_\eta N_\eta d_\eta).\] Taking \(L_\eta\asymp 1\) and balancing \[N_\eta^{-4\gamma_\eta^\star} \qquad\text{and}\qquad \frac{N_\eta^2\log N_\eta\log n}{n}\] gives \[N_\eta\asymp n^{1/(4\gamma_\eta^\star+2)}\] up to logarithmic factors. Therefore, \[\|\widehat\eta(\hat{\boldsymbol{f}}_X)-\eta(\nu(\boldsymbol{f}_X))\|_{L_2(P)}^2 = O_p\left( n^{-2\gamma_\eta^\star/(2\gamma_\eta^\star+1)} (\log n)^a + \delta_{n,f}^2 \right).\] Equivalently, \[\|\widehat\pi_1(\hat{\boldsymbol{f}}_X)-\pi_1(\boldsymbol{f}_X)\|_{L_2(P)}^2 = O_p\left( n^{-2\gamma_\eta^\star/(2\gamma_\eta^\star+1)} (\log n)^a + \delta_{n,f}^2 \right).\]

13.3 Rate for \(\widehat \mu_{10}\) under the transformed factor representation↩︎

We now prove the rate for \(\mu_{10}(\boldsymbol{f}_X) :=\bbE[\mu_1(\boldsymbol{f}_X,\boldsymbol{f}_M)\mid A=0,\boldsymbol{f}_X]\). Define the relevant nuisance parameter \(\eta := \mu_{10}\circ \nu_X^{-1}\) so that \(\eta(\nu_X(\boldsymbol{f}_X))=\mu_{10}(\boldsymbol{f}_X)\). Let us also define \[\cI_0:=\{i\in\cI_{\rm nuis}:A_i=0\}, \qquad n_0:=|\cI_0|.\] The estimator \(\widehat\eta\) is obtained by regressing the generated outcome \(\widehat\mu_1(\widehat\boldsymbol{f}_{X,i},\widehat\boldsymbol{f}_{M,i})\) on \(\widehat\boldsymbol{f}_{X,i}\): \[\widehat\eta \in \arg\min_{g\in\mathcal{G}_{\eta,n}} \frac{1}{n_0} \sum_{i\in\cI_0} \left\{ \widehat\mu_1(\widehat\boldsymbol{f}_{X,i},\widehat\boldsymbol{f}_{M,i}) - g(\widehat\boldsymbol{f}_{X,i}) \right\}^2 .\] Assume \[\|\widehat\boldsymbol{f}_X-\nu_X(\boldsymbol{f}_X)\|_{L_2(P_0)} = O_p(\delta_{n,f}),\] where \(P_0\) denotes the distribution conditional on \(A=0\). Assume also that \(\eta\) admits a bounded Lipschitz HCM extension to a compact set containing both \(\nu_X(\boldsymbol{f}_X)\) and \(\widehat\boldsymbol{f}_X\) with probability tending to one, and \(\eta\in\cH(d_\eta,\ell_\eta,\cP)\) with DAS \(\gamma^*_\eta\). As before, let \[\mathcal{G}_{\eta,n} := \cG(L_\eta,d_\eta,1,N_\eta,M_\eta,B_\eta)\] be the truncated ReLU class. By Theorem 4 of [32], there exists \[\eta^\dagger \in \cG(C_1,d_\eta,1,C_2N_\eta,\infty,C_3N_\eta^{C_4})\] such that \[\|\eta^\dagger-\eta\|_\infty \le C N_\eta^{-2\gamma_\eta^\star}.\] After truncation, we may regard \(\eta^\dagger\in\mathcal{G}_{\eta,n}\) and retain \[\|\eta^\dagger-\eta\|_\infty^2 \le C N_\eta^{-4\gamma_\eta^\star}.\] As in the proof of \(\widehat\mu_1\), the relevant oracle error is \[\|\eta^\dagger(\widehat\boldsymbol{f}_X)-\eta(\nu_X(\boldsymbol{f}_X))\|_{L_2(P_0)}^2.\] Using the same decomposition as before and Lipschitzness of \(\eta\), \[\label{eq:mu1095oracle95approx} \|\eta^\dagger(\widehat\boldsymbol{f}_X)-\eta(\nu_X(\boldsymbol{f}_X))\|_{L_2(P_0)}^2 = O_p\left( N_\eta^{-4\gamma_\eta^\star} + \delta_{n,f}^2 \right).\tag{24}\] Now, for \(i\in\cI_0\), define \[\begin{align} \xi_i & := \mu_1(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i})-\mu_{10}(\boldsymbol{f}_{X,i}), \\ \Delta_i &:=\widehat\mu_1(\widehat\boldsymbol{f}_{X,i},\widehat\boldsymbol{f}_{M,i})-\mu_1(\boldsymbol{f}_{X,i},\boldsymbol{f}_{M,i}). \end{align}\] Then the generated outcome satisfies \[\widehat\mu_1(\widehat\boldsymbol{f}_{X,i},\widehat\boldsymbol{f}_{M,i}) = \eta(\nu_X(\boldsymbol{f}_{X,i})) + \xi_i + \Delta_i.\] Here \(\xi_i\) is the intrinsic second-stage regression noise, while \(\Delta_i\) is the first-stage error from estimating \(\mu_1\). By independent of \((\boldsymbol{f}_X, \boldsymbol{f}_M)\) with \((\bu_X, \bu_M)\), we have: \[\bbE[\xi_i\mid \boldsymbol{f}_{X,i},\bu_{X,i},A_i=0]=0,\] so that, conditional on the representation-learning fold, \[\bbE[\xi_i\mid \widehat\boldsymbol{f}_{X,i},A_i=0]=0.\] For any \(g\in\mathcal{G}_{\eta,n}\), define \[r_g(i) := g(\widehat\boldsymbol{f}_{X,i}) - \eta(\nu_X(\boldsymbol{f}_{X,i})).\] By empirical risk minimization, \[\frac{1}{n_0}\sum_{i\in\cI_0} \{\xi_i+\Delta_i-r_{\widehat\eta}(i)\}^2 \le \frac{1}{n_0}\sum_{i\in\cI_0} \{\xi_i+\Delta_i-r_{\eta^\dagger}(i)\}^2 .\] Expanding both sides and canceling common terms gives \[\begin{align} \label{eq:mu1095basic} \|r_{\widehat\eta}\|_{n,0}^2 \le \|r_{\eta^\dagger}\|_{n,0}^2 + \frac{2}{n_0}\sum_{i\in\cI_0} \xi_i \left\{ \widehat\eta(\widehat\boldsymbol{f}_{X,i}) - \eta^\dagger(\widehat\boldsymbol{f}_{X,i}) \right\} + \frac{2}{n_0}\sum_{i\in\cI_0} \Delta_i \left\{ \widehat\eta(\widehat\boldsymbol{f}_{X,i}) - \eta^\dagger(\widehat\boldsymbol{f}_{X,i}) \right\}. \end{align}\tag{25}\] The first stochastic term is handled exactly as in the proof of \(\widehat\mu_1\), using Lemma 4 of [32]. The second term is new and comes from the generated outcome.

By Cauchy’s inequality and Young’s inequality, \[\begin{align} \left|\frac{2}{n_0}\sum_{i\in\cI_0}\Delta_i\left\{\widehat\eta(\widehat\boldsymbol{f}_{X,i})-\eta^\dagger(\widehat\boldsymbol{f}_{X,i})\right\}\right| & \le 2\|\Delta\|_{n,0}\|\widehat\eta(\widehat\boldsymbol{f}_X)\eta^\dagger(\widehat\boldsymbol{f}_X)\|_{n,0}\\ &\qquad \le \zeta \|r_{\widehat\eta}\|_{n,0}^2+C_\zeta \|r_{\eta^\dagger}\|_{n,0}^2+C_\zeta \|\Delta\|_{n,0}^2. \end{align}\] Combining this with 25 , and applying the same empirical process bound as in the proof of \(\widehat\mu_1\), we obtain \[\|r_{\widehat\eta}\|_{n,0}^2 \lesssim \|r_{\eta^\dagger}\|_{n,0}^2 + \frac{V_n\log n}{n} + \|\Delta\|_{n,0}^2 + \frac{t}{n},\] where \[V_n:=\operatorname{Pdim}(\mathcal{G}_{\eta,n}).\] By Bernstein’s inequality for the fixed oracle and the already established rate for \(\widehat\mu_1\) under \(P_0\), \[\|\Delta\|_{n,0}^2 = O_p(r_{\mu_1,0,n}^2),\] where \[r_{\mu_1,0,n}^2 := \|\widehat\mu_1(\widehat\boldsymbol{f}_X,\widehat\boldsymbol{f}_M) - \mu_1(\boldsymbol{f}_X,\boldsymbol{f}_M)\|_{L_2(P_0)}^2 .\] Therefore, \[\label{eq:mu1095pre95balance} \|\widehat\eta(\widehat\boldsymbol{f}_X)-\eta(\nu_X(\boldsymbol{f}_X))\|_{L_2(P_0)}^2 = O_p\left( N_\eta^{-4\gamma_\eta^\star} + \delta_{n,f}^2 + \frac{V_n\log n}{n} + r_{\mu_1,0,n}^2 \right).\tag{26}\]

Finally, as before, \[V_n \lesssim \left( L_\eta^2N_\eta^2+d_\eta L_\eta N_\eta \right) \log(L_\eta N_\eta d_\eta).\] Taking \(L_\eta\asymp 1\) and balancing \[N_\eta^{-4\gamma_\eta^\star} \qquad\text{and}\qquad \frac{N_\eta^2\log N_\eta\log n}{n}\] gives \[N_\eta\asymp n^{1/(4\gamma_\eta^\star+2)}\] up to logarithmic factors. Hence, \[\|\widehat\eta(\widehat\boldsymbol{f}_X)-\eta(\nu_X(\boldsymbol{f}_X))\|_{L_2(P_0)}^2 = O_p\left( n^{-2\gamma_\eta^\star/(2\gamma_\eta^\star+1)} (\log n)^a + \delta_{n,f}^2 + r_{\mu_1,0,n}^2 \right).\] Equivalently, \[\|\widehat\mu_{10}(\widehat\boldsymbol{f}_X)-\mu_{10}(\boldsymbol{f}_X)\|_{L_2(P_0)}^2 = O_p\left( n^{-2\gamma_\eta^\star/(2\gamma_\eta^\star+1)} (\log n)^a + \delta_{n,f}^2 + r_{\mu_1,0,n}^2 \right).\]

14 A sufficient condition for factor recovery↩︎

15 Proof of Proposition 14↩︎

By condition 27 and 28 , we have: \[d_{\mathcal{V}}^2(\hat{z}) \le C\{\mathcal{L}_\lambda(\hat{z})-\mathcal{L}_\lambda^\star\} + C(\delta_{X, \rm app}^2+\delta_{M, \rm app}^2) = O_p(r_n^2+\delta_{X, \rm app}^2+\delta_{M, \rm app}^2).\] Therefore, \[d_{\mathcal{V}}(\hat{z}) = O_p(r_n+\delta_{X, \rm app}+\delta_{M, \rm app}).\] By the definition of \(d_{\mathcal{V}}\), for every \(\varepsilon>0\) there exists \(\nu_\varepsilon=(\nu_{X,\varepsilon},\nu_{M,\varepsilon})\in\mathcal{V}\) such that \[\begin{align} &\|\hat{g}_X(X)-\nu_{X,\varepsilon}(\boldsymbol{f}_X)\|_{L_2(P)}^2 + \|\hat{g}_M(M)-\nu_{M,\varepsilon}(\boldsymbol{f}_M)\|_{L_2(P)}^2 \\ &\qquad\le d_{\mathcal{V}}^2(\hat{z})+\varepsilon . \end{align}\] Taking, for instance, \(\varepsilon=n^{-1}\) and writing \(\nu_n:=\nu_{1/n}\) gives \[\|\hat{g}_X(X)-\nu_{X,n}(\boldsymbol{f}_X)\|_{L_2(P)}^2 + \|\hat{g}_M(M)-\nu_{M,n}(\boldsymbol{f}_M)\|_{L_2(P)}^2 = O_p(r_n^2+\delta_{X, \rm app}^2+\delta_{M, \rm app}^2).\] Taking square roots and using \(\sqrt{x+y+z}\le \sqrt x+\sqrt y+\sqrt z\) yields \[\left( \|\hat{g}_X(X)-\nu_{X,n}(\boldsymbol{f}_X)\|_{L_2(P)}^2 + \|\hat{g}_M(M)-\nu_{M,n}(\boldsymbol{f}_M)\|_{L_2(P)}^2 \right)^{1/2} = O_p(r_n+\delta_{X, \rm app}+\delta_{M, \rm app}).\] Consequently, \[\|\hat{g}_X(X)-\nu_{X,n}(\boldsymbol{f}_X)\|_{L_2(P)} + \|\hat{g}_M(M)-\nu_{M,n}(\boldsymbol{f}_M)\|_{L_2(P)} = O_p(r_n+\delta_{X, \rm app}+\delta_{M, \rm app}),\] where the last step follows because the sum of two nonnegative terms is bounded by \(\sqrt 2\) times their Euclidean norm.

Finally, by the definition of \[\hat{\boldsymbol{f}}:=(\hat{g}_X(X),\hat{g}_M(M)), \qquad \nu_n(\boldsymbol{f}):=(\nu_{X,n}(\boldsymbol{f}_X),\nu_{M,n}(\boldsymbol{f}_M)),\] we obtain \[\|\hat{\boldsymbol{f}}-\nu_n(\boldsymbol{f})\|_{L_2(P)} = O_p(r_n+\delta_{X, \rm app}+\delta_{M, \rm app}).\] Thus, if \(r_n+\delta_{X, \rm app}+\delta_{M, \rm app}=o(n^{-1/4})\), then \[\|\hat{\boldsymbol{f}}-\nu_n(\boldsymbol{f})\|_{L_2(P)} = o_p(n^{-1/4}),\] which is the desired factor-recovery condition.

16 Additional Real Data Details↩︎

We apply our proposed method to investigate how depressive symptoms influence cognitive decline in Alzheimer’s disease, and whether this relationship is mediated by DNA methylation. Understanding this pathway is important because prior studies have established associations between depression and cognitive impairment, but the biological mechanisms—particularly epigenetic mediation—remain less clear. Here, we present the detailed data preprocessing procedure.

16.0.0.1 Dataset.

We conduct our analysis using data from the Alzheimer’s Disease Neuroimaging Initiative (ADNI), a longitudinal, multi-center observational study designed to validate biomarkers for Alzheimer’s disease (AD) clinical trials [33]. The dataset is accessible via https://adni.loni.usc.edu/. Our treatment variable \(A\) is the Geriatric Depression Scale (GDS), a self-reported measure ranging from 0 to 15, with higher values indicating more severe depressive symptoms. We binarize GDS using the clinically standard threshold (e.g., GDS \(> 5\) indicates depression). The outcome \(Y\) is the Alzheimer’s Disease Assessment Scale–Cognitive Subscale (ADAS-Cog), ranging from 0 to 85, where higher scores reflect greater cognitive impairment [35], [36].

16.0.0.2 Preprocessing.

We include clinically relevant covariates such as demographic and health-related variables, denoted by \(X\). Missing values in \(X\) are imputed using mice [52]. To enhance model flexibility, we construct higher-order polynomial features and interaction terms, resulting in a covariate dimension of \(\dim(X) = 171\). The mediator \(M\) consists of DNA methylation measurements obtained using the Illumina HumanMethylationEPIC BeadChip array. We preprocess the methylation data using the minfi package [53], and we select CpG sites that have been reported to be significantly associated with Alzheimer’s disease in large-scale epigenome-wide association studies (EWAS), including a meta-analysis of Braak stage and re-analysis studies of publicly available GEO datasets [54]–[56].

based on missingness and variability, retaining only those with no missing values. This yields \(\dim(M) = 3{,}206\) mediators. After preprocessing, we retain \(n = 649\) subjects with complete data across all variables.

16.0.0.3 Choice of \(\tilde{p}\) and \(\tilde{q}\).

We apply our proposed estimation procedure to estimate the natural indirect effect (NIE), natural direct effect (NDE), and total effect (TE). The projection dimensions \((\tilde{p}, \tilde{q})\) are selected based on principal component analysis (PCA) where \(\tilde{p} = 5\) and \(\tilde{q} = 40\) number of coodinates is able to explain 80% variance.

17 More on Simulation and Ablation Study↩︎

All experiments were conducted in a cluster computing environment on a CPU compute node. The computational environment used Python 3.10.12 with 32 CPU cores. A typical full experimental run with 500 replications took approximately 84 hours to complete. Tables 3–9 and Tables 4–10 report the performance of different methods and ablation study compared with MediEncoder across varying sample sizes and covariate/mediator dimensions. The ablation study further evaluates the performance of MediEncoder after removing the alignment component in the loss function 8 .

5pt

5pt

5pt

5pt

References↩︎

---

1. \(^\dagger\) Corresponding authors.↩︎