ODE Discovery for Longitudinal Heterogeneous Treatment Effects Inference
Authors: Krzysztof Kacprzyk
Samuel Holt
Jeroen Berrevoets
Zhaozhi Qian
Mihaela van der Schaar
Venue: ICLR 2024 spotlight (2024)
Keywords: Treatment Effects over Time
Primary area: causal reasoning

Abstract: Inferring unbiased treatment effects has received widespread attention in the machine learning community. In recent years, our community has proposed numerous solutions in standard settings, high-dimensional treatment settings, and even longitudinal settings. While very diverse, the solution has mostly relied on neural networks for inference and simultaneous correction of assignment bias. New approaches typically build on top of previous approaches by proposing new (or refined) architectures and learning algorithms. However, the end result—a neural-network-based inference machine—remains unchallenged. In this paper, we introduce a different type of solution in the longitudinal setting: a closed-form ordinary differential equation (ODE). While we still rely on continuous optimization to learn an ODE, the resulting inference machine is no longer a neural network. Doing so yields several advantages such as interpretability, irregular sampling, and a different set of identification assumptions. Above all, we consider the introduction of a completely new type of solution to be our most important contribution as it may spark entirely new innovations in treatment effects in general. We facilitate this by formulating our contribution as a framework that can transform any ODE discovery method into a treatment effects method.

Published as a conference paper at ICLR 2024
ODE DISCOVERY FOR LONGITUDINAL HETEROGE- NEOUS TREATMENT EFFECTS INFERENCE
Krzysztof Kacprzyk' University of Cambridge
Samuel Holt"
Jeroen Berrevoets" University of Cambridge
University of Cambridge
Zhaozhi Qian University of Cambridge
Mihaela van der Schaar University of Cambridge
ABSTRACT
Inferring unbiased treatment effects has received widespread attention in the ma- chine learning community. In recent years, our community has proposed numerous solutions in standard settings, high-dimensional treatment settings, and even lon- gitudinal settings. While very diverse, the solution has mostly relied on neural networks for inference and simultaneous correction of assignment bias. New approaches typically build on top of previous approaches by proposing new (or refined) architectures and learning algorithms. However, the end result-a neural- network-basee inference machine- remains unchallenged. In this paper, we intro- duce different type of solution in the longitudinal setting: closed-form ordinary differential equation (ODE). While we still rely on continuous optimization learn an ODE, the resulting inference machine is no longer neural network. Doing so yields several advantages such as interpretability, irregular sampling, and a different set of identification assumptions, Above all, we consider the introduction completely new lype solution to be our most important contribution as it may spark entirely new innovations in treatment effects in general. We facilitate this by formulating our contribution as a framework that can transform any ODE discovery method into treatment effects method.
INTRODUCTION
Inferring treatment effects over time has received a lot of attention from the machine-learning community (Lim, 2018; Schulam & Saria, 2017: Gwak et al., 2020; Bica et al., 2020b). A major reason for this is the wide range of applications one can apply such a longitudinal treatment effects model. Consider for example the important problem of constructing treatment plan for cancer patients or even a training schedule to combat unemployment.
The increased attention from the machine learning community resulted in many methods relying on novel neural network architectures and learning algorithms. Once trained, these neural nets are used as inference machines, with innovation focusing on new architectures and learning strategies. This type of approach has indeed yielded many successes, but we believe that for some situations, one may want to consider an entirely different type of model: the ordinary differential equation (ODE).
This point is illustrated in fig. 1, where we show a standard treatment effects (TE) model on the left, and our new approach on the right. Essentially, standard TE model will first build a representation to adjust the covariate shift-e.g through propensity weighting (Robins et al., 2000) or adversarial learning (Bica et al., (220b)-used to model the outcome. Conversely, our new approach does not adjust the dataspace in any way, and instead discovers global ODE which we refine per patient, the goal of ODE discovery is finding underlying closed-form concise ODEs from observed trajectories.
Doing so yields advantages ranging from interpretability to irregular sampling (Lok, 2008; Saarela & Liu, 2016; Ryalen et al., 2019), and even modifying certain assumptions relied upon by contemporary techniques (Pearl, 2009; Bollen & Pearl, 2013). Most of all, we believe that proposing this new strategy may spark a radically different approach to inferring treatment effects over time. Moreover, given that the resulting model is an equation, one may use the discovered solution to engage in further
"qqual contribution; authors listed in reverse alphahetic order.

---

Published as a conference paper at ICLR 2024
TE Data: D
Standard TE: (D) + y
Our method: D - X)
p(o(X)
P(X (A))
Fx
Figure I: Conceptual overview. Left: A longitudinal treatment effects dataset with biased samples. Middle: A standard treatment effects (TE) approach will learn representation of the data, D, and will use the representation for inference. Right: Our approach, which learns ODE, refined for each specific patient in the dataset.
research and data collection, given that we can now understand certain behaviours the treatment in the environment. The latter is, of course, not possible when using black-box neural network models (Rudin, 2019; Angrist, 1991; Kraemer et al., 2002; Bica et al., 2020a; 2021).
Our contribution is usable framework that allows to translate any ODE discovery method (Brunton et al., 2016) into the treatment effects problem formulation. While other TE inference methods have used ideas from ODE literature, none have focused on ODE discovery, we have devoted appendix C.I to explain this subtle but important difference, Hence, in this paper, we first explain the differences between ODE discovery (and ODEs in general), and treatment effects inference, before connecting them through our proposal. Using our framework, we build an example method (called INSITE)', tested in accepted benchmark settings used throughout the literature. Transforming an ODE discovery method into a TE method results in new set of identification assumptions. This need not be a limitation, instead, can be seen as an extension as the typical TE identification assumptions may not always hold in ODE discovery settings.
HETEROGENEOUS TREATMENT EFFECTS OVER TIME
To reformulate the longitudinal heterogeneous treatment effects problem as an ODE discovery problem, in this section, we first introduce longitudinal treatment effects with its assumptions and then introduce ODE discovery in section 3. Let random variables X (t) E Rdd AA E A and y, E y be the 7th E [N] individual's features, treatment, and outcome at time horizon. We also denote the static covariates of the 7th individual as V R In the treatment effects literature, typically A = {0, 1} and y R. Unless required for clarity, will drop the individual indicator i. We use the following notation Xt,tt to denote the observed features in the time window 42). The observed dataset contains the realizations of the random variables above, D [[v, (x,, Wt: (.)) : [N], 1}.
T). with the time
Of interest is estimating the expected potential outcome Yes+r (Qnttr)) for some т > 0 under hypothetical future treatments ur.r+7 given the historical features Xo and the previous treatments Aua (Neyman, 1923; Rubin, 1980):
([ Xoss Aoz)
By definition, the potential outcome defined above includes multiple hypothetical scenarios that cannot be simultaneously observed in the real world. This problem is often referred to as the fundamental challenge treatment effects and causal inference (Holland, 1986).
As such, treatment effects (over time) literature typically introduces a set of assumptions to link the (unobservable) potential outcomes the observable quantities Y. X, and A, to correctly estimate (attr((ttt-)) in eq. (I). These assumptions are:
Assumption 2.1 (Consistency) For an observed treatment process Aurrte a, the potential out- come is the same the factual outcome Y Y(a) Yoro.,
We provide code http:s////..tooooocoo///atttttts Het e enteterogeneee www.erouss--ueoeeeooooooooorueee
tud
wtw.effffeeeeeeeeeeeee.... Inference.
2

---

Published as a conference paper at ICLR 2024
Assumption 2.2 (Overlap) The treatment intensity process (() is not deterministic given any filtration St (Klenke, 2008) and time point E [., i.e.,
p(Ar+8t
St Ae 0 < - 7: with Y E (0.1)
7 >< )() 56-+0
lim
öt
Assumption 2.3 (ggnorability) The intensity process A(t) given the filtration Be is equal to the intensity process that is generated by the filtration 3 U (Y,) t that includes the o-algebras generated by future outcomes (a(Y.) >>>> 1).
The above generalizes the standard identification assumptions made in static treatment effect literature (Rosenbaum & Rubin, 1983; Seedat et al., 2022). This generalization is largely based on previous extensions to continuous-timm causal effects (Lok, 2008; Saarela Liu, 2016; Ryalen et al., 2019).
Assum. 2.2 and assum. 2.3 rely on a treatment intensity process, 3() which can be considered a generalization of the propensity score in continuous time (Robins, 1999). Essentially, assum. 2.2 allows that any treatment can be chosen at time t, given the past observations in the filtering 3t. Furthermore, assum. 2.3 ensures is sufficient to condition on patient's 's past observed trajectory to block any backdoor paths to future potential outcomes (eq. (1)).
Given the above, we are allowed the following equality, which identifies the potential outcome (LHS) to be estimated using the observed variables (RHS):
E Y VV Xo:t Aax] = B Xo:e Ag:s Act++ = Cut+rl.
(2)
which is exploited by most works the treatment effects literature (albeit by first regularizing models such that they respect assum. 2.1 to 2.3) (Bica et al., 2020b; Lim, 2018; Melnychuk et al., 2022). A thorough overview of related works is presented appendix C.
3 UNDERPINNINGS OF ORDINARY DIFFERENTIAL EQUATION DISCOVERY
We now propose model this treatment effect time problem as dynamical system, where their temporal evolution can often be well represented by ODEs (Hamilton, 2020). That means assume the time-varying features and outcomes of the hh individual (z, (i) Vt [0, (+)) discrete (and possibly noisy) measurements of underlying continuous trajectories of observed features (0))) [0. 7] + RRd and potential outcomes (0)) [0, T] - TR, where T IR called the time horizon (Birkhoff, 1927). To make our formalism consistent, we also assume there is treatment trajectory a [0, 2| + # such that ! either constitute snapshots of the underlying continuous treatment a or if the treatments are administered at discrete times then a is step function whose values corresponds to the currently administered treatment.
There are many fields which already use ODEs as formal language to express time dynamics. One such example is pharmacology, where large portion of the literature is dedicated recovering an ODE from observational data. The found ODE is then used to reason about possible treatments and disease progression (Geng et al., 2017: Butner et al., 2020). We further assume that this system modelled by a system of ODEs which describes the time derivative of x as function of æ, and The outcome trajectory, y. depends on the features æ. particular, we assume dx(t) (t) F(v, at)) and y(t) g(æ(t)),
dt
where (t) the differential of x. g is prespecified by the user and often models the outcome as one of the features (g(x)= = zy)) e.g., the tumour volume. This is general formulation that may also include direct effects of a on y. We further discuss eq. (3) appendix The goal of ODE discovery is to recover the underlying system of ODEs F based on the observed dataset D (Brunel, 2008; Brunel et al., 2014). Of course, the reliance on such a dataset implies that, while the trajectories are defined in continuous time, they are observed in discrete time. Moreover, methods for ODE discovery make the following assumptions:
Assumption 3.1 (Existence and Uniqueness) The underlying process be modelled by system of ODEs at) = F(u, att), and for every initial condition Eg. v and treatment plan a
We further explain treatment intensity processes and filtrations in the assumptions in appendix B. Current ODE discovery methods are not designed to work with static features but they can be adapted to this setting by considering them time-varying features that constant throughout the trajectory.
3

---

