SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Hao Wang<sup>1</sup> Zhichao Chen<sup>1</sup> Jiajun Fan<sup>2</sup> Haoxuan Li<sup>3</sup> Tianqiao Liu<sup>4</sup> Weiming Liu<sup>1</sup> Quanyu Dai<sup>5*</sup> Yichao Wang<sup>5</sup> Zhenhua Dong<sup>5</sup> Ruiming Tang<sup>5</sup>  
<sup>1</sup>Zhejiang University <sup>2</sup>Tsinghua University <sup>3</sup>Peking University <sup>4</sup>Purdue University <sup>5</sup>Huawei Noah's Ark Lab <sup>haohaow@zju.edu.cn quanyu.dai@connect.polyu.hk
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.
The estimation of the conditional average treatment effect (CATE) through randomized controlled trials serves as a cornerstone in causal inference, finding applications in diverse fields such as health care [63], e-commerce [2, 43, 81], and education [15]. Although randomized controlled trials are often considered the gold standard for CATE estimation [38, 53], the associated financial and ethical constraints can make them infeasible in practice [39, 79]. Therefore, alternative approaches that rely on observational data have gained prominence. For instance, drug developers could utilize post-marketing monitoring reports to assess the efficacy of new medications rather than conducting costly clinical A/B trials. With the growing access to observational data, estimating CATE from observational data has attracted intense research interest [41, 80, 84].
Estimating CATE with observational data has two main challenges: (1) missing counterfactuals, i.e., only one factual outcome out of all potential outcomes can be observed; (2) treatment selection bias, i.e., individuals have different propensities for treatment, leading to non-random treatment assignments and a resulting covariate shift between the treated and untreated groups [77, 79]. Traditional meta-learners [36] handle the issue of missing counterfactuals by breaking down the CATE estimation into more tractable sub-problems of factual outcome estimation. However, the covariate shift makes it difficult to generalize the factual outcome estimators trained within respective treatment groups to the entire population and thus biases the derived CATE estimator.
Beginning with counterfactual regression [67] and its remarkable performance, there are various attempts that handle the selection bias by minimizing the distribution discrepancy between the
(a) Handling treatment selection bias with  $\psi (\cdot)$ .
(b) Architecture of ESCFR.
Figure 1: Overview of handling treatment selection bias with ESCFR. The red and blue colors signify the treated and untreated groups, respectively. (a) The treatment selection bias manifests as a distribution shift between  $X_{1}$  and  $X_{0}$ . The scatters and curves represent the units and the fitted outcome mappings. (b) ESCFR mitigates the selection bias by aligning units from different treatment groups in the representation space:  $R = \psi(X)$ , which enables  $\phi_{1}$  and  $\phi_{0}$  to generalize across groups.
treatment groups in the representation space [13, 27, 85, 86]. However, two critical issues with these methods have long been neglected, which impedes them from handling the treatment selection bias. The first issue is the mini-batch sampling effects (MSE). Specifically, current representation-based methods compute the distribution discrepancy within mini-batches instead of the entire data space, making it vulnerable to bad sampling cases. For example, the presence of an outlier in a mini-batch can falsely increase the estimate of the distribution discrepancy, thereby misleading the update of the estimator. The second issue is the unobserved confounder effects (UCE). Specifically, current approaches mainly operate under the unconfoundedness assumption [48] and ignore the pervasive influence of unobserved confounders, which makes the resulting estimators biased given the existence of unobserved confounders.
Contributions and outline. In this paper, we propose an effective CATE estimator based on optimal transport, namely Entire Space CounterFactual Regression (ESCFR), which handles both the MSE and UCE issues with a generalized sinkhorn discrepancy. Specifically, after presenting preliminaries in Section 2, we redefine the problem of CATE estimation within the framework of stochastic optimal transport in Section 3.1. We subsequently showcase the MSE issue faced by existing approaches in Section 3.2 and introduce a mass-preserving regularizer to counter this issue. In Section 3.3, we explore the UCE issue and propose a proximal factual outcome regularizer to mitigate its impact. The architecture and learning objectives of ESCFR are elaborated upon in Section 3.4, followed by the presentation of our experimental results in Section 4.
2.1 Causal inference from observational data
This section formulates preliminaries in observational causal inference. We first formalize the fundamental elements in Definition 2.1 following the general notation convention<sup>2</sup>.
Definition 2.1. Let  $X$  be the random variable of covariates, with support  $\mathcal{X}$  and distribution  $\mathbb{P}(x)$ ; Let  $R$  be the random variable of induced representations, with support  $\mathcal{R}$  and distribution  $\mathbb{P}(r)$ ; Let  $Y$  be the random variable of outcomes, with support  $\mathcal{Y}$  and distribution  $\mathbb{P}(y)$ ; Let  $T$  be the random variable of treatment indicator, with support  $\mathcal{T} = \{0,1\}$  and distribution  $\mathbb{P}(T)$ .
In the potential outcome framework [61], a unit characterized by the covariates  $x$  has two potential outcomes, namely  $Y_{1}(x)$  given it is treated and  $Y_{0}(x)$  otherwise. The CATE is defined as the conditionally expected difference of potential outcomes as follow:
$$
\tau (x) := \mathbb {E} \left[ Y _ {1} - Y _ {0} \mid x \right], \tag {1}
$$
2We use uppercase letters, e.g.,  $X$  to denote a random variable, and lowercase letters, e.g.,  $x$  to denote an associated specific value. Letters in calligraphic font, e.g.,  $\mathcal{X}$  represent the support of the corresponding random variable, and  $\mathbb{P}()$  represents the probability distribution of the random variable, e.g.,  $\mathbb{P}(X)$ .
where one of these two outcomes is always unobserved in the collected data. To address such missing counterfactuals, the CATE estimation task is commonly decomposed into outcome estimation subproblems that are solvable with supervised learning method [36]. For example, T-learner models the factual outcomes  $Y$  for units in the treated and untreated groups separately; S-learner regards  $T$  as a covariate, and then models  $Y$  for units in all groups simultaneously. The CATE estimate is then the difference of the estimated outcomes when  $T$  is set to treated and untreated.
Definition 2.2. Let  $\psi : \mathcal{X} \to \mathcal{R}$  be a mapping from support  $\mathcal{X}$  to  $\mathcal{R}$ , i.e.,  $\forall x \in \mathcal{X}$ ,  $r = \psi(x) \in \mathcal{R}$ . Let  $\phi : \mathcal{R} \times \mathcal{T} \to \mathcal{Y}$  be a mapping from support  $\mathcal{R} \times \mathcal{T}$  to  $\mathcal{Y}$ , i.e., it maps the representations and treatment indicator to the corresponding factual outcome. Denote  $Y_{1} = \phi_{1}(R)$  and  $Y_{0} = \phi_{0}(R)$ , where we abbreviate  $\phi(R, T = 1)$  and  $\phi(R, T = 0)$  to  $\phi_{1}(R)$  and  $\phi_{0}(R)$ , respectively, for brevity.
TARNet [67] combines T-learner and S-learner to achieve better performance, which consists of a representation mapping  $\psi$  and an outcome mapping  $\phi$  as defined in Definition 2.2. For an unit with covariates  $x$ , TARNet estimates the CATE as:
$$
\hat {\tau} _ {\psi , \phi} (x) = \hat {Y} _ {1} - \hat {Y} _ {0}, \quad \hat {Y} _ {1} = \phi_ {1} (\psi (x)), \quad \hat {Y} _ {0} = \phi_ {0} (\psi (x)), \tag {2}
$$
where  $\psi$  is trained over all individuals,  $\phi_1$  and  $\phi_0$  are trained over the treated and untreated group, respectively. Finally, the quality of CATE estimation is evaluated with the precision in estimation of heterogeneous effect (PEHE) metric:
$$
\epsilon_ {\mathrm {P E H E}} (\psi , \phi) := \int_ {\mathcal {X}} \left(\hat {\tau} _ {\psi , \phi} (x) - \tau (x)\right) ^ {2} \mathbb {P} (x) d x. \tag {3}
$$
However, according to Figure 1(a), the treatment selection bias causes a distribution shift of covariates across groups, which misleads  $\phi_{1}$  and  $\phi_0$  to overfit their respective group's properties and generalize poorly to the entire population. Therefore, the CATE estimate  $\hat{\tau}$  by these methods would be biased.
2.2 Discrete optimal transport and Sinkhorn divergence
Optimal transport (OT) quantifies distribution discrepancy as the minimum transport cost, offering a tool to quantify the selection bias in Figure 1(a). Monge [49] first formulated OT as finding an optimal mapping between two distributions. Kantorovich [34] further proposed a more applicable formulation in Definition 2.3, which can be seen as a generalization of the Monge problem.
Definition 2.3. For empirical distributions  $\alpha$  and  $\beta$  with  $n$  and  $m$  units, respectively, the Kantorovich problem aims to find a feasible plan  $\pi \in \mathbb{R}_+^{n\times m}$  which transports  $\alpha$  to  $\beta$  at the minimum cost:
$$
\mathbb {W} (\alpha , \beta) := \min  _ {\boldsymbol {\pi} \in \Pi (\alpha , \beta)} \left\langle \mathbf {D}, \boldsymbol {\pi} \right\rangle , \Pi (\alpha , \beta) := \left\{\boldsymbol {\pi} \in \mathbb {R} _ {+} ^ {n \times m}: \boldsymbol {\pi} \mathbf {1} _ {m} = \mathbf {a}, \boldsymbol {\pi} ^ {\mathrm {T}} \mathbf {1} _ {n} = \mathbf {b} \right\}, \tag {4}
$$
where  $\mathbb{W}(\alpha, \beta) \in \mathbb{R}$  is the Wasserstein discrepancy between  $\alpha$  and  $\beta$ ;  $\mathbf{D} \in \mathbb{R}_+^{n \times m}$  is the unit-wise distance between  $\alpha$  and  $\beta$ , which is implemented with the squared Euclidean distance;  $\mathbf{a}$  and  $\mathbf{b}$  indicate the mass of units in  $\alpha$  and  $\beta$ , and  $\Pi$  is the feasible transportation plan set which ensures the mass-preserving constraint holds.
Since exact solutions to (4) often come with high computational costs [5], researchers would commonly add an entropic regularization to the Kantorovich problem as follow:
$$
\mathbb {W} ^ {\epsilon} (\alpha , \beta) := \left\langle \mathbf {D}, \boldsymbol {\pi} ^ {\epsilon} \right\rangle , \quad \boldsymbol {\pi} ^ {\epsilon} := \underset {\boldsymbol {\pi} \in \Pi (\alpha , \beta)} {\arg \min } \left\langle \mathbf {D}, \boldsymbol {\pi} \right\rangle - \epsilon \mathrm {H} (\boldsymbol {\pi}), \quad \mathrm {H} (\boldsymbol {\pi}) := - \sum_ {i, j} \boldsymbol {\pi} _ {i, j} (\log (\boldsymbol {\pi} _ {i, j}) - 1), \tag {5}
$$
which makes the problem  $\epsilon$ -convex and solvable with the Sinkhorn algorithm [19]. The Sinkhorn algorithm only consists of matrix-vector products, making it suited to be accelerated with GPUs.
In this section, we present the Entire Space CounterFactual Regression (ESCFR) approach, which leverages optimal transport to tackle the treatment selection bias. We first illustrate the stochastic optimal transport framework for distribution discrepancy quantification across treatment groups, and demonstrate its efficacy for improving the performance of CATE estimators. Based on the framework, we then propose a relaxed mass-preserving regularizer to address the sampling effect, and a proximal factual outcome regularizer to handle the unobserved confounders. We finally open a new thread to summarize the model architecture, learning objectives, and optimization algorithm.
3.1 Stochastic optimal transport for counterfactual regression
Representation-based approaches mitigate the treatment selection bias by calculating distribution discrepancy in the representation space and then minimizing it. OT is a preferred method to quantify the discrepancy due to its numerical advantages and flexibility over competitors. It is numerically stable in the cases where other methods, such as  $\phi$ -divergence (e.g., Kullback-Leibler divergence), fails [64]. Compared with adversarial discrepancy measures [7, 33, 87], the it can be calculated efficiently and integrated naturally with the traditional supervised learning framework.
Formally, we denote the OT discrepancy between treatment groups as  $\mathbb{W}\left(\mathbb{P}_{\psi}^{\mathrm{T}=1}(r), \mathbb{P}_{\psi}^{\mathrm{T}=0}(r)\right)$ , where  $\mathbb{P}_{\psi}^{\mathrm{T}=1}(r)$  and  $\mathbb{P}_{\psi}^{\mathrm{T}=0}(r)$  are the distributions of representations in treated and untreated groups, respectively, induced by the mapping  $r = \psi(x)$ . The discrepancy is differentiable with respect to  $\psi$  [24], and thus can be minimized by updating the representation mapping  $\psi$  with gradient-based optimizers.
Definition 3.1. Let  $\hat{\mathbb{P}}^{T = 1}(x) \coloneqq \{x_i^{T = 1}\}_{i = 1}^n$  and  $\hat{\mathbb{P}}^{T = 0}(x) \coloneqq \{x_i^{T = 0}\}_{i = 1}^m$  be the empirical distributions of covariates at a mini-batch level, which contain  $n$  treated units and  $m$  untreated units, respectively;  $\hat{\mathbb{P}}_{\psi}^{T = 1}(r)$  and  $\hat{\mathbb{P}}_{\psi}^{T = 0}(r)$  be that of representations induced by the mapping  $r = \psi(x)$  in Definition 2.2.
However, since prevalent neural estimators mainly update parameters with stochastic gradient methods, only a fraction of the units is accessible within each iteration. A shortcut in this context is to calculate the discrepancy at a stochastic mini-batch level:
$$
\hat {\mathbb {W}} _ {\psi} := \mathbb {W} \left(\hat {\mathbb {P}} _ {\psi} ^ {\mathrm {T} = 1} (r), \hat {\mathbb {P}} _ {\psi} ^ {\mathrm {T} = 0} (r)\right). \tag {6}
$$
The effectiveness of this shortcut is investigated by Theorem 3.1 (refer to Appendix A.4 for proof), which demonstrates that the PEHE can be optimized by iteratively minimizing the estimation error of factual outcomes and the mini-batch group discrepancy in (6).
Theorem 3.1. Let  $\psi$  and  $\phi$  be the representation mapping and factual outcome mapping, respectively;  $\hat{\mathbb{W}}_{\psi}$  be the group discrepancy at a mini-batch level. With the probability of at least  $1 - \delta$ , we have:
$$
\epsilon_ {\mathrm {P E H E}} (\psi , \phi) \leq 2 \left[ \epsilon_ {\mathrm {F}} ^ {T = 1} (\psi , \phi) + \epsilon_ {\mathrm {F}} ^ {T = 0} (\psi , \phi) + B _ {\psi} \hat {\mathbb {W}} _ {\psi} - 2 \sigma_ {Y} ^ {2} + \mathcal {O} \left(\frac {1}{\delta N}\right) \right], \tag {7}
$$
where  $\epsilon_{\mathrm{F}}^{T = 1}$  and  $\epsilon_{\mathrm{F}}^{T = 0}$  are the expected errors of factual outcome estimation,  $N$  is the batch size,  $\sigma_Y^2$  is the variance of outcomes,  $B_{\psi}$  is a constant term, and  $\mathcal{O}(\cdot)$  is a sampling complexity term.
3.2 Relaxed mass-preserving regularizer for sampling effect
Although Theorem 3.1 guarantees that the empirical OT discrepancy (6) bounds the PEHE, the sampling complexity term  $\mathcal{O}(\cdot)$  inspires us to investigate the potential risks of bad cases caused by stochastic sampling. Precisely, the term  $\mathcal{O}(\cdot)$  results from the discrepancy between the entire population and the sampled mini-batch units (see (30) and (32) in Appendix A) which is highly dependent on the uncontrollable sampling quality. Therefore, a reliable discrepancy measure should be robust to bad sampling cases, otherwise the resulting vulnerability would impede us from quantifying and minimizing the actual discrepancy.
(a) Ideal
(b) Imbalance
(c) Outlier
We consider three sampling cases in Figure 2, where the transport strategy is reasonable and applicable in the ideal sampling case in Figure 2(a). However, the transport strategy calculated with (6) is vulnerable to disturbances caused by non-ideal sampling conditions. For instance, Figure 2(b) reveals false matching of units with disparate outcomes in the
sampled mini-batch where the outcomes of the two groups are imbalanced; Figure 2(c) showcases false matching of the mini-batch outliers with normal units, causing a substantial disruption of the
Figure 2: Optimal transport plan (upper) and its geometric interpretation (down) in three cases, where the connection strength represents the transported mass. Different colors and vertical positions indicate different treatments and outcomes, respectively.
transport strategy for other units. Therefore, the vanilla OT in (6) fails to quantify the group discrepancy for producing erroneous transport strategies in non-ideal mini-batches, thereby misleading the update of the representation mapping  $\psi$ . We term this issue the mini-batch sampling effect (MSE).
Furthermore, this MSE issue is not exclusive to OT-based methods but is also prevalent in other representation-based approaches. Despite this, OT offers a distinct advantage: it allows the formalization of the MSE issue through its mass-preservation constraint, as indicated in (5). This constraint mandates a match for every unit in both groups, regardless of the real-world scenarios, complicating the transport of normal units and the computation of true group discrepancies. This issue is further exacerbated by small batch sizes. Such formalizability through mass-preservation provides a lever for handling the MSE issue, setting OT apart from other representation-based methods [85, 86].
An intuitive approach to mitigate the MSE issue is to relax the marginal constraint, i.e., to allow for the creation and destruction of the mass of each unit. To this end, inspired by unbalanced and weak transport theories [14, 65], a relaxed mass-preserving regularizer (RMPR) is devised in Definition 3.2. The core technical point is to replace the hard marginal constraint in (4) with a soft penalty in (9) for constraining transport strategy. On the basis, the stochastic discrepancy is calculated as
$$
\hat {\mathbb {W}} _ {\psi} ^ {\epsilon , \kappa} := \mathbb {W} ^ {\epsilon , \kappa} \left(\hat {\mathbb {P}} _ {\psi} ^ {\mathrm {T} = 1} (r), \hat {\mathbb {P}} _ {\psi} ^ {\mathrm {T} = 0} (r)\right), \tag {8}
$$
where the hard mass-preservation constraint is removed to mitigate the MSE issue. Building on Lemma 1 in Fatras et al. [23], we derive Corollary 3.1 to investigate the robustness of RMPR to sampling effects, showcasing that the effect of mini-batch outliers is upper bounded by a constant.
Definition 3.2. For empirical distributions  $\alpha$  and  $\beta$  with  $n$  and  $m$  units, respectively, optimal transport with relaxed mass-preserving constraint seeks the transport strategy  $\pi$  at the minimum cost:
$$
\mathbb {W} ^ {\epsilon , \kappa} (\alpha , \beta) := \left\langle \mathbf {D}, \pi \right\rangle , \pi := \underset {\pi} {\arg \min } \left\langle \mathbf {D}, \pi \right\rangle - \epsilon \mathrm {H} (\pi) + \kappa \left(\mathrm {D} _ {\mathrm {K L}} \left(\pi \mathbf {1} _ {m}, \mathbf {a}\right) + \mathrm {D} _ {\mathrm {K L}} \left(\boldsymbol {\pi} ^ {\mathrm {T}} \mathbf {1} _ {n}, \mathbf {b}\right)\right) \tag {9}
$$
where  $\mathbf{D} \in \mathbb{R}_+^{n \times m}$  is the unit-wise distance, and  $\mathbf{a}, \mathbf{b}$  indicate the mass of units in  $\alpha$  and  $\beta$ , respectively.
Corollary 3.1. For empirical distributions  $\alpha, \beta$  with  $n$  and  $m$  units, respectively, adding an outlier  $a'$  to  $\alpha$  and denoting the disturbed distribution as  $\alpha'$ , we have
$$
\mathbb {W} ^ {0, \kappa} \left(\alpha^ {\prime}, \beta\right) - \mathbb {W} ^ {0, \kappa} (\alpha , \beta) \leq 2 \kappa \left(1 - e ^ {- \sum_ {b \in \beta} \left(a ^ {\prime} - b\right) ^ {2} / 2 \kappa}\right) / (n + 1), \tag {10}
$$
which is upper bounded by  $2\kappa / (n + 1)$ .  $\mathbb{W}^{0,\kappa}$  is the unbalanced discrepancy as per Definition 3.2.
In comparison with alternative methods [9, 83] to relax the marginal constraint, the RMPR implementation in Definition 3.2 enjoys a collection of theoretical properties [65] and can be calculated via the generalized Sinkhorn algorithm [14]. The calculated discrepancy is differentiable w.r.t.  $\psi$  and thus can be minimized via stochastic gradient methods in an end-to-end manner.
3.3 Proximal factual outcome regularizer for unobserved confounders
Existing representation-based methods fail to eliminate the treatment selection bias due to the unobserved confounder effects (UCE). Beginning with CFR [67], the unconfoundedness assumption A.1 (see Appendix A) is often taken to circumvent the UCE issue [75]. In this context, given two units  $r_i \in \mathbb{P}_{\psi}^{T=1}(r)$  and  $r_j \in \mathbb{P}_{\psi}^{T=0}(r)$ , for instance, OT discrepancy in Definition 3.2 calculates the unit-wise distance as  $\mathbf{D}_{ij} := \| r_i - r_j \|^2$ . If Assumption A.1 holds, it eliminates the treatment selection
bias since it blocks the backdoor path  $X \to T$  in Figure 3(a) by balancing confounders in the latent space. However, Assumption A.1 is usually violated due to the existence of unobserved confounders as per Figure 3(b), which hinders existing methods from handling treatment selection bias since the backdoor path  $X' \to T$  is not blocked. The existence of unobserved covariates  $X'$  also makes the transport strategy in vanilla OT unidentifiable, which invalidates the calculated discrepancy.
(b)
Figure 3: Causal graphs with (a) and w/o (b) the unconfoundedness assumption. The shaded node indicates the hidden confounder  $X'$ .
To mitigate the effect of  $X'$ , we introduce a modification to the unit-wise distance. Specifically, we observe that given balanced  $X$  and identical  $T$ , the only variable reflecting the variation of  $X'$  is the
outcome  $Y$ . As such, resonating with Courty et al. [16], we calibrate the unit-wise distance  $\mathbf{D}$  with potential outcomes as follow:
$$
\mathbf {D} _ {i j} ^ {\gamma} = \left\| r _ {i} - r _ {j} \right\| ^ {2} + \gamma \cdot \left[ \left\| y _ {i} ^ {T = 0} - y _ {j} ^ {T = 0} \right\| ^ {2} + \left\| y _ {j} ^ {T = 1} - y _ {i} ^ {T = 1} \right\| ^ {2} \right], \tag {11}
$$
where  $\gamma$  controls the strength of regularization. The underlying regularization is straightforward: units with similar (observed and unobserved) confounders should also have similar potential outcomes. As such, for a pair of units with similar observed covariates, i.e.,  $\| r_i - r_j\|^2 \approx 0$ , if their potential outcomes given the same treatment  $t = \{0,1\}$  differ greatly, i.e.,  $\| y_i^t - y_j^t\| >> 0$ , their unobserved confounders should likewise differ significantly. The vanilla OT technique in (6) where  $\mathbf{D}_{ij} = \| r_i - r_j\|^2$  would incorrectly match this pair, generate a false transport strategy, and consequently misguide the update of the representation mapping  $\psi$ . In contrast, OT based on  $\mathbf{D}_{ij}^\gamma$  would not match this pair as the difference of unobserved confounders is compensated with that of potential outcomes.