Paper: ODE Discovery for Longitudinal Heterogeneous Treatment Effects Inference (ICLR-2024)

Reviewer comment (Writing):
In Section 1, please clearly state what the ODE discovery task is and what its goal is. Otherwise, the readers cannot understand why it yields interpretability.

Author response:
1. In Section 1, we now explicitly define the goal of ODE discovery as "finding underlying closed-form concise ODEs from observed trajectories," clarifying how these transparent equations lead to interpretability.

Reviewer comment (Writing):
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 included in index set \([T]\). However, in Assumption 2.2, “time point \( t \in [0, T] \)” seems to indicate that it is continuous.

Author response:
2. To clarify the nature of time in our model, Section 2 has been updated to specify time as a continuous variable, represented as $ t \in [0, T] $, to eliminate any confusion. We also clarify explicitly that $T$ is the time horizon.

Reviewer comment (Writing):
In Section 3, “we assume that time-varying features and outcomes … are discrete measurements of underlying continuous trajectories…” is unclear. It is confusing 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. This makes the paper extremely difficult to follow.

Author response:
3. In Section 3, we have refined our explanation: trajectories are indeed defined in continuous time, while observations are made at discrete time points (at possibly varying intervals). We also clarified the concept of "discrete (or continuous) treatment plan" in relation to the set of allowed treatments $\mathcal{A}$.

Reviewer comment (Theory):
In Eq. (3), function \( g \) is introduced, and the authors say “it is known and is often assumed to be the identity.” The reason for this assumption is not explained.

Author response:
4. The role of $g$ in Eq. (3) is to describe the treatment outcome as a function of the observed features. Thus, knowing $g$ is the same as just specifying the outcome variable. It is often assumed to correspond to one of the features, e.g., the tumor volume. We have clarified this point. For a detailed discussion about Eq. (3), please take a look at Appendix C.

Reviewer comment (Presentation):
Please separate the paragraph at “The goal of ODE discovery is …” in Section 3. Do not mix the proposed problem setting and the existing one in the same paragraph, as this makes it hard to follow.

Author response:
5. We have separated the paragraphs in Section 3 to distinctly delineate our proposed problem setting from existing ones, enhancing readability.

Reviewer comment (Writing):
Although the proposed method employs the existing ODE discovery method (called SINDy), it is unclear what this method does and how it works. The advantages and limitations of this method are not explained, nor is there discussion on how to choose among existing methods.

Author response:
_The soundness of employing ODE discovery methods_ - We have added an additional appendix M that describes SINDy, how it works, and how to choose other existing methods. This further expands on the original method description of SINDy in Appendix F.

**_Actions taken:_** We include a new Appendix M to explain SINDy and other possible ODE learners.

Reviewer comment (Novelty):
The difference from existing ODE-based treatment effect estimation methods is unclear in the main paper. 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 the paper clearer.

Author response:
**\(C\) ODE-based TE inference and ODE-discovery for TE inference.** Thank you for your suggestion. We have now included a summary of the relevant descriptions from Appendix B.1 in the main paper, outlining the distinctions between our method and existing ODE-based treatment effect estimation methods. This addition aims to clarify the motivation and unique contributions of our work.

**_Actions taken:_** A summary of B.1 is included in the main text.

Reviewer comment (Theory):
Determining the correct form of ODE is highly challenging.

Author response:
**Learning ODEs.** You are correct. ODE learning is an incredibly challenging task and, in certain scenarios, infeasible. For these scenarios, we recommend a practitioner rely on the tried-and-tested pure inference machines for treatment effects inference (such as those built on top of neural networks or random forests and beyond).

However, there are also multiple scenarios where ODE learning actually _is_ feasible, and moreover- **_preferred_**. These scenarios include pharmacology, physics, and even economics. The reason is that existing literature on these topics (outside of machine learning) is already heavily driven by the need to discover and recover underlying ODE systems. 

**_Actions taken:_** We include pharmacology as an example in Section 3 to further emphasize where ODE discovery makes sense. We also explicitly list the limitations associated with INSITE (and our general framework) in our ethics statement. Hopefully, guiding a practitioner on when (and when not) to use our ODE-learner-based approach to TE inference.

Reviewer comment (Reproducibility):
The discussion lacks guidance on learning a feature library, a crucial aspect for this method.

Author response:
**Learning feature libraries.** Setting a good feature library requires two important considerations:

1. Which ODE method are we using?
2. In which scenario are we applying the method?

For 1, we have included a new Appendix M, which explains SINDy (our underlying ODE learner) in detail, as well as points the interested reader to alternative ODE learners, should SINDy be insufficient (wrt to feature sets or other reasons). 

