Counterfactual Generative Models for Time-Varying Treatments

Thanks much for your valuable and insightful comments. We express our gratitude to the reviewer for your positive comments regarding the importance of addressing high-dimension outcomes, the use of generative models, and the positive experimental results. We also thank the reviewer for emphasizing the potential extension of our framework, the comparison to current literature, and the meaning of the treatment variables. We have updated the manuscript according to your comments. Please see our point-by-point responses below.

How do readers use this model to solve causal issues, such as ITE estimation?

We appreciate the reviewer's reference to Individual Treatment Effects (ITE). We would like to clarify that our approach to estimating the counterfactual distribution across a population, rather than individual-focused estimates (ITE), is driven by the practicality of our research's applications, such as in public health policy. Implementing policies on an individual basis can be costly and sometimes unfeasible. By focusing on the population level (or a sub-population), we capture individual differences, which is essential for effective policy-making. In contrast, the variation within individual counterfactuals often stems from observational noise and might need substantial data for precise estimation. Consider, for instance, the COVID-19 data: our objective is to analyze the counterfactual outcome of a state-wide policy (mask mandate) and provide policymakers with a range of possible counterfactual outcomes that reflect variability in county-level confirmed cases. In such contexts, policymakers might be less concerned with the estimation noise at the individual level and more focused on the broader variability across different individuals.

