# RSINET: INPAINTING REMOTELY SENSED IMAGES USING TRIPLE GAN FRAMEWORK

###### Abstract

We tackle the problem of image inpainting in the remote sensing domain. Remote
sensing images possess high resolution and geographical variations, that
render the conventional inpainting methods less effective. This further
entails the requirement of models with high complexity to sufficiently capture
the spectral, spatial and textural nuances within an image, emerging from its
high spatial variability. To this end, we propose a novel inpainting method
that individually focuses on each aspect of an image such as edges, colour and
texture using a task specific GAN. Moreover, each individual GAN also
incorporates the attention mechanism that explicitly extracts the spectral and
spatial features. To ensure consistent gradient flow, the model uses residual
learning paradigm, thus simultaneously working with high and low level
features. We evaluate our model, alongwith previous state of the art models,
on the two well known remote sensing datasets, Open Cities AI and Earth on
Canvas, and achieve competitive performance. The code can be referred here:
https://github.com/advaitkumar3107/RSINet.

Index Terms—  Image inpainting, remote sensing, generative adversarial
networks

##  1 Introduction

Image inpainting is the process of filling in the missing part or conserving
the damaged and deteriorated image (which can be physical or digital). Digital
image inpainting is an important problem statement in the field of computer
vision which has applications in various domains such as restoring damaged
images/videos, remote sensing, object removal, text removal, automatic
modifications of images/videos, image compression and super resolution [1].
The drawback of the traditional mathematical methods is their low PSNR on
complex images such as remote sensing images [2]. Remote sensing images have a
variety of boundaries, objects and colours which makes it a semantically
tougher problem. Hence, deep learning has been introduced earlier to tackle
this challenge such as in [2], which solved three missing information tasks in
remote sensing data using a deep convolutional network combined with spatio-
temporal information. Similarly, [3] used a generative approach for image
inpainting. These images looked quite similar to the original version.

There have been several image inpainting approaches as well, such as [4],
which used a UNet based encoder-decoder structure. Furthermore, [4] also used
residual learning to compute the gradients effectively for image restoration.
[5] built on the above idea and used guidance loss for inpainting. [6]
proposed a new type of layer called ‘partial convolution’ (PartialConv) to
improve the current best performing image inpainting models. However, all the
above models used simple baselines, which resulted in lower PSNR on complex
satellite images. [7] proposed the EdgeConnect (EC) model, that initially
builds up the outline of the image to be restored and then fills it with
details. Furthermore, to ensure the relationship between the neighbouring
areas and entire image as a whole, [8] proposed a multi GAN approach, that
trains a global as well as a local discriminator (GLCIC) on top of their VGG-
type model [9]. Moreover, [10] proposed a global and local attention based
model to tackle the issue of image irregularities such as holes. In addition,
to preserve the semantic style of the original images, additional loss
functions have been explored. For instance, [11] introduces style loss, where
the textural features are synthesised from the images, which assist in the
style transfer. Similarly, [12] introduces perceptual loss for super-
resolution, (instead of pixelwise loss) and get high resolution images.

In the aforementioned methods, even though the methods seemed to perform
better individually on the task specific images, they could not prove much
effective on the images with high spatial variation such as those from remote
sensing domain. This is because each of the individual models was designed to
focus on a specific aspect of the image such as colour, edge or texture.
Hence, there arises a need to have a common model that simultaneously works
with all the image aspects and leads to more accurate image restoration.
Inspired from this notion, we propose an image inpainting approach that
utilizes multiple GANs to effectively capture the different aspects of the
image. In addition, to handle irregularities of the image defects,
Convolutional Block Attention Module (CBAM) [13] and the Gated Attention
Layers [14] based attention modules are added. The GANs are also reinforced
with skip connection to ensure a consistent gradient flow. Moreover, to
preserve the semantic style of the images and get more robust representation,
we incorporate a cocktail of adversarial, style and perceptual loss. Our
approach can be summarised as follows:

  1. 1.

We introduce a model that combines the characteristics of edge detection GAN,
colour filling GAN and an global GAN to focus on the different image aspects.

  2. 2.

We introduce attention layers for guided backpropagation to get more robust
spectral-spatial representation and increase the sharpness of the images,
while simultaneously incorporating style and perceptual losses with the
adversarial loss to capture the semantic information of the images.

  3. 3.

We train our model on the Open Cities AI [15] and the Earth on Canvas [16]
datasets where our model outperforms the state of the art deep inpainting
models by at least 0.06 % and 2.24 % respectively for the two datasets.

##  2 Model Description

Fig. 1: Architecture of our proposed GAN. The Generator consists of an
encoder-residual blocks-decoder structure. The encoder has a skip connection
as well as feature maps that serve as a gate for the gated attention layers.
The Discriminator consists of 6 convolutional layers that successively
downsample the input into a 28 $\times$ 28 probabilistic grid.

In this section we discuss our proposed model. It consists of three GANs
trained on top of each other, namely, the edge completion GAN
$\mathcal{G}_{1}$, colour filling GAN $\mathcal{G}_{2}$ and the global GAN
$\mathcal{G}_{3}$. We describe the loss functions and the architecture used
for the basic GAN unit in detail in the subsequent sections. The GAN
architecture used in RSINet is presented in Fig. 1. All the GANs follow the
similar architecture.

Let $I_{gt}$ be the groundtruth image and, $C_{gt}$ and $I_{gray}$ be the
groundtruth edge map (after applying a canny edge detector [17] to the true
image) and grayscale image respectively. In the generator for
$\mathcal{G}_{1}$, we use the masked grayscale image
$\tilde{I}_{gray}=I_{gray}\odot(1-M)$, its masked edge map
$\tilde{C}_{gt}=C_{gt}\odot(1-M)$ ($\odot$ being the Hadamard product), and
the image masks ${M}$ as the pre-condition (proposed in [18] as input to
pix2pix GAN, where 1 and 0 denote the masked region and the background
respectively). The generator produces the completed edge map for the image
(see Eqn. 1).

$C_{pred}=\mathcal{G}_{1}(\tilde{I}_{gray},\tilde{C}_{gt},M)$ (1)

We use $C_{gt}$ and $C_{pred}$ conditioned on $I_{gray}$ as the input to the
discriminator, which predicts whether or not the edge map is real. This
predicted edge map is then passed on to the colour filling model. The colour
filling model uses the masked RGB image, $\tilde{I}_{gt}=I_{gt}\odot(1-M)$,
conditioned on the completed edge map $C_{comp}=C_{gt}\odot(1-M)+C_{pred}\odot
M$, as taken from the previous model. The model outputs the completed RGB
image with the right colours filled at the right places.

$I_{pred}=\mathcal{G}_{2}\left(\tilde{I}_{gt},C_{comp}\right)$ (2)

The previous two models gave blurry outputs on the tougher satellite images.
Hence, we propose to train a global GAN on top of the output from the colour
filling model which helps in refining the output and makes the image visually
sharper. This was based upon the suggestion [19] that using the same encoder-
decoder architectures for image inpainting were observed to give better
results since both models learnt the same features semantically.

$I_{refined}=\mathcal{G}_{3}(I_{comp},\tilde{I}_{gt})$ (3)

###  2.1 Loss Functions

The edge completion model is trained the adversarial loss [18] and the feature
matching loss [20]. Adversarial training is thought of as a min-max game
between two players. Here, the generator tries to ‘fool’ the discriminator,
i.e. make the discriminator predict with a high probability that the
generator’s output belongs to the input data, while the discriminator is
trained to differentiate between the original and the generated samples. The
adversarial loss is given in Eqn. 4.

$\begin{split}L_{adv}=\operatorname*{\mathbb{E}}_{C_{gt},I_{gray}}[\log\;\mathcal{D}_{1}(C_{gt},I_{gray})]+\\\
\operatorname*{\mathbb{E}}_{I_{gray}}[\log(1-\mathcal{D}_{1}(C_{pred},I_{gray}))]\end{split}$
(4)

In Eqn. 4, $\mathcal{D}_{1}$ is the discriminator for edge completion GAN.

The feature matching loss, $L_{FM}$ compares the activation maps in the
intermediate layers of the discriminator. This is similar to perceptual loss
[12] (used in the colour filling model), where activations are compared with
those from a pre-trained VGG network. However, a VGG network is not trained to
produce the edge maps, so we use the $L_{FM}$ instead of the $L_{perc}$ here.
The $L_{FM}$ is defined as :

$L_{FM}=\operatorname*{\mathbb{E}}\left[\sum\limits_{i=1}^{L}\;\frac{1}{N_{i}}\left\lVert\mathcal{D}_{1}^{(i)}\;(C_{gt})-\mathcal{D}_{1}^{(i)}\;(C_{pred})\right\rVert_{1}\right]$
(5)

Fig. 2: Ablation study on Open Cities AI dataset with rectangular masks to
compare the models by removing different components. As expected, the
performance is better with all the components.

Here $L$ is the number of layers in the discriminator. $N_{i}$ is the number
of elements in the $i^{th}$ convolutional layer and $\mathcal{D}_{1}^{(i)}$ is
the activation in the $i^{th}$ layer of the $\mathcal{D}_{1}$. Applying it to
both the $\mathcal{G}_{1}$ and $\mathcal{D}_{1}$ gave us better results, as
suggested in [7].

The colour filling model is trained on 4 losses. The adversarial loss, L1 loss
[21], style loss [11] and perceptual loss [12]. To ensure proper scaling the
L1 loss is normalized by the mask size. The adversarial loss is similar to the
previous model. The perceptual loss gives the L1 distance between the
activations from a few specific layers of a pre-trained network as well as our
model. We choose VGG-19 network (pretrained on ImageNet [9]) since its
architecture is similar to ours.

$L_{perc}=\operatorname*{\mathbb{E}}\left[\sum\limits_{i=1}\;\frac{1}{N_{i}}\left\lVert\phi_{i}\;(I_{gt})-\phi_{i}\;(I_{pred})\right\rVert_{1}\right]$
(6)

where $\phi_{i}$ represents the activation map of the $i^{th}$ layer of the
pre-trained VGG19 network. The style loss returns the differences in
covariances of these activation maps by constructing the gram matrix from the
activation maps. Given the feature maps of sizes $C_{j}\times H_{j}\times
W_{j}$, style loss is computed by

$L_{style}=\operatorname*{\mathbb{E}}_{j}\left[\left\lVert
G_{j}^{\phi}\left(\tilde{I}_{pred}\right)-G_{j}^{\phi}\left(\tilde{I}_{gt}\right)\right\rVert_{1}\right]$
(7)

where $G_{j}^{\phi}$ is the gram matrix constructed from activation maps
$\phi_{j}$. The global GAN model is also trained on the weighted sum of the 4
losses as discussed above, namely, $L_{adv}$, $L_{perc}$, $L1$ and
$L_{style}$. The final loss for RSINet is given in Eqn. 8. The value of
${\lambda}_{1}$, ${\lambda}_{2}$ and ${\lambda}_{3}$ is empirically fixed to 1
during implementation.

$L_{final}=L_{adv}+{\lambda}_{1}L1+{\lambda}_{2}L_{perc}+{\lambda}_{3}L_{style}$
(8)

##  3 Experiments

The model has a VGG type architecture [9] (an encoder followed by a decoder).
The encoder part consists of 3 convolutional layers each successively halving
the image dimensions and doubling the number of channels, followed by 8
residual blocks and finally 3 convolution layers again. This converts the
feature representation of the residual blocks back to the image size (the
output). A skip connection (consisting of 2 layers) has been added to the
second convolutional layer of the network. We have also added 3 gated
attention layers in between, for further refinement of the feature maps.

###  3.1 Datasets

We worked primarily with two famous satellite imagery datasets. The first one
is the Open Cities AI challenge dataset [15], which consists 500 randomly
sampled images size of 1024$\times$1024, divided into 16 images of size
256$\times$256\. The second one is the Earth on Canvas dataset [16] which
consists of 1400, 256$\times$256 images from 14 classes. It is highly
uncorrelated and a good test for our model’s performance. Both the datasets
have been divided into 60%, 20% and 20% ratio for train, validation and test
sets respectively.

###  3.2 Protocol

