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 ϕ via contrastive lower bounds on I(ϕ;(Y∗,S)) and I(ϕ;S) (Equation 6 ). On the RCT we feed (ϕ(X),T,Y~) to a CATE meta-learner, with Y~=h(S,ϕ(X)) when Y∗ is unobserved (Case 1) or 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(ϕ)=10 . Curves show means over 10 runs and the surrounding shaded regions show ±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
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 α as more treatment-effect heterogeneity depends on features that cannot be learned from the historical data.
Figure 4: Sufficiency-(ii) violation, DML — normalized top-20% policy value. The policy advantage declines gradually under partial violation and disappears at α=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; ϕ(X) is trained to recover Z . Bottom row: the two violations. Left, surrogacy violation: a direct T→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 ϕ was trained on, so sufficiency-(ii) holds; at T=1 (right sub-panel) a violating X→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 .
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.
Figure 7: Surrogacy violation, DML – normalized top-20% policy value. The policy advantage declines as the direct effect grows, but remains positive at δ=2 because the surrogate-mediated component remains informative.
Figure 8: Synthetic data results using the X-learner. Rows: error metric (PEHE; normalized top-20% policy value). Columns: CATE target ( Y^ vs. Y ). Curves show means over 10 historical and experimental resamples from the same DGP; shaded regions show ±1 standard error.
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 ±1 standard error.
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 ± 0.048
0.818 ± 0.052
0.587 ± 0.028
0.500 ± 0.012
0.474 ± 0.008
1.167 ± 0.058
0.789 ± 0.048
0.508 ± 0.028
0.395 ± 0.012
0.357 ± 0.008
MI-MINE
1.138 ± 0.044
0.999 ± 0.039
0.841 ± 0.010
0.829 ± 0.009
0.816 ± 0.007
1.167 ± 0.058
0.978 ± 0.038
0.836 ± 0.013
0.818 ± 0.013
0.806 ± 0.010
Encoder
1.141 ± 0.057
1.000 ± 0.040
0.853 ± 0.020
0.833 ± 0.018
0.824 ± 0.018
1.167 ± 0.058
0.789 ± 0.054
0.437 ± 0.012
0.364 ± 0.006
0.329 ± 0.007
PLS
1.145 ± 0.046
0.832 ± 0.041
0.534 ± 0.022
0.485 ± 0.021
0.423 ± 0.012
1.167 ± 0.058
0.809 ± 0.042
0.498 ± 0.017
0.441 ± 0.018
0.393 ± 0.008
PCA
1.159 ± 0.052
1.313 ± 0.060
1.542 ± 0.047
1.529 ± 0.039
1.454 ± 0.019
1.167 ± 0.058
1.314 ± 0.059
1.564 ± 0.049
1.551 ± 0.031
1.469 ± 0.022
Appendix
Table 1: Normalized PEHE on the synthetic data, X-learner ( dim(ϕ)=10 ). Entries are means ± 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) .
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 ± 0.085
0.611 ± 0.062
0.851 ± 0.015
0.890 ± 0.010
0.911 ± 0.005
0.012 ± 0.084
0.659 ± 0.042
0.881 ± 0.020
0.931 ± 0.010
0.944 ± 0.005
MI-MINE
-0.117 ± 0.109
0.451 ± 0.053
0.594 ± 0.022
0.590 ± 0.023
0.613 ± 0.013
0.084 ± 0.111
0.496 ± 0.048
0.611 ± 0.024
0.610 ± 0.020
0.634 ± 0.010
Encoder
-0.075 ± 0.140
0.321 ± 0.075
0.632 ± 0.079
0.709 ± 0.063
0.773 ± 0.060
0.060 ± 0.102
0.691 ± 0.060
0.908 ± 0.008
0.928 ± 0.005
0.950 ± 0.003
PLS
-0.144 ± 0.105
0.618 ± 0.070
0.858 ± 0.009
0.894 ± 0.009
0.915 ± 0.008
0.072 ± 0.084
0.637 ± 0.066
0.872 ± 0.011
0.902 ± 0.012
0.923 ± 0.006
PCA
-0.032 ± 0.081
0.055 ± 0.059
0.076 ± 0.029
0.086 ± 0.023
0.101 ± 0.021
-0.031 ± 0.050
0.049 ± 0.057
0.080 ± 0.031
0.088 ± 0.029
0.120 ± 0.011
Appendix
Table 2: Normalized top-20% policy value on the synthetic data, X-learner ( dim(ϕ)=10 ). Entries are means ± 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) .
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 ± 0.059
0.650 ± 0.059
0.486 ± 0.037
0.423 ± 0.019
0.408 ± 0.011
1.326 ± 0.053
0.586 ± 0.056
0.364 ± 0.032
0.274 ± 0.012
0.252 ± 0.010
MI-MINE
1.616 ± 0.069
0.928 ± 0.046
0.822 ± 0.013
0.825 ± 0.008
0.811 ± 0.008
1.598 ± 0.059
0.916 ± 0.036
0.806 ± 0.008
0.791 ± 0.008
0.781 ± 0.007
Encoder
1.776 ± 0.151
0.977 ± 0.036
0.861 ± 0.023
0.837 ± 0.017
0.830 ± 0.019
1.563 ± 0.091
0.550 ± 0.060
0.357 ± 0.013
0.312 ± 0.008
0.291 ± 0.005
PLS
1.919 ± 0.123
0.902 ± 0.047
0.486 ± 0.036
0.387 ± 0.019
0.306 ± 0.020
1.901 ± 0.120
0.911 ± 0.048
0.454 ± 0.028
0.313 ± 0.014
0.238 ± 0.011
PCA
2.188 ± 0.089
1.796 ± 0.037
1.469 ± 0.067
1.271 ± 0.028
1.177 ± 0.016
2.243 ± 0.094
1.806 ± 0.042
1.487 ± 0.063
1.296 ± 0.034
1.180 ± 0.015
Appendix
Table 3: Normalized PEHE on the synthetic data, DML ( dim(ϕ)=10 ). Entries are means ± 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) .
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 ± 0.140
0.788 ± 0.043
0.904 ± 0.015
0.936 ± 0.008
0.944 ± 0.006
0.581 ± 0.074
0.849 ± 0.040
0.945 ± 0.012
0.966 ± 0.005
0.976 ± 0.002
MI-MINE
0.136 ± 0.136
0.503 ± 0.060
0.620 ± 0.027
0.595 ± 0.019
0.623 ± 0.013
0.251 ± 0.125
0.540 ± 0.040
0.628 ± 0.018
0.622 ± 0.013
0.638 ± 0.013
Encoder
-0.025 ± 0.163
0.344 ± 0.070
0.585 ± 0.090
0.750 ± 0.039
0.793 ± 0.064
0.388 ± 0.135
0.890 ± 0.023
0.928 ± 0.008
0.944 ± 0.006
0.953 ± 0.004
PLS
0.384 ± 0.116
0.712 ± 0.044
0.879 ± 0.021
0.933 ± 0.009
0.956 ± 0.007
0.426 ± 0.107
0.750 ± 0.035
0.910 ± 0.012
0.948 ± 0.008
0.975 ± 0.002
PCA
0.041 ± 0.059
-0.023 ± 0.049
0.028 ± 0.021
0.040 ± 0.023
0.022 ± 0.010
0.084 ± 0.049
-0.031 ± 0.048
0.032 ± 0.024
0.034 ± 0.024
0.033 ± 0.018
Appendix
Table 4: Normalized top-20% policy value on the synthetic data, DML ( dim(ϕ)=10 ). Entries are means ± 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) .
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.
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 (XY,S) , T=0 / T=1
0.940 / 0.951
Representation sufficiency
Representation
Outcome Δ log loss
S(0)ΔR2
S(1)ΔR2
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(ϕ)=10 . Points show means over 10 runs; error bars show ±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 ± 0.076
1.247 ± 0.047
1.124 ± 0.031
1.133 ± 0.020
1.090 ± 0.019
MI-MINE
1.340 ± 0.062
1.186 ± 0.052
1.099 ± 0.035
1.152 ± 0.039
1.082 ± 0.022
Encoder
1.470 ± 0.068
1.173 ± 0.034
1.106 ± 0.026
1.119 ± 0.028
1.059 ± 0.017
PLS
1.790 ± 0.102
1.322 ± 0.062
1.118 ± 0.030
1.135 ± 0.031
1.055 ± 0.015
PCA
1.659 ± 0.093
1.372 ± 0.061
1.193 ± 0.049
1.216 ± 0.036
1.123 ± 0.022
ICA
1.668 ± 0.102
1.356 ± 0.044
1.151 ± 0.046
1.192 ± 0.023
1.130 ± 0.017
Appendix
Table 6: Normalized PEHE on the perinatal semi-synthetic dataset at dim(ϕ)=10 . Entries are means ± standard errors over 10 runs; lower is better.
Method
n=100
n=250
n=500
n=750
n=1000
MI-InfoNCE
0.245 ± 0.081
0.344 ± 0.040
0.381 ± 0.043
0.353 ± 0.033
0.417 ± 0.024
MI-MINE
0.286 ± 0.072
0.422 ± 0.033
0.413 ± 0.028
0.439 ± 0.032
0.484 ± 0.017
Encoder
0.357 ± 0.044
0.413 ± 0.028
0.485 ± 0.029
0.464 ± 0.031
0.549 ± 0.015
PLS
0.295 ± 0.048
0.341 ± 0.047
0.456 ± 0.027
0.386 ± 0.049
0.517 ± 0.016
PCA
0.121 ± 0.043
0.209 ± 0.036
0.255 ± 0.023
0.175 ± 0.032
0.293 ± 0.025
ICA
0.118 ± 0.042
0.193 ± 0.034
0.260 ± 0.033
0.207 ± 0.015
0.268 ± 0.025
Appendix
Table 7: Normalized top-20% policy value on the perinatal semi-synthetic dataset at dim(ϕ)=10 . Entries are means ± standard errors over 10 runs; higher is better.
Figure 13: The nine methods used in the main comparison at dim(ϕ)=5 . The prediction encoder gives similar PEHE and policy value at dimensions 5 and 10.
Normalized PEHE
Policy@20%
n
dim(ϕ)=5
dim(ϕ)=10
dim(ϕ)=5
dim(ϕ)=10
100
1.471 ± 0.105
1.470 ± 0.068
0.325 ± 0.083
0.357 ± 0.044
250
1.153 ± 0.025
1.173 ± 0.034
0.425 ± 0.039
0.413 ± 0.028
500
1.114 ± 0.029
1.106 ± 0.026
0.510 ± 0.028
0.485 ± 0.029
750
1.148 ± 0.029
1.119 ± 0.028
0.404 ± 0.056
0.464 ± 0.031
1000
1.084 ± 0.025
1.059 ± 0.017
0.555 ± 0.018
0.549 ± 0.015
Appendix
Table 8: Representation-dimension sensitivity for the prediction encoder on the perinatal semi-synthetic dataset. Entries are means ± standard errors over 10 runs.
Figure 14: The eight distinct methods at dim(ϕ)=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 ±1 standard error.
Method
n=100
n=250
n=500
n=750
n=1000
MI-InfoNCE
1.952 ± 0.117
1.522 ± 0.077
1.282 ± 0.041
1.096 ± 0.026
1.104 ± 0.043
MI-MINE
1.717 ± 0.108
1.339 ± 0.051
1.224 ± 0.059
1.121 ± 0.037
1.088 ± 0.023
Encoder
2.094 ± 0.153
1.420 ± 0.078
1.183 ± 0.032
1.102 ± 0.023
1.063 ± 0.023
PLS
2.594 ± 0.088
1.686 ± 0.140
1.292 ± 0.058
1.116 ± 0.023
1.059 ± 0.039
PCA
2.315 ± 0.172
1.698 ± 0.093
1.400 ± 0.031
1.214 ± 0.027
1.157 ± 0.035
ICA
2.451 ± 0.168
1.624 ± 0.096
1.340 ± 0.039
1.237 ± 0.031
1.181 ± 0.037
Appendix
Table 9: Normalized PEHE on the perinatal semi-synthetic dataset at dim(ϕ)=10 using true experimental Y . Entries are means ± standard errors over 10 runs; lower is better.
Method
n=100
n=250
n=500
n=750
n=1000
MI-InfoNCE
0.303 ± 0.047
0.313 ± 0.041
0.293 ± 0.032
0.354 ± 0.046
0.336 ± 0.022
MI-MINE
0.307 ± 0.060
0.283 ± 0.057
0.352 ± 0.034
0.356 ± 0.052
0.368 ± 0.040
Encoder
0.321 ± 0.064
0.302 ± 0.056
0.411 ± 0.038
0.392 ± 0.056
0.438 ± 0.032
PLS
0.273 ± 0.049
0.296 ± 0.045
0.399 ± 0.032
0.398 ± 0.042
0.399 ± 0.043
PCA
0.162 ± 0.038
0.173 ± 0.022
0.216 ± 0.028
0.201 ± 0.023
0.253 ± 0.024
ICA
0.169 ± 0.040
0.169 ± 0.036
0.211 ± 0.056
0.193 ± 0.021
0.215 ± 0.026
Appendix
Table 10: Normalized top-20% policy value on the perinatal semi-synthetic dataset at dim(ϕ)=10 using true experimental Y . Entries are means ± standard errors over 10 runs; higher is better.
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.
Shi Bo, Debarghya Mukherjee, AmirEmad Ghassami
Department of Mathematics and Statistics, Boston University
Predicting the effect of interventions with many possible variations, e.g., therapeutic content that affects mental health outcomes or an earnings call transcript that drives movement in share price, is useful across several domains. However, classical causal estimators tend to assume that all possible interventions are observed, which is infeasible when interventions vary widely, for instance, in the space of all text strings. We adapt a well-known approach of recasting causal inference as a learning problem, to address high-dimensional treatment spaces. Specifically, under standard assumptions like no unobserved confounding, we show that causal error decomposes into a series of moment-balancing errors of increasing order, and design objectives that directly improve causal estimation. We also show how to project the effect of a high-dimensional treatment onto lower-dimensional treatment attributes, which allows a single model to answer several causal questions without additional attribute-specific training. We empirically evaluate our estimators in settings with high-dimensional continuous, discrete, and text treatments, the last of which used a semi-synthetic dataset of Amazon Reviews. Our experiments demonstrate the benefit of higher-order balance error optimization and competitive performance of projected causal estimates with attribute-specific estimators.
Nikita Dhawan, Arnav Paruthi, Andrew Kim +3
University of Toronto, Vector Institute · University of Toronto · Vanguard
We propose a two-stage estimator for structural mediation parameters that combines deep representation learning with G-estimation under the "no essential heterogeneity" (NEH) assumption. We call the method UNIT. In the first stage,TARNet estimates the heterogeneous effect of a randomized treatment on a mediator by learning a shared covariate representation across treatment arms.The resulting conditional average treatment effect (CATE) estimate provides a plug-in approximation to the heterogeneity-dependent component of the weight function entering the G-estimating equation of Zheng and Zhou (2015), which identifies the structural parameters even in the presence of unmeasured mediator-outcome confounding. We show that more accurate first-stage representation learning can yield a more informative plug-in weight and thereby improve the precision of the structural parameter estimator. In simulations with non-Gaussian covariates and nonlinear mediator effects, TARNet weights reduce the Stage-2 standard error of the mediation coefficient by a factor of 1.45 to 1.51 (median across replications, n≥2000) relative to the classical approach, at no cost to bias or coverage.