Supervised Learning Objectives for High-Dimensional Treatment Effects

The gist The authors adapt causal inference by recasting it as a learning problem to address challenges in estimating average potential outcomes for high-dimensional treatments, such as those found in text strings, by decomposing causal error into moment-balancing errors and projecting high-dimensional estimates onto lower-dimensional attributes.

Recasting Causal Inference as Learning

The paper focuses on estimating the average potential outcome (APO) of a high-dimensional treatment, defined as the average outcome that would be obtained under that treatment in a given population in the Neyman–Rubin model The traditional causal estimation approaches struggle when faced with high-dimensional treatments because randomized controlled trials are too expensive or infeasible to conduct while covering the entire space of possible treatments Naïve applications of classical observational estimators also face challenges, such as Inverse Propensity Score Weighting (IPW) being undefined for any string not observed in the training data The approach adapted is to recast causal estimation as a learning problem so that generalization handles previously unseen treatments It is known that causal estimation can be recast as weighted risk minimization, which the authors call Causal Risk Minimization (CRM) throughout.

Decomposition of APO Error

The key contribution addresses two challenges faced by CRM with high-dimensional treatments First, APO estimation via risk minimization requires targets that must themselves be estimated from data, and empirical or trained models of these weights often run the risk of producing extreme weights which in turn harm APO training and prediction. The authors show that the APO estimation error decomposes into a series of moment-balancing errors — errors in balancing the moments of different orders of the confounder distribution across treatment groups This decomposition motivates learning importance weights with explicit balance regularization up to a chosen order. Specifically, under Assumption 1, where the conditional outcome function f(t, X) is a polynomial of degree at most K for all t in T or confounders X are supported on a finite set, the error decomposes as gˆ(t) − g(t) = X/K sum k=0 ck(t) · εk(t), where εk(t) = E[wˆ(T, X) · Xk T = t] − E[Xk].

Learning Stabilized Weights and Moment Balancing

The decomposition suggests a training scheme for learning stabilized weights to minimize all balancing errors up to order K. For IPW-CRM, this is added as a regularization term to the standard propensity score training objective. The balancing property of stabilized weights states that they balance any function h of the confounders X across treatments: E[w(t, X) h(X)] = E[h(X)] for all t in T. Perfect APO estimation via the optimization in eq. (6) requires the learned wˆ(t, X) to satisfy this for h(X) = f(t, X), leading to gˆSW-CRM(t) = E[f(t, X)].

Lower-Dimensional Projections

The paper also shows how to efficiently project the effect of a high-dimensional treatment onto lower-dimensional treatment attributes post hoc, given access to a model of the relationship between the attributes and text. Under Assumption 2, which states that T' is conditionally independent of X and Y given T, the lower-dimensional APO g'(t') is given by g'(t') = ETT'=t' [g(T)] ≈ Pn i=1 I T'i = t' gˆ(Ti). This allows a single learned gˆ to simultaneously answer causal questions about any lower-dimensional attribute that can be labeled given the high-dimensional treatment, without retraining a separate estimator per attribute.

Empirical Evaluation and Results

The authors evaluated their CRM estimators across three settings of increasing complexity: a simple linear setting with continuous treatments, a synthetic dataset with high-dimensional discrete treatments, and a semi-synthetic dataset of real Amazon Reviews. In the simple linear setting, optimizing higher-order momentbalancing errors consistently improved performance of their weighted CRM estimator, outperforming all other variants. For the Amazon Reviews text treatments, SW-CRM avoided training instability and performed competitively with OI and OI-CRM even at K = 1, with higher-order moments (K = 2) further improving performance. Furthermore, projections of the Amazon Reviews APO estimator onto attributes like sentiment and review length matched or outperformed bespoke estimators that were trained separately per attribute.

Conclusion

In this work, they recast average potential outcome (APO) estimation with high-dimensional treatments as causal risk minimization (CRM) over the input space of treatments, overcoming challenges of generalization to unseen treatments They presented an APO error decomposition into momentbalancing errors to learn accurate targets for risk minimization and showed how to efficiently answer multiple causal questions about treatment attributes by projecting a single high-dimensional estimator onto lower-dimensional attributes, with no additional training cost. When the required assumptions are satisfied, their results empirically verify the analytic APO error decomposition and demonstrate performance trends with scaled dataset and model sizes that are consistent with standard risk minimization. The semi-synthetic Amazon Reviews experiments show successful application of CRM estimators to real text treatments and post hoc lower-dimensional projections that are competitive with attribute-specific retraining, unlocking the potential for realworld causal estimation in domains with truly high-dimensional treatments.

Limitations

The work has several limitations, including the strong ignorability assumption being necessary though untestable in practical scenarios, the computational costs of optimizing higher-order moment-balancing errors which are linear in K, generalization relying on inductive biases of the trained model and training data, and the requirement for access to the relationship between high-dimensional treatment and its attributes. The lower-dimensional projected APOs require access to the relationship between the high-dimensional treatment and its attributes, pT'|T, along with the conditional independence assumption that requires T' to depend only on T.

Appendix Notation

The appendix provides notation defining variables such as X Random variable corresponding to confounders, T Random variable corresponding to high-dimensional treatment, and g(t) Average potential outcome (APO) given by E[Y (t)], where the expectation is over some defined population. It also defines eˆ(t, x) Propensity score estimator, to approximate pTX(t x), and wˆ(t, x) Stabilized weights estimator, to approximate pT (t)pT X(tx). The order-k balancing error is denoted as εk(t) Order-k balancing error corresponding to Xk.

Derivations

The derivations show that the APO estimation error of the risk minimizer can be written in terms of moment balancing errors of the confounders. For continuous X, this leads to gˆ(t) − g(t) = X sum k=0 f(k)(t, 0) k! · εk(t), where f(k)(t, 0) = ∂f/(∂x k)x=0. The derivation for lower-dimensional projections shows g'(t') = ETT'=t' [g(T)] ≈ Pn i=1 I T'i = t' gˆ(Ti).

Dataset Details

The paper details three settings of increasing complexity: a linear continuous setting, a synthetic discrete dataset, and the semi-synthetic Amazon Reviews text dataset. The Amazon Reviews text dataset uses n = 10, 000 reviews with X ∈ 0,..., 7 being the product’s total rating count.

1. textbfProjected Causal Inference for High-Dimensional Text Treatments: Capabilities for Multi-Attribute Causal Query Answering. The improved system can simultaneously answer several lower-dimensional causal questions about any lower-dimensional attribute that can be labeled given the high-dimensional treatment, without retraining a separate estimator per attribute. This allows users to compare the estimated APOs of two or more texts or investigate the effects of some lower dimensional attribute of the text, such as length or sentiment, in order to deduce general strategies that favor an outcome.

2. textbfHigher-Order Moment Balancing for Robust APO Estimation: Improved Accuracy with Confounder Knowledge. The system can achieve better causal estimation by optimizing higher-order momentbalancing errors, which consistently improved performance of our weighted CRM estimator, outperforming all other variants. This optimization is crucial for achieving the goal of minimizing the APO estimation error when dealing with complex confounders where the true conditional outcome function has a polynomial structure.

3. textbfReal-World Text Causal Prediction: Competitiveness Against Attribute-Specific Models. The system can perform effectively on real, high-dimensional text data by demonstrating that projections of our Amazon Reviews APO estimator onto attributes like sentiment and review length matched or outperformed bespoke estimators that were trained separately per attribute. This capability allows for the deployment of a single, computationally efficient model for diverse causal queries across different textual features.

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.