# Auto-differentiable Ensemble Kalman Filters

Yuming Chen University of Chicago, Chicago, IL<EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS>Daniel Sanz-
Alonso11footnotemark: 1 Rebecca Willett11footnotemark: 1

###### Abstract

Data assimilation is concerned with sequentially estimating a temporally-
evolving state. This task, which arises in a wide range of scientific and
engineering applications, is particularly challenging when the state is high-
dimensional and the state-space dynamics are unknown. This paper introduces a
machine learning framework for learning dynamical systems in data
assimilation. Our auto-differentiable ensemble Kalman filters (AD-EnKFs) blend
ensemble Kalman filters for state recovery with machine learning tools for
learning the dynamics. In doing so, AD-EnKFs leverage the ability of ensemble
Kalman filters to scale to high-dimensional states and the power of automatic
differentiation to train high-dimensional surrogate models for the dynamics.
Numerical results using the Lorenz-96 model show that AD-EnKFs outperform
existing methods that use expectation-maximization or particle filters to
merge data assimilation and machine learning. In addition, AD-EnKFs are easy
to implement and require minimal tuning.

##  1 Introduction

Time series of data arising across geophysical sciences, remote sensing,
automatic control, and a variety of other scientific and engineering
applications often reflect observations of an underlying dynamical system
operating in a latent state-space. Estimating the evolution of this latent
state from data is the central challenge of data assimilation (DA) [39, 28,
75, 49, 68]. However, in these and other applications, we often lack an
accurate model of the underlying dynamics, and the dynamical model needs to be
learned from the observations to perform DA. This paper introduces auto-
differentiable ensemble Kalman filters (AD-EnKFs), a machine learning (ML)
framework for the principled co-learning of states and dynamics. This
framework enables learning in three core categories of unknown dynamics: (a)
parametric dynamical models with unknown parameter values; (b) fully-unknown
dynamics captured using neural network (NN) surrogate models; and (c)
inaccurate or partially-known dynamical models that can be improved using NN
corrections. AD-EnKFs are designed to scale to high-dimensional states,
observations, and NN surrogate models.

In order to describe the main idea behind the AD-EnKF framework, let us
introduce briefly the problem of interest. Our setting will be formalized in
Section 2 below. Let $x_{0:T}:=\\{x_{t}\\}_{t=0}^{T}$ be a time-homogeneous
_state process_ with transition kernel $p_{\theta}(x_{t}|x_{t-1})$
parameterized by a vector $\theta.$ For instance, $\theta$ may contain unknown
parameters of a parametric dynamical model or the parameters of a NN surrogate
model for the dynamics. Our aim is to learn $\theta$ from observations
$y_{1:T}:=\\{y_{t}\\}_{t=1}^{T}$ of the state, and thereby learn the unknown
dynamics and estimate the state process. The AD-EnKF framework learns $\theta$
iteratively. Each iteration consists of three steps: (i) use EnKF to compute
an estimate $\mathcal{L}_{{\footnotesize\text{EnKF}}}(\theta)$ of the data
log-likelihood $\mathcal{L}(\theta):=\log p_{\theta}(y_{1:T});$ (ii) use auto-
differentiation (“autodiff”) to compute the gradient
$\nabla_{\theta}\mathcal{L}_{{\footnotesize\text{EnKF}}}(\theta);$ and (iii)
take a gradient ascent step. Filtered estimates of the state are obtained
using the learned dynamics.

The EnKF, reviewed in Section 3, estimates the data log-likelihood using an
ensemble of particles. Precisely, given a transition kernel
$p_{\theta}(x_{t}|x_{t-1})$, the EnKF generates particles $x_{0:T}^{1:N};$
that is, each particle $x_{t}^{n}\in\mathbbm{R}^{d_{x}}$ and
$x_{0:T}^{1:N}:=\\{x_{t}^{n}\\}_{{t=0,\ldots T\atop n=1,\ldots,N}}.$ The log-
likelihood estimate $\mathcal{L}_{{\footnotesize\text{EnKF}}}(\theta)$ depends
on $\theta$ through these particles and also through the given transition
kernel. Differentiating the map
$\theta\mapsto\mathcal{L}_{{\footnotesize\text{EnKF}}}(\theta)$ in step (ii)
of AD-EnKF involves differentiating _both_ the map $\theta\mapsto
x_{0:T}^{1:N}$ from parameter to EnKF particles and the map
$(\theta,x_{0:T}^{1:N})\mapsto\mathcal{L}_{{\footnotesize\text{EnKF}}}(\theta)$
from parameters and EnKF particles to EnKF log-likelihood estimate. A key
feature of our approach is that
$\theta\mapsto\mathcal{L}_{{\footnotesize\text{EnKF}}}(\theta)$ can be auto-
differentiated using the reparameterization trick ([44] and Section 4.1) and
autodiff capabilities of NN software libraries such as PyTorch [64], JAX [11],
and Tensorflow [1]. Automatic differentiation is different from numerical
differentiation in that derivatives are computed exactly through compositions
of elementary functions whose derivatives are known, as opposed to finite
difference approximations that cause discretization errors.

The AD-EnKF framework represents a significant conceptual and methodological
departure from existing approaches to blend DA and ML based on the
expectation-maximization (EM) framework, see Fig. 1 below. Specifically, at
each iteration, EM methods that build on the EnKF [65, 12, 9] employ a
surrogate likelihood $\mathcal{L}_{{\footnotesize\text{EM-
EnKF}}}\bigl{(}\theta\,;x_{0:T}^{1:N}\bigr{)}$ where the particles
$x_{0:T}^{1:N}$ are generated by EnKF and _fixed_. Importantly, EM methods
compute gradients used to learn dynamics by differentiating only through the
$\theta$-dependence in $\mathcal{L}_{{\footnotesize\text{EM-EnKF}}}$ that does
not involve the particles. In particular, in contrast to AD-EnKF, the map
$\theta\mapsto x_{0:T}^{1:N}$ from parameter to EnKF particles is _not_
differentiated. Moreover, the performance of EM methods is sensitive to the
specific choice of EnKF algorithm in use, and the tuning of additional
algorithmic parameters of EM can be challenging [12, 9]. Our numerical
experiments suggest that, even when optimally tuned, EM methods underperform
AD-EnKF in high-dimensional regimes. The outperfomance of AD-EnKF may be
explained by the additional gradient information obtained by differentiating
the map $\theta\mapsto x_{0:T}^{1:N}$.

$\theta$$y_{1:T}$$x_{0:T}^{1:N}$$\mathcal{L}_{{\footnotesize\text{EnKF}}}$ (a)
AD-EnKF

$\theta$$y_{1:T}$$x_{0:T}^{1:N}$$\mathcal{L}_{\text{EM-EnKF}}$ (b) EM-EnKF

With AD-EnKF, parameters $\theta$ and observations $y_{1:T}$ are used to
generate EnKF particles $x_{0:T}^{1:N}$; the particles together with $\theta$
and $y_{1:T}$ are used to compute the likelihood
$\mathcal{L}_{{\footnotesize\text{EnKF}}}$, and the gradient
$\nabla_{\theta}\mathcal{L}_{{\footnotesize\text{EnKF}}}$ explicitly accounts
for the map from $\theta$ to the particles $x_{0:T}^{1:N}$. In contrast, with
EM-EnKF, the likelihood $\mathcal{L}_{{\footnotesize\text{EM-EnKF}}}$ is a
function of $\theta$ and fixed particles $x_{0:T}^{1:N}$ generated by EnKF, so
that computing the gradient $\nabla_{\theta}\mathcal{L}_{\text{EM-EnKF}}$ does
not account for the map from $\theta$ to the particles $x_{0:T}^{1:N}$.

Figure 1:  Computational graph of AD-EnKF and EM-EnKF. Dashed squares
represent computations performed by the EnKF. Gray arrows in (b) indicate that
the construction of $\mathcal{L}_{\text{EM-EnKF}}$ is performed in two steps:
(1) obtain $x_{0:T}^{1:N}$ from $\theta$ and $y_{1:T}$ (gray arrows); and (2)
use $\theta$ and $x_{0:T}^{1:N}$ (no longer seen as a function of $\theta$) to
define $\mathcal{L}_{\text{EM-EnKF}}.$ In contrast, those lines are black in
(a), indicating that in AD-EnKF the particles $x_{0:T}^{1:N}$ in
$\mathcal{L}_{{\footnotesize\text{EnKF}}}$ are seen as varying with $\theta$.

The AD-EnKF framework also represents a methodological shift from existing
differentiable particle filters [59, 55, 51]. Similar to AD-EnKF, these
methods rely on autodiff of a map
$\theta\mapsto\mathcal{L}_{{\footnotesize\text{PF}}}(\theta),$ where the log-
likelihood estimate $\mathcal{L}_{{\footnotesize\text{PF}}}(\theta)$ depends
on $\theta$ through weighted particles $(w_{0:T}^{1:N},x_{0:T}^{1:N})$
obtained by running a particle filter (PF) with transition kernel
$p_{\theta}(x_{t}|x_{t-1})$. However, the use of PF suffers from two caveats.
First, it is not possible to auto-differentiate directly through the PF
resampling steps [59, 55, 51]. Second, while the PF log-likelihood estimates
are consistent, their variance can be large, especially in high-dimensional
systems. Moreover, their _gradient_ , which is the quantity used to perform
gradient ascent to learn $\theta$, is not consistent [17].

###  1.1 Contributions

This paper seeks to set the foundations and illustrate the capabilities of the
AD-EnKF framework through rigorous theory and systematic numerical
experiments. Our main contributions are:

  * •

We develop new theoretical convergence guarantees for the large sample EnKF
estimation of log-likelihood gradients in linear-Gaussian settings (Theorem
3.2).

  * •

We combine ideas from online training of recurrent networks (specifically,
Truncated Backpropagation Through Time – TBPTT) with the learning of AD-EnKF
when the data sequence is long, i.e. $T$ is large.

  * •

We provide numerical evidence of the superior estimation accuracies of log-
likelihoods and gradients afforded by EnKF relative to PF methods in high-
dimensional settings. In particular, we illustrate the importance of using
localization techniques, developed in the DA literature, for EnKF log-
likelihood and gradient estimation, and the corresponding performance boost
within AD-EnKF.

  * •

We conduct a numerical case study of AD-EnKF on the Lorenz-96 model [54],
considering parameterized dynamics, fully-unknown dynamics, and correction of
an inaccurate model. The importance of the Lorenz-96 model in geophysical
applications and for testing the efficacy of filtering algorithms is
highlighted, for instance, in [56, 48, 47, 12]. Our results show that AD-EnKF
outperforms existing methods based on EM or differentiable PFs. The
improvements are most significant in challenging high-dimensional and
partially-observed settings.

###  1.2 Related Work

The EnKF algorithm was developed as a state estimation tool for DA [27] and is
now widely used in numerical weather prediction and geophysical applications
[84, 89]. Recent reviews include [38, 41, 70]. The idea behind the EnKF is to
propagate $N$ equally-weighted particles through the dynamics and assimilate
new observations using Kalman-type updates computed with empirical moments.
When the state dimension $d_{x}$ is high and the ensemble size $N$ is
moderate, traditional Kalman-type methods require $O(d_{x}^{2})$ memory to
store full covariance matrices, while storing empirical covariances in EnKFs
only requires $O(Nd_{x})$ memory. The use of EnKF for joint learning of state
and model parameters by _state augmentation_ was introduced in [3], where EnKF
is run on an augmented state-space that includes the state and parameters.
However, this approach requires one to design a pseudo-dynamic for the
parameters which needs careful tuning and can be problematic when certain
types of parameters (e.g., error covariance matrices) are involved [79, 21] or
if the dimension of the parameters is high. In this paper, we employ EnKFs to
approximate the data log-likelihood. The use of EnKF to perform derivative-
free maximum likelihood estimation (MLE) is studied in [81, 65]. An empirical
comparison of the likelihood computed using the EnKF and other filtering
algorithms is made in [13]; see also [35, 58]. The paper [26] uses EnKF
likelihood estimates to design a pseudo-marginal Markov chain Monte Carlo
(MCMC) method for Bayesian inference of model parameters. The works [79, 80]
propose online Bayesian parameter estimation using the likelihood computed
from the EnKF under a certain family of conjugate distributions. However, to
the best of our knowledge, there is no prior work on state and parameter
estimation that utilizes gradient information of the EnKF likelihood.

The embedding of EnKF and ensemble Kalman smoothers (EnKS) into the EM
algorithm for MLE [22, 7] has been studied in [85, 86, 25, 65], with a special
focus on estimation of error covariance matrices. The expectation step
(E-step) is approximated with EnKS under the Monte Carlo EM framework [88]. In
addition, [12, 60] incorporate deep learning techniques in the maximization
step (M-step) to train NN surrogate models. The paper [9] proposes Bayesian
estimation of model error statistics, in addition to an NN emulator for the
dynamics. On the other hand, [87, 16] consider online EM methods for error
covariance estimation with EnKF. Although gradient information is used during
the M-step to train the surrogate model [12, 60, 9], these methods do not
auto-differentiate through the EnKF (see Fig. 1), and accurate approximation
of the E-step is hard to achieve with EnKF or EnKS.

Another popular approach for state and parameter estimation are particle
filters (PFs) [33, 24] that approximate the filtering step by propagating
samples with a kernel, reweighing them with importance sampling, and
resampling to avoid weight degeneracy. PFs give an unbiased estimate of the
data likelihood [19, 5]. Based upon this likelihood estimate, a particle MCMC
Bayesian parameter estimation method is designed in [5]. Although PF
likelihood estimates are unbiased, they suffer from two important caveats.
First, their variance can be large, as they inherit the weight degeneracy of
importance sampling in high dimensions [78, 10, 2, 73, 76]. Second, while the
propagation and reweighing steps of PFs can be auto-differentiated, the
resampling steps involve discrete distributions that cannot be handled by the
reparameterization trick. For this reason, previous differentiable PFs omit
autodiff of the resampling step [59, 55, 51], introducing a bias.

An alternative to MLE methods is to optimize a lower bound of the data log-
likelihood with variational inference (VI) [7, 44, 66]. The posterior
distribution over the latent states is approximated with a parametric
distribution and is jointly optimized with model parameters defining the
underlying state-space model. In this direction, variational sequential Monte
Carlo (VSMC) methods [59, 55, 51] construct the lower bound using a PF
algorithm. Moreover, the proposal distribution of the PF is parameterized and
jointly optimized with model parameters defining the state-space model.
Although VSMC methods provide consistent data log-likelihood estimates, they
suffer from the same two caveats as likelihood-based PF methods. Other works
that build on the VI framework include [45, 67, 29]. An important challenge is
to obtain suitable parameterizations of the posterior, especially when the
state dimension is high. For this reason, a restrictive Gaussian
parameterization with a diagonal covariance matrix is often used in practice
[45, 29].

### Outline

This paper is organized as follows. Section 2 formalizes our framework and
reviews a characterization of the likelihood in terms of normalizing constants
arising in sequential filtering. Section 3 overviews EnKF algorithms for
filtering and log-likelihood estimation. Section 4 contains our main
methodological contributions. Numerical experiments on linear-Gaussian and
Lorenz-96 models are described in Section 5. We close in Section 6.

### Notation

We denote by $t\in\\{0,1,\dots,T\\}$ a discrete time index and by
$n\in\\{1,\dots,N\\}$ a particle index. Time indices will be denoted with
subscripts and particles with superscripts, so that $x_{t}^{n}$ represents a
generic particle at time $t.$ We denote
$x_{t_{0}:t_{1}}:=\\{x_{t}\\}_{t=t_{0}}^{t_{1}}$ and
$x^{n_{1}:n_{2}}:=\\{x^{n}\\}_{n=n_{0}}^{n_{1}}$. The collection
$x_{t_{0}:t_{1}}^{n_{0}:n_{1}}$ is defined in a similar way. The Gaussian
density with mean $m$ and covariance $C$ evaluated at $x$ is denoted by
$\mathcal{N}(x;m,C)$. The corresponding Gaussian distribution is denoted by
$\mathcal{N}(m,C)$. For square matrices $A$ and $B$, we write $A\succ B$ if
$A-B$ is positive definite, and $A\succeq B$ if $A-B$ is positive semi-
definite. For $A\succeq 0$, we denote by $A^{1/2}$ the unique matrix $B\succeq
0$ such that $B^{2}=A$. We denote by $|v|$ the 2-norm of a vector $v$, and by
$|A|$ the Frobenius norm of a matrix $A$.

##  2 Problem Formulation

Let $x_{0:T}$ be a time-homogeneous Markov chain of hidden _states_
$x_{t}\in\mathbbm{R}^{d_{x}}$ with transition kernel
$p_{\theta}(x_{t}|x_{t-1})$ parameterized by
$\theta\in\mathbbm{R}^{d_{\theta}}.$ Let $y_{1:T}$ be observations of the
state. We seek to learn the parameter $\theta$ and recover the state process
$x_{0:T}$ from the observations $y_{1:T}$. In Section 2.1, we formalize our
problem setting, emphasizing our main goal of learning unknown dynamical
systems for improved DA. Section 2.2 describes how the log-likelihood
$\mathcal{L}(\theta)=\log p_{\theta}(y_{1:T})$ can be written in terms of
normalizing constants arising from sequential filtering. This idea will be
used in Section 3 to obtain EnKF estimates for $\mathcal{L}(\theta)$ and
$\nabla_{\theta}\mathcal{L}(\theta)$, which are then employed in Section 4 to
learn $\theta$ by gradient ascent.

###  2.1 Setting and Motivation

We consider the following state-space model (SSM)

(transition) $\displaystyle\quad\quad x_{t}$
$\displaystyle=F_{\alpha}(x_{t-1})+\xi_{t},\quad\quad$
$\displaystyle\xi_{t}\sim\mathcal{N}(0,Q_{\beta}),\quad\quad$ $\displaystyle
1\leq t\leq T,$ (2.1) (observation) $\displaystyle\quad\quad y_{t}$
$\displaystyle=Hx_{t}+\eta_{t},\quad\quad$
$\displaystyle\eta_{t}\sim\mathcal{N}(0,R),\quad\quad$ $\displaystyle 1\leq
t\leq T,$ (2.2) (initialization) $\displaystyle\quad\quad x_{0}$
$\displaystyle\sim p_{0}(x_{0}).$ (2.3)

The initial distribution $p_{0}$ and the matrices $H\in R^{d_{y}\times d_{x}}$
and $R\succ 0$ are assumed to be known. Nonlinear observations can be dealt
with by augmenting the state. We further assume independence of all random
variables $x_{0},$ $\xi_{1:T}$ and $\eta_{1:T}.$ Finally, the transition
kernel
$p_{\theta}(x_{t}|x_{t-1})=\mathcal{N}(x_{t};F_{\alpha}(x_{t-1}),Q_{\beta}),$
parameterized by $\theta:=\\{\alpha,\beta\\},$ is defined in terms of a
deterministic map $F_{\alpha}$ and Gaussian additive noise. This kernel
approximates an unknown state transition of the form

$x_{t}=F^{*}(x_{t-1})+\xi_{t},\quad\quad\xi_{t}\sim\mathcal{N}(0,Q^{*}),\quad\quad
1\leq t\leq T,$ (2.4)

where $Q^{*}=0$ if the true evolution of the state is deterministic. The
parameter $\beta$ allows us to estimate the possibly unknown $Q^{*}.$ We
consider three categories of unknown state transition $F^{*}$, leading to
three types of learning problems:

  1. (a)

_Parameterized dynamics_ : $F^{*}=F_{\alpha^{*}}$ is parameterized, but the
true parameter $\alpha^{*}$ is unknown and needs to be estimated.

  2. (b)

_Fully-unknown dynamics_ : $F^{*}$ is fully unknown and $\alpha$ represents
the parameters of a NN surrogate model $F_{\alpha}^{\footnotesize\text{NN}}$
for $F^{*}.$ The goal is to find an accurate surrogate model
$F_{\alpha}^{\footnotesize\text{NN}}$.

  3. (c)

