Paper: DBRNet: Advancing Individual-Level Continuous Treatment Estimation through Disentangled and Balanced Representation (ICLR-2024)

Reviewer comment (Experiments):
It is not clear if authors have addressed my concern about under-training/over-training due to training all algorithms for fixed 800 epochs (and not using something like early stopping).

Author response:
Regarding your concern about the use of 800 epochs for every model, we agree with your point. To address this, we have implemented an early stopping technique across all models to prevent overfitting and underfitting in this ICLR version. This modification is detailed in our code. Furthermore, we have specifically highlighted this change in blue on page 13 in the Implementation Details section of our revised manuscript.

Reviewer comment (Experiments):
Can you confirm if you have addressed the under-training/over-training issue caused by training for fixed 800 epochs?

Author response:
Regarding your concern about the use of 800 epochs for every model, we agree with your point. To address this, we have implemented an early stopping technique across all models to prevent overfitting and underfitting in this ICLR version. This modification is detailed in our code. Furthermore, we have specifically highlighted this change in blue on page 13 in the Implementation Details section of our revised manuscript.

Reviewer comment (Evaluation):
The paper is mainly motivated by learning the **individualized** dose response function. This is a different task than learning the **average** dose response function. Hence, the main metric for evaluation should be the MISE and not the AMSE (even though it is a nice add-on but not necessary). Thus, the MISE should also be the displayed metric e.g., in Fig. 3 and Table 3, and the main interpretation of the results should focus on the MISE.

Author response:
We are grateful for your insightful suggestion regarding the emphasis on MISE. Initially, we had prioritized AMSE as the primary metric due to its popularity. However, after carefully considering your valuable suggestion, we agree that focusing on MISE is indeed more appropriate for our paper. To reflect this, we have thoughtfully reordered the tables to give precedence to MISE over AMSE. Additionally, we have revised the displayed metric in Fig. 3 according to your suggestion. We have also refined the description of our results to highlight the MISE findings more prominently. All these modifications have been highlighted in blue for ease of review.

Reviewer comment (Evaluation):
The sensitivity analysis in Fig. 3 states in the description: “Sensitivity analysis with different values of \(\alpha, \beta\)”, but the axis of the plots are \(\beta\) value and \(\gamma\) value which is confusing. Also, the plots for \(\beta\) and \(\gamma\) do not show clear performance gains outside of the standard deviation bounds. This may indicate that the (quite substantial) performance gains reported in Table 2 are not significant because of the high variance of the method, or may be quite prone to hyperparameter tuning and model selection. Clarification here would be much appreciated.

Author response:
We appreciate your detailed review and have corrected the typo. **Furthermore, we have expanded the test range of hyperparameters and included all hyperparameters in the sensitivity analysis. The new results can be found in the revised paper.** The figure shows that the re-weighting function and the discrepancy loss ($\beta$) have more substantial influence on the model's performance, which consists with the findings from the previous ablation study. The performance improves notably when $β$ is larger and the re-weighting factor is set to 1.

Reviewer comment (Experiments):
The paper argues that existing methods using selection bias adjustment “fail to accurately adjust for selection bias, as they do not make any adjustments to eliminate the bias, or resort to a simple and sophisticated approach to balance the entire representation.” While there is work on binary treatments showing this, it is important to benchmark against such methods in a continuous setting to support this argument, e.g., using the balancing applied in [4] as a baseline.

Author response:
Thank you for your suggestion. We would love to add this baseline; however, the code is not open-sourced, and despite our efforts to contact the author, we have not received a response. Should the code be released, we will certainly do our best to include it as a baseline.

Reviewer comment (Presentation):
Minor: There are a many spacing issues with missing spaces around references.

Author response:
We appreciate your detailed review and **have already corrected the missing space.**

Reviewer comment (Experiments):
I assume the “_TR” in the baseline methods refers to targeted regularization. However, TR is designed for ATE / ADRF estimation. Motivation about why these baselines are included (and still oftentimes outperform methods without TR even in MISE) would be appreciated.

Author response:
As you rightly pointed out, TR represents targeted regularization. 
As you mentioned, TR is designed for ATE / ADRF estimation. Furthermore, to obtain ATE and ADRF, the papers [2-3] (which proposed TR) actually first calculate the individual treatment effect and then average them to derive the ATE and ADRF. We conjecture that thoes approaches may also lead to better individual treatment effect estimation. Otherwise, the papers [2-3] might not achieve improved ATE/ADRF results through averaging if individual treatment effect were inaccurately estimated. Therefore, we believe it would be interesting to compare these with our methods. Moreover, including more baseline methods also makes the results more convincing.

