Representation Learning for Sample-Efficient CATE Estimation by Leveraging Multiple Outcomes

Maitreyi Swaroop; Shikha Bhat; Samantha Rodriguez; Tamar Krishnamurti; Bryan Wilder

cs.LGSep 5, 2026

Representation Learning for Sample-Efficient CATE Estimation by Leveraging Multiple Outcomes

Authors: Maitreyi Swaroop, Shikha Bhat, Samantha Rodriguez, Tamar Krishnamurti, Bryan Wilder

Organizations: Machine Learning Department, Carnegie Mellon University · Department of General Internal Medicine, University of Pittsburgh

Abstract

Estimating conditional average treatment effects (CATE) enables efficient targeting of interventions, but many applications have limited experimental samples, making it difficult to estimate heterogeneous effects from high-dimensional covariates. In such settings, policymakers and medical practitioners often succumb to the curse of dimensionality or apply off-the-shelf dimension reduction methods that may not preserve treatment heterogeneity. Yet these domains often come with large historical datasets measuring a wide range of outcomes -- a source of supervision that is rarely exploited in practice. Following causal representation learning, we hypothesize that such domains with high-dimensional covariates have lower-dimensional underlying dynamics. We can thus leverage the diverse outcomes measured in historical data to learn a lower-dimensional representation of the covariates. Theoretically, we prove that when the auxiliary outcomes satisfy a set of surrogacy conditions and the representation retains relevant covariate information, the original CATE is identified when the high-dimensional covariates are replaced by the learned representation. Combined with existing dimension-dependent rates for CATE estimation, the result implies greater sample-efficiency on the same experimental sample. Additionally, we characterize the bias-variance tradeoff when the assumptions do not hold perfectly, and show that the representation-based estimator can still achieve lower error when the reduction in estimator variance outweighs the bias due to compression. Empirically, we evaluate the method on synthetic data and semi-synthetic medical data.

Figures & tables

Figure 1: Method overview. On historical data we train $\phi$ via contrastive lower bounds on $I(\phi;(Y^{*},S))$ and $I(\phi;S)$ (Equation 6 ). On the RCT we feed $(\phi(X),T,\tilde{Y})$ to a CATE meta-learner, with $\tilde{Y}=h(S,\phi(X))$ when $Y^{*}$ is unobserved (Case 1) or $\tilde{Y}=Y^{*}$ when it is observed (Case 2).

Figure 2: Normalized PEHE and normalized top-20% policy value on the synthetic and perinatal semi-synthetic datasets at

\dim(\phi)=10

. Curves show means over 10 runs and the surrounding shaded regions show

\pm 1

standard error. The two raw-

X

curves differ only in whether the historical outcome model uses

S

Appendix figures & tables22 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 3 — Figure 3: Sufficiency-(ii) violation, DML — normalized PEHE. True-outcome-aware representations retain a clear advantage under partial violations, but their normalized PEHE increases with $\alpha$ as more treatment-effect heterogeneity depends on features that cannot be learned from the historical data.

Panel of Figure 4 — Figure 4: Sufficiency-(ii) violation, DML — normalized top-20% policy value. The policy advantage declines gradually under partial violation and disappears at $\alpha=1$ , when the representation no longer contains the information that determines treatment-effect rankings.

Figure 5: Causal structure: default and two identification violations. Top row, left to right: notation; the historical DAG, where no treatment is present; the experimental DAG, where

T

acts on

S

and

S

mediates the effect on

Y^{*}

Z

is the ground-truth latent representation of

X

on which

S

and

Y^{*}

depend;

\phi(X)

is trained to recover

Z

. Bottom row: the two violations. Left, surrogacy violation: a direct

T\!\to\!Y^{*}

edge (dashed red) breaks the mediation through

S

; analyzed in Appendix D.1 . Right, sufficiency-(ii) violation: the violation is specific to the treated group. At

T=0

(left sub-panel) the structure matches what

\phi

was trained on, so sufficiency-(ii) holds; at

T=1

(right sub-panel) a violating

X\!\to\!S

edge (dashed red) appears, encoding heterogeneity in the treatment effect on

S

that depends on features of

X

unrecoverable from untreated data. Analyzed in Appendix C.3 .

Panel of Figure 6 — Figure 6: Surrogacy violation, DML – normalized PEHE. The outcome-aware representations lose accuracy as more of the treatment effect bypasses $S$ , but remain more accurate than unsupervised compression and raw $X$ throughout the sweep.

