ODE Discovery for Longitudinal Heterogeneous Treatment Effects Inference

We thank the reviewer for their time and helpful comments. Our rebuttal below follows the structure of your review. To facilitate further discussion, we have also stated how each specific comment changed our paper

(A) Problem setting. Thank you for highlighting areas where our problem setting needed clarification. We have made the following revisions:

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.

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.

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}$.

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.

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

These revisions aim to make the problem setting clearer and the paper easier to follow.

Actions taken: Revisions 1-5 (outlined above) are applied to our paper.

rebuttal continues in next comment...

Thank you once again for your invaluable insights and time dedicated to the review process!

If there were any leftover concerns, we would sincerely appreciate the opportunity to clarify them—before the discussion period for the authors ends in one day. We believe our responses (Nov 15) have addressed in detail the full set of questions you raised:

• Improving the clarity of the problem setting; see Response (A)
• Exposition on technical soundness and clarifications; see Response (B)
• Clarifications between existing ODE-based TE methods; see Response (C)

These are paired with corresponding updates to the submission (Nov 15), including extensive new additional experimental results, discussions, and clarifications—which include two new experimental appendices appendix J.7 and J.8, entitled "INSITE also empirically works for irregularly sampled data" and "INSITE can also work with more challenging observation functions $g(x)$" respectively.

We would appreciate it if the reviewer kindly let us know if there were any further questions in the very limited time remaining of one day. We are eager to do our utmost to address them!

Thank you,

Paper 5372 authors

Thank you once again for your constructive comments that have ultimately improved the paper, and your time dedicated to the review process!

Our initial change highlighting inadvertently compared the current version of the paper to an old git commit, leading to an overestimation of the changes made. We have now corrected this by highlighting only the text changed during the author discussion period, compared to the original submission (Sep 28). These changes can be verified using the OpenReview pdf difference tool. We find our changes to the paper to be minimal, and inline with addressing the reviewers comments—to be comprehensive we list out all these changes below, and their type:

• Abstract: Clarified ODE discovery methods discover an equation, by changing "human-readable equation" to "equation".
• Introduction: Clarified the goal of ODE discovery with the addition of “the goal of ODE discovery is finding underlying closed-form concise ODEs from observed trajectories.”
• Introduction: Clarified ODE discovery methods discover an equation, by changing "human-readable equation" to "equation".
• Introduction: Clarified contribution using ODE discovery methods by “While other TE inference methods have used ideas from ODE literature, none have focused on ODE discovery, we have devoted appendix B.1 to explain this subtle but important difference.”
• Section 2. Clarified time to be continuous “$t \in [0,T]$, with $T$ the time horizon”.
• Section 3. Renamed section to “Underpinnings of ODE discovery”
• Section 3. Added additional clarification and example of how and where ODEs are used by “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 literature is dedicated to recover 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).”
• Section 3. Clarified $g$ as “$g$ is prespecified by the user and often models the outcome as one of the features ($g(\mathbf{x}) = x_j$), e.g., the tumor volume.”
• Section 3. Clarified ODE trajectories by “Of course, the reliance on such a dataset implies that, while the trajectories are defined in continuous time, they are observed in discrete time”.
• Section 4.1. Clarified the four different treatment types by “In summary, there are 4 different treatment types: binary, categorical, multiple, and continuous treatments (Bica et al., 2020). Each of them can be a static or dynamic treatment, which are either constant or change during a trajectory, respectively. This results in 8 scenarios, which we 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 6 in appendix D.”
• Section 7. Added a future work of “learning piece-wise continuous ODE systems (Wang et al., 2023)”
• Ethics statement. Clarified the limitations with listing three limitations that a user should take into account before using the method in practice.
• New appendix J.7 on empirical results to show that “INSITE also empirically works for irregularly sampled data”.
• New appendix J.8 on empirical results to show that “INSITE can also work with more challenging observation functions $g(x)$”.
• New appendix M on “How to choose an ODE discovery method”.
• New appendix N on “Numerical Solver and Parameters”

In summary the total changes to the paper during the author rebuttal discussion period are only additional clarification changes or the addition of additional empirical results for special settings. Each change directly addresses feedback from the review process, enhancing the paper’s clarity, depth, and practical relevance. We assure you that the core arguments, methodology, and results remain unchanged. The fundamental nature and validity of our work have not been altered; rather, the paper has been refined to better communicate its value.

