Set Transformer | Bohan Chen's Personal Webpage

Learning Probabilistic Filters with Strictly Proper Scoring Rules

Thu, 25 Jun 2026 00:00:00 +0000

The proper scoring ensemble filter (PSEF) learns an ensemble analysis operator that targets the complete Bayesian filtering distribution rather than only its conditional mean. The operator takes a forecast ensemble and a new observation as input and returns an analysis ensemble. A permutation-invariant transformer ensures that the result respects the exchangeability of ensemble members and can be evaluated at different ensemble sizes.

The key training device is a strictly proper scoring rule, implemented with the energy score. Its expected value is uniquely minimized by the true filtering distribution, yet its empirical form requires only simulated state–observation trajectories—not direct access to the generally intractable filtering distribution. Under a realizability assumption, the population objective recovers the Bayesian filter, and a mean-field time-averaging argument connects the practical single-trajectory objective to this population target.

Experiments cover a linear–Gaussian diagnostic, the strongly non-Gaussian doubling-angle model, Lorenz-63, and Lorenz-96. Energy-score training captures multimodal distributions missed by Gaussian filters and mean-based learning objectives. The results also reveal a useful architectural distinction: an EnKF correction provides an effective inductive bias in close-to-Gaussian problems, while a flexible end-to-end update is better suited to strongly non-Gaussian posterior geometry.

The manuscript has been submitted to the Journal of Machine Learning Research.

Learning the Whole Filtering Distribution with Proper Scoring Rules

Thu, 25 Jun 2026 00:00:00 +0000

Table of Contents

Our new paper, “Learning Probabilistic Filters with Strictly Proper Scoring Rules,” is now available on arXiv and has been submitted to the Journal of Machine Learning Research (JMLR).

Paper: ·
Code:
Publication page and BibTeX:
Original announcement:

Learning a Distribution Without Seeing It

Bayesian filtering aims to infer the evolving distribution of a hidden state from partial and noisy observations. That full distribution—not only its mean—is what represents uncertainty, multimodality, and the range of plausible system states.

This creates a basic learning problem: the true filtering distribution is generally intractable, so it is not available as a supervised target. A simulator can readily produce state trajectories and their corresponding noisy observations, but it does not hand us the conditional distribution we would like a filter to learn.

Our solution is to change the training objective. A strictly proper scoring rule evaluates a predictive distribution against a realized state, and its expected score is uniquely optimized by the true data-generating distribution. We use the energy score, which lets us train from simulated state–observation trajectories while still targeting the entire Bayesian filtering distribution.

The Proper Scoring Ensemble Filter

We introduce the proper scoring ensemble filter (PSEF). Its learned analysis map takes a forecast ensemble and the latest observation, then outputs an updated analysis ensemble. The map is implemented with a permutation-invariant transformer, respecting the fact that an ensemble represents an unordered empirical distribution. The same parameterization can also operate across different ensemble sizes, with lightweight fine-tuning where needed.

The paper studies two ways to parameterize the update:

Correction terms: retain the EnKF update as an inductive bias and learn how to correct it.
End to end: learn a more flexible analysis map without constraining it to an EnKF-style correction.

The distinction matters. In close-to-Gaussian settings, the EnKF structure is useful and the correction-based approach performs best. When the observation creates a strongly non-Gaussian or multimodal posterior, the end-to-end map has the flexibility needed to represent its geometry.

Why the Loss Function Matters

The header figure gives a controlled example on the doubling-angle model. The true posterior is bimodal. Training the same learning architecture with the energy score captures both modes across ensemble sizes, including ensembles much smaller than those used by the classical baselines. A normalized mean-squared objective and Kalman-type filters instead miss the distributional structure, because matching a conditional mean is not enough to learn uncertainty.

The theory mirrors this empirical difference. Under a realizability assumption, the population proper-scoring objective is minimized by the true Bayesian filtering distribution. We also connect the finite-ensemble, single-trajectory loss used in training to the population objective through a mean-field consistency and time-averaging argument.

