[论文解读] Optimal Transport for Treatment Effect Estimation

Hao Wang, Zhichao Chen+8 more|arXiv (Cornell University)|Oct 27, 2023

Advanced Causal Inference Techniques被引用 5

一句话总结

该论文提出ESC-FR，一种基于最优传输的新型条件平均处理效应（CATE）估计方法，通过松弛的质量保持正则化器和近端事实结果正则化器，分别解决小批量采样效应（MSE）和未观测混杂效应（UCE）。该方法通过在整个数据空间中利用随机最优传输对齐处理组，实现了观测数据上的最先进性能。

ABSTRACT

Estimating conditional average treatment effect from observational data is highly challenging due to the existence of treatment selection bias. Prevalent methods mitigate this issue by aligning distributions of different treatment groups in the latent space. However, there are two critical problems that these methods fail to address: (1) mini-batch sampling effects (MSE), which causes misalignment in non-ideal mini-batches with outcome imbalance and outliers; (2) unobserved confounder effects (UCE), which results in inaccurate discrepancy calculation due to the neglect of unobserved confounders. To tackle these problems, we propose a principled approach named Entire Space CounterFactual Regression (ESCFR), which is a new take on optimal transport in the context of causality. Specifically, based on the framework of stochastic optimal transport, we propose a relaxed mass-preserving regularizer to address the MSE issue and design a proximal factual outcome regularizer to handle the UCE issue. Extensive experiments demonstrate that our proposed ESCFR can successfully tackle the treatment selection bias and achieve significantly better performance than state-of-the-art methods.

研究动机与目标

• 为解决从观测数据中进行CATE估计时因处理组与对照组之间协变量偏移而产生的处理选择偏差。
• 克服小批量采样效应（MSE），即现有方法中异常值或不平衡小批量导致的分布差异估计偏差。
• 缓解未观测混杂效应（UCE），即在标准无混杂性假设下导致差异计算偏差的问题。
• 构建一个基于随机最优传输的严谨因果推断框架，该框架作用于整个数据空间而非局部小批量。
• 通过整合考虑潜在未观测混杂因子的近端正则化器，改进反事实结果估计。

提出的方法

• 提出一种基于随机最优传输的全空间反事实回归（ESCFR）框架，用于从观测数据中估计CATE。
• 引入一种松弛的质量保持正则化器，通过降低对小批量采样伪影（如异常值和不平衡）的敏感性，稳定分布对齐。
• 设计一种近端事实结果正则化器，通过建模未观测混杂因子对结果的单调影响，来考虑其潜在影响。
• 采用广义Sinkhorn差异来计算整个数据空间上的分布差异，避免依赖可能有偏的小批量估计。
• 采用对偶优化框架，联合学习表示映射和反事实结果，确保分布对齐与结果保真度。
• 应用未观测混杂因子与结果之间的单调性假设，以缩小可能的传输策略集合，提升估计精度。

实验结果

研究问题

• RQ1 如何通过缓解小批量采样效应对分布差异估计的影响，来提升观测数据中CATE估计的性能？
• RQ2 未观测混杂因子在多大程度上会偏差现有基于最优传输的CATE估计器？该偏差如何纠正？
• RQ3 能否设计一种基于原理的最优传输框架，使其作用于整个数据空间而非局部小批量，以提升鲁棒性？
• RQ4 在存在未观测混杂的情况下，引入事实结果的近端正则化器如何提升反事实预测的准确性？
• RQ5 能否利用未观测混杂因子与结果之间的单调性假设，来减少传输策略选择中的模糊性？

主要发现

• ESC-FR在多个基准数据集上显著优于现有最先进方法，在CATE估计中展现出对分布偏移更强的鲁棒性。
• 松弛的质量保持正则化器有效降低了小批量中异常值和结果不平衡的影响，稳定了训练过程并提升了对齐精度。
• 近端事实结果正则化器减少了未观测混杂因子带来的偏差，尤其在单调性假设成立时，显著提升了反事实预测的准确性。
• 实证结果表明，ESC-FR在CATE估计中相比现有基线方法实现了更低的均方误差（MSE），尤其在高未观测混杂场景下表现更优。
• 该方法在多样化数据集上均表现出一致的性能提升，包括具有复杂非线性结果关系和隐藏混杂因子的数据集。
• 消融研究证实，质量保持正则化器与近端正则化器均对性能有显著贡献，验证了其设计合理性。