As the author discussion period draws to a close, we remain open and responsive to any further questions or concerns. Your feedback has been invaluable, and we are committed to addressing any remaining issues promptly.

We hope this clarifies your concern and you will kindly consider accepting the paper. Thank you once again for your invaluable insights and the opportunity to improve our manuscript.

Kind regards,

Paper 5372 authors

We thank the reviewer for their time and helpful comments. Our rebuttal below follows the structure of your review. To facilitate further discussion, we have also stated how each specific comment changed our paper

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.

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.

We thank the reviewer once more for their insightful and positive review. We believe our paper is improved as a direct result of your comments. Should there be any more questions, please ask.

We thank the reviewer for their time and helpful comments. Our rebuttal below follows the structure of your review. To facilitate further discussion, we have also stated how each specific comment changed our paper

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.

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.

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.

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.

rebuttal continues in next comment...

Thank you once again for your invaluable insights and time dedicated to the review process!

If there were any leftover concerns, we would sincerely appreciate the opportunity to clarify them—before the discussion period for the authors ends in one day. We believe our responses (Nov 15) have addressed in detail the full set of questions you raised:

• Improving the presentation of the paper; see Response (1)
• Exposition of the theoretical contributions; see Response (2)
• Facilitating decision-making and insights; see Response (3)
• Corrections; see Response (4)
• Relaxing assumption 3.1; see Response (Questions - Assumption 3.1)
• Clarification on how ODEs are solved; see Response (Questions - Solving ODEs)

These are paired with corresponding updates to the submission (Nov 15), including extensive new additional experimental results, discussions, and clarifications—which include two new experimental appendices, appendix J.7 and J.8.

We would appreciate it if the reviewer kindly let us know if there were any further questions in the very limited time remaining of one day. We are eager to do our utmost to address them!

Thank you,

Paper 5372 authors

Thank you once again for your constructive feedback and for dedicating your time to the review process!

Regarding the term 'human-readable equation,' we initially used this phrase to emphasize the interpretability of the discovered closed-form ODE. However, we acknowledge that this term might have been ambiguous without a detailed explanation. Thus, we have revised the manuscript, replacing 'human-readable equation' with 'equation' in the abstract and introduction to avoid any misunderstanding—therefore, this term does not appear in the paper anymore.

Concerning your valuable insights on numerical problems and solver details, we have now added an additional (new) appendix N titled “Numerical Solver and Parameters”, which details in one place the following.

The numerical method (solver) used for INSITE and the other ODE discovery methods (e.g., SINDy) is that of the Euler method, which is a first-order numerical algorithm for solving ODEs given an initial value. This uses a time step size of $\delta_t = 0.166$ seconds, which was sufficient for our 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 an 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 Euler numerical solver to estimate future values at future times. We also kindly highlight further implementation details can be found in appendices E, F and H.

As the author discussion period draws to a close, we remain open and responsive to any further questions or concerns. Your feedback has been invaluable, and we are committed to addressing any remaining issues promptly.

We hope this clarifies your concern and you will kindly consider accepting the paper. Thank you once again for your invaluable insights and the opportunity to improve our manuscript.

Kind regards,

Paper 5372 authors

We thank the reviewer for their time and helpful comments. Our rebuttal below follows the structure of your review. To facilitate further discussion, we have also stated how each specific comment changed our paper

1. Identification. When the complete causal system is known (including the underlying structural equations), one can use the causal system to solve problems such as counterfactual regression, without resorting to regression techniques. Chapter 3.4. "Counterfactual analysis in structural model" in Pearl (2009) [A], makes this point explicitly. There, a counterfactual query is resolved using the complete causal system (being the causal graph and structural equations).

Of course, most practical scenarios do not have access to the full causal system. Performing counterfactual regression in these settings then relied on regression strategies, which are only approximate. Using regression techniques to perform counterfactual regression then requires at least some proof of identification. The burden of proof is only necessary because we cannot access the complete causal system.

In our setting, we actually learn the complete causal system (including the structural equations). Papers proposing ODE learners claim to recover the true underlying ODE. If there is such one true underlying equation (cfr. Assumption 3.1), then it is only logical that this equation corresponds to the one true underlying (causal) structural equation. Hence, by learning the true ODE, one has access to the full causal system, which in turn alleviates the need for an identification proof.

