Revisiting Covariate and Hypothesis Roles in ITE Estimation: A New Approach Using Laplacian Regularization

Response to Reviewer UdZR

We appreciate the reviewer’s thoughtful feedback and constructive comments, which have provided valuable insights and helped us refine our work. Below, we address the specific points raised in detail.

Theoretical Framework

We appreciate the reviewer’s feedback and concerns regarding the constraint term in Theorem 3. The theorem provides an upper bound on the ITE error that is fundamentally different from the classical covariate shift-based bounds. Instead of focusing on aligning covariates, our upper bound explicitly depends on the difference between the expected loss predictions for the potential outcomes. This shift in focus moves away from directly minimizing the covariate shift and emphasizes the role of prediction discrepancies in bounding the ITE error. The proposed upper bound explicitly incorporates the hypothesis functions, advocating for their integration in ITE estimation under the covariate shift without direct minimization as applied by traditional methods. This reflects our assertion that covariate shifts in ITE estimation are often rooted in highly predictive features, which should be utilized rather than minimized.

Following the theoretical framework, our approach employs Laplacian regularization to ensure smooth predictions under the covariate shift. The Laplacian graph encodes the covariate shift geometry through its edges, enabling smooth predictions across the learned representation space. Crucially, we do not directly regularize the covariate shift in the representation space. This decision preserves highly imbalanced features, which are often predictive and essential for accurate ITE estimation.

We have also revised the simplification in equation (22) to maintain a tighter proposed bound. Our updated formulation is as follows: $(1-u) \int_{\mathcal{R}} p^{t=0}_{\Phi}(r) \bigl| \ell\_{h, \Phi}(\Psi(r), 1) - \ell\_{h, \Phi}(\Psi(r), 0) \bigr| dr + u \int\_{\mathcal{R}} p^{t=1}\_{\Phi}(r) \bigl| \ell\_{h, \Phi}(\Psi(r), 0) - \ell\_{h, \Phi}(\Psi(r), 1) \bigr| dr = \int\_{\mathcal{R}} p\_{\Phi}(r, t) \bigl| \ell\_{h, \Phi}(\Psi(r), 1) - \ell\_{h, \Phi}(\Psi(r), 0) \bigr| dr,$ where the joint distribution $p\_{\Phi}(r, t)$ satisfies $p\_{\Phi}(r, t) = (1-u) \cdot p^{t=0}\_{\Phi}(r) + u \cdot p^{t=1}\_{\Phi}(r)$. This modification still demonstrates that the counterfactual upper bound reflects the difference between the expected loss predictions—now weighted by the joint distribution over $t$ and $r$. The bound is still dependent on the hypothesis functions, distinguishing it from classical upper bounds.

Section 3.2 Title Alignment with Content

We aimed to explain selection bias in the context of ITE estimation, emphasizing the challenges arising from having only a partial view of the joint distribution. Following the reviewer’s comment, we will revise the title to "ITE estimation under the selection bias" to better reflect the content of the section.

Response to Reviewer Wr16

We sincerely thank the reviewer for their detailed feedback and critical evaluation. While we understand the concerns raised, we believe they provide an opportunity to clarify the significant contributions and address any misunderstandings about the novelty and impact of our work. Below, we respond to the specific points raised.

Contribution Significance

We appreciate the reviewer’s feedback and understand the concerns regarding the perceived novelty of Equation (13) and Algorithm 1. However, our work, particularly in the context of Individual Treatment Effect (ITE) estimation, extends beyond these elements. Our approach revisits the role of covariate shifts and hypothesis functions within ITE estimation, presenting a novel perspective that distinguishes it from existing methods. We depart from the classical paradigm, which has been widely used for the past nine years across many methods. Instead of directly minimizing the covariate shift, we emphasize the importance of estimating ITE under the covariate shift. We posit that the covariate shift is rooted in highly predictive features, as doctors assign treatments based on those features, in contrast to classical domain adaptation, where covariate shifts are treated as artifacts to be mitigated. To support this, we provide an upper bound on ITE error that depends on the hypothesis function rather than the classical covariate shift. Following this theorem and arguments, we advocate for hypothesis functions that are designed to estimate potential outcomes under the covariate shift rather than minimizing it. We consider this new concept as a contribution in its own right.