We kept the parameters for the global GAN similar to those of colour filling
GAN. Adam optimizer [22] is used for training all the models, with learning
rates of 10-3 and 10-4 for generator and discriminator respectively. We also
tried adding jigsaw training [23], that refers to splitting the input image
into square patches and then randomly shuffling those patches to get a new
transformed image. This image would be fed to the model and it would try to
reconstruct the original image. For evaluation, peak signal-to-noise ratio
(PSNR) is used. For the GLCIC model, its 160$\times$160 output was
interpolated to 256$\times$256 while computing the PSNR.

Fig. 3: Inpainting results on Earth on Canvas dataset with (a-b) rectangular masks, (c-d) salt-pepper noise, (e-f) irregular masks. For each example, the image on the left is the original image, the image in the centre has been corrupted by a mask while the one in the right is the reconstructed image after inpainting. Table 1:  Accuracy analysis for irregular masks on Open Cities AI and Earth on Canvas datasets for different methods | PSNR (DB)
---|---
Model Name |  Open Cities AI |  Earth on Canvas
EC original [7] |  37.113 |  33.145
EC (our model) |  37.433 |  33.784
Only colour filling GAN (our model) |  30.347 |  31.198
PartialConv [6] |  35.912 |  31.246
EC+Global GAN (our model) |  37.456 |  34.542

###  3.3 Discussions

As observed in Fig. 2, the baseline EC model (i.e. with edge completion GAN
and colour filling GAN) along with the global GAN trained on top of it,
outperforms the original EC as well as the GLCIC models on the rectangular
masks type ablation on both the datasets. The performance for all the models
is presented in Table 1. It is visible that our model outperforms all the
models for Open Cities AI dataset (highest PSNR of 37.456) and Earth on Canvas
dataset (highest PSNR of 34.542).

###  3.4 Critical Analysis

####  3.4.1 Ablations

We primarily dealt with 3 types of ablations; rectangular masks wherein a
random rectangle of the input image was replaced with white pixels, covering
between 5%-30% of the image area during training, salt and pepper masks where,
random image pixels as sampled from a gaussian distribution were whitened,
covering anywhere between 5%-95% of the image area and irregular masks, as
shown in [6]. Furthermore, we also present the visualizations of inpainting
for the different masks in Fig. 3 for the three kinds of masks. The method has
shown a good performance in reconstructing the image for salt and pepper, and
irregular masks. However, in case of rectangular mask, in Fig. 3 (b), the
model could not construct the aeroplane completely. This is because, the
method exploits the locally available information (using convolutions) for
reconstruction.

Fig. 4: Performance of various loss functions presented on Open Cities AI
dataset. It can be observed that relatively better performance is observed
when the model is trained on all the losses.

####  3.4.2 Effect of Loss Function

As can be seen from Fig. 4, a mixture of losses works best for inpainting
(presented on Open Cities AI). This hybrid loss function consisting of
perceptual, style, L1 and adversarial losses for the generator, and feature
matching and adversarial losses for the discriminator performs well for this
semantically tough image inpainting task. Perceptual and style losses help in
improving the quality of the reconstructed image by capturing image semantics
and minimising the difference of activations in all the layers of the target
and the output images. L1 loss improves the model performance due to its high
noise robustness. Feature matching loss used in the discriminator helps to
keep the GAN stable and prevent mode collapse. To show the diverse
performance, the ablation study is performed for different percentage of
rectangular masks.

##  4 Conclusion

We presented a novel approach to tackle the problem of image inpainting in
remote sensing domain. Our approach builds on the existing EdgeConnect model
and strengthens it by incorporating attention mechanism and global GAN, that
gives realistic results by combining global and local features. Furthermore,
we incorporate multiple losses such as perceptual loss and style loss, which
further boost our model’s performance. We evaluate our model on Open Cities AI
and Earth on Canvas datasets, where our approach gives competitive results
with respect to previous benchmarks. In future, we would consider more
semantically challenging datasets and explore the problem in a multimodal
scenario.

## Acknowledgement

The authors would like to acknowledgement the IITB-ISRO Grant
RD/0120-ISROC00-005.

## References

  * [1] Omar Elharrouss, Noor Almaadeed, Somaya Al-Maadeed, and Younes Akbari,  “Image inpainting: A review,”  Neural Processing Letters, pp. 1–22, 2019.
  * [2] Junyu Dong, Ruiying Yin, Xin Sun, Qiong Li, Yuting Yang, and Xukun Qin,  “Inpainting of remote sensing sst images with deep convolutional generative adversarial network,”  IEEE Geoscience and Remote Sensing Letters, vol. 16, no. 2, pp. 173–177, 2018.
  * [3] A. Kuznetsov and M. Gashnikov,  “Remote sensing image inpainting with generative adversarial networks,”  in 2020 8th International Symposium on Digital Forensics and Security (ISDFS), 2020, pp. 1–6.
  * [4] Yang Liu, Jinshan Pan, and Zhixun Su,  “Deep blind image inpainting,” 2017.
  * [5] Zhaoyi Yan, Xiaoming Li, Mu Li, Wangmeng Zuo, and Shiguang Shan,  “Shift-net: Image inpainting via deep feature rearrangement,” 2018.
  * [6] Guilin Liu, Fitsum A. Reda, Kevin J. Shih, Ting-Chun Wang, Andrew Tao, and Bryan Catanzaro,  “Image inpainting for irregular holes using partial convolutions,” 2018\.
  * [7] Kamyar Nazeri, Eric Ng, Tony Joseph, Faisal Z. Qureshi, and Mehran Ebrahimi,  “Edgeconnect: Generative image inpainting with adversarial edge learning,” 2019.
  * [8] Satoshi Iizuka, Edgar Simo-Serra, and Hiroshi Ishikawa,  “Globally and locally consistent image completion,”  ACM Trans. Graph., vol. 36, no. 4, July 2017.
  * [9] Karen Simonyan and Andrew Zisserman,  “Very deep convolutional networks for large-scale image recognition,” 2014.
  * [10] S. M. Nadim Uddin and Yong Jung,  “Global and local attention-based free-form image inpainting,”  Sensors, vol. 20, pp. 3204, 06 2020.
  * [11] Leon A. Gatys, Alexander S. Ecker, and Matthias Bethge,  “Image style transfer using convolutional neural networks,”  2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 2414–2423, 2016.
  * [12] Justin Johnson, Alexandre Alahi, and Li Fei-Fei,  “Perceptual losses for real-time style transfer and super-resolution,” 2016.
  * [13] Sanghyun Woo, Jongchan Park, Joon-Young Lee, and In So Kweon,  “Cbam: Convolutional block attention module,” 2018.
  * [14] Ozan Oktay, Jo Schlemper, Loic Le Folgoc, Matthew Lee, Mattias Heinrich, Kazunari Misawa, Kensaku Mori, Steven McDonagh, Nils Y Hammerla, Bernhard Kainz, Ben Glocker, and Daniel Rueckert,  “Attention U-Net: Learning where to look for the pancreas,” 2018\.
  * [15] Milind Naphade, Zheng Tang, Ming-Ching Chang, David C Anastasiu, Anuj Sharma, Rama Chellappa, Shuo Wang, Pranamesh Chakraborty, Tingting Huang, Jenq-Neng Hwang, et al.,  “The 2019 ai city challenge.,”  in CVPR Workshops, 2019, pp. 452–460.
  * [16] Ushasi Chaudhuri, Biplab Banerjee, Avik Bhattacharya, and Mihai Datcu,  “A zero-shot sketch-based intermodal object retrieval scheme for remote sensing images,”  IEEE Geoscience and Remote Sensing Letters, 2021.
  * [17] Paul Bao, Lei Zhang, and Xiaolin Wu,  “Canny edge detection enhancement by scale multiplication,”  IEEE transactions on pattern analysis and machine intelligence, vol. 27, no. 9, pp. 1485–1490, 2005.
  * [18] Phillip Isola, Jun-Yan Zhu, Tinghui Zhou, and Alexei A. Efros,  “Image-to-image translation with conditional adversarial networks,” 2016\.
  * [19] Suriya Singh, Anil Batra, Guan Pang, Lorenzo Torresani, Saikat Basu, Manohar Paluri, and CV Jawahar,  “Self-supervised feature learning for semantic segmentation of overhead imagery.,”  in BMVC, 2018, vol. 1, p. 4.
  * [20] Tim Salimans, Ian Goodfellow, Wojciech Zaremba, Vicki Cheung, Alec Radford, and Xi Chen,  “Improved techniques for training gans,” 2016.
  * [21] Hang Zhao, Orazio Gallo, Iuri Frosio, and Jan Kautz,  “Loss functions for neural networks for image processing,” 2015.
  * [22] Diederik P Kingma and Jimmy Ba,  “Adam: A method for stochastic optimization,”  arXiv preprint arXiv:1412.6980, 2014.
  * [23] Mehdi Noroozi and Paolo Favaro,  “Unsupervised learning of visual representations by solving jigsaw puzzles,” 2016.


# Learning to Assimilate in Chaotic Dynamical Systems

Michael McCabe
Department of Computer Science
University of Colorado Boulder
<EMAIL_ADDRESS>
&Jed Brown
Department of Computer Science
University of Colorado Boulder
<EMAIL_ADDRESS>

###### Abstract

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

##  1 Introduction

Forecasting in the geosciences is an initial value problem of enormous
practical significance. In high value domains like numerical weather
prediction [1], climate modeling [2], atmospheric chemistry [3], seismology
[4], and others [5, 6, 7], forecasts are produced by estimating the current
state of the dynamical system of interest and integrating that state forward
in time using numerical models based on the discretization of differential
equations. These numerical models are derived from physical principles and
possess desirable extrapolation and convergence properties. Yet despite the
efficacy of these models and the vast amount of compute power used in
generating these forecasts, obtaining highly accurate predictions is non-
trivial. Discretization introduces numerical errors and it is often too
computationally expensive to directly simulate the system of interest at the
resolution necessary to capture all relevant features.

Further complicating matters in geoscience applications is that many systems
of interest are chaotic, meaning that small errors in the initial condition
estimates can lead to significant forecasting errors over relatively small
time frame [8]. This can be problematic when the initial condition for a
forecast is current state of the Earth, a large-scale actively evolving system
that cannot be directly controlled. Modern numerical weather prediction (NWP)
models utilize hundreds of millions of state variables scattered over a three-
dimensional discretization of the Earth’s atmosphere [9]. To perfectly
initialize a simulation, one would need to know the true value of all of these
state variables simultaneously. This data is simply not available. In reality,
what is available are noisy, partial measurements scattered non-uniformly in
time and space.

The inverse problem of estimating the true state of a dynamical system from
imperfect observations is the target of our work and of the broader field
known as data assimilation. Data assimilation techniques produce state
estimates by systematically combining noisy observations with the numerical
model [10, 9]. The data assimilation problem differs from other state
estimation problems in that the numerical model acts both as an evolution
equation and as a mechanism for transporting information from regions where
observations are dense to regions where observations are sparse [11]. Thus for
accurate estimates of the full system state, it is important to utilize both
the model and observations.

One recent trend in deep learning research has been the push to develop neural
network models to replace expensive, oft repeated processes. This incurs a
large upfront cost to train the network in exchange for a faster solution to
subsequent iterations of the problem. These amortized methods were initially
developed for statistical inference [12] but have also been used in the
context of numerical simulation [13], and meta-learning [14]. Our work extends
this approach to variational data assimilation [15].

In this work, we introduce a self-supervised framework for learning to
assimilate which we call amortized assimilation. Here, we use the term self-
supervised to indicate that our method learns to assimilate entirely from
trajectories of noisy observations without the use of ground truth data during
training. Our design incorporates the objective flexibility of variational
assimilation but amortizes the expense of solving the nonlinear optimization
problem inherent to variational methods into the training of a neural network
that then behaves as a sequential filter.

Our contributions are both theoretical and empirical. In Section 3.1, we
introduce an amortized assimilation architecture based on the Ensemble Kalman
Filter [16]. We then develop the theory of amortized assimilation in Section
3.2 by extending the powerful self-supervised denoising results of Batson and
Royer [17] to the dynamical systems setting through the use of differentiable
simulation. We then support this theory through a set of numerical
experiments111Code available at
https://github.com/mikemccabe210/amortizedassimilation. in Section 5, where we
see that amortized assimilation methods match or outperform conventional
approaches across several benchmark systems with especially strong performance
at smaller ensemble sizes.

##  2 Preliminaries

#### Problem Setting

