Optimal Transport for Reducing Bias in Causal Inference without Data Splitting

Thank you for the valuable and insightful comments. We will revise the submission according to the comments. Our responses are as follows.

Q1: Comparison with [1,2]

A1:

1. Both [1] and [2] consider the binary treatment setting. Different from them, we propose a unified framework for both binary and continuous treatment settings.

2. [1] reduce the distribution shift between the treatment group and the control group, i.e., $\mathcal{W}(q_1, q_0)$. Different from it, we reduce the shift between each conditional distribution and the marginal distribution, i.e., $\sum_{t} \mathcal{W}(q_t, q)$.

3. [2] considers an optimal transport problem between the factual and counterfactual distributions, which is different from our optimal transport model. In addition, [2] relies on a linear potential outcome model. We employ the neural network to train a nonlinear outcome model.

4. We have conducted the methods in [1,2] as the compared methods in the revision. Our method achieves better performance. The results are as follows.

Q2: Difference from CFR-Wass.

A2:

Apart from data splitting, there are two major difference between our method and CFR-Wass.

1. CFR-Wass reduces the Wasserstein distance between the treatment group and the control group, i.e., $\mathcal{W}(q_1, q_0)$. Different from it, we reduce the Wasserstein distance between each conditional distribution and the marginal distribution, i.e., $\sum_{t} \mathcal{W}(q_t, q)$.

2. CFR-Wass assigns equal weights to the samples in one group. Different from it, we introduce the generalized propensity score into the optimal tranpsort model as the sample weights.

Yes, we conduct CFR with the Wasserstein distance as a compared method in Section 5.2.

Q3: Parameter sensitivity.

A3: The hyperparameter $\lambda$ between the outcome prediction loss and the Wasserstein discrepancies range from $[0.5, 1.3]$. The trade-off hyperparameter of the entropy regularization $\gamma$ is tuned in the range $[0.0002, 0.0018]$, and the sensitivity results on synthetic data $(\beta = 0.25)$ of $\sqrt{AMSE}$ are reported as follows.

Q4: The efficiency of the method.

A4: Running time results are given as follows. Although our method takes more time to solve optimal transport, our method achieves better performance compared with others.

1. Continuous treatment setting on synthetic data $(\beta = 0.25)$ in one realization

2. Binary treatment setting on the IHDP-1000 data in one realization

Thank you for the valuable and insightful comments. We will revise the submission according to the comments. Our responses are as follows.

Q1: Explanation regarding data splitting.

A1:

1. Theoretically, more samples are helpful for better distribution estimation and alignment, which is supported by Eq. (15) in Theorem 3, in which the term $O(1/\sqrt{\delta n})$ indicates that more samples can induce a tighter upper bound.

2. In practice, less samples bring higher variances. This issue becomes even worse when multiple or continuous treatments are considered. More possible treatment values will induce fewer samples in each subset, resulting in higher variances and inaccurate distribution estimation and alignment. Therefore, it is important to improve the data efficiency to leverage more samples [1].

Q2: Regarding AMSE, Eq. (16), and data splitting.

A2:

1. AMSE is unbiased since the prediction loss is measured on unbiased distribution $q(x)$. AMSE is usually adopted as the target to minimize in causal inference for continuous treatments [2]. Since all the combinations of $x$ and $t$ are considered and the integral is conducted on $q(x)p(t) dx dt$, $p(t=1) \gg p(t=0)$ does not affect AMSE as long as $p(t=0) > 0$, which is satisfied because of the positive assumption.

2. AMSE involves counterfactual outcomes and is intractable. Different from AMSE, the factual loss $\mathcal{L}$ in Eq. (16) is tractable but is conducted on the conditional distribution $q\_t(x) = p(x|t)$ rather than the marginal distribution $q(x)$. We propose to balance the distribution of $q\_t(x)$ and $q(x)$ by minimizing the Wasserstein distance between them, so that the model trained on balanced $q\_t(x)$ can be well generalized to $q(x)$. As shown in Theorem 3, minimizing the Wasserstein distances and $\mathcal{L}$ can reduce the upper bound of AMSE.