_Model correction_ : $F^{*}$ is unknown, but an inaccurate model
$F_{\text{approx}}\approx F^{*}$ is available. Here $\alpha$ represents the
parameters of a NN $G_{\alpha}^{\footnotesize\text{NN}}$ used to correct the
inaccurate model. The goal is to learn $\alpha$ so that
$F_{\alpha}:=F_{\text{approx}}+G_{\alpha}^{\footnotesize\text{NN}}$
approximates $F^{*}$ accurately.

In some applications, the map $F^{*}$ may represent the flow between
observations of an autonomous differential equation driving the state, i.e.

$\frac{\text{d}x}{\text{d}s}=f^{*}(x),\quad\quad F^{*}:x(s)\mapsto
x(s+\Delta_{s}),$ (2.5)

where $f^{*}$ is an unknown vector field and $\Delta_{s}$ is the time between
observations. Then, the map $F_{\alpha}$ in (2.1) (resp.
$F_{\alpha}^{\footnotesize\text{NN}}$, $F_{\text{approx}}$,
$G_{\alpha}^{\footnotesize\text{NN}}$) will be similarly defined as the
$\Delta_{s}$-flow of a differential equation with vector field $f_{\alpha}$
(resp. $f_{\alpha}^{\footnotesize\text{NN}}$, $f_{\text{approx}}$,
$g_{\alpha}^{\footnotesize\text{NN}}$). Once $\theta=\\{\alpha,\beta\\}$ is
learned, the state $x_{0:T}$ can be recovered with a filtering algorithm using
the transition kernel $p_{\theta}(x_{t}|x_{t-1})$. We will illustrate the
implementation and performance of AD-EnKF in these three categories of unknown
dynamics in Section 5 using the Lorenz-96 model to define the vector field
$f^{*}$. We remark that learning NN surrogate models for the dynamics may be
useful even when the true state transition $F^{*}$ is known, since
$F_{\alpha}^{\footnotesize\text{NN}}$ may be cheaper to evaluate than $F^{*}$.

###  2.2 Sequential Filtering and Data Log-likelihood

Suppose that $\theta=\\{\alpha,\beta\\}$ is known. We recall that, for $1\leq
t\leq T,$ the _filtering distributions_ $p_{\theta}(x_{t}|y_{1:t})$ of the SSM
(2.1)-(2.2)-(2.3) can be obtained sequentially, alternating between _forecast_
and _analysis_ steps:

(forecast) $\displaystyle\quad\quad
p_{\theta}(x_{t}|y_{1:t-1})=\int\mathcal{N}(x_{t};F_{\alpha}(x_{t-1}),Q_{\beta})p_{\theta}(x_{t-1}|y_{1:t-1})\text{d}x_{t-1},$
(2.6) (analysis) $\displaystyle\quad\quad
p_{\theta}(x_{t}|y_{1:t})=\frac{1}{Z_{t}(\theta)}\mathcal{N}(y_{t};Hx_{t},R)p_{\theta}(x_{t}|y_{1:t-1}),$
(2.7)

with the convention $p_{\theta}(\cdot|y_{1:0}):=p_{\theta}(\cdot)$. Here
$Z_{t}(\theta)$ is a normalizing constant which does not depend on $x_{t}$. It
can be easily shown that

$Z_{t}(\theta)=p_{\theta}(y_{t}|y_{1:t-1})=\int\mathcal{N}(y_{t};Hx_{t},R)p_{\theta}(x_{t}|y_{1:t-1})\text{d}x_{t},$
(2.8)

and therefore the data log-likelihood admits the characterization

$\mathcal{L}(\theta):=\log p_{\theta}(y_{1:T})=\sum_{t=1}^{T}\log
p_{\theta}(y_{t}|y_{1:t-1})=\sum_{t=1}^{T}\log Z_{t}(\theta).$ (2.9)

Analytical expressions of the filtering distributions
$p_{\theta}(x_{t}|y_{1:t})$ and the data log-likelihood $\mathcal{L}(\theta)$
are only available for a small class of SSMs, which includes linear-Gaussian
and discrete SSMs [40, 62]. Outside these special cases, filtering algorithms
need to be employed to approximate the filtering distributions, and these
algorithms can be leveraged to estimate the log-likelihood.

##  3 Ensemble Kalman Filter Estimation of the Log-likelihood and its Gradient

In this section, we briefly review EnKFs and how they can be used to obtain an
estimate $\mathcal{L}_{{\footnotesize\text{EnKF}}}(\theta)$ of the log-
likelihood $\mathcal{L}(\theta)$. As will be detailed in Section 4, the map
$\theta\mapsto\mathcal{L}_{{\footnotesize\text{EnKF}}}(\theta)$ can be readily
auto-differentiated to compute
$\nabla_{\theta}\mathcal{L}_{{\footnotesize\text{EnKF}}}(\theta)$, and this
gradient can be used to learn the parameter $\theta$. Section 3.1 gives
background on EnKFs, Section 3.2 shows how EnKFs can be used to estimate
$\mathcal{L}(\theta)$, and Section 3.3 contains novel convergence guarantees
for the EnKF estimation of $\mathcal{L}(\theta)$ and
$\nabla_{\theta}\mathcal{L}(\theta).$

###  3.1 Ensemble Kalman Filters

Given $\theta=\\{\alpha,\beta\\},$ the EnKF algorithm [27, 28] sequentially
approximates the filtering distributions $p_{\theta}(x_{t}|y_{1:t})$ using $N$
equally-weighted particles ${x}_{t}^{1:N}.$ At forecast steps, each particle
${x}_{t}^{n}$ is propagated using the state transition equation Eq. 2.1, while
at analysis steps a Kalman-type update is performed for each particle:

(forecast step)
$\displaystyle\quad\quad\widehat{x}_{t}^{\,n}=F_{\alpha}({x}_{t-1}^{n})+{\xi}_{t}^{n},$
$\displaystyle\quad\quad{\xi}_{t}^{n}\overset{\text{i.i.d.}}{\sim}\mathcal{N}(0,Q_{\beta}),$
(3.1) (analysis step)
$\displaystyle\quad\quad{x}_{t}^{n}=\widehat{x}_{t}^{\,n}+\widehat{K}_{t}(y_{t}+{\gamma}_{t}^{n}-H\widehat{x}_{t}^{\,n}),$
$\displaystyle\quad\quad{\gamma}_{t}^{n}\overset{\text{i.i.d.}}{\sim}\mathcal{N}(0,R).$
(3.2)

Note that the particles $x_{0:T}^{1:N}$ depend on $\theta,$ and (3.1)-(3.2)
implicitly define a map $\theta\mapsto x_{0:T}^{1:N}.$ The Kalman gain
$\widehat{K}_{t}:=\widehat{C}_{t}H^{\top}(H\widehat{C}_{t}H^{\top}+R)^{-1}$ is
defined using the empirical covariance $\widehat{C}_{t}$ of the forecast
ensemble ${\widehat{x}}_{t}^{1:N},$ namely

$\widehat{C}_{t}=\frac{1}{N-1}\sum_{n=1}^{N}(\widehat{x}_{t}^{\,n}-\widehat{m}_{t})(\widehat{x}_{t}^{\,n}-\widehat{m}_{t})^{\top},\quad\text{where}\quad\widehat{m}_{t}=\frac{1}{N}\sum_{n=1}^{N}\widehat{x}_{t}^{\,n}.$
(3.3)

These empirical moments provide a Gaussian approximation to the _forecast
distribution_

$p_{\theta}(x_{t}|y_{1:t-1})\approx\mathcal{N}(\widehat{m}_{t},\widehat{C}_{t}).$
(3.4)

Several implementations of EnKF are available, but for concreteness we only
consider the “perturbed observation” EnKF defined in (3.1)-(3.2). In the
analysis step (3.2), the observation $y_{t}$ is perturbed to form
$y_{t}+{\gamma}_{t}^{n}.$ This perturbation ensures that in linear-Gaussian
models the empirical mean and covariance of ${x}_{t}^{1:N}$ converges as
$N\rightarrow\infty$ to the mean and covariance of the filtering distribution
[53, 50].

###  3.2 Estimation of the Log-Likelihood and its Gradient

Note from (2.9) that in order to approximate $\mathcal{L}(\theta)=\log
p_{\theta}(y_{1:T})$, it suffices to approximate $p_{\theta}(y_{t}|y_{1:t-1})$
for $1\leq t\leq T$. Now, using (2.8) and the EnKF approximation (3.4) to the
forecast distribution, we obtain

$p_{\theta}(y_{t}|y_{1:t-1})\approx\int\mathcal{N}(y_{t};Hx_{t},R)\mathcal{N}(x_{t};\widehat{m}_{t},\widehat{C}_{t})\text{d}x_{t}=\mathcal{N}\big{(}y_{t};H\widehat{m}_{t},H\widehat{C}_{t}H^{\top}+R\big{)}.$
(3.5)

Therefore, we have the following estimate of the data log-likelihood:

$\mathcal{L}_{{\footnotesize\text{EnKF}}}(\theta):=\sum_{t=1}^{T}\log\mathcal{N}\big{(}y_{t};H\widehat{m}_{t},H\widehat{C}_{t}H^{\top}+R\big{)}\approx\mathcal{L}(\theta).$
(3.6)

Notice that the forecast empirical moments
$\\{\widehat{m}_{t},\widehat{C}_{t}\\}_{t=1}^{T}$, and hence
$\mathcal{L}_{{\footnotesize\text{EnKF}}}(\theta),$ depend on $\theta$ in two
distinct ways. First, each forecast particle $\widehat{x}_{t}^{\,n}$ in (3.1)
depends on a particle ${x}_{t}^{n},$ which indirectly depends on $\theta.$
Second, each forecast particle depends on $\theta=\\{\alpha,\beta\\}$ directly
through $F_{\alpha}$ and $Q_{\beta}$. The estimate
$\mathcal{L}_{{\footnotesize\text{EnKF}}}(\theta)$ can be computed online with
EnKF. The whole procedure is summarized in Algorithm 3.1, which implicitly
defines a map $\theta\mapsto\mathcal{L}_{{\footnotesize\text{EnKF}}}(\theta).$
Before discussing the autodiff of this map and learning of the parameter
$\theta$ in Section 4, we establish the large ensemble convergence of
$\mathcal{L}_{{\footnotesize\text{EnKF}}}(\theta)$ and
$\nabla_{\theta}\mathcal{L}_{{\footnotesize\text{EnKF}}}(\theta)$ towards
$\mathcal{L}(\theta)$ and $\nabla_{\theta}\mathcal{L}(\theta)$ in a linear
setting.

Algorithm 3.1 Ensemble Kalman Filter and Log-likelihood Estimation

1:Input: $\theta=\\{\alpha,\beta\\},y_{1:T},x_{0}^{1:N}.$ (If $x_{0}^{1:N}$ is
not specified, draw $x_{0}^{n}\overset{\text{i.i.d.}}{\sim}p_{0}(x_{0})$.)

2:Initialize $\mathcal{L}_{{\footnotesize\text{EnKF}}}(\theta)=0.$

3:for $t=1,\ldots,T$ do

4: Set $\widehat{x}_{t}^{\,n}=F_{\alpha}({x}_{t-1}^{n})+{\xi}_{t}^{n}$, where
${\xi}_{t}^{n}\overset{\text{i.i.d.}}{\sim}\mathcal{N}(0,Q_{\beta})$.
$\triangleright$ Forecast step

5: Compute $\widehat{m}_{t},\widehat{C}_{t}$ by Eq. 3.3 and set
$\widehat{K}_{t}=\widehat{C}_{t}H^{\top}(H\widehat{C}_{t}H^{\top}+R)^{-1}.$

6: Set
${x}_{t}^{n}=\widehat{x}_{t}^{\,n}+\widehat{K}_{t}(y_{t}+{\gamma}_{t}^{n}-H\widehat{x}_{t}^{\,n})$,
where ${\gamma}_{t}^{n}\overset{\text{i.i.d.}}{\sim}\mathcal{N}(0,R)$.
$\triangleright$ Analysis step

7: Set
$\mathcal{L}_{{\footnotesize\text{EnKF}}}(\theta)\leftarrow\mathcal{L}_{{\footnotesize\text{EnKF}}}(\theta)+\log\mathcal{N}\big{(}y_{t};H\widehat{m}_{t},H\widehat{C}_{t}H^{\top}+R\big{)}.$

8:end for

9:Output: EnKF particles $x_{0:T}^{1:N}$. Log-likelihood estimate
$\mathcal{L}_{{\footnotesize\text{EnKF}}}(\theta)$.

###  3.3 Large Sample Convergence: Linear Setting

In this section we consider a linear setting and provide large $N$ convergence
results for the log-likelihood estimate
$\mathcal{L}_{{\footnotesize\text{EnKF}}}(\theta)$ and its gradient
$\nabla_{\theta}\mathcal{L}_{{\footnotesize\text{EnKF}}}(\theta)$ towards
$\mathcal{L}(\theta)$ and $\nabla_{\theta}\mathcal{L}(\theta)$ for any given
$\theta$, for a fixed data sequence $y_{1:T}$. The mappings $\mathcal{L}$ and
$\mathcal{L}_{{\footnotesize\text{EnKF}}}$ are defined in Eq. 2.9 and Eq. 3.6,
respectively. For notation convenience, we drop $\theta$ in the function
argument since the main dependence will be on $N$ in this section. Similar to
[53, 46], we study $L^{p}$ convergence for any $p\geq 1$.

######  Theorem 3.1.

Assume that the state transition Eq. 2.1 is linear, i.e.,

$x_{t}=A_{\alpha}x_{t-1}+\xi_{t},\quad\quad\xi_{t}\sim\mathcal{N}(0,Q_{\beta}),\quad\quad
A_{\alpha}\in\mathbbm{R}^{d_{x}\times d_{x}},$ (3.7)

and that the initial distribution $p_{0}$ is Gaussian. Then, for any
$\theta=\\{\alpha,\beta\\}$ and for any $p\geq 1,$
$\mathcal{L}_{{\footnotesize\text{EnKF}}}$ converges to $\mathcal{L}$ in
$L^{p}$ with rate $1/\sqrt{N}$, i.e.,

$\bigl{(}\operatorname{\mathbb{E}}\big{|}\mathcal{L}_{{\footnotesize\text{EnKF}}}-\mathcal{L}\big{|}^{p}\bigr{)}^{1/p}\leq
cN^{-1/2},$ (3.8)

where $c$ does not depend on $N$.

The linearity of the flow $F_{\alpha}(\cdot)$ is equivalent to the linearity
of the vector field $f_{\alpha}(\cdot)$. Although the convergence of EnKF to
the KF in linear settings has been studied in DA [53, 50, 46, 20] and in
filtering approaches to inverse problems [77, 14], there are no existing
convergence results for EnKF log-likelihood estimation. An exception is [42],
where the authors provide a heuristic argument for convergence in the case
$T=1$. Most of the theoretical analysis of EnKF is based on the _propagation
of chaos_ statement [57, 83]: EnKF defines an interacting particle system,
where the interaction is through the empirical mean $\widehat{m}_{t}$ and
covariance matrix $\widehat{C}_{t}$ of the forecast ensemble
${\widehat{x}}_{t}^{1:N}$. As $N\rightarrow\infty$, one hopes that these
empirical moments can be replaced by their deterministic limits, and that the
particles will hence evolve independently. The large $N$ limits of
$\widehat{m}_{t},\widehat{C}_{t}$ turn out to be the mean and covariance
matrix of the KF forecast distribution. We will leave the construction of the
propagation of chaos statement as well as the proof of Theorem 3.1 to Appendix
A.

Since this paper focuses on gradient based approaches to the learning of
$\theta=\\{\alpha,\beta\\}$, it is thus interesting to compare the gradient
$\nabla_{\theta}\mathcal{L}_{{\footnotesize\text{EnKF}}}$ to the true gradient
$\nabla_{\theta}\mathcal{L}$, as $N\rightarrow\infty$, if both of them exist.
The intuition is that if
$\nabla_{\theta}\mathcal{L}_{{\footnotesize\text{EnKF}}}$ is an accurate
estimate of $\nabla_{\theta}\mathcal{L}$, then one can perform gradient-based
optimization over $\mathcal{L}_{{\footnotesize\text{EnKF}}}$ as if one was
directly optimizing over the true log-likelihood $\mathcal{L}$. For the
gradient w.r.t. $\beta$ to be well-defined, we write
$S_{\beta}=Q_{\beta}^{1/2}$ in the following statement so that $\beta$ does
not appear in the stochasticity of the algorithm. This is also known as the
“reparameterization trick,” which will be discussed later in Section 4.1.

######  Theorem 3.2.

Assume that the state transition Eq. 2.1 is linear, i.e.,

$x_{t}=A_{\alpha}x_{t-1}+S_{\beta}\xi_{t},\quad\quad\xi_{t}\sim\mathcal{N}(0,I_{d_{x}}),\quad\quad
A_{\alpha}\in\mathbbm{R}^{d_{x}\times d_{x}},$ (3.9)

and that the initial distribution $p_{0}$ is Gaussian. Assume the
parameterizations $\alpha\mapsto A_{\alpha}$ and $\beta\mapsto S_{\beta}$ are
differentiable. Then, for any $\theta=\\{\alpha,\beta\\}$, both
$\nabla_{\theta}\mathcal{L}_{{\footnotesize\text{EnKF}}}$ and
$\nabla_{\theta}\mathcal{L}$ exist and, for any $p\geq 1,$
$\nabla_{\theta}\mathcal{L}_{{\footnotesize\text{EnKF}}}$ converges to
$\nabla_{\theta}\mathcal{L}$ in $L^{p}$ with rate $1/\sqrt{N}$, i.e.,

$\bigl{(}\operatorname{\mathbb{E}}\big{|}\nabla_{\theta}\mathcal{L}_{{\footnotesize\text{EnKF}}}-\nabla_{\theta}\mathcal{L}\big{|}^{p}\bigr{)}^{1/p}\leq
cN^{-1/2},$ (3.10)

where $c$ does not depend on $N$.

An important observation is that $\theta$ only enters the objective function
$\mathcal{L}_{{\footnotesize\text{EnKF}}}$ through the empirical mean
$\widehat{m}_{t}$ and covariance matrix $\widehat{C}_{t}$ of the forecast
ensemble. As $N\rightarrow\infty$, one hopes that these empirical moments can
be replaced by their deterministic limits, and gradients based on these
empirical moments can be replaced by gradients based on their deterministic
limits. The gradients taken in the limits turn out to be those of the true
log-likelihood $\mathcal{L}$. Again, the proof relies on the propagation of
chaos statement and is left to Appendix B.

##  4 Auto-differentiable Ensemble Kalman Filters

This section contains our main methodological contributions. We introduce our
AD-EnKF framework in Section 4.1. We then describe in Section 4.2 how to
handle long observation data, i.e., large $T$, using TBPTT. In Section 4.3, we
highlight how various techniques introduced for EnKF in the DA community,
e.g., localization and covariance inflation, can be incorporated into our
framework. Finally, Section 4.4 discusses the computational and memory costs.

###  4.1 Main Algorithm

Algorithm 4.1 Auto-differentiable Ensemble Kalman Filter (AD-EnKF)

1:Input: Observations $y_{1:T}$. Learning rate $\eta$.

2:Initialize SSM parameter $\theta^{0}$ and set $k=0.$

3:while not converging do

4:
$x_{0:T}^{1:N},\mathcal{L}_{{\footnotesize\text{EnKF}}}(\theta^{k})=\textsc{EnsembleKalmanFilter}(\theta^{k},y_{1:T}).$
$\triangleright$ Alg. 3.1

5: Compute
$\nabla_{\theta}\mathcal{L}_{{\footnotesize\text{EnKF}}}(\theta^{k})$ by auto-
differentiating the map
$\theta^{k}\mapsto\mathcal{L}_{{\footnotesize\text{EnKF}}}(\theta^{k}).$

