Bayesian Neural Controlled Differential Equations for Treatment Effect Estimation
Authors: Konstantin Hess
Valentyn Melnychuk
Dennis Frauen
Stefan Feuerriegel
Venue: ICLR 2024 poster (2024)
Keywords: treatment effect estimation
neural differential equation
variational Bayes
medicine
Primary area: causal reasoning
TL;DR: We develop a novel variational Bayesian method for uncertainty-aware treatment effect estimation in continuous time using neural controlled differential equations and neural stochastic differential equations.

Abstract: Treatment effect estimation in continuous time is crucial for personalized medicine. However, existing methods for this task are limited to point estimates of the potential outcomes, whereas uncertainty estimates have been ignored. Needless to say, uncertainty quantification is crucial for reliable decision-making in medical applications. To fill this gap, we propose a novel Bayesian neural controlled differential equation (BNCDE) for treatment effect estimation in continuous time. In our BNCDE, the time dimension is modeled through a coupled system of neural controlled differential equations and neural stochastic differential equations, where the neural stochastic differential equations allow for tractable variational Bayesian inference. Thereby, for an assigned sequence of treatments, our BNCDE provides meaningful posterior predictive distributions of the potential outcomes. To the best of our knowledge, ours is the first tailored neural method to provide uncertainty estimates of treatment effects in continuous time. As such, our method is of direct practical value for promoting reliable decision-making in medicine.

Published as a conference paper at ICLR 2024
BAYESIAN NEURAL CONTROLLED DIFFERENTIAL EQUATIONS FOR TREATMENT EFFECT ESTIMATION
Konstantin Hess, Valentyn Melnychuk, Dennis Frauen & Stefan Feuerriegel
Munich Center for Machine Learning
LMU Munich
(k. hess, me me lnychuk, frauen, feuerrieeeeeee
ABSTRACT
Treatment effect estimation in continuous time is crucial for personalized medicine. However, existing methods for this task are limited to point estimates the potential outcomes, whereas uncertainty estimates have been ignored. Need- less to say, uncertainty quantification is crucial for reliable decision-making in medical applications. To fill this gap, we propose a novel Bayesian neural con- trolled differential equation (BNCDE) for treatment effect estimation in contin- uous time. our BNCDE, the time dimension is modeled through coupled system of neural controlled differential equations and neural stochastic differen- tial equations, where the neural stochastic differential equations allow for tractable variational Bayesian inference. Thereby, for an assigned sequence of treatments, our BNCDE provides meaningful posterior predictive distributions of the potential outcomes. To the best of our knowledge, ours is the first tailored neural method to provide uncertainty estimates of treatment effects in continuous time. As such, our method is of direct practical value for promoting reliable decision decision-making in medicine.
INTRODUCTION
Personalized medicine seeks to choose treatments that improve patient's future health trajectory. To this end, reliable estimates of treatment effects over time are needed (Allam et al.] 2021; Bica (et al. 202I). For example, in cancer therapy, a physician may base the decisions applying chemotherapy on whether or not the expected health trajectory will improve after treatment.
In medicine, there a growing interest in estimating treatment effects from patient trajectories using observational data (c.g., electronic health records) (Allam et al. 2021 Bica et al., 2021 Feuerriegel
ct al. 20244). Methods for this task should fulfill two requirements: (1) Existing methods typically
require a patient's ss health trajectory be recorded in regular time steps (e.g., Bica al.) 2020 Mel- nyccuk et al. 2022). However, medical practice is highly volatile and dynamic, as patients may need immediate treatment. Hence, methods are needed that model patient trajectories not in discrete time (e.g., fixed daily or hourly time steps) but in continuous time (i.c., actual timestamps). reliable decision-making, medicine ing uncertainty (e.g., credible intervals) (Zampieri et al. 2021 Banerji et al. 2023). For example, rather than saying that treatment is expected to reduce the size of a tumor by x, one is interested in whether the size of tumor will reduce by with 95% probability, given the evidence of available data. Hence, methods for treatment effect estimation must allow for uncertainty quantification. To the best of our knowledge, tailored method that accounts for both (1) and (2) is still missing.
(2) To ensure only interested in point estimates but also the correspond-
Several neural methods have been developed for individualized treatment effect estimation from ob- servational data over time (see Section 2 for overview). Here, methods often focus on simplified settings in discrete time (e.g., Lim et al. 2018 Bica et al.) 2020; Kuzmanovic et al. 2021 Li et al. 2021 Melnychuk et al. 2022) but not in continuous time. In contrast, there is only single neural method that operates in continuous time, namely. TE-CDE (Seedat et al., 2022). Yet, this method
lacks rigorous uncertainty quantification.

---