Figure 1: Observations are assumed to be generated from a Hidden Markov Model.

In this work, we consider a dynamical system with state variable
$\textbf{x}(t)\in C\subset\mathbb{R}^{d}$ where $C$ is a compact subset of
$\mathbb{R}^{d}$ and system dynamics defined by the differential equation:

$\displaystyle\frac{d\textbf{x}}{dt}=g(\textbf{x}(t))$ (1)

with $g$ as a time-invariant, Lipschitz continuous map from
$\mathbb{R}^{d}\to\mathbb{R}^{d}$. The state is observed at a sequence of $K$
discrete time points $\\{\tau_{0},\tau_{1},\dots,\tau_{K}\\}$ generating the
time-series
$\textbf{y}_{0:K}=\\{\textbf{y}_{0},\textbf{y}_{1},\dots,\textbf{y}_{K}\\}$
where $0\leq k\leq K$ is an index corresponding to the time point $\tau_{k}$.
These observations are imperfect representations of the system states
$\textbf{x}_{0:K}=\\{\textbf{x}_{0},\textbf{x}_{1},\dots\textbf{x}_{K}\\}$
generated from observation operators of the form:

$\displaystyle\textbf{y}_{k}$
$\displaystyle=\mathcal{H}_{k}(\textbf{x}_{k})+\bm{\eta}_{k}$ (2)

where $\mathcal{H}_{k}$ is an arbitrary potentially nonlinear function and
$\bm{\eta}_{k}\sim\mathcal{N}(0,\sigma_{k}^{2}\textbf{I})$. This noise model
is often referred to as a Hidden Markov Model (Figure 1). In the classical
assimilation setting, the observation operators are assumed to be known a
priori though our proposed method is capable of learning a subset of
observation operators. For a problem like numerical weather prediction, these
operators may represent measurements taken of certain atmospheric quantities
at a particular measurement location.

Data assimilation is then the inverse problem of estimating the true
trajectory $\textbf{x}_{0:K}$ from our noisy observations $\textbf{y}_{0:K}$
and our model of the system evolution:

$\displaystyle\textbf{x}(\tau_{k+1})=\mathcal{M}_{k:k+1}(\textbf{x}(\tau_{k}))=\textbf{x}(\tau_{k})+\int_{\tau_{k}}^{\tau_{k+1}}g(\textbf{x}(t))\
dt.$ (3)

Figure 2: Ensemble filters model the time-evolution of uncertainty under
nonlinear dynamics by maintaining a set of samples from the state distribution
and simulating their trajectories forward in time. Observations are then used
to refine these forecast estimates in what is called the analysis step.

#### Sequential Filtering

In the absence of ground truth, it is practical to present data assimilation
from a statistical perspective. In particular, the quantity of interest that
we are concerned with in this paper is the filtering distribution
$p(\textbf{x}_{k}\>|\>{\textbf{y}_{0:k}})$. Efficient algorithms for
estimating the filtering distribution are derived from exploiting dependence
relationships in the generative model to produce the two step cycle:

Forecast: $\displaystyle
p(\textbf{x}_{k}\>|\>{\textbf{y}_{0:k-1}})=\int_{\mathbb{R}^{d}}p(\textbf{x}_{k}\>|\>{\textbf{x}_{k-1}})\
d\textbf{x}_{k-1}$ (4a) Analysis: $\displaystyle
p(\textbf{x}_{k}\>|\>{\textbf{y}_{0:k}})\propto
p(\textbf{y}_{k}\>|\>{\textbf{x}_{k}})p(\textbf{x}_{k}\>|\>{\textbf{y}_{0:k-1}})$
(4b)

When evaluating data assimilation methods, the objective function is often
evaluated with respect to the “analysis” estimates as these act as these are
the initial conditions used for the next forecast phase. In subsequent
sections, we will use $\textbf{x}^{f}_{k}$ and $\textbf{x}^{a}_{k}$ to refer
to forecast and analysis estimates of $\textbf{x}_{k}$ respectively.

In the case of linear dynamics with Gaussian uncertainties, the Kalman filter
[18] provides both the forecast and analysis distributions exactly in closed
form. However, the general case is significantly more complicated. Under
nonlinear dynamics, computing the forecast distribution requires the solution
of the Fokker-Planck partial differential equation [19]. While finding an
exact solution to these equations is computationally infeasible, one can
produce what is called an ensemble estimate of the forecast distribution using
Sequential Monte Carlo methods [20].

Ensemble filters (Figure 2) replace an exact representation of the uncertainty
over states with an empirical approximation in the form of a small number of
samples called particles or ensemble members. The forecast distribution can
then be estimated by integrating the dynamics forward in time for each
ensemble member independently. Allowing for full generality leads to particle
filter approaches while enforcing Gaussian assumptions leads to the Ensemble
Kalman filter (EnKF) [16]. The former is rarely used in data assimilation
settings as particle filters can become degenerate in high dimensions [21]
while the latter has been quite successful in practice. Like the classical
Kalman filter, the EnKF assumes that both the state distribution and
observation likelihoods are Gaussian. The latter is often reasonable, but the
former is a rather strong assumption under nonlinear dynamics.

Ensemble methods possess an inherent trade-off in that larger ensembles more
accurately model uncertainty, but each ensemble member requires independently
integrating the dynamics forward in time which can be significantly more
expensive than the assimilation step itself. Bayesian filters rely on robust
covariance estimates to maintain stability which can be difficult to obtain
when using a small ensemble. Heuristics such as covariance inflation [22] and
localization [23] have been developed to improve stability at smaller ensemble
sizes, but there remains value in developing methods that improve upon the
accuracy of these approaches.

#### Variational Assimilation

The alternative to sequential filters is variational assimilation. Variational
assimilation methods, like 4D-Var [24], directly solve for the initial
conditions of a system by finding the state $\textbf{x}_{0}$ which minimizes
the negative log posterior density through nonlinear optimization. This is
expensive, but gradient-based optimization utilizing the adjoint of the
numerical model avoids the persistent Gaussian state assumptions necessary for
tractable Bayesian filters. Furthermore, these methods can be augmented by
auxiliary objectives promoting concordance with physical properties like
conservation laws.

##  3 Amortized Assimilation

Figure 3: The AmEnF replaces the analysis step of a conventional ensemble
filter with a neural network trained using historical data. The model uses an
augmented state which includes both standard inputs as well as recurrent
memory.

###  3.1 Amortized Ensemble Filters

Our proposed approach attempts to combine the flexibility of variational
methods with the efficiency of sequential filters by training a neural network
to act as an analysis update mechanism. Before diving into the theory, we
present the general amortized assimilation framework from the perspective of a
specific architecture which we call the Amortized Ensemble Filter (AmEnF). The
AmEnF (Figure 3) replaces the EnKF analysis equations with a parameterized
function in the form of a neural network whose weights are optimized for a
specific dynamical system during training. By unrolling a fixed number of
assimilation cycles as first proposed in Haarnoja et al. [25], the AmEnF can
be trained efficiently using backpropagation through time [26]. The loss used
for this training will be discussed in more detail in Section 3.2.

One advantage of decoupling from the statistical model in this way is that it
allows us to incorporate more information into our analysis process. Like the
EnKF, the analysis ensemble members are computed using the forecast values
$\textbf{x}^{f}_{i,k}$ for each ensemble member $i\in[m]$ with $m$ denoting
the number of ensemble members, the ensemble covariance $P^{f}_{k}$, and the
observations $\textbf{y}_{k}$. However, we augment these features with the
dynamics $g(\textbf{x}^{f}_{i,k})$ and coordinate-wise recurrent memory
$\textbf{c}_{i,k-1}$ resulting in the following update equations:

$\displaystyle\textbf{x}_{i,k}^{f}$
$\displaystyle=\mathcal{M}_{k-1:k}(\textbf{x}_{i,k-1}^{a})\quad$
$\displaystyle\quad P_{k}^{f}$
$\displaystyle=Cov(\\{\textbf{x}^{f}_{i,k}\\}_{k=1}^{m})$ (5)
$\displaystyle\textbf{z}_{x_{i}},\textbf{z}_{c_{i}}$
$\displaystyle=f_{L}(\textbf{x}_{i,k}^{f},P_{k}^{f},\textbf{y}_{k},g(\textbf{x}^{f}_{i,k}),\textbf{c}_{i,k-1})\quad$
$\displaystyle\quad\bm{\lambda}_{x_{i}},\bm{\lambda}_{c_{i}}$
$\displaystyle=f_{N}(\textbf{x}_{i,k}^{f},P_{k}^{f},\textbf{y}_{k},g(\textbf{x}^{f}_{i,k}),\textbf{c}_{i,k-1})$
$\displaystyle\textbf{x}_{i,k}^{a}$
$\displaystyle=\bm{\lambda}_{x_{i}}\odot\textbf{x}_{i,k}^{f}+\textbf{z}_{x_{i}}$
$\displaystyle\textbf{c}_{i,k}$
$\displaystyle=\bm{\lambda}_{c_{i}}\odot\textbf{c}_{i,k-1}+\textbf{z}_{c_{i}}$

where $f_{L}$ and $f_{N}$ are neural networks with linear and sigmoid final
activations respectively. We implement these as a single network which is
split only at the final layer. We observed notable gains to training stability
through the inclusion of this sigmoid-gated final layer. For scalability, we
do not use the full covariance matrix as an input to the model. Rather, we
include local covariance entries as additional feature channels where each
channel represents the covariance with the state value at a fixed relative
spatial position. For instance in a 1D system, the relative covariance
channels may contain the variance at the given point and the covariances with
the points to the left and right.

The inclusion of recurrent memory may seem out of place when the dynamics are
assumed to be Markovian. However, while the dynamics themselves are Markovian,
the data assimilation process is not. The inclusion of memory allows the model
to learn from the past without explicitly including previous states and
observations. This is a significant benefit in practice. The other non-
standard inclusion, the local dynamics, led to a small performance boost as
well. As our method is not reliant on generative models of the inputs,
additional features could be added trivially.

###  3.2 Self-Supervised Assimilation

The major obstacle to training an assimilation model is the lack of ground
truth data. In this section, we derive theoretical motivation for the self-
supervised training procedure that lies at the heart of amortized
assimilation. Recent work in image denoising has shown that self-supervised
methods are a promising alternative to denoising methods with explicit noise
models. Lehtinen et al. [27] initially demonstrated the ability to denoise
images without clean ground truth images by using multiple noise samples.
Batson and Royer [17] expanded this idea to the more general
$\mathcal{J}$-invariance framework which exploits independence assumptions
between noise in different partitions of the state vector to derive a valuable
decomposition of the denoising loss. Techniques developed using this approach
have exhibited comparable performance to explicitly supervised denoising
across multiple applications [28, 29, 30]. In our work, we extend this
approach to the dynamical systems setting using differentiable simulation. As
we do not need the full generality of the $\mathcal{J}$-invariance framework
to derive our results, we begin by restricting Proposition 1 from Batson and
Royer [17] to our setting.

######  Lemma 1 (Noise2Self – Restricted).

Suppose concatenated noisy observation vector
$\textbf{y}=[\textbf{y}_{k};\textbf{y}_{k+1}]$ is an unbiased estimator of
concatenated state vector $\textbf{x}=[\textbf{x}_{k};\textbf{x}_{k+1}]$ and
that the noise in $\textbf{y}_{k}$ is independent from the noise in
$\textbf{y}_{k+1}$. Now let
$\textbf{z}=f(\textbf{y})=[\textbf{z}_{k};\textbf{z}_{k+1}]$. If $f$ is a
function such that $\textbf{z}_{k+1}$ does not depend on the value of
$\textbf{y}_{k+1}$ then:

$\mathbb{E}_{\textbf{y}}\|f(\textbf{y})_{k+1}-\textbf{y}_{k+1}\|^{2}=\mathbb{E}_{\textbf{y}}[\|f(\textbf{y})_{k+1}-\textbf{x}_{k+1}\|^{2}+\|\textbf{y}_{k+1}-\textbf{x}_{k+1}\|^{2}]$
(6)

While the proof of this lemma can be found in the supplementary materials
(C.1), the main takeaway is that if we are able to define a search space for
our denoising function such that all $f$ in the space have the stated
independence property, then the expected self-supervised denoising loss can be
decomposed into the expected supervised loss and the noise variance. As the
noise variance term is irreducible, this decomposition implies that minimizing
the expected self-supervised loss is equivalent to minimizing the expected
supervised loss.

