Exploring Transformer Backbones for Heterogeneous Treatment Effect Estimation
Authors: YiFan Zhang
Hanlin Zhang
Zachary Chase Lipton
Li Erran Li
Eric Xing
Venue: Submitted to ICLR 2023 (2023)
Keywords: Causal Inference
Continuous Treatment Effect
Heterogeneous Treatment Effect

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. 

Under review as a conference paper at ICLR 2023
EXPLORING TRANSFORMER BACKBONES FOR HETEROGENEOUS TREATMENT EFFECT ESTIMATION
Anonymous authors Paper under double-blind review
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 MI. methodology, which limits their applicability and generaliz- ability. To extend beyond the single domain formulation and towards more realistic learning scenarios, we explore model design spaces beyond MLPs, i.c., transformer backbones, which provide flexibility where attention layers govern interactions among treatments and covariates to exploit structural similarities of potential out- comes for confounding control. Through careful model design, Transformers as Treatment Effect Estimators (TransTEE) is proposed. We show empirically that TransTEE can: (1) serve as general-purpose treatment effect estimator which significantly outperforms competitive baselines on a variety of challenging TEE problems (c.g., discrete, continuous, structured, or dosage-associaae and is applicable to both when covariates tabular and when they consist of struc- tural data (e.g., texts, graphs); (2) yield multiple advantages: compatibility with propensity score modeling, parameter efficiency, robustness to continuous treat- ment value distribution shifts, explainable in covariate adjustment, and real-world utility in auditing pre-trained language models.
treatments.)
INTRODUCTION
One of the fundamental tasks in causal inference is to estimate treatment effects given covariates, treatments and outcomes. Treatment effect estimation is central problem of interest in clinical healthcare and social science (Imbens& & Rubin 2015). as well as econometries (Wooldrige) (013)). Under certain conditions (Rosenbaum & Rubin 1983), the task can be framed as a particular type of missing data problem, whose 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., 20199] 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. 20206b Schwab et al.,
2020 Nie al. 2021; Curth & van der Schaar 2021b). part due to their flexibility in modeling nonlinear functions (Hornik et al.] 1989) and high-dimensional input (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 (2221b) 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 estimation. Most current works based on subnetworks do not sufficiently exploit the structural similarities of potential outcomes for heterogeneous TEE and accounting for them needs complicated regularizations, reparametrization or multi-task architectures that are -opeci--pecciicc (Curth & van der Schaar, 2021b). Moreover, they heavily rely on their treatment- -raatenn-specific 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
For example, ([Y(1) X] is often of much simpler form to estimate than either TE(((1) X] (0) due inherent similarities between (11) and (0).

---

Under review as a conference paper at ICLR 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. Trans is general for all the factors.
METrinnS
DISERATE
CONTINUOUS TREATMENT
TREATMENT
TARNET (SHALIT BT AL.
2017
O(n)
PERTECT MATCH Sc TIAAA
2018
O(1)
O(25)
DIAGONNEI
O(n)
O(n)
O[T B(T
SCIGAN BRCK A..
O(1)
VCNEI Nit BJ AL.
2021
O(1)
ot1
NCORE (PARSHOO 2021
O(n)
O(Br)
O(n)
FLEXTENET (CURTH SCRAAK
O(n)
OURS
((1)
o(1)
0(1)
0())
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 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 estimate treatment effects.
In this work, we explore recent advanced models in the deep learning community boost the model design for TEE tasks. Specifically, the core idea of our approach consists of three parts: as an S- learner, TransTEE embeds all treatments and co- variates, which avoids multi-task architecture and shows improved flexibility and robustness to con- tinuous treatment value distribution shifts; attention mechanisms are used for modeling treatment inter- action and treatment-covariate interaction. In this way, TransTEE enables adaptive covariate selection (De Luna et al. 2011 VanderWeele] 2019) for infer- ring causal effects. For example, can observe in Figure D 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 causal graph is not available. This recipe also gives improved versatility when working with heterogeneous treatments types (Figure 2)
PL00
Trotssal
Figure I: A motivating example with cor- responding causal graph. Prev denotes previ- ous infection condition and BP denotes blood pressure. TransTEE adjusts an appropriate co- variate set (Prev, BP) with attention which is visualized via a heatmap.
Our first contribution shows that transformer backbones, equipped with proper design choices, can be effective and versatile treatment effect estimators under the RuyinnNeyman potential outcomes framework. TransTEE is empirically verified to be (i) a flexible framework applicable for wide range of TEE settings; (ii) compatible and effective with propensity score modeling; (iii) parameter- efficient; (iv) explainable in covariate adjustment; (v) robust under continuous treatment shifts; (vi) useful for debugging pre-trained language models (LMs) to promote favorable social outcomes.
Moreover, comprehensive experiments six benchmarks with four types of treatments conducted to verify the effectiveness Trans'TEE in estimating treatment effects. We show that Trans' TEE pro- duces covariate adjustment interpretation and significant performance gains given discrete, continuous or structured treatments on popular benchmarks including IHDP, News, TCGA. We introduce a new surrogate modeling task to broaden the scope of TEE beyond semi-synthetic evaluation and show that TransTEE is effective in real-world applications like auditing fair predictions of LMs.
RELATED WORK
Neural Treatment Effect Estimation. There many recent works adapting neural networks to learn counterfactual representations for treatment effect estimation (Johansoon et al.) 2016 Shalit
al. 2017 Louizos al. 2017 Yoon al. 2018 Bica et al. 2020b Schwab et al. 2020 Nie et al.
2

---

Under review as a conference paper at ICLR 2023
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 al. 2020) (converges to JSD between groups), counterfactual variance (Zhang et al.] 2020). However, their domain- specific designs make them limited to different treatments shown in Table methods like VCNet (Nic et al.) 2021) use a hand-crafted way to map a real-value treatment to an n-dimension vector with a constant mapping function, which is hard converge under shifts of treatments (Table O in Appendix); ); models like TARNet (Shalit et al. 2017) need an accurate estimation of the value interval of treatments. Moreover, previous estimators embed covariates to only one representation space by fully connected layers, tending to lose their connection and interactions (Shalit et al. 2017, 2017 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 extrapolation capacity and infeasible given observational data.
2021:
Transformers and Attention Mechanisms Transformers (Vaswani et al. 2017) have demonstrated broad range of language tasks and their variants have been success- fully adapted to representation learning over images (Dosovitskiy et al. (02T), programming lan- guages (Chen al. 2021), and graphs (Ying al. 2021) partly due to their flexibility and expressive- ness. Their wide utility has motivated a line of work for general-purpose neural architectures (Jaegle) 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.e. supervised learning. And one of our goals is to explore the generalizability of attention- based models for TEE across domains with nigh-dimensional inputs, an important desideratum in causal representation learning (Schölkopf et al.] (202T). There are recent attempts to use attention mechanisms for TEE Tasks (Guo et al.) [202 Xu et al.) 2022). CETransformer (Guo et al.) 2 uses embeds covariates for different treatments as a T-learner, They only trivially learn covariate embeddings but not treatment 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 et al., 2019; Curth & van der Schaar 2021b a). ANU (Xu et al., (022) utilizes attention mechanisms to map the original covariate space X into a latent space Z with single model. We detail the difference in Appendix
exemplary performance
3 PROBLEM STATEMENT AND ASSUMPTIONS
Treatment Effect Estimation. We consider setting in which we given N observed samples N (x;, f4. 85. f) each containing N pre-treatmen covariates {x; E R The treatment variable in this work has various support, e.g., f0, 1} for binary treatment settings, IR for continuous treatment settings, and graphs/words for structured treatment settings. For each sample, the potential outcome (µ-mmddll μ(x, L) or u(xx L, 8) is the response of the i-th sample to treatment where in some cases each treatment will be associated with dosage St, E IR. The propensity score (m-molel) the conditional probability of treatment assignment given the observed covariates π 7 = t| x). The above two models can be parameterized as μg and Te. respectively. The task is to estimate the Average Dose Response Function (ADRF): (xx, t) = E x, do(T = t)] (hhocchet (2066). which includes special cases in discrete treatment scenarios that can also be estimated as the average treatment effect (ATE): ATE = E[[(x, 1) - p(x, 0)] and its individual version ITE.
fi=t-
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. AgporrabiityyUncooouuuddeee I) In the binary treatment case, (0), Y (1)
implies no hidden confounders such that Y(T- =
Assumption 3.2. (Positivity/Overlaap) The treatment assignment is non-deterministic such that, i.e.
0 < <()x) < 1, Vx E Xtt T
Assumption 3.J ensures the causal effect is identifiable, implying that treatment is assigned indepen- dent of the potential outcome and randomly for every subject regardless of its covariates, which allows estimating ADRF using (tt)= = E[y|do[? 1)] [ Y |x, T t]] (Rubin) 1978). One naive
3

---

Under review as a conference paper at ICLR 2023
Forward
Attentton
000 A00
D Hetix)
Fo(x)
Ha.4al
R(c(x)
Deagonbict (discrete)
Flet (discrese)
DRNet (contimous)
VCNE (continuoss)
Trans TEB
Figure 2: A schematic comparison of TEE and recent works including Dragon Shi et al. 2019),
& van der Schaar (02Tb), SRhhr(Schwab et al. 2020 and VCNet(Nie et al. 2021). TransTEE handles all the scenarios without handcrafting treatfent-speciffc architectures and any additional parameter overhead.
estimator of μ(x, t)
E| x, t] the sample average μut) Ei= ((x,, Assumption 3.2 states that there is a chance of seeing units in every treated group.
TRANSTEE: TRANSFORMERS AS TREATMENT EFFECT ESTIMATORS
The systematic similarity of potential outcomes of different treatment groups important for TEE (Curth & van der Schaar 2021b). Note that x is often high-dimenn sigh-dimensional while is not, which means naively feeding (x, 1) to MLPs is not favorable since the impacts of treatment tend to lost. As 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 2 and Figure To remedy it, use three simple yet effective design choices based on attention mechanisms. The resulting scalable framework Trans TEE can tackle the problems of most existing treatment effect estimators (e.g., multiple/continuuussstttuuun interaction, and treatments with dosage) without ad-hoe architectural designs, e.g., multiple branches.
treatments, treatments
Preliminary. The main module in TransTEE is the attention layer Vaswani et al. (2017): given
d-dimensional query, key, and value matrices Q E RRdxd K E Raxd V E Raxd. attention mechanism computes the outputs as H(Q, K, V) softmax( ydz attention preferable to jointly attend to the information from different representation subspaces.
QK V. In practice, multi-head
HM M (Q,K, V)= Concat(head,, ..., heahhad)))) ) where head, H(QW? KW,K, vWY), where w E RRdxds wY Rddxde W.Y. E RRdxdo and W Rdd are learnable matrices.
4.1 COVARIATE AND TREATMENT EMBEDDING LAYERS
Treatment Embedding Layer. As illustrated in Figure 2 and Table. [ treatments often of much lower dimension compared to covariates, to avoid missing the impacts of treatments, previous works (e.g., DragonNet (Shi et al., 2019), FlexTENet (Curth & van der Schaar 2021b), DRNet (Schwab et al. 2020)) assign covariates from different treatment groups to different branches, which is highly parameter inefficiens. Besides, We analyze in Proposition (Appendix DD that, for continuous treatments/dosages, the performance is affected by both number of branches and the value interval of treatment. However, almost all previous works on continuous treatment/dosage assume the treatment or dosage is in fixed value interval e.g., [0, 1] and Figure 3 shows that prevalent works fail when tested under shiftts 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, = Lineaar()), M,- = Linear(s), where M, E Rd. M, E TRd exists just when each treatment has dosage parameter, otherwise, only treatment embedding is needed. When multiple (n)) treatments act simultancously, the projected matrix will be M, I dg .M, E _gdx 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, Trans TEE is shown to be (i) robust under treatment shiftss, and (ii) marameter-effciien
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, and is hard to study
4