Published as a conference paper at ICLR 2024
Continuous
Outcome uncertainty
Model uncertainty
RMSNs (Lm et al.
2018
X
CRN [Bica et al. 2020
X
G-Net
(4)
CT Meinychuk
2022
X
TE-CDE (Seedat et al.
2022
(X)
BNCDE (ours)
(3): Authors use only MC dropout for model uncertainty
Table 1: Overview of key neural methods for treatment effect estimation over time. Model uncer- tainty (epistemic) refers to the uncertainty with respect to the optimal model parameters. Outcome uncertainty (aleatoric) is the uncertainty that is inherent to the data-gent gentraanng process. While model uncertainty decreases with increasing sample size, outcome uncertainty does not.
In this work, we develop novel neural method for treatment effect estimation from observational data in continuous time that allows for Bayesian uncertainty quantification. For this purpose, we propose the Bayesian neural controlled differential equation (called BNCDE). In our BNCDE, we follow the Bayesian paradigm to account for both model uncertainty and outcome uncertainty. To the best of our knowledge, ours is the first tailored neural method for uncertainty-awarr treatment effect estimation in continuous time.
In our BNCDE, the time dimension modeled through coupled system of neural controlled differ- ential equations and neural stochastic differential equations, where the neural stochastic differential equations allow for tractable variational Bayesian inference. Specifically, we use latent neural SDEs to parameterize the posterior distribution of the weights in the neural controlled differential equa- tions. By design, the solutions to our SDEs are stochastic weight processes based on which we then compute the Bayesian posterior predictive distribution of the potential outcomes in continuous time.
We contribute to three different streams of the literature D (1) We contribute to the literature on treatment effect estimation: We develop novel method for treatment effect estimation from obser- vational data continuous time with uncertainty quantification. (2) We contribute the the neural differential equation literature: To the best of our knowledge, version of neural controlled differential equations. (3) We contribute to the medical literature: We show empirically that our method yields state-of-the-an performance, paving the way for reliable, uncertainty- medical decision making.
the first propose Bayesian
RELATED WORK
We discuss methods related to (i) individualized treatment effect estimation over time and (ii) neural ordinary differential equations, well as Bayesian methods for (i) and (ii). Thereby, show that Bayesian method for uncertainty quantification in our setting is missing (see T).le We emphasize that we focus on methods for individualized treatment effect estimation over time and for average treatment effect estimation (c.g. Robins et al. 2000 Robins & Hernán 2009 Rytgaard et al. 2022
Frauen al. 2023).
Treatment effect estimation over time: Many works focus on estimating heterogeneous treatment effects from observational data the static setting (e.g., Johansson et al. 2016 Alaa & van der Schaar 2017: Louizos et al.) 2017] 2SaliSh et al.] 2017 Yoon et 20118] Zhang et al.
chuk et 2023). In contrast, only a few works consider individualized treatment effect estimation in the dynamic setting, that is, over time F We focus neural methods for this task, as existing
non parammamet methods (Xu et al.) 2016 Schulam & Sarial 2017 Solcimani et al. (017)) impose strong assumptions on the outcome distribution, are not designed for multi-dimensional outcomes and static covariate data, and scalability limited.
(1) Some methods operate in discrete time. Examples the recurrent marginal structural net- works (RMSNs) (Lim et al. 2018), counterfactual recurrent network (CRN) (Bica et al. 2020).
https: / thub, coalb..cco//ie comm/onsstan inhess/ hases///aaeeiiaannaaaaa CCDE
Brouwer al. 2022) focus different setting with single, static treatment (see Supplement
2

---

Published as a conference paper at ICLR 2024
2021). and causal transformer (CT) Melnychuk et al. 2022). However, these meth- ods are all limited to discrete time (e.g., regular recordings such as in daily or hourly time steps), which is often unrealistic in medical practice.
G-Net (Li et al.,
(2) A more realistic approach is to estimate treatment effects in continnous time. To the best of our knowledge, the only neural method for that purpose is the treatment effect neural controlled dif- ferential equation (TE-CDE) (Seedat et al.) 2022). TE-CDE leverages neural controlled differential equations (CDEs) capture treatment effects in continuous time (see Supplement [ [or details on TE-CDE). However, unlike our method, TE-CDE does not allow for uncertainty quantification.
In the continuous time setting, the timestamps of observation may be informative about the potential outcomes and bias their estimates. A general framework to address informative sampling has been developed in (Vanderschueren et al. 2023) and applied to TE-CDE. We later also apply this approach
to our BNCDE; see Supplement K
Uncertainty quantification for treatment effect estimation: Jesson et al. (2020) highlight the importance of Bayesian uncertainty quantification for treatment effect estimation in the static setting. In the time-varying setting, existing methods are limited that they use Monte Carlo (MC) dropout
(Gal & Ghahramani 2016) as an ad-hoc solution. However, MC dropout relies upon mixtures of Dirac distributions in parameter space, which leads to approximations of the true posterior which questionable and not faithful (Le Folgoc et al. 2021). In contrast, tailored neural method for Bayesian uncertainty quantification in the continuous time setting is still missing.
Neural ordinary differential equations (ODEs) can be seen as infinitel--deep residual neural networks (Chen et al., 2018; Haber & Ruthotto 2018 Lu et al.) 2018), where infinitesimal small hidden layer transformations correspond to the dynamics of a time- evolving latent vector field. Neural CDEs (Kidger et al. 2020; Morrill et al.) 202T) extend neural ODEs in order process time series data in continuous time (for more details, Supplement D))
Neural differential equations:
Uncertainty quantification for neural ODEs: There are only a few existing approaches for this task. (1) Variants of Markov chain Monte Carlo (MCMC) have directly been applied to neural ODEs Dandekar et al.) (022).. However, MCMC methods are known to scale poorly to high dimensions. Fast approximate Bayesian inference can be achieved through Laplace approximation (Ott al., 2023 (3) VAE (Yildız al. 2019) uses Gaussian distributions for variational inference in neural ODEs. Yet, both (2) and (3) rely on unimodal distributions with limited expressiveness. (4) The posterior can be approximated through neural stochastic differential equations (SDEs) (Li) ctt al. (020). A key benefit of this is that neural SDEs as a variational family are both scalable and arbitrarily expressive (Tzen Raginsky 2019 Xu et al. 2022). However, no work has integrated Bayesian uncertainty quantification into neural CDEs.
Research gap: As shown above, there is no method for treatment effect estimation in continuous time with Bayesian uncertainty quantification. To fill this gap, we propose our BNCDE. To the best of our knowledge, our BNCDE also the first Bayesian neural CDE.
3 PROBLEM FORMULATION
Setup: Let L [) be the observation time and let E (0, A] be the prediction window. We then consider n patients (i.i.d.) as realizations of the following variables: (1) Oulcomes over time are given by Y, e E (c.g.. tumor volume). (2) Covariates X, E ddd include additional patient informa- tion such as comorbidity. (3) Treatments are given by A, E {0, 13dd. where multiple treatments can be administered at the same time. The treatment assignment is controlled by multivariate counting process N₂ E No with intensity X(t), where N₂ corresponds to the number of treatments assigned up to a specific point in time t. This is consistent with medical practice where the same treatment oftentimes applied multiple times. For example, cancer patients may receive multiple cycles of chemotherapy (Curran et al. 201T).
Importantly, we focus on a setting in continuous time. That is, the outcomes and covariates of each patient i are recorded at timestamps {ti, (ti,ti..... .i,...,,t,, Gn, with 0 and et F. where m. is the number of timestamps and 8 is the latest observation time for patient i Observation times
W...o.g. we assume that e 0 for all
3

---

Published as a conference paper at ICLR 2024
assumed follow another, possibly diseeey-dependent, intensity process {(t). This is a crucial difference from a setting in discrete time, since, in our setting, the timestamps are arbitrary and thus non-regular. Further, the timestamps may differ between patients.
We have access to observational data hi = foo d ( "j..g E [0, which we refer to as patient trajectories. Formally, the observed outcomes time are given by vfo.1 Uega, (vig) and the observed covariates by 2/0.4] Uega, (x where ti I t is the latest observation time up to time with k; S my. In contrast, we have full knowledge of the history of assigned treatments, i.e., d0.1] = Usge U. We are interested in the potential outcome for a new patient time 1' + A. Yi+ (( for an arbitrary future sequence of treatments d/z- [+++] = Ui df hτ U (a(p)44))
given patient's observed trajectory hi_. For notation, write (((+2)
when we include the factual future sequence of treatments. We write h ++ = M(0+2)- U (
when we further include the realized observed outcome for this future sequence of treatments.
Our objective is to predict the potential outcome for a future sequence of assigned treatments, given the observed patient history. In the following, we adopt the potential outcomes framework [Neyman| 1923] Rubin 1978) and its extensions to the time-varying setting (Lok 2008
Estimation Task:
Robins & Hernán, 2009 Saarela & Liu 2016 Rytgaard et al. 2022).
Unique to our setting that we do not focus on simple point estimates but perform Bayesian un- certainty quantification. In the Bayesian framework, model parameters E S are assigned prior distribution p(w). Further, for an individual and a future sequence of treatments d/z- + the po- tential outcome +r+ *r+alo(r= (((r=,+++)) has a likelihood p(Yi++a + i ++)]] hi,, w), given a parameter realization and patient trajectory hi.. We thus aim estimate the posterior predictive distribution
p(Y++ ( P(Y
+aj) h H) = Ep(w #0) [P(Yr MoYY+taldf lali- a +aj] his, []]:
(1)
which is the weighted average likelihood under the parameter posterior distribution p(w HH) given the training data H U h - Of note, our setting is different from other works such as TE-CDE Seedat et al. 2022), where only point estimates such as IE(i-+ala'' (++a|] hi.1 are computed but without uncertainty quantification. Instead, we estimate the full distribution of the potential outcomes.
The above estimation task is challenging due to the fundamental problem of causal inference (Imbens & Rubin 2015) in that only factual but never counterfactual patient trajecto- ries are observed. To ensure identifiability from observational data, we make the following three assumptions that are standard in the literature (Lok) 2008, Robins & Hernán 2009 [aarela & Liu 2016 Rytgaard et al. 2022 Ying. 2022): (I) Consistency: Given a sequence of treatments Ap,,+r] a'u,ttri. f 2 0 and T [0, A), the potential outcome Y=+~|00o, coincides with the observed outcome Y++r- (2) Overlap: For any realization of patient history H, = h4, there is pos- itive probability of receiving treatment any point t time. That is, the intensity process satisfies 0 < X(L he)) < 1. Similar to estimating the propensity score in the static setting (Schweisthal et al. 2023), the intensity can be estimated from data (Leemis) 1991}. (3) Unconfoundedness: The treatment assignment probability is independent of future outcomes and unobserved information.
Identifiability:
This means that the intensity process satisfies X(t fu) = X(t hr F(Y((l) : > t)), where
F(Y.ld/ t) is the filtration generated by future potential outcomes.
4 BAYESIAN NEURAL CONTROLLED DIFFERENTIAL EQUATION
4.1 ARCHITECTURE
Our BNCDE builds upon an encoder-decodes architecture (see Fig. D with three components. The A encoder receives the patient trajectory H, in continuous time and encodes into a hidden rep- resentation Z, up to time 7. The B decoder then takes the hidden representation Zi together with future sequence of treatments a' and transforms it into new hidden representation Ż, for 0- < T < A, where A is the desired prediction window. Both the encoder and the decoder each consist of (i) neural CDE and (ii) latent neural SDE (see Supplement D for overview).
4

---

Published as a conference paper at ICLR 2024
Inpelt nutput
connectoon
Noure COEK
Ax
Decoder
Prodiction
Input
Neure 50
Prediction
head
Erooder variable
Vesodervanitlle
Powterior
Dats
Aa
Te-a
Time
Figure 1: Our BNCDE consists of an encoder, a decoder, and prediction head.
(i) The neural CDEs compute hidden representations of the patient trajectories continuous time. The hidden representations in the neural CDEs allow model non-linearities in the data, reduce dimensionality, and capture dependencies between observed variables (Y,, Xt, A.). (ii) The latent neural SDEs approximate the posterior distribution of the neural CDE weights. Finally, the C pre- diction head receives Z and parameterizes the likelihood the potential outcome. We explain the different components in the following.
Encoder: The encoder has two components, namely a neural CDE and a latent neural SDE: (i The neural CDE encodes hidden representation that is driven by the patient history. (ii) The latent neural SDE approximates the posterior distribution of the neural CDE weights.
The neural CDE consists of trainable embedding network n RI+d, +d. Rdd that encodes the initial outcome Yo. covariates Xo. and treatments Ao into hidden state Zo. The hidden state Zo serves as the initial condition of the neural CDE. Formally, it is given by
Ze = ( (Z W.) dl X,, As), t (0, Zo = no Xu. An),
(2)
where f' wi): Rd.+1 (dexd is the neural vector field controlled by [Y, X,, Ar] given the network weights w/. In particular, f is Bayesian neural network. encodes information about the patient patient history into the hidden states Z, any time & [". Given the observed data, the posterior distribution the random weights We is modeled through the latent neural SDE.
Hence, the weights not static but varying. The output Zp is then passed to the decoder.
The neural SDE is used to model the neural CDE weights Wi E Rd,,, Here, we assume that the weights evolve over time according to stochastic process, where the stochastic process is the lution to the latent neural SDE. Specifically, we optimize the neural SDE to approximate the joint posterior distribution p( HH) of the weights we [0, T) given the training data. We refer to the joint variational distribution (w(o.))) Then, the neural SDE is formally defined by
f (ws, s) ds + ( o dBB,, t (0, Wo E
where 98 Rd=+1 Rdw is the drift network, is constant diffusion coefficient, B, RRd is standard Brownian motion, and eg *N(2,, .) is the initial condition.
The decoder also consists of a neural CDE and a latent neural SDE. (i) The neural CDE is controlled by future sequence of treatments. (ii) The latent neural SDE approximates the posterior distribution of the neural CDE weights.
B
Decoder:
The neural CDE takes the hidden representation ZT from the encoder as the initial condition. The neural CDE then given by
2, 1/' ('(Z,, (2)) da',, TE (0, Zo ZT.
5

---

Published as a conference paper at ICLR 2024
(.)): Rd + ((+((H++tdd)) is the Bayestan neural vector field controlled by future sequence of treatments al given the network weights W,. The hidden state Z is then passed to the prediction head.
where (z,
The neural SDE is defined in the same way as in the encoder. It is used make variational approximation d (s) of the joint posterior distribution p( H, w)) of the weights w, on [0, A] given the training data and the neural CDE weights from the encoder. Formally, the neural SDE is given by
f.=[['0 (( I/ o TE (0, Wp- = Es:
W -,= =
(5)
where 98 Rd Rd is the drift network, B, Rd is another standard Brownian motion, and cØ (N 0) the initial condition.
C Prediction head: The prediction head is a trainable mapping Te. Rd - RR? It takes the hidden state Z as input and then predicts the expected value JUT+A and the outcome uncertainty ~ of the potential outcome at time 7 + A via
(UF-A, ΣT+A) = (2a)
(6)
4.2 COUPLED SYSTEM OF NEURAL DIFFERENTIAL, EQUATIONS
We can write our BNCDE coupled system of differential equations, which allows our BNCDE to learn in an end-to-end manner. Formally, the system of neural differential equations is given by
98 (W):T) (2r (+))
Br
d
(c)) od
)
2r
0
W:
(Z -p(l2,t wsw
Be
d
d (e:
od
(8)
Ze
0
with initial conditions
= Zr, Wo 5/ Zo (V,, Xo, Ao)
and wi
Eg
(9)
and where E (0, 7 and (0, A]. The prediction head then outputs (UT+ Ef+a) 7/8 (Z)) The embedding network n/g the prediction head 7/g the initial conditions Eg and Eg and the drift networks O and L learned via gradient-basse methods.
For our model, the use of neural SDEs has three main benefits. (1) Neural SDEs allow us to approx- imate parameter posterior distributions that are arbitrarily complex. We emphasize that, although o is constant, prior research (Tzen & Raginsky] 2019, Xu et al.) [022] shows that given sufficiently expressive family of drift networks, neural SDEs are able to approximate the true posterior distri- butions with arbitrary accuracy. (2) Our neural SDEs are non-autonommuss, which means that they directly incorporate dependency on time. This way, we make sure that the weights learn to adjust their dynamics over time. (3) Our neural SDEs are learned through variational Bayes and enable inference in seconds.
4.3 POSTERIOR PREDICTIVE DISTRIBUTION
Our model generates the full posterior predictive distribution of the potential outcome, given patient history and training data. For a potential outcome at time f* + A, our variational approximation of the posterior predictive distribution in Eq. is given by
[p(Y+ Aoresaldee e + wioss , w/o,r-1, n(-) R (w).))) E Egg(oio.s)) |P(Y++a (oftr,-+4l d+j)] /o. wp. hīz
(10)

---

Published as a conference paper at ICLR 2024
Hence, our posterior predictive distribution can be thought of Gaussian mixture model with infinitely many mixture components parameterized by the weights W|,,T-1 and WpoA).. We provide a formal derivation in Supplement
We emphasize that our model incorporates both (1) model uncertainty (epistemic) and (2) outcome uncertainty (aleatoric): (1) Model uncertainty is represented by the variance of the latent neural SDEs. The variances of the marginals L (wpr) and L (wi,) are larger for patient histories h2. and future sequences of treatments d/e- FF that are unlike any observations in the training data H. Correspondingly, this leads to a larger variance of the posterior predictive distribution. Our BNCDE therefore informs about uncertainty in predictions due to lack data support. (2) Outcome uncer- lainty is determined by the likelihood variance Er-+A A larger likelihood variance implies that for a given patient history hü- and a future sequence of treatments d(p ((+))))) there a higher degree of variability in the potential outcome at future time E + A.
4.4 TRAINING
Our BNCDE trained by maximizing the evidence lower bound (ELBO). which further requires the specification of prior distributions of the neural CDE weights and a likelihood.
Priors: Our chosen priors for both the encoder weights wt and the decoder weights W, are indepen- dent Ornstein- Brreen--hhlebeek (OU) processes due to their finite variance in the time limit. We set the prior drifts to h" ("())) = (-wi) and h ((,) (-w,) respectively and the diffusion coefficients to σ.
Likelihood:
The likelihood in Eq. 10 is a normal distribution parameterized by the prediction head in Eq. 6 That is, model the likelihood as
YT+ Yf+ald/,,4+)) (Wo.A)) Hy) ~ ~N(T++ ET+A).
(11)
For training, we optimize the variational posterior distribution of the weights using the data H. Specifically, maximize the ELBO for observation given by
ELBO:
log MB+4) ) Bogtoooo (agieno log p(uf + W).,) W.. 1 (( }}}
- -DKLlg ll p(W(),)))] - DEL 1 (w)))) lI p(w(,,))) (12)
where DKL is the Dallback-Leibles divergence, (()() and p(w)o.ry) are the joint distributions of the OU priors on [0, A] and [o, 7') respectively. We provide a full derivation of the ELBO loss in Supplement F A summary of the numerical ELBO approximation provided in Supplement Further implementation details of our model are provided in Supplement
5
NUMERICAL EXPERIMENTS
5.1
SETUP
tumor growth model by Geng et al. (2017). Variants of this model are the standard used to assess the performance of time-var varying treatment effect models (Lim et al., 2018 Bica et al. 2020; Li al. Melnychuluuuu et al. 2022, Seedat et al.) [2022, Vanderschueren et al., 2023). ). In the tumor growth model, the outcome Y, is the tumor volume Geng et al. (2017). At time I, a treatment may be applied and, if so, includes chemotherapy, radiotherapy, or both. We employ the same model variant in Vanderschueren et al. (2023), where observations are sampled non-regularly, that is, in continuous time. We provide more details in Supplement
Data: We benchmark our method using the established hatpmaoookknnte-hhaaaooooyooeo
Baselines: Due to the novelty of our setting, there are no existing baselines for uncertainty quan- tification in continuous time (see Table 1). The only comparable method is TE-CDE with Monte Carlo (MC) dropout (Seedat et al. 2022 Implementation details are in Supplement
Performance metrics: We use different metrics assess whether the uncertainty estimates of our method are faithful and sharp. To this end, we examine the posterior predictive credible intervals (Crls) of the potential outcomes generated by the different methods. For each patient, compute
7

---

Published as a conference paper at ICLR 2024
the individual equal-tailed (1 a) posterior predictive CrI with [0.01, 0.05]. This choice is motivated by medical practice, where the treatment effectiveness is typically evaluated based on the 95% and 99% Crls. We assess the faithfulness of the Cris by computing the empirical coverage, that is, the proportion of times the Cris contain the outcomes in the test data. We further assess the sharpness of the computed Cris by reporting the median width of the Crls between 95% and 99% over all patients in the test data. We report the results for the prediction windows A = 1, 3. Results for additional prediction windows in Supplement We report the averaged results along with the standard deviation over five different seeds.
We further evaluate the enror in the point estimates. Specifically, we compute the mean squared error (MSE) of the Monte Carlo mean estimates of the observed outcomes.
5.2 RESULTS
Faithfulness: We first evaluate whether the estimated Cris are faithful. For this, we show the empirical coverage across different quantiles (1 a). The results are in Fig. 2 For our BNCDE, the (1 (1-a) Crl almost always contains at least (1 ax) of the outcome, implying that the estimated Crls are generally faithful. In contrast, the estimated Crls for TE-CDE are generally not faithful. In particular, the CrIs from TE-CDE often contain fewer outcomes than the (1 - a) level should guarantee, implying that the uncertainty estimates from TE-CDE are overconfident. This is in line with the literature, where MC dropout is found to produce poor approximations of the posterior (Le) Folgoc al.) 2021). For longer prediction windows A, the advantages from method greater. In sum, the results demonstrate that our method is clearly superior.
1538
-30
(a)
(b)
(c)
Figure 2: Faithfulness: Empirical coverage across different Crl quantiles. Shown are different prediction windows A 1,2,3. 3. Areas in green (red) indicate that the CrIs are faithful (not faithful).
Sharpness: We compare the median width of estimated CrIs (see Fig. The Crls from our BNCDE are significantly sharper than those of TE-CDE. This holds for all prediction windows, which further demonstrates the effectiveness of our method.
TECTE
Pearrer
(a) A =
(b) A = 2
(c) A = 3
Figure 3: Sharpness: Width of the Cris (median) for different quantiles a.
Error in point estimates: We further compare the errors in the Monte Carlo mean estimates of the treatment effects. The rationale is that our BNCDE may be better at computing the posterior distribution, yet the more complex architecture could hypothetically let the point estimates deterio- rate. However, this is not the case, and we see that our BNCDE is clearly superior (see Fig. 4. We further find that the performance gain from our method is robust across different levels of noise in the data-generating process (i.e., larger Var(e:)).
8

---

Published as a conference paper at ICLR 2024
642
(a) A
(b) A
(c)
Figure 4: Error in point estimates: Reported is the median over the mean squared errors (MSE) of the point estimates of the outcomes. The results are based on test data that is generated with varying levels of noise, i.e., Var(cz).
5.3 COMPARISON OF MODEL UNCERTAINTY
We now provide deep-dive comparing the uncertainty estimates for model uncertainty only (and thus without outcome uncertainty). We do so for two reasons: (1) We can directly compare the model uncertainty from both our BNCDE and TE-CDE with MC dropout as the latter is limited to model uncertainty.
(2) We can better understand the role of the neural SDE in our method (i.e., before the variational approximations are passed to the prediction head). For TE-CDE we thus use the variance in the MC dropout samples
measure of model uncertainty. For our BNCDE, we capture the model uncertainty with the Monte Carlo variance in the means of the likelihood Mr++2 under the neural SDE weights wju [-] and W.. However, both measures
not directly comparable. Hence, inspired by analyses for static settings (Jesson et al. 2020: 2021: Oprescu et al,, 2023), compare the deferral rate. That is, report errors in the MSE of the treatment effect as we successively withhold samples with the largest Monte Carlo variance.
FLLLES
221L dicinal
(a) A=
(b) =2
(c) (A)
Figure 5: To compare model uncertainty, the normalized MSE of estimated treatment effects be- tween the treatment arms versus deferral rate is shown.
Fig. 5 shows the normalized MSE across different proportions of withheld samples. For TE-CDE with MC dropout, the MSE decreases almost randomly, implying that the MC estimates mea- sure of model uncertainty are highly uninformative (and close to random deferral). For our BNCDE, larger MSE in the treatment effect estimates corresponds to higher model uncertainty, thus im- plying that the model uncertainty in our method is more informative.
We repeat experiments for informative sampling (see Supplement K). For this, extend both our BNCDE and TE-CDE using the inverse intensity weighting method from Vander- schueren et al. (2023). We find that our method remains clearly superior.
Extension:
Our results show the following: (1) Our BNCDE produces posterior predictive Cris that are faithful.. In contrast, the Crls from TE-CDE can be overconfident, which, in medicine, may lead to treatment decisions that are dangerous. (2) Our BNCDE further provides approximations of the CrIs that are sharper. (3) Our BNCDE generates more informative estimates of model uncer- tainty compared to TE-CDE with MC dropout. (4) Our BNCDE is further more robust against noise than TE- EDE with MC dropout and, also, (5) highly effective for longer prediction windows. We discuss the applicability of our BNCDE in medical scenarios in Supplement A
Conclusion:

---

Published as a conference paper at ICLR 2024
REFERENCES
Ahmed M. Alaa and Mihaela van der Schaar. Bayesian inference of individualized treatment effects using multi-task Gaussian processes. In NeuriPS, 2017.
Ahmed M. Alaa, Scott Hu, and Mihaela van der Schaar. Learning from clinical judgments: Semi- Markov-modulated marked Hawkes processes for risk prognosis. In ICML, 2017.
Ahmed Allam, Stefan Feuerriegel, Michael Rebhan, and Michael Krauthammer. Analyzing patient trajectories with artificial intelligence. Journal of Medical Internet Research, 23(12):e29812, 2021.
Christopher R. S. Banerji, Tapabrata Chakraborti, Chris Harbron, and Ben D. MacArthur. Clinical AI tools must convey predictive uncertainty for each individual patient. Nature medicine, 2023,
Ioana Bica, Ahmed M. Alaa, James Jordon, and Mihaela der Schaar. Estimating counterfactual treatment outcomes over time through adversarially balanced representations. In ICLR, 2020.
Ioana Bica, Ahmed M. Alaa, Craig Lambert, and Mihaela van der Schaar. From real-world patient data individualized treatment effects using machine learning: Current and future methods to address underlying challenges. Clinical Pharmacology and Therapeutics, 109(1):87-100, 2021.
Edward De Brouwer, Javier González Hernández, and Stephanie Hyland. Predicting the impact of
treatments over time with uncertainty aware neural differential equations. AISTATS, 2022.
Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary dif- ferential equations. In NeurIPS, 2018.
KyungHyun Cho, Bart van Merrienboer, Dzmitry Bahdanau, and Yoshua Bengio. On the properties of neural machine translation: encoder-decoder approaches. arXiv. 1409. 1259, 2014.
Walter J. Curran, Rebecca Paulus, Corey J. Langer, Ritsuko Komaki, Jin S. Lee, Stephen Hauser, Benjamin Movsas, Todd Wasserman, Seth A. Rosenthal, Elizabeth Gore, Mitchell Machtay, William Sause, and James Cox. Sequential vs. concurrent chemoradiation for stage IIII non- small cell lung cancer: Randomized phase III trial RTOG 9410. Journal of the National Cancer Institute, 103(19):1452-146) 2011.
Alicia Curth and Mihaela van der Schaar. Nonparametric estimation of heterogeneous treatment effects: From theory to learning algorithms. In AISTATS, 2021.
Raj Dandekar, Karen Chung, Vaibhay Dixit, Mohamed Tarek, Aslan Garcia- Valadez, Krishna Vishal Vemula, and Chris Rackauckas. Bayesian neural ordinary differential equations. arXiv preprint, 2012.07244, 2022.
Stefan Feuerriegel, Dennis Frauen, Valentyn Melnychuk, Jonas Schweisthal, Konstantin Hess, Ali- cia Curth, Stefan Bauer, Niki Kilbertus, Isaac S. Kohane, and Mihacla van der Schaar. Causal machine learning for predicting treatment outcomes. Nature Medicine, 2024.
Dennis Frauen, Tobias Hatt, Valentyn Melnychuk, and Stefan Feuerriegel. Estimating average causal effects from patient trajectories. In AAAÍ, 2023.
Yarin Gal and Zoubin Ghahramani. Dropout as a Bayesian approximation: Representing model uncertainty in deep learning. In ICML, 2016.
Changran Geng, Harald Paganetti. and Clemens Grassberger. Prediction of treatment response for combined chemo- and radiation therapy for non-small cell lung cancer patients using bio- mathematical model. Scientific Reports, 7(1):13542, 2017.
Eldad Haber and Lars Ruthotto. Stable architectures for deep neural networks. Inverse Problems, 34( 34(1)4014004, 2018
Stephanie L. Hyland, Martin Faltys, Matthias Hüser, Xinrui Lyu, Thomas Gumbsch, Cristóbal Es- teban, Christian Bock, Max Hor, Michael Moor, Bastian Rieck, Mare Zimmermann, Dean Bo- denham, Karsten Borgwardt, Gunnar Rätsch, and Tobias M. Merz. Early prediction circulatory failure in the intensive care unit using machine learnnn.. Nature medicine, 26(3):364-373 2020.
10

---

Published as a conference paper at ICLR 2024
Guido W. Imbens and Donald B. Rubin. Causal inference for statistics, social, and biomedical sciences: An introduction. Cambridge books online. Cambridge University Press, Cambridge, 2015. ISBN 978113902
Andrew Jesson, Sören Mindermann, Uri Shalit, and Yarin Gal. Identifying causal-effect inference failure with uncertainty-awiin models. In NeurlPS, 2020.
Andrew Jesson, Sören Mindermann, Yarin Gal, and Uri Shalit. Quantifying ignorance in individual- level causal-effect estimates under hidden confounding. In ICML, 2021.
Fredrik Johansson, Uri Shalit, and David Sonntag. Learning representations for counterfactual inference. In ICML, 2016.
Alistair E. W. Johnson, Tom Pollard, Lu Shen, Li-wei H. Lehman, Mengling Feng, Mohammad Ghassemi, Benjamin Moody, Peter Szolovits, Leo Anthony Celi, and Roger G. Mark. MIMIC-III, freely accessible critical care database. Scientific Data, 311::600335, 2016.
Patrick Kidger, James Morrill, James Foster, and Terry Lyons. Neural controlled differential equa- tions for irregular time series. In NeurIPS, 2020.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In ICLR, 2015.
Milan Kuzmanovic, Tobias Hatt, and Stefan Feuerriegel. Deconfounding temporal autoencoder: Estimating treatment effects over time using noisy proxies. In ML4HI 2021.
Loic Le Folgoc, Vasileios Baltatzis, Sujal Desai, Anand Devaraj, Sam Ellis, Octavio E. Martinez Manzanera, Arjun Nair, Huaqi Qiu, Julia Schnabel, and Ben Glocker. Is MC dropout Bayesian? arXiv preprint, 21 10.04286, 2021.
Lawrence M. Leemis. Nonparametric estimation of the cumulative intensity function for nonho- mogeneous Poisson process. Management Science, 37(7):886 88009008, 1991.
Rui Li, Stephanic Hu, Mingyu Lu, Yuria Utsumi, Prithwish Chakraborty, Daby M. Sow, Piyush Madan, Jun Li, Mohamed Ghalwash, Zach Shahn, and Li-wei Lehman. G-Net: A recurrent net- work approach G-computation for counterfactual prediction under a dynamic treatment regime. In ML4H, 2021.
Xuechen Li, Ting-Kam Leonard Wong, Ricky T. Q. Chen, and David Duvenaud. Scalable gradients for stochastic differential equations. In AISTATS, 2020.
Bryan Lim, Ahmed M. Alaa, and Mihaela van der Schaar. Forecasting treatment responses over time using recurrent marginal structural networks. In NeurlPS, 2018.
Judith J. Lok. Statistical modeling of causal effects in continuous time. Annals of Statistics, 36(3), 2008.
Christos Louizos, Uri Shalit, Joris Mooij, David Sontag, Richard Zemel, and Max Welling. Causal effect inference with deep latent-variable models. In NeurIPS, 2017.
Yiping Lu, Aoxiao Zhong, Quanzheng Li, and Bin Dong. Beyond finite layer neural networks: Bridging deep architectures and numerical differential equations. In ICML, 2018
Valentyn Melnychuk, Dennis Frauen, and Stefan Feuerriegel. Causal transformer for estimating counterfactual outcomes. ICML, 2022.
Valentyn Melnychuk, Dennis Frauen, and Stefan Feuerriegel. Normalizing flows for interventional density estimation. ICML, 2023.
Valentyn Melnychuk, Dennis Frauen, and Stefan Feuerriegel. Bounds on representation-induuee confounding bias for treatment effect estimation. In ICLR, 2024.
James Morrill, Patrick Kidger, Lingyi Yang, and Terry Lyons. Neural controlled differential equa-
tions for online prediction tasks. arXiv preprint, 2106.1 11028, 2021.

---

Published as a conference paper at ICLR 2024
Elizabeth Murray, Eric B. Hekler, Gerhard Andersson, Linda M. Collins, Aiden Doherty, Chris Hol- lis, Daniel E. Rivera, Robert West, and Jeremy C. Wyatt. Evaluating Digital Health Interventions: Key Questions and Approaches. American journal of preventive medicine, 31(5):443-835 2016.
Jerzy Neyman. On the application of probability theory to agricultural experiments. Annals of Agriculural Sciences, 10:1-51, 1923.
Miruna Oprescu, Jacob Dorn, Marah Ghoummaid, Andrew Jesson, Nathan Kallus, and Uri Shalit. B-Learner: Quasi-oracle bounds on heterogeneous causal effects under hidden confounding, In ICML, 2023.
Katharina Ott, Michael Tiemann, and Philipp Hennig. Uncertainty and structure neural ordinary differential equations. arXiv preprint, 2305. 2023.
Yilmazcan Özyurt, Mathias Kraus, Tobias Hatt, and Stefan Feuerriegel. ArDMM: An attentive deep Markov model for risk scoring in intensive care units. In 2021.
Judea Pearl. Causality. Cambridge University Press, New York City, 2009. ISBN 189566.
Tom Pollard, Alistair E. Johnson, Jesse D. Raffa, Leo A. Celi, Roger Mark, and Omar Badawi. The eICU Collaborative Research Database, a freely available multi-center database for critical care research. Scientific Data, 5: 5:180178, 2018.
James M. Robins and Miguel A. Hernán. Estimation the the causal effects of time-varying exposures. Chapman & Hall/CRC handbooks of modern statistical methods. CRC Press, Boca Raton, 2009. ISBN 9781588886587.
James M. Robins. Miguel Hernán, and Babette Brumback. Marginal structural models and causal inference in epidemiology. Epidemiology, 1 11(5):5005660 2000.
Niklas Rodemund, Bernhard Wernly, Christian Jung, Crispiana Cozowicz, and Andreas Koköfer.
The Salzburg Intensive Care database (SICdb): an openly available critical care dataset. Intensive care medicine, 49(6):700-77) 7022 2023.
Donald B. Rubin. Bayesian inference for causal effects: The role of randomization. Annals of Statistics, 6(1):34-58, 1978.
Helene C. Rytgaard, Thomas A. Gerds, and Mark J. van der Laan. Continuous-tine targeted min- imum loss-based estimation of intervention-specte -specffcc mean outcomes. The Annals of Statistics, 2022.
Olli Saarela and Zhihui Amy Liu. A flexible parametric approach for estimating continuous-time inverse probability of treatment and censoring weights. Statistics in Medicine, 35(23):423--421,, 2016.
Peter Schulam and Suchi Saria. Reliable decision support using counterfactual models. In NeurIPS, 2017.
Jonas Schweisthal, Dennis Frauen, Valentyn Melnychuk, and Stefan Feuerriegel. Reliable off-policy learning for dosage combinations. In NeurIPS, 2023.
Nabeel Seedat, Fergus Imrie, Alexis Bellot, Zhaozhi Qian, and Mihaela der Schaar. Continuous- time modeling of counterfactual outcomes using neural controlled differential equations. In ICML, 2022.
Uri Shalit, Fredrik D. Johansson, and David Sontag. Estimating individual treatment effect: Gener- alization bounds and algorithms. In ICML, 2017.
Hossein Soleimani, Adarsh Subbaswamy, and Suchi Saria. Treaoment-response models for counter- factual reasoning with continuous-time, continuous-valued interventions. In UAI, 2017.
12

---

Published as a conference paper at ICLR 2024
Patrick J. Thoral, Jan M. Peppink, Ronald H. Driessen, Eric J. G. Sijbrands, Erwin J. O. Kompanje, Lewis Kaplan, Heatherlee Bailey, Jozef Kesecioglu, Maurizio Cecconi, Matthew Churpek, Gilles Clermont, Mihacla van der Schaar, Ari Ercole, Armand R. J. Girbes, and Paul W. G. Elbers. Shar- ing ICU patient data responsibly under the Society of Critical Care Medicine/Furopent Society of Intensive Care Medicine Joint Data Science Collaboration: The Amsterdam University Medical Centers Database (AmsterdamUMCCa) example. Critical Care Medicine, 49(6):563-577, 2021.
Belinda Tzen and Maxim Raginsky, Neural stochastic differential equations: Deep latent Gaussian models in the diffusion limit. arXiv preprint, 1905.09883. 2019.
Toon Vanderschueren, Alicia Curth, Wouter Verbeke, and Mihaela van der Schaar. Accounting for informative sampling when learning to forecast treatment outcomes over time. In ICML, 2023.
Kerstin N. Vokinger, Stefan Feuerriegel, and Aaron S. Kesselheim. Mitigating bias in machine learning for medicine. Communications Medicine, 1:25, 2021.
Shirly Wang, Matthew B.A. McDermott, Geeticka Chauhan. Marzych Ghassemi, Michael C. Hughes, and Tristan Naumann. MIMIC-extract: A data extraction, preprocessing, and repre- sentation pipeline for MIMIC-III In CHIL, 2020.
Winnie Xu, Ricky T. Q. Chen, Xuechen Li, and David Duvenaud. Infinitely deep Bayesian neural networks with stochastic differential equations. In AISTATS, 2022.
Yanbo Xu, Yanxun Xu, and Suchi Saria. non-parametric bayesian approach for estimating treatment-response curves from sparse time series. In MI.4H, 2016.
Çağatay Yıldiz, Markus Heinonen, and Harri Lähdesmäki. ODE?VAE: Deep generative second order ODEs with Bayesian neural networks. In NeurlPS, 2019.
Andrew Ying. Causality of functional longitudinal data, arXiv preprint, 12525, 2022.
Jinsung Yoon, James Jordon, and Mihaela der Schaar. GANITE: Estimation of individualized treatment effects using generative adversarial nets. In ICLR, 2018.
Fernando G. Zampieri, Jonathan D. Casey, Manu Shankar- Hari, Frank E. Harrell, and Michael O. Harhay. Using Bayesian methods augment the interpretation of critical care trials. An overview of theory and example reanalysis of the alveolar recruitment for acute respiratory distress syn- drome trial. American Journal of Respiratory and Critical Care Medicine, 203(5):543-35) 2021.
Yao Zhang, Alexis Bellot, and Mihaela van der Schaar. Learning overlapping representations for the estimation of individualized treatment effects. In AISTATS, 2020.
13

