ODE Discovery for Longitudinal Heterogeneous Treatment Effects Inference

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 a 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, a 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., 2020b)—used to model the outcome. Conversely, our new approach does not adjust the dataspace in any way, and instead discovers a global ODE which we refine per patient, as 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 research and data collection, given that we can now understand certain behaviours of 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 a usable framework that allows us 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 section˜C.1 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)¹¹1We provide code at https://github.com/samholt/ODE-Discovery-for-Longitudinal-Heterogeneous-Treatment-Effects-Inference., tested in accepted benchmark settings used throughout the literature. Transforming an ODE discovery method into a TE method results in a new set of identification assumptions. This need not be a limitation, instead, it can be seen as an extension as the typical TE identification assumptions may not always hold in ODE discovery settings.

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 $\bm{X}^{(i)}_{t}\in\mathbb{R}^{d}$, $\bm{A}^{(i)}_{t}\in\mathcal{A}$ and $Y^{(i)}_{t}\in\mathcal{Y}$ be the $i^{\textrm{th}}\in[N]$ individual’s features, treatment, and outcome at time $t\in[0,T]$, with $T$ the time horizon. We also denote the static covariates of the $i^{\textrm{th}}$ individual as $\bm{V}^{(i)}\in\mathbb{R}^{m}$. In the treatment effects literature, typically $\mathcal{A}=\{0,1\}$ and $\mathcal{Y}=\mathbb{R}$. Unless required for clarity, we will drop the individual indicator $i$. We use the following notation $\bm{X}_{t_{1}:t_{2}}$ to denote the observed features in the time window $[t_{1},t_{2}]$. The observed dataset contains the realizations of the random variables above, $\mathcal{D}=\{(\bm{v},\langle\bm{x}_{t},y_{t},\bm{a}_{t}\rangle)^{(i)}:i\in[N],t\in[0,T^{(i)}]\}$.

Of interest is estimating the expected potential outcome $Y_{t:t+\tau}(\bar{\bm{a}}_{t:t+\tau})$, for some $\tau>0$ under hypothetical future treatments $\bar{\bm{a}}_{t:t+\tau}$ given the historical features $\bm{X}_{0:t}$ and the previous treatments $\bm{A}_{0:t}$ (Neyman, 1923; Rubin, 1980):

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 to treatment effects and causal inference (Holland, 1986).

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-time causal effects (Lok, 2008; Saarela & Liu, 2016; Ryalen et al., 2019).

Assum.˜2.2 and assum.˜2.3 rely on a treatment intensity process, $\lambda(t|\mathfrak{F}_{t})$ 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 $\mathfrak{F}_{t}$. Furthermore, assum.˜2.3 ensures it is sufficient to condition on a patient’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):

which is exploited by most works in the treatment effects literature (albeit by first regularizing models such that they respect assum.˜2.1, 2.2 and 2.3) (Bica et al., 2020b; Lim, 2018; Melnychuk et al., 2022). A thorough overview of related works is presented in appendix˜C.

We now propose to model this treatment effect over time problem as a dynamical system, where their temporal evolution can often be well represented by ODEs (Hamilton, 2020). That means we assume the time-varying features and outcomes of the $i^{\text{th}}$ individual ($\bm{x}^{(i)}_{t}\ \forall t\in[0,T^{(i)}]$) are discrete (and possibly noisy) measurements of underlying continuous trajectories of observed features $\bm{x}^{(i)}:[0,T]\rightarrow\mathbb{R}^{d}$ and potential outcomes $y^{(i)}:[0,T]\rightarrow\mathbb{R}$, where $T\in\mathbb{R}$ is called the time horizon (Birkhoff, 1927). To make our formalism consistent, we also assume there is a treatment trajectory $\bm{a}:[0,T]\rightarrow\mathbb{R}^{k}$ such that $\bm{a}^{(i)}_{t}$’s either constitute snapshots of the underlying continuous treatment $\bm{a}$ or if the treatments are administered at discrete times then $\bm{a}$ is a step function whose values corresponds to the currently administered treatment.