Moving back over to data assimilation, our ultimate goal is to train a neural
network to act as our denoising agent. Letting $f_{\theta}$ denote a neural
network parameterized by weights $\theta$, we can formalize the objective that
we’d actually like to minimize as the supervised analysis loss:

$\mathcal{L}^{a}(\theta)=\frac{1}{K-1}\sum_{k=1}^{K-1}\bigg{\|}\bigg{(}\frac{1}{m}\sum_{i=1}^{m}f_{\theta}(\textbf{x}_{i,k}^{f},P_{k}^{f},\textbf{y}_{k},g(\textbf{x}^{f}_{i,k}),\textbf{c}_{i,k-1})\bigg{)}-\textbf{x}_{k}\bigg{\|}^{2}.$
(7)

This supervised loss reflects that our goal is to produce the best initial
condition estimates at the start of a forecast window. Unfortunately,
$\textbf{x}_{k}$ is unknown so this objective cannot be used to train a
denoising model. In the amortized assimilation framework, we instead minimize
what we call the self-supervised forecast loss:

$\mathcal{L}^{ssf}(\theta)=\frac{1}{K-1}\sum_{k=1}^{K-1}\bigg{\|}\bigg{(}\frac{1}{m}\sum_{i=1}^{m}\mathcal{H}_{k+1}(\mathcal{M}_{k:k+1}\circ
f_{\theta}(\textbf{x}_{i,k}^{f},\cdot)\bigg{)}-\textbf{y}_{k+1}\bigg{\|}^{2}.$
(8)

The self-supervised approach (Figure 4) has the obvious advantage that it uses
readily available noisy observations as the training target. However, we can
also show that it is theoretically well-motivated under Lemma 1.
Differentiable simulation acts as a bridge between our denoised estimate and
future observations which are assumed to have independently sampled noise
under the generative model described in Section 2. While Lemma 1 is stated in
terms of one step ahead forecasting, similar to prior work [31], we observe
significantly improved performance by unrolling the procedure over multiple
steps.

The power of this approach lies in the fact that the forecast objective can
actually be used to bound the analysis objective so that minimizing the
forecast objective is equivalent to minimmizing an upper bound on the analysis
objective. To show this, apart from the previously specified generative model
assumptions, we rely on two additional assumptions:

######  Assumption 1 (Uniqueness of Initial Value Problem (IVP) Solution).

The autonomous system evolution equation $x^{\prime}(t)=g(x(t))$ is
deterministic and $L$-Lipschitz continuous in $x$ thus admitting a unique
solution to the initial value problem under the Picard-Lindelöf theorem [32].

Figure 4: The self-supervised training process is unrolled across multiple
filtering steps. Loss is evaluated between “pseudo-observations” and noisy
observations obtained from the system.

######  Assumption 2 (Unbiased Observation Operators).

For all observation operators $\mathcal{H}_{k}$, the observation $y_{k}$ is an
unbiased estimator of some subset of $x_{k}$, ie. if $x_{k}^{S}$ is the
restriction of $x_{k}$ to some subset of features $S$ then
$\mathbb{E}[y_{k}\>|\>{x_{k}^{S}}]=x_{k}^{S}$.

We also define an additional loss, the supervised forecast loss
$\mathcal{L}^{f}$ as the mean squared error between the forecast state and the
true state. For the purposes of our analysis, we’ll assume that $S$ is the
full state such that $\mathbb{E}[y_{k}\>|\>{x_{k}}]=x_{k}$ though this will
not be the case in our numerical experiments. Our first result focuses on the
case where the supervised loss is driven to zero.

######  Proposition 1 (Zero Loss Regime).

Under the stated assumptions, let $f_{\theta}$ denote a family of functions
parameterized by $\theta$ for which $\min_{\ \theta}\
\mathcal{L}^{f}(\theta)=0$. For any $\theta$ which achieves this minimum, it
is also true that $\mathcal{L}^{a}(\theta)=0$ and that $\theta$ is in set of
minimizers of $\mathbb{E}_{\textbf{y}_{0:K}}[\mathcal{L}^{ssf}(\theta)]$.

The proof can be found in the supplementary materials (C.2). This case
provides important motivation by showing that if a perfect denoiser exists and
is contained within our search space then it is also in the set of minimizers
for our expected self-supervised objective. This does not necessarily mean
that such a function will be discovered as the objective is non-convex and in
practice, we’re minimizing the self-supervised loss over a finite sample
rather than over the expectation of the noise distribution.

For the more general case, with sufficiently smooth dynamics, one might expect
that over short time horizons, a function which minimizes the forecast loss
should perform well for the analysis loss, this is not guaranteed and depends
on the properties of the system under study. For our bound, we only consider
the supervised losses as per the decomposition in Lemma 1, minimizing the
expected self-supervised forecast loss is equivalent to minimizing the
expected supervised forecast loss.

######  Proposition 2 (Non-zero Loss Regime).

Under the previously stated assumptions, the supervised analysis loss can be
bounded by the supervised forecast loss as:

$\frac{1}{K-1}\sum_{k=1}^{K-1}\bigg{\|}\bigg{(}\frac{1}{m}\sum_{i=1}^{m}f_{\theta}(\cdot))\bigg{)}-\textbf{x}_{k}\bigg{\|}\leq\frac{e^{L(\tau_{k+1}-\tau_{k})}}{K-1}\sum_{k=1}^{K-1}\max_{i\in[m]}\bigg{\|}\mathcal{M}_{k:k+1}\circ
f_{\theta}(\textbf{x}^{f}_{i,k},\cdot)-\textbf{x}_{k}\bigg{\|}$ (9)

The proof (found in C.3) simply uses the Lipschitz coefficient to bound the
system Lyapunov exponents which govern the evolution of perturbations. While
the tightness of the bound is system specific, large Lyapunov exponents can be
offset by more frequent assimilation cycles. In Section 5, our numerical
results suggest that training under this loss is effective in practice for a
number of common test systems designed to simulate real-world phenomena.

###  3.3 Implementation Details

#### Avoiding Ensemble Collapse

Ensemble-based uncertainty estimates are an elegant solution to the problem of
computing the evolution of uncertainty under nonlinear dynamics, but without
assumptions on the noise model, evaluating the analysis uncertainty is non-
trivial. A well-known problem in deep learning-based uncertainty
quantification is feature collapse [33, 34], a phenomenon in which the hidden
representations resulting from regions of the input space are pushed
arbitrarily closely together by the neural network. In our ensemble-based
uncertainty quantification scheme, this poses an issue as ensemble members can
quickly converge over a small number of assimilation steps resulting in
ensemble collapse.

We address this by combining ensemble estimation with MC Dropout [35] which
interprets a pass through a dropout-regularized neural network as a sample
from a variational distribution over network weights. Thus, for standard usage
of MC Dropout one can estimate uncertainty in the predictive distribution by
using multiple dropout passes to integrate out weight uncertainty. In our
case, we further have to account for uncertainty over the input distribution
which is represented by our forecast ensemble. We employ a relatively simple
Monte Carlo scheme in which the forecast ensemble members are assimilated by a
single pass through a dropout-regularized network with independent dropout
masks giving us $m$ samples from the predictive distribution as our analysis
ensemble. While this is a crude approximation, we found it resulted in more
consistent performance with less tuning compared to the sensitivity
constraints found in recent uncertainty quantification methods [33, 36, 37].

#### Incomplete Observations

One important concern in data assimilation is handling a variety of
observation operators as the full state will rarely be observed at one time
point. We discuss two processes for addressing this in amortized filters. The
first, which is far more flexible but less effective at exploiting prior
knowledge, is to initialize an network with independent weights corresponding
to each observation type. This can be used with any observation type, even
ones with unknown relationships to the system state. These unknown observation
operators can be trained by evaluating the loss at points with future, known
observation operators and passing the gradient signal back in time through the
differentiable simulation connecting the time points.

The second, which we implement in the AmEnF, is to use our knowledge of the
spatial distribution of the observation operators to assign the observed
values corresponding to certain locations as channels in the input
representation. Coordinates which are not observed can be masked and an
additional channel is appending indicating whether a particular coordinate was
observed during the given assimilation cycle. This is depicted by the shading
in the observation channels of Figure 3. More details on the masking process
can be found in the supplementary material.

##  4 Related Work

Work on the intersection of deep learning and data assimilation has exploded
in recent years. Some early efforts explicitly used reanalysis data produced
by traditional assimilation methods as a target for supervised training [38,
39]. These approaches are limited as they cannot reasonably expect to
outperform the method used to produce the reanalysis data in terms of
accuracy. More recently, significant effort has gone into using traditional
assimilation alongside learned surrogate models to stabilize training, learn
augmented dynamics, or accelerate simulation with larger ensembles [40, 41,
42, 43].

The work most comparable to our own is that which uses machine learning
methods to learn a component of the assimilation mechanism directly in
continuous nonlinear systems [44, 45]. Our method differs from Grönquist et
al. [44] in that we use differentiable simulation to enable self-supervised
learning. While Frerix et al. [45] incorporates differentiable simulation,
their Learned Inverse Operator framework acts as a component of the 4D-Var
algorithm and still requires solving an expensive optimization problem at each
assimilation step rather than using using a closed form update. Outside of the
field of data assimilation specifically, there has also been significant work
into neural parameterization of Bayesian filters [25, 46, 47, 48, 49, 50, 51,
52, 53, 54].

##  5 Experiments

(a) Mean ensemble spread over assimilation cycles

(b) Ensemble trajectories with observation noise $\sigma=1.0$

(c) Ensemble trajectories with observation noise $\sigma=2.5$

Figure 5: MC Dropout introduces an additional stochastic element that
stabilizes the ensemble resulting in realistic relationships between the
ensemble spread and system uncertainty.

We examine the empirical performance of amortized assimilation from two
perspectives. In the first, we derive empirical justification for our design
choices, comparing the performance of the various objective functions and
qualitatively examining the evolution of uncertainty in the ensemble.

The second objective is to compare the performance of our amortized model, the
AmEnF, against standard data assimilation methods over multiple common
benchmark dynamical systems. All methods tested in this section have explicit
knowledge of the system dynamics. Due to the challenge of learning chaotic
dynamics, we found that methods which rely on learning the dynamics were
uncompetitive.

For all experiments, we generate a training set consisting of 6000 sequences
consisting of 40 assimilation steps each. The validation set consists of a
single sequence of an additional 1000 steps and the test set is a further
10,000 steps. AmEnF models are developed in PyTorch [55] using the torchdiffeq
[56] library for ODE integration. Models are trained on a single GTX 1070 GPU
for 500 epochs using the Adam [12] optimizer with initial learning rate $8e-4$
with a warm-up over 50 iterations followed by halving the learning rate every
200 iterations. All experiments are repeated over ten independent noise
samples and error bars indicate a single standard deviation.

###  5.1 Qualitative Evaluation

Experiments in this section are performed using the Lorenz 96 system [57].
Lorenz 96 is a system of coupled differential equations intended to emulate
the evolution of a single atmospheric quantity across a latitude circle.

$\dot{x}_{j}=(x_{j+1}-x_{j-2})x_{j-1}-x_{j}+F$ (10)

The system is defined to have periodic boundary conditions $x_{D+1}=x_{1}$,
$x_{0}=x_{D}$, and $x_{-1}=x_{D-1}$ where $D$ is the number of system
dimensions. We set the number of dimensions to 40 and the forcing value to
$F=8$. This is perhaps the most widely used test system in data assimilation
and is what we will use the explore the behavior of the AmEnF model. Unless
explicitly stated otherwise, all experiments in this section were performed
with $\sigma$=1 and $m=10$.

#### Ensemble Behavior

Figure 6: Analysis RMSE evaluated on the validation set across models trained
using different losses for $\sigma=1$.

Figure 5 demonstrates the benefit of incorporating MC Dropout. In the left
figure, we see that the ensemble variance quickly collapses to zero using
deterministic networks while it varies along a stable range using MC Dropout.
More interesting are the central and right figures which compare the ensemble
member trajectories (light blue) with the true trajectory (black) for a single
spatial coordinate across several noise settings. We can see that the
uncertainty patterns in our ensemble behave largely as one would hope. Less
certainty in the system results in a larger ensemble spread.

