Counterfactual Data Augmentation with Contrastive Learning

Dear Reviewer AQb9,

Thank you for your feedback. We appreciate your attention to detail and have carefully considered your comments. We updated our manuscript accordingly (see the corrected proof of theorem 1 in Appendix B). Please find below a detailed answer for the rest of the questions.

Question: “I do not agree that individuals from both groups are identical...”

In the context of causal inference, when referring to the distribution of the treatment group and the control group we intend the marginals of the features $p(X|T=1)$ and p(X|T=0), not $p(Y,X|T=1)$ and $p(Y,X|T=0)$. Our characterization is accurate in illustrating that if $p(X|T=0) = p(X|T=1)$, standard learning theory guarantees that good performance on $p(X|T=1)$ implies a good performance on prediction $Y_1$ on $p(X|T=0)$ after learning the potential outcome $Y_1$ on $p(X|T=1)$.Thus, the challenge in estimating counterfactual outcomes well lies in the disparity between the treatment groups, specifically in their feature distribution.

Question: “On page 3, it stated that "The counterfactual distribution, denoted by pCF, is defined as the sampling distribution”

By inverting the treatment assignment mechanism we indeed change the structural equation relating the treatment variable T to the features X. Hence the factual outcomes will change. If the reviewer means by completely determined that there is a deterministic function between the observed factual outcome $Y = T Y1 + (1-T)Y0$ and the counterfactual outcome random variable $Y_{CF} = (1-T) Y1 + T Y0$ then the answer is clearly yes. If he means that there is no information between both then the answer is clearly no.It is not true that $p_{F}(y|x,t) = p_{CF}(y|x,t)$ instead it is true that $p_{F}(y|x,1-t) = p_{CF}(y|x,t)$ (follows from the provided identities above).

Question: :”On page 4 above (5), it stated that "Let ℓ : {0, 1} × {0, 1} → R". Later, function g serves as the first arguement which has range [0,1] (the probability). Hence I wonder if it should be ℓ : [0, 1] × {0, 1} → R”

Thank you for catching this typo, we will update the manuscript accordingly.

Question: “Solving equation (5) can have two challenges.

The computational challenge and imbalance issue in solving equation (5) is indeed significant. While it is true that the number of pairs is increasing at a rate $\frac{N*(N-1)}{2}$, we only need to select a subset of all possible pairs for the training. For the imbalanced problem, we only select a subset of individuals for which we perform well, that is why we don’t augment every point in the dataset. Moreover, the regions where we have high overlap between the different treatment groups are the regions where we will do most of the imputation and these are the regions where there is the least imbalance between the different treatment groups. In practice, to overcome the imbalance problem, it is easy to sample a mini-batch with an equal number of positive and negative pairs during the learning of the contrastive learning module.

Question: think this consistent property is not just a property of pF alone, but also a joint property of both pF and pCF

The mentioned identity $p_{CF}(x, 1 − t) = p_{F}(x, t)$ is always true for any binary treatment assignment while consistency only holds when the treatment assignment is independent of the features $X$. There is no direct relationship between unconfoundedness, a property about observing all the features that are causally affecting $T$ and $Y$, and consistency, as defined in Definition 5, a property that is equivalent to the factual and the counterfactual distribution being identical. As an example, we can imagine a scenario where the treatment is assigned by flipping a coin without knowing any of the patient's information. We know that unconfoundedness doesn’t hold in this case (since the features are not observed) but consistency holds. Also, we can have a case where we observe all the features affecting T and Y but there is still a bias and unconfoundedness doesn’t hold. Similarly, we can construct scenarios where they both hold and both don’t.

Question: While this theorem guarantees reliable imputation of its counterfactual outcome, it does not establish the asymptotic consistency of COCOA (which was claimed to be the case). It only says that COCOA can estimate the counterfactual distribution well, but for consistency you need to check Definition 5. 6.2 Theorem 2 is very trivial and the proof is very elementary. I would not call this a theorem.

We will label it as a proposition instead. While being easy to prove, it highlights the important point that the likelihood of encountering similar individuals under positivity is increasing, which is crucial for local imputation of missing outcomes. We will clarify this in the manuscript.

Dear Reviewer NLwi,

Thank you for your feedback.

1. Theoretical Results Significance: We respectfully disagree with the reviewer, the provided theoretical analysis establishes a foundation for understanding data augmentation methods in causal inference. Mainly Theorem 3, postulates that the error in estimating treatment effects is bounded by the loss on the training augmented data, the distance between the features distribution in the augmented data and that of a Randomized Controlled Trial, and the average error for the local approximator. It is important to note that the first term is guaranteed to converge to $0$ by standard learning theory, the second term is guaranteed to converge to 0 as soon as we select all the individuals to impute their outcomes, and the third term depends on the quality of imputation but is also guaranteed to converge to 0 as the number of data points increase. This being said we agree with the reviewer that further investigation of using contrastive learning to identify smooth local regions for local approximation is an open research question that requires a more theoretical understanding of contrastive learning in general. We will be happy if the reviewers can point us to good theoretical works for contrastive learning that will help us in this research direction. We will also change the wording in our updated manuscript to make the limitations of our preliminary theoretical analysis clearer.

