1 Introduction↩︎

Effective decision-making in real-world high-stakes domains often relies on predicting the effects of a combination of many different choices. For instance, the success of a financial advisor’s client call may depend on the products recommended, the order in which information is presented, the style of language used, and how these interact. Understanding the effect of all of these elements of a call on a single outcome, e.g., customer satisfaction, is a causal estimation problem with a large space of possible interventions (or treatments), which in this case is the space of all strings. Other examples include a counselor’s response to a patient in distress [1], [2], the phrasing of a public health message [3], [4], word choices in advertisements [5], and the framing of an earnings disclosure [6].

In this work, we focus on estimating the average potential outcome (APO) of a high-dimensional treatment, which is defined as the average outcome that would be obtained under that treatment in a given population in the Neyman–Rubin model [7], [8]. Text strings as treatments is the key motivating setting, since understanding the effect of a whole text intervention can enable many types of decision-making: selecting optimal wording combinations, ranking candidate messages, or identifying linguistic strategies to achieve an outcome. Specifically, our setting is challenging due to the dimensionality of the treatment because each observed data point may have received a unique treatment and test-time treatments may be unobserved during training.

The traditional causal estimation approaches struggle when faced with high-dimensional treatments. 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 to this setting also face several challenges. Inverse propensity score weighting (IPW) [9], [10] computes an empirical sum that is restricted to units that received the exact treatment of interest, and is undefined for any string not observed in the training data. Outcome imputation [11] can sidestep this by learning a model of the conditional outcome given treatment and confounders and marginalizing over the confounders, but this requires inference-time access to individual-level confounder values, which may be impractical when those variables are sensitive (e.g., income or health markers) and prohibitively expensive when the confounders themselves are high-dimensional (e.g., textual confounders).

The approach we adapt 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 [12]–[14], a general approach we call Causal Risk Minimization (CRM) throughout, but works in the CRM literature tend to focus on binary treatments, which obscures some of the challenges with high-dimensional treatments.

Our key contribution is to address two of the challenges faced by causal risk minimization with high-dimensional treatments. First, APO estimation via risk minimization requires targets that must themselves be estimated from data. In the case of IPW, the targets are weighted outcomes, such that APO estimation quality depends directly on the correctness of the learned weights. Empirical or trained models of these weights often run the risk of producing extreme weights which in turn harm APO training and prediction. We address this by showing 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, derived for binary, discrete, and continuous confounders, motivates learning importance weights with explicit balance regularization up to a chosen order. While prior works have used low-order balance regularization to improve propensity score training [15], our error decomposition provides a principled strategy to choose the order upto which balancing confounder moments improves APO estimation itself.

Second, practical use-cases of an APO estimator may include asking lower-dimensional causal questions about the effects of specific text attributes on an outcome. Prior work using texts as treatments has largely circumvented high-dimensionality issues by collapsing the text into a low-dimensional, often binary, summary before classical estimation [5], [16]–[18]. However, this approach may sacrifice valuable information, requires committing to the textual attributes of interest a priori, and involves training a separate estimator per attribute. Instead, we demonstrate how to efficiently project a single high-dimensional APO estimator trained over the full (text) treatment space onto lower-dimensional attributes post hoc, given access to a model of the relationship between the attributes and text. This allows us to simultaneously answer several lower-dimensional causal questions, at no additional training cost per attribute.

We evaluated our 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 [19] with 10,000 unique text treatments. For text treatments, we finetuned a pretrained language model on the CRM loss function that is aligned with common pretraining objectives and hence, can harness the priors that these models encode over natural language [20], [21]. Our experiments empirically verify the analytic APO error decomposition, demonstrate generalization to unseen treatments, and tease apart the effect of different experimental parameters on prediction quality. In particular, we found that optimizing higher-order moment-balancing errors consistently improved performance of our weighted CRM estimator, outperforming all other variants, without access to confounders at inference-time. Finally, 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.

2 Preliminaries↩︎

Let \(\mathcal{T}\) be the space of possible treatments and let \(Y(t)\) for \(t \in \mathcal{T}\), \(X\), and \(T\) be random variables corresponding to the potential outcomes, the confounders, and the treatment, respectively, all drawn from some joint distribution. We assume \(Y(t)\) takes values in either \(\mathbb{R}\) or \(\{0, 1\}\) and \(X\) takes values in \(\mathcal{X}\). Our goal is to estimate the average potential outcome (APO) under treatment \(t\), \[g(t) = \mathbb{E}[Y(t)],\] as defined in the Neyman–Rubin model [7]. We assume access to \(n\) i.i.d. samples \((Y_i, X_i, T_i)\) from an observational distribution, where we abuse notation and let \(Y = Y(T)\) be the random potential outcome realized under the observed treatment \(T\). Since the treatment assignment may depend on \(X\), the naïve conditional mean \(\mathbb{E}[Y \mid T = t]\) is in general biased for \(g(t)\), as illustrated in 1. Throughout, we assume strong ignorability (i.e. no unmeasured confounding): \(\{Y(t)\}_{t \in \mathcal{T}} \perp T \mid X\), positivity: \(0 < p_{T \mid X}( t \mid x ) < 1\), and the stable unit treatment value assumption [22], under which the following classical estimators are valid. We provide a complete list of notation in 7.