Of particular note is the behavior near outlier observations. We can see that
while the ensemble trajectories occasionally approach such observations, when
the ensemble spread is tight the learned assimilator trusts the current
forecast trajectory and the ensemble members do not move onto a worse path to
match the observation.

#### Training Objective

Figure 6 shows a comparison between the test-time supervised analysis loss
across AmEnF models trained using several objective functions. The problem
used for comparison is the fully observed Lorenz 96 system with observation
noise $\sigma=1$. As one might anticipate, the self-supervised analysis
objective, the only objective available without differentiable simulation,
results in the AmEnF learning an identity mapping from the observation so that
the error is exactly the observation variance. The supervised analysis (which
we are able to use since this is simulated data) and self-supervised forecast
training approaches result in learning a legitimate denoising function. What
is intriguing is that the self-supervised forecast loss actually outperforms
the supervised analysis loss for this system. The optimization problem is
nonconvex, so neither approach is guaranteed to reach a global minimizer. This
is not definitive as the performance gap is within the margin of error;
however, it does suggest that there may be some value to including
differentiable simulation in the training pipeline beyond simply enabling the
denoising property.

###  5.2 Method Comparison

Figure 7: Results from numerical comparisons across multiple ensemble sizes
and observation noises.

We evaluate performance on three benchmark systems. Loss is reported as the
time-averaged RMSE between the mean analysis state estimate and the true state
on the test state. Each system is evaluated across a variety of ensemble sizes
(denoted by $m$) and isotropic Gaussian observation noise levels (whose
standard deviation is denoted by $\sigma$). All results are reported in Figure
7.

We compare performance against a set of widely used filtering methods for data
assimilation implemented in the Python DAPPER library [58]. These methods
include 4D-Var [24], the Local Ensemble Transform Kalman Filter (LETKF) [59],
and the Iterative Ensemble Kalman Smoother (iEnKS) [60]. 4D-Var and the iEnKF
are used in a filtering capacity, using only the latest observation rather
than a multi-step trajectory. Full experiment settings including
hyperparameters such as inflation and localization settings searched are
included in the supplementary materials.

#### Lorenz 96

For method comparisons, we examine the more challenging case of a partially
observed system. At every assimilation step, we observe only a rotating subset
consisting of every fourth spatial dimension so that only $25\%$ of the system
is observed at any time. These experiments use an RK-4 [61] integrator with
step sizes of .05 and partial observations every two integration steps.

#### Kuramoto-Shivashinsky

While our theoretical results focus on ordinary differential equations, the
method can also be used with partial differential equations like the KS
equation. The KS equation [62] is a fourth-order partial differential equation
known to exhibit chaotic behavior. In one spatial dimension, the governing
equation can be written as:

$u_{t}+u_{x}+u_{xxxx}+uu_{x}=0$ (11)

with periodic boundary conditions. These experiments were fully observed. The
system was integrated with an ETD-RK4 method with step sizes of $.5$ with
observations every two steps.

#### Vissio-Lucarini 20

The system described by Vissio and Lucarini [63] is a coupled system intended
to represent a minimal model of the earth’s atmosphere. It augments the Lorenz
system with “temperature" variables allowing for more complicated behavior:

$\displaystyle\frac{dX_{j}}{dt}=X_{j-1}(X_{j+1}-X_{j-2})-\alpha\theta_{j}-\gamma
X_{j}+F$ (12)
$\displaystyle\frac{d\theta_{j}}{dt}=X_{j+1}\theta_{j+2}-X_{j-1}\theta_{j-2}+\alpha
X_{j}-\gamma\theta_{j}+G$

We set $F=10$, $G=10$, $\gamma=1$, $\alpha=1$, and use 36 spatial points.
Experiments using this system are also partially observed with every fourth
spatial coordinate being observed at one time. The system is integrated using
RK-4 with steps of length $.05$ and partial observations every two steps.
4D-Var was excluded in this comparison as performance on prior experiments did
not justify further examination.

#### Results

As we can see in Figure 7, the AmEnF models match or outperform the
competitive methods with more significant improvements coming at smaller
ensemble sizes and higher noise levels which are associated with stronger
nonlinear effects [60]. The only system in which the AmEnF does not notably
outperform the comparison models is the KS equation. This is interesting as
our theoretical results rely on assumptions about ODEs while the KS equation
is a PDE. Furthermore, the bound we derived has a time component and the KS
experiments occur at the largest time scales. Nonetheless, despite these
potential theoretical snags, the AmEnF model near matches the the top
performing models in accuracy while exhibiting stronger stability under
repeated experiments.

##  6 Discussion and Conclusion

#### Limitations

Despite the successful numerical results, the method is not without
challenges. The largest consideration is related to scaling. First, while the
use of differentiable simulation allowed our method to outperform competitive
methods, the availability of adjoints in numerical forecasting is not certain.
Forecasts in domains like NWP often use legacy code where adjoints need to be
maintained by hand. We expect this to change over time, but this will limit
the adoption of similar approaches in the near future. Assuming the presence
of adjoints, training cost is also a concern as it requires the repeated
simulation of the dynamical systems. However, numerical models tend to change
slowly and applications like NWP are among the largest users of compute in the
world [1] with assimilation cycles occurring multiple times per day, thus we
expect the amortization to pay off over time.

#### Conclusion

We theoretically and experimentally motivate the use of a self-supervised
approach for learning to assimilate in chaotic dynamical systems. This
approach uses a hybrid deep learning-numerical simulation architecture which
outperforms or matches several standard data assimilation methods across
several numerical experiments. The performance gains are especially profound
at low ensemble sizes, which could enable higher resolution, more accurate
simulation for assimilation problems leading to stronger predictions for high
value problems like numerical weather prediction in the future.

## Acknowledgments and Disclosure of Funding

This material is based upon work supported by the National Science Foundation
under Grant No. 1835825.

