cs.LGJun 25, 2026

Learning Probabilistic Filters with Strictly Proper Scoring Rules

Authors: Eviatar BachRicardo BaptistaJochen BröckerBohan ChenAndrew Stuart

Organizations: Department of Meteorology, University of Reading Brian Hoskins Building, Reading, RG6 6ET, UK · Department of Mathematics and Statistics, University of Reading Pepper Lane, Reading, RG6 6AX, UK · National Centre for Earth Observation Brian Hoskins Building, University of Reading, Reading, RG6 6ET, UK · Department of Statistical Sciences, University of Toronto · Jun 700JochenUniversityBr\"ockerAve., Toronto, ON M5G 1Z5, Canada · Department of Meteorology, University of Reading 25 Brian Hoskins Building, Reading, RG6 6ET, UK · The Computing + Mathematical Sciences Department, California Institute of Technology 1200 E. California Blvd., Pasadena, CA 91125, USA[cs.LG] · The Computing + Mathematical Sciences Department, California Institute of Technology 1200 E. California Blvd., Pasadena, CA 91125, USA

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.

Explore similar work

Sep 23, 2026stat.ML

Improving Ensemble Filters with Flow Matching

Data assimilation estimates a dynamical state from partial and noisy observations. Classical ensemble filters are efficient but restrict analysis updates through finite sample covariance and affine Gaussian distribution. We introduce the Flow Ensemble Filter (FlowEF), which uses conditional flow matching to transport the forecast ensemble from a classical baseline filter to an analysis ensemble. FlowEF uses a localized Gaussian source during training, transports forecast ensemble members from a baseline filter at deployment, and conditions its velocity field on ensembles from that baseline filter and the observation. The proposed model therefore learns a nonlinear update while mapping each baseline ensemble independently. For sparsely observed dynamical systems, FlowEF improves both deterministic and probabilistic metrics over all four classical ensemble filters. It also achieves the best performance among the state-of-the-art generative data assimilation models.
Haoyuan Chen, Alexandre Thiéry
May 15, 2026eess.SY

The Score Kalman Filter

A central obstacle in nonlinear Bayesian filtering is representing the belief distribution. Moment-based filters address this by propagating polynomial moments and reconstructing a density from them. Recent work completes the predict-update loop via the maximum-entropy (MaxEnt) principle, but each step requires the partition function and its gradient, both nn-dimensional integrals whose cost scales exponentially, restricting the demonstrated MaxEnt moment filtering to n4n \le 4. We avoid the partition function entirely by combining score matching with Stein's identity. In our setting, score matching reduces the density fit to a single linear solve whose coefficients are assembled directly from the propagated moments. The same parameters then drive Stein's identity to close the moment hierarchy during prediction and to recover posterior moments after each Bayesian update, keeping the full predict-update loop free of partition function evaluation. The resulting Score Kalman Filter (SKF) reduces to the classical information-form Kalman filter as a special case and performs every step through linear algebra. On nonlinear coupled-oscillator networks, the SKF runs through n=20n=20 and reports lower RMSE than the EKF, UKF, EnKF, and particle-filter baselines on the tested synthetic benchmarks.
Kaito Iwasaki, Anthony Bloch, Taeyoung Lee +1
Jun 18, 2026cs.LG

ΩΩ: Operator-based Mixture Ensemble for Generative Assimilation

Characterizing non-Gaussian posterior distributions in partially observed high-dimensional nonlinear systems remains a fundamental challenge in data assimilation. Ensemble Kalman filters rely on Gaussian approximations that can be inaccurate for strongly non-Gaussian posteriors, whereas particle filters suffer from severe scalability limitations. Recent score-based generative approaches improve posterior characterization but typically require supervised training with ground-truth posterior samples, which are unavailable in most practical applications. We introduce ΩΩ (Operator-based Mixture Ensemble for Generative Assimilation), a scalable framework that integrates conditional Gaussian surrogate modeling, unsupervised score learning, and generative sampling. The conditional Gaussian surrogate provides a nonlinear non-Gaussian baseline approximation while admitting closed-form conditional posterior distributions for the unresolved variables. First, ΩΩ exploits these closed-form conditional distributions to analytically recover the high-dimensional unobserved component, reducing computational cost and mitigating the curse of dimensionality. Second, ΩΩ learns only the residual discrepancy beyond an analytical baseline through denoising score matching using ensemble trajectories alone, eliminating the need for ground-truth posterior samples and substantially reducing the learning burden. Third, ΩΩ reconstructs the full non-Gaussian posterior distribution of both observed and unobserved variables via a Gaussian mixture representation, capturing multimodal, skewed, and heavy-tailed statistics. Finally, ΩΩ employs annealed Langevin sampling to iteratively refine ensemble members from the baseline toward the target posterior. ΩΩ is validated on several turbulent models with intermittency and extreme events, consistently improving posterior accuracy.
Pouria Behnoudfar, Nan Chen