Causal Graph ODE: Continuous Treatment Effect Modeling in Multi-agent Dynamical Systems

Zijie Huang, Jeehyun Hwang, Junkai Zhang, Jinwoo Baik, Weitong Zhang, Dominik Wodarz, Yizhou Sun, Quanquan Gu, Wei Wang

Keywords: Dynamical System, NeuralODE, Causal Inference

Abstract:

Real-world multi-agent systems are often dynamic and continuous, where agents interact over time and undergo changes in their trajectories. For example, the COVID-19 transmission in the U.S. can be viewed as a multi-agent system, where states act as agents and daily population movements between them are interactions. Estimating the counterfactual outcomes in such systems enables accurate future predictions and effective decision-making, such as formulating COVID-19 policies.
However, existing methods fail to model the continuous dynamic effects of treatments on the outcome, especially when multiple treatments are applied simultaneously.
To tackle this challenge, we propose Causal Graph Ordinary Differential Equations (CAG-ODE), a novel model that captures the continuous interaction among agents using a Graph Neural Network (GNN) as the ODE function. The key innovation of our model is to learn time-dependent representations of treatments and incorporate them into the ODE function, enabling precise predictions of potential outcomes. To mitigate confounding bias, we further propose two domain adversarial learning-based objectives, which enable our model to learn balanced continuous representations that are not affected by treatments or interference. Experiments on two datasets demonstrate the superior performance of CAG-ODE.

Submission Number: 29