## References

  * Bauer et al. [2015] Peter Bauer, Alan Thorpe, and Gilbert Brunet.  The quiet revolution of numerical weather prediction.  _Nature_ , 525(7567):47–55, 2015.
  * Randall et al. [2007] David A Randall, Richard A Wood, Sandrine Bony, Robert Colman, Thierry Fichefet, John Fyfe, Vladimir Kattsov, Andrew Pitman, Jagadish Shukla, Jayaraman Srinivasan, et al.  Climate models and their evaluation.  In _Climate change 2007: The physical science basis. Contribution of Working Group I to the Fourth Assessment Report of the IPCC (FAR)_ , pages 589–662. Cambridge University Press, 2007.
  * Bocquet et al. [2015] M. Bocquet, H. Elbern, H. Eskes, M. Hirtl, R. Žabkar, G. R. Carmichael, J. Flemming, A. Inness, M. Pagowski, J. L. Pérez Camaño, P. E. Saide, R. San Jose, M. Sofiev, J. Vira, A. Baklanov, C. Carnevale, G. Grell, and C. Seigneur.  Data assimilation in atmospheric chemistry models: current status and future prospects for coupled chemistry meteorology models.  _Atmospheric Chemistry and Physics_ , 15(10):5325–5358, 2015.  doi: 10.5194/acp-15-5325-2015.  URL https://acp.copernicus.org/articles/15/5325/2015/.
  * Fichtner [2010] Andreas Fichtner.  _Full seismic waveform modelling and inversion_.  Springer Science & Business Media, 2010.
  * Bresch et al. [2010] Didier Bresch, Thierry Colin, Emmanuel Grenier, Benjamin Ribba, and Olivier Saut.  Computational modeling of solid tumor growth: the avascular stage.  _SIAM Journal on Scientific Computing_ , 32(4):2321–2344, 2010.
  * Bertoglio et al. [2012] Cristóbal Bertoglio, Philippe Moireau, and Jean-Frederic Gerbeau.  Sequential parameter estimation for fluid–structure problems: application to hemodynamics.  _International Journal for Numerical Methods in Biomedical Engineering_ , 28(4):434–455, 2012.
  * Engl et al. [2009] Heinz W Engl, Christoph Flamm, Philipp Kügler, James Lu, Stefan Müller, and Peter Schuster.  Inverse problems in systems biology.  _Inverse Problems_ , 25(12):123014, 2009.
  * Wiggins et al. [1990] Stephen Wiggins, Stephen Wiggins, and Martin Golubitsky.  _Introduction to applied nonlinear dynamical systems and chaos_ , volume 2.  Springer, 1990.
  * Carrassi et al. [2018] Alberto Carrassi, Marc Bocquet, Laurent Bertino, and Geir Evensen.  Data assimilation in the geosciences: An overview of methods, issues, and perspectives.  _WIREs Climate Change_ , 9(5):e535, 2018.  URL https://wires.onlinelibrary.wiley.com/doi/abs/10.1002/wcc.535.
  * Penny and Hamill [2017] Stephen G. Penny and Thomas M. Hamill.  Coupled Data Assimilation for Integrated Earth System Analysis and Prediction.  _Bulletin of the American Meteorological Society_ , 98(7):ES169–ES172, 2017.  doi: 10.1175/BAMS-D-17-0036.1.
  * Kalnay [2003] Eugenia Kalnay.  _Atmospheric modeling, data assimilation and predictability_.  Cambridge university press, 2003.
  * Kingma and Welling [2014] Diederik P. Kingma and Max Welling.  Auto-Encoding Variational Bayes.  In _2nd International Conference on Learning Representations, ICLR 2014, Banff, AB, Canada, April 14-16, 2014, Conference Track Proceedings_ , 2014.
  * Xue et al. [2020] Tianju Xue, Alex Beatson, Sigrid Adriaenssens, and Ryan Adams.  Amortized Finite Element Analysis for Fast PDE-Constrained Optimization.  In Hal Daumé III and Aarti Singh, editors, _Proceedings of the 37th International Conference on Machine Learning_ , volume 119 of _Proceedings of Machine Learning Research_ , pages 10638–10647. PMLR, 13–18 Jul 2020.  URL http://proceedings.mlr.press/v119/xue20a.html.
  * Ravi and Beatson [2019] Sachin Ravi and Alex Beatson.  Amortized Bayesian Meta-Learning.  In _International Conference on Learning Representations_ , 2019.  URL https://openreview.net/forum?id=rkgpy3C5tX.
  * Asch et al. [2016] Mark Asch, Marc Bocquet, and Maëlle Nodet.  _Data assimilation: methods, algorithms, and applications_ , volume 11.  SIAM, 2016.
  * Evensen [1994] G. Evensen.  Sequential data assimilation with a nonlinear quasi-geostrophic model using Monte Carlo methods to forecast error statistics.  _Journal of Geophysical Research_ , 1994.  ISSN 01480227.  doi: 10.1029/94jc00572.
  * Batson and Royer [2019] Joshua Batson and Loic Royer.  Noise2Self: Blind Denoising by Self-Supervision.  In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, _Proceedings of the 36th International Conference on Machine Learning_ , volume 97 of _Proceedings of Machine Learning Research_ , pages 524–533. PMLR, 09–15 Jun 2019.  URL https://proceedings.mlr.press/v97/batson19a.html.
  * Kalman [1960] R. E. Kalman.  A new approach to linear filtering and prediction problems.  _Journal of Fluids Engineering, Transactions of the ASME_ , 1960.  ISSN 1528901X.  doi: 10.1115/1.3662552.
  * Jazwinski [1972] A H Jazwinski.  Stochastic processes and filtering theory.  _IEEE Transactions on Automatic Control_ , 1972.
  * Doucet et al. [2001] Arnaud Doucet, Nando Freitas, and Neil Gordon.  An Introduction to Sequential Monte Carlo Methods.  In _Sequential Monte Carlo Methods in Practice_. 2001.  doi: 10.1007/978-1-4757-3437-9_1.
  * Snyder et al. [2008] Chris Snyder, Thomas Bengtsson, Peter Bickel, and Jeff Anderson.  Obstacles to high-dimensional particle filtering.  _Monthly Weather Review_ , 2008.  ISSN 00270644.  doi: 10.1175/2008MWR2529.1.
  * Anderson and Anderson [1999] J. L. Anderson and S. L. Anderson.  A Monte Carlo implementation of the nonlinear filtering problem to produce ensemble assimilations and forecasts.  _Mon. Wea. Rev._ , 127:2741–2758, 1999.
  * Hamill et al. [2001] 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.  URL https://journals.ametsoc.org/view/journals/mwre/129/11/1520-0493_2001_129_2776_ddfobe_2.0.co_2.xml.
  * Rabier et al. [2000] F Rabier, H Järvinen, E Klinker, J.-F. Mahfouf, and A Simmons.  The ECMWF operational implementation of four-dimensional variational assimilation. I: Experimental results with simplified physics.  _Quarterly Journal of the Royal Meteorological Society_ , 126(564):1143–1170, 2000.  URL https://rmets.onlinelibrary.wiley.com/doi/abs/10.1002/qj.49712656415.
  * Haarnoja et al. [2016] Tuomas Haarnoja, Anurag Ajay, Sergey Levine, and Pieter Abbeel.  Backprop KF: Learning discriminative deterministic state estimators.  In _Advances in Neural Information Processing Systems_ , 2016.
  * Werbos [1990] P.J. Werbos.  Backpropagation through time: what it does and how to do it.  _Proceedings of the IEEE_ , 78(10):1550–1560, 1990.  doi: 10.1109/5.58337.
  * Lehtinen et al. [2018] Jaakko Lehtinen, Jacob Munkberg, Jon Hasselgren, Samuli Laine, Tero Karras, Miika Aittala, and Timo Aila.  Noise2Noise: Learning image restoration without clean data.  In _35th International Conference on Machine Learning, ICML 2018_ , 2018.  ISBN 9781510867963.
  * Laine et al. [2019] Samuli Laine, Tero Karras, Jaakko Lehtinen, and Timo Aila.  High-quality self-supervised deep image denoising.  In _Advances in Neural Information Processing Systems_ , volume 32. Curran Associates, Inc., 2019.  URL https://proceedings.neurips.cc/paper/2019/file/2119b8d43eafcf353e07d7cb5554170b-Paper.pdf.
  * Fadnavis et al. [2020] Shreyas Fadnavis, Joshua Batson, and Eleftherios Garyfallidis.  Patch2self: Denoising diffusion mri with self-supervised learning.  In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, _Advances in Neural Information Processing Systems_ , volume 33, pages 16293–16303. Curran Associates, Inc., 2020.  URL https://proceedings.neurips.cc/paper/2020/file/bc047286b224b7bfa73d4cb02de1238d-Paper.pdf.
  * Krull et al. [2019] Alexander Krull, Tim-Oliver Buchholz, and Florian Jug.  Noise2Void - Learning Denoising From Single Noisy Images.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , June 2019.
  * Um et al. [2020] Kiwon Um, Robert Brand, Yun (Raymond) Fei, Philipp Holl, and Nils Thuerey.  Solver-in-the-Loop: Learning from Differentiable Physics to Interact with Iterative PDE-Solvers.  In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, _Advances in Neural Information Processing Systems_ , volume 33, pages 6111–6122. Curran Associates, Inc., 2020.  URL https://proceedings.neurips.cc/paper/2020/file/43e4e6a6f341e00671e123714de019a8-Paper.pdf.
  * Coddington and Levinson [1955] Earl A Coddington and Norman Levinson.  _Theory of ordinary differential equations_.  Tata McGraw-Hill Education, 1955.
  * van Amersfoort et al. [2021] Joost van Amersfoort, Lewis Smith, Andrew Jesson, Oscar Key, and Yarin Gal.  Improving Deterministic Uncertainty Estimation in Deep Learning for Classification and Regression.  _CoRR_ , abs/2102.11409, 2021.  URL https://arxiv.org/abs/2102.11409.
  * Papyan et al. [2020] Vardan Papyan, X. Y. Han, and David L. Donoho.  Prevalence of neural collapse during the terminal phase of deep learning training.  _Proceedings of the National Academy of Sciences_ , 117(40):24652–24663, 2020.  ISSN 0027-8424.  doi: 10.1073/pnas.2015509117.  URL https://www.pnas.org/content/117/40/24652.
  * Gal and Ghahramani [2016] Yarin Gal and Zoubin Ghahramani.  Dropout as a Bayesian approximation: Representing model uncertainty in deep learning.  In _33rd International Conference on Machine Learning, ICML 2016_ , 2016.  ISBN 9781510829008.
  * Van Amersfoort et al. [2020] Joost Van Amersfoort, Lewis Smith, Yee Whye Teh, and Yarin Gal.  Uncertainty estimation using a single deep deterministic neural network.  In Hal Daumé III and Aarti Singh, editors, _Proceedings of the 37th International Conference on Machine Learning_ , volume 119 of _Proceedings of Machine Learning Research_ , pages 9690–9700. PMLR, 13–18 Jul 2020.  URL https://proceedings.mlr.press/v119/van-amersfoort20a.html.
  * Rosca et al. [2020] Mihaela Rosca, Theophane Weber, Arthur Gretton, and Shakir Mohamed.  A case for new neural network smoothness constraints.  In Jessica Zosa Forde, Francisco Ruiz, Melanie F. Pradier, and Aaron Schein, editors, _Proceedings on "I Can’t Believe It’s Not Better!" at NeurIPS Workshops_ , volume 137 of _Proceedings of Machine Learning Research_ , pages 21–32. PMLR, 12 Dec 2020.  URL https://proceedings.mlr.press/v137/rosca20a.html.
  * Härter and de Campos Velho [2008] Fabrício P. Härter and Haroldo Fraga de Campos Velho.  New approach to applying neural network in nonlinear dynamic model.  _Applied Mathematical Modelling_ , 32(12):2621–2633, dec 2008.  ISSN 0307-904X.  doi: 10.1016/J.APM.2007.09.006.
  * Cintra et al. [2016] Rosângela Cintra, Haroldo Campos Velho, and Steven Cocke.  Tracking the model: Data assimilation by artificial neural network.  pages 403–410, 2016.  doi: 10.1109/IJCNN.2016.7727227.
  * Brajard et al. [2019] J. Brajard, A. Carrassi, M. Bocquet, and L. 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.  _Geoscientific Model Development Discussions_ , 2019:1–21, 2019.  doi: 10.5194/gmd-2019-136.  URL https://gmd.copernicus.org/preprints/gmd-2019-136/.
  * Ouala et al. [2018] Said Ouala, Ronan Fablet, Cédric Herzet, Bertrand Chapron, Ananda Pascual, Fabrice Collard, and Lucile Gaultier.  Neural network based Kalman filters for the spatio-temporal interpolation of satellite-derived sea surface temperature.  _Remote Sensing_ , 2018.  ISSN 20724292.  doi: 10.3390/rs10121864.
  * Bocquet et al. [2019] 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.
  * Chattopadhyay et al. [2021] A. Chattopadhyay, M. Mustafa, P. Hassanzadeh, E. Bach, and K. Kashinath.  Towards physically consistent data-driven weather forecasting: Integrating data assimilation with equivariance-preserving spatial transformers in a case study with ERA5.  _Geoscientific Model Development Discussions_ , 2021:1–23, 2021.  doi: 10.5194/gmd-2021-71.  URL https://gmd.copernicus.org/preprints/gmd-2021-71/.
  * Grönquist et al. [2021] Peter Grönquist, Chengyuan Yao, Tal Ben-Nun, Nikoli Dryden, Peter Dueben, Shigang Li, and Torsten Hoefler.  Deep learning for post-processing ensemble weather forecasts.  _Philosophical Transactions of the Royal Society A: Mathematical, Physical and Engineering Sciences_ , 379(2194):20200092, Feb 2021.  ISSN 1471-2962.  doi: 10.1098/rsta.2020.0092.
  * Frerix et al. [2021] Thomas Frerix, Dmitrii Kochkov, Jamie Smith, Daniel Cremers, Michael Brenner, and Stephan Hoyer.  Variational data assimilation with a learned inverse observation operator.  In Marina Meila and Tong Zhang, editors, _Proceedings of the 38th International Conference on Machine Learning_ , volume 139 of _Proceedings of Machine Learning Research_ , pages 3449–3458. PMLR, 18–24 Jul 2021.  URL https://proceedings.mlr.press/v139/frerix21a.html.
  * Coskun et al. [2017] Huseyin Coskun, Felix Achilles, Robert Dipietro, Nassir Navab, and Federico Tombari.  Long Short-Term Memory Kalman Filters: Recurrent Neural Estimators for Pose Regularization.  In _Proceedings of the IEEE International Conference on Computer Vision_ , 2017.  ISBN 9781538610329.  doi: 10.1109/ICCV.2017.589.
  * Karl et al. [2017] Maximilian Karl, Maximilian Soelch, Justin Bayer, and Patrick Van der Smagt.  Deep Variational Bayes Filters: Unsupervised learning of state space models from raw data.  In _Proceedings of the International Conference on Learning Representations (ICLR)_ , 2017.
  * Krishnan et al. [2015] Rahul G. Krishnan, Uri Shalit, and David Sontag.  Deep Kalman Filters, 2015.
  * Rangapuram et al. [2018] Syama Sundar Rangapuram, Matthias W Seeger, Jan Gasthaus, Lorenzo Stella, Yuyang Wang, and Tim Januschowski.  Deep State Space Models for Time Series Forecasting.  In S Bengio, H Wallach, H Larochelle, K Grauman, N Cesa-Bianchi, and R Garnett, editors, _Advances in Neural Information Processing Systems 31_ , pages 7785–7794. Curran Associates, Inc., 2018.  URL http://papers.nips.cc/paper/8004-deep-state-space-models-for-time-series-forecasting.pdf.
  * [50] Xiao Ma, Péter Karkus, David Hsu, and Wee Sun Lee.  Particle Filter Recurrent Neural Networks.  In _The Thirty-Fourth AAAI Conference on Artificial Intelligence, AAAI, 2020_ , pages 5101–5108.
  * Gu et al. [2015] Shixiang Gu, Zoubin Ghahramani, and Richard E. Turner.  Neural adaptive Sequential Monte Carlo.  In _Advances in Neural Information Processing Systems_ , 2015.
  * Le et al. [2018] Tuan Anh Le, Maximilian Igl, Tom Rainforth, Tom Jin, and Frank Wood.  Auto-encoding Sequential Monte Carlo.  In _6th International Conference on Learning Representations, ICLR 2018 - Conference Track Proceedings_ , 2018.
  * Mei et al. [2019] Hongyuan Mei, Guanghui Qin, and Jason Eisner.  Imputing missing events in continuous-time event streams.  In _36th International Conference on Machine Learning, ICML 2019_ , 2019.  ISBN 9781510886988.
  * Becker et al. [2019] Philipp Becker, Harit Pandya, Gregor Gebhardt, Cheng Zhao, C. James Taylor, and Gerhard Neumann.  Recurrent Kalman Networks: Factorized Inference in High-Dimensional Deep Feature Spaces.  In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, _Proceedings of the 36th International Conference on Machine Learning_ , volume 97 of _Proceedings of Machine Learning Research_ , pages 544–552. PMLR, 09–15 Jun 2019.  URL https://proceedings.mlr.press/v97/becker19a.html.
  * Paszke et al. [2019] Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary Devito Facebook, A I Research, Zeming Lin, Alban Desmaison, Luca Antiga, Orobix Srl, and Adam Lerer.  Automatic differentiation in PyTorch.  In _Advances in Neural Information Processing Systems 32_ , 2019.
  * Chen et al. [2018] Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud.  Neural ordinary differential equations.  _Advances in Neural Information Processing Systems_ , 2018.
  * Lorenz [2006] Edward N. Lorenz.  Predictability-a problem partly solved.  In _Predictability of Weather and Climate_. 2006.  ISBN 9780511617652.  doi: 10.1017/CBO9780511617652.004.
  * Raanes and Others [2018] Patrick N Raanes and Others.  nansencenter/DAPPER: Version 0.8, 2018.  URL https://doi.org/10.5281/zenodo.2029296.
  * Bocquet [2011] M. Bocquet.  Ensemble Kalman filtering without the intrinsic need for inflation.  _Nonlinear Processes in Geophysics_ , 2011.  ISSN 10235809.  doi: 10.5194/npg-18-735-2011.
  * Sakov et al. [2012] Pavel Sakov, Dean S. Oliver, and Laurent Bertino.  An iterative EnKF for strongly nonlinear systems.  _Monthly Weather Review_ , 2012.  ISSN 00270644.  doi: 10.1175/MWR-D-11-00176.1.
  * Runge [1895] C. Runge.  Ueber die numerische Auflösung von Differentialgleichungen.  _Mathematische Annalen_ , 1895.  ISSN 00255831.  doi: 10.1007/BF01446807.
  * Kuramoto [1978] Yoshiki Kuramoto.  Diffusion-Induced Chaos in Reaction Systems.  _Progress of Theoretical Physics Supplement_ , 1978.  ISSN 0375-9687.  doi: 10.1143/ptps.64.346.
  * Vissio and Lucarini [2020] Gabriele Vissio and Valerio Lucarini.  Mechanics and thermodynamics of a new minimal model of the atmosphere.  _The European Physical Journal Plus_ , 135(10), Oct 2020\.  ISSN 2190-5444.  doi: 10.1140/epjp/s13360-020-00814-w.
  * Zupanski [1997] Dusanka Zupanski.  A general weak constraint applicable to operational 4dvar data assimilation systems.  _Monthly Weather Review_ , 125(9):2274–2292, 1997\.