Published as a conference paper at ICLR 2024
at to, there exists a unique continuous solution [to. T] R satisfying the ODEs for all t E (to, (LLindef.. 1894; Ince, 1956).
Assumption 3.2 (Observability) All dimensions of all variables in F are observed for all individ- wals, ensuring sufficient data to identify the system dynamics and infer the ODE's structure and parameters (Kailath, 1980).
Assumption 3.3 (Functional Space) Each ODE in F belongs some subspace of closed-form ODEs. These are equations that can be represented as mathematical expressions consisting of binary operations {+,-, x, }. input variables, some well-known functions (e.g., {log, exp, sin)), and numeric constants (e.g., 1-0.2..... 5.2} E IR) (Schmidt & Lipson, 2009).
Identification. Assum. 3.1 and 3.2 play a crucial role in ODE discovery. Essentially, we require making such assumptions in order correctly identify the underlying equation. The assumptions made in the treatment effects literature (assum. 1 to 2.3) serve a similar purpose as they allow us to interpret the estimand causal effect, i.e., they are necessary for identification (Imbens & Rubin, 2015; Rosenbaum & Rubin, 1983; Imbens, 2004).
Assum. 3.1 ensures that the discovered ODE has unique solution, which is essential for making reliable predictions, assum. 3.2 is necessary such that the observed data can be used to accurately identify the underlying ODE. Finally, assum. 3.3 defines the space of equations for the optimization algorithm to consider. We review methods for ODE discovery in appendix C.
CONNECTING TREATMENT EFFECTS INFERENCE AND ODE DISCOVERY
-THE FRAMEWORK
From eqs. (2) and (3), one can see that both treatment effects and ODE discovery involve learning a func- tion that can issue predictions about future states. Al- though the connection is natural, there remain some discrepancies between these two fields. To apply ODE discovery methods for treatment effects infer- ence, we have resolve these discrepancies. Essen- tially, we establish framework one can use to apply any ODE discovery technique in treatment effects.
treatments
Cont.
Bin.
Cat.
BSV (table 2)
section
B
+
+
C
We identify three discrepancies between the treatment effects literature and the ODE discovery literature: (I) different assumptions, (2) discrete (not continuous) treatment plans, and (3) variability across subjects. Each discrepancy is explained and reconciled in dedicated subsection below with actionable steps.
D
+
INSITE
Figure 2: Dimensions of our framework. The x-axis shows different treatment types (cfr. sec- tion 4.2)
y-axis shows between-subject variability in increasing difficulty (cfr. section 4.3). Vnndicates "no adapting needed", and $ shows our framework. In green shows what ODE discov-
Figure 2 shows the areas ODE discovery methods can be expanded to, using our framework. Ranging from simple adaption, to proposing completely new method (in section 5), our framework significantly in- creases the reach of existing ODE discovery methods. Our new method-INSSTE: shuld be considered the result after implementing our practical framework we present in the remainder of this section.
ery methods can do out of the box. Blue shows INSITE's possibilities, encompassing all settings, including complex BSV.
4.1 DECIDING ON ASSUMPTIONS (DISCR.
Discrepancy 1 (Different assumptions.) In sections 2 and 3 we listed the most common assump- tions made in treatment effects and ODE discovery literature, respectively. While solving similar problems, these assumptions do not correspond one-to-one (table 1). However, the fact that they do not can be seen as major advantage-in in some scenarios it may be more appropriate to assume assum. 2.1 to 2.3 versus assum. 3.1 to 3.3 or vice versa, which expands the application domain.
Increasingly, the treatment effects literature is considering settings that violate the overlap assumption assum. 2.2 (D' Amour et al., 2021). It has been shown that correct model specification can weaken or replace the overlap assumption (Gelman & Hill, 2006, Chapter 10). The same holds true for
4

---

Published as a conference paper at ICLR 2024
Table Comparing assumptions. This table lists the typical assumptions made in ODE discovery and in treat- ment effects. While some seem to correspond (or are at least similar), we have shaded green where we could relax some assumptions with others. This a powerful idea stemming directly from connecting these two fields. We show the robustness of these assumptions in violating settings section 6. We note that the correspondence of similarity between these assumptions is not formal but indicates the discrepancy between the formalisms of treatment effect inference and ODE discovery that needs to be considered as part of the framework.
ODE discovery
Treatment effects
Explanation
ref
assumption
assumption
ref
3.1
existence & uniqueness consistency
2.1
2.1 is implicit through 3.2.
3.2
observability o
- overlap
2.2
2.2 can be relaxed by 3.1 and 3.3
3.3
functional spaces a
ignorability
2.3
2.3 is similar as 3.2.
the ODE discovery methods, where overlap in assum. 2.2 can be relaxed with assum. 3.1 and 3.3. Consider an example where the true and specified models are both linear. Here, overlap can be violated as we can safely extrapolate outside the support of either treatment distribution. In contrast, the recent neural network-based treatment effects models can rarely satisfy correct model specification (Lim, 2018; Bica et al., 2020b; Berrevoets et al., 2021; Melnychuk et al., 2022). As a result, the overlap assumption plays key role for these methods. We show this empirically in section 6.
Framework step 1:
Accept ODE discovery assumptions (assum. 3.1 to 3.3).
4.2 INCORPORATING DIVERSE TREATMENT TYPES (DISCR. 2)
Discrepancy 2 (Discrete treatment plans.) In section 3 we assume that a is a trajectory like x and with a; snapshot observed similarly to observing covariates and outcomes. This implies that, like covariates and outcomes, the treatment plan is defined continuous time, and the treatment is itself continuous value. While there exist scenarios where this could be possible (e.g., in settings where treatment is constantly administered), most settings violate this. Hence, we need to reconcile modelling treatments in continuous time and values to settings violating these basic setups.
To connect ODE discovery with the treatment effects literature, we require incorporating different types of treatment plans. Continuous valued treatment that is administered over time is adopted quite naturally a differential equation. However, the treatment effects literature is typically focused on other types of treatment (Bica et al., 2021): binary treatments, categorical treatments, multiple simultaneous treatments, and dosage. Each of these treatments can be static (i.e., constant throughout the trajectory) or dynamic. With such diverse treatments, it might be impossible to express F closed-form expression. For instance, where the treatment is a categorical variable (i.e., A (l,..., k.)). This violates assum. 3.3 of continuous solutions made by ODE discovery methods.
To reconeile it, we need to decide how to incorporate treatment a into F. This can done two ways: either we discover different (piecewise) closed-form ODEs for different treatments (Trefethen et al., 2017; Jianwang et al., 2021), or we incorporate the action variable a into the closed-form ODE, Depending the type of treatment one both of these approaches can be chosen. In appendix E we outline the ways of incorporating the treatment plan a in according to the treatment types listed above. In summary, there are 4 different treatment types: binary, categorical, multiple, and continuous treatments (Bica et al., 2020b). Each of them can be a static or dynamic treatment, which either constant changes during a trajectory, respectively. This results in scenarios, which model in two ways: either the ODE changes for each treatment option; or the treatment is part of the starting condition. These are all outlined in table appendix
In section 6 we demonstrate the effectiveness of discovering ODEs per categorical treatment.
Framework step 2:
Incorporate treatment plans as dynamical systems using appendix E.
4.3 MODELLING BETWEEN-SURJERE VARIABILITY (DISCR. 3)
Diserepancy 3 (Between-subject variability.) In a classical ODE discovery setting, we usually have only one kind of trtween-subjed variability corresponding unexplained variability (RUV). However, there are a few other sources of variability that need to be accounted for before considering it as an ODE discovery problem.
noisy measurement residual
5

---

Published as a conference paper at ICLR 2024
Table 2: Layers of Between-Subject Variability (BSV). We show four levels of between- subject variability and explain the different sources of where variability is coming from, In the table, have: ODE (i), RUV (ii),
Covariates (iii), Parameter distributions (iv) and ((a) awp + wi. detailed appendix
(i)
iii)
(iii)
(iv)
Causal graph
Example v(1)
ODF
+RUV
+Cov
+Dist
Parameters (eq. (5))
A
X
x(t)
Co = co. C1
B
V
/
a(1) +
Co - co. m
C
/
/
X
a(0)+
Co ((r)), e(e))
/
d
x(t) + Co (og((o)): ~ ((1(e))
As pharmacological models play a prominent role in treatment effect literature (Geng et al., 2017; Bica et al., 2020b; Berrevoets et al., 2021; Lim, 2018; Seedat et al., 2022) we employ formalism from pharmacology literature to discuss between-subject variability (BSV).
There are two types of BSV (Mould Upton, 2012): (i) Unexplained variability which includes RUV and parameter distributions; and (ii) Explained variability includes covariate models where we model the impact of static covariates on the equation's 's parameters.
We base the model of our causal dynamical system on the formalism introduced in Peters et al. (2022). partieular, they use the term Deterministic Casual Kinetic Model (DCKM), which can be summarised as system of first-order autonomous ODEs. This is the simplest pharmacological model with no BSV (setting A in Table 2). In light of realism, however, we can add different layers of BSV, making the model more complex and thus more difficult to discover by current methods.
Table 2 outlines different types of BSV as graphical models, each with increasing complexity. In section 6, we relate the parameter columns in table 2 to the underlying ground-truth equations used in our experiments. We now detail these layers in increasing complexity.
B: RUV noisy measurements. Noisy measurements discovery (Brunton et al., 2016). This is because noisy measurements cause derivative estimation to be very challenging. Recently, methods based on the weak formulation of ODEs (Qian al., 2022; Messenger & Bortz, 2021) have managed to circumvent the derivative estimation step, making them more robust noisy measurements. We represent DCKMs with noisy measurements graphically similar to the standard DCKM but will now explicitly add noise-related nodes (table 2 row B).
one of the biggest challenges of ODE
C: Covariate models. One way we can incorporate explainable variability into the model is by modelling the impact of the covariates on the model parameters. The covariates are incorporated into the model by closed-form expressions. Usually simple ones such as linear, exponential, or power (Mould & Upton, 2012). Sometimes other equations are used, such as complex closed-form expressions or piece-wise functions (Chung et al., 2021). We add a node corresponding to static covariates to represent covariate models graphically (table 2 row C).
D: Parameter distributions. Although the covariate models let us calculate the group parameters- parameters for specific group of people based, e.g., on their age weight, the actual parameter for an individual might deviate from this value in some random way. We can model this by assuming a the covariates. We can depict this graphi-
distribution of parameter values with its mean depending cally by adding the nodes corresponding to the parameters and the associated noise nodes coming into them. Similarly, ODE discovery methods are not designed to discover such models (table 2 row D).
Framework step 3: Decide on the BSV using table 2 and, if type D, adjust as in section 5.
5 INSITE
-TE FRAMEWORK IN PRACTICE
Having introduced our general framework of applying ODE discovery techniques to treatment effects inference (in section 4), we now turn to apply it. Following the framework steps (FS) as described in section 4, we have to: (FS1) accept new set of assumptions (discr. 1), (FS2) incorporate treatment plans (discr. 2), and (FS3) decide on the BSV type (diser. 3). For our purposes, we use SINDy (Brunton et al.. 2016), one of the most widely adopted and used methods in ODE discovery. One can use any underlying ODE discovery methods as long as they: (1) model ODEs that include a set of numeric constants (3 Rm with m numeric constants), and (2) we accept that BSV is modelled

---

Published as a conference paper at ICLR 2024
Algorithm 1 Individualized Nonlinear Sparse Identification Treatment Effect (INSITE)
1: Input: Patient data D cfr. section 2; deterministic DE discovery method, DE.
2: Output: Patient-specific DEs F(;; population DEs F.
3: F <DE (D)
Step Discover Population Differential Equation b. Step 2: Discover Patient- Differential Equations
4: for patient (i) in DD do
5: Fine-tune numeric constants F using patient (i) eq. (4)
6: Obtain patient-specifee equation F(i)
7:
return Patient- Patient-specific F(i) and F
through these numeric constants alone, i.e., two varying subjects can only differ by having different model numeric constants.*
Individualized Nonlinear Sparse Identification Treatment Effect (INSITE) consists of two main steps, reminiscent of the remaining discr. 2 and 3: (1) discover the population (global) differential equation F for all patients, and (2) discover the patient-specit differential equation F(i) by fine-tuning the population equation's 's numeric constants- outlined in algorithm 1. Having set of numeric constants (one for each patient) we derive a population differential equation where each individual's numeric constants are represented by sample from this population differential equation distribution.
Step (cfr. framework step 2): Discovering the Population Differential Equation. Given an observed dataset ID of patients, we aim to discover the population differential equation governing the patient covariates' interaction with their potential outcomes. We can use any deterministic differential equation discovery method such as SINDy (Brunton et al., 201 6) to recover the population ODE F. We adapt it to handle treatments as outlined in section 4.2 and appendix (as per discr. 2), where for each separate k categorical (or binary) treatment, we discover separate ODE-wil which is only active for that categorical treatment, i.c., F ===
Step 2 (cfr. framework step 3): Discovering Patient-Specife Differential Equations. Once the population differential equation is discovered, we then fine-tune the numeric constants in the population equation to obtain patient-specific differential equations F(). By keeping the same functional form but allowing unique numeric constants for each patient to be refined, we can model the individualized treatment effect and account for patient heterogeneity (between-subee -ubjectt variability). Crucially, to avoid overfitting and allow only small deviations away from the initial population numeric constants B, we fit the observed patient trajectory up the current time by minimizing.
1 T
((
( ( A (1) ) +1113 - 3
Mean Squared Error (MSE) of the predicted trajectory evolution y for the observed trajectory. Additionally, we also require a regularization term in eq. (4). Specifically, for each individual tra- jectory, (i), we refine the population set of parameters (3) to obtaín ( while still using them to regularize to prevent overfitting. The latter is an insight borrowed from transfer learning (Tommasi et al., 2010; Tommasi & Caputo, 2009; Takada & Fujisawa, 2020). Interestingly, without such regularization, find INSITE to underperform significantly (cfr. appendix K.1). We use Broy- den-Fletcher Goodearr---oobbbboaaaaaaa
(BFGS) algorithm (Fletcher, 2013) to minimize eq. (4) as is standard in symbolic regression (Petersen et al., 2020), and provide full inference details in appendix G.I.
Parameter distributions. Unique to our method is that after obtaining the patient-specific differential equations, we derive population differential equation where each numeric constant is represented by a distribution, such as normal distribution or mixture of distributions distiibutions-rrcoovrnnng probabilistic interpretation of the underlying data generating process. This is explored in appendix K.2.
6 EXPERIMENTS AND EVALUATION
To allow for robust and systematic comparison of longitudinal treatment effect (LTE) methods and ODE discovery techniques, we create synthetic testbed designed for treatment effect prediction
Essenttialyy we cannot have mathematical expression changed across subjects (such as exp changed to log. In fact, scenario such this would violate assum. 3.3).
7

---

Published as a conference paper at ICLR 2024
across different synthetic datasets generated from different classes of pharmacological models. We note that using synthetic datasets is common in benchmarking LTE methods, as for a real dataset the counterfactual outcomes unknown (Lim, 2018; Bica et al., 2020b). Our testbed consists of:
Diverse underlying pharmacological models F. The synthetic datasets are generated by sampling pharmacological model F with given treatment assignment policy. To ensure robustness across diverse pharmacological scenarios, we include the model classes outlined in section 4.3, from A to D-which differ between noise, static covariates, and parametric distributions of parameters. We analyze two standard Pharmacokinetic- Saetiii-aaaaaaaaaaaaiiaa
(PKPD) models for each model class (section 4.3). First, a compartaattt PKPD model (eq. (5)) with a binary static action. Second, a state-of-the-an biomedical PKPD model of tumor growth, used to simulate the combined effects of chemotherapy and radiotherapy in lung cancer (Geng al., 2017) (eq. (6))-this has been extensively used by other works (Seedat et al., 2022; Bica et al., 2020b; Melnychuk et al., 2022). Here, to explore continuous types of treatments, we use a continuous chemotherapy treatment c(t) and a binary radiotherapy treatment d(t), both changing over time. For both models z, is the volume of the tumor & days after diagnosis, modeled separately as:
dr(() dt
Cp x(t), if = 0 V
(5)
(r(t),
= a
da(()] (plog xxt)
K
BeC(L)
(aa,((t) + B, e )=(()
(6)
dt
Chemocherapy
Radiotherapy
Noise
Tumorgrowth
where the parameters Co, C1, V, a, B, are sampled according to the different layers of between-subject variability (table 2) forming variations of A-D, with parameter distributions following that as described in Geng et al. (2017) or otherwise detailed in appendix F. We also compare to the standard implementation of eq. (6) (labelled as Cancer PKPD) to ensure standard comparison to existing state-of-ther art methods, where the treatments are both binary. We further detail dataset generation for each in appendix F.
Action assignment policy. We introduce time-dependent confounding by making the treatment assignment vary from a purely random treatment assignment to purely deterministic based on threshold of the outcome predictor value, controlled by a scalar 7 E..-such that y = 0 corresponds to no time-dependent confounding and larger values correspond to increasing time- dependent confounding, Further details of this treatment policy are in appendix F.
Benehmark methods. We seek to compare against the existing state-of-the-ar (SOTA) methods from (I) longitudinal treatment effect models and (2) ODE discovery. First, longitudinal treatment effect models often consist of black-box neural network- interpretable. A significant theme in these works is the development of methods to mitigate time- dependent confounding. Many mitigation methods have been proposed that we use as benchmarks, such as propensity networks in Marginal Structural Models (MSM) (Robins et al., 2000) and Recurrent Marginal Structural Networks (RMSN) (Lim, 2018), Gradient Reversal in Counterfactual Recurrent Networks (CRN) (Bica et al., 2020b), Confusion in Causal Transformers (CT) (Melnychuk et al., 2022) and g g-computation in G-Net (G-Net) (Li et al., 2021). Second, ODE discovery methods aim to discover an underlying closed-form mathematical equation F that best fits the underlying controlled ODE of the observed state and action trajectories dataset, 'D. In these methods, the state derivative is unobserved; therefore, they have to approximate the derivative using finite diffcrences, as in Sparse Identification of Nonlinear Dynamics (Brunton et al., 2016), or employ variational loss as the objective, as in Weak Sparse Identification of Nonlinear Dynamics (Messenger & Bortz, 2021)². To make these two ODE discovery methods more competitive, we adapt them to model individual ODEs per categorical treatment (section 4.2), (termed A-SINDy, A- A-WSINDy respectively). We further diseuss why we chose these SOTA methods and their implementation details in appendix G.
approaches, i.e., not closed-form or easily
We evaluate against the standard longitudinal treatment effect metrics of test counterfac- tual r-step ahead prediction normalized root mean squared error (RMSE), where = (1, 4, 5, 6). for a sliding treatment plan (Meinychuk et al., 2022). Unless otherwise specified, each result is run for five random seed runs with time-dependent confounding of 7 = 2-see appendix H for details.
Evaluation.
'We only include results for WSINDy when it is possible use it, as it does not support sparse short. trajectories. Detailed further in appendix G.
8

---

Published as a conference paper at ICLR 2024
MSM
RMSN
CRN A-WSINDy
G-Net INSITE
CT
A-SINDy
2.5
n-RMSE
0.5
1.5
0.1
0.02
1
2
6
2
(a) T-Steps
(b) 7
Figure 3: (a) Counterfactual 7-step ahead prediction error (7 = 2). on eq. (5).D from table 3. (b) Counterfactual 6-step ahead prediction error (for increasing time-dependent confounding, 7), on the standard Cancer PKPD dataset. INSITE maintains a low normalized RMSE (high performance) across long time horizons and increasing time- dinedeepeneent confounding. Further results are appendix K.
Table 3: Counterfactual normalized RMSE, for 6-step ahead prediction of the benchmarks on each synthetic dataset detailed in appendix F. We quote 95% confidence intervals with each value, and all metries averaged over ten random seed runs. Each PKPD underlying pharmacological model simulated with different layers of
BSV (table 2) for A-D. Other 7-step results in the appendix K.3. Our contribution is shaded below.
Method
15.A
(5,C
(5.D
(6.B
H9LC
eq.OoLD
Cancer PKPD
MSM
0:99 1.777777
0.90...00
2,05 0.00
2.55.1115
2.5510.13
2,06 1.15
2.11.004
2.301.0.12
ODDED EITE
RMSN
CRN
1.05:0.00
1:92:0:24
14550.10
1,94:015
0.2:10099
16811.19
19110.11
L99:0.14
1,23,0015
1005.000
105:00
1.2510.55
103
1.000.1
1.03:00 10
1.10.n
0.92:0.05
1 1.0400117
G.Net
0.91:50.99
0.9120.20
0.7250.14
0.97=0.15
1.33:027
13349027
L02s0.11
1,25 15
1.22a0.14
CT
0,90
0.90s0118
0.7550114
L00:0.14
1.29:30.7
12960.10
103
1,14 10
1/0760077
A-SINDy
0:31:0000
0.11:0000
8.13:229-11
0.15:0000
1.45:0088
1,55000
1,02001
1.51:1000
1.2350.13
A WSINND
0.11=7.246.05
0.11:2440044
0.1251447-00
0.10:761:-04
NA
NA
NA
NA.
NA
INSITE
0.82 #2.00-15
0.10 =0:00
0J94 20.00
0.05 #5
2005
0.34:3005
0.54
2004
0.87 2000
0.79 #7.77e117
Main results.
The test counterfactual normalized RMSE for 6-step ahead prediction for each benchmark dataset is tabulated in table 3. Our method, INSITE, achieves the lowest test counterfactual normalized RMSE across all methods. We provide additional experimental results in appendix K.
Interpretable equations. Unique the proposed framework and INSITE, is that the discovered differential equation is fully interpretable; however, it relies on strong assumptions. Some of these discovered equations shown in appendix J. It is clear that even with binary actions, INSITE is able to discover a useful equation that performs well and is similar in form to the true underlying equation, even discovering exactly in some simple settings of eq. (5). The method can even do this in the presence of noise (BSV layers of B-D) and extrapolate well for future 7-step predictions, fig. 3 (a).
Model misspecification. INSITE and the adapted ODE discovery methods (A-SINDy, A-WSINDy) are only correctly specified for eq. (5).A-D datasets, as they use the feature library of LINSITE = {I, zg. 21, xg Crucially, the ODE discovery methods are misspecified for eq. (6).A-D, and the Cancer PKPD datasets, as their underlying equation would require the feature library of fCancer = {1, To, 31, Ig 2bxr1, 20 (og(zo) (og())) to be correctly specified. Although this misspecification persists, as noticeable from the increased order of magnitude increase in error, we still observe INSITE achieves a lower error than the longitudinal treatment effect models (app. K.6).
Relaxing the overlap assumption. We can relax the overlap assumption (assum. 2.2) by increasing the time-dependent confounding of the treatment given, increasing 7. We observe, in fig. (b), that INSITE can still discover good approximating underlying equation even in the presence of high time-dependent confounding- hence where there is low overlap.
Table 4: INSITE Ablation
INSITE Ablation. INSITE's gain in performance derives from our
framework's ' discr. 2, to discover individual ODEs per categorical treatment, well as the ability to discover individualized (fine- tuned) ODEs. Notably, the improvement arising from the ability to
ODE per Cat. (section 4.2
Fine Tuning
(5LD n-RMSE
0,85 23c
71ce--17
discover individual ODEs per categorical treatment can be widely applied to existing ODE discovery methods using our framework for treatment effects over time. Furthermore, we conduct additional insight experiments to evaluate the methods against datasets of smaller sample sizes and increasing observation noise in appendix K.
10.00
7 CONCLUSION AND FUTURE WORK
In conclusion, we presented first framework that connects longitudinal heterogeneous treatment effects with ODE discovery methods, enabling improved interpretability and performance. Naturally, this is just the beginning. We hope that building on framework (appendix A), the connection between treatment effects and ODE discovery is further solidified and perhaps extended to other types of differential equations and dynamical systems in general (see e.g. appendix C).
9

---

Published as a conference paper at ICLR 2024
Ethics statement. This paper's novel approach of integrating ODE discovery methods into treatment effects inference can revolutionize precision medicine by enabling personalized, effective treatment strategies. However, misuse or application in inappropriate contexts could lead to incorrect treatment decisions. Moreover, despite the potential for individualized treatment, privacy concerns arise. Thus, proper data governance, clear communication of model limitations, and rigorous validation using diverse datasets are vital for responsible application.
Regarding limitations of our method, INSITE, and our framework in general, list major limitations one should take into account before applying our work in practice:
A correct set of candidate functions (tokens) necessary for correct model recovery. To show the importance of this, we include experiments where we have wrongly specified this token library (see appendix K.6) and observe degrading performance as a result.
2. ODE discovery works best in sparse settings. The reason is two-fold: from a technical point of view, sparse equations are much less complex and simply easier to recover; from usability point of view, the usefulness of non-sparse equations is limited as interpretabilityy is negatively affected by non-sparse (or non-parsimonious) equations (Crabbe et al., 2020).
3. ODEs are noise free. Since we recover ODEs, the found equations do not model source of noise as typically the case in structural equation modelling. To model noise terms explicitly, our framework should be extended into stochastic DEs, as is stated in our future works paragraph above.
Reproducibility statement. We provide all code https: http:/hwww...o///ooooo/////////t Discovery- for-Longi loongituding studinal-Biti Reedinal-Hetergennnn Heterogeneous Teetenntt------eeeeett---e-e Efecss fefff-ettaaeeeeee Inf erence. To ensure this paper is fully reproducible, we include an extensive appendix with all implementation and experimental details to recreate the method and experiments. These are outlined as the following: for benchmark dataset details, see appendix F; benchmark method implementation details, including those of the treatment effects baselines and adapted ODE discovery methods which include the proposed INSITE method see G; evaluation metric details see appendix H; dataset generation and model training see appendix I.
The authors would like to acknowledge and thank their corresponding funders, where SH is funded by AstraZeneca. JB is funded by W.D. Armstrong Trust, KK is funded by Roche, ZQ is funded by the Office of Naval Research (ONR). Moreover, we would like to warmly thank all the anonymous reviewers, alongside research group members of the van der Schaar lab (www. vander r rschaar ab. com), for their valuable input, comments, and suggestions as the paper was developed.
Acknowledgements.
REFERENCES
Ahmed M Alaa and Mihaela der Schaar. Attentive state-space modeling of disease progression. Advances in neural information processing systems, 32, 2019.
Mauricio A Alvarez and Neil D Lawrence. Sparse convolved multiple output gaussian processes.
arXiv preprint arXiv:091 107, 2009.
Joshua Angrist. Instrumental variables estimation of average treatment effects econometries and epidemiology, 1991.
Joshua D Angrist, Guido W Imbens, and Donald B Rubin. Identification of causal effects using instrumental variables. Journal of the American statistical Association, 91(434):4144555, 1996.
Susan Athey and Guido W Imbens. Identification and inference in nonlinear diffeereneeinn-diffeeeness models. Econometrica, 74(2):331449,, 2006.
Jeroen Berrevoets, Alicia Curth, Ioana Bica, Eoin McKinney, and Mihaela van der Schaar. Disen- tangled counterfactual recurrent networks for treatment effect inference over time. arXiv preprint 1r201:2112.03811, 2021.
Jeroen Berrevoets, Krzysztof Kacprzyk, Zhaozhi Qian, and Mihaela van der Schaar. Causal deep learning. arXiv preprint arXiv:2303.021 026 2023.
10

---

Published as a conference paper at ICLR 2024
Ioana Bica, Ahmed Alaa, and Mihaela Van Der Schaar. Time series deconfounder: Estimating treatment effects over time in the presence of hidden confounders. In International Conference on Machine Learning, pp. 884 895. PMLR, 2020a.
Ioana Bica, Ahmed M. Alaa, James Jordon, and Mihaela van der Schaar. Estimating counterfactual treatment outcomes over time through adversarially balanced representations. In International Conference on Learning Representations, 2020b.
Ioana Bica, Ahmed M. Alaa, Craig Lambert, and Mihaela Schaar. From Real-World Patient Data to Individualized Treatment Effects Using Machine Learning: Current and Future Methods to Address Underlying Challenges. Clinical Pharmacology & Therapeutics, 109(1):87-100, January 2021. ISSN 0009-9236, 1532-6535. doi: 10.1002/cpt.1907.
George David Birkhoff. Dynamical systems, volume 9. American Mathematical Soc., 1927.
Kenneth A Bollen and Judea Pearl. Eight myths about causality and structural equation models. Handbook of causal analysis for social research, pp. 301-328, 2013.
Nicolas JB Brunel. Parameter estimation of ode's via nonparametric estimators. 2008
Nicolas JB Brunel, Quentin Clairon, and Florence d' Alché Buc. Parametric estimation of ordinary dif- ferential equations with orthogonality conditions. Journal of the American Statistical Association, (09(505):133-18) 2014.
Steven I. Brunton, Joshua I. Proctor, and Nathan Kutz. Discovering governing equations from data by sparse identification of nonlinear dynamical systems. Proceedings of the national academy of sciences, 1 113(15):3922-3937 2016.
Joseph D Butner, Dalia Elganainy, Charles Wang, Zhihui Wang, Shu-Hsia Chen, Nestor Esnaola, Renata Pasqualini, Wadih Arap, David S Hong, James Welsh, et al. Mathematical prediction of clinical outcomes in advanced cancer patients treated with checkpoint inhibitor immunotherapy. Science advances, 6 ((18))8a86629, 2020.
Gustau Camps- Valls, Andreas Gerhardus, Urmi Ninad, Gherardo Varando, Georg Martius, Emili Balaguer- Balaguer-Ballesterr Ricardo Vinuesa, Emiliano Diaz, Laure Zanna, and Jakob Runge. Discovering causal relations and equations from data. arXiv preprint arXiv:2305. 13341, 2023.
Kyunghyun Cho, Bart Van Merriënboer, Caglar Gulcehre, Dzmitry Bahdanau, Fethi Bougares, Holger Schwenk, and Yoshua Bengio. Learning phrase representations using encoder-decoder for statistical machine translation. arXiv preprint 1406. 1078, 2014.
Erin Chung, Jonathan Sen, Priya Patel, and Winnie Seto. Population Pharmacokinetic Models of Vancomycin in Paediatric Patients: A Systematic Review. Clinical Pharmacokinetics, 60(8): 985-1001, August 2021. ISSN 0312-5963 1179-1926. doi: 0000777/402222-22-0002777
Jonathan Crabbe, Yao Zhang, William Zame, and Mihaela van der Schaar. Learning outside the black-box: The pursuit of interpretable models. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and Lin (eds.), Advances Neural Information Processing Systems, volume 33, pp- 17838-17888 Curran Associates, Inc., 2020. URL ht http:///ooooeiiiieee neurips cc pdf.
Miles Cranmer, Alvaro Sanchez Gonzalez, Peter Battaglia. Rui Xu, Kyle Cranmer, David Spergel, and Shirley Ho. Discovering symbolic models from deep learning with inductive biases. Advances in Neural Information Processing Systems, 33:17429-17444 2020.
Kanti Bhushan Datta and Bosukonda Murali Mohan. Orthogonal functions in systems and control, volume 9. World Scientific, 1995.
Alexander D' Amour, Peng Ding, Avi Feller, Lihua Lei, and Jasjeet Sekhon. Overlap in observational studies with high-dimensional covariates. Journal of Econometries, 221(2):644 654, 2021.
William A Falcon. Pytorch lightning. GitHub, 3, 2019.

---

Published as a conference paper at ICLR 2024
Roger Fletcher. Practical methods of optimization. John Wiley & Sons, 2013.
Jean-Pierre Florens, James J Heckman, Costas Meghir, and Edward Vytlacil. Identification of treatment effects using control functions in models with continuous, endogenous treatment and heterogeneous effects. Econometrica, 1655):119111200 1206, 2008
Michele Jonsson Funk, Daniel Westreich, Chris Wiesen, Til Stürmer, M Alan Brookhart, and Marie Davidian. Doubly robust estimation of causal effects. American joneriaanjournall of epidemiology, 173(7): 761-767, 2011.
Yaroslav Ganin and Victor Lempitsky. Unsupervised domain adaptation by backpropagation. In International conference on machine learning, pp. 1180-1189. PMLR, 2015.
Armengol Gasull, Antoni Guillamon, and Jordi Villadelprat. The period function for second-order
quadratic odes is monotone. Qualitative Theory Dynamical Systems, 4(2):329-355) 2004.
Andrew Gelman and Jennifer Hill. Data analysis using regression and multteeve/hieerarccicc models. Cambridge university press, 2006.
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 Pio-Mattemmatical Model. Scientific Reports, 131):1354,, October 2017. ISSN 2045-2322. doi:
10.1038/s41598-011 010-41159-----7-11 13646-7.
Pawan Goyal and Peter Benner. Discovery of nonlinear dynamical systems using a runge- kutta inspired dictionary- sparse regression approach. Proceedings of the Royal Society A, 478 (2262):20210883 2022.
Daehoon Gwak, Gyuhyeon Sim, Michael Poli, Stefano Massaroli, Jaegul Choo, and Edward Choi. Neural ordinary differential equations for intervention modeling. arXiv preprint arXiv:2010.08304, 2020.
James Douglas Hamilton. Time series analysis. Princeton university press, 2020.
Miguel A Hernán, Babette Brumback, and James M Robins. Marginal structural models to estimate the joint causal effect of nonrandomized treatments. Journal of the American Statistical Association, 96(454):440 448, 2001.
Çağlar Hızh, ST John, Anne Tuulikki Juuti, Tuure Tapani Saarinen, Kirsi Hannele Pietiläinen, and Pekka Marttinen. Joint point process model for counterfactual treatment-outcomm trajectories under policy interventions. In NeurlPS 2022 Workshop Learning from Time Series for Health, 2022.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8): 1735-1780, 1997.
Paul W Holland. Statistics and causal inference. Journal of the American statistical Association, 81 (396):9459966) 1986.
Samuel Holt, Alihan Hüyük, Zhaozhi Qian, Hao Sun, and Mihaela van der Schaar. Neural laplace control for continuous-time delayed systems. In International Conference Artificial Intelligence and Statistics, pp. 1747-1778. PMLR, 2023a.
Samuel Holt, Zhaozhi Qian, and Mihaela van der Schaar. Deep generative symbolic regression. arXiv preprint arXiv:2401.0022 2023b.
Samuel Holt, Alihan Hüyük, and Mihaela der Schaar. Active observing in continuous-time control. Advances in Neural Information Processing Systems, 36, 2024.
Samuel Holt, Zhaozhi Qian, and Mihacla van der Schaar. Neural laplace: Learning diverse classes of differential equations in the laplace domain. In International Conference on Machine Learning, pp. 8811-8833 PMLR, 2022.
Guido W Imbens. Nonparametric estimation of average treatment effects under exogeneity: A review. Review of Economics and statistics, 86(1) 36(1):4-29, 2004.
12

---

Published as a conference paper at ICLR 2024
Guido W Imbens and Donald B Rubin. Causal inference in statistics, social, and biomedical sciences. Cambridge University Press, 2015.
Edward L Ince. Ordinary differential equations. Courier Corporation, 1956.
Rakshit Jha, Mattijs De Paepe, Samuel Holt, James West, and Shaun Ng. Deep learning for digital asset limit order books. arXiv preprint arXiv:2010.01241, 2020.
Hong Jianwang, Ricardo A Ramirez-Mendozz, and Xiang Yan. Statistical inference for piecewise affine system identification. Mathematical Problems in Engineering, 2021:1-9, 2021.
Fredrik Johansson, Uri Shalit, and David Sontag. Learning representations for counterfactual inference. In International conference on machine learning, pp. 3020-3029. PMLR, 2016.
Kaitlyn E Johnson, Grant Howard, William Mo, Michael K Strasser, Ernesto ABF Lima, Sui Huang, and Amy Brock. Cancer cell population growth kinetics at low densities deviate from the exponential growth model and suggest allee effect. PLoS biology, 17(8):e3000399, 2019.
Thomas Kailath. Linear systems, volume 156. Prentice-Hall Englewood Cliffs, NJ, 1980.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv: 42260690, 2014.
Achim Klenke. Probability theory. universitext, 2008.
Helena Chmura Kraemer, G Terence Wilson, Christopher G Fairburn, and W Stewart Agras. Mediators and moderators of treatment effects in randomized clinical trials. Archives of general psychiatry, 59(10):877-883, 2002.
Milan Kuzmanovie, Tobias Hatt, and Stefan Feuerriegel. Estimating conditional average treatment effects with missing treatment information. In International Conference Artificial Intelligence and Statistics, pp. 746-766. PMLR, 2023.
Mathias Lechner and Ramin Hasani. Learning long-term dependencies in iregularly-sampled time series. arXiv preprint arXiv:2006.0444/ 2020.
Rui Li, Stephanie Hu, Mingyu Lu, Yuria Utsumi, Prithwish Chakraborty, Daby M Sow, Piyush Madan, Jun Li, Mohamed Ghalwash. Zach Shahn, et al. G-net: a recurrent network approach to g-computation for counterfactual prediction under dynamic treatment regime. In Machine Learning for Health, pp. 282-299. PMLR, 2021.
Bryan Lim. Forecasting treatment responses over time using recurrent marginal structural networks. Advances in neural information processing systems, 31, 2018.
Ernest Lindelöf. Sur l'application de la méthode des approximations successives aux équations différentielles ordinaires du premier ordre. Comptes rendus hebdomadaires des séances de l'Académie des sciences, 116(3):454-57 1894.
Lennart Ijuug. System identification. Springer, 1998.
Judith Lok. Statistical modeling of causal effects in continuous time. 2008.
Valentyn Melnychuk, Dennis Frauen, and Stefan Feuerriegel. Causal transformer for estimating counterfactual outcomes. In International Conference on Machine Learning, pp. 15293-15329, PMLR, 2022.
Daniel A. Messenger and David M. Bortz. Weak SINDy: Galerkin- -aseed Data-Don Data-Driven Model Selection. Multiscale Modeling & Simulation, 19(3): January 2021. ISSN 1540-3459, 1540-3467. doi: 10.1 137/20M 333000033331666
Dr Mould and Rn Upton. Basic Concepts in Population Modeling, Simulation, and Model-Based Drug Development. CPT: Pharmacometriic & Systems Pharmacology, 1(9):6, 2012. ISSN 2163-8306. doi: 10.1024/psp.20122..
13

---

Published as a conference paper at ICLR 2024
S Chandra Mouli, Muhammad Ashraful Alam, and Bruno Ribeiro. Metaphysica: Ood robustness in physics-informed machine learning. arXiv preprint arXiv:2303.03/81,,22 2023
Jersey Neyman. Sur les applications de théorie des probabilités aux experiences agricoles: Essai des principes. Roczniki Nauk Rolniczych, 10(1):1-51, 1923.
Kimia Noorbakhsh and Manuel Rodriguez. Counterfactual temporal point processes. Advances in Neural Information Processing Systems, 35:24810-24888 2022.
Judea Pearl, Causal inference in statistics: An overview. 2009.
Jonas Peters, Stefan Bauer, and Niklas Pfister. Causal models for dynamical systems. In Probabilistic and Causal Inference: The Works of Judea Pearl, pp. 671-690. ACM, 2022.
Brenden K Petersen, Mikel Landajuela Larma, Terrell Mundhenk, Claudio Prata Santiago, Soo Kyung Kim, and Joanne Taery Kim. Deep symbolic regression: Recovering mathemat- ical expressions from data via risk-seeking policy gradients. In International Conference on Learning Representations, 2020.
Zhaozhi Qian, Ahmed M Alaa, and Mihaela der Schaar. When and how lift the lockdown? global covid-19 scenario analysis and policy assessment using compartmental gaussian processes. Advances in Neural Information Processing Systems, 33: 13729-00100740, 2020.
Zhaozhi Qian, Krzysztof Kacprzyk, and Mihaela van der Schaar. D-code: Discovering closed-form odes from observed trajectories. In International Conference Learning Representations, 2022.
Carl Edward Rasmussen, Christopher KI Williams, et al. Gaussian processes for machine learning, volume Springer, 2006.
James Robins. Correcting for non-compliance in randomized trials using structural nested mean models. Communications in Statistics-Theon and methods, 23(8):2379-22 2412, 1994.
James M Robins. Association, causation, and marginal struetural models. Synthese, 121(1/2): 151-179, 1999.
James Robins and Miguel Heránn. Estimation of the causal effects of time-varying exposures. Longitudinal data analysis, 553:599, 2009.
James M Robins, Miguel Angel Hernan, and Babette Brumback. Marginal structural models and causal inference in epidemiology. Epidemiology. pp. 550-560, 2000.
Paul R Rosenbaum and Donald B Rubin. The central role of the propensity score in observational studies for causal effects. Biometrika, 70(1):41-55, 1983.
Donald B. Rubin. Comment on "randomization analysis of experimental data: The fisher ran- domization test". Journal of the American Statistical Association, 75(371):591, 1980. doi: 10.2307/2287653. URL https: /tw.:/ii.ii//// ort. 230.02230000 2287753.
Cynthia Rudin. Stop explaining black box machine learning models for high stakes decisions and use interpretable models instead. Nature machine intelligence, 1(5):206-215 2150 2019.
Pål C Ryalen, Mats J Stensrud, and Kjetil Røysland. The additive hazard estimator is consistent for comtinuous-time marginal structural models. Lifetime data analysis, 25:611-638,
1-638, 2019.
Päl Christie Ryalen, Mats Julius Stensrud, Sophie Fosså, and Kjetil Røysland. Causal inference in continuous time: an example on prostate cancer therapy. Biostatisttics, 21( 2020.
Olli Saarela and Zhihui Liu. A flexible parametric approach for estimating continuous-time inverse probability of treatment and censoring weights. Statistics in medicine, 35(23):4238-4251, 2016.
Michael Schmidt and Hod Lipson. Distilling free-form natural laws from experimental data. science, 324(5923):81-85, 2009.
Peter Schulam and Suchi Saria. Reliable decision support using counterfactual models. Advances in neural information processing systems, 30, 2017.
14

---

Published as a conference paper at ICLR 2024
Nabeel Seedat, Fergus Imrie, Alexis Bellot, Zhaozhi Qian, and Mihaela der Schaar. Continuous- time modeling of counterfactual outcomes using neural controlled differential equations. arXiv preprint arXiv:2206.0831/ 2022.
Alex Sherstinsky. Fundamentals of recurrent neural network (rnn) and long short-term memory (Istm) network. Physica D: Nonlinear Phenomena, 404:132306, 2020.
Hossein Soleimani, Adarsh Subbaswamy, and Suchi Saria. Treameentressonse models for coun- terfactual reasoning with continuous-time, continuous-values interventions. arXiv preprint arXiv:/704.02038, 2017.
Marshall H Stone. The generalized weierstrass approximation theorem. Mathematics Magazine, 21 (5):237-254, 1948.
Masaaki Takada and Hironori Fujisawa. Transfer learning via lell regularization. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances Neural Information Processing Systems, volume 33, pp. 14266-14277. Curran Associates, Inc., 2020. URL fi pdf.
dings www.neniiii.o httpp/ppeee.eleeepppeeeeeeeeeeee c c72bb17996243-
C. D. Thron. Linearity and Superposition in Pharmacokineties. Pharmacological Reviews, 26(1):
3-31, March 1974. ISSN 0031-6997, 1521-0081.
Tatiana Tommasi and Barbara Caputo. The more you know, the less you learn: from knowledge transfer to one-shot learning of object categories. In Proceedings of the British Machine Vision Conference, number CONF, 80-1, 2009.
Tatiana Tommasi, Francesco Orabona, and Barbara Caputo. Safety in numbers: Learning categories from few examples with multi model knowledge transfer. In 2010 IEEE Computer Society Conference on Computer Vision and Pattern Recognition, pp. 3081-3088. IEFF, 2010.
Lloyd N Trefethen, Ásgeir Birkisson, and Tobin Driscoll. Exploring ODEs. SIAM, 2017.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
Bochen Wang, Liang Wang, Jiahui Peng, Mingyue Hong, and Wei Xu. The identification of piecewise non-linear dynamical system without understanding the mechanism. Chaos: An Interdisciplinary Journal of Nonlinear Science, 33(6), 2023.
Ronald J Williams and David Zipser. A learning algorithm for continually running fully recurrent neural networks. Neural computation, 1(2):270-280, 1989.
Yanbo Xu, Yanxun Xu, and Suchi Saria. A bayesian nonparametric approach for estimating individu- alized treatment-responsse curves. In Machine learning for healthcare conference, pp. 282-300. PMLR, 2016
15

---

Published as a conference paper at ICLR 2024
Appendix
Table of Contents
Future Work
17
B Explanation of Treatment Intensity Processes and Filtrations in the Assumptions
17
C Extended Related Work
18
C.I. Related work approaches
18
C.2 Related work methods
19
D ODE formalism for treatment effect estimation
20
E Treatments
20
F Benchmark Dataset Details
21
F.1 Cancer PKPD
22
F22 Synthetic Equation Classes
22
G Benchmark Method Implementation Details
23
G.1 INSITEs inference procedure
26
H Evaluation Metrics
27
Dataset Generation and Model Training
27
Interpretable Equations
28
K Additional Experiments
28
K.1
Sensitivity À
28
K.2
Parametric distribution of numeric constants
28
K.3
Additional main results
29
K.4
Additional results with varying sample sizes
29
K.5
Additional results with incrcasing observation noise
31
K.6 Model Misspecification
31
K.7 INSITE also empirically works for irregularly sampled data
34
K.8 INSITE can also work with challenging observation functions g(x)
34
L Reconciling Assumptions in Table 1
35
M Scale of the outcomes
35
How to choose an ODE discovery method
Numerical Solver and Parameters
36
Code. All code is available http://///athhaaoo//oo/oooooooor or have broader research group codebase http://..oooooaaro/oe/////aaaa/a /CDE Di scovery for wwwperrr-foo--ooiitooooe orwpre/woo..ogggiiaaaaaaeeeeeeee ts Io-feeereeee
We
16

---

Published as a conference paper at ICLR 2024
AUTHOR CONTRIBUTIONS
All authors provided valuable contributions to the paper from the idea, writing, and editing, and managing the project's progress.
KK wrote the formalism section 3, proposed (and described) how to adapt ODE discovery methods to include discrete treatment plans (section 4.2) and different layers of between-subject variability (section 4.3), constructed the initial experimental settings in table 8, and provided the idea for fig. 2.
SH wrote parts of the introduction, all the ODE assumptions in section 3; proposed and wrote the method in section 5, developed and wrote all the experiments in section 6, and the conclusion in section 7; designed and came up with the method INSITE; solely coded up the method and all baseline implementations, initial exploratory experiments for INSITE; all coding and experiment design, implementation and running for all nine datasets, seven baselines, evaluation; led all results within the paper in section 6, and eight additional new experimental settings and ablations of all the baselines across all datasets in appendix K; solely wrote 20 appendix sections, of additional experiments, ablations, explanations of key parts of the paper, additional related work, baseline descriptions, and synthetic dataset descriptions.
JB connected causal treatment effect estimation with ordinary differential equation discovery leading to the framework presented in this paper.
A FUTURE WORK
In our paper we have been very explicit about the settings in which we can use ODE discovery settings (see section 3 in particular). As these settings do not correspond one-to-one with standard TE settings, our paper can be viewed more to be done! In particular, future endeavours could include: discovering stochastic differential equations, leveraging latent variable models for handling unobserved variables, learning piece-wise continuous ODE systems (Wang et al., 2023), and developing strategies to transfer population-level ODEs other populations where there is limited data. Furthermore, INSITE could be considered first link in connecting equation discovery over time with causality (Berrevoets et al., 2023), we hope many more works of this type will follow.
expansion of typical TE application areas. However, there is still
B EXPLANATION OF TREATMENT INTENSITY PROCESSES AND FILTRATIONS IN THE ASSUMPTIONS
The concept of treatment intensity processes has its roots in the generalization for treatment effects in continuous time, introduced in seminal works such Lok (2008) and later embraced by Saarela & Liu (2016). Intensity processes serve as an extension of propensity scores, traditionally denoted as P(]]),, continuous time framework.
In treatment effects literature, propensity scores provide the probability of treatment assignment conditioned observed covariates. However, when dealing with cantinuous-time data, the traditional propensity scores pose a challenge. Specifically, propensities in this setting would govern entire treatment trajectories, considering not just the treatment at specifie time point T, but all potential future treatments. Given the vast space of potential trajectories in discrete time, this complexity becomes unmanageable in continuous time.
To address this challenge, filtrations are employed to restrict when new treatment can be sampled. Filtrations essentially offer controlled space for the random processes play. Using the terminology from probability theory:
For a stochastic process (X,)neN, the filtration F, oXxikk S n)) a alagebrr.. This means the filtration IF = (Fx) )ee limits the stochastic process to only those random variables X, for which S n, even though the original probability space might encompass variables with n.
Such a structure becomes increasingly important in continuous time applications, as it enables us to understand and manage the dependencies among observations over time. Furthermore, intensity
17

---

Published as a conference paper at ICLR 2024
processes in this framework be analogous to well-known processes in other contexts, such as the Hawkes point processes (that model excitation) or the Poisson point processes (which assume no influence from past observations).
Now, relating back to assum. 2.2 (Overlap) and assum. 2.3 (Ignorability): The intensity process, represented by ((f(3e), captures the propensity of individual receiving treatment at time f given all the information up to that point, encapsulated in the filtration 8,- The overlap assumption ensures that the intensity process is never deterministic, implying that there's 's always some randomness in treatment assignment irrespective of past information. On the other hand, the ignorability assumption implies that the intensity process conditioned on past information (up time the same as the intensity process that considers even the future outcomes beyond time t. This ensures that treatment assignment is independent of potential outcomes, thereby satisfying a key requirement for unbiased causal inference.
C EXTENDED RELATED WORK
C.I RELATED WORK APPROACHES
This research brings together two previously separated fields, namely temporal treatment effect estimation and ODE discovery, and there exists a plethora of works within each domain. Here, we will focus on the most relevant and recent ones.
Treatment Effects over Time. There has been substantial research on the problem of estimating treatment effects. One of the pioneering works in this area is the potential outcomes framework introduced by Neyman (1923) and further developed by Rubin (1980).
Traditional methods for estimating treatment effects include propensity score matching (Rosenbaum Rubin, 1983), instrumental variables (Angrist al., 1996; Angrist, 1991), and difference-in- differences (Athey & Imbens, 2006). However, these methods are primarily suited for static treatment settings and may have limitations when extended handle treatment effects over (continuous) time.
The treatment effects literature is well represented in discrete time settings. With more traditional work stemming from epidemiology, based on g g-computation, Structural Nested Models, Gaussian Processes, and Marginal Structural Models (MSMs) (Robins, 1994; Robins & Hernán, 2009; Xu et al., 2016; Robins et al., 2000); and more recent work relying on advances in deep learning architectures (Bica et al., 2020b; Melnychuk et al., 2022; Berrevoets et al., 2021: Lim, 2018). Some works have expanded the static treatment effect literature to the continuous time setting (Lok, 2008; Saarela & Liu, 2016; Ryalen al., 2019), introducing Nevertheless, these works still face some critical challenges. For example, they often assume constant treatment effects and do not naturally handle irregular sampling and continuous treatments.
assumptions help identify potential outcomes.
ODE discovery. The field of ODE discovery has also seen significant progress over the years. One of the primary goals of ODE discovery is learn model of a system's 's behaviour from observational data. This involves discovering the underlying system of ODEs that governs the system.
The methods that solve this problem range from linear ODE discovery techniques (Ljung, 1998) to more recent machine learning approaches like Recurrent Neural Networks (RNN) and Long Short- Term Memory (LSTTM) (Hochreiter & Schmidhuber, 1997; Sherstinsky, 2020). However, these later methods often produce black-box models that lack the interpretability required for many applications, including those in regulated domains.
Recent efforts in the MI. community have focused on finding interpretable models using techniques such as Sparse Identification of Nonlinear Dynamics (SINDy) (Brunton et al., 2016), which discovers sparse dynamical systems from data. However, traditional ODE discovery methods largely focus on homogeneous systems where the system dynamics are the same across different systems or individuals. They do not readily handle the heterogeneity of treatment effects (between-subject variability), which is a major concern in real-world applications like personalized medicine.
We consider ODE discovery to be sub-field of the wider Symbolic Regression field, which is concerned with discovering a more general class of equation (beyond ODEs) (Camps-Vare Valls et al., 2023; Mouli et al., 2023; Cranmer al., 2020).
18

---

Published as a conference paper at ICLR 2024
Dynamical systems, ODEs, and treatment effects. We want to explicitly state that, ODE discovery is a subfield of the broader area of dynamical systems identification. We believe this to be necessary as often work is presented as an identification method for dynamical systems, while they are only applicable to ODEs which is only one particular type of dynamical system. That means that works such Schulam & Saria (2017); Hizh et al. (2022); Noorbakhsh & Rodriguez (2022); Ryalen et al. (2020); Alaa & van der Schaar (2019); Qian et al. (2020); Soleimani et al. (2017) may appear related as they model treatment effects in type of dynamical system. However, is crucial note that none of these works model treatment effects as an ODE, nor are they concerned with actually uncovering the underlying ground truth dynamical system in general. Instead, they model treatment effects as a parameterized dynamical system (such as, for example, a point process (Noorbakhsh & Rodriguez, 2022; Alaa & van der Schaar, 2019)), which is learned purely for accurate inference. In contrast, while we are also interested in accurate inference, the resulting model (the ODE) is additionally important.
Fusing Treatment Effects and ODE discovery. Our work is the first, to our knowledge, to combine these two fields for the problem of treatment effect estimation. We seek to harness the strengths of both fields to address their respective challenges, such as the lack of interpretability and robustmess in the treatment effects literature and the assumptions about system behaviours in the ODE discovery community. In particular, our INSITE approach is unique in its ability to address the challenges from both communities and thus provides a promising new direction for these intertwined fields.
There are some ideas that seem to correspond, but only first glance. For example Florens et al. (2008) connects control (a field which is very much related to ODE discovery), yet differs one some crucial points: (i) time is not considered; (ii) no equations are discovered; (iii) the assumption sets differ widely.
C.2 RELATED WORK METHODS
Longitudinal Treatment Effect Methods. Longitudinal treatment effect methods focus on estimating the outcome of treatment at given time t,, often using black-box neural network-basee approaches, iee. not closed-form or easily interpretable. A significant theme in these works is the development of methods to mitigate time-dependent confounding. Many mitigation methods have been proposed that we use as benchmarks. such as propensity networks in Marginal Structural Models (MSM) (Robins et al., 2000) and Recurrent Marginal Structural Networks (RMSN) (Lim, 2018), Gradient Reversal in Counterfactual Recurrent Networks (CRN) (Bica et al., 2020b), Confusion in Causal Transformers (CT) (Melnychuk et al., 2022) and g-computation in G-Net (G-Net) (Li et al., 2021).
ODE Discovery Methods. ODE discovery methods aim to discover an underlying closed-fome mathematical equation F that best fits the underlying controlled ODE of the observed state and action trajectories dataset, D. In these methods, the state derivative is unobserved; therefore, they have to approximate the derivative using finite differences, as in Sparse Identification of Nonlinear Dynamics (Brunton et al., 2016), or employ variational loss as the objective, in Weak Sparse Identification of Nonlinear Dynamics (Messenger & Bortz, 2021). Underlying these methods is the broader class of methods for symbolic regression (Holt et al., 2023b). To make these two ODE discovery methods more competitive, we adapt them to model individual ODEs per categorical treatment (section 4.2), (termed A-SINDy, A-WSINDY respectively).
Other Related Approaches. In addition to the aforementioned categories, several other methods have been proposed model longitudinal data and estimate treatment effects. For instance, recurrent neural networks (RNNs) (Cho et al., 2014) and long short-term memory (LSTM) networks (Hochreiter & Schmidhuber, 1997) are commonly used for handling sequential data, as well as other sequential time series models (Holt et al., 2022), and they have been applied to estimate treatment effects in longitudinal settings (Lim, 2018). Gaussian processes (Rasmussen et al., 2006) and continuous- time Gaussian processes (Alvarez. & Lawrence, 2009) have also been employed for modelling and learning from longitudinal data. Furthermore, counterfactual regression (Johansson et al., 2016) and doubly robust estimation (Funk et al., 2011) are statistical methods used to estimate causal effects in observational studies, and they can be adapted to handle -meedeepnndett confounding in longitudinal settings. Additionally, once model is learned, it can be used for optimal planning (Holt al., 2023a; 2024), and such time-series models can also applied to other related domains such as in finance (Jha et al., 2020).
19

---

Published as a conference paper at ICLR 2024
A more detailed comparison of the related approaches discussed above can be seen in table 5. This table compares various aspects of these methods, such as interpretability. eetween-subjec variability, continuous-timg time modelling, robustness to observation noise and the ability to handle categorical and continuous treatments, and data size requirements.
Table 5: Comparison of related works. Columns: Interpretable? can it provide closed-fore form equation? BSV? can between-subjed variability be modeled?? Continuous timtussstimemm--- is it a continuous time model (i.c., able to learn from irregularly sampled observation and treatment trajectories naturally)? Noise? is it robust to observation noise? Categorical E z?- can it model categorical treatments (inclusive of binary treatments)? Continuous E R?-can can model continuous treatments?
Approach
Ref
Interpeetable?
BSV?
Continoou- time?
Noisc?
Categorical
E 2?
Continuous
R?
Longiudinal Treatment Effect Methods
MSM
(Robins al., (Lim, 2018)
RMSN
CRN
(Bica et al., 2020b)
x
G-Net
(Li al., 2021)
(Melnychuk al., 2022)
ODE discovery Methods
A-WSINDy
A-SINDy
(Messenger Bartz, 2021)
(Brunton 2016)
-population only
populaation only
x /
/
INSITE
(Ours)
/
D ODE FORMALISM FOR TREATMENT EFFECT ESTIMATION
In this section, we want to justify our ODE formalism for treatment effect estimation, eq. (3), and explain its generality.
We assume that the treatment effect is deterministic function of the modelled covariates (often the identity). This is inspired by examples considered in treatment effect papers, where we equate the treatment effect the tumour volume or the response drug given a particular concentration (through a dose-response curve) (Bica et al., 2020b; Berrevoets et al., 2021; Lim, 2018; Seedat et al., 2022).
We want to emphasize that direct link between a and y can be included through extended covariates x' (t) (zxt), y()) and the function that links covariates to the outcome as g g(x, y) = y. Then the new treatment effect variable is defined as y (l) g( (()) v(t). This is the same as treating one of the covariates as the outcome. Please note that our formalism extends this setting and allows for more complicated treatment outcome that depends several covariates.
E TREATMENTS
In the following. we summarize the ways of incorporating the treatment plan in F according to the treatment types of binary treatments, categorical treatments, multiple simultancous treatments, and continuous treatments. This is necessary so that can simplify the complex and possibly not closed-form F into simpler closed-form f's that we can discover using current methods. In this assumed be closed-form.
subsection, every function denoted by the letter f
As described in Bica et al. (2021) there are four main types of treatments considered in the treatment effect estimation literature
* Binary treatment (treatment or no treatment)
Categorical treatment (one treatment out of multiple possible treatment options)
Multiple treatments assigned simultaneously
Single treatment or multiple treatments with associated dosage
Each of these treatments can be either static (constant throughout the trajectory) or dynamic. As we are interested in the discovery of closed-form ODEs, the treatment can be incorporated into F in two ways.
We learn different closed-form ODEs for different treatments
20

---

Published as a conference paper at ICLR 2024
We incorporate the action variable into the closed-form ODE.
Depending on the type of treatment one or both of these approaches can be chosen. We summarize the ways of treating in F based on the type of treatment in Table 6.
Binary treatment considers only two kinds of treatments (usually "treatment" and "no treatment") The action trajectory is one-dimensional with values in {0, 1}. In static setting, is constant, in dynamic setting it is piece-wise constant-that means that we are allowed to change the treatment throughout the trajectory. This kind of action can be incorporated in F in two ways. We can either learn two different closed-form systems ODEs (fo and fi) or incorporate a directly in the equation, e.g., fo + fi
Categorical treatment considers one treatment out of multiple possible treatments. The action trajectory is deeedimensional with values in K). As previously, is constant in static setting and piece-wise constant a dynamic setting. We no longer can easily incorporate a into the closed-form expression (as it attains K discrete values), so the only option is learn K separate closed-form systems of ODES fI...., JK-
Multiple treatments considers scenarios where multiple treatments be assigned simultaneously. We represent a as a vector (0, 1) where (t) = the 7th treatment was assigned at time t. As previously, constant in static setting and piece-wise constant in dynamic setting. We can consider this case as just having 2 treatments (all possible subsets) and reduce to categorical treatment where we learn 2 separate equations. terms for each treatment and combine them using principle of superposition (Thron, 1974), We can then represent the system of ODEs at time t as Σ ft(t) v). However, no current ODE discovery algorithm leverages this representation.
alternative approach would be learn separate
Continuous treatment is usually considered when want model the dosage the strength of the treatment. a(t) is represented as a K dimeensiionl real vector where (t is the dose/strength of the 7th treatment. In static setting, consider to be continuous. As (t(() can have any real value, the only way to incorporate it into the equations have it directly in the closed-form ODE (=((), v, (()))
be constant. In dynamic setting, consider it
Table How different treatment types can represented an ODE formalism. Here corresponds to a static or dynamic treatment.
Treatment
S/D
Domain of a
Constant
F(z(t),
Yes
(x()))(())
Binary
a(t) {0, 1}
D
Piece-wise
fo(x(1), v) + (z()))
Yes
Categorical
a(t) [L
fa(o(z(t).
D
wise
S
faco(a((),,)) Yes
Multiple
D
a(t) E 1)
Piece-wise
or
ΣK (()((t))(()(tte.))))))))))))))))
S
Yes
Continuous
att) E R
s(x(1), v, a()(
No
F BENCHMARK DATASET DETAILS
In the following we outline the standard Cancer PKPD equation, then outline the between-subject variability layers that equation models. We also provide the parameter distributions used to generate the different equation classes. For dataset generation and model training see appendix I.
to generate the different equation classes from the two base PKPD
21

---

Published as a conference paper at ICLR 2024
F.1 CANCER PKPD
This is a state-of-the-ar biomedical aatiaaooknneii--Paaaaadddddaee (PKPD) model of tumour growth, used to simulate the combined effects of chemotherapy and radiotherapy in lung cancer (Geng et al., 2017) (eq. (7))-this has been extensively used by other works (Seedat et al., 2022: Bica al., 2020b; Melnychuk al., 2022). Specifically, this models the volume of the tumour xx() for days after the cancer diagnosis- -hhree the outcome is one- mne-dimensional. The model has two binary treatments: (1) radiotherapy at and (2) chemotherapy oi.
dxx() (plog dt
K
B.C(L)
(a,,d(1)) + 3,d(1)3) et )a(()
x(1)
Chemotherapy
Radiotherapy
Noise
Tumorgrowth
Where the parameters K, P. Beg Or Br for each simulated patient are sampled distributions detailed in Geng et al. (2017), which are also described in table 7. Here e, ~ (O0 0.012) is random noise, modelling randomness in the tumour growth.
Table 7: Cancer PKPD parameter distrihutions.
Model
Variable
Parameter
Distribution
Parameter Value (p, σ))
Growth parameter
Normal
7.00 x 10-5, 7.23 x 10 10-3 30
Tumor growth
k
Carrying capacity
Constant
Radiotherapy
Radio kill (a) Radio cell kill (3)
a
Normal
0.0398, 0.168 Set s.t. a B-1010
B,
Chemotherapy Chemo kill
Bo
Normal
0.028, 0.0007
Furthermore, we incorporate heterogeneous responses, following Bica et al. (2020b); Lim (2018); Melnychuk al. (2022), where the means are modified for 3e and a, by creating three groups of patients (i.e., to represent three types of patients with heterogeneity in treatment response).
modify the mean of a, so that u(a,)) = 1.1 p(a,) and for patient group we modify the mean of a, so that μ(az) = 1.1 x p(az).
For patient group
Additionally, the chemotherapy drug concentration e(t) follows an exponential decay relationship with half-life of one day:
dc(t) -.55c(t)
(8)
dt
where the chemotherapy binary action represents increasing the c(t) concentration by 5.0mg/ m mmn Vinblastine given at time t.
of
Whereas the radiotherapy concentration d(() represents 2.0Gy fractions of radiotherapy given at timestep f, where Gy is the Gray ionizing radiation dose.
Time-dependent confounding, We introduce vire-ig confounding into the data generation process. This is accomplished by characterizing the allocation of chemotherapy and radiotherapy as Bernoulli random variables. The associated probabilities, Pe and pr, are determined by the tumor diameter as follows:
(1(() o (
D ( (D(1 (()))))))
Pr p+(1)
= p=(1==0
( 8,)
(Dmax
(9)
where Dmax 13cm represents the largest tumor diameter, 0, - O, Dmax/2 and D(I) signifies the mean tumor diameter. The parameters To and Yr manage the extent of time-varying confounding. Higher values of Yerr amplify the influence of this confounding factor over time.
F.2 SYNTHETIC EQUATION CLASSES
We outline the two synthetic equations, which then use in the four layers of between- subject variability settings A-D table which differ between noise, static covariates, and parametric distributions of parameters. The two synthetic equations study are standard PKPD models,
22

---

Published as a conference paper at ICLR 2024
the first being one-comparttmental PKPD model with binary static action. The second is the same state-of-thh-arr art biomedical PKPD model of tumor growth from the Cancer PKPD dataset (Geng et al., 2017), where the model has continuous chemotherapy treatment c(t) and a binary radiotherapy treatment d(t). Therefore the volume of the tumor days after diagnosis for each model is given by:
dx(t) dt
Ger(t), if = 0 [ r(t),
(10)
i
dx(t) dt
K
(plog
B.C(t)
(a,d(t) + B,d(t)3) a )=(t)
(11)
x(t))
Chemotherapy
Radiotherapy
Noise
Tumorgrowth
where the parameters Co. C1, V, K,,,a, 8, et are sampled according to the different layers of between-subject variability (table 2) forming variations of A-D, following the setup in table 8,
Specifically, for the one-compartmental PKPD model of eq. (10), co (0.5, 0.05), e1 ~ (0.5, 0.05). V = 1 and has additive observation noise of E (.000.01). Moreover, wp- 1.0, w1 = 0.05, w2 = 1.0, W/y 0.15.
For the tumour growth PKPD model, eq. (11), we ablate the full model to create a homogencous version of one patient type (UPaiient Type 1) for the noise BSV layer (A) and noise BSV layer (B), then reintroduce patient heterogeneity (Upatiem Type E {1. 2, 3)) and restrict B. to only be the mean value for the BSV layer of covariates (C). Finally, the full BSV layer (D) of noise, covariates and parametric distribution of parameters is the full Cancer PKPD model. Here unless otherwise noted the other parameters are still sampled from their respective defined distributions, outlined in table Here we follow the same treatment assignment policy as the Cancer PKPD model above, and as the LTE methods are only designed to handle binary treatments, we provide the LTE methods with the binary treatment assignments, and the corresponding ODE discovery methods with the continuous value of c(t) and the binary value of d(t).
demeedeeendentt confounding Similarly to the Cancer PKPD model, we also introduce time- dependent confounding the data-generating process, following a similar setup. We achieve this by characterizing each treatment assignment as Bernoulli random variables. Each associated probability p depends on the outcome normalized value. Here, we can vary the scalar value 7 E [0, inf) to increase time-dependent confounding.
ptt) a (7 (ū(() 0.5))
(12)
(tt) Bern(p(t))
Where g(() is rolling window average of the previous outcomes, with window size of w = 15.
G BENCHMARK METHOD IMPLEMENTATION DETAILS
We seek to compare against the existing tfate-of-the-an the-art art (SOTA) methods from (1) longitudinal treatment effect models and (2) ODE discovery. First, longitudinal treatment effect models often consist of black-box neural network-based approaches, i.e., not closed-form or easily interpretable. significant theme in these works is the development of methods mitigate time-dependent confounding. Many mitigation methods have been proposed that we use as benchmarks, such as propensity networks in Marginal Structural Models (MSM) (Robins et al., 2000) and Recurrent Marginal Structural Networks (RMSN) (Lim, 2018), Gradient Reversal in Counterfactual Recurrent Networks (CRN) (Bica et al., 2020b). Confusion in Causal Transformers (CT) (Melnychuk et al., 2022) and g-computation in G-Net (G-Net) (Li et al., 2021). Second, ODE discovery methods aim to discover an underlying closed-form mathematical equation F that best fits the underlying controlled ODE of the observed state and action trajectories dataset, D. In these methods, the state derivative is unobserved; therefore, they have to approximate the derivative using finite differences, as in Sparse Identification of Nonlinear Dynamics (Brunton al., 2016), employ variational loss as the objective, as in Weak Sparse Identification of Nonlinear Dynamics (Messenger & Bortz, 2021). To
'We only include results for WSINDy when it is possible use it, as it does not support sparse short. trajectories.
23

---

Published as a conference paper at ICLR 2024
Table 8: Parameter values for the different layers of between-subject variability (A-D) for the two PKPD models.
Daraset
Outcome (y)
Covariates
Parameters
Treatment
Co Co
oq. (5)AA
y(t) = =(f)
cg.61
C
a() a(0) e (0.1)
Co Cn
oq. (5)B
u(t) = x +6
a() (0.
C G
Co 4g W1
oq. (5)C
M(t) = x(1) +
Ca.
a0)
C ~ Mroig wn 0.25)
oq. (5).D
M(t) = x(1) +
00.
C. M(e ws. 0.25)
= 6 (0.1)
m(o.) = 1.1 ((.))
og. (6).
v(1) = 2(()
Ploment
8. p().)
(c(t).
m(a,)
1.1 x ((a.)
oq. (6.B
w(t) = x(1) +e
Ulstient Type
8. p(3.)
a= (c(t).
1: ((a,)
((r)
eq (6).
v(t) xl1
OPaiicat Type (1,
Opatica Type
3: p(a.) = 1.1 p(a.)
(e(t), d())
2: No modificution
Be p(3.)
m(a,) ((ae)
eq. (6)D
y(t) = a
Opsiin Typ E 21,
Oputient
3: =(a,)= x p(o,)
(c())
2: No modificution
8. ~ -((((3).
make these two ODE discovery methods more competitive, we adapt them to model individual ODEs per categorical treatment (section 4.2), (termed A-SINDy, A-WSINDy respectively). We further discuss why we chose these state-of-the-ar methods and their associated implementation details in the following.
To make the LTE methods competitive, the hyperparameters are tuned following Melnychuk et al. (2022), the Cancer PKPD dataset-and and then kept constant for all experiments. Specifically, use Melnychuk et al. (2022)'s 's tuned hyperparameters, which are detailed below for each LTE method. also detail hyperparameter tuning appendix I.
LONGITUDINAL TREATMENT EFFECT MODELS
Marginal Structural Models (MSMs) Robins al. (2000); Hernán al. (2001) are popular methods in epidemiology designed for the estimation of counterfactual outcomes with inverse probability of treatment weights (IPTW) through linear modelling. This approach utilizes stabilized weights to eliminate time-narying confounding bias.
Upon estimation the stabilized weights, they are normalized and truncated their 1st and 99th percentiles to align with Lim (2018). Outcome regressions are applied separately for each prediction horizon. For any given horizon т, the dataset is split into smaller segments using a rolling origin. and stabilized weights computed for each segment.
No hyperparameters are required for MSMs; hence, in all experiments, training and validation subsets combined.
Recurrent Marginal Structural Networks (RMSNs) Lim (2018) use a gequence-to-seuuemee architecture composed of four LSTM subnetworks. It manages multiple binary treatments by re- weighting the objective with the IPTW during training, creating a pseudo-population that emulates a randomized controlled trial.
24

---

Published as a conference paper at ICLR 2024
The propensity networks are initially trained to estimate the stabilized weights. Subsequently, the encoder is trained using a mean squared error weighted with the stabilized weights for one-step-ahead predictions. Lastly, the decoder is trained by minimizing the loss using the fully stabilized weights. The dataset is processed into smaller chunks using rolling origins for this stage of training.
Here the hyperparameters are: the propensity treatment model has 8 sequential hidden units, dropout rate of 0.1, one layer, uses a batch size of 64, with a max grad norm of 2.0, and is optimized with the Adam optimizer (Kingma & Ba, 2014) with a learning rate of 0.001. The propensity history model has 16 sequential units, with a dropout rate of 0.3, one layer, uses a batch size of 256, with max grad norm of 1.0, and an Adam optimizer with learning rate of 0.01. The encoder model has 12 sequential hidden units, a dropout rate of 0.1, one layer, a batch size of 64, a max grad norm of 2.0, and an Adam optimizer with learning rate of 0.001. Moreover, the decoder model uses 64 sequential hidden units, with dropout rate of 0.2, one layer, a batch size of 256, a max grad norm of 1.0, and an Adam optimizer with a learning rate of 0.001.
Counterfactual Recurrent Network (CRN) Bica et al. (2020b) utilizes an encoder-decoder archi- tecture, applying an adversarial learning technique, gradient reversal Ganin & Lempitsky (2015), to create balanced representations non-predictive of the treatment assignment.
CRN's encoder and decoder are composed of a single L.STM-layer each. Training involves minimizing a loss function that applies gradient reversal to minimize cross-entropy between the predicted and current treatment, while simultaneously maximizing the entropy of the built representations.
Here the hyperparameters are, the encoder model has balancing representation size of 6, 18 fully connected hidden units, with a dropout rate of 0.2, 24 sequential hidden units, uses a batch size of 64 and is optimized with Adam optimizer with learning rate of 0.01. Moreover, the decoder model has a balancing representation size of 3, 9 fully connected hidden units, with a dropout rate of 0.2, 24 sequential hidden units, uses batch size of 512 and is optimized with Adam optimizer with learning rate of 0.001. To make this baseline as competitive as possible, we use the domain confusion balancing loss from Melnychuk et al. (2022).
G-Net Li et al. (2021) is a recent method for estimating time-varying treatment effects, based on the G-computation formulation. It leverages a recurrent architecture that directly models the evolution of treatments and outcomes over time.
The network designed to simultaneously model the auto-regressive dynamics of outcomes and the time-varying effects of treatments. For each time step, the model uses LSTM cells to capture the auto-regressive outcome dynamics and a separate LSTT layer to model the time-varying treatment effects.
The model trained using a mean squared error loss for outcome prediction and binary cross-entropy loss for treatment assignment prediction.
Here the hyperparameters are, 24 sequential hidden units, 48 fully connected hidden units, one layer, r size of 3, with a dropout rate of 0.1, uses a batch size of 128, 25 Monte Carlo samples and is optimized with an Adam optimizer with learning rate of 0.01.
Causal Transformer (CT) Melnychuk et al. (2022) is recent state-of-the-ar method for treatment effects over time. is based on the Transformer architecture Vaswani et al. (2017), which is effective in modelling rongerange dependencies in sequential data. CT uses causal attention mechanism to model the temporal dynamics of the treatment effects. It is composed of three transformer subnetworks with separate inputs for time-varying covariates, previous treatments, and previous outcomes into a joint network that has in-between cross-attentions.
To encourage balanced representations that are predictive of the next outcome but predpiedie of the current treatment assignment, CT optimizes a loss that is composed of two terms. The first term supervised loss, such as Mean Squared Error (MSE), which measures the difference between the predicted and actual outcomes. The second term is a domain confusion loss, which encourages the model to minimize the difference between the distributions of representations conditioned on different treatments.
Here the hyperparameters are 16 sequential hidden units, balancing representation size of 16, 32 fully connected units, a dropout rate of 0.1, a batch size of 256, and is optimized with an Adam optimizer with a learning rate of 0.01.
25

---

Published as a conference paper at ICLR 2024
ODE DISCOVERY METHODS
Sparse Identification of Nonlinear Dynamics (SINDy) (Brunton al., 2016), data-driven frame- work that aims to discover the governing dynamical system equations directly from time-series data. The algorithm works by iteratively performing sparse regression library of candidate functions to identify the sparsest yet most accurate representation of the dynamical system.
use a polynomial library of order two, which feature library of == {1, 2o, I ). Finite difference approximations are used to compute time derivatives from the input time-series data, of order one. Here the alpha parameter is kept constant at 0.5 across all experiments, and the sparsity threshold is set to 0.001 for all experiments, apart from the eq. (5) datasets where it is to 0.1.
In our implementation,
Weak Sparse Identification of Nonlinear Dynamics (WSINDy) (Messenger & Bortz, 2021) is a variant of SINDY designed to handle noisy sampled time-series data. Instead of directly differentiating the time-series data to obtain derivatives, WSINDy formulates a variational optimization problem that simultaneously identifies the governing equations and estimates the derivatives.
The implementation details are largely similar to those for SINDy, with the significant difference being the formulation of the loss function for the optimization problem. Here, we use 100 domain centres and polynomial library of order two. Specifically, WSINDy requires multiple domain centres in the variational formulation, which precludes the application of WSINDy to sparse and short sequence trajectories- leading the inclusion of this baseline only where it is possible to apply (when the minimum trajectories are long enough that it can be applied).
To form the Adapted SINDy (A-SINDy) and Adapted WSINDy (A AA-WSINDy), modified the SINDY and WSINDy methods to handle categorical treatments by learning individual ODEs for each treatment group. This adaptation makes these models more competitive for the benchmark comparisons in this study.
Individualized Nonlinear Sparse Identification Treatment Effect (INSITE), our proposed method, builds on top of SINDy that discovers population (global) differential equation for all patients in the training dataset. Therefore uses the same hyperparameters as defined above for SINDy.
Moreover, INSITE fine-tunes the discovered population ODE to existing observed patient trajectories by minimizing the MSE loss between the predicted and observed outcomes. Here the regularization parameter À is set to À 10.0 across all experiments. It was set following the same hyperparameter tuning procedure on the validation dataset on the Cancer PKPD dataset and then kept constant for all experiments.
G.1 INSITES INFERENCE PROCEDURE
INSITE principally consists training step and an inference step. First, for the training step, discovers a population (global) differential equation for all patients in the training dataset. The exact form of this global differential equation that is discovered depends on the treatment type, as following the types outlined in appendix E. Commonly, the treatment will be a categorical treatment, where the action trajectory is one- -ie-imneniional with values in [1, K1, therefore the form of the global ODEs that are learned are K separate dooed-forrm systems of ODEs fi....,fk:.
For the inference step at run-time, a new (unseen) patient is observed with covariates and an initial treatment history following a specified treatment plan. INSITE fine-tunes the active K discovered population ODEs to existing observed patient trajectories by minimizing the MSE loss between the predicted and observed outcomes, for the active th separate closed-form ODE as determined by the existing observed treatments. Given this, there can exist case where a discrete treatment type may not be observed in the existing patient observed treatment history, in that case, the separate ODE that corresponds that discrete treatment is not fine-tuned, and instead is left as the underlying population (global) differential equation for all patients in the training dataset with that specific discrete treatment.
26

---

Published as a conference paper at ICLR 2024
H EVALUATION METRICS
We evaluate against the standard longitudinal treatment effect metrics of test counterfactual 7-step ahead prediction normalized root mean squared error (RMSE), where {1, 3, 4, 5, 6), for sliding treatment plan (Bica et al., 2020b; Meinychuk et al., 2022). Unless otherwise specified, each result run for five random seed runs with time-dependent confounding of = 2.
To evaluate each baseline, we use the same setup as Melnychuk et al. (2022). That is for each dataset, for each patient in the test set and for each time step, we simulate multiple counterfactual trajectories by setting the treatment to counterfactual possible treatment values, depending on т. First, for one-step ahead prediction T = 1, we simulate all the possible counterfactual treatment values (e.g., in eq. (5) we simulate both = 0 and = 1, and eq. (6) & Cancer PKPD we simulate all four combinations of {(aī = 0, aj = 0), (a, = 1, af 0), (aî 0, at = 1), (a: = 1, of 1)}) for the next counterfactual outcome yt + 1). This corresponds to the PKPD model under all the feasible treatment assignments. Second, for multi-step ahead prediction, the number of possible potential outcomes yt + 2. ...., + mmax) grows exponentially with the projection horizon maax- Therefore we use the sliding window treatment assignment (Bica et al., 2020b; Melnychuk et al., 2022). This sliding window treatment assignment can help test that the correct timing of treatment is chosen. It is implemented by simulating trajectories with a single treatment, however, the treatments are iteratively moved over a window ranging from to + Tmax 1, this essentially results in [[Tmax 1) trajectories.
For each random seed run, we generate new train, validation and test dataset independently. With 1000 training trajectories, 100 validation trajectories and 100 test trajectories, unless otherwise noted. We then train each baseline on the training dataset and tune the hyperparameters on the validation dataset. We then evaluate the performance of each baseline on the test dataset. We repeat this process for each random seed run, for a total of five random seed runs, unless otherwise noted. We then report the mean and 95% confidence interval of the normalized RMSE between the predicted and observed outcomes for each of these counterfactual trajectories. We report the normalized RMSE, as is standard (Bica et al., 2020b; Melnychuk et al., 2022), where we normalize by the maximum outcome value for that dataset, where for eq. (5) ymax = 50.0 and for eq. (6) and Cancer PKPD Ymax 1150em3.
I DATASET GENERATION AND MODEL TRAINING
Dataset generation D. We generate a dataset for each underlying pharmacological model F and a given action policy. This is achieved by sampling a set of N initial conditions 2o E X, v E V from the models specified domains X, V. These initial values and the action policy simulate the covariate trajectory up to a defined end time T, using a numerical ODE solver. This forms a dataset as described in Section 3. We follow this process independently sample train, validation, and test dataset D.
Model training. We train each baseline on the training dataset Perain We follow the same training setup as Melnychuk et al. (2022), where all the baselines are implemented in Py lightning (Falcon, 2019) and trained with the Adam optimizer (Kingma & Ba, 2014). Following Melnychuk et al. (2022), we train all LTE baselines using the teacher forcing technique (Williams & Zipser, 1989) when training the models for multi-step ahead prediction. During the evaluation of multi-step ahead prediction, we switch off teacher forcing and autoregressively feed model predictions. Furthermore, we train each baseline for 100 epochs, for the experimental evaluation setup, see appendix H. We perform all experiments and training using single Intel Core i9- 12900K CPU @ 3.20GHz, 64GB RAM with an Nvidia RTX3090 GPU 24GB.
Hyper parameter tuning. We followed the hyperparameter tuning setup of Melnychuk et al. (2022) and used their tuned hyperparameters for all the baselines to the Cancer PKPD dataset across all datasets and all experiments. Therefore, only tuned the hyperparameters of the ODE discovery methods to the validation dataset of the Cancer PKPD dataset and fixed them throughout all experiments all other datasets.
27

---

Published as a conference paper at ICLR 2024
J INTERPRETABLE EQUATIONS
Unique to the proposed framework and INSITE is that the discovered differential equation is fully interpretable. Some of these discovered equations are shown in table 9. clear that even with binary actions, INSITE is able to discover a useful equation that performs well and is similar in form the true underlying equation, even discovering exactly in some simple settings of eq. (5) (e.g.. we discretize the numeric constants, rounding to one decimal place for eq. (5).A and eq. (5).B we recover the true underlying ODE).
Table Interpretable discovered ODEs. We show discovered equations alongside the true underlying data generating ODE for the one-compartment PKPD model in eq. (5), with increasing levels of BSV complexity.
Dataset
True ODE (Data generating process)
Discovered Population ODE (F)
og. (5).A
Caz(t), if a 0 Cx((t), if
de((1)
10085466G2()),,
if a
dt
dt
(008550C12()),
if
Coz(t), if O
1. (.005599g2()),, if a if 1.019177Coz(1) -00455776z()), ifa 1.021360C;a(() - 001464457z(), ifa=
eq. (5).B
drie)
dn(t)
Cix(t), a
1.085445(a(()).
eq. (5)C
(x(())
Cor(() -0.05a(1),, i da(t)
de
-(x((() -0.15x(1), if
de
(ta(t) -0.05r(t), if
(.020334Coz(t) 0.1076062(t), if 0
eq. (5)D
C(x(t) 0.15x((). if de
(. 0.047018m(t), if
K ADDITIONAL EXPERIMENTS
K.1 SENSITIVITY TO À
In the following, we explore the sensitivity of INSITE to the regularization hyperparameter constant A. As tabulated in table 10, we see that INISTE does indeed depend on the correct choice of A. We follow the hyperparameter tuning setup of tuning À on the validation test dataset, as is standard for other LTE methods, such as Causal Transformer (CT) (Melnychuk et al., 2022). We further detail this tuning strategy in appendix I; specifically, we chose À = 10, tuned on the validation set for dataset Eq. 5.D. and kept the same value across all datasets, and all runs (unless explicitly, as noted this sensitivity analysis).
Table 10: Counterfactual normalised RMSE, for G-step ahead prediction of the benchmarks on each synthetic dataset detailed in appendix F, with INSITF varying its regularization hyperparameter X. We quote 95% confidence intervals with each value, and all metries are averaged over ten random seed runs. Each PKPD underlying pharmacological model is simulated with different layers of BSV (table 2) for A-D. INSITE sensitive to its regularization hyperparameter X.
Mathed ISLA
ISVC
(5LD
(B1-A
(51.8
Canoer PRPD
0.0011000
05728-37617
200:0000
2.55 13
2.55
2.11a00
3.3020112
OH1--b ETR
RASNN
L055000
L052k000
1555514
1.25:000
Libesti
LI055L10
0.9221U8
6-N41
09118130
0.918331
0.728814
1.332121
1.34
L0251111
1
CT
0.0020155
07550114
1.005:14
1,20 p1
1.25
103350111
LI4e5100
1,072007
A-WIIINDY
0111-72665
0:11-2.0-034
B1 12511-7--3
11:07700
NA
NA
NA
INSITT
0.0
632 2000
0JD -3000
BLM 10
RJ8 =0
3.07 -000
1:17 417 86
290 2000
019 20.00
1.22 u0
INSITE-A
10.0
2310
RLM4 06
IL06
0.94 2005
10m5
10ml
037 50
K.2 PARAMETRIC DISTRIBUTION OF NUMBRIC CONSTANTS
Unique our method, INSITE is that after obtaining the patient-specific differential equations, we can derive a population differential equation where each numeric constant is represented by a distribution, such normal distribution or mixture of i probabilistic interpretation of the underlying data generating process.
We explore a further insight experiment, where we discover the parametric distribution of the numeric constants of a further synthetic equation. We extended our existing eq. (5), to have an offset term Ao that is sampled from bimodal Gaussian mixture distribution. We follow same experimental
28

---

Published as a conference paper at ICLR 2024
setup and discover individualized differential equations for each patient. To explore if we can recover the underlying bimodal distribution, we plot the distribution of the numeric constants Su for each patient. We observe that the distribution of the numeric constants for each patient is bimodal and that the two modes are the two modes of the underlying bimodal distribution. This is shown in fig. 4, with kernel distribution estimation plot.
0.8
0.7
Dendty
0.5
0.3
0.1
0.0
-3
8o.
Figure 4: Kernel distribution estimation plot of the individualized numeric constants So for cach patient, for further synthetic equation that extends eq. (5). have offset term 3o that is sampled from bimodal Gaussian mixture distribution. We observe that the distribution of the numeric constants 8o for each patient is bimodal the two modes are the two modes the underlying bimodal distribution. This verifies that INSITE can recover probabilistic interpretation of the underlying data-generating process, with a population differential equation where each numeric constants can be represented by distribution.
K.3 ADDITIONAL MAIN RESULTS
In the following, provide full additional results for all counterfactual 7-step ahead prediction errors of the benchmarks on each symthetic dataset. Where we tabulate: 1-step in table 11, 2-step in table 12, 3-step in table 13, 4-step in table 14, 5-step in table 15 and 6-step in table 16.
Table 11: Counterfactual normalized RMSE, for 1 1-step ahead prediction of benchmarks on cach synthetic dataset detailed in appendix F. We quote 95% confidence intervals with each value, and all metries are averaged over ten random seed runs. Each PKPD underlying pharmacological model is simulated with different layers of
BSV (table 2) for A-D. Our contribution is shaded below.
Melbed
1614
ChearPKFO
224-7
BABSN
2,0001117
1704817
2
25244
01774811
2.9788111
01241210
109814
CT
6.00mm
0.394895
8.70000
0 3360 65
0.87408
0830000
051:6400
00660
A WSSINNN
011-7906-5
033144.95 01
0 10 65 01
NA
NA
NA.
NA
NA
24STTE
9.8 332600
1282-11
0.51177M077
Table 12: Counterfactual normalized RMSE, for 2-step ahead prediction of the benchmarks on each synthetic dataset detailed in appendix F. We quote 95% confidence intervals with each value, and all metrics averaged over ten random seed runs. Each PKPD underlying pharmacological model simulated with different layers of BSV (table 2) for A-D. Our contribution is shaded below.
Matted
4% BIAA
0y C5LI
CSLD
RL. (6LA
eg P9-00LI
4G. HILC"
Pg/iLDD
ChporFPKTo
N05M
L3541577-18
E3443
20000115
2.00B:15
1:54.:.39
1.558111
033
108ee 15
8.7
1/554016
127-4111
111-0
G-N61
BASLAU
02318-4
Q431E 39
0.54000
QATABES
$59LCC99
00600
A-SINDs
011--186177
0.111830
147,000
4471855
1.31000
1.0_0.15
12210014
INSITE
882 Juve
0A8
8.83
0.22 1882
0.2 J0c0
091332
8:00 2X12
0.901200
K.4 ADDITIONAL RESULTS WITH VARYING SAMPLE SIZES
In the following, we provide additional results where we vary the number training sample sizes, in the range of {10, 100, 1, 00000 10, 000). For each sample size, we provide the counterfactual
29

---

Published as a conference paper at ICLR 2024
Table 13: Counterfactual normalized RMSE, for 3-step ahead prediction of the benchmarks on each synthetic dataset detailed in appendix F. We quote 95% confidence intervals with each value, and metries are averaged over ten random seed runs. Each PKPD underlying pharmacological model is simulated with different layers of BSV (table 2) for A-D. Our contribution is shaded below.
Meiheel
eq. IRA
rq 1FLB
IRVC
(Q-(5LD
eg. 16.A
eg. (ALB
(FAC
HSLD
Cinorr PKPD
RMSN
22120014
2:22:011
2.125014
1.0eetre
1.0
1.045:11
1.34s0us
0.56e0:0
CBN
05002207
0.86 20.53
0.86 2x003
R&2 L1.94
0 2111
0,73 ase
13.
077522116
87748411
10842411
1.00a6.11
1.00..cun
119950.11
01002912
0.58,55.00
0.5010000
1:223000
0.88-5oe
A-SINity
0.3540000
0116:200-11
148000
1.610013
1
A. A.SSS
0. 0.3147790-43
0.11=250p04
0.12
0. 0.1057 100
NA
NA.
NA
NA.
NA
INSITE
1L01 20000
0.83 20.00
0.08 21.03
0.85 20100
0,83 H0L29
0.53 =000
0,90 SALLE
0,78 20000
Table 14: Counterfactual normalized RMSE, for 4-step ahead prediction of the benchmarks on each synthetic dataset detailed in appendix F. We quote 95% confidence intervals with cach value, and all metrics are averaged over ten random seed runs. Each PKPD underlying pharmacological model is simulated with different layers of BSV (table 2) for A-D. Our contribution is shaded below.
Metbel
15LA
(73B
IBLC
LPLD
IBLA
(RLB
(BAC
IBLD
Canorr PXPD
MSM
12scu0
1:12:300
1.12::477:16
2.5542.16
2.55ac.15
2.06:8.10
2:15a2ue
2.2488:12
RMSSS
1:6es::1
1.22213
1191254M
BS2::316
066555111
0,59
G-Net
015025.15
0.9022.18
0.5625.33
0.54:8812
1.3428.16
LJ442.15
1.2185.11
1.002a00
CT
1L7840114
0.78:8114
0.6685122
0.5Re5.111
1,22:006
1.22:6400
L.4280n
1.3485.3
0.1147.56645
011:2.51e0x
0.124)
NA
NA
NA
INSITE
1000 52403-18 18
0083 15
R.64 18
BL04 0:000
0.86 20.39
I.96
:083
PUST LA 53600
0.88 20.1
IL716
20:00
Table 15: Counterfactual normalized RMSE, for 5-step ahead prediction of the benchmarks on each synthetic dataset detailed in appendix F. We quote 95% confidence intervals with each value, and all metrics are averaged over ten random seed runs. Each PKPD underlying pharmacological model is simulated with different layers of BSV (table 2) for A-D. Our contribution is shaded below.
eq.
MSM
1:05=000
2.60c2.11
2.50et.11 15
2.09c8:43
2.50e0::1 12
RMSN
201:2919
1.T7esis
Steesur
/Bea
1.19a6
1.00:815
1.115542
0196
0.562019
1.2428:21
1.2486
1.24-001
1.1740:11
0085e0.16
0.6628177
0.715645
0.95e0.14
1.2k:00e
1.38a6000
1.04s00
1. 1.16S5.11
1.00scm
OWKCD
A-SINITy
0111:000
12
14516510m
1
INSITT
(L02 10:00
.0.30
IL04 15776-18
I.64 -.533
(L90
0.90 1000
FUS .0.99
0.87 A4:04
IL7K
Table 16: Counterfactual normalized RMSE, for 6-step ahead prediction of the benchmarks on cach synthetic dataset detailed in appendix F. We quote 95% confidence intervals with each value, and all metries averaged over ten random seed runs. Each PKPD underlying pharmacological model simulated with different layers of BSV (table 2) for A-D. Our contribution is shaded below.
Methoel
eq. (SAA
eq. (5LB
P1-I5CC
Mg. 15LD
eg. IBLA
ei-[6LB
Pg. 1B-IBCC
Pgg 19-I5LD
Coner PAPD
MSM
119048 STo
1550-5000
2:00:0000
2,666613
2:0600
2.1150m
2:5040 13
008-D 117
RMSN
10540.10 10
196-523
105---10
B8240000
19
1 191:0 15
10-0011
105-0m
1.2546 15
100--00
1:B0-6
1. 1.1045:11
1.8525.00
0.97-40:8
1.06-6 IT
0.7550.1
1.20:p00
1.2055100
1.0:5:11
1.1420100
A-SINEDy
ILI1ACAD
0.11:000
8.1562406-17
0.188810
1.451000
LAULONII
1.511309
1,211117
A-WSINEL
01
0.1132.40p04
411263477444
NA
NA
normalized RMSE, for 6-step ahead prediction of the benchmarks on each synthetic dataset. These tabulated in: table 17, table 18, table 19 and table 20 respectively.
Table 17: Counterfactual normalized RMSE, for 6-step ahead prediction of the benchmarks on each synthetic dataset detailed in appendix F, with each method trained training dataset with 10 samples. We quote 95% confidence intervals with each value, and metries are averaged over five random seed runs. Each PKPD underlying pharmacological model is simulated with different layers of BSV (table 2) for A-D. Our contribution shaded below.
Metheel
rq ISLA
rg: 15LM
(5LD
rg (BLR
(EMC
(6LD
Conorr PAPD
011--1 131
2.191010
CWN
2 2.561466
2.562046
19961045
3.355860
250:200
2.00C000
14448a0
20845.19
2 2.9948.27
13
3474030
34422289
3554083
2-4812488
24835
1.50 20.37
2084263
2.91 $1.00
2544344
25251.0
2.091.200
248.1100
3.01:006
2.956611
20110337
2344211
3.4840-12
A-SINiiy
[1 0.314000
0:001.000
0.34-1500
0.141000
1.251000
1,511.00
1801000
24410000
2,111000
A. W.WSSNES
0. 0.118)130--44
0.10 14c408
A.1248.770
041:5442000
NA
NA.
NA.
NA.
NA
INSITE
1000 20000
0.8 20000
8J66 20102
0.85 211
009 20000
0L92 50000
1.16 20,0000
1,27 50,00
2.15 =0.00
30

---

Published as a conference paper at ICLR 2024
Table 18: Counterfactual normalized RMSE, for 6-step ahead prediction of the benchmarks on each synthetic dataset detailed in appendix F, with each method trained on a training dataset with n 100 samples. We quote 95% confidence intervals with each value, and all metrics are averaged over five random seed runs. Each PKPD underlying pharmacological model is simulated with different layers of BSV (table 2) for A-D. Our contribution
is shaded below.
Methol
USAA
IFAB
:BVC
16LA
Cancer PXPD
MSM
0.50:500
1.21=200
1.54:5:0
2.80e2000
2.005500
221:0000
2:25-200
3,45acu0
RMSS
1.758833
1.7589.5
1.35
2.77s1.29
2.0006:11
2.00c.111
141:841
1.418812
2.11:6:00
Gi-Nex
LM63541
1.561041
1.35,5.30
23368064
1.5622:19
1.6.6.10
2.0.8.19
2.
CT
LS24699
1.521962
1,2812/00
E4S4831
2.548331
2.53622
204:1016
3,2010.00
OnE-D
A-WSINIT
10 0.3112.5600
0111:725000
0.121521000
B 0.10-355000
NA
NA
NA
NA
PANITE:
1L02 10.00
1U85 10:30
B.86 (4103
RL88 00.00
IL96
10:30
1.15
(0:0
IU87
.0.00
0,99 10:00
1,19
(0:00
Table 19: Counterfactual normalized RMSE, for 6-step ahead prediction of the benchmarks on each synthetic dataset detailed in appendix F, with each method trained on a training dataset with n = 1, 000 samples. We quote 95% confidence intervals with each value, and all metrics averaged over five random seed runs. Each PKPD underlying pharmacological model is simulated with different layers of BSV (table 2) for A-D. Our
contribution is shaded below.
oq (SLA
<q. (EAE
<9 (5LC
<g.5AD
<q. IRLA
cq. (GLI
<g. BRABCCC
cq.
MSM
0.502210
0.91a154e-10
2,00:100
261:200
2512CI0
1.992202211 to
2:182200
2.362C00
001-- 1TT
RMSS
L00
L872833
10090:12
1955246
0.778812
1872530
LS1:823
LSS5833
1.302214
1.3586114
1.97cca
180 23.29
1:022.17
1.062288
1.07188.12
08336631
08550053
1,161619
10000000
1.11:014
1,001111 11
A-SINID
0.3125000
0:11:0000
0:158800
1,460000
1,47:6000
1,00:0000
1,9665000
1,29:0000
0.314377700
0.11-3.9e
NA
Table 20: Counterfactual normalized RMSE, for 6-step ahead prediction of the benchmarks on each synthetic dataset detailed in appendix F, with each method trained on a training dataset with n = 10, 000 samples. We quote 95% confidence intervals with each value, and all metrics are averaged five random seed runs. Fach PKPD underlying pharmacological model simulated with different layers of BSV (table 2) for A-D. Our contribution is shaded below.
Metloal
09. (SLA
cq. (FLB
<5.15LC
(g-(5LD
<9. (RLA
cq, (BLB
6g. (BLB
cq. HSLD
Cuacer PAPD
(100-00 6:T
BMSN
101 1244
0.70 11:30
0.41 J017
0.54.028
(L99 J1.27
0.95 11133
[.91 .0.30
186 (1.51
14431632
CRN
483 =9:9
6,30 -12.0
2.31 5462
3.59 29.39
0.93 =0.21
0.91 =0:30
I.71 =0.07
1,56 22.41
4.33 ms
G.Not
1.05-946
019920177
0.9 -1:00
1-14 -0.30
1,16 =0.43
8.97:9255
1.0925.13
1.2546 11
=0.50
051-992
=0.17
0.74
08550012
1.06e0
A-SINDY
0.112000
0.1285000
0.155840
1.395511
1.3080.18
1.305931
14585520
A-WSINITY
01111441144
0.11:1:77000
8:101750-00
NA
NA
NA
NA
NA
INSETT:
(U8 .0.30
IL94 12316-17
0.85 1030
0.96 0775500
0.93 111
(L78 P4.
0.58 1017
K.S ADDITIONAL RESULTS WITH INCREASING OBSERVATION NOISE
In the following, we provide additional results where we increase the observation noise of the BSV datasets, in the range of ε = (0.01, 0.1, 1.0]. For each sample size, we provide the counterfactual normalized RMSE, for 6-step ahead prediction of the benchmarks on each synthetic dataset. These are tabulated in: table 21, table 22 and table 23.
We note that for high observational noise settings, although the performance of SINDy degrades, the performance of WSINDY remains competitive. In settings where WSINDY is supported (i.e., trajectory lengths are not short and not sparse variable length). INSITE can WSINDy to discover the population (global) differential equation model instead of SINDy. This arises as INSITE and more broadly the proposed framework can fine-tune any differential equation discovery method that produces a population (global) closed-form differential equation that has numeric constants.
K.6 MODEL. MISSPECIFICATION
INSITE and the other ODE discovery methods (e.g., A-SINDy) are only correctly speci- fied for eq. A5) datasets, they use the feature library of LINSITE {1, 21, }. Crucially, the ODE discovery methods are misspecified for eq. (6).A-D, and the Cancer PKPD datasets,
their underlying equation would require the feature library of LCaner {1, Zo. 11 21 2, 2n log(z)). log(21)} to be correctly specified. Although this misspec- ification persists (in the main experimental table in table 3), as noticeable from the increased order
31

---

Published as a conference paper at ICLR 2024
Table 21: Counterfactual normalized RMSE, for 6-step ahead prediction of the benchmarks on each synthetic dataset detailed in appendix F, with BSV observation noise of 0.01. We quote 95% confidence intervals with each value, and all metrics are averaged over five random seed runs. Each PKPD underlying pharmacological model is simulated with different layers BSV (table 2) for A-D. Our contribution is shaded below.
Meileed
eq. IRA
(51.8
eg C5CC
eq. (SLD
rg (6LJ
mq. HBLC
S4
0.7at 14
RMSN
181:0025
1,30sht14
1.35
3.16
CBN
L
1951832
E.7721077
1.07/0
8/06.0
13
0772:0
143 2845
1.45 2i45
1UN6 us
3:39::133 26
0.636031
0.8314322
0/4400
1.340115
1.16:0319
1.20::000
OHW-4
A-SINity
0.3140000
0.11 1001
0:33:000
151000
14460001
1.47-000
144000000
1.56-003
A. J.SSIINN
0.31
0.1 0.11:6-72:64
0,1122 15-48
0. 1.1021200-33
8A
NA.
NA.
N5.
INSITE
1000 20000
0.3 2010
0,84 20000
1L06 20100
0.96 2000
0.96 =0000
0,32 146-16
8:99 2000
Table 22:
Counterfactual normalized RMSE, for 6-step ahead prediction of the benchmarks on cach synthetic dataset detailed in appendix F, with BSV observation noise of = 0.1, We quote 95% confidence intervals with each value, and all metrics are averaged over five random seed runs. Each PKPD underlying pharmacological model is simulated with different layers of BSV (table 2) for A-D. Our contribution is shaded below.
Mellodi
cq. 15IA
15LB
(BVC
(SLD
69.A
H6LC
(9.D
MSM
[19943 Ste-16
[1:-4000
1.17.a
2:11000
2
2:61::014 60
2:13-u4
11546.15
1 166015
0.97-0
1540
E 2019
1:00-:006
11.96
(2-N0
11905645
D (M5-0-33
0.70.0 15
1. 116-0%
1 18m
1.49 -n-4
(U:6 -0:23
1:41:013 37
CT
0.5548531
0.91
0.82:0030 30
019011022
1.34
1.36::0:5
110020014
1.24 06
0DE-D
0:24:000
147::000
A. A-WSINES
0.23 22455-45
0.23 15:40
0.23 222 94-10
SA
NA
NA
NA
INSITE
1000 SAEDS
0.23:000
8,23 1100)
(1241000
1422000
0.55 21000
0.82
2010
8,90 210M
Table 23:
Counterfactual normalized RMSE, for 6-step ahead prediction of the benchmarks on each synthetic dataset detailed in appendix F, with BSV observation noise of e = 1.0. We quote 95% confidence intervals with each value, and all metrics are averaged over five random seed runs. Each PKPD underlying pharmacological model is simulated with different layers of BSV (table 2) for A-D. Our contribution is shaded helow.
Method
CPWC
MSM
(L5022 24c-6
2.0024000
3,03200:00 00
4.1022000
2.0120000
2:51::000
1.5022000
2.13 20
0007- LTR.
BMSN
3.01
3:21.00
3.724:.25
1:15-:0:5
10:13
(2-Nex
11901044
2.51
2.35,0
2 196.0.25
1.48 10W
1.48 LIVE5
00.96 10.23
1/42.01
L.
08551031
2.17
2:12.000
2:3010006
1:54.912
1:34.:033
1.24..05
J.WSINE
0.1113.77:44
245
NA
INSITE
0.02 50002
2185000
2.12:0000
2.182000
LA22000
0.96 2000
0.83 20.00
691 2000
of magnitude increase in error, we still observe INSITE achieves a lower error than the longitudinal treatment effect models.
To investigate how INSITE and the baseline ODE discovery methods compare under different levels of model misspecification, we performed a complete re-run of our main experimental table with varying feature libraries. Here we went from using an overly-restricted feature library of L == (1) to a still misspecified feature library (for eq. (6).A-D, and the Cancer PKPD datasets) that overly expressive of L = {1,
(1, xo. 1 xb, T031, 27, xn, -p. x xf). The results are tabulated {1, To. 31, T0 } in table 26, L = {1, zg "1, " 11 " in table 27, L (1 Ii, " 11 .I. af T at) in table 28. We observe the same pattern of low error when the feature library is correctly specified (that is the underlying dataset was generated with features within the searching or discovery library set) and an order of magnitude increase in error or higher when the feature library is misspecified (for eq. (-).A-,, and the Cancer PKPD datasets).
as: L = {1} in table 24, L = [1, Tn, r1} in table 25, L
An actionable insight is that the ODE discovery method should have a sufficiently expressive feature library to search over. Moreover, even the underlying data-generating equation contains unique feature library terms, such as log(so), log(z1), that are not explicitly given, employing polynomial feature library can serve as a robust approximation to the unidentified feature library set. This approach is supported by the Stone-War Weierstrass theorem in function approximation theory, which posits that any continuous function defined over a closed interval can be uniformly approximated by polynomial function with arbitrarily close degree of precision (Stone, 1948).
32

---

Published as a conference paper at ICLR 2024
Table 24: Counterfactual normalized RMSE, for 6-step ahead prediction of the benchmarks on each synthetic dataset detailed appendix F, with all ODE discovery methods using a restricted feature library of LISSIE = (1}. We quote 95% confidence intervals with each value, and all metrics are averaged over ten random seed runs. Each PKPD underlying pharmacological model is simulated with different layers of BSV (table 2) for A-D. Our contribution is shaded below.
Method
dRAA
(6LB
(6LD
Canoee PKPP
MSM
0:9:0120
0.7788976-17
2.00a::00
2.55-00:3
2.5588113
2.06a2-16
2.13::004
2.50atu12
1.23-00:5
ORN
1.355.10
0 8822400
1.9858.14
L0555005
1.052516
1.03:0
0.92auus
1.Nex
0.721011
0.9.aC.15
1.33::222
LS511227
1.225511
8.2bu0155 15
1.221814
CV
0900u835
0.7840114
1:0000.14
1.29-0.0
1.201000
1.0445111
1:04:430
LIW7a1077
INSITE:
28.43 16
28.45
38J18 17.00
28.53
38.51 10.00
28,51 1010)
1833 10:00
18:38 10:00
17.36 17.0011
Table 25: Counterfactual normalized RMSE, for 6-step ahead prediction of the benchmarks on each synthetic dataset detailed in appendix F, with all ODE discovery methods using restricted feature library of LINSITE = [1, 31 We quote 95% confidence intervals with each value, and all metries are averaged over ten random seed runs. Each PKPD underlying pharmacological model simulated with different layers of BSV (table 2) for
A-D. Our contribution is shaded below.
Method
(S.A
(5L.
15,ISLC
<q.I5LD
00. (0.691A
c9. IBAC
eq. (S1.)
Caace PKPD
HMSS
1:
1944:8.23
1.918C118
1.23::115
12518115
11012.11
1.00ud3:
LINAR117
CRN
1.0666.0 to
1,05-0.0
1 1.96-5 14
1 1:05-089
Lutsom
10545566
1:33:0 90
G-Nox
0:01-020
0 77:014
0.9740
1:31-027
1.0943:11
1:25:41 15
A SISIND INSITE
0.5240.00
388200.3
0 0.6820000
0.5752000
0:98:0:00
D.752000
0.73-0:00
8.47 20.00
0.47 2000
0.30 2900
0.40 277700017
0.50 21:30015
0.80 546-16
0.63 20.00
0,61 2000
0.86 20000
Table 26: Counterfactual normalized RMSE, for 6-step ahead prediction of the benchmarks on cach synthotic dataset detailed in appendix F, with all ODE discovery methods using restricted feature library of LINSITE (1, 20, 11, 20211 ). We quote 95% confidence intervals with each value, and all metries are averaged random seed runs. Each PKPD underlying pharmacological model is simulated with different layers of BSV (table 2) for A-D. Our contribution is shaded below. This table is identical table and is included here for completeness.
Method
(ELA
(5.
(5LC
(5LD
(6).A
(6,,
(6LC
(6).D
Canoea PKPO
MSM
17-117
0.9925000
0.7728377-11
2.0920000
2:5580.13
2.55:0113
2.0660.16
2.11=004
2.30a0012
RMSN
1.9220024
15
1.6540.19
19120.18 18
1.25e0.15
1.25-0.15
L10k5.11
L/0s0.11
CRN
1,08:0110 10
1005:000 10
0.2200000
1 19980
1055003
105:000
1.0560.0
10050010
0,92=005
G-Net
0091=0220
0:91::220
0. 0.72:4014
159200
1.35+027
13449227
1.226.11
1,28:0155
1.22-0
CT
0,903:011 18
0,903-115 15
0.75:4014
10040 14
12290007
1299:000
10546.11
1.14:010
1.00-007
ODE
A-SINLy
0.11:0000
0:11:0000
01151000
1.455003
14550003
1.40:0011
1.51:009
1,232013
INSTEE
0.02 33
0.03 -0:00
6.04 20.00
0.05 15
0.94:000
20.0
3000
0.84 20:4
0.87 10.0
0.79 54.376-17
Table 27: Counterfactual normalized RMSE, for 6-step ahead prediction of the benchmarks on each synthetic dataset detailed in appendix F, with all ODE discovery methods using restricted feature library of LINSITE = [1, To, x1, .u. To .z).. We quote 95% confidence intervals with each value, and all metrics are averaged over ten random seed runs. Fach PKPD underlying pharmacological model simulated with different layers of BSV
(table 2) for A-D. Our contribution is shaded below.
Method
<q (5LA
(51.B
0g I5UC
cq. (5.D
0g. (91.A
ON HRLB
C6. IELC
e5 (9..)
Caaoel PKPD
MSM
1000
0.9218-576-17
2.00AC00
2.55-113
2.SSBLIS
20h15.16
2.102:044
2,3028112
HMSSN
1.91ac.1K
1:23-015
1-1012.18
LIMEA117
CRN
1.866
1005-0.00
0 032:000
1.98-0 14
1.05-099
10550003
1:05-306
1/33-0 90
01920008
1.31-037
0,906
0:90-0
1775001
1.0060
1
1.0565111
1.14:0
1074007
A-SINDS
0.11:1996177
0.12:1000117
011
0:0000.0
0.822200
0.85
L71
INSITE:
8002 10000
6/33 LEOM
BJH 10100
6.04 10003
0,67 110M
14000
Table 28: Counterfactual normalized RMSE, for 6-step ahead prediction of the benchmarks on each synthetic dataset detailed appendix F, with all ODE discovery methods using restricted feature library of LISSIR 2u, 31, 28, T1 2ř, æ0, 2fæ1, Zozf, }.). We quote 95% confidence intervals with each value, and all metrics are averaged over ten random seed runs. Each PKPD underlying pharmacological model is simulated with different layers of BSV (table 2) for A-D. Our contribution is shaded below.
Method 15AA
15-15DD
(61.4
16L.5
BBLC
(61.1)
Canore (9PP
MSM
17
0490u1 30
0.9788 Wh-17
2:0940:00
2.55--0 13
5.0603 16
2:13:0044
E3R
RMSN
19440123
1.91a0.1
1
LINE4117
0325509)
LISasous
1.08221
0.118220
0.2 0.228114
0.9722.15
1.33
1.025511
1.25
2220114
CT
01905518
0.7588114
1.0005.14
1229:000
12200000
10052211
L:04
L177
19
INSITE:
8J00
6413 5009
B. 2910
0.66
6,55 51.56/16
0085
S46.1
33

---

Published as a conference paper at ICLR 2024
K.7 INSITE ALSO EMPIRICALLY WORKS FOR IRREGULARLY SAMPLED DATA
INSITE also supports irregularly sampled time series data, as its underlying ODE discovery method of SINDY supports discovering ODEs from irregularly sampled time series data (Brunton et al., 2016; Goyal & Benner, 2022). We further empirically verify this by re-running SINDy and INSITE across all datasets, now with the irregularly sampled datasets, where we randomly sub-sampled imrgullall, excluding 10% of the original observations along the time dimension. The subsequent results, conducted across ten random seeds, are outlined in table table 29 below. We observe that INSITE and SINDy still achieve good performance (hence low prediction error).
Table 29; Counterfactual normalized RMSE, for 6-step ahead prediction of the benchmarks on cach synthetic dataset detailed in appendix F. We quote 95% confidence intervals with each value, and all metries are averaged over ten random seed runs. Datasets sub-sampled irregularly, excluding 10% of the original observations along the time dimension. Each PKPD underlying pharmacological model is simulated with different layers of
BSV (table 2) for A-D. Our contribution is shaded below.
Method eg (5LA
eq. ($).)
eq. C51C
eq. (5LD
eg. (6.A
eq. (6.
eq (6)C
ng. (6AD
Cancer PXPD
ODE
A-SINDy
0:1111
0.12.1 Vie11
0.111000
0.001 54-18
0.901 14-18
1.68 1.0k-16
INSITE
2000
.000
0.04
0J14 1000
10.00
IL75
10:00
0.70 10:00
1,00 11:34e-15
K.8 INSITE CAN ALSO WORK WITH MORE CHALLENGING OBSERVATION FUNCTIONS g(x)
Although the observation function that describes the treatment outcome as function of the observed features is often prespecified by the user as often models the outcome as one of the features (g(x) 2j), e.g., the tumor volume, other more complex forms of g that are not prespecified are also of interest. To investigate whether INSITE can work under these more challenging settings of having. be non-linear function, we empirically investigated this by implementing three more challenging forms of g(x) representing more complex settings. These are, (1) an exponential function g(x) exp(x,) that often models rapid changes growth, commonly found in many biological and economic systems (Johnson et al., 2019), (2) a quadratic function g(x) = 2² that models a non-linea relationship in systems (Gasull et al., 2004), and (3) trigonometric function g(x) sin(2,) that often models periodic components within oscillatory dynamics, which are commonly found in natural and engineering systems (Datta & Mohan, 1995). For each of these new observation functions we re-ran SINDy and INSITE across all datasets, conducted across five random seeds, with the results tabulated in table 30 below. We observe that INSITE and SINDy can still achieve acceptable performance. To implement this experimentally. before applying the observation function after sampling the dataset initially, we normalized the dataset using min-max scaler (normalizer) so the min-max was from 0 to I, applied g(x) then un-normalized to return the min-max back to the original range of the dataset- we did to help numerical stability, and such normalization step is common when pre- pre-processing datasets.
Table 30: Counterfactual normalized RMSE, for 6-step ahead prediction of the benchmarks on cach synthetic dataset detailed in appendix F. We quote 95% confidence intervals with each value, and metrics are averaged over five random seed runs. We modify the observation function g(r), to be one of the following g(x) = {æz, exp(x;), ), az. sin(2;)) Each PKPD underlying pharmacological model is simulated with different layers of BSV (table for A-D. Our contribution is shaded below.
olm)
Mothed eg (5)A
eg. (5LB
eg (55.D
(q.(
eq (61.8
eg (6VC
eq.
Cancer PKPD
A-SINDs
0.111000
0.11.0000
0.1310000
0.1510000
1,46.0000
1:47.0.00
1,29.000
INSITE
10.0
0.0 5000
0.04 20:0
10000
0.96 5009
0.96 090:2000
11.54e
0.90 10000
0.79 10:0
A-NINDy INSITS
0:75-0:0
0266:006
=0.34
1.02 =0.57
=002
061
60.30
003
61.31
$0.59
0.78
23
A-SENDs INSETE
0:1015885-05
#1
0.09
=2.67e
01121596-01
0.09
42
0.121001
+6
1.12 #1.04
1,16 1105
1.15 11.6
=1.00
1.63 #2.05
1,99
0.57
40.00
03
A-SINDs
2.4210077
2.42 0.09
2.31.10.7
2.42 10:40
1,62 2044
146 20.4
2,30 11.96
1.99 .0.99
2.00 10.46
INSITE
2.09 31.10 10
2.09 20.00
2.06 2006
2.22 10:25
1,46 2033
144 30.42
1.96 21.09
1.72 200/6
1.59 10.29
34

---

Published as a conference paper at ICLR 2024
L RECONCILING ASSUMPTIONS IN TABLE 1
Let us discuss in detail how the assumptions set for treatment effects (assum. 2.1 to 2.3) and ODE discovery (assum. 3.1 to 3.3) compare as reported in table in section 4.
existence & uniqueness (assum. 3.1) - --- overlap (assum. 2.2) and functional spaces (assum, 3.3)
overlap (assum. 2.2).
These assumptions do not map one-to-one. However, they do serve a similar purpose. Essentially, overlap allows to assume that there is always very similar sample (in terms of covariates), such that we have access full range of potential outcomes such that we can estimate eq. (1) for every treatment and covariates. When violated, Gelman & Hill (2006) pointed out that one has to rely much more on correct model specification, i.e., rather than assuming overlap, one may resort to the correct model specification instead. Assum. 3.1 and 3.3 tell us exactly this: the former assumes that there exists one discoverable function (i.c.. one that can be correctly specified), while the latter tells us that the correct function is comprised of the token set that is provided to the ODE discovery method. Combined, these assumptions should result a correctly specified model, which can then be used to relax an overlap assumption.
observability (assum. 3.2) o consistency (assum. 2.1) and
observability (assum. 3.2) While the previous assumptions did not map one-to-one, we find that assum. 3.2 does map one-to-one with assum. 2.1 and 2.3. With eq. (3) we combine covariates, (potential) outcomes and treatments into one dynamical system. Given this, assum. 3.2 states that every variable of the complete ODE observed the dataset. These variables include the potential outcomes (which corresponds to assum. 2. as well as all the covariates (corresponding to assum. 2.3). One could make the argument that assum. 3.2 even stricter than combining assum. 2.1 and 2.3 the latter does not mention observing treatment indicators. While there is work missingness in treatment indicators (such as (Kuzmanovic et al., 2023)) it still an exception to the rule.
overlap (assum. 2.3).
M SCALE OF THE OUTCOMES
In the following Table 31, we provide the scale of outcomes for each of the datasets used in all experiments.
Table 31: Scale of outcomes for each of the datasets used in all experiments.
Dataset
mean
std
min
25%
50%
75%
max
Cancer PKPD
6,93553
49.099786
0.0
0.000263
0.011018
0.179058
1150.34651
Eq. (5).A
5,225302
7.873385
0.006063
0.48826
1.723761
6.418705
49.480554
Eq. (5)BB
5,225425
7.873301
-0.00044
0.489809
1.726515
6,415482
49.481121
Eq. (5).C
4.376515
7.459998
-0.010424
0.232287
1.039811
4.889878
49,481 121
Eq. (5).)
4.992425
8.127416
-0.012978
0.300557
1.308329
5,801 108
49. 49,848292
Eq. (6).A
7.772205
53,479708
0.000157
0.012505
0.226499
1144.486013
Eq. (6).B
7.772162
53.479744
-0.034542
0.001075
0.016917
0,22768
1144.490737
Eq. (6).
7.31387
51.156823
-0.034717
0.000713
0.015301
0.189101
1 1150.33633
Eq. (6).D
6.935527
49.099739
-0.034994
0,001075
0,01645
0.178607
1150,334351
N HOW TO CHOOSE AN ODE DISCOVERY METHOD
In the following, we first explain how Sparse Identification of Nonlinear Dynamics (SINDy) (Brunton et al., 2016) works and the necessary assumptions it imposes, then second, a discussion on alternative ODE discovery methods and how to choose suitable one for given problem setting.
Sparse Identification of Nonlinear Dynamics (SINDy) (Brunton et al., 2016), a data-driven frame- work that aims to discover the governing dynamical system equations directly from time-series data. The algorithm works by iteratively performing sparse regression on a library of candidate functions to identify the sparsest yet most accurate representation of the dynamical system.
35

---

Published as a conference paper at ICLR 2024
SINDy makes explicit sparse assumption that the dynamics governed by only a few important terms so that the discovered equation is sparse in the space of possible functions. Such sparse parsimonious ODEs balance accuracy with model complexity, which aids in avoiding overfitting. SINDy considers a simpler problem of modelling
dx((t) - = f(x(t)) dt
(13)
where the function f is an equation with a few numbers of terms, that is a sparse equation in the space of possible functions. SINDy uses sparse regression to identify the terms in f. which also helps it mitigate noise. It collects trajectories and approximates their derivative ((t) numerically, from which it then iteratively solves eq. (13) through sparse regression, where it constructs a feature library of candidate non-linear functions (XX), ), where X = [a' (41), ())))) xx (12)..... ...., xx (lm)) For example, the candidate lonelr functions could consist of constant, trigonometric polynomial terms:
(
X
Pa
X Pa
sin
cos X) sin(2X) cox
(14)
Given this feature library, we can then formulate sparse vectors of coefficients [S] & ... En] which determine which features are active:
E
X = O
(15)
Such method is poorly suited for large state dimension due to the factorial growth of O (Brunton et al., 2016). Furthermore, such an approach relies on choosing the right coordinates and basis for the best representation of the dynamies. Moreover, this also assumes that given the non-linear candidate function library, we can recover an equation that linearly decomposes into these candidate functions.
Alternative ODE discovery methods. range of alternative methods exist, including symbolic regression (Schmidt & Lipson, 2009) and other ODE discovery methods that better handle noise, such as WSINDy (Messenger & Bortz, 2021).
Symbolic Regression. It can similarly discover an underlying ODE, however, it performs combinato rial search (Schmidt & Lipson, 2009), which can be slower; however, it can discover more non-linear relationships amongst the candidate library terms, such one feature divided by another.
Feature Learning Methods in ODE Discovery. Feature learning methods, particularly those employing neural networks like Recurrent Neural Networks (RNNs) and Long Short-Term Memory (LSTM) networks, offer data-driven approach for ODE discovery. These methods automatically extract relevant features from time-series data, adapting to complex dynamics without predefined candidate functions (Lechner & Hasani, 2020).
Unlike Sparse Identification of Nonlinear Dynamics (SINDy), which uses predetermined library of functions, feature learning methods learn these features directly from data, offering flexibility in modelling non-linear relationships in systems with unknown complicated dynamics.
However, the use of neural networks in ODE discovery often leads to models that lack interpretability and require substantial data for effective training, posing challenges in extracting physical laws or governing equations.
Considerations for Use. The choice between feature learning methods and others like SINDy or Symbolic Regression depends on system complexity, data availability, and interpretability needs. Feature learning is advantageous in complex, data-rich environments but less so in scenarios where interpretability and simplicity are priorities. When employing methods such as SINDy, user could start with feature library of polynomial terms and slowly increase the polynomial order till the empirical generalization accuracy was suitable for a given treatment effect setting. Other suitable terms commonly found in PKPD models include exp, log terms and non- non-linear combinations of these terms (Geng et al., 2017).
O NUMERICAL SOLVER AND PARAMETERS
The numerical method (solver) used for INSITE and the other ODE discovery methods (c.g.. SINDy) is that of the Fuler method, which is first-order numerical algorithm for solving ODEs given
36

---

Published as a conference paper at ICLR 2024
an initial value. This uses time step size of 6, 0.166 seconds, which sufficient for experiments. This numerical method is used with the same settings across all methods, including INSITE. We kindly note that first, the ODE discovery methods discover ODE, and then at inference time, when we wish to determine future values, we forward simulate the initial observed value using the discovered ODE and this Fuler numerical solver to estimate future values future times. We also kindly highlight further implementation details can be found in appendices appendices F, G and
37


Reviews:
[{"review_id": "2YfPz3tUQ5", "reviewer": "Reviewer_6MV1", "summary": "The paper proposes the Individualized Nonlinear Sparse Identification Treatment Effect (INSITE) framework for estimating heterogeneous time-varying treatment effects given dynamic and statist covariates. INSITE leverages deterministic ODE discovery methods to infer the population differential equation, which is fine-tuned to recover patient-specific differential equations. Experimental results on two synthetic datasets demonstrate that INSITE is competitive in counterfactual prediction relative to baselines.", "questions": "- For the BSV scenario A (eqns 5 and 6), shouldn't we expect ODE methods to do as well as INSITE\n- Fig 3(b): Shouldn't we expect INSITE to be monotonically increasing with $\\gamma$", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper outlines a comprehensive framework for leveraging ODE discovery methods in time-varying treatments/covariates setup\n- The paper is relatively well-written and easy to follow\n- Experimental results demonstrate that the proposed approach is competitive relative to baselines\n- The learned ODES are interpretable, unlike previously proposed neural network approaches", "weaknesses": "- I encourage the authors to include a discussion section focused on the limitations of the proposed approach. The paper makes strong assumptions about the underlying dynamics, e.g., \n1) The success of the INSITE is dependent on the choice of the library of candidate functions\n2) INSITE assumes that the system is sparse, which rarely holds in high-dimensional settings\n3) INSITE assumes that the system is deterministic and noise-free. However, real-world systems often have noise and other stochastic elements, which can significantly affect the model's accuracy.\n- I encourage the authors to include a complete description of how INSITE handles different treatment types (categorical and continuous ). Given that this setup is distinct from the previous ODE discovery methods, complete details should be provided in the main paper\n\n**Underwhelming Experiments**\n- While the paper claims to handle irregularly sampled data, which is typical in real-world settings (i.e., patient covariates are not measured continuously over time), the synthetic experiments don't seem to explore this scenario\n- The synthetic experiments are too simplistic, $g(x)$ is an identity function, $x$ is 1-dimensional. I encourage the authors to explore more challenging scenarios, including semi-synthetic datasets", "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": "ew2SxnTrxL", "reviewer": "Reviewer_aGwv", "summary": "This paper studies longitudal treatment effect estimation from a new prospective, using ordinary differential equation (ODE). This potentially opens up a new paradigm of treatment effect estimation methods that does not rely on neural network. The ODE approach naturally offers interpretability and requires slightly different assumptions for identification. The authors first discussed the difference and similarities between longitudal treatment effect estimation and ODE, then proposed a framework for bridging the problem and solution. The proposed framework is then compared with several neural network based longitudal effect estimators.", "questions": "What are the \"strong assumptions\" mentioned in this statement? \" Unique to the proposed framework and INSITE, is that the discovered\ndifferential equation is fully interpretable; however, it relies on strong assumption\"\n\nThe data generation process in the synthetic benchmark in Eq 5 seems to closely follow with the ODE assumption in Eq 3, is there experiment results when such the data is generated in different ways?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The ODE perspective is very novel for longitudal treatment effect estimation.\n2. The paper is well-written and well-organized, with extensive additional information for reproducibility.\n3. The ODE framework can relax the overlap assumption in causal effect estimation.", "weaknesses": "1. There seems to lack a formal identification result in terms of identifying the causal effects from the ODE perspective.\n2. The experiments are conducted on synthetic datasets only. Results on real-world datasets could further strengthen the evaluation.", "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": "iyhOljAiFf", "reviewer": "Reviewer_WcGm", "summary": "The authors present a framework that can transform any ODE discovery method into a treatment effects method. This involves reformulating the longitudinal heterogeneous treatment effects problem as an ODE discovery problem, with the goal of recovering the underlying system of ODEs based on observed datasets. A proposed model called INSITE (section 5) is built using this framework and tested in accepted benchmark settings.", "questions": "- The assumption 3.1 is not only about existence and uniqueness (these are very mild assumptions and acceptable), but also about the continuous trajectory. I wonder if this is too strict and it may limit the scope of application for the method you've proposed. Could it be relaxed to weaker conditions, such as piecewise continuity? As time changes, different patterns of $\\boldsymbol{x}(t)$ may emerge over the interval from $t_0$ to $T$.\n\n- In practice, what method do you use for solving ODEs? How about the complexity of gradient query / solving for this numerical system? Is your method faster than the neural network-based approaches?", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- The exposition of the method is quite clear. The proposed INSITE framework is very flexible, as it can be used on top of many other time varying TE methods.\n\n- It's reasonable and natural to model this treatment effect over time problem as a dynamical system with ODEs representation. This research brings together two previously separated fields, i.e., temporal TE estimation and ODE discovery, which set up a bridge for these communities.\n\n- The experimental simulation, especially in the appendix, is comprehensive and appears to be reproducible for me.", "weaknesses": "1. The presentation of this paper, such as the related work, should be further improved before acceptance. e.g., the $T$ appears at section 2 but is introduced in section 3 (\"is called the time horizon\"). Another example is ODE discovery paragraph in B.1; I didn't understand why RNN and LSTM are considered ODE discovery methods.\n\n2. Compared to existing black-box TE methods (like neural-network based architectures), the authors claim that the main contribution of this work is a human-readable, interpretable framework. But if the current form of the manuscript only offers empirical validation, I feel that its contributions might not yet meet the threshold of ICLR conf. I wonder that have authors investigated the inherent theoretical properties of the ODE-based framework (which is a white-box model compared to deep network)? e.g., for certain nonlinear differential equations, the solution might be chaotic, meaning it's highly sensitive to initial conditions $\\boldsymbol{x}_0, \\boldsymbol{v}$ and treatment plan $\\boldsymbol{a}$ at $t_0$. \n\n3. Given the claimed emphasis on 'human-readability', how does the proposed method facilitate better decision-making or insights for practitioners in the field is still unclear. Addressing these concerns with more discussions on closed-form ODE would further enhance this work.\n\n4. Minor issues don't affect the rating:\n\n- page 8: $\\gamma = 0$ corresponds no time-dependent -> add 'to' for consistency\n\n- Table 5: word 'Intrepetible' -> 'Interpretable'\n\n- there's one paragraph exceeds the 9 pages limit of the main body\n\n- the full name of ODE/ODEs appears many times in the text, which seems unnecessary", "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": "tZKnWzGddM", "reviewer": "Reviewer_h9Ux", "summary": "The author investigated longitudinal treatment effect estimation from the perspective of ODE discovery. They introduced a framework capable of converting any ODE discovery method into a treatment effect estimation approach. They explored the distinctions between ODE discovery and treatment effect estimation. Furthermore, they presented a versatile solution capable of accommodating both continuous and categorical treatments. Importantly, their method yields interpretable equations, a crucial requirement in the healthcare field.", "questions": "Is it possible to combine any feature learning method within this framework to acquire the necessary dictionary? My primary concern revolves around the challenge of determining the equation's form.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- Method offers interpretable solutions for better understanding.\n- Handles various treatments through defined continuous or piecewise equations.\n- Applicable for estimating treatment effects at both population and individual levels.\n- Manages diverse between-subjects variability effectively.\n- Yields improved results even in cases of model misspecification.", "weaknesses": "- Determining the correct form of ODE is highly challenging.\n- The discussion lacks guidance on learning a feature library, a crucial aspect for this method.\n-Assuming y=x poses the question: what if this assumption is invalid? Would ODE or neural networks be more appropriate to model the outcome function? How might this framework integrate with a neural network or other machine learning techniques as g function for the outcome?", "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": "31tvdDjQ3j", "reviewer": "Reviewer_pnuM", "summary": "The authors propose to solve a task of ODE discovery in the treatment effect estimation setting. Although the task is not sufficiently explained in the main paper, I believe it is to construct an ODE model with simple expressions (e.g., log, exp, sin) from time series data at continuous time stamps. Thus, the paper seems to tackle a novel task. However, it is completely unclear because the motivation is not sufficiently described, the paper is not well structured, the problem setting is not clearly presented, and many important descriptions are moved to Appendix.", "questions": "See all questions in (A), (B), and (C), in particular, (B).", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 1, "contribution": 2, "strengths": "- Introducing the ODE discovery in the context of treatment effect estimation seems novel.", "weaknesses": "(A) Problem setting is unclear\n\n- First of all, in Section 1, please clearly state that what the ODE discovery task is and what its goal is. Otherwise, the readers cannot understand why it yields interpretability. \n\n- It is unclear whether time $t$ is defined as continuous or discrete. For instance, in Section 2, “at $t \\in [T]$” seems to say that the value of time $t$ is discrete and it is included in index set $[T]$. However, for example in Assumption 2.2, “time point $t \\in [0, T]$” seems to say that it is continuous.\n\n- In Section 3, “we assume that time-varying features and outcomes … are discrete measurements of underlying continuous trajectories…” is unclear. I could not understand what is discrete and what is continuous. Do you imply that trajectories are defined at continuous time and their observations are measured at discrete time? Or do you mean that the values of features and outcomes are continuous in ODE models, but their observed values are discrete? Similarly, “discrete (or continuous) treatment plan” is unclear to me. This makes the paper extremely difficult to follow.\n\n- Eq. (3) introduces function $g$, and the authors say “it is known and is often assumed to be the identity.” Why? \n\n- Please separate the paragraph at “The goal of ODE discovery is …” in Section 3. Please do not mix the proposed problem setting and the existing one in the same paragraph, which is hard to follow.\n\n\n(B) The technical soundness is unclear\n\n- **Assumptions**: I could not understand why the authors say that “Assumption 2.2 can be relaxed with Assumptions 3.1 and 3.3.”  I could not see what the point of Discrepancy 1 is. It simply says that “the assumptions in treatment effect estimation do not correspond one-to-one to those in ODE discovery,” which seems trivial. It is hard to follow.\n \n- **The soundness of employing ODE discovery methods**: Although the proposed method employs the existing ODE discovery method (called SINDy), it is unclear what this method does and how it works. What is the advantages and the limitations of this method? How can we choose the existing methods? \n\n- **Ability to express the variability of each subject**: The authors seem to simply change the constant values in simple ODE models to express the variability across subjects. Can we really always express the treatment and the outcome curves with such a simple strategy, given that the continuous time series data often involve complex dynamics?\n\n(C) The difference from existing ODE-based treatment effect estimation methods is unclear in the main paper\n\n- Please briefly summarize the description of “Dynamical systems, ODEs, and treatment effects” in Appendix B.1 in the main paper. This will make the motivation of this paper clearer.", "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 (spotlight)