Inverse Propensity Score Weighting (IPW). The IPW estimator [9], [10] uses an estimate \(\hat{e}(t, x)\) of the propensity score \(p_{T \mid X}(t \mid x)\) to estimate APOs as: \[\label{eq:ipw95classic} \hat{g}_{\text{IPW}}(t) = \frac{1}{n} \sum_{i=1}^n \frac{Y_i \, \mathbb{I} \Bigl(T_i = t\Bigr)}{\hat{e}(T_i, X_i)}.\tag{1}\] Notice that the indicator \(\mathbb{I}(T_i = t)\) restricts the sum above to units that received treatment \(t\). In high-dimensional settings like those with text treatments, we may only ever observe one datapoint per treatment and for treatments not seen in the training data, this estimator is undefined.

Outcome Imputation (OI). The outcome imputation estimator [11] learns an estimate \(\hat{f}(t, x)\) of the expected conditional outcome \(\mathbb{E}[Y \mid T=t, X=x]\) to estimate APOs as: \[\label{eq:oi95classic} \hat{g}_{\text{OI}}(t) = \frac{1}{n} \sum_{i=1}^n \hat{f}(t, X_i).\tag{2}\] While \(\hat{f}\) may generalize to unseen \(t\) with a correctly specified model class, the Monte Carlo average over \(n\) confounder values for each treatment \(t\) can be computationally expensive at inference time. Moreover, when confounders \(X\) are sensitive (e.g., income or health status), access to individual-level \(X_i\) at inference time may not be feasible for privacy reasons.

These observations motivate the development of estimators that

3 Causal Risk Minimization↩︎

Many causal estimation problems reduce to risk minimization problems [12]–[14], approaches that we collectively call Causal Risk Minimization (CRM). 1 demonstrates with a toy example how weighted risk minimization can recover true APOs given correct weights. The CRM approach that we take is to derive a supervised learning problem whose risk minimizer is the APO function. The critical quantity in this problems is a weight function, whose estimation is the focus of one of our contributions. Our second key contribution is a simple Monte Carlo method for computing APO functions of treatments projected to lower dimensions. For text treatments in particular, the cross-entropy loss variant derived below allows us to finetune pretrained language models on an objective that aligns with their next-token-prediction pretraining objective.

3.1 Deriving CRM Targets↩︎

IPW-CRM. The IPW estimator in 1 is limited by the requirement that treatment \(t\) appear in the observed data. Note that, under strong ignorability and with true propensity scores, \[\label{eq:ipw95crm} \begin{align} g(t) &= \mathbb{E}\left[\frac{Y \mathbb{I} \Bigl(T = t\Bigr)}{p_{T \mid X}(T \mid X)} \right] = \mathbb{E}\left[\mathbb{E}\left[\frac{Y \mathbb{I} \Bigl(T = t\Bigr)}{p_{T \mid X}(T \mid X)} \, \middle| \, T \right]\right] \\ &= \int p_T(s) \mathbb{E}\left[\frac{Y \mathbb{I} \Bigl(T = t\Bigr)}{p_{T \mid X}(T \mid X)} \, \middle| \, T=s \right] ds = \mathbb{E}\left[\frac{Y p_T(t)}{p_{T \mid X}(t \mid X)} \, \middle| \, T=t \right], \end{align}\tag{3}\] where the first three equalities follow, respectively, from the standard IPW estimator, the law of total expectation, and the definition of expectation with respect to \(T\). The final equality holds because in the integral over \(s \in \mathcal{T}\), the indicator variable \(\mathbb{I} (T=t)\) evaluates to \(1\) when \(s=t\) and \(0\) otherwise. 3 motivates fitting a parametric IPW-CRM model \(\hat{g}(t)\) with the following objective: \[\label{eq:ipw95loss} \hat{g}_{\text{IPW-CRM}} = \mathop{\mathrm{arg\,min}}_g \sum_{i=1}^n \mathcal{L}\left(g(T_i), \frac{\hat{p}_T(T_i)}{\hat{e}(T_i, X_i)} \, Y_i \right),\tag{4}\] where \(\hat{e}(T_i, X_i)\) is a learned propensity score, \(\hat{p}_T = \sum_{i=1}^n \hat{e}(T_i, X_i)\) is the marginal treatment probability estimate, and \(\mathcal{L}\) is a suitable loss function. For real \(Y\), we can take \(\mathcal{L}\) as squared error, and for binary \(Y\), we take it as cross-entropy loss with targets \([1 - Y_i \cdot \frac{\hat{p}_T(T_i)}{\hat{e}(T_i, X_i)}, Y_i \cdot \frac{\hat{p}_T(T_i)}{\hat{e}(T_i, X_i)}]\).

SW-CRM. The weighted-outcome targets in 4 are computed by first learning a model for propensity scores. However, we may also directly learn a model of the stabilized weights \(\hat{w}(T_i, X_i)\) [23] to then fit a Stabilized-Weighted APO estimator (SW-CRM) given by: \[\label{eq:sw95loss} \hat{g}_{\text{SW-CRM}} = \mathop{\mathrm{arg\,min}}_g \sum_{i=1}^n \mathcal{L}\left(g(T_i), \hat{w}(T_i, X_i) \, Y_i \right),\tag{5}\] where the loss function \(\mathcal{L}\) is chosen similarly to that for IPW-CRM.