##  Appendix A Notation Table

Table 1: Notation used in paper. Symbol | Domain | Description
---|---|---
$\textbf{x}(t)$ | $\mathbb{R}^{d}$ | System state evaluated at time $t$
$g$ | $\mathbb{R}^{d}\to\mathbb{R}^{d}$ |  $\mathbf{\dot{x}}$ or equivalently the governing dynamics of system
$L$ | $\mathbb{R}$ | Lipschitz constant of $g$
$\tau_{k}$ | $\mathbb{R}$ | Time observation $k$ was recorded
$\mathcal{M}_{k:k+1}$ | $\mathbb{R}^{d}\to\mathbb{R}^{d}$ | $\textbf{x}_{k}+\int_{\tau_{k}}^{\tau_{k+1}}g(\textbf{x}(t))\ dt$
$\mathcal{H}_{k}$ | $\mathbb{R}^{d}\to\mathbb{R}^{o}$ | Observation operator used at time index $k$
$\textbf{y}_{k}$ | $\mathbb{R}^{o}$ | Observation value recorded at time index $k$
$\sigma_{k}^{2}$ | $\mathbb{R}$ | Variance of isotropic Gaussian observation noise at time index $k$
$\textbf{x}_{k}$ | $\mathbb{R}^{d}$ | $\textbf{x}(\tau_{k})$
$m$ | $\mathbb{Z}$ | Ensemble size
$K$ | $\mathbb{Z}$ | Number of observations in assimilation window
$\textbf{x}^{f}_{i,k}$ | $\mathbb{R}^{d}$ | Forecast estimate of $\textbf{x}_{k}$ from ensemble member $i$
$P^{f}_{k}$ | $\mathbb{R}^{d\times d}$ | Forecast estimate of state covariance at time index $k$
$\textbf{x}^{a}_{i,k}$ | $\mathbb{R}^{d}$ | Analysis estimate of $\textbf{x}_{k}$ from ensemble member $i$
$P^{a}_{k}$ | $\mathbb{R}^{d\times d}$ | Analysis estimate of state covariance at time index $k$
$f_{\theta}$ | $\mathbb{R}^{d}\to\mathbb{R}^{d}$ | Learned analysis update parameterized by weights $\theta\in\mathbb{R}^{p}$
$\mathcal{L}^{a}$ | $\mathbb{R}^{d}\to\mathbb{R}$ | Supervised analysis loss
$\mathcal{L}^{f}$ | $\mathbb{R}^{d}\to\mathbb{R}$ | Supervised forecast loss
$\mathcal{L}^{ssf}$ | $\mathbb{R}^{o}\to\mathbb{R}$ | Self-supervised forecast loss
$\mathcal{L}^{ssa}$ | $\mathbb{R}^{o}\to\mathbb{R}$ | Self-supervised analysis loss

##  Appendix B Connections to Classical Assimilation Methods

The Amortized Ensemble Filters (AmEnF) used in our experiments have strong
connections to both ensemble-based sequential filtering methods and
variational assimilation. We discuss both here.

###  B.1 Connections to the EnKF

As we mentioned in the main text of the paper, the AmEnF architecture can be
seen as directly replacing the analysis update equations from the EnKF with an
optimized neural network. There are a large number of EnKF variants, but the
most straightforward to describe is the stochastic EnKF, which “samples" a
unique observation corresponding to each ensemble member from the observation
noise distribution. Letting $H_{k}$ be the linearization of our observation
operator $\mathcal{H}_{k}$, the update equations for the stochastic EnKF are:

Forecast: $\displaystyle\textbf{x}^{f}_{i,k}$
$\displaystyle=\mathcal{M}_{k-1:k}(\textbf{x}^{f}_{i,k}),$ $\displaystyle\quad
i=1,2,\dots,m$ (13) Analysis: $\displaystyle P^{f}_{k}$
$\displaystyle=Cov(\\{\textbf{x}^{f}_{i,k}\\}_{i=1}^{m})$
$\displaystyle\mathbf{\tilde{y}}_{i,k}$
$\displaystyle=\textbf{y}+\bm{\epsilon}_{i},$
$\displaystyle\quad\bm{\epsilon}_{i}\sim\mathcal{N}(0,\sigma_{k}^{2}I)$
$\displaystyle K_{k}$
$\displaystyle=P_{k}^{f}H^{T}_{k}(H_{k}P^{f}_{k}H^{T}_{k}+\sigma_{k}^{2}I)^{-1}$
$\displaystyle\textbf{x}^{a}_{i,k}$
$\displaystyle=\textbf{x}_{i,k}^{f}+K_{k}(\mathbf{\tilde{y}}_{i,k}-H_{k}\textbf{x}^{f}_{i,k}).$

We emphasize that this formulation does not lead to an efficient
implementation and use it only for explanatory purposes.

As with most ensemble-based filters, the forecast step is identical between
the AmEnF and the EnKF. The difference comes from the analysis step. The EnKF
update is the Kalman filter update with the Kalman gain $K$ computed using the
empirical covariance of the ensemble. The Kalman filter formula can be derived
both from Bayesian statistics and through a variational principle as the
optimal update in the case that both the observation and state distributions
are Gaussian [15].

As the image of a Gaussian random variable under nonlinear transformation is
not guaranteed to be Gaussian, this update is no longer optimal under all
nonlinear dynamics. Amortized assimilation replaces this pre-specified update
with a learned update which we optimize using historical data. The amortized
model uses the same basic inputs (the ensemble members, the observation, and
the covariance) with a small number of additional features included due to the
ease of doing so.

###  B.2 Connections to Weak-Constraint 4D-Var

Weak-constraint 4D-Var [64] is a variation of classical, or strong-constraint,
4D-Var which accounts for model inaccuracy by performing optimization over the
full state trajectory $\textbf{x}_{0:K}$ within an assimilation window rather
than just finding initial conditions $\textbf{x}_{0}$ and using the time
evolution of those initial conditions for all likelihood calculations. This
results in the following objective function:

$\displaystyle\mathcal{J}(\textbf{x}_{0:K})=$
$\displaystyle(\textbf{x}_{0}-\textbf{x}_{b})^{T}B^{-1}(\textbf{x}_{0}-\textbf{x}_{b})$
(14)
$\displaystyle+\sum_{k=1}^{K}(\mathcal{H}_{k}(\textbf{x}_{k})-\textbf{y}_{k})^{T}R_{k}^{-1}(\mathcal{H}_{k}(\textbf{x}_{k})-\textbf{y}_{k})$
$\displaystyle+\sum_{k=0}^{K-1}(\textbf{x}_{k+1}-\mathcal{M}_{k:k+1}(\textbf{x}_{k}))^{T}Q_{k+1}^{-1}(\textbf{x}_{k+1}-\mathcal{M}_{k:k+1}(\textbf{x}_{k}))$

where the subscript $b$ indicates a prior or background estimate and the
matrices $B,R_{k},Q_{k}$ indicate the background, observation, and model error
covariances respectively. On the surface, this looks quite different from the
objective used in our work. However, it actually only differs by the inclusion
of a prior term prior to unrolling the self-supervised amortized assimilation
objective. If we included a prior based on the ensemble distribution at each
step of the amortized assimilation objective (which was done in an earlier
iteration of our approach, though we typically observed worse performance
compared to the presented model) and assume $R_{k}=I$, the objective could be
re-written as:

$\displaystyle\mathcal{J}(\textbf{x}^{a}_{0:K})=$ $\displaystyle\
\sum_{k=1}^{K}(\mathcal{H}_{k}(\textbf{x}_{k}^{a})-\textbf{y}_{k})^{T}R_{k}^{-1}(\mathcal{H}_{k}(\textbf{x}_{k}^{a})-\textbf{y}_{k})$
(15)
$\displaystyle+\sum_{k=0}^{K-1}(\textbf{x}_{k+1}^{a}-\mathcal{M}_{k:k+1}(\textbf{x}_{k}))^{T}{P^{f}_{k+1}}^{-1}(\textbf{x}_{k+1}-\mathcal{M}_{k:k+1}(\textbf{x}_{k}))$

with all $\textbf{x}_{k}$ generated through the sequential filtering process.
This objective differs from the weak-constraint 4D-Var objective with
$Q_{k}=P^{f}_{k}$ only by the initial prior. $Q_{k}=P^{f}_{k}$ is not an
especially sensible choice for model error covariance and the mechanics of
generating the predicted trajectories are quite different (direct optimization
vs. sequential updates), but the connection is nonetheless interesting as it
indicates that our method may be interpreted as a direct amortization of the
weak-constraint 4D-Var procedure under certain conditions.

##  Appendix C Proofs

###  C.1 Lemma 1

######  Lemma 2 (Noise2Self – Restricted).

Suppose concatenated noisy observation vector
$\textbf{y}=[\textbf{y}_{k};\textbf{y}_{k+1}]$ is an unbiased estimator of
concatenated state vector $\textbf{x}=[\textbf{x}_{k};\textbf{x}_{k+1}]$ and
that the noise in $\textbf{y}_{k}$ is independent from the noise in
$\textbf{y}_{k+1}$. Now let
$\textbf{z}=f(\textbf{y})=[\textbf{z}_{k};\textbf{z}_{k+1}]$. If $f$ is a
function such that $\textbf{z}_{k+1}$ does not depend on the value of
$\textbf{y}_{k+1}$ then:

$\mathbb{E}_{\textbf{y}}\|f(\textbf{y})_{k+1}-\textbf{y}_{k+1}\|^{2}=\mathbb{E}_{\textbf{y}}[\|f(\textbf{y})_{k+1}-\textbf{x}_{k+1}\|^{2}+\|\textbf{y}_{k+1}-\textbf{x}_{k+1}\|^{2}]$
(16)

_Proof:_

$\displaystyle\mathbb{E}_{\textbf{y}}\|f(\textbf{y})_{k+1}-\textbf{y
}_{k+1}\|^{2}$
$\displaystyle=\mathbb{E}_{\textbf{y}}\|(f(\textbf{y})_{k+1}-{\textbf{x}_{k+1}})-(\textbf{y}_{k+1}-\textbf{x}_{k+1})\|^{2}]$
(17)
$\displaystyle=\mathbb{E}_{\textbf{y}}[\|f(\textbf{y})_{k+1}-\textbf{x}_{k+1}\|^{2}+\|\textbf{y}_{k+1}-\textbf{x}_{k+1}\|^{2}-2\langle
f(\textbf{y})_{k+1}-\textbf{x}_{k+1},\textbf{y}_{k+1}-\textbf{x}_{k+1}\rangle]$
(18)
$\displaystyle=\mathbb{E}_{\textbf{y}}[\|f(\textbf{y})_{k+1}-\textbf{x}_{k+1}\|^{2}+\|\textbf{y}_{k+1}-\textbf{x}_{k+1}\|^{2}]-2\sum_{j=1}^{d}\mathbb{E}_{\textbf{y}}[(f(\textbf{y})_{k+1}^{(j)}-\textbf{x}_{k+1}^{(j)})(\textbf{y}_{k+1}^{(j)}-\textbf{x}_{k+1}^{(j)})]$
(19)
$\displaystyle=\mathbb{E}_{\textbf{y}}[\|f(\textbf{y})_{k+1}-\textbf{x}_{k+1}\|^{2}+\|\textbf{y}_{k+1}-\textbf{x}_{k+1}\|^{2}]-2\sum_{j=1}^{d}\mathbb{E}_{\textbf{y}}[f(\textbf{y})_{k+1}^{(j)}-\textbf{x}_{k+1}^{(j)}]\mathbb{E}_{\textbf{y}}[\textbf{y}_{k+1}^{(j)}-\textbf{x}_{k+1}^{(j)}]$
(20)
$\displaystyle=\mathbb{E}_{\textbf{y}}[\|f(\textbf{y}_{k+1})-\textbf{x}_{k+1}\|^{2}+\|\textbf{y}_{k+1}-\textbf{x}_{k+1}\|^{2}].$
(21)

