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

• Small randomized experiments and high-dimensional covariates make CATE estimation difficult, and existing practice may substitute manually specified subgroups for data-driven heterogeneity discovery.Source: §1 · p. 1
• Historical covariates and outcome indicators can reflect latent factors such as social support and health conditions that may drive treatment-effect heterogeneity.Source: §1 · p. 2
• Historical auxiliary outcomes train a lower-dimensional representation that can identify heterogeneous effects under surrogacy conditions and improve sample-efficiency on the experimental sample.The authors also characterize gains using existing CATE error rates that depend on covariate dimension.Source: §1 · p. 2
• When assumptions are imperfect, compression coarsens the estimand, but reduced experimental-stage variance can outweigh the resulting bias and lower total error.Source: §1 · p. 2
• The method shows substantial sample-efficiency improvements for CATE estimation on synthetic and semi-synthetic medical data.Source: §1 · p. 2

• High-dimensional covariates worsen convergence for standard CATE estimators, while existing dimension-reduction remedies may require manual selection or linearity assumptions.Source: §2 · p. 2
• Causal representation-learning research either learns treatment-effect representations within one dataset or seeks latent causal variables, rather than addressing both multi-outcome historical data and treatment-effect estimation.Source: §2 · pp. 2–3
• Prior surrogacy work studies identification or variance reduction, whereas this method uses surrogate structure to transfer learned representations for sample-efficient CATE estimation, even when the target outcome is observed.Source: §2 · p. 3
• Unlike data-fusion and causal-transfer approaches, this work targets underpowered experiments with historical and experimental covariates from the same population, without requiring the treatment to appear in historical data.Source: §2 · p. 3

• The setting uses a large historical sample and a smaller randomized experiment, with covariates and auxiliary outcomes in both and a target outcome in the historical data.Source: §3 · p. 3
• Historical units need not have received treatment, and observing the target outcome in the experiment is not required for identification under the stated surrogacy conditions.When the target outcome is observed experimentally, reliance on those conditions weakens.Source: §3 · p. 3
• The overview trains ϕ and an outcome model on historical data, then uses a CATE meta-learner on the RCT with either a predicted or observed target outcome.Source: Figure 1 · p. 4
• The goal is to learn a representation ϕ: R^d → R^m with m ≪ d from historical observations and use it to estimate the target-outcome CATE in the experiment.Source: §3 · p. 4
• The setup assumes randomized treatment assignment in the experiment, treatment overlap across covariates, and experimental covariate support contained within historical support.Source: §3 · p. 4
• Surrogacy requires treatment independence from the target outcome conditional on surrogates and covariates, while comparability requires the target-outcome relationship given surrogates and covariates to agree across populations.Source: §3 · p. 4

4.1 Representation requirements

• ϕ must retain outcome-relevant information given S and surrogate-relevant covariate information across treatment; randomization means compression need not preserve confounders.Source: §4 · p. 4, §4.1 · p. 5

4.2 Identification result

• Under Assumptions 3.1–4.1, ϕ(X) identifies τ(x) by contrasting expected h(S, ϕ(x)) across randomized treatment arms.The outcome model h is learned on historical data and applies across arms under surrogacy and comparability.Source: §4.2 · p. 5
• The two-stage procedure learns ϕ and h from historical data, then fits CATE pseudo-outcomes on ϕ(X) in the experiment rather than on X.The historical outcome relationship and randomized shifts in surrogates jointly identify treatment effects; unit-level prediction needs covariates, not post-treatment surrogates.Source: §4.2 · p. 5
• If surrogate sufficiency fails, the estimand remains a causal effect but coarsens to the average treatment effect among units sharing ϕ(X).Heterogeneity along compressed directions is averaged away.Source: §4.2 · p. 6
• Compression adds squared bias, while estimating over ϕ(X) instead of X can reduce experimental-stage variance; MSE falls when the bias is smaller than the variance reduction.The tradeoff is especially relevant when m ≪ d.Source: §4.2 · p. 6