OI-CRM. The OI estimator of 2 required fitting a conditional outcome model \(\hat{f}\) and then estimating APOs with \(\hat{g}(t) = \frac{1}{n}\sum_{i=1}^n \hat{f}(t, X_i)\), assuming access to individual-level confounder values and requiring multiple (\(n\)) forward passes during inference. Given a trained \(\hat{f}\), we can lift this requirement by training an OI-CRM estimator: \[\label{eq:oi95loss} \hat{g}_{\text{OI-CRM}} = \mathop{\mathrm{arg\,min}}_g \sum_{i=1}^n \mathcal{L}\left(g(T_i), \frac{1}{n}\sum_{i=1}^n \hat{f}(T_i, X_i) \right),\tag{6}\] where the loss function \(\mathcal{L}\) is squared error. Note that all estimators learned via 4 5 6 learn a direct mapping from treatments to APOs, do not assume access to confounder values at inference time, and require only a single forward pass to predict the APO for a given treatment.

3.2 APO Error Decomposition↩︎

The solutions of 4 5 6 rely crucially on the targets used in the loss function, which are themselves computed using a learned model of propensity scores, stabilized weights, or conditional outcomes, respectively. However, even small errors in calibration of these learned models, a common pitfall of neural networks, can lead to large errors in APO estimation. In the case of propensity scores and stabilized weights, we can decompose the error in APO estimation into a series of balancing error terms. Confounder balancing is a well-known necessary condition for correct propensity scores [15]. The error decomposition below shows the sufficiency of minimizing balance error upto some order for minimizing APO estimation error via risk minimization. We present the decomposition in terms of the stabilized weights \(w(t, x)\), while the same analysis holds for propensity scores \(e(t, x)\).

The balancing property of stabilized weights states that they balance any function \(h\) of the confounders \(X \in \mathcal{X}\) across treatments: \(\mathbb{E}[ w(t, X) \, h(X)] = \mathbb{E}[h(X)]\) for all \(t \in \mathcal{T}\). Perfect APO estimation via the optimization in 5 requires the learned \(\hat{w}(t, X)\) to satisfy this for \(h(X) = f(t, X)\). \[\label{eq:f95balance} \begin{align} \hat{g}_{\text{SW-CRM}}(t) &= \mathbb{E}[\hat{w}(T, X) \, Y \mid T = t] = \mathbb{E}[ \mathbb{E}[ \hat{w}(T, X) \cdot Y \mid T=t, X ] \mid T = t ] \\ &= \mathbb{E}[ \hat{w}(t, X) \mathbb{E}[ Y \mid T=t, X ] \mid T = t ] = \mathbb{E}[ \hat{w}(t, X) \, f(t, X) \mid T = t ].\\ \hat{g}_{\text{SW-CRM}}(t) &= g(t) = \mathbb{E}[f(t, X)] \iff \mathbb{E}[ \hat{w}(t, X) \, f(t, X) \mid T = t ] = \mathbb{E}[f(t, X)]. \end{align}\tag{7}\]

In other words, the APO estimation error depends directly on how well the weights \(\hat{w}(t, X)\) balance \(f(t, X)\) for each treatment \(t\). Without access to the true \(f\), existing work attempts to learn weights that balance either arbitrary functions of \(X\) [24] or, more commonly, the first moment of \(X\) [25]. Under the following assumption on the structure of \(f(t, \cdot)\), we can relate APO estimation error to moment-balancing errors, such that minimizing balance errors directly minimizes APO error.

Under 1, we can write \(f(t,X) = \sum_{k=0}^K c_k(t)\,X^k)\) for constants \(c_k(t)\), and 7 gives: \[\hat{g}(t) - g(t) = \sum_{k=0}^K c_k(t) \cdot \varepsilon_k(t), \text{ where } \varepsilon_k(t) = \mathbb{E} \left[\hat{w}(T,X) \cdot X^k \mid T=t\right] - \mathbb{E}[X^k], \label{eq:general-decomp}\tag{8}\] is the order-\(k\) balancing error corresponding to \(X^k\). We derive this decomposition for binary, discrete, and continuous confounders in 8.1 This error decomposition suggests a training scheme for learning stabilized weights to minimize all balancing errors upto order \(K\): \[\hat{w}(T_i, X_i) = \mathop{\mathrm{arg\,min}}_w \sum_{i=1}^n \sum_{k=0}^K \lVert \mathbb{E}[w(T_i, X_i) \cdot X^k] - \mathbb{E}[X^k] \rVert^2,\] For IPW-CRM, we add this as a regularization term to the standard propensity score training objective.

3.3 Lower-dimensional projections of high-dimensional APO estimates↩︎