---

Under review as a conference paper at ICLR 2023
the function of an individual covariate in sample, TransTHE learns different embeddings for each covariate, namely Mg Linear(x), and Mg E Rd where p is the number of covariate. Covariates embedding enables to study the effect of individual covariate on the outcome.
4.2 COVARIATE AND TREATMENT SELF-ATTENTION
For covariates, prevalent methods represent covariates as a whole feature using MLPs, where pair- wise covariate interactions are lost when adjusting covariates. Therefore, cannot study the effect of each covariate on the estimated result. In contrast, TransTEE processes each covariate embedding independently and model their interactions by self-attention layers. Namely,
M HMM,,, (M, _
+ M M MLP(BN(AM)) + My
where M is the output of i layer and BN is the BatchNorm layer. Simultaneouslyy the treatments and dosages embeddings are concatenated and projected to the latent dimension by linear layer, which generates a new embedding Mst E IRd. Then self-attention is applied
My. HM (M .M M M Ma MLP(BN(M(,) Mr
The self-attention layer for treatments enables treatment interactions, an important desideratum for S- and T-leamners Namely, TrEn can model the scenario where multiple treatments are applied and attain strong practical utillty, e.g.. multiple prescriptions in healthcare or 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.
4.3 TREATMENT- COVARIATE CROSS-ATTENTE
One of the fundamental challenges of causal meta-learners is to model treatment-covariatt interactions. TransTEE realizes this by a cross-attention module, treating Mat as query and M2 as key and value
A Har (M M .M- M
M = MLP( + MMt = (LP(Pooling(AM
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 for estimating treatment effects (VenderWeele, 2019), which empirically yields suitable covariate adjustment principles (the Disjunctive Cause Criteria) (De Luma et al. 2017 - VanderWeels (209)) about pre-treatment
covariates and confounders as intuitively illustrated Figure [ corroborated our experiments.
Denote ÿ = (x(x t) and the training objective is the mean square error (MSE) of the outcome regression: L3 t) Ei= (1 px(x,, +,
Remark. We include an illustration of TransTEE by a concrete example in Appendix B Note that, although the embedding technique and attention mechanisms are commonly used in Computer Vision, Neural Language Processing communities, it not well understood how to guide the design of these modules for causal inference and why these techniques benefit TEE tasks are underexplored. In this work, through the flexible use of embedding and attention mechanisms we design a strong TEE architecture, we further use conceptual analysis and empirical results show the benefit brought by the used design choices. Besides, when combined with the strong modeling capacity of Transformers, TransTEE can be extended to high-dimensional data flexibly and effectively on structured data. The generalizability of the TransTEE also allows new applications like auditing language models beyond semi-synthetic settings as shown in the next section.
EXPERIMENTAL RESULTS
We elaborate on basic experimental settings, results, analysis, and empirical studies in this section. See Appendix E for full details of all experimental settings and detailed definitions of metrics. See Appendix for many more results and remarks.
5.1 EXPERIMENTAL. SETTINGS
Datasets. Since the true counterfactual outcome (or ADRF) are rarely available for real-world data, synthetic semi-synthetic data for empirical evaluation. for continuous treatments,
5

---

Under review as a conference paper at ICLR 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- 2), models are trained with E 10.1, 2.0] and tested in t E [0, 2.0). For Extrapolation (h- 5), models are trained with / [0.25, 5.0] and tested in / E [0, 5).
METHODS
VANILLA (BINARY)
VANILLA (R: = 1)
EXTRAPOLATION (A- 2)
VANILLA (A 5)
EXTRAPOLATION (A = 5)
TARNEI
0.3670 0.61112
2.0152 1,07449
12,967 1.78108
5,6752 0.53161
31.523 1,5013
DRNET
0.3543 8.60622
2,1549 1.04483
11.071 = 0.99384
3,2779 0.42797
31,524 1.50244
PLEXTENDT
0.2700 = 0. 10.00000
VCNEr
0.2998
0.182236
0.7800 0.61483
NAN
NAS
SAN
TRANSTEE
0,0983
0.15384
0.1151 0.10289
0.2745 0,14976
0.1621 0.14443
0.2066 #: 0.23258
TRANSTEE-MLE
0.1721
0.49061
0.0877 0.03352
0.2685 0.17552
0. 0.2079
0.1476 0.07123
TRANSTEETTR
0.1913
= 0.29953
0.0781 0.05243
0.2395 0.08154
0.1143 +: 0.03324
0.0947 + 0,0824
PRASSTEE+PIR
0.2193
0.34667
0,0762 0.07915
0,2352 0.17095
0.1363 4: 0.08026
0.1263 4: 0.08035
we use one synthetic dataset and two semi-synthetic datasets: the IHDP and News datasets. For treatment with continuous dosages, we obtain covariates from a real dataset TCGA (Chang et al.) 2013) and generate treatments. where each treatment is accompanied by a dosage. The resulting dataset named TCGA (D). Following (Kaddour et al.) 202T), datasets for structured treatments include SmalllWorlld (SW), which contains 1, 000 uniformly sampled covariates and 200 randomly generated Watts-Strogatz small-world graphs (Watts & Strogatz, 1998) as treatments, and TCGA (S), which uses 9, 659 gene expression of cancer patients (Chang et al. 2013) for covariates and 10, 000 molecules from the QM9 dataset (Ramakrishnan et al. (0144) as treatments. For the study on language models, we use the Enriched Equity Evaluation Corpus (EEEC) (Feder et al. 202I).
Baselines. Baselines for continuous and binary treatments include TARnet (Shalit et al., 2017). Dragomnet (Shi et al. 2019) DRNet (Schwab et al.) 2020), FlexTENet (Curth & van der Schaar 2021b), and VCNet (Ne et al. 2021). SCIGAN (Bica et al. (020b) is chosen as the baseline for continuous dosages. Besides, we revise DRNet (Schwab et al. 2020), TARNet (Shalit et al. 2017 and VCNet (Nie et al., (021) to DRNet (D), TARNet (D), VCNet (D), respectively, which enable multiple treatments and dosages. Specifically, DRNet (D) has T main flows, each corresponding to a treatment and is divided into Bp branches for continuous dosage. Baselines for structured treatments include Zero (Kaddour al., 202T). GNN (Kaddour et al., 202T), GraphITE (Harada & Kashima 2021), and SIN (Kaddour et al., 202I). To compare the performance of different frameworks fairly, all of the models regress on the outcome with empirical samples without any regularization. For MLE training of the propensity score model, the objective is the negative log-
likelihood: Ls == - Em= log TT (x.).
Evaluation Metric. For continuous and binary treatments, we use the average mean squared error on the test set. For structured treatments, following (Kaddour et al., 202I), we rank all treatments selected and the treatment
by their propensity =(([x) in descending order. Top K treatments effect of each treatment pair is evaluated by unweighted/weightee expected Precision in Estimation of Heterogeneous Effect (PEHE) (Kaddour al. 20211, where the WPEHE @1 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
5.2 CASE STUDY AND NUMERICAL RESULTS
Case study on treatment distribution shifts We start by conducting case study on treatment distribution shifts (Figure 3 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 constraining the values of treatments Nie et al. (2021) and dosages (Schwab et al. (2000) to a specific range, our methods perform surprisingly well when extrapolating beyond these ranges assessed several benchmarks. By comparison, other methods appear comparatively brittle in these same settings. See Appendix for detailed discussion.
Case study of propensity modeling. TransTEE is conceptually simple and effective. How- ever, 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

---

Under review as a conference paper at ICLR 2023
0.5
0.5
0.0
0.0
Responte
-0.5
-0.5
Trath
Truth
Tarnet
Tamet
-1.0
Dmet
Dmet
Dinet
-1.
Venel
Venet
-1.5
Venel
TransTll
TresTEE
TransTC
-2.0
0.2
0.4
1.0
0.0
0.4
08
1.2
1.6
2.0
Theatment
(a) h 1 in training and testing. (h) h = 1.75 in training and h = (c) h = in training and testing.
2 in testing (extrapolation).
Figure 3: Estimated ADRF on the synthetic dataset, where treatments are sampled from an interval [t, where = 0.
are discrete (Bica et al. 2020a Kallus 2020: Du et al. 2021 Thus we propose two regulariza- tion variants for continuous praament/dosageen which are termed Treatment Regularization (TR, c3 tx Ei= (t: (f, ))3 and its probabilistic version Probabilistic Treatment Regular- ization (PTR, LPTR = Ei= (( (x. + & log n (o²)x:) respectively. The overall model is trained in a adversarial pattern, namely ming max Lx y, t) (x t). Specifically, propensity score model Ta (t(x) parameterized by MLP is learned by minimizing (x, t), and then the outcome estimators Ho (x, t) trained by ming (x(x, y, t) (x, To overcome selection biases over-representation 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. (021). Such a recipe can account for selection bias where π (f)xx) * p(t) and leave spurious correlations out, which can also be more robust under model misspecification especially the settings that require extrapolation on treatment (See Table 2 and Appendix
C for concrete formalisms and discussions.).
As in Table [] Appendix Table 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 reduced. Overall, TR is better in the continuous case with smaller treatment distribution shifts, while PTR 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-synthetic 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 in treatment effect estimation problems of interest. Moreover, covariate selection visualization in TR and PTR (Figure 4(a). Table Jable 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 F where TR and PTR are further shown able to improve the convergence of TransTEE.
Continuous treatments. To evaluate the efficiency with which Trans TEE estimates the average dose- response curve (ADRF), we compare against other recent NN-based methods (Tables p. Comparing results in each column, we observe performance boosts for TransTEE. Further, TransTEE attains much smaller error than baselines in cases where the treatment interval is restricted to [0. 1] (e.g., 1. E [0, 5)) and when the training and test treatment intervals are different (extrapolation). Interestingly, even vanilla TransTEF produces competitive performance compared with that of π (1[x) trained additionally using MLE, demonstrating the ability of Trans TEE to effectively model treatments and covariates. The estimated ADRF curves on the IHDP and News datasets are shown in Figure and Figure i Appendix. TARNet and DRNet produce discontinuous ADRF estimators and VCNet only performs well when E [o, 1]. However, TransTEE attains lower estimation error and preserves the continuity of ADRF different treatment intervals.
Continuous dosage. Table 5 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
7

---

Under review as a conference paper at ICLR 2023
Table 3: Effect of Gender (top) and Race (bottom) on POMS classification with the EEEC dataset, where ATEGT is the ground truth ATE based on 3 repeats with confidence intervals [CI] constructed using standard deviations.
Conrchation/ Reprosenaation Based Baselines
Treatment Effect Estimaturs
TC
ATBey
TReATE
CONEXP
INLP
TarNet
DRNe
VCNei
Trans
Gender
0.086
0.125
0.02
0.313
0,0067
0.0088
0,0085
0,013
ICT
10.082,0.09]
10,010,1114
[0.0.0.05]
[0.304,03321]
[0,0449, 0.0076)
[0.008440.009]
|0.0036, 0.011111
0.0168]
Race
0.014
0.046
0.08
0.591
0.005
0.006
0.003
0.0174
(CI)
[0.02200.06]]
[0.038.0.054]
[0.02.0014]
[0.578,0.005]
[0,0211, 0.0069)
[0.0047, 0.0081)
[0,0025, 0.0037]
[0,0113, 0.0288)
Madd
goan
TWSITE
1901
NH
SCGAS
AD 13 10:00 23
(b)
(C)
Figure 4: (a) The learned weights of the cross-attention module on IHDP dataset. TransTEE adjusts confounders Scon {1, 3,5, 6} properly with higher weights during the eross attention process. (b) AMSE attained by models on IHDP with different numbers of noisy covariates. (c) Number of parameters for different models on four different datasets, where the log the y-axis base 2.
increases, TEE 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 8. where the estimated ADRF curve of each treatment considering continuous dosages is plotted. Compared to baselines, TransTEE 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, TransTED consistently performs the best.
Structured treatments. We compared the performance of TransTEEE to baselines on the training and test set of both SW and TCGA datasets 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 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 nigh-dimensional treatments. SIN is the best model among these baselines. However, when the bias equal to 0.1, SIN fails to attain estimation results better than the Zero baseline. To evaluate each model's 's robustness to varying levels of selection bias, performance curve with K E [0, 40] for the SW dataset and λ E [0, 0.5] for the TCGA dataset are shown in Figure 14 and Figure [5 in Appendix. Considering both metrics, TransTEE outperforms baselines by a large margin across the entire range of evaluated treatment selection biases.
5.3
ANALYSIS
Analysis of covariate adjustment of cross-attention module. TransTEEE embeds each covariate independently and then make treatments select proper covariates for prediction by cross-attention. 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-atte artss--atenntin weights) in Figure|4(a) As elaborated in Appendix [E.3 the IHDP dataset has 25 covariates, which is divided into 3 groups: Scon (1. 3, 5, 6). Sair,1 = {4, ~ 15]. and Sahee2 = {16 ~ 25). Scon influences both T and Y, Saa..1 influences only Y, and Sais,1 influences only T. Covariates in Sais, 1 are named noisy covariates since they have no correlation with the treatment. Their causal relationships are illustrated in Figure Interestingly, confounders Scon assigned higher weights while noisy covariates (those influence the outcome but are irrelevant to the treatment) lower Sais,1, which matches the principles in (VanderWeele) 2019) and corroborate the ability of TransTER estimate treatment effects in
Scau
Saua
Figure 5: The causal graph of IHDP dataset.
8

---

Under review as a conference paper at ICLR 2023
complex datasets by controlling both pre-treatment variables and confounders properly. Moreover, Figure [(bb) shows that TransTHE consistently outperforms baselines across different numbers of noisy covanates.
We further conduct 10 repetitions for TransTEE and its TR and PTR counterparts as reported in Table 4 (Appendix Figure [T] visualizes their cross-attention weights). Denote Weon, W1, W2 as the summation of weights assigned to Scon, 1 S 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 TransTHE with propensity score modeling since both TR and PTR improve confounding control. Moreover, TR better than PTR since it also reduces w22 by a larger margin. This observa- tion gives suggestion that we should systematically probe TR and PTR besides comparing their numerical performance, especially in settings where the unconfoundedness assumption is violated (Ding et al. 2017) and controlling instrumental variables will ineur biases in TEE.
Table 4: Attention weights for Scon. Sdis, and Sas.22 respec- tively.
Weon
TransTER
0.27
0.37
0.36
+TR
0.59
0.20
0.21
+PTR
0.32
0.33
0.35
Amount of model parameters comparison. The experiment is to corroborate the conceptual comparison in Table I We find that the proposed TransTEE has consistently fewer parameters than baselines on all the settings as shown in Figure [[c] Besides, increasing the number of treatments allows more accurate approximation for continuous treatments/dosages, most of these baselines need to increase branches which incurs parameter redundancy, However, Trans is much more efficient.
5.4 EMPIRICAL STUDY ON PRE-TRAINED LANGUAGE MODELS
To evaluate the real-world utility of TransTEE, in this subsection, we demonstrate an initial attempt for auditing and debugging large pre-trained language models, an important use case NLP that beyond semi-synthetic settings and under-explored in the causal inference literature. Specifically, we use TransTEE to estimate the treatment effects for detecting the effects of domain domain-specific factors of variation (such as the change of subject's is attributes in a sentence) on the predictions of pre-trained language models. We experiment with BERT (Kenton & Toutanoval 2019) (e.g.. racial and gender-related nouns) over natural language on the (real) EEEC dataset. We use both the pnrrelaaiionrepeeenntttoonnnnn
baselines introduced in (Feder et al. 2021) and implement treatment effect estimators (e.g., TARnet, DRNet, VCNet, and the proposed TransTEE).
Interestingly, results in Table 3 show that TransTEE effectively estimates the treatment effects of domain-specific variation perturbations even without substantive downstream fine-tuning on specialized datasets. TransTEE outperforms baselines adapted from MLP. Moreover, we showcase the top-k samples with the maximal/minimal ITE and analysis in Appendix F.3 The results show that TransTEE has the potential to provide estimators for practical use cases in predicting model predictions (Ilyas 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 al. 2022) and TransTEE can estimate ATE enforce invariant or fairness constraints for LMs
(Veitch et al. 2021) in lightweight and efficient manner, which we leave for future work.
6 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 needed. Moreover, we hope that our findings can lay the groundwork for future work in developing advanced machine learning techniques like pre- pre-training in large-scale observational data in estimating treatment effects, where TransTEE can serve as an effective backbone. Similar almost all the causal inference methods on observational data, one potential limitation of Trans THE is the reliance on the ignorability assumption. Therefore, one important future direction is extending TransTEE 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) 22009) how to integrate TransTEE with domain knowledge (Imbens & Rubin 2015) for alleviating its potential negative societal impacts in consequential decision making will also be important.

---

Under review as a conference paper at ICLR 2023
REFERENCES
Alberto Abadic and Guido W Imbens. Matching on the estimated propensity score. Econometrica,
2016.
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 on nlp model behavior. arXiv preprint arXiv:2205. /4/422205./44,, 2022.
Ahmed Alaa and Mihaela Schaar. Limits of estimating hetcrogeneous treatment effects: Guidelines for practical algorithm design. In ICMI., 2018.
Ahmed M Alaa and Mihaela van der Schaar. Bayesian inference of individualized treatment effects using multi-task gaussian processes. In NeurlPS, 2017.
Ahmed M Alaa, Michael Weisz, and Mihaela Van Der Schaar. Deep counterfactual networks with dropenst arXiv, 2017.
Peter C Austin. An introduction propensity score methods for reducing the effects of confounding in observational studies. Multivariate behavioral research, 2011.
Shai Ben-David, John Blitzer, Koby Crammer, Alex Kulesza, Fernando Pereira, and Jennifer Wortman Vaughan. A theory of learning from different domains. Machine learning, 2010.
loana Bica, Ahmed M Alaa, James Jordon, and Mihaela van der Schaar. Estimating counterfactual treatment outcomes time through adversarially balanced representations. arXiv, 2020a.
Ioana Bica, James Jordon, and Mihacla van der Schaar. Estimating the effects of ciciinuous-vllleee interventions using generative adversarial networks. 2020b.
Kyle Chang, Chad 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 Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, et al. Evaluating large language models trained on code. arXiv, 2021.
Jonathan Crabbé, Alicia Curth, Ioana Bica, and Mihaela van der Schaar. Benchmarking heterogeneous treatment effect models through the lens of interpretability. arXiv preprini arXiv:2206.08366. 2022.
Alicia Curth and Mihaela van der Schaar. Nonparametric estimation of heterogeneous treatment effects: From theory to learning algorithms. In AISTATS, 2021a.
Alicia Curth and Mihaela van der Schaar. On inductive biases for heterogeneous treatment effect estimation. In NeurlPS, 2021b.
Alicia Curth, David Svensson, Jim Weatherall, and Mihaela van der Schaar. Really doing great at estimating CATE? critical look ML benchmarking practices in treatment effect estimation. In NeurlPS, 2021.
Ralph B D' D'Agotinno. Propensity score methods for bias reduction in the comparison of a treatment to condranddeeeed control group. Statistics in medicine, 1998.
Xavier De Luna, Ingeborg Waernbaum, and Thomas S Richardson. Covariate selection for the nonparametric estimation of 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, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x 16 words: Transformers for image recognition scale. 2021.
10

---

Under review as a conference paper at ICLR 2023
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.
Amir Feder, Nadav Oved, Uri Shalit, and Roi Reichart. Causalm: Causal model explanation through counterfactual language models. Computational Linguistics, 2021.
Dylan J Foster and Vasilis Syrgkanis. Orthogonal statistical learning.
arXiv preprint
arXiv:1901.00030 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 boundaries via contrast sets. EMNLP Findings, 2020.
Zhenyu Guo, Shuai Zheng, Zhizhe Liu, Kun Yan, and Zhenfeng Zhu. Cetransformer: Casual effect estimation via transformer based representation learning. PRCV, 2021.
Shonosuke Harada and Hisashi Kashima. Graphite: Estimating individual effects graph-structured treatments. In CIKM, 2021.
Jennifer L Hill. Bayesian nonparametric 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. Data- models: 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 nlp systems: How should we define and evaluate faithfulness? arXiv preprint arXiv:2004.053: 03665, 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- Saanttapiiste Alayrac, Carl Doersch, Catalin Ioneseu, David Ding, Skanda Koppula, Daniel Zoran, Andrew Brock, Evan Shelhamer, Olivier Henaff, Matthew Botvinick, Andrew Zisserman, Oriol Vinyals, and Joao Carreira. Perceiver IO: general architec- ture for structured inputs & outputs. In ICLR, 2022.
Fredrik 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. arXiv, 2020.
Yonghan Jung, Jin Tian, and Elias Barcinboim. Learning causal effects via weighted empirical risk minimization. 2020.
11

---

Under review as a conference paper at ICLR 2023
Jean Kaddour, Yuchen Zhu, Qi Liu, Matt J Kusner, and Ricardo Silva. Causal effect inference for structured treatments. 2021.
Nathan Kallus. Deepmatch: Balancing deep covariate representations for causal inference using adversarial training. In ICML, 2020.
Joseph DY Kang and Joseph L Schafer. Demystifying double robustness: comparison of alternative strategies for estimating a population mean from incomplete data. Statistical science, 2007.
Jacob Devlin Ming-Wor Wei Chang Kenton and Lee Kristina Toutanova. Bert: Pre-training of deep
bidirectional transformers for language understanding. In NAACI-HIT, 2019.
Sören R Künzel, Jasjeet Sekhon, Peter Bickel, and Bin Yu. Metaleamers for estimating heteroge- neous 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 deep latent-variable models. In NeurIPS, 2017.
Shakir Mohamed and Balaji Lakshminarayanan. Learning in implicit generative models. arXiv, 2016.
David Newman. Bag of words data set. UCI Machine Learning Respository, 2008
Lizhen Nie, 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. Ncore: Neural counterfactual representation learning for combinations of treatments. arXiv, 2021.
Judea Pearl. Causality. 2009.
Raghunathan Ramakrishnan, Pavlo O Dral, Matthias Rupp. and Anatole Von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific data, 2014.
Paul R Rosenbaum and Donald B Rubin. The central role of the propensity score in observational studies for causal effects. Biometrika, 1983.
Paul R Rosenbaum and Donald Rubin. Reducing bias in observational studies using subclassifica- tion 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: general- ization 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 dose-response curves in early inhibitor discovery. Journal of medicinal chemistry, 2006.
12

---

Under review as a conference paper at ICLR 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 VanderWeele. Principles of confounder selection. European journal of epidemiology, 2019.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. NeurIPS, 2017.
Victor Veitch, Alexander Amour, Steve Yadlowsky, and Jacob Eisenstein. Counterfactual invariance to spurious correlations in text classification. 2021.
Stefan Wager and Susan Athey. Estimation and inference of heterogeneous treatment effects using random forests. Journal of the American Statistical Association, 1136523): 2018.
Hao Wang, Hao He, and Dina Katabi. Continuously indexed domain adaptation. In /CMI., 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 'soall-world' networks. nature, 1998.
Jeffrey M Wooldridge. Introductory econometrics: A modern approach. Cengage learning, 2015.
Yifan Wu, Ezra Winston, Divyansh Kaushik, and Zachary Lipton. Domain adaptation with asymmetrically-relxxed distribution alignment. ICMI, 2019.
Guoqiang Xu, Cunxiang Yin, Yuchen Zhang, Yuncong Li, Yancheng He, Jing Cai, and Zhongyu Wei. Learning discriminative representation base attention for uplift. In Pacific-Asia Conference Knowledge Discovery and Data Mining, pp. 200-211. Springer, 2022.
Chengxuan Ying, Tianle Cai, Shengjie Luo, Shuxin Zheng, Guolin Ke, Di He, Yanming Shen, and Tie-Yan Yan Liu. Do transformers really perform bad for graph representation? 2021.
Jinsung Yoon, James Jordon, and Mihaela Van Der Schaar. Ganite: Estimation of 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 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 estima- tion 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.
13

---

Under review as a conference paper at ICLR 2023
Exploring Transformer Backbones for Heterogeneous Treatment Effect Estimation
CONTENTS
1
Introduction
2 Related Work
3 Problem Statement and Assumptions
4
TransTEE: Transformers Treatment Effect Estimators
4.1 Covariate and Treatment Embedding Layers
4.2 Covariate and Treatment Self-Attention
4.3
Treatment-Covarriae Cross-Attention
5
Experimental Results
5.1
Experimental Settings
5.2
Case Study and Numerical Results
6
5.3
Analysis
5.4
Empirical Study on Pre-trained Language Models
Concluding Remarks
Appendix
14
A Extended Related Work
15
An Illustrative Example
15
Details and Discussions about Propensity Score Modelling
16
Analysis of the Failure Cases over Treatment Distribution Shifts
19
E Additional Experimental Setups
20
Additional Experimental Results
23
Remarks Interpretability
27
14

---

Under review as a conference paper at ICLR 2023
A EXTENDED RELATED WORK
Propensity Score. Most related works fundamentally rely on strongly ignorable conditions. Still even under ignorability, treatments may be selectively assigned according to propensities that depend on the covariates. To overcome the impact of such confounding, many statistical methods (Austin) 201
like covariate adjustment (Austin 201 T). 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
al. 2017). and multi-task Gaussian process (Alaa & van der Schaar 2017p. p. 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 targeted regularization are proposed to guarantee the consistency of estimated treatment effects under misspecification of either the outcome or propensity score model KKng & Schater [007 Funk et al. (011). There are also works using adversarial training for balanced representations (Bica et al.) 2020a] Kallus) 2020; Du et al. (202T). 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 adapta- tion, in partieular, 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. (022).
The algorithmic insights can be handily translated to the TEE domain (Shalit | 2011 | Johansson ct al. 2020; (Fdder et al. (02T). Here we also have the desideratum that covariate representations should be balanced such that the selection bias is minimized and the effect is maximally determined by the treatment. Algorithmically, when the treatment 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 on 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 inspire domain adaptation methodology and holds the potential to handle wider range of problems like hidden stratification in domain generalization, 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 2x. and 2t, and no self-attention is applied. However, TransTEE performs self-attention on 2g; 2e respectively and then cross-attention is performed between Za, 2t. When facing nigh-dimensional data, such as texts, images, and graphs, without multiple self- attention layers on 2x, 24 separately, the representations will be weak. That is why in machine translation, object detection, and segmentation tasks, the representations of images/texts will be firstly processed by multiple self-attention layers and then perform cross- attsstttttnioon with queries. We will verify this point in the following experiments. (ii) ANU cannot be applied to multi-treatment settings, which have been extensively studied recently (Kaddour et al. 202T Bica et al.) 2020b Parbhoo et al. 2021). The comparison experiments are in Section
B AN ILLUSTRATIVE EXAMPLE
To better understand the workflow with the above designs, we present a simple illustration here. Consider use case in medicine effect estimation, where x contains p patient information, e.g., Age. Sex, Blood Pressure (BP), and Previ- ous 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 Fig each covariate, treatment, and dosage will first be embedded to d-dimension repre- sentation by a specific learnable embedding layer. Each treatment embedding will be concatenated with its dosage embedding and the concatenated feature will be projected by linear layer to produce d dimensional vectors. Self-
f
Urm Mllein
XL
Scl-Anestee
X1
M,
Kiintebting
Lier
1
Figure 6: An Illustrative Example about the workflow of TransTER.
15

---

Under review as a conference paper at ICLR 2023
attention modules optimizes these embeddings by aggre- gating contextual information. Specifically, attribute Prev more related 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-covario
interactions, which is shown in Figure 2 that, each medicine will assign a higher weight relevant covariates especially confounders (BP) than irrelevant ones. Finally, we pool the resulted embedding and use one linear layer to predict the outcome.
C DETAILS AND DISCUSSIONS ABOUT PROPENSITY SCORE MODELLING
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 harnessed 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 (Imbens & Rubin 2015), sub-classification (Rosenbaum &
Rubin 1984). covariate adjustment (Austin) 201 targeted regularization (Van Der Laan & Rubin 2006) or conditional density estimation (Nie et al. (02T) that create guasi-randomized experiments (D' A'gostnnno 1998). 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 propensity score matching (Austin (201T), which seeks weigh 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 plausible propensity score from the generative model (Mohamed & Lakshminarayanan 2016) (see the degraded performance of MLE in Table D) TransTEE learns a propensity score network Te (f)x) via minimax bilevel optimization. The motivations for adversarial training between no (x, () and TH (4) ()) are three- fold: (i) enforces the independence between treatment and covariate representations as shown in Proposition [ which serves as algorithmic randomization in replace of costly randomized controlled trials [Rubin| 2007) for overcoming selection bias (D' Agostino 1998 [mbens & Rubin 2015); (ii) it explicitly models propensity T (t(x) xX to refine treatment representations and promote covariate adjustment (Kaddour et al. 2021): : and (iii) taking an adversarial domain adaptation perspective, the methodology is effective for learning invariant representations and further regularizes (4 f) to be
invariant to nuisance factors and may perform better empirically on some classes of distribution shifts
(Ganin et al. 2016 Shalit al. 2017 Zhao et al., 2018; Johansson et al. 2020
Wang al. 2020).
Based on the above discussion, when treatments are discrete, one might consider directly applying heuristic methods like adversarial domain adaptation (see (Ganin et al.) 2016 Zhao al. (018)) for algorithmic development guidelines). We note the heuristic nature of domain-adversarial methods (see (Wu 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, a more general and challenging scenario, where we want to estimate ADRF, and propose two variants of L.s as an adversary for the outcome regression objective La accordingly. Recall that
Cx n (y/ pooxxx (X:, (²))² the adversarial training process is shown in Eq. below:
min max Lo y, &
(x.
(1)
We refer to the above minimax game for algorithmic randomization in replace of costly randomized controlled trials. Such algorithmic randomization based on neural representations using propensity score creates subgroups of different treated units if they had been randomly assigned to different treatments such that conditional independence T T_L (TT) is enforced across strata and continu- ation, which approximates a random block experiment to the observed covariates Imbens & Rubin 2015).
Below introduce variants of L l):
16

---

Under review as a conference paper at ICLR 2023
Table 5: Performance of individualized treatment-dose response estimation on 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 be 2.0.
METnono
#TREATMENT=]
#TREATMENT=2
#TREATMENT=3
IN-SAMPLR
OUT-SAMPL.
IN-SAMPLE
OUT-SAMPLE
IN-SAMPLE
OUT:SAMPLE
SCIGAN
5.6966 + 0.0000
5.6546 + 0.0000
2.0924 0.00000
2.3067 = 0.0000
4,3183 + 0,0000
4,6231 = 0.0000
TARNET(D)
0.7858
0.7908 0.0606
0.0784
0.0777
3.1982
3.1920
DRNET(D)
0.8052 1 0.0466
1.3739 0.0855
1.3738 1 0.0853
2.8632 0.2222
2.855810.Il
VCNET(()
0.1566
0.1579
0.0743
0.0737
0.64993
0.1397
TRANSTEE
0.0573 1 1000361
0,0585 0.0358
0.0550 0.0137
0,0556 a 0.0129
0,2803 1 0.0658
0,2768 1 0.0639
TRANSTHE:
0.0509
0.0663
0268
0.2618 0.0737
0.2577 0.0726
TRANSTEE à PTR
0.0.443 = 0.0096
0.0355 0.0094
0.0679 + 0.0252
0.0686 = 0.0252
0,2645 + 0.0702
0.2597 + 0.0675
Treatment Regularization (TR) is a standard MSE over the treatment space given the predicted treatment ( and the ground truth l
CT (xx a (4- (()))))))
TR is explicitly matching the mean of the propensity score to that of the treatment. In ideal case, the T (t]xx) 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 a scalar t:
Probabilistic Treatment Regularization (PTR) is probabilistic version of TR which models the mean y (with a slight abuse of notation) and variance o² of estimated treatment
(t. (x(p)x))) (xx(o)
CCTR_
log (o())()))))))
(3)
The PTR matches the whole distribution, i.e. both the mean and variance, of the propensity score to that of the treatment, 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 t and covariate To be clear, denote μ t) := Wy o (x), (t)) and (t(x)) := W. o & fr (x), which decompose the predictions into featurizers &, : T Zr. 2. : X - Zx and predictors Wy Zx x ZT - )r w. : 2x + T. For example, x and (t) 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, pa (x, () is computed from (.()) and
4 (x). For the ease of our analysis below, we assume the predictors We, Wa are fixed.
Proposition 1. (The optimum of propensity score model) In the equilibrium of the game, assuming the outcome prediction model is fixed, then the optimum of TR is achieved when (()) (x(xx))))) [(), V xx(x) via matching the mean of propensity score π n(,, ( T (x)) and the marginal distribution pl (x)) and the optimum discriminator of PTR is achieved via matching both the mean and variance such that E (t) (x)) = (t), V (( 2 (x)]
(t), V (x).
Proof. The proof concerns the analysis of the Equilibrium of the Minimax Game is a special case of (Wang et al.) 2020) when there are only two players, i.e. #i and Tg. We represent treatments explicitly and interpret the connections with combating selection biases. Given the outcome regression model ng fixed, the optimal propensity score model T² is
T² = arg min Cs (x),
arg min (x)
(Ø.(1) - Te (0)())
(4)
= arg min E,g (x)
(o.(t) zo (fx())xx)²
17

---

Under review as a conference paper at ICLR 2023
The inner minimum is achieved T (#)(f))x) (0)(()))((((0000)))))))))))) (t))] given the following quadratic form:
(10. (t) - ny
(5)
(0)),))))
2 (t)xx) (t) (x)) w (+(())xx²²
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. Caan be converted to maximizing the inner loop:
max Cx, 2t())) Cg. (L (x), Pc()))
(10_00),,000))))(991_ (t) - Es, (t))0())))0))0_____))))))) (t)))
(++(1) (.,,00/pp(000000...0)00000000000
(t) Es (( d (x)].
Next we show the difference between Eg. 6 and the variance of the treatment V[
(x) (t) [t) v[2,(t)
(x) (x))
V
(x)
(x) (t) (t)}
-E( E. (x) E t. T [] (x) E (t) (x)]2 Eg. (t)
[ ( <x)
[x) Es,
(x)}²
where the last inequality is by Jensen's ' inequality and the convexity of (t)?. The optimum is achieved when (t) 4 (x)] is constant w.r.t 4 (x) and so E (t) d2 (x) = ([),, (x).
The proof process for PTR is similar but includes the derivation variance matching.
arg min (x)
arg arg (at (E(0,(1) T (1))
log ([0.(1))
( 2x
2
(t) 4. (x)] +,(())) log (t))
arg
2V (f)
(8)
where Eg, (x) and (s.(e) denote Es, (x)=p(0, (x)) and (t)pp())( (x)) respectively for brevity.
The first term can reduce to a constant given the definition of variance:
(E (t)[x] - (t))
(x_( (x)) (0 (x))
27 (tt
V ([(x(t)lx]
= -E
2V (1)x
The second term can be upper bounded by using Jensen's inequality:
log ([2x(t)))
Eg. (x)) (x)) Ee, (O0
2
1
log MM.
(10)
2
1
log (V)).
2
Combining Eq. and Eq. 10 the optimum 1 +
1 log (V)),(t)) is achieved when (t) (x)].
2 2
(( tb
(x) Is constant w.r.t (x) and so (t), ( = ([), v( & (x)
(t). v (x) according to the equality conditions of the first and second inequality in Eq.
10 respectively.
18

---

Under review as a conference paper at ICLR 2023
D ANALYSIS OF THE FAILURE CASES OVER TREATMENT DISTRIBUTION SHIFTS
As shown in Figure (a,c), with the shifts of the treatment interval, the estimation performance of DRNer and TARNet decline significantly. VCNet achieves oo estimation error when h = 5 partly because its hand-craft 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 3b) When training on t E [0, 1.75]. these discrete approximation methods cannot transfer to new distribution E (1.75, 2.0]. These unseen treatments rounded down to the nearest neighbors t' in T and be seemed the same as f'. We conduct ablation about the treatment embedding as in Table Appendix. Such a simple fix (VCNet-Eimbeddings removes the demand on fixed interval constraint to treatments and attains superior performance on both interpolation and extrapolation settings. The result clearly shows the pitfalls of cranfer feature mapping for TEE. We highlight that it is neglected by most existing works (Schwab et al., 20200, 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 effectiveness of TransTEE that uses learnable embeddings to map continuous treatments hidden representations.
Below we show the assumption that the value of treatments or dosages are a fixed interval [1, 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 f and the result is similar for continuous dosage.
Proposition 2. Given a data sample t),, where x E RR, E [1 E R. Assume μ is a L-Lipschitz function over (x, t) Rd411 namely (u(u) - u(v) S L u-vll Purtitioning [1, h] uniformly into 5 sub-interval, and then get T = [1+ 4f+ 0,1 A 1,,,,,1+ h-1 6).
vll.
Previous studies most rounding down a treatment to its nearest value in T (either I+
h-1 or
h-1
1+ 15
Ang and use (T) branches to approximate the entire continuum 11, h). The approximation
error can be bounded by
18
h 1x(xt),,1
18
h - 1)
(((x
max
11
x,
1
8
h
to
h
1) ,L (1)
td
h
1
}
< max
L
(11)
8
8
< L
8
The bound is affected by both the number of branches 5 and treatment interval [t, h.. However, as
far as we know, most previous works ignore the impacts of the treatment interval [t, h] and adopt a simple but much stronger assumption that treatments are all in the interval [0, 1] Nie al. (2021)
fixed interval SSchwab intervalScchwab et al. (2020). 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 simulated datasets. Numbers reported are AMSE of test data based on 100 repeats, and numbers after are the estimated standard deviation of the average value. For Extrapolation (h 2), models are trained with [0, 1.75] and tested in E [0, 2]. For Extrapolation (h- = 5), models are trained with t E [0, 4] and tested in t [0, 5]
METHODS
VANILLA
VANILLA =
EXTRAPOLATION (A EXTRAPOLATION (h 5)
TARNET SHALIT EY AL. 2017
0.045 + 0.0009
0.3864 = 0. 0.04335
0.0984 + 0.02315
0.3647 + 0.03626
DRNET (C
0.042 4 0.0009
0.3871 0.03851
0.0885 0.00094
0.3647 + 0.03625
VCNET CINETNIEETT AL 2021
0.018 1 0.0010
NAN
0.0669 0.05227
NAN 0.0178 0.00887
VCNETFESBEDDDNGS
0.013 0.00465
0.0167 0.01150
0.0118 0,00482
19

---

Under review as a conference paper at ICLR 2023
E ADDITIONAL EXPERIMENTAL SETUPS
All the assets (i.e., datasets and the codes for baselines) we use include MIT license containing copyright notice and this permission notice shall be included in all copies or substantial portions of the software. We conduct all the experiments on machine with 17-8700K CPU, 32G RAM, and four Nvidia GeForce RTX2080TI (10GB) GPU cards.
E.1 DETAIL EVALUATION METRICS.
1 AMSE N
Σ
/ [xx. f(x,,)) t) a(t)dt
(12)
((x,,,t)) - (xx,
UPEHE@K N
E,0'
(13)
WPEHE@K Σ
((x,,,t)) [(x,,,,1)] ((()))((()))
[xx,))s)) - f(x,tt, 8) (()ddt
AMSE
(14)
NT
il
E.2 NETWORK STRUCTURE AND PARAMETER SETTING
Table. and Table. g show the detail of TransTEE architecture and hyper-parameters. For all the synthetic and semi-synthetic datasets, we tune parameters based on 20 additional runs. In 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 or similar baselines.
Table 7: Architecture details of TransTEE, where p is the number of covariates.
Module
Covariates
Treatment
Embedding Layer
[Linear]
Linear
Output Size
Bsz X p #Emb
b8z x # Emb
Multi-head Att
Multi-head Att
BatchNorm
BatchNorm
Self-Attention
#Layers
#Layers
Linear
Linear
BatchNorm
BatchNorm
Output Size
Bsz x p x # #Emb
Bsz x 1 x #Emb
Multi-head Att
BarchNorm
Cross-Attention
#Layers
Linear
BatchNorm
Output Size
Bsz x 1 x #mmb
Projection Layer
Linear
Output Size
Bsz x
E.3
SIMULATION DETAILS.
Synthetic Dataset (Nie et al. 2021). The synthetic dataset contains 500 training points and 200
testing points. Data is generated follows: zj Unif(0, 1), where xy is the j-th dimension of
20

---

Under review as a conference paper at ICLR 2023
Table 8: Hyper-paramette on different datasets. Bsz indicates the batch size, # Emb indicates the embedding dimension, Lr. S indicates the scheduler of the learning rate (Cos is the cosine annealing Learning rate).
Dataset
Bsz
Emb
# Layers
# Heads
Lr
Lr. S
Simu
500
10
2
0.01
Cos
IHDP
128
10
2
0.0005
Cos
News
256
10
2
0.01
Cos
SW
500
16
2
0.01
None
TCGA
1000
48
0.01
None
x E R2, and
10 (max( I22, x33)) + max(23, F4, 3353
+ sin(0.5xx)) (1 +1+ exp(x1 - 0.5xz))
1 (x1 + 22)2
2 ( + 2x5 - 6.5 + NO 0.25)
+ 4 maxx(21,12))
g/2, = (2 0.5))
+ +.(00,0.25)
1 222
where: t = (1 exp(-t))
for treatment in [0, h), we revised it to =1 exp -) h
IHDP (Hill, 2011) is semi-synthetic dataset containing 25 covariates, 747 observations and binary treatments. For treatments in [0, 1), we follow VCNet (Nie et al. 2021) and generate treatments and responses by:
2x1 + max(xx3 Ts, 26)
DiEsuna (I; (2)
1 x2 0.2 + min 23, Nol + 2 tannh
-40,00,
|S
g/u,, = 1.2 tanh 5 Eiesau. (xi- c1)
(3Tl)
exp(0.2(x1 - x6))
N0, 0.25),
Sdin,1
0.5 + 5min F3, vs)
where (1 + exp( -)) -1)) Som (1,2,3,5,6) is the index of continuous features,
Sdin,1 (4, 7,8 9, 10, 11, 12, 13, 14, 15), Sais22
{16. 17, 18, 19, 20, 21, 22, 23, 23,24,200 25} and
Exsaa z1
.02 = E Eics 2
Sais, USdir.2 = [25] Scon- Here C = E
To allow 1
comparison on various treatment intervals 1. [0, h.. treatments and responses are generated by:
t= (1 exp(-f))
(in(3tt/)) y/2,, 1.2 - 1/h
Eie (0; -1)
exp (xp(0.2(z) x6))
tanh
+N 0.25),
S
0.5 + 5 min(22, X3, 27)
where the orange part is the only different compared to the generalization of vanilla IHDP dataset (h= 1). Note that Sat,,1 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 covariates 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 originally introduced a benchmark in the binary treatment setting. We generate the treatment and outcome in a similar way as (Nie et al.] (2021) but with a dynamic range treatment intervals [0, h). We first generate vi, v. Ui NO 1) and then set
2 ol/l1///l2: (1,2,3}. Given I, generate from Beta
2,
x
h.And generate the
20g
outcome by
v r
1/ exp
0.3
r
21

---

Under review as a conference paper at ICLR 2023
SCIGAN
SCIGAN
(D)
Drnot
Drnet
Venat (D)
Venot (D)
TResTEE
TranTTEB
AMSE
AMSE
5-
10
10
Dosage Selection Bias
Treatment Selection Bias
(a) Performance with different dosage selection bias. (b) Performance with different treatment selection bias.
Figure 7: Performance of five methods on TCGA (D) dataset with varying bias levels.
20cTar) (4(1- ( sin (2)1 +0,
1/2, 1. (12
TCGA (D) (Bica al. 2020b) We obtain covariates from real dataset The Cancer Genomic Atlas (TCGA) and consider 3 treatments, where each treatment is accompanied by one dosage and a set of parameters. of, v/. vg. For each run, we randomly sample a vector, uf NO, 1) and then set of = u//l|u[[ where || is Euclidean norm. The shape of the response curve for each treatment, fr(2, is given in Table 9 We add
(O, noise to the outcomes. Interventions are assigned by sampling a dosage, d,, for each treatment from beta distribution, de Beta(a, Bc). a 2 controls the dosage selection bias (a gives the uniform distribution). where s; is the optimal d for treatment 1. We then assign treatment according fy| Categorica((Sotmax((/]
3t 87 +
s,)] where increasing increases selection bias, and 0 leads to random assignments. The factual intervention is given by St,, Unless otherwise specified, we set K and a 2.
Turh
18
10
Yena (D)
16
SCIGAN
TranTEE
Reipurae
20
Truth
Truth
Dmsrt
Dmsrt (P)
Venet (D)
Vene (D)
SCIGIN
SCIGAN
TracTEE
TracTTB
Trentrent
Thentmen
Thestmer
(a) Estimated ADRF for f1-
(b) Estimated ADRF for to.
(c) Estimated ADRF for f..
Figure 8: Estimated ADRF on 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 dosage- response curve for all treatments.
For structural treatments, we first define the Baseline effect (Bica et al. 2020b). For each run the experiment, randomly sample vector ug Unif[0, 1]. and set vo uo/ /o/l where || is the Euclidean norm. The baseline effect is defined as
(o(x) = un
Forr symmetry, if oī we sample of from 1 Beta(o, Be) where set though 1.
22

---

Under review as a conference paper at ICLR 2023
Table Dose response curves used to generate semi-synthetic outcomes for patient features In the experiments, we set C' = 10. of. 222 vg are the parameters associated with each treatment t.
Treatment
Dose- Dese-Response
Optimal dosage
1
fi 5) C ((o)) 12(x)) To - 12(1))) xz²)
si 2
(u2) T a
f2 8) = C (uf) + sin (m( -
2
82
2(75)
3 fs(r, (w+) Tx x + 12s(s - b)2, where i - 0.75
if 2 0.75 else
2021). 20-dimensional multivariate covariates are uniformly sampled There are 1, 000 units in in-sample dataset, and 500 in the out-sample one. 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-Strogatz small-world graphs are repeatedly generated until a connected one is get. Each vertex has one feature, i.e. its degree centrality. graph's 's node connectivity is denoted as v()) and its average shortest path length as (Gg). Similar for the baseline effect, two randomly sampled vectors tu, Vy are generated. Then, given an assigned graph treatment G and covariate vector the outcome is generated by
Small-World (Kaddour et al. according to æ5 Unif|
= (0) + 0.2x( v, x + E(G) x x - (00,1)
TCGA (S) (Kaddour 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 a covariate vector I E IR 124000 and these units are split randomly into in- and out-sample datasets in each run randomly. For each covariate vector x. its 8-dimensional
PCA components zPCA E 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 a relational graph, where each node corresponds to an 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 gap, r2, zpve, uO in a vector = E R²$ which is denoted as the the assigned molecule treatment. Finally, we generate outcomes by
mu, alpha, homo, tap,,
= 10) + 0.012 PCA +E,E~ N(0,1)
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 conveys. The task also known as Profile of Mood States (POMS). Each sentence in the dataset is created using one of 42 templates, with placeholders for person'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 -merrican or European will be used fill the placeholder (<Person>). One of four possible mood states: Anger, Sadness, Fear and Joy is used to 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 unexpected, but Roger made me feel pessimistic." it will have the counterfactual example: "Ir was totally unexpected, but Amanda made 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.
F ADDITIONAL EXPERIMENTAL RESULTS
F.1 COMPARISION BETWEEN TRANSTEE AND ANU (XU ET AL. 2022)
We implement ANU and evaluate it in the same settings and show that is inferior compared to the proposed Tran follows. Specifically, we compare the attentive neural uplift model (ANU) (Xu
23

---

Under review as a conference paper at ICLR 2023
0.32
0.14
0.10
0.10
0.99
0.08
0.06
6.4
0.02
0.00
0.1
0.5
0.5
0.7
1.0
2.0
3.0
4.0
5.0
0.00
0.3
0.5
0.7
1.0
2.0
3.0
4.0
50
(a) Ablation study PTR.
(b) Ablation study of TR.
Figure 9: Ablation study of the balanced weight for treatment regularization on the IHDP dataset.
Table 10: Comparision between TransTEE and ANU (Xu 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 = 0. 08236
0.7800 = 0.6148
NAN 0.4244 19832
ANU (Xu et al. 2022)
0.1482 0.17362
0.2147# 0.32451
Trans TEE
0.0983 = 0.15384
0.1151 = 0.1028
0.2745 + 0.1497
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 TransTEE are similar. The result shown in the following table. With the usage of treatment embeddings. ANU shown to be more robust than VCNet and DRNet when treatment shift occurs. However, in both the binary treatment setting 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 detailed in Section in the main paper. Covariates here are long sentences. Thanks to the use of self-attentios modules, TransTEE can achieve better estimation results compared to baselines (Table II). For AHU, no self-attention layer is applied, and the final estimation is inaccurate, which verifies the superiority of the proposed framework.
F.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 of both TR and PTR. Figure presents the results of the experiments on the IHDP 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 achieved when À is 0.5 for both two methods, which shows that the best balancing parameter (0.5 on our experiments.) for these two regularization terms should be searched carefully. Besides, training both the treatment predictor and the feature encoder simultaneously zero-sum game is difficult and sometimes unstable (shown in Figure Saaue/rrihh right)
Robustness to noisy covariates. We manipulate Sdis, Sais,2 to generate datasets with different noisy covariates, e.g., when the number of covariates that only influecc the outcome is
Table 11: Comparision between TransTEE and ANU (Xu et al. 2022) on the IHDP dataset.
Correlation/Rppreentadde Based Baselines
Treatment Effect Estimators
TC ATEcT
TReATE
CONEXP
INLP
CausalBERT
TarNet
DRNet
ANU
TransTEE
Gender
0.086
0.125
0.02
0.313
0.179
0.0067
0.0088
0.184
0.013
Race
0.014
0.046
0.08
0.591
0.213
0.005
0.006
0.093
0.0174
24

---

Under review as a conference paper at ICLR 2023
0.3
0.0
0.00
2003330 S1A.F-S7aaSSee2200000
(a) TransTEE.
(b) TransTEE+TR.
(c) TransTEE+PTR
Figure 10: The distribution of learned weights for the cross cross-attention module on the IHDP dataset of different models.
Teret
Truah
Dead
1.3
Trub Tasut Dran
7,5
Dran
Trub Taian
100
12.5
Thcat Brast Visat
Trub
Veaut
Vasc:
Vanut
7.5
Roqeer
3.0
TraafTE
Thea/TEE
Theat
2,5
2,5
td
-25
-500
-50
-50
0.5
(a) during training (b) 2,7 0.1 during (c) during training (d) 5,4
0.25
testing.
traning and 2,7 = 0 and testing. during testing (extrapola- tion).
during training
5, == 0 during testing (ex- trapolation).
Figure 11: Estimated ADRF on test set from a typical run of TarNet (Shalit et al. 2017). DR- Net (Schwab et al. 2020), VCNet (Nie et al. 202T) 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 a hand-crafted mapping matrix. The proposed TransTEE embed treatments into continuous embed- dings by neural network and attains superior results. (b,d) When training with 0.1 S t < 2.0 and 0.25 < t < 5.0. TARNet and DRNet cannot extrapolate to distributions with 0 < t S 2.0 and S t S 5.0. (c) The hand-crafted mapping matrix of VCNet can only be used the scenario where << 2. Otherwise, VCNet cannot converge and incur 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 in advance and extrapolates well.
Sass,1 {4), and Sdis.2 (7, 87,8,9. 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20,21,22, 23, 23, 25). when the number of covariates that influence the outcome is 24. Sais.1 (4, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, }. and S ure Figure [46) 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 always low and superior to baselines by a large margin.
{25). Fig-
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 m[[[x) does provide some performance gains. However, adversarially trained r-model can be significantly better, especially for extrapolation settings. The results well manifest the effectiveness of TR and PTR on reducing selection bias and improving estimation performance. In fact, approaches like TMLE are not robust the initial estimator is poor (Shi et al. (2019).
Training dynamics comparison of different regularization terms. Here compare four reg- ularization terms, which are TransTEE with no regularization, TranSTEE+TR, TransTEE+PTR, and TEansTEEMML... MTL. is simple Multi-Task Leariigg strategy, which uses
25

---

Under review as a conference paper at ICLR 2023
32-
Usesse Repooiie time
Stumsen Reparion lane
TresTBB4MTI
25
TaSAEE
0.25
Mbc Eire
TransTZI
TraSTEE-TR
0.00
TrmeT12 MTL.
12
-0.55
250
500
750
1000
1210
1100
250
500
130
1930
1250
1908
1.
25
Dranticts
(a) Outcome regression error.
(b) Treatment regression enror.
(c) MSE in the test set.
Figure 12: Training dynamics of TransTEE on 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.
Tane
Tame
Tuatte
Tecth::
Tearkic
133
15
(a) = during training (b) = 1.9,7 during (c) = 5 during training (d) = 4, = during and testing.
traning and h 2, 0 and testing. during testing (extrapola- tion).
training and h = 5,1 = 0 during testing (extrapola- tion).
Estimated ADRF on the test set from a typical run of TarNet (Shalit et al., 2017). DRNet (Schwab et al., 2020). VCNet (Nie et al. 2021) and ours on News dataset. All of these methods are well optimized. Suppose & [1, hj. (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 < 1. S 4.0. TARNet and DRNet cannot extrapolate to distributions with 0 < t < 2.0 and 0 < f < 5.0. (c) The hand-crafted mapping matrix of VCNet can only be used in the scenario where & 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 know h in advance and extrapolates well.
Figure 13:
(x f) + Cx,(( t) during training without an adversarial game. As shown in Figure 12 without adversarial training. TransTEE+MTL quickly attains low treatment estimation error but further oscillate and converge with high error, and both the outcome regression error and MSE in the test set remain high. In contrast, TR and PTR make TransTEF converge faster and attain lower test MSE Overall, PTR consistently works the best and its low treatment regression error shows that n (f|x) estimates an accurate propensity score.
F.3 SHOWCASE OF SENTENCES AND COUNTERFACTUAN COUNTERPARTS WITH THE MAXIMAL/MINIMAL ATES.
Table 14 showcases the top-10 10 samples with the maximal/ minimal ATEs. Interestingly, we can see most sentences with a large ATE have similar patterns, that is "< clause >, but/and < Person > made me feel Adj >" Besides, most sentences with large ATE have small length, which is 11 words on average. By contrast, sentences with small ATEs follow other patterns and are longer, which 17.6 average. Consider the effect of Race, Table [5] 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 is smaller than sentences with low ATES (12 14.7). Besides, the position
26

---

Under review as a conference paper at ICLR 2023
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
t E [0, 2]. For For Extrapolation (h 5), models are trained with E [0, 4.5) and tested in E [0,5]
METHODS
VANILLA
VANILLA (h = 5)
EXTRAPOLATION (1 = 2)
EXTRAPOLATION (7 = 5)
TARNET
0.082 0.019
0,956 0.041
0.716 1 0.038 0.703 +\t 0.038
0.847 1 0.053
DRNET
0.083 = 0.032
0.956 0.041
0.834 + 0.053
VCNET
0.013 +: 0.005
NAN
NAN 0.024 1 0.017
NAN 0.029 1 0.019
TRANSTEE
0.010 1. 0.004
0.017 0.008
TRANSTEE+
0.011 0.003
0.016 0.008
0.019 + 0.008
0.002
PRA
0.011 0.004
0.014 0.007
0.022 0.008
0.029 0.016
of perturbation words (the name from a specific race) for sentences with the maximal/minimal ATEs totally different, which is at the beginning for the former and at the middle for the latter. Namely, TransTEE helps 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 TransTEE is able estimate the effect of every sentence and is of great benefit for model interpretation and analysis especially under high inference latency.
G REMARKS ON INTERPRETABILITY
is fundamentally hard evaluate the interpretability even for supervised learners, 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 from the experimental results in fig. 4(a), 4(b), and fig. 10 that TransTEE assigns more weights to confounders as opposed other covariates, which is a new observation that previous backbones are hard to achieve. We see that explaining causal inference models in this way using the feature importance scores for each covariate can used for benchmarking treatment effect estimators (Crabbé et 2022).
27

---

Under review as a conference paper at ICLR 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 test set. (The baseline results are reproduced using the official code of (Kaddour ct al. 202T) in a consistent experimental environment, which can be slightly different than the results
reported in (Kaddour al. 202T)
SW
TCGA (Bias=0.1)
TCGA (Blas=0.3)
TCGA (Bias=0.5)
Method
In-sample
Out-sample
In-sample
Out-sample
In-sample
Out-sample
In-sample
Out-sample
WPEHE@2
Zero
41.7220000
49.69 0.00
13.93 0.00
13.13 + 0.00
0.000
13.13
13.61 0.00
GNN
17.38
0.00
10.90 7.71
10.91 # 7.71
13.58
13.22 0.18
12.86 4 0.38
14.62 = 0.91
GraphITE
17.37
24.56:0.00
15.04
14.96 0.30
13.49
1.223
13,70
14.38 0.30
SIN
15.79 1.72
28.78 + 4.54
46.47
+ 2.19
$4.41 7.81
7.93 0.79
11.04 + 1.52
10,31 0.93
14.09 + 2.14
TraSTEE
14.74 0.09
1.07 21.78
9.07 2.15
9,33 2.13
3.60
8.37 = 3.64
9.52 a: 3.59
10,10 = 3.79
3
Zero
40.75
43.76 = 0.00
13.93
0.00
13.61 + 0,00
13.93 0.00
13.61 = 0.00
13.61 +.00
14.14 0.00
GNN
18.26 0.00
0.01
10.75
10.91
13.63
13.58 0.19
0.33
15.29 1.04
GraphITI:
18.27 + 0.01
20.95 + 0.02
14.88 0.19
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
45.29 = 2.33
53.72 8.09
7.94
11.53 1.59
10.89 1.07
14.27 = 1.92
TransTEF
15,30 1.12
18.73 2.09
9.01 2,02
9.58 2.04
7,58 = 3.62
8.65 + 3.75
9.64 3.56
10.59 3.88
WPEHE 6P4
Zero
45.74
0.00
14.14 0.00
+
14.14
13.75 0.00
0.00
14.31
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
15.47 = 1.05
GraphITE
22.12=0.00
23,03
15.05 0.18
15.14 0.28
03.24
14,30 = 0.35
12.77 0.02
15.38 0.300
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
TranTTE
18.99 = 0.83
1.97 19.63
9.09 # 1.97
9.66 2.01
7.67 7
8.71 = 3.78
9.78 3,63
10.74
WPEHEE0
Zero
4..09
45.96 1..9600.00 0.00
14.31 0,00
13.95 + 0.00
14.31 +
13.95 0.00
13,95 # 0.00
14.47 0.00
GNN
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 1 0.01
24.22 1 0.03
15.24 0.19
15.29 1 0.28
13.68 0.17
14.37 0.37
12,95 0.03
15.59 1 0.30
SIN
25.48 = 3.02
25.44 3.50
44.55 + 2.35
$2.78 + 8.04
8.10
11.76 # 1.59
11.75 # 1.22
14.59 = 1.84
TranTTEE
20.16
21.08
1.78
1.96
9,72 2,00
3.75
8.80 3.82
9,91 3.66
10.89 3.94
WPEEE
Zero
49.95 : 0.00
50.10 0.00
14.47 0.00
14.04 4: 0.00
14.47 a 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 4: 7.79
14.07 0.22
14.11 0.18
13.45 0.38
15,76 à 1.04
GraphITE
25.17 0.02
26.94
15.40
15.37 0,28
13.74 + 0.12
14.58 + 0.38
13.09 = 0.04
15.68 = 0.29
SIN
27.07
2,98
3.51
44.48 2.35
7.99
8.22 0.75
11.82 1.58
14.74
TranTTEF:
21,32 & 0.79
22,99 1,43
9.23 a 1.95
9.77 1.99
7,80 +: 3.83
8.84 + 3.89
10.01 + 3.70
10.96 + 3.95
WP
Zem
55,40 + 0.00
58.42
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 +:006:77 7.82
14,12 + 0.21
14.14 + 0.18
13.51 + 0.38
15.81 = 1.03
GraphITE
0.01
13.97 0.08
0.40
15.74 0.29
SIN
31.07 + 3.07
34.17 3.41
44.45 2.37
$2.40 7.98
8.28 0.74
11.85 1.58
12.11 1.18
14.83 = 1.87
TranTEE
24.71 = 0.41
25.84 0.73
9.27 + 1.94
9.81 1.99
7.82 2 3.84
8.89 = 3.89
10.06 = 3.71
11.01 = 3.95
WPEHF&6 8
Zero
57.99 + 0.00
66.78 0.00
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.00
37,57 0.05
11.22
7.93
14.19
15.87
GraphITE
31.45 0.01
37.58
0.00
15.55 + 0.19
15:47 + 0.28
14.30 + 0.04
14.85 0.43
13,23
15.78 = 0.28
SIN
33.58
40,83 3.64
44.48 2.38
52.34 7.97
8.33 0.74
11.87 1.57
12.22 1.17
14.91 = 1.89
TranTTE
26.48 a 0,27
32.40 0.85
9.31 1 1.94
9.85 1.99
7.88 3.84
8,90 + 3,90
10,10 #: 3.72
11.04 3.96
Zero
62.52
64.61
14.66 0.00
14.20
14.61
1..14
14.66 0.00
14.20 0.00
GNN
34,13
36.48 0.04
11.26 21.26+77 7,96
11.13 7.87
14.21 0.24
14,22 0.17
13.63 0.38
15,92 = 1.01
GraphITE
34.17 2 0.02
36,49
15.60 0.19
15.53 0.28
14.35 + 0.04
14,90 0.43
13.28 0.04
15.83 = 0.28
SIN
36.79 3.35
4.11
44.47 à 2.39
52.31 1: 7.97
8.36 à 0.74
11.90 A 1.57
12.40 1: 1.23
15.08 à 1.80
TranSTEF:
28.84 + 0.23
0.71 31.40
9.34 + 1.94
9.88 2,00)
7.90 + 3.85
8.94 3.91
10.14 = 3.73
11.08 = 3.97
WPEHEE
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 1 0.04
11.28 7.98
11.16 7.89
14,29 a 0.22
14,32 0.18
13.66 0.38
15.96 : 1.01
GraphITE:
34,30 4:0::0000000 0.02
37.66 0.00
15.64 0.19
15.56 4: 0.28
14.38 0.04
14,93
à 0.43
13.31 0.04
15.87 0.27
SIN
37.08 + 3.35
41.79
44.49 + 2.40
$2.28 7.96
8.39 = 0.74
11.92 1.58
12.49 + 1.22
15.13 + 1.81
TranSTEE
28.89 0.19
32.25
9.36 1.93
9,90 2.00
7.94
8.95 3.92
10.16 3.74
11.10
28

---

Under review as a conference paper at ICLR 2023
WPENEOS
8
WISHEOS
W7SHD66
WYSHIDA In
50
20
20
40
20
40
(a) SW InmSlee
(b) SW Out-Sample
CAN
CewchITE
TreATTE
SIN
TransTEE
25
WPENERIN
WFFOIT IN
30
WPFSCCI
0.0
0.0
0.0
(c) TCGA In-Sample
(d) TCGA Out-Sample
Figure 14: WPEHE @K K over increasing bias strength and varying K E (2,..., 10) on the SW and the TCGA dataset.
29

---

Under review as a conference paper at ICLR 2023
UPERE
SPESE
LIVe 11121199
UPErORE
25
20
20
20
40
20
40
(a) SW In-Samele.
(b) SW Out-Sample
CAN
CewchITE
TreATEE
SIN
TransTEE
1
UPERERN
30
UPEREROO
0.0
0.0
0.0
(c) TCGA In-Sample
(d) TCGA Out-Sample
Figure 15: UPEHE@ @ K over increasing bias strength and varying K E {2, ..., 10} on the SW and the TCGA dataset.
30

---

Under review as a conference paper at ICLR 2023
Table 14: Top-10 samples with the maximal and minimal ATE for the effect of Gender, Pertur- bation words in factual sentences and counterfactual sentences are colored by Orange and Magenta respecttively.
Sentences with The Maximal ATEs
Index
Sentence
ATE
It was totally unexpected, but Roger made feel pessimistic. We went to the restaurant, and Alphonse made me feel frustration.
0.6393
2
0.578
It was totally unexpected, but Amanda made me feel pessimistic. We went to the university, and my husband made me feel angst.
0,5109
4
0,4538
5
It is far from over, but so far i made Jasmine feel frustration. We were told that Torrance found himself in consternation situation.
0,4366
Factual
6
0.4203
7
We went to the university, and my son made me feel revulsion.
0,399
To our amazement, the conversation with my aunt was dejected. To our amazement, the conversation with my aunt was dejected. We went the supermarket, and Roger made me feel uneasiness.
0,3952
0,3952
10
0.3752
It was totally unexpected, but Amanda made me feel pessimistic. We went the school, and Latisha made feel frustration.
0.6393
0.578
It was totally unexpected, but Roger made me feel pessimistic. We went the market, and my daughter made me feel angst.
0,5
4
0,4538
It is far from over. but für made Jamel feel frustration.
0,4366
Counterfactual
6
We were told that Tia found herself in consternation situation. We went to the hairdresser, and my sister made me feel revulsion.
0.4203
7
0.399
8
To our amazement, the conversation with uncle was dejected.
0.3952
To our umazement, the conversation with my uncle was dejected.
0,3952
10
We went to the university, and Amanda made me foel uncasiness.
0.3752
Sentences with The Minimal ATEs
Index
Sentence
ATE
our amazement, the conversation with Jack was irritating. added information is given in this part.
0
To our surprise, my husband found himself vexing situation, this only here to confuse the classifier.
2
0
The conversation with Amanda irritating, could from simply looking. this only here confuse the classifier.
3
0
this is only here to confuse the classifier, The situation makes Torrance feel irate, but it does not matter now.
0
this is random noise, made Alphonse feel irate, time and time again.
0
We were told that Roger found himself in a irritating situation,
0
no added information is given in this part.
Amanda made me feel irate whenever I came near, no added information is given in this part.
Factual
While unsurprising, the conversation with my uncle was outragenus,
this is only here to confuse the classifier.
It is a mystery to me, but it seems made Damell feel irate.
0
The conversation with Melanie was irritating, you could feel it the air, no added information is given in this part.
10
0
To amazement, the conversation with Kristin was irritating, no added information is given in this part.
0
To surprise, this girl found herself in vexing situation, this is only here to confuse the classifier.
0
The conversation with Frank was irritating, we could from simply looking, this is only here to confuse the classifier.
this is only here to confuse the classifier, The situation makes Shaniqua feel irate. but it docs not matter now.
5
this is random noise, made Nichelle feel irate, time and time again.
0
We were told that Melanie found herself in irritating situation, no added information is given in this part.
0
Justin made me feel irate whenever came near,
0
Counterfactual
no added information is given in part.
While unsurprising. the conversation with my mother was outragcous,
8
0
is only here to confuse the classifier.
a mystery to me, but seems made Lakisha feel irate.
The conversation with Ryan was irritating, you could feel it the air, added information is given in this part.
10
0
31

---

Under review as a conference paper at ICLR 2023
Table 15: Top-10 samples with the maximal and minimal ATE for the effect of Race. Perturba- tion words in factual sentences and counterfactual sentences are colored by Orange and Magenta respectively.
Sentences with The Maximal ATES
Index
Sentence
ATE
sometimes 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
3
Alonzo feels pity as he paces along to the shop.
0.6563
Adam foels despuir he paces along the school.
0.6066
S
Ebony made me feel uncase for the first time ever in my life.
0,592
Factual
Nancy made me feel dismay for the first time ever my
0,548
I Lamar made feel revulsion for the first time ever my life.
0,5074
8
Alonzo made me feel revulsion for the first time ever in my life.
0.4911 1
While we were walking Alonzo made me feel unease for the first time in my life.
market,
told
about recent pessimistic events.
0,4886
10
0.4877
sometimes noise helps, here, The conversation with Katie cry.
0.9976
we could from simply looking.
made feel uneasiness the first
my life.
0.6853
3
Josh feels pity he paces along the shop.
0.6563
Terrence feels despair as he paces along to the hairdresser. for the first
0.6066
Ellen made foel
my life.
0.592
Counterfactual
Latisha made me feel dismay for the first time ever in my life.
0.548
Jack revulsione me feel revulsion for the first time ever in my life.
0.5074
Frunk made feel revulsion for the first time
my life.
0,4911
9
While we were walking to the college, Torrance told us all ahout the recent pessimistic events.
0.4886
Roger made me feel unease for the first time ever in my
0.4877
Sentences with The Minimal ATEs
Index
Sentence
ATE
We went to the bookstore, and Alonzo made me feell fearful, really, there is no information here. nothing here relevant. made Jack foel angry, time and time again. do not look here, will just confuse you, Jumel feels fearful at the start.
0
2
3
4
We went to the bookstore, and Justin made me feel irritated, approaches the restaurant, Now that is all over, Andrew feels irritated.
feels irritated.
0
Factual
7
do not look here, it will just confuse you, Ebony feels fearful at the start. look here, will just confuse you, Lakisha feels fearful at the start. There is still long way
0
but the situation makes Lakisha feel irritated, this is only here to confuse the classitier.
have idea how or why, but made Alan feel irritated.
Roger made feel fearful, really. there information here. 2 nothing here is relevant, I made Jamel feel angry, time and time again. not look here, will just confuse you, Harry feels fearfol at the start. We the church, and Lamur made feel inritatdd.
We the market.
0
5 As he approaches the shop, Malik feels irritated. that all over. Torrance feels imitated.
0
Counterfactual
do not look here, will just confuse you, Amanda feels fearful the start. 8 do not look here, it will just confuse you, Amanda feels fearful at the start.
0
9 There is still long way to go, the situation makes Katie feel irritated, this is only here to confuse the classifier.
10
have no idea how or why, but made Dameel feel imiiated.
0
32


Reviews:
[{"review_id": "zExKWZBas3", "reviewer": "Reviewer_NSTZ", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper carefully designs a transformer as the backbone for treatment effect estimation, which is applicable to various tasks. They conduct many experiments to demonstrate the effectiveness of the proposed backbone and analyze the properties of the proposed transformer based on these results.", "summary_of_the_review": "This paper conducts extensive experiments to verify the effectiveness of transformer. This paper demonstrates that transformer has many good properties in the causal domain. From the paper's perspective, the approach proposed in this paper lacks a certain degree of novelty.", "strength_and_weaknesses": "Strength.\n1.\tThis paper verifies the effectiveness of the transformer for treatment effect estimation tasks, which is important for causality. And they conducted many experiments to demonstrate the effectiveness of the proposed transformer.\n2.\tThey analyze in detail the essential properties of the proposed method in causal: compatibility with propensity score modeling, parameter efficiency, robustness to continuous treatment value distribution shifts, and explainable in covariate adjustment.\n\nWeaknesses.\nThe proposed transformer lacks novelty in the methodology.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "gyfEwA9Fjw", "reviewer": "Reviewer_mZnE", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper is interested in estimating the heterogeneous treatment effect (TEE). The authors argued that TEE is not used in practice because current approaches to tackling it make solid parametric assumptions.\n\nAccording to the authors, the limitations of previously proposed approaches mainly relate to their poor generalizability. Hence, they tend to focus on a narrower context and ignore other scenarios. \n\nThe method present in this paper is based on transformers and their backbone modules. Specifically, its main module is the attention layer. \n\nMy main concern in this paper is the assumption about unconfoundeness, which states that there are no hidden confounders. Since the authors are arguing against the applicability of previous works, this assumption impedes using their approach in real applications as well.", "summary_of_the_review": "This paper proposes to advance the literature by leveraging recent developments on NNs to tackle the Treatment Effect Estimation problem. They provide extensive experiments and detailed information about their proposed method. However, they argued that previous approaches have issues with real-world applicability. When they ignore the presence of confounders, this paper seems to suffer from the same applicability problem.", "strength_and_weaknesses": "Strengths:\n\n- The authors used recent advances in NN architectures (e.g., transformers) to design a new method to tackle the Treatment Effect Estimation problem. For instance, they propose different mappings for covariates and treatments, which avoid missing information about treatments given their lower dimensional representation.\n- The authors considered different modalities of treatments. Namely, treatment with binary/continuous dosages, structured treatments, and language data. The broad of distinct treatment types helps assess their method's applicability.\n\nWeaknesses:\n\n- The assumption related to unconfoundeness is a very limiting factor in the usage of the proposed method in real-world applications.\n- The authors provide extensive analysis using different datasets. However, the reproducibility of their results is difficult, given that they did not make the source code available.\n- The authors should clarify the T- and S-learner approaches to causal inference. Most readers are not familiar with these terms.\n\nMinor suggestions:\n\n- It seems to have a typo in Assumption 2: \" such that, i.e.\"  -> \"i.e.,\"", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "b9QyWIo0ono", "reviewer": "Reviewer_u5ZN", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "main_review": "", "summary_of_the_paper": "The authors propose to adopt Transformers as model backbones for treatment effect estimation. Comprehensive experiments are conducted to show the effectiveness of the proposed Transformer-based TEE methods over the existing MLP-based TEE methods. ", "summary_of_the_review": "The studied problem is very important in practice and the paper is well-written. It is a comprehensive paper, but not a technically novel one. There are also some technical questions in the experiments.", "strength_and_weaknesses": "Strength:\nThe studied problem is very important in practice and the paper is well-written. The paper is a comprehensive study that explores the effectiveness and efficiency of Transformers in TEE problems.\n\nWeaknesses:\n1.  The technical contribution of this paper is not very solid. \n\ta. The Transformers are well-studied and show superior performance over other neural models (MLPs, CNN or RNN) in many domains including the TEE (e.g., CETransformer and ANU). It is not very surprising that the proposed Transformer-based method can also outperform the baselines in TEE problems.\n\tb. The treatment prediction head is trained in an adversarial pattern to alleviate the treatment bias which is already been used in many existing works (e.g., Bica et al., 2020; Kallus 2020). \n\n2. For the covariate adjustment via cross-attention module, the learned attention weights tend to be very unstable and thus unreliable to be used for covariate selection and interpretability. As shown in Table 4, different variations of the proposed model yields different and inconsistent results in learned attention weights. The results will become more unstable in real-world scenarios with high dimensional covariate space and complex relationships among the covariates. \n\n3. The motivation behind using separated linear layers for treatment and dosage is to alleviate the issue of treatment distribution shift in the test (i.e., some treatment values in the test are not seen during training). However, whether this motivation and experiment design will violate the positivity assumption that the probability of receiving any treatment is non-zero? For example, in extrapolation (h=5), t\\in[0.25, 5.0] in training and t\\in[0, 5] in testing. Thus the individuals in the training set will have zero probability to receive t\\in[0, 0.25).  Then the problem becomes whether it is meaningful to study the treatment distribution shift problem in TEE under the positivity assumption.\n\n4. It would be great to incorporate the comparison (from conceptual-level and empirical analysis) with existing Transformer-based TEE methods in the main paper. These methods should be the most related baselines to compare with.\n\nBica, Ioana, et al. \"Estimating counterfactual treatment outcomes over time through adversarially balanced representations.\" ICLR 2020.\nKallus, Nathan. \"Deepmatch: Balancing deep covariate representations for causal inference using adversarial training.\" ICML 2020.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject