SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Moving forward, since  $y_{i}^{T = 0}$  and  $y_{j}^{T = 1}$  in (11) are unavailable due to the missing counterfactual outcomes, the proposed proximal factual outcome regularizer (PFOR) uses their estimates instead. Specifically, let  $\hat{y}_i$  and  $\hat{y}_j$  be the estimates of  $y_{i}^{T = 0}$  and  $y_{j}^{T = 1}$ , respectively, PFOR refines (11) as
$$
\mathbf {D} _ {i j} ^ {\gamma} = \left\| r _ {i} - r _ {j} \right\| ^ {2} + \gamma \cdot \left[ \left\| \hat {y} _ {i} - y _ {j} \right\| ^ {2} + \left\| \hat {y} _ {j} - y _ {i} \right\| ^ {2} \right], \quad \hat {y} _ {i} = \phi_ {0} (r _ {i}), \quad \hat {y} _ {j} = \phi_ {1} (r _ {j}), \tag {12}
$$
Additional justifications, assumptions and limitations of PFOR are discussed in Appendix D.3.
3.4 Architecture of entire space counterfactual regression
The architecture of ESCFR is presented in Figure 1(b), where the covariate  $X$  is first mapped to the representations  $R$  with  $\psi(\cdot)$ , and then to the potential outcomes with  $\phi(\cdot)$ . The group discrepancy  $\mathbb{W}$  is calculated with the optimal transport equipping with the RMPR in (9) and PFOR in (12).
The learning objective is to minimize the risk of factual outcome estimation and the group discrepancy. Given mini-batch distributions  $\hat{\mathbb{P}}^{T = 1}(x)$  and  $\hat{\mathbb{P}}^{T = 0}(x)$  in Definition 3.1, the risk of factual outcome estimation following [68] can be formulated as
$$
\mathcal {L} _ {\mathrm {F}} (\psi , \phi) := \mathbb {E} _ {x _ {i} \in \hat {\mathbb {P}} ^ {T = 1} (x)} \| \phi_ {1} (\psi (x _ {i})) - y _ {i} \| ^ {2} + \mathbb {E} _ {x _ {j} \in \hat {\mathbb {P}} ^ {T = 0} (x)} \| \phi_ {0} (\psi (x _ {j})) - y _ {j} \| ^ {2}, \tag {13}
$$
where  $y_{i}$  and  $y_{j}$  are the factual outcomes for the corresponding treatment groups. The discrepancy is:
$$
\mathcal {L} _ {\mathrm {D} ^ {\gamma}} ^ {\epsilon , \kappa} (\psi) := \mathbb {W} ^ {\epsilon , \kappa , \gamma} \left(\hat {\mathbb {P}} _ {\psi} ^ {\mathrm {T} = 1} (r), \hat {\mathbb {P}} _ {\psi} ^ {\mathrm {T} = 0} (r)\right), \tag {14}
$$
which is in general the optimal transport with RMPR in Definition 3.2, except for the unit-wise distance  $\mathbf{D}^{\gamma}$  calculated with the PFOR in (12). Finally, the overall learning objective of ESCFR is
$$
\mathcal {L} _ {\mathrm {E S C F R}} ^ {\epsilon , \kappa , \gamma , \lambda} := \mathcal {L} _ {\mathrm {F}} (\psi , \phi) + \lambda \cdot \mathcal {L} _ {\mathrm {D} \gamma} ^ {\epsilon , \kappa} (\psi), \tag {15}
$$
where  $\lambda$  controls the strength of distribution alignment,  $\epsilon$  controls the entropic regularization in (5),  $\kappa$  controls RMPR in (9), and  $\gamma$  controls PFOR in (12). The learning objective (15) mitigates the selection bias following Theorem 3.1 and handles the MSE and UCE issues.
The optimization procedure consists of three steps as summarized in Algorithm 3 (see Appendix B). First, compute  $\pi^{\epsilon ,\kappa ,\gamma}$  by solving the OT problem in Definition 3.2 with Algorithm 2 (see Appendix B), where the unit-wise distance is calculated with  $\mathbf{D}^{\gamma}$ . Second, compute the discrepancy in (14) as  $\langle \pi^{\epsilon ,\kappa ,\gamma},\mathbf{D}^{\gamma}\rangle$ , which is differentiable to  $\psi$  and  $\phi$ . Finally, calculate the overall loss in (15) and update  $\psi$  and  $\phi$  with stochastic gradient methods.
Datasets. Missing counterfactuals impede the evaluation of PEHE using observational benchmarks. Therefore, experiments are conducted on two semi-synthetic benchmarks [67, 85], i.e., IHDP and ACIC. Specifically, IHDP is designed to estimate the effect of specialist home visits on infants' potential cognitive scores, with 747 observations and 25 covariates. ACIC comes from the collaborative perinatal project [51], and includes 4802 observations and 58 covariates.
Table 1: Performance (mean±std) on the PEHE and AUUC metrics. “*” marks the baseline estimators that ESCFR outperforms significantly at p-value  $< 0.05$  over paired samples t-test.
Dataset Model	ACIC (PEHE)	IHDP (PEHE)	ACIC (AUUC)	IHDP (AUUC)
In-sample	Out-sample	In-sample	Out-sample	In-sample	Out-sample	In-sample	Out-sample
OLS	3.749±0.080*	4.340±0.117*	3.856±6.018	5.674±9.026	0.843±0.007	0.496±0.017*	0.652±0.050	0.492±0.032*
R.Forest	3.597±0.064*	3.399±0.165*	2.635±3.598	4.671±9.291	0.902±0.016	0.702±0.026*	0.736±0.142	0.661±0.259
S.Learner	3.572±0.269*	3.636±0.254*	1.706±1.600*	3.038±5.319	0.905±0.041	0.627±0.014*	0.633±0.183	0.702±0.330
T.Learner	3.429±0.142*	3.566±0.248*	1.567±1.136*	2.730±3.627	0.846±0.019	0.632±0.020*	0.651±0.179	0.707±0.333
TARNet	3.236±0.266*	3.254±0.150*	0.749±0.291	1.788±2.812	0.886±0.046	0.662±0.014*	0.654±0.184	0.711±0.329
C.Forest	3.449±0.101*	3.196±0.177*	4.018±5.602*	4.486±8.677	0.717±0.005*	0.709±0.018*	0.643±0.141	0.695±0.294
k-NN	5.605±0.168*	5.892±0.138*	2.208±2.233*	4.319±7.336	0.892±0.007*	0.507±0.034*	0.725±0.142	0.668±0.299
O.Forest	8.094±4.669*	4.148±2.224*	2.605±2.418*	3.136±5.642	0.744±0.013	0.699±0.022*	0.664±0.157	0.702±0.325
PSM	5.228±0.154*	5.094±0.301*	3.219±4.352*	4.634±8.574	0.884±0.010	0.745±0.021	0.740±0.149	0.681±0.253
BNN	3.345±0.233*	3.368±0.176*	0.709±0.330	1.806±2.837	0.882±0.033	0.645±0.013*	0.654±0.184	0.711±0.329
CFR-MMD	3.182±0.174*	3.357±0.321*	0.777±0.327	1.791±2.741	0.871±0.032	0.659±0.017*	0.655±0.183	0.710±0.329
CFR-WASS	3.128±0.263*	3.207±0.169*	0.657±0.673	1.704±3.115	0.873±0.029	0.669±0.018*	0.656±0.187	0.715±0.311
ESCFR	2.252±0.297	2.316±0.613	0.502±0.252	1.282±2.312	0.796±0.030	0.754±0.021	0.665±0.166	0.719±0.311
Baselines. We consider three groups of baselines. (1) Statistical estimators: least square regression with treatment as covariates (OLS), random forest with treatment as covariates (R.Forest), a single network with treatment as covariates (S.learner [36]), separate neural regressors for each treatment group (T.learner [36]), TARNet [67]; (2) Matching estimators: propensity score match with logistic regression (PSM [60]), k-nearest neighbor (k-NN [18]), causal forest (C.Forest [73]), orthogonal forest (O.Forest [73]); (3) Representation-based estimators: balancing neural network (BNN [31]), counterfactual regression with MMD (CFR-MMD) and Wasserstein discrepancy (CFR-WASS) [67].
Training protocol. A fully connected neural network with two 60-dimensional hidden layers is selected to instantiate the representation mapping  $\psi$  and the factual outcome mapping  $\phi$  for ESCFR and other neural network based baselines. To ensure a fair comparison, all neural models are trained for a maximum of 400 epochs using the Adam optimizer, with the patience of early stopping being 30. The learning rate and weight decay are set to  $1e^{-3}$  and  $1e^{-4}$ , respectively. Other settings of optimizers follow Kingma and Ba [35]. We fine-tune hyperparameters within the range in Figure 5, validate performance every two epochs, and save the optimal model for test.
Evaluation protocol. The PEHE in (3) is the primary metric for performance evaluation [67, 85]. However, it is unavailable in the model selection phase due to missing counterfactuals. As such, we use the area under the uplift curve (AUUC) [2] to guide model selection, which evaluates the ranking performance of the CATE estimator and can be computed without counterfactual outcomes. Although AUUC is not commonly used in treatment effect estimation, we report it as an auxiliary metric for reference. The within-sample and out-of-sample results are computed on the training and test set, respectively, following the common settings [44, 45, 67, 85].
Table 1 compares ESCFR and its competitors. Main observations are noted as follows.
- Statistical estimators demonstrate competitive performance on the PEHE metric, with neural estimators outperforming linear and random forest methods due to their superior ability to capture non-linearity. In particular, TARNet, which combines the advantages of T-learner and S-learner, achieves the best overall performance among statistical estimators. However, the circumvention to treatment selection bias results in inferior performance.
- Matching methods such as PSM exhibit strong ranking performance, which explains their popularity in counterfactual ranking practice. However, their relatively poor performance on the PEHE metric limits their applicability in counterfactual estimation applications where accuracy of CATE estimation is prioritized, such as advertising systems.
- Representation-based methods mitigate the treatment selection bias and enhance overall performance. In particular, CFR-WASS reaches an out-of-sample PEHE of 3.207 on ACIC, advancing most statistical methods. However, it utilizes the vanilla Wasserstein discrepancy, wherein the MSE and UCE issues impede it from solving the treatment selection bias. The proposed ESCFR
Table 2: Ablation study (mean±std) on the ACIC benchmark. “*” marks the variants that ESCFR outperforms significantly at p-value  $< 0.01$  over paired samples t-test.
In-sample	Out-sample
SOT	RMPR	PFOR	PEHE	AUUC	PEHE	AUUC
✘	✘	✘	3.2367±0.2666*	0.8862+0.0462	3.2542+0.1505*	0.6624+0.0149*
✓	✘	✘	3.1284±0.2638*	0.8734+0.0291	3.2073+0.1699*	0.6698+0.0187*
✓	✓	✘	2.6459+0.2747*	0.8356+0.0286	2.7688±0.4009	0.7099+0.0157*
✓	✘	✓	2.5705±0.3403*	0.8270±0.0341	2.6330±0.4672	0.7110±0.0287*
✓	✓	✓	2.2520±0.2975	0.7968±0.0307	2.3165+0.6136	0.7542±0.0202
achieves significant improvement over most metrics compared with various prevalent baselines<sup>3</sup>. Combined with the comparisons above, we attribute its superiority to the proposed RMPR and PFOR regularizers, which accommodate ESCFR to the situations where MSE and UCE exist.
In this section, to further verify the effectiveness of individual components, an ablation study is conducted on the ACIC benchmark in Table 2. Specifically, ESCFR first augments TARNet with the stochastic optimal transport to align the confounders in the representation space, as described in Section 3.1, which reduces the out-of-sample PEHE from 3.254 to 3.207. Subsequently, it mitigates the MSE issue with RMPR as per Section 3.2 and the UCE issue with PFOR as per Section 3.3, which reduces the out-of-sample PEHE to 2.768 and 2.633, respectively. Finally, ESCFR combines the RMPR and PFOR in a unified framework as detailed in Section 3.4, which further reduces the value of PEHE and advances the best performance of other variants.
4.4 Analysis of relaxed mass-preserving regularizer
Most prevalent methods fail to handle the label imbalance and mini-batch outliers in Figure 2(b-c). Figure 4 shows the transport plan generated with RMPR in the same situations, where RMPR alleviates the MSE issue in both bad cases. Initially, RMPR with  $\kappa = 50$  presents similar matching scheme with Figure 2, since in this setting the loss of marginal constraint is strong, and the solution is thus similar to that of the vanilla OT problem in Definition 2.3; RMPR with  $\kappa = 10$  further looses the marginal constraint and avoids the incorrect matching of units with different outcomes; RMPR with  $\kappa = 2$  further gets robust to the outlier's interference and correctly matches the remaining units. This success is attributed to the relaxation of the mass-preserving constraint according to Section 3.2.
Notably, RMPR does not transport all mass of a unit. The closer the unit is to the overlapping zone in a batch, the greater the mass is transferred. That
is, RMPR adaptively matches and pulls closer units that are close to the overlapping region, ignoring outliers, which mitigates the bias of causal inference methods in cases where the positivity assumption does not strictly hold. Current approaches mainly achieve it by manually cleaning the data or dynamically weighting the units [32], while RMPR naturally implements it via the soft penalty in (9).
We further investigate the performance of RMPR under different batch sizes and  $\kappa$  in Appendix D.2 to verify the effectiveness of RMPR more thoroughly.
(a)  $\kappa = 50$
Figure 4: Geometric interpretation of OT plan with RMPR under the outcome imbalance (upper) and outlier (down) settings. The dark area indicates the transported mass of a unit, i.e., marginal of the transport matrix  $\pi$ . The light area indicates the total mass.
(b)  $\kappa = 10$
${}^{3}$  An exception would be the within-sample AUUC,which is reported over training data and thus easy to be overfitted. This metric is not critical as the factual outcomes are typically unavailable in the inference phase. We mainly rely on out-of-sample AUUC instead to evaluate the ranking performance and perform model selection.
Figure 5: Parameter sensitivity of ESCFR, where the lines and error bars indicate the mean values and  $90\%$  confidence intervals, respectively. (a) Impact of alignment strength  $(\lambda)$ . (b) Impact of entropic regularization strength  $\epsilon$ . (c) Impact of PFOR strength  $\gamma$ $(\times 10^{3})$ . (d) Impact of RMPR strength  $\kappa$ .
4.5 Parameter sensitivity study
In this section, we investigate the role of four critical hyperparameters within the ESCFR framework—namely,  $\lambda$ ,  $\epsilon$ ,  $\kappa$ , and  $\gamma$ . These hyperparameters have a profound impact on learning objectives and the performance, as substantiated by the experimental results presented in Figure 5.
First, we vary  $\lambda$  to investigate the efficacy of stochastic optimal transport. Our findings suggest that a gradual increase in the value of  $\lambda$  leads to consistent improvements. For instance, the out-of-sample PEHE diminishes from 3.22 at  $\lambda = 0$  to roughly 2.85 at  $\lambda = 1.5$ . Nonetheless, overly emphasizing distribution balancing within a multi-task learning framework can result in compromised factual outcome estimation and, consequently, suboptimal CATE estimates. The next parameter scrutinized is  $\epsilon$ , which serves as an entropic regularizer. Our study shows that a larger  $\epsilon$  accelerates the computation of optimal transport discrepancy [24]. However, this computational advantage comes at the cost of a skewed transport plan, as evidenced by an increase in out-of-sample PEHE with oscillations.
Following this, we explore the hyperparameters  $\gamma$  and  $\kappa$ , associated with PFOR and RMPR, respectively. Extending upon vanilla optimal transport, we incorporate unit-wise distance via (12) and assess the impact of PFOR. Generally, the inclusion of PFOR positively influences CATE estimation. However, allocating excessive weight to proximal outcome distance compromises the performance, as it dilutes the role of the unit-wise distance in the representation space. Concurrently, we modify the transport problem with (9) to examine the implications of RMPR. Relieving the mass-preserving constraints through RMPR significantly enhances CATE estimation, but an excessively low value of  $\kappa$  hampers performance as it fails to ensure that the representations across treatment groups are pulled closer together in the optimal transport paradigm.
Current research aims to alleviate treatment selection bias through the balancing of distributions between treated and untreated groups. These balancing techniques can be broadly categorized into three categories: reweighting-based, matching-based, and representation-based methods.
Reweighting-based methods mainly employ propensity scores to achieve global balance between groups. The core procedure comprises two steps: the estimation of propensity scores and the construction of unbiased estimators. Propensity scores are commonly estimated using logistic regression [11, 20, 37, 88]. To enhance the precision of these estimates, techniques such as feature selection [69, 76, 77], joint optimization [88, 89], and alternative training techniques [92] have been adopted. The unbiased estimator is exemplified by the inverse propensity score method [59], which inversely re-weights each units with the estimated propensity scores. While theoretically unbiased, it suffers from high variance with low propensity and bias with incorrect propensity estimates [39, 40, 74]. To alleviate these drawbacks, doubly robust estimators and variance reduction techniques have been introduced [26, 42, 58]. Nevertheless, these methods remain constrained by their reliance on propensity scores, affecting their efficacy in real-world applications.
Matching-based methods aim to match comparable units from different groups to construct locally balanced distributions. The key distinctions between representative techniques [8, 46, 60, 79] lie in their similarity measures. A notable exemplar is the propensity score matching approach [60], which
computes unit (dis)similarity based on estimated propensity scores. Notably, Tree-based methods [73] can also be categorized as matching approaches, but use adaptive similarity measures. However, these techniques are computationally intensive, limiting their deployment in large-scale operations.
Representation-based methods aim to construct a mapping to a feature space where distributional discrepancies are minimized. The central challenge lies in the accurate calculation of such discrepancies. Initial investigations focused on maximum mean discrepancy and vanilla Wasserstein discrepancy [31, 67], later supplemented by local similarity preservation [85, 86], feature selection [13, 27], representation decomposition [27, 80] and adversarial training [87] mechanisms. Despite their success, they fail under specific but prevalent conditions, such as outlier fluctuations [23] and unlabeled confounders [90], undermining the reliability of calculated discrepancies.
The recent exploration of optimal transport in causality [71] has spawned innovative reweighting-based [21], matching-based [10] and representation-based methods [44, 45]. For example, Li et al. [45] utilize OT to align factual and counterfactual distributions; Li et al. [44] and Ma et al. [48] use OT to reduce confounding bias. Despite these advancements, they largely adhere to the vanilla Kantorovich problem that corresponds to the canonical Wasserstein discrepancy, akin to Shalit et al. [67]. Adapting OT problems to meet the unique needs of treatment effect estimation remains an open area for further research.
Due to the effectiveness of mitigating treatment selection bias, representation learning has been the primary approach to estimating individual treatment effect. However, existing methods neglect the mini-batch sampling effects and unobserved confounders, which hinders them from handling the treatment selection bias. A principled approach named ESCFR is devised based on a generalized OT problem. Extensive experiments demonstrate that ESCFR can mitigate MSE and UCE issues, and achieve better performance compared with prevalent baseline models.
Looking ahead, two research avenues hold promise for further exploration. The first avenue explores the use of normalizing flows for representation mapping, which offers the benefit of invertibility [12] and thus aligns with the assumptions set forth by Shalit et al. [67]. The second avenue aims to apply our methodology to real-world industrial applications, such as bias mitigation in recommendation systems [74], an area currently dominated by high-variance reweighting methods.
This work is supported by National Key R&D Program of China (Grant No. 2021YFC2101100), National Natural Science Foundation of China (62073288, 12075212, 12105246, 11975207) and Zhejiang University NGICS Platform.
[77] Haotian Wang, Kun Kuang, Long Lan, Zige Wang, Wanrong Huang, Fei Wu, and Wenjing Yang. Out-of-distribution generalization with causal feature separation. IEEE Trans. Knowl. Data Eng., 2023.
[78] Christopher KI Williams and Carl Edward Rasmussen. Gaussian processes for machine learning, volume 2. MIT press Cambridge, MA, 2006.
[79] Anpeng Wu, Kun Kuang, Ruoxuan Xiong, Bo Li, and Fei Wu. Stable estimation of heterogeneous treatment effects. In ICML, volume 202 of Proceedings of Machine Learning Research, pages 37496-37510. PMLR, 2023.
[80] Anpeng Wu, Junkun Yuan, Kun Kuang, Bo Li, Runze Wu, Qiang Zhu, Yueting Zhuang, and Fei Wu. Learning decomposed representations for treatment effect estimation. IEEE Trans. Knowl. Data Eng., 35(5):4989-5001, 2023.
[81] Peng Wu, Haoxuan Li, Yuhao Deng, Wenjie Hu, Quanyu Dai, Zhenhua Dong, Jie Sun, Rui Zhang, and Xiao-Hua Zhou. On the opportunity of causal learning in recommendation systems: Foundation, estimation, prediction and challenges. In IJCAI, pages 23-29, 2022.
[82] Zeke Xie, Issei Sato, and Masashi Sugiyama. A diffusion theory for deep learning dynamics: Stochastic gradient descent exponentially favors flat minima. In ICLR, 2020.
[83] Renjun Xu, Pelen Liu, Yin Zhang, Fang Cai, Jindong Wang, Shuoying Liang, Heting Ying, and Jianwei Yin. Joint partial optimal transport for open set domain adaptation. In *IJCAI*, pages 2540–2546, 2020.
[84] Mengyue Yang, Xinyu Cai, Furui Liu, Weinan Zhang, and Jun Wang. Specify robust causal representation from mixed observations. In KDD, pages 2978-2987. ACM, 2023.
[85] Liuyi Yao, Sheng Li, Yaliang Li, Mengdi Huai, Jing Gao, and Aidong Zhang. Representation learning for treatment effect estimation from observational data. In NeurIPS, pages 2638-2648, 2018.
[86] Liuyi Yao, Sheng Li, Yaliang Li, Mengdi Huai, Jing Gao, and Aidong Zhang. ACE: adaptively similarity-preserved representation learning for individual treatment effect estimation. In ICDM, pages 1432-1437, 2019.
[87] Jinsung Yoon, James Jordon, and Mihaela van der Schaar. GANITE: estimation of individualized treatment effects using generative adversarial nets. In ICLR, 2018.
[88] Wenhao Zhang, Wentian Bao, Xiao-Yang Liu, Keping Yang, Quan Lin, Hong Wen, and Ramin Ramezani. Large-scale causal approaches to debiasing post-click conversion rate estimation with multi-task learning. In WWW, pages 2775-2781, 2020.
[89] Yang Zhang, Dong Wang, Qiang Li, Yue Shen, Ziqi Liu, Xiaodong Zeng, Zhiqiang Zhang, Jinjie Gu, and Derek F Wong. User retention: A causal approach with triple task modeling. In IJCAI, 2021.
[90] Jiajing Zheng. Sensitivity Analysis for Causal Inference with Unobserved Confounding. University of California, Santa Barbara, 2021.
[91] Jiajing Zheng, Alexander D'Amour, and Alexander Franks. Copula-based sensitivity analysis for multi-treatment causal inference with unobserved confounding. CoRR, abs/2102.09412, 2021.
[92] Ziwei Zhu, Yun He, Yin Zhang, and James Caverlee. Unbiased implicit recommendation and propensity estimation via combinational joint learning. In RecSys, pages 551-556, 2020.