3. For the Wasserstein distance $\mathcal{W} (c, \hat{q}\_t, \hat{q})$ involving $\hat{q}\_t(x) = \hat{p}(x|t)$, we do not conduct data splitting. Instead of considering samples receiving $t$ only, we model $\hat{q}\_t = \sum_{i=1}^{n} \hat{q}\_t(x_i) \delta\_{x\_i}$ in Eq. (14), where $\hat{q}\_t(x\_i)$ is estiamted based on the generalized propensity scores. As a result, all the training samples are involved in the empirical distributions, which avoids the issue of data splitting and enhances the performance of distribution estimation.

Q3: Computational complexity.

A3: Running time results are given as follows. Although our method takes more time to solve optimal transport, our method achieves better performance compared with others.

1. Continuous treatment setting on synthetic data $(\beta = 0.25)$ in one realization

2. Binary treatment setting on the IHDP-1000 data in one realization

Q4: Implementation of $\theta$

A4: We are sorry for the confusion. The implementation of $\theta$ is based on [3] and described as follows. Assuming that the conditional distribution of treatment given covariates is Gaussian, i.e., $P(t \mid x_i) \sim \mathcal{N}(\theta(\phi(x_i)), \sigma^2)$. We can estimate the parameters by maximizing the likelihood:

$\max_{\theta,\sigma} L(\hat{\theta},\hat{\sigma};t,x):=\prod_{i=1}^{n} \frac{1}{\sqrt{2\pi \sigma^2}} \exp\left(-\frac{1}{2\sigma^2}(t_i - \theta(\phi(x_i)))^2\right)$.

After that, the estimated generalized propensity score is given by:

$\hat{p}(t \mid x_i) = \frac{1}{\sqrt{2\pi \hat{\sigma}^2}} \exp\left(-\frac{1}{2\hat{\sigma}^2}(t - \hat{\theta}(\phi(x_i)))^2\right)$.

We will revise the submission to make it clear.

Thank you for the valuable and insightful comments. We will revise the submission according to the comments. Our responses are as follows.

Q1: Regarding the assumption of RKHS.

A1:

1. The kernel method requires to calculate kernel functions between all the pairs of two samples, which has high computational complexity.