Panel of Figure 7 — Figure 7: Surrogacy violation, DML – normalized top-20% policy value. The policy advantage declines as the direct effect grows, but remains positive at $\delta=2$ because the surrogate-mediated component remains informative.

Panel of Figure 8 — Figure 8: Synthetic data results using the X-learner. Rows: error metric (PEHE; normalized top-20% policy value). Columns: CATE target ( $\hat{Y}$ vs. $Y$ ). Curves show means over 10 historical and experimental resamples from the same DGP; shaded regions show $\pm 1$ standard error.

Panel of Figure 9 — Figure 9: Synthetic data results using the DML learner using predicted and true outcomes as CATE targets. Curves show means over the same 10 historical and experimental resamples as the main analysis; shaded regions show $\pm 1$ standard error.

	$\widehat{Y}$ target					$Y$ target
Method	$n=50$	$n=100$	$n=250$	$n=500$	$n=750$	$n=50$	$n=100$	$n=250$	$n=500$	$n=750$
MI-InfoNCE	1.141 $\pm$ 0.048	0.818 $\pm$ 0.052	0.587 $\pm$ 0.028	0.500 $\pm$ 0.012	0.474 $\pm$ 0.008	1.167 $\pm$ 0.058	0.789 $\pm$ 0.048	0.508 $\pm$ 0.028	0.395 $\pm$ 0.012	0.357 $\pm$ 0.008
MI-MINE	1.138 $\pm$ 0.044	0.999 $\pm$ 0.039	0.841 $\pm$ 0.010	0.829 $\pm$ 0.009	0.816 $\pm$ 0.007	1.167 $\pm$ 0.058	0.978 $\pm$ 0.038	0.836 $\pm$ 0.013	0.818 $\pm$ 0.013	0.806 $\pm$ 0.010
Encoder	1.141 $\pm$ 0.057	1.000 $\pm$ 0.040	0.853 $\pm$ 0.020	0.833 $\pm$ 0.018	0.824 $\pm$ 0.018	1.167 $\pm$ 0.058	0.789 $\pm$ 0.054	0.437 $\pm$ 0.012	0.364 $\pm$ 0.006	0.329 $\pm$ 0.007
PLS	1.145 $\pm$ 0.046	0.832 $\pm$ 0.041	0.534 $\pm$ 0.022	0.485 $\pm$ 0.021	0.423 $\pm$ 0.012	1.167 $\pm$ 0.058	0.809 $\pm$ 0.042	0.498 $\pm$ 0.017	0.441 $\pm$ 0.018	0.393 $\pm$ 0.008
PCA	1.159 $\pm$ 0.052	1.313 $\pm$ 0.060	1.542 $\pm$ 0.047	1.529 $\pm$ 0.039	1.454 $\pm$ 0.019	1.167 $\pm$ 0.058	1.314 $\pm$ 0.059	1.564 $\pm$ 0.049	1.551 $\pm$ 0.031	1.469 $\pm$ 0.022

Appendix

Table 1: Normalized PEHE on the synthetic data, X-learner (

\dim(\phi)=10

). Entries are means

\pm

standard errors over 10 paired data resamples; lower is better. For the

Y

target, the two raw-

X

variants coincide, so we report a single value under

h(X)

	$\widehat{Y}$ target					$Y$ target
Method	$n=50$	$n=100$	$n=250$	$n=500$	$n=750$	$n=50$	$n=100$	$n=250$	$n=500$	$n=750$
MI-InfoNCE	-0.006 $\pm$ 0.085	0.611 $\pm$ 0.062	0.851 $\pm$ 0.015	0.890 $\pm$ 0.010	0.911 $\pm$ 0.005	0.012 $\pm$ 0.084	0.659 $\pm$ 0.042	0.881 $\pm$ 0.020	0.931 $\pm$ 0.010	0.944 $\pm$ 0.005
MI-MINE	-0.117 $\pm$ 0.109	0.451 $\pm$ 0.053	0.594 $\pm$ 0.022	0.590 $\pm$ 0.023	0.613 $\pm$ 0.013	0.084 $\pm$ 0.111	0.496 $\pm$ 0.048	0.611 $\pm$ 0.024	0.610 $\pm$ 0.020	0.634 $\pm$ 0.010
Encoder	-0.075 $\pm$ 0.140	0.321 $\pm$ 0.075	0.632 $\pm$ 0.079	0.709 $\pm$ 0.063	0.773 $\pm$ 0.060	0.060 $\pm$ 0.102	0.691 $\pm$ 0.060	0.908 $\pm$ 0.008	0.928 $\pm$ 0.005	0.950 $\pm$ 0.003
PLS	-0.144 $\pm$ 0.105	0.618 $\pm$ 0.070	0.858 $\pm$ 0.009	0.894 $\pm$ 0.009	0.915 $\pm$ 0.008	0.072 $\pm$ 0.084	0.637 $\pm$ 0.066	0.872 $\pm$ 0.011	0.902 $\pm$ 0.012	0.923 $\pm$ 0.006
PCA	-0.032 $\pm$ 0.081	0.055 $\pm$ 0.059	0.076 $\pm$ 0.029	0.086 $\pm$ 0.023	0.101 $\pm$ 0.021	-0.031 $\pm$ 0.050	0.049 $\pm$ 0.057	0.080 $\pm$ 0.031	0.088 $\pm$ 0.029	0.120 $\pm$ 0.011