---

Published as a conference paper at ICLR 2024
A DISCUSSION ON THE APPLICABILITY IN MEDICAL SCENARIOS
In our main paper, we demonstrated that BNCDE represents significant step toward reliable, uncertainty-
cerrisioommaking in medicine. To this end, we improve over existing methods by leveraging the Bayesian paradigm and by introducing a novel neural approach for uncertainty- treatment effect estimation. Specifically, we provide strong empirical evidence that BNCDE based on a combination of neural CDEs and latent neural SDEs outperforms the existing ad-hoc baselines, namely TE-CDE (Seedat et al. 2022) with MC dropout (Gal & Ghahramani] 2016). Ad- ditionally, we demonstrate numerically that our method is compatible with other extensions, such as the informative sampling framework by Vanderschueren al. 2023) and balanced representations
Bica et al. 2020; Melnychuk et al. 2022 Seedat et al. 20222 (see Supplements
K and D.
Our experiments further demonstrate that our method holds great relevance in clinical settings. First, our method allows for treatment effect estimation in continuous time. That is, it does not rely on the unrealistic assumption of regular recording times in electronic health records (Özyurt et al. 2021 but can instead account for highly irregularly sampled recording times. In electronic health records, irregular observation times are common and quasi-standard (Alaa et al. 2017 Allam et al.) 2021) Because of this, the assumptions in existing methods for discrete time are directly violated, while a particular strength of our method is that we explicitly account for irregular sampling. Second, our method does not only provide point estimates of the potential outcomes but estimates the full posterior predictive distribution. This is crucial in personalized medicine, as physicians may want imperative for all
to base their decisions uncertainty estimates. In fact, uncertainty estimates medical applications to avoid patient harm (Banerji et al. 2023). To account for that, our posterior predictive distribution provides insights into both (i) model uncertainty (see Fig. 5 and (ii) outcome uncertainty (see Supplement]] D of the potential outcomes in continuous time, and is the first tailored neural method for this task.
Combining outcome estimation in continuous time and rigorous uncertainty quantification in a sin- gle, neural end-to-end architecture is a major improvement over all existing neural baselines (see Section 2) Furthermore, existing non-parametric approaches (Xu et al. 2016, [Schulam & Saria, 2017 Soleimani et al. 2017) impose strong assumptions on the outcome distribution, they are not designed for multi-dimensional outcomes and static covariate data, and they face scalability issues. For instance, Gaussian Process methods as in (Schulam & Saria 2017) require a matrix inversion of the covariance matrix involving all patients in the training data. Due to the cubic scaling of this op- eration, this is completely impractical in medical scenarios with constrained time and computational resources.
However, while our BNCDE presents a first and important step toward uncertainty-aware medical decision-making with neural networks, our method also has limitations any other. We discuss them in the following:
(i Our BNCDE remains a black-box method that is not directly explainable. While explainability is often desired in medical scenarios for reliable decision-making, this issue is not unique to our method but applies to all previously mentioned neural approaches (Lim et al., 2018 Bica et al., 2020 Li et al. 2021 MeInychuk et al., 2022] Scedat et al. 2022). As such, we suggest a careful use of our method high-stakes decisions. Instead, we recommend the use of our method in automated or routine tasks. For example, there is a growing number of smartphone apps with digital health interventions (Murray et al. 2016), which could benefit from personalization through our method. Further, we argue that recognizing highly non-linear, complex, and inexplicable patterns in data is what neural networks are designed for; this, of course, contradicts direct explainability.
(ii) As reliability is of high importance, we emphasize that tuning the diffusion hyperparameter has noticeable impact on the approximation quality of the latent neural SDEs. In particular, during our experimental studies, increasing the diffusion constant directly led to stronger conservatism in the predictive predictive estimates.
(iii) Training our BNCDE involves solving high-dimensional stochastic differential equations. This a numerically challenging task and is the computational bottleneck of our method. However, we emphasize that, validated in Supplement the training time scales only linearly with the dimen- sion of the latent neural SDEs. This is not different from the scaling of any other fully connected neural network. Furthermore, we emphasize that, once the network is trained, inference for new
14

---

Published as a conference paper at ICLR 2024
patients can be achieved in seconds. Hence, this makes our method suitable for deployment decision support tool in clinical settings.
(iv) Our work focuses approximate Bayesian inference for neural CDEs of moderate size as countered in medical practice and has not been designed for very high-dimensional settings (c.g., inference on neural CDEs with billions of parameters). While such high-dimensional settings out of scope for this work, uncertainty quantification is typically needed for moderately sized (rather than very large) datasets and neural networks in practice. The reason is that, for moderately sized datasets, we cannot expect to learn the true data-generating mechanisms. Instead, we expect that there is uncertainty how medical treatments affect patient health and, to deal with this, it is im- perative in medicine to make treatment choices by incorporating the underlying uncertainty (Banerji et al. 2023).
(v) Due to the highly complex structure in electronic health records, sufficient data is required to obtain robust estimation results. However, we are optimistic that this will not be an issue in the near future due to the increasing amount of available data and data sharing agreements (e.g.. Johnson ct al. 2016, Pollard et al. 2018 Hyland et al. 2020 Thoral et al. 2021 Rodemund et al. 2023
(vi) As with any other method for treatment effect estimation, ours relies upon mathematical assump- tions. Importantly, the above implementation of our BNCDE does not account for confounding bias. Confounding bias distorts the true causal relationship between treatment and outcome, stem- ming from the influence of an unaccounted-for variable. This occurs when a variable is associated with both the treatment and the outcome, introducing a potential source of bias. For instance, pa- tients with more severe health conditions may be more likely to receive treatment A, while those with better health status may be given treatment B (Vokinger et al., 2021). Consequently, when estimating the effect of treatment A from data, bias towards a worse outcome emerges due to the association between health condition and treatment assignment. Addressing confounding bias in discrete-timg -ire literature involves adjustments like inverse propensity weighting (IPW) (RMSNs (Lim al. 2018)) and G-computation (G-Net (Li et al. (02T)). However, these adjustments lead very high estimation variance. Other neural methods such as (Bica et al.] 2020) Melnychuk et al.) 2022 Seedat et al. 2022) try to reduce confounding bias through balanced representations. However, balancing is approach for reducing estimation variance and not for mitigating confounding bias (Shalit et al. 2017). Furthermore, it may even introduce an infinite data bias (Melnychuk et al.) 2024) (see Supplement N for a detailed discussion balanced representations). As in Vander- schueren et al. 2023). we therefore decided against balanced representations note, our BNCDE is generally compatible with balancing, standing, while estimation bias due to confounding is still an open issue, we emphasize that this is not unique our method but general problem in the literature for time-varying treatment effect estimation.
our main paper. Of we show in Supplement L Notwith-
Finally, we acknowledge that treatment effect estimation in the continuous-time setting challeng- ing and remains subject to ongoing research. Our BNCDE faces the same lack of proper adjustments as the only other neural baseline (Seedat al.] 2022) and does not offer unbiasedness guarantees. Nevertheless, we believe that our BNCDE is a major improvement over existing methods and direct step towards reliable decision-makinn medicine.
15

