Source-linked AI summary
Generalization Bounds and Representation Learning for Estimation of Potential Outcomes and Causal Effects
Fredrik D. Johansson, Uri Shalit, Nathan Kallus, David Sontag
TL;DR
Estimating individual causal effects from observational records is difficult because only the outcome under the treatment taken is observed, and treatment groups may differ. The paper derives distribution-distance bounds and representation-learning methods with re-weighting; it shows that expected risk can still be bounded without assuming treatment-group overlap.
Problem
Method
The paper derives risk bounds based on treatment-group distribution distances and uses re-weighting and learned representations to minimize bound-motivated objectives.
Source: §1 · pp. 2–3, §5 · p. 22Results
The bounds do not require treatment-group overlap and can still bound the expected error of hypotheses in a given class.
Source: §8 · p. 36Takeaways & Limitations
Even without overlap, the framework can bound expected risk, although causal effects may remain nonparametrically unidentified.
Source: §8 · p. 36Without overlap, the bound is not generally tight, though it may be informative when the lack of overlap is small.
Source: §4.4 · p. 20
Abstract
from arXiv · showhide
Practitioners in diverse fields such as healthcare, economics and education are eager to apply machine learning to improve decision making. The cost and impracticality of performing experiments and a recent monumental increase in electronic record keeping has brought attention to the problem of evaluating decisions based on non-experimental observational data. This is the setting of this work. In particular, we study estimation of individual-level causal effects, such as a single patient's response to alternative medication, from recorded contexts, decisions and outcomes. We give generalization bounds on the error in estimated effects based on distance measures between groups receiving different treatments, allowing for sample re-weighting. We provide conditions under which our bound is tight and show how it relates to results for unsupervised domain adaptation. Led by our theoretical results, we devise representation learning algorithms that minimize our bound, by regularizing the representation's induced treatment group distance, and encourage sharing of information between treatment groups. We extend these algorithms to simultaneously learn a weighted representation to further reduce treatment group distances. Finally, an experimental evaluation on real and synthetic data shows the value of our proposed representation architecture and regularization scheme.
1 Introduction
The paper studies individual-level potential outcomes and causal effects from observational data, where only the outcome under the intervention actually taken is observed. It develops distance-based generalization bounds and representation-learning methods to address treatment-group differences.
1 Introduction
- Generalization guarantees for flexible causal models must account for confounding and systematic differences between treatment groups, which can undermine estimation when groups have limited overlap.The paper assumes confounders are measured and can be adjusted for, but notes that partial overlap may prevent consistent estimation.Source: §1 · p. 2
- The paper bounds potential-outcome and CATE prediction risk using distributional distances between treatment groups, with sample re-weighting and connections to unsupervised domain adaptation.The bounds apply to finite samples and misspecified hypothesis classes.Source: §1 · pp. 2–3
- Weighting makes a bias–variance trade-off explicit: more uniform, potentially biased weights increase treatment-group distance while reducing variance.The propensity score is typically unknown, and estimating it can introduce bias.Source: §1 · p. 3
- The authors extend the bounds to learned representations and develop regularized algorithms that combine representation learning with re-weighting.Under appropriate assumptions, the bounds can be tight with large samples and overlap, and the algorithms can consistently estimate causal effects; the framework is evaluated on synthetic and real-world benchmarks.Source: §1 · p. 3
2 Potential outcomes and effects of interventions
The potential-outcomes framework defines causal effects by comparing outcomes under alternative treatments, although only the outcome under the treatment received is observed. Identification relies on assumptions about confounding and treatment availability across contexts.
2 Potential outcomes and effects of interventions
- The individual treatment effect is the difference between a unit’s two potential outcomes, but the outcome under the treatment not received is an unobserved counterfactual.The CATE is the expected treatment effect conditional on context and is the paper’s feature-based target for individualized effect estimation.Source: §2 · p. 5
- Observational treatment assignment can create context-distribution differences between treated and control groups, producing confounding bias and inflated variance if unadjusted.Figure 2 illustrates that group-average outcomes can reverse the direction of the causal effect present at every context.Source: §2 · p. 6, Figure 2 · p. 6
2.1 Identifying assumptions
- Under ignorability, overlap, and SUTVA, expected potential outcomes and CATE are identifiable from observed data.Ignorability requires conditional independence of treatment and potential outcomes given context; overlap requires each treatment to have nonzero probability in every context.Source: §2 · p. 5, §2.1 · p. 7
- Ignorability and SUTVA are untestable from observational data alone, while overlap becomes harder to satisfy and check as context dimensionality grows.Without additional assumptions, disjoint treatment-group contexts can prevent identification across the population.Source: §2.1 · pp. 6–7
3 Related work
Related work spans outcome regression, weighting, representation learning, and causal-estimation guarantees. This paper situates its bounds and algorithms among these approaches, emphasizing finite-sample distribution differences and partially satisfied overlap.
3 Related work
- Outcome-regression approaches include treatment-specific T-learners, single-model S-learners, direct effect models, and meta-learners, with bias and regularization concerns in high-dimensional, strongly confounded, small-sample settings.Under appropriate assumptions, several methods achieve asymptotic consistency or optimal rates when their regression functions are well specified or nonparametrically estimated.Source: §3 · p. 8
- Prior representation-learning methods share outcome information across treatment groups and regularize their distributional distance, while this work combines representation learning and re-weighting to address destructive representations.The related work also considers preserving context information, disentangling causes of treatment and outcomes, and estimating counterfactual uncertainty.Source: §3 · pp. 8–9
- Propensity weighting can require estimating unknown propensities and can have high variance when propensities are small; learned weights instead minimize distributional distance in a prespecified representation.This paper develops weighted risk minimization for potential outcomes and CATE in fixed and learned representations.Source: §3 · p. 9
- The paper relates its learned representations to prognostic scores, trading approximate outcome sufficiency against a balancing objective for finite-sample CATE estimation.It also focuses on settings where overlap is only partially satisfied, which prior work has addressed less extensively.Source: §3 · pp. 9–10
- Unlike asymptotic minimax results, the paper’s bounds address finite-sample risk affected by treatment-group differences and restricted hypothesis classes.Prior work finds that the best asymptotic and low-sample estimators need not coincide in general.Source: §3 · p. 10
4 Generalization bounds for estimating potential outcomes and CATE
This section develops finite-sample generalization bounds for potential-outcome and CATE estimates, accounting for treatment-group differences through reweighting and function-class complexity. It also examines when the bounds can be tight or loosened, including that the main bound does not require treatment-group overlap, though consistent nonparametric estimation does.
4 Generalization bounds for estimating potential outcomes and CATE
- The fully observable risk bounds combine reweighting and finite-sample guarantees, and the main generalization bound does not require treatment-group overlap.Consistent nonparametric estimation still requires overlap.Source: §4 · p. 11
4.1 Risk for hypotheses of potential outcomes and CATE
- CATE risk is bounded through the risks of separate potential-outcome hypotheses, but that decomposition may be suboptimal when the CATE is simpler than either outcome function.CATE risk is not directly computable from observed data, making counterfactual risk the central difficulty.Source: §4.1 · pp. 12–13
4.2 Importance-weighting hypotheses & propensity scores
- Under overlap, balancing weights based on propensity scores can make weighted factual risk match population risk; partial overlap and unobserved confounding limit what is identified.Weights based on propensity scores may also have poor finite-sample behavior, motivating bounds for arbitrary weights.Source: §4.2 · pp. 14–15
4.3 Bounds on the risk of a re-weighted estimator
- Assuming expected pointwise losses belong to a specified function family, risk differences are bounded by an integral probability metric, with tightness depending on the weighting and overlap conditions.The loss-family assumption can encode boundedness or functional properties such as Lipschitz continuity or RKHS membership.Source: §4.3 · pp. 15–16
4.4 Bounds based on finite samples
- Finite-sample guarantees apply to any valid weighting function and, with the reweighting bound, separate the bias and variance introduced by weighting.The finite-sample step bounds empirical reweighted risk estimates against their expected counterparts.Source: §4.4 · pp. 17–18
- Wasserstein distance learning has O(d) dimension-dependent sample complexity, while kernel MMD has O(1) dimension dependence; the Wasserstein-based bound has worse sample complexity.The MMD sample-complexity bound applies to universal kernels.Source: §4.4 · pp. 18–19
- The bound remains informative when the loss-function family is larger than the hypothesis class, provided the hypothesis pseudo-dimension and loss-family kernel norm are bounded.The example uses fixed-size deep neural networks for hypotheses and an RBF-kernel RKHS for losses.Source: §4.4 · p. 19
- With overlap and known propensity scores, importance weighting makes the bound tight as sample size grows; without overlap, it may remain informative but is not generally tight.Without overlap, density-ratio weights may be undefined where one treatment group has zero density, leaving a nonzero IPM.Source: §4.4 · p. 20
- Reweighting creates a bias-variance trade-off: balancing dissimilar treatment groups can reduce bias but increase variance, and biased weights may yield a smaller population-risk bound.This can hold even when the true treatment propensity is known.Source: §4.4 · p. 20
- Treatment-predictive features that do not predict outcomes, such as instrumental variables, can inflate the IPM without improving risk, motivating bounds on learned representations.Ignoring such information may tighten the bound.Source: §4.4 · p. 20
5 Generalization bounds for representation learning
Section 5 extends marginal-risk bounds to hypotheses operating on learned representations, while Section 5.1 relates the framework to unsupervised domain adaptation and distinguishes its bounds from prior results.
5 Generalization bounds for representation learning
- As input dimension grows, treatment groups tend to diverge and bounds loosen; representations can change which covariate differences contribute to measured distance.The factor CΦ counteracts bound reductions from simply scaling down the representation.Source: §5 · pp. 21, 23
- Invertible representations yield finite-sample generalization bounds on marginal risk using reweighted factual risk and treatment-group distance in representation space.Source: §5 · pp. 22–23
- The representation theorem assumes twice-differentiability and invertibility, and its function-class condition is not guaranteed after changing spaces, though overlap makes it verifiable with infinite data.Source: §5 · pp. 23–24
- A similar bound using the Wasserstein distance is available, but it has worse sample complexity.Source: §5 · p. 23
5.1 Relation to unsupervised domain adaptation
- Prior domain-adaptation bounds do not allow reweighting and may be unnecessarily loose even with overlap; this work addresses learned representations by requiring invertibility.Counterfactual outcome estimation under ignorability is mathematically equivalent to unsupervised domain adaptation under covariate shift.Source: §5.1 · p. 24
6 Estimation
Section 6 turns the theoretical results into practical objectives for estimating potential outcomes and CATE, then describes shared-representation estimators and ways to regularize treatment-group distance. It also discusses implementations of the distributional distances used in that regularization.
6 Estimation
- Practical objectives target potential outcomes and CATE, and the proposed parameterizations draw on theoretical insights even when they are not guaranteed to minimize the bounds.The section also reports empirical results for these parameterizations.Source: §6 · p. 25
6.1 A first learning objective and asymptotic consistency
- Objective (21) combines weighted factual risk with distributional regularization; balancing weights recover inverse-propensity-weighted regression asymptotically, while finite-sample IPM remains nonzero from sampling variance.The finite-sample IPM can remain nonzero even when the treatment-group distributions match.Source: §6.1 · p. 25
- Jointly learned sample weights make the bias–variance trade-off explicit, with representation-dependent weights restricting variance relative to freely chosen sample weights.The objective regularizes the norm of the weights.Source: §6.1 · pp. 25–26
- Under the stated assumptions, minimizers of objective (22) converge to the representation and hypothesis that minimize counterfactual risk as sample size grows.The convergence is to the counterfactual-risk minimizer in the infinite-sample limit.Source: §6.1 · p. 26
6.2 Practical representation learning and learning shared representations
- TARNet shares a representation across treatment groups; CFR adds treatment-group distance regularization, while RCFR also learns weights to trade off selection bias and variance.The shared representation lets the treatment groups share information; the CFR distance penalty is intended to support better counterfactual generalization.Source: §6.2 · pp. 27–28
6.3 Implementing regularization of distributional distance
- CFR and RCFR use empirical kernel MMD and Wasserstein distance as integral probability metrics; Sinkhorn distances approximate costly Wasserstein linear programs by fixed-point iteration.A differentiable kernel supports gradient-based MMD learning; pair sampling reduces quadratic cost but yields a higher-variance unbiased estimator.Source: §6.3 · p. 29
7 Experiments
The experiments evaluate causal-effect methods on synthetic, semi-synthetic, and real-world data, separating the impacts of architecture and balancing regularization.
7.1 Experimental setup & baselines
- Models were assessed on held-out data using CATE MSE when labels were available or policy risk when propensity scores were known; nearest-neighbor pseudo-labels served as a hyperparameter-selection surrogate and may mildly bias selection.Comparators included OLS S- and T-learners, k-NN, TMLE, random forests, BART, causal forests, and prior balancing estimators.Source: Evalutation criteria & hyperparameter… · pp. 30–31, §7.1 · p. 30
7.2 Synthesized outcome: IHDP
- On IHDP, neural estimators outperformed selected baselines, while CFR improved only marginally over TARNet for CATE estimation.Table 1 reports CATE and ATE mean squared errors within- and out-of-sample.Source: Table 1 · p. 32, §7.2 · p. 33
- IHDP uses synthesized outcomes on a confounded, overlap-limited cohort; its selected setting guarantees ignorability and produces nonlinear CATE with sparse coefficients.The cohort has 747 subjects after biased removal of treated children with nonwhite mothers; the selected outcome model has 60% coefficient sparsity.Source: §7.2 · p. 32
- On IHDP, CATE error decreased as α rose to about 1 and worsened above 2; learned non-uniform weights reduced sensitivity to large α in the smaller architecture.The best result in the learned-weight experiment used a larger α and a small penalty on weight non-uniformity.Source: Figure 8 · p. 33, §7.2 · pp. 33–34
7.3 Partially randomized study: National Supported Work program.
- On Jobs, nuanced nonlinear policies had lower policy risk than treat-all logistic regression, but the out-of-sample difference narrowed; k-NN generalized poorly and had the worst baseline ATT estimate.The table reports policy risk and ATT error; Figure 9 plots policy risk across treatment thresholds.Source: Table 2 · p. 35, Figure 9 · p. 35, §7.3 · p. 36
8 Discussion
The discussion connects representation learning and sample re-weighting for causal estimation, including settings without treatment-group overlap. It also identifies binary treatments and limited benchmark representativeness as important boundaries.
8 Discussion
- Without treatment-group overlap, the bounds still constrain expected hypothesis risk, although nonparametric causal effects are not identifiable.The discussion notes that overlap may be absent when covariates strongly affect treatment but weakly affect outcomes, including possible high-dimensional settings; this is presented as a conjecture.Source: §8 · p. 36
- Representation learning and sample re-weighting offer a unified perspective on treatment-group shift and when either or both approaches may be useful.The discussion highlights finite-sample behavior and overlap violations as settings where the problem is less understood than asymptotic estimation under overlap and ignorability.Source: §8 · p. 36
- The bounds cover binary treatments; potential-outcome risk can extend one-vs-all, but causal-effect bounds need another approach.An all-pairs treatment-group distance is proposed for causal effects, though it may be inefficient with many treatments; the appropriate approach depends on the estimand.Source: §8 · p. 37
- IHDP and Jobs are not necessarily representative of application problems: IHDP is small and relatively simple, while Jobs has small samples and few measured covariates.The authors say comparisons are most informative among similar neural-network models and call for a broader set of benchmarks reflecting real-world observational studies.Source: §8 · p. 37
A Proof of Theorem 3
The proof establishes convergence of the weighted learning objective’s minimizers to a representation and hypothesis minimizing counterfactual risk as sample size grows.
A Proof of Theorem 3
- Under the stated assumptions, minimizers of objective (22) converge to the representation and hypothesis minimizing counterfactual risk as sample size tends to infinity.The theorem statement specifies sufficiently large objective parameters and a minimizer based on n samples.Source: A Proof of Theorem 3 · p. 46
- The proof bounds a stochastic term C by O(1/√n), then obtains the final result from Theorem 2.The proof invokes Markov’s inequality for this bound.Source: A Proof of Theorem 3 · p. 48
- The theorem setup assumes a bounded-kernel RKHS and weak overlap, including a finite second moment for the policy-density ratio in Corollary 3.The marginal risk is defined as expected loss for predicting a single potential outcome over the population.Source: A Proof of Theorem 3 · pp. 46, 49
B Experiment details
The experiment-details section directs readers to Table 3 for hyperparameters and search ranges for TARNet, CFR Wass, and CFR MMD.
B Experiment details
- Table 3 is referenced for hyperparameters and search ranges for TARNet, CFR Wass, and CFR MMD.The supplied table caption identifies it as a hyperparameters-and-ranges table.Source: B Experiment details · p. 49
C Architecture for joint learning of sample weights
The section points to Figure 10 for an illustration of the re-weighed CFR estimator.
C Architecture for joint learning of sample weights
- Figure 10 illustrates the re-weighed CFR estimator.The passage identifies the figure’s subject but does not state a result to compare.Source: C Architecture for joint learning of… · p. 49
D Minimization of approximate Wasserstein distances
The section approximates Wasserstein distances with entropic regularization and Sinkhorn iterations, then describes how to compute gradients for learning representations. The procedure uses treatment-control minibatches and can be implemented with neural-network training frameworks.
D Minimization of approximate Wasserstein distances
- Entropic regularization makes Wasserstein distance approximations computable with Sinkhorn-Knopp matrix scaling, at orders of magnitude faster speed than solving a linear program.The resulting Sinkhorn distance is computed through fixed-point iterations involving repeated multiplication by a kernel matrix.Source: D Minimization of approximate… · p. 49
- The stochastic-gradient procedure samples treatment and control minibatches, builds their pairwise representation-distance matrix, and computes an approximate optimal transport matrix.The minibatch contains m treated and m′ control units; pairwise distances form an m × m′ matrix.Source: D Minimization of approximate… · p. 50
- Differentiating through Sinkhorn iterations computes the gradient; disregarding the gradient of T* amounts to minimizing an upper bound on Sinkhorn transport.The passage identifies this as the computation of gradient g1.Source: D Minimization of approximate… · p. 50
- The framework is agnostic to the representation parameterization, while experiments use neural networks and may represent Sinkhorn fixed-point iterations as a recurrent neural network.The objective can be minimized with commonly used neural-network training frameworks out of the box.Source: D Minimization of approximate… · pp. 50–51
- The kernel matrix corresponds to a metric such as Euclidean distance, with entries defined from distances between represented inputs.The passage also identifies n_c and n_t as the control- and treatment-group sizes.Source: D Minimization of approximate… · p. 50