Note that the superscript $j$ indicates vector element indices while the
subscript corresponds to block entries corresponding to time indices.
Statement (7) uses Fubini’s theorem to move the expectation integral into the
summation. We then use the assumed independence of $f(y)_{k+1}^{(j)}$ and
$y_{k+1}^{(j)}$ to convert the expectation of the product of independent
random variables into the product of expectations. Finally,
$\mathbb{E}_{\textbf{y}}[\textbf{y}_{k+1}^{(j)}-\textbf{x}_{k+1}^{(j)}]=0$ as
$\textbf{y}_{k+1}^{(j)}$ is an unbiased estimator of $\textbf{x}_{k+1}^{(j)}$,
resulting in the final term vanishing and giving us the stated identity.

###  C.2 Proposition 1

######  Proposition 3 (Zero Loss Regime).

Under the stated assumptions, let $f_{\theta}$ denote a family of functions
parameterized by $\theta$ for which $\min_{\ \theta}\
\mathcal{L}^{f}(\theta)=0$. For any $\theta$ which achieves this minimum, it
is also true that $\mathcal{L}^{a}(\theta)=0$ and that $\theta$ is in set of
minimizers of $\mathbb{E}_{\textbf{y}_{0:K}}[\mathcal{L}^{ssf}(\theta)]$.

_Proof:_ Plug in $\mathcal{M}_{k:k+1}\circ
f_{\theta}(\textbf{x}_{k}^{f},\textbf{y}_{k})$ as the function used in Lemma
2. This obeys the assumptions that $f(\textbf{y}_{k})$ is independent from
$\textbf{y}_{k+1}$ as observation noise is assumed to be uncorrelated in time
and the composition predicts $\textbf{y}_{k+1}$ based entirely on prior
information. Since the supervised component of the loss is zero, the expected
self-supervised loss consists only of the irreducible variance term and thus
is also minimized.

For the analysis loss, assumption 1 gives us that the solution to the initial
value problem is unique and therefore the preimage of $\mathcal{M}_{k:k+1}$
contains a single point for all inputs in the range of $\mathcal{M}_{k:k+1}$.
Since our forecast is produced by applying the operator, the forecast is
clearly in the range of the operator. Thus, if $\mathcal{M}_{k:k+1}\circ
f_{\theta}(\cdot)=x_{k+1}$ then $f_{\theta}(\cdot)=x_{k}$ resulting in an
analysis loss of zero as well.

###  C.3 Proposition 2

######  Proposition 4 (Non-zero Loss Regime).

Under the previously stated assumptions, the supervised analysis loss can be
bounded by the supervised forecast loss as:

$\frac{1}{K-1}\sum_{k=1}^{K-1}\bigg{\|}\bigg{(}\frac{1}{m}\sum_{i=1}^{m}f_{\theta}(\cdot))\bigg{)}-\textbf{x}_{k}\bigg{\|}\leq\frac{e^{L(\tau_{k+1}-\tau_{k})}}{K-1}\sum_{k=1}^{K-1}\max_{i\in[m]}\bigg{\|}\mathcal{M}_{k:k+1}\circ
f_{\theta}(\textbf{x}^{f}_{i,k},\cdot)-\textbf{x}_{k}\bigg{\|}$ (22)

_Proof:_ We construct the bound using the Lypanuov spectrum of the system
defined by dynamics $g$. The Lyapunov exponents of the system characterize the
rate of separation of perturbed system states. Since we are bounding the
separation at an earlier time point by the separation at a later time, the
relevant quantity is the the minimal Lyapunov exponent $\lambda_{1}$, which
assuming the value is negative, bounds the rate at which perturbations in the
initial conditions can shrink:

$\bigg{\|}\mathcal{M}_{k:k+1}\circ
f_{\theta}(\textbf{x}^{f}_{i,k},\cdot)-\textbf{x}_{k+1}\bigg{\|}\geq
e^{\lambda_{1}(\tau_{k+1}-\tau_{k})}\bigg{\|}f_{\theta}(\textbf{x}^{f}_{i,k},\cdot)-\textbf{x}_{k}\bigg{\|}$
(23)

The Lyapunov exponents are inherently local properties, but as $g$ is
$L$-Lipschitz, we can globally bound the Lyapnuov exponents $|\lambda_{i}|\leq
L$ for all $i=1,\dots,d$. Plugging this in and rearranging, we get:

$\bigg{\|}f_{\theta}(\textbf{x}^{f}_{i,k},\cdot)-\textbf{x}_{k}\bigg{\|}\leq
e^{L(\tau_{k+1}-\tau_{k})}\bigg{\|}\mathcal{M}_{k:k+1}\circ
f_{\theta}(\textbf{x}^{f}_{i,k},\cdot)-\textbf{x}_{k+1}\bigg{\|}$ (24)

which using linearity and the fact that the maximum distance from an ensemble
member cannot exceed the distance from the mean of the ensemble yields the
stated bound.

##  Appendix D Incomplete or Unknown Observations

###  D.1 Incomplete Observations

Figure 8: Partial observations were included in the convolutional architecture
by replacing missing values with forecast values from a shuffled ensemble.

We experimented with two approaches for incorporating partial observations
with known observation operators into convolutional architectures. In both, we
incorporated an additional channel with indicators of .1 if the coordinate was
observed and -.1 if the coordinate was not observed. The simpler of the two
approaches was simply to pass a value of zero in the observation channel for
all unobserved coordinate.

The second approach, which was used in our experiments, used partial
observations as an opportunity to cross-pollinate the update with information
from other ensemble members. In this approach (Figure 8), each observation
channel consisted of the actual observed values in the coordinates that were
observed while the missing information was replaced with the forecast values
from other ensemble members. The information from the ensemble was treated as
though it was part of the observation in that the imputed values were detached
from the computation graph so that no gradient information flowed backwards
through this process.

###  D.2 Unknown Observation Operators

Figure 9: The process for learning to use unknown observation operators. In
this example $\mathcal{H}_{2}$ is learned from data.

While not the focus of our work, amortized assimilation does admit a natural
approach for learning observation operators which we describe here (Figure 9).
We do not evaluate the loss based on unknown observation operators. Instead,
the unknown observation type is processed by an additional independent neural
network whose architecture may depend on the data type of the unknown input.
This learned observation operator is then trained using only gradient
information flowing backwards through the differentiable simulation from
losses evaluated using future, known observation types. This enables direct
assimilation of data types without direct mapping to state variables like
hyperspectral imaging or other satellite-derived information.

##  Appendix E Experiment Settings

Figure 10: Outline of neural network architecture used for learned
assimilation step in all experiments.

###  E.1 L96 and VL20

All experiments used the same architecture outline as Figure 10. For L96 and
VL20, all convolutional filters had width $5$ with circular padding of 2. Each
layer, apart from the final readout layer, used 64 convolutional filters. Six
channels of memory were used. Dropout was set to .2. We used a covariance
width of 3, indicating that channels were included for the coordinate variance
and covariances with the coordinates to the left and right of the current
entry. Training was done using mini-batches of size 64.

###  E.2 KS Equation

The KS experiments used the same architecture as the L96 and VL20 experiments
but with slightly different hyperparameter settings. For these experiments, we
used a 64 filters per layer with a filter size of $7$ and circular padding of
size $3$. The dropout rate was set to .2 and the covariance width was set to
5. Memory depth was set to 4. The system was discretized uniformly into 128
spatial coordinates and training used mini-batches of size 32.

###  E.3 Hyperparameter Search

We tuned the comparison methods on validation sets generated from 5 different
initial conditions through a grid search. This grid search was performed
independently for each tested dynamical system, ensemble size, and noise
level. For testing, the parameters giving the highest average RMSE over the
starting conditions were used. The parameters searched are listed in table 2.

The AmEnF was tuned in a more ad-hoc manner with variables moved up or down
the ladder of available options when instabilities occurred. This primarily
affected the initial learning rate as we found the approach was largely
agnostic to sensible settings of the dropout rate while performance generally
improved with longer sequences during training at the expense of additional
memory consumption. Note: for the AmEnF, this only includes parameters
explored while tuning the presented architecture. Decisions on number of
filters and batch sizes were made to fully utilize available GPU memory
subject to the state size of the dynamical system, the number of unrolled
steps, and the ensemble size. All AmEnF models were trained with an ensemble
size of $m=10$ and tested using ensembles of various sizes.

Table 2: Hyperparameter search space.  Method | Parameter | Range
---|---|---
4D-Var | B | {Historical Covariance, $I$}
| B-Scale | $\\{.1,.2,.4,1,1.25,1.5,.75,2,2.25,2.5,3\\}$
LETKF | Inflation | $\\{1.00,1.01,1.02,1.04,1.07,1.1,1.2,1.4,1.7,2.00\\}$
| Rotation | {True, False}
| Localization Radius | $\\{0.1,0.2,0.4,0.7,1.0,2.0,4.0,7.0,10.0,20.0,40.0,70.0\\}$
iEnKS | Inflation | $\\{1.00,1.01,1.02,1.04,1.07,1.1,1.2,1.4,1.7,2.00\\}$
| Rotation | {True, False}
AmEnF | Dropout Rate | $\\{.1,.2,.3,.5\\}$
| Initial Learning Rate | $\\{1e-2,5e-3,1e-3,8e-4,5e-4,3e-4,1e-4\\}$
| Steps per Training Seq | {1, 10, 20, 40}

##  Appendix F Paper Checklist

  1. 1.

For all authors…

    1. (a)

Do the main claims made in the abstract and introduction accurately reflect
the paper’s contributions and scope? [Yes]

    2. (b)

Did you describe the limitations of your work? [Yes]

    3. (c)

Did you discuss any potential negative societal impacts of your work? [N/A]

    4. (d)

Have you read the ethics review guidelines and ensured that your paper
conforms to them? [Yes]

  2. 2.

If you are including theoretical results…

    1. (a)

Did you state the full set of assumptions of all theoretical results? [Yes]
They’re listed immediately before the theoretical results are discussed.

    2. (b)

Did you include complete proofs of all theoretical results? [Yes] In the
supplementary materials.

  3. 3.

If you ran experiments…

    1. (a)

Did you include the code, data, and instructions needed to reproduce the main
experimental results (either in the supplemental material or as a URL)? [Yes]
Yes, the code is included in a linked repository

    2. (b)

Did you specify all the training details (e.g., data splits, hyperparameters,
how they were chosen)? [Yes] though moved to the supplement for space reasons

    3. (c)

Did you report error bars (e.g., with respect to the random seed after running
experiments multiple times)? [Yes] Yes, all comparisons to other methods
include error bars

    4. (d)

Did you include the total amount of compute and the type of resources used
(e.g., type of GPUs, internal cluster, or cloud provider)? [Yes] The hardware
used was listed.

  4. 4.

If you are using existing assets (e.g., code, data, models) or
curating/releasing new assets…

    1. (a)

If your work uses existing assets, did you cite the creators? [Yes] Yes. All
existing libraries are cited under experiments

    2. (b)

Did you mention the license of the assets? [No]

    3. (c)

Did you include any new assets either in the supplemental material or as a
URL? [No]

    4. (d)

Did you discuss whether and how consent was obtained from people whose data
you’re using/curating? [N/A] All data is simulated by the authors

    5. (e)

Did you discuss whether the data you are using/curating contains personally
identifiable information or offensive content? [N/A] All data is simulated
from governing equations

  5. 5.

If you used crowdsourcing or conducted research with human subjects…

    1. (a)

Did you include the full text of instructions given to participants and
screenshots, if applicable? [N/A]

    2. (b)

Did you describe any potential participant risks, with links to Institutional
Review Board (IRB) approvals, if applicable? [N/A]

    3. (c)

Did you include the estimated hourly wage paid to participants and the total
amount spent on participant compensation? [N/A]