Paper: Bayesian Neural Controlled Differential Equations for Treatment Effect Estimation (ICLR-2024)

Reviewer comment (Novelty):
With respect to novelty, it appears that the main formulations of Stochastic Differential Equations (SDE) and Controlled Differential Equations (CDE) are derived from prior works, and the variational inference methods employed are consistent with those utilized in existing latent SDE literature [1]. This gives an impression that the current work may be an amalgamation of these previous methodologies, adapted to extend the CDE framework.

Author response:
**Response to (1a) neural SDEs and neural CDEs:**

Thank you for giving us the opportunity to clarify the novelty of our method and how we add to the existing literature. In particular, our work makes non-trivial contributions along **three literature streams**. 


1. We contribute to the literature on neural differential equations. Neural CDE [1] and neural SDE [2] have been proposed previously in isolation, while ours is the **first** method to combine them into a single method. Further, we are not aware of any prior work that has coupled neural SDEs in an encoder-decoder framework. More importantly, to estimate treatment effects, a simple ‘plug-n-play’ approach combining both is typically **not** sufficient. In contrast, our method extends them in a non-trivial way with additional techniques from the causal machine learning literature (e.g., informative sampling in **Supplement L**, balanced representations in **new Supplement M**). Moreover, there does not exist a single tailored Bayesian neural CDE method.
2. We contribute to the literature on treatment effect estimation. Existing methods for treatment effect estimation in continuous time (TE-CDE [3]) are limited to point estimates. So far, no method for uncertainty quantification in this setting existed. Ours is the **first** method that allows for rigorous, Bayesian uncertainty quantification. 
3. We contribute to the medical literature. We show that the MC dropout [4] (as used in our baselines [3]) gives _incorrect_ uncertainty estimates, which may lead to _harmful_ decisions in medical practice. As a remedy, we offer a new method to support reliable medical decision-making over time by providing **reliable** uncertainty estimates. As such, we study an important and underexplored problem in medicine.  

**Action:** We carefully checked that we clearly spelled out the novelty of our work. In particular, we revised **Section 1** and emphasized our contribution to the three different literature streams.

Reviewer comment (Novelty):
There are noticeable similarities between this work and [2]. Despite the authors’ footnote indicating a substantial divergence from [2], there is a potential for their framework to be adaptable for time series treatment and to incorporate CDE elements. A more comprehensive discussion of the distinctions between the current work and [2] could enhance the clarity on this matter.

Author response:
**Response to (1b) distinction from CF-ODE:**

Thank you for giving us the opportunity to highlight the differences between our BNCDE and the CF-ODE [5]. The are **six important differences**: 



1. CF-ODE is inherently designed to forecast the treatment effect over time of a **single**, **static** treatment. In contrast, our BNCDE is designed to estimate treatment effects for a **sequence of multiple** treatments in the future. This is a much more complex task.
2. Accordingly, CF-ODE builds upon a **static identifiability** framework. Here, adjusting for the measurements of the patient history $S_{t^*}(X)$ at the time of treatment, $t^*$ is assumed to be sufficient to adjust for all confounders (Assumption 2.3).
3. CF-ODE directly models the latent states in the decoder with a neural SDE. We do not see how this can be effectively combined with neural CDEs to model arbitrary, future sequences of multiple treatments.
4. All patient information in CF-ODE is contained **deterministically** in the initial state of the ODE. Hence, all **uncertainty** needs to be incorporated in a **deterministic**, initial state. We argue that this is very **inefficient**. Instead, our BNCDE captures patient information in the random variable $Z_T$. Here, all the Monte Carlo variance is propagated through the encoded trajectory $Z_{[0,T]}$, accounting for uncertainty at each point in time $t \in [0,T]$. Therefore, model uncertainty is directly incorporated into the patient encoding.
5. The use of a GRU in the CF-ODE encoder implies **regular sampling** of the patient trajectories. This assumption is a **crucial violation of medical reality**, where it is standard that measurements are taken at **irregular** points in time. Therefore, CF-ODE does **not** account for uncertainty due to irregular sampling. 
6. The single treatment decision in CF-ODE is limited to **binary treatments**. Instead, our BNCDE can deal with any treatment (in our experiments, treatments are discrete).