---

Published as a conference paper at ICLR 2024
B CANCER SIMULATION
The tumor data was simulated according to the lung cancer model by Geng et al. (2017), which has previously been used in (Lim et al. 2018 Bica et al. 2020 Li et al. [022] Melnychuk al.)
2022 Seedat et al.] [022]] Vanderschueren et al. 2023) In particular, we adopt the simulation setting by Vanderschueren
(2023), which includes irregularly sampled observations and avoids introducing confounding bias in the treatment assignment mechanism.
The tumor volume, which refer as the outcome variable in the main paper, evolves according to an ordinary differential equation:
K Y
dy
+ plog
Ct,CC
(arde + B, (
Ee
Y dt,
(13)
Chemotherapy
Radiotherapy
Noise
Tumor growth
where o is growth parameter, K is the carrying capacity, Oles Or and 3, control the chemo and radio cell kill, respectively, and 6, introduces randomness into the growth dynamics. The parameters were sampled according Geng al. (2017) and are summarized in Table The variables o and de are set following the work by Lim et al. (2018); Bica et al. (2020): Seedat et al. (2022). The time t is measured in days.
Variable
Parameter
Distribution
Value (H, o²)
Growth parameter
Normal
(7.00 x 10 7.23 x 10 30
Tumor growth
Carrying capacity
K
Constant
Radiotherapy
Radio cell kill
Or
Normal
(0.0398, 0.168)
Radio call kill
Br
Set to 8. 10 x or
Chemotherapy
Chemo cell kill
Normal
(0.028, 7.00 10
Noise
Normal
(0, 0.01.)
Table 2: Sampling details of parameters used in the tumor simulation model.
We make the following, additional adjustments as informed by prior literature:
As in (Lim heterogeneity between patients is introduced by modeling different subgroups. Each sub- group differs in their average response to treatment, expressed by the mean values of the normal distributions. For subgroup 1, we increase µ(a,)) by 10%, and, for subgroup 2, we increase (a(ce) by 10%.
2018 Bica al., 2020; Seedat et al. 2022 Vanderschueren al. 2023
Following Kidger al. (2020). add multivariate counting variable that counts how often each treatment has been administered up to specific time.
As in (Vanderschueren et al., 2023), the observation process is governed by an intensity process with history- intensity
[y ( - })
ç = sigmoid
(14)
where " controls the sampling informativeness, D 13 cm, and D is the average tumor diameter over the last 15 days. For our informative sampling experiments in Section and for the additional prediction windows Supplement]] we set 7 1. For our experiments in Supplement
K we increased the sampling informativeness to = 2.
Treatments are assigned according to either a concurrent or a sequential treatment arm (Curran et al., 2011). That is, patients receive treatment either (i) weekly chemotherapy for five weeks and then radiotherapy for another five weeks or (ii) biweekly chemotherapy and radiotherapy for ten weeks.
Patients are randomly divided between the two treatment regimes. For any patient in the test data, simulate both the factual outcome under the assigned treatment arm and the counterfactual outcome under the unassigned treatment arm.
16

---

Published as a conference paper at ICLR 2024
For training, validation, and testing, observed time window is 55 days with additional predic- tion window of up to days. For training and validation, we simulated 10, 000 and 1000 observa- tions, respectively. For testing, simulated the trajectories and the outcomes for 10, 000 patients under both treatment arms, respectively. We standardized our data with the training set.
17

---

Published as a conference paper at ICLR 2024
C TE-CDE
In this section, we give a brief overview the so-called treatment effect neural controlled differential equation (TE-CDE) (Seedat al. 2022). We present the main ideas the original TE-CDE model but refer to the original paper for more details. We build upon our notation to avoid ambiguities.
TE-CDE uses neural controlled differential equations predict potential outcomes in continuous time. For this, TE-CDE assumes that hidden state variable Ze E Rd is driven by the outcome history, covariate history, and treatment history. The control path is interpolated in continuous time from observational, irregularly sampled data. The hidden state variable Ze encodes information about the potential outcomes for a given treatment assignment. In (Seedat et al. 2022). TE-CDE trained with an adversarial loss term to enforce balanced representations (see Supplement NN for discussion on balanced representations). However, balanced representations are not the focus of our paper, and we thus omit them for better comparability.
As is common for neural CDEs, TE-CDE encodes the initial outcome variable Yo, covariates Xo and treatment Ao into the hidden state Zo through embedding network 1/S Rltd, - R The hidden state then serves as the initial condition of the controlled differential equation given by
Z / f5(Zz) dY, X,, A.).
E (0, and Zo = nYo, Xo, An)-
(15)
Eq. is the encoder network of TE-CDE, The last hidden state Z is then passed through a decoder to predict the potential outcome fr+ald/r (d/ (A)) for an arbitrary but fixed time window in the future. The decoder is given by the neural CDE
YT+alde: ~ (Za), Z, = L' f (Z)) TE (0, and Zo =o= (16)
where ng is the read-out mapping. In particular, for new patient *, TE-CDE seeks to approximate the expected value of the potential outcome, given the patient trajectory H;.. That is, TE-CDE approximates the quantity
EYp+ald(a---a) Hz-).
(17)
Reassuringly, highlight the following key differences between TE-CDE and BNCDE:
1. TE-CDE targets Eq. 17 and thus makes point estimate of the potential outcome at time Ft + A. In contrast, our BNCDE estimates the full posterior predictive distribution of the potential outcome.
2. TE-CDE directly optimizes the neural vector fields fÿ and fg of the CDE. Instead, our BNCDE parameterizes the distribution of the neural CDE weights with latent neural SDEs. This neural SDE is optimized shape the posterior distribution of the neural CDE weights given the training data. As such, the architectures of TE-CDE and our BNCDE are vastly different.
3. TE-CDE may employ MC dropout to approximate model uncertainty. However, as show in our paper, the tailored neural SDEs provide more informative model uncertainty quantification.
4. TE-CDE does provide estimate of the outcome uncertainty quanrification. On the other hand, our BNCDE incorporates outcome uncertainty through the likelihood variance Ei+ (see Supplement[] I
5. The training objective of TE-CDE is against the mean squared error (MSE). In contrast, our method optimizes the evidence lower bound (ELBO). The ELBO seeks to balance the expected likelihood under the posterior distribution with the Kullback-Leibler divergence between prior and posterior SDEs. Therefore, our loss naturally incorporates weight regu- larization.
18

---

Published as a conference paper at ICLR 2024
D NEURAL DIFFERENTIAL EQUATIONS
We provide a brief overview of (1) neural ODEs, (2) neural CDEs, and (3) neural SDEs.
Neural ODEs: Neural ordinary differential equations (ODEs) (Haber & Ruthotto 2018 Lu et al. 2018) combine neural networks with ordinary differential equations. Recall that, in a traditional residual neural network with t I...., T residual layers, the hidden states are defined
(18) where fe.() (-) is the t-th residual layer with trainable parameters 6 and Zo- X is the imput data. These update rules correspond to an Euler discretization of a continuous transformation. That is, stacking infinitely many infinitesimal small residual transformations, fz(+) defines vector field induced by the initial value problem
Z Zi-1 + fe. (Zz-1),
Z, 2. [-
Zo X.
(19)
Instead of specifying discrete layers and their parameters, neural ODE describes the continuous changes of the hidden states over an imaginary time scale. A neural ODE thus learns a continuous flow of transformations:
input X = Zo is passed through ODE solver produce output Y Zr. The vector field fo is then updated either by propagating the error backward through the solver or using the adjoint method (Chen et al. 2018).
Neural controlled differential equations (CDEs) (Kidger al. 2020 extend the above architecture in order to handle time series data in continuous time. In neural ODE, all information has to be captured in the input Zo- X, i.e., at time hand, are able to process time series data. They can thus be seen as the continuous-time analog of recurrent neural networks. For path of covariate data X, E Rd== & E [0, T). a neural CDE given by an embedding network n (-), readout mapping 1i (-), and neural vector field fo that satisfy
Neural CDEs:
Neural CDEs, on the other
Y \zn)(ZT), Z, S// fo(Z.. dX,, EE (0, and Zo 7 (Xo),
(20)
where Z, E Rd. and ZZ, E Rd,xd, The integral is a Riemann- Stieltjes integral, where fz s) dX, refers to vecror multiplication. Here, the neural differential equation is said to be controlled by X,. Under mild regularity conditions, it can be computed as
Z z= L.' fo( dx. ds ds, EE (0,
(21)
Computing the derivative with respect to time requires a representation of the data stream X, for any E [0, T.. Hence, irregularly sampled observations ((to, 2o), (1).2....... (fn, = (f. x.,)) need to be interpolated over time, yielding continuous time representation X₁. The neural vector field fa. the embedding network vig. and the readout network 7/ are then optimized analogous to the neural ODE, For data interpolation. Morril et al. (202T) suggest different interpolation schemes for different tasks such as rectilinear interpolation for online prediction.
Neural SDEs: Neural stochastic differential equations (SDEs) (Li et al. 2020 learn stochastic differential equations from data. A neural SDE is given by drift network Is and diffusion network og that satisfy
We S. go(wars) 9) ds I' "olwws 5) dB., t 0 and wo Es
(22)
where B, a standard Brownian motion and Eo NVV,, .8) is the initial condition.
For two SDEs with shared diffusion coefficient, their Kullback-Leible (KL. divergence can be com- puted on path space. This makes neural SDEs particularly suited for variational Bayesian inference.
Let 9 and h be the drifts of the variational neural SDE and the prior SDE respectively. Let further-
more a (wr, t) be their shared diffusion coefficient. Tzen & Raginsky (2019) Li et al. 2020 show that the KL-divergence on path space then satisfies
DKL. (9% (w) p(W]).))) = Bagg(oo.n) Lo' (1o((es))) h(ws, 8) s))/o(ws, 8))² ds
(23)
19

---

Published as a conference paper at ICLR 2024
Of note, even for constant shared diffusion parameter a (wr, t) variational posterior distribu- tion parameterized this way can approximate the true posterior arbitrarily closely using a sufficiently expressive family of drift functions 9 (Tzn & Raginsky 2019 Xu et al. 2022)
20

---

Published as a conference paper at ICLR 2024
E MOTIVATION FOR USING NEURAL CDES (AND NOT NEURAL ODES)
In the following, we provide a discussion why we used neural CDEs instead of neural ODEs to capture patient observations and to model future sequences of treatments for personalized decision- making. We also highlight the main differences to CF-ODE (Brouwer et al. 2022), which builds on neural ODEs to estimate treatment effects in a different setting than ours.
Why did we opt for neural CDEs instead of neural ODEs?
Neural ODEs can be thought of as residual neural networks with infinitely many, infinitesimal small hidden layer transformations (see Supplemen()) Thus, neural ODEs are designed to describe how deterministic system evolves, given its initial conditions (e.g., a patient's initial health condition, an initial treatment decision). However, it is impossible to change these dynamics over time, once the neural ODE is learned. This is different from neural CDEs, which adjust for sequentially incoming information (Kidger et al.] 20200 Morrill et al., 2021). Further, neural CDEs make use of latent representations, which are able to capture non-linearities and complex dependencies in the data.
Why is single neural ODE suitable for task?
There are clear theoretical reasons why using a single neural ODE to model the dynamics of (Y, Xe, At) is not suitable. If we used neural ODE to directly model the dynamics of (Y,, Xt, At), this would mean that we learn an ODE that describes the deterministic evolution of (Y,, Xx,AA) At). given the initial conditions (Yo, Xo, Ao). In particular, we would then make the following assump- tions that would conflict with our task:
(i) The health conditions of a patient are completely captured in her initial state (Yo, Ao). The initial state is the only variable that influences the evolution of (neural) ODE. However, we want the patient encoding to be updated over time.
(ii) Given these initial conditions, the outcome and covariate evolution are deterministic. It is not possible "correct" an ODE future point in time to account for the randomness of future observations. However, we want to take into account the information from the observations in the future when they are measured.
(iii) The treatment assignment plan is fixed a priori and cannot be changed. That means, we would assume that sequences of treatments evolve like deterministic process. This makes modeling of arbitrary future sequences of treatments impossible.
(iv) Future observations cannot change the dynamics of the system and, importantly, cannot update beliefs about the potential outcomes. In particular, future observations cannot cor- rect the posterior distribution that captures the uncertainty. However, we intend to update the posterior distribution when patient information is recorded at later points in time.
How our architecture based neural CDEs address the above issues?
Our architecture solves all the above issues. The reason is that neural CDEs have much larger flexibility that is beneficial for our task:
(i) Neural CDEs allow for patient personalization at any point in time t, and do not assume that all information is captured at time O.
(ii) Future observations (Y,, Xx, Ar) can update the neural CDE. This crucial, because not all patient trajectories follow the same deterministic evolution.
(iii) Using neural CDE in our decoder, we model arbitrary future sequences of mul- tiple treatments.
(iv) Incoming observations after time t = can update our beliefs. This makes the neural CDEs particularly suited for the Bayesian paradigm.
Importantly, the hidden states Z, in a neural CDE also allow us to model non-linearities in the data, reduce dimensionality, and capture dependencies between observed variables YY, In summary, using neural ODEs to directly model the dynamics of (Y,, X,, A.) has significantly lower modeling capacities, limits patient personalization, and imposes prohibitive assumptions on the real world. To this end, argue that it is difficult if not impossible to adequately model the input
21

---

Published as a conference paper at ICLR 2024
through a joint neural ODE in our task, and, remedy, opted for approach based on neural
CDEs instead.
Of note, the Counterfactual ODE (CF-ODE) Brouwer al., 2022 leverages neural ODEs for treatment-effect estimation in a different setting to ours. In the following, we demonstrate why CF-ODE is not applicable our setting. We highlight the limitations of this method and clarify important differences to our BNCDE:
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 sequence of multiple treatments the future.
(ii) Accordingly, CF-ODE builds upon static identifiability framework. Here, adjusting for the measurements of the patient history at the time of treatment is assumed to be suffi- cient to adjust for all confounders. However, the static identifiability form CF-ODE does not hold true for time-varying settings as ours.
(iii) CF-ODE directly models the latent states in the decoder with neural SDE. In contrast, our neural SDEs model the distribution of the neural CDE weights.
(iv) All patient information in CF-ODE is contained deterministically in the initial state of the ODE. Hence, all uncertainty needs to be incorporated a deterministic, initial state. On the other hand, our BNCDE captures patient information in the random variable Z7 = Zo- Here, the Monte Carlo variance is propagated through the encoded trajectory Z1o.7). ac- counting for uncertainty at each point in time t E [0, T]. Therefore, model uncertainty is directly incorporated into the patient encoding.
(iv) The use of gated recurrent unit Cho et al. (014) in the CF-ODE encoder implies regular sampling of the patient trajectories. This assumption is crucial violation of med- ical 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.
(vi) The single treatment decision in CF-ODE limited binary treatments. Instead, our BNCDE can deal with any treatment.
22

---

Published as a conference paper at ICLR 2024
F VARIATIONAL APPROXIMATIONS
We provide a detailed derivation of (1) our approximate posterior predictive distribution and (2) our evidence lower bound.
Notation: We slightly abuse notation for better readability. We denote the expectation with respect to the stochastic process wi on path space as
f () = Bruporp)))
(24)
The analogue applies for wr.
Approximate posterior predictive distribution: In the Bayesian framework, the posterior predic- tive distribution of the outcome of interest Yi++al((z f +-j given observed inputs hț. and training
data H is the expectation of the outcome likelihood under the weight posterior distribution. Hence, we approximate the posterior predictive distribution via
p (Vo+slde .e +A)) h
(25)
p(Ye+alde x +
(1) H)
(26)
/ P(Yesald Mi. dp d w|u H) dp(w),,7 H)
(27)
/ NVVe+ald F Wr.
(w).)))
(28)
= -Fng(woo
[ ( p(Vo+aldrr.. +4)) |1. 4 worr hz.
(29)
Evidence lower bound: Recall that, in notation, training
sample is the target.
144+ = M(T+A)- U (1b4a) where M(+4))- is the training input and Bf4s
Our objective is an ELBO maximization, which we derive via
logp(u/+4 -
(30)
=log
plvita M(8+2) op0,
(31)
= =log
piui+al (1.) dp(
(32)
=log|
=
p( )
(w)o do (w)),
(33)
/ logpluh44 | (1
Lfd ((())))))))
- log 18(410001)
- log
dg (2)) do
(34)
p(
p(
logp(u Wj, Wo, h(F+a) }}}
= Egg (2 (..)
E (g
- -DKLl9 (w)p.)) p(W(
- DKL bas (w)).:)) || p(w
(35)
where Eq. 31 follows from the standard pairwise independence assumption of inputs and weights, Eq. 32 due to the independence of the OU priors, and Eq. 34 follows from Jensen's inequality.
The ELBO objective is maximized for (a))s)) = p(Wjo.a) h a w).)1) and
of (1))))))) p(w70,7 4)) We this by substituting in the true weight posteriors
23

