Paper: Self-Distilled Disentanglement for Counterfactual Prediction (ICLR-2024)

Reviewer comment (Theory):
Incorrect Definiton about $I(A; B \mid C)$, which should not denote conditional mutual information. If $I(Z ; Y\mid T)$ represents conditional mutual information, then in Eqs. (2,3,4), it is evident that the conditional mutual information $I(Z ; Y \mid T) ≠ 0$, as the open of the collider structure $Z → T ← \{C, U\}$ will make $Z$ dependent on $\{C, U\}$ when $T$ is fixed as a condition, which consequently results in the dependence of Z and $Y$ through $\{C, U\}$.

Author 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!

Reviewer comment (Theory):
The authors use the conditional mutual information in the chain rule of mutual information again. The authors should clarify this point and differentiate the definition $I(Z ; Y \mid T)$ in Eqs. (2,3,4) from the definition of conditional mutual information. The relevant content may need to be restated.

Author 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!

Reviewer comment (Experiments):
Is it necessary to use a shallow network for $Q^z_T$ and $Q^c_T$? Why not use a network of the same size as the reference network?

Author 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.

We intentionally opted for a relatively shallow neural framework with 2 or 3 layers in each network (deep, retain, shallow). This choice prioritized computational efficiency and addressed concerns related to overfitting, which is crucial for tasks involving large datasets or limited computational resources.

Reviewer comment (Theory):
The optimization directions of the losses $L\left(Q_T^z, T\right) + L\left(Q_T^c, T\right)$ and $L\left(Q_T^c, Q_T^z\right)$ may be different or conflicting because the former aims to maximize the predictive abilities of $z$ and $c$ on $T$ (the predictive abilities of $c$ and $z$ on $T$ are different), while the latter actually implies forcing a better-performing model to reduce its predictive abilities. This can be achieved by modifying either the predictive network or the representation network. In essence, it is a non-zero-sum game problem. Only when all three are equal to 0 does it mean that $D[\mathcal{P}_T^{R_z} \| \mathcal{P}_T^{R_c}]=0$. Otherwise, minimizing the loss of self-distilled disentanglement does not necessarily mean minimizing $D[\mathcal{P}_T^{R_z} \| \mathcal{P}_T^{R_c}]$. I am not sure if I have missed any important parts, but it clearly requires further clarification.

Author response:
- **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.

Reviewer comment (Novelty):
The Teacher network only aims at enhancing information prediction abilities and does not seem to be the focus of this paper? I will adjust my scores based on the author's response.

Author response:
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 **S**elf-**D**istilled **D**isentanglement, i.e., $SD^2$.  Therefore, the teacher model is utilized to enhance the network's prediction capability, thereby facilitating the distillation and disentanglement process.

Our distillation network is designed to improve the predictive capabilities of the shallow network. Let's delve into the distillation framework aimed at minimizing the loss $L_c^z$​. The teacher network consists of both the retain network and the deep network, utilizing Z and C as joint input features. In contrast, the student network, known as the shallow network, is only half the size of the teacher network and takes either Z or C as its input feature. Although the teacher network may not be large in terms of absolute parameter count, its prediction heads contain a greater amount of predictive information. This information is substantial enough to guide the training of the student network, thereby enhancing its predictive capabilities.

Reviewer comment (Presentation):
Some captions are not descriptive enough of the contents of the figures. E.g., Fig. 1(b) and (c).

Author 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.

Reviewer comment (Presentation):
Use of inline equations should be avoided if possible.

Author 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.

Reviewer comment (Presentation):
It’s best to state each finding in a separate bullet-point.

Author 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.

Reviewer comment (Presentation):
When referring to the appendix, it’s best to also indicate its section number, to make it easier to find.

Author 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.

Reviewer comment (Writing):
I’m not quite sure about the method’s name, specifically, what is being “distilled” here? Is this referring to dissected factors from X?

Author 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 **S**elf-**D**istilled **D**isentanglement, i.e., $SD^2$.

Reviewer comment (Evaluation):
What is being measured in the radar charts in Figure 3(b)? The caption states it shows “the contribution of actual and other variables to the decomposed representations”; how is this measured?

Author response:
- **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.

[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.

Reviewer comment (Evaluation):
It is stated in the paper that “the IV-based methods perform worse than the non-IV-based ones under the continuous scenario” but no reference or discussion is included. Please elaborate why.

Author response:
- **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.

Reviewer comment (Theory):
Eq 3: both parts confuse me. On the left: it seems like the constraint I(Z; Y | T) = 0 contradicts the statement from the bottom of p3 about this inducing dependence through a collider. On the right: I'm not sure where this inequality constraint comes from, or why it's necessary.

Author 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!

Reviewer comment (Writing):
Corollary 4.3: I don't understand this notation (10a-c, 11a-c) — does this mean minimizing all 3? Minimizing the smallest? Again, I think clearer notation could be used around statements around optimization.

Author response:
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.

Reviewer comment (Writing):
Sec 4.2: I don't follow a lot of these architectural choices. I think the "retain network" and "teacher networks" could have their roles explained more, as well as the relationship between the deep and shallow networks.

Author response:
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).

Reviewer comment (Evaluation):
What is the metric in 3b representing "contribution of actual variables?"

Author response:
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.

Reviewer comment (Writing):
Top of p5: you say A is independent of C and Z — do you mean here that C is independent of A and Z? That seems to align more closely with the topic of the paragraph.

Author response:
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.

Reviewer comment (Writing):
I find the notation mk, mz, etc. to be confusing — these look like products to me rather than single variables. Maybe consider using m_k etc.

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

Reviewer comment (Theory):
Is there an implicit assumption in this method that T, Y are causally downstream of X? What other assumptions are required for this method to work well?

Author 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.

Reviewer comment (Theory):
An explicit counterexample shows that the mutual information is not zero in general:  
\( U \sim Bern(0.5 )\)  
\( Z \sim Bern(0.5 )\)  
\(T = XOR(Z, U)\)  
\(Y = XOR(T,U)\)  
where \(XOR\) is the exclusive or function that evaluates to 1 if either argument is 1 but not both. Here, \(H(Y | T)\) is 1, since all the randomness in \(Y\) comes from \(U\) which has entropy 1 by construction. More importantly, if you work through all the possible outcomes of the above system, you will see that, conditional on \(T\), \(Z\) is always equal to \(Y\) (i.e. \(H(Y | T, Z) = 0\)), but that they are conditionally independent when you condition on both \(T\) and \(U\). For example, if \(T = 1\), then when \(z=0\) we know \(U=1\) (since otherwise \(T\) would not be 1), and hence \(Y = XOR(1, 1) = 0 = z\).

Author response:
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.

Reviewer comment (Theory):
This counterexample highlights that the claim \(I(Z;Y | T) = 0\) is not true in general.

Author 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!

Reviewer comment (Theory):
The authors make no assumptions on the functional form of the structural equation that determines \(Y\), and from the LATE framework and Pearl's work on bounding treatment effects, it is known that it is impossible to non-parametrically identify the treatment effect with access to an instrument. Therefore, there are surely more necessary assumptions for the method to work in general.

Author 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.

Reviewer comment (Theory):
Question — Can you give sufficient conditions under which this method identifies the treatment effect?

Author 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.