Learning and using scores for efficient Bayesian filtering

Here we present a new algorithm for Bayesian data assimilation problems where we are interested in sampling from the probability distribution of the state of a chaotic system given past observations. This sampling problem is an example of a recurring theme across the geosciences and data science, where the target distribution is high-dimensional, non-Gaussian, and only specified up to a normalizing constant. The score of a probability distribution is defined as the gradient of its log density, and it is known when, e.g., the density is available up to a normalization constant. In this work, we observe that the score of the target density conditioned on unstable manifolds can be estimated with accuracy that increases exponentially in trajectory length; that is, we estimate conditional scores directly from the chaotic model and timeseries data, without having access to the unnormalized target. We develop a novel algorithm -- Score-operator Newton or SCONE method to transport samples from any known reference distribution to the target using these computed conditional scores. This Newton method finds a transport map for sampling the target as a fixed point of a discrete-time dynamical system in infinite dimensions. We show how the SCONE algorithm can be restricted to the unstable manifold to enable sampling of target densities on the unstable manifold. This allows sampling of singular probability distributions that have absolutely continuous conditionals on the unstable manifold such as targets encountered in Bayesian data assimilation. We demonstrate that, unlike other flow-based sampling approaches, our Newton updates, which relies on elliptic PDE solves, tend to effectively tackle multimodality.