6: Set
$\theta^{k+1}=\theta^{k}+\eta\nabla_{\theta}\mathcal{L}_{{\footnotesize\text{EnKF}}}(\theta^{k})$
and $k\leftarrow k+1.$

7:end while

8:Output: Learned SSM parameter $\theta^{k}$ and EnKF particles
$x_{0:T}^{1:N}$.

Our core method is shown in Alg. 4.1, and our PyTorch implementation is at
https://github.com/ymchen0/torchEnKF. The gradient of the map
$\theta^{k}\mapsto\mathcal{L}_{{\footnotesize\text{EnKF}}}(\theta^{k})$ can be
evaluated using autodiff libraries [64, 11, 1]. More specifically, reverse-
mode autodiff can be performed for common matrix operations like matrix
multiplication, inverse, and determinant [32]. We use the “reparameterization
trick” [44, 69] to auto-differentiate through the stochasticity in the EnKF
algorithm. Specifically, in Alg. 3.1 line 4, we draw ${\xi}_{t}^{n}$ from a
distribution $\mathcal{N}(0,Q_{\beta})$ that involves a parameter $\beta$ with
respect to which we would like to compute the gradient. For this operation to
be compatible with the autodiff, we reparameterize

$\widehat{x}_{t}^{n}=F_{\alpha}(x_{t}^{n})+{\xi}_{t}^{n}\quad{\xi}_{t}^{n}\overset{\text{i.i.d.}}{\sim}\mathcal{N}(0,Q_{\beta})\quad\Longleftrightarrow\quad\widehat{x}_{t}^{n}=F_{\alpha}(x_{t}^{n})+Q_{\beta}^{1/2}{\xi}_{t}^{n}\quad\xi_{t}^{n}\overset{\text{i.i.d.}}{\sim}\mathcal{N}(0,I_{d_{x}}),$
(4.1)

so that the gradient with respect to $\beta$ admits an unbiased estimate. In
contrast to the EnKF, the resampling step of PFs cannot be readily auto-
differentiated [59, 55, 51].

###  4.2 Truncated Gradients for Long Sequences

If the sequence length $T$ is large, although
$\mathcal{L}_{{\footnotesize\text{EnKF}}}(\theta)$ and its gradient
$\nabla_{\theta}\mathcal{L}_{{\footnotesize\text{EnKF}}}(\theta)$ can be
evaluated using the aforementioned techniques, the practical value of Alg. 4.1
is limited for two reasons. First, computing these quantities requires a full
filtering pass of the data, which may be computationally costly. Moreover, for
the gradient ascent methods to achieve a good convergence rate, multiple
evaluations of gradients are often needed, requiring an equally large number
of filtering passes. The second reason is that, like recurrent networks, Alg.
4.1 may suffer from exploding or vanishing gradients [63] as the derivatives
are multiplied together using chain rules in the backpropagation.

Our proposed technique can address both of these issues by borrowing the ideas
of TBPTT from the recurrent neural network literature [90, 82] and the
recursive maximum likelihood method from the hidden Markov models literature
[52]. The idea is to divide the sequence into subsequences of length $L$.
Instead of computing the log-likelihood of the whole sequence and then
backpropagating, one computes the log-likelihood of each subsequence and
backpropagates within that subsequence. The subsequences are processed
sequentially, and the EnKF output of the previous subsequence (i.e., the
location of particles) are used as the input to the next subsequence. In this
way, one performs $\lceil T/L\rceil$ gradient updates in a _single_ filtering
pass, and since the gradients are backpropagated across a time span of length
at most $L$, gradient explosion/vanishing is more unlikely to happen. This
approach is detailed in Alg. 4.2.

Algorithm 4.2 AD-EnKF with Truncated Backprop (AD-EnKF-T)

1:Input: Observations $y_{1:T}$. Learning rate $\eta$. Subsequence length $L$.

2:Initialize SSM parameter $\theta^{0}$ and set $k=0.$

3:while not converging do

4: Set $x_{0}^{n}\overset{\text{i.i.d.}}{\sim}p_{0}(x_{0})$.

5: for $j=0,\ldots,T/L-1$ do

6: Set $t_{0}=jL$, $t_{1}=\min\\{(j+1)L,T\\}$.

7:
$x_{t_{0}:t_{1}}^{1:N},\mathcal{L}_{{\footnotesize\text{EnKF}}}(\theta^{k})=\textsc{EnsembleKalmanFilter}(\theta^{k},y_{(t_{0}+1):t_{1}},x_{t_{0}}^{1:N}).$
$\triangleright$ Alg. 3.1

8: Set
$\theta^{k+1}=\theta^{k}+\eta\nabla_{\theta}\mathcal{L}_{{\footnotesize\text{EnKF}}}(\theta^{k})$
and $k\leftarrow k+1.$

9: end for

10:end while

11:Output: Learned SSM parameter $\theta^{k}$ and EnKF particles
$x_{0:T}^{1:N}$.

###  4.3 Localization for High State Dimensions

In practice, the state often represents a physical quantity that is
discretized in spatial coordinates (e.g., numerical solution to a time-
evolving PDE), which leads to a high state dimension $d_{x}$. In order to
reduce the computational and memory complexity, EnKF is often run with
$N<d_{x}$. A small ensemble size $N$ causes rank deficiency of the forecast
sample covariance $\widehat{C}_{t}$, which may cause spurious correlations
between spatial coordinates that are far apart. In other words, for $(i,j)$
such that $|i-j|$ is large, the $(i,j)$-th coordinate of $\widehat{C}_{t}$ may
not be close to 0, although one would expect it to be small since it
represents the correlation between spatial locations that are far apart. This
problem can be addressed using localization techniques, and we shall focus on
_covariance tapering_ [36]. The idea is to “taper” the forecast sample
covariance matrix $\widehat{C}_{t}$ so that the nonzero spurious correlations
are zeroed out. This method is implemented defining a $d_{x}\times d_{x}$
matrix $\rho$ with 1’s on the diagonal and entries smoothly decaying to $0$
off the diagonal, and replacing the forecast sample covariance matrix
$\widehat{C}_{t}$ in Alg. 3.1 by $\rho\circ\widehat{C}_{t}$, where $\circ$
denotes the element-wise matrix product. Common choices of $\rho$ were
introduced in [31]. Covariance tapering can be easily adopted within our AD-
EnKF framework. We find that covariance tapering not only stabilizes the
filtering procedure, which had been noted before, e.g., [37, 34], but it also
helps to obtain low-variance estimates of the log-likelihood and its gradient
— see the discussion in Section 5.1.2. Localization techniques relying on
local serial updating of the state [37, 61, 72] could also be considered.

Another useful tool for EnKF with $N<d_{x}$ is _covariance inflation_ [4],
which prevents the ensemble from collapsing towards its mean after the
analysis update [30]. In practice, this can be performed by replacing the
forecast sample covariance matrix $\widehat{C}_{t}$ in Alg. 3.1 by
$(1+\zeta)\widehat{C}_{t}$, where $\zeta>0$ is a small constant that needs to
be tuned. Although not considered in our experiments, covariance inflation can
also be easily adopted within our AD-EnKF framework.

###  4.4 Computation and Memory Costs

Autodifferentiation of the map
$\theta^{k}\mapsto\mathcal{L}_{{\footnotesize\text{EnKF}}}(\theta^{k})$ in
Alg. 4.1 does not introduce an extra order of computational cost compared to
the evaluation of this map alone. Thus, the computational cost of AD-EnKF is
at the same order as that of a standard EnKF. The computation cost of EnKF can
be found in, e.g., [70]. Moreover, AD-EnKF can be parallelized and speeded up
with a GPU.

Like a standard EnKF, when no covariance tapering is applied, AD-EnKF has
$O(Nd_{x})$ memory cost since it does not explicitly compute the sample
covariance matrix $\widehat{C}_{t}$111Note that
$\widehat{C}_{t}H^{\top}=\frac{1}{N-1}\sum_{n=1}^{N}(\widehat{x}_{t}^{\,n}-\widehat{m}_{t})(H\widehat{x}_{t}^{\,n}-H\widehat{m}_{t})^{\top}$
and
$H\widehat{C}_{t}H^{\top}=\frac{1}{N-1}\sum_{n=1}^{N}(H\widehat{x}_{t}^{\,n}-H\widehat{m}_{t})(H\widehat{x}_{t}^{\,n}-H\widehat{m}_{t})^{\top}$,
which require $O(d_{x}\max\\{N,d_{y}\\})$ and $O(d_{y}\max\\{N,d_{y}\\})$
memory respectively. Both of them are less than $O(d_{x}^{2})$ if $d_{y}\ll
d_{x}$ and $N\ll d_{x}$. . With covariance tapering, the memory cost is at
most $O(\max\\{N,r\\}d_{x})$, where $r$ is the tapering radius, if the
tapering matrix $\rho$ is sparse with $O(rd_{x})$ nonzero entries. This
sparsity condition is satisfied when using common tapering matrices [31]. In
terms of the time dimension, the memory cost of AD-EnKF can be reduced from
$O(T)$ to $O(L)$ with the TBPTT in Section 4.2. Unlike previous work on EM-
based approaches [12, 9, 65], where the locations of all particles
$x_{t}^{1:N}$ across the whole time span of $T$ need to be stored, AD-EnKF-T
only requires to store the particles within a time span of $L$ to perform a
gradient step.

If the transition map $F_{\alpha}$ is defined by the flow map of an ODE with
vector field $f_{\alpha}$, we can use adjoint methods to differentiate
efficiently through $F_{\alpha}$ in the forecast step Eq. 3.1. Use of the
adjoint method is facilitated by NeuralODE autodiff libraries [15] that have
become an important tool to learn continuous-time dynamical systems [6, 71,
18]. Instead of discretizing $F_{\alpha}$ with a numerical solver applied to
$f_{\alpha}$ and differentiating through solver’s steps as in [12, 9], we
directly differentiate through $F_{\alpha}$ by solving an adjoint differential
equation, which does not require us to store all intermediate steps from the
numerical solver, reducing the memory cost. More details can be found in [15],
and the PyTorch package provided by the authors can be incoporated within our
AD-EnKF framework with minimal effort.

##  5 Numerical Experiments

###  5.1 Linear-Gaussian Model

In this section, we focus on parameter estimation in a linear-Gaussian model
with a banded structure on model dynamic and model error covariance matrix.
This experiment falls into the category of “parameterized dynamics” in Section
2.1. We first illustrate the convergence results of the log-likelihood
estimate $\mathcal{L}_{{\footnotesize\text{EnKF}}}$ and gradient estimate
$\nabla_{\theta}\mathcal{L}_{{\footnotesize\text{EnKF}}}$ presented in Section
3.3, since the true values $\mathcal{L}$ and $\nabla_{\theta}\mathcal{L}$ are
available in closed form. We also show that the localization techniques
described in Section 4.3 lead to a more accurate estimate when the ensemble
size is small. Finally, we show that having a more accurate estimate,
especially for the gradient, improves the parameter estimation.

We compare the EnKF to PF methods. Similar to the EnKF, the PF also provides
an estimate of the log-likelihood and its gradient. Different from [59, 55,
51], we adopt the PF with optimal proposal [24] as it is implementable for the
family of SSMs considered in this paper [23, 75], and we find it to be more
stable than separately training a variational proposal. To compute the log-
likelihood gradient for the PF, we follow the same strategy as in [59, 55, 51]
and do not differentiate through the resampling step. The full algorithm,
which we abbreviate as AD-PF, is presented in Appendix D.

We consider the following SSM, similar to [91, 80]

$\displaystyle x_{t}$ $\displaystyle=A_{\alpha}x_{t-1}+\xi_{t},\quad\quad$
$\displaystyle\xi_{t}\sim\mathcal{N}(0,Q_{\beta}),\quad\quad$ $\displaystyle
1\leq t\leq T,$ (5.1a) $\displaystyle y_{t}$
$\displaystyle=Hx_{t}+\eta_{t},\quad\quad$
$\displaystyle\eta_{t}\sim\mathcal{N}(0,0.5I_{d_{y}}),\quad\quad$
$\displaystyle 1\leq t\leq T,$ (5.1b) $\displaystyle x_{0}$
$\displaystyle\sim\mathcal{N}(0,4I_{d_{x}}),$ (5.1c) where
$A_{\alpha}=\begin{bmatrix}\alpha_{1}&\alpha_{2}&&0\\\
\alpha_{3}&\alpha_{1}&\ddots&\\\ &\ddots&\ddots&\alpha_{2}\\\
0&&\alpha_{3}&\alpha_{1}\\\
\end{bmatrix},\quad\quad[Q_{\beta}]_{i,j}=\beta_{1}\exp(-\beta_{2}|i-j|).$
(5.1d)

Here $[Q_{\beta}]_{i,j}$ denotes the $(i,j)$-th entry of $Q_{\beta}$.
Intuitively, $\beta_{1}$ controls the scale of error, while $\beta_{2}$
controls how error is correlated across spatial coordinates. We set
$\alpha=(\alpha_{1},\alpha_{2},\alpha_{3}),$ $\beta=(\beta_{1},\beta_{2}),$
and $\theta=\\{\alpha,\beta\\}$.

####  5.1.1 Estimation Accuracy of $\mathcal{L}_{{\footnotesize\text{EnKF}}}$
and $\nabla_{\theta}\mathcal{L}_{{\footnotesize\text{EnKF}}}$

As detailed above, a key idea proposed in this paper is to estimate
$\mathcal{L}(\theta)$, $\nabla_{\alpha}\mathcal{L}(\theta)$ and
$\nabla_{\beta}\mathcal{L}(\theta)$ with quantities
$\mathcal{L}_{{\footnotesize\text{EnKF}}}(\theta)$,
$\nabla_{\alpha}\mathcal{L}_{{\footnotesize\text{EnKF}}}(\theta)$ and
$\nabla_{\beta}\mathcal{L}_{{\footnotesize\text{EnKF}}}(\theta)$ obtained by
running an EnKF and differentiating through its computations using autodiff.
Since these estimates will be used by AD-EnKF to perform gradient ascent, it
is critical to assess their accuracy. We do so in this section for a range of
values of $\theta.$

We first simulate observation data $y_{1:T}$ from the true model with
$d_{x}=d_{y}\in\\{20,40,80\\}$, $T=10$, $H=I_{d_{x}}$,
$\alpha^{*}=(0.3,0.6,0.1)$ and $\beta^{*}=(0.5,1)$. Given data $y_{1:T}$, the
true data log-likelihood $\mathcal{L}(\theta)=p_{\theta}(y_{1:T})$ and
gradient $\nabla_{\theta}\mathcal{L}(\theta)$, which can be decomposed into
$\nabla_{\alpha}\mathcal{L}(\theta)$, $\nabla_{\beta}\mathcal{L}(\theta)$, can
be computed analytically. We perform $P=50$ EnKF runs, and report a Monte
Carlo estimate of the relative $L^{2}$ errors of the log-likelihood and
gradient estimates (see Appendix F for their definition) as the ensemble size
$N$ increases. Fig. 2 shows the results when $\theta$ is evaluated at the true
parameters $\\{\alpha^{*},\beta^{*}\\}$. Intuitively, this $\theta$ is close
to optimal since it is the one that generates the data. We also show in Fig. 9
in Appendix C the results when $\theta$ is evaluated at a parameter that is
not close to optimal: $\alpha=(0.5,0.5,0.5),\beta=(1,0.1)$. Both figures
illustrate that the relative $L^{2}$ estimation errors of the log-likelihood
and its gradient computed using EnKF converge to zero at a rate of
approximately $N^{-1/2}$. Moreover, the state dimension $d_{x}$ has small
empirical effect on the convergence rate. On the other hand, those computed
using PF have a slower convergence rate or barely converge, especially for the
gradient (see the third plot in Fig. 9). We recall that the resampling parts
are discarded from the autodiff of PFs, which introduces a bias. Moreover, the
empirical convergence rate is slightly slower in higher state dimensions.
Comparing the estimation error of EnKF and PF under the same $d_{x}$ choice,
we find that when the number of particles is large ($>500$), EnKF gives a more
accurate estimate than PF. However, when the number of particles is small,
EnKF is less accurate, but we will show in the next section how the EnKF
results can be significantly improved using localization techniques.

Figure 2: Relative $L^{2}$ estimation errors of the log-likelihood (left) and
its gradient w.r.t. $\alpha$ (middle) and $\beta$ (right), computed using EnKF
and PF, as a function of $N$, for the linear-Gaussian model (5.1). State
dimension $d_{x}\in\\{20,40,80\\}$. $\theta$ is evaluated at the true
parameters $\\{\alpha^{*},\beta^{*}\\}$. (Section 5.1.1)

####  5.1.2 Effect of Localization

In practice, for computational and memory concerns, the number of particles
used for EnKF is typically small ($<100$), and hence it is necessary to get an
accurate estimate of log-likelihood and its gradients using a small number of
particles. We use the covariance tapering techniques discussed in Section 4.3,
where $\widehat{C}_{t}$ is replaced by $\rho\circ\widehat{C}_{t}$ in Alg. 3.1,
and $\rho$ is defined using the fifth order piecewise polynomial correlation
function of Gasperi and Cohn [31]. The detailed construction of $\rho$ is left
to Appendix F, with a hyperparameter $r$ that controls the tapering radius.

Fig. 3 shows the estimation results when the state dimension is set to be
$d_{x}=80$ and $\theta$ is evaluated at $(\alpha^{*},\beta^{*})$, while
different tapering radii $r$ are applied. The plots of EnKF with no tapering
and the plots of PF are the same as in Fig. 2. We find that covariance
tapering can reduce the estimation error of the log-likelihood and its
gradient when the number of particles is small. Moreover, having a smaller
tapering radius leads to a better estimation when the number of particles is
small. As the number of particles grows larger, covariance tapering may worsen
the estimation of both log-likelihood and its gradient. This is because the
sampling error and spurious correlation that occurs in the sample covariance
matrix in EnKF will be overcome by large number of particles, and hence
covariance tapering will only act as a modification to the objective function
$\mathcal{L}_{{\footnotesize\text{EnKF}}}$, leading to inconsistent estimates.
However, there is no reason for using localization when one can afford a large
number of particles. When computational constraints require fewer particles
than state dimension, we find that covariance tapering is not only beneficial
to the parameter estimation problems but is also beneficial to learning of the
dynamics in high dimensions, as we will show in later sections. Results when
$\theta$ is evaluated at parameters that are not optimal
($\alpha=(0.5,0.5,0.5)$, $\beta=(1,0.1)$) are shown in Fig. 10 in Appendix C,
where the beneficial effect of tapering is evident.

Figure 3: Relative $L^{2}$ estimation errors of log-likelihood (left) and its
gradient w.r.t. $\alpha$ (middle) and $\beta$ (right), computed using EnKF and
PF, with different covariance tapering radius applied to EnKF for the linear-
Gaussian model (5.1). State dimension $d_{x}=80$. $\theta$ is evaluated at the
true parameters $\\{\alpha^{*},\beta^{*}\\}$. (Section 5.1.2)

####  5.1.3 Parameter Learning

Here we illustrate how the estimation accuracy of the log-likelihood and its
gradient, especially the latter, affect the parameter learning with AD-EnKF.
Since our framework relies on gradient-based learning of parameters,
intuitively, the less biased the gradient estimate is, the closer our learned
parameter will be to the true MLE solution.

