S
D ISCREPANCY
FOR
C OUNTERFACTUAL
G ENERALIZATION
Anonymous authors
Paper under double-blind review
A BSTRACT
Estimating individual treatment effects from observational data is highly challenging due to the existence of treatment selection bias.
Most prevalent approaches mitigate this issue by aligning distributions of different treatment groups in the representation space.
However, there are two critical problems circumvented: (1) mini-batch sampling effects (MSE), where the alignment easily fails due to the outcome imbalance or outliers at a mini-batch level; (2) unobserved confounder effects (UCE), where the unobserved confounders damage the correct alignment.
To tackle these problems, we propose a principled approach named E ntire S pace
C ounter F actual R egression (ESCFR) based on a generalized sinkhorn discrepancy for distribution alignment within the stochastic optimal transport framework.
Based on the 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 our proposed ESCFR can successfully tackle the treatment selection bias and achieve significantly better performance than state-of-the-art methods.
I NTRODUCTION
Estimating individual treatment effect (ITE) with randomized controlled trials is a common practice in causal inference, which has been widely used in e-commerce (Betlei et al., 2021), education (Cordero et al., 2018), and health care (Schwab et al., 2020).
For example, drug developers would conduct clinical A/B tests to evaluate the drug effects.
Although randomized controlled trials are the gold standard (Pearl & Mackenzie, 2018) for causal inference, it is often prohibitively expensive to conduct such experiments.
Hence, observational data that can be acquired without intervention has been a tempting shortcut.
For example, drug developers tend to assess drug effects with post-marketing monitoring reports instead of clinical A/B trials.
With the growing access to observational data, estimating ITE from observational data has attracted intense research interest.
Estimating ITE 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 their preferences for treatment selection, making units in different treatment groups heterogeneous.
To handle missing counterfactuals, meta-learners (K¨unzel et al., 2019) decompose the
ITE estimation task into solvable factual outcome estimation subproblems.
However, the treatment selection bias makes it difficult to generalize the factual outcome estimators trained within respective treatment groups to the entire population; consequently, the derived ITE estimator is biased.
Beginning with counterfactual regression (Shalit et al., 2017) and its revolutionary performance, most prevalent methods handle the selection bias by minimizing the distribution discrepancy between groups in the representation space (see Liuyi et al., 2018; Hassanpour & Greiner, 2020; Cheng et al., 2022).
However, two critical issues with these methods have long been neglected, which significantly impedes them from handling the treatment selection bias.
The first problem is the mini-batch sampling effects (MSE).
Specifically, current representation-based methods (Shalit et al., 2017; Liuyi et al., 2018) compute distribution discrepancy within mini-batches instead of the entire data space, making it vulnerable to bad sampling cases.
For example, given two aligned distributions, if a mini-batch outlier exists in the sampled distribution, the mini-batch discrepancy will be significant, making the training process noise-filled.
The second problem is the unobserved confounder effects
Y
Xψ
R
WW
YY
LL
TT ^ ^ Rϕ
Yϕ
Xϕ 0 (R 0 ) R 0 =
ψ (X 0 ) ϕ
1 (R 1 )
F a c t ua l ou t c o m e
Covariate
Treatment indicator
Factual outcome
Factual error
Representation Outcome estimate
X X
R (
Y
Xψ
R
WW
YY
LL
TT ^ ^ Rϕ
Yϕ
Xϕ 0 (R 0 ) R 0 =
ψ (X 0 ) ϕ
1 (R 1 )
F a c t ua l ou t c o m e
Covariate
Treatment indicator
Factual outcome
Factual error
Representation Outcome estimate
X
[Figure] [Figure]
(UCE).
Specifically, current approaches directly assume unconfoundedness Ma et al.
(2022), while the unobserved confounders widely exist in real scenarios and make the resulting estimators biased.
Contributions and outline.
In this paper, we propose an effective ITE estimator based on optimal transport, Entire Space CounterFactual Regression (ESCFR), which tackles both the MSE and UCE issues with a generalized sinkhorn discrepancy.
Specifically, after preliminaries in Section 2, we first reformulate the ITE estimation problem as a stochastic optimal transport problem in Section 3.1.
We next showcase the MSE issue faced by existing approaches in Section 3.2 and propose a relaxed mass-preserving regularizer to mitigate this issue.
We further investigate the UCE issue in Section 3. and propose a proximal factual outcome regularizer to solve it.
We finally formulate the architecture and learning objectives of ESCFR in Section 3.4, and report the experimental results in Section 4.
P RELIMINARIES 2. C AUSAL I NFERENCE FROM O BSERVATIONAL D ATA
This section formulates basic definitions and models in observational causal inference.
We first formalize the fundamental elements in Definition 2.1 following the general notation convention 1 .
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 ) .
Following the potential outcome framework (Rubin, 1974), an individual with covariates x has two potential outcomes, namely Y 1 ( x ) given it is treated and Y 0 ( x ) otherwise.
The ground-truth individual treatment effect (ITE) is always formulated as the expected difference of potential outcomes:
[Equation]
where one of these two outcomes is always unobserved.
To address such missing counterfactuals, the
ITE estimation task is commonly decomposed into potential outcome estimation subproblems that are solvable with any supervised learning method (K¨unzel et al., 2019).
For example, T-learner models the factual outcomes Y for units in the treated and untreated groups separately; S-learner regards the treatment indicator T as one of the covariates X , and models Y for all units simultaneously.
The ITE 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.
For example, Y 1 = ϕ 1 ( R ) , Y
= ϕ 0 ( R ) , where we abbreviate ϕ ( R,T = 1 ) and ϕ ( R,T = 0 ) to ϕ 1 ( R ) and ϕ 0 ( R ) , respectively, for brevity.
TARNet (Shalit et al., 2017) obtains better performance by absorbing the advantages of both T-learner and S-learner, consisting of a representation mapping ψ and an outcome mapping ϕ as defined in
Definition 2.2.
For an individual with covariates X , TARNet estimates ITE as:
[Equation]
where
[Equation]
where ψ is trained over all individuals, ϕ 1 and ϕ 0 are trained over the treated and untreated group, respectively, to minimize the factual error.
Finally, the performance of ITE estimators is mainly evaluated with the precision in estimation of heterogeneous effect (PEHE):
[Equation]
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 ITE estimate ˆ τ by these methods would be biased.
2. D ISCRETE O PTIMAL T RANSPORT AND S INKHORN D IVERGENCE
Optimal transport (OT) instantiates distribution discrepancy as the minimum transport cost, which provides a grip for quantifying the treatment selection bias in Figure 1(a).
Monge (1781) first formulated OT as finding an optimal mapping between two distributions.
However, this formulation cannot guarantee the existence and uniqueness of solutions.
Kantorovich (2006) proposed a more applicable formulation in Definition 2.3, which can be seen as a generalization of 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 minimum cost:
[Equation]
where W ( α,β ) ∈ R is the Wasserstein discrepancy between α and β ; D ∈ R n × m + is the unit-wise distance 2 between α and β ; a and b indicate the mass of units in α and β , and Π is the feasible transportation plan set which ensures the mass-preserving constraint holds.
However, exact solutions (Bonneel et al., 2011) to (4) always come with high computational costs.
As such, researchers would always add an entropic regularization to the Kantorovich problem:
[Equation]
making the problem ϵ -convex and solvable with the Sinkhorn algorithm (Cuturi, 2013).
The Sinkhorn algorithm only consists of matrix-vector products, making it suited to be accelerated with GPUs.
P ROPOSED M ETHOD
In this section, we present the proposed Entire Space CounterFactual Regression (ESCFR) approach based on optimal transport to tackle the treatment selection bias.
We first illustrate the stochastic optimal transport framework for distribution discrepancy minimization across treatment groups.
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 summarize the model architecture, learning objectives, and optimization algorithm.
3. S TOCHASTIC O PTIMAL T RANSPORT FOR C OUNTERFACTUAL R EGRESSION
Representation-based approaches mitigate the treatment selection bias by calculating the distribution discrepancy in the representation space and then minimizing it.
We select optimal transport to compute the discrepancy since it has shown compelling advantages over its competitors.
Specifically, it accounts for the distribution’s geometry and thus works in the cases where ϕ -divergence ( e .
g .,
Kullback-Leibler divergence) fails (Seguy et al., 2018).
In addition, the calculated discrepancy can be optimized with the traditional supervised learning framework instead of the adversarial learning framework, and is therefore easier to optimize than adversarial-based methods (Kallus, 2020).
Optimal transport calculates the group discrepancy as W ( P T = 1 ψ ( r ) , P T = 0 ψ ( r )) , where P T = 1 ψ ( r ) and
P T =
ψ ( 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 ψ (Flamary et al., 2021), thus can be minimized by updating the representation mapping ψ with gradient-based optimizers.
Definition 3.1.
Let ˆ P T = 1 ( x ) and ˆ P T = 0 ( x ) 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 =
ψ ( r ) be that of representations induced by the representation mapping r = ψ ( x ) defined 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 group discrepancy at a stochastic mini-batch level:
[Equation]
To further investigate the effectiveness of this shortcut, Theorem 3.1 demonstrates that PEHE can be optimized by iteratively minimizing the factual outcome estimation error and the mini-batch group discrepancy (6).
The proof of the theorem can be found in Appendix A.3.
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:
[Equation]
where ϵ T = 1 F and ϵ T =
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. R ELAXED M ASS -P RESERVING R EGULARIZER FOR S AMPLING E FFECT
[Figure]
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, 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, the discrepancy measure should be robust to bad sampling cases, otherwise the resulting huge variance will impede it from computing and reducing the actual discrepancy.
The OT discrepancy in (6) can be easily disturbed by many sampling cases, except for the ideal case in Figure 2(a) where the transport strategy is reasonable and applicable.
For example, according to
Figure 2(b), it falsely matches units with unrelated factual outcomes in the mini-batch where the outcomes between groups are imbalanced; according to Figure 2(c), it falsely matches the mini-batch outliers to normal units, causing a substantial disruption of the transportation strategy.
In short, the vanilla OT technique in (6) fails to quantify the group discrepancy for producing erroneous transport strategies in non-ideal mini-batches, and thus misleads the update of the representation mapping ψ .
We summarize this phenomenon as the mini-batch sampling effect (MSE) issue
This issue is attributed to the mass-preservation constraint in (5), which requires that all units in both groups match each other, regardless of the actual situation.
Mini-batch outliers, for instance, would be compelled to be transported according to Figure 2, which impedes the transport of normal units and the computation of the actual group discrepancy.
A small batch size would exacerbate this defect.
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:
[Equation]
where D ∈ R n × m + is the unit-wise distance, and a ; b indicate the mass of units in α and β , respectively.
An intuitive approach to mitigate MSE is to relax the marginal constraint and allow for the creation and destruction of unit’s mass.
To this end, a relaxed mass-preserving regularizer (RMPR) is devised in Definition 3.2, which replaces the hard marginal constraint in (4) with a soft penalty in (8) for deriving transport strategy.
In this context, the stochastic discrepancy is calculated as
[Equation] [Equation]
where the hard mass-preservation constraint is removed to mitigate the MSE issue.
Inspired by Fatras et al.
(2021), the robustness of RMPR to sampling effects can be further theoretically investigated in
Theorem 3.2, where the effect of mini-batch outliers is upper bounded by a constant.
Theorem 3.2.
For empirical distributions α,β with n and m units, respectively, adding an outlier a ′ to α and denoting the disturbed distribution as α ′ , we have
[Equation]
which is upper bounded by 2 κ / n .
W 0 ,κ is the unbalanced discrepancy as per Definition 3.2.
In addition, compared with alternatives (Xu et al., 2020; Chapel et al., 2021) to relax the marginal constraint, the approach in Definition 3.2 has better metric properties (S ´ ejourn ´ e et al., 2019) and can be accelerated via the generalized Sinkhorn algorithm (Chizat et al., 2018).
It is differentiable w .
r .
t .
ψ and thus can be minimized via stochastic gradient methods in an end-to-end manner.
3. P ROXIMAL F ACTUAL O UTCOME R EGULARIZER FOR U NOBSERVED C ONFOUNDERS
[Figure]
Existing representation-based methods fail to eliminate the treatment selection bias due to the unobserved confounder effects (UCE).
Beginning with
CFR (Shalit et al., 2017), the unconfoundedness assumption A.1 (see Appendix A) is often taken to circumvent the UCE issue (Ma et al., 2022).
Given two units r i ∈ P T = 1 ψ ( r ) and r j ∈ P T = 0 ψ ( r ) , for instance, optimal transport in Definition 3.2 calculates the unit-wise distance as D ij ∶= ∥
r i − r j ∥ 2 .
If
Assumption A.1 holds, this approach mitigates the treatment selection bias since it blocks the backdoor path X
→ T in Figure 3(a) by balancing the confounders across groups in a latent space.
However,
Assumption A.1 is usually violated in practice as per Figure 3(b), which hinder existing methods including OT from handling treatment selection bias since the backdoor path X ′
→ T is not blocked.
According to Figure 3(b), given balanced X and identical T , the only variable reflecting the variation of X ′ is the outcome Y .
As such, inspired by the joint distribution transport technique (Courty et al., 2017a), we propose to calibrate the unit-wise distance D with the potential outcomes as follow:
[Equation]
where γ controls the strength of regularization.
The underlying regularization is: units with similar (both observed and unobserved) confounders should 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 ti − y tj ∥ >> 0 , their unobserved confounders should likewise differ significantly.
The vanilla OT technique in (6) with D ij = ∥
r i − r j ∥ 2 would incorrectly match this pair because ∥ r i − r j ∥ 2 ≈ 0 , 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.
Moreover, 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 =
i and y T = 1 j , respectively, PFOR refines (11) as
[Equation]
Additional justifications, assumptions and limitations of PFOR are discussed in Appendix D.3.
3. A RCHITECTURE OF E NTIRE S PACE C OUNTERFACTUAL R EGRESSION
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 (8) 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 Shi et al.
(2019) can be formulated as
[Equation]
where y i and y j are the factual outcomes for the corresponding treatment groups.
The discrepancy is:
[Equation]
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
[Equation]
where λ controls the strength of distribution alignment, ϵ controls the entropic regularization in (5), κ controls RMPR in (8), 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.
First, compute πππ ϵ,κ by solving the linear programming problem in Definition 3.2 with Algorithm 2, where the unit-wise distance is calculated with D γ .
Second, compute the discrepancy in (14) as ⟨ πππ ϵ,κ , D γ ⟩ , where D γ is differentiable to ψ , making it feasible to minimize this discrepancy with its gradient w.r.t.
ψ .
Finally, calculate the overall loss in (15) and update ψ and ϕ with stochastic gradient methods.
E XPERIMENTS 4. E XPERIMENTAL S ETUP
Datasets.
Missing counterfactuals impedes the evaluation of PEHE with observational benchmarks.
Following Liuyi et al.
; Shalit et al.
(2017), experiments are conducted on two semi-synthetic benchmarks.
Specifically, the IHDP benchmark aims to estimate the effect of specialist home visits on infants’ potential cognitive scores, with 747 observations and 25 covariates; the ACIC dataset comes from the collaborative perinatal project, with 4802 observations and 58 covariates.
Baselines.
We consider three groups of baselines.
(1) Statistical estimators: least square regression with the treatment as covariates (OLS)
[Table]
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 400 epochs with Adam optimizer, with the learning rate and weight decay both set to 0.001.
Other settings of optimizers follow Kingma & Ba (2015).
We fine-tune hyperparameters within the range in
Figure 5, validate performance every two epochs, and save the optimal model for test.
Evaluation protocol.
Following Liuyi et al.
(2018), PEHE in (3) is used as the precision metric for performance evaluation.
However, it is unavailable in the model selection phase due to missing counterfactuals.
As such, we use the area under the uplift curve (AUUC)
(Betlei et al., 2021) to guide model selection, which measures the counterfactual ranking performance of the model and can be computed without counterfactual outcomes.
The within-sample and out-of-sample results are reported on the training and test data, respectively, following Shalit et al.
4. P ERFORMANCE C OMPARISON
Table 1 reports the performance of ESCFR and its competitors over ten runs.
Statistical estimators exhibit competitive performance on the PEHE metric.
Due to the superiority to depict non-linearity, neural estimators outperform linear and random forest methods.
In particular, TARNet that absorbs the advantage of T-learner and S-learner achieves the best overall performance in statistic estimators.
However, the circumvention to treatment selection bias leads to inferior performance.
Matching methods, e .
g ., PSM exhibit compelling ranking performance, which explains why they are favored in counterfactual ranking practice (Betlei et al., 2021).
However, their poor performance on
PEHE hinders their application in counterfactual estimation applications such as advertising systems that place more emphasis on the accuracy of treatment effect estimation.
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, significantly outperforming most statistical methods.
However, the MSE and UCE issues impede these methods from solving the treatment selection bias.
The proposed ESCFR achieves significant improvement over most metrics compared with various prevalent baselines 4 .
Combined with the comparisons above, we attribute its superiority to the proposed RMPR and PFOR regularizers, which makes it robust to
MSE and UCE.
See Appendix D.5 for additional comparison results.
4. A BLATION S TUDY
To verify the effectiveness of individual components, an ablation study is conducted on the ACIC benchmark in Table 2.
Specifically, ESCFR first augments TARNet with stochastic optimal transport in Section 3.1, which effectively reduces the out-of-sample PEHE from 3.254 to 3.207.
Then, it
[Table]
mitigates the MSE issue with RMPR in Section 3.2 and the UCE issue with PFOR in Section 3.3, reducing the out-of-sample PEHE to 2.768 and 2.633, respectively.
Finally, ESCFR combines the
RMPR and PFOR in a unified framework in Section 3.4, reducing the out-of-sample PEHE to 2.316.
4. A NALYSIS OF R ELAXED M ASS -P RESERVING R EGULARIZER
[Equation]
κ =
Most prevalent methods fail to cope the label
- ance and mini-batch outliers in Figure 2(b-c).
Figureshows the transport plan generated with RMPR in the same situations, where RMPR alleviates the MSE issue in both bad cases.
RMPR with κ = 10 , for instance, avoids the incorrect matching of units with different outcomes; RMPR with κ = 2 gets robust to the outlier’s interference and correctly matches the remaining units.
We attribute the success to the relaxed mass-preserving constraint in Section
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 (Johansson et al., 2020), while RMPR naturally implements it via the soft penalty in (8).
We further investigate the performance of RMPR under different batch sizes and κ in Appendix D.2 to verify the effectiveness of RMPR more thoroughly.
4. P ARAMETER S ENSITIVITY S TUDY
We discuss four critical hyperparameters in ESCFR, i .
e
., λ , ϵ , κ , γ , which are the weights in the learning objective and influence the final performance significantly.
We firstly vary λ to investigate the influence of stochastic optimal transport.
Specifically, increasing λ consistently improves the precision of ITE estimates.
For example, the out-of-sample PEHE reduces from 3.22 at λ = 0 to approximately 2.85 at λ = 1 . 5 .
However, assigning a higher focus on distribution balancing in a multi-task learning framework leads to difficulties in factual outcome estimation, leading to sub-optimal ITE estimates.
Then, we vary ϵ to investigate the entropic regularizer.
It does accelerate the computation of optimal transport discrepancy with a large ϵ corresponding to a faster computation speed (Flamary et al., 2021).
However, it comes with a biased transport plan, evidenced by the out-of-sample PEHE increasing to 2.95 with fluctuations.
We further vary γ and κ to showcase the influence of PFOR and RMPR, respectively.
PFOR benefits
ITE estimation, while assigning considerable weight to the proximal outcome distance is detrimental to ITE estimates, because the unit-wise distance calculated with representations (covariates) would be neglected.
Relaxing the mass-preserving constraint via RMPR could significantly improve the ITE estimates; however, too-small κ is always detrimental as we cannot guarantee that the representations across treatment groups are drawn closer together in the optimal transport approach.
[Figure] [Figure]
In
[Figure]
Out
[Figure]
sample
[Figure]
(c)
[Figure]
In
[Figure]
sample Out
[Figure]
sample
[Figure]
(d)
[Figure]
P E
[Figure]
E
In-sample Out
[Figure]
sample
R ELATED W ORKS
Current works mitigate the treatment selection bias by balancing the distributions of treated and untreated groups, which can be divided into three groups according to the methods used for balancing.
Reweighting-based methods weight individuals with balanced scores to achieve globally balanced distributions, represented by the inverse propensity score (IPS) approach (Rosenbaum & Rubin, 1983a) and its doubly robust variant (Robins et al., 1994).
Imai & Ratkovic (2014) and Fong et al.
(2018) propose calculating the balancing score by solving an optimization problem.
Kuang et al.
(2017b) and Kuang et al.
(2017a) consider additional non-confounding factors in covariates.
However, these methods are susceptible to non-overlapping units and suffer from a high variance issue.
Matching-based methods match similar units from different groups to construct locally balanced distributions.
Representative methods (Rosenbaum & Rubin, 1983b; Chang & Dy, 2017; Li et al., 2016) mainly differ in terms of similarity measures.
Notably, Tree-based methods (Wager & Athey, 2018) can also be considered as matching methods with adaptive neighborhood metrics.
However, the computational cost hinders the application of these methods in large-scale scenarios.
Begining with BNN (Johansson et al., 2016) and CFR (Shalit et al., 2017), representation-based methods minimize the group discrepancy in the latent space.
Liuyi et al.
and Hassanpour &
Greiner (2020) further augment CFR with local similarity and non-confounding factors, respectively.
Kallus (2020) and Yoon et al.
(2018) propose to balance the distributions of representations with adversarial training.
Due to its scalability and avoidance of the high variance issue, representationbased methods have been predominant for handling the treatment selection bias.
It is necessary to distinguish ourselves from emerging OT-based causal inference approaches.
Dunipace (2021) augments the IPS method with a propensity score estimator based on OT; however, it is limited by the aforementioned high variance issue.
Torous et al.
(2021) uses the push-forward operator to improve change-in-change models; however, they are designed for multi-phase data which is not available in our case.
Li et al.
(2022) has a similar setup to to us, while it focuses on the decomposition of latent variables and is identical to Shalit et al.
(2017) in terms of alignment technology.
Our work is a new take on OT under the CFR framework, alleviating the MSE and UCE issues that have been long neglected by the causal inference community until this year.
C ONCLUSION
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-match 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
Sinkhorn discrepancy.
Extensive experiments demonstrate that ESCFR can largely mitigate MSE and
UCE issues, and achieve better performance compared with other baseline models.
There are two directions of future work that we intend to pursue.
The first direction attempts to construct the representation mapping with normalizing flows (Chen et al., 2018), which is invertible and thus satisfies the assumption to the representation mapping by Shalit et al.
Another direction seeks to extend our methodology to industrial applications, e .
g ., debias in recommenders (Wang et al., 2022), which has long been dominated by reweighting methods with the high variance issue.