Twinned Interventional Flows

Thank you for your feedback. We address your questions and concerns below.

1. Coherence and motivation behind the combination of ideas.

   We appreciate your feedback and understand the importance of clarity in our writing. Our objective is to address a combination of real-world problems, each of which is carefully chosen based on its relevance to applications such as medical decision-making. The selection of ideas is motivated by the need to integrate knowledge from different fields to tackle complex systems effectively.

   The issue of stiffness, in particular, arose during the development of our method. Despite searching existing literature, we found no satisfactory solutions. Consequently, we devised our own approach, which demonstrated significant improvements in our experiments. Please see the visualization in Section 4.1, if it's sufficient to explain the problem.

   We acknowledge the importance of conveying these motivations more explicitly, and we will make adjustments to the text to provide a clearer explanation of why these specific problems are included in the paper. If you have additional suggestions or specific areas you would like us to address, please let us know.

2. Regarding stiffness in ODEs.

   There is no clear measure of stiffness in the literature and our definition, "Stiff ODEs are numerically difficult to invert," is derived from Ernst Hairer’s observation that "stiff equations are problems for which explicit methods do not work" [1], combined with insights from [2] that errors can significantly magnify when ODEs are inverted.

   To address your concern, we will refine our wording and explicitly reference [1] to provide a clearer foundation for our definition. If you have any further suggestions or specific areas you would like us to enhance, please let us know.

   Regarding the visualization of stiffness, it is presented in Figure 4f compared to the performance when the penalty term is employed (4e).

   How would you like us to further visualize these problems?

3. Focus of related work and background section.

   In our paper, we approach causality from a specific angle, choosing not to delve deeply into causal discovery. We mention it mainly due to its relevance to the intersection of reinforcement learning (RL) and causal fields. Instead, our primary focus lies in recognizing the causal properties of the system, particularly addressing challenges like unobserved confounders. We achieve this by incorporating these considerations into the model design, ensuring flexibility to tackle such complexities.

   To accomplish our goals, we amalgamate tools from both the causal inference and RL domains. As a result, our background section delves into these topics, expecting a foundational understanding from the readers. We acknowledge the absence of exhaustive technical background knowledge in the paper, given its assumption of familiarity with these fields.

   Nevertheless, we are open to adjusting our focus if you deem any particular aspect irrelevant or have specific topics in mind that should receive more attention. Your insights are valuable in shaping the paper to meet the expectations of our audience.

4. The balance between appendixes and main body.

   We share this concern with you, however the strict page limit forces us to make some compromises.

   If there are specific sections in the appendix that you believe could be integrated into the main body or should be discarded, please let us know. We are open to reorganizing the content to enhance the overall flow while acknowledging to the page limit.

5. The setting summarized in a sentence. Our primary objective is to develop a predictive model that accurately captures the system dynamics under various interventions and possible (unobserved) confounders. This model should be versatile enough to estimate the likelihood of events and make use of both observational and interventional data. We believe this framework is applicable to various real-world scenarios, particularly in fields like medicine.

6. Wording and additions. Thank you for your detailed suggestions on the wording, (e.g. wording in paragraph 3, and adding the term causal sufficiency). We appreciate your feedback and will incorporate the necessary changes in the upcoming version, which we aim to upload shortly. Your understanding of the delay is much appreciated. We look forward to continuing the discussion.

7. Definition of stiff regions.

   We will add this definition to the text: By stiff regions we mean the regions of weight space where the corresponding NODE becomes stiff.

The response will be continued in the following message.

Many thanks for your thoughtful feedback. We address all your comments and suggestions below:

1. Causal reinforcement learning: connection between RL and treatment effect estimation. Thank you for your comment. The connection between causal inference and RL fields is well studied by Bareinboim [1]. In our setting, the reinforcement learning set up (POMDP) is closely tied with causal treatment effect prediction; and our work can be seen as addressing two of the six main tasks in causal reinforcement learning, namely, (a) Generalized policy learning and (b) Counterfactual decision-making. Additionally, deciding when and where to intervene could be addressed with a good predictor that is continuous in nature.

   Specifically, each action $a_t$ serves the role of a treatment, and time-varying confounders can be included in $l_t$. Effectively predicting treatment effects implies that the problem can be cast as an RL task, where the objective is to optimize the treatment based on the predictions. This perspective highlights the interplay between causal treatment effect prediction and RL, where a good predictor with continuous properties can inform decision-making in RL. Our POMDP formulation serves this purpose.

   Hopefully this clarified the connections between these topics and we will make the according changes to the paper to highlight these topics.

