Closed-form conditional diffusion models for data assimilation

TL;DR

This paper introduces a novel closed-form approach to data assimilation using conditional diffusion models, leveraging analytical score functions and kernel density estimation to improve over traditional filtering methods.

Contribution

It develops a closed-form conditional diffusion model for data assimilation that operates in black-box settings and outperforms ensemble Kalman and particle filters in nonlinear problems.

Findings

Outperforms ensemble Kalman and particle filters with small to moderate ensembles.

Effectively models complex, non-Gaussian distributions in nonlinear systems.

Operates without explicit knowledge of system or measurement models.

Abstract

We propose closed-form conditional diffusion models for data assimilation. Diffusion models use data to learn the score function (defined as the gradient of the log-probability density of a data distribution), allowing them to generate new samples from the data distribution by reversing a noise injection process. While it is common to train neural networks to approximate the score function, we leverage the analytical tractability of the score function to assimilate the states of a system with measurements. To enable the efficient evaluation of the score function, we use kernel density estimation to model the joint distribution of the states and their corresponding measurements. The proposed approach also inherits the capability of conditional diffusion models of operating in black-box settings, i.e., the proposed data assimilation approach can accommodate systems and measurement processes without their explicit knowledge. The ability to accommodate black-box systems combined with the superior capabilities of diffusion models in approximating complex, non-Gaussian probability distributions means that the proposed approach offers advantages over many widely used filtering methods. We evaluate the proposed method on nonlinear data assimilation problems based on the Lorenz-63 and Lorenz-96 systems of moderate dimensionality and nonlinear measurement models. Results show the proposed approach outperforms the widely used ensemble Kalman and particle filters when small to moderate ensemble sizes are used.