Reviewer comment (Theory):
The discrepancy loss in Eq (5.) leverages KL divergence between representations. However, as I understand, the learned representations are not probabilistic so how is the KL divergence approximated which is defined as a measure between two probability distributions? Also, if there are some implicit assumptions here which are necessary to ensure statistical independence between the representations they should be stated.

Author response:
Dear Reviewer,

Thank you for acknowledging our updates. We would like to further clarify the implementation of the independent loss. In our paper, the intent of the independent loss is to **make the three representations more separate from each other**. We were inspired by the KL divergence and the KL divergence loss in PyTorch to design this loss. After carefully considering your comments and those of Reviewer cz1b, we agree that referring to it as 'KL divergence' is not appropriate. As you mentioned, KL divergence is typically used to calculate the discrepancy between probability distributions, which is indeed different from ours. Therefore, we have decided to rename this loss as 'divergence loss' $L_D$ and include a full definition of this loss in our revised paper to avoid confusion (the modifications are highlighted in blue). However, we kindly argue that **our loss is still effective for separating the representations**.  for example, $L_D(\Gamma(x_i), \Delta(x_i))$ is calculated as follows :

$L_{temp}(\Gamma(x_i), \Delta(x_i)) = \Delta(x_i)\log\left(\frac{\Delta(x_i)}{\Gamma(x_i)}\right) = \Delta(x_i)(\log(\Delta(x_i)) - \log(\Gamma(x_i)))$

$L_D(\Gamma(x_i), \Delta(x_i)) = \frac{1}{n}\sum_{n,m} L_{temp}(\Gamma(x_i), \Delta(x_i)),$

where $n$ and $m$ represent the batch size and the latent factor dimension, respectively.

Using this formula, if the representations $\Gamma(x_i)$ and $\Delta(x_i)$ are identical, the result will be 0. Conversely, if $\Gamma(x_i)$ and $\Delta(x_i)$ are distinctly different, the value will be larger. Therefore, we utilize the reverse of this loss to ensure that the representations are **separate** from each other, thereby guaranteeing that each component learns only the information relevant to its specific functionality. To demonstrate the effectiveness of this loss in our model, we present t-SNE plots of the model trained with and without $L_{disc}$ on a synthetic dataset in Fig.2. The incorporation of the discrepancy loss results in more distinct representations and a higher KL-divergence score, validating our loss approach.

We would like to inquire if there are any points that remain unaddressed in your view, as we believe we have responded to **all nine questions**. Your advice and suggestions are invaluable to us.

Reviewer comment (Theory):
In the proof citing the law of total probability for Equation 9, it seems to assume binary treatment. This appears inconsistent with the claim that the method is capable of handling continuous treatments.