4.3 Objective for obtaining ϕ

• The Information Bottleneck objective maximizes I(Y*; ϕ(X)|S) for outcome relevance and I(S; ϕ(X)) as a historical-data proxy for surrogate relevance.Source: §4.3 · p. 6
• The finite-sample objective weights these terms by λ; feedforward encoders use mutual-information estimator heads or BCE/MSE prediction-loss substitutes.The prediction losses lower-bound the corresponding mutual-information terms.Source: §4.3 · p. 7

5 Empirical Results

• All methods use a two-stage procedure: historical data learns representations where applicable and trains outcome models; experimental data supplies predicted outcomes as CATE targets, estimated with cross-fitted DML.They report mean normalized PEHE and normalized top-20% policy value across 10 random data draws, using the same draws for every method.Source: §5 · pp. 7–8

5.1 Synthetic dataset

• The synthetic data use 1,000 covariates and a 10-dimensional latent vector, with outcomes and treatment effects depending on a five-dimensional component under the main setup.The experiment varies experimental training size from 50 to 750, uses a 10-dimensional learned representation, and separately studies violations by allowing the other latent component to affect outcomes and treatment effects.Source: §5.1 · p. 8
• At n = 100, MI-InfoNCE reaches normalized PEHE 0.650 ± 0.059, versus 0.902 ± 0.047 for PLS and 3.883 ± 0.319 for raw X.Unsupervised baselines reach at least 1.750; MI-InfoNCE’s top-20% policy value is 0.788 ± 0.043 versus 0.091 ± 0.043 for raw X with h(X).Source: §5.1 · p. 8
• On the synthetic data, MI-InfoNCE and PLS are effectively tied by n = 250, and PLS is modestly ahead at larger sample sizes.The linear data-generating process favors PLS, while MI-InfoNCE leads at the smallest sample sizes and substantially outperforms unsupervised baselines throughout.Source: §5.1 · p. 8

5.2 Semi-synthetic medical dataset

• The perinatal semi-synthetic study uses real covariates from 60,248 historical and 11,747 experimental pregnancies, with 37 auxiliary diagnosis features summarized into five principal components.Treatment assignment and potential outcomes are synthetic, and the data construction makes randomization, surrogacy, and outcome-model comparability hold by construction.Source: §5.2 · p. 9
• At n = 100, conditional MINE achieves normalized PEHE 1.340, compared with more than 4.2 for either raw-X baseline.At n = 1000, the prediction encoder and PLS each achieve approximately 1.06, versus 1.70 for h(X) and 2.10 for h(X, S); policy results follow a similar relative trend.Source: §5.2 · p. 9

6 Discussion

• The approach targets internally valid but underpowered experiments by using multi-outcome historical data to improve the statistical power of CATE estimation.The authors also hope their findings motivate more deliberate collection of auxiliary outcomes in administrative and operational data.Source: §6 · p. 9
• Identification relies on surrogacy and sufficiency, although empirical gains are reported as robust to imperfect satisfaction of these conditions.Sufficiency violations coarsen the estimand to an average effect among units sharing ϕ(X); observing Y* in the experiment can sidestep surrogacy violations, while performance under more general failures remains open.Source: §6 · pp. 9–10

A Proof of Theorem 4.3

• Under Assumptions 3.1–4.1, the CATE τ(x) is identified by Equation (3).The proof establishes the result separately for each treatment arm before taking their difference.Source: A Proof of Theorem 4.3 · p. 15
• For each arm t, the experimental potential-outcome mean equals the expectation of h(S, ϕ(x)) conditional on ϕ(x), T = t, and the experimental population.Here h(s, z) is the historical-population conditional expectation of Y given S = s and ϕ(X) = z.Source: A Proof of Theorem 4.3 · p. 15
• The CATE is the difference between those representation-conditioned expectations under treatment and control in the experimental population.The identified expression uses the same historical outcome model h for both treatment arms.Source: A Proof of Theorem 4.3 · p. 15