2. Identifiability and our assumptions. We understand your concerns about identifiability. Indeed, identifying the number of confounders or recovering a full causal graph in general is tedious and often infeasible. We also acknowledge that there are scenarios where data limitations may hinder learning of the true causal structure due to identifiability issues.

   Our primary goal here is to develop a predictive model that accurately captures the system dynamics under various interventions and possible (unobserved) confounders in our setting that treats the size of the latent space as a tunable parameter, and abstracts many practical situations with confounders. We therefore appeal to an approach that does not constrain the causal hierarchy within observed and unobserved variables and can be described in terms of a continuous POMDP.

   Importantly, if desired, identifiability can be achieved with appropriate additional assumptions in our setting. Indeed, one way is identification via (weighted) time-independent sampling by adapting Tennenholtz et al (2020) [2]. This, however, is susceptible to the curse of dimensionality. Instead, one can invoke a (more amenable) reduction to proximal causal inference as illuminated by Bennett and Kallus [3].

   The set of assumptions provided in the paper are also consequence of this approach. The Overlap assumption by Seedat et al. [4] is listed in our set of assumption in Appendix C, as it necessary for deriving our theoretical results. However the two other assumptions by [4] are more related to the identifiability of the problem, which is not our focus here, and so were not listed. However we are happy to add these assumptions with a discussion on identifiability based on your feedback.

3. Regarding the experimental evidence. Our primary focus in the experimental section was to showcase the various capabilities of the proposed methood TIF, but not to imply or assert its superiority over low-dimensional embeddings like TE-CDE [4]. We emphasize the flexibility afforded by TIF due to its many strengths inherited from continuous normalizing flows, particularly in density estimation. In this context, it is crucial to note that while density estimation is primarily a theoretical framework, the stability of the Monte Carlo estimation algorithm, as outlined in Appendix D, requires further validation. Consequently, we refrained from making comparisons with, for example, Bayesian methods.

   We are grateful to you for bringing the two new papers [5], [6] to our attention, which are important and highly relevant to our topic. Unfortunately these works were not published when we worked with the background study of TIF. We are happy the see that the field is evolving and we will cite and position these papers appropriately in the final version of the paper. However, due to time constraints, benchmarking against them is unfortunately not feasible during the response period.

   We believe that this work already makes significant contributions, and lays the foundation to investigate aspects such as those pointed by you as future work. Based on your feedback, we will also be sure to underscore the motivation underlying each of our experiments.

The response will be continued in the following message.

Dear reviewer, we are grateful for your constructive feedback that helped us elucidate some important aspects of this work. We hope our response has sufficiently addressed your concerns, and hope it translates into an increased score. We are also happy to discuss further if you have any further questions before the discussion window ends tomorrow. Many thanks!

1. Thank you - we'll structure the related section better based on your comments. Our main contribution is to use flow models in an augmented space that couples observations and latent variables including unknown confounders, to enable exact (not just lower bounds on the) log-likelihood estimates as well as predict treatment effects. We're not aware of any previous work in literature that can handle both these tasks simultaneously in the presence of unknown confounders. Since we do not observe the unknown variables only part of the information is visible and the state evolves continuously over time, so the problem can be investigated in a POMDP setting. We would like to emphasize that our goal is not to cherrypick any RL tasks, but to motivate this work as a generic framework that can also be leveraged in RL settings. In our RL task, we already compare with Latent-ODE which is a strong normalizing flow-based baseline.
2. Identifiability is certainly important, however, identifiability results in our POMDP setting can be derived using e.g., proximal causal inference with additional assumptions. Since such theoretical results on identifiability are already known, and these additional assumptions do not play any role otherwise in the entire setup (e.g., for the purpose of likelihood estimation), we did not elaborate on these to keep the exposition focused. We will make this clear based on your comments.
3. Thank you for pointing these papers to us. We will appropriately position the contributions of these works. Please note that neither of these papers allow for exact log-likelihood estimation in the presence of unknown confounders
4. The fine tuning of the learning rate for the tumor effect prediction was done with values 0.01- 0.0001 for TIF and the values reported in literature were used for TE-CDE. Based on your comment, we have now also shared the code of this experiment for reproducibility.
5. To our understanding the assumption 3 by Seedat et.al. means that the current history is sufficient for making accurate counterfactual predictions without incorporating the information about the potential factual outcomes. This does not mean that there is no confounding effect in the system, but rather that the history observations are sufficient to accurately model them. But we will re visit this and make corrections if needed.

9. Thank you for elaborating this. We see it as: there is unobserved confounding in causal diagram (see $l$ in in Figure 1 and observe that red arrow is present only in the observational data) however the markovian assumption in POMDP is satisfied in respect to the full state $s=[l,o]$, thus if we obtain accurate estimate of $l$ we can make accurate predictions based on the current state. The parental markovian condition follows from the structure of of the DAG, that is, each variable is independent on all its nondescendants, given its parents (Pearl & Verma 1995). However, we will revisit our wording and make sure that there is no ambiguity and the terms are conveyed correctly.