Author response:
Thank you for pointing this out. We realized that the equation and term were confusing, so we have carefully revised them. A new paragraph has been added to Appendix A, which includes our revised equation for continuous settings, highlighted in blue. We have also included additional explanations in the proof section, similarly highlighted. Typically,  $t'$ represents all counterfactual treatments of $x$, defined as $t'= \\{t'| t'\in \mathcal{T}, t'\neq t \\}$. The union of $t$ and $t'$ constitutes the universal set of treatments. Therefore, given the probability density function $f(x,t)$ we have, $ \mathbb{P}(x,t)+\mathbb{P}(x,t')=\int_{t}f(x,t)dt+\int_{\mathcal{T} / \{t\}}f(x,t)dt
=\int_{\mathcal{T}}f(x,t)dt
=\mathbb{P}(x)$

Reviewer comment (Reproducibility):
The process of identifying and estimating the counterfactual distribution $p(x, t')$ from observed data, especially when $t'$ is a continuous variable, is not sufficiently detailed. More explicit steps or examples would greatly benefit readers in understanding this crucial aspect of the proposed methodology.

Author response:
In our paper, we do not estimate the counterfactual distribution $p(x, t')$; instead, our goal is to estimate the IDRF as shown in Equation 1 $\mu(t, x) = \mathbb{E}[Y(T=t)|X=x, T=t]$. This involves estimating the corresponding $y$ for any given $x$ and $t$. We acknowledge, as you rightly pointed out, that addressing this continuous treatment estimation is a crucial aspect of our work. Typically, we adopt the varying coefficient network [2] (as introduced in Section 3.2 in factual loss) to estimate treatment effects under continuous settings. **Thank you for pointing out the confusion of the function of the varying coefficient network. We have added a sentence highlighted in blue in the methods section, to elucidate that.**

In binary settings, as demonstrated in the paper by Hassanpour and Greiner [2019b], the common approach is to learn a deep representation of $x$ and then construct a neural network for each treatment to estimate $\mu(t, x)$. However, this method cannot be extended to continuous settings due to the infinite number of treatments, making it impractical to build a neural network for each value of $t$. Therefore, we utilize the varying coefficient network [2] (introduced in factual loss in section 3.2), which enables us to estimate treatment effects under continuous settings. Specifically, the varying coefficient structure uses a function $g_{\theta(t)}(\Delta(x_i), \Upsilon(x_i))$ with varying parameters $\theta(t)$ rather than fixed parameters to predict outcomes $y$ given the treatment $t$ and the corresponding deep representation of covariates $\Delta(x_i), \Upsilon(x_i)$. In this model, the parameters are determined by $t$, allowing us to incorporate the continuous $t$ in outcome estimation and estimate $\mu(t, x)$ without the need to construct a neural network for each value of $t$. 

Particularly, we employ a B-spline of degree $p$ with $q$ knots, resulting in $k=p+q+1$ basis functions. Let $\textbf{B} =[b_1, b_2, ..., b_k] \in \mathbb{R}^{n\times k}$ denote the spline basis for the treatment $T \in \mathbb{R}^{n\times 1}$. For a single-layer feedforward network with $p$ inputs and $q$ outputs, the function is given by $f_{\theta(t)} = \sum_{i=1}^{k} (b_i \cdot (\textbf{X}\textbf{W}))$, where $\textbf{W} \in \mathbb{R}^{p\times q\times k}$ is the optimizable weight matrix.

Reviewer comment (Theory):
The authors use the independent loss for disentanglement, but this does not necessarily imply full identifiability of the real underlying causal factors. Given the known challenges in achieving full identifiability in nonlinear independent component analysis (up to a component linear transformation), it would be helpful to discuss this limitation and its effect in the paper. A more rigorous proof or argument regarding the identifiability and independence of representations in the model would be beneficial.

Author response:
We respectfully argue that our method differs from nonlinear independent component analysis, and our latent factors are fully identifiable through supervised learning. We acknowledge your point that achieving full identifiability in independent component analysis is challenging, as it is inherently unsupervised. While it can separate the latent factors, identifying the semantic meaning of each latent factor is difficult. In contrast, our method leverages **supervised learning** to ensure that each deep representation learns its corresponding information effectively.

In our paper, we describe three types of deep representations: 
1. Instrumental factors $\Gamma(x)$, which are associated with the treatment but not with the outcome except through the treatment.
2. Confounder factors $\Delta(x)$, which are associated with both the treatment and the outcome.
3. Adjustment factors $\Upsilon(x)$, which are predictive of the outcome but not associated with the treatment.

We use instrumental and confounder factors to predict treatment, and confounder and adjustment factors to predict the outcome. The ground truth of the treatment and the outcome, along with the loss function (Equation 2), are used to optimize the learned representations (supervised learning). This approach ensures that different representations only encode relevant information for their respective roles, thereby achieving identifiability. In addition, we attempt to explore whether the representations capture the corresponding factors by utilizing t-SNE to visualize the three deep representations. The results in section 4.5 show that our method has the ability to identify the corresponding factors.

Reviewer comment (Theory):
The proposed method relies heavily on the causal graph shown in Figure 1(a). However, the authors haven't provided a sufficient explanation for the soundness of this causal graph. The authors should at least provide several specific examples (such as Johansson et al., 2016) about what $\Gamma, \Delta, \Upsilon$ respectively represents.

Author response:
1.**Explanation and soundness of the causal graph**

We agree with you that utilizing a concrete example can help readers better understand the causal graph and enhance its soundness, and thus we **have added a section in Appendix H (An instantiation of causal graph) and the detailed motivation and intuition for the disentanglement in the method section.(both are highlighted in blue)**.

The motivation and intuition behind the disentanglement in our study are rooted in causal inference theory, which typically classifies covariates (input features) into three categories: instrumental variables (predictive only of treatment), confounder variables (influencing both treatment and outcome), and adjustment variables (predictive solely of outcome). Furthermore, disentangling $X$ into these three factors is a commonly used approach in causal inference research, as evidenced in [1, 2, 3, 4, 5], allowing for the precise extraction of relevant information.

We presented a detailed example in a medical scenario, as proposed in [1]: We might collect extensive historical data from each patient, including the patients' features $𝑋$ (e.g., age, gender, living environment, doctor-in-charge), treatment of patients $𝑇$ (taking a particular medicine or not), and the final outcome $𝑌$ (cured or not). Among these features, age and gender simultaneously affect the treatment (as a doctor would consider these factors when choosing a treatment) and the outcome (since they can also impact the patient’s recovery rate), hence they are confounding factors $𝐶$. In contrast, the doctor-in-charge would influence only the treatment decision, without affecting the outcome, thus being an instrumental factor $𝐼$. Environment, which only affects the outcome but not the treatment, falls into the category of adjustment factors $𝐴$.

Reviewer comment (Theory):
In Section 3.2, the authors formulate the loss function of their proposed DBRNet, it is too sophisticated and contains too many hyper-parameters. As the authors claimed, the discrepancy loss encourages $\Gamma, \Delta, \Upsilon$ to be independent of each other, and the independent loss encourages $\Upsilon, T$ to be independent of each other. However, based on the causal graph shown in Figure 1, why not DBRNet contain another term to encourage $\Gamma, Y$ to be independent of each other given $T, \Delta$? Besides, there are some questions about the soundness of the discrepancy loss and independent loss, please see Questions.

Author response:
2.**Clarification on KL divergence**

Thank you for prompting us to clarify our approach to KL divergence. The aim of our causal graph is to demonstrate that covariates can be decomposed into three factors: $Γ(x_i), Δ(x_i), Υ(x_i)$, where $Γ(x_i)$ and $Δ(x_i)$ should be used to estimate $t_i$, while $Δ(x_i)$ and $Υ(x_i)$ should be used to estimate $Y$. Hence, our model structure is already fulfilled this causal graph, and the rest of loss is mainly aim to adjust selection bias.

The disentanglement loss $L_{dis}$ is designed to ensure that these factors are more separated and encode only the relevant information. The purpose of adding an independent loss $L_{ind}$ is to eliminate the selection bias caused by biased treatment assignment. Since we aim to prevent the distribution of the adjustment factor $Υ(x_i)$ from varying across different treatments, so we introduce this term (make the adjustment factor to be independent of $t_i$) to make necessary adjustments. Although many previous papers have proposed similar ideas, most of them focus on binary settings and may therefore adopt different methods [5-7]. This loss function also allows us to encode all information about $t_i$ in $Γ(x_i)$ and $Δ(x_i)$, rather than in $Υ(x_i)$, which aids in making the estimation of $w$ more precise and accurate. Hence, this loss further enhances bias elimination through re-weighting. Secondly, from a practical standpoint, since $Y$ is unknown and is the variable we aim to predict, it is not practical to include a term to measure and minimize the discrepancy loss between two distributions when one of them is unknown. We agree with you using the causal graph structure to justify the $L_{ind}$ may be confusing **and we have already modified that part (in section 3.2 Independent loss, highlighted in blue) to clarify the motivation and aim. Thank you again for your suggestion!**

[1] Wu A, Kuang K, Yuan J, et al. Learning decomposed representation for counterfactual inference[J]. arXiv preprint arXiv:2006.07040, 2020.
[2] Yao, Liuyi, et al. "A survey on causal inference." ACM Transactions on Knowledge Discovery from Data (TKDD) 15.5 (2021): 1-46.

[3] Kuang, Kun, et al. "Treatment effect estimation with data-driven variable decomposition." Proceedings of the AAAI Conference on Artificial Intelligence. Vol. 31. No. 1. 2017.

[4] Liu, Dugang, et al. "Mitigating confounding bias in recommendation via information bottleneck." Proceedings of the 15th ACM Conference on Recommender Systems. 2021.

[5] Hassanpour, Negar, and Russell Greiner. "Learning disentangled representations for counterfactual regression." International Conference on Learning Representations. 2020.


[6] Shalit, Uri, Fredrik D. Johansson, and David Sontag. "Estimating individual treatment effect: generalization bounds and algorithms." International Conference on Machine Learning. PMLR, 2017.

[7] Bellot A, Dhir A, Prando G. Generalization bounds and algorithms for estimating conditional average treatment effect of dosage[J]. arXiv preprint arXiv:2205.14692, 2022.

Reviewer comment (Experiments):
I'm happy that the authors have performed sensitivity analysis w.r.t. $\alpha, \beta$, so why not perform that w.r.t. $\gamma, \lambda$?

Author response:
4.**More results for sensitivity analysis**

We are grateful for your constructive suggestion regarding the addition of more results to enhance the robustness of our experiments. In line with your advice, **we have added an ablation study for $\alpha$ (Table 2) and a sensitivity analysis for $\alpha$ and $\lambda$ based on both MISE and AMSE in our revised paper (Figure 3&6)**. Due to space limitations, we have chosen to only display the figures based on MISE, while the results for AMSE are included in Appendix I.  The results show that despite $\alpha$ also contributes to IDRF estimation, its values have a minimal impact on model performance.  As for $λ$, which controls the regularization, excluding it leads to a decrease in performance. However, when included, the specific value of $λ$ does not significantly affect the results.

Reviewer comment (Theory):
Discrepancy loss minimizes KL divergence between $\Gamma(x_i)$ and $\Delta(x_i)$ and the authors claimed that if $\Gamma(x_i)$ and $\Delta(x_i)$ are independent of each other, the KL divergence is 1. First, considering $x_i$ is a realization of the random variable X (rather than a random variable), what does "$\Gamma(x_i)$ and $\Delta(x_i)$ are independent of each other" mean? Second, in my experience, KL Divergence can only measure the similarity between two distributions but not the independence of two random variables. The similar question also exists for independence loss. If I'm wrong, please provide some references.

Author response:
5.**Meaning of $Γ(x_i)$ and $Δ(x_i)$ and clarification on KL divergence**

In our paper, $X$ represents the entire dataset of covariates, encompassing multiple features, rather than a singular random variable. In the previous given example, within the medical historical data of all patients, $X$ represents all features of all patients. For a specific patient $x_i$, the the patient’s age and gender are the confounder factors $Δ(x_i)$, while the patient’s doctor-in-charge is the instrumental factor $Γ(x_i)$. It is apparent that the doctor-in-charge is independent of the patient's age and gender. We believe that adding this example, as per your suggestion, will elucidate this point more clearly. Moreover, $Γ(x_i)$ and $Δ(x_i)$ generated by our model are not random variables; they are deep representations with a dimension of 50, as generated by three representation networks in our model (Figure 2). We agree with your observation that KL divergence is employed to measure the discrepancy between two distributions of representations. In our work, we specifically use the KL divergence loss, with its detailed implementation available in PyTorch: [https://pytorch.org/docs/stable/generated/torch.nn.KLDivLoss.html](https://pytorch.org/docs/stable/generated/torch.nn.KLDivLoss.html). **Your suggestion is valuable. In response, we have added a section in the Appendix D, highlighted in blue, to clarify KL divergence. Additionally, we have included a sentence in the methods section to direct readers to this new Appendix section.**

Reviewer comment (Novelty):
The motivation of this paper is convincing but lacks innovation because neither the disentangled representation nor the re-weighting technique is originally proposed in this paper. I hope the authors can provide deeper insights into why disentangled covariates should be considered in the context of continuous treatment settings.

Author response:
1.**Motivation of using disentanglement in the context of continuous settings**

We are grateful for your question. The motivation for using disentanglement in continuous settings is to **precisely** adjust selection bias. Existing methods often resort to balancing the entire representation, a simple yet brute approach. However, we argue that not all information in the latent representation should be balanced to adjust for selection bias. For instance, while confounder factors in the representation bring selection bias, they also contribute to outcome predictions. Balancing instrumental factors is theoretically implausible since they relate to treatment assignment and should not be identical across treatments. Therefore, indiscriminate balancing of the entire representation of input covariates is not appropriate. Hence, we apply disentanglement techniques to disentangle the three factors for the further precise bias adjustment. **We appreciate your suggestion and have added more details about the motivation (highlighted in blue) in both the introduction and the methods sections of our paper.**

Reviewer comment (Experiments):
In the experimental setting, it would be beneficial to compare the proposed method with more disentanglement-based baselines, such as DR-CFR[1] and DeR-CFR[2].  
[1] [ICLR'20] Learning Disentangled Representations for CounterFactual Regression  
[2] [ArXiv'20] Learning decomposed representation for counterfactual inference

Author response:
3.**Comparison with other methods**

Thank you for sharing these interesting papers! After a careful review, one of them [1] is extensively discussed in the related work and section 3.4 of our paper. The limitations of [1] actually motivated our approach. Specifically, [1] suggests building regression networks on the top of disentangled embeddings for each treatment value in binary settings. However, in continuous settings, building a network for each treatment value is impractical due to the infinite nature of $T$. Moreover, upon revisiting their code according to your suggestion, we note that [1] requires prior knowledge of the probability of each treatment to compute the re-weighting function, making it unsuitable for continuous treatment scenarios. Similarly, the methods in [2] are limited, as they binarize the treatment. The lack of open-source code for [2] further complicates direct comparison. Overall, while these methods employ disentanglement, they address different questions in different settings. We would be eager to compare our method with theirs if they release their code, especially for continuous settings.