Supervised Learning Objectives for High-Dimensional Treatment Effects
Abstract
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 such settings, assuming realizability of the learned function and no unobserved confounding. We show that causal error decomposes into moment-balancing errors of increasing order and exploit structure of the conditional outcome in the confounders to design objectives that directly improve causal estimation. We also show how to project the effect of a high-dimensional treatment onto coarser 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. Our experiments demonstrate generalization to held-out treatments, benefits of balancing higher-order moments when required, and projected causal estimates that are competitive with attribute-specific ones.
Related
- Deformed Q-learning for High-Dimensional Dynamic Treatment Regimes
- Neural Supervised Clustering for Learning Individualized Treatment Rules with Many Treatments
- Different representation learning objectives recover distinct latent structures from the same psychometric data
- Energetic Variational Deep Learning for Solving High-Dimensional Stationary Fokker-Planck Equations
- Diversity-Aware Active Learning Builds Materials Datasets for Current and Future Objectives
BibTeX
@misc{dhawan2026supervised,
title = {{Supervised Learning Objectives for High-Dimensional Treatment Effects}},
author = {Nikita Dhawan and Arnav Paruthi and Andrew Kim and Lovedeep Gondara and Jekaterina Novikova and Chris J. Maddison},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2605.27281},
eprint = {2605.27281},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2605.27281},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - Supervised Learning Objectives for High-Dimensional Treatment Effects
AU - Nikita Dhawan
AU - Arnav Paruthi
AU - Andrew Kim
AU - Lovedeep Gondara
AU - Jekaterina Novikova
AU - Chris J. Maddison
PY - 2026
DA - 2026/10/15
PB - arXiv
DO - 10.48550/arxiv.2605.27281
UR - https://arxiv.org/abs/2605.27281
AB - 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 such settings, assuming realizability of the learned function and no unobserved confounding. We show that causal error decomposes into moment-balancing errors of increasing order and exploit structure of the conditional outcome in the confounders to design objectives that directly improve causal estimation. We also show how to project the effect of a high-dimensional treatment onto coarser 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. Our experiments demonstrate generalization to held-out treatments, benefits of balancing higher-order moments when required, and projected causal estimates that are competitive with attribute-specific ones.
ER -