Searching for new physics: Using explainable AI to understand deep learned parameterizations of turbulent heat fluxes

Machine learning (ML), and particularly deep learning (DL), for geophysical research has shown dramatic successes in recent years. However, these models are primarily geared towards better predictive capabilities, and are generally treated as black box models, limiting researchers’ ability to interpret and understand how these predictions are made. As these models are incorporated into larger models and pushed to be used in more areas it will be important to build methods that allow us to reason about how these models operate. This will have implications for scientific discovery that will ensure that these models are robust and reliable for their respective applications. Recent work in explainable artificial intelligence (XAI) has been used to interpret and explain the behavior of machine learned models.

Here, we apply new tools from the field of XAI to provide physical interpretations of a system that couples a deep-learning based parameterization for turbulent heat fluxes to a process based hydrologic model. To develop this coupling we have trained a neural network to predict turbulent heat fluxes using FluxNet data from a large number of hydroclimatically diverse sites. This neural network is coupled to the SUMMA hydrologic model, taking imodel derived states as additional inputs to improve predictions. We have shown that this coupled system provides highly accurate simulations of turbulent heat fluxes at 30 minute timesteps, accurately predicts the long-term observed water balance, and reproduces other signatures such as the phase lag with shortwave radiation. Because of these features, it seems this coupled system is learning physically accurate relationships between inputs and outputs. 

We probe the relative importance of which input features are used to make predictions during wet and dry conditions to better understand what the neural network has learned. Further, we conduct controlled experiments to understand how the neural networks are able to learn to regionalize between different hydroclimates. By understanding how these neural networks make their predictions as well as how they learn to make predictions we can gain scientific insights and use them to further improve our models of the Earth system.