---

Published as a conference paper at ICLR 2024
into Eq. 34
This yields
/
log logwlubis M10+0)- (10,)), w(o,))) - log p(((1o.a) ( I wo.F1)
p(w)o.)))
p(w1o,b) ( dp(ăto.a) (0,1) dp(wyo, m2+2)
log
(36)
p(w(o.F1))
plui Mi++)) wI0, (1. dp(01o.)).00..)) h24s))
/ log
(37)
p( 11 M(++))- vi(a)
= f log
PlU (10.), w0.89) Min+a)
de ()) (1+ta)
(38)
p( W M(B+2) 38++a)
f log
dp(#)o,a)) 410.1
(
W(0.1 h/0+ Vo+4
(39)
M ) "p( 1 = log plyb+ llggp(vt+
(40)
log p(ui M(0+)));;
(41)
follows again from prior independence of input data and weights and the fact that h #4 (D4A)- U (.
where Eq.
24

---

Published as a conference paper at ICLR 2024
G NUMERICAL COMPUTATION OF THE EVIDENCE LOWER BOUND
The objective of our BNCDE is the evidence lower bound (ELBO), which approximated after the forward pass of the numerical SDE solver.
Recall that, for an observation i, the ELBO is given by
log logr ((++))) 2 E Bog E (
log p(uf++ (10,0), War, Mo+2y-)BB })]
DRL. D (H)oo)) |l p(
DRL (a (1)) p( }]- (42)
We approximate the expectations with respect to o (4(0.)) and (w)))))) with Monte Carlo sam- of the weight trajectories from the numerical SDE solver. That is, compute
ples wfo.21 and w [..A] the expected likelihood via
{ B ()))
logp(vi-s (10. w0,F) Mi +)))))
Byg (o))))
(43)
Σ
Σ
logp(yh+4 (ومر) aar h M
(44)
0.84
where 10, 8 and [0, A are grid approximations of [0, 0' and [0, A], respectively.
The Kullhack-Leibler (KI.) divergences between priors and variational posteriors on path space can also be computed through Monte Carlo approximations. Here, focus on the KL-divergence in the encoder component; the decoder part follows analogously, Recall that we seek to compute the KL-divergence between the solutions of the SDEs
(g((r dt + a dB,
dw,
(45)
hh (w,) dt + a dB,.
Previously, Tzen & Raginsky (2019) and Li et al. (2020) have shown that for two stochastic differ- ential equations with shared diffusion coefficient a, the KL-divergence on path space satisfies
D (w)) p(W),, Egg(w my) [,,0 ((og(we)) - (w.))/oo) dt
(46)
Therefore, it can be approximated with Monte Carlo samples from the numerical SDE solver. That is, we compute
Eug(o [.'" ((95(we) he (wrr))/))²
1
Σ
1
((o((W)) (
(w(())) dt
10.09
W D..
2
(47)
where Ju (-) dt denotes the integral approximation from the SDE solver. In other words, for each Monte Carlo particle j. we pass the KL-divergence term as additional state through the SDE solver and then take the Monte Carlo average at time J..
25

---

Published as a conference paper at ICLR 2024
H IMPLEMENTATION DETAILS
In the following, we summarize the implementation details of our BNCDE and TE-CDE. Experi- ments were carried out on 1x NVIDIA A100-PCIE-40GB. Training of BNCDE and TE-CDE took approximately 34 hours and 22 hours, respectively.
The hyperparameters of TE-CDE training and testing. the dropout probabilities in the prediction head were set to = 0.1. As in our BNCDE, we included the time covariate an input to TE-CDE, this improved performance.
taken from the original work (Seedat et 2022 For both
We optimized the hidden layers of the neural SDEs drift over {16, (16, 16), (16, 64, 16), (16, 64, 64, 16), (16, 64, 64, 64, 16)} and the diffusion coefficient a over f0.01, 0.001, 0.0001). We used the same neural SDE configuration for both the encoder and the decoder. Importantly, we no- ticed that higher diffusions generally lead to more conservative estimates of the posterior predictive distribution, which may be desirable in medical practice. Our performance metric was the ELBO objective on the validation set. We only optimized the hyperparameters for the prediction window A and then used the same configurations for A = 2....... 5. For better comparability, our BNCDE had the same hyperparameter specifications as TE-CDE where possible.
For both our BNCDE and TE-CDE, we increased the learning rates of the linear embedding and prediction networks compared to the original TE-CDE implementation for more efficient training (Kidger et al.] (020)). Further,
used cubic Hermite splines with backward differences in the neural CDES, which have been shown to be superior to linear interpolation (Morrill et al., 2021). As we want understand the true source of performance gain, we removed the balancing loss term in TE- CDE for comparability. However, we emphasize that the balancing loss term could easily be added our BNCDE, as we show in our extended analysis in Supplement
Table 3 lists the selected hyperparameters.
Component
Hyperparameser
(ouns)
Scedat
Batch size
64
64
Oplimizer
Kingra
2015
Max. number spuchs
500
500
General
Pañence
10
10
MC samples traiming
10
10
sumples prediction
Hilden state size d
Inierpolatinn mathed
Cubiic Hermite splinos
Culic Hermite splines
Differential equation selver
Solver
Malcsaaroyama
Faler Adaptive
Step size
Adapeiee
Leurning layers Output activation
10
10
Embodding
Leuming rate
10
Hidden layers dimessions Hidden activations
Neural CDFs
(128, 128)
128)
ReLU
ReLU
Output actication
tanh
tanh
1 Leainng rate
10
Hidden layers
Neural SDEs
Hidden activations
Hidden dimensions
Rel.U
(16, 64, 64, 16)
Outpul activation
Diffusion coeflicient
0.001
Leaninng
Prediction network
Hidden layers
Chipon activation
sc0plos)
Dropout probability
Leanning layers Output activation
Inteuityy Conly Supplement
sigmoid
sigmoid
Leamiing rate
Ralancing network (only for Supplemen
Hiddza loyers activation
Balancing pppeparamettrr
Table Hyperparameters of BNCDE and TE-CDE. Hyperparameters of TE-CDE were taken from Seedat et al. 2022). Our BNCDE used the same hyperparameters where possible.
26

---

Published as a conference paper at ICLR 2024
I EVALUATION OF OUTCOME UNCERTAINTY
We now assess the quality of the estimated outcome (aleatoric) uncertainty. In our BNCDE, out- come uncertainty captured in the likelihood variance M++ via
YA+a (+))))))++))) "
H+) ~ NNUT+A, ET+A)
(48)
Outcome uncertainty is due to the randomness that is inherent to the data-generating process. It is irreducible in the sense that, even if we had infinite training data and hence zero model uncertainty, the potential outcome would still be random.
For given patient history Hj and future sequence of treatments ( +A) computing the true out- come uncertainty is intractable. Instead, we can check whether our BNCDE produces informative estimates of the outcome uncertainty. Variational Bayesian methods may heavily over- or underes- timate the outcome uncertainty, yielding uninformative estimates. Hence, through our analysis, we provide insights that our BNCDE produces meaningful estimates outcome uncertainty for differ- ent patient histories and future sequences of treatments. In particular, we show that there is direct relation between estimated outcome uncertainty and prediction error. Intuitively. prediction errors should be larger for outcomes with higher outcome uncertainty. This means that if the potential out- come for particular patient history and future sequence of treatments subject to a large variance, the prediction errors should increase accordingly.
The results in Fig. show the relation between prediction errors and estimated outcome uncertainty for varying prediction windows. For each patient, compute the Monte Carlo average of the esti- mated outcome uncertainty ET+ as well as the average point prediction error. We notice that our BNCDE provides varying levels of outcome uncertainty for different patients and future sequences of treatments. Further, prediction errors are highly correlated with estimated outcome uncertainties. Therefore, our BNCDE generates meaningful estimated outcome uncertainties. Importantly. this is another key advantage of our BNCDE over approaches such as TE-CDE (Seedat et al.] (022) with MC dropout (Gal & Ghahramani 2016), which are not capable of estimating potential outcome uncertainties.
(a) 1
(b) 2
A 3
(d) A = 4
(e) A =5
Figure 6: We analyze the association between estimated outcome uncertainty and prediction error. The prediction errors are highly correlated with the estimated outcome uncertainties.
27