We first consider the setting where the state dimension is set to be
$d_{x}=80$. We run AD-EnKF for 1000 iterations with gradient ascent under the
following choices of ensemble size and tapering radius: (1) $N=1000$ with no
tapering; (2) $N=50$ with no tapering; and (3) $N=50$ with tapering radius 5.
We also run AD-PF with $N=1000$ particles. Throughout, one “training
iteration” corresponds to processing once the whole data sequence. Additional
implementation details are available in the appendices. Fig. 4 and 5 show a
single run of parameter learning under each setting, where we include for
reference the MLE obtained by running gradient ascent until convergence with
the true gradient $\nabla_{\theta}\mathcal{L}$ (denoted with the red dashed
line). The objective function, i.e., the likelihood estimates
$\mathcal{L}_{{\footnotesize\text{EnKF}}}$ and
$\mathcal{L}_{{\footnotesize\text{PF}}}$ are also plotted as a function of
training iterations. Results with other choices of state dimension $d_{x}$ are
summarized in Table 5.1, where we take the values of $\alpha$ at the final
iteration and compute their distance to the true MLE solution. The procedure
is repeated 10 times, and the mean and standard deviations are reported. The
results all show a similar trend: AD-EnKF with $N=1000$ particles performs the
best (small errors and small fluctuations) for all settings, while AD-EnKF
with $N=50$ particles and covariance tapering performs second best. AD-EnKF
with $N=50$ without covariance tapering comes at the third place, and AD-PF
method performs the worst, indicating the superiority of AD-EnKF method to the
AD-PF method for high-dimensional linear-Gaussian models of the form (5.1) and
the utility of localization techniques. Importantly, the findings here are
consistent with the plots in Fig. 3. This behavior is in agreement with the
intuition that the estimation accuracy of the log-likelihood gradient
determines the parameter learning performance.

Figure 4: Learned parameter $\alpha$ as a function of training iterations for the linear-Gaussian model (5.1). State dimension $d_{x}=80$. Red dashed lines are the MLE solutions to the true data log-likelihood $\mathcal{L}$. (Section 5.1.3) Figure 5: Learned parameter $\beta$, and training objective $\mathcal{L}_{{\footnotesize\text{EnKF}}}$, $\mathcal{L}_{{\footnotesize\text{PF}}}$ as a function of training iterations for the linear-Gaussian model (5.1). Red dashed lines are the MLE solutions to the true data log-likelihood $\mathcal{L}$ (left and middle), and the maximum value attained by $\mathcal{L}$ (right). (Section 5.1.3) |  | $d_{x}=20$
---
$N=50$
| $d_{x}=20$
---
$N=1000$
| $d_{x}=40$
---
$N=50$
| $d_{x}=40$
---
$N=1000$
| $d_{x}=80$
---
$N=50$
| $d_{x}=80$
---
$N=1000$
AD-EnKF (no taper) |  1.65 $\pm$0.30 |  0.07 $\pm$0.06 |  4.12$\pm$0.73 |  0.17$\pm$0.09 |  4.14$\pm$0.67 |  0.20$\pm$0.14
AD-EnKF(taper=5) |  0.53$\pm$0.18 | $-$ |  0.35$\pm$0.27 | $-$ |  1.05$\pm$0.38 | $-$
AD-PF |  7.75$\pm$0.37 |  3.51$\pm$0.35 |  8.58$\pm$0.25 |  5.59$\pm$0.31 |  9.28$\pm$0.49 |  6.77$\pm$0.24
Table 5.1: Euclidean distance ($\times 10^{-2}$) from the learned parameter
$\alpha$ at the final iteration to the true MLE solution, under varying
dimensional settings for the linear-Gaussian model (5.1) (Section 5.1.3).

###  5.2 Lorenz-96

In this section, we illustrate our AD-EnKF framework in the three types of
learning problems mentioned in Section 2.1: parameterized dynamics, fully-
unknown dynamics, and model correction. We will compare our method to AD-PF,
as in Section 5.1. We will also compare our method to the EM-EnKF method
implemented in [9, 12], which we abbreviate as EM, and is detailed in Appendix
E. We emphasize that the gradients computed in the EM are different from the
ones computed in AD-EnKF, and in particular do not auto-differentiate through
the EnKF.

The reference Lorenz-96 model [54] is defined by (2.5) with vector field

$f^{*(i)}(x)=-x^{(i-1)}(x^{(i-2)}-x^{(i+1)})-x^{(i)}+8,\quad\quad 0\leq i\leq
d_{x}-1,$ (5.2)

where $x^{(i)}$ and $f^{*(i)}$ are the $i$-th coordinate of $x$ and component
of $f^{*}$. By convention $x^{(-1)}:=x^{(d_{x}-1)},x^{(-2)}:=x^{(d_{x}-2)}$
and $x^{(d_{x})}:=x^{(0)}$. We assume there is no noise in the reference state
transition model, i.e., $Q^{*}=0$. The goal is to recover the reference state
transition model with
$p_{\theta}(x_{t}|x_{t-1})=\mathcal{N}(x_{t};F_{\alpha}\big{(}x_{t-1}),Q_{\beta}\big{)}$
from the data $y_{1:T},$ where $F_{\alpha}$ is the flow map of a vector field
$f_{\alpha}$, and then recover the states $x_{1:T}$. The parameterized error
covariance $Q_{\beta}$ in the transition model is assumed to be diagonal,
i.e., $Q_{\beta}=\text{diag}(\beta)$ with $\beta\in\mathbbm{R}^{d_{x}}$. The
parameterized vector field $f_{\alpha}$ is defined differently for the three
types of learning problems, as we lay out below. We quantify performance using
the forecast error (RMSE-f), the analysis/filter error (RMSE-a), and the test
log-likelihood. These metrics are defined in Appendix F.

####  5.2.1 Parameterized Dynamics

We consider the same setting as in [8], where

$\begin{split}f_{\alpha}^{(i)}(x)=\big{[}&1,x^{(i-2)},x^{(i-1)},x^{(i)},x^{(i+1)},x^{(i+2)},\\\
&\big{(}x^{(i-2)}\big{)}^{2},\big{(}x^{(i-1)}\big{)}^{2},\big{(}x^{(i)}\big{)}^{2},\big{(}x^{(i+1)}\big{)}^{2},\big{(}x^{(i+2)}\big{)}^{2},\\\
&x^{(i-2)}x^{(i-1)},x^{(i-1)}x^{(i)},x^{(i)}x^{(i+1)},x^{(i+1)}x^{(i+2)},\\\
&x^{(i-2)}x^{(i)},x^{(i-1)}x^{(i+1)},x^{(i)}x^{(i+2)}\big{]}^{\top}\alpha,\quad\quad
0\leq i\leq d_{x}-1,\end{split}$ (5.3)

and $\alpha\in\mathbbm{R}^{18}$ is interpreted as the coefficients of some
“basis polynomials” representing the governing equation of the underlying
system. The parameterized governing equation of the $i$-th coordinate depends
on its $N_{1}=5$ neighboring coordinates, and the second order polynomials
only involve interactions between coordinates that are at most $N_{2}=2$
indices apart. The reference ODE Eq. 5.2 satisfies $f^{*}=f_{\alpha^{*}}$,
where $\alpha^{*}\in\mathbbm{R}^{18}$ has nonzero entries

$\alpha^{*}_{0}=8,\quad\alpha^{*}_{3}=-1,\quad\alpha^{*}_{11}=-1,\quad\alpha^{*}_{16}=1,$
(5.4)

and zero entries otherwise. Here the dimension of $\theta=\\{\alpha,\beta\\}$
is $d_{\theta}=18+d_{x}$.

We first consider the specific case with $d_{x}=d_{y}=40$, $H=I_{40}$. We set
$R=I_{40}$ and $x_{0}\sim\mathcal{N}(0,50I_{40})$. We generate four sequences
of training data with the reference model for $T=300$ with time between
consecutive observations $\Delta_{s}=0.05$. Both flow maps $F^{*}$ and
$F_{\alpha}$ are integrated using a fourth-order Runge Kutta (RK4) method with
step size $\Delta_{s}^{\text{int}}=0.01$, with adjoint methods implemented for
backpropagation through the ODE solver [15].

We use AD-EnKF-T (Alg. 4.2) with $L=20$ and covariance tapering Eq. F.3 with
radius $r=5$. We compare with AD-PF-T (see Appendix D) with $L=20$ and EM (see
Appendix E). The implementation details, including the choice of learning
rates and other hyperparameters, are discussed in Appendix F.

Comparison of the three algorithms is shown in Fig. 6. Our AD-EnKF-T recovers
$\alpha^{*}$ better than the other two approaches. The EM approach converges
faster, but has a larger error. Moreover, EM tends to converge to a higher
level of learned model error $\sigma_{\beta}$ (defined in Eq. F.4), while our
AD-EnKF-T shows a consistent drop of learned error level. Note that
$Q_{\beta}$ in the learned transition kernel acts like covariance inflation,
which is discussed in Section 4.3, but is “learned” to be adaptive to the
training data rather than manually tuned; therefore, having a nonzero error
level $\sigma_{\beta}$ may still be helpful. The plot of the log-likelihood
estimate during training indicates that AD-EnKF-T searches for parameters with
a higher log-likelihood than the EM approach, which is not surprising as AD-
EnKF-T directly optimizes $\mathcal{L}_{{\footnotesize\text{EnKF}}}$, while EM
does so by alternatively optimizing a surrogate objective. Also, the large
discrepancy between the optimized $\mathcal{L}_{{\footnotesize\text{EnKF}}}$
and $\mathcal{L}_{{\footnotesize\text{PF}}}$ objective may be due to
$\mathcal{L}_{{\footnotesize\text{PF}}}$ being a worse estimate for the true
log-likelihood $\mathcal{L}$ than that of
$\mathcal{L}_{{\footnotesize\text{EnKF}}}$. Note that PFs may not be suitable
for high-dimensional systems like the Lorenz-96 model. Even with knowledge of
the true reference model and a large number of particles, the PF is not able
to capture the filtering distribution well due to the high dimensionality —
see, e.g., Figure 5 of [10].

Figure 6:  Learning parameterized dynamics of Lorenz-96 (5.3), with $d_{x}=40$
and $H=I_{40}$. Learned value of the 18 coefficients of $\alpha$ (upper left
for nonzero entries and upper right for zero entries, where the truth
$\alpha^{*}$ is plotted in red dashed lines), averaged diagnosed error level
$\sigma_{\beta}$ Eq. F.4 (lower left) and log-likelihood
$\mathcal{L}_{{\footnotesize\text{EnKF}}}$/$\mathcal{L}_{{\footnotesize\text{PF}}}$
during training (lower right), as a function of training iterations.
Throughout, shaded area corresponds to $\pm 2$ std over 5 repeated runs.
(Section 5.2.1)

We also consider varying the state dimension $d_{x}$ and observation model
$H$. (The parameterization in Eq. 5.3 is valid for any choice of $d_{x}$.) We
measure the Euclidean distance between the value of learned $\alpha$ at the
final training iteration (at convergence) to $\alpha^{*}$. The training
procedure is repeated 5 times and the results are shown in Table 5.2. We vary
$d_{x}\in\\{10,20,40,80\\}$ and consider two settings for $H$: fully observed
at all coordinates, i.e., $H=I_{d_{x}}$, and partially observed at every two
out of three coordinates [74], i.e.,
$H=[e_{1},e_{2},e_{4},e_{5},e_{7},\cdots]^{\top}$, where
$\\{e_{i}\\}_{i=1}^{d_{x}}$ is the standard basis for $\mathbbm{R}^{d_{x}}$.
The number of particles used for all algorithms is fixed at $N=50$, and
covariance tapering Eq. F.3 with radius $r=5$ is applied to the EnKF. For both
AD-EnKF-T and AD-PF-T, $L=20$. We find that AD-EnKF-T is able to consistently
recover $\alpha^{*}$ regardless of the choice of $d_{x}$ and $H$, and is able
to perform well in the important case where $N<d_{x}$, with an accuracy that
is orders of magnitude better than the other two approaches. The EM approach
is able to recover $\alpha^{*}$ consistently in fully observed settings, but
with a lower accuracy. In partially-observed settings, EM does not converge to
the same value in repeated runs, possibly due to the existence of multiple
local maxima. AD-PF-T is able to converge consistently in fully observed
settings but with the lowest accuracy, and runs into filter divergence issues
in partially-observed settings, so that the training process is not able to
complete. Moreover, we observe that the error of AD-PF-T tends to grow with
the state dimension $d_{x}$, while the two approaches based on EnKF do not
deteriorate when increasing the state dimension. This is further evidence that
EnKF is superior in high-dimensional settings.

|  | $d_{x}=10$
---
(full)
| $d_{x}=20$
---
(full)
| $d_{x}=20$
---
(partial)
| $d_{x}=40$
---
(full)
| $d_{x}=40$
---
(partial)
| $d_{x}=80$
---
(full)
| $d_{x}=80$
---
(partial)
EM |  0.308$\pm$ 0.026 |  0.289$\pm$ 0.0114 |  2.28$\pm$ 4.92 |  0.268$\pm$ 0.0103 |  7.754$\pm$ 8.057 |  0.231$\pm$ 0.0209 |  7.382$\pm$ 4.812
AD-PF-T |  0.262$\pm$ 0.020 |  0.711$\pm$ 0.0291 | $-$ |  1.557$\pm$ 0.0422 | $-$ |  2.079$\pm$ 0.0275 | $-$
AD-EnKF-T | 0.217$\pm$ 0.027 | 0.0325$\pm$ 0.0128 | 0.0835$\pm$ 0.0189 | 0.0283$\pm$ 0.0022 | 0.0930$\pm$ 0.0098 | 0.0540$\pm$ 0.0065 | 0.0813$\pm$ 0.0083

Table 5.2: Lorenz-96, learning parameterized dynamics with varying $d_{x}$ and
observation models. The table shows recovery of learned $\alpha^{*}$ for each
algorithm at the final training iteration, in terms of its distance to the
truth $\alpha^{*}$ Eq. 5.4. “Full” corresponds to full observations, i.e.,
$H=I_{d_{x}}$. “Partial” corresponds to observing 2 out of 3 coordinates,
i.e., $H=[e_{1},e_{2},e_{4},e_{5},e_{7},\cdots]^{\top}$. The “-” indicates
that training cannot be completed due to filter divergence. (Section 5.2.1)

####  5.2.2 Fully Unknown Dynamics

We assume no knowledge of the reference vector field $f^{*}$, and we
approximate it by a neural network surrogate
$f_{\alpha}^{\footnotesize\text{NN}}:\mathbbm{R}^{d_{x}}\rightarrow\mathbbm{R}^{d_{x}},$
where here $\alpha$ represents the NN weights. The structure of the NN is
similar to the one in [12] and is detailed in Appendix F. The number of
parameters combined for $\alpha$ and $\beta$ is $d_{\theta}=9317$. The
experimental results are compared to the model correction results, and hence
are postponed to Section 5.2.3.

####  5.2.3 Model Correction

We assume $f^{*}$ is unknown, but that an inaccurate model $f_{\text{approx}}$
is available. We make use of the parametric form Eq. 5.3, and define
$f_{\text{approx}}$ via a perturbation $\widetilde{\alpha}$ of the true
parameter $\alpha^{*}$:

$f_{\text{approx}}:=f_{\widetilde{\alpha}},\quad\quad\text{where
}\widetilde{\alpha}_{i}\sim\begin{cases}\mathcal{N}(\alpha_{i}^{*},1),\quad&\text{if
}i=0,\\\ \mathcal{N}(\alpha_{i}^{*},0.1),\quad&\text{if
}i\in\\{1,\dots,5\\},\\\ \mathcal{N}(\alpha_{i}^{*},0.01),\quad&\text{if
}i\in\\{6,\dots,17\\}.\end{cases}$ (5.5)

The coefficients of a higher order polynomial have a smaller amount of
perturbation. $\widetilde{\alpha}$ is _fixed_ throughout the learning
procedure. We approximate the residual $f^{*}-f_{\text{approx}}$ by a NN
$g_{\alpha}^{\footnotesize\text{NN}}$, where $\alpha$ represents the weights,
and $g_{\alpha}^{\footnotesize\text{NN}}$ has the same structure and the same
number of parameters as in the fully unknown setting. The goal is to learn
$\alpha$ so that
$f_{\alpha}:=f_{\text{approx}}+g_{\alpha}^{\footnotesize\text{NN}}$
approximates $f^{*}$.

We set $d_{x}=40$ and consider two settings for $H$: fully observed with
$H=I_{40}$, $d_{y}=40$, and partially observed at every two out of three
coordinates with $d_{y}=27$ (see Section 5.2.1). Eight data sequences are
generated with the reference model for training and four for testing, each
with length $T=1200$. Other experimental settings are the same as in Section
5.2.1.

For the setting where training data is fully observed, we compare AD-EnKF-T
with AD-PF-T and the EM approach. The results are plotted in Fig. 7. The
number of particles used for all algorithms is fixed at $N=50$, and covariance
tapering Eq. F.3 with radius $r=5$ is applied to EnKF. The subsequence length
for both AD-EnKF-T and AD-PF-T is chosen to be $L=20$. We find that, whether
$f^{*}$ is fully known or an inaccurate model is available, AD-EnKF-T is able
to learn the reference vector field $f^{*}$ well, with the smallest forecast
RMSE among all methods. Applying a filtering algorithm to the learned model,
we find that the states recovered by the AD-EnKF-T algorithm at the final
iteration have the lowest error (filter RMSE) among all methods, indicating
that AD-EnKF-T also has the ability to learn unknown states well. Moreover,
the filter RMSE of AD-EnKF-T is close to the one computed using a filtering
algorithm _with_ known $f^{*}$ and $Q^{*}$. The test log-likelihood
$\mathcal{L}_{{\footnotesize\text{EnKF}}}$ of the model learned by AD-EnKF-T
is close to the one evaluated with the reference model. We also find that
having an inaccurate model $f_{\text{approx}}$ is beneficial to the learning
of AD-EnKF-T. The performance metrics are boosted compared to the ones with a
fully unknown model. EM has worse results, where we find that the forecast
RMSE does not consistently drop in the training procedure and the states are
not accurately recovered. This might be because the smoothing distribution
used by EM cannot be approximated accurately. AD-PF-T has the worst
performance, possibly because PF fails in high dimensions.

Figure 7: Learning the Lorenz-96 model from fully unknown dynamics (Section
5.2.2) v.s. model correction (Section 5.2.3), with full observations
($H=I_{d_{x}}$). All performance metrics are evaluated after each training
iteration. Red dashed lines correspond to metric values obtained with the
reference model $f^{*}$ and $Q^{*}$.

We repeat the learning procedure in the setting where training data is
partially observed at every two out of three coordinates. The results are
shown in Fig. 8. Those for AD-PF-T are not shown since training cannot be
completed due to filter divergence. We find that AD-EnKF-T is still able to
recover $f^{*}$ consistently as well as the unknown states for all
coordinates, including the ones that are not observed, and has a filter RMSE
close to the one computed with knowledge of $f^{*}$. However, the performance
metrics of the EM algorithm in the model correction experiment deteriorate as
training proceeds, indicating that it may overfit the training data. In
addition, we find that the EM algorithm does not converge to the same point in
repeated trials, particularly so in the setting of fully unknown dynamics. All
of these results indicate that AD-EnKF is advantageous when learning from
partial observations in high dimensions.

The ability to recover the underlying dynamics and states even with incomplete
observations and fully unknown dynamics is most likely due to the
convolutional-type architecture of the NN
$f_{\alpha}^{\footnotesize\text{NN}}$, which implicitly assumes that each
coordinate only interacts with its neighbors, and that this interaction is
spatially invariant.

Figure 8: Learning Lorenz-96 from fully unknown dynamics (Section 5.2.2) v.s.
model correction (Section 5.2.3) with partial observations
($H=[e_{1},e_{2},e_{4},e_{5},e_{7},\cdots]^{\top}$). All performance metrics
are evaluated after each training iteration. Red dashed lines correspond to
metric values obtained with the reference model $f^{*}$ and $Q^{*}$. The
absence of lines for EM in the fully unknown setting is due to its low and
unstable performance.

##  6 Conclusions and Future Directions

This paper introduced AD-EnKFs for the principled learning of states and
dynamics in DA. We have shown that AD-EnKFs can be successfully integrated
with DA localization techniques for recovery of high-dimensional states, and
with TBPTT techniques to handle large observation data and high-dimensional
surrogate models. Numerical results on the Lorenz-96 model show that AD-EnKFs
outperform existing EM and PF methods to merge DA and ML.

