Tightening Causal Bounds via Covariate-Aware Optimal Transport

We highly appreciate the reviewers' summar and comments.

Several other PI methods have been published in recent years that do not necessarily rely on optimal transport theory, ...

The other PI methods deal with the case where there is unobserved confounder/leaky IV, such that the observed marginal distributions of outcomes do not equal the true ones, thus needing PI. In contrast, we are interested in estimating a functional of the joint coupling of the outcomes, which is also partially identifiable and thus we seek PI through OT (a coupling theory). As a result, our PI and the other PI methods that the reviewer lists deal with different cases of partial identifiability, thus are not very relevant and cannot be compared in the our setting.

It appears from the code that this is fixed at 1. Unless I'm missing something?

In the code, the input $d_Z$ for function vip_estimate can be set to any integer, for computing the true value of Vc or Vip, in the simulation result, we make the examples using $d_Z = 1$.

The Neyman CI approach is not really an apples-to-apples comparison.

The Neyman CI approach is basically about the variance estimation of the outcomes from the sample, which is a standard estimand in causal inference. We use this as an example to show that the quadratic objective h is a practical one. Specifically, in [1], the tightest Neyman variance estimator without covariate information is the same as the optimal objective value of the associated OT problem, and our approach incorporates the covariate information in this formulation on top of their work.

Some aspects of this manuscript are clearly quite thoroughly researched ... However, this is not the case when it comes to partial identification (even including Appx A1) …

Our PI and the other PI methods that the reviewer lists deal with different cases of partial identifiability.

In the reviewer’s list, due to hidden confounder / leaky instrumental variables, the observed marginal data distributions are different from the target marginal distribution. (i) hidden confounder: https://ojs.aaai.org/index.php/AAAI/article/view/17437 https://proceedings.mlr.press/v162/zhang22ab.html https://proceedings.mlr.press/v213/padh23a.html (ii) leaky instrumental variable: https://proceedings.mlr.press/v235/jiang24b.html https://openreview.net/forum?id=OaJLMx2nwS

In contrast, our PI is due to the unobserved joint coupling distribution of the outcomes in different treatment groups (i.e. cross-world dependence). Therefore, we utilize the framework of optimal transport to analyze the possible couplings of the observed cross-world outcomes. As a result, our PI approach and the other PI approaches the reviewer lists target different types of partial identifiability. In future work, we will investigate the combination of our PI and the other PIs, i.e., dealing with scenarios involving two sources of partial identifiabilities.

In the literature related to the PI relevant to our approach (PI on joint couplings), we have discussed several papers in the Appendix A. In the revised version, we will discuss the difference between our PI and the other PI that is more relevant to the reviewer’s list (e.g. [2]).

I'm intrigued by this target estimand ... between an ATE and a CATE …

Our target estimand is different from ATE and CATE in that, in our randomized experiment setting, both ATE and CATE are identifiable ($E[Y(1) - Y(0)]$), while our target estimand $E[(Y(0) - Y(1))^2]$ is partially identifiable because it not only relies on the marginal distribution of $Y(0), Y(1)$, but more importantly depends on their unobserved joint distribution. Therefore, the nature of our target estimand and ATE (CATE) are different, thus our PI targets on the partial identifiability of coupling of outcomes, while other PI, particularly ATE, CATE, may focus on the the partial identifiability resulted by hidden confounders.

(a) whether the user is supposed to know the true parents/children of Z and (b) which of the following graphs are fair game ...

We set $X$ to be treatment, $Y$ to be outcome, and $Z$ to be covariate.

In view of a causal diagram, our randomized experiment setting is: (3) $X \rightarrow Y \leftarrow Z$. Our approach can be extended to the case (1) $Z \rightarrow X \rightarrow Y \leftarrow Z$, i.e., allowing the covariates to influence the treatment assignment, by incorporating the propensity scores. We leave this direction for future research.

According to the DAG, we have: Z has no parents, and Y is the children of Z. Actually, for the vector variable Z, we allow an arbitrary structure inside the vector Z. In the paper,, we assume there are no latent confounders (Assp. 3.1).

