Self-Distilled Disentanglement for Counterfactual Prediction

Weaknesses - Response:

Firstly, we acknowledge there is an issue with the interpretation of the disentanglement process for variable Z within the context of the associated loss function in Formula 4. Our initial concept was grounded in the understanding that when using the treatment label T as the feature to predict outcome Y, the information in Z might become redundant. Therefore, we sought to disentangle Z from X by minimizing the KL-divergence between the predicted distribution of Y conditioned on both Rz and T, and the predicted distribution of Y conditioned solely on T, i.e., ${D}_{KL}\left[ p(Y \mid {Z},{T}) \Vert p(Y \mid T) \right]$. This served as the motivation for the design of the loss function used in subsequent experiments, and empirical studies have demonstrated its effectiveness. However, we recognize the inadequacy of the previous explanation. Therefore, in the updated version, we have removed the misleading description of the conditional mutual information $I (Z; Y|T)$. We appreciate your insightful feedback!

Questions - Response:

• The Necessity of the Shallow Network: A shallow network design for Z and C is necessary. As shown in Equation 11, we aim to directly obtain the predicted distribution of T based on Z and C, facilitating the disentanglement of C. This predicted distribution is precisely achieved by designing two shallow networks for Z and C. The retain network consolidates shared information from Z and C, and therefore, it should incorporate more parameters. A deep network is subsequently introduced to output predictions. This is the rationale behind our design of different shallow and deep networks.

• Clarification on the Loss Function: We didn't specify whether it's a zero-sum or non-zero-sum problem. In theory, to minimize the mutual information between $R_z$ and $R_c$, we theoretically only need to minimize either $D_{KL}[P^{R_z}\_{T}|P_T]$ or $D_{KL}[P^{R_c}\_{T}| P_T]$ (we elaborate on choosing $D_{KL}[P^{R_z}\_{T}|P_T]$ for clarity) and $D_{KL}[P^{R_z}\_{T}|P ^{R_c}\_T]$. The corresponding loss functions are $L (Q^Z\_T, T)$ and $L (Q^Z\_T, Q^C\_T)$. However, in practice, including only these two components in the model loss makes convergence challenging, resulting in unsatisfactory performance. We speculate that this is due to a lack of sufficient supervised information guiding the updating direction of the prediction network for $R_c$. Therefore, we introduce the minimization of $D_{KL}[P^{R_c}\_{T}| P_T]$ (corresponding loss function $L (Q^C_T, T)$) to expedite loss convergence. Empirical results indicate that this loss function design improves disentanglement and counterfactual prediction performance, suggesting that, while not precise, it helps the model approach the optimum more quickly. We acknowledge that the original theoretical statements were not rigorous and have modified formulas 10 and 11 in section 4.2, further explaining loss functions in Section 4.3 in the updated version.

  Your understanding of the teacher network is correct. We can set up a deep prediction network and and use its supervision information to train the shallow prediction models. From this perspective, our approach is essentially a self-distillation model and that’s why we name our method as Self-Distilled Disentanglement, i.e., $SD^2$. Therefore, the teacher model is utilized to enhance the network's prediction capability, thereby facilitating the distillation and disentanglement process.

Weaknesses - Response:

Thank you for your advice. We have incorporated your suggestions into the updated version of the paper by adding more captions for Figure 1, transcribing all inline equations into display equations, breaking down the findings into separate bullet points and stating specifying the section numbers in the appendix when referenced, aiming to improve the paper's readability.

Questions - Response:

• Method's name: Here, "distillation" refers to a framework design that utilizes deep classifiers (for binary scenarios) or regressors (for continuous scenarios) to guide shallow ones in knowledge extraction, thereby promoting disentanglement performance.

  Technically, the optimization of equation (8a, 8b, 9a, 9b) can be realized by training prediction networks about $Y$. Additionally, we can set up a deep prediction network from $T$, $A$, $C$ to $Y$ and use its supervision information to train the shallow prediction models. From this perspective, our approach is essentially a self-distillation model. Therefore, we name our method as Self-Distilled Disentanglement, i.e., $SD^2$.

