Source-linked AI summary
Learning Representations for Counterfactual Inference
Fredrik D. Johansson, Uri Shalit, David Sontag
TL;DR
The paper studies how to answer counterfactual questions from observational data, where only one intervention outcome is observed for each context. It combines representation learning with domain adaptation, using balanced representations and discrepancy-based regularization, and reports stronger counterfactual inference performance, while noting limitations for linear outcome models without treatment interactions.
Problem
Counterfactual inference from observational studies must estimate outcomes for interventions not observed for each context, making causal effects unavailable directly from the data.
Method
The method learns representations that balance treatment populations while supporting factual and counterfactual outcome prediction, using domain adaptation and discrepancy-based regularization.
Results
The balancing neural network BNN-2-2 performs best on both IHDP and News for estimating ITE and PEHE and remains competitive for ATE.
Takeaways & Limitations
Learning balanced representations improves counterfactual inference in the paper’s theoretical analysis and experiments across two datasets.
Takeaways & Limitations
Linear models and neural networks with a single linear outcome layer cannot capture individual treatment-effect differences because they lack interactions between Φ(x) and t.
Abstract
from arXiv · showhide
Observational studies are rising in importance due to the widespread accumulation of data in fields such as healthcare, education, employment and ecology. We consider the task of answering counterfactual questions such as, "Would this patient have lower blood sugar had she received a different medication?". We propose a new algorithmic framework for counterfactual inference which brings together ideas from domain adaptation and representation learning. In addition to a theoretical justification, we perform an empirical comparison with previous approaches to causal inference from observational data. Our deep learning algorithm significantly outperforms the previous state-of-the-art.
1. Introduction
The paper addresses counterfactual inference in observational studies by learning representations that combine domain adaptation with representation learning. It proposes balanced representations and evaluates them theoretically and empirically across simulated and real-world tasks.
- Counterfactual inference asks what outcome would have occurred under an alternative intervention, using observational data from fields including healthcare, education, employment, and ecology.
- The empirical evaluation covers simulated and real-world tasks, while the paper also provides a theoretical justification for the proposed approach.
- The method connects counterfactual inference to domain adaptation and covariate shift, treating differing intervention populations as distributions requiring balanced representations.
- The proposed framework regularizes learned representations by encouraging similarity between populations receiving different interventions, reducing variance when transferring predictions across distributions.
- The paper develops linear and deep-learning representation algorithms and argues that balanced representations improve counterfactual inference compared with sample re-weighting approaches.
2. Problem setup
Counterfactual inference estimates outcomes for interventions not observed for each context, creating a distribution shift between factual data and the counterfactual population. The paper frames this as covariate shift and uses representations, prediction functions, and discrepancy minimization to address it.
- For each context, only one potential outcome is observed, so the outcome under the alternative treatment must be inferred.
- The individualized treatment effect is Y1(x) − Y0(x), while the average treatment effect averages this difference over a population distribution.
- Estimating the individualized treatment effect requires predicting outcomes for the counterfactual set {(x_i, 1−t_i)} rather than only the observed factual samples.
- Because factual and counterfactual distributions can differ through the treatment assignment mechanism, counterfactual prediction is a covariate-shift problem.
- Figure 1 depicts representation-based response prediction with treatment-group information while minimizing distributional imbalance measured by discrepancy.
3. Balancing counterfactual regression
The method learns representations and outcome hypotheses that balance treated and control populations while predicting factual and counterfactual outcomes. It combines prediction losses, counterfactual-neighbor penalties, and domain-adaptation discrepancy minimization, with linear and neural implementations.
- Core objective: The learned representation balances treatment populations while supporting prediction of observed and unobserved outcomes.It uses a representation Φ and hypothesis h, optimizing factual prediction, counterfactual prediction, and distributional balance.
- Core objective: Counterfactual predictions are regularized toward the nearest observed outcome from the corresponding treated or control group.The nearest opposite-treatment neighbor is computed once in the input space and remains fixed during representation learning.
- Core objective: The method minimizes discrepancy between treated and control representations to improve generalization from factual to counterfactual distributions.The discrepancy is a hypothesis-class-dependent distance from domain adaptation; for complex hypothesis spaces it generally lacks an exact closed form.
- Linear representations: Balancing Linear Regression learns sparse feature re-weightings, trading predictive accuracy against similarity between treatment-group distributions.Features differing substantially between groups receive smaller weights, and the objective is optimized by alternating sub-gradient descent.
- Neural representations: The neural approach learns Φ with initial hidden layers, then combines the representation with treatment assignment to predict outcomes.The representation layers support discrepancy computation, while subsequent layers receive treatment as additional input.
- Neural representations: Models with a single linear outcome layer cannot capture individual treatment-effect differences because they omit interactions between Φ(x) and treatment.Adding nonlinear post-concatenation layers can introduce interactions but removes the closed-form discrepancy expression.
4. Theory
The theory bounds relative counterfactual generalization error using measurable fit, discrepancy, and overlap-related quantities. It motivates learning representations that balance treated and control groups while preserving predictive accuracy.
- Theory: Theorem 1 bounds relative counterfactual error for a ridge model trained on factual outcomes and evaluated on counterfactual data.The bound compares this error with ridge regression trained on the unobserved counterfactual outcomes.
- Theory: The bound includes a discrepancy term measuring differences between factual and counterfactual representation distributions.For linear hypotheses, this discrepancy is related to differences between treated and control representation means.
- Theory: The unknown counterfactual-fitting term is bounded using factual fitting, nearest-neighbor counterfactual fitting, and an overlap-dependent distance term.The overlap term is smaller when potential-outcome functions are smooth and treated and control groups have nearby counterparts.
- Theory: The proposed procedure first minimizes discrepancy and fitting terms over the representation, then performs ridge regression on factual outcomes and treatment assignments.This can reduce reliance on features strongly associated with treatment but weakly associated with outcomes.
- Linear discrepancy: For linear hypotheses, minimizing discrepancy is equivalent to matching treated and control means in representation space.The discrepancy is expressed through second-order moments and reduces to a weighted difference between treated and control means.
5. Related work
Prior causal-inference methods include non-parametric matching and weighting, parametric outcome models, doubly robust estimators, and covariate-balancing approaches.
- Non-parametric methods: Non-parametric methods include nearest-neighbor matching, propensity score matching, and propensity score re-weighting.These methods do not model the relation between context, intervention, and outcome parametrically.
- Parametric methods: Parametric methods model the relation between context, intervention, and outcome using regressions, random forests, or regression trees.Examples include linear and logistic regression.
- Doubly robust methods: Doubly robust methods combine propensity-score-weighted regression with aspects of parametric and non-parametric approaches.Their efficacy may decline when treatment assignment probabilities must be estimated, as in most observational studies.
- Balancing approaches: Earlier covariate-balancing work transformed or selected covariates while modeling treatment-covariate interactions.This represents one of the few prior approaches explicitly achieving balance through covariate transformation or selection.
6. Experiments
The experiments evaluate balancing linear and neural representation methods on IHDP and a newly simulated News dataset, comparing them with established counterfactual-inference baselines. Across both datasets, balancing neural networks perform best for individual treatment-effect estimation, while balancing and variable-selection methods show advantages over standard neural networks or strong baselines in specific comparisons.
- Experimental setup: The evaluation compares BLR and BNN variants with linear regression, doubly robust regression, LASSO + RIDGE, BART, and a standard neural network.The experiments use IHDP and News data, with results reported in Tables 1 and 2.
- Metrics: The experiments report ITE RMSE, absolute ATE error, and PEHE, with smaller values indicating better performance.PEHE requires accurate estimation of both factual and counterfactual responses, and the experiments use simulated outcomes with access to counterfactual samples.
- Datasets: The simulated News dataset represents news items by word counts, assigns desktop or mobile viewing as treatment, and models outcomes using topic representations.It contains 5000 news items based on 50 LDA topics and a 3477-word vocabulary, with 50 evaluation realizations.
- Results: BNN-2-2 performs best on both datasets for estimating ITE and PEHE and remains competitive for ATE.The authors attribute the result to balancing regularization helping avoid overfitting representations to factual outcomes.
- Results: For linear methods, BLR and LASSO + RIDGE perform best for ITE estimation, with LASSO + RIDGE described as a strong baseline competitive with BLR.On News, the two methods perform equally well despite selecting different variables.
- Results: On News, BNN-4-0, BLR, and LASSO + RIDGE outperform the standard neural network, while BART may be disadvantaged by dataset dimensionality.The authors note that BART could improve with hyperparameter tuning.
7. Conclusion
The paper treats counterfactual inference as a domain adaptation problem and learns representations with similar treated and untreated distributions. It leaves extensions to multiple treatments, improved optimization, and richer discrepancy measures as open questions.
- Counterfactual inference is formulated as a domain adaptation problem, motivating representations suited to this task.
- Balanced representations have similar distributions among treated and untreated populations.
- Trading off representation balance with standard data-fitting and regularization terms is presented as practically and theoretically prudent.
- Open questions include extending the method beyond two treatments, improving optimization algorithms, and using richer discrepancy measures.
A. Proof of Theorem 1
The proof connects factual and counterfactual prediction through domain adaptation bounds, representation control, and assumptions on potential outcomes. It uses nearest opposite-treatment neighbors and Lipschitz continuity to bound counterfactual error.
- The proof treats the factual distribution as the source domain and the counterfactual distribution as the target domain.
- The analysis explicitly bounds the approximation term because the linear hypothesis space is weaker than the RKHS space used in prior work.
- The factual outcome uses the received treatment, while the counterfactual outcome uses the opposite treatment.
- Potential outcomes are assumed to be Lipschitz continuous on a metric space, with constants K0 and K1 controlling outcome changes with distance.
- The theorem analyzes ridge-regression solutions for factual and counterfactual empirical distributions within a linear hypothesis class.
- For each unit, the proof selects a nearest neighbor from the group receiving the opposite treatment and uses that distance in the bound.