[Paper Review] Learning to Assimilate in Chaotic Dynamical Systems

Michael T. McCabe, Jed Brown|arXiv (Cornell University)|Nov 1, 2021

Meteorological Phenomena and Simulations58 references4 citations

TL;DR

This paper introduces amortized assimilation, a self-supervised deep learning framework that learns to estimate initial states in chaotic dynamical systems using only noisy, partial observations and differentiable simulation—without requiring ground truth data. It outperforms traditional data assimilation methods, especially at low ensemble sizes and high noise levels, by amortizing variational assimilation's computational cost through neural network training.

ABSTRACT

The accuracy of simulation-based forecasting in chaotic systems is heavily dependent on high-quality estimates of the system state at the time the forecast is initialized. Data assimilation methods are used to infer these initial conditions by systematically combining noisy, incomplete observations and numerical models of system dynamics to produce effective estimation schemes. We introduce amortized assimilation, a framework for learning to assimilate in dynamical systems from sequences of noisy observations with no need for ground truth data. We motivate the framework by extending powerful results from self-supervised denoising to the dynamical systems setting through the use of differentiable simulation. Experimental results across several benchmark systems highlight the improved effectiveness of our approach over widely-used data assimilation methods.

Motivation & Objective

• To address the challenge of accurate initial condition estimation in chaotic dynamical systems, where small errors amplify rapidly.
• To develop a data assimilation method that operates without ground truth data, relying solely on noisy, partial observations.
• To reduce the computational burden of variational data assimilation by learning a neural network policy through amortization.
• To extend self-supervised denoising principles from static to dynamical systems using differentiable simulation.
• To improve forecasting accuracy in high-stakes applications like numerical weather prediction, especially under resource constraints such as small ensemble sizes.

Proposed method

• The method uses a neural network architecture inspired by the Ensemble Kalman Filter (EnKF), trained via a self-supervised objective derived from differentiable simulation.
• It leverages differentiable simulation to backpropagate through the system dynamics, enabling end-to-end training of the assimilation network.
• The framework treats the data assimilation problem as a sequence-to-sequence learning task, where the network maps noisy observation sequences to improved state estimates.
• The training objective is derived from extending Batson and Royer's self-supervised denoising theory to dynamical systems, using the model's own predictions as pseudo-ground truth.
• The method is implemented using the DAPPER library and supports both ODEs and PDEs, including the Lorenz-96, Kuramoto-Shivashinsky, and Vissio-Lucarini 20 systems.
• It enables amortized inference: once trained, the network performs fast, sequential filtering without solving costly optimization problems at inference time.

Experimental results

Research questions

• RQ1 Can a self-supervised deep learning framework effectively learn to assimilate in chaotic systems without access to ground truth initial states?
• RQ2 How does amortized assimilation compare to established variational and ensemble-based data assimilation methods in terms of accuracy and robustness?
• RQ3 Does the method maintain performance under challenging conditions such as low ensemble sizes and high observation noise?
• RQ4 To what extent can the theoretical foundations of self-supervised denoising be extended to dynamical systems with time-evolving states?
• RQ5 Can the framework generalize across different types of chaotic systems, including both ODEs and PDEs?

Key findings

• Amortized assimilation (AmEnF) matches or exceeds the performance of 4D-Var, LETKF, and iEnKS across all benchmark systems, including Lorenz-96, Kuramoto-Shivashinsky, and Vissio-Lucarini 20.
• The method shows the most significant performance gains at low ensemble sizes and high noise levels, where traditional methods struggle due to nonlinearities and poor ensemble representation.
• On the Kuramoto-Shivashinsky PDE, AmEnF achieves accuracy comparable to the best-performing models despite theoretical assumptions being based on ODEs.
• The method demonstrates strong stability across repeated experiments, indicating reliable generalization and robustness.
• Training cost is high due to repeated differentiable simulations, but the amortized inference is fast and scalable, making it suitable for operational forecasting.
• The approach is compatible with legacy numerical models if adjoints are available, though this remains a current limitation for adoption.