---

Published as a conference paper at ICLR 2024
J ADDITIONAL PREDICTION WINDOWS
We repeat the experiments from Section for the prediction windows A = and 5. We report the results Figures and respectively. The results are consistent with our main paper: (i) The credible intervals generated by our BNCDE remain faitffa and sharp. In contrast, TE-CDE produces unfaithful Crls. (ii) The point estimates of the outcome by our BNCDE are more stable under increasing noise in the data generation. (iii) The model uncertainty of our BNCDE is better calibrated.
CH mninos on emspttil cowerage
0.00
Creditie Laternl wahth
0.35
TE CDE
BNCDE
0.30
-0.02
TECDE
DNCDE
0.00
Overecddddout
-0.00
Conservalive
0.00
0.85
0.90
0.97
0g8
0.90
0.05
0.900
0.97
0.95
0.90
Posterice prodictive croditble Interval
Posterioe prodictive creditble Interval
(a) CrI coverage
(b) CrI width
10a
MSE Moute Cacoo eslannee
0.033
10
0.015
10-2
0.120
TR-CDE ENCDE (oun)
10-
Normalied
10-1
0.005
30-1
TE CDE BNCDE
10-1
0.005
Rraadom deferal
/arter]
Proportion of most uncertade somples withheld
(c) Error point estimates
(d) Comparison model uncertainty
Figure 7: We repeat our experiments from Section Reported are the results for the prediction
window A
28