How do we use high-dimensional APO estimates in practice? In the case of text treatments, for instance, we may compare APOs of two texts or rank some number of text candidates to select the best one for an outcome of interest. We may also be interested in 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. We denote this attribute as \(T'\) and assume it satisfies the following.

This assumption holds trivially for deterministic properties of the treatment, for instance, mention of particular words in the text or the length of the text. Under 2, the lower-dimensional APO \(g'(t')\) is given by the expectation of high-dimensional APOs \(g(\cdot)\) over samples where \(T' = t'\), and can be computed via a Monte Carlo estimate as: \[\label{eq:proj} g'(t') = \mathbb{E}_{T| T'=t'} [g(T)] \approx \frac{\sum_{i=1}^n \mathbb{I} \Bigl(T'_i=t'\Bigr) \, \hat{g}(T_i)}{\sum_{i=1}^n \mathbb{I} \Bigl(T'_i=t'\Bigr)}.\tag{9}\] See 8.2 for a complete description of the lower-dimensional projections of treatments and derivation of their APOs. 9 enables the use of a single learned \(\hat{g}\) to simultaneously answers causal questions about any lower-dimensional attribute that can be labeled given the high-dimensional treatment, without retraining a separate estimator per attribute.

4 Related Work↩︎

High-dimensional and text treatments. Causal inference with complex, multi-valued, or high-dimensional treatments is a well-recognized challenge [26], [27]. [28] defined the generalized propensity score (GPS) to extend propensity-score-based techniques to continuous treatments. [29] studied high-dimensional settings (upto 200 treatments) using bag-of-words representations, with access to multiple observations per treatment value. A growing literature applies these ideas when the treatment is expressed in natural language [30], though prior work has largely circumvented the high-dimensionality of text by first extracting lower-dimensional representations [31]. [5] learned lexical features as a treatment proxy, [16] and [17] inferred a latent treatment variable from the text, and [32] fine-tuned a BERT-based model to predict outcomes from (synthetic) text with underlying binary treatment variables. [18] similarly reduced counselor conversations to a handful of attributes (e.g. length, speed, sentiment). [33] highlighted the challenges of using text classifiers for such dimensionality reduction. Large language models are also increasingly being used in causal pipelines where text plays roles other than treatment, for structuring patient self-reports [34] or controlling for it as a confounder [35]. Our work departs from the above by treating each unique text string as a distinct treatment value. CRM estimators learn APOs in the high-dimensional treatment space and can be projected onto lower-dimensional attributes as needed.

The study of text treatments is motivated by a wide range of applications. In finance, earnings call transcripts have been used to estimate the effect of expressed political risk or sentiment on stock volatility and movement [36], [37], and the effect of COVID-19 exposure language on firm capital structure [6], [38]. In consumer settings, the causal effect of review sentiment on purchase behavior [17], [37] and complaint politeness on resolution time [17] have been studied. In mental health applications, text-treatment APO estimation has potential in evaluating therapist response quality [1], [2], [39] and counseling interventions [40]. Across all these settings, prior work has relied on manual annotation or classification to reduce treatments to a lower dimensions a priori.

Covariate balancing. Correct propensity scores balance functions of covariates across treatments, a property that has inspired several methods to learn balancing weights for binary [15], [41], [42] or continuous [43] treatments. [25] combined first-moment balancing weights with regularized regression adjustment for high-dimensional linear outcome models. [44] developed a generalized optimal matching framework that subsumes covariate balance and establishes a balance-variance trade-off. Other approaches include adversarial balancing [24], [45]–[47], end-to-end stable weight learning [48], and joint propensity-outcome training with targeted regularization [49], all for either binary or continuous treatments. Our contribution connects APO estimation error directly to moment-balancing errors and applies to the high-dimensional discrete treatment setting, which allows the handling of text treatments in the space of possible strings.

5 Empirical Evaluation↩︎

Our experiments investigated the following questions:

Datasets. We evaluated different methods on three settings of increasing complexity. First, we verified our error decomposition in a simple linear setting with continuous Gaussian confounders, treatments, and outcomes, where the true APO is a linear function of confounders and the treatment. Next, we constructed a synthetic setting with discrete confounders, binary outcome, high-dimensional discrete treatments to mirror natural language tokens, and nonlinear relationships between different variables. Finally, we evaluated different methods on a real-world dataset of Amazon Reviews [19] in a semi-synthetic setting, where each product review is a different treatment, total number of ratings received by the product is the \(8\)-dimensional confounder, and the purchase of the product by a user reading a given review is the binary outcome of interest. Outcomes were synthetically generated by using GPT-5.1 [50] as a conditional outcome model, which we prompted for a realistic probability of purchase, given each possible confounder value and treatment review. All datasets were split into training, validation, and test subsets such that treatments are non-overlapping across the subsets. In the case of Amazon Reviews, each treatment (or review) is observed only once. Complete details for all datasets and experimental parameters are provided in [app:dataset_details] [app:exp_details], respectively.

Evaluation. Across all settings we evaluated APO estimation quality via the relative mean absolute error (MAE) and correlation between estimated and true APOs over held-out treatments. For estimators that measure and optimize balance errors, we also report the average of balance errors upto some order \(K\), where \(K\) is specific to each setting.

5.1 Does optimizing balance errors reduce APO error in a linear continuous setting?↩︎

