Uncertainty quantified discovery of chemical reaction systems via Bayesian scientific machine learning
Emily Nieves (1, 2), Raj Dandekar (1), Chris Rackauckas (1)
(1) Department of Computational Science and Engineering, Massachusetts Institute of Technology; (2) Department of Biological Engineering, Massachusetts Institute of Technology
Objectives: Mechanistic models in quantitative systems pharmacology (QSP) offer more interpretability than purely data-driven machine learning approaches, but practitioners often lack complete knowledge of the underlying systems. Methods of scientific machine learning (SciML) account for this epistemic uncertainty by mixing neural network techniques with mechanistic modeling forms to allow for the automated discovery of missing components in mechanistic models. In this work we build upon recent results in SciML to demonstrate how incorporating prior structural knowledge into deep learning can improve interpretability and prediction accuracy in QSP modeling. We build upon the previously proposed chemical reaction neural network (CRNN) [1] which is a structured neural network which directly encodes the mathematical form of mass-action kinetics into the neural architecture. In this work, we extend the CRNN to represent the uncertainty in the learned network using a Bayesian inference framework.
Methods: We combined the CRNN with a preconditioned Stochastic Gradient Langevin Descent (pSGLD) Markov Chain Monte Carlo (MCMC) stepper [2] of the Bayesian neural ODEs [3] with derivatives computed using stiffly-stabilized adjoint sensitivity analysis [4] via Julia’s SciML ecosystem to more efficiently and accurately obtain an estimate of the posterior probability of the parameters of the learned reaction network. We tested our algorithm with a toy reaction system containing four reactions and then extended it to a simplified version of the EGFR-STAT3 pathway.
Results: Our Bayesian CRNN was able to directly learn the complete reaction network purely from data and provide the trained result in a mechanistic interpretable form for both the simple case and the EGFR-STAT3 pathway. We demonstrate that our pSGLD MCMC is required for increasing the robustness of the Bayesian training process, and showcase the noise robustness of the Bayesian CRNN with this improved stepping method. When compared to training with a standard SGLD optimizer, our preconditioned form was more efficient and entered the sampling phase ~6,500 epochs before the SGLD. It also better represented the uncertainty in the learned reaction rates by achieving an average percent deviation from the true reaction rates of 18.5%± 21.3 compared to 117.8% ± 219.7 for the SGLD.
We also compared our pSGLD-trained CRNN to a purely data-driven neural ODE system that did not include prior knowledge about the structure of chemical reactions and found that in addition to being interpretable to mechanistic predictions, our CRNN also obtains a higher extrapolation accuracy when simulated for a time span 25% longer than was used for training (mean absolute percent error 9.0%±3.2 for CRNN, 254.7%± 109.8 for neural ODE).
Conclusions: The combination of a preconditioned SGLD optimizer with the CRNN allows for efficient training and posterior sampling as well as reliable estimates of parametric uncertainty while accounting for epistemic uncertainty in a way that builds confidence in the learned reaction network and can help determine if more data is needed. This work demonstrates that knowledge-embedded machine learning techniques via SciML approaches may greatly outperform purely deep learning methods in a small-medium data regime that is common in Quantitative Systems Pharmacology (QSP) and demonstrates viable techniques for the automated discovery of QSP models directly from timeseries data.
References:
[1] W. Ji and S. Deng, Autonomous discovery of unknown reaction pathways from data by chemical reaction neural network, 2021
[2] C. Li, et al., Preconditioned stochastic gradient langevin dynamics for deep neural networks, 2015
[3] R. Dandekar, et al., Bayesian neural ordinary differential equations, 2020
[4] S. Kim, et al., Stiff neural ordinary differential equations, 2021