Bayesian Neural Controlled Differential Equations for Treatment Effect Estimation

TL;DR

This work tackles continuous-time treatment effect estimation with uncertainty by introducing Bayesian Neural Controlled Differential Equations (BNCDE), a neural architecture that combines neural controlled differential equations with latent neural stochastic differential equations to model posterior weight dynamics and yield posterior predictive distributions for outcomes under future treatment sequences. By jointly training an encoder that processes irregular histories and a decoder that conditions on future treatments, BNCDE captures both model (epistemic) and outcome (aleatoric) uncertainty in a principled Bayesian framework. Empirical results on a non-regular tumor-growth simulator show that BNCDE provides faithful and sharp credible intervals, more informative model uncertainty, and robustness to noise and longer prediction horizons, outperforming the continuous-time TE-CDE baseline. The method holds promise for uncertainty-aware, personalized medical decision-making in settings with irregular sampling and time-varying treatments, albeit with caveats around explainability and computational demands intrinsic to Bayesian neural differential equations.

Abstract

Treatment effect estimation in continuous time is crucial for personalized medicine. However, existing methods for this task are limited to point estimates of the potential outcomes, whereas uncertainty estimates have been ignored. Needless to say, uncertainty quantification is crucial for reliable decision-making in medical applications. To fill this gap, we propose a novel Bayesian neural controlled differential equation (BNCDE) for treatment effect estimation in continuous time. In our BNCDE, the time dimension is modeled through a coupled system of neural controlled differential equations and neural stochastic differential equations, where the neural stochastic differential equations allow for tractable variational Bayesian inference. Thereby, for an assigned sequence of treatments, our BNCDE provides meaningful posterior predictive distributions of the potential outcomes. To the best of our knowledge, ours is the first tailored neural method to provide uncertainty estimates of treatment effects in continuous time. As such, our method is of direct practical value for promoting reliable decision-making in medicine.

Figures (12)

• Figure 1: Our BNCDE consists of an encoder, a decoder, and a prediction head.
• Figure 2: Faithfulness: Empirical coverage across different CrI quantiles. Shown are different prediction windows $\Delta = 1,2,3$. Areas in green (red) indicate that the CrIs are faithful (not faithful).
• Figure 3: Sharpness: Width of the CrIs (median) for different quantiles $\alpha$.
• Figure 4: Error in point estimates: Reported is the median over the mean squared errors (MSE) of the point estimates of the outcomes. The results are based on test data that is generated with varying levels of noise, i.e., $\text{Var}(\epsilon_t)$.
• Figure 5: To compare model uncertainty, the normalized MSE of estimated treatment effects between the treatment arms versus deferral rate is shown.