The Ensemble Kalman Inversion Race

TL;DR

This study tackles calibration of climate-model parameters from noisy time-averaged statistics using derivative-free ensemble Kalman inversion methods. By racing variants of TEKI, ETKI, UKI, and IEKF against a derivative-based Levenberg-Marquardt baseline on Lorenz63 and Lorenz96 models (including grid and neural-network parameterizations), it quantifies computational efficiency via forward-model evaluations to a target RMSE. The results show no single winner across all problems; UKI excels in low-dimensional settings with informative priors, IEKF performs well with informative priors, and TEKI/ETKI offer robust performance even with uninformative priors, while LM consistently fails due to noisy loss landscapes. These findings yield practical guidelines for choosing ensemble variants in GCM calibration and highlight the enduring value of derivative-free methods for optimizing models driven by statistics of chaotic systems.

Abstract

Ensemble Kalman methods were initially developed to solve nonlinear data assimilation problems in oceanography, but are now popular in applications far beyond their original use cases. Of particular interest is climate model calibration. As hybrid physics and machine-learning models evolve, the number of parameters and complexity of parameterizations in climate models will continue to grow. Thus, robust calibration of these parameters plays an increasingly important role. We focus on learning climate model parameters from minimizing the misfit between modeled and observed climate statistics in an idealized setting. Ensemble Kalman methods are a natural choice for this problem because they are derivative-free, scalable to high dimensions, and robust to noise caused by statistical observations. Given the many variants of ensemble methods proposed, an important question is: Which ensemble Kalman method should be used for climate model calibration? To answer this question, we perform systematic numerical experiments to explore the relative computational efficiencies of several ensemble Kalman methods. The numerical experiments involve statistical observations of Lorenz-type models of increasing complexity, frequently used to represent simplified atmospheric systems, and some feature neural network parameterizations. For each test problem, several ensemble Kalman methods and a derivative-based method "race" to reach a specified accuracy, and we measure the computational cost required to achieve the desired accuracy. We investigate how prior information and the parameter or data dimensions play a role in choosing the ensemble method variant. The derivative-based method consistently fails to complete the race because it does not adaptively handle the noisy loss landscape.

Figures (9)

• Figure 1: Flow chart of a forward model $\mathcal{F}(\cdot)$ whose outputs are statistics of the system state. For a given set of parameters $u$, the forward model first simulates the dynamics for a time period $\tau$, starting from a random initial condition. The statistics $f=\mathcal{F}(u)$ are subsequently computed from the simulation output after a spin-up period. (Image of dynamics and subgrid processes simulation is from ST17.)
• Figure 2: Random samples taken from the prior distributions constructed for the L'96 model with a neural network parameterization. The purple lines are the functions generated by neural networks whose weights and biases are draws from a Gaussian prior. The black dashed lines are the prior means. (a) Samples from the "uninformative" prior. (b) Samples from the "informative" prior.
• Figure 3: Average number of forward model runs required to reach a target RMSE of one as a function of ensemble size for the L'63 test problem. Shown are results for TEKI (blue), ETKI (yellow) and IEKF (green). For UKI, the ensemble size is fixed (Section \ref{['sec:UKI']})
• Figure 4: Summary of results for estimating two model parameters from statistics of L'63 dynamics. (a) Average number of forward model runs required to reach target accuracies of $\text{RMSE}=1$, $\text{RMSE}=1.1$, or $\text{RMSE}=1.2$. (b) Average number of iterations required to reach the target accuracies. The error bars in panels (a) and (b) are derived from the 5$^{\text{th}}$ and 95$^{\text{th}}$ percentiles. (c) Optimal ensemble sizes for three target accuracies. (d) Parameter combinations in the $\rho-\beta$ plane. The light gray contours are a 2D histogram of an averaged posterior distribution (darker grays represent regions of higher posterior probability and lighter grays represent regions of lower posterior probability). The ensemble means of TEKI, ETKI, UKI and IEKF are shown as blue, yellow, pink, and green dots. The true parameter pair is a dark gray dot and the expected posterior mode is a dark gray cross. Slices of the loss function obtained by varying $\rho$ or $\beta$ independently and separately are shown at the top and right sides of panel (d).
• Figure 5: Summary of results for the model parameter estimation problem with L'96 dynamics under a constant forcing. (a) Average number of forward model runs required to reach target accuracies of $\text{RMSE}=1$, $\text{RMSE}=1.1$, or $\text{RMSE}=1.2$. (b) Average number of iterations required to reach the target accuracies. The error bars in panels (a) and (b) are derived from the 5$^{\text{th}}$ and 95$^{\text{th}}$ percentiles. Note the truncated $y$-axis in panel (a). (c) Optimal ensemble sizes for three target RMSEs. (d) Loss function and ensemble means of TEKI, ETKI, UKI and IEKF for one of our experiments. Also shown is the result of a derivative-based optimization via LM, which fails to converge to a useful solution (defined as $\text{RMSE} \leq 1.2$).