🤖 AI Summary
This work addresses the challenge of accurately quantifying uncertainty in the conditional distribution for online state inference of dynamical systems under noisy observations. The authors propose an ensemble data assimilation method trained on synthetic trajectories, leveraging a permutation-invariant Transformer architecture to learn the Bayesian filtering analysis map. For the first time, strictly proper scoring rules—specifically the energy score—are employed to guide training, ensuring global optimality of the learned probability distributions under realizability assumptions. Both theoretical analysis and empirical evaluations demonstrate that the proposed approach significantly outperforms conventional methods and mean squared error–based learning strategies in complex scenarios involving nonlinearity, non-Gaussianity, and multimodal posteriors: it effectively corrects the ensemble Kalman filter (EnKF) in near-Gaussian settings and achieves superior performance through end-to-end learning in strongly non-Gaussian regimes.
📝 Abstract
Bayesian filtering of partially and noisily observed dynamical systems seeks to infer the evolving conditional distribution of the state of a dynamical system, given observations, in an online fashion. This Bayesian filtering distribution is the natural object for uncertainty quantification, but it is rarely available as a supervised learning target. However, one can often use the forecast model to generate synthetic system trajectories, along with synthetic observations. We introduce the proper scoring ensemble filter (PSEF), an ensemble data assimilation method based on training an analysis map to approximate the filtering distribution using only synthetic state--observation trajectories. The analysis step is represented as a permutation-invariant, transformer-based map that takes as input a forecast ensemble and observations, producing an analysis ensemble. Training is based on strictly proper scoring rules -- with the energy score used in our implementation -- so that probabilistic accuracy is rewarded over the whole probability distribution. We prove that, under a realizability assumption, the population objective is minimized by the true Bayesian filtering distribution. We also derive the finite-ensemble empirical objective used in training and relate its single state--observation trajectory form to the population objective, using a mean-field consistency argument. Numerical experiments show that the learned filter accurately approximates challenging filtering distributions, including nonlinear, non-Gaussian, and multi-modal posteriors, and achieves stronger performance in data assimilation tasks than classical methods or learning-based methods with mean-squared-error objectives. For close-to-Gaussian problems, learning a correction to the EnKF is the best approach, while for highly non-Gaussian problems an end-to-end approach that discards this inductive bias is superior.