Training Physics‐Based Machine‐Learning Parameterizations With Gradient‐Free Ensemble Kalman Methods

Abstract

Most machine learning applications in Earth system modeling currently rely on gradientbased supervised learning. This imposes stringent constraints on the nature of the data used for training (typically, residual time tendencies are needed), and it complicates learning about the interactions between machine-learned parameterizations and other components of an Earth system model. Approaching learning about process-based parameterizations as an inverse problem resolves many of these issues, since it allows parameterizations to be trained with partial observations or statistics that directly relate to quantities of interest in long-term climate projections. Here we demonstrate the effectiveness of Kalman inversion methods in treating learning about parameterizations as an inverse problem. We consider two different algorithms: unscented and ensemble Kalman inversion. Both methods involve highly parallelizable forward model evaluations, converge exponentially fast, and do not require gradient computations. In addition, unscented Kalman inversion provides a measure of parameter uncertainty. We illustrate how training parameterizations can be posed as a regularized inverse problem and solved by ensemble Kalman methods through the calibration of an eddy-diffusivity mass-flux scheme for subgrid-scale turbulence and convection, using data generated by large-eddy simulations. We find the algorithms amenable to batching strategies, robust to noise and model failures, and efficient in the calibration of hybrid parameterizations that can include empirical closures and neural networks.

Updates on Model Hierarchies for Understanding and Simulating the Climate System: A Focus on Data‐Informed Methods and Climate Change Impacts

“…High‐resolution models can provide training data for dynamics‐driven SGS processes in the atmosphere and oceans, such as gravity waves, turbulence, convection, and cloud cover. These processes lie in the “gray zone,” with scales O (1–100 km) which are under‐resolved in typical climate models but are largely resolved in computationally intensive sub‐kilometer scale models (e.g., atmospheric subgrid momentum fluxes: Yuval & O’Gorman, 2020; Yuval et al., 2021; Wang et al., 2022; ocean momentum forcing: Guillaumin & Zanna, 2021; Perezhogin et al., 2023; convection: Brenowitz & Bretherton, 2019; Gentine et al., 2018; clouds: Rasp et al., 2018; gravity waves: Sun et al., 2023) or large eddy simulations (e.g., eddy‐diffusivity momentum flux: Lopez‐Gomez et al., 2022; Shen et al., 2022). However, other coupled processes such as atmospheric chemistry, sea ice cover, and vegetation dynamics are not modeled explicitly at any resolution and may make use of observational data sets (e.g., ozone: Nowack et al., 2018; sea‐ice: Andersson et al., 2021; vegetation: Chen et al., 2021).…”

EnsembleKalmanProcesses.jl: Derivative-free ensemble-based model calibration

“…It does so by defining transformation maps under-the-hood from the constrained space to an unconstrained space where the optimization problem can be suitably defined. Constrained optimization using this framework has been successfully demonstrated in a variety of settings (Dunbar et al, 2022;Lopez-Gomez et al, 2022;Schneider, Dunbar, et al, 2022).…”

“…• EnsembleKalmanProcesses.jl has been used to train physics-based and machine-learning models of atmospheric turbulence and convection, implemented using Flux.jl and TurbulenceConvection.jl (Lopez-Gomez et al, 2022). In this application, the available model outputs are not differentiable with respect to the learnable parameters, so gradientbased optimization was not an option.…”

Gradient Flows for Sampling: Mean-Field Models, Gaussian Approximations and Affine Invariance

“…This problem is fundamental in many applications arising in computational science and engineering and is widely studied in the applied mathematics, machine learning and statistics communities. A particular application is Bayesian inference for large-scale inverse problems; such problems are ubiquitous, arising in applications from climate science [65,118,63,91], through numerous problems in engineering [136,38,21] to machine learning [111,100,28,32]. These applications have fueled the need for efficient and scalable algorithms which employ noisy data to learn about unknown parameters θ appearing in models and perform uncertainty quantification for predictions then made by those models.…”