Our proposed LITE method, while simple, is effective and provides state-of-the-art results on two leading benchmarks. It outperforms 21 competing methods, including state-of-the-art methods from recent years. Additionally, our regularization term (LITE) is not utilized only to align the learned representation with the hypothesis function outcomes, but also to dynamically learn the graph structure itself through the optimization process. The latent representation and the resulting graph are continuously refined through optimization, aligning the geometry of the learned representation with the intrinsic geometry of the data.

Graphs and Similarity Kernels

The common practice for representing a graph involves using a similarity kernel. Based on this graph, we define the (graph-)Laplacian operator, which enables us to construct the quadratic Laplacian form. Minimizing this quadratic form promotes predictions that align with the span of smooth eigenvectors of the Laplacian graph. As a result, the counterfactuals are explicitly inferred from the factual predictions, grounded in observations, by leveraging the smoothness requirement over the learned geometry of the covariate shift.

Response to Reviewer x8oW

We thank the reviewer for their valuable comments, which have helped us better articulate the key points of our contribution.

Q1: The Covariate Shift

Please note that the purpose of LITE is not to minimize the covariate shift, and any connections between $h_1$ and $h_0$ are not assumed in ITE estimation. The covariate shift in ITE estimation stems from the non-random treatment assignment that promotes imbalanced features between the control and treated groups. Conventional methods for ITE estimation directly try to reduce the covariate shift in the representation space to allow estimation of both potential outcomes from more "reliable" aspects of the data and thus allow better ITE estimation. We assert that as opposed to classical domain adaptation, where covariate shift is treated as an artifact to be mitigated, in ITE estimation, those highly imbalanced features are also highly predictive (as doctors, for example, assign treatment based on those features). LITE aims to allow predictions of both potential outcomes under the natural covariate shift rather than minimizing it. To this end, LITE utilizes the Laplacian, which encodes the covariate shift in the representation space within the graph edges. The quadratic Laplacian form regularization allows smooth prediction over the covariate shift geometry and thus allows more reliable estimations in regions where factual labels are limited.

Q2: Theoretical Framework

We derive a new upper bound of the expected ITE loss and show that it explicitly depends on hypothesis functions rather than the classic covariate shift bound. Following this general theorem, we depart from conventional approaches that minimize the covariate shift and call for the incorporation of the hypothesis function to estimate the potential outcomes under the covariate shift. We acknowledge that our theory does not directly relate to the proposed Laplacian method but advocates for incorporating functions in ITE estimation instead of solely relying on covariate shift minimization. The Laplacian regularization allows the hypothesis function to estimate under the covariate shift by promoting smoothness of the factual and counterfactual predictions over the geometry of the learned representation encoded within the Laplacian graph.

Q3: Simulation Study

In simulation 5.1, our aim was to show the effectiveness of our term as an ablation study. Therefore, this section was focused on the LITE method only. Across the simulations on News and IHDP leading datasets, we compared our method to 21 methods, including state-of-the-art methods. Our superior empirical results demonstrate the effectiveness of considering covariate shift rather than minimizing it.

Q4: Lack of Ground Truth and Hyperparameter Tuning in ITE Estimation

The lack of ground truth counterfactual for hyperparameter tuning is inherent in ITE estimation and poses a challenge to all methods. There are multiple criteria for hyperparams in ITE estimation, and those methods are actually essential for real-world data and may affect the performance of ITE estimation. For more details on hyperparameter tuning in ITE, see [1] and [2]. To disentangle the effect of hyperparams criteria mechanism and the performance of the ITE estimation method, as we note in the appendix, we follow the commonly used practice in the literature for hyperparameter selection [3] based on the PEHE metric on the validation set, while the hyperparameters are fixed across multiple realizations (1000 for IHDP and 50 for News). Although not applicable to real-world data, this approach validates the robustness of the selected parameters and prevents overfitting.

References:

[1] Machlanski, D., Samothrakis, S., and Clarke, P. Hyperparameter tuning and model evaluation in causal effect estimation. arXiv preprint arXiv:2303.01412, 2023.

[2] Saito, Y., and Yasui, S. Counterfactual cross-validation: Stable model selection procedure for causal inference models. In International Conference on Machine Learning, pages 8398--8407, 2020.

[3] Johansson, F., Shalit, U., and Sontag, D. Learning representations for counterfactual inference. In International Conference on Machine Learning, pages 3020--3029, PMLR, 2016.

Response to Reviewer TWu7

We appreciate the reviewer’s insightful comments, which have allowed us to better clarify and emphasize the key aspects of our contribution.