[1] Aronow, Peter M., Donald P. Green, and Donald KK Lee. "Sharp bounds on the variance in randomized experiments." (2014): 850-871.

[2] Imbens, Guido W., and Donald B. Rubin. Causal inference in statistics, social, and biomedical sciences. Cambridge university press, 2015.

We highly appreciate the positive comments of the reviewer and are happy to answer any question if needed.

We greatly appreciate the reviewer's summary and questions. In the following, we address each question individually.

In Assp 3.3, why is it important to assume that $\nabla_y h(y, \cdot)$ is injective for all $y \in \mathcal{Y}$ ?

This is because we define the $V_{ip}(\eta)$ by the expectation of $h$ under a coupling $\pi$ that is computed through our mirror-relaxed OT problem. To make the $\pi$ well-defined, a standard assumption of to let the objective to have an injective gradient everywhere, e.g., a quadratic function.

In Prop 4.1, do we need the ratio $\frac{n}{m}$ to converge to a positive constant?

For Prop 4.1, as long as $m,n \rightarrow \infty$, the estimator is consistent. The convergence rate may depend on the ratio $\frac{n}{m}$, but we do not invetigate in this direction. Alternatively, in Thm 4.1, we let the convergence rate to depend on $\min (m,n)$.

For Table 3, can the authors somehow obtain the true value of $\rho$ in this example?

Unfortunately, in table 3, we are analyzing a real dataset and the true value of $\rho$ is unidentifiable. We thank the reviewer for their suggestion on making a simulation such that $\rho > 0$ and the lower bound is meaningful. As a response to this, we can actually twist the synthetic experiments in Fig 3 to set the estimand to be $\rho > 0$, since estimating $\rho$ is essentially estimating for quadratic $h$. In the real dataset, the sample appears to be noisy and the correlation between covariate and outcome is not as significant as that in the synthetic dataset (this is why the lower bound is negative, indicating that we cannot reject the possibility of negative $\rho$); In Fig 3, we have larger sample, the estimation error is relatively small, and we set the correlation to be significantly positive, therefore, in that case, the lower bound will be larger than zero.

One important point that the author may consider addressing is to design a valid statistical inference procedure on the partial identification bounds.

The OT estimation may struggle when the dimension of covaraite plus outcome is larger than three, thus a useful inference procedure would be challanging to construct. However, when the dimension is smaller than three, there are some CLT limiting results that can be used to build an inference procedure. We will leave this for future research.

Can the authors generalize the results to convex cost functions?

This is also in our scope of the future research. In the present version, the results depends on the convergence rate of the Brenier's map, which are studied under quadratic loss function. An extension to convex cost functions would be very nontrivial, and we are investigating alternative ways to reduce the effect of large $\eta$ on the magnitude of the estimator.

For Figure 3 (c,f), I wonder why the $L_1$ error of $\hat V_{ip}(\eta)$ will eventually go up as $\eta$ increases.

This is because even though $V_{ip}(\eta) \leq V_{c}$ in the population level, the estimator $\hat V_{ip}(\eta)$ may overestimate $V_{ip}(\eta)$ to be larger than $V_{c}$ based on a finite sample. Therefore, for the $\eta$ such that $\hat V_{ip}(\eta)$ is larger than $V_{c}$, the $L_1$ error increases since $\hat V_{ip}(\eta)$ is increasing with respect to $\eta$.

On Line 727: did you define the operation $\circ$ somewhere in the paper?

Sorry for the confusion, $\circ$ denotes the composition of mapping, which will be added in the revised version.

We highly appreciate the positive comments of the reviewer and are happy to answer any questions if needed.

We address potential limitations.

The penalty tuning.

We provide a data-driven selection method for Q1, which works well in Fig 3, 4.

Efficiency loss relative to COT.

No direct estimator of COT has been established (see Q2). Although there is a gap of Vip and Vc, we maintain consistency and explicit convergence rate based on a finite sample, which is not known for COT.

Multi-valued or continuous treatments.

Vip can be extended to multi-valued treatment using multi-marginal OT. Continuous treatments are not applicable directly for classic OT framework.

1. Penalty:

