Twinned Interventional Flows

Linda Hemmann, Vikas Garg

Primary Area: generative models

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Keywords: counterfactual predictions, normalizing flows, causal reinforcement learning, causal effects, neural differential equations, partially observed MDPs

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2024/AuthorGuide.

Abstract:

Real-world problems in continuously evolving settings, such as predicting the efficacy of medical treatment, often require estimating the causal effects of interventions. Issues such as irregularly-sampled and missing data, unobserved factors, and ethical concerns make such settings especially challenging. The existing methodology relies on low-dimensional embeddings, potentially incurring information loss.

We circumvent this limitation with a novel approach ``twinning" that augments the partial observations with additional latent variables and appeals to conditional continuous normalizing flows to model the system dynamics, obtaining accurate density estimates. We also introduce a new approach to overcome a key technical challenge, namely, mitigating stiffness of the underlying neural ODE. The model provably benefits from auxiliary non-interventional data during training. We showcase the flexibility of the proposed method with tasks like anomaly detection and counterfactual prediction, and benchmark on standard reinforcement learning (Half-Cheetah) and treatment effect prediction (tumor growth) contexts.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors' identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 7202