There are many fields which already use ODEs as a formal language to express time dynamics. One such example is pharmacology, where a large portion of the literature is dedicated to 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 is modelled by a system of ODEs which describes the time derivative of $\bm{x}$ as a function of $\bm{x},\bm{v}$, and $\bm{a}$. The outcome trajectory, $y$, depends on the features $\bm{x}$. In particular, we assume

where $\bm{\dot{x}}(t)$ is the differential of $\bm{x}$. $g$ is prespecified by the user and often models the outcome as one of the features ($g(\bm{x})=x_{j}$), e.g., the tumour volume. This is a general formulation that may also include direct effects of $\bm{a}$ on $y$. We further discuss eq.˜3 in appendix˜D. The goal of ODE discovery is to recover the underlying system of ODEs $\bm{F}$ based on the observed dataset $\mathcal{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:

Identification. Assum.˜3.1 and 3.2 play a crucial role in ODE discovery. Essentially, we require making such assumptions in order to correctly identify the underlying equation. The assumptions made in the treatment effects literature (assum.˜2.3, 2.2 and 2.1) serve a similar purpose as they allow us to interpret the estimand as a 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 a 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.

From eqs.˜2 and 3, one can see that both treatment effects and ODE discovery involve learning a function that can issue predictions about future states. Although the connection is natural, there remain some discrepancies between these two fields. To apply ODE discovery methods for treatment effects inference, we have to resolve these discrepancies. Essentially, we establish a framework one can use to apply any ODE discovery technique in treatment effects.

We identify three discrepancies between the treatment effects literature and the ODE discovery literature: (1) different assumptions, (2) discrete (not continuous) treatment plans, and (3) variability across subjects. Each discrepancy is explained and reconciled in a dedicated subsection below with actionable steps.

Figure˜2 shows the areas ODE discovery methods can be expanded to, using our framework. Ranging from simple adaption, to proposing a completely new method (in section˜5), our framework significantly increases the reach of existing ODE discovery methods. Our new method—INSITE—should be considered the result after implementing our practical framework we present in the remainder of this section.

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 a new set of assumptions (discr.˜1), (FS2) incorporate treatment plans (discr.˜2), and (FS3) decide on the BSV type (discr.˜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 ($\beta\in\mathbb{R}^{m}$, with $m$ numeric constants), and (2) we accept that BSV is modelled through these numeric constants alone, i.e., two varying subjects can only differ by having different model numeric constants.⁴⁴4Essentially, we cannot have a mathematical expression changed across subjects (such as an $\exp$ changed to a $\log$. In fact, a scenario such as this would violate assum. 3.3).

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 $\bar{\bm{F}}$ for all patients, and (2) discover the patient-specific differential equation $\bm{F}^{(i)}$ by fine-tuning the population equation’s numeric constants—as outlined in algorithm˜1. Having a set of numeric constants (one for each patient) we derive a population differential equation where each individual’s numeric constants are represented by a sample from this population differential equation distribution.

◆ Step 1 (cfr. table˜1): Discovering the Population Differential Equation.  Given an observed dataset $\mathcal{D}$ 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., 2016) to recover the population ODE $\bar{\bm{F}}$. We adapt it to handle treatments as outlined in section˜4.2 and appendix˜E (as per discr.˜2), where for each separate $k$ categorical (or binary) treatment, we discover a separate ODE—which is only active for that categorical treatment, i.e., $\bar{\bm{F}}=\{\bar{\bm{F}}_{1},\dots,\bar{\bm{F}}_{k}\}$.

◆ Step 2 (cfr. discr.˜3): Discovering Patient-Specific 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 $\bm{F}^{(i)}$. 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-subject variability). Crucially, to avoid overfitting and allow only small deviations away from the initial population numeric constants $\bar{\bm{\beta}}$, we fit the observed patient trajectory up to the current time by minimizing,

a Mean Squared Error (MSE) of the predicted trajectory evolution $\hat{\bm{Y}}$ for the observed trajectory. Additionally, we also require a regularization term in eq.˜4. Specifically, for each individual trajectory, $(i)$, we refine the population set of parameters ($\bar{\bm{\beta}}$) to obtain $\bm{\beta}^{(i)}$ 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, we find INSITE to underperform significantly (cfr. section˜K.1). We use Broyden–Fletcher–Goldfarb–Shanno (BFGS) algorithm (Fletcher, 2013) to minimize eq.˜4 as is standard in symbolic regression (Petersen et al., 2020), and provide full inference details in section˜G.1.