• Metrics in radar charts: Following [2] and [3], we use the first (second) slice to denote the weight matrix that connects the variables in X belonging (not belonging) to the actual variables. The polygons’ radii in Figure 3b are scaled between 0: 0.09 and quantify the average weights of the first slice (in red) and the second slice (in blue). For better readability, we have added the description for this metric in Section 5.3.2. Thanks for the question.

• Discussion under the continuous scenario: This outcome highlights the significant issue of confounding bias stemming from the treatment regression stage in IV-based methods. It emphasizes the essential need to disentangle confounding factors ($C$) from observed variables ($X$) and subsequently address confounding bias led by $C$. We appreciate your question, and we have integrated this discussion into the updated version.

[2]. Negar Hassanpour and Russell Greiner. Learning disentangled representations for counterfactual regression. In International Conference on Learning Representations, 2019.

[3]. Learning decomposed representation for counterfactual inference. arXiv preprint arXiv:2006.07040, 2020.

I thank the authors for their responses.

After carefully reading the other reviews and the authors’ responses, I think the point raised by reviewer FUp5 is very important and must be properly addressed in the paper. It is necessary that their assumptions are tight enough that renders that counterexample not applicable.

In light of this, I have decided to reduce my score.

Weaknesses-Response:

• Eq 3: Firstly, we acknowledge there is an issue with the interpretation of the disentanglement process for variable Z within the context of the associated loss function in Formula 4. Our initial concept was grounded in the understanding that when using the treatment label T as the feature to predict outcome Y, the information in Z might become redundant. Therefore, we sought to disentangle Z from X by minimizing the KL-divergence between the predicted distribution of Y conditioned on both Rz and T, and the predicted distribution of Y conditioned solely on T, i.e., ${D}_{KL}\left[ p(Y \mid {Z},{T}) \Vert p(Y \mid T) \right]$. This served as the motivation for the design of the loss function used in subsequent experiments, and empirical studies have demonstrated its effectiveness. However, we recognize the inadequacy of the previous explanation. Therefore, in the updated version, we have removed the misleading description of the conditional mutual information $I (Z; Y|T)$. We appreciate your insightful feedback!

• Eq 4: In equation (4), we aim to minimize the KL-divergence between the predicted distribution of Y from $R_z$, T and the predicted distribution of Y from T, i.e., ${D}_{KL}\left[ p(Y \mid {Z},{T}) \Vert p(Y \mid T) \right]=-\int p(Y \mid {Z},{T})log \frac{p(Y \mid {Z},{T})}{p(Y \mid T)} dY$. In the revised version, we have removed the relevant statements about $I (Z; Y | T)$. The reasons are detailed in the response to Q1.

• During the training phase, to obtain reasonable representations for C and A, we set up two separate prediction networks from R_c and R_a to Y, which means we use Y as supervisory information to guide the training process of R_c and R_a. This is the only information shared between R_c and R_a during the training process. Therefore, the mutual information between R_a and R_c is all related to Y, and the information contained in Ra that is related to Rc but unrelated to Y is 0 at this stage.

• Corollary 4.3: Thanks for your advice. Here, the notation used in Equation 10 and 11 refers to minimizing all three KL-Divergence together, and we have added related statements in Corollary 4.3 based on your and Reviewer Yekp’s suggestions for better clarity.

• Sec 4.2: We acknowledge that the definitions for these networks need to be sufficiently clear. Below is a brief explanation of these modules, which we have updated in Sec 4.2. “Retain network” here represents the neural network for retaining the information from both Z and C. “Deep networks” and “Shallow networks” are named based on their relative proximity to $R_a$ and $R_c$. Essentially, both are classifiers (for binary scenarios) or regressors (for continuous scenarios).

• Eq. 12: Thank you for your correction. $W$ represents the model weight parameters involved in the $L_2$ regularization loss. As you comprehend, $\alpha, \beta, \gamma$ are global scalars, and $w$ should be example-wise weights. Thus, we modify Eq. 12 in the revised version.

• Metric for “contribution of actual variables”: Following [2] and [3], we use the first(second) slice to denote the weight matrix that connects the variables in X belonging (not belonging) to the actual variables. The polygons’ radii in Figure 3b are scaled between 0: 0.09 and quantify the average weights of the first slice (in red) and the second slice (in blue). For better readability, we have added the description for this metric in Section 5.3.2.

