Ultra-fast machine learning solver for chemical kinetics in atmospheric models

Modeling of atmospheric chemistry is critical to major environmental problems including air pollution, ozone hole, and climate change. The evolution of hundreds of chemical components is governed by a large and stiff system of ordinary differential equations (ODEs) describing complex kinetic mechanisms. The computational cost of solving such a system has been a major barrier in the development of atmospheric chemistry models and has hindered the inclusion of atmospheric chemistry in Earth system models.

Here we use neural networks (NNs) as a fast alternative to traditional ODE solvers, speeding up the calculation by orders of magnitude with comparable accuracy. NNs can efficiently approximate the low-dimensional manifolds present in the chemical space and thus remedy the curse of dimensionality arising from the large number of chemical components. The training data generated by the traditional chemical solver are noise-free and smoothly-distributed, which allows the use of second-order quasi-Newton optimization algorithms for much faster convergence without the worry of over-fitting. Changes in the chemical mechanism requires re-training a new NN, but the training time can be brought down to several minutes using the PyTorch library with GPU acceleration on the Amazon cloud platform. We present first applications of our NN solver in the GEOS-Chem chemical transport model and the NASA GEOS-5 Earth system model.

The lack of error control on individual predictions is a prominent problem with machine-learning-based solvers. We discuss potential approaches to detect "bad predictions" and ensure accuracy, such as ensemble methods, confidence interval estimate, and conformal prediction techniques.