Appendix

Table 2: Normalized top-20% policy value on the synthetic data, X-learner (

\dim(\phi)=10

). Entries are means

\pm

standard errors over 10 paired data resamples; higher is better. For the

Y

target, the two raw-

X

variants coincide, so we report a single value under

h(X)

	$\widehat{Y}$ target					$Y$ target
Method	$n=50$	$n=100$	$n=250$	$n=500$	$n=750$	$n=50$	$n=100$	$n=250$	$n=500$	$n=750$
MI-InfoNCE	1.371 $\pm$ 0.059	0.650 $\pm$ 0.059	0.486 $\pm$ 0.037	0.423 $\pm$ 0.019	0.408 $\pm$ 0.011	1.326 $\pm$ 0.053	0.586 $\pm$ 0.056	0.364 $\pm$ 0.032	0.274 $\pm$ 0.012	0.252 $\pm$ 0.010
MI-MINE	1.616 $\pm$ 0.069	0.928 $\pm$ 0.046	0.822 $\pm$ 0.013	0.825 $\pm$ 0.008	0.811 $\pm$ 0.008	1.598 $\pm$ 0.059	0.916 $\pm$ 0.036	0.806 $\pm$ 0.008	0.791 $\pm$ 0.008	0.781 $\pm$ 0.007
Encoder	1.776 $\pm$ 0.151	0.977 $\pm$ 0.036	0.861 $\pm$ 0.023	0.837 $\pm$ 0.017	0.830 $\pm$ 0.019	1.563 $\pm$ 0.091	0.550 $\pm$ 0.060	0.357 $\pm$ 0.013	0.312 $\pm$ 0.008	0.291 $\pm$ 0.005
PLS	1.919 $\pm$ 0.123	0.902 $\pm$ 0.047	0.486 $\pm$ 0.036	0.387 $\pm$ 0.019	0.306 $\pm$ 0.020	1.901 $\pm$ 0.120	0.911 $\pm$ 0.048	0.454 $\pm$ 0.028	0.313 $\pm$ 0.014	0.238 $\pm$ 0.011
PCA	2.188 $\pm$ 0.089	1.796 $\pm$ 0.037	1.469 $\pm$ 0.067	1.271 $\pm$ 0.028	1.177 $\pm$ 0.016	2.243 $\pm$ 0.094	1.806 $\pm$ 0.042	1.487 $\pm$ 0.063	1.296 $\pm$ 0.034	1.180 $\pm$ 0.015

Appendix

Table 3: Normalized PEHE on the synthetic data, DML (

\dim(\phi)=10

). Entries are means

\pm

standard errors over 10 paired data resamples; lower is better. For the

Y

target, the two raw-

X

variants coincide, so we report a single value under

h(X)

	$\widehat{Y}$ target					$Y$ target