We used continuous Gaussian data with a linear true APO function to empirically verify the APO error decomposition into the order-0 and order-1 balance errors, as analytically derived in 9.1. Notice that the continuous setting requires generalization to unseen treatments, as in the case of text treatments. We compared IPW-CRM estimators with and without balance regularization against an oracle outcome imputation estimator (OI) that trains a correctly specified linear network to model the conditional outcomes. 2 shows relative APO MAE, alongside the corresponding order-0 and order-1 balance errors for the IPW-CRM estimators, all of which reduce as sample size increases for all estimators. Without any balance error regularization, IPW-CRM has larger balance errors that persist even as sample size grows, with correspondingly high APO error. In contrast, the balance error regularizer used in IPW-CRM with \(K=1\) substantially reduced balancing errors, which in turn reduced APO error, as expected, and matched the oracle performance of OI. This empirically verifies that minimizing moment-balancing errors directly translates to reduced APO error.

5.2 How does CRM perform with discrete treatments and varying experimental parameters?↩︎

As a step towards high-dimensional text treatments, we used a synthetic dataset of high-dimensional discrete treatments and trained small transformer models to compare the performance of different estimators. As shown in 1, SW-CRM (\(K=2\)) outperforms vanilla OI and other variants. Including higher-order moment-balancing regularization improves its performance. For IPW-CRM, we report results for the best performing \(K=1\). For larger \(K\), IPW-CRM suffered from training instability and hence, poorer performance, a pitfall avoided by SW-CRM.

For the best-performing method, SW-CRM, we conducted further ablation studies to tease apart the effects of different experimental conditions on the correlation between predicted and true APOs, which are visualized in 3. Performance steadily improved with sample size as well as model size, yielding better correlation and lower standard errors. To systematically increase model size, we increased the transformer embedding dimension and its feedforward hidden dimension in proportion. Performance remained stable with increasing confounding strength, suggesting that the estimator was able to correct for large amounts of confounding as long as all confounders were known and the assumption of no unmeasured confounding was satisfied. Finally, we studied the impact of the maximum moment order balanced, \(K\), on APO correlation and the total balance error. 4 shows that as moments of higher orders are included in the balance-error regularizer, APO correlation monotonically increased and correspondingly, the total sum of balance errors upto order-\(4\) decreased. This is consistent with our APO error decomposition for non-binary discrete covariates.

5.3 Do our methods scale to text treatments and large language models?↩︎

To evaluate our methods with treatments in the very high-dimensional and discrete space of all strings, we used the Amazon Reviews dataset, treating each individual review as a treatment, and finetuned Gemma-3-270M to train different components of our estimators. Note that in this more realistic setting, each treatment is observed in only one datapoint and APO estimation relies on the language model’s ability to generalize to unseen treatments. 1 shows similar trends in estimators’ performance to those of the synthetic dataset. IPW-CRM failed more dramatically in this setting, which is explained by extreme propensity scores in the large treatment space of texts and resulting training instability. SW-CRM avoided this instability and performed competitively with OI and OI-CRM even at \(K=1\), with higher-order moments (\(K=2\)) further improving performance. Notably, SW-CRM and OI-CRM achieved high correlations with the true APOs without requiring access to confounder values at test time, which is necessary for inference with standard OI.

5.4 Can we repurpose a single high-dimensional APO estimator to answer several lower-dimensional causal questions efficiently?↩︎

Finally, we explored how a single high-dimensional APO estimator may be operationalized and used in practice. A straightforward use-case is to compare estimated APOs of two or more texts to predict which one would lead to the best outcome of interest. It is also possible to ask lower dimensional causal questions about specific attributes of the text and answer them by projecting the high-dimensional APO estimate down from the space of texts to the space of the attribute in question, using 9 . For instance, in the case of Amazon Reviews, we maybe interested in the effects of the review rating (\(1, \dots, 5\)), sentiment (positive or negative), or length (greater or less than 100 words) on the outcome of purchase. For each estimator, we computed the projected APOs for these three treatment spaces and compared them against their separately retrained attribute-specific counterparts.

2 shows that the projected estimates from a single high-dimensional APO estimator have competitive or better performance than their retrained counterparts, with SW-CRM (\(K=2\)) outperforming other variants. This performance is consistent across different estimators for the Rating and Sentiment treatments, with no additional training cost for each attribute. The case of Length provides a challenging task with two possible treatments (i.e. possible correlations of \(-1\) or \(1\)) with near-constant true APOs. This means that small errors in predictions can reverse the treatment ordering (i.e. \(-1\) correlation), which was the case with projected estimates of OI and OI-CRM. Overall, these results show that a single high-dimensional APO estimator can serve as an efficient tool for answering multiple downstream causal queries simultaneously, without retraining separate estimators for each, as long as the assumption of no unmeasured confounding is satisfied and the lower-dimensional attributes satisfy 2.

6 Limitations and Conclusions↩︎

In this work, we 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. We presented an APO error decomposition into moment-balancing 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.

Our work has several limitations that may impact its application, including those shared with any observational study:

When the required assumptions are satisfied, our 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 real-world causal estimation in domains with truly high-dimensional treatments.

Appendix↩︎

7 Notation↩︎

8 Derivations↩︎

