Rethinking Forward Processes for Score-Based Nonlinear Data Assimilation in High Dimensions

TL;DR

This paper introduces a novel score-based filtering method, MASF, tailored for high-dimensional nonlinear data assimilation, improving accuracy and speed over existing methods by transforming states toward measurement space.

Contribution

The paper proposes a measurement-aware forward process and develops MASF, a new score-based filter that enhances performance and computational efficiency in high-dimensional nonlinear data assimilation.

Findings

MASF outperforms existing score-based and ensemble Kalman filters.

MASF achieves up to 28.2 times speedup with amortized pretraining.

MASF demonstrates improved accuracy on high-dimensional fluid dynamics benchmarks.

Abstract

Data assimilation is the process of estimating the state of a dynamical system over time by combining model predictions with measurements. This task becomes challenging when the system is nonlinear and high-dimensional. To address this, score-based Bayesian filters have recently emerged. However, these methods still show unsatisfactory performance in certain cases, particularly under spatially sparse measurements. Such degradation stems from heuristic approximations of the likelihood score, whose errors can accumulate over time. This limitation arises because the methods simply adopt a classical forward process for generative modeling that transforms a data distribution toward a Gaussian distribution, which is independent of the measurement equation. Here, we propose a forward process tailored for filtering that transforms the system state toward the measurement space, enabling a theoretically sound formulation of the likelihood score. Based on this, we develop the Measurement-Aware Score-based Filter (MASF). We evaluate MASF on Kolmogorov flow, a high-dimensional fluid benchmark with up to $\mathcal{O}(10^5)$ dimensions, under diverse measurement operators, including nonlinear cases with a dimensional mismatch between the state and the measurements. MASF shows improved performance over existing score-based filters and ensemble-type Kalman filters. Notably, MASF achieves up to a $28.2\times$ wall-clock speedup compared with the baselines when using amortized pretraining. Our implementation is available at \texttt{https://github.com/tcnllab-oss/masf}.