Parameter distributions. Unique to our method 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 as a normal distribution or a mixture of distributions—recovering a probabilistic interpretation of the underlying data generating process. This is explored in section˜K.2.

To allow for a robust and systematic comparison of longitudinal treatment effect (LTE) methods and ODE discovery techniques, we create a synthetic testbed designed for treatment effect prediction 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 are unknown (Lim, 2018; Bica et al., 2020b). Our testbed consists of:

Diverse underlying pharmacological models $\bm{F}$. The synthetic datasets are generated by sampling a pharmacological model $\bm{F}$ with a 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-Pharmacodynamic (PKPD) models for each model class (section˜4.3). First, a one-compartmental PKPD model (eq.˜5) with a binary static action. Second, a state-of-the-art biomedical PKPD model of tumor growth, used to simulate the combined effects of chemotherapy and radiotherapy in lung cancer (Geng et 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 $y=x$, is the volume of the tumor $t$ days after diagnosis, modeled separately as:

where the parameters $C_{0},C_{1},V,\rho,K,\gamma,\alpha,\beta,e_{t}$ 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-the-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 a threshold of the outcome predictor value, controlled by a scalar $\gamma\in\mathbb{R_{+}}$—such that $\gamma=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.

Benchmark methods. We seek to compare against the existing state-of-the-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. 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). Second, ODE discovery methods aim to discover an underlying closed-form mathematical equation $\bm{F}$ that best fits the underlying controlled ODE of the observed state and action trajectories dataset, $\mathcal{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 a variational loss as the objective, as in Weak Sparse Identification of Nonlinear Dynamics (Messenger & Bortz, 2021)⁵⁵5We only include results for WSINDy when it is possible to use it, as it does not support sparse short trajectories. Detailed further in appendix G.. 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). We further discuss why we chose these SOTA methods and their implementation details in appendix˜G.

Evaluation. We evaluate against the standard longitudinal treatment effect metrics of test counterfactual $\tau$-step ahead prediction normalized root mean squared error (RMSE), where $\tau=\{1,2,3,4,5,6\}$, for a sliding treatment plan (Melnychuk et al., 2022). Unless otherwise specified, each result is run for five random seed runs with time-dependent confounding of $\gamma=2$—see appendix˜H for details.

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 to 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 are 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 it 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 $\tau$-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 $\mathcal{L}_{\text{INSITE}}=\{1,x_{0},x_{1},x_{0}x_{1}\}$. 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 $\mathcal{L}_{\text{Cancer}}=\{1,x_{0},x_{1},x_{0}x_{1},x_{0}^{2}x_{1},x_{0}x_{1}^{2},\log(x_{0}),\log(x_{1})\}$ 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 $\gamma$. We observe, in fig.˜3 (b), that INSITE can still discover a good approximating underlying equation even in the presence of high time-dependent confounding—hence where there is low overlap.

INSITE Ablation. INSITE’s gain in performance derives from our framework’s discr.˜2, to discover individual ODEs per categorical treatment, as well as the ability to discover individualized (fine-tuned) ODEs. Notably, the improvement arising from the ability to 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.

In conclusion, we presented a 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 our 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).

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, we list 3 major limitations one should take into account before applying our work in practice:

1. 1.

   A correct set of candidate functions (tokens) is necessary for correct model recovery. To show the importance of this, we include experiments where we have wrongly specified this token library (see section˜K.6) and observe degrading performance as a result.

2. 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 a usability point of view, the usefulness of non-sparse equations is limited as interpretability is negatively affected by non-sparse (or non-parsimonious) equations (Crabbe et al., 2020).

3. 3.

   ODEs are noise free. Since we recover ODEs, the found equations do not model a source of noise as is 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 at https://github.com/samholt/ODE-Discovery-for-Longitudinal-Heterogeneous-Treatment-Effects-Inference. 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.

Acknowledgements. 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.vanderschaar-lab.com), for their valuable input, comments, and suggestions as the paper was developed.