Non-Gaussian Filtering in Lorenz Systems

Lorenz-63 with a partial square observation. The end-to-end, energy-score-trained PSEF most closely follows the particle-filter reference, including its cross-shaped and bimodal structures.

We evaluate PSEF on a linear–Gaussian diagnostic, the doubling-angle model, Lorenz-63, and Lorenz-96. In Lorenz-63, squaring the observed coordinate introduces sign ambiguity and produces complex non-Gaussian filtering distributions. The end-to-end PSEF trained with the energy score reproduces these structures more faithfully than mean-based learned filters and classical Gaussian ensemble filters.

Lorenz-96 results across identity, square, and arctangent observations. Proper-score-trained filters remain among the strongest methods, while the best architecture depends on the posterior geometry.

The Lorenz-96 results reinforce that there is no single architectural bias that is optimal for every problem. Correction terms are effective when the filtering distribution remains approximately unimodal, whereas the end-to-end energy-score model is strongest when the observation induces multimodality. Across both cases, probabilistic training is what allows the learned ensemble to target distributional accuracy rather than only state-estimation error.

Collaboration and Reproducibility

I led this project, including the implementation and numerical experiments, in collaboration with Eviatar Bach, Ricardo Baptista, Jochen Bröcker, and Andrew Stuart. The contains the code and experiment scripts.

Figures on this page were extracted from the paper, which is available under a CC BY 4.0 license.

Learning Enhanced Ensemble Filters

Sun, 15 Feb 2026 00:00:00 +0000

This work introduces the measure neural mapping enhanced ensemble filter (MNMEF), a learning-based data-assimilation method derived from a mean-field formulation of the filtering problem. Measure neural mappings extend neural operators to maps acting on probability measures; their finite-ensemble implementation uses a permutation-invariant set transformer.

The mean-field construction allows most learned parameters to be shared across ensemble sizes. A model can therefore be trained efficiently using a small ensemble and deployed at other ensemble sizes, while lightweight fine-tuning adapts a limited number of parameters such as inflation and localization.

Experiments on Lorenz-63, Lorenz-96, and Kuramoto–Sivashinsky systems show improved filtering accuracy relative to optimized classical ensemble methods across both small and larger ensembles.

Learning Enhanced Ensemble Filters Published in JCP

Thu, 11 Dec 2025 00:00:00 +0000

Our paper, “Learning Enhanced Ensemble Filters,” has been published in the Journal of Computational Physics.

Journal article:
Paper: ·
Code:
Publication page and BibTeX:
Original announcement:

The paper introduces the measure neural mapping enhanced ensemble filter (MNMEF). Starting from a mean-field formulation of filtering, we use a set transformer to learn corrections to ensemble filtering updates while respecting the permutation symmetry of an ensemble. The mean-field perspective allows a model trained at one ensemble size to be deployed at another, with lightweight fine-tuning for ensemble-size-dependent parameters.

Lorenz-96 results across ensemble sizes and two observation-noise levels. The pretrained and fine-tuned MNMEF models remain competitive as the ensemble size changes.

Across Lorenz-63, Lorenz-96, and Kuramoto–Sivashinsky experiments, MNMEF improves on optimized classical ensemble filters over a range of ensemble sizes.

Kuramoto–Sivashinsky filtering with only every eighth dimension observed. With an ensemble of ten members, MNMEF produces a more accurate state estimate than the optimized LETKF benchmark while maintaining a comparable spread.

The method also learns ensemble-size-dependent inflation and localization during the lightweight fine-tuning stage. Although the localization function is not constrained to have a prescribed shape, it learns a distance-decay profile similar to the widely used Gaspari–Cohn function.

Learned localization weights across ensemble sizes. The red curves adapt across assimilation steps, while the blue Gaspari–Cohn reference remains fixed.

I led this work in collaboration with Eviatar Bach, Ricardo Baptista, Edoardo Calvello, and Andrew Stuart.

The final article appears in Journal of Computational Physics, Volume 547, Article 114550: .