PaperlayerView exact revision on arXiv

Source-linked AI summary

Exploring Transformer Backbones for Heterogeneous Treatment Effect Estimation

Yi-Fan Zhang, Hanlin Zhang, Zachary C. Lipton, Li Erran Li, Eric P. Xing

arXiv:2202.01336v5cs.LG

TL;DR

Treatment effect estimators often use treatment-specific designs that do not sufficiently exploit similarities among potential outcomes, limiting their applicability across settings. TransTEE combines treatment and covariate embeddings with attention, and the paper reports gains across diverse benchmarks and treatment types, including utility in auditing language-model predictions.

  • Problem

    Existing treatment effect methods do not sufficiently exploit structural similarities among potential outcomes and often rely on treatment-specific designs that generalize poorly across settings.

    Source: §1 · pp. 1–2

  • Method

    TransTEE is an S-learner that embeds treatments and covariates and uses attention to model treatment and treatment-covariate interactions.

    Source: §1 · p. 2

  • Results

    TransTEE shows significant gains across discrete, continuous, and structured treatment settings on benchmarks including IHDP, News, and TCGA, and is effective in auditing fair language-model predictions.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The paper reports TransTEE as a flexible estimator with compatibility with propensity score modeling, parameter efficiency, explainable covariate adjustment, and robustness to continuous treatment shifts.

    Source: §1 · p. 2

    TransTEE, like almost all causal inference methods on observational data, relies on the ignorability assumption.

    Source: §6 · p. 11

Abstract

from arXiv · show

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 work uses 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 that significantly outperforms competitive baselines in 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

1 Introduction

Treatment effect estimation is difficult to use in practice because existing approaches do not adequately exploit similarities among potential outcomes and often depend on treatment-specific designs. TransTEE uses treatment and covariate embeddings with attention, and is reported to perform well across varied treatment settings and benchmarks.

1 Introduction

  • The paper frames treatment effect estimation as a central problem in healthcare, social science, and econometrics, with covariate and treatment representation choices that differ from supervised learning.The introduction notes that the task can be framed as a missing-data problem under certain conditions.Source: §1 · p. 1
  • Existing methods do not sufficiently exploit structural similarities among potential outcomes, and treatment-specific designs limit generalizability to settings such as text and graph inputs.The introduction also identifies parameter inefficiency and brittleness to treatment-distribution shifts as limitations of MLP-based models.Source: §1 · pp. 1–2
  • TransTEE embeds treatments and covariates as an S-learner and uses attention to model their interactions, enabling adaptive covariate selection and flexibility across heterogeneous treatment types.The motivating example describes higher attention weights for pre-treatment covariates and confounders.Source: §1 · p. 2, Figure 1 · p. 2
  • TransTEE shows significant performance gains on discrete, continuous, or structured treatments across benchmarks including IHDP, News, and TCGA, and is effective in auditing fair language-model predictions.The experiments cover six benchmarks and four treatment types.Source: §1 · p. 2

2 Related Work

Prior neural estimators use domain-specific architectures and covariate representations, while attention-based approaches to treatment effect estimation have explored different modeling choices. The paper positions TransTEE as an S-learner that represents treatments as well as covariates.

2 Related Work

  • Neural estimators use domain-specific designs that limit transfer across treatment settings; for example, VCNet [47] maps real-valued treatments with a constant function and can struggle under treatment shifts.TARNet [60] requires an accurate estimate of the treatment-value interval.Source: §2 · p. 3
  • Prior estimators can lose covariate interactions by embedding covariates as a single representation, while DRNet [58] may require n×m branches for n treatments and m associated dosages.The passage describes the branch requirement as a rigid constraint on extrapolation and potentially infeasible with observational data.Source: §2 · p. 3
  • Among attention-based TEE approaches, CETransformer [27] uses a T-learner and does not represent treatments, while ANU maps covariates into a latent space.The paper contrasts these approaches with TransTEE's S-learner formulation.Source: §2 · pp. 3–4

3 Problem Statement and Assumptions

