Exploring Transformer Backbones for Heterogeneous Treatment Effect Estimation
Authors: YiFan Zhang
Hanlin Zhang
Zachary Chase Lipton
Li Erran Li
Eric Xing
Venue: Accepted by TMLR ()

Abstract: Previous works on Treatment Effect Estimation (TEE) are not in widespread use because they are predominantly theoretical, where strong parametric assumptions are made but untractable for practical application. Recent works use Multilayer Perceptron (MLP) for modeling casual relationships, however, MLPs lag far behind recent advances in ML methodology, which limits their applicability and generalizability. To extend beyond the single domain formulation and towards more realistic learning scenarios, we explore model design spaces beyond MLPs, i.e., transformer backbones, which provide flexibility where attention layers govern interactions among treatments and covariates to exploit structural similarities of potential outcomes for confounding control. Through careful model design, Transformers as Treatment Effect Estimators (TransTEE) is proposed. We show empirically that TransTEE can: (1) serve as a general-purpose treatment effect estimator which significantly outperforms competitive baselines on a variety of challenging TEE problems (e.g., discrete, continuous, structured, or dosage-associated treatments.) and is applicable to both when covariates are tabular and when they consist of structural data (e.g., texts, graphs); (2) yield multiple advantages: compatibility with propensity score modeling, parameter efficiency, robustness to continuous treatment value distribution shifts, explainable in covariate adjustment, and real-world utility in auditing pre-trained language models.

Hanlin Zhang" Published in Transactions on Machine Learning Research (11/2023)
Exploring Transformer Backbones for Heterogeneous Treatment Effect Estimation
Yi-Fan Zhang
yifanshong.cs es es@gmail com
Institute of Automation University of Chinese Academy of Sciences
Horvard University
Zachary C. Lipton
Cornegie Mellon University
Li Erran Lit Amazon
Eric P. Xing Cornegie Mellon University
Reviewed on OpenReview: https: /ttp://ppppe openreview. net/ forum? id= YM247P
//
Abstract
Previous works on Treatment Effect Estimation (TEE) can be limited to carefully examined use cases because they are predominantly theoretical, where strong parametric assumptions are made but untractable for practical application. Recent works use Multilayer Perceptron (MLP) for modeling casual relationships, however, MLP's lag far behind recent advances in ML methodology in handling unstructured, high-dimensional, non-tabular data, which limits their applicability and generalizability. To extend beyond the single domain formulation and towards more realistic learning scenarios, we explore model design spaces beyond MLPs, i.e., transformer backbones, which provide flexibility where attention layers govern interactions among treatments and covariates to exploit structural similarities of potential outcomes for confounding control. Through careful model design, Transformers as Treatment Effect Estimators (TransTEE) is proposed. We show empirically that TrausTEE can: (1) serve as purposs treatment effect estimator which significantly outperforms competitive basclines on variety of challenging TEE problems (c.g., discrete, continuous, structured, or dosage-associated treatments.) and is applicable to both when covariates are tabular and when they consist of structural data (e.g., texts, graphs); (2) yield multiple advantages: compatibility with propensity score modeling, parameter efficiency, robustness to continuous treatment value distribution shifts, explainability in covariate adjustment, and real-world utility in auditing fair predictions of pre-trained language models.
1 Introduction
One of the fundamental tasks in causal inference is to estimate treatment effects given covariates, treatments and outcomes. Treatment effect estimation is a central problem interest in clinical healthcare and social science fImbens & Rubin 2015), as well as econometrics (Wooldridge) 20015). Under certain conditions RRosenbaum & Rubin 1983), the task can be framed as particular type of missing data problem, whose
"Equal Contribution.
Work done outside A Amazon.

---

Published in Transactions on Machine Learning Research (11/2023)
structure is fundamentally different in key ways from supervised learning and entails a more complex set of covariate and treatment representation choices.
Previous works in statistics leverage parametric models (Imbens & Rubin 2015 Wager & Athey 2018
Künzel et al. 2019 Foster & Syrgkanis 2019) to estimate heterogeneous treatment effects. To improve their utilities, feed-forward neural networks have been adapted for modeling causal relationships and estimating
treatment effects (Yoon et al. 2018 Bica et al. 2020b [Scwwb et al. 2020 Nie et al.] [0.]2021]] Curth & van der Schaar 2021b), in part due to their flexibility in modeling nonlinear functions (Hornik et al.) 1989] and high-dimensional imput (Johansson et al,, 2016). Among them, the specialized NN's architecture plays key
role in learning representations for counterfactual inference (Alaa & Schaar 2018 Curth & van der Schaar 2021bb such that treatment variables and covariates are well distinguished (Shalit et al. 2017].
Despite these encouraging results, several key challenges make it difficult to adopt these methods as standard tools for treatment effect cstimation. Most current works based on subnctworks do not sufficiently exploit the structural similarities of potential outcomes for heterogeneous TEE and accounting for them needs complicated regularizations, reparametrization multi-task taali architectures that are problem-specific (Curth & van der Schaar 2021b). Moreover, they heavily rely on their treatment-specifte designs and cannot be easily extended beyond the narrow context in which they are originally. For example, they have poor practicality and generalizability when high-dimensional structural data (e.g., texts and graphs) are given as input (Kaddour et al. 2021.. Besides, those MLP-based models currently lag far behind recent advances in machine learning methodology, which are prone to issues of scale, expressivity and flexibility. Specifically, those side limitations include parameter inefficiency (Table D and brittleness under different scenarios, such as when treatments shift slightly from the training distribution. The above limitations clearly show a pressing need for an effective and practical framework to estimate treatment effects.
In this work, we explore recent advanced models in the deep learning community to boost the model design for TEE tasks. Specifically, the core idea of our approach consists of three parts: as an S-learner, TransTEF embeds all treatments and covariates, which avoids multi-task ar- chitecture and shows improved flexibility and robustness to treatment value distribution shifts; attention mecha- nisms are used for modeling treatment interaction and treatment-covariall interaction. In this way, Trans TEE en- ables adaptive covariate selection (De Luna et al. 2011 VanderWeele Weele, 2019) for inferring causal effects. For exam- ple, one can observe in Figure [ that both pre-treatment covariates and confounders are appropriately adjusted with higher weights, which recovers the "disjunctive cause criterion" (De Luna et al. 2011 that accounts for those two kinds of covariates and is helpful for ensuring the plausibility of the conditional ignorability assumption when complete knowledge of a causal graph is not available. This recipe also gives improved versatility when working with heterogeneous treatments (Figure
Civariate X Sex
Gr) Mafa)
IX
BY
ED
1X
Cansel- Effeet
prees
ap
Age
Scx
TEE
Treatment
Figure 1: A motivating example with a corre- sponding causal graph. Prev denotes previous infection condition and BP denotes blood pres- sure. TranSTEE adjusts an appropriate covariate set (Prev, BP} with attention which is visualized via a heatmap.
Our first contribution shows that transformer, equipped with proper design choices, can be effective and versatile treatment effect estimators under the Rubin-Neyman potential outcomes framework. TransTEF is empirically verified to be (i) flexible framework applicable for a wide range of TEE settings; (ii) compatible and effective with propensity score modeling; (ii) parameter-effiient;; (iv) explainable in covariate adjustment; (v) robust under treatment shifts; (vi) useful for debugging pre-trained language models (T.Ms) to promote favorable social outcomes.
Morcover, comprehensive experiments on six benchmarks with four types of treatments are conducted verify the effectiveness of Trans We show that TransTEF produces covariate adjustment interpretation and significant performance gains given discrete, contimuous or structured treatments on popular benchmarks
I For example, E|Y(1) - -Y0 is often of a much simpler form estimate than either ([((I)X] or ([)())) due to inherent similarities between Y()) and (()).
2

---

Published in Transactions on Machine Learning Research (11/2023)
Table 1: Comparison of existing works and TransTEE in terms of parameter complexity. n is the number of treatments. Br, BD are the number of branches for approximating continuous treatment and dosage. Treatment interaction means explicitly modeling collective effects of multiple treatments. TransTEE is general for all the factors.
MRTHODS
DISCRETR TREATMENT
CONTINUOUS TREATMENT
TREATMENT INTERACTION
DOSAGR
TARNET SHALIT ET AL.
2017
(n)
PERPECT MATCH
SCHWAR FT
2018
((n)
((2")
DRAGUNNET
SHI B/F Al.
(()
DRNET SCHWAB
((n)
O(TBa) ((TBa)
SCIGAN BICA Ri Al.
202UR
(())
VCNET
NiB ET AL
2021
((1)
O(1)
NCORE. PARRBOO
2021
O(n)
O(Br)
O(x)
PLENTENRE CURRr
St
2021
O(n)
OURS
O(1)
O(1)
0(1)
0(1)
including IHDP, News, TCGA. We introduce a new surrogate modeling task to broaden the scope of TEE beyond semi-synihhtlic evaluation and show that TransTEE. is effective in real-world applications like auditing fair predictions of LMs.
2
Related Work
Neural Treatment Effect Estimation. There are many recent works on adapting neural networks to
learn countorfactual representations for treatment effect estimation (Johansson al. 2016 Shalit al.
2017 Louizos et al. 2017: Yoon et. al.] 2018 Bica et. al., 2020b Schwab et al.] 2020 Nie et al.) 2021
Curth & van der Schaar 2021b). To mitigate the imbalance of covariate representations across treatment
groups, various approaches are proposed including optimizing distributional divergence (e.g. IPM including
MMD, Wasserstein distance), entropy balancing (Zeng et al. 2020) (converges to JSD between groups), counterfactual variance (Zhang et al. 2020). However, their domain-specific designs make them limited to different treatments as shown in Table ! methods like VCNet (Nie et al., 2021) use a hand-crafted way to map a real-value treatment to an n-dimensioon vector with a constant mapping function, which is hard converge under shifts of treatments (Table in Appendix); models like TARNet (Shalit et al., 2017] need an accurate estimation of the value interval of treatments. Morcover, previous estimators embed covariates to only one representation space by fully connected layers, tending to lose their connection and interactions (Shalit et al. 20177 Johansson et al. 2020). And it is non-trivial to adapt to the wider settings given existing ad hoc designs on network architectures. For example, the case with n treatments and m associated dosage requires n x m branches for methods like DRNet (Schwab et al., 2020), which put a rigid requirement on the rxtoattatiion capacily and infeasible given observational data.
Transformers and Attention Mechanisms Transformers Vaswani et al. 2017) have demonstrated broad range of language tasks and their variants have been successfully adapted to representation learning over images (Dosovitskiy et al. 2021), programming languages (Chen et. al. 2021 and graphs (Ying et al. 2021) partly due to their Hlexibility and expressiveness. Their wide utility has motivated a line of work for general-purpose neural architectures (Jaegle et al. 2021 2022) that can be trained to perform tasks across various modalities like images, point clouds, audios and videos. But causal inference is fundamentally different from the above models' focus, i.c. supervised learning. And one of our goals is to explore the generalizability of attention-baser models for TEE across domains with high-dimenssonal inputs, an important desideratum in causal representation learning (Schölkoof et al. 202Ib. There are recent attempts use attention mechanisms for TEE Tasks (Guo et al. 2021 Xu et al. 2022). CETransformer (Guo ct al. 2021) uses embeds covariates for different treatments as T-Icarner, They only trivially Iearn covariate embeddings but not treat ment embedding, while the latter is shown more important for TEE tasks. In contrast, TransTEE is an S-learner, which is more well-suited to account for causal heterogeneity (Künzel ct al.] 2019 Curth & van der Schaar 2021bla). ANU (Xu et al., (2022) utilizes attention mechanisms map the original covariate space X into a latent space Z with a single model. We detail the difference in A Appendix A
exemplary performance
3

---

Published in Transactions on Machine Learning Research (11/2023)
Feed Forward
Embedding Layer
Mutti-Hfed Attention
(x)
(xx,
A(s,,t
100
[fs.fa]
(x(
a(t(x)
DmgonNett (discrete)
FlexTENet (discrete)
DRNet (continuous)
VCNet (contimaous)
TransTEF
Figure 2: A schematic comparison of TransTEL and recent works including DragonNet (Shi et al. 2019 FlexTENer (Crrth & van der Schaar 2021b), DRNet (Schwab et al. 2020) and VCNet (Nie et al. TransTEE handles all the scenarios without handcrafting treatment-specifco architectures and any additional parameter overhead.
2021
3 Problem Statement and Assumptions
Treatment Effect Estimation. We consider setting in which we are given N observed samples (x,, ty St b) each containing N pre-treatment covariates {x, E R' The treatment variable t, in this work has various support, e.g,, {0, 1) for binary treatment scttings, IE for continuous treatment settings, and
words for structured treatment settings. For each sample, the potential outcome (n-model) p(x,0 1) or p(x, 1) 8) is the response of the i-th sample to a treatment. 4, where in some cases each treat will be associated with a dosage St, E R. The propensity score (T-model) is the conditional probability of treatment assignment given the observed covariates т(T t|X x). The above two models can be parameterized as pi and ne, respectively. The task is estimate the Average Dose Response Function (ADRF): u(x,t) t) = E x, do(T t)] (Shoichet 2006), which includes special cases in discrete treatment scenarios that can also be estimated as the average treatment effect (ATE): ATE Biu(x, 1) u(x, 0)] and its individual version ITE.
What makes the above problem more challenging than supervised learning is that we never see the missing counterfactuals and ground truth causal effects in observational data. Therefore, we first introduce the required fundamentally important assumptions that give the strongly ignorable condition such that statistical estimands can be interpreted causally.
Assumption 3.1. (Ignorabilityy watpraaillty///uoooouuunee Unconfoundedness) implies no hidden confounders such that Y(T = I In the binary treatment case, Y(0), (1))
Assumption 3.2. (Positivity/Overrap)) The treatment assignment. is non-deterministic such that, ie. 0< <(t|z) < 1, V2 E X,t T
Assumption 3.1 ensures the causal effect is identifiable, implying that treatment is assigned independent of the potential outcome and randomly for every subject regardless of its covariates, which allows estimating ADRF using p(t) := [[ = t)] = =[E[[ x T = t]] (Rubin 1978). One naive estimator of μ(x, t) = E| x, T t] is the sample avcrage v²t) En= ()x;, Assumption 3.2 states that there is a chance of seeing units in every treated group.
4 Trans Transformers as Treatment Effect Estimators
The systematic similarity of potential outcomes of different treatment groups is important for TEE (Curth & van der Schaar (0211).. Note that x is often high-dimensional while t is not, which means naively leeding (x, () to MLPs is not. favorable since the impacts of meett lend to be lost. As a result, various architectures and regularizations have been proposed to enforce structural similarity and differences among treatment groups. However, they are limited to specific use cases as shown in Section P and Figure [ To remedy it, we use three simple yet effective design choices based on attention mechanisms. The resulting scalable framework TransTEF can tackle the problems of most existing treatment effect estimators (e.g.,
4

---

Published in Transactions on Machine Learning Research (11/2023)
treatments, treatments interaction, and treatments with dosage) without ad-hoc architectural designs, e.g,, multiple branches.
Preliminary. The main module in TranSTEE is the attention layer Vaswani et al. 2017 given query, key, and value matrices Q E Raxdy, K E Rdxd, V E Rdxd, attention mechanism computes the outputs as H(Q, K, V)= softmax( QK- vd, )V. In practice, multi-head attention is preferable to jointly attend to the information from different representation subspaces.
HM,Q, K, V) = Concat(headi, ..., head, WWo
where head, H(QW? KW." VW.),
where woO E Rdxdi wY E Rdddd W. E Rddd, and wo Ridde are learnable matrices.
4.1 Covariate and Treatment Embedding Layers
Treatment Embedding Layer. As illustrated in Figure 2 and Table. [ as treatments are often of much lower dimension compared to covariates, to avoid missing the impacts of treatments, previous works (e.g., DragonNet (Shi el al. 2019), FlexTENee (Curth & van der Schaar 2021b), DRNet (Schwab et al. 2020
assign covariates from different treatment groups to different branches, which highly parameter inefficient. Besides, We analyze in Proposition 3 (Appendix C that, for continuous treatments /dosagss, the performance is affected by both number of branches and the value interval of treatment. However, almost all previous works continuous treatment /dosage assume the treatment or dosage is in fixed value interval e.g., [0, 1] and Figure ! shows that prevalent works foil when tested under shifts of treatments. These two observations motivate us to use two learnable linear layers to project scalar treatments and dosages to d-dimension vectors separately: M, = Linear((), M,, = Linear(s), where Mt E Rd. M, E Rd exists just when each treatment has needed. When multiple (n) treatments act
a dosage parameter, otherwise, only treatment embedding simultaneously, the projected matrix will be M, E Rdxa, .M, E Rdxx and when facing structural treatments (languages, graphs), the treatment embedding will be projected by language models and graph neural networks respectively. By using the treatment embeddings, TransTEE is shown to be (i) robust under treatment shifts, and (ii) perameter-efficient.
Covariates Embedding Layer. Different from previous works that embed all covariates by one fully connected layer, where the differences between covariates tend to be lost. Therefore, they are hard to study the function of an individual covariate. TransTEE learns different embeddings for each covariate, namely M,, = Linear(x), and M2 E Rdxp, where p is the number of covariate. Covariates embedding enables us to study the effect of individual covariate on the outcome.
4.2 Covariate and Treatment Self-Attention and Attss-Attention
For covariates, prevalent methods represent covariates a whole fea- ture using MLPs, where pair-wise covariate interactions are lost when adjusting covariates. Therefore, we cannot study the effect of each covariate on the estimated result. In contrast, Trans processes each covariate embedding independently and model their interactions by self-attention layers. Finally, one of the fundamental challenges of causal meta-learners is to model treatment-covariate interactions. TransTEE realizes this by a cross-attention module, which enables us to model the scenario where mudliple trealments are applied and allain strong practical utility, e.g., multiple prescriptions in healthcare different financial measures in economics. This is an effective remedy for existing methods which are limited to settings where various treatments are not used simultaneously. Refer to Appendix A.1 for a formal definition and explanation.
Jwliy Lizer
Crus Amertce xL
Ni
M2
Kinbeddig
1 1 1
x.
Figure 3: An Illustaatvve Example about the workflow of TransTEB.
An Illutrative Example To better understand the workflow with the above designs, we present simple illustration here. Consider use case in medicine effect estimation,
5

---

Published in Transactions on Machine Learning Research (11/2023)
where x contains p patient information, e.g., Age, Sex, Blood Pressure (BP), and Precious infection condition (Prev) with a corresponding causal graph (Figure D. n medicines (treatments) are applied simultaneously and each medicine has a corresponding dosage. As shown in Figure 3 each covariate, treatment, and dosage will first be embedded to d-dimension representation by a specific learnable embedding layer. Each treatment embedding will be concatenated with its dosage embedding and the concatenated feature will be projected by a linear layer to produce d dimensional vectors. Self-attention modules optimizes these embeddings by aggregating contextual information. Specifically, attribute Prev is more related to age than sex, hence the attention weight of Prev feature to age feature is larger and the update of Prev feature will be more dependent on the age feature. Similarly, the interaction of multi-medicines is also attained by the self-attention module. The last Cross-attention module enables treatment-covariate interactions, which is shown in Figure 2 that, each medicine will assign a higher weight to relevant covariates especially confounders (BP) than irrelevant. ones. Finally, we pool the resulted embedding and use one linear layer to predict the outcome.
5 Propensity Score Modeling
We first discuss the fundamental differences and common goals between our algorithm and traditional ones:
as general approach to causal inference, TransTEE can be directly harnessod with traditional methods
that estimate propensity scores by including hand-crafted features of covariates (Imbens & Rubin 2015) to
reduce biases through covariate adjustment (Austin 2011), matching (Rubin & Thomas 1996 Abadie &
Imbens 2016), stratification (Frangakis & Rubin 2002), reweighting (Hirano et al. 2003); ) g-computation
[Imeess& & Rubin 2015), sub-classification (Rosenbaum & Rubin 1984], covariate adjustment (Austin 2011
largeted regularization (Van Der Laan & Rubin 2006) conditional density estimation (Nie et al. 2021
that create quasi- raaasi-anddmized experiments (D (''otoo [998). It is because the general framework provides an advantage to using an off-the-shelf propensity score regularizer for balancing covariate representations. Similar to the goal of traditional methods like inverse probability weighting and propensily score matching (Austin 2011), which seeks to weigh a single observation to mimic the randomization effects with respect to the covariate from different treatment groups of interest.
Unlike previous works that use hand-crafted features or directly model the conditional density via maximum likelihood training, which is prone to high variance when handling high-dimensional, structured treatments (Singh et al. 2019) and can be problematic when we want to estimate a plausible propensity score from the generative model (Mohaeed & Lakshminarayanan 2016) (see the degraded performance of MLE in Table 2 Trans learns propensity score network Ta (4|x via minimax bilevel optimization. The motivations for adversarial training between no (x, t) and a (f|x) are three-fold: (i) it enforces the independence between treatment and covariate representations as shown in Proposition [ which serves as algorithmic randomizat tin in place of costly randomized controlled trials (Rubin (207)) for overcoming selection bias (D'Agostino) 1998 Imbens& & Rubin 2015); (ii) it explicitly models the propensity (t(x))) to refine treatment representations and promote covariate adjustment (Kaddour et al., 2021
Based on the above discussion, when treatments are discrete, one might consider directly applying heuristic mothods such as adversarial domain adaptation (see (Ganin et al. 2016 Zhao ct al., 2018) for algorithmic development guidelines). We note the heuristic nature of domain- adversarial methods (see (Wu et al.) 2019
for clear failure cases) and a debunking of the common claim that (Ben-David et al.] 2010) guarantees the robustness of such methods. Here, we focus on continuous TEE, more general and challenging scenario, where we want to estimate ADRF, and propose two variants of Lss as an adversary for the outcome regression objective Lo in Eq. 4 accordingly. To be clear, we denote [x () = wy o ( (x), (x), and no (+)x) == Wt T (x), which decompose the predictions into featurizers $, T + Zr, $. - Zx and predictors wy Zx x ZT + 3), w; Zx + For example, (x) and &e(() can be the linear embedding layer and attention modules in our implementation. The propensity is computed on (x),,, an intermediate feature representation of x. Similarly, (x t) is computed from (t) and (x) (x, y, 2 denote either a raw or representation space). The process is shown in Eq. below:
min max Lo(x, v, t) - (x,
(1)
6

---

Published in Transactions on Machine Learning Research (11/2023)
We refer to the above minimax game for algorithmic randomization in place of costly randomized controlled trials. Such algorithmic randomization based on neural representations using propensity score creates subgroups of different treated units as if they had been randomly assigned to different treatments such that conditional independence T (T is enforced across strata and contimationn, which approximates a random block experiment to the observed covariates (Imbens & Rubin 2015).
Below, we introduce two variants of (x,
Treatment Regularization (TR) is standard MSE over the treatment space given the predicted treatment & and the ground truth 1
a
R 1) = Σ (4 - (4:x)((((x;)))
(2)
TR explicitly matching the mcan of the propensity score to that of the treatment. In an ideal case, <(t|x) should be uniformly distributed given different x. However, the above treatment regularization procedure only provides matching for the mean of the propensity score, which can be prone to bad equilibriums and treatment misalignment (Wang et al. 2020). Thus, we introduce the distribution of t and model the uncertainty rather than predicting scalar t:
Probabilistic Treatment Regularization (PTR) is a probabilistic version of TR that models the mean p (with a slight abuse of notation) and the variance o² of the estimated treatment f
12
(1+ - no To(x)xx))² (x2(0) + log n (0))
CPTR_ &
Σ
(3)
2
The PTR matches the whole distribution, i.c. both the mean and variance, of the propensity score to that of the trealment, which can be preferable in certain cases.
Equilibrium of the Minimax Game We analyze that TR and PTR can align the first and second moment of continuous treatments at equilibrium respectively, and thus promote the independence between treatment and covariale x. For the ease of our analysis below, we assume the predictors We, wx are fixed.
(The optimum of propensity score model) In the equilibrium of the game, assuming the out- come prediction model is fixed, then the optimum of TR is achieved when (() [[), Y __(x)) via matching the mean of propensity score ((b((t) (x)) and the marginal distribution p( (æ)) and the oplimum discriminator of PTR is achieved via matching both the mean and variance such that ([),(t)), V (x).. See Appendix B for the proof.
Proposition 1.
(
IE [E), (
6 Experimental Results
We elaborate on basic experimental settings, results, analysis, and empirical studies in this section. See Appendix D for full details of all experimental settings and detailed definitions of metrics. See Appendix for many more results and remarks.
6.1 Experimental Settings
Datasets. Since the true counterfactual outcome (or ADRF) are rarely available for real-world data, we use syuthetic or semi-synthetic data for empirical evaluation. for continuous treatments, we use one synthetic dataset and two semi-synthetic datasets: the IHDP and News datasets. For treatment with continuous dosages, we obtain covariates from real dataset TCGA (Chang et al., 2013) and generate treatments, where each treatment is accompanied by a dosage. The resulting datasct is named TCGA (D). Following (Kaddour et al. 2021), datasets for structured treatments include Small- World (SW), which contains 1, 1,000 uniformly sampled covariates and 200 randomly generated Watts SStrssattrrzzz small-world graphs (Watts & SSrogatzz 1998) as treatments, and TCGA (S), which uses 9, 659 gene expression of cancer patients (Chang al. [0033) for covariates and 10, 000 molecules from the QM9 dataset (Ramakrishnan et al., 2014 as treatments. For the study language models, we the Fnriched Equily Fvaluation Corpus (EFEC) Feder et al. 2021
7

---

Published in Transactions on Machine Learning Research (11/2023)
Table 2: Experimental results comparing NN-based methods on the IHDP datasets, where
means the model is not suitable for continuous treatments. We report the results based on 100 repeats, and numbers after + are the estimated standard deviation of the average value. For the vanilla setting with binary treatment, we report the mean absolute difference between the estimated and true ATE. For Extrapolation (h 1 2), models are trained with & E [0.1, 2.0] and tested in t E [0, For Extrapolation (h = 5), models are trained with t E [0.25, 5.0] and tested in E [0, 5].
METHOODS
VANILLA (BISARY)
VANILLA (h 1)
EXTRAPOLATION (h
VANILLA (h 5)
EXTRAPOLATION (A 5)
TARNET
0.3870 = 0.61112
2.0152 + 1.07449
12.967 + 1.78108 11,071 + 0.99384
5,6752 + 0,53161
31,523 = 1.5013 31,524 + 1,50264
DRNET
0.3543 + 0 1.60622
2.1549 1 1.04483
3.2779 + 0.42797
FLEXTENET
0.2700 + 0.10000
VONET
0.2095 + 0. 18236
0,7800 + 0.61483
XAN 0.2745 0.14976
NAN
XAN
TRANSTEE
0.0983 0.15384
1010289
0.1621 0.14443
0:2066 0:23258 0.1417 + 0.07123
TRANSTEE TEE--EEE
0.1721- = 0. J0061
0,0877 + 0,03352
0,2685 + 0,17552
0,2079 + 0.17637
TRANSTEE-TR
0.1913 + 0. 29953
0,0781
+ 0.03243
0,2393 + 0.08154
0.1148 0.08224
0.0947 + 0.0824
TRANSTER. PTR
0.2193 + 0.34667
0.0762
0.07915
0.2352 0.17095
0.1363 +: 0.08036
0.1363 + 0.08035
Table 3: Performance of individualized treatment-dose response estimation the TCGA (D)
dataset with different numbers of treatments. We report AMSE and standard deviation over 30 repeats. The selection bias on treatment and dosage are both set to be 2.0.
Miernoos
#TREATMENT=1
# #TREATMENT=2
#3
IN-SAMPLE
OUT-SAMPLR
PN-SAMPLE
OUT-SANPUR
IN-A
OUT-SAMPLR
SCIGAN
5.6960 =: 0. 0.0000
5.0546 1 0.0000
2,0924 +: 0.0000
2.3067 1 0.0080
4.3183 1: 0.0000
4,6231 = 0.0000
(D)
0.7888
0.0009
0.7908
0.0606
1.4207 +: 0.0784
1.4206
0.0777
3.1982 4: 0.5847
3.1920
0.5746
DRNET(D)
0,8034
= 0.0469
0.8052 1 0.0466
1.3739 1: 0.0858
1.3738
1 0.0853
2,8632 1: 0.4227
2.8558 = 0.4143
(D
0.1506
0. 03033
0.1579
0.0301
0.2919 + 0.0743
0.2918
0.0737
0.6459 # 0.1387
0.6493
0.1397
TRANSTEE
0.0585
0.0358
0,0350 +
0.0556
0.0129
0,2803
0.0639
THANSTER TR
0.0495
0.0176
0.0500
0.0180
0.0863
+ 0.0268
0.0671
0.0268
0.2618
+: 0.0737
0.2577
0.0726
TRANSTEE PTR
0,0343 0.0096
0.0355
0,0094
0,0679 4: 0.0252
0.0686
0.0252
0,2645
0,0702
0.2597
0.0075
Baselines. Baselines for continuous and binary treatments include TARnet (Shalit et al.
2017 Drag-
onnet (Shi et al., 2019). DRNet (Schwab et al.
FlexTENet (Curth & van der Schaar 2021b and 2020b is chosen as the baseline for continuous dosages.
2020).
VCNet (Nie et al.
2021 SCIGAN [Bica et al.
Besides, we revise DRNet (Schwab et al. 2020). TARNet (Shalit et al. 2017), and VCNet (Nie et al. 2021
to DRNet (D), TARNet (D), VCNet (D), respectively, which enable multiple treatments and dosages. Spocifi- cally, DRNel (D) has T main flows, each corresponding to a treatment, and is divided into Bn branches for
continuous dosage. Baselines for structured treatments include Zero (Kaddour et al. 2021]. GNN (Kaddour
ct al. 2021), GraphITE (Harada & Kashima 2021p, and SIN (Kaddour et al. 2021). To compare the
performance of different frameworks fairly, all of the models regress on the outcome with empirical samples
without any regularization. For ML.E training of the propensity score model, the objective is the negative log-likelln loglliie:hood:: Lo- = - E log T (t.(x.)..
Evaluation Metric. For continuous and binary treatments, we use the average mean squared error the test set. For structured treatments, following (Kaddour et al., 2021 we rank all treatments by their propensity m(t[x) in descending order. Top K treatments are selected and the treatment effect of each treatment pair is evaluated by unweighted/ /eeigghueddweegee weighted expected Precision in Estimation of Heterogeneous Effect (PEHE) (Kaddour et al.) 2021), where the WPEHE@K accounts for the fact that treatment pairs that are less likely to have higher estimation errors should be given less importance. For multiple treatments and dosages, AMSE is calculated over all dosage and treatment pairs, resulting in AMSE
6.2 Case Study and Numerical Results
Case study on treatment distribution shifts We start by conducting a case study on treatment
distribution shifts (Figure and exploring an extrapolation setting in which the treatment may subsequently be administered at values never seen before during training. Surprisingly, we find that while standard
results rely on constraining the values of treatments Nie et al. (2021) and dosages Schwab et al. (2020) to a specific range, our methods perform surprisingly well when extrapolating beyond these ranges as assessed on
8

---

Published in Transactions on Machine Learning Research (11/2023)
0.5
0.5
0.0
Resporss
Resporse
Resporte
00.5
-0.5
Truth Tamntt
Truth
Truth
-1.0
2
Tamet
-1.0
Tamet
Dieet
Dmet
Dmeet
Venet
Venet
-1.5
Venet
TransTEE
Trans TEE
TrunsTEE
0.0
02
0.4
0.6
0.8
1.0
0.4
0.8
1.2
1.6
2.0
2
Treatment
Treatment
Treatment
(a) = training and testing.
(b) h 1.75 in training and à in testing (extrapolation).
(c) h training and testing.
Figure 4: Estimated ADRF the synthotic dataset, where treatments sampled from interval h), where 1 = 0.
several benchmarks. By comparison, other methods appear comparatively brittle in these same settings. See Appendix
C for detailed discussion.
Case study of propensity modeling. TranSTEE conceptually simple and effective. However, when the sample size is small, it becomes important to account for selection bias (Alaa & Schaar 2018 However, most
existing regularizations can only be used when the treatments are discrete (Bica et al. 2020a Kallus 2020 Du et al. 2021). Thus we propose two regularization variants for continuous treatment/ (oltmmenn//oagge,, which are termed Treatment Regularization (TR, C2 CTR( 1) = E (1; To (f.]xx) ?) and its probabilistic version Probabilistic Treatment Regularization (PTR, C
[o + 1 log ग (021x)))) respectively. The overall (x
(x, Specifically, propensity score model T (4(x) parameterized by an MLP is learned by minimizing (x, 1), and then the outcome estimators μe (x, is trained by ming Lo(x, t) (x(x,t), To overcome selection biases overrrepresnntation space, the bilevel optimization enforces effective treatment. effect estimation while modeling the discriminative propensity features to partial out parts of covariates that cause the treatment but not the outcome and dispose of nuisance variations of covariates (Kaddour et al. 2021). Such recipe account for selection bias where [[{]x) * p(t) and leave spurious correlations out, which can also be more robust under model misspecification especially in the scttings that require extrapolation on treatment (See Table and Appendix B for concrete formalisms and discussions.).
model trained in a adversarial pattern, namely ming max Lo y.
As in Table 2 Table 3 and Table 12 with the addition of adversarial training as well as TR and PTR, TransTEE's estimation error with continuous treatments can be further roduced. Overall, TR is better in the continuous case with smaller treatment distribution shifts, while PTR. is preferable when shifts are greater. Both TR and PTR cannot: bring performance gains over discrete cases. The superiority of TR and PTR in combination with TranSTEE over comprehensive existing works, especially in semi-synthetie benchmarks like IHDP that may systematically favor some types of algorithms over others (Curth et al.) 2021). also calls for more understanding of NNs' inductive biases treatment effect estimation problems of interest. Moreover, covariate selection visualization TR and PTR (Figure 6(a) Table 5 and Appendix F. supports the idea that modeling the propensity score effectively promotes covariate adjustment and partials out the effects from the covariates the treatment features. We also compare the training dynamic of different regularizations in Appendix where TR and PTR are further shown able to improve the convergence of TransTER.
Continuous treatments. To evaluate the elficiency with which TransTEE estimates the average dose- response curve (ADRF), we compare against other recent NN-based methods (Tables 2). Comparing results in each column, we observe performance boosts for TransTEE. Further, TransTEF attains much smaller error than baselines in cases where the treatment interval is not restricted [0, 1] (c.g., E [o, 5]) and when the training and test treatment intervals are different (extrapolation). Interestingly, even vanilla
9

---

Published in Transactions on Machine Learning Research (11/2023)
Trll Dood Teeete
Treuihy
Teartey
Teafte
Kegpoote
Reepoose
120
175
Teestment
(a) h- during training and (b) 1.9,/ = during tran- (c) = 5 during training and (d) h 4 during training
testing.
ing and h = 2,7 during testing.
and h 5, = 0 during testing (extrapolation).
testing (extrapolation).
Figure 5: Estimated ADRF on the test set from a typical run of TarNet (Shalit al. 2017p. DRNet (Schwab) et al. p. VCNet (Ne et al., 2021) and ours on News dataset. All of these methods are well-optimized. Suppose E [1. hl. (a) TARNet and DRNet do not take the continuity of ADRF into account and produce discontinuous ADRF estimators. VCNet produces continuous ADRF estimators through hand-crafted mapping matrix. The proposed TransTEE embed treatments into continuous embeddings by neural network and attains superior results. (b,d) When training with 0 $ < 1.9 and 0 I t S 4.0. TARNet and DRNet cannot extrapolate to distributions with 0 < 2.0 and 0 S 5.0. (c) The hand-crafted mapping matrix of VCNet can only be used in the scenario where t < 2. Otherwise, VCNet cannot converge and incur an infinite loss. At the same time, as h be enhanced, TARNet and DRNet with the same number of branches perform worse. TransTEE needs not know h in advance and extrapolates well.
Table 4: Effect of Gender (top) and Race (bottom) on POMS classification with the EEEC dataset, where ATEgT the ground truth ATE based on 3 repeats with confidence intervals [CI] constructed using standard deviations.
Correlation/ Representation Based Baselines
Treatment Effect Estimators
TC
ATEgT
TReATE
CONEXP
INLP
TarNet
DRNet
VCNet
Trans' T'E
Gender
0.086
0.125
0.02
0.313
0.0067
0.0088
0.0085
0.013
[CI]
[0.082,0.09]
00.110,0114]
[0.000055]
[0.30103321]
[0.0019, 0.0076]
[0.0081.0.009]
[0.0036, 0.0111]
10008, 0.0168]
Race
0.014
0.046
0.08
0.591
0.005
0.006
0.003
0.0174
[C]]
[0.012,0.01]]
[0.038.0.054]
[0.02.0.01.]
[0.578.0.605]
[0.0021, 0.0069]
[0.0047, 0.0081]
[0.0025, 0.0037]
13, 0.0238)
TEansTEE produces competitive performance compared with that of (([xx) trained additionally using MLE, demonstrating the ability of TransTEF. effectively model treatments and covariates. The estimated ADRF curves on the IHDP and News datasets are shown in Figure [2] and Figure TARNet and DRNet produce discontimoous ADRF estimators and VCNet only performs well when E [0, 1]. However, TransTEE attains lower estimation error and preserves the continuity of ADRF different treatment intervals.
Continuous dosage. In Table 3 we compare TranSTEE against baselines on the TCGA (D) dataset with default treatment selection bias 2.0 and dosage selection bias 2.0. As the number of treatments increascs, TransTEE and its variants (with regularization term) consistently outperform the baselines by a large margin on both training and test data. TransTEE's effectiveness is also shown in Appendix Figure g where the estimated ADRF curve of each treatment considering continuous dosages plotted. Compared to baselines, Trans attains better results over all treatments. Stronger selection bias in the observed data makes estimation more difficult because it becomes less likely to see certain treatments or particular covariates. Considering different dosages and treatment selection bias, Appendix Figure shows that as biases increase, Trans consistently performs the best.
Structured treatments. We compared the performance of TransTEE baselines on the training and test set of both SW and TCGA datascts with varying degrees of treatment selection bias. The numerical results are shown in Appendix Table 13 The performance gain between GNN and Zero indicates that taking
10

---

Published in Transactions on Machine Learning Research (11/2023)
Medel TIANTTE VIDSA URKNE TARNa
Tursiees
0000
630
0.00
10 19 17 1-
Covarates
NEWS
TOGA/21
TOGNS)
(b)
(c)
Figure 6: (a) The learned weights of the cross-attention module on IHDP dataset. TransTEF adjusts confounders Scon (1, 3,5,, 5, properly with higher weights during the cross attention process.
(b) AMSE (c) Number of parameters for
attained by modcls on IHDP with different mumbers of noisy covariates. different models on four different datasets, where the log on the y-axis is base
into account graph information significantly improves estimation. The results suggest that, overall, the performance of TransTEE is the best due to the strong modeling capability and advanced model structure for processing high-dimensional treatments. STN is the best model among these baselines. However, when the bias is equal to 0.1, SIN fails to attain estimation results better than the Zero baseline. To evaluate each model's robustness to varying levels of selection bias, performance curve with E [0, 40) for the SW dataset and & E [0, 0.5) for the TCGA dataset are shown in Figure 14 and Figure 15 in Appendix. Considering both metrics, TransTEE outperforms baselines by a large margin across the entire range of evaluated treatment. selection biases.
6.3 Analysis
Analysis of covariate adjustment of cross-attention module. TransTEE embeds cach covariate independently and then make treatments select proper covariates for prediction by cross-atlention. The resulting interpretability of the covariate adjustment process using attention weights is one clear advantage over existing works. Thus we visualize the covariate selection results (cross-attention weights) in Fig- ure ((a) As elaborated in Appendix D.3 the IHDP dataset has 25 covariates, which is divided into 3 groups: Scon 21,2,3, 3, 5, 6), Sais,1 [4.7 15), and Sais,2 (16 25). Scon influences both T and Y, Sais,1 influences only Y, and Sais,1 influences only T. Covariates in Sais,1 named noisy variates since they have correlation with the treatment. Their causal relationships are illustrated in Figure 7 Interestingly, confounders Scon are assigned higher weights while noisy covariates (those in- fluence the outcome but are irrelevant to the treatment) lower Sdis,1, which matches the principles in Vander Weele 2019) and corroborate the ability of TransTEE to estimate treatment effects in complex datasets by controlling both pre-treatment variables and confounders properly. Moreover, Figure 6(b) shows that TransTEE consistently outperforms baselines across different numbers of noisy covariates.
We further conduct 10 repetitions for Trans TEE and its TR and PTR counterparts as reported in Table (Appendix Figure II visualizes their cross-attention weights). Denote Weon, W1, W2 as the summation of weights assigned to Scon: Sase. Saie.2 re- spectively. We can see that, incorporated with both TR and PTR regularization, TransTEE assigns more weights to confounding covariates (Scon) and fewer weights on noisy covariates, which further verifies the compatibility of TransTTE with propensity score modeling since both TR and PTR. improve confounding control. Moreover, TR. is better than PTR since it also reduces by a larger margin. This observation gives a suggestion that we should systematically probe TR and PTR besides compar- ing their numerical performance, especially in settings where the unconfoundedness Ding et al., 2017) and controlling instrumental variables will incur biases in TEE.
Sauz
Saisa
T
Figure 7: The causal graph of IHDP dataset.
assumption is violated
Amount of model parameters comparison. The experiment to corroborate the conceptual comparison in Table 1 We find that the proposed TransTEF has consistently fewer parameters than baselines all
11

---

Published in Transactions on Machine Learning Research (11/2023)
the settings as shown in Figure 6(c) Besides, increasing the number of treatments allows more accurate approximation for continuous treatments/dossggen most of these baselines need to increase branches which incurs parameter redundancy. However, TransTEE is much more efficient.
6.4 Empirical Study on Pre-trained Language Models
To evaluate the real-world utility of TrEnSTEE, in this subsection, we demon- strate an initial attempt for auditing and debugging large pre-trained lan- guage models, an important use case in NLP that is beyond semi-synthetic settings. The goal of our stylized analyses is to stress test the utility of TransTEF on high-dimeessional settings where assumptions such ignora- bility might not necessarily hold, which can be under-explored in the causal inference literature. Specifically, we use TransTEE to estimate the treatment effects for detecting the effects of domain-specific factors of variation (such as the change of subject's attributes in a sentence) on the predictions of pre- trained language models. We experiment with BERT (Kenton & Toutanova (e.g., racial and gender-related nouns) over natural language on the
Table 5: Attention weights for Scon Saie, and S re- spectively.
Weon
W1
W/2
TransTEE +TR +PTR
0.27
0.37
0.36
0.59
0.20
0.21
0.32
0.33
0.35
2019
(real
EEEC dataset. We use both the correlation /rrpreseentaaioonaaseee baselines introduced in (Foder
ct al.
2021 and implement treatment effect estimators (c.g., TARnet, DRNet, VCNet, and the proposed
Trans
Interestingly, results in Table show that TransTER effectively estimates the treatment effects of domain- specific variation perturbations even without substantive downstream fine-tuning on specialized datasets. TransTEF. outperforms baselines adapted from MT.P. Moreover, we showcase the top-k samples with the
maximal/mining ITE and analysis in Appendix E.3 The results show that TransTEE has the potential to
provide estimators for practical use cases in predicting model predictions (llyas et al. 2022). For example, those identified samples can provide actionable insights like function as contrast sets for analyzing and
understanding LMs (Gardner et al. 2020; Abraham et al. 2022) and TransTEE can estimate ATE to enforce
invariant or fairness constraints for LMs (Viichh et al. 2021 in a lightweight and efficient manner, which we leave for future work.
7 Concluding Remarks
In this work, we show attention mechanisms can be effective and versatile design choices for TEE tasks. Extensive experiments well verify the effectiveness and utility of the proposed TransTEE, which also imply that more challenging and unified evaluation alternatives of TEE are nceded. Morcover, we hope that our findings can lay the groundwork for future work developing advanced machine learning techniques like pre-training in large-scale observational data in estimating treatment effects, where TransTEE can serve effective backbone. Similar to almost all the causal inference methods on obscrvational data, potential limitation of TransTEE is the reliance on the ignorability assumption. Therefore, one important future direction is extending TransTEF to settings with more complex causal graphs and generate identifiable causal functionals tractable for optimization ІJung et al. 2020) supported by identification theory. Since adjusting covariates without accounting for the causal graph might yield inaccurate or biased estimates of the causal effect (Pearl 2009), how to integrate TransTER with domain knowledge (Imbens & Rubin 2015) for alleviating its potential negative societal impacts in consequential decision making will also be important.
Acknowledgment
The authors would like to thank Edward H. Kennedy and Hongyi Wang for their helpful discussions. This roscarch has been generously funded by NSF IIS1955532, NSF HM04762010002, NIH R014M140467.
References
Alberto Abadie and Guido W Imbens. Matching on the estimated propensity score. Econometrion, 2016.
12

---

Published in Transactions on Machine Learning Research (11/2023)
Eldar David Abraham, Karel D'Oosterlinck, Amir Feder, Yair Ori Gat, Atticus Geiger, Christopher Potts, Roi Reichart, and Zhengxuan Wu. Cebab: Estimating the causal effects of real-world concepts nlp model behavior. arXiv preprint arXiv:2205.141440 2022.
Ahmed Alaa and Mihaela Schaar. Limits of estimating heterogeneous treatment effects: Guidelines for practical algorithm design. In ICML, 2018.
Ahmed M Alaa and Mihacla van der Schaar. Bayesian inference of individualized treatment effects using multi-task gaussian processes. In NeurIPS, 2017.
Ahmed M Alaa, Michael Weisz, and Mihaela Van Der Schaar. Deep counterfactual networks with propensity- dropout. arXiu, 2017.
Peler C Austin. An introduction to propensity score methods for reducing the effects of confounding observational studies. Multivariate behavioral research, 2011.
Shai Ben- nen-David, John Blitzer, Koby Crammer, Alex Kulesza, Fernando Pereira, and Jennifer Wortman Vaughan. A theory of learning from different domains. Machine learning, 2010.
Toana Bica, Ahmed M Alaa, James Jordon, and Mihaela der Schaar. Estimating counterfactual treatment. outcomes over time through adversarially balanced representations. arXiv, 2020a.
loana Bica, James Jordon, and Mihaela van der Schaar. Estimating the effects of continuous-valued interventions using generative adversarial networks. 2020b.
Kyle Chang, Chad J Creighton, Caleb Davis, Lawrence Donehower, Jennifer Drummond, David Wheeler, Adrian Ally, Miruna Balasundaram, Inanc Birol, Yaron SN Butterfield, et al. The cancer genome atlas pan-cancer analysis project. Nat Genet, 2013.
Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde de Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, et al. Evaluating large language models trained on code. Xiv, 2021.
Jonathan Crabbé, Alicia Curth, Ioana Bica, and Mihaela van der Schaar. Benchmarking heterogeneous treatment effect models through the lens of interpretability. arXiv preprint orXiw:2206.08338 2022.
Alicia Curth and Mihaela van der Schaar. Nonparametric estimation of heterogeneous treatment effects: From theory to learning hmgo In AISTATS, 2021a.
Alicia Curth and Mihaela van der Schaar. On inductive biases for heterogeneous treatment effect estimation. In Newr/PS, 2021b.
Alicia Curth, David Svensson, Jim Weatherall, and Mihaela van der Schaar. Really doing great at estimating CATE? critical look at ML benchmarking practices treatment effect estimation. In NeurIPS, 2021.
Ralph B D' D'Agostino. Propensity score methods for bias reduction in the comparison of a treatment to a non-randomized control group. Statistics in medicine, 1998.
Xavier De Luna, Ingeborg Waernbaum, and Thomas S Richardson. Covariate selection for the nonparametric estimation of an average treatment effect. Biometrika, 2011.
Peng Ding, TJ VanderWeele, and James M Robins. Instrumental variables as bias amplifiers with general outcome and confounding. Biometrika, 2017.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Marthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. 2021.
Xin Du, Lei Sun, Wouter Duivesteijn, Alexander Nikolaev, and Mykola Pechenizkiy. Adversarial balancing- based representation learning for causal effect inference with observational data. DMKD, 2021.
13

---

Published in Transactions on Machine Learning Research (11/2023)
Amir Feder, Nadav Oved, Uri Shalit, and Roi Reichart. Causalm: Causal model explanation through
counterfactual language models. Compulational linguistics, 2021.
Dylan J Foster and Vasilis Syrgkanis. Orthogonal statistical learning. aiviv preprint 1rx::11911.000 2019.
Constantine E Frangakis and Donald B Rubin. Principal stratification in causal inference. Biometrics, 2002.
Michele Jonsson Funk, Daniel Westreich, Chris Wiesen, Til Stürmer, M Alan Brookhart, and Marie Davidian. Doubly robust estimation of causal effects. American journal of epidemiology, 2011.
Yaroslav Ganin, Evgeniya Ustinova, Hana Ajakan, Pascal Germain, Hugo Larochelle, François Laviolette, Mario Marchand, and Victor Lempitsky. Domain-adversarial training of neural networks. JMLR, 2016.
Matt Gardner, Yoav Artzi, Victoria Basmov, Jonathan Berant, Ben Bogin, Sihao Chen, Pradeep Dasigi, Dheeru Dua, Yanai Elazar, Ananth Gottumukkala, et al. Evaluating models' local decision boundarics via contrast sets. In EMNLP Findings, 2020.
Zhenyu Guo, Shuai Zheng, Zhizhe Liu, Kun Yan, and Zhenfeng Zhu. Cetransformer: Casual effect estimation via transformer based representation learning. In PRCV, 2021.
Shonosuke Harada and Hisashi Kashima. Graphite: Estimating individual effects of graph-structured treatments. In CIKM, 2021.
Jennifer L Hill. Bayesian nonparametrie modeling for causal inference. Journal of Computational and Graphical Statistics, 2011.
Keisuke Hirano, Guido W Imbens, and Geert Ridder. Efficient estimation of average treatment effects using the estimated propensity score. Econometrica, 2003.
Kurt Hornik, Maxwell Stinchcombe, and Halbert White. Multilayer feedforward networks are universal approximators. Neural networks, 1989.
Andrew Ilyas, Sung Min Park, Logan Engstrom, Guillaume Leclere, and Aleksander Madry, Datamodels: Predicting predictions from training data. arXiv, 2022.
Guido W Imbens and Donald B Rubin. Causal inference in statistics, social, and biomedical sciences. Cambridge University Press, 2015.
Alon Jacovi and Yoav Goldberg. Towards faithfully interpretable ulp systems: How should we define and evaluate faithfulness? arXiv preprint arXiv:2004.036655, 2020.
Andrew Jaegle, Felix Gimeno, Andy Brock, Oriol Vinyals, Andrew Zisserman, and Joao Carreira. Perceiver: General perception with iterative attention. In ICML, 2021.
Andrew Jaegle, Sebastian Borgeaud, Jean- Dapt-sap Alayrac, Carl Doersch, Catalin lonescu, David Ding, Skanda Koppula, Daniel Zoran, Andrew Brock, Evan Shelhamor, Olivier Henaff, Matthow Botvinick, Andrew Zisserman, Oriol Vinyals, and Joao Carreira. Perceiver IO: general architecture for structured inputs & outputs. In ICLR, 2022.
Fredrík Johansson, Uri Shalit, and David Sontag. Learning representations for counterfactual inference. In ICML, 2016.
Fredrik D Johansson, Uri Shalit, Nathan Kallus, and David Sontag. Generalization bounds and representation learning for estimation of potential outcomes and causal effects. arXie, 2020.
Yonghan Jung, Jin Tian, and Elias Barcinboim. Learning causal effects via weighted empirical risk minimiza- tion. 2020.
Jean Kaddour, Yuchen Zhu, Qi Liu, Matt Kusner, and Ricardo Silva. Causal effect inference for structured treatments. 2021.
14

---

Published in Transactions on Machine Learning Research (11/2023)
Nathan Kallus. Deepmatch: Balancing deep covariate representations for causal inference using adversarial training. In TCML, 2020.
Joseph DY Kang and Joseph L Schafer. Demystifying double robustness: A comparison of alternative strategies for estimating population mean from incomplete data. Statistical science, 2007.
Jacob Devlin Ming-Was Wei Chang Kenton and Lee Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In NAACL-HLT, 2019.
Sören R. Künzel, Jasjeet S Sekhon, Peter Bickel, and Bin Yu. Metalearners for estimating heterogeneous treatment effects using machine learning. PNAS, 2019.
Christos Louizos, Uri Shalit, Joris M Mooij, David Sontag, Richard S Zemel, and Max Welling, Causal effect inference with deop latent-variable models. In NeurIPS, 2017.
Shakir Mohamed and Balaji Takshminarayanan. Learning in implicit generative models. arxiv, 2016.
David Newman. Bag of words data set. UCI Machine Learning Respository. 2008.
Lizhen Nic, Mao Ye, Qiang Liu, and Dan Nicolae. Venet and functional targeted regularization for learning causal effects of continuous treatments. 2021.
Sonali Parbhoo, Stefan Bauer, and Patrick Schwab. Neore: Neural counterfactual representation learning for combinations of treatments. arv, 2021.
Judea Pearl. Causalily. 2009.
Raghunathan Ramakrishnan, Pavlo Dral, Matthias Rupp. and O Anatole Von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific dala, 2014.
Paul Rosenbaum and Donald Rubin. The central role of the propensity score in observational studies for causal effects. Biometrika, 1983.
Paul R Rosenbaum and Donald B Rubin. Reducing bias in observational studies using subclassification on the propensity score. JASA, 1984.
Donald B Rubin. Bayesian inference for causal effects: The role of randomization. The Annals of statistics, 1978.
Donald B Rubin. The design versus the analysis of observational studies for causal effects: parallels with the design of randomized trials. Statistics in medicine, 2007.
Donald B Rubin and Neal Thomas. Matching using estimated propensity scores: relating theory to practice. Biometrics, 1996.
Bernhard Schölkopf, Francesco Locatello, Stefan Bauer, Nan Rosemary Ke, Nal Kalchbrenner, Anirudh Goyal, and Yoshua Bengio. Toward causal representation learning. Proceedings of the IEEE, 2021.
Patrick Schwab, Lorenz Linhardt, and Walter Karlen. Perfect match: A simple method for learning representations for counterfactual inference with neural networks. arxiv 2018
Patrick Schwab, Lorenz Linhardt, Stefan Bauer, Joachim M Buhmann, and Walter Karlen. Learning counterfactual representations for estimating individual dose-response curves. In AAAI, 2020.
Uri Shalit, Fredrik D Johansson, and David Sontag. Estimating individual treatment effect: generalization bounds and algorithms. In ICML, 2017.
Claudia Shi, David Blei, and Victor Veitch. Adapting neural networks for the estimation of treatment effects. 2019.
Brian K Shoichet. Interpreting steep doso-nesponse curves in early inhibitor discovery. Journal of medicinal chemistry, 2006.
15

---

Published in Transactions on Machine Learning Research (11/2023)
Rahul Singh, Maneesh Sahani, and Arthur Gretton. Kernel instrumental variable regression. 2019.
Mark J Van Der Laan and Daniel Rubin. Targeted maximum likelihood learning. The international journal of biostatistics, 2006.
Tyler J VanderWeele. Principles of confounder selection. European journal of epidernioloyy, 2019.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In NeurIPS, 2017.
Victor Veitch, Alexander D' D'Amour, Steve Yadlowsky, and Jacob Eisenstein. Counterfactual invariance spurious correlations text classification. 2021.
Stefan Wager and Susan Athey. Estimation and inference of heterogeneous treatment effects using random forests. Journal of the American Statistical Association, 113(523): (3(1223):222811222 2018.
Hao Wang, Hao He, and Dina Katabi. Continuously indexed domain adaptation. In 1CML, 2020.
Haohan Wang, Zeyi Huang, Hanlin Zhang, Yong Jae Lee, and Eric Xing. Toward learning human-aligned cross-domain robust models by countering misaligned features. 2022.
Duncan Watts and Steven H Strogatz. Collective dynamics of small-worldd networkss nature, 1998.
Jeffrey M Wooldridge. Introductory econometrics: A modern approach. Cengage learning, 2015.
Yilan Wu, Ezra Winston, Divyansh Kaushik, and Zachary Lipton. Domain adaptation with asymmetrically- relaxed distribution alignment. In JCML, 2019.
Guoqiang Xu, Cunxiang Yin, Yuchen Zhang, Yuncong Li, Yancheng He, Jing Cai, and Zhongyu Wei. Learning discriminative representation base on attention for uplift. In Pacific-Asia Conference on Knowledge Discovery and Data Mining, pp. 200 211. Springer, 2022.
Chengxuan Ying, Tianle Cai, Shengjie Luo, Shuxin Zheng, Guolin Ke, Di Hc, Yanming Shen, and Tin Liu. Do transformers really perform bad for graph representation? 2021.
Jinsung Yoon, James Jordon, and Mihaela Van Der Schaar. Ganite: Estimation individualized treatment effects using generative adversarial nets. In ICLR, 2018.
Shuxi Zeng, Serge Assaad, Chenyang Tao, Shounak Datta, Lawrence Carin, and Fan Li. Double robust representation ion learning for counterfactual prediction, 2020.
Hanlin Zhang, Yi-Fan Zhang, Weiyang Liu, Adrian Weller, Bernhard Schölkopf, and Eric P Xing. Towards principled disentanglement for domain generalization. CVPR, 2022.
Yao Zhang, Alexis Bellot, and Mihaela Schaar. Learning overlapping representations for the estimation of individualized treatment effects. In AISTATS, 2020.
Han Zhao, Shanghang Zhang, Guanhang Wu, José MF Moura, Joao P Costeira, and Geoffrey J Gordon Adversarial multiple source domain adaptation. 2018
16

---

Published in Transactions on Machine Learning Research (11/2023)
Exploring Transformer Backbones for Heterogeneous Treatment Effect Estimation
Contents
A Extended Related Work
18
A.1 Detail information for Covariate and Treatment Sclf-Attention and Cross-Attention
18
A.2 Treatment-Covariate Cross-Ateentinn
19
B Details and Discussions about Propensity Score Modelling
19
C Analysis of the Failure Cases over Treatment Distribution Shifts
20
D Additional Experimental Setups
21
D.1 Detail Evaluation Metrics.
D.2 Network Structure and Parameter Setting
22
D.3 Simulation details.
23
E Additional Experimental Results
26
E.1 Comparision between TransTEE and ANU (Xu et al., 2022)
26
E.2 Additional Numerical Results and Ablation Studies..
27
E.3 Showcase of sentences and counterfactual counterparts with the maximal mokimal/minimal ATEs.
28
F Remarks on Interpretability
29
17

---

Published in Transactions on Machine Learning Research (11/2023)
A Extended Related Work
Most related works fundamentally rely on strongly ignorable conditions. Still even under ignorability, treatments may be sclectively assigned according to propensities that depend on the covariates. To overcomne the impact of such confounding, many statislical methods (Austin 2011 like covariate adjustment
Propensity Score.
(Austin 2011p, matching (Rubin & Thomas
1996; Abadie & Imbens 2016), stratification (Frangakis & Rubin 2002). reweighting (Hirano et al. 2003). g-computation (Imbens & Rubin 2015), have been proposed.
More recent approaches include propensity dropout (Alaa et al., 2017), and multi-takk Gaussian process
(Alaa & van der Schaar, 2017). Explicitly modeling the propensity score, which reflects the underlying policy for assigning treatments to subjects, has also shown to be effective in reasoning about the unobserved counterfactual outcomes and accounting for confounding. Based upon it, double robust estimators and largeted regularization are proposed to guarantee the consistency of estimated mreat effects under misspecification of either the outcome or propensity score model (Kang & Schafer 20077 Funk et al., 2011
There are also works using adversarial training for balanced representations (Bica et al.) 2020a Kallus 2020:
Du et al. 2021). However, most traditional approaches are restricted to binary treatments and the capacity of NNs for such problems have not been fully leveraged.
Domain Adaptation There are some close connections between causal inference and domain adaptation, in particular, out-of-distribution robustness. Intuitively, traditional domain adversarial training learns representations that are indistinguishable by the domain classifier by minimizing the worst-domain empirical error (Ganin et al. 2016 Zhao et al., 2018 Wang et al.,
2022
Zhang et al. 2022). The algorithmic insights can
be handily translated to the TEE domain Shalit et al.
2017 Johansson et al. 2020 Feder et al. (021). Here
we also have the desideratum that covariate representations should be balanced such that the selection bias is minimized and the effect maximally determined by the treatment. Algorithmically, when the trcatment is continuous, we connect our method to continuously indexed domain adaptation (Wang et al. 2020). Our formulation and algorithm also serve to build connections to diverse set of statistical thinking causal inference and domain adaptation, of which much can be gained by mutual exchange of ideas (Johansson et al. 2020). Explicitly modeling the propensity score also seeks to connect causal inference with transfer learning to inspire domain adaptation methodology and holds the potential to handle a wider range of problems like hidden stratification in domain genoralization, which we leave for future work.
Comparision between TransTEE and ANU (Xu al., 2022). (i) The model structure is different. ANU performs cross-attention between 2g, and 25, and no self-attention is applied. However, TransTEE performs self-attention 2z, respectively and then cross-attention is performed between 2n, 24. When facing high-dimensional data, such as lexts, images, and graphs, without multiple self-attention layers on Zz separately, the representations will be weak. That is why in machine translation, object detection, and segmentation tasks, the representations of imagcs/texts will be firstly processed by multiple self-attention layers and then perform cross- attention with queries. We will verify this point in the following experiments. (ii) ANU cannot be applied multi-treatment settings, which have been extensively studied recently (Kaddour 2021 Bica et al.) 2020b Parbhoo et al. 2021) The comparison experiments are in Section
et al.
A.1 Detail information for Covariate and Treatment Self-Attention and Cross-Attention
For covariates, prevalent methods represent covariates as a whole feature using MLPs, where pair- wise covariate interactions are lost when adjusting covariates. Therefore, we cannot study the effect of cach covariate on the estimated result. In contrast, TransTEE processes each covariate embedding independently and model their interactions by self-attention layers. Namely,
Mr (M M SI--T,ME-1 MM-1) M!-1 +M--) M = MLP(BN(M))) + M..
where M the output of I layer and BN is the BatchNorm layer. Simultancously, the treatments and dosages embeddings concatenated and projected to the latent dimension by linear layer, which generates
18

---

Published in Transactions on Machine Learning Research (11/2023)
new embedding Msr E R Then self-attention is applied
M - Har(Mar' Maa', !, M)r') + Mar1
M MLP(BN(M)) + M
Note that we don't need to explicitly include the key valuuequeryy weight matrices. Because the definition of M MM-1 (MM-, 1) includes the projection of the query, key, and value inputs. The solf-attention layer for treatments emables trealment interactions, an important desideratum for S- and T-learnens. Namely, TransTEE can model the scenorio where multiple treatments are applied and attain strong practical utility, e.g., multiple prescriptions healthcare or different financial measures in economics. This is an effective remedy for existing methods which are limited settings where various treatments are not used simultaneously.
A.2 Sreatment-Covariite
One of the fundamental challenges of causal meta-learners is to model treatment-covariite interactions. TransTEE realizes this by cross-attention module, treating Mat as query and Mg as key and value
M Ha (M " M +
M = MLP((( + Mt', 9 = (Poolloliig(M
where M is the output of the last cross-attention layer and M = M The above interactions are particularly important for adjusting proper covariate or confounder sets estimating trealment elfects WanderWeeley 2019), which empirically yields suitable covariate adjustment principles (the Disjunctive Cause Criteria) (De Luna et al.) 2011: Vander VanderWeele, 2019) about pre-treatment covoriates and confounders intuitively illustrated in Figure and corroborated in our experiments.
Denote g = [x(x t) and the training objective is the mean square error (MSE) of the outcome regression:
Lo(x, Σ (ys -o(x)xx)))))
(4)
B Details and Discussions about Propensity Score Modelling
(The optimum of propensity score model) In the equilibrium of the game, assuming the out- come prediction model is fixed, then the oplimum of TR is achieved when (z) B[O,(0)), La(x)) via matching the mean of propensity score (t) (xx))) and the marginal distribution p(2x(x)) and the optimum discriminator of PTR is achieved via matching both the mean and variance such that E ((), 2, (x)) ((), (1) (x)) ((),(1), (x)).
Proposition 2.
Proof. The proof concerns the analysis of the equilibrium of the Minimax Game. It is a special case of (Wang et al. (020)) when there are only two players, i.e. μ4 and Tg. We represent treatments explicitly and interpret the connections with combating selection biases. Given the fixed outcome regression model pos the optimal propensity score model x is
arg min (D,,, (x), (.))))
arg min (x),
(2,(t) - To (f)x)))))
(5)
arg min Eg, (x)
(L) ma (f)x)))))
The inner minimum is achieved at TH (f)xx) 1000000) given the following quadratic
form:
E (000 (1) - To (1) (x)))
(6)
[a, ( (t) - 2mg (f) (x) (e, (+) + no (0)
19

---

Published in Transactions on Machine Learning Research (11/2023)
We assume the above optimum condition of the propensity score model always holds with respect to the outcome regression model during training, then the minimax game in Eq. can be converted to maximize the inner loop:
max Lo (t)) La- 2t()))
(1))
(t) Eg, (0.0 (00
(7)
E4, (x.) Ex (00
(0, (.0.00000e0r00..000..0.0)0.00001
(x) (x)) V (ttt) (4) Ex,
Next we show the difference between Eq. and the variance of the treatment V[1,(t)]:
(x.(x))p(0=xx) Vo. (000:0/0ee.kky10e000000000000000 ((),
(8)
-E[+,(t)] Ex (x) Ex (t) (x) (x)
(x) (t) (x) Ox f (1) (x)
where the last. inequality is determined by Jensen's 's inequality and the convexity of (1)2.. The optimum is
achieved when ( is constant w.r.t (x) and so (t))(t))+_(x)) = (t),, (x,
The proof process for PTR is similar but includes the derivation of variance matching.
arg min (x), (t))
(E (t))
log (()
arg min E (,0
(t) (x)]
2
(9)
(E) 0 (x)) +,(1))² log V( (1)
arg min Egg (x) (.)
2V (t) [x)]
where (x) and Es, (1) denote (x)=xp((( (x)) and Es, (t respectively for brevity. The first term can be reduced to constant given the definition of variance:
(E( x +a(())²
Es, (x))p(0 (xx) Es (t))
(t)x)
(10)
(x)
2V (t))x(
2
The second term can be upper bounded by using Jensen's inequality:
log (t)xx
E
2
log ((
(11)
log (V]),
Combining Eq.
10 and Eq. 11 the optimum } + 1 log (t)). is achieved when ((),
(t), (x)] is constant w.r.t D (x) and so (t)x (x)] = (t),, (t) v[), vx)(x)
according to the equality conditions of the first and second incquality in Eq. respectively.
C Analysis of the Failure Cases over Treatment Distribution Shifts
As shown in Figure 4 (a.c), with the shifts of the treatment interval, the estimation performance of DRNet and TA RNtt decline significantly. VCNet achieves oo estimation error when h 5 partly because its hand-craft.
20

---

Published in Transactions on Machine Learning Research (11/2023)
projection matrix can only process values near [0, 1]. Another problem brought by this assumption is the extrapolation dilemma, which can be seen in Figure [b). When training on E [O, 1.75,, these discrete approximation methods cannot transfer to new distribution & E (1.75, 2.0]. These unseen treatments are rounded down to the nearest neighbors t' in T and be secmed the same as t'. We conduct ablation about the ment embedding as in Table G in Appendix. Such a simple fix EVCNeetEEnbeddingg removes the demand on a fixed interval constraint to treatments and attains superior performance on both interpolation and extrapolation settings. The result clearly shows the pitfalls of hand-crafted feature mapping for TEE. We highlight that it is neglected by most existing works (Schwab et al. 2020 Nie et al. 2021 Shi et al. 2019 Guo et al. 2021). Extrapolation is still a challenging open problem. We can see that no existing work does well when training and test treatment intervals have big gaps. However, the empirical evidence validates the improved elfectiveness of TransTEE that uses learnable embeddings to map continuous treatments to hidden representations.
Below we show the assumption that the value of treatments or dosages are in fixed interval [t. h] is sub-optimal and thus these methods get poor extrapolation results. For simplicity, we only consider a data sample has only one continuous treatment t and the result is similar for continuous dosage.
Proposition 3. Given a data sample (x, t),, where x E Rd, ,t, [1, y E R. Assume μ is a L-Lipschitz function over (x.t) E Rd+1 namely Iạ(u) [v(v] S L|/u u v/l. Partitioning h] uniformly into 5 sub- interval, and then get T [ h-t _ 0, 1 h-d
I h-! 5]. Previous studies most rounding down
a treatment to its nearest value in T (either ++ approzimate the entire conlinuum [1. h]. The approssmation error can be bounded by
65 A or l I+
15 h-1
1 and use IT') branches
h-1
max
1
x,
tô
h 1
t(x,tt,,1
t)
x,
tô
h
8 ) m(xx,)) 1)
1)
l
8
h
1o
h
1).
(1)
h
1
}
max
L
(12)
5
1
8
-1 l
< T.
5
The bound is affected by both the number of branches d and treatment interval [t, h]. However, as far as we know, most previous works ignore the impacts of the treatment interval [7. h] and adopt simple but much stronger assumption that treatments are all in the interval [0, 1] Nie et al. (2021) or a fixed interval [Schwab] et al. (020)). These observations well manifest the motivation of our general framework for TEE without the need for treatment-specific architectural designs.
Table 6: Experimental results comparing NN-based methods on simulated datasets. Numbers reported are AMSE of test data based 100 repeats, and numbers after + are the estimated standard deviation of the average value. For Extrapolation (h = 2), models are trained with & E [0, 1.75) and tested in E [0. 2]. For Extrapolation (h = 5), models are trained with t E [0,4) and tested in 1 E [0.5]
METHODS
VANILLA
VANILLA (h 5)
EXTRAPOLATION (A 2)
EXTRAPOLATION (h 5)
TARNET
SHALIT ETT AL..
2017
0.045 + 0.0009
0.3864
0.04335
0.0984 + 0.02315
0.3647 + 0.03626
DRNET SCHWAR ET AL.
2020
0.042 + 0.0009
0.3871
0.03851
0.0885 t 0.00094
0.3647 + 0.03625
VCNET NiR FT At..
202T
0.018 + 0.0010
NAN
0.0669 + 0.05227
NAN
VCNET EMBEDDINGS
0.013 + 0.00465
0.0167
+ 0.01150
0.0118 + 0.00482
0.0178 E 0.00887
D
Additional Experimental Setups
All the assets (i.e., datasets and the codes for baselines) we use include a MIT license containing a copyright notice and this permission notice shall be included in all copies or substantial portions of the software. We conduct all the experiments on a machine with 17-8700K CPU, 32G RAM, and four Nvidia GeForce RTX2080Ti (10GB) GPU cards.
21

---

Published in Transactions on Machine Learning Research (11/2023)
D.1 Detail Evaluation Metrics.
N
AMSET 1 N
Σ
- (xx, +()dtt
(13)
UPEHE@K
f(x,ttt) (f(n))-((sss tx
C 1,!!
(14)
WPEHEOK
-(x,hht))((xx,11)) - ((x,,,1))) ((10))))((())pp)pp) p(11x))((1xx))
N
T
AMSE
1
f(x,tt f(xm.,, a(s)dt
(15)
NT
D.2 Network Structure and Parameter Setting
Table. 7 and Table. show the detail of TransTEE architecture and sspea-parameters. For all the synthetic and semi-syuthetic datasets, we tune parameters based on additional runs. each run, we simulate data, randomly split into training and testing, and use AMSE on testing data for evaluation. For fair comparisons, in all experiments, the model size of Trans is less than similar to baselines.
Table Architecture details of TEE, where the number of covariates.
Module
Covariates
Treatment
Embedding Layer
[Linear]
[Linear]
Output Size
Bsz x x # Emb
baz x 1 x # Emb
Multi-head Att.
Multi-head Att
BatchNorm
BatchNorm
AetenAttention
X # Layers
Layers
Linear
Linear
BatchNorm
BatchNorm
Output Size
Bsz #Emb
Bsz x 1 #Emb
Multi-head Att
BatchNorm
Cross-Attention
x # Layerss
Lincar
BatchNorm
Output Size
Bsz x #Emb
Projection Layer
[Linear]
Output Size
Bsz
Table 8: plaaer-parameters on different datasets. Bsz indicates the batch size, Emb indicates the embedding dimension, Lr. S indicates the schoduler of the learning rate (Cos is the cosine anncaling Learning
rate).
Dataset
Bsz
Emb
# Layers
Heads
Lr
Lr. S
Simu
500
1
0.01
Cos
IHDP
128
0.0005
Cos
News
256
10
1
2
0.01
Cos
SW
500
16
1
2
0.01
None
TOGA
1000
48
3
4
0.01
None
22

---

Published in Transactions on Machine Learning Research (11/2023)
D.3 Simulation details.
Synthetic Dataset (Nie el al. 2021 The synthelic dataset contains 500 training points and 200 testing points. Data is generated as follows: 2j ~ Unif[O, 1], where zÿ is the j-th dimension of x E R² and
z sinin (max( J2, 23)) + max(23, 25x
sin(0.5xz) (1 exp(24 - 0.5zg)) +
1+ (x) + zz)²
23 + 2 (in(( + 2m. - 6.5 (0,,0.25)
( 4 4max(x11 T6)
(2 0.5))
+.((0,0.25)
1 203
where (1 exp(-
for treatment in [0, hl, we revised it to (1 + exp -t)- 1-
IHDP Hill 2011 is a semi-synthetic dataset containing 25 covariates, 747 observations and binary treatments.
For treatments in [0, 1}, we follow VCNet. (Nee et al., 2021 and generale treatments and responses by:
2x1 + mmx((13s 25)
Σ Diesun (I. (2) -4+ 44
Z = 1 x2
0.2 + min min(13, 25. xo) + 2 tanh
|
y/2, t = sin(3rt)
Eie Eiesar (xt c1)
cxp(0.2((x) - x6))
tanh
+ (00,0.25)
1.2 -
Sdia.1
0.5 + 5 min (ain(22, x3, x35)
where +1+ exp( 1 S Sem
(1,2,3,5,6} is the index set of continuous features, Sdis,1
74, 8,9,8,9, 8, 10, 11, 12, 13, 14, 15), S
{16,17,88, 19, 20, 21, 23, 25} and Saie.1 Sas,, [25] Scon
Σ
Here C1
E
.2 E
To allow comparison various treatment intervals
1
E [0. hl. treatments and responses are generated by:
t= =1xxxp(--))-
(in(3mt/h) y/x,t t- 1.2 - 1/h
taanh 5 Diesuinn (06--1)
exp(0.2(x1 - mi6)
+ +N 0.25),
Saia.1
0.5 + 5 Z3, 23, 275)
where the orange part is the only different compared to the generalization of vanilla IHDP dataset (h 1-=10 1). Note that S only impacts outcome that serves to be noisy covariates; Saie.2 contains pre-treatment covariates that only impact treatments, which also serves to be instrumental variables. This allows us to observe the improvement using TransTEE when noisy covariates exist. Following (Hill 2011 I covariates are standardized with mean and standard deviation 1.
News. The News dataset consists of 3000 randomly sampled news items from the NY Times corpus Newman 2008) and was originally introduced as a benchmark in the binary treatment setting. We generate the treatment and outcome in a similar way as (Nie et al. 2021 but with dynamic range or treatment intervals [0, h) We first generate uf, 0/g, o/ ~ ~N(00, and then set v, = o//llu//l; e (1, Given we generate from
Beta
2,
o. x
A.dnd we generate the outcome by
12 0.3
v|zz,t= =ex==xxp exp
1gg
g/2, t = 2(max min(2. 20x 5²) 4(t - 0.5) + sin (2t) + +.(,0.555
TCGA (D)
[Bia et al. 2020bb We obtain covariates from real dataset The Cancer Genomic Atlas (TCGA) and consider 3 treatments, where each treatment accompanied by one dosage and set of
23

---

Published in Transactions on Machine Learning Research (11/2023)
8
SCIGAN
SCIGAN
20
(D)
7-
Tamet (D)
Drnet (D)
Drnet (D)
Venet (D)
Venet (D)
TranTTEE
TransTEE
AMSE
AMSE
10
2-
1
0
2
10
10
Dosage Scloction Bias
Treatment Selection Bias
(a) Performance with different dosage selection bias.
(b) Performance with different treatment selection bias.
Figure 8: Performance of five methods on TCGA (D) dataset with varying bias levels.
parameters, of, u². v/. For each run, randomly sample vector, ut (0, and then set of = w [u]
where ||- 1/ is Euclidean norm. The shape of the response curve for each treatment, (z, s) is given in Table g We add e ~ .V(0, noise to the outcomes. Interventions are assigned by sampling dosage, d,, for each treatment from a beta distribution, d/x Beta(a, 3t). a 2 1 controls the dosage selection bias (a = gives the uniform distribution). Br = e_ع + a, where si is the optimal dosage for treatment t. We then assign treatment according to tyl Cateeorrcal(Softmaxxffff s;)) where increasing increases selection bias, and 0 leads to random assignments. The factual intervention is given by (4,, 8ty)) Unless otherwise specified, we set = 2 and = 2.
Tuuhh
18
30
Dmntt (D)
Venet (D)
SCIGAN
25
13
14
Respone
12
Response
12
Respone
20
10-
10
Drnet
Drnet (D)
Vendt
Venet (D)
SCIGAN
SCIGAN
Tra
TransTER TEE
0.0
0.4
0.8
0.2
0.5
0.8
1.0
0.0
0.4
1.0
Treatment
Treatment
Treatmeni
(a) Estimated ADRF for
(b) Estimated ADRF for t2.
(c) Estimated ADRF for f3.
Figure 9: Estimated ADRF the test set from a typical run of DRNet (D), TARNet (D), VCNet (D), and SCIGAN. All of these methods are well optimized. TransTEE can well estimate the desage-response curve for all treatments.
For structural treatments, we first define the Baseline effect (Bica el al. 2020b). For each run of the experiment, we randomly sample vector up Unif(0, 1), and set up ug/ /oll where || is the Euclidean norm. The baseline effect is defined
[o(x) = u. x
For symmetry, if sf O, sample of from 1-Beta(a,89) (4) where By though of
24

---

Published in Transactions on Machine Learning Research (11/2023)
Table 9: Dose response curves used to generate semi-synthetic outcomes for patient features z.
In the experiments, we set C 10. rf, 122 1/2 are the parameters associated with cach treatment t.
Treatment
Dose Respones
Optimal dosage
1
f1 8) = C ((+)) x 12(mg) x8 -
$1 2( 1 x
(e))
2
fe(a, = C (uf)TT x sin (x())==))
$2 2( x
(og) x
3 8s(2, = C (wf) Tx 12s(s - 0)2, where b 0.75 (1))
if 2 0.75 else 1
Small- World (Kaddour et al., to x; ~ 1n,ff-1,1]. There are 0,000 units in in-sample datasct, and 500 in the out-sample onc. Graph interventions For each graph intervention, a number of nodes between 10 and 120 are uniformly sampled, the number of neighbors for each node is between 3 and 8, and the probability of rewiring each edge is between 0.1 and 1. Watts-Strogan vertex has one feature, i.c. its degree centrality. A graph's node connectivity is denoted v(g) and its average shortest path length as ( Similar for the baseline elfect, two randomly sampled vectors V2, Uy are generated. Then, given an assigned graph treatment G and a covariate vector x, the outcome is generated by
2021 20-dimensional multivariate covariates are uniformly sampled according
small-world graphs are repeatedly generated until a connected one is get. Each
y = (x) + 0.2x())² vr,x + (()) of te, -,0+-~(0011)))
TCGA (S) [addour et al. 2021 We use 9, 659 gene expression measurements of cancer patients for covariates. The in-sample and datasets consist of 5, 000 units and the out-sample one of 4, 659 units, respectively. Each unit is covariate vector x E 124000 and these units are split randomly into in- and out-sample datasets in each run randomly. For each covariate vector æ; its 8-dimensional PCA components PCA R is computed. Graph interventions We randomly sample 10, 000 molecules from the Quantum Machine 9 (QM9) dataset (Ramakrishnan et al., 2014) (with 133k molecules in total) in each run. We create relational graph, where each node corresponds to atom and consists of 78 atom features. We label each edge corresponding to the chemical bond types, e.g., single, double, triple, and aromatic bonds. We collect 8 molecule properties mu, alpha, homo, lumo, gap, r2, zpve, u0 in vector 2 E R$, which is denoted as the the assigned molecule treatment. Finally, we generate outcomes by
y = (0)))()) + 0.012 x. PCA +c,6~N(00.
Enriched Equity Evaluation Corpus (EEEC) Feder et al. 2021 consists of 33, 738 English sentences and the label of each sentence is the mood state it conveys. The task is also known as Profile of Mood States (POMS). Each sentence in the dataset is created using one of 42 templates, with placeholders for a person's 's name and the emotion, e.g., "<Person>> made me feel <emotional state word>.". ". A list of common names that are tagged as male or female, and as African-American or European will be used to fill the placeholder (< (<erson>>). One of four possible mood states: Anger, Sadness, Fear and Joy is used fill the emotion placeholder. Hence, EEEC has two kinds of counterfactual examples, which are Gender and Race. For the Gender case, it changes the name and the Gender pronouns in the example and switches them, such that for the original example: 'It was totally unerpected, but Roger made me feel pessimistic. it will have the counterfactual example: was lolally unespecled, bul Amanda made me feel pessimistic.." " For the Race concept, it creates counterfactuals such that for the original example "Josh made me feel uneasiness for the first time ever in my life.", ", the counterfactual example is: "Darnell made me feel uneasiness for the first time ever in my life.". For each counterfactual example, the person's 's name is taken at random from the pre-existing list. corresponding to its type.
25

---

Published in Transactions on Machine Learning Research (11/2023)
0.12
0.14
0.12
0.10
0.10
AMSE
0.08
0.06
0.06
0.04
0.04
0.02
0.00
0.00
1.0
4.0
(a) Ablation study of PTR.
(b) Ablation study of TR.
Figure 10: Ablation study of the balanced weight for treatment regularization on the IHDP dataset.
Table 10: Comparision between TransTEE and ANU (Xu et al. 2022) on the IHDP dataset.
Methods
Vanilla (Binary)
Vanilla (h = 1)
Extrapolation (h 2)
DRNet
0.3543 + 0.6062
2.1549 + 1.04483
11.071 + 0.9938
VCNet:
0.2098 + 1236
0.7800 + 0.6148
NAN 0.4244 + 0.19832
ANU
Xu et al. 2022)
0.1482 + 0.17362
0.2147# 0.32451
TransTEE)
0.0983 = 0.15384
0.1151 0.1028
0.2745 0.1497
E
Additional Experimental Results
E.1
Comparision between TransTEE and ANU (Xu et al. 2022
We implement ANU and evaluate in the same settings and show that inferior compared the proposed TransTEF follows. Specifically, we compare the attentive neural uplift model (ANU) (Xu et al.) 2022) with ours in the following two settings. (1) IHDP dataset in Table 10 in the main manuscript. We adjust the layers of ANU such that the total parameters of ANU and Trans TEE are similar. The result is shown in the following table. With the usage of treatment embeddings, ANU is shown to be more robust than VCNet. and DRNet when a treatment shift occurs. However, in both the binary treatment sctting and continuous treatment settings, TransTEE performs better than ANU.
(2) We further evaluate the real-world utility of ANU (Xu et al. 2022) and the experimental setting is detailed in Scction 6..4] in the main paper. Covariates here are long sentences. Thanks to the use of self-attention modules, Trans TEE can achieve better estimation results compared to baselines (Table 11.. For AHU, no self-attention layer is applied, and the final estimation is inaccurate, which verifies the superiority of the proposed framework.
Table 11: Comparision between TransTEE, and ANU (Xu et al.
2022 on the IHDP dataset.
Correlation/ Reairisentteeeeeettttttiinnn Based Baselines
Treatment Effect Estimators
TC Gender
ATEgT
TReATE
CONEXP
INLP
CausalBERT
TarNet
DRNet
ANU
TransTEE
0.086
0.125
0.02
0.313
0.179
0.0067
0.008s
0.184
0.013
Race
0.014
0.046
0.08
0.591
0.213
0.005
0.006
0.093
0.0174
26

---

Published in Transactions on Machine Learning Research (11/2023)
023
6.2
1.3
6.5
5.5
685
10
5.0
001103322112111111000 23732225
029
Covorares
Covaniats
Covoriates
(a) TransTEF.
(b) TransTER+TR.
(c) TransTER+
Figure 11: The distribution of learned weights for the cross-altention module on the IHDP dataset of different models.
125
Tame
Tomnet
Tamat
10.0
Tarnel
7.5
7,5
Diset Veatt
Veset
7.5
Reqonts
5.0
TrausTBE
Resp000
TRESTTEE
Bergone
ThaSTEE
2.5
0.0
-2.5
Treatment
Treatmert
Treaimm
(a) h during training and (b) 2,/ 0.1 during tran- (c) 5 during training and (d) h 5, 0.25 during train- testing.
ing and 2,1 during testing.
ing and 5, during testing (extrapolation).
testing (extrapolation).
Figure 12: Estimated ADRF on test from a typical run of TarNet (Shalit al. 2017) DRNet (Schwab et al. 2020), VCNet (Nie et al. 2021) and ours on IHDP dataset. All of these methods are well optimized. (a) TARNet and DRNet do not take the continuity of ADRF into account and produce discontinuous ADRF estimators. VCNet. produces continuous ADRF estimators through hand-crafted mapping matrix. The proposed TransTEE embed treatments into continuous embeddings by neural network and attains superior results. (b,d) When training with 0.1 < 1 < 2.0 and 0.25 < < 5.0. TARNet and DRNet cannot extrapolate to distributions with 0 < < 2.0 and 0 < < 5.0. (c) The hand-crafted mapping matrix VCNet can only be used in the scenario where the same time, h be enhanced, TARNet and DRNet with the same number of branches perform worse. Trans needs not to know h in advance and extrapolates well.
Otherwise, VCNet cannot converge and incur an infinite loss. At
E.2 Additional Numerical Results and Ablation Studies
Choice of the balancing weight for treatment regularization. To understand the effect of propensity score modeling, we conduct an ablation study on the balancing weights both TR and PTR. Figure I0 presents the results of the experiments on the THDP dataset. The main observation is that both TR and PTR with a proper regularization strength consistently improve estimation compared to TransTEE without regularization. The best performers are achieved when 0.5 for both two methods, which shows that the best balancing parameter (0.5
experiments.) for these two regularization terms should be searched carefully. Besides, training both the treatment predictor and the feature encoder simultancously in a zero-sum game is difficult and sometimes unstable (shown in Figure 10 right)
Robustness to noisy covariates. We manipulate Saie.1 Saie,2 to generate datasets with different noisy covariates, c.g., when the number of covariates that only influence the outcome is 6, Sā = (4), and Sass22 {7,8,9) 9, 10. 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25), when the number of covariates that influence the oulcome is 24, Sais.1 7,,7,8, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24,), and
27

---

Published in Transactions on Machine Learning Research (11/2023)
Trin
Ourcoms Regrosion Ervor
Trans PTX
25
Trans E22
Tam
Repezssion
0.25
Mac Eitor
TIaIsTES
TresSEE-TR
TraISTER*PRR
15
TransEES+PT
TrasTTi MTL
-0.25
M+MM..
-0.90
5
0.15
250
500
750
1000
1250
1500
300
750
1000
1250
1500
5
15
20
25
30
Berations
Berations
(a) Outcome regression
(b) Treatment regression
(c) MSE in the test set.
Figure 13: Training dynamics of TransTEE IHDP dataset with various regularization terms, where the total training iteration is 1, 500 and (c) is evaluated on the test set per 50 training iterations.
Table 12: Experimental results comparing neural network based methods on the News datasets. Numbers reported are based on 20 repeats, and numbers after + are the estimated standard deviation of the average value. For Extrapolation (h = 2), models are trained with E [0, 1.9] and tested in E [0, 2). For For Extrapolation (h - 5), models are trained with & E [0, 4.5] and tested in 5)
METHODS
VANILLA
VANILLA (h 5)
EXTRAPOLATION (h EXTRAPOLATION (h
TARNET
0.082 0.019
0.956 0.041
0.716 1 0.038
0.847 + 0.053
DRNET
0.083 + 0.032
0.956 + 0.041
0.703 1: 0.038
0.834 + 0.053
VCNET
0.013 0.005
NAN 0.017 + 0.008
NAN
NAN
TRANSTEE
0.010 + 0.004
0.024 1 0.017
0.029 + 0.019
TRR
0.011 0.003
0.016 0.008
0.019 1 0.008
0.028 +\2 0.002
TRANSTEE+PPR
0.011 + 0.004
0.014 + 0.007
0.022 + 0.008
0.029 + 0.016
Sas,2 (25). Figure Figure 6(b) shows that, as the number of covariates that only influence the outcome increases, both TARNet and DRNet become better estimators, however, VCNet performs worse and even inferior to TARNet and DRNet when the number is large than 16. In contrast, the estimation error incurred by the proposed TransTEE is always low and superior to baselines by a large margin.
Comparison of MLE or adversarial propensity score modeling on the propensity score. Seeing results in Table 2 additionally combine TransTEE with maximum likelihood training of x(t)x) does provide some performance gains. However, an adversarially trained m-model be significantly better, especially for extrapolation settings. The results well manifest the effectiveness of TR and PTR reducing selection bias and improving estimation performance. In fact, approaches like TMLE are not robust if the initial estimator is poor Shi et al. 2019)..
Training dynamics comparison of different regularization terms. Here we compare four regularization terms, which are TransTEE with no regularization, TranSTEE+TR, TransTEE-PTR, and MTL. MTL is simple Multi-Task Learning strategy, which uses Lo(x, y, t) CTR t) during training without adversarial game. As shown in Figure 13 without adversarial training, TransTEE+MTI. quickly attains low trcatment estimation error but further oscillate and converge with high crror, and both the outcome regression error and MSE in the test set remain high. In contrast, TR and PTR make TransTEE converge faster and attain lower test MSE. Overall, PTR consistently works the best and its low treatment regression error shows that n (t(x) estimates an accurate propensity score.
E.3 Showcase of sentences and counterfactual counterparts with the maximal/minimal ATEs.
Table 14 showcases the top-10 samples with the maximal/ minimal ATEs. Interestingly, we can see most sentences with large ATE have similar patterns, that is "< clause >, bul/ond Person made me feel <
28

---

Published in Transactions on Machine Learning Research (11/2023)
Adj >". Besides, most sentences with large ATE have small length, which is 11 words average. By cantrast, sentences with small ATFs follow other patterns and are longer, which is 17.6 average. Consider the effect of Race, Table 15 showcases the top-10 samples. Similarly, there are also some dominant patterns that have pretty high or low ATEs and the average length of sentences with high ATEs smaller than sentences with low ATEs (12 ws 14.7). Besides, the position perturbation words (the name from a specific race) for sentences with the maximal//miiime former and the middle for the latter. Namely, TransTEE helps us mitigate spurious correlations that exist in model prediction, e.g., length of sentences, the position of perturbation words, certain sentence patterns and is useful in mitigating undesirable bias ingrained in the data. Besides, well-optimized TransTEF is able to estimate the effect of every sentence and is of great benefit for model interpretation and analysis especially under high inference latency.
ATEs totally different, which at the beginning for the
F Remarks on Interpretability
It is fundamentally hard to evaluate the interpretability even for supervised learners, as the evaluation crucially depends on specific models, tasks, and input spaces (Jacovi & Goldberg 2020 TransTEE provide an initial step to promote causal inference model interpretability. We can see from the experimental results in fig. 4(a), 4(b), and fig. 10 that TransTEE assigns more weights to confounders as opposed to other covariates, which is causal inference models in this way using the feature importance scores for each covariate can be used for benchmarking mrettment effect estimators (Crabbé et al. 2022).
observation that previous backbones are hard to achieve. We that explaining
29

---

Published in Transactions on Machine Learning Research (11/2023)
Table 13: Error of CATE estimation for all methods, measured by WPEHE@2-10. Results are averaged over trials, + denotes std error. In-Sample means results in the training set and Out-sample means results in the lest set. (The baseline results are reproduced using the official code of (Kaddour el al. 2021 in a consistent experimental environment, which can be slightly different than the results reported
(Kaddour al. )
Sw
TCGA (Bias=0.1)
TOGA (Bins=0.3)
TOGA (Bias=0.5)
Method
In-semple
Out-sample
In-sample
Out-sample
In-sample
Out-sample
In-sample
Out-eample
WPEHE
Zero
41.72 + 0.00
49.69 + 0.00
13.93 = 0.00
13.13 = 0.00
13.93 + 0.00
13.13
0.00
13.93 + 0.00
13.61 = 0.00
GNN
17.38
24.53 +:
10.90
7.71
7.71
13.58
13.22
+: 0.38
0.91
GraphITEE
17.37
0.01
24.56 + 0.02
15.04
0.20
14.96
0.30
13.49
0.23
13:70
0.52
12.41 + 0.02
14.38
0.30
SIN
15.79 + 1.72
28.78
+ 4.54
46.47
2.19
54.41
7.81
7.93 + 0.79
11.04
1.52
10.31
+ 0.93
14.09
2.14
Trans TEE
14.74 + 0.09
21.78 + 1.07
9.07 + 2.15
9.33 + 2.13
7.54 + 3.60
8.37 + 3.64
9.52 + 3.59
10.10 + 3.79
WPEHER
Zero
40.75 + 0.00
43.76 + 0.00
13.93 + 0.00
13.61 + 0.00
13:93 + 0.00
13.61
0.00
13.61 + 0.00
14.14 - 0.00
GNN
18:26 +: 0.00
20.91 + 0.01
10.75 = 7.60
10.91 = 7.72
13.63 + 0.18
13.58
0.19
12.92 + 0.33
15.29 = 1.04
GrahiTE
18.27 +: 0.01
20.95 +: 0.02
14.88
0.19
15.12
0.29
13:49 +: 0.22
14.19
0.43
12.56
+ 0.01
15.18 = 0.31
SIN
18.15 = 1.97
23.62
+ 3.93
29 2.33
53.72 = 8.09
7.94 + 0.75
11.53
1.59
10.89
+ 1.07
14.27 = 1.92
TraEsTEE
15.30 1.12
18.73
d: 2.09
2.02 9.07
9.58
2.04
7.58
3.62
de 3.75
9.64
3.56
10.59 +
WPENEE3
Zero
45.74
0.00
0.00
14.14 0.00
0.00
14.14
13.75
+ 0.00
= 0.00
GNN
22.09 + 0.01
23.01 + 0.01
10.87 7.69
10.88 7.69
13.87
0.18
13.71
0.19
13.13 + 0.34
15.47 1.05
GraphITE
22.12 + 0.00
23.03 + 0.02
15.05 0.18
15.14 0.28
13.64 + 0.20
14.30
0.35
12.77 + 0.02
15.38
0.30
SIN
22.14
2.30
23.70
+ 3.67
44.72
2.35
53.12 8.09
7,99 + 0.73
11.66
1.59
11.38 + 1.04
14.37
1.83
TransTEF
Trans TEE
18.99 + 0.83
19.65 + 1.97
9.09 + 1.97
9.66 + 2.01
7.67 + 3.70
8.71 + 3.78
9.78 + 3.63
10.74 + 3.91
WPPHE
Zero
49.19 + 0.00
45.96 + 0.00
14.31 = 0.00
13.95 = 0.00
14:31 + 0.00
13.95
0.00
13:95 + 0.00
14.47 = 0.00
GNN
24.18 +: 0.01
24.20 + 0.01
10.99
7.77
10.97 = 7.76
13.98 + 0.17
13.92
0.18
13.31
+ 0.37
15.67 = 1.05
GraphITE
24.22 + 0.01
21.22 + 0.03
15.21
0.19
15.29 = 0.28
13.68 + 0.17
14.37
0.37
12.95
+ 0.03
15.59 = 0.30
SIN
25.48
25.44
3.50
44.55
2.35
52.78
8.04
0.75
11.76
1.22
1.84
TrausTEE
20.16 0.42
21.08 + 1.78
9.17
1.96
9.72
2.00
7.76 + 3.75
8.80 + 3.82
9.91
3.66
10.89 + 3.94
WPEL1E66
Zero
49.95 +: 0.00
50.10 + 0.00
14.47 0.00
14.04 = 0.00
14.47 + 0.00
14.04 0.00
14.04 + 0.00
14.53
0.00
GNN
25.13 : 0.00
26.93 + 0.01
11.11 7.86
11.02 = 7.79
11.07 + 0.22
14.11
0.18
13:45 + 0.38
15.76
1.01
GraphITE
25.17 + 0.02
26.94 + 0.02
15.40
0.19
15.37 = 0.28
13.74 + 0.12
14.58 + 0.38
13.00 + 0.04
15.68
0.29
SIN
27.07 +: 2.98
28.11
+ 3.51
44.48 = 2.35
52.54 = 7.99
8.22 + 0.75
11.82 + 1.58
11.97
+ 1.19
14.74 = 1.86
TransTE TEE
21.32 + 0.79
22.99 + 1.43
9.23 + 1.93
9.77 + 1.99
7.80 + 3.83
8.84 + 3.89
10.01 + 3.70
10.96 1 3.95
WPEHE337
Zero
55.40 +: 0.00
58.42 + 0.00
14.53 = 0.00
14.09 = 0.00
14.53 + 0.00
14.09
0.00
14.53 + 0.00
14.09 = 0.00
GNN
29.30 + 0.03
32.15 + 0.03
11.16 7.89
11.06 = 7.82
14.12 0.21
14.14
0.18
13.51 + 0.38
15.81 = 1.03
GraphTTE
29.34
15.47 0.19
0.28
13.97
14.69
0.04
15.74 0.29
SIN
31.07
3.07
34.17
3.41
44.45 2.37
52.40 = 7.98
8.28
0.74
11.85
1.58
12.11
1.18
14.83 = 1.87
TransTEL
24.71
0.41
25.84 1: 0.73
9.27
1.94
9.81 1.99
7.82
3.84
8.89
1: 3.89
10.06
3.71
11.01 1: 3.95
WPE.HESS
Zero
57.99 : 0.00
66.78 + 0.00
14.61 0.00
14.14 = 0.00
14.60 + 0.00
14.12
0.00
14.61 + 0.00
14.14 0.00
GNN
31.41 + 0.03
33.57 + 0.05
11.22
7.93
11.09
7.85
14.19 + 0.25
14.20
0.18
13.58 + 0.38
15.87
1.02
GapLITE
31.45 +: 0.01
37.58 + 0.00
15.55 = 0.19
15.47 = 0.28
14.30 1 0.04
14.85 + 0.43
13.23 + 0.04
15.78
0.28
SIN
33.58 + 3.37
10.83
+ 3.64
11.8 = 2.38
$2.31 = 7.97
8.33 + 0.74
11.87 + 1.57
12.22
+ 1.17
14.91 = 1.89
TraISTEE
26.48 = 0.27
32.40 + 0.85
9.31 + 1.94
9.85
1.99
7.88 + 3.84
8.90 + 3.90
10.10 = 3.72
11.04 = 3.96
WPEHE
Zero
62.52 + 0.00
61.61 + 0.00
14.66 0.00
14:20 = 0.00
14.61 0.00
14.14 0.00
14.66 + 0.00
14.20 = 0.00
GNN
34.13
36.48
0.04
11.26
7.96
7.87
14.21
14.22
+ 0.38
GraphITE
34.17
0.02
36.49
0.01
15.60
0.19
15.53
0.28
14.35
0.01
14.90
0.43
13.28
0.04
15.83 = 0.28
SIN
36.79 + 3.35
40.99 + 5.14
44.47
2.39
52.31
7.97
8.36 + 0.74
11.90
1.57
12.40
+ 1.23
15.08
1.80
TransTEL
28.84 # 0.23
31.40 + 0.71
9.34 + 1.94
9.88 + 2.00
7.90 1 3.85
8.94 + 3.91
10.14 + 3.73
11.08 1 3.97
WPEHER10
Zero
62.65 + 0.00
65.59 + 0.00
14.69 + 0.00
14.23 + 0.00
14.69 + 0.00
14.23
0.00
14.69
+ 0.00
14.23
0.00
GNN
34.26 :: 0.04
37.65
+ 0.04
11.28 = 7.98
11.16 = 7.89
14.29 + 0.22
14.32 4 0.18
13.66
+ 0.38
15.96 = 1.01
GraphITE
34.30 :: 0.02
37.66
0.00
15.64 = 0.19
15.56 = 0.28
14.38 1 0.04
14.93 + 0.43
13.31
+ 0.04
15.87 = 0.27
SIN
37.08 1 3.35
11.79
+ 5.21
11.19 2.40
52.28
= 7.96
8.39
- 0.74
11.92
1.58
12.49
+ 1.22
15.13 = 1.81
TransTE: TEE
0.19
32.25 4: 0.69
9.36
1.93
9.90
2.00
7.94 ::
4: 3.92
10.16
+: 3.74
11.10 +
30

---

Published in Transactions on Machine Learning Research (11/2023)
GNN
GraphITE
Zero
TransTEE
SIN
GraphITE
Zero
TransTEE
WPEHESS
WPPNDe2
36
15
25
WPEEEES
50
WFBEDDe
50
WFEEDEE
23
WFEHEBS
50-
WPPL 12ge
WE-EHDD 10
25
20
20
40
20
40
20
(a) SW In-Sample.
(b) SW Out-Sample
GraphITE
TransTEE
GraphITE
TransTEE
WPEREBS
40
WPERReN
WPCmme
50
20
25
WPSRERS
40
WP IPEEEE
20
WPEHEBS
40
WPO 11G8
WP WPPOFBBE O
50
WP MPPUIZE O
20-
25
(c) TCGA In-Sample
(d) TOGA Out-Sample
Figure 14: WPEHR@K over increasing bias strength and varying K E (2...., 10} on the SW and the TCGA
dataset.
31

---

Published in Transactions on Machine Learning Research (11/2023)
GNN
GraphITE
Zero
TransTEE
SIN
GraphITE
Zero
TransTEE
UPSRE
UPTHEW
UP211084
50
S0
23
25
UPENEBS
10
UPD Piiiig
UPEREE
UPENEES
UPENEO
UPP UPETEE
25
25
UPSRE99
UPCt18g9
UPERIBeO
UPEN220
UPEtmeee
50
25
20
20
40
20
40
20
(a) SW In-Sample.
(b) SW Out-Sample
GraphITE
TransTEE
GraphITE
TransTEE
UPSHE82
50
UPSLIC03
UPST
50
UPERE22
UPSLICA
25
25
75
UPERERS
UPENDD
UPENEEI
UPCIee
UPERB:09
UPEN189
UPCNER 10
UPCH
50
UPCIE20 10
50
25
25
(c) TCGA In-Sample
(d) TOGA Out-Sample
Figure 15: UPEHE@K over increasing bias strength and varying K E (,,
... 10) on the SW and the TCGA
dataset.
32

---

Published in Transactions on Machine Learning Research (11/2023)
Top-10 samples with the maximal and minimal ATE for the effect of Gender. Perturbation words in factual sentences and counterfactual sentencos are colored by Orange and Magenta respecttively.
Table 14:
Sentences with The Maximal ATEs
Index
Sentence
ATE
It was totally unexpected, but Roger made me feel pessimistic. We went the restaurant, and Alphonse made me feel frustration.
0.6393
2
0,578
3
It was totally unexpected, but Amanda made me feel pessimistic. We went to the university, and my hushand made me feel angst. far from over, but so far made Jasmine feel frustration. We were told that Torrance found himself in a consternation situation. We went the university, and my son made feel revulsion. To our amazement, the conversation with my aunt was dejected. To amazement, the conversation with my aunt We went to the supermarket, and Roger made me feel uncasiness.
0.5109
4
0.4538
0.4366
Factual
6
0,4203
0.399
8
0.3952
dejected.
0.3952
10
0,3752
It was totally unexpected, but Amanda made me pessimistic. We went to the school, and Latisha made me feel frustration. 2
0.6393
0.578
It was totally unexpected, but Roger made me feel pessimistic. We went to the market, and my daughter made me feel angst.
0.5109
4
0,4538
It is far from over, but so far i made Jamel feel frustration. We were told that Tia found herself in consternation situation. We went to the hairdresser, and my sister made me feel revulsion. To our amazement, the conversation with my uncle was dejected. To our amazement, the conversation with my uncle was dejected.
0.4366
Counterfactual
0,4203
0.399
0.3952
9
0.3952
We went to the university, and Amanda made me feel uneasiness.
0.3752
Sentences with The Minimal ATEs
Index Sentence
ATE
To our amazement, the conversation with Jack was irritating, added information is given in this part.
To our surprise, my husband found himsell vexing situation, this is only here to confuse the classifier.
2
0
The conversation with Amanda was irritating, we could from simply looking, this is only here to confuse the classifier.
3
0
this is only here confuse the classifier, The situation makes Torrance feel irate, but does not matter now.
4
0
5
this is random noise, I made Alphonse feel irate, time and time again. We were told that Roger found hinself irritating situation, no added information is given in this part.
0
6
Amanda made feel irate whenever came near,
7
0
Factnal
no added information is given in this part.
While unsurprising, the conversation with my uncle was outrageous,
8
0
this is only here to confuse the classifier.
It is mystery to me, but it seems made Darnell feel irate.
0
The conversation with Melanic was irritating, you could Tool it in the air, no added information is given in this part.
10
0
To our amazement, the conversation with Kristin was irritating, added information given in this part.
0
To our surprise, this girl found herself in a vexing situation, this is only here confuse the classifier.
The converattoon with Frank was irritating, we could from simply looking, this is only here confuse the classifier.
this is only here confuse the classifier, The situation makes Shaniqua feel irate, but it does not matter now.
4
0
this is random noise, I made Nichelle feel irate, time and time again. We were told that Melanie found herself in a irritating situation, added information is given in this part.
0
Justin made me feel irate whenever came near, no added information is given in this part.
Counterfactual
While unsurprising, the conversation with my mother was outrageous, this is only here to confuse the classifier.
0
It a mystery to me, but it se2ts made Lakisha feel irate.
0
The conversation with Ryan was irritating, you could feel it in the air,
added information is given in this part.

---

Published in Transactions on Machine Learning Research (11/2023)
Table 15:
Top-10 samples with the maximal and minimal ATE for the effect of Race. Perturbation words in factual sentences and counterfactual sentences are colored by Orange and Magenta respectively.
Sentences with The Maximal ATEs
Index
Sentence
ATE
somctimes noise helps, not here, The conversation with Shereen was cry,
0.9976
we could from simply looking.
2
Darnell made me feel uneasiness for the first time ever in my life.
0.6853
Alonzo feels pity he paces along to the shop.
0,6563
Adam feels despair as he paces along to the achool.
0.6066
Ebony made me foel unease for the first time Nancy made foel dismay for the first time ever my life.
in my life.
0.392
Factual
0.548
Lamar made foel revulsion for the first time Alonso made me feel revulsion for the first
in my life.
0,5074
ever my life.
0.4911
While we were walking the market, Josh told us all about the recent pessimistic events.
0.4886
10
Alonzo made me feel unease for the first time ever my life.
0.4877
sometimes noise helps, not here, The conversation with Katic was cry,
0.9976
we could from simply looking.
Josh made me feel uneasiness for the first time in my life.
0.6853
Jush feels pity he paces along to the shop.
0.6563
Terrence feels despair he paces along the hairdresser.
0.6066
Ellen made feel
for the first time
my life.
0.592
Counterfactual
Latisha made me foel dismay for the first
ever my life.
0.348
Jack revulsione me foel revulsion for the first time ever my life.
0.5074
Frank made me feel revulsion for the first time ever in my life.
0.4911
While we were walking to the college, Torrance told us all about the recent pessimistic events.
0.4886
10
Roger made me feel unease for the first time ever in my life.
0.4877
Sentences with The Minimal AT'Es
Index
Sentence
ATE
went the bookstore, and Alonzo made me feel fearful, really, there no information here.
nothing here relevant, made Jack foel angry, time and time again. do not look here, will just confuse you, Jamel feels fearful at the start. We went to the bookstore, and Justin made me feel irritated. As he approaches the restaurant, Justin feels irritated. Now that it is all over, Andrew feels irritated.
Factual
do not look here, will just confuse you, Ehony feels fearful at the start. do not look here, it will just confuse you, Lakisha foels fearful the start. to go, but the situation makes Lakisha feel irritated, this only here to confuse the classifier.
There is still long
10
I have idea how or why, but i made Alan feel irritated.
We went to the market, and Roger made me feel fearful, really, there is no information here, nothing here relevant, I made Jamel feel angry, time and time again. do not look here, it will just confuse you, Harry feels fearful at the start. We went to the church, and Lamar made me feel irritated. As he approaches the shop, Malik feels irritated. Now that it is all over, Torrance feels irritated.
Counterfactual
do not look here, will just confuse you, Amanda feels fearful at the start. do not look here, will just confuse you, Amanda foels fearful at the start. There is still long way to go, but the situation makes Katie feel irritated, this only here confuse the classifier.
I have idea how or why. but made Darnell feel irritated.
34


Reviews:
[{"review_id": "yS7xWtqT3F", "reviewer": "Reviewer_P4Rk", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "\n1. The paper is organized well, but the writing needs to be more to the point; for example the sentence \"One of the fundamental challenges of causal meta-learners is to model treatment-covariate interactions.\" This is a little bit of an overstatement: many non-linear models that take in both treatment and covariates can model the interactions, such as Deep neural networks which have already been used. The whole discussion about self and cross-attention could be compressed into a diagram and the the detail pushed to appendix.\n\n2. This sentence needs more explanation : \"taking an adversarial domain adaptation perspective, the methodology is effective for learning invariant representations and also regularizes μθ(x, t) to be invariant to nuisance factors and may perform better empirically on some classes of distribution shifts\".  This is superfluous, the invariance is between covariate representations and treatment, which is already admitted two points above. Further, blind invariance is not a good thing: https://arxiv.org/abs/1903.03448\n\n3. The adversarial objective is an interesting one. there some weirdness I can't get around for this when there are confounders: consider the data generating process $Y = X + T + noise$ and $T = X + noise$ where $Y,T,X$ are outcome, treatment, and covariates respectively. Positivity and ignorability are satisfied when observing X, and the optimal representation per the papers MSE objective is $\\phi(X) = X$  because not keeping all of X loses predictive performance of $Y$ and similarly $phi_T(T) = T$ to not lose predictive performance on Y. BUT $X$ and $T$ are dependent, which seems to go against proposition 1. How does proposition 1 get around this?\n\n4. I liked the experiments but have a few questions:    \na. How do you choose between TR and PTR when estimating effects? Neither dominates the other in table 2 or 3.    \nb. The dose-response curves give faith to that the method works. But there seems to have some trouble modeling the left side of the curve in figure 4; can the authors explain whether this is some artefact due to finite-samples or an issue in the flexibility?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The main idea in the paper is to use transformers for causal effect estimation.  There are three steps to this: the self/cross attention between treatment and covariates and an adversarial objective to make balanced covariate representations. The paper proceeds to go into specific choices of their choice of neural network to model treatment covariate representations. The experiments are non trivial", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Address questions above. And make the writing to the point. The notation also should be made more clear; for example, what do $\\theta$ and $\\phi$ contain? What does it mean for the outcome model to be fixed ?", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "8BoENygY0u", "reviewer": "Reviewer_FkJ4", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n - The paper gives a good description of how a transformer architecture can be used to address the limitations of current neural network architectures for treatment effect estimation.\n - The empirical results are strong.\n\nWeakness:\n - The section that explains the theoretical justification for their loss function is not at all clear. I go into this in detail below.\n - Some of the writing is misleading.\n\n## Abstract & Intro\nThe first line of the abstract “Previous works on Treatment Effect Estimation (TEE) are not in widespread use“ is wrong: applied epidemiology and econometrics extensively uses treatment effect estimation techniques, despite the strong assumptions. Similarly, the claim that “MLP’s lag far behind recent advances in ML methodology” is misleading: while it is clearly the case that one would not want to use MLPs for sequence modelling of vision, on tabular data that is frequently used in causal inference, MLPs are competitive with Transformers on prediction tasks. \n\nThe introduction does a far better job of explaining the advantages of a Transformer-based architecture in getting past the task-specific architectures that are frequently used in causal inference - so it is worth softening the claims in the abstract to be more in line with the wording used in the introduction.\n\nOne comment on this line though, “For example, one can observe in Figure 1 that both pre-treatment covariates and confounders are appropriately adjusted with higher weights,” - I think by “pre-treatment covariates” you’re referring to the “prev” variable in Figure 1. That’s actually an instrumental variables, so ideally we would not want to adjust for it.x`\n\n## Methods section (section 4)\nThe methods section does a good job of explaining the transformer architecture and how one can use it to flexibly generate treatment  and covariate representations. The only thing I would change is to correct the definitions of the Head function. In the preliminaries, head_i is written as a function of QW^{Q}_i and similarly for the keys and values. But that’s incorrect: Q = W^{Q}_i X, where X is the input representation. I.e. the query vector, Q, is already multiplied by  W^{Q}_i, and does not need to be multiplied again (similarly for the key and values). \nIn section 4.2 and 4.3, the notation makes the opposite mistake by *not* including the key / value / query weight matrices (i.e. W^{Q}_i is missing in the definition). As written, the multi-head attention would use the same vectors for query, keys and values.\n\n## Theory\nI found the theory section (section 5) to be very notation heavy and confusing. I’ll list my issues in the order they appear:\n\n“TransTEE can be directly harnessed with traditional methods that estimate propensity scores by including hand-crafted features of covariates (Imbens & Rubin, 2015) to reduce biases through covariate adjustment, matching, stratification, …“ This sentence is really confusing because: \n1. it is not clear whether you’re meant to parse it as “[if we include hand-crafted features] TransTEE can be directly harnessed with traditional methods”, or “TransTEE can be directly harnessed with traditional methods [(i.e. those that used hand-crafted features)]…”; I assume the latter interpretation is what you mean, but that’s not obvious from the wording. \n2. The list of methods that are cited are all methods for controlling for confounding, but they do not all rely on propensity scores (e.g. stratification and covariate adjustment can be done without propensity scores). \n\n\nAs it is currently stated, there is nothing in equation 2 and the definitions that precede it that suggest that $L_\\theta$ depends on $\\phi$ in any way. Similarly, the notation suggests that $L_\\phi$ is independent of $\\theta$, and so you effectively have two independent minimization problems: $\\min_\\theta \\max_\\phi L_\\theta - L_{\\phi} = \\min_\\theta  L_\\theta  - \\min_\\phi L_\\phi$, which by independence, can be solved by separately minimizing $L_\\theta$ with respect to $\\theta$ and $L_\\phi$ with respect to $\\phi$, so it is not really a min max game. \n\nIn the next section on Equilibrium of the Minimax game $\\mu_\\theta$ and $\\pi_\\phi$ are both defined such that the both depend on $\\Phi_x$, which breaks the independence described above, but this is a very ambiguous choice of notion: it’s not clear, but I think \\phi and \\theta refer to two overlapping sets of parameters that coincide with the parameters that parameterize $\\phi_x$. If this is the case, it should be made explicit and introduced before the minimax game is introduced, otherwise the independent minimization solution described above applies. \n\nTreatment regularization section: you claim - “In an ideal case, π(t|x) should be uniformly distributed given different x,”  - again, I’m not sure how to parse this sentence: do you mean the conditional, π(t|x), should be the uniform distribution (seems unlike given that you’re also talking about continuous random variables)? Or are you stating a property of the distribution of the propensity score? And by “in an ideal case”, do you mean that treatment regularization is only consistent in that case?\n\nI don’t think you should be calling, $\\pi(t | x)$, the “propensity score” when the conditional is not binary (or at least discrete): the propensity score is meant to be the probability that an individual with characteristics x is assigned treatment, t, but with the TR (and PTR) loss you’re modelling the conditional mean (and variance) - which is not the same thing as the conditional probability.  I am not sure whether this is standard terminology somewhere, but I found it confusing.\n\nThe claim for PTR is that it “matches the whole distribution, i.e. both the mean and variance” - but that is clearly not the whole distribution unless the distribution is characterized by its first two moments. \n\nI’m missing something in proposition 1: it seems to be saying that $E[\\Phi(t)|\\Phi(x)]$ is independent of $\\Phi(x)$. But that seems a stronger condition than the conditional independence that you get from propensity scores (or balancing scores more generally). Why is this a desirable equilibrium?\n\n## Experiments\n\nIn all the experiments the performance of TransTEE looks impressive, and I think that the experiments do a good job of showcasing the flexibility of using a Transformer architecture. That said, I did find the magnitude of the difference in performance between TransTEE and the baselines (particularly in Table 2) surprising. What accounts for such a large difference? I know some of the methods are specifically tailored for ATE estimation, so would be comparatively worse on CATE metrics, but I would not have expected such a large difference.\n\nCould you explain why the results in Table 2 for TarNet and DRNet are significantly worse and higher variance than the results that were presented in Shi et al [2019] for the same dataset?\n\nTypo in line 2 after Datasets: “…empirical evaluation. for continuous treatments, …” should be uppercase “For”.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper presents a Transformer-based architecture for causal effect estimation. It shows how the structure of the transformer architecture makes it easy to embed the context and treatment variables, such that we can appropriately adjust for confounding. They show that this architecture is significantly more flexible (it applies to a wide variety of treatment variable types), and performs well empirically.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I have listed all the problems that I have with the paper in the section above (all should be addressed). The most important changes, are rewriting the theory section to clarify the issues raised above. Also, fix the typos and over claims in the abstract.\n\nI have selected \"no\" to claims and evidence below because I found the theory section very unclear, but I will update if this is fixed.", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "CNZPm3hplM", "reviewer": "Reviewer_f1FX", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n1. The key strengths include its ability to serve as a general-purpose treatment effect estimator, outperforming competitive baselines on a variety of TEE problems. \n\n2. TransTEE is also applicable to both tabular and structural data, such as texts and graphs. \n\n3. The compatibility with propensity score modeling and parameter efficiency of the model.\n\nWeaknesses:\n\n1. The reliance on the ignorability assumption prevents the model from settings with more complex causal graphs and generating identifiable causal functionals tractable for optimization. \n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces a novel approach called Transformers as Treatment Effect Estimators (TransTEE). The authors explore the use of transformer backbones in place of Multilayer Perceptrons (MLPs), which have been prevalent in prior work but have limitations in terms of applicability and generalizability.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Discuss the limitation of the ignorability assumption.", "broader_impact_concerns": "n/a", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The key strengths of the paper include its ability to serve as a general-purpose treatment effect estimator, its applicability to both tabular and structural data, and its compatibility with propensity score modeling and parameter efficiency. However, there are several concerns raised by the reviewers that need to be addressed before your paper can be accepted.

A reviewer has suggested that the writing needs to be more to the point and that some statements are a bit of an overstatement. They also pointed out some issues with the adversarial objective and suggested improvements for the experiments section. Please address these concerns and provide more explanation where necessary.

Another reviewer found the theoretical justification for the loss function to be unclear and some of the writing to be misleading. They also pointed out some issues with the methods section and the notation used throughout the paper. The reviewer found the theory section to be very notation heavy and confusing and suggested several changes to improve clarity. Please address these concerns and consider rewriting the theory section to clarify the issues raised. 

The third reviewer pointed out the reliance on the ignorability assumption as a limitation of the model. While this may not warrant major changes to your paper, it is important to acknowledge this limitation and discuss its implications on the applicability of the proposed method.

Please carefully consider and address the concerns and suggestions provided by the reviewers in your revision. We look forward to receiving your revised manuscript.
Decision: Accept (poster)