8.1 APO error decompositions into balancing errors↩︎

Recall that under 1, we can write \(f(t,X) = \sum_{k=0}^K c_k(t)\,X^k)\) for constants \(c_k(t)\). Then, using 7 , we have: \[\hat{g}(t) - g(t) = \sum_k c_k(t) \cdot \varepsilon_k(t), \text{ where } \varepsilon_k(t) = \mathbb{E} \left[\hat{w}(T,X) \cdot X^k \mid T=t\right] - \mathbb{E}[X^k], \label{appeq:general-decomp}\tag{10}\] is the order-\(k\) balancing error corresponding to \(X^k\). We now instantiate this decomposition for confounders in different settings: binary, discrete, and continuous.

8.1.0.1 Binary X (warm-up).

If \(X \in \{0,1\}\), we have \(f(t,X) = f(t,0) \cdot 1 + (f(t,1) - f(t,0)) \cdot X\). Substituting into 10 , the APO error depends on two balance error terms: \[\hat{g}(t) - g(t) = f(t,0) \cdot \varepsilon_0(t) + \big[f(t,1) - f(t,0)\big] \cdot \varepsilon_1(t), \label{eq:binary-decomp}\tag{11}\] where \(\varepsilon_0(t) = \mathbb{E}[\hat{w} \mid T=t] - 1\) is a normalization error and \(\varepsilon_1(t) = \mathbb{E}[\hat{w}(t, X)\cdot X \mid T=t] - \mathbb{E}[X]\) is the common first moment balance error. In this case, the APO error of the risk minimizer vanishes whenever the weights are properly normalized and balance the first moment of \(X\).

8.1.0.2 Discrete finite \(X\).

For \(X \in \{0, 1, \ldots, K\}\), the Newton forward difference formula [51], \(f(t, \cdot)\) admits the finite expansion, \[f(t, X) = \sum_{k=0}^{K} \frac{\Delta^k f(t,0)}{k!} \prod_{j=0}^{k-1} (X-j),\] where \(\Delta^k f(t,0) = \sum_{j=0}^k (-1)^{k-j}\binom{k}{j}f(t,j)\) is the \(k\)-th forward difference at zero and the product \(\prod_{j=0}^{k-1} (X-j)\) is also known as the falling factorial. Notice that we may collect terms to rewrite this expansion in terms of powers of \(X\): \(f(t, X) = \sum_{k=0}^{K} c_k(t) X^k\) for some \(c_k(t)\). Hence, \[\hat{g}(t) - g(t) = \sum_{k=0}^{K} c_k(t) \cdot (\mathbb{E} [\hat{w}(t, X) \cdot X^k \mid T=t ] - \mathbb{E} [X^k ]).\] In other words, balancing the moments \(\mathbb{E}[X^k]\) for \(k = 0, \dots, K\) drives the APO error to zero.

8.1.0.3 Continuous \(X\).

Finally, when \(X \in \mathbb{R}\) and \(f\) is a continuous polynomial of order at most \(K\), the Taylor expansion of the true conditional outcome function \(f(t,x) = \mathbb{E}[Y \mid T=t, X]\) around \(X = 0\) yields for each \(t\): \[f(t,X) = \sum_{k=0}^{K}\frac{f^{(k)}(t,0)}{k!}\,X^k \implies g(t) = \mathbb{E}[f(t, X)] = \sum_{k=0}^K\frac{f^{(k)}(t,0)}{k!}\,\mathbb{E}[X^k],\] where \(f^{(k)}(t,0) = \frac{\partial^k f(t,x)}{\partial x^k}\big|_{x=0}\). Similarly, the estimated APO is: \[\hat{g}(t) = \mathbb{E}[\hat{w}(T,X) \cdot f(t, X) \mid T = t] = \sum_{k=0}^{\infty}\frac{f^{(k)}(t,0)}{k!}\,\mathbb{E}\!\left[\hat{w}(T,X)\cdot X^k \mid T=t\right].\] Let the order-\(k\) balancing error be \(\varepsilon_k(t) = \mathbb{E}[\hat{w}(T,X)\cdot X^k \mid T=t] - \mathbb{E}[X^k]\). Then, \[\hat{g}(t) - g(t) = \sum_{k=0}^{\infty}\frac{f^{(k)}(t,0)}{k!}\cdot\varepsilon_k(t).\]

Across all cases, the APO estimation error of the risk minimizer can be written in terms of moment balancing errors of the confounders.

8.2 Deriving lower-dimensional projections of high-dimensional APOs↩︎