The paper formulates treatment effect estimation using observed covariates, treatments, and outcomes, with the goal of estimating dose-response quantities and their discrete-treatment special cases. Causal interpretation depends on ignorability and overlap assumptions.

3 Problem Statement and Assumptions

  • The target is the Average Dose Response Function, µ(x, t) = E[Y |X = x, do(T = t)], with ATE and ITE as discrete-treatment quantities.The treatment may be binary, continuous, or structured, and can include an associated dosage.Source: §3 · p. 4
  • Observational data omit counterfactual outcomes and ground-truth causal effects, motivating assumptions that permit causal interpretation of statistical estimands.The paper introduces strongly ignorable conditions for this purpose.Source: §3 · p. 4
  • Ignorability rules out hidden confounders, while positivity requires a nonzero chance of observing each treatment for each covariate profile.The paper states that ignorability supports identification and positivity ensures a chance of seeing units in every treatment group.Source: §3 · p. 4

4 TransTEE: Transformers as Treatment Effect Estimators

TransTEE uses treatment and covariate embeddings with attention layers to model interactions without treatment-specific architectures. Its design includes treatment and covariate self-attention and treatment-to-covariate cross-attention for covariate adjustment.

4 TransTEE: Transformers as Treatment Effect Estimators

  • TransTEE applies attention-based design choices to multiple, continuous, structured, interacting, and dosage-associated treatments without ad hoc architectures such as multiple branches.The schematic comparison says it handles the scenarios without additional parameter overhead.Source: §4 · p. 4, Figure 2 · p. 5

4.1 Covariate and Treatment Embedding Layers

  • Treatment and dosage values are projected into learnable embeddings, while each covariate receives its own embedding to preserve information about individual covariates.Structural treatments use language models or graph neural networks for their embeddings.Source: §4.1 · p. 5

4.2 Covariate and Treatment Self-Attention

  • TransTEE processes covariate embeddings independently and uses self-attention to model covariate interactions; treatment self-attention enables interactions among simultaneously applied treatments.The paper identifies simultaneous prescriptions or economic measures as examples of such treatment combinations.Source: §4.2 · pp. 5–6

4.3 Treatment-Covariate Cross-Attention

  • Cross-attention uses the treatment representation to attend to covariates, supporting adjustment of covariates and confounders; the outcome-regression objective is mean squared error.The paper reports that the resulting adjustment empirically reflects the Disjunctive Cause Criteria.Source: §4.3 · p. 6
  • The paper presents its embedding and attention design as a flexible route to structured high-dimensional data and applications such as auditing language models.It notes that guidance on designing these modules for causal inference remains underexplored.Source: §4.3 · p. 6

5 Experimental Results

Experiments evaluate TransTEE across treatment types and datasets, finding strong performance under continuous-treatment extrapolation, structured-treatment selection bias, and dosage variation. Analyses also examine covariate adjustment, parameter efficiency, propensity regularization, and auditing language-model predictions.

5.1 Experimental Settings

  • Experiments cover synthetic and semi-synthetic continuous-treatment data, TCGA dosage and structured-treatment datasets, SW structured treatments, and EEEC language-model auditing.Baselines vary by treatment type, and evaluation uses test-set AMSE for continuous and binary treatments and PEHE-based measures for structured treatments.Source: §5.1 · p. 7

5.2 Case Study and Numerical Results

  • For continuous treatments, TR performs better under smaller distribution shifts, whereas PTR is preferable under greater shifts; neither improves discrete-treatment performance.Adding adversarial training and TR or PTR can further reduce TransTEE's continuous-treatment estimation error.Source: §5.2 · p. 8
  • TransTEE achieves lower estimation error while preserving continuous ADRF continuity across treatment intervals, including extrapolation settings where several baselines are brittle.Vanilla TransTEE is competitive with additional propensity-model MLE; TARNet and DRNet produce discontinuous ADRF estimators, while VCNet performs well only for t ∈[0, 1].Source: §5.2 · p. 9
  • Considering both metrics, TransTEE outperforms baselines by a large margin across the entire evaluated treatment-selection-bias range on the SW and TCGA datasets.The evaluated ranges were κ ∈ [0, 40] for SW and κ ∈ [0, 0.5] for TCGA.Source: §5.2 · pp. 9–10