2. The kernel function between two samples $x$ and $x'$ is the inner product of their mapped representations, i.e., $k(x, x') = \langle \phi(x), \phi(x') \rangle$. In practice, we employ a neural network to approximate the mapping function $\phi$. A similar approach has also been used in [1].

Q2: Explanation of the performance improvement.

A2: The performance improvement mainly comes from two perspectives.

1. Existing distribution alignment methods usually split samples into multiple subsets, each of which only includes the samples receiving a specific treatment. Different from them, we do not split data into smaller subsets. This means that more samples are leveraged for learning. This advantage is theoretically supported by Eq. (15) in Theorem 3, in which the term $O(1/\sqrt{\delta n})$ indicates that more samples induce a tighter upper bound.

2. For balanced representation learning achieved by minimizing the Wasserstein distances, we further introduce the generalized propensity scores into the model. This approach can adaptively assign weights for samples, making samples with higher generalized propensity scores contribute more to the model.

Q3: Ablation study on the loss function involving Wasserstein distances.

A3: We have conducted the ablation study on the loss function involving Wasserstein distances in the revision. The results are reported as follows.

Q4: Comparison with other re-weighting approaches.

A4:

1. Propensity score-based methods employ propensity scores to re-weight samples. Different from them, we leverage propensity scores to model the conditional distributions, and incorporate propensity scores into the optimal transport model for balanced representation learning.

2. Some re-weighting methods learn weights for samples to balance the distributions of different groups, where the distribution shift is measured by a predefined metric. However, for distributions with a very large shift, re-weighting cannot well reduce the distribution shift since the support sets are fixed. Different from them, we characterize the confounding bias by considering the balancing error between the conditional distribution and the marginal distribution as shown in Eqs. (6) and (7), and connect the confounding bias with the optimal transport, which motivates us to learn balanced representations to reduce the confounding bias.

Q5: Regarding unobserved confounders.

A5: Since we characterize the confounding bias by measuring the discrepancy between $q_t(x)$ and $q(x)$, the ignorability assumption is required. If unobserved confounders exist, the confounding bias can not be fully captured by considering $q_t(x)$ and $q(x)$ only. We will investigate the situation with unobserved confounders in the future.

Q6: How to estimate the probability mass $q(x_i)$ for high-dimensional variables?

A6: The probability mass $q_t(x)$ is estimated by generalized propensity scores, which can be estimated by a regression model. For high-dimensional variables, we can reduce the dimension by PCA or neural networks, or introduce the $\ell_1$ norm to select informative variables.

[1] Shalit U, Johansson F D, Sontag D. Estimating individual treatment effect: generalization bounds and algorithms. International conference on machine learning. PMLR, 2017: 3076-3085.

Thank you for the valuable and insightful comments. We will revise the submission according to the comments. Our responses are as follows.

Q1: Explanation of the method.

A1:

1. Following [1], we aim to reduce AMSE that is the unbiased outcome estimation error for all the treatment values. Based on Lines 289-296 and Eq. (12), AMSE is upper-bounded by the Wasserstein distances between $q_t$ and $q$, as well as the feasible biased loss $\varepsilon_{q_t}(h_t)$. This analysis motivates us to minimize the upper bound of AMSE, i.e., the sum of the feasible loss and the Wasserstein distances, and design the algorithm in Section 4.3.

2. Generalized propensity scores obtained by $\theta$ are modeled as $\hat{q}_t$ and are used in the minimization of the Wasserstein distances, as shown in Eqs. (18) and (19), which tends to obtain balanced representations by learning $\phi$.

3. Without balanced representations, the outcome prediction model $\psi$ will be trained on biased data distributions, which means that the prediction model $\phi(\cdot,t)$ could fit the conditional distribution $q_t$ (i.e., samples receiving the treatment $t$) while not being well generalized to other subsets (i.e., samples receiving $t' \neq t$). By minimizing the Wasserstein distances based on $\phi$ and generalized propensity scores, we can obtain balanced representations, and $\phi$ can be well generalized to all the samples rather than a subset receiving specific treatments.

4. Setting $\phi$ to be identically zero will make the Wasserstein distance zero. On the other hand, setting $\phi$ to zero will induce a very large outcome prediction loss $\hat{\mathcal{L}}$ in Eq. (17). Therefore, minimizing $\hat{\mathcal{L}}$ can avoid the situation where $\phi$ being zero. This is discussed in Lines 290-296, which motivates us to propose Eq. (12).

Q2: Algorithm description.

A2: We are sorry for the confusion. Since Problem (20) involves multiple blocks of variables to optimize, we cannot minimize the objective in one iteration. Instead, in Step 9, we update the parameters of the model based on the gradient descent method. We will revise the submission to make it clear.

Q3: Description of the existing literature.

A3: Thank you for the valuable comments. In this paper, we mainly focus on the machine learning method for causal inference. We will revise the submission accordingly.

Q4: The performance of classical or non neural-network methods.

A4: We have conducted non neural-network methods in the revision. The results are reported as follows.

Q5: The reason why loss $\mathcal{\hat{L}}$ does not incorporate the generalized propensity score $\theta$.

A5: Generalized propensity scores obtained by $\theta$ are modeled as $\hat{q}_t$ and are used in the minimization of the Wasserstein distances, as shown in Eq.s (18) and (19), which tends to obtain balanced representations by learning $\phi$. Based on balanced representations, the outcome prediction model $\psi$ can be trained on unbiased data distributions, making the prediction model $\phi(\cdot,t)$ generalized to all the samples rather than only the subset receiving the treatment $t$.

[1] Nie L, Ye M, Nicolae D. VCNet and Functional Targeted Regularization For Learning Causal Effects of Continuous Treatments. International Conference on Learning Representations.