For 2, our method heavily relies on domain knowledge. We believe domain knowledge should, of course, take into account point 1 when making a decision on the ODE learner. While this could be problematic, the same is also true for _any_ ML method to be applied in practice. 

**_Actions taken:_** We include Appendix M, which explains, in detail, SINDy and provides alternatives.

Reviewer comment (Presentation):
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.

Author response:
**1. Presentation of the paper.** Thank you for pointing us to these necessary corrections. We have taken the following steps in our revision to address your concerns:

1. _On $T$:_ We have revised the paper to introduce the concept of the time horizon earlier in Section 2.
2. _On RNNs and LSTMs for ODE discovery:_ In Appendix B.1, we have amended the text to replace references to RNN and LSTM with a more accurate discussion of Neural ODEs. This modification provides clarity on our approach and aligns better with the context of ODE discovery.

We believe these changes enhance the overall presentation and coherence of the paper, and we hope that they address your concerns effectively.

**_Actions taken:_** Adjusted the paper where necessary to incorporate the above suggestions.

Reviewer comment (Theory):
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$.

Author response:
**2. Theoretical Contributions.** Our main contribution lies in creating a _usable framework_ that translates ODE discovery methods into the treatment effects problem formulation, bridging two previously distinct fields. 

Our theoretical contribution is the alignment of two utterly different problem formulations and assumptions between treatment effects inference and ODE discovery. As shown in Figure 2, ODE discovery methods only work out of the box in the simplest settings (i.e., continuous treatments and only one kind of between-subject variability (BSV) corresponding to a noisy measurement). We demonstrate how to extend ODE discovery methods to allow for discrete treatment plans (both dynamic and static) and two additional sources of BSV: covariate models and parameter distributions (BSV levels C and D). Identification of these two dimensions (axes in Figure 2) and subsequent reconciliation of discrepancies (as shown in Section 4) constitute a major theoretical milestone.

While the theoretical properties of ODE discovery are not the central focus of this work, we acknowledge their importance and direct readers to references A and B for foundational background in this area.

**_Actions taken:_** We have included references [A] and [B] in the main text of our paper, referring the interested reader to the many theoretical underpinnings of ODE discovery.

Reviewer comment (Evaluation):
Given the claimed emphasis on 'human-readability', how the proposed method facilitates 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.

Author response:
**3. Facilitating decision-making and insight.** We have included Appendix I, where we demonstrate that our framework, INSITE, successfully identifies equations that are not only predictive but also closely resemble or exactly match the true underlying equations, facilitating further analyses, such as stability analysis, impact analysis, and optimal treatment analysis.

As we have now elaborated in Section 3, many fields 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 ([D],[E]).

Furthermore, in Section 6.1, we illustrate that our method remains effective even when the common overlap assumption is not met. This aspect of INSITE enhances decision-making in complex scenarios where traditional assumptions may not hold.

Reviewer comment (Writing):
Minor issues don't affect the rating: page 8: $\gamma = 0$ corresponds no time-dependent -> add 'to' for consistency.

Author response:
**4. Corrections:** Thank you for pointing us to these errors; they have now been implemented. We kindly note that the ethics and reproducibility statement do not count toward the page limit, as per the official ICLR author guide instructions: https://iclr.cc/Conferences/2024/AuthorGuide.

Reviewer comment (Writing):
Minor issues don't affect the rating: Table 5: word 'Intrepetible' -> 'Interpretable'.

Author response:
**4. Corrections:** Thank you for pointing us to these errors; they have now been implemented. We kindly note that the ethics and reproducibility statement do not count toward the page limit, as per the official ICLR author guide instructions: https://iclr.cc/Conferences/2024/AuthorGuide.

Reviewer comment (Writing):
Minor issues don't affect the rating: the full name of ODE/ODEs appears many times in the text, which seems unnecessary.

Author response:
**4. Corrections:** Thank you for pointing us to these errors; they have now been implemented. We kindly note that the ethics and reproducibility statement do not count toward the page limit, as per the official ICLR author guide instructions: https://iclr.cc/Conferences/2024/AuthorGuide.

Reviewer comment (Reproducibility):
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?

Author response:
**Solving ODEs.** To solve the discovered ODE, we use a standard numerical ODE solver of the Euler method, although other solvers can also be used, such as Runge-Kutta methods. At inference time, this approach is comparable to the inference time with a neural network to make a prediction. Interestingly, discovering ODEs, such as with the underlying sparse identification of nonlinear dynamics (SINDy) method, is faster than training neural network-based TE methods.

Reviewer comment (Experiments):
The experiments are conducted on synthetic datasets only. Results on real-world datasets could further strengthen the evaluation.