Because of the above, CF-ODE is **not** applicable to our setting.

**Action:** We added a new **Supplement F** where we spell out the differences between CF-ODE and our method. In sum, CF-ODE is **not** applicable to our setting. We also refer to our new **Supplement E,** where we further highlight the limitations of neural ODE modeling.

Reviewer comment (Experiments):
In terms of scalability, the neural network weights are modeled through an SDE with a neural network drift function. This configuration implies that the input dimensionality of the SDE drift scales with the dimensionality of the weights, which could potentially escalate to millions. This raises pertinent questions regarding the computational cost and the stability of the training process under such high-dimensional circumstances.

Author response:
**Response to (2) scalability:**

This is an excellent question. Increasing the depth of the neural CDE networks directly affects the dimension of the latent neural SDEs. Accordingly, the differential equation solver must solve a higher dimensional SDE. Dimensionality is, however, not a real problem for our method for several reasons: 



* (i) In **Supplement I**, we **report the average training time** of our method. Overall, the average time of our BNCDE was **comparable** to the  TE-CDE [3] baseline (approx. 34h vs 22h).
* (ii) We provide new experiments in **Supplement N** which show that the training time increases only **linearly** with the depth of the neural CDE networks (i.e., the dimension of the neural SDE). Linear scaling is not different from any other, fully connected neural network architecture.
* (iii) Furthermore, using multiple Monte Carlo trajectories at the same time reduces Monte Carlo variance during the training and, thus, stabilizes training. In the new **Supplement N,** we show that increasing the number of Monte Carlo samples comes at very little extra cost. 
* (iv) We did not encounter any instability issues for our method. Our results are very consistent over different seeds (see our experiments in Section 5). 
* (v) Importantly, discretization errors in SDE solvers do not increase with the dimension, but only depend on the order of the solver and the chosen time grid. Modern solvers can adaptively choose the grid size and provide error guarantees.

**Action:** We added new experiments in **Supplement N**, where we show that the training time only increases linearly with the depth of the neural CDEs, i.e., dimension of the neural SDEs. Further, we demonstrate that using multiple Monte Carlo samples at the same time comes at very little extra cost.

Reviewer comment (Theory):
Since the potential outcome is modelled through the latent embeddings, does this affect the identifiability?

Author response:
**Response to (2) latent embeddings**

This is a great question. We use latent embeddings to capture non-linear dependencies in the patient data and reduce dimensionality. This is completely in line with previous works on treatment effect estimation over time [7,8,4,9] and even with works in the static setting [10,11,12].

However, when encoding into latent representations, information on confounders in the original data may be lost and, indeed, affect identifiability. However, as we train our model to be predictive of the potential outcomes, we will not lose relevant information on confounders if the dimension of our latent representation is not too small. 

**Action:** We elaborate on the use of latent representations and the implications for representation-induced bias in **Section 4.1**. We point to an analysis between representation learning and representation-induced bias as future research, which, even in the static setting, has not yet been explored.

Reviewer comment (Evaluation):
The prediction also estimates the outcome uncertainty. But typically, this can potentially lead to a naive model with large uncertainty. I wonder about the quantification quality for the aleatoric uncertainty.

Author response:
**Response to (3) quality of aleatoric uncertainty**

Thank you. We are happy to demonstrate the **quantification quality of aleatoric uncertainty** in our BNCDE (see new **Supplement J**). To this end, we performed **new experiments** where we plot estimation errors against estimates of aleatoric uncertainty, finding that our method can adequately capture uncertainty in the data-generating process.  

**Action:** We add a new experiment that demonstrates the quantification quality of aleatoric uncertainty (see new  **Supplement J**). Thereby, we show that our BNCDE provides meaningful estimates of aleatoric uncertainty.

Reviewer comment (Theory):
Considering medical applications in mind, it is important to mention: 1) possible downsides of the proposed approach; and 2) which assumptions need to be satisfied for the method to be robust.