---

Published as a conference paper at ICLR 2024
0.25
0.00
engleial covttng)
0.20
Crodable inservl wilch
0.02
0.35
TECDD
0.00
--- RNCOR (ouns
0.00
TR-CDE
-0.0%
BNCDE (uuns
0.05
Openceefident
0.08
Consevvatove
Posterior predictive credible interval
Pustetior predictive crediblee intervol
(a) Crl coverage
(b) Crl width
10-1
0.015
MSE
00-2
MSS of Monte Cado 2288
0.004
TE:
(ouns
0.012
30
0.001
30-1
TE-CDE BNCDF (uum)
0,009
Randum
10-7
0.02
0.04
0.05
0.08
0.0
0.2
0.4
0.6
0.8
1.0
v /Varie)
Proportion wacertain somples withheld
(c) Error in point estimates
(d) Comparison of model uncertainty
Figure 8: We repeat our experiments from Section window A 5.
Reported are the results the prediction
29

---

Published as a conference paper at ICLR 2024
K EXTENSION TO INFORMATIVE SAMPLING
We show that our BNCDE is superior over TE-CDE, even when used together with the informative sampling framework from Vanderschueren et al. 2023). The motivation for this that the obser- vation times of an outcome may themselves contain valuable information about the health status of a particular patient. For example, patients in a more severe state of health are likely be visited more often, leading to a correspondingly higher observation intensity. To address such informative sampling, Vanderschueren et al.] (202)) propose to weight the training objective with the inverse observation intensity. To achieve this, they introduce an additional prediction head for TE-CDE, which estimates the observation intensities. The inverse estimated observation intensities are then used to weight the training objective. which has been shown to improve performance in outcome estimation. This extension of TE-CDE is called the TESAR-CDE.
Changes to tumor growth simulator: We extend our simulation setup follows. In the tumor growth simulator (see Supplement B). the observation times are determined by an intensity process with intensity function
Ç sigmoid [ ( D (
2 1
(49)
where y controls the sampling informativeness, D 13em and D the average tumor diameter over the last 15 days. For higher values of % the total number of observed outcomes reduces but the informativeness of each observation timestamp increases. We later change ? = 2.0 to introduce informative sampling.
Changes to our BNCDE: We can easily incorporate the inverse intensity weighting into our BNCDE to account for informative sampling. The final hidden representation ğ of patient is additionally passed through second prediction head n to estimate the observation intensity via
C++A The (1))
(50)
The evidence lower bound (ELBO) for patient observation is then weighted with the inverse estimated observation intensity, i.e.,
ELBO' /S844
3000//$444
(51)
Experiments: We repeat the experiments from Section 5 where we increase the level of infor- mativeness y from 1.0 2.0. We focus on the prediction window A with inverse intensity weighting enabled for both methods.
We benchmark our BNCDE with TESAR-CDE. For RESAR-CDE, we enable MC dropout in the outcome prediction head vi but not in the intensity prediction head v/s. We provide implementation details of nh. in Supplement For both TESAR-CDE and our extended BNCDE, optimization of n is performed as in the original work by Vanderschueren al. (2023). That is, the error in the estimated observation intensity is not propagated through the whole computational graph but only through nh
Results: Fig 9 shows, that in the informative sampling setting, method remains superior. Our BNCDE (i) produces more faithful and sharper credible interval approximations: (ii) generates more stable point estimates of the outcome under increasing noise; and (iii) has better calibrated model uncertainty. In sum, this demonstrates the effectiveness of our BNCDE even for informative sam- pling.
30

---

Published as a conference paper at ICLR 2024
0.00
0.30
TESAECDE ENCEE (mur))
enpirioal coeetage
Uvtrosalidaud
0.25
Conservntive
Creditee inteval width
0.20
TESAR CDE
(ouns
0.20
1000
0.05
-0.0
0.00
0.85
0.90
0.97
3:88
0.85
0.90
0.97
0.85
0.90
predictive
predictive
(a) CrI coverage
(b) Crl width
0.040
TESAR-CDE BNCDE (omzs)
MSE Motte Catto estansse
0.035
10
0.000
0.075
Normalbedd
0.070
10-1
0.015
TESAR CDE
10-1
INCDE (mum)
Rendom deferal
10-7
0.02
0.04
0.05
0.08
0.30
0.0
0.2
0.4
0.5
10
Proportion of most uncertala samples withheld
(c) Error in point estimates
(d) Comparison model uncertainty
Figure 9: We repeat the experiments for settings with informative sampling. For this, increase the informativeness parameter in the tumor growth simulator to = 2.0. Following (Vanderschueren al. 2023), weight the training objective with the estimated inverse observation intensity. We benchmark our extended BNCDE against TESAR-CDI (Vanderschueren et al. 2023). for which we enable MC dropout the outcome prediction head.
31

---

Published as a conference paper at ICLR 2024
L REAL-WORLD AND SEMI-SYNTHETIC DATA
We provide further insights that our BNCDE remains reliable on different data sets. For this, we use the MIMIC-extract (Wang al. 2020). In particular, we use the same preprocessing of MIMIC-III (Johnson et al.) 2016) as in earlier studies (Melnychuk et al. 2022). The MIMIC-extract provides intensive care unit data at irregularly recorded times measured units of hours.
We also examine whether our BNCDE is compatible with balanced representations. While it is often used the literature for mitigating confounding bias (Bica et al. 2020 Melnychuk et al. 2022 SSeddtt et al. 22022) it actually reduces estimation variance (Shalit et al. 2017). e.g., due to overlap violations, which is particularly relevant the time-varying setting.
L. EXPERIMENTAL. SETUP
Real-world data:
We consider patients with 30 observation times. For each patient, we select time- varying patient covariates heart rate, red blood cell count, sodium, mean blood pressure, systemic vascular resistance, glucose, chloride urine, glascow coma scale total, and hematocrit and positive end-expiratory pressure set. Additionally, we use gender and age as static covariates. We have additional access discrete treatments vasopressors and mechanical ventilation. We seek predict the outcome diastolic blood pressure. The remaining covariates serve as unobserved confounders.
We 60/15/25 split for training, validation, and testing on N 000 observations.
Semi-synthetic data generation: In order to generate treatment assignments and outcomes, we proceed as follows:
(i) We only consider patient histories with 30 to 50 observation times. For each patient, we then select time-varying covariates (1) heart rate, (2) red blood cell count, (3) sodium, (4) mean blood pressure and (5) systemic vascular resistance, and static covariates (6) gender and (7) age. Impor- tantly, these covariates are observed irregularly at different hourly levels. Let K 7 be the number of covariates. We write X (k)) when referring to covariate & E 21, .......... 7) of patient at hour j.
(ii) We forward-fill each patient covariate in time. Note that this is done only for data generation, for training and testing, as our model is designed handle irregularly sampled data.
(iii) Then, we simulate outcomes of patient for a burn-in period of H = O...., 4 hours via
H
Σ
1
tanh
By X is (b)) vin + ct (v),
(52)
h +1
2
h=0
k=1
where vi, for all i, Bx NO, 1) are randomly sampled covariate coefficients, and ef, (y) (0 0.01²) are i.i.d. outcome noise terms.
(iv) Next, we sample the treatment assignments and patient outcomes at hours j > 5. At each time ty. the binary treatment mechanism depends on the patient history in order introduce confounding. Specifically, the treatment probabilities are parameterized by
pi. (a) sigmoid ttanh
Σ
3xx: (k)) - tanh (vi,..) + ci (a)())
(53)
K
k=1
where et, (a) 0.00) are i.i.d. noise terms. Treatment assignments are then sampled via
at Ber(pii(a))
(54)
We simulate outcomes via U
2 Ba Baxi (x)
vi, cos (vi,ci,) + (),_,())
tanh
(55)
h
2
(v) Since our BNCDE is designed for irregularly sampled data, we do not use the forward-filled covariates but the original, irregularly sampled covariates for training and testing. Further,
32

---