Let \(T'\) denote the lower-dimensional attribute of the high-dimensional treatment \(T\), where the relationship between the two variables is specified through a known conditional distribution \(p_{T' \mid T}\), so that \(p_{T'}(t') = \mathbb{E}_T [p_{T'|T}(t'|T)]\). We assume \(T'\) is independent of \(X\) and \(Y(\cdot)\) given \(T\).

To relate the APOs \(g(\cdot)\) over the support of \(T\) to the APOs \(g'(\cdot)\) over the support of \(T'\), consider the ideal completely randomized experiment (CRE) where the assignment distribution over \(T\) induces an assignment distribution over \(T'\). The randomized distribution over all variables is given by \[p_{\text{cre}}(y, x, t, t') = p_{Y(t), X}(y, x) p_{T}(t) p_{T' | T}(t' | t).\] Under this distribution, we can estimate joint-treatment APOs under \(T=t\) and \(T'=t'\) as: \[\begin{align} g(t,t') &= \frac{1}{n} \sum_{i=1}^n \frac{Y_i\mathbb{I} \Bigl(T_i = t, T'_i = t'\Bigr)}{p_{T, T'}(T_i, T'_i)} \\ &\approx \sum_{x,s,s'} \int\frac{y \mathbb{I} \Bigl(s=t,s'=t'\Bigr)}{p_{T,T'}(s,s')}p_{\text{cre}}(y, x, s, s') \, dy \\ &= \mathbb{E}_{X,Y|T=t, T'=t'} [Y | T=t, T'=t'] \\ &= \mathbb{E}_{X,Y|T=t} [Y | T=t] \\ &= g(t), \end{align}\] where the second-to-last equality follows from the conditional independence of \(T'\) and \((X, Y(\cdot))\) given \(T\). Hence, under the randomized distribution over \(T'\), we have APOs given by: \[\begin{align} g'(t') &= \mathbb{E}_{X,Y|T'=t'} [Y | T'=t'] \\ &= \mathbb{E}_{T| T'=t'} [ \mathbb{E}_{X,Y|T'=t', T} [Y | T'=t', T] \\ &= \mathbb{E}_{T| T'=t'} [g(T, t')] \\ &= \mathbb{E}_{T| T'=t'} [g(T)]. \end{align}\] Our experiments estimate \(g\) with different APO estimators \(\hat{g}\) and the expectation above with the Monte Carlo estimate: \[\hat{g}'(t') = \frac{\sum_{i=1}^n \mathbb{I} \Bigl(T'_i=t'\Bigr) \, \hat{g}(T_i)}{\sum_{i=1}^n \mathbb{I} \Bigl(T'_i=t'\Bigr)}.\]

9 Dataset Details↩︎

9.1 Linear Continuous Setting↩︎

We generated data according to: \[\begin{align} X_i &\sim \mathcal{N}(0, 1), \\ T_i \mid X &\sim \mathcal{N}(X, 1), \\ Y_i &= 1 + 2 T + 3 X + \epsilon, \end{align}\] where \(\epsilon \sim \mathcal{N}(0, 1)\) is independent noise. The true APOs are given by \(g(t) = 1 + 2 t + 3 \, \mathbb{E}[X]\).

The APO error decomposition into the order-\(0\) and order-\(1\) moment-balancing error follows from the continuous case analyzed in 3.2, or can be directly derived for this data generating process. \[\begin{align} \hat{g}_{\text{SW-CRM}}(t) - g(t) &= \mathbb{E}[\hat{w}(T, X) Y \mid T=t] - g(t) \\ &= \mathbb{E}[\hat{w}(T, X) \mid T{=}t] + 2 t \mathbb{E}[\hat{w}(T, X) \mid T=t] \\ &\quad + 3 \mathbb{E}[\hat{w}(T, X) X \mid T=t] + \mathbb{E}[\hat{w}(T, X) \epsilon \mid T=t] \\ &\quad - (1 + 2t + 3 \, \mathbb{E}[X]) \\ &= (1+2t) (\mathbb{E}[\hat{w}(T, X) \mid T{=}t] - 1) + 3 (\mathbb{E}[\hat{w}(T, X) X \mid T=t] - \mathbb{E}[X]) \\ &\quad + \mathbb{E}[\hat{w}(T, X) \mid T=t] \, \mathbb{E}[\epsilon] \\ &= (1 + 2 t) \varepsilon_0(t) + 3 \varepsilon_1(t). \end{align}\]

9.2 Synthetic Discrete Dataset↩︎

We generated discrete sequence data with length-\(4\) confounders and length-\(3\) treatments according to: \[\begin{align} X_i^{(k)} &\sim \mathrm{Uniform}\bigl({0,\ldots,V^X_k{-}1}\bigr), \quad k = 0,1,2,3, \\ T_i^{(j)} \mid X_i &\sim \mathrm{Categorical}\bigl(p_j(\cdot \mid X_i)\bigr), \quad j = 0,1,2, \\ Y_i \mid T_i, X_i &\sim \mathrm{Bernoulli}\bigl(\sigma(\mu(T_i, X_i))\bigr), \end{align}\] where “vocabulary” sizes were \(V^X = (5,4,2,3)\), such that \(|\mathcal{X}| = 120\), and \(V^T = (4,2,2)\), such that \(|\mathcal{T}| = 16\). For the model size ablation study, we used \(V^T = (16,2,2)\), such that \(|\mathcal{T}| = 64\), to demonstrate the benefits of larger models. The true propensity factorizes across the treatment sequence, \(p_{T|X}(t \mid x) = \prod_{j=0}^{2} p_j(t^{(j)} \mid x)\), where \[p_j(t^{(j)} \mid x) = \frac{\exp(\ell_j(x,t^{(j)}))}{\sum_{t^{(j)}=0}^{V^T_j-1}\exp(\ell_j(x,t^{(j)}))},\] and the logits are \[\begin{align} \ell_0(x,t^{(0)}) &= t^{(0)}\bigl(x^{(0)} + (x^{(0)})^2 + 2(x^{(0)})^3\bigr), \\ \ell_1(x,t^{(1)}) &= -t^{(1)}\bigl(x^{(3)} + (x^{(3)})^2 + 2(x^{(3)})^3\bigr), \\ \ell_2(x,t^{(2)}) &= t^{(2)}\bigl(x^{(1)} + (x^{(1)})^2 + 2(x^{(1)})^3\bigr). \end{align}\] The true conditional outcome model is \[\begin{align} \mu(t,x) &= 0.3t^{(0)} + 0.2t^{(1)} + 0.15t^{(2)} \\ &\quad + 0.4x^{(0)} + 0.10x^{(1)} + 0.25x^{(2)} + 0.15x^{(3)} \\ &\quad + t^{(0)} x^{(0)} + t^{(0)}(x^{(0)})^2 + t^{(0)}(x^{(0)})^3. \end{align}\] The true APOs are computed by exact enumeration: \[g(t) = \mathbb{E}[Y(t)] = \frac{1}{120}\sum_{x \in \mathcal{X}} \sigma(\mu(t,x)).\]