A.1 Discussion of proof steps

• Randomization identifies each arm’s potential-outcome mean with its observed-outcome mean, and the tower property introduces S into the expectation.The proof then applies surrogacy, comparability, and sufficiency to reach the representation-based expression.Source: §A.1 · p. 15
• Surrogacy removes treatment from the outcome expectation conditional on S and X, and comparability transfers that conditional outcome model from the experimental to historical population.The assumptions are T ⊥ Y | S, X, P = E and Y ⊥ P | S, X, respectively.Source: §A.1 · p. 15
• Sufficiency makes X redundant for predicting Y given ϕ(X) and S, and redundant for predicting S given ϕ(X) and treatment.These conditions ensure the inner outcome model and outer surrogate distribution depend on X only through the representation.Source: §A.1 · p. 15
• Historical data fit h(S, ϕ(X)), while experimental data estimate how treatment shifts the surrogate distribution conditional on ϕ(X).The CATE is obtained by contrasting h’s expected predictions across treatment arms.Source: §A.1 · p. 15

B.1 Ignorability under compression

• Under simple randomization, any measurable function of covariates, including ϕ(X), preserves ignorability for treatment-effect estimation.Randomization makes treatment independent of covariates and potential outcomes, so the representation need not separately retain information for confounding adjustment.Source: §B.1 · p. 16

B.2 Extension to stratified randomizations

• For stratified randomization, concatenating the known strata variables V with ϕ(X) preserves conditional ignorability while keeping the representation dimension m + |V| ≪ d.The stratification variables are retained whether or not the encoder learns them.Source: §B.2 · p. 16

C.1 Identification under violation

• When sufficiency-(ii) fails, the identified target is the ϕ-averaged CATE, which averages treatment-effect heterogeneity over covariates compressed away while preserving heterogeneity along retained directions.It equals the original CATE when that effect is constant on the representation’s level sets.Source: §C.1 · pp. 16–17

C.2 Bias-variance tradeoff under imperfect sufficiency

• Under imperfect sufficiency, representation-based estimation has lower error when its variance reduction exceeds the squared bias from coarsening.Figure 3 shows normalized PEHE rising with α under partial violations, while true-outcome-aware representations retain an advantage.Source: §C.2 · p. 17, Figure 3 · p. 18

C.3 Empirical analysis of sufficiency violation

• The α sweep shifts treatment-effect heterogeneity toward features unlearnable from untreated historical data; α=0 satisfies sufficiency-(ii), while α=1 places all heterogeneity on those features.At α=1, the representation identifies the coarsened CATE rather than the original CATE.Source: §C.3 · p. 18
• At α=1, MI-InfoNCE’s normalized PEHE reaches 1.031 and policy value falls to 0.041, versus 0.423 PEHE and 0.936 policy value at α=0 for n=500.At α=0.5, its normalized PEHE is 0.549; compression can still beat raw X, but cannot recover heterogeneity absent from the representation.Source: §C.3 · pp. 18–19, Figure 4 · p. 19

D Surrogacy Violation

• A direct treatment-to-target path violates full mediation through S; it introduces bias unless Y* is observed experimentally, when randomization identifies the representation-conditional CATE directly..Source: D Surrogacy Violation · pp. 19–20

D.1 Empirical analysis of surrogacy violation (δ)

• The scalar δ controls the strength of treatment’s direct effect on Y* that bypasses S, and this direct term is absent from untreated historical data.At δ=0, the treatment effect on Y* is fully mediated by S; at δ=2, the direct component accounts for 80% of CATE variance.Source: §D.1 · p. 20
• At δ=2 with true experimental outcomes, MI-InfoNCE’s normalized PEHE is 0.898 versus raw X’s 1.961; outcome-aware representations remain more accurate across the sweep.At n=500, MI-InfoNCE’s normalized PEHE rises from 0.423 at δ=0 to 0.767 at δ=1, while policy value falls from 0.936 to 0.665.Source: Figure 6 · p. 21, §D.1 · p. 21
• The evaluation varies δ over {0, 0.5, 1, 2} and sample size n over {50, 100, 250, 500, 750}, using three-fold DML and ten historical-experimental sample pairs.Methods use either historical outcome-model predictions or true experimental outcomes as the CATE target.Source: §D.1 · p. 21