Author response:
**2. Experiments.** Note that we adopt the experimental settings from the related literature on ODE discovery. While in TE literature, it is common to rely on semi-synthetic data, truly testing our method requires more than just access to the unobserved potential outcome (which is the reason TE literature resorts to semi-synthetic testbeds). Specifically, we also require access to the underlying causal system, which is only possible with fully synthetic data. 

We also kindly refer to Appendix J, which includes many more experiments that may interest the reviewer. In particular, we included experiments on:
* Sensitivity on $\lambda$
* Parametric distributions of numeric constants
* Varying sample sizes and higher noise
* Model misspecification

Reviewer comment (Theory):
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., the success of the INSITE is dependent on the choice of the library of candidate functions.

Author response:
**Limitations.** We agree with the reviewer, limitations should be clearly stated:
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 Appendix J.6) and observe degrading performance as a result. We will explicitly refer to Appendix J.6 from our concluding section.

**_Actions taken:_** The above limitations (1-3) are now stated in our ethics statement. We include it there as it should be considered explicitly before applying our method (and framework) in practical settings.

Reviewer comment (Theory):
The paper makes the strong assumption that INSITE assumes the system is sparse, which rarely holds in high-dimensional settings.

Author response:
**Limitations.** We agree with the reviewer, limitations should be clearly stated:
2. Sparsity is an assumption stemming from the type of function we wish to recover. Only when the function is sparse (that is, it contains not too many terms) can we deem it interpretable. There is nothing technical that limits INSITE (or the underlying SINDy method) to sparse functions; however, relaxing sparsity would also result in less useful recovered equations.

**_Actions taken:_** The above limitations (1-3) are now stated in our ethics statement. We include it there as it should be considered explicitly before applying our method (and framework) in practical settings.

Reviewer comment (Theory):
The paper makes the strong assumption that INSITE assumes 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.

Author response:
**Limitations.** We agree with the reviewer, limitations should be clearly stated:
3. ODEs are noise-free. Since we recover ODEs, we therefore recover noiseless (deterministic) functions. In case one is interested in stochastic differential equations (SDEs), one needs to extend our work to encompass this type of differential equation. 

**_Actions taken:_** The above limitations (1-3) are now stated in our ethics statement. We include it there as it should be considered explicitly before applying our method (and framework) in practical settings.

Reviewer comment (Reproducibility):
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.

Author response:
**Treatment systems.** For a complete description of incorporating different treatment types, we refer to Appendix D. Due to the extent of our discussion, we moved this to our supplementary materials. In particular, Table 6 (in Appendix D) shows how a non-continuous treatment is transformed into a dynamical system. However, on page 5, we have included a brief description of the table presented in Appendix D.

**_Actions taken:_** Include a summary of Table 6 in Appendix D in the main text.

Reviewer comment (Experiments):
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.

Author response:
### **(A) 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 [B, C]. 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 irregularly, excluding 10\% of the original observations along the time dimension. The subsequent results, conducted across ten random seeds, are outlined in table 29 below. We observe that INSITE and SINDy still achieve good performance (hence low prediction error).

**Table 29**
| Method | eq.(5).A | eq.(5).B | eq.(5).C | eq.(5).D | eq.(6).A | eq.(6).B | eq.(6).C | eq.(6).D | Cancer PKPD |
|--------|--------------------------------|--------------------------------|--------------------------------|--------------------------------|----------------|----------------|----------------|----------------|-------------|
| A-SINDy | 0.11±1.93e-17 | 0.11±1.93e-17 | 0.12±1.93e-17 | 0.13±0.00 | 0.90±1.54e-16 | 0.90±1.54e-16 | 0.80±0.00 | 0.96±1.54e-16 | 1.68±3.08e-16 |
| INSITE | **0.02±0.00** | **0.03±0.00** | **0.04±0.00** | **0.04±0.00** | **0.85±1.54e-16** | **0.85±0.00** | **0.75±0.00** | **0.70±0.00** | **1.00±1.54e-16** |


***Actions taken:*** We now include this new experimental result in a **(new) Appendix J.7**, entitled **"INSITE also empirically works for irregularly sampled data"**.

Reviewer comment (Experiments):
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.

Author response:
### **(B) INSITE can also work with more challenging observation functions $g(x)$**

Although the observation function $g$ that describes the treatment outcome as a function of the observed features is often prespecified by the user as it often models the outcome as one of the features ($g(\mathbf{x})=x_j$), 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 $g$ be a 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(\mathbf{x})=\exp(x_j)$ that often models rapid changes or growth, commonly found in many biological and economic systems [D], (2) a quadratic function $g(\mathbf{x})=x_j^2$ that models a non-linear relationship in systems [E], and (3) a trigonometric function $g(\mathbf{x})=\sin(x_j)$ that often models periodic components within oscillatory dynamics, which are commonly found in natural and engineering systems [F]. 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.