Under-penalizing may get an estimator to be below the COT value. Over-penalizing leads to a slow convergence on the $\eta$ in Thm 4.3, (thus minimizing empirical PI width is not ideal). Notably, $V_{ip}(\eta)$ with arbitrary $\eta$ already improves over Vu by using covariate, Vc has no known direct estimator. We use the elbow method to select $\eta$ as in [6], which is common in unsupervised setups, such as determining the number of principal components in PCA.

2. Loss to COT:

Exp 3.8 shows the gap between $Vip(\eta)$ and Vc for Gaussian model. In general, we do not have a bound. In Sec 6, we discuss a variant to address this issue. We aim to utilize the covariate to improve over Vu. (Prop. 3.6, Sec 5). COT has no direct estimator, Vip possesses an explicit finite-sample convergence rate (Thm 4.3).

3. High Dim:

The convergence of OT value is known to depend on the dimension, like a rate of $O(n^{-2/d})$ for squared Wasserstein [7], which is sharp [4]. Note that, our convergence rate in Thm 4.3 is aligned with this result. If we have a reliable parametric distribution of Y given Z, the Vip estimator is compatible with the models. But we should note the risk of model misspecification as in Fig 3. (b) (e). If no parametric model is available, a variable selection method for covariate would help. See Q10.

4. Overfitting:

Fig 3 and Fig 4 show that the elbow method is stable as the sample size goes from 500 to 1500.

5. Extension:

For multi-level, the Vip estimator can be directly extended using the multi-marginal OT [5]. For continuous value, our causal bound is out of the scope of a standard OT framework.

6. Sensitivity to h:

The cost function h in the causal estimand is typically chosen by the user, e.g. Sec 5.3. The Vc is a minimum of a linear functional of h, an error of order $O(a)$ will cause an error of order at most $O(a)$.

7. Real-World:

Vip is the OT problem with objective h with encouraging the alignment of Z in the two groups by regularization. Regularization is popular in classic statistics and machine learning approaches, e.g. LASSO.

8. Semi-Parametric

Semi-para estimates identifiable causal quantities but we do partially identified ones. Bounds based on variance inequalities are not consistent but ours is consistent. Fan (2023) considers inference for PI with moment equalities. They model the joint distribution of Y using copula, which is suited for univariate outcomes, but our approach can handle vectors. They use Bernstein copulas, but we have no restrictions on the coupling. They require estimating conditional CDFs, we don’t.

9. Alters:

E.g. [3]. But our mirror method has pros: (i) the mirror relaxation $\leq$ COT, and for causal bound, an underestimation beats overestimation since a lower bound still forms a valid causal bound. (ii) we guarantee an improvement over Vu.

10. Var-Select:

$\sup_{\eta} \inf_{\pi} \int h(Y(0), Y(1)) + \sum_{i} \eta_i |Z^{(i)}(0) - Z^{(i)}(1)| d\pi - C \sum_{i} |\eta_i|.$

We use the L1 norm of the $\eta$ vector to encourage sparsity.

11. Data Perturb:

The robustness property of $V_{\ip}$ against data perturbation inherits from the stability of the OT with respect to its marginals. When the objective function $h$ is Lipschitz, then a distributional shift of order $O(a)$ measured in 1-Wasserstein distance induces a bias of at most order $O(a^{2/15})$ on the OT map used in computing $Vip(\eta)$, also the estimator [8].

12. Hidden Confound:

In this case, we can relax the OT. Suppose that the bias of observed conditionsal law is measured by a distance D and smaller than a constant, then a relaxed OT problem (or robust OT) can be used [1,2].

[1] "Optimal transport with relaxed marginal constraints."

[2] "On robust optimal transport: Computational complexity and barycenter computation."

[3] "Consistent optimal transport with empirical conditional measures."

[4] "Sharp convergence rates for empirical optimal transport with smooth costs."

[5] "Bridging multiple worlds: multi-marginal optimal transport for causal partial-identification problem."

[6] “Determining the number of clusters/segments in hierarchical clustering/segmentation algorithms.”

[7] “On the rate of convergence in Wasserstein distance of the empirical measure."

[8] "Quantitative stability of optimal transport maps and linearization of the 2-Wasserstein space."