9.3 Amazon Reviews Text Dataset↩︎

We constructed a semi-synthetic dataset from the Amazon Reviews 2023 corpus [19], Electronics category, with \(n = 10,000\) reviews. The treatment \(T\) is the full text of a product review, along with the rating given. The confounder \(X \in {0, \dots, 7}\) is the product’s total rating count, discretized into 8 bins with thresholds \((10, 50, 200, 1000, 5000, 20000, 100000)\). The outcome was the binary variable corresponding to whether or not a customer purchases the product after reading a given review. Since counterfactual outcomes are unobservable, we used GPT-5.1 to generate the true conditional purchase probabilities for the purpose of evaluation. For each review and each rating count bin, we queried GPT-5.1 with the system prompt:

You are a customer shopping for an electronic product online on Amazon. You get to see the number of ratings the product has and one complete user review, with their rating. Based on both these things, what is the probability (0.0 to 1.0) that you will purchase the product? Reason about how the number of ratings and the particular review affect your probability of purchase:

Larger number of ratings indicate a popular product, making you more likely to purchase it.
The more positive the review sentiment is, the more likely you are to purchase it.
A detailed and informative review makes your purchase probability more aligned with the review sentiment.
A large number of ratings also reinforces the review sentiment and accordingly your purchase probability.
Respond with only a single float between 0.0 and 1.0, nothing else.

and a user prompt of the form:

Number of Ratings: <bin label>
Rating: <rating>/5
Review: <review text>

yielding \(\mu(t, k) = P(Y{=}1 \mid T{=}t, X{=}k)\) for each \((t, k)\) pair. The observed outcome was sampled as \(Y_i \sim \mathrm{Bernoulli}(\mu(T_i, X_i))\). The true APO was computed as: \[g(t) = \mathbb{E}[Y(t)] = \sum_{k=0}^{7} P(X{=}k) \, \mu(t, k),\] where \(P(X{=}k)\) is the empirical bin frequency.

For the lower-dimensional projections of text APOs, we annotated each review with the following three attributes:

1. Rating: the rating out of \(5\), e.g."4/5",

2. Sentiment: "very positive" if rating \(> 4\), else "negative or neutral",

3. Length: "more than 100 words" or "less than 100 words".

10 Experimental Details↩︎

All hyperparameters were chosen on a validation split for each dataset and model.

10.1 Linear Continuous Setting↩︎

We used linear networks for the conditional outcome model and the APO model that maps treatments to APO targets, each trained with the Adam optimizer at learning rate \(10^{-2}\) for 2000 steps. The propensity model was a \(3\)-layer multilayer perceptron with hidden dimension 64 and ReLU activations, mapping \(X \in \mathbb{R}\) to two scalar outputs \(\mu_\theta(X)\) and \(\log \sigma_\theta(X)\), which parameterized the Gaussian propensities \(T \mid X \sim \mathcal{N}(\mu_\theta(X),\, \sigma_\theta(X)^2)\). It was trained via maximum likelihood, with and without balance-error regularization, using Adam with learning rate \(10^{-3}\) for 2000 steps.

10.2 Synthetic Discrete Dataset↩︎

All trained models (propensity, stabilized weights, conditional outcome, and APO) were \(2\)-layer transformers with embedding dimension \(8\), \(2\) attention heads, feedforward dimension \(32\), trained with AdamW and learning rate \(10^{-4}\) on \(n = 10,000\) samples. Propensity and stabilized models were trained for 2000 epochs each and APO and outcome models for 3000 epochs with the cross-entropy loss derived in 3.1 for binary outcomes.

10.3 Amazon Reviews Text Dataset↩︎

For all models, we finetuned a pretrained Gemma-3-270M model with AdamW on their corresponding objectives, using the Transformers library [52] on 2 NVIDIA L40 GPUs. Hyperparameters vary per estimator and are summarized in 3.

References↩︎