Method	$n=50$	$n=100$	$n=250$	$n=500$	$n=750$	$n=50$	$n=100$	$n=250$	$n=500$	$n=750$
MI-InfoNCE	0.371 $\pm$ 0.140	0.788 $\pm$ 0.043	0.904 $\pm$ 0.015	0.936 $\pm$ 0.008	0.944 $\pm$ 0.006	0.581 $\pm$ 0.074	0.849 $\pm$ 0.040	0.945 $\pm$ 0.012	0.966 $\pm$ 0.005	0.976 $\pm$ 0.002
MI-MINE	0.136 $\pm$ 0.136	0.503 $\pm$ 0.060	0.620 $\pm$ 0.027	0.595 $\pm$ 0.019	0.623 $\pm$ 0.013	0.251 $\pm$ 0.125	0.540 $\pm$ 0.040	0.628 $\pm$ 0.018	0.622 $\pm$ 0.013	0.638 $\pm$ 0.013
Encoder	-0.025 $\pm$ 0.163	0.344 $\pm$ 0.070	0.585 $\pm$ 0.090	0.750 $\pm$ 0.039	0.793 $\pm$ 0.064	0.388 $\pm$ 0.135	0.890 $\pm$ 0.023	0.928 $\pm$ 0.008	0.944 $\pm$ 0.006	0.953 $\pm$ 0.004
PLS	0.384 $\pm$ 0.116	0.712 $\pm$ 0.044	0.879 $\pm$ 0.021	0.933 $\pm$ 0.009	0.956 $\pm$ 0.007	0.426 $\pm$ 0.107	0.750 $\pm$ 0.035	0.910 $\pm$ 0.012	0.948 $\pm$ 0.008	0.975 $\pm$ 0.002
PCA	0.041 $\pm$ 0.059	-0.023 $\pm$ 0.049	0.028 $\pm$ 0.021	0.040 $\pm$ 0.023	0.022 $\pm$ 0.010	0.084 $\pm$ 0.049	-0.031 $\pm$ 0.048	0.032 $\pm$ 0.024	0.034 $\pm$ 0.024	0.033 $\pm$ 0.018

Appendix

Table 4: Normalized top-20% policy value on the synthetic data, DML (

\dim(\phi)=10

). Entries are means

\pm

standard errors over 10 paired data resamples; higher is better. For the

Y

target, the two raw-

X

variants coincide, so we report a single value under

h(X)

Panel of Figure 10 — Figure 10: Representation dimension sweep, DML — normalized PEHE. Increasing $m$ through the outcome-relevant latent dimension substantially improves the best supervised methods. Beyond that point, additional dimensions provide no uniform benefit and can increase error for some methods.

Panel of Figure 11 — Figure 11: Representation dimension sweep, DML — normalized top-20% policy value. MI-InfoNCE and PLS recover most of the available ranking signal once $m$ is large enough to represent the outcome-relevant latent structure. Unsupervised methods remain near random across dimensions.

Cohort overlap
Cohort-classification AUC	0.752
NN coverage in $(X_{Y},S)$ , $T=0$ / $T=1$	0.940 / 0.951
Representation sufficiency
Representation	Outcome $\Delta$ log loss	$S(0)$ $\Delta R^{2}$	$S(1)$ $\Delta R^{2}$
Prediction encoder	0.0065	0.0146	0.0379
Conditional MINE	0.0089	0.0501	0.0772

Appendix

Table 5: Diagnostics for cohort overlap and representation sufficiency. The final three columns report the cross-validated improvement from appending raw

X

to the representation; smaller values indicate that less predictive information remains beyond the representation.

Figure 12: The nine methods used in Section 5.2 on the perinatal semi-synthetic dataset at $\dim(\phi)=10$ . Points show means over 10 runs; error bars show $\pm 1$ standard error. Supervised representations learned from historical outcomes outperform both raw- $X$ variants and generally outperform unsupervised compression.

Method	$n=100$	$n=250$	$n=500$	$n=750$	$n=1000$
MI-InfoNCE	1.451 $\pm$ 0.076	1.247 $\pm$ 0.047	1.124 $\pm$ 0.031	1.133 $\pm$ 0.020	1.090 $\pm$ 0.019
MI-MINE	1.340 $\pm$ 0.062	1.186 $\pm$ 0.052	1.099 $\pm$ 0.035	1.152 $\pm$ 0.039	1.082 $\pm$ 0.022
Encoder	1.470 $\pm$ 0.068	1.173 $\pm$ 0.034	1.106 $\pm$ 0.026	1.119 $\pm$ 0.028	1.059 $\pm$ 0.017
PLS	1.790 $\pm$ 0.102	1.322 $\pm$ 0.062	1.118 $\pm$ 0.030	1.135 $\pm$ 0.031	1.055 $\pm$ 0.015
PCA	1.659 $\pm$ 0.093	1.372 $\pm$ 0.061	1.193 $\pm$ 0.049	1.216 $\pm$ 0.036	1.123 $\pm$ 0.022
ICA	1.668 $\pm$ 0.102	1.356 $\pm$ 0.044	1.151 $\pm$ 0.046	1.192 $\pm$ 0.023	1.130 $\pm$ 0.017