Several research directions stem from this work. First, gradient and Hessian
information of $\mathcal{L}_{{\footnotesize\text{EnKF}}}$ obtained by autodiff
can be utilized to design optimization schemes beyond the first-order approach
we consider. Second, the convergence analysis of EnKF estimation of the log-
likelihood and its gradient may be generalized to nonlinear settings. Third,
the idea of AD-EnKF could be applied to auto-differentiate through other
filtering algorithms, e.g. unscented Kalman filters, and in Bayesian inverse
problems using iterative ensemble Kalman methods. Finally, the encouraging
numerical results obtained on the Lorenz-96 model motivate the deployment and
further investigation of AD-EnKFs in scientific and engineering applications
where latent states need to be estimated with incomplete knowledge of their
dynamics.

## Acknowledgments

YC was partially supported by DMS-2027056 and NSF OAC-1934637. DSA is grateful
for the support of DMS-2027056. RW is grateful for the support of DOD
FA9550-18-1-0166, DOE DE-AC02-06CH11357, NSF OAC-1934637, DMS-1930049, and
DMS-2023109.

## References

  * [1] Martín Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, et al.  Tensorflow: A system for large-scale machine learning.  In 12th $\\{$USENIX$\\}$ Symposium on Operating Systems Design and Implementation ($\\{$OSDI$\\}$ 16), pages 265–283, 2016.
  * [2] S. Agapiou, O. Papaspiliopoulos, D. Sanz-Alonso, and A. M. Stuart.  Importance sampling: Intrinsic dimension and computational cost.  Statistical Science, 32(3):405–431, 2017.
  * [3] Jeffrey L Anderson.  An ensemble adjustment Kalman filter for data assimilation.  Monthly Weather Review, 129(12):2884–2903, 2001.
  * [4] Jeffrey L Anderson and Stephen L Anderson.  A Monte Carlo implementation of the nonlinear filtering problem to produce ensemble assimilations and forecasts.  Monthly Weather Review, 127(12):2741–2758, 1999.
  * [5] Christophe Andrieu, Arnaud Doucet, and Roman Holenstein.  Particle Markov chain Monte Carlo methods.  Journal of the Royal Statistical Society: Series B (Statistical Methodology), 72(3):269–342, 2010.
  * [6] Ibrahim Ayed, Emmanuel de Bézenac, Arthur Pajot, Julien Brajard, and Patrick Gallinari.  Learning dynamical systems from partial observations.  arXiv preprint arXiv:1902.11136, 2019.
  * [7] Christopher M Bishop.  Pattern Recognition and Machine Learning.  Springer, 2006.
  * [8] Marc Bocquet, Julien Brajard, Alberto Carrassi, and Laurent Bertino.  Data assimilation as a learning tool to infer ordinary differential equation representations of dynamical models.  Nonlinear Processes in Geophysics, 26(3):143–162, 2019.
  * [9] Marc Bocquet, Julien Brajard, Alberto Carrassi, and Laurent Bertino.  Bayesian inference of chaotic dynamics by merging data assimilation, machine learning and expectation-maximization.  arXiv preprint arXiv:2001.06270, 2020.
  * [10] Marc Bocquet, Carlos A Pires, and Lin Wu.  Beyond Gaussian statistical modeling in geophysical data assimilation.  Monthly Weather Review, 138(8):2997–3023, 2010.
  * [11] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang.  JAX: composable transformations of Python+NumPy programs, 2018\.
  * [12] Julien Brajard, Alberto Carrassi, Marc Bocquet, and Laurent Bertino.  Combining data assimilation and machine learning to emulate a dynamical model from sparse and noisy observations: a case study with the lorenz 96 model.  Journal of Computational Science, 44:101171, 2020.
  * [13] Alberto Carrassi, Marc Bocquet, Alexis Hannart, and Michael Ghil.  Estimating model evidence using data assimilation.  Quarterly Journal of the Royal Meteorological Society, 143(703):866–880, 2017.
  * [14] Neil K Chada, Yuming Chen, and Daniel Sanz-Alonso.  Iterative ensemble Kalman methods: a unified perspective with some new variants.  arXiv preprint arXiv:2010.13299, 2020.
  * [15] Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud.  Neural ordinary differential equations.  arXiv preprint arXiv:1806.07366, 2018.
  * [16] Tadeo J Cocucci, Manuel Pulido, Magdalena Lucini, and Pierre Tandeo.  Model error covariance estimation in particle and ensemble Kalman filters using an online expectation–maximization algorithm.  Quarterly Journal of the Royal Meteorological Society, 147(734):526–543, 2021.
  * [17] Adrien Corenflos, James Thornton, Arnaud Doucet, and George Deligiannidis.  Differentiable particle filtering via entropy-regularized optimal transport.  arXiv preprint arXiv:2102.07850, 2021.
  * [18] Edward De Brouwer, Jaak Simm, Adam Arany, and Yves Moreau.  Gru-ode-bayes: Continuous modeling of sporadically-observed time series.  arXiv preprint arXiv:1905.12374, 2019.
  * [19] Pierre Del Moral.  Feynman-Kac Formulae: Genealogical and Interacting Particle Systems with Applications.  Series: Probability and Applications, Springer-Verlag, New York, 2004\.
  * [20] Pierre Del Moral, Julian Tugaut, et al.  On the stability and the uniform propagation of chaos properties of ensemble Kalman–Bucy filters.  The Annals of Applied Probability, 28(2):790–850, 2018.
  * [21] Timothy DelSole and Xiaosong Yang.  State and parameter estimation in stochastic dynamical models.  Physica D: Nonlinear Phenomena, 239(18):1781–1788, 2010.
  * [22] Arthur P Dempster, Nan M Laird, and Donald B Rubin.  Maximum likelihood from incomplete data via the EM algorithm.  Journal of the Royal Statistical Society: Series B (Methodological), 39(1):1–22, 1977.
  * [23] Arnaud Doucet, Simon Godsill, and Christophe Andrieu.  On sequential Monte Carlo sampling methods for Bayesian filtering.  Statistics and Computing, 10(3):197–208, 2000.
  * [24] Arnaud Doucet and Adam M Johansen.  A tutorial on particle filtering and smoothing: Fifteen years later.  Handbook of nonlinear filtering, 12(656-704):3, 2009.
  * [25] Denis Dreano, Pierre Tandeo, Manuel Pulido, Boujemaa Ait-El-Fquih, Thierry Chonavel, and Ibrahim Hoteit.  Estimating model-error covariances in nonlinear state-space models using Kalman smoothing and the expectation–maximization algorithm.  Quarterly Journal of the Royal Meteorological Society, 143(705):1877–1885, 2017.
  * [26] Christopher Drovandi, Richard G Everitt, Andrew Golightly, Dennis Prangle, et al.  Ensemble MCMC: accelerating pseudo-marginal MCMC for state space models using the ensemble Kalman filter.  Bayesian Analysis, 2021.
  * [27] Geir Evensen.  Sequential data assimilation with a nonlinear quasi-geostrophic model using Monte Carlo methods to forecast error statistics.  Journal of Geophysical Research: Oceans, 99(C5):10143–10162, 1994\.
  * [28] Geir Evensen.  Data assimilation: the ensemble Kalman filter.  Springer Science & Business Media, 2009.
  * [29] Marco Fraccaro, Simon Kamronn, Ulrich Paquet, and Ole Winther.  A disentangled recognition and nonlinear dynamics model for unsupervised learning.  arXiv preprint arXiv:1710.05741, 2017.
  * [30] Reinhard Furrer and Thomas Bengtsson.  Estimation of high-dimensional prior and posterior covariance matrices in Kalman filter variants.  Journal of Multivariate Analysis, 98(2):227–255, 2007.
  * [31] Gregory Gaspari and Stephen E Cohn.  Construction of correlation functions in two and three dimensions.  Quarterly Journal of the Royal Meteorological Society, 125(554):723–757, 1999.
  * [32] Mike B Giles.  Collected matrix derivative results for forward and reverse mode algorithmic differentiation.  In Advances in Automatic Differentiation, pages 35–44. Springer, 2008.
  * [33] Neil J Gordon, David J Salmond, and Adrian FM Smith.  Novel approach to nonlinear/non-Gaussian Bayesian state estimation.  IEE proceedings F (radar and signal processing), 140(2):107–113, 1993.
  * [34] Thomas M Hamill, Jeffrey S Whitaker, and Chris Snyder.  Distance-dependent filtering of background error covariance estimates in an ensemble Kalman filter.  Monthly Weather Review, 129(11):2776–2790, 2001.
  * [35] Alexis Hannart, Alberto Carrassi, Marc Bocquet, Michael Ghil, Philippe Naveau, Manuel Pulido, Juan Ruiz, and Pierre Tandeo.  Dada: data assimilation for the detection and attribution of weather and climate-related events.  Climatic Change, 136(2):155–174, 2016.
  * [36] Peter L Houtekamer and Herschel L Mitchell.  Data assimilation using an ensemble Kalman filter technique.  Monthly Weather Review, 126(3):796–811, 1998.
  * [37] Peter L Houtekamer and Herschel L Mitchell.  A sequential ensemble Kalman filter for atmospheric data assimilation.  Monthly Weather Review, 129(1):123–137, 2001.
  * [38] Peter L Houtekamer and Fuqing Zhang.  Review of the ensemble Kalman filter for atmospheric data assimilation.  Monthly Weather Review, 144(12):4489–4532, 2016.
  * [39] Andrew H Jazwinski.  Stochastic processes and filtering theory.  Courier Corporation, 2007.
  * [40] Rudolph Emil Kalman.  A new approach to linear filtering and prediction problems.  Transactions of the ASME–Journal of Basic Engineering, 82(Series D), 1960.
  * [41] Matthias Katzfuss, Jonathan R Stroud, and Christopher K Wikle.  Understanding the ensemble Kalman filter.  The American Statistician, 70(4):350–357, 2016.
  * [42] Matthias Katzfuss, Jonathan R Stroud, and Christopher K Wikle.  Ensemble Kalman methods for high-dimensional hierarchical dynamic space-time models.  Journal of the American Statistical Association, 115(530):866–885, 2020.
  * [43] Diederik P Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  arXiv preprint arXiv:1412.6980, 2014.
  * [44] Diederik P. Kingma and Max Welling.  Auto-Encoding Variational Bayes.  In 2nd International Conference on Learning Representations, 2014\.
  * [45] Rahul Krishnan, Uri Shalit, and David Sontag.  Structured inference networks for nonlinear state space models.  In Proceedings of the AAAI Conference on Artificial Intelligence, 2017.
  * [46] Evan Kwiatkowski and Jan Mandel.  Convergence of the square root ensemble Kalman filter in the large ensemble limit.  SIAM/ASA Journal on Uncertainty Quantification, 3(1):1–17, 2015\.
  * [47] Kody JH Law, Daniel Sanz-Alonso, Abhishek Shukla, and Andrew M Stuart.  Filter accuracy for the Lorenz 96 model: Fixed versus adaptive observation operators.  Physica D: Nonlinear Phenomena, 325:1–13, 2016.
  * [48] Kody JH Law and Andrew M Stuart.  Evaluating data assimilation algorithms.  Monthly Weather Review, 140(11):3757–3782, 2012.
  * [49] Kody JH Law, Andrew M Stuart, and Kostas Zygalakis.  Data Assimilation.  Cham, Switzerland: Springer, 214, 2015.
  * [50] Kody JH Law, Hamidou Tembine, and Raul Tempone.  Deterministic mean-field ensemble Kalman filtering.  SIAM Journal on Scientific Computing, 38(3):A1251–A1279, 2016.
  * [51] Tuan Anh Le, Maximilian Igl, Tom Rainforth, Tom Jin, and Frank Wood.  Auto-encoding sequential Monte Carlo.  arXiv preprint arXiv:1705.10306, 2017.
  * [52] François Le Gland and Laurent Mevel.  Recursive identification in hidden markov models.  In Proceedings of the 36th Conference on Decision and Control, San Diego 1997, volume 4, pages 3468–3473, 1997.
  * [53] François Le Gland, Valérie Monbet, and Vu-Duc Tran.  Large sample asymptotics for the ensemble Kalman filter.  PhD thesis, INRIA, 2009.
  * [54] Edward N Lorenz.  Predictability: A problem partly solved.  In Proc. Seminar on Predictability, 1996.
  * [55] Chris J Maddison, Dieterich Lawson, George Tucker, Nicolas Heess, Mohammad Norouzi, Andriy Mnih, Arnaud Doucet, and Yee Whye Teh.  Filtering variational objectives.  arXiv preprint arXiv:1705.09279, 2017.
  * [56] Andrew J Majda and John Harlim.  Filtering Complex Turbulent Systems.  Cambridge University Press, 2012.
  * [57] Henry P McKean.  Propagation of chaos for a class of non-linear parabolic equations.  Stochastic Differential Equations (Lecture Series in Differential Equations, Session 7, Catholic Univ., 1967), pages 41–57, 1967\.
  * [58] Sammy Metref, Alexis Hannart, Juan Ruiz, Marc Bocquet, Alberto Carrassi, and Michael Ghil.  Estimating model evidence using ensemble-based data assimilation with localization–The model selection problem.  Quarterly Journal of the Royal Meteorological Society, 145(721):1571–1588, 2019.
  * [59] Christian Naesseth, Scott Linderman, Rajesh Ranganath, and David Blei.  Variational sequential Monte Carlo.  In International Conference on Artificial Intelligence and Statistics, pages 968–977. PMLR, 2018.
  * [60] Duong Nguyen, Said Ouala, Lucas Drumetz, and Ronan Fablet.  Em-like learning chaotic dynamics from noisy and partial observations.  arXiv preprint arXiv:1903.10335, 2019.
  * [61] Edward Ott, Brian R Hunt, Istvan Szunyogh, Aleksey V Zimin, Eric J Kostelich, Matteo Corazza, Eugenia Kalnay, DJ Patil, and James A Yorke.  A local ensemble Kalman filter for atmospheric data assimilation.  Tellus A: Dynamic Meteorology and Oceanography, 56(5):415–428, 2004\.
  * [62] Omiros Papaspiliopoulos, Matteo Ruggiero, et al.  Optimal filtering and the dual process.  Bernoulli, 20(4):1999–2019, 2014.
  * [63] Razvan Pascanu, Tomas Mikolov, and Yoshua Bengio.  On the difficulty of training recurrent neural networks.  In International conference on machine learning, pages 1310–1318. PMLR, 2013.
  * [64] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala.  Pytorch: An imperative style, high-performance deep learning library.  In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems 32, pages 8024–8035. Curran Associates, Inc., 2019.
  * [65] Manuel Pulido, Pierre Tandeo, Marc Bocquet, Alberto Carrassi, and Magdalena Lucini.  Stochastic parameterization identification using ensemble Kalman filtering combined with maximum likelihood methods.  Tellus A: Dynamic Meteorology and Oceanography, 70(1):1–17, 2018\.
  * [66] Rajesh Ranganath, Sean Gerrish, and David Blei.  Black box variational inference.  In Artificial intelligence and statistics, pages 814–822. PMLR, 2014.
  * [67] Syama Sundar Rangapuram, Matthias Seeger, Jan Gasthaus, Lorenzo Stella, Yuyang Wang, and Tim Januschowski.  Deep state space models for time series forecasting.  In Proceedings of the 32nd international conference on neural information processing systems, pages 7796–7805, 2018.
  * [68] Sebastian Reich and Colin Cotter.  Probabilistic Forecasting and Bayesian Data Assimilation.  Cambridge University Press, 2015.
  * [69] Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra.  Stochastic backpropagation and approximate inference in deep generative models.  In International conference on machine learning, pages 1278–1286. PMLR, 2014.
  * [70] Michael Roth, Gustaf Hendeby, Carsten Fritsche, and Fredrik Gustafsson.  The Ensemble Kalman filter: a signal processing perspective.  EURASIP Journal on Advances in Signal Processing, 2017(1):1–16, 2017.
  * [71] Yulia Rubanova, Ricky TQ Chen, and David Duvenaud.  Latent ordinary differential equations for irregularly-sampled time series.  2019\.
  * [72] Pavel Sakov and Laurent Bertino.  Relation between two common localisation methods for the EnKF.  Computational Geosciences, 15(2):225–237, 2011.
  * [73] Daniel Sanz-Alonso.  Importance sampling and necessary sample size: an information theory approach.  SIAM/ASA Journal on Uncertainty Quantification, 6(2):867–879, 2018\.
  * [74] Daniel Sanz-Alonso and Andrew M Stuart.  Long-time asymptotics of the filtering distribution for partially observed chaotic dynamical systems.  SIAM/ASA Journal on Uncertainty Quantification, 3(1):1200–1220, 2015.
  * [75] Daniel Sanz-Alonso, Andrew M Stuart, and Armeen Taeb.  Inverse Problems and Data Assimilation.  arXiv preprint arXiv:1810.06191, 2019.
  * [76] Daniel Sanz-Alonso and Zijian Wang.  Bayesian update with importance sampling: Required sample size.  Entropy, 23(1):22, 2021.
  * [77] Claudia Schillings and Andrew M Stuart.  Analysis of the ensemble Kalman filter for inverse problems.  SIAM Journal on Numerical Analysis, 55(3):1264–1290, 2017.
  * [78] Chris Snyder, Thomas Bengtsson, Peter Bickel, and Jeff Anderson.  Obstacles to high-dimensional particle filtering.  Monthly Weather Review, 136(12):4629–4640, 2008.
  * [79] Jonathan R Stroud and Thomas Bengtsson.  Sequential state and variance estimation within the ensemble Kalman filter.  Monthly Weather Review, 135(9):3194–3208, 2007.
  * [80] Jonathan R Stroud, Matthias Katzfuss, and Christopher K Wikle.  A Bayesian adaptive ensemble Kalman filter for sequential state and parameter estimation.  Monthly Weather Review, 146(1):373–386, 2018.
  * [81] Jonathan R Stroud, Michael L Stein, Barry M Lesht, David J Schwab, and Dmitry Beletsky.  An ensemble Kalman filter and smoother for satellite data assimilation.  Journal of the American Statistical Association, 105(491):978–990, 2010.
  * [82] Ilya Sutskever, Oriol Vinyals, and Quoc V Le.  Sequence to sequence learning with neural networks.  arXiv preprint arXiv:1409.3215, 2014.
  * [83] Alain-Sol Sznitman.  Topics in propagation of chaos.  In Ecole d’été de probabilités de Saint-Flour XIX—1989, pages 165–251. Springer, 1991.
  * [84] Istvan Szunyogh, Eric J Kostelich, Gyorgyi Gyarmati, Eugenia Kalnay, Brian R Hunt, Edward Ott, Elizabeth Satterfield, and James A Yorke.  A local ensemble transform Kalman filter data assimilation system for the NCEP global model.  Tellus A: Dynamic Meteorology and Oceanography, 60(1):113–130, 2008\.
  * [85] Pierre Tandeo, Manuel Pulido, and François Lott.  Offline parameter estimation using EnKF and maximum likelihood error covariance estimates: Application to a subgrid-scale orography parametrization.  Quarterly journal of the royal meteorological society, 141(687):383–395, 2015.
  * [86] Genta Ueno and Nagatomo Nakamura.  Iterative algorithm for maximum-likelihood estimation of the observation-error covariance matrix for ensemble-based filters.  Quarterly Journal of the Royal Meteorological Society, 140(678):295–315, 2014.
  * [87] Genta Ueno and Nagatomo Nakamura.  Bayesian estimation of the observation-error covariance matrix in ensemble-based filters.  Quarterly Journal of the Royal Meteorological Society, 142(698):2055–2080, 2016.
  * [88] Greg CG Wei and Martin A Tanner.  A Monte Carlo implementation of the EM algorithm and the poor man’s data augmentation algorithms.  Journal of the American statistical Association, 85(411):699–704, 1990.
  * [89] Jeffrey S Whitaker, Thomas M Hamill, Xue Wei, Yucheng Song, and Zoltan Toth.  Ensemble data assimilation with the ncep global forecast system.  Monthly Weather Review, 136(2):463–482, 2008.
  * [90] Ronald J Williams and David Zipser.  Gradient-based learning algorithms for recurrent.  Backpropagation: Theory, architectures, and applications, 433:17, 1995.
  * [91] Ke Xu and Christopher K Wikle.  Estimation of parameterized spatio-temporal dynamic models.  Journal of Statistical Planning and Inference, 137(2):567–588, 2007\.