In the meantime, a related extension of our work might be to infer conditional counterfactual outcomes (related to the conditional average treatment effect, CATE). This corresponds to looking at the distribution of the outcome under a specific subpopulation. Since our covariates are assumed to be time-varying, a common approach is to introduce a set of static baseline covariates, $V$ (Robins et al, 1999). The baseline covariate $V$ denotes the static feature (such as a patient's gender or age) that will influence both the time-varying covariates $X$ and the outcome $Y$. We can then draw counterfactual samples from a specific sub-population by conditioning on the values of the $V$. In this framework, we observe ($Y_t^i$, $\overline{A}_t^i$, $\overline{X}_t^i$, $V^i$), where $V \in \mathbb{R}^{ \nu}$ is a static baseline variable that varies by individual. Accordingly, the generator will have an additional input: $g\_{\theta}(z, \overline{a},v): \mathbb{R}^r \times \mathcal{A}^d \times \mathbb{R}^{\nu} \rightarrow \mathcal{Y}$ , and the IPTW weights will become $w\_\phi(\overline{a},\overline{x},v) = \frac{1}{\prod\_{\tau=t-d+1}^{t} f\_\phi(a\_{\tau}|\overline{a}\_{\tau-1},\overline{x}\_{\tau},v)}$. Correspondingly, the objective in Proposition 1 will become $\mathbb{E}\_{v} ~\mathbb{E}\_{\overline{a}} ~\left[\mathbb{E}\_{y \sim f\_{\overline{a},v}} ~\log f\_\theta(y|\overline{a},v)\right] \approx \frac{1}{N}\sum\_{(y,\overline{a},\overline{x},v) \in \mathcal{D}}w\_\phi(\overline{a},\overline{x},v) \log f\_\theta(y|\overline{a},v).$

We have conducted additional experiments using the fully synthetic dataset, where the baseline variable, $V$, was divided into two groups. The $V$ was uniformly drawn from $[-1,0]$ in the first group and from $[0,1]$ in the second group. We then looked at the performance of the generative samples corresponding to two sub-groups. We included the results in Appendix I (Table 4, Fig.11, and Fig.12) of the revised manuscript. Our method (MSCVAE) still consistently outperforms other baseline methods in the simplified results below (Wasserstein distance, smaller the better):

• $V\in[-1,0]$

• $V\in[0,1]$

Thanks much for your valuable and insightful comments. We express our gratitude to the reviewer for your positive comments regarding the technical soundness of our method. We also thank the reviewer for the careful read and the mention of the confusion from the notations. We have updated the manuscript according to your comments. Please see our point-by-point responses below.

The authors may not give detailed explanations about the causal assumptions in the main paper, but at least mention the names of the causal assumptions in the paper.

Thanks for the suggestion. In the revised manuscript, we mentioned the three causal assumptions in the methodology section.

In Algorithm 1, the authors suggest that we should draw a sample epsilon from $N(0,I), where$I$ is the total number of individuals. What is the point of setting epsilon as a realization with large variance?

We would like to clarify that $I$ should stand for the identity matrix, and $N$ represents the total number of samples. We have fixed it in the revised paper.

Why do you model epsilon as normally distributed?

The Gaussian assumption for the VAE is a common assumption which enables a closed form expression for the ELBO (Kingma and Welling, 2013). To elaborate, our generative framework essentially involves identifying a transformation that converts random noise into the desired target distribution, which in this case is the counterfactual distribution. It's a standard practice in the field of generative models to use the normal distribution as the basis for our source distribution (Kingma and Welling, 2013; Sohn et al, 2015; Higgins et al, 2016).

What should be the objective value of training?

We would like to clarify that, the objective function in Proposition 1 is evaluated over the entire $N$ data samples, where $N$ represents the total number of samples across individuals and time (see below). Similarly, in the algorithm box, where batch optimization can be used, the objective is then evaluated over the mini-batch.

In the revised manuscript, we have clarified how data were indexed. Consistent with much of the literature on longitudinal causal inference, such as (Lim et al, 2018; Bica et al, 2020; Li et al, 2021), our approach assumes that data is observed across individuals, with each individual having data for $T$ time points. Importantly, given its history $\overline{X}$ and $\overline{A}$, the outcome, $Y$, can be treated as conditionally independent samples. Therefore, for simplicity, we could 'chop' the time series of each individual, leading to a total of $N$ data tuples across individuals and time. For example, if there are $1000$ individuals each with $100$ time points and $d=3$, then $N = 1000\times (100-3+1)=98,000$ (we truncate the first $d-1$ time points because they do not have a full history). Therefore, we denote the data tuples simply as ($y^i,\overline{x}^i,\overline{a}^i$), where $i$ denotes the index of the sample and ranges from $1$ to $N$.

It is strange that t=d, …, T in Algorithm 1. Is it a typo mistake?

Following the previous point, we ranged $t$ from $d$ instead of $1$ because the first $d-1$ time points do not have full-length ($=d$) history.

How to determine the value of d?

As a standard in longitudinal causal inference (Robins et al, 1994; Lim et al, 2018; Bica et al, 2020; Li et al, 2021; Bica et al 2021), $d$ is pre-determined as prior knowledge of the data generating process. This approach is validated by our numerical findings, proving to be an effective method. For instance, in our COVID-19 studies, we set $d=3$ following insights from epidemiological research, which indicates that a mask mandate typically requires two to three weeks to become effective in a population.

Intuitively, $d$ can be considered as a causal structure that determines how many arrows should point to the current variable of interest from history. Therefore, if one is interested in estimating $d$ instead of pre-determining it, a potential method is causal discovery which automatically infer the DAG (Figure 8) from data.

We thank the reviewer for the comments. We appreciate your input and regret any confusion that may have arisen if our methodology wasn’t adequately elucidated. We would like to clarify that our generative framework aims to draw high-quality counterfactual outcomes, which can be viewed as an implicit distribution estimator, bypassing the need for direct density estimation. We want to emphasize that this objective differs drastically from the majority of the existing literature, which only focus on estimating the mean of the counterfactual distribution. Please see our point-by-point responses below.

How could $\overline{a}$, a fixed value, be drawn from $\mathcal{D}$? Wouldn’t the RHS of (4) be the same for all $\overline{a}$ while the LHS is supposed to be different?

We thank the reviewer for pointing it out. Our generator approximates the underlying counterfactual for any $\overline{a}$,, so there should be an expectation over $\overline{a}$, for the objective function (Equation 2): $\hat{\theta} = argmin_{\theta \in \Theta}~\mathbb{E}_{\overline{a}} ~\left[D_f(f\_{\overline{a}}, f\_\theta(\cdot|\overline{a})) \right],$

where $D$ represents the distributional distance between $f\_{\overline{a}}$ (true counterfactual distribution) and $f\_\theta$ (the proxy conditional distribution represented by our proposed counterfactual generator $g\_\theta$), such as KL divergence. The resulting weighted loss function in Proposition 1 becomes:

$$\mathbb{E}\_{\overline{a}} [\mathbb{E}\_{y \sim f\_{\overline{a}}} ~\log f\_\theta(y|\overline{a})] \approx \frac{1}{N}\sum\_{(y,\overline{a},\overline{x}) \in \mathcal{D}}w\_\phi(\overline{a},\overline{x}) \log f\_\theta(y|\overline{a})$$

Why is it a sum instead of an average?

We note that, in Proposition 1, we used the sum for simplicity, which is numerically equivalent to the average in practice as we fixed the batch size in our training.

And why is the index in (5) from $t-d$ instead of from $t-d-1$?

The index in (5) should be from $t-d+1$ instead of $t-d$, which is a typo. We have fixed in the revised paper.

Where does the expectation over $\overline{a}$ come from?

The$\overline{a}$in the proof for Proposition 1 is due to that we were taking an expectation over $\overline{a}$ (as noted above). We have modified our notations in the revised manuscript.

I don’t seem to understand the comment in Remark 1 ...

We would like to clarify that we did not suggest that the DR method lacks robustness under model misspecification than IPTW methods. In Remark 1, we intended to convey that it was challenging to extend our IPTW framework to a DR method. This is because in doing so, we need to augment the IPTW estimator with an outcome model for approximating the counterfactual distribution. Such an extension would require accurately approximating the conditional outcome distribution $f(Y_t|\overline{X}_t,\overline{A}_t)$ and the conditional covariate distribution $f(X_t|\overline{X}\_{t-1},\overline{A}\_{t-1})$. This becomes particularly challenging when $Y$ is high-dimensional, a difficulty that is further exacerbated in the time-varying case with the need to accurately approximate $f(X_t|\overline{X}\_{t-1},\overline{A}\_{t-1})$. The G-Net is the only outcome-based method known to us that addresses such problem, and we have therefore included it as a baseline in our study. Our analysis revealed that G-Net's performance is suboptimal due to the aforementioned challenges, leading us to opt for the IPTW-only framework. We have clarified these points in the Appendix E of the revised paper.

Furthermore, we would like to emphasize our main contribution is to use a generative model to draw high-quality counterfactual samples that follow the underlying counterfactual distribution. We opted for IPTW due to its seamless integration with our generative framework (the weighted loss in Proposition 1) and suitability for the high-d time-varying density approximation (it's easier to estimate the binary propensity model).

One pivotal assumption is that d, the length of history dependence, is finite and known...

In line with established practices in longitudinal causal inference, as referenced in works like (Robins et al, 1994; Lim et al, 2018; Bica et al, 2020; Li et al, 2021; Bica et al 2021), we selected $d$ based on our pre-existing understanding of the data generation process. This approach is validated by our numerical findings, proving to be an effective method. For instance, in our COVID-19 studies, we set $d=3$ following insights from epidemiological research, which indicates that a mask mandate typically requires two to three weeks to become effective in a population.

(Continued from the last response)

In the meantime, a related extension of our work might be to infer conditional counterfactual outcomes (related to the conditional average treatment effect, CATE). This corresponds to looking at the distribution of the outcome under a specific subpopulation. Since our covariates are assumed to be time-varying, a common approach is to introduce a set of static baseline covariates, $V$ (Robins et al, 1999). The baseline covariate $V$ denotes the static feature (such as a patient's gender or age) that will influence both the time-varying covariates $X$ and the outcome $Y$. We can then draw counterfactual samples from a specific sub-population by conditioning on the values of the $V$. In this framework, we observe ($Y_t^i$, $\overline{A}_t^i$, $\overline{X}_t^i$, $V^i$), where $V \in \mathbb{R}^{ \nu}$ is a static baseline variable that varies by individual. Accordingly, the generator will have an additional input: $g\_{\theta}(z, \overline{a},v): \mathbb{R}^r \times \mathcal{A}^d \times \mathbb{R}^{\nu} \rightarrow \mathcal{Y}$ , and the IPTW weights will become $w\_\phi(\overline{a},\overline{x},v) = \frac{1}{\prod\_{\tau=t-d+1}^{t} f\_\phi(a\_{\tau}|\overline{a}\_{\tau-1},\overline{x}\_{\tau},v)}$. Correspondingly, the objective in Proposition 1 will become $\mathbb{E}\_{v} ~\mathbb{E}\_{\overline{a}} ~\left[\mathbb{E}\_{y \sim f\_{\overline{a},v}} ~\log f\_\theta(y|\overline{a},v)\right] \approx \frac{1}{N}\sum\_{(y,\overline{a},\overline{x},v) \in \mathcal{D}}w\_\phi(\overline{a},\overline{x},v) \log f\_\theta(y|\overline{a},v).$

We have conducted additional experiments using the fully synthetic dataset, where the baseline variable, $V$, was divided into two groups. The $V$ was uniformly drawn from $[-1,0]$ in the first group and from $[0,1]$ in the second group. We then looked at the performance of the generative samples corresponding to two sub-groups. We included the results in Appendix I (Table 4, Fig.11, and Fig.12) of the revised manuscript. Our method (MSCVAE) still consistently outperforms other baseline methods in the simplified results below (Wasserstein distance, smaller the better):

• $V\in[-1,0]$

• $V\in[0,1]$

Can you provide a proof for equation 2?

We have included a proof for Eq 2 (Eq 3 in the revised manuscript) in Appendix C (pasted below). The objective function for training the counterfactual generator minimizes the difference between $f\_\theta(\cdot|\overline{a})$ and the true counterfactual distribution $f\_{\overline{a}}$ with respect to a distributional difference $D_f(\cdot,\cdot)$ over all treatment combinations, $\overline{a}$. When the distance measure is the KL-divergence, this equals maximizing the log-likelihood of the conditional distribution when data is sampled from the counterfactual distribution: $\hat{\theta} = argmin\_{\theta \in \Theta} \mathbb{E}\_{\overline{a}} [{\rm KL}(f\_{\overline{a}}(\cdot) || f\_\theta(\cdot|\overline{a})) ]$ $= argmin\_{\theta \in \Theta} \mathbb{E}\_{\overline{a}} [\int \log \left(\frac{f\_{\overline{a}}(y)}{f\_\theta(y|\overline{a})}\right)f\_{\overline{a}}(y) dy]$ $= argmax\_{\theta \in \Theta} \mathbb{E}\_{\overline{a}} [ \int \log \left(f\_\theta(y|\overline{a})\right)f\_{\overline{a}}(y) dy ]$ $= argmax\_{\theta \in \Theta} \mathbb{E}\_{\overline{a}} [\mathbb{E}_{y \sim f\_{\overline{a}}} ~\log f\_\theta(\cdot|\overline{a}) ]$

What would be the difference in objective if we were to define the distance in terms of Maximum Mean Discrepancy (MMD) or Wasserstein distance?

We specifically study KL-divergence because it is well formulated for majority of generative models, including CVAE and many others, and has been focused in related work on counterfactual density estimation (Kennedy 2021, Melnychuk, 2023). Replacing KL with (kernel) MMD or Wasserstein-1 distance may lead to technical difficulties, due to the min-max formulation in MMD and Wasserstein distance. In particular, we have the following unified objective function for kernel MMD and Wasserstein-1 distance:

\hat{\theta} = argmin\_{\theta \in \Theta}~ \mathbb{E}\_{\overline{a}} [\sup\_{\phi \in \Phi} | \mathbb{E}\_{y \sim f\_\theta(\cdot | \overline{a})}[\phi(y)] -\mathbb{E}\_{y \sim f\_{\overline{a}}(\cdot)}[\phi(y)] \| \].

Here $\Phi$ is a class of functions. For kernel MMD, $\Phi$ is class of kernel mappings and for Wasserstein-1 distance, $\Phi$ consists of all $1$-Lipschitz functions. As can be seen, the objective function now involves an inner maximization operation, which often leads to numerical instability. This phenomenon is well recognized in generative adversarial networks' training (Mescheder et al., 2018; Goodfellow et al., 2020). Therefore, we would like to emphasize the computational efficiency brought by Proposition 1 for the KL case. This proposition allows us to bypass the direct estimation of $f_{\overline{a}}$, significantly enhancing computational scalability in high-dimensional cases.

Considering your framework, it appears straightforward to expand it to the individual outcome level. What led to the decision to overlook that possibility?

We appreciate the reviewer's reference to Individual Treatment Effects (ITE). We would like to clarify that our approach to estimating the counterfactual distribution across a population, rather than individual-focused estimates (ITE), is driven by the practicality of our research's applications, such as in public health policy. Implementing policies on an individual basis can be costly and sometimes unfeasible. By focusing on the population level (or a sub-population), we capture individual differences, which is essential for effective policy-making. In contrast, the variation within individual counterfactuals often stems from observational noise and might need substantial data for precise estimation. Consider, for instance, the COVID-19 data: our objective is to analyze the counterfactual outcome of a state-wide policy (mask mandate) and provide policymakers with a range of possible counterfactual outcomes that reflect variability in county-level confirmed cases. In such contexts, policymakers might be less concerned with the estimation noise at the individual level and more focused on the broader variability across different individuals.

Thanks much for your valuable and insightful comments. Our sincere thanks to the reviewer for recognizing the underlying motivation and flexibility of our approach. Please see our point-by-point responses below.

IPTW values can be notably small and, as highlighted by the author, require precise definition.

Thank you for your comments regarding our use IPTW. We agree that IPTW can be unstable and might lead to less reliable results in certain contexts. Nevertheless, it's important to highlight that integrating IPTW with our novel generative framework represents a significant and novel methodological advancement. This combination is particularly effective in capturing complex, high-dimensional patterns in counterfactual outcomes, an area our study uniquely addresses to the best of our knowledge.

To address the stability concerns in a practical context, our experiments incorporated established stabilization strategies, such as quantile truncation and standardization, as extensively discussed in the literature (Xiao et al, 2010; Chesnaye et al, 2022). The results from our experiments also show that incorporating these techniques into the design of our marginal structure model greatly results in promising and dependable numerical outcomes.

Given we're handling a treatment sequence, it's important to note that the counterfactual treatment is not unique.

We would like to clarify that our generator approximates the counterfactual distribution for any $\overline{a}$, so there should be an expectation over $\overline{a}$ for the objective function: $\hat{\theta} = argmin_{\theta \in \Theta}~\mathbb{E}_{\overline{a}} ~\left[D_f(f\_{\overline{a}}, f\_\theta(\cdot|\overline{a})) \right],$

where $D$ represents the distributional distance between $f\_{\overline{a}}$ (true counterfactual distribution) and $f\_\theta$ (the proxy conditional distribution represented by our proposed counterfactual generator $g\_\theta$), such as KL divergence. We genuinely appreciate the reviewer's careful read of these details, and have updated the manuscript (Equation 2, 3 and Proposition 1) to address this issue.

I believe the following two papers could also serve as baseline references ...

We thank the reviewer for mentioning these baseline methodologies. However, it's important to note that these methods are tailored for mean prediction, which differs from our objective of distribution estimation. Additionally, these methods primarily focus on predicting outcomes at the individual level, whereas our research is geared towards assessing the outcomes at the population level. These fundamental differences might make it challenging for a fair comparison between their methods and ours. In this regard, we could incorporate these methods to compare with our method in terms of mean estimate, should the reviewer consider it necessary.

If x isn't utilized in the generator, what's the rationale for calculating weights based on x? Why not solely estimate treatment probability based on the treatment sequence?

We would like to clarify that, estimating treatment probability based solely on the treatment sequence will introduce bias. This is because the propensity score, $f(A_t|\overline{A}\_{t-1},\overline{X}\_{t})$, depends the history of both treatments and covariates. Therefore, IPTW can vary for the same treatment sequence, $\overline{a}$, owing to the variability in $X$. Importantly, the covariates $X$ play a crucial role in the standard approaches for IPTW estimation (Robins et al, 1999; Lim et al, 2018). As a result, while we do not have to directly use $X$ during training our conditional generative model, we still access them through the IPTW weights in Proposition 1.

Dear reviewers,

We hope that our responses have addressed your concerns. Please let us know if you have any additional questions or comments. We would be more than happy to follow up with additional details. Thanks much for taking the time to review our work.