E Prediction losses as mutual information lower bounds

• BCE and MSE losses lower-bound mutual information, while BCE’s bound is tight at convergence with sufficient capacity and MSE’s can remain looser.The MSE bound is tight when the conditional outcome is Gaussian and its conditional variance is constant in X; under additive Gaussian noise, both gaps close.Source: E Prediction losses as mutual… · pp. 21–22
• The policy advantage declines as the direct effect grows but remains positive at δ = 2.The surrogate-mediated component remains informative at that violation level.Source: Figure 7 · p. 22

F Empirical results (cont.)

• The empirical objective uses prediction heads, conditional InfoNCE, or conditional MINE, with MI variants also including a marginal surrogate term weighted by λS − 1.The reported experiments set λS = 2.Source: F Empirical results (cont.) · p. 23

F.1.1 Description

• Synthetic comparisons use X-learner and DML with predicted or observed outcomes as targets, alongside supervised, unsupervised, and raw-feature baselines.They report normalized PEHE and top-k policy value across five experimental sample sizes, averaging the same 10 historical and experimental resamples.Source: §F.1.1 · p. 24

F.1.2 Additional results: figures

• Figures 8–9 cover X-learner and DML results across both outcome targets, while Tables 1–4 tabulate their PEHE and policy-value results.The supplied passages specify the visual scopes and metrics but do not provide table cell values.Source: §F.1.2 · p. 24, §F.1.3 · p. 24, Table 1 · p. 24

F.2 Empirical analysis of representation dimension (m)

• At n = 500, increasing m from 2 to 5 lowers MI-InfoNCE normalized PEHE from 0.961 to 0.538 and raises normalized top-20% policy value from 0.341 to 0.940.Beyond m = 5, gains are not uniform; MI-InfoNCE remains consistently good at m = 10 and 20, while PLS degrades at m = 50.Source: §F.2 · pp. 26–27, Figure 10 · p. 27, Figure 11 · p. 28

F.3.1 Construction

• The semi-synthetic study uses historical EHR data and a distinct experimental-pool cohort, with auxiliary diagnosis trajectories forming five surrogate components.The first five PCA components explain 41.3% of historical auxiliary-outcome variance; treatment affects the target outcome only through the surrogate by construction.Source: §F.3.1 · pp. 28–29

F.3.2 Overlap and representation diagnostics

• Supervised representations outperform raw-X variants and generally unsupervised compression, while diagnostics indicate approximate sufficiency in this experiment.Adding X reduces outcome log loss by at most 0.0104; supervised-representation improvements for predicting potential surrogates range from 0.0146 to 0.0837.Source: Figure 12 · p. 30, §F.3.2 · p. 30

F.3.3 Results and representation dimension

• For the prediction encoder, dimensions 5 and 10 give similar results: dimension 10 has lower normalized PEHE at four of five sample sizes, but neither has uniformly higher policy value.The study retains dimension 10 for consistency with the synthetic experiment and uses dimension 5 as a sensitivity analysis.Source: §F.3.3 · p. 30, Figure 13 · p. 31, Table 8 · p. 32

F.3.4 True-outcome sensitivity

• With true experimental Y as the CATE target, low-dimensional supervised representations retain an advantage over raw X.At n = 100, conditional MINE has normalized PEHE 1.717 ± 0.108 versus 6.833 ± 0.323 for raw X; at n = 1000, PLS and the prediction encoder are nearly tied at 1.059 ± 0.039 and 1.063 ± 0.023, respectively.Source: §F.3.4 · pp. 30–31, Figure 14 · p. 32, Table 9 · p. 32