5.3 Analysis

  • Cross-attention assigns higher weights to IHDP confounders and lower weights to noisy covariates, while TransTEE outperforms baselines across noisy-covariate counts.With TR and PTR, it assigns more weight to confounders and less to noisy covariates; TR also reduces weights on treatment-only covariates more than PTR.Source: §5.3 · p. 10, Figure 4 · p. 10, Table 4 · p. 10
  • The analysis flags a boundary: when unconfoundedness is violated, controlling instrumental variables can incur bias in treatment-effect estimation.Source: §5.3 · p. 11
  • TransTEE uses fewer parameters than baselines across the evaluated settings, while adding treatment branches to many baselines incurs parameter redundancy.Figure 4(c) reports parameter counts on four datasets with a base-2 logarithmic y-axis.Source: Figure 4 · p. 10, §5.3 · p. 11

5.4 Empirical Study on Pre-trained Language Models

  • On EEEC, TransTEE estimates effects of domain-specific perturbations to BERT predictions and outperforms MLP-adapted baselines without specialized downstream fine-tuning.The authors present this as a potential surrogate-modeling use, including actionable contrast sets and lightweight estimates for fairness constraints.Source: §5.4 · p. 11

6 Concluding Remarks

The conclusion presents attention mechanisms as versatile design choices for treatment-effect estimation and identifies reliance on ignorability as a limitation. The related work connects propensity-score modeling and domain adaptation to causal inference, while the illustration traces TransTEE from embedded inputs through attention to outcome prediction.

6 Concluding Remarks

  • TransTEE relies on the ignorability assumption; future work includes more complex causal graphs and integrating domain knowledge to mitigate potential societal impacts.The paper cautions that covariate adjustment without accounting for the causal graph can yield inaccurate or biased effects.Source: §6 · p. 11

A Extended Related Work

  • Propensity-score modeling reflects treatment-assignment policy and supports confounding adjustment; double-robust estimators and targeted regularization address misspecification of either the outcome or propensity model [41].The discussion notes that many traditional approaches are restricted to binary treatments.Source: A Extended Related Work · p. 18
  • Causal inference and domain adaptation share interests in out-of-distribution robustness and balanced covariate representations; for continuous treatments, TransTEE connects to continuously indexed domain adaptation [69].Explicit propensity-score modeling also connects causal inference with transfer learning, while hidden stratification in domain generalization is left for future work.Source: A Extended Related Work · p. 18
  • Unlike ANU, TransTEE applies self-attention separately to covariate and treatment representations before cross-attention, and supports multi-treatment settings.The authors argue that separate self-attention layers strengthen representations for high-dimensional data such as texts, images, and graphs.Source: A Extended Related Work · p. 18

B An Illustrative Example

  • The illustrative workflow embeds covariates, treatments, and dosages, then uses self-attention to model contextual and multi-treatment interactions.Cross-attention models treatment-covariate interactions, after which pooled embeddings feed a linear outcome predictor; the example gives confounders higher attention than irrelevant covariates.Source: B An Illustrative Example · pp. 18–19

C Details and Discussions about Propensity Score Modelling

TransTEE’s propensity-score modeling combines a propensity network with minimax training and treatment regularization to address selection bias, including for continuous treatments. The analysis describes how its regularizers align treatment-feature moments across covariate representations.

