Non-Gaussian smoothing improves data assimilation in complex simulators

PR-Smoother: Simulator-Preserving Non-Gaussian Smoothing for Data Assimilation

Machine Learning

Summary

Data assimilation is a way to combine physical models with real-world observations to better understand systems like weather or fluid flow. Traditional methods often assume simple (Gaussian) errors and struggle with complex, high-dimensional systems. The authors propose PR-Smoother, a method that keeps the original simulator intact but learns corrections to better capture uncertain, complicated outcomes. This method can handle tricky cases where observations are ambiguous and works well even for very large systems.

What this means in practice

  • For meteorological forecast teams: Improve weather forecasts by better representing uncertain atmospheric states using complex physical models with PR-Smoother's non-Gaussian smoothing technique.
  • For computational fluid dynamics engineers: Enhance simulation calibration and state estimation in turbulent flow problems with high-dimensional models using PR-Smoother's ability to jointly infer states and parameters from observations.

Authors

Yuta Tarumi

Abstract

Many physical data assimilation (DA) workflows require smoothing methods that represent non-Gaussian posteriors over physical state variables, scale to high-dimensional simulators, train from observation windows alone, and remain compatible with calibration of the prescribed simulator. We introduce PR-Smoother, a simulator-preserving amortized smoother designed for this prescribed-simulator DA regime. Its key design principle is to keep the prescribed simulator explicit in both the evidence lower bound and the variational family: rather than learning replacement dynamics or a learned trajectory prior, PR-Smoother learns only future-conditioned corrections around the prescribed rollout. This yields an explicit non-Gaussian smoothing distribution over physical trajectories and supports joint state, parameter, and sensor-bias learning from observations alone. The variational family contains the exact smoother in deterministic and linear-Gaussian limits. Empirically, PR-Smoother captures multimodal posteriors in 4-dimensional Lorenz-96, remains accurate under ambiguous nonlinear observations and process noise in 40-dimensional Lorenz-96, and scales to joint state-parameter-bias inference in 16,384-dimensional Kolmogorov flow.