Published as a conference paper at ICLR 2024
apply an observation mask to the outcomes with observation probabilities
sigmoid vi,)
(56)
Further, we mask all outcomes for which no covariates are observed in the future.
We use 60/15/25 split for training, validation and testing N 32, 673 observations.
I.2 BNCDE WITH BALANCED REPRESENTATIONS
For both the real-world data and the semi-synthetic data, our BNCDE is trained with balanced rep- resentations (Bica al., 2020 Melnychuk et al. [022] Scedat et al. 2022) in order reduce finite- sample estimation variance. We provide discussion on balanced representations in Supplement N To this end, we follow the implementation as in (Seedat al. 2022). That is, we add second prediction head 7% as in (Seedat et al. 2022) to make the hidden representations Ze and Z, non- predictive of the administered treatments. The treatment prediction head vg estimates the probability of a future treatment P², (a) at timestamp Tj. For a sequence treatments with treatment decisions at irregular timestamps [TI...., Jm the treatment prediction head is trained to maximize the binary cross-entropy given by
1 ar log(p", (a)) + (1 ar,) log(1 1op5(1--,.11 p (a)).
BCE
Σ 3=1
(57)
The overall objective is then to maximize the weighted sum of evidence lower bound and binary cross entropy. That is, we maximize
ELBO BCE,
(58)
where a is hyperparameter.
We choose a 0.01 and a single linear transformation for n2. As proposed the original TE-CDE paper (Seedat et al. 2022). also use balanced representations for our baseline, TE-CDE with MC dropout. Importantly, we only use MC dropout in the outcome prediction head, not in the treatment prediction head. For both our BNCDE and TE-CDE, all hyperparameters are kept as in the main paper (see Supplement
L.3 RESULTS
For both datasets, we repeat our main experiments on the reliability of the posterior predictive dis- tributions for A = 1 hour ahead prediction.
(i) Real-world data: Fig. 10 shows the empirical coverage for the posterior predictive credible intervals along with the width. Here, our BNCDE tends to have wider intervals. However, we clearly see that estimates are much more reliable. Importantly, the credible intervals from our BNCDE coincide very accurately with the frequentist outcome quantiles. On the other hand, TE-CDE with MC dropout completely fails to obtain reliable credible intervals. In fact, the Cris from TE-CDE with MC dropout out are way too overconfident (i.e., too narrow), which could lead harmful decisions. In sum, our BNCDE therefore strongly outperforms the baseline on real-world data with balancing.
(ii) Semi-synthetic data: Fig. shows the empirical coverage for the posterior predictive credible intervals along with the CrI width. Further, we increase the variance in the outcome noise e, from 0.01 to 0.1 and repeat our robustness studies. Both methods generate estimates that are conservative and have sufficient predictive coverage. In particular, all of the reported posterior predictive credible intervals of our BNCDE contain all of the outcomes in the test set. We noticed that using balancing on this semi-synthetic dataset increased the conservatism of our estimates. This is desirable and shows in particular why our BNCDE is compatible with balancing. Further, compared with the TE- CDE baseline with MC dropout, our BNCDE has much sharper credible intervals, while providing more coverage than the baseline. Finally, we see that the point estimates of our BNCDE again much more resistant to increasing noise in the outcome distribution. In sum, this demonstrates the effectiveness of proposed BNCDE.
33

---

Published as a conference paper at ICLR 2024
Chefiile Itherll vddd
C2l artiattee erqueeeee
644
2.
-4- DNcmr (oar)
Overostiinal
BNCDE (wan)
0.0
-01
0.56
0.56
00:00
0.98
1.00
0.55
0.95
0.9
6.68
3.00
Prstenor produttoon interva
Protemsrr producton croiille
(a) CrI coverage
(b) Crl width
Figure 10: Real-world data: repeat our main analysis of the reliability of the posterior predictive credible intervals with balanced representations. We benchmark extended BNCDE against TE- CDE with MC dropout in the outcome prediction head. We include balanced representations as in the original work (Seedat al. 2022). We report the mean and standard deviation over five different prediction runs.
0.14
6000
0.10
09040
6006
0.66
TLC BNCIDE
-000
UNCUU
-008
TE-CD6 BNCDE
0.14
Cyeoedilas
610
0000
Conemarhe
0.50
11.66
I437
200
1123
1496
0.00
0.50
0002
11.04
2100
11.55
3.30
Pastecior proiotive coolidde inteval
Photense prodictive crollile nntervi
(a) CrI coverage
(b) Crl width
(c) Error point estimates
Figure 11: Semi-synthetic data: We repeat our main analysis of the reliability of the posterior predictive credible intervals with balanced representations. Further, we increase the variance of the outcome noise in Eq. (55 from 0.012 to 0.12 in the test data. We benchmark our extended BNCDE against TE-CDE with MC dropout the outcome prediction head. We include balanced represen- tations as in the original work (Seedat et al. 2022). We report the mean and standard deviation over five different prediction runs.
34

---

Published as a conference paper at ICLR 2024
M RUNTIME
Estimating treatment effects over time from electronic health records is notoriously difficult and may require complex neural architectures. Importantly, this is a common issue the existing baselines well in our method. Below, we provide discussion of why we expect that our runtime is still reasonable in practical applications.
Arguably, the computational bottleneck in training our BNCDE is solving the neural SDEs. If the neural CDE networks are very large, high dimensional SDE has to be solved in the forward pass of the training. To better understand this, we examine empirically how the training time of our method scales with the dimension of the SDE. Thereby, we validate that training time only scales linearly with the dimension of the neural CDE.
For this, we train our BNCDE for a single epoch on the cancer simulation data. As detailed in Supplement B, there are N 10, 000 training observations and we train on batches of size 64. All experiments were carried out on x NVIDIA A100-PCIE-40GBB
The
Stmder viggee seasd CDE modic
Stabe Marre Clrll
(a) We increase the number of hidden (b) We increase the number of Monte layers in the encoder neural CDE. Run- Carlo samples in the SDE solver. Mul- time increases approximately linearly. tiple samples can be used at little cost.
Figure 12: Runtime evaluation of our BNCDE.
First, we use the same parameterization as provided in Supplement H but change the number of hidden layers in the neural CDE of the encoder. That is, we increase the number of hidden layers of size 128 in the neural CDE and, hence, the dimension of the latent neural SDE. Fig. [2]ff)) reports the average training time in minutes along with the standard deviation over three seeds against the number of weights neural CDE of the encoder. The training time only increases linearly with the number of weights, which mimics the usual computational complexity for any fully connected neural network.
Second, investigate how training time scales with the number of Monte Carlo samples of the weight trajectories w|o,2] and (/..]]- Both the approximations of (i) the expected likelihood and (ii) the Kullback-Leibler divergence in the evidence lower bound rely on Monte Carlo samples of Wio.r] and W1o.A] (see Supplement G). Hence, training our BNCDE stabilizes for higher number of Monte Carlo samples of W[o.r] and W[o,A]- In order to gain insight how training time of our BNCDE scales under multiple draws Monte Carlo draws of Wj,,7 and Wlo at once, we use the same hyperparameters as in our main experiments (see Supplement(H]) but we increase the number of Monte Carlo draws from to 100. Fig. [2] (right) reports the average training time in minutes along with the standard deviation over three seeds against the number of Monte Carlo draws for both wj0,21 and W|o,A]- Training time barely increases, which shows that Monte Carlo variance in the training can be reduced at little cost.
As side note, we emphasize that neither increasing the dimension of the latent neural SDEs nor the number of Monte Carlo samples affects the discretization error of the SDE solver. our ex- periments, we found that adaptive solvers suitable, e.g., the adaptive Euler-Maruyama scheme. Following (Chen et al.) 2018), we argue that discretization errors are in practice not an issue, as modern solvers provide guarantees accuracy.
In sum, we conclude that there are no drawbacks in using very deep neural CDE networks that are specific our BNCDE. Still, is the case with other approximate Bayesian methods, there
35

---

Published as a conference paper at ICLR 2024
stability guarantees when training our method, that is, for learning very high-dimensional neural SDEs (c.g., in the billions). Rather, our work focuses on Bayesian inference for neural CDEs of moderate size encountered in medical practice.
36

---

Published as a conference paper at ICLR 2024
N DISCUSSION ON BALANCED REPRESENTATIONS
Some prior works on estimating heterogeneous treatment effects propose to learn balanced repre- sentations that are non-predictive of the treatment (Bica et al., 2020 Melnychuk et al. 2022] Seedat itt al. 2022). The idea is to mimic randomized clinical trials and reduce finite-sample error due to estimation variance (Shalit et al. 2017). However, for uanevagying treatment effects, this approach has several drawbacks:
Learning guarantees only exist in static (i.e., non con-time-varynng settings (Shalit et al. 2017).
Thus, balanced representations do not help with reducing bias due identifiability issues induced by time-varying confounders. For this, proper adjustment methods such as G-
computation prverse-rropensity weighting would be necessary (Pearl 2009 Robins & Hernán 2009).
2. Even in static settings, methods based on balanced representations impose invertibility as- sumptions the learned representations (Shalit et al.) (017). This is highly unrealistic in tare-varying settings, where representations not only incorporate information about the confounders but also about the full patient history.
3. Because invertibility is difficult to ensure, balanced representations may actually increase bias. We refer to Curth & van der Schaar (2021) and Melnychuk et al. (2024) for a detailed discussion this issue. Again, this is particularly detrimental in time mare-iariing settings.
In our main paper. we therefore decided not to incorporate balanced representations due to the rea- sons described above. This consistent with previous works in literature (Curth & van der Schaar 2021 Vanderschueren et al., 2023). Nevertheless, we emphasize that balanced representa- tions can be casily integrated into our BNCDE, see Supplement L
heuristic approach for reducing finite-sample variance. Importantly, in order to avoid confounding bias, proper adjustments such as G-computation or inverse- proseesttppesttty weighting needed.
In summary, balanced representations
37


Reviews:
[{"review_id": "ZCGWhQCxLC", "reviewer": "Reviewer_CatB", "summary": "This paper proposed to estimate the treatment effect and quantify the uncertainty over continuous time. To achieve this goal, a VAE-based Bayesian method was proposed to learn ODE (resp. SDE) over latent variables (resp. encoder parameters), as well as the posterior distribution of potential outcomes. An evidence lower bound was derived for optimization. Experiments on semi-synthetic data were conducted.", "questions": "Please see the weakness above.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The treatment effect over continuous time is a very important and interesting topic since it can model the instantaneous effect that is beyond the scope of traditional causal modeling such as Granger causality. Besides, this paper conducted a thorough experimental analysis regarding the uncertainty quantification. The paper is well-written and easy to understand.", "weaknesses": "This paper fails to discuss the motivation for learning latent representations, which was commonly used for unstructured data. However, the experimental setting in this paper is still structured data. In this regard, why do not directly model the ODE over (treatment, covariates, and outcome)? Moreover, the introduction of latent variable $Z$ may violate the unconfoundness assumption. Specifically, the implicit causal assumption behind latent representation is $Z_t$ affects $(A_t, X_t, Y_t)$ and $Z_t \\to Z_{t'} (t' > t)$ which also effects $(A_{t'}, X_{t'}, Y_{t'})$. In this regard, an unblocked path will open from $Y_t$ and $Y_{t'}$. Besides, since $Y_t \\perp A_t | X_t$, we should implement backdoor adjustment for $X_t$, which is necessary to eliminate the confounding bias. However, I fail to see such an adjustment in the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "zt8jNbfTYt", "reviewer": "Reviewer_UGox", "summary": "The paper introduces BNCDE which combines neural controlled differential equations with neural stochastic differential equations to model the time dimension. It leverages the Bayesian paradigm to account for both model and outcome uncertainty, allowing for uncertainty-aware treatment effect estimation in continuous time.", "questions": "Please see section Weaknesses", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The paper focuses on the crucial task of estimating treatment effects over time in the context of personalized medicine. \n\nThe incorporation of uncertainty quantification is a critical aspect of the proposed methodology. It allows for probabilistic estimates of treatment effects, which is essential for making informed decisions in medical contexts.\n\nThe problem setting appears to be innovative and unique.", "weaknesses": "1/ 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?\n\n2/ 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?\n\n3/ Experimental results are only on one synthetic data. Is there any real-life data or semi-synthetic data that is applicable to this model?\n\n4/ Page 7: \"we the maximize ELBO\"", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Rx4FldYV15", "reviewer": "Reviewer_6Zs2", "summary": "The paper proposes an estimation of the treatment effects with uncertainty in continuous time by exploiting Bayesian neural controlled differential equations. The proposed method (BNCDE) is based on an encoder-decoder architecture and allows for estimating posterior predictive distributions of the potential outcomes. In the proposed method the time dimension is modeled through a coupled system of neural-controlled differential equations and neural stochastic differential equations.", "questions": "-\tThe conclusion states: “(4) Our BNCDE is further fairly robust against noise”. It would be nice to be more explicit about what “fairly” means.\n\n-\tPlease discuss possible failure modes of the method. Considering that the aim is to apply it in the medical domain, this would be relevant for anyone trying to use it. \n\n-\tIn 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?\n\n-\tThe 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?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "-\tThe paper is clearly written and identifies the gap in the literature with clear motivation (estimating treatment effects in continuos time with uncertainty quantification). The paper is self-contained and easy to follow.\n \n-\tThe implementation details are well documented in both the paper/appendix and the provided code.", "weaknesses": "-\tConsidering 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. \n\n-\tThe 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "PNmfZ3Agf2", "reviewer": "Reviewer_kEvk", "summary": "This paper presents a Bayesian Controlled Neural Differential Equations (BNCDE) framework aimed at estimating treatment effects in time-series data within a continuous time. The approach involves employing a Bayesian Neural Network to model the drift function of the CDE, with the posterior distribution of the neural network weights defined by the solution of a stochastic differential equation.\nThe authors outline a training regimen for the framework by optimizating the Evidence Lower Bound (ELBO) in an end-to-end manner. \nOn the empirical side, the author includes experiments conducted on synthetic tumor growth data, with a focus on evaluating the uncertainty estimation quality of the proposed model. The results from these experiments are compared to non-Bayesian counterparts, highlighting the performance differences between the two approaches.", "questions": "1. In section 3.3, the author briefly talked about the 3 assumptions required for identifiability. To make it even more clear, the author can consider adding explicit reference or discussions regarding why those 3 assumptions can lead to identifiability. \n\n2. Since the potential outcome is modelled through the latent embeddings, does this affect the identifiability?\n\n3. The prediction also estimate the outcome uncertainty. But typically, this can potential lead to the a naive model with large uncertainty. I wonder the quantification quality for the aleatoric uncertainty.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "This paper proposes a treatment effect estimation framework for time series in continuous time, demonstrating a better uncertainty estimation compared to the non-Bayesian counterpart with MC dropout. I briefly checked the derivation, which seems to be correct. The presentation of method is clear in most places but there are still some ambiguities that requires further clarification.", "weaknesses": "There are two aspects that warrant further discussion, particularly in terms of its novelty and scalability. 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. Furthermore, 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.\n\nIn 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. An exploration of these aspects would contribute to a more thorough understanding of the framework’s practical applicability and limitations.\n\n[1] Li, Xuechen, et al. \"Scalable gradients for stochastic differential equations.\" International Conference on Artificial Intelligence and Statistics. PMLR, 2020.\n[2] De Brouwer, Edward, Javier Gonzalez, and Stephanie Hyland. \"Predicting the impact of treatments over time with uncertainty aware neural differential equations.\" International Conference on Artificial Intelligence and Statistics. PMLR, 2022.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)