[论文解读] Exploring Transformer Backbones for Heterogeneous Treatment Effect Estimation

Yifan Zhang, Hanlin Zhang+3 more|arXiv (Cornell University)|Feb 2, 2022

Advanced Causal Inference Techniques被引用 12

一句话总结

本文提出 TransTEE，一种基于 Transformer 主干网络的新型治疗效应估计器，通过自注意力机制建模治疗与协变量之间的交互作用，实现对离散、连续及结构化治疗的灵活、稳健且可解释的异质性治疗效应估计。TransTEE 在多种基准测试中表现优于强基线模型，并支持可解释的协变量调整与预训练语言模型的审计。

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 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

研究动机与目标

• 为解决现有基于 MLP 的治疗效应估计器在分布偏移下泛化能力差、参数效率低且易失效的问题。
• 探索使用 Transformer 主干网络作为 MLP 的灵活替代方案，以建模异质性治疗效应估计中复杂的治疗-协变量交互作用。
• 开发一种与倾向得分建模兼容、参数高效且对连续治疗分布偏移具有鲁棒性的统一框架。
• 通过基于注意力的特征重要性实现可解释的协变量调整，支持现实场景下的因果推断。
• 在预训练语言模型的公平性与社会影响审计中展示其实用价值。

提出的方法

• TransTEE 采用单一 Transformer 主干网络联合嵌入所有治疗与协变量，避免多任务架构，支持端到端学习。
• 自注意力机制动态建模治疗与协变量之间的交互作用，实现自适应协变量选择与混杂控制。
• 模型利用注意力权重识别并优先处理混杂因子与相关协变量，与条件不可忽略性的析取因果准则保持一致。
• 设计为与倾向得分建模兼容，提升观察性研究中模型的鲁棒性与可解释性。
• 框架支持表格数据与结构化输入（如文本、图），具备广泛适用性。
• 引入代理建模任务，以评估在合成基准之外的真实世界实用性。

实验结果

研究问题

• RQ1 Transformer 主干网络能否有效建模涵盖离散、连续及结构化治疗等多种治疗类型的异质性治疗效应？
• RQ2 与最先进的基于 MLP 的基线模型相比，TransTEE 在性能、鲁棒性与参数效率方面表现如何？
• RQ3 TransTEE 中的注意力机制在多大程度上能提供可解释且有意义的协变量调整，以支持因果推断？
• RQ4 TransTEE 是否可有效用于审计预训练语言模型的公平性与社会影响？
• RQ5 TransTEE 对连续治疗值的分布偏移具有多强的鲁棒性？

主要发现

• TransTEE 在涵盖四种治疗类型的六个基准测试中显著优于竞争性基线模型，包括 IHDP、News 和 TCGA，展现出卓越的泛化能力与性能。
• 在半合成与真实世界数据集上，TransTEE 在估计平均治疗效应（ATE）方面达到最先进水平，代理任务中的最大 ATE 达到 0.9976。
• TransTEE 对连续治疗值的分布偏移具有鲁棒性，性能稳定，而基于 MLP 的模型则出现明显退化。
• TransTEE 中的注意力权重能有效识别并优先处理混杂因子与关键协变量，提供与因果识别原则一致的可解释特征重要性评分。
• TransTEE 有效支持对预训练语言模型的审计，通过扰动分析揭示社会偏见，极端 ATE 值达到 0.9976 和 0.0，证明其在公平性评估中的实用价值。
• 该框架具有参数高效特性，且与倾向得分建模兼容，显著提升其在现实世界因果推断场景中的实际部署能力。