##### Notation

We denote by $c$ a constant that does not depend on $N$ and may change from
line to line. We denote by $\|U\|_{p}$ the $L^{p}$ norm of a random
vector/matrix $U$:
$\|U\|_{p}:=\big{(}\operatorname{\mathbb{E}}|U|^{p}\big{)}^{1/p}$, where
$|\cdot|$ is the underlying vector/matrix norm. (Here we use 2-norm for
vectors and Frobenius norm for matrices.) For a sequence of random
vectors/matrices $U_{N}$, we write

$U_{N}\xrightarrow{1/2}U$

if, for any $p\geq 1$, there exists a constant $c$ such that

$\|U_{N}-U\|_{p}\leq cN^{-1/2},\quad\quad\forall N\geq 1.$

For a scalar valued function $f(U)$ that takes a vector/matrix $U$ as input,
we denote by $\partial_{U}f$ the derivative of $f$ w.r.t. $U$, which collects
the derivative of $f$ w.r.t. each entry of the vector/matrix $U$. When $U$ is
a vector, the notations $\partial_{U}f$ and $\nabla_{U}f$ are equivalent.

For a vector/matrix valued function $U(a)$ that takes a scalar $a$ as input,
we denote by $\partial_{a}U$ the derivative of $U$ w.r.t. $a$, which collects
the derivative of each entry of the vector/matrix $U$ w.r.t. $a$.

##  Appendix A Proof of Theorem 3.1

We first recall the propagation of chaos statement. Notice that in the EnKF
algorithm Alg. 3.1, we compute ${x}_{t}^{1:N}$ sequentially, based on the
forecast ensemble $\widehat{x}_{t}^{1:N}$ and its empirical mean and
covariance $\widehat{m}_{t},\widehat{C}_{t}$. We build “substitute particles”
${\mathsf{\bf{x}}}_{t}^{1:N}$ in a similar fashion, except that at each step
the population mean and covariance
$\widehat{\mathsf{\bf{m}}}_{t},\widehat{\mathsf{\bf{C}}}_{t}$ are used instead
of their empirical versions. Starting from
$\mathsf{\bf{x}}_{0}^{1:N}=x_{0}^{1:N}$, the update rules of substitute
particles are listed below, with a side-by-side comparison to the EnKF update
rules:

$\begin{array}[]{rl|rl}\lx@intercol\hfil\text{EnKF
particles}\hfil\lx@intercol\vrule\lx@intercol&\lx@intercol\hfil\text{Substitute
particles}\hfil\lx@intercol\\\
\hline\cr\widehat{x}_{t}^{\,n}&=F_{\alpha}({x}_{t-1}^{n})+{\xi}_{t}^{n}&\widehat{\mathsf{\bf{x}}}_{t}^{\,n}&=F_{\alpha}({\mathsf{\bf{x}}}_{t-1}^{n})+{\xi}_{t}^{n}\\\
\widehat{m}_{t}&=\frac{1}{N}\sum_{n=1}^{N}\widehat{x}_{t}^{\,n}&\widehat{\mathsf{\bf{m}}}_{t}&=\operatorname{\mathbb{E}}\bigl{[}\widehat{\mathsf{\bf{x}}}_{t}^{\,n}\bigr{]}\\\
\widehat{C}_{t}&=\frac{1}{N-1}\sum_{n=1}^{N}(\widehat{x}_{t}^{\,n}-\widehat{m}_{t})(\widehat{x}_{t}^{\,n}-\widehat{m}_{t})^{\top}&\widehat{\mathsf{\bf{C}}}_{t}&=\operatorname{\mathbb{E}}\bigl{[}(\widehat{\mathsf{\bf{x}}}_{t}^{\,n}-\widehat{\mathsf{\bf{m}}}_{t})(\widehat{\mathsf{\bf{x}}}_{t}^{\,n}-\widehat{\mathsf{\bf{m}}}_{t})^{\top}\bigr{]}\\\
\widehat{K}_{t}&=\widehat{C}_{t}H^{\top}(H\widehat{C}_{t}H^{\top}+R)^{-1}&\widehat{\mathsf{\bf{K}}}_{t}&=\widehat{\mathsf{\bf{C}}}_{t}H^{\top}(H\widehat{\mathsf{\bf{C}}}_{t}H^{\top}+R)^{-1}\\\
{x}_{t}^{n}&=\widehat{x}_{t}^{\,n}+\widehat{K}_{t}(y_{t}+{\gamma}_{t}^{n}-H\widehat{x}_{t}^{\,n})&{\mathsf{\bf{x}}}_{t}^{n}&=\widehat{\mathsf{\bf{x}}}_{t}^{\,n}+\widehat{\mathsf{\bf{K}}}_{t}(y_{t}+{\gamma}_{t}^{n}-H\widehat{\mathsf{\bf{x}}}_{t}^{\,n})\\\
\end{array}$ (A.1)

Notice that the substitute particles use the _same_ realization of random
variables as the EnKF particles, including initialization of particles
$x_{0}^{1:N}$, forecast simulation error ${\xi}_{t}^{n},$ and noise
perturbation ${\gamma}_{t}^{n}$. As $N\rightarrow\infty$, one can show that
the EnKF particles ${x}_{t}^{1:N}$ (resp. ${\widehat{x}}_{t}^{1:N}$) are close
to the substitute particles ${\mathsf{\bf{x}}}_{t}^{1:N}$ (resp.
${\widehat{\mathsf{\bf{x}}}}_{t}^{1:N}$), and hence the law of large numbers
guarantees that $\widehat{m}_{t},\widehat{C}_{t}$ are close to
$\widehat{\mathsf{\bf{m}}}_{t},\widehat{\mathsf{\bf{C}}}_{t}$. We summarize
the main results from [53] (see also [46]):

######  Lemma A.1.

Under the same assumption of Theorem 3.1:
(1) For each $t\geq 1$, the substitute particles ${\mathsf{\bf{x}}}_{t}^{1:N}$
are i.i.d., and each of them has the same law as the true filtering
distribution $p(x_{t}|y_{1:t})$. Similarly,
$\widehat{\mathsf{\bf{x}}}_{t}^{\,1:N}$ are i.i.d., and each of them has the
same law as the true forecast distribution $p(x_{t}|y_{1:t-1})$. In
particular,

$p(x_{t}|y_{1:t-1})=\mathcal{N}(x_{t};\widehat{\mathsf{\bf{m}}}_{t},\widehat{\mathsf{\bf{C}}}_{t}).$
(A.2)

(2) For each $t,n,p\geq 1,$ the EnKF particle ${x}_{t}^{n}$ converges to the
substitute particle ${\mathsf{\bf{x}}}_{t}^{n}$ in $L^{p}$ with convergence
rate $N^{-1/2}$, and the substitute particle ${\mathsf{\bf{x}}}_{t}^{n}$ has
finite moments of any order. The same holds for forecast particles
$\widehat{x}_{t}^{\,n}$:

${x}_{t}^{n}\xrightarrow{1/2}{\mathsf{\bf{x}}}_{t}^{n},\quad\quad\widehat{x}_{t}^{\,n}\xrightarrow{1/2}\widehat{\mathsf{\bf{x}}}_{t}^{\,n},\quad\quad\|{\mathsf{\bf{x}}}_{t}^{n}\|_{p}\leq
c,\quad\quad\|\widehat{\mathsf{\bf{x}}}_{t}^{\,n}\|_{p}\leq c.$ (A.3)

In particular, $\widehat{m}_{t}$, $\widehat{C}_{t}$ converge to
$\widehat{\mathsf{\bf{m}}}_{t}$, $\widehat{\mathsf{\bf{C}}}_{t}$ in $L^{p}$
with convergence rate $N^{-1/2}$:

$\widehat{m}_{t}\xrightarrow{1/2}\widehat{\mathsf{\bf{m}}}_{t},\quad\quad\quad\quad\widehat{C}_{t}\xrightarrow{1/2}\widehat{\mathsf{\bf{C}}}_{t}.$
(A.4)

###### Proof.

Eq. A.2 corresponds to Lemma 2.1 of [53]. Eq. A.3 corresponds to Proposition
4.4 of [53]. Eq. A.4 is a direct corollary of Theorem 5.2 of [53]. ∎

###### Proof of Theorem 3.1.

By Eq. A.2, using the Gaussian observation assumption Eq. 2.2:

$\mathcal{L}(\theta)=\sum_{t=1}^{T}\log
p(y_{t}|y_{1:t-1})=\sum_{t=1}^{T}\log\mathcal{N}(y_{t};H\widehat{\mathsf{\bf{m}}}_{t},H\widehat{\mathsf{\bf{C}}}_{t}H^{\top}+R).$
(A.5)

By Eq. 3.6,

$\mathcal{L}_{{\footnotesize\text{EnKF}}}(\theta)=\sum_{t=1}^{T}\log\mathcal{N}(y_{t};H\widehat{m}_{t},H\widehat{C}_{t}H^{\top}+R).$
(A.6)

Define

$\begin{split}h_{t}(m,C)&:=\log\mathcal{N}(y_{t};Hm,HCH^{\top}+R)\\\
&=-\frac{1}{2}\log\det(HCH^{\top}+R)-\frac{1}{2}(y_{t}-Hm)^{\top}(HCH^{\top}+R)^{-1}(y_{t}-Hm)+{\text{const}}.\end{split}$
(A.7)

It suffices to show that, for each $t\geq 1$,

$h_{t}(\widehat{m}_{t},\widehat{C}_{t})\xrightarrow{1/2}h_{t}(\widehat{\mathsf{\bf{m}}}_{t},\widehat{\mathsf{\bf{C}}}_{t}).$
(A.8)

We denote by $\mathbb{S}_{+}^{d_{x}}\subset\mathbbm{R}^{d_{x}\times d_{x}}$
the space of all positive semi-definite matrices equipped with Frobenius norm.
Notice that $h_{t}$ is a continuous function on
$\mathbbm{R}^{d_{x}}\times\mathbb{S}_{+}^{d_{x}}$, since $HCH^{\top}+R\succeq
R\succ 0$. To show convergence in $L^{p}$, intuitively one would expect a
Lipschitz-type continuity to hold for $h_{t}$, in a suitable sense. We inspect
the derivatives of $h_{t}$ w.r.t. $m$ and $C$, which will also be useful for
later developments:

$\begin{split}\partial_{m}h_{t}(m,C)&=-H^{\top}(HCH^{\top}+R)^{-1}(y_{t}-Hm),\\\
\partial_{C}h_{t}(m,C)&=-\frac{1}{2}H^{\top}(HCH^{\top}+R)^{-1}H\\\
&\quad\quad\quad+\frac{1}{2}H^{\top}(HCH^{\top}+R)^{-1}(y_{t}-Hm)(y_{t}-Hm)^{\top}(HCH^{\top}+R)^{-1}H.\end{split}$
(A.9)

Since $R^{d_{x}}\times\mathbb{S}_{+}^{d_{x}}$ is convex, by the mean value
theorem, triangle inequality and Cauchy-Schwarz, and define
$m(\chi):=\chi\widehat{\mathsf{\bf{m}}}_{t}+(1-\chi)\widehat{m}_{t}$,
$C(\chi):=\chi\widehat{\mathsf{\bf{C}}}_{t}+(1-\chi)\widehat{C}_{t}$,

$\begin{split}\big{|}h_{t}(\widehat{m}_{t},\widehat{C}_{t})-h_{t}(\widehat{\mathsf{\bf{m}}}_{t},\widehat{\mathsf{\bf{C}}}_{t})\big{|}&\leq\sup_{\chi\in[0,1]}\big{|}\partial_{m}h_{t}\big{(}m(\chi),C(\chi)\big{)}\big{|}|\widehat{m}_{t}-\widehat{\mathsf{\bf{m}}}_{t}|\\\
&\quad\quad\quad+\sup_{\chi\in[0,1]}\big{|}\partial_{C}h_{t}\big{(}m(\chi),C(\chi)\big{)}\big{|}|\widehat{C}_{t}-\widehat{\mathsf{\bf{C}}}_{t}|.\end{split}$
(A.10)

Taking $L_{p}$ norm on both sides,

$\begin{split}\big{\|}h_{t}(\widehat{m}_{t},\widehat{C}_{t})-h_{t}(\widehat{\mathsf{\bf{m}}}_{t},\widehat{\mathsf{\bf{C}}}_{t})\big{\|}_{p}&\leq\sup_{\chi\in[0,1]}\big{\|}\partial_{m}h_{t}\big{(}m(\chi),C(\chi)\big{)}\big{\|}_{2p}\|\widehat{m}_{t}-\widehat{\mathsf{\bf{m}}}_{t}\|_{2p}\\\
&\quad\quad\quad+\sup_{\chi\in[0,1]}\big{\|}\partial_{C}h_{t}\big{(}m(\chi),C(\chi)\big{)}\big{\|}_{2p}\|\widehat{C}_{t}-\widehat{\mathsf{\bf{C}}}_{t}\|_{2p},\end{split}$
(A.11)

where we have used the triangle inequality and the $L^{p}$ Cauchy-Schwarz
inequality $\||U||V|\|_{p}\leq\|U\|_{2p}\|V\|_{2p}$, see e.g., Lemma 2.1 of
[46]. Also, by plugging in Eq. A.9, for each $\chi\in[0,1]$,

$\begin{split}\big{\|}\partial_{m}h_{t}\big{(}m(\chi),C(\chi)\big{)}\big{\|}_{2p}&\leq\big{\|}|H||(HC(\chi)H^{\top}+R)^{-1}|(|y_{t}-\chi
H\widehat{\mathsf{\bf{m}}}_{t}-(1-\chi)H\widehat{m}_{t}|)\big{\|}_{2p}\\\
&\leq|H||R^{-1}|\big{(}|y_{t}|+|H||\widehat{\mathsf{\bf{m}}}_{t}|+|H|\|\widehat{m}_{t}\|_{2p}\big{)}\leq
c,\end{split}$ (A.12)

where we have used that $|(HC(\chi)H^{\top}+R)^{-1}|\leq|R^{-1}|$, that
$\widehat{\mathsf{\bf{m}}}_{t}$ is deterministic, and that all moments of
$\widehat{m}_{t}$ are finite, by Eq. A.4. Similarly,

$\begin{split}\big{\|}\partial_{C}h_{t}\big{(}m(\chi),C(\chi)\big{)}\big{\|}_{2p}&\leq\frac{1}{2}|H|^{2}|R^{-1}|+\frac{1}{2}|H|^{2}|R^{-1}|^{2}\|y_{t}-\chi
H\widehat{\mathsf{\bf{m}}}_{t}-(1-\chi)H\widehat{m}_{t}\|_{4p}^{2}\\\
&\leq\frac{1}{2}|H|^{2}|R^{-1}|+\frac{1}{2}|H|^{2}|R^{-1}|^{2}(|y_{t}|+|H||\widehat{\mathsf{\bf{m}}}_{t}|+|H|\|\widehat{m}_{t}\|_{4p})^{2}\leq
c,\end{split}$ (A.13)

where we have used that $|vv^{\top}|=|v|^{2}$ for vector $v$. Thus, combining
Eqs. A.11, A.12, A.13, and A.4 gives

$\|h_{t}(\widehat{m}_{t},\widehat{C}_{t})-h_{t}(\widehat{\mathsf{\bf{m}}}_{t},\widehat{\mathsf{\bf{C}}}_{t})\|_{p}\leq
cN^{-1/2},$ (A.14)

which concludes the proof. ∎

##  Appendix B Proof of Theorem 3.2

Without loss of generality, we assume that $\theta\in\mathbbm{R}$ is a scalar
parameter, since in general the gradient w.r.t. $\theta$ is a collection of
derivatives w.r.t. each element of $\theta$. We will use the following lemma
repeatedly:

######  Lemma B.1.

For sequences of random vectors/matrices $U_{N}$, $V_{N}$:
(1) If $U_{N}-V_{N}\xrightarrow{1/2}0$ and $V_{N}\xrightarrow{1/2}V,$ then

$U_{N}\xrightarrow{1/2}V.$ (B.1)

(2) If $U_{N}\xrightarrow{1/2}U$, $V_{N}\xrightarrow{1/2}V$ and $U,V$ have
finite moments of any order, then

$U_{N}V_{N}\xrightarrow{1/2}UV.$ (B.2)

More generally, the result holds for multiplication of more than two
variables.

###### Proof.

(1) Using $U_{N}-V=(U_{N}-V_{N})+(V_{N}-V)$, the proof follows from the
triangle inequality.
(2) Applying triangle inequality and $L^{p}$ Cauchy-Schwarz inequality,

$\begin{split}\|U_{N}V_{N}-UV\|_{p}&\leq\|(U_{N}-U)V_{N}\|_{p}+\|U(V_{N}-V)\|_{p}\\\
&\leq\|U_{N}-U\|_{2p}\|V_{N}\|_{2p}+\|U\|_{2p}\|V_{N}-V\|_{2p}\\\ &\leq
cN^{-1/2}\Bigl{(}\|V\|_{2p}+cN^{-1/2}\Bigr{)}+\|U\|_{2p}cN^{-1/2}\\\ &\leq
cN^{-1/2}.\end{split}$ (B.3)

∎

The following result, which we will use repeatedly, is an immediate corollary
of Lemma A.1:

######  Lemma B.2.

Under the same assumption of Theorem 3.1,

$(H\widehat{C}_{t}H^{\top}+R)^{-1}\xrightarrow{1/2}(H\widehat{\mathsf{\bf{C}}}_{t}H^{\top}+R)^{-1}.$
(B.4)

###### Proof.

Using the identity $A^{-1}-B^{-1}=A^{-1}(B-A)B^{-1}$ for invertible matrices
$A$, $B$:

$\begin{split}&\phantom{{}={}}\|(H\widehat{C}_{t}H^{\top}+R)^{-1}-(H\widehat{\mathsf{\bf{C}}}_{t}H^{\top}+R)^{-1}\|_{p}\\\
&=\|(H\widehat{C}_{t}H^{\top}+R)^{-1}H(\widehat{\mathsf{\bf{C}}}_{t}-\widehat{C}_{t})H^{\top}(H\widehat{\mathsf{\bf{C}}}_{t}H^{\top}+R)^{-1}\|_{p}\\\
&\leq|R^{-1}|^{2}|H|^{2}\|\widehat{C}_{t}-\widehat{\mathsf{\bf{C}}}_{t}\|_{p}\\\
&\leq cN^{-1/2},\end{split}$ (B.5)

where we have used the $L^{p}$ convergence of $\widehat{C}_{t}$ to
$\widehat{\mathsf{\bf{C}}}_{t}$ Eq. A.4, and the fact that
$|(HCH^{T}+R)^{-1}|\leq|R^{-1}|$ for $C\succeq 0$. ∎

######  Lemma B.3.

Under the same assumption of Theorem 3.2, for each $t\geq 1$, both
$\partial_{\theta}\widehat{x}_{t}^{\,n}$ and
$\partial_{\theta}{\widehat{\mathsf{\bf{x}}}_{t}^{\,n}}$ exist, and
$\partial_{\theta}\widehat{x}_{t}^{\,n}$ converges to
$\partial_{\theta}\widehat{\mathsf{\bf{x}}}_{t}^{\,n}$ in $L^{p}$ for any
$p\geq 1$ with convergence rate $N^{-1/2}$. Moreover,
$\partial_{\theta}\widehat{\mathsf{\bf{x}}}_{t}^{\,n}$ has finite moments of
any order:

$\displaystyle\partial_{\theta}\widehat{x}_{t}^{\,n}\xrightarrow{1/2}\partial_{\theta}\widehat{\mathsf{\bf{x}}}_{t}^{\,n},\quad\quad\|\partial_{\theta}\widehat{\mathsf{\bf{x}}}_{t}^{\,n}\|_{p}\leq
c,\quad\quad\forall n.$ (B.6)