Appendix

Table 6: Normalized PEHE on the perinatal semi-synthetic dataset at

\dim(\phi)=10

. Entries are means

\pm

standard errors over 10 runs; lower is better.

Method	$n=100$	$n=250$	$n=500$	$n=750$	$n=1000$
MI-InfoNCE	0.245 $\pm$ 0.081	0.344 $\pm$ 0.040	0.381 $\pm$ 0.043	0.353 $\pm$ 0.033	0.417 $\pm$ 0.024
MI-MINE	0.286 $\pm$ 0.072	0.422 $\pm$ 0.033	0.413 $\pm$ 0.028	0.439 $\pm$ 0.032	0.484 $\pm$ 0.017
Encoder	0.357 $\pm$ 0.044	0.413 $\pm$ 0.028	0.485 $\pm$ 0.029	0.464 $\pm$ 0.031	0.549 $\pm$ 0.015
PLS	0.295 $\pm$ 0.048	0.341 $\pm$ 0.047	0.456 $\pm$ 0.027	0.386 $\pm$ 0.049	0.517 $\pm$ 0.016
PCA	0.121 $\pm$ 0.043	0.209 $\pm$ 0.036	0.255 $\pm$ 0.023	0.175 $\pm$ 0.032	0.293 $\pm$ 0.025
ICA	0.118 $\pm$ 0.042	0.193 $\pm$ 0.034	0.260 $\pm$ 0.033	0.207 $\pm$ 0.015	0.268 $\pm$ 0.025

Appendix

Table 7: Normalized top-20% policy value on the perinatal semi-synthetic dataset at

\dim(\phi)=10

. Entries are means

\pm

standard errors over 10 runs; higher is better.

Figure 13: The nine methods used in the main comparison at $\dim(\phi)=5$ . The prediction encoder gives similar PEHE and policy value at dimensions 5 and 10.

	Normalized PEHE		Policy@20%
$n$	$\dim(\phi)=5$	$\dim(\phi)=10$	$\dim(\phi)=5$	$\dim(\phi)=10$
100	1.471 $\pm$ 0.105	1.470 $\pm$ 0.068	0.325 $\pm$ 0.083	0.357 $\pm$ 0.044
250	1.153 $\pm$ 0.025	1.173 $\pm$ 0.034	0.425 $\pm$ 0.039	0.413 $\pm$ 0.028
500	1.114 $\pm$ 0.029	1.106 $\pm$ 0.026	0.510 $\pm$ 0.028	0.485 $\pm$ 0.029
750	1.148 $\pm$ 0.029	1.119 $\pm$ 0.028	0.404 $\pm$ 0.056	0.464 $\pm$ 0.031
1000	1.084 $\pm$ 0.025	1.059 $\pm$ 0.017	0.555 $\pm$ 0.018	0.549 $\pm$ 0.015

Appendix

Table 8: Representation-dimension sensitivity for the prediction encoder on the perinatal semi-synthetic dataset. Entries are means

\pm

standard errors over 10 runs.

Figure 14: The eight distinct methods at $\dim(\phi)=10$ using true experimental $Y$ as the CATE target. The $h(X,S)$ raw- $X$ variant is omitted because it is identical to raw $X$ when the true outcome is used directly. Points show means over 10 runs; error bars show $\pm 1$ standard error.

Method	$n=100$	$n=250$	$n=500$	$n=750$	$n=1000$
MI-InfoNCE	1.952 $\pm$ 0.117	1.522 $\pm$ 0.077	1.282 $\pm$ 0.041	1.096 $\pm$ 0.026	1.104 $\pm$ 0.043
MI-MINE	1.717 $\pm$ 0.108	1.339 $\pm$ 0.051	1.224 $\pm$ 0.059	1.121 $\pm$ 0.037	1.088 $\pm$ 0.023
Encoder	2.094 $\pm$ 0.153	1.420 $\pm$ 0.078	1.183 $\pm$ 0.032	1.102 $\pm$ 0.023	1.063 $\pm$ 0.023
PLS	2.594 $\pm$ 0.088	1.686 $\pm$ 0.140	1.292 $\pm$ 0.058	1.116 $\pm$ 0.023	1.059 $\pm$ 0.039
PCA	2.315 $\pm$ 0.172	1.698 $\pm$ 0.093	1.400 $\pm$ 0.031	1.214 $\pm$ 0.027	1.157 $\pm$ 0.035
ICA	2.451 $\pm$ 0.168	1.624 $\pm$ 0.096	1.340 $\pm$ 0.039	1.237 $\pm$ 0.031	1.181 $\pm$ 0.037