**Table 30**
| $g(x)$    | Method | eq.(5).A | eq.(5).B | eq.(5).C | eq.(5).D | eq.(6).A | eq.(6).B | eq.(6).C | eq.(6).D | Cancer PKPD |
|-------------|------------|-------------------------------|-------------------------------|-------------------------------|-------------------------------|----------------|----------------|----------------|----------------|-----------------|
| **$x$**       | A-SINDy    | 0.11 ±0.00                    | 0.11 ±0.00                    | 0.13 ±0.00                    | 0.15 ±0.00                    | 1.46 ±0.00     | 1.47 ±0.00     | 1.40 ±0.00     | 1.56 ±0.00     | 1.29 ±0.00      |
|             | INSITE     | 0.02 ±0.00                    | 0.03 ±0.0                     | 0.04 ±0.00                    | 0.05 ±0.00                    | 0.96 ±0.00     | 0.96 ±0.00     | 0.82 ±0.00     | 0.90 ±0.00     | 0.79 ±0.00      |
| **$\exp(x)$**  | A-SINDy    | 0.75 ±0.02                    | 0.75 ±0.02                    | 0.74 ±0.03                    | 0.76 ±0.06                    | 0.76 ±0.34     | 0.76 ±0.34     | 1.21 ±1.34     | 1.02 ±0.57     | 0.94 ±0.20      |
|             | INSITE     | 0.59 ±0.03                    | 0.59 ±0.02                    | 0.61 ±0.02                    | 0.62 ±0.05                    | 0.67 ±0.30     | 0.68 ±0.32     | 1.08 ±1.31     | 0.90 ±0.59     | 0.78 ±0.13      |
| **$x^2$**     | A-SINDy    | 0.10 ±0.01                    | 0.12 ±0.00                    | 0.12 ±0.00                    | 0.12 ±0.01                    | 1.16 ±1.03     | 1.15 ±1.03     | 2.11 ±3.58     | 1.69 ±2.02     | 0.57 ±0.10      |
|             | INSITE     | 0.03 ±0.00                    | 0.09 ±0.00                    | 0.09 ±0.00                    | 0.09 ±0.01                    | 1.12 ±1.04     | 1.11 ±1.04     | 2.04 ±3.57     | 1.63 ±2.03     | 0.52 ±0.08      |
| **$\sin(x)$**  | A-SINDy    | 2.42 ±0.07                    | 2.42 ±0.07                    | 2.31 ±0.07                    | 2.42 ±0.40                    | 1.62 ±0.43     | 1.62 ±0.43     | 2.30 ±1.86     | 1.99 ±0.69     | 2.00 ±0.46      |
|             | INSITE     | 2.09 ±0.10                    | 2.09 ±0.10                    | 2.06 ±0.08                    | 2.22 ±0.25                    | 1.46 ±0.36     | 1.41 ±0.42     | 1.96 ±1.78     | 1.72 ±0.76     | 1.59 ±0.29      |

***Actions taken:*** We now include this new experimental result in a **(new) Appendix J.8**, entitled **"INSITE can also work with more challenging observation functions $g(x)$"**.

Reviewer comment (Evaluation):
For the BSV scenario A (eqns 5 and 6), shouldn't we expect ODE methods to do as well as INSITE.

Author response:
**BSV scenario A.** We agree that for the BSV scenario A (eq.5 and eq.6), we expect ODE discovery methods to perform as well as INSITE. This discrepancy could likely occur due to the accuracy to which the parameters are optimized internally within the SINDy method. We observe that an alternative optimization method that is more accurate than that of WSINDy (Messenger et al., 2021) achieves a better performance than SINDy. Whereas INSITE fine-tunes the ODEs parameters to the initially observed trajectory of a patient at inference time using a higher-accuracy optimization method of the Broyden–Fletcher–Goldfarb–Shanno (BFGS) algorithm, a second-order nonlinear optimization method.

Reviewer comment (Evaluation):
Fig 3(b): Shouldn't we expect INSITE to be monotonically increasing with \(\gamma\).

Author response:
**INSITE and $\gamma$.** We agree that INSITE's error is monotonically increasing with $\gamma$. However, we note that it still achieves a more accurate prediction compared to existing methods.

The fact that INSITE's error increases with $\gamma$ is also expected. As $\gamma$ increases, overlap gets more violated. While INSITE is _more robust_ against overlap violations, it is not immune. The reason is that the hypothesis space associated with ODEs is still large enough that we cannot confidently extrapolate to low-sample regions (such as regions of low/non-overlap). On the other hand, the ODE hypothesis space is still smaller than the neural networks used in other TE inference techniques.