In addition, all derivatives $\partial_{\theta}\widehat{m}_{t}$,
$\partial_{\theta}\widehat{\mathsf{\bf{m}}}_{t}$,
$\partial_{\theta}\widehat{C}_{t}$,
$\partial_{\theta}\widehat{\mathsf{\bf{C}}}_{t}$,
$\partial_{\theta}\widehat{K}_{t}$ and
$\partial_{\theta}\widehat{\mathsf{\bf{K}}}_{t}$ exist, and

$\displaystyle\partial_{\theta}\widehat{m}_{t}\xrightarrow{1/2}\partial_{\theta}\widehat{\mathsf{\bf{m}}}_{t},\quad\quad\partial_{\theta}\widehat{C}_{t}\xrightarrow{1/2}\partial_{\theta}\widehat{\mathsf{\bf{C}}}_{t},\quad\quad\partial_{\theta}\widehat{K}_{t}\xrightarrow{1/2}\partial_{\theta}\widehat{\mathsf{\bf{K}}}_{t}.$
(B.7)

###### Proof.

We will prove this by induction. For $t=1$, since
$\widehat{x}_{1}^{\,n}=Ax_{0}^{n}+S\xi_{0}^{n}=\widehat{\mathsf{\bf{x}}}_{1}^{\,n}$,

$\partial_{\theta}\widehat{x}_{1}^{\,n}=(\partial_{\theta}A)x_{0}^{n}+(\partial_{\theta}S)\xi_{0}^{n}=\partial_{\theta}\widehat{\mathsf{\bf{x}}}_{1}^{\,n},$
(B.8)

and both derivatives $\partial_{\theta}\widehat{x}_{1}^{\,n}$ and
$\partial_{\theta}\widehat{\mathsf{\bf{x}}}_{1}^{\,n}$ exist. Also,

$\|\partial_{\theta}\widehat{\mathsf{\bf{x}}}_{1}^{n}\|_{p}\leq|\partial_{\theta}A|\|x_{0}^{n}\|_{p}+|\partial_{\theta}S|\|\xi_{0}^{n}\|_{p}\leq
c,$ (B.9)

since $x_{0}^{n}$ and $\xi_{0}^{n}$ are drawn from Gaussian distributions,
which have finite moments of any order. So Eq. B.6 holds for $t=1$.

Assume Eq. B.6 holds for step $t$. Then, using the definition for
$\widehat{m}_{t}$:

$\partial_{\theta}\widehat{m}_{t}=\frac{1}{N}\sum_{n=1}^{N}\partial_{\theta}\widehat{x}_{t}^{\,n}\xrightarrow[\text{\scriptsize1⃝}]{1/2}\frac{1}{N}\sum_{n=1}^{N}\partial_{\theta}\widehat{\mathsf{\bf{x}}}_{t}^{\,n}\xrightarrow[\text{\scriptsize2⃝}]{1/2}\operatorname{\mathbb{E}}[\partial_{\theta}\widehat{\mathsf{\bf{x}}}_{t}^{\,n}]\underset{\text{\scriptsize3⃝}}{=}\partial_{\theta}\operatorname{\mathbb{E}}[\widehat{x}_{t}^{\,n}]=\partial_{\theta}\widehat{\mathsf{\bf{m}}}_{t}.$
(B.10)

Convergence \scriptsize1⃝ follows from induction assumption Eq. B.6.
Convergence \scriptsize2⃝ follows from law of large numbers in $L^{p}$, since
$\partial_{\theta}\widehat{\mathsf{\bf{x}}}_{t}^{\,n}$ are i.i.d. and the
moments of $\partial_{\theta}\widehat{\mathsf{\bf{x}}}_{t}^{\,n}$ are finite
by induction assumption Eq. B.6. The swap of differentiation and expectation
in \scriptsize3⃝ is valid since the expectation is taken over a distribution
that is independent of $\theta$. Both derivatives
$\partial_{\theta}\widehat{m}_{t}$ and
$\partial_{\theta}\widehat{\mathsf{\bf{m}}}_{t}$ exist. Similarly,

$\begin{split}\partial_{\theta}\widehat{C}_{t}&\hskip 5.0pt=\hskip
5.0pt\frac{1}{N-1}\sum_{n=1}^{N}\partial_{\theta}(\widehat{x}_{t}^{\,n}(\widehat{x}_{t}^{\,n})^{\top})-\frac{N}{N-1}\partial_{\theta}(\widehat{m}_{t}\widehat{m}_{t}^{\top})\\\
&\hskip 5.0pt=\hskip
5.0pt\frac{1}{N-1}\sum_{n=1}^{N}\big{(}(\partial_{\theta}\widehat{x}_{t}^{\,n})(\widehat{x}_{t}^{\,n})^{\top}+\widehat{x}_{t}^{\,n}(\partial_{\theta}\widehat{x}_{t}^{\,n})^{\top}\big{)}-\frac{N}{N-1}\big{(}(\partial_{\theta}\widehat{m}_{t})\widehat{m}_{t}^{\top}+\widehat{m}_{t}(\partial_{\theta}\widehat{m}_{t})^{\top}\big{)}\\\
&\xrightarrow[\text{\scriptsize1⃝}]{1/2}\frac{1}{N-1}\sum_{n=1}^{N}\big{(}(\partial_{\theta}\widehat{\mathsf{\bf{x}}}_{t}^{\,n})(\widehat{\mathsf{\bf{x}}}_{t}^{\,n})^{\top}+\widehat{\mathsf{\bf{x}}}_{t}^{\,n}(\partial_{\theta}\widehat{\mathsf{\bf{x}}}_{t}^{\,n})^{\top}\big{)}-\frac{N}{N-1}\big{(}(\partial_{\theta}\widehat{\mathsf{\bf{m}}}_{t})\widehat{\mathsf{\bf{m}}}_{t}^{\top}+\widehat{\mathsf{\bf{m}}}_{t}(\partial_{\theta}\widehat{\mathsf{\bf{m}}}_{t})^{\top}\big{)}\\\
&\xrightarrow[\text{\scriptsize2⃝}]{1/2}\mathbb{E}[\partial_{\theta}\big{(}\widehat{\mathsf{\bf{x}}}_{t}^{\,n}(\widehat{\mathsf{\bf{x}}}_{t}^{\,n}\big{)}^{\top})]-\partial_{\theta}(\widehat{\mathsf{\bf{m}}}_{t}\widehat{\mathsf{\bf{m}}}_{t}^{\top})\\\
&\hskip 5.0pt\underset{\text{\scriptsize3⃝}}{=}\hskip
5.0pt\partial_{\theta}(\operatorname{\mathbb{E}}[\widehat{\mathsf{\bf{x}}}_{t}^{\,n}(\widehat{\mathsf{\bf{x}}}_{t}^{\,n})^{\top}]-\widehat{\mathsf{\bf{m}}}_{t}\widehat{\mathsf{\bf{m}}}_{t}^{\top})\\\
&\hskip 5.0pt=\hskip
5.0pt\partial_{\theta}\widehat{\mathsf{\bf{C}}}_{t}.\end{split}$ (B.11)

For \scriptsize1⃝ we have used the $L^{p}$ convergence of
$\widehat{x}_{t}^{\,n}$ to $\widehat{\mathsf{\bf{x}}}_{t}^{\,n}$,
$\partial_{\theta}\widehat{x}_{t}^{\,n}$ to
$\partial_{\theta}\widehat{\mathsf{\bf{x}}}_{t}^{\,n}$, $\widehat{m}_{t}$ to
$\widehat{\mathsf{\bf{m}}}_{t}$ and $\partial_{\theta}\widehat{m}_{t}$ to
$\partial_{\theta}\widehat{\mathsf{\bf{m}}}_{t}$ with rate $N^{-1/2}$, and the
fact that $\widehat{\mathsf{\bf{x}}}_{t}^{\,n}$ and
$\partial_{\theta}\widehat{\mathsf{\bf{x}}}_{t}^{\,n}$ have finite moments of
any order, followed by Lemma B.1. Convergence \scriptsize2⃝ follows from law
of large numbers in $L^{p}$ since
$(\partial_{\theta}\widehat{\mathsf{\bf{x}}}_{t}^{\,n})(\widehat{\mathsf{\bf{x}}}_{t}^{\,n})^{\top}$
are i.i.d. with finite moments, by the Cauchy-Schwarz inequality.
\scriptsize3⃝ is valid since the expectation is taken over a distribution that
is independent of $\theta$. Both derivatives
$\partial_{\theta}\widehat{C}_{t}$ and
$\partial_{\theta}\widehat{\mathsf{\bf{C}}}_{t}$ exist. Similarly,

$\begin{split}\partial_{\theta}\widehat{K}_{t}&\hskip 5.0pt=\hskip
5.0pt\partial_{\theta}\big{(}\widehat{C}_{t}H(H\widehat{C}_{t}H^{\top}+R)^{-1}\big{)}\\\
&\hskip 5.0pt\underset{\text{\scriptsize1⃝}}{=}\hskip
5.0pt(\partial_{\theta}\widehat{C}_{t})H(H\widehat{C}_{t}H^{\top}+R)^{-1}\\\
&\quad\quad\quad-\widehat{C}_{t}H(H\widehat{C}_{t}H^{\top}+R)^{-1}(H(\partial_{\theta}\widehat{C}_{t})H^{\top}+R)(H\widehat{C}_{t}H^{\top}+R)^{-1}\\\
&\xrightarrow[\text{\scriptsize2⃝}]{1/2}(\partial_{\theta}\widehat{\mathsf{\bf{C}}}_{t})H(H\widehat{\mathsf{\bf{C}}}_{t}H^{\top}+R)^{-1}\\\
&\quad\quad\quad-\widehat{\mathsf{\bf{C}}}_{t}H(H\widehat{\mathsf{\bf{C}}}_{t}H^{\top}+R)^{-1}(H(\partial_{\theta}\widehat{\mathsf{\bf{C}}}_{t})H^{\top}+R)(H\widehat{\mathsf{\bf{C}}}_{t}H^{\top}+R)^{-1}\\\
&\hskip 5.0pt\underset{\text{\scriptsize1⃝}}{=}\hskip
5.0pt\partial_{\theta}\big{(}\widehat{\mathsf{\bf{C}}}_{t}H(H\widehat{\mathsf{\bf{C}}}_{t}H^{\top}+R)^{-1}\big{)}\\\
&\hskip 5.0pt=\hskip
5.0pt\partial_{\theta}\widehat{\mathsf{\bf{K}}}_{t}.\end{split}$ (B.12)

Here equalities \scriptsize1⃝ and \scriptsize3⃝ follow from chain rule. For
\scriptsize2⃝ we have used the $L^{p}$ convergence of $\widehat{C}_{t}$ to
$\widehat{\mathsf{\bf{C}}}_{t}$, $\partial_{\theta}\widehat{C}_{t}$ to
$\partial_{\theta}\widehat{\mathsf{\bf{C}}}_{t}$ and
$(H\widehat{C}_{t}H^{\top}+R)^{-1}$ to
$(H\widehat{\mathsf{\bf{C}}}_{t}H^{\top}+R)^{-1}$ with rate $N^{-1/2}$ (by Eq.
B.4), followed by Lemma B.1. Both derivatives
$\partial_{\theta}\widehat{K}_{t}$ and
$\partial_{\theta}\widehat{\mathsf{\bf{K}}}_{t}$ exist since $R\succ 0$.

To show Eq. B.6 holds for step $t+1$, we need to investigate the derivatives
of the analysis ensemble $\partial_{\theta}{x}_{t}^{n}$, by plugging in the
EnKF update formula:

$\begin{split}\partial_{\theta}{x}_{t}^{n}&\hskip 5.0pt=\hskip
5.0pt\partial_{\theta}(\widehat{x}_{t}^{\,n}+\widehat{K}_{t}(y_{t}+{\gamma}_{t}^{n}-H\widehat{x}_{t}^{\,n})\\\
&\hskip 5.0pt\underset{\text{\scriptsize1⃝}}{=}\hskip
5.0pt(I-\widehat{K}_{t}H)\partial_{\theta}\widehat{x}_{t}^{\,n}+(\partial_{\theta}\widehat{K}_{t})(y_{t}+{\gamma}_{t}^{n}-H\widehat{x}_{t}^{\,n})\\\
&\hskip 5.0pt=\hskip
5.0pt\big{(}I-\widehat{C}_{t}H^{\top}(H\widehat{C}_{t}H^{\top}+R)^{-1}H\big{)}\partial_{\theta}\widehat{x}_{t}^{\,n}+(\partial_{\theta}\widehat{K}_{t})(y_{t}+{\gamma}_{t}^{n}-H\widehat{x}_{t}^{\,n})\\\
&\xrightarrow[\text{\scriptsize2⃝}]{1/2}(I-\widehat{\mathsf{\bf{C}}}_{t}H^{\top}(H\widehat{\mathsf{\bf{C}}}_{t}H^{\top}+R)^{-1}H)\partial_{\theta}\widehat{\mathsf{\bf{x}}}_{t}^{\,n}+(\partial_{\theta}\widehat{\mathsf{\bf{K}}}_{t})(y_{t}+{\gamma}_{t}^{n}-H\widehat{\mathsf{\bf{x}}}_{t}^{\,n})\\\
&\hskip 5.0pt=\hskip
5.0pt(I-\widehat{\mathsf{\bf{K}}}_{t}H)\partial_{\theta}\widehat{\mathsf{\bf{x}}}_{t}^{\,n}+(\partial_{\theta}\widehat{\mathsf{\bf{K}}}_{t})(y_{t}+{\gamma}_{t}^{n}-H\widehat{\mathsf{\bf{x}}}_{t}^{\,n})\\\
&\hskip 5.0pt\underset{\text{\scriptsize3⃝}}{=}\hskip
5.0pt\partial_{\theta}(\widehat{\mathsf{\bf{x}}}_{t}^{\,n}+\widehat{\mathsf{\bf{K}}}_{t}(y_{t}+{\gamma}_{t}^{n}-H\widehat{\mathsf{\bf{x}}}_{t}^{\,n})\\\
&\hskip 5.0pt=\hskip
5.0pt\partial_{\theta}{\mathsf{\bf{x}}}_{t}^{n}.\end{split}$ (B.13)

Equalities \scriptsize1⃝ and \scriptsize3⃝ follow from chain rule. For
\scriptsize2⃝ we have used the $L^{p}$ convergence of $\widehat{x}_{t}^{\,n}$
to $\widehat{\mathsf{\bf{x}}}_{t}^{\,n}$,
$\partial_{\theta}\widehat{x}_{t}^{\,n}$ to
$\partial_{\theta}\widehat{\mathsf{\bf{x}}}_{t}^{\,n}$, $\widehat{C}_{t}$ to
$\widehat{\mathsf{\bf{C}}}_{t}$, $\partial_{\theta}\widehat{K}_{t}$ to
$\partial_{\theta}\widehat{\mathsf{\bf{K}}}_{t}$, and
$(H\widehat{C}_{t}H^{\top}+R)^{-1}$ to
$(H\widehat{\mathsf{\bf{C}}}_{t}H^{\top}+R)^{-1}$, with convergence rate
$N^{-1/2}$, and the fact that $\widehat{\mathsf{\bf{x}}}_{t}^{\,n}$,
$\partial_{\theta}\widehat{\mathsf{\bf{x}}}_{t}^{\,n}$ and the Gaussian random
variable ${\gamma}_{t}^{n}$ have finite moments of any order, followed by
Lemma B.1. Both derivatives $\partial_{\theta}{x}_{t}^{n}$ and
$\partial_{\theta}{\mathsf{\bf{x}}}_{t}^{n}$ exist since $R\succ 0$. We also
have the moment bound on $\partial_{\theta}{\mathsf{\bf{x}}}_{t}^{n}$:

$\|\partial_{\theta}{\mathsf{\bf{x}}}_{t}^{n}\|\leq|I-\widehat{\mathsf{\bf{K}}}_{t}H|\|\partial_{\theta}\widehat{\mathsf{\bf{x}}}_{t}^{\,n}\|_{p}+|\partial_{\theta}\widehat{\mathsf{\bf{K}}}_{t}|(|y_{t}|+\|{\gamma}_{t}^{n}\|_{p}+|H|\|\widehat{\mathsf{\bf{x}}}_{t}^{\,n}\|_{p})\leq
c,$ (B.14)

since $\widehat{\mathsf{\bf{x}}}_{t}^{\,n}$,
$\partial_{\theta}\widehat{\mathsf{\bf{x}}}_{t}^{\,n}$ and the Gaussian random
variable ${\gamma}_{t}^{n}$ have finite moments of any order. Then,

$\begin{split}\partial_{\theta}\widehat{x}_{t+1}^{\,n}&\hskip 5.0pt=\hskip
5.0pt\partial_{\theta}(A{x}_{t}^{n}+S{\xi}_{t}^{n})\\\ &\hskip
5.0pt\underset{\text{\scriptsize1⃝}}{=}\hskip
5.0pt(\partial_{\theta}A){x}_{t}^{n}+A(\partial_{\theta}{x}_{t}^{n})+(\partial_{\theta}{S}){\xi}_{t}^{n}\\\
&\xrightarrow[\text{\scriptsize2⃝}]{1/2}(\partial_{\theta}A){\mathsf{\bf{x}}}_{t}^{n}+A(\partial_{\theta}{\mathsf{\bf{x}}}_{t}^{n})+(\partial_{\theta}S){\xi}_{t}^{n}\\\
&\hskip 5.0pt\underset{\text{\scriptsize3⃝}}{=}\hskip
5.0pt\partial_{\theta}(A{\mathsf{\bf{x}}}_{t}^{n}+S{\xi}_{t}^{n})\\\ &\hskip
5.0pt=\hskip
5.0pt\partial_{\theta}\widehat{\mathsf{\bf{x}}}_{t+1}^{\,n}.\end{split}$
(B.15)

Here equalities \scriptsize1⃝ and \scriptsize3⃝ follow from chain rule. For
\scriptsize2⃝ we have used the $L^{p}$ convergence of ${x}_{t}^{n}$ to
${\mathsf{\bf{x}}}_{t}^{n}$ and $\partial_{\theta}{x}_{t}^{n}$ to
$\partial_{\theta}{\mathsf{\bf{x}}}_{t}^{n}$. Both derivatives
$\partial_{\theta}\widehat{x}_{t+1}^{\,n}$ and
$\partial_{\theta}\widehat{\mathsf{\bf{x}}}_{t+1}^{\,n}$ exist since both
$\partial_{\theta}{x}_{t}^{n}$ and
$\partial_{\theta}{\mathsf{\bf{x}}}_{t}^{n}$ exist. We also have the moment
bound:

$\begin{split}\|\partial_{\theta}\widehat{\mathsf{\bf{x}}}_{t+1}^{\,n}\|_{p}&\leq|\partial_{\theta}A|\|{\mathsf{\bf{x}}}_{t}^{n}\|_{p}+|A|\|\partial_{\theta}{\mathsf{\bf{x}}}_{t}^{n}\|_{p}+|\partial_{\theta}S|\|{\xi}_{t}^{n}\|_{p}\leq
c,\end{split}$ (B.16)

since ${\mathsf{\bf{x}}}_{t}^{n}$,
$\partial_{\theta}{\mathsf{\bf{x}}}_{t}^{n}$ and Gaussian random variable
${\xi}_{t}^{n}$ have finite moments of any order. Thus Eq. B.6 is proved for
step $t+1$ and the induction step is finished, which concludes the proof of
the lemma. ∎

###### Proof of Theorem 3.2.

Recall the definition of $h_{t}$ Eq. A.7. It suffices to show that, for each
$t\geq 1$,

$\partial_{\theta}\Big{(}h_{t}(\widehat{m}_{t},\widehat{C}_{t})\Big{)}\xrightarrow{1/2}\partial_{\theta}\Big{(}h_{t}(\widehat{\mathsf{\bf{m}}}_{t},\widehat{\mathsf{\bf{C}}}_{t})\Big{)}.$
(B.17)

We first investigate the convergence of derivatives of $h_{t}$ w.r.t
$\widehat{m}_{t}$ and $\widehat{C}_{t}$. The derivatives are computed in Eq.
A.9:

$\begin{split}\partial_{m}h_{t}(\widehat{m}_{t},\widehat{C}_{t})&\hskip
5.0pt=\hskip
5.0pt-H^{\top}(H\widehat{C}_{t}H^{\top}+R)^{-1}(y_{t}-H\widehat{m}_{t})\\\
&\xrightarrow{1/2}-H^{\top}(H\widehat{\mathsf{\bf{C}}}_{t}H^{\top}+R)^{-1}(y_{t}-H\widehat{\mathsf{\bf{m}}}_{t})\\\
&\hskip 5.0pt=\hskip
5.0pt\partial_{m}h_{t}(\widehat{\mathsf{\bf{m}}}_{t},\widehat{\mathsf{\bf{C}}}_{t}),\end{split}$
(B.18)

and

$\begin{split}\partial_{C}h_{t}(\widehat{m}_{t},\widehat{C}_{t})&\hskip
5.0pt=\hskip 5.0pt-\frac{1}{2}H^{\top}(H\widehat{C}_{t}H^{\top}+R)^{-1}H\\\
&\quad\quad\quad\quad+\frac{1}{2}H^{\top}(H\widehat{C}_{t}H^{\top}+R)^{-1}(y_{t}-H\widehat{m}_{t})(y_{t}-H\widehat{m}_{t})^{\top}(H\widehat{C}_{t}H^{\top}+R)^{-1}H\\\
&\xrightarrow{1/2}-\frac{1}{2}H^{\top}(H\widehat{\mathsf{\bf{C}}}_{t}H^{\top}+R)^{-1}H\\\
&\quad\quad\quad\quad+\frac{1}{2}H^{\top}(H\widehat{\mathsf{\bf{C}}}_{t}H^{\top}+R)^{-1}(y_{t}-H\widehat{\mathsf{\bf{m}}}_{t})(y_{t}-H\widehat{\mathsf{\bf{m}}}_{t})^{\top}(H\widehat{\mathsf{\bf{C}}}_{t}H^{\top}+R)^{-1}H,\end{split}$
(B.19)

where we have used the $L^{p}$ convergence of $\widehat{m}_{t}$ to
$\widehat{\mathsf{\bf{m}}}_{t}$ and $(H\widehat{C}_{t}H^{\top}+R)^{-1}$ to
$(H\widehat{\mathsf{\bf{C}}}_{t}H^{\top}+R)^{-1}$ by Eq. B.4, followed by
Lemma B.1. Then, by chain rule,

$\begin{split}\partial_{\theta}\Big{(}h_{t}(\widehat{m}_{t},\widehat{C}_{t})\Big{)}&\hskip
5.0pt=\hskip
5.0pt\Big{(}\partial_{m}h_{t}(\widehat{m}_{t},\widehat{C}_{t})\Big{)}^{\top}\partial_{\theta}\widehat{m}_{t}+\text{Tr}\bigg{(}\Big{(}\partial_{C}h_{t}(\widehat{m}_{t},\widehat{C}_{t})\Big{)}^{\top}\partial_{\theta}\widehat{C}_{t}\bigg{)}\\\
&\xrightarrow{1/2}\Big{(}\partial_{m}h_{t}(\widehat{\mathsf{\bf{m}}}_{t},\widehat{\mathsf{\bf{C}}}_{t})\Big{)}^{\top}\partial_{\theta}\widehat{\mathsf{\bf{m}}}_{t}+\text{Tr}\bigg{(}\Big{(}\partial_{C}h_{t}(\widehat{\mathsf{\bf{m}}}_{t},\widehat{\mathsf{\bf{C}}}_{t})\Big{)}^{\top}\partial_{\theta}\widehat{\mathsf{\bf{C}}}_{t}\bigg{)}\\\
&\hskip 5.0pt=\hskip
5.0pt\partial_{\theta}\Big{(}h_{t}(\widehat{\mathsf{\bf{m}}}_{t},\widehat{\mathsf{\bf{C}}}_{t})\Big{)}.\end{split}$
(B.20)

Both derivatives exist since $\partial_{\theta}\widehat{m}_{t}$,
$\partial_{\theta}\widehat{\mathsf{\bf{m}}}_{t}$,
$\partial_{\theta}\widehat{C}_{t}$ and
$\partial_{\theta}\widehat{\mathsf{\bf{C}}}_{t}$ exist, by Lemma B.3. We have
used Eq. B.18 and Eq. B.19 above, the $L^{p}$ convergence of
$\partial_{\theta}\widehat{m}_{t}$ to
$\partial_{\theta}\widehat{\mathsf{\bf{m}}}_{t}$ and
$\partial_{\theta}\widehat{C}_{t}$ to
$\partial_{\theta}\widehat{\mathsf{\bf{C}}}_{t}$ with rate $N^{-1/2}$ by Lemma
B.3, followed by Lemma B.1. ∎

######  Remark B.4.

We again emphasize that all the derivatives and chain rule formulas do _not_
need to be computed by hand in applications, but rather through the modern
autodiff libraries. We list them out only for the purpose of proving
convergence results.

##  Appendix C Additional Figures

See Fig. 9 and 10 for additional figures to the linear Gaussian experiment
Section 5.1:

Figure 9: Relative $L^{2}$ estimation errors of the log-likelihood (left) and
its gradient w.r.t. $\alpha$ (middle) and $\beta$ (right), computed using EnKF
and PF, as a function of $N$, for the linear-Gaussian model (5.1). State
dimension $d_{x}\in\\{20,40,80\\}$. $\theta$ is evaluated at
$\alpha=(0.5,0.5,0.5),\beta=(1,0.1)$. (Section 5.1.1) Figure 10: Relative
$L^{2}$ estimation errors of log-likelihood (left) and its gradient w.r.t.
$\alpha$ (middle) and $\beta$ (right), computed using EnKF and PF, with
different covariance tapering radius applied to EnKF. State dimension
$d_{x}=80$. $\theta$ is evaluated at $\alpha=(0.5,0.5,0.5),\beta=(1,0.1)$.

##  Appendix D Auto-Differentiable Particle Filters

Here we review the auto-differentiable particle filter methods introduced in,
e.g. [59, 55, 51]. A particle filter (PF) propagates $N$ weighted particles
$({w}_{t}^{1:N},{x}_{t}^{1:N})$ to approximate the filtering distribution
$p_{\theta}(x_{t}|y_{1:t})$. It first specifies a proposal distribution
$r(x_{t}|x_{t-1},y_{t}).$ Then it iteratively samples new particles from the
proposal, followed by a reweighing procedure:

$\widehat{x}_{t}^{\,n}\sim
r(\cdot|{x}_{t-1}^{n},y_{t}),\quad\quad{\widehat{w}}_{t}^{n}=\frac{p_{\theta}(\widehat{x}_{t}^{\,n}|{x}_{t-1}^{n})\mathcal{N}(y_{t};H\widehat{x}_{t}^{\,n},R)}{r(\widehat{x}_{t}^{\,n}|{x}_{t-1}^{n},y_{t})}{w}_{t-1}^{n}.$
(D.1)

A resampling step is performed when necessary, e.g., if the effective sample
size drops below a certain threshold:

$({w}_{t}^{n},{x}_{t}^{n})=\Bigl{(}\frac{1}{N},\widehat{x}_{t}^{{a}_{t}^{n}}\Bigr{)},\quad\quad\text{where
}{a}_{t}^{n}\sim\text{Categorical}\bigl{(}\cdot|{\widehat{w}}_{t}^{1:N}\bigr{)}.$
(D.2)

We refer the reader to [24] for a more general introduction to PF. Here we
consider PF with the optimal proposal
$r(\widehat{x}_{t}^{\,n}|{x}_{t-1}^{n},y_{t}):=p_{\theta}(\widehat{x}_{t}^{\,n}|{x}_{t-1}^{n},y_{t})$,
as it is readily available for the family of SSMs that is considered in this
paper [23, 75]. The unnormalized weights ${\widehat{w}}_{t}^{n}$ also provides
an approximation to the data log-likelihood. Define

$\mathcal{L}_{{\footnotesize\text{PF}}}(\theta)=\sum_{t=1}^{T}\log\Big{(}\sum_{n=1}^{N}{\widehat{w}}_{t}^{n}\Big{)}.$
(D.3)

Note that the weights ${\widehat{w}}_{t}^{n}$ depends implicitly on $\theta$.
It is a well-known result [19, 5] that the data likelihood
$\exp\mathcal{L}(\theta)$ can be unbiasedly estimated, and a lower bound for
$\mathcal{L}(\theta)$ can be derived:

$\operatorname{\mathbb{E}}[\exp\mathcal{L}_{{\footnotesize\text{PF}}}(\theta)]=\exp\mathcal{L}(\theta),\quad\quad\operatorname{\mathbb{E}}[\mathcal{L}_{{\footnotesize\text{PF}}}(\theta)]\leq\mathcal{L}(\theta),$
(D.4)

where the last inequality follows from Jensen’s inequality. We summarize the
optimal PF algorithm and parameter learning procedure (AD-PF) in Alg. D.1 and
Alg. D.2 below. AD-PF with truncated backprop (AD-PF-T) can be built in a
similar fashion as AD-EnKF-T (Alg. 4.2), and is omitted. For a derivation of
update rules of optimal PF, see e.g. [75]. For AD-PF where the proposal is
learned from data, see e.g. [59, 55, 51]. Following the same authors, the
gradients from the resampling steps are discarded. Therefore,
$\nabla_{\theta}\mathcal{L}_{{\footnotesize\text{PF}}}$ is biased and
inconsistent as $N\rightarrow\infty$ (see [17]).

Algorithm D.1 Optimal Particle Filtering and Log-likelihood Estimation

1:Input: $\theta=\\{\alpha,\beta\\},y_{1:T},(w_{0}^{1:N},x_{0}^{1:N}).$ (If
$(w_{0}^{1:N},x_{0}^{1:N})$ is not specified,
$x_{0}^{n}\overset{\text{i.i.d.}}{\sim}p_{0}(x_{0})$, $w_{0}^{n}=1/N$.)

2:Initialize $\mathcal{L}_{{\footnotesize\text{PF}}}(\theta)=0.$

3:Set $S=HQ_{\beta}H^{\top}+R$ and $K=Q_{\beta}H^{\top}S^{-1}.$

4:for $t=1,\ldots,T$ do

5: Set
$\widehat{x}_{t}^{\,n}=(I-KH)F_{\alpha}({x}_{t-1}^{n})+Ky_{t}+{\xi}_{t}^{n}$,
where ${\xi}_{t}^{n}\sim\mathcal{N}(0,Q_{\beta})$.

6: Set
$\widehat{w}_{t}^{\,n}=\mathcal{N}\bigl{(}y_{t};HF_{\alpha}({x}_{t-1}^{n}),S\bigr{)}$

7: Set
$({w}_{t}^{n},{x}_{t}^{n})=\Bigl{(}\frac{1}{N},\widehat{x}_{t}^{{a}_{t}^{n}}\Bigr{)}$,
where
${a}_{t}^{n}\sim\text{Categorical}\bigl{(}\cdot|{\widehat{w}}_{t}^{1:N}\bigr{)}$
$\triangleright$ Resampling step

8: Set
$\mathcal{L}_{{\footnotesize\text{PF}}}(\theta)\leftarrow\mathcal{L}_{{\footnotesize\text{PF}}}(\theta)+\log\Big{(}\sum_{n=1}^{N}\widehat{w}_{t}^{\,n}\Big{)}.$

9:end for

10:Output: PF particles $(w_{0:T}^{1:N},x_{0:T}^{1:N})$. Log-likelihood
estimate $\mathcal{L}_{{\footnotesize\text{PF}}}(\theta)$.

Algorithm D.2 Auto-differentiable Particle Filtering (AD-PF)

1:Input: Observations $y_{1:T}$. Learning rate $\eta$.

2:Initialize SSM parameter $\theta^{0}$ and set $k=0.$

3:while not converging do

4:
$(w_{0:T}^{1:N},x_{0:T}^{1:N}),\mathcal{L}_{{\footnotesize\text{PF}}}(\theta^{k})=\textsc{ParticleFilter}(\theta^{k},y_{1:T}).$
$\triangleright$ Alg. D.1

5: Compute $\nabla_{\theta}\mathcal{L}_{{\footnotesize\text{PF}}}(\theta^{k})$
by auto-differentiating the map
$\theta^{k}\mapsto\mathcal{L}_{{\footnotesize\text{PF}}}(\theta^{k}).$

6: Set
$\theta^{k+1}=\theta^{k}+\eta\nabla_{\theta}\mathcal{L}_{{\footnotesize\text{PF}}}(\theta^{k})$
and $k\leftarrow k+1.$

7:end while

8:Output: Learned SSM parameter $\theta^{k}$ and PF particles
$(w_{0:T}^{1:N},x_{0:T}^{1:N})$.

##  Appendix E Expectation-Maximization with Ensemble Kalman Filters

Here we review the EM methods introduced in, e.g., [65, 9]. Consider the log-
likelihood expression

$\displaystyle\log p_{\theta}(y_{1:T})$ $\displaystyle=\log\int
p_{\theta}(y_{1:T},x_{0:T})\text{d}x_{0:T}$ (E.1)
$\displaystyle=\log\operatorname{\mathbb{E}}_{q(x_{0:T}|y_{1:T})}\Big{[}\frac{p_{\theta}(y_{1:T},x_{0:T})}{q(x_{0:T}|y_{1:T})}\Big{]}$
(E.2)
$\displaystyle\geq\operatorname{\mathbb{E}}_{q(x_{0:T}|y_{1:T})}\big{[}\log
p_{\theta}(y_{1:T},x_{0:T})-\log
q(x_{0:T}|y_{1:T})\big{]}:=\mathcal{L}(q,\theta).$ (E.3)

The algorithm maximizes $\mathcal{L}(q,\theta)$, which alternates between two
steps:

  1. 1.

(E-step) Fix $\theta^{k}$. Find $q^{k}$ that maximizes
$\mathcal{L}(q,\theta^{k})$, which can be shown to be the exact posterior

$q^{k}(x_{0:T}|y_{1:T})=p_{\theta^{k}}(x_{0:T}|y_{1:T}).$

  2. 2.

(M-step) Fix $q^{k}$. Find $\theta^{k+1}$ that maximizes
$\mathcal{L}(q^{k},\theta)$:

$\theta^{*}=\operatorname*{arg\,max}_{\theta}\operatorname{\mathbb{E}}_{q^{k}(x_{0:T}|y_{1:T})}\big{[}\log
p_{\theta}(y_{1:T},x_{0:T})\big{]}.$ (E.4)

The exact posterior in the E-step is not always available, so one can run an
Ensemble Kalman Smoother (EnKS) to approximate the distribution
$p_{\theta^{k}}(x_{0:T}|y_{1:T})$ by a group of equally-weighted particles
$\\{x_{0:T}^{k,n}\\}_{n=1}^{N}$, and replace the expectation in the M-step by
an Monte-Carlo average among these particles.

Consider again the nonlinear SSM defined by equations (2.1)-(2.2)-(2.3) with
dynamics $F_{\alpha}$ and noise covariance $Q_{\beta}.$ The key to the methods
proposed in, e.g., [65, 9] is that, for the optimization problem Eq. E.4 in
the M-steps, optimizing w.r.t. $\alpha$ is hard, but optimizing w.r.t. $\beta$
is easy. If we write $\\{x_{0:T}^{n}\\}_{n=1}^{N}$ as an approximation of
$q(x_{0:T}|y_{1:T})$ (we drop $k$ here for notation convenience), the
optimization problem Eq. E.4 can be rewritten as (using Monte-Carlo
approximation)

$\displaystyle\operatorname*{arg\,max}_{\theta}\frac{1}{N}\sum_{n=1}^{N}\log
p_{\theta}(y_{1:T},x_{0:T}^{n})$
$\displaystyle=\operatorname*{arg\,max}_{\alpha,\beta}\frac{1}{N}\sum_{n=1}^{N}\sum_{t=1}^{T}\log\mathcal{N}\bigl{(}{x}_{t}^{n};F_{\alpha}({x}_{t-1}^{n}),Q_{\beta}\bigr{)}.$

When $\alpha$ is given, optimizing w.r.t. $\beta$ becomes the usual MLE
computation for covariance matrix of a multivariate Gaussian. We summarize the
EM algorithm in Algorithm E.1. Note that this special type of M-step update
only applies to SSMs with transition kernel of the form Eq. 2.1. For more
general SSMs, updates for $\beta$ may not admit a closed form expression, and
one may have to resort to a joint optimzation of $(\alpha,\beta)$ in the
M-steps.

Algorithm E.1 Expectation Maximization (EM) with Ensemble Kalman Filter

1:Input: Observations $y_{1:T}$. Learning rate $\eta$. Smoothing parameter
$S$. Inner-loop gradient ascent steps $J$.

2:Initialize SSM parameter $\theta^{0}=(\alpha^{0},\beta^{0})$. Initialize
$k=0.$

3:while not converging do

4: $x_{0:T}^{1:N}=\textsc{EnsembleKalmanFilter}(\theta^{k},y_{1:T})$.
$\triangleright$ Alg. 3.1

5: Perform smoothing for lag $S$ (see e.g. [42]) and update $x_{0:T}^{1:N}$.

6: Set
$Q_{\beta^{k+1}}=\frac{1}{NT}\sum_{n=1}^{N}\sum_{t=1}^{T}\Big{(}{x}_{t}^{n}-F_{\alpha^{k}}({x}_{t-1}^{n})\Big{)}\Big{(}{x}_{t}^{n}-F_{\alpha^{k}}({x}_{t-1}^{n})\Big{)}^{\top}.$

7: Initialize $\alpha^{k,0}=\alpha^{k}.$

8: for $j=0,\dots,J-1$ do

9:
$\alpha^{k,j+1}=\alpha^{k,j}+\eta\nabla_{\alpha}\mathcal{L}_{{\footnotesize\text{EM-
EnKF}}}(\alpha^{k,j},x_{0:T}^{1:N}),$ where
$\mathcal{L}_{{\footnotesize\text{EM-
EnKF}}}(\alpha,x_{0:T}^{1:N}):=\frac{1}{N}\sum_{n=1}^{N}\sum_{t=1}^{T}\log\mathcal{N}({x}_{t}^{n};F_{\alpha}({x}_{t-1}^{n}),Q_{\beta^{k+1}}).$

10: end for

11: Set $\alpha^{k+1}=\alpha^{k,J}$, $k\leftarrow k+1.$

12:end while

13:Output: Learned SSM parameter $\theta^{k}=(\alpha^{k},\beta^{k})$ and
particles $x_{0:T}^{1:N}.$

##  Appendix F Implementation Details and Additional Performance Metrics

###  F.1 Linear-Gaussian Model

The relative $L^{2}$ errors of the log-likelihood and gradient estimates in
Section 5.1.1 are given by

$\frac{\operatorname{\mathbb{E}}[(\mathcal{L}_{{\footnotesize\text{EnKF}}}(\theta)-\mathcal{L}(\theta))^{2}]^{1/2}}{|\mathcal{L}(\theta)|},$
(F.1)

and

$\frac{\operatorname{\mathbb{E}}[|\nabla_{\alpha}\mathcal{L}_{{\footnotesize\text{EnKF}}}(\theta)-\nabla_{\alpha}\mathcal{L}(\theta)|^{2}]^{1/2}}{|\nabla_{\alpha}\mathcal{L}(\theta)|},\quad\quad\frac{\operatorname{\mathbb{E}}[|\nabla_{\beta}\mathcal{L}_{{\footnotesize\text{EnKF}}}(\theta)-\nabla_{\beta}\mathcal{L}(\theta)|^{2}]^{1/2}}{|\nabla_{\beta}\mathcal{L}(\theta)|}.$
(F.2)