Appendix

Table 9: Normalized PEHE on the perinatal semi-synthetic dataset at

\dim(\phi)=10

using true experimental

Y

. Entries are means

\pm

standard errors over 10 runs; lower is better.

Method	$n=100$	$n=250$	$n=500$	$n=750$	$n=1000$
MI-InfoNCE	0.303 $\pm$ 0.047	0.313 $\pm$ 0.041	0.293 $\pm$ 0.032	0.354 $\pm$ 0.046	0.336 $\pm$ 0.022
MI-MINE	0.307 $\pm$ 0.060	0.283 $\pm$ 0.057	0.352 $\pm$ 0.034	0.356 $\pm$ 0.052	0.368 $\pm$ 0.040
Encoder	0.321 $\pm$ 0.064	0.302 $\pm$ 0.056	0.411 $\pm$ 0.038	0.392 $\pm$ 0.056	0.438 $\pm$ 0.032
PLS	0.273 $\pm$ 0.049	0.296 $\pm$ 0.045	0.399 $\pm$ 0.032	0.398 $\pm$ 0.042	0.399 $\pm$ 0.043
PCA	0.162 $\pm$ 0.038	0.173 $\pm$ 0.022	0.216 $\pm$ 0.028	0.201 $\pm$ 0.023	0.253 $\pm$ 0.024
ICA	0.169 $\pm$ 0.040	0.169 $\pm$ 0.036	0.211 $\pm$ 0.056	0.193 $\pm$ 0.021	0.215 $\pm$ 0.026

Appendix

Table 10: Normalized top-20% policy value on the perinatal semi-synthetic dataset at

\dim(\phi)=10

using true experimental

Y

. Entries are means

\pm

standard errors over 10 runs; higher is better.

Explore similar work

MediEncoder: Nonlinear Representation Learning for High-Dimensional Causal Mediation Analysis
Jun 5, 2026Shi Bo, Debarghya Mukherjee, AmirEmad GhassamiCausal Effect Estimation Representation Learning
Department of Mathematics and Statistics, Boston University
Causal Risk Minimization for High-Dimensional Treatments
May 26, 2026Nikita Dhawan, Arnav Paruthi, Andrew Kim +3Causal Effect Estimation Causal Inference
University of Toronto, Vector Institute · University of Toronto · Vanguard
Representation Learning for Semiparametric Causal Mediation Analysis under No Essential Heterogeneity
Jul 12, 2026Roberto Faleh, Sofia Morelli, Holger BrandtSemiparametric Inference Parameter Estimation
Methods Center University of Tübingen, Germany

Jun 5, 2026J/K move · Enter open · S save

MediEncoder: Nonlinear Representation Learning for High-Dimensional Causal Mediation Analysis

Shi Bo, Debarghya Mukherjee, AmirEmad Ghassami

Department of Mathematics and Statistics, Boston University

Causal Effect Estimation Representation Learning Causal Representation Learning Causal Mediation Analysis Causal Inference

Open paper AskPDF arXiv

Causal mediation analysis decomposes a treatment effect into indirect pathways through mediators and direct pathways not operating through them. Modern biomedical studies often involve high-dimensional covariates and mediators that are noisy proxies for lower-dimensional latent biological processes. Existing methods typically rely on sparsity, linear factor models, or ignore the connection among variables in the learned representations, which can be restrictive when measurements are nonlinear and covariate and mediator factors are structurally dependent. We propose MediEncoder, a representation-learning framework for nonlinear high-dimensional mediation analysis. MediEncoder jointly learns low-dimensional covariate and mediator representations using a coupled encoder-decoder architecture with a cross-factor network that links treatment and covariate representations to mediator representations. The learned features are then used in a cross-fitted efficient influence function-based estimator of natural direct and indirect effects. The resulting estimator is multiply robust and asymptotically normal under suitable regularity conditions. Simulations show that MediEncoder improves estimation accuracy over competing dimension-reduction approaches, and an application to Alzheimer's Disease Neuroimaging Initiative data illustrates its utility in high-dimensional biomedical causal mediation analysis.

Abstract

Explore similar work

MediEncoder: Nonlinear Representation Learning for High-Dimensional Causal Mediation Analysis

Causal Risk Minimization for High-Dimensional Treatments

Representation Learning for Semiparametric Causal Mediation Analysis under No Essential Heterogeneity