Author response:
Thank you. We added a **new discussion** around possible downsides and the assumptions that need to be satisfied to ensure robustness (see **new Supplement A**). Therein, we especially focus on potential limitations (e.g., explainability, scalability, etc.) that are relevant for medical practice. 

**Action:** We added a new discussion on the applicability in medical practice (see **new Supplement A**).

Reviewer comment (Experiments):
The method is only compared against other neural methods. I understand those are the most natural competitors, but this ignores a large body of literature. Are there other methods which could deal with the same setup? Or not at all?

Author response:
Thank you. We are more than happy to explain why non-neural methods are **not** applicable (see Table 1). 

* To the best of our knowledge, the **only** applicable baseline (both neural and non-neural) is TE-CDE [1]. TE-CDE builds upon the same setting (i.e., treatment effect estimation in continuous time) as our work. However, it does **not** capture aleatoric and epistemic uncertainty. Hence, we extended TE-CDE with MC dropout [5] to obtain a baseline that is applicable to our setting.
* There are some non-neural methods that focus on a related setting [2,3,4] but still are **not** applicable. (We cited all of them in our related work section.) There are different reasons for that: they impose strong assumptions on the outcome distribution of the potential outcomes, they are not designed for multi-dimensional outcomes and static covariate data, and scalability is limited. For example, the methods in [4] can only deal with continuous treatments but **not** discrete treatments as in our work. The Gaussian process from [3] can **not** deal with static covariates because of which all patient-level information is ignored. Further, similar to [2], it scales cubically because of which practical application in medicine is prohibited.

**Action:** We have carefully checked our related work section to explain that existing, non-neural baselines are not applicable. Furthermore, we reiterate the downsides of existing baselines and the strengths of our proposed method in our new discussion (see **new Supplement A**).

Reviewer comment (Writing):
The conclusion states: “(4) Our BNCDE is further fairly robust against noise”. It would be nice to be more explicit about what “fairly” means.

Author response:
Thank you. Upon reading your comment, we found that we should remove “fairly” in our conclusion. Rather, we now point to our experiments in Figure 4 where we show that our method is more robust to noise in the data-generating process as compared to the baseline. 

**Action:** We improved our wording in the conclusion to be more precise (see the conclusion revised **Section 5**).

Reviewer comment (Novelty):
In the abstract, it states: “However, existing methods for this task are limited to point estimates of the potential outcomes, whereas uncertainty estimates have been ignored.”. I believe this is an incorrect assumption. There are methods which take uncertainty into account. Are you specifically referring here to neural methods?

Author response:
Thank you. We fixed our statement. You are correct: we were referring to uncertainty-aware, **neural** methods for time-varying treatment effect estimation.

**Action:** We clarified the abstract by adding the term “neural”.

Reviewer comment (Theory):
The confounding problem is not discussed in the proposed approach. As far as I understand, the baseline method is able to deal with confounders through a balancing term in the loss. Is it possible to address the problem of confounders within this framework?

Author response:
Thank you for your suggestion. The key motivation of our work was to highlight the effectiveness of combining neural CDEs and neural SDEs for uncertainty-aware prediction of the potential outcomes. 

To address your point, we added an extension of our method where we used balanced representations to deal with confounding. Thereby, we are consistent with the baseline (i.e., TE-CDE [1]), which also uses balanced representation. Our results demonstrate that our extended method outperforms the baseline with balanced representations. 

**Action:** We added an extension of our method with balanced representations (see **new** **Supplement M**). We find that our proposed method outperforms the baseline with balanced representations on a new real-world and a new semi-synthetic data set.

Reviewer comment (Theory):
It appears that the timestamps follow a point process. Is there a method to verify if the intensity satisfies the Overlap assumption: 0 < λ(t | H_t^i) < 1? I understand that intensities in point processes can exceed 1. I believe that the number of events would be less frequent with the constraint λ(t | H_t^i) < 1. Could you please provide some insights on this?