• top of p5: Thank you for pointing out our mistakes. The correct statement is, indeed, "C is independent of A and Z." We have corrected this statement in the new version.

• Maybe consider using m_k etc: Thank you for your suggestion. We have made corrections based on your advice.

• More information of DRCFR: Thank you for your suggestion. We have added relevant statement “ a baseline that ensures the disentanglement of A by proposing Prop 3.1, yet does not provide disentangling solutions for Z and C” in Section 5.3.2 in the new version.

Questions - Response:

Your understanding is correct. As illustrated in Figure 1, T and Y are causally downstream components of different underlying factors in X in the context of our study.

We share the same assumptions for identifying the treatment effects with [1]. Specifically, suppose the generated IVs satisfy three conditions (Relevance, Exclusion, Unconfoundedness) of valid IVs, additional assumptions are necessary to identify the average causal effect (ACE) of T on Y. Adequate assumptions encompass homogeneity in the causal impact of T on Y, uniformity in the relationship between Z and T, and the absence of effect modification. On the basis of [1], we did not introduce additional conditions for identifiability, as our primary contribution lies in designing an algorithm to identify instrumental variables that satisfy the given conditions. We have also incorporated the description regarding identifiability into Section 3 of the revised paper.

Weaknesses-Response:

Firstly, we acknowledge there is an issue with the interpretation of the disentanglement process for variable Z within the context of the associated loss function in Formula 4. Our initial concept was grounded in the understanding that when using the treatment label T as the feature to predict outcome Y, the information in Z might become redundant. Therefore, we sought to disentangle Z from X by minimizing the KL-divergence between the predicted distribution of Y conditioned on both Rz and T, and the predicted distribution of Y conditioned solely on T. This served as the motivation for the design of the loss function used in subsequent experiments, and empirical studies have demonstrated its effectiveness. However, we recognize the inadequacy of the previous explanation. Therefore, in the updated version, we have removed the misleading description of the conditional mutual information $I (Z; Y|T)$. We appreciate your insightful feedback!

Questions-Response:

We share the same assumptions for identifying the treatment effects with [1]. Specifically, suppose the generated IVs satisfy three conditions (Relevance, Exclusion, Unconfoundedness) of valid IVs, additional assumptions are necessary to identify the average causal effect (ACE) of T on Y. Adequate assumptions encompass homogeneity in the causal impact of T on Y, uniformity in the relationship between Z and T, and the absence of effect modification. On the basis of [1], we did not introduce additional conditions for identifiability, as our primary contribution lies in designing an algorithm to identify instrumental variables that satisfy the given conditions. We have also incorporated the description regarding identifiability into Section 3 of the revised paper.

[1]. Fernando Pires Hartwig, Linbo Wang, George Davey Smith, and Neil Martin Davies. Average causal effect estimation via instrumental variables: the no simultaneous heterogeneity assumption. Epidemiology, 34(3):325–332, 2023.

Thanks for your comment! To ensure the disentanglement of Z, we minimize the KL-divergence between the predicted distribution of Y conditioned on both Rz and T, and the predicted distribution of Y conditioned solely on T, i.e., ${D}_{KL}\left [ p (Y \mid {Z}, {T}) \Vert p (Y \mid T) \right]$. Our initial concept was grounded in the understanding that when using the treatment label T as the feature to predict outcome Y, the information in Z might become redundant. In addition, we set up the representation network of Z and make it align with the causal structure depicted in Figure 1(a) to ensure the generated IVs satisfy three conditions of valid IVs. Identifiability of disentangled representation is challenging, but empirical progress has been made in this direction, for example, [2]. Our method further advances the disentanglement performance. The experimental results in the main paper, Appendix validate the effectiveness of our generated IVs. Thus, following [1], the only assumption in our work is that we can recover the valid IVs from observed features.

[1] Louizos, C.; Shalit, U.; Mooij, J. M.; Sontag, D.; Zemel, R.; and Welling, M. 2017. Causal effect inference with deep latent-variable models. Advances in neural information processing systems, 30.

[2] Wu, A.; Kuang, K.; Xiong, R.; Zhu, M.; Liu, Y.; Li, B.; Liu, F.; Wang, Z.; and Wu, F. 2023. Learning Instrumental Variable from Data Fusion for Treatment Effect Estimation. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 37, 10324–10332.