Title: Optimal Transport for Treatment Effect Estimation
Venue: NeurIPS 2023
Accepted: yes
Keywords: treatment effect estimation

optimal transport

wasserstein

causal inference

counterfactual
Abstract: Estimating individual treatment effects from observational data is challenging due to treatment selection bias. Prevalent methods mainly mitigate this issue by aligning different treatment groups in the latent space, the core of which is the calculation of distribution discrepancy. However, two issues that are often overlooked can render these methods invalid:
(1) mini-batch sampling effects (MSE), where the calculated discrepancy is erroneous in non-ideal mini-batches with outcome imbalance and outliers;
(2) unobserved confounder effects (UCE), where the unobserved confounders are not considered in the discrepancy calculation.
Both of these issues invalidate the calculated discrepancy, mislead the training of estimators, and thus impede the handling of treatment selection bias.
To tackle these issues, we propose Entire Space CounterFactual Regression (ESCFR), which is a new take on optimal transport technology in the context of causality.
Specifically, based on the canonical optimal transport framework, 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 ESCFR estimates distribution discrepancy accurately, handles the treatment selection bias effectively, and outperforms prevalent competitors significantly.

Optimal Transport for Treatment Effect Estimation

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 E ntire S pace C ounter F actual R egression (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 masspreserving 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.

1 Introduction

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

Y0

ψ

R1

W

Y

L

T

^

^

R0

ϕ

1

Y1

ϕ

0

X1

ϕ

0

(R0

)

R0=

ψ

(X0

)

ϕ

1

(R1

)

Factual outcome

Covariate

Treatment

indicator

Factual

outcome

Factual

error

Representation

Outcome estimate

X1

R

(a) Handling treatment selection bias with ψ (⋅) .

Y0

ψ

R1

W

Y

L

T

^

^

R0

ϕ

1

Y1

ϕ

0

X1

ϕ

0

(R0

)

R0=

ψ

(X0

)

ϕ

1

(R1

)

Factual outcome

Covariate

Treatment

indicator

Factual

outcome

Factual

error

Representation

Outcome estimate

X1

R

(b) Architecture of ESCFR.

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 Preliminaries

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

Definition 2.1. Let X be the random variable of covariates, with support X and distribution P ( x ) ; Let R be the random variable of induced representations, with support R and distribution P ( r ) ; Let Y be the random variable of outcomes, with support Y and distribution P ( y ) ; Let T be the random variable of treatment indicator, with support T = { 0 , 1 } and distribution 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:

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 ψ ∶ X → R be a mapping from support X to R , i.e., ∀ x ∈ X , r = ψ ( x ) ∈ R . Let ϕ ∶ R×T → Y be a mapping from support R×T to Y , i.e., it maps the representations and treatment indicator to the corresponding factual outcome. Denote Y 1 = ϕ 1 ( R ) and Y 0 = ϕ 0 ( R ) , where we abbreviate ϕ ( R,T = 1 ) and ϕ ( R,T = 0 ) to ϕ 1 ( R ) and ϕ 0 ( R ) , respectively, for brevity.

TARNet [67] combines T-learner and S-learner to achieve better performance, which consists of a representation mapping ψ and an outcome mapping ϕ as defined in Definition 2.2. For an unit with covariates x , TARNet estimates the CATE as:

where ψ is trained over all individuals, ϕ 1 and ϕ 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:

However, according to Figure 1(a), the treatment selection bias causes a distribution shift of covariates across groups, which misleads ϕ 1 and ϕ 0 to overfit their respective group's properties and generalize poorly to the entire population. Therefore, the CATE estimate ˆ τ 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 α and β with n and m units, respectively, the Kantorovich problem aims to find a feasible plan π ∈ R n × m + which transports α to β at the minimum cost:

where W ( α,β ) ∈ R is the Wasserstein discrepancy between α and β ; D ∈ R n × m + is the unit-wise distance between α and β , which is implemented with the squared Euclidean distance; a and b indicate the mass of units in α and β , and Π 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:

which makes the problem ϵ -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.

3 Proposed method

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 ϕ -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 W ( P T = 1 ψ ( r ) , P T = 0 ψ ( r )) , where P T = 1 ψ ( r ) and P T = 0 ψ ( r ) are the distributions of representations in treated and untreated groups, respectively, induced by the mapping r = ψ ( x ) . The discrepancy is differentiable with respect to ψ [24], and thus can be minimized by updating the representation mapping ψ with gradient-based optimizers.

Definition 3.1. Let ˆ P T = 1 ( x ) ∶= { x T = 1 i } n i = 1 and ˆ P T = 0 ( x ) ∶= { x T = 0 i } i m = 1 be the empirical distributions of covariates at a mini-batch level, which contain n treated units and m untreated units, respectively; ˆ P T = 1 ψ ( r ) and ˆ P T = 0 ψ ( r ) be that of representations induced by the mapping r = ψ ( 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:

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 ψ and ϕ be the representation mapping and factual outcome mapping, respectively; ˆ W ψ be the group discrepancy at a mini-batch level. With the probability of at least 1 -δ , we have:

where ϵ T = 1 F and ϵ T = 0 F are the expected errors of factual outcome estimation, N is the batch size, σ 2 Y is the variance of outcomes, B ψ is a constant term, and O(⋅) 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 O(⋅) inspires us to investigate the potential risks of bad cases caused by stochastic sampling. Precisely, the term O(⋅) 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.

Weconsider 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

(a) Ideal

(b) Imbalance

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

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 α and β with n and m units, respectively, optimal transport with relaxed mass-preserving constraint seeks the transport strategy π at the minimum cost:

where D ∈ R n × m + is the unit-wise distance, and a , b indicate the mass of units in α and β , respectively. Corollary 3.1. For empirical distributions α,β with n and m units, respectively, adding an outlier a ′ to α and denoting the disturbed distribution as α ′ , we have

which is upper bounded by 2 κ /( n + 1 ) . W 0 ,κ 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 . ψ 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 ∈ P T = 1 ψ ( r ) and r j ∈ P T = 0 ψ ( r ) , for instance, OT discrepancy in Definition 3.2 calculates the unitwise distance as D ij ∶= ∥ r i -r j ∥ 2 . If Assumption A.1 holds, it eliminates the treatment selection

X

Y

T

(a)

X

Y

T

X

'

(b)

bias since it blocks the backdoor path X → 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 ′ → T is not blocked. The existence of unobserved covariates X ′ also makes the transport strategy in vanilla OT unidentifiable, which invalidates the calculated discrepancy.

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 D with potential outcomes as follow:

where γ 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 ≈ 0 , if their potential outcomes given the same treatment t = { 0 , 1 } differ greatly, i . e ., ∥ y t i -y t j ∥ >> 0 , their unobserved confounders should likewise differ significantly. The vanilla OT technique in (6) where 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 ψ . In contrast, OT based on D γ ij would not match this pair as the difference of unobserved confounders is compensated with that of potential outcomes.

Moving forward, since y T = 0 i and y T = 1 j in (11) are unavailable due to the missing counterfactual outcomes, the proposed proximal factual outcome regularizer (PFOR) uses their estimates instead. Specifically, let ˆ y i and ˆ y j be the estimates of y T = 0 i and y T = 1 j , respectively, PFOR refines (11) as

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 ψ (⋅) , and then to the potential outcomes with ϕ (⋅) . The group discrepancy 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 ˆ P T = 1 ( x ) and ˆ P T = 0 ( x ) in Definition 3.1, the risk of factual outcome estimation following [68] can be formulated as

where y i and y j are the factual outcomes for the corresponding treatment groups. The discrepancy is:

which is in general the optimal transport with RMPR in Definition 3.2, except for the unit-wise distance D γ calculated with the PFOR in (12). Finally, the overall learning objective of ESCFR is

where λ controls the strength of distribution alignment, ϵ controls the entropic regularization in (5), κ controls RMPR in (9), and γ 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 π ϵ,κ,γ by solving the OT problem in Definition 3.2 with Algorithm 2 (see Appendix B), where the unit-wise distance is calculated with D γ . Second, compute the discrepancy in (14) as ⟨ π ϵ,κ,γ , D γ ⟩ , which is differentiable to ψ and ϕ . Finally, calculate the overall loss in (15) and update ψ and ϕ with stochastic gradient methods.

4 Experiments

4.1 Experimental setup

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.

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 ψ and the factual outcome mapping ϕ 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 1 e -3 and 1 e -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].

4.2 Overall performance

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

achieves significant improvement over most metrics compared with various prevalent baselines 3 . 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.

4.3 Ablation study

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(bc). 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 κ = 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 κ = 10 further looses the marginal constraint and avoids the incorrect matching of units with different outcomes; RMPR with κ = 2 further gets robust to the outlier's interference and correctly matches the remaining units. This success is attributed to the relaxation of the masspreserving 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

(a)

κ

=

50

(b)

κ

=

10

(c)

κ

=

2

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 κ in Appendix D.2 to verify the effectiveness of RMPR more thoroughly.

0.0

0.5

1.0

1.5

2.0

3.0

λ

(a)

2.6

2.8

3.0

3.2

PEHE

In-sample

Out-sample

0.5

1.0

1.5

2.0

2.5

3.0

ε

(b)

2.8

3.0

3.2

PEHE

In-sample

Out-sample

0

0.1

0.5

1.0

5.0 10.0

γ

(c)

2.6

2.7

2.8

2.9

3.0

PEHE

In-sample

Out-sample

5.0

10.0 15.0 20.0 25.0

κ

(d)

2.6

2.7

2.8

2.9

3.0

PEHE

In-sample

Out-sample

4.5 Parameter sensitivity study

In this section, we investigate the role of four critical hyperparameters within the ESCFR framework-namely, λ , ϵ , κ , and γ . 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 λ to investigate the efficacy of stochastic optimal transport. Our findings suggest that a gradual increase in the value of λ leads to consistent improvements. For instance, the out-of-sample PEHE diminishes from 3.22 at λ = 0 to roughly 2.85 at λ = 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 ϵ , which serves as an entropic regularizer. Our study shows that a larger ϵ 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 γ and κ , 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 κ hampers performance as it fails to ensure that the representations across treatment groups are pulled closer together in the optimal transport paradigm.

5 Related works

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 reweightingbased [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.

6 Conclusion

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.

References

Jason M. Altschuler, Jonathan Weed, and Philippe Rigollet. Near-linear time approximation algorithms for optimal transport via sinkhorn iteration. In NeurIPS , pages 1964-1974, 2017.

Artem Betlei, Eustache Diemert, and Massih-Reza Amini. Uplift modeling with generalization guarantees. In SIGKDD , pages 55-65, 2021.

Jose H. Blanchet, Arun Jambulapati, Carson Kent, and Aaron Sidford. Towards optimal running times for optimal transport. CoRR , abs/1810.07717, 2018.

François Bolley, Arnaud Guillin, and Cédric Villani. Quantitative concentration inequalities for empirical measures on non-compact spaces. Probability Theory and Related Fields , 137(3): 541-593, 2007.

Nicolas Bonneel, Michiel van de Panne, Sylvain Paris, and Wolfgang Heidrich. Displacement interpolation using lagrangian mass transport. ACM Trans. Graph. , 30(6):158, 2011.

MAlan Brookhart, Sebastian Schneeweiss, Kenneth J Rothman, Robert J Glynn, Jerry Avorn, and Til Stürmer. Variable selection for propensity score models. American journal of epidemiology , 163(12):1149-1156, 2006.

Aaron Chadha and Yiannis Andreopoulos. Improved techniques for adversarial discriminative domain adaptation. IEEE Trans. Image Process. , 29:2622-2637, 2020.

Yale Chang and Jennifer G. Dy. Informative subspace learning for counterfactual inference. In AAAI , pages 1770-1776, 2017.

Laetitia Chapel, Rémi Flamary, Haoran Wu, Cédric Févotte, and Gilles Gasso. Unbalanced optimal transport through non-negative penalized linear regression. In NeurIPS , pages 2327023282, 2021.

Arthur Charpentier, Emmanuel Flachaire, and Ewen Gallic. Optimal transport for counterfactual estimation: A method for causal inference. arXiv preprint arXiv:2301.07755 , 2023.

Jiawei Chen, Hande Dong, Yang Qiu, Xiangnan He, Xin Xin, Liang Chen, Guli Lin, and Keping Yang. Autodebias: Learning to debias for recommendation. In SIGIR , pages 21-30, 2021.

Tian Qi Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary differential equations. In NeurIPS , pages 6572-6583, 2018.

Mingyuan Cheng, Xinru Liao, Quan Liu, Bin Ma, Jian Xu, and Bo Zheng. Learning disentangled representations for counterfactual regression via mutual information minimization. In SIGIR , pages 1802-1806, 2022.

Lenaïc Chizat, Gabriel Peyré, Bernhard Schmitzer, and François-Xavier Vialard. Scaling algorithms for unbalanced optimal transport problems. Math. Comput. , 87(314):2563-2609, 2018.

José M Cordero, Víctor Cristóbal, and Daniel Santín. Causal inference on education policies: A survey of empirical studies using pisa, timss and pirls. J. Econ. Surv. , 32(3):878-915, 2018.

Nicolas Courty, Rémi Flamary, Amaury Habrard, and Alain Rakotomamonjy. Joint distribution optimal transportation for domain adaptation. In NeurIPS , pages 3730-3739, 2017.

Nicolas Courty, Rémi Flamary, Devis Tuia, and Alain Rakotomamonjy. Optimal transport for domain adaptation. IEEE Trans. Pattern Anal. Mach. Intell. , 39(9):1853-1865, 2017.

Richard K Crump, V Joseph Hotz, Guido W Imbens, and Oscar A Mitnik. Nonparametric tests for treatment effect heterogeneity. Rev. Econ. Stat. , 90(3):389-405, 2008.

Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In NeurIPS , pages 2292-2300, 2013.

Quanyu Dai, Haoxuan Li, Peng Wu, Zhenhua Dong, Xiao-Hua Zhou, Rui Zhang, Rui Zhang, and Jie Sun. A generalized doubly robust learning framework for debiasing post-click conversion rate prediction. In SIGKDD , pages 252-262, 2022.

Eric Dunipace. Optimal transport weights for causal inference. CoRR , abs/2109.01991, 2021.

Pavel E. Dvurechensky, Alexander V. Gasnikov, and Alexey Kroshnin. Computational optimal transport: Complexity by accelerated gradient descent is better than by sinkhorn's algorithm. In ICML , volume 80, pages 1366-1375, 2018.

Kilian Fatras, Thibault Séjourné, Rémi Flamary, and Nicolas Courty. Unbalanced minibatch optimal transport; applications to domain adaptation. In ICML , volume 139, pages 3186-3197, 2021.

Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z. Alaya, Aurélie Boisbunon, Stanislas Chambon, Laetitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, Léo Gautheron, Nathalie T. H. Gayraud, Hicham Janati, Alain Rakotomamonjy, Ievgen Redko, Antoine Rolet, Antony Schutz, Vivien Seguy, Danica J. Sutherland, Romain Tavenard, Alexander Tong, and Titouan Vayer. POT: python optimal transport. J. Mach. Learn. Res. , 22:78:1-78:8, 2021.

Sergio Garrido, Stanislav Borysov, Jeppe Rich, and Francisco Pereira. Estimating causal effects with the neural autoregressive density estimator. Journal of Causal Inference , 9(1):211-228, 2021.

Siyuan Guo, Lixin Zou, Yiding Liu, Wenwen Ye, Suqi Cheng, Shuaiqiang Wang, Hechang Chen, Dawei Yin, and Yi Chang. Enhanced doubly robust learning for debiasing post-click conversion rate estimation. In SIGIR , pages 275-284, 2021.

Negar Hassanpour and Russell Greiner. Learning disentangled representations for counterfactual regression. In ICLR . OpenReview.net, 2020.

Guido W Imbens and Donald B Rubin. Causal inference in statistics, social, and biomedical sciences . Cambridge University Press, 2015.

Arun Jambulapati, Aaron Sidford, and Kevin Tian. A direct o ( 1 / epsilon ) iteration parallel algorithm for optimal transport. In NeurIPS , pages 11355-11366, 2019.

Chi Jin, Rong Ge, Praneeth Netrapalli, Sham M Kakade, and Michael I Jordan. How to escape saddle points efficiently. In ICML , pages 1724-1732, 2017.

Fredrik D. Johansson, Uri Shalit, and David A. Sontag. Learning representations for counterfactual inference. In ICML , pages 3020-3029, 2016.

Fredrik D. Johansson, Uri Shalit, Nathan Kallus, and David A. Sontag. Generalization bounds and representation learning for estimation of potential outcomes and causal effects. CoRR , abs/2001.07426, 2020.

Nathan Kallus. Deepmatch: Balancing deep covariate representations for causal inference using adversarial training. In ICML , volume 119, pages 5067-5077, 2020.

Leonid V Kantorovich. On the translocation of masses. J. Math. Sci. , 133(4):1381-1382, 2006.

Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In ICLR , 2015.

Sören Reinhold Künzel, Jasjeet Sekhon, Peter Bickel, and Bin Yu. Metalearners for estimating heterogeneous treatment effects using machine learning. Proc. Natl. Acad. Sci. U. S. A. , 116 (10):4156-4165, 2019.

Jae-woong Lee, Seongmin Park, and Jongwuk Lee. Dual unbiased recommender learning for implicit feedback. In SIGIR , pages 1647-1651, 2021.

Haoxuan Li, Yanghao Xiao, Chunyuan Zheng, and Peng Wu. Balancing unobserved confounding with a few unbiased ratings in debiased recommendations. In WWW , pages 305-315, 2022.

Haoxuan Li, Quanyu Dai, Yuru Li, Yan Lyu, Zhenhua Dong, Xiao-Hua Zhou, and Peng Wu. Multiple robust learning for recommendation. In Brian Williams, Yiling Chen, and Jennifer Neville, editors, AAAI , pages 4417-4425, 2023.

Haoxuan Li, Yanghao Xiao, Chunyuan Zheng, Peng Wu, and Peng Cui. Propensity matters: Measuring and enhancing balancing for recommendation. In ICML , volume 202 of Proceedings of Machine Learning Research , pages 20182-20194. PMLR, 2023.

Haoxuan Li, Chunyuan Zheng, Yixiao Cao, Zhi Geng, Yue Liu, and Peng Wu. Trustworthy policy learning under the counterfactual no-harm criterion. In ICML , volume 202 of Proceedings of Machine Learning Research , pages 20575-20598. PMLR, 2023.

Haoxuan Li, Chunyuan Zheng, and Peng Wu. Stabledr: Stabilized doubly robust learning for recommendation on data missing not at random. In ICLR , 2023.

Haoxuan Li, Chunyuan Zheng, Peng Wu, Kun Kuang, Yue Liu, and Peng Cui. Who should be given incentives? counterfactual optimal treatment regimes learning for recommendation. In KDD , pages 1235-1247. ACM, 2023.

Qian Li, Zhichao Wang, Shaowu Liu, Gang Li, and Guandong Xu. Deep treatment-adaptive network for causal inference. VLDBJ , pages 1-16, 2022.

Qian Li, Zhichao Wang, Shaowu Liu, Gang Li, and Guandong Xu. Causal optimal transport for treatment effect estimation. IEEE Trans. Neural Networks Learn. Syst. , 34(8):4083-4095, 2023.

Sheng Li, Nikos Vlassis, Jaya Kawale, and Yun Fu. Matching via dimensionality reduction for estimation of treatment effects in digital marketing campaigns. In IJCAI , pages 3768-3774, 2016.

Tianyi Lin, Nhat Ho, and Michael I. Jordan. On efficient optimal transport: An analysis of greedy and accelerated mirror descent algorithms. In ICML , volume 97, pages 3982-3991, 2019.

Jing Ma, Mengting Wan, Longqi Yang, Jundong Li, Brent Hecht, and Jaime Teevan. Learning causal effects on hypergraphs. In SIGKDD , page 1202-1212, 2022.

Gaspard Monge. Mémoire sur la théorie des déblais et des remblais. Mem. Math. Phys. Acad. Royale Sci. , pages 666-704, 1781.

X Nie and S Wager. Quasi-oracle estimation of heterogeneous treatment effects. Biometrika , 108(2):299-319, 2020.

Kenneth R Niswander and Myron Gordon. The women and their pregnancies: the Collaborative Perinatal Study of the National Institute of Neurological Diseases and Stroke , volume 73. National Institute of Health, 1972.

Elizabeth L Ogburn and Tyler J VanderWeele. On the nondifferential misclassification of a binary confounder. Epidemiology , 23(3):433, 2012.

Judea Pearl and Dana Mackenzie. The book of why: the new science of cause and effect . Basic books, 2018.

Ofir Pele and Michael Werman. Fast and robust earth mover's distances. In ICCV , pages 460-467, 2009.

Gabriel Peyré and Marco Cuturi. Computational optimal transport. Found. Trends Mach. Learn. , 11(5-6):355-607, 2019.

Khiem Pham, Khang Le, Nhat Ho, Tung Pham, and Hung Bui. On unbalanced optimal transport: An analysis of sinkhorn algorithm. In ICML , volume 119, pages 7673-7682, 2020.

Ievgen Redko, Amaury Habrard, and Marc Sebban. Theoretical analysis of domain adaptation with optimal transport. In PKDD , volume 10535, pages 737-753, 2017.

James M Robins, Andrea Rotnitzky, and Lue Ping Zhao. Estimation of regression coefficients when some regressors are not always observed. J. Am. Stat. Assoc. , 89(427):846-866, 1994.

Paul Rosenbaum and Donald B Rubin. The central role of the propensity score in observational studies for causal effects. Biometrika , 70(1):41-55, 1983.

Paul Rosenbaum and Donald B Rubin. The central role of the propensity score in observational studies for causal effects. Biometrika , 70(1):41-55, 1983.

Donald B Rubin. Estimating causal effects of treatments in randomized and nonrandomized studies. J. Educ. Psychol. , 66(5):688, 1974.

Brian C Sauer, M Alan Brookhart, Jason Roy, and Tyler VanderWeele. A review of covariate selection for non-experimental comparative effectiveness research. Pharmacoepidemiology and drug safety , 22(11):1139-1145, 2013.

Patrick Schwab, Lorenz Linhardt, Stefan Bauer, Joachim Buhmann, and Walter Karlen. Learning counterfactual representations for estimating individual dose-response curves. In AAAI , pages 5612-5619, 2020.

Vivien Seguy, Bharath Bhushan Damodaran, Rémi Flamary, Nicolas Courty, Antoine Rolet, and Mathieu Blondel. Large scale optimal transport and mapping estimation. In ICLR , 2018.

Thibault Séjourné, Jean Feydy, François-Xavier Vialard, Alain Trouvé, and Gabriel Peyré. Sinkhorn divergences for unbalanced optimal transport. CoRR , abs/1910.12958, 2019.

Thibault Séjourné, François-Xavier Vialard, and Gabriel Peyré. Faster unbalanced optimal transport: Translation invariant sinkhorn and 1-d frank-wolfe. In AISTATS , volume 151, pages 4995-5021, 2022.

Uri Shalit, Fredrik D. Johansson, and David Sontag. Estimating individual treatment effect: generalization bounds and algorithms. In ICML , pages 3076-3085, 2017.

Claudia Shi, David M. Blei, and Victor Veitch. Adapting neural networks for the estimation of treatment effects. In NeurIPS , pages 2503-2513, 2019.

Susan M Shortreed and Ashkan Ertefaie. Outcome-adaptive lasso: variable selection for causal inference. Biometrics , 73(4):1111-1122, 2017.

Tamar Sofer, David B Richardson, Elena Colicino, Joel Schwartz, and Eric J Tchetgen Tchetgen. On negative outcome control of unobserved confounding as a generalization of difference-indifferences. Stat. Sci. , 31(3):348, 2016.

Ruibo Tu, Kun Zhang, Hedvig Kjellström, and Cheng Zhang. Optimal transport for causal discovery. In ICLR . OpenReview.net, 2022.

Cédric Villani. Optimal transport: old and new , volume 338. Springer, 2009.

Stefan Wager and Susan Athey. Estimation and inference of heterogeneous treatment effects using random forests. J. Am. Stat. Assoc. , 113(523):1228-1242, 2018.

Hao Wang, Tai-Wei Chang, Tianqiao Liu, Jianmin Huang, Zhichao Chen, Chao Yu, Ruopeng Li, and Wei Chu. Escm 2 : Entire space counterfactual multi-task model for post-click conversion rate estimation. In SIGIR , 2022.

Haotian Wang, Wenjing Yang, Longqi Yang, Anpeng Wu, Liyang Xu, Jing Ren, Fei Wu, and Kun Kuang. Estimating individualized causal effect with confounded instruments. In KDD , pages 1857-1867. ACM, 2022.

Haotian Wang, Kun Kuang, Haoang Chi, Longqi Yang, Mingyang Geng, Wanrong Huang, and Wenjing Yang. Treatment effect estimation with adjustment feature selection. In KDD , pages 2290-2301. ACM, 2023.

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.

Christopher KI Williams and Carl Edward Rasmussen. Gaussian processes for machine learning , volume 2. MIT press Cambridge, MA, 2006.

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.

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.

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.

Zeke Xie, Issei Sato, and Masashi Sugiyama. A diffusion theory for deep learning dynamics: Stochastic gradient descent exponentially favors flat minima. In ICLR , 2020.

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.

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.

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.

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.

Jinsung Yoon, James Jordon, and Mihaela van der Schaar. GANITE: estimation of individualized treatment effects using generative adversarial nets. In ICLR , 2018.

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.

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.

Jiajing Zheng. Sensitivity Analysis for Causal Inference with Unobserved Confounding . University of California, Santa Barbara, 2021.

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.

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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this paper, authors have applied optimal transport to address selection-bias issue in binary treatment setting for individualized treatment effect estimation. They have proposed a relaxed mass-preserving regularizer to address the mis-alignment generated from outcome imbalance and outliers in non-ideal mini-batches. They also proposed a proximal factual outcome regularizer to relax the unconfoundedness assumption. They have done extensive evaluation of the proposed approach on IHDP and ACIC, two semi-synthetic datasets.\n\nSTRENGTHS:\n- A novel method is proposed for treatment effect estimation in binary setting based on optimal transport (OT) that uses a relaxed mass-preserving regularizer to address the mis-alignment generated from outcome imbalance and outliers in non-ideal mini-batches, as well as a proximal factual outcome regularizer to relax the unconfoundedness assumption.\n- Extensive evaluations is performed to study the proposed model that provides necessary analysis and ablation studies.\n- Theoretical results are provided to bound the error in terms of the proposed components to show the convergence (I have not verified the proof).\n- Idea was clearly stated with sufficient details.\n\nWEAKNESSES:\n- Baselines used in the paper seem old. There are very large number of methods proposed for treatment effect (TE) estimation. It would have been nice to see comparison with recent techniques. Hoever, even if authors don't add new baselines, I think paper has sufficient novelty to be above the acceptance threshold.\n- The paper solves (equation (1)) Conditional average treatment effect estimation also called as heterogeous treatment effects estimation, as called in machine learning literature, but paper did not even mention these names. Please update the paper so that readers don't have any confusion. There is already a lot of confusion in ML regarding the terminology.\n- Causal inference methods are typically based on some assumptions -- but I did not see any. Please clearly state the assumptions required to apply these methods.\n\nQUESTIONS:\n- Training all networks for 400 epochs do not look fair to me. Are all the methods getting trained before 400 epochs or not? In my experience this affects the performance in causal inference. All the networks should be trained using early stopping and by setting large number of max-epochs so that every learner gets chance to train properly.\n- Have you used two layers each in representation learning and outcome mapping or total two layers for both?\n- First line of Introduction section does not look right to me. Doesn't RCT mostly answer population-level or heterogeneous-level questions but not the individual-level treatment effects?\n- In figures 1(a) and 2, please clearly state what are the different axis, e.g., for Figure 2, it is not clear what is x-axis or what is y-axis?\n- Why does title says 'a new perspective' when the idea is very widely used, including in causal inference, as authors have themselves discussed that?\n\nLIMITATIONS:\nAuthors did not discuss the assumptions and limitations of thier work.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-06-16"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose to minimize a (relaxed) Sinkhorn distance between representations under the two treatment groups in order to balance the latent space and better predict the treatment effects. They highlight two problems with existing TARNet-based approaches that do not regularize with optimal-transport cost: mini-batch sampling effects (MSE) and unobserved confounder effects (UCE). Theorem 3.1 & Corollary 3.1 formally characterize how the proposed method suffers less from MSE. The authors also introduce additional regularization terms that are meant to deal with UCE. Finally, they show improved performance with two classical semi-synthetic benchmarks.\n\nSTRENGTHS:\n* The problem of estimating heterogeneous treatment effects is important, the paper is written well, and the solution is creative.\n * The motivation for the relaxed Sinkhorn distance (Sec. 3.2) is great. \n * The benchmark results are solid.\n * The additional empirical explorations (Sec. 4.3 - 4.5) are interesting as well.\n\nWEAKNESSES:\nMy biggest problem is with the UCE part of the paper.\n * Sec. 3.3: the explanation for the hidden-confounding regularizer makes intuitive sense. However, I do not see a formal connection. Once the authors substitute the actual potential outcomes with the model's predictions (Eq. 11 -> Eq. 12), one begins to wonder if this can actually help with hidden confounding.\n * A major confounder with the issue above is that the authors do not appear to include hidden confounders in their semi-synthetic benchmarks. Therefore, it is impossible to tell if this regularization approach really helps with UCE.\n * If the authors still wish to make the case for UCE mitigation, they need to reference the growing literature on partial identification / sensitivity analysis for hidden confounders and how their regularization relates to sensitivity models.\n\nBesides that, I think the authors should spend more time addressing what the theoretical results mean for the MSE issue that they highlighted in the introduction.\n\nSmall nitpicks:\n* Figure 1a, the orange annotation looks like it is referring to the X_1 part of the figure.\n * line 107-108: \"Optimal transport is a preferred method due to its advantages over competitors.\" is overly vague.\n * line 108-109, \"It accounts for the distribution’s geometry, making it effective [where the KL-divergence] fails.\" is a strange statement. Why does the geometry help here? No explanation is provided. \n * The beginning of the paper is a bit verbose.\n\nQUESTIONS:\nHow do you know that you address UCE? As the paper stands, I recommend that you refrain from discussing hidden confounders in your motivation for the method and stick with what really *can* be shown for your proposed regularization. I am willing to raise my score once this is addressed.\n\nIs MSE really an issue with previous approaches? Can you show that clearly?\n\nLIMITATIONS:\nLimitations or potential societal impacts are not discussed.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-06-18"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe objective of this research paper is to tackle two primary concerns: (1) the effects of mini-batch sampling (MSE), which lead to misalignment in non-ideal mini-batches characterized by outcome imbalance and outliers, and (2) the impacts of unobserved confounders (UCE), resulting in inaccurate discrepancy calculation due to the neglect of these unobserved factors. To address these challenges, this paper presents a novel approach based on optimal transport in the context of causality.\n\nMore specifically, the proposed approach builds upon the stochastic optimal transport framework. It introduces a relaxed mass preserving regularizer to mitigate the MSE issue and devises a proximal factual outcome regularizer to address the UCE problem. By incorporating these techniques, the proposed method achieves a principled solution for treatment selection bias.\n\nThe effectiveness of the proposed method is extensively evaluated through a series of experiments. The results demonstrate that the approach successfully overcomes treatment selection bias and outperforms existing state-of-the-art methods by a significant margin.\n\nSTRENGTHS:\nThe paper is clearly written and well-presented. It delves into the important question of improving the estimation of individualized treatment effects, a crucial topic within the field of causal inference. Optimal transport, an emerging field in machine learning, is effectively employed to measure treatment selection bias. Furthermore, the two proposed regularizers are technically intriguing and efficient, as supported by both theoretical and empirical evidence.\n\nWEAKNESSES:\nI have two main concerns: \n\n1. The potential applicability of the proposed relaxed mass-preserving regularizer extends beyond causal effect estimation and encompasses a wide range of application settings. It appears to be independent of the estimation of causal effects. The authors acknowledge that the OT discrepancy can be easily affected by various sampling cases, highlighting that it is not solely limited to the distribution discrepancy across treatment groups. Consequently, the main contribution of this regularizer may be seen as an incremental addition to the field of causal inference, raising concerns about the originality and novelty of the paper.\n\n2. The existing literature on the subject appears to be considerably understudied. For example, important references like \"Optimal Transport for Counterfactual Estimation: A Method for Causal Inference\" by Arthur Charpentier, Emmanuel Flachaire, and Ewen Gallic are missing and remain unaddressed in the related work section. It would be beneficial for the authors to conduct a more comprehensive investigation of the related literature and consider incorporating these references into their study.\n\nQUESTIONS:\nSee above\n\nLIMITATIONS:\nSee above", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-26"}

comments: SUMMARY:
The paper introduces a new approach called Entire Space CounterFactual Regression (ESCFR) to estimate individual treatment effects from observational data, addressing two critical problems that existing methods fail to solve. The first problem is mini-batch sampling effects (MSE), which lead to misalignment in non-ideal mini-batches with outcome imbalance and outliers. The second problem is unobserved confounder effects (UCE), which result in inaccurate discrepancy calculation due to the neglect of unobserved confounders. ESCFR combines the principles of optimal transport and causality to overcome these issues.

STRENGTHS:
- The theoretical evidence is solid and convincing.
- The topic, mentioned in the paper as the mini-batch misalignment and UCE error, are interesting and practical issues of ITE estimation.
- Experimental results are solid.

WEAKNESSES:
Is there some theoretical guarantee for the UCE regularizer proposed in this paper? If so, I think the quality can be further improved.

QUESTIONS:
See in Weaknesses

LIMITATIONS:
See in Weaknesses
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-07-05