C Details and Discussions about Propensity Score Modelling

  • TransTEE’s general framework can be combined with traditional propensity-score methods, including covariate adjustment, matching, stratification, reweighting, and g-computation.The paper attributes this compatibility to using off-the-shelf propensity-score regularization to balance covariate representations.Source: C Details and Discussions about… · p. 19
  • Minimax training of the propensity network is presented as a way to reduce selection bias by encouraging independence between treatment and covariate representations.The paper also describes the network as refining treatment representations and supporting covariate adjustment.Source: C Details and Discussions about… · p. 19
  • For continuous treatments, Treatment Regularization (TR) matches the propensity score’s mean to the treatment, while PTR matches both mean and variance.PTR models the mean and variance of the estimated treatment; matching the whole distribution may be preferable in some cases.Source: C Details and Discussions about… · p. 20
  • At equilibrium, TR’s optimum aligns conditional and marginal treatment-feature means; PTR’s optimum aligns both means and variances across covariate representations.The analysis assumes the outcome prediction model is fixed.Source: C Details and Discussions about… · pp. 21–22

D Analysis of the Failure Cases over Treatment Distribution Shifts

The section examines why fixed treatment intervals can undermine extrapolation and reports that learnable treatment embeddings improve performance in both interpolation and extrapolation settings. Large gaps between training and test treatment intervals remain challenging for existing methods.

D Analysis of the Failure Cases over Treatment Distribution Shifts

  • Learnable treatment embeddings improve TransTEE’s performance in both interpolation and extrapolation, although large gaps between training and test intervals remain challenging.The section identifies fixed-interval treatment mappings as a source of poor extrapolation.Source: D Analysis of the Failure Cases over… · p. 22
  • For h = 5, VCNet has infinite estimation error, partly because its hand-crafted projection matrix processes only values near [0, 1].With training on t ∈[0, 4] and testing on t ∈[0, 5], the comparison evaluates extrapolation beyond the training interval.Source: D Analysis of the Failure Cases over… · p. 22, Table 6 · p. 23
  • When training covers t ∈[0, 1.75] and testing includes t ∈(1.75, 2.0], discrete approximation methods round unseen treatments to the nearest represented value.The paper describes this as an extrapolation dilemma caused by fixed treatment mappings.Source: D Analysis of the Failure Cases over… · p. 22
  • The fixed-interval assumption is described as suboptimal for extrapolation, and prior methods are said to adopt intervals such as [0, 1] or another fixed range.The section motivates a general TEE framework that does not require treatment-specific architectural designs.Source: D Analysis of the Failure Cases over… · pp. 22–23

E Additional Experimental Setups

The experimental assets are described as MIT-licensed, and the experiments used a specified workstation configuration.

E Additional Experimental Setups

E.1 Detail Evaluation Metrics.

The setup describes additional-run tuning and testing with AMSE, constrains model size relative to baselines, and points to a table of TransTEE architecture details.

E.1 Detail Evaluation Metrics.

  • Hyperparameters for synthetic and semi-synthetic datasets were tuned using 20 additional runs.Each run simulated data, randomly split it into training and testing sets, and evaluated test AMSE.Source: §E.2 · p. 24
  • TransTEE’s model size was less than or similar to the baselines in all experiments.Source: §E.2 · p. 24
  • Table 7 provides TransTEE architecture details, with p denoting the number of covariates.Source: Table 7 · p. 24

E.3 Simulation details.

The simulations cover tabular, dosage-associated, structural, and text-based treatment settings, using both semi-synthetic and real-world covariates. Dataset-specific setups vary treatment ranges, dosage assignment, and structural interventions.

IHDP

  • IHDP supplies 25 covariates, 747 observations, and binary treatments, with generated continuous treatments considered over intervals [0, h].The section also describes an IHDP treatment-generation procedure following VCNet [47].Source: §E.3 · p. 25

News dataset

  • The News benchmark uses 3,000 NY Times news items and varies continuous treatment intervals across [0, h].Treatment and outcome generation follows VCNet [47].Source: §E.3 · p. 25

TCGA (D)

  • TCGA (D) uses real cancer covariates and three dosage-associated treatments, with dosage selection controlled by a beta distribution and treatment assignment bias by κ.The stated default settings are κ = 2 and α = 2.Source: §E.3 · p. 26, Table 9 · p. 26

Small-World

  • Small-World experiments use 20-dimensional covariates and graph interventions with sampled node counts, degrees, and edge-rewiring probabilities.The in-sample and out-sample sets contain 1,000 and 500 units, respectively.Source: §E.3 · p. 27