Author response:
Thank you for your question. Upon reading your questions, we realized that we should have spelled out more explicitly how point processes are used. You are correct: the observation times of the covariates $X_t$ and outcomes $Y_t$ follow a point process with a, possibly history-dependent, intensity function, say $\zeta(t)$. However, the point process for the observation times is **different** from the point process in the identifiability assumption. The latter is a point process with intensity $\lambda(t | H_t)$, where $\lambda(t | H_t)$ is the intensity function that governs the treatment assignment. Intuitively, this can be thought of as the analog of the propensity score in the static setting. For any given time $t$, we want the probability of treatment to be greater $0$ to ensure identifiability. In the continuous-time setting, this means that the intensity function $\lambda(t|H_t)$ is non-zero. If the intensity function was hypothetically zero, then a patient with history $H_t$ would have had no opportunity to receive treatment at time $t$. Instead, if  $\lambda(t|H_t)>0$, there is a chance that treatment is administered. 

We can attempt to **verify** this assumption in practice empirically. In the static setting, one would simply estimate the propensity score from data (see, e.g., [1]). We can also adopt the same strategy in our continuous-time setting. Here, we can simply estimate the intensity function from data, for example, via maximum likelihood or nonparametric approaches [2]. Thereby, we can verify whether the intensity satisfies the overlap assumption.  

**Action:** We clarified that the observation times of $X_t$ and $Y_t$ follow a point process $\zeta(t)$ that is different from the treatment assignment process $\lambda(t)$ (see our revised **Section 3**). We further explained how the treatment intensity can be estimated from data, so that the overlap assumption can be verified (see our revised **Section 3**).

Reviewer comment (Theory):
Understand that the variance of the variational posterior is often smaller than that of the true posterior. Could this potentially lead to overconfidence in uncertainty predictions?

Author response:
Thank you for asking this important question. In theory, the variational weight posterior variance could also be smaller than the true weight posterior variance. In such cases, uncertainty estimates will be overconfident. This has two implications: if practitioners rely on uncertainty estimates as a decision rule to administer a treatment, they may be misled. Further, the posterior predictive distribution will become overconfident.

However, our empirical findings show that this is **not** an issue for our BNCDE. Across all experiments, we find that the posterior predictive distribution of our BNCDE tends to be **conservative** (see, e.g., Figure 2 and the additional studies in **Supplements K, L, M**). That is, it tends to provide **more coverage** than the predictive posterior credible intervals should guarantee from a frequentist point of view. Further, although we cannot assess the true, underlying posterior variance, we highlight that Figure 4 clearly shows the **informativeness** of our estimated weight posterior variance. That is, we observe meaningful differences in magnitude, which demonstrates the effectiveness of our method. 

During hyperparameter tuning of our BNCDE, we further noticed that a higher diffusion hyperparameter in the latent SDEs generally led to **more conservativeness** in the potential outcome estimation. We therefore suggest that a careful choice of the diffusion hyperparameter is one of the most important hyperparameters in our method for applications in medical scenarios. 

**Action:** We added the above recommendations for choosing the diffusion hyperparameter to our implementation details (see revised **Supplement I**). Furthermore, we discuss that a larger diffusion hyperparameter leads to more conservative estimates of the posterior distribution (see new **Supplement A**).

Reviewer comment (Experiments):
Experimental results are only on one synthetic data. Is there any real-life data or semi-synthetic data that is applicable to this model?

Author response:
Thank you for giving us the opportunity to demonstrate the effectiveness of our BNCDE for other datasets. In particular, we followed your suggestions and evaluated our method using both a real-world dataset and a semi-synthetic dataset from medicine. Our real-world dataset is MIMIC III [4], which includes patients from intensive care units. Further, we followed [3] and developed a semi-synthetic dataset based on MIMIC III. The results show that our method is clearly suitable for other datasets and outperforms the baseline.  

**Action:** We added **new experiments** **using real-world data and semi-synthetic data** (see **Supplement M**). Again, our experiments show that our BNCDE consistently outperforms the baseline and yields very reliable results.

Reviewer comment (Writing):
Page 7: "we the maximize ELBO"

Author response:
Thank you very much for pointing this out. We fixed the typo.