11. Thank you. Finding simpler components that could be alternatives to concatsquash block is certainly an interesting direction, but requires further research. In general, beyond Occam's razor (as often implemented e.g., by instantiating the minimum description length principle - which loosely speaking amounts to choosing a simple hypothesis from a set of candidate hypotheses that all fit the training data well), we know from learning theory foundations that there is an inherent tension between expressivity and generalization when considering different families of architectures (or class of hypotheses) and the sweet-spot lies where the two are roughly balanced. Unfortunately, the generative ML community is still trying to work out the generalization bounds for flow based models, so a detailed principled analysis of the impact of different components is very much an open question that this work inherits. Please do note that several works have shown the empirical merits of using concatsquash block in similar settings, i.e., for embedding in dynamic systems, so we adopted this choice to be consistent with the prevailing paradigm.

Based on your comments, we'll be happy to include a discussion on this in the Supplementary section as well. Thank you again for your thoughtful comments, and we hope our response here satisfactorily addresses your concerns and translates into an improved score.

Many thanks for your constructive feedback. We address all your concerns and questions below:

1. Improving presentation and clarity.

   Thank you for your inputs. We appreciate your insights, and are currently working on enhancing readability, including, by incorporating more intuitive reasoning into the methods section and elucidating the experimental section as suggested. The revised version will be uploaded shortly.

2. Time complexity and scaling.

   Thank you for the opportunity to shed light on the computational aspects of the proposed method. We would like to emphasize that TIF does not necessitate full matrix inversion (which would indeed have been computationally expensive with a complexity of $O(D^3)$, where $D$ denotes the data dimension), instead inheriting the more tractable computation needs of continuous normalizing flows that can be solved more efficiently by integrating backward in time [1].

   Furthermore TIF inherits from the continuous flows normalizing a more efficient way of estimating the log likelihood.

   If the cost of evaluating TIF is $O(DH)$ where $D$ is the dimensionality of the data, and $H$ is the largest hidden dimension in the concatsquach block, then the time complexity of one transformation using a naive normalizing flow approach would be $O(DH + D^3)$

   The cubic format arises from using the change of variables to obtain complex densities $z \sim p_z(z)$ from a simple distribution $u \sim p_u(u)$ with transformation $T(u) =x$, as described in [1,2]

   $log(p_z(z))=log(p_u(u)) - log(|det J_T(u)|) ~,$

   where $J_T$ is the corresponding Jacobian matrix ($n \times n$), and the time complexity of computing the log determinant is $O(n^3)$.

   With continuous normalizing flows the complexity can be reduces to $O(DH + D^2)$, as the change in log probabilities can be written as [3, 2]

   $log (p(z_{t_1}))= log (p(z_{t_0})) - \int_{t_0}^{t_1} Tr \{J_{T}(z_t) \} dt~,$

   Where the time complexity of obtaining the exact Trace is $O(n^2)$ [1].

   Instead of exact Trace computation, we further reduce the time complexity to $O(DH + D)$ using Hutchinson’s trace estimator to obtain an estimate of the trace [1,2,3].

   $Tr \{ J_{T}(z_t) \} = E_v [ v^T J_{T}(z_t) v ]$

   where $v$ is a random vector with zero mean and unit covariance. We draw binary values for $v$ from a Rademacher Distribution. Reducing the time complexity to $O(n)$ [2].

   However the use of numerical ODE solvers introduces an additional cost as the transformation needs to be performed $L$ times, leading to an overall time complexity of $O(L(DH + D))$. The impact of $L$ is significant in TIF, especially as the discontinuities of the derivative at the intervention points increase the number of steps taken by the numeric solver. Additionally, for the log likelihood sample, the model needs to be run forwards and backwards in time, i.e., $O(2L(DH + D))$. Finally, to obtain the Monte Carlo estimate (Appendix D), several samples are needed, which can, however, be computed in parallel.

Thank you once again for your valuable comments. We will incorporate the discussed topics into the updated version, scheduled for upload by Monday at the latest. We appreciate your engagement and are open to continuing this discussion with you.

References

[1] Papamakarios, G., Nalisnick, E. T., Rezende, D. J., Mohamed, S. & Laksh minarayanan, B. (2021), ‘Normalizing flows for probabilistic modeling and inference.’, Journal of Machine Learning Research 22(57), 1–64. URL: https: // jmlr. org/ papers/ v22/ 19-1028. html

[2] Will Grathwohl, Ricky T. Q. Chen, Jesse Bettencourt, Ilya Sutskever, and David Duvenaud. Ffjord: Free-form continuous dynamics for scalable reversible generative models. In International Conference on Learning Representations, ICLR 2019, 6-9 May, New Orleans, USA, Conference Track Proceedings, 2019. URL https://openreview.net/forum?id=rJxgknCcK7.

[3] Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. In Advances in Neural Information Processing Systems, NeurIPS 2018, 2-8 December, Montr´eal, Canada, volume 31. Curran Associates, Inc., 2018. URL https://proceedings.neurips.cc/paper_files/paper/2018/ file/69386f6bb1dfed68692a24c8686939b9-Paper.pdf