2. Hyperparameters $\epsilon$ and $K$: We agree with the reviewer that these hyperparameters are very important for the efficacy of the method, which is why we conducted extensive ablation studies on them presented in the main text and in Appendix C.3. We conducted our experiments in benchmark causal inference datasets (IHDP, Twins, and News) as well two synthetic datasets (a Linear and Nonlinear dataset). Our experimental results suggest that our method is robust to a wide range of choices of hyper-parameters: it is not challenging to find a set of hyper-parameters that consistently improve the performance. This being said, we believe developing principal ways to tune these hyperparameters (as well as other widely used hyperparameters in Machine Learning in general and causal inference in particular) is a very important research direction that is beyond the scope of the current presentation.

3. Sample Size Effect: Semi-synthetic datasets such as IHDP already have a limited number of individuals, and reducing the dataset size further complicates the learning process, particularly for data augmentation methods. As the dataset shrinks, finding close neighbors becomes increasingly challenging. Consequently, when the number of data points is very small, the method may struggle to identify suitable neighbors, potentially leading to no data points being selected for data augmentation. However, we would like to note that compared to other methods, our algorithm will not perform any data augmentation in this case. This is one of the key properties of our algorithm that it only imputes with high confidence, thus never worsening the problem.

We thank the reviewer again for his feedback and we will be happy to answer any further questions.

Dear Reviewer Lde6,

Thank you for your constructive feedback and for the relevant resources you shared with us. We will include them in our revised related work section.

1. The contrastive loss: Contrastive Learning is used as a heuristic (proven to be effective in various applications) to identify similar individuals within the same treatment group. To clarify the potential misunderstanding, as in the provided example the pair $\{X = \text{“young”}, T=0, Y=0\}$ and $\{X = \text{“old”}, T=0, Y=0\}$ will be in the positive dataset. And the pair $\{X = \text{“young”}, T=1, Y=10\}$ and $\{X = \text{“old”}, T=1, Y=0\}$ will be in the negative dataset. In this case, no data augmentation will happen. This is indeed the key to the success of our algorithm: it does not impute unreliable outcomes. Assuming that we have more points. Therefore if we get individual $a=\{X = \text{“old”}, T=1, Y=0\}$ (Y value is hidden from us), we can test that it will be similar to the individuals $b=\{X = \text{“old”}, T=0, Y=0\}$ and we can hence impute his outcome and have $Y \approx 0$ (a has the same $X$ value as $b$ hence $b$ is a close neighbor of $a$, and we use $b$’s factual outcome as the imputation results for $a$).

2. The place of the theory: The proposed theoretical analysis applies to general data augmentation in causal inference. While local approximators are guaranteed to converge (there has been a considerate amount of work that studies the statistical consistency of local approximation methods as in [1], hence we believe it is unnecessary to repeat these results in this work), the contrastive learning part is a heuristic that we use to identify similar individuals. The other natural choice is using Euclidean distance but that doesn’t generalize well for high dimensional data. We will add a remark to our revised manuscript outlining this.

3. Comparison to ReiszNet: To the best of our understanding the provided reference builds a method to estimate the average treatment effect (ATE). Our proposed method is mainly built to estimate the conditional average treatment effect (CATE) at an individual level. Hence, we fail to see a direct way of applying Reisznet to CATE estimation. We would appreciate it if the reviewer could share advice on how to adapt ReiszNet to CATE estimation. This being said, we will include this method for completeness in the related work when discussing ATE estimation methods.

4. Selection Criteria: Yes we use the same threshold $\epsilon$ in practice, hence if $\|g(x_i,x_j)\|\leq \epsilon$ we classify the individuals as similar otherwise we classify them as non-similar. We will add this detail to the algorithm section for enhanced clarity.

Thank you again for your constructive feedback.

[1] Tjøstheim, D., Otneim, H. and Støve, B., 2021. Statistical Modeling Using Local Gaussian Approximation. Academic Press.

Dear Reviewer FUqm,

Thank you for your valuable feedback.

1. Computational Complexity: The computational complexity of identifying the similar and the non-similar individuals is $\frac{N^2}{2}$ in the worst-case scenario. The complexity of the rest of the learning algorithm is on par with that of the other modules. For very large datasets, we can select a subset of the possible similar and non-similar pairs. In practice, when the pairs of positive samples are moderate, we can select only a small set of negative samples for every gradient update of the contrastive learning module. Please note that when the dataset contains lots of pairs of similar individuals, data augmentation is not needed as performance is guaranteed by our error bound (simply setting $\alpha = 0$).

2. Binary Outcomes: Since we assume that the functions are locally smooth and can be approximated locally by a linear model or a Gaussian Process, the method currently does not generalize well to classification tasks. We will incorporate this clarification in our revised manuscript.

3. Definition 4: Definition 4 establishes the counterfactual and factual losses, which are crucial in the proposed theoretical analysis in Section 4. Please note that $t$ in the definition is the treatment assignment (that is discrete in our case). This analysis introduces a notion of consistency for data augmentation in causal inference.

We hope these clarifications address your concerns and reinforce the value of our work. Thank you once again for your insightful comments.