Auto-differentiable data assimilation: Co-learning of states, dynamics, and filtering algorithms

1. Problem Statement

Data assimilation (DA) aims to estimate the state of a dynamical system from partial, noisy observations. Traditional DA methods (e.g., Kalman Filter, 3DVar, EnKF) face two critical challenges:

1. Model Error: The dynamics governing the system are often unknown, expensive to simulate, or imperfectly modeled.
2. Parameter Tuning: DA algorithms require extensive manual tuning of hyperparameters (e.g., background error covariance, covariance inflation, mixing parameters) to balance trust between the forecast model and observations.

Existing machine learning approaches often address these issues in isolation (e.g., learning only the dynamics or only the filter gain). There is a lack of a unified framework that can jointly learn the system states, the underlying dynamics model, and the filtering algorithm parameters simultaneously from noisy data.

2. Methodology: Auto-Differentiable Filtering

The authors propose a general framework called Auto-Differentiable Filtering. This approach treats the entire data assimilation pipeline as a differentiable computational graph, allowing for end-to-end optimization via gradient descent.

A. General Framework

The framework assumes a hidden Markov model where:

• Dynamics: $x^*_t = F_{\theta_1}(x^*_{t-1}) + \text{noise}$. The true dynamics $F_{\theta^*_1}$ are unknown; an imperfect model $F_{\theta_1}$ (potentially a neural network) is used.
• Observations: $y_t = H_t x^*_t + \eta_t$.
• Filtering: The analysis step is defined generally as:
  $$x^n_t = (I - K_\phi H_t) M_\theta(x^n_{t-1}) + K_\phi Y(y_t; R_t)$$
  Where $M_\theta$ is the forecast model, $K_\phi$ is the Kalman gain (parameterized by $\phi$), and $Y$ incorporates observations.

B. The Learning Process (Algorithm 1)

The core innovation is the use of auto-differentiation to update parameters $\theta$ (dynamics) and $\phi$ (filtering) based on a loss function derived from the observations.

1. Forward Pass: The filter runs through the time series $y_{1:T}$ using current parameters to generate state estimates (analyses).
2. Loss Function: The authors define a loss function $L(\theta, \phi)$ based on the negative log-likelihood of the observations:
   $$L(\theta, \phi) = -\sum_{t=1}^T \log \mathcal{N}(y_t; H_t \hat{m}^\theta_t, H_t C^{\theta,\phi}_t H_t^\top + R_t)$$
   This measures the discrepancy between observed data and the predicted observation distribution (mean and covariance) generated by the filter.
3. Backward Pass: Gradients $\nabla_\theta L$ and $\nabla_\phi L$ are computed via backpropagation through the time series.
4. Optimization: Parameters are updated using gradient descent. To handle long time series and avoid vanishing/exploding gradients, the authors employ Truncated Backpropagation Through Time (TBPTT).

C. Special Cases Derived

The framework generalizes and extends several existing methods:

• AD-3DVar-C: Learns a static forecast error covariance matrix $C_\phi$ (parameterized as $B_\phi B_\phi^\top$) rather than the Kalman gain directly. This allows handling of time-varying observation matrices $H_t$.
• AD-Ens3DVar: A hybrid approach learning a convex combination of static and ensemble-based covariances.
• AD-EnKF (Modified): Extends the existing AD-EnKF to jointly learn the covariance inflation factor alongside dynamics.
• AD-3DVar-K (Comparison): A baseline method from prior work (Levine & Stuart, 2022) that learns the Kalman gain directly but cannot handle time-varying $H_t$.

3. Key Contributions

1. Unified Framework: Formulated a general, theoretically motivated framework for the co-learning of states, dynamics, and filtering parameters.
2. New Algorithms: Derived AD-3DVar-C and AD-Ens3DVar, which offer specific advantages over existing methods (e.g., handling time-varying observation models).
3. Enhanced AD-EnKF: Modified the existing AD-EnKF to co-learn filtering parameters (covariance inflation), not just dynamics.
4. Comprehensive Evaluation: Validated the framework across three distinct scientific domains with varying complexities (linear vs. nonlinear, static vs. time-varying observations).

4. Experimental Results

The framework was tested on three dynamical systems:

A. Astrodynamics (Clohessy-Wiltshire Equations)

• Task: Linear system; learning a scalar orbital rate parameter from partial position observations.
• Result: Ensemble-based methods (AD-EnKF, AD-Ens3DVar) significantly outperformed 3DVar methods. They successfully recovered the true parameter and accurately estimated unobserved velocity states. 3DVar methods struggled to recover unobserved components when $H_t$ was static.

B. Atmospheric Science (Lorenz-96 System)

• Task: Nonlinear, chaotic system; learning a residual neural network to correct an imperfect physics model.
• Result: Ensemble methods were robust to high state dimensions ($d_x$), high noise levels, and sparse observations (low $d_y/d_x$ ratio). 3DVar methods degraded rapidly as noise increased or observability decreased. Ensemble methods achieved roughly 50% lower filtering error compared to 3DVar variants.

C. Systems Biology (Generalized Lotka-Volterra)

• Task: Nonlinear population dynamics; learning interaction matrices with time-varying observation matrices ($H_t$ changes over time).
• Result: This experiment highlighted the superiority of AD-3DVar-C and ensemble methods over AD-3DVar-K (which fails here). All methods handled the time-varying $H_t$, but ensemble methods provided more accurate trajectory reconstructions for sparse, noisy data.

5. Significance and Guidelines for Practitioners

The paper provides a "menu" of algorithms with trade-offs between accuracy, computational cost, and flexibility:

| Algorithm | Best Use Case | Pros | Cons |
| :--- | :--- | :--- | : |
| AD-EnKF / AD-Ens3DVar | High accuracy required; Nonlinear systems; Sparse/Noisy data. | Highest accuracy; Robust to unobserved states; Learns complex dynamics. | High memory/compute cost (scales with ensemble size $N$). |
| AD-3DVar-C | Time-varying observation models ($H_t$); High-dimensional systems with limited memory. | Handles time-varying $H_t$; Lower memory than EnKF; Scales well with $d_x$. | Struggles to estimate unobserved states if $H$ is static; Lower accuracy than EnKF. |
| AD-3DVar-K | Fixed observation models ($H$); Extreme computational constraints. | Lowest memory/compute cost. | Cannot handle time-varying $H$; Struggles with unobserved states; Lower accuracy. |

Conclusion:
This work bridges the gap between machine learning and classical data assimilation. By making the filtering process auto-differentiable, it eliminates the need for manual tuning of covariance matrices and inflation factors, allowing the system to "learn" the optimal balance between model and data. While ensemble-based methods offer superior accuracy, the framework provides lighter-weight alternatives (3DVar variants) for scenarios where computational resources are constrained or observation models change dynamically.