Of course, the accuracy and identification guarantees in our setting (above) are only as strong as those provided by the underlying ODE learner. In case one uses a faulty ODE learner (with, for example, poor guarantees), INSITE would suffer as a result. In this case, our framework is at least explicit about which assumptions and settings one can rely on ODE discovery for TE inference.

Actions taken: We include Appendix M, which explains which identification guarantees (and under which necessary assumptions) SINDy (our example underlying ODE learner) provides. Furthermore, Appendix M also includes alternative ODE learners, in case SINDy is unsuitable for the practical setting at hand.

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

We thank the reviewer once more for their insightful and positive review. We believe our paper is improved as a direct result of your comments. Should there be any more questions, please ask.

References

[A] Judea Pearl. Causal inference in statistics: An overview. 2009.

Thanks for the discussion. They addressed most of my concerns.

On the public comment regarding Symbolic Regression, I think it will be beneficial to discuss the position of this work in SR literature. I am familiar with TE literature and unfamiliar with SR/ODE, and I would assume many readers of this paper would also be the same.

I increased my score based on the rebuttal. But I also want to emphasize that my score comes with a low confidence as I am unfamiliar with ODE and SR literature.

We thank the reviewer for their time and helpful comments. Our rebuttal below follows the structure of your review. To facilitate further discussion, we have also stated how each specific comment changed our paper

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.

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.

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.

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.

Experiments. The underlying method SINDy implemented supports irregularly sampled time series data [B, C], therefore INSITE also supports learning ODEs from irregularly sampled time series. For the function introduced in Eq. (3) and its common assumption as the identity function, please refer to related works [A] for a detailed background.

Actions taken: Included a new Appendix M to detail SINDy and other ODE discovery learners. We will also provide, before the end of the discussion period (due to high running times), a set of experiments on irregular sampling and alternative functions $g(x)$.

Questions

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.

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.

We thank the reviewer once more for their insightful and positive review. We believe our paper is improved as a direct result of your comments. Should there be any more questions, please ask.

References

[A] Edward L Ince. Ordinary differential equations. Courier Corporation, 1956.

[B] Brunton, Steven L., Joshua L. Proctor, and J. Nathan Kutz. "Discovering governing equations from data by sparse identification of nonlinear dynamical systems." Proceedings of the national academy of sciences 113.15 (2016): 3932-3937.

[C] Goyal, Pawan, and Peter Benner. "Discovery of nonlinear dynamical systems using a Runge–Kutta inspired dictionary-based sparse regression approach." Proceedings of the Royal Society A 478.2262 (2022): 20210883.

Dear Shurui, Thank you very much for your interest in our paper! We truly appreciate this, as well as the comments you made. Hence, we will, as we have for our reviewers, answer in the structure of your comment below.

1. Symbolic regression. We agree with your statement:

ODE discovery, is fundamentally a part of Symbolic Regression.

Indeed, A large part of the ODE discovery literature employs techniques such as symbolic regression. However, symbolic regression also encompasses equations beyond ODEs. Since we do not claim to also connect other types of equation systems to TEs, we did not want to overclaim: "a connection between TEs and symbolic regression". More work is required to make such a claim.

However, we do agree that further context is required. As such, we will use the references you have provided to correctly position ODE discovery within the larger symbolic regression community.

2. Motivation. Our TE community's focus of the last few years has mostly been: to mitigate bias in favor of increased accuracy in counterfactual regression. This is a useful cause, of course, but it is unlikely that practitioners will adopt the state-of-the-art in TE inference (implying a complete retraining of their systems), for a possible tiny increase in accuracy.

Hence, with ODEs, one gets additional benefits beyond accuracy. Indeed, by mapping ODE discovery methods (such as SINDy, or perhaps alternative symbolic regressors), we get accuracy while also getting interpretability. The latter is incredibly important as most application areas for TE inference would benefit hugely from a fully interpretable model.

Beyond the above, we hope our paper also engages researchers (and practitioners) from the ODE discovery field (or perhaps, as you suggest, the wider symbolic regression community). In this way, both TEs inference and ODE discovery could build upon ideas from both sides, pushing research even further.

Thank you once more for commenting on our paper! We believe your comments make our manuscript better. Hopefully, more people will find our paper as thought provoking as you have.