TCGA (S)

  • TCGA (S) combines cancer gene-expression covariates with molecular graph treatments drawn from QM9 and eight molecule properties.The passage describes 9,659 gene-expression measurements and 10,000 sampled molecules per run.Source: §E.3 · p. 27

EEEC

  • The EEEC task uses counterfactual sentences that alter names and pronouns to represent gender, or names to represent race.The corpus contains 33,738 sentences constructed from 42 templates and four mood states.Source: §E.3 · pp. 27–28

F.1 Comparision between TransTEE and ANU [74]

The comparison evaluates TransTEE against ANU in IHDP treatment settings and a real-world sentence-based application. TransTEE is reported to perform better than ANU in both binary and continuous treatment settings and in the sentence-based estimation task.

F.1 Comparision between TransTEE and ANU [74]

  • TransTEE performs better than ANU in both binary and continuous treatment settings on IHDP.ANU is adjusted to have a similar parameter count; treatment embeddings make ANU more robust than VCNet and DRNet under treatment shift, but TransTEE still performs better.Source: §F.1 · p. 28, Table 11 · p. 29
  • On the real-world long-sentence task, TransTEE achieves better estimation results than baselines, while ANU without self-attention produces inaccurate estimates.The comparison is reported with Table 11.Source: §F.1 · p. 29, Table 11 · p. 29

F.2 Additional Numerical Results and Ablation Studies

The additional experiments examine robustness, propensity-score regularization, and training dynamics, alongside sentence-level ATE examples and covariate-importance interpretations. Results report gains from suitable regularization and robustness to noisy covariates, while TransTEE extrapolates across treatment ranges and highlights confounders in its learned weights.

F.2 Additional Numerical Results and Ablation Studies

  • Both TR and PTR improve estimation over unregularized TransTEE at suitable strengths, with the best performance at λ = 0.5 for both methods.The passage cautions that training the treatment predictor and feature encoder together in a zero-sum game can be difficult and unstable.Source: §F.2 · p. 29
  • TransTEE’s estimation error remains low as outcome-only noisy covariates increase, while VCNet worsens and TARNet and DRNet improve.TransTEE is reported to outperform the baselines by a large margin across the tested covariate counts.Source: §F.2 · p. 29
  • TransTEE extrapolates across tested treatment ranges without knowing the upper bound h in advance, whereas TARNet and DRNet fail to extrapolate and VCNet’s mapping is limited to t < 2.The figure descriptions report these patterns on IHDP and News; TARNet and DRNet worsen as h increases with the same number of branches.Source: Figure 11 · p. 30, Figure 13 · p. 31, Table 12 · p. 31
  • Adversarial propensity-score modeling can outperform maximum-likelihood modeling, especially in extrapolation settings.The passage attributes reduced selection bias and improved estimation performance to TR and PTR.Source: §F.2 · p. 30
  • TR and PTR help TransTEE converge faster and achieve lower test MSE than the compared regularization alternatives, with PTR performing best overall.TransTEE+MTL quickly reaches low treatment-estimation error but later oscillates and converges with high error; PTR’s low treatment-regression error indicates accurate propensity-score estimates.Source: §F.2 · p. 31

F.3 Showcase of sentences and counterfactual counterparts with the maximal/minimal ATEs.

  • High-ATE sentences tend to be shorter than low-ATE sentences for gender effects, averaging 11 versus 17.6 words, and for race effects, averaging 12 versus 14.7.The passage also notes differing sentence patterns and perturbation-word positions for maximal and minimal race ATEs.Source: §F.3 · p. 32, Table 14 · p. 36, Table 15 · p. 37

G Remarks on Interpretability

  • TransTEE assigns more weight to confounders than to other covariates in the reported experiments, offering feature-importance scores for benchmarking treatment-effect estimators.The paper frames this as an initial step toward interpretability and notes that evaluation depends on models, tasks, and input spaces.Source: G Remarks on Interpretability · p. 32

Related papers

Loading 2202.01336v5…