Terminology

We acknowledge the distinction between ITE and CATE as highlighted by the reviewer. In our manuscript, the term CATE is used when formally defining ITE. While the use of ITE throughout the paper may not strictly conform to its traditional definition, we will include a note in the ITE definition clarifying that we are actually estimating CATE, which is the expectation of ITE, while throughout the paper, we refer to CATE as ITE interchangeably. The CATE captures a causal effect which is as specific to an individual, rather than a population-level causal effect, and therefore, it is sometimes called individual treatment effect (ITE). This choice aligns with the terminology adopted in prior works and is discussed in [1].

Proposed Theorem 3

To highlight our novel contributions, Theorems 1 and 2 are now defined as Propositions. We have also revised the simplification in equation (23) to maintain a tighter proposed bound. Our updated formulation is as follows: $(1-u) \int_{\mathcal{R}} p^{t=0}_{\Phi}(r) \bigl| \ell\_{h, \Phi}(\Psi(r), 1) - \ell\_{h, \Phi}(\Psi(r), 0) \bigr| dr + u \int\_{\mathcal{R}} p^{t=1}\_{\Phi}(r) \bigl| \ell\_{h, \Phi}(\Psi(r), 0) - \ell\_{h, \Phi}(\Psi(r), 1) \bigr| dr = \int\_{\mathcal{R}} p\_{\Phi}(r, t) \bigl| \ell\_{h, \Phi}(\Psi(r), 1) - \ell\_{h, \Phi}(\Psi(r), 0) \bigr| dr,$ where the joint distribution $p\_{\Phi}(r, t)$ satisfies $p\_{\Phi}(r, t) = (1-u) \cdot p^{t=0}\_{\Phi}(r) + u \cdot p^{t=1}\_{\Phi}(r)$. This modification still demonstrates that the counterfactual upper bound reflects the difference between the expected loss predictions—now weighted by the joint distribution over $t$ and $r$. The bound is dependent on the hypothesis function, distinguishing it from classical upper bounds. Following the new Theorem 3, our approach deviates from traditional methods that minimize the covariate shift. Instead, we advocate for integrating the hypothesis function into the estimation of potential outcomes under the covariate shift. We posit that the treatment assignment, which contributes to the covariate shift, is rooted in a predictive feature and should not be reduced in the representation space as done in classical methods.

Why Imposing Smoothness on Prediction Helps Bound the Distance Term in Theorem 3

In scenarios with random treatment assignment, we would anticipate that predictions for both potential outcomes would typically exhibit similar levels of expected loss. Consequently, the difference, which constitutes our upper bound, should be minimal. However, in cases with selection bias, at each point $r$ (i.e., the covariate representation) the presence of factual labels may inherently favor one potential outcome over another due to the treatment assignment policy (i.e., the covariate shift). This results in a larger expected loss at point $r$ for the outcome with fewer factual labels available compared to the other, thus leading to a larger difference between both expected loss predictions.

LITE is designed to facilitate reliable predictions for both potential outcomes within the inherent covariate shift, thereby yielding smaller upper bounds. We first calculate both potential outcome predictions for each sample—both factual and counterfactual. Subsequently, LITE evaluates the quadratic Laplacian form for each potential outcome. In particular, for each potential outcome, this term includes both factual and counterfactual predictions, in addition to the Laplacian operator, which quantifies the distances among all samples (representing the covariate shift in the representation space). The quadratic form minimization promotes predictions that live in the span of smooth eigenvectors of the Laplacian graph. Thus, the counterfactuals are explicitly inferred from the factual predictions, which are grounded in observations, using the smoothness requirement over the covariate shift geometry. Not requiring smoothness in regions with sparse factual labels—i.e., omitting our proposed term—leads to unreliable predictions, resulting in larger upper bounds.

Continuity of Potential Outcomes

In tasks where potential outcomes fall into categorical predictions, we represent these categories using one-hot vectors, while the predictions are generated by the sigmoid function, yielding continuous values between 0 and 1. This approach allows us to apply smoothness to these continuous outputs. We note that each categorical decision fundamentally stems from a natural "continuous decision" before being discretized, making these continuous predictions more natural.

References:

[1] Johansson, F. D., Shalit, U., Kallus, N., et al. (2022). Generalization bounds and representation learning for estimation of potential outcomes and causal effects. JMLR, 23(166), 1-50.