Learning Dynamical Characteristics with Neural Operators for Data Assimilation
Authors: Yi Xiao
Wei Xue
Venue: Submitted to ICLR 2023 (2023)
Keywords: AI for science
data assimilation
generative models

Abstract: Data assimilation refers to a group of algorithms that combines numerical models with observations to obtain an optimal estimation of the system's states. In areas like earth science, numerical models are usually formulated by differential equations, also known as the prior dynamics. It is a great challenge for neural networks to properly exploit the dynamical characteristics for data assimilation, because first, it is difficult to represent complicated dynamical characteristics in neural networks, and second, the dynamics are likely to be biased. The state-of-the-art neural networks borrow from the traditional method to introduce dynamical characteristics by optimizing the 4D-Var objective function in which the dynamics are inherently quantified, but the iterative optimization process leads to high computational cost. In this paper, we develop a novel deep learning framework with neural operators for data assimilation. The key novelty of our proposed approach is that we design a so-called flow operator through self-supervised learning to explicitly learn dynamical characteristics for reconstructed states. Numerical experiments on the Lorenz-63 and Lorenz-96 systems, which are the standard benchmarks for data assimilation performance evaluation, show that the proposed method is at least three times faster than state-of-the-art neural networks, and reduces the dynamic loss by two orders of magnitude. It is also demonstrated that our method is well-adapted to biases in the prior dynamics. 

Under review as a conference paper at ICLR 2023
LEARNING DYNAMICAL CHARACTERISTICS WITH NEURAL OPERATORS FOR DATA ASSIMILATION
Anonymous authors Paper under double-blind review
ABSTRACT
Data assimilation refers to a group of algorithms that combine numerical mod- els of system with observations obtain optimal estimation of the system's s states. In domains like earth science, numerical models are usually formulated by differential equations, also known as prior dynamics. neural networks to properly exploit the dynamical characteristics for data assim- ilation, because first, it is difficult to represent complicated dynamical character- istics in neural networks, and second, the dynamics are likely to be biased. The state-of-the-an neural networks emulate traditional approaches to introduce dy- namical characteristics by optimizing an objective function in which the dynamics are inherently quantified, but the iterative optimization process leads to high com- putational cost. In this paper, develop novel deep learning framework with neural operators for data assimilation. The key novelty of our proposed approach is that we design a so-called flow operator to explicitly learn dynamical charac- teristics for reconstructing sequences of physical states. Numerical experiments the Lorenz-63 and Lorenz-96 systems, which the standard benchmarks for data assimilation performance evaluation, show that the proposed method is at least three times faster than oftte-of-theaar neural networks, and reduces the dy- namic loss by two orders of magnitude. It is also demonstrated that our method is well-adapted biases in the prior dynamics.
great challenge for
INTRODUCTION
Data assimilation can be essentially defined as statistical technique to combine the prior dynamics with a sequence of noisy and aregularly-samplee observations, which plays an important role in utilizing observational data, especially for numerical weather forecasting systems. For instance, the European Centre for Medium- Range Weather Forecasts (ECMWF) employs the variational assimi- lation algorithms in its operational systems to take full advantage of both in situ and -atellite-derived data, as well as state-of-the-ar numerical models (Rabier et al. 2000).
Classical Assimilation Methods Classical data assimilation methods range from early empirical analysis (Bergthorsson et al. 1955) and optimal interpolation (Gandin) 1963) to later variational- based assimilation algorithms (Sasaki 1970) and the filtering methods based on statistical estima- tion theory, such as Kalman filter (Welch et al. 1995), and particle filter (Carpenter et al.) 1999). Although classical methods have been essential tools for improving the capability of major numer- ical weather prediction centers worldwide (Kalnay 2003), they generally do not account for the propagation of observation information along time and therefore are unable to effectively utilize observations at different times (Song et al. 2017).
Time-dependent Assimilation Method In recent years, new data assimilation methods have been developed (Evensen 2003 Lorenc & Rawlins, 2005 Hunt et al.) 2007). Unlike previous methods, they consider the revolution observation information along time, other words, the time-dependent dynamical characteristics (Song et al. 2017). Among them, the four-dimensional variational (4D-ar) assimilation algorithm is a cutting-edge one. Experimental results prove the advantage of 4D-Var Var assimilation methods in utilizing observational information (Lorene & Rawl- ins] 2005). However, due to the complicated modeling and solving process of the 4D. Var algorithm, is considered computationally expensive, especially for high resolution cases (Fisher 1998).

---

Under review as a conference paper at ICLR 2023
Machine Learning for data assimilation Prior to 2020, extensive research into developing ma- chine learning for data assimilation had been successful in reducing computational costs (Hsich &
Tang 1998; Vijaykumar al., 2002, de Campos Velho et al. 2002; Härter & de Campos Velho Cintra & de Campos Velho 2018; Maak et al. 2020). However, these works base their the- ory on classical assimilation methods, thus suffering from low accuracy brought by not considering time-dependent dynamical characteristics. From 2020, a lot of neural network-basee methods con- sider time-dependent assimilation. They can be generally divided into three categories: learning the inverse observation operator, learning model biases, and learning the optimization algorithm.
2008
Learning Inverse Observation Operators The aim of learning inverse observation operators is to use neural networks to construct a mapping from observations to reconstructed states. The representative work is Frerix et al. (202T), which the observational data is exploited by a neural operator to provide better initial states for the optimization of 4Dr objective function. This work has at least two flaws. The first is that the neural network is used only to exploit observational information, while the integration of observation and dynamical characteristics is still implemented by Var, leading to a limited improvement in computational efficiency. The second that the strict constraints of dimendepeneent dynamical characteristics which prevent it from generalizing to the cases in which the prior dynamics are biased.
the prerequisites for this work,
Learning Model Biases As mentioned above, the prior dynamics can be biased, and another group of works uses neural networks to learn the bias the prior dynamic models for data assim- ilation (Brajard al.) 2020) Arcueci et al. 2022 Farchi et al. 2021 [Bonavita & Laloyaux] 2020). These works generally achieve a balance between accuracy and efficiency, but the improvement in computational efficiency is insignificant, because the assimilation of dynamical characteristics and observations is still done in an iterative manner.
Learning Optimization Algorithm Since the main computational overhead of the 4D-Va algo- rithm lies in the process of iteratively optimizing the 4D-Var objective function, some studies have focused on replacing traditional optimization algorithms with neural networks. Among them, the representative one is Fablet et al. (2021b a). This work simultaneously learns both the dynamical characteristics of the model and the optimization algorithm of the 4D-Var objective function through neural networks. It not only takes into consideration the weak constraints of time-dependent dynam- ical characteristics, but also successfully reduces the number iterations compared with traditional 4D- -ar algorithms, subsequently improving computational efficiency. We consider this work to be the state-of-ther f--th-e-rttt art method for comparison. Despite the achievements of the work, there is still room for improvement in terms of both accuracy and efficiency.
We develop new approach of combing data-driven strategies with model- driven methods for time-dependent data assimilation. This framework consists of three operators implemented by neural networks, namely, the inverse observation operator, the perturbator, and the flow operator. Among them, the cooperation of the perturbator and the flow operator decouples the two goals of learning observations and learning dynamical characteristics, and they succeed in continuously refining reconstructed states.
Our Contributiinss
To our knowledge, this is the first work that directly uses neural networks to blend the dynam- ical prior with sequence of observations without explicit formulations of the 4D-Var objective function for assimilation.
We test the proposed framework on the Lorenz-63 and Lorenz-96 systems. Experimental results support the effectiveness of our framework. The proposed architecture is at least three times faster than the state-of-thed -ott--rrrt neural network Fablet et al. (202Tb). reduces the dynamic loss by two orders of magnitude and is comparable to the state-of-the-art neural networks in terms of the reconstruction error.
We design experiments in which the prior dynamies different from ground truth to simulate real- woaldworld scenarios. Experimental results prove that our method is well-adapted to deviations of the prior dynamic model.
2 PROBLEM STATEMENT
Suppose that there is a dynamic system with d dimensions to be observed and estimated, and we denote the state variables by x(t) IRd. Consider N evenly distributed time points, +
2

---

Under review as a conference paper at ICLR 2023
1 i - where At is short enough for the integration step. We are interested in the assimilation problem over a fixed-lengee Lendt period. [to ty- and we call the period of interest "assimilation window" (Ir emolet 2006, Fablet et al.] 2021b). Our goal is to estimate the states of the system on these N time points, x; x(4:). All the information we have is (1) the a priori knowledge of the dynamics and (2) the observations of the system states during this period. We note that in our work, the guess of the initial states X is not considered, and we focus on the integration of the prior knowledge and the observational information, as done in Frerix al. 2021 and Fablet et al. 2).1b).
The Prior Knowledge real-world scenarios, scientists develop physical models to estimate the dynamic system. We regard the physical model as the prior knowledge, and name it "the prior dynamies". The prior dynamics are usually represented by the following differential equation,
dx(t) = Mpr (x(t)),
(I)
d/
where Mpr a mapping from Rdd to RR The prior dynamics can be either biased or unbiased.
The Observations The observations are indispensable for the reconstruction of the physical states. In our problem setting, the observations are assumed be discrete and partial, with white Gaussian noise. This assumption is widely considered in the data assimilation community, espe- cially for developing novel data assimilation algorithms (Brajard et al. 2020; Farchi et al. 2021 Arcucci et al.) (021). Suppose that all the observations are obtained at time 1, to RAA,
J = 0, I, ..., N/ R and only variables whose indexes belong to S2, are observed at time ty. The formulation can be written
yp(t) (t) c(t), Vt E (to, to + RAt, (N/R 1) Vp Q
(2)
where c(t) represents the white Gaussian noise.
3 LEARNING FRAMEWORK WITH THREE OPERATORS
The major difficulty of the problem lies in how to properly utilize the prior dynamic model. Numer- ous studies have developed methods based on Bayesian inference to deal with the time series states estimation problem by formulating the transition process in probabilistic manner (Krishnan al.) 2017 Karl al. 2016 Marino et al. 2018 Kingma & Welling 2013). Despite their success, is still difficult to properly deal with the biases of the prior dynamics for data assimilation.
Our work approaches this issue from a different angle. Instead of explicitly formulating the transition process, simultancously generate all physical states within the time window at once by an inverse observation operator. The prior dynamics are fed into our model via a flow operator that follows behind, which works with perturbers to progressively refine the reconstructed state. The overall structure of our framework is shown in Figure
3.1 TWo METRICES
Before proceeding to the detail of our framework, let first illustrate two important metries for training and evaluation.
Reconstruction Loss free The reconstruction loss is defined the mean square error between the generated sequence of states and the ground truth states. Denote Xgt = ((xx))) (x.r)..... Nx1)) the discrete sequence of true states, x [X., x...... x x quence of generated states, then
the se-
(rea(x) | Xgx xpel1 1
Dynamic Loss Cdyn considered dynamics. It is an unsupervised metric, formulated by
The dynamic loss measures how well the given sequence of states fits the
Ldyn (x) = 1b (x) - x||
(4)
where + defined by
{
o(x)o = Xx0 ((x); Mp 1
(5)
3

---

Under review as a conference paper at ICLR 2023
Rough Physical States
Sparse Observations
sp00e
90322
Inverse Observation Operator
Flow Operator
amo
Refined Physical States
Flow
Flow
Perturbator
Perturbator
Operttor
Operator
2B blocks
Figure 1: Proposed deep learning framework. The inverse observation operator (orange box) and the flow operator (blue box) aim to convert sparse observations to complete physical states and bring them closer to the prior dynamics, respectively, while the second row demonstrates the process of refining reconstructed physical states with the help of perturbators (green box) and flow operators.
Mpr (x;-x;;1: t.) represents integrating X;-1 step ahead from t;-1 to t4, using the prior dy- namic model Mpr- (xx) is also sequence of states. The output at the ith time point, denoted by ((x);, is obtained by integrating X;-1 step forward. The RK45 integration scheme (Dormand & Prince 1980) is adopted throughout our paper. The smaller the dynamic loss, the better the sequence of states fits the dynamics. the sequence of states perfectly matches the dynamic system, then (xx x holds (ignore the error of integration scheme), and the dynamic loss is zero.
3.2 How THE THREE OPERATORS WORK
Inverse Observation Operator Our framework starts with the inverse observation operator A., which aims to learn lower-dimensional observations and convert them into physical states (Frerix et al. 202T). With the help of the inverse observation operator, we are already able to construct the sequence of physical states, but with relatively low accuracy. Instead of seeking help from traditional 4D- Var algorithms to further improve accuracy (Frerix et al. 2021), we construct the so-called perturbator and the flow operator for such refining purposes.
Flow Operator The main function of the flow operator is to reduce the dynamic loss of the reconstructed sequence of states and output sequence that better conforms to the prior dynamics. Suppose that x is reconstructed sequence of states, and we require that the dynamic loss after the operator is smaller than given value Lo. which is usually orders of magnitude smaller than the original dynamic loss. Denote the flow operator by F, and the requirement can be formulated as
(F(x)) < Co < Cx)
(6)
We emphasize that although it is generally infeasible for neural networks to learn how to make arbi- trary inputs satisfy prior dynamics, in our case, the input already conforms the previous dynamics to some extent, which makes the goal achievable.
The role of the flow operator to guarantee that the reconstructed sequence con- forms to the dynamics, but this is insufficient for data assimilation, because there are infinitely many sequences of states that conform to the dynamics, and we need to find the one that also fits the ob- servations well. On the other hand, thanks to the flow operator's s ability to guarantee that the output consistent with the prior dynamics, we can add a new neural network, called the perturbator, in front of the flow operator, which has more free space utilize observational information by learning from ground truth labels. can use the learned knowledge to revise the reconstructed sequence of states without considering the problem of conforming to the prior dynamics.
Perturbator
In this way, we decouple the two goals conforming to observations and conforming to prior dynamics, and both the flow operator and the perturbator can perform their duties. Specifically,

---

Under review as a conference paper at ICLR 2023
the perturbator uses what it has learned from the observations and labels to perturb the reconstructed states away from the original flow, and then the flow operator corrects the deviated flow to make it a new flow that conforms to the prior dynamics.
How the Perturbator and the Flow Operator Work Together The perturbator and the flow operator always work in pairs. Given sequence of states with relatively low dynamic loss Ldyn and relatively high reconstruction loss Cree, assuming that we have trained a flow operator F, the function of the perturbator P can be formulated as
Crea (Fe((F(P())) S (x) Cdyn (F(((x()))) ~ Layn (x)
The first formulation requires that the reconstruction loss is reduced, and the second requires that the dynamic loss should remain in the same order of magnitude. Ideally, more pairs of perturbators and flow operators are applied to the input, the reconstruction loss will continuously decrease while maintaining the dynamic loss.
We emphasize that the output of the perturbator P(x) is likely to have both high reconstruction loss and high dynamic loss, but it does not matter because we only require that the losses after the flow operator F (P(x)) are small.
3.3 NETWORK TRAINING
Our framework trained using sequences of true states via supervised learning. Since the overall framework adopts a checkerboard style for refining the reconstructed states, we design a specific training procedure for all the three operators work properly. We first use the reconstruction loss function to train single inverse observation operator; then flow operator is added after the inverse observation operator and trained with the dynamic loss. Afterwards, we gradually add pairs of perturbators and flow operators, and train and fine-tune them with hybrid loss. More details are shown in Algorithm We embed the unobserved regions with zeros to make the input observations y the same size the output physical states
Stopping Condition In the case which the prior dynamics are biased, we stipulate a stopping criterion for the training procedure in order to alleviate overfitting, which is formulated as follows:
(8) where Ex,, [Cdyn (x is the expectation of the prior dynamic loss of ground truth states, which measures the difference between the prior dynamics and ground truth. can estimated in advance by averaging the dynamic loss over all samples in the training set. Each time the dynamic loss of the outputs Cdyn (x) is calculated, we compare it with Ex,, [Cayn (xg:))) if the incquality i not satis- fied, we end the current training phase and move on to the next one. Inequality guarantees that the dynamic loss of the reconstructed states is no smaller than that of the true states. is necessary but not sufficient condition to avoid overfitting in terms of dynamics, and any reconstructed states with smaller dynamic loss than true states indicate that more information than expected has been learned from the prior dynamics. Specifically, if the prior dynamics accurate, then Ex,, [Ldyn (X.t)] O. inequality triiiillly holds, and we don't 't need to consider the stopping condition.
Layn (x) 2 Ex,, (Xgt))
3.4 NETWORK STRUCTURES FOR OPERATORS
The network structures are based the Residual U-Net structure (Ronneberger al.) 2013 Sid-
dique et al. 202T) for all three operators, as shown in the Appendix. The hierarchical structure of O-Net allows to capture both the large-scale and small-scal information of observation states and
is ubiquitously used in the reconstruction domain (Guan et al. 2019. Siddique et al.) 2021 Punn
Agarwal 2022 Feng et al., 2020). Moreover, the residual structure helps alleviate the vanishing gradient problem (He et al. 2016), allowing U-Net models with deeper layers to be designed.
In addition, this paper also modifies the network according to the specificity of the data assimilation issue. Unlike images, of which the three dimensions are length, height, and channels, the inputs of data assimilation are sequences of observations (states), whose dimensions are time, space, and channels. The length and height of images are symmetric, while time and space generally not. To tackle this difference, we carry out a dimensionality reduction by maxpooling only in the time dimension for the U-Net.
5

---

Under review as a conference paper at ICLR 2023
Algorithm Training procedure of our proposed framework (without stopping condition)
Require:
Ground Truth Dataset D = {y, x) observationsy E IRdx.N states x E R Initialize the inverse observation operator Ae, the flow operator Fg. and the perturbator Pv.
Choose the number of blocks 2B + 2, the hyperparameter À > 0, and different epoch numbers eA,ep,
for epoch from 1 to 6A do for each batch X) E do do
Phase I: training the inverse observation operator.
(X corresponds to the data batch (y, x).
X Ag(Y), La +- (
Update 0 with the Adam optimizer (Kingma Ba 2014
end for end for for epoch from 1 to do for each batch (Y, X) do X - F-Fo(Ao Ls + Ldyn (X) Update with the Adam optimizer. end for end for
Phase 2: training the flow operator.
The inverse observation operator is fixed.
for from 1 to B do
b Phases 3 ~ B + 2: training the perturbator and fine-tuning. epoch 1: training; epoch 1: fine-tuning
for epoch from 1 Ep do for each batch X) do
X + (F& o Pv) o Fo. AY()), Las +- (xre((X) + A (X)
Update & with the Adam optimizer.
The flow operator is fixed.
end for end for for epoch from 1 to ep do
epoch 1: fine-tuning
for each batch (Y D do o o Fg Ay(Y), Ls + Lrec(X) + Addy Update @ with the Adam optimizer.
The perturbator fixed.
end for end for end for
As for the inverse observation operator, add skip layer from the input to the bottom hidden variables for better information transfer. As for the flow operator, we add residual block from the input the output to make the training process converge faster. More details can be found in the Appendix.
noted that our framework is limited to the U-Net. More advanced networks, such as the work of Wang et al. (2020), can also be adopted in the future to better handle other more difficult chaotic systems.
4 NUMERICAL EXPERIMENTS
In this section, we design the numerical experiments to test the performance of our framework in different scenarios.
4.1 DESIGN OF EXPERIMENTS
Dynamic Systems The first step is to select dynamic systems to test on. Lorenz systems are widely used for the evaluation of various data assimilation algorithms, owing to their chaotic prop- erties and similarities with meteorological systems, especially in data-driven and machine learning case-studies (Lguensat et al. 2017 Raissi 2018: Fablet et al.) 2021b Frerix et al. (021). In our work, we choose the Lorenz-63 63 and Lorenz 96 systems for evaluation. The dimension of the Lorenz- 63 system is three: the dimension of the Lorenz-96 system modifiable, and we set it to 40.

---

Under review as a conference paper at ICLR 2023
The Prior Dynamics training of the flow operator. We construct the prior dynamics by varying parameters in the ordinary differential equations (ODE).
The second step is to determine the prior dynamics, which is crucial for the
The Lorenz-63 and Lorenz-96 dynamic systems are governed by the ordinary differential equations (ODE) expressed as Equation 9 and Equation T0 respectively (Lorenz 1963 1996).
dX₃
= 0 - X1)
dt
dX₂
pXx X2 X1
(9)
dXs
= X2X22 - BX8
de
dX; = (Ii+1 24-2) 24-1 Ijj + F, = 1,2,...,40
(10)
dt
As for the ground truth dynamics, we assign 10 p 28 8 and F for the Lorenz-96 system. To construct the equations for the prior dynamics, we increase by Ao in Equation 9 for the Lorenz-63 system and F by AF in Equation I0 for the Lorenz-96 system.
8/3 for the Lorenz-63 system
Dataset Construction Before constructing the dataset, parameters like N, At, R, Sr should be determined in advance. Refer to the Appendix for these experimental settings. We do a long inte- gration over time using the ground truth dynamics and divide the time sequence into three sections. The first section is for building the training set, the second is for the validation set, and the last is for the test set. For each section, continuously randomly intercept sequences of length N dataset samples. Also refer to the Appendix for detailed settings.
4.2 UNBIASED PRIOR DYNAMICS
First, we consider the cases in which the prior dynamics are identical the ground truth dynamics. We use total of eight (B = 3) and ten blocks (B 4) to train the deep learning framework for the Lorenz-63 and the Lorenz-96 systems, respectively, as shown Figure
We compare our framework with two baselines. The first Fablet et al. (2021b), which is considered as the state-of-the-arr. The second framework, but follows simple training procedure: (1) train the inverse observation operator and the flow operator in the front with the hybrid loss; (2) add one pair of the perturbator and the flow operator the back, and then train the entire network with the hybrid loss; (3) repeat (2) until the number of blocks meets the setting of our models.
U-Net baseline, which shares the same structure of our
Apart from the reconstruction loss and the dynamic loss defined above, evaluate the efficiency of our framework by test time. Test time refers to the time for the neural network to finish evaluation for round the test dataset. Shorter test time indicates higher computational efficiency.
The experimental results are demonstrated in Table I Our proposed method reduces the reconstruc- tion error by 44% and 3% for the Lorenz-63 and the Lorenz-96 systems, respectively, and reduces the dynamic loss by two orders of magnitude, compared with Fablet et al. (2021b). In terms of test time, our proposed framework four times and three times faster than |ablltt al. (2021b) in both dynamic systems. Our method also outperforms the U-Net baseline, in terms of both Eree and Ldyns which proves the advantage of our training procedure.
We draw an example from the test dataset of the Lorenz-63 system to further illustrate how our architecture works, as shown in Figure 2 Focus on the blue oval box area the upper right corner, where we find how the cooperation of the perturbator and the flow operator succeeds in correcting the rough reconstructed sequence of states. In the beginning, the physical states are clearly inconsistent with the ground truth in the upper right region. After perturbator is added, disturbances with high frequencies are involved in the misaligned area to distinguish them from well-constructed areas. flow operator is further added behind the perturbator, the freeuuegyyncy noise disappears, and the inconsistency is corrected to some extent. Such a perturbation misalignment becomes smaller with repetition.
perturbationn-refiiemmn refinement step is repeated and the
7

---

Under review as a conference paper at ICLR 2023
Dynamic System
Model
of Blocks
Lree
Layn
AVG (s)
SD (8)
Lorenz-63
Our work
2
1.28
1.10e-3
3.1
0.65
Our work
4
0.923
3.82e-4
4.5
0.48
Our work
6
0.799
2.08e-4
5.8
0.63
Our work
8
0.751
1.52e-4
6.4
0.40
U-Net baseline
8
0.849
3.99c-4
7.2
0.56
Fablet et al. (2021b)
20
1.34
1.00e-2
27.1
3.80
Lorenz-96
Our work
0,571
1.40e-3
7.8
0.60
Our work
4
0,450
5.09c-4
9.8
0.97
Our work
6
0.397
3. 3.14e-4
11.5
0.93
Our work
8
0.373
2.66e-4
13.5
0.77
Our work
10
0.367
2.36e-4
15.8
0.72
U-Net baseline
8 0.599
2.24e-3
12.8
0.88
Fablet et al. (2021b)
20
0.38
>1.00e-2
47.3
1.46
Table Results of comparative experiments. We report the performance of the proposed framework and two baselines for the Lorenz-63 and Lorenz-96 systems. The results of Fablet et al. (2021b) draw directly from the original paper (Fablet et al. 2021b). Bold values refer to the best score. All of the models tested on a
with 2000 samples for Lorenz-63 and 3000 samples for Lorenz- 96. The batch size is set to eight, and all the evaluations are performed on one TITAN RTX GPU. Test time may vary across experiments, and we provide the average (AVG) and standard deviation (SD) of test time from 10 independent tests.
-38
-10
-25
223
325
159
106
125
156
175
206
330
725
350
779
Coarse
Refined Reconstructed States
Flow Operator
Reconstructed
Perturbator
Perturbator
Operator
Stales
10
Legend:
ground truth
reconstructed
15
106
100 155
186
116
206
An intuitive presentation of how the refining structure works. A sample is randomly drawn from the test dataset of the Lorenz-63 system for demonstration. The rough reconstructed states are generated by the first two operators (the inverse observation operator and the flow operator). We use two pairs of perturbators and flow operators refine the input, and extract the intermediate results generated by each block to demonstrate the detailed procedure.
Figure 2:
4.3 BIASED PRIOR DYNAMICS
Next, we consider the cases in which the prior dynamics are biased. Specifically, we select Ao from the set {0.25, 0.50, 0.75, 1.00, 1.50, 2.00, 2.50. 3.00, 3.50, 4.00, 4.50, 5.00) for the Lorenz-63 system, select AF from the set (0.50, 1.00, 1.50, 2.00. 2.50, 3.00, 3.50, 4.00, 4.50, 5.00] for the Lorenz-96 system, and conduct numerical experiments accordingly.
8

---

Under review as a conference paper at ICLR 2023
We report results in Figure Take the Lorenz-63 system an example. As illustrated in the upper-right figure, when Ao is smaller than 1.0, the increase in Lree is smaller than 20%, which shows that our framework is robust to small biases. When Ao is smaller than 2.00, the Free increases as Ao increases, whereas if Ao is larger than 2.00, the Free fluctuates without monotonic changes. Such phenomenon can be explained by looking into how network structure works: When Ao is small, the difference between the prior dynamics and the ground truth is small, and the information of the prior dynamics can be successfully exploited by the flow operator. Therefore, as Ao gets larger, useful information provided by the prior dynamics decreases and the performance of our framework worsens. However, if Ao too large, no longer necessary for the flow operator to learn from the bad prior dynamics; it is cheaper to learn directly from ground truth labels. Hence, in this case, the value of Ao does not affect the performance of the neural network. The results shown in the upper-left figure can further confirm our explanation. When Ao < 2, the difference between Layn (x) and Cagn (xyt) is close to zero, which is consistent with our stopping criterion (Inequality 8), so the network works as expected. However, if Ao 2. the bias of prior dynamics is too large for the flow operator to learn, and solely learning from the ground truth is sufficient for the neural network to generate states whose dynamic loss is smaller than Ldya (Xgz). The findings above also hold for the Lorenz-96 system, as shown in the lower figures. The difference is that the threshold changes from 2.0 around 2.5. They provide further evidence that our network learns features from prior dynamics rather than simply remembering ground-trute labels, and that our framework works properly despite biases in prior dynamics.
Dynamic Loss of the Lorenz-63 system
Reconstruction Loss of the Lorenz-63 system
0.012
EsyO0i
1.4
0.010
CarlKg)
1.3
1.2
0.005
L.1
10.
0.004
0.002
0.8
Dut
Dynamic
Lorenz-96 system
Reconstruction Loss the Lorenz-96 system
0.05
0.05
CorlKy) Lgali)
0 04
0.03
0.5
0.01
0.50
0.4
2
Figure 3: The left panels show the dynamic loss of reconstructed states (blue line), the biases of prior dynamics (orange line), and their differences (green dotted line) against Ao or AF, The right panels show the reconstruction loss (Crec) in the test set against Ao AF. The upper and lower panels represent the results of the Lorenz-63 and Lorenz-96 systems, respectively.
5 CONCLUSIONS
We have introduced deep learning framework for data assimilation. The key novelty of our work is that we develop the flow operator learn from the prior dynamics and continuously refine the recon- structed sequence of states. The results of numerical experiments on the Lorenz-63 and Lorenz-96 systems support the relevance of our framework. In terms of assimilation accuracy, our framework has reduced the dynamic loss by two orders of magnitude, and at least maintained the reconstruc- tion accuracy, compared with the state-of-the oftte-of-hheaase art method. In terms of computational efficiency, our framework has been found be four times and three times faster than the tfate-of-the-ar method on the Lorenz-63 and Lorenz-96 systems, respectively. It has also been found to be well-adapted to biases in the prior dynamics.

---

Under review as a conference paper at ICLR 2023
REFERENCES
Rossella Areucci, Jiangcheng Zhu, Shuang Hu, and Yi-Ke Guo. Deep data assimilation: integrating deep learning with data assimilation. Applied Sciences, 11(3):114, 2021.
Bergthorsson, BR Döös, S Fryklund, O Haug, and R Lindquist. Routine forecasting with the barotropic model. Tellus, 7(2):272-274, 1955.
Massimo Bonavita and Patrick Laloyaux. Machine learning for model error inference and correction. Journal of Advances in Modeling Earth Systems, I2(D))c2020MSS00223,, 2020.
Julien Brajard, Alberto Carrassi, Mare Bocquet, and Laurent Bertino. Combining data assimilation and machine learning to emulate dynamical model from sparse and noisy observations: A case study with the lorenz 96 model. Journal of Computational Science, 44:101171, 2020.
James Carpenter, Peter Clifford. and Paul Fearnhead. Improved particle filter for nonlinear problems. IEE Proceedings-Rulee Sonar and Navigation, 146(1) ):2-7, 1999.
Rosangela Saher Cintra and Haroldo F de Campos Velho. Data assimilation by artificial neural net- works for atmospheric general circulation model. Advanced upplications for artificial neural networks, 265, 2018.
HF de Campos Velho, S Stephany, AJ Preto, NL Vijaykumar, and AG Nowosad. A neural network implementation for data assimilation using mpi. WIT Transactions on Information and Commu- nication Technologies, 27, 2002.
John R Dormand and Peter J Prince. family of embedded runge-kutta formulae. Journal of computational and applied mathematics, 6(1):19-26, 1980.
Geir Evensen. The ensemble kalman filter: Theoretical formulation and practical implementation.
Ocean dynamics, 53(4):343-367, 2003.
Ronan Fablet, Mohamed Mahmoud Amar, Quentin Febvre, Maxime Beauchamp, and Bertrand Chapron. End-to-end physics-infomae informed representation learning for satellite ocean remote sensing data: Applications satellite altimetry and sea surface currents. ISPRS Annals of Photogranme- try, Remote Sensing & Spatial Information Sciences, (3), 2021a.
Ronan Fablet, Bertrand Chapron, Lucas Drumetz, Etienne Mémin, Olivier Pannekoucke, and François Rousseau. Learning variational data assimilation models and solvers. Journal of Ad- vances Modeling Earth Systems, (3(10)::202 2021b.
Alban Farchi, Patrick Laloyaux, Massimo Bonavita, and Mare Bocquet. Using machine learning to correct model error in data assimilation and forecast applications. Quarterly Journal of the Royal 2021.
Meteorological Society, 147(739):3067
Jinchao Feng, Jianguang Deng, Zhe Li, Zhonghua Sun, Huijing Dou, and Kebin Jia. End-to-end res-unet based reconstruction algorithm for photoacoustic imaging. Biomedical optics express, 11(9):5321-5340, 2020.
M Fisher. Minimization algorithms for variational data assimilation. In Proc. of ECMWF Seminar on Recent Development in Numerical Methods for Aumospheric Modelling (7-11 September 1998, Reading, UK). ECMWF, 1998.
Thomas Frerix, Dmitrii Kochkov, Jamie Smith, Daniel Cremers. Michael Brenner, and Stephan Hoyer. Variational data assimilation with a learned inverse observation operator. In International Conference on Machine Learning, pp. 3458. PMLR, 2021.
Lev Semenovich Gandin. Objective analysis of meteorological field. Giddomeetcoollogiceekkee Izdate stvo, 286, 1963.
Steven Guan, Amir A Khan, Siddhartha Sikdar, and Parag V Chitnis. Fully dense unet for 2-d sparse photoacoustic tomography artifact removal. IEEE journal of biomedical and health informatics, 24(2):568-576 2019.
10

---

Under review as a conference paper at ICLR 2023
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, 2008.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp- 770-778, 2016.
William W Hsieh and Benyang Tang. Applying neural network models to prediction and data anal- ysis in meteorology and oceanography. Bulletin of the American Meteorological Society, 79(9): 1855-1870, 1998.
Brian Hunt, Eric J Kostelich, and Istvan Szunyogh. Efficient data assimilation for spatiotemporal chaos: A local ensemble transform kalman filter. Physica D: Nonlinear Phenomena, 230(1-2): 112-126, 2007.
Eugenia Kalnay. Amospheric modeling, data assimilation and predictability. Cambridge university press, 2003.
Maximilian Karl, Maximilian Soelch, Justin Bayer, and Patrick Van der Smagt. Deep varia- tional bayes filters: Unsupervised learning state space models from raw data. arXiv preprint 1rXi::1605.06432, 2016.
Diederik Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint arXiv:1412.6988, 2014.
Diederik Kingma and Max Welling. arXiv:1312.6114.222 2013.
Auto-encoding variational bayes.
arXiv preprint
Rahul Krishnan, Uri Shalit, and David Sontag. Structured inference networks for nonlinear state space models. In Proceedings of the AAAI Conference Artificial Intelligence, volume 31, 2017.
Redouane Lguensat, Pierre Tandeo, Pierre Ailliot, Manuel Pulido, and Ronan Fablet. The analog data assimilation. Monthly Weather Review, 145(10):4093-4107 2017.
Andrew C Lorene and F Rawlins. Why does 4d-var beat 3d-var? Quarterly Journal of the Royal Meteorological Society: A journal of the atoosphrric sciences, applied meteorology and physical 2005.
oceanography, 131(613))3247-3223
Edward N Lorenz. Deterministic nonperiodic flow. Journal of atmospheric sciences, 200211130-141, 1963.
Edward N Lorenz. Predictability: A problem partly solved. In Proc. Seminar on predictability, volume 1996.
Julian Mack, Rossella Areucci, Miguel Molina-Solana and Yi-Ke Guo. Attention-based convolu- tional autoencoders for 3d- variational data assimilation. Computer Methods in Applied Mechanics and Engineering, 372:113291 13291, 2020.
Joseph Marino, Milan Cvitkovic, and Yisong Yue. A general method for amortizing variational filtering. Advances in neural information processing systems, 31, 2018.
Narinder Singh Punn and Sonali Agarwal. Modality specific u-net variants for biomedical image segmentation: survey. Artificial Intelligence Review, pp. 1-45, 2022.
Florence Rabier, Heikki Järvinen, E Klinker, J-F Mahfouf, and A Simmons. The ecmwf opera- tional implementation of four-dimensional variational assimilation. i: Experimental results with simplified physics. Quarterly Journal of the Royal Meteorological Society, 1463644) 2000)
Maziar Raissi. Deep hidden physics models: Deep learning of nonlinear partial differential equa- tions. The Journal of Machine Learning Research, 19(1):93-9955 2018.
11

---

Under review as a conference paper at ICLR 2023
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedi- cal image segmentation. In International Conference on Medical image computing and computer- assisted intervention, pp. 234-241. Springer. 2015.
Yoshikazu Sasaki. Some basic formalisms in numerical variational analysis. Monthly Weather Review, 98(12):875-883, 1970.
Nahian Siddique, Sidike Paheding, Colin P Elkin, and Vijay Devabhaktuni. U-net and its variants for medical image segmentation: review theory and applications. Ieee Access, 9:82031-82057,
2021.
Hyo-Jong Song, Seoleun Shin, Ji-Hyun Ha, and Sujeong Lim. The advantages of hybrid 4denvar in the context of the forecast sensitivity initial conditions. Journal of Geophysical Research: Atmospheres, 122(22): 2017.
Yannick Tr' emooe.. Accounting for imperfect model in 4dr Quarterly Journal of the Royal Meteorological Society: A journal of the atosppericc sciences, applied meteorology and physical oceanography, 132(621):24833 2504, 2006.
Nandamudi Lankalapalli Vijaykumar, Stephan Stephany, HF Campos Velho, AJ Preto, and AG Nowosad. Optimized neural network code for data assimilation. In Brazilian Congress on Meteorology, volume 12, pp. 38413849, 2002.
Gabriele Vissio and Valerio Lucarini. Mechanics and thermodynamics of new minimal model of the atmosphere. The European Physical Journal Plus, 135(10):1211, 2020.
Rui Wang, Karthik Kashinath, Mustafa Mustafa, Adrian Albert, and Rose Yu. Towards physics- informed deep learning for turbulent flow prediction. In Proceedings of the 26th ACM SIGKDD International Conference Knowledge Discovery & Data Mining, pp. 1457-1466, 2020.
Greg Welch, Gary Bishop, al. An introduction to the kalman filter. 1995.
A DETAILED NETWORK STRUCTURE FOR OPERATORS
We take the three-layer U-Net as an example to show the detailed network structures for the inverse observation operator (Figure 4 and the perturbator (Figure 5 The structure of the flow operator almost the same as that of the perturbator, except that skip layer is added from the input to the output.
B EXPERIMENTAL SETTINGS
B. DYNAMIC SYSTEMS SETUP
In the Lorenz-63 system, t = 0.01, that the time step for integration is 0.01, and the number of time points in the assimilation window 200. Only the first component of the physical states is observed. The observations are sampled every eight steps, and we also add Gaussian noise with variance of 2.0 them.
200, R 8, Var[e(t)] = 2.0, Ω = (1}. This means
In the Lorenz-96 system, t 0.05, N = 200, R = 4, Var[e(t)] = 2.0, Ωr (sample 20 random integers between I and 40 for each t). This means that the time step for inte- gration is 0.05, and the number of time points in the assimilation window is 200. The observations are sampled every four steps, and each time we observe 20 random components of the physical states. Gaussian noise with a variance of 2.0 is also added to the observations.
B.2 DATASET INFORMATION
In the Lorenz-63 system, the training data comprises 10000 sequences sampled from time series with 1 12000 steps; the validation data comprises 2000 sequences from time series with 5000 steps and the test data comprises with 3000 sequences from 10000 time steps. In the Lorenz-96 system, the training data comprises 2000 sequences sampled from time series with 7500 steps; the validation
12

---

Under review as a conference paper at ICLR 2023
h+h
KxN
2h+2h
KxNN
Conv 3x3 SILU
4h
4h4h
Conv 3x3
KxN,
Max Pool
TrunsConv
Obs Time
KXIN,
Concat
Figure 4: An example of the 3-layer U-Net architecture of the inverse observation operator. We denote K the spatial dimension of the physical states, the number of channels (the total dimension of the system is d = с K), N the total sequence length, / the hidden dimension, N, the sequence length with valid observations, N₁, the intermediate hidden sequence length. The input of the inverse observation operator is observations. Before input, the sparse observations are padded to the same size the output sequence physical states by padding unobserved regions with zeros.
h
h+h
h
KxN
2h
2h+2h
2h
Conv SiLU
KxNN
Conv
MaxPool
4h
TransConv
KxN,
Concat
Figure 5: An example of the 3-yeyrr U-Net architecture of the perturbator. We let readers refer to Figure 4 for details. Both the input and output of the perturbator are a sequence of physical states.
data comprises 300 sequences from time series with 2500 steps and the test data comprises with 2000 sequences from 10000 time steps. Settings for the training and validation set correspond to those in [Fbblet et al. (2021b), and we increase the size test data to increase the credibility of experimental results.
13

---

Under review as a conference paper at ICLR 2023
B.3 NETWORK SETTINGS
In the Lorenz-63 experiments, the input is reshaped as e x K x N, where e = 3, K = 1, and N 200. four-layer U-Net adopted, with hidden dimension experiments, the input is reshaped as с x K x N, where с = K = 40, and N = 200. A three- layer U-Net is adopted, with hidden dimension = 32. In both experiments, = 100, C = 20, er 200, 20 for training according to Algorithm [
32. In the Lorenz-96
C DETAILED RESULTS OF TRAINING
C.I TRAINING COST
The training cost of our framework is affordable one TITAN RTX GPU. The training speed is around three times that of Fablet et al. (2021b). Details are demonstrated Table 2 and Table for the Lorenz-63 and Lorenz-96 systems, respectively.
Model
Phase
Time per epoch
# of epochs
Phase time
Total time
Fablet et al. (2021b)
Phase I block))
m 38 s
20
32 m 40 175
Phase 2 (10 blocks)
45 s
100
Total time
207 m 40 s
Our work
Phase
7
20
2 m 20
Phase 2
10 s
200
33 m 20
Phase 3 (4 blocks)
24 s
20
8 m
Phase 4 (6 blocks)
36
20
12 m
Phase 5 (8 blocks)
44
20
14 m 40
Total time
70 m 20
Table 2: Training time for the Lorenz-63 system. If the validation loss does not continue decrease, we consider the network to converge and end the current phase. The training procedure of Fablet phases: the first phase is trained with 10 blocks and the second phase is trained with 20 blocks. We should at least allocate 20 epochs for the first phase and 100 epochs for the second phase to make the training loss converge well. We trained our framework up to eight blocks, strictly following the training procedure of Algorithm []
al. (2021b) is divided into
The phases of our work correspond to those defined Algorithm [ Both models are trained on a TITAN RTX GPU with batch size set to 8. The statistics of time per epoch, number of epochs for each phase, phase time and total time are shown in the table.
Model Fablet et al.
Phase
Time per epoch
# of epochs
Time
Phase 1 (10 blocks)
2 m 34 s
20
51 m 20
(2021b)
Phase 2 (10 blocks)
m 45 s
100
475 m
Total time
526 m 20
Our work
Phase I
9
20
3 3
Phase
33 s
200
110
Phase 3 (4 blocks)
33
20
11 m
Phase 4 (6 blocks)
48
20
16
Phase 5 (8 blocks)
63 s
20
21 m
Total time
161 m
Table 3: Training time for the Lorenz-96 system. We let readers refer to Table 2 for details.
C.2 LOss CHANGE DURING TRAINING
The variations of the reconstruction loss and the dynamic loss against the number of epochs for the Lorenz-63 and Lorenz-96 systems are reported in Figure 6 and Figure 7 respectively. In the first phase, the inverse observation operator learns a mapping from the observational space to the physical states. Since shift between two different spaces takes place in this period, we see steep decrease
14

---

Under review as a conference paper at ICLR 2023
Phose Train the inverse observation operator
Phases Train the perturbator Eine-tane
Phuse Triin the flow operator
1011
10-2
train Cyee
10-3-
20
50
100
150
200
250
15
epochs
10
5-
50
100
150 of epochs
200
250
0.5
epochs
Figure 6: The variation of the reconstruction loss free and the dynamic loss Ldy for training and validation with the number of epochs for the Lorenz-63 system. Two blue vertical dashed lines divide the training procedure into three periods: training the inverse observation operator, training the flow operator, and phases 3~5. The green vertical further divide the phases 3~5 to two periods: training the perturbator and fine-tuning. The left panel demonstrates the changes in the reconstruction loss during training, the details of which are enlarged and shown in the bottom right panel. The red and blue dashed lines in the left panel and the bottom right panel show how continuing to train the inverse observation operator without adding new operators would result in change in the reconstructed loss. The top right panel demonstrates the variation of dynamic loss during training.
1.9
1.6
1.0
10-3
0.4
train Cve
train Gan
0.2
of epochs
epochs
Figure 7: The variation of the reconstruction loss Cree and the dynamic loss Layn with the number of epochs for the Lorenz-96 system. The left panel demonstrates the changes in the reconstruction loss during training, while the right panel demonstrates the variation of the dynamic loss.
of loss within the first 20 epochs. After the training of the inverse observation operator is finished, we start the training of the flow operator to correct the flow dynamics of the reconstructed states. The training of the flow operator is relatively slow because an unsupervised loss is employed in this procedure. We claim that the flow operator is indispensable in further decreasing the reconstruction loss and dynamic loss, for the reason that without the flow operator, the inverse observation operator would quickly lead to overfitting, in other words, the difference between the training loss and vali- dation loss will diverge, shown in the red and blue dashed lines in the figure. There is sudden
15

---

Under review as a conference paper at ICLR 2023
increase in the loss function after the second blue vertical line because a perturbator with untrained parameters enters our framework. During the fine-tuning procedure, the reconstruction loss drops steadily while the dynamic loss fluctuates. This is not surprising considering that the fine-tuning procedure alternates between the perturbators and the flow operators.
Reconstructed examples and associated absolute errors
work blocks)
work (6 blocks)
work blocks)
Fablet (2021b)
Figure 8: The sequences reconstructed states and their corresponding absolute error maps of the Lorenz-96 96 system, with the horizontal axis being the time axis and the vertical axis being the spatial axis. the displayed figures in the same row share the same colormap.
D SAMPLE STUDY OF THE LORENZ-96 SYSTEM
Figure 8 demonsttrttes the physical states reconstructed by different neural networks for the Lorenz- 96 system and their associated absolute error maps. Visually, the dark regions in the error map becomes smaller and lighter as the number of block grows, which confirms the effects of the collab- oration of the perturbators and the flow operators in this example.
E FURTHER EXPERIMENTS
E. COMPARATIVE STUDY ON BIASED PRIOR DYNAMICS
Reconstruction Loss of the Lorenz-63 system
Reconstruction Loss of the Lorenz-96 system
1.4
work
1.0
0.8
0.9
0.8
Ib)
0.4
Figure 9: Results of our work and Fablet et al. (2021b) under biased prior dynamics. The left panel shows the reconstruction loss (Cree) against Ao for the Lorenz-63 system; the right panel shows the reconstruction loss against AF for the Lorenz-96 system.
Here, compare the performance of our work with Fablet et al. (021bb under biased prior dynam- ics. Since Fablet et al. (2021b) did not realize their algorithms under this experimental setting, we implement their algorithms based on best understanding of their work. The results are shown in worth noting that the results of the case in which bias equals zero are close to those presented in their paper, which confirms the credibility of our experiments.
Figure
As for the Lorenz-63 system, the reconstruction loss their model is larger than that of our model the prior bias is smaller than 2.5, and converges around 1. 1.20~ 1.35 as the prior bias gets larger
16

---

Under review as a conference paper at ICLR 2023
than 2.5, which is similar to the performance of our model. As for the Lorenz-96 system, the recon- struction loss of their model seems to diverge as the prior bias gets larger, while the reconstruction loss of our model converges to around 0.7. These comparative experiments show that proposed model more adaptive to the prior bias than the baseline model. The difference performance of their model under different dynamical systems be attributed the use of two different neural networks in the two cases.
E.2 ABLATION STUDY
We add two ablation experiments to show that 1) the refining process by perturbators and flow op- erators greatly contributes to reducing the reconstruction error the states generated by the inverse observation operator; 2) when the experimental setting for observations is changed, only the inverse observation operator needs retraining, and the network parameters for refining process can remain unchanged.
Experiment 1 one single inverse observation operator. It has been found that the checkerboard refining pro- cess can further reduce the reconstruction error by at least 50%, compared with one single inverse observation operator. The detailed results are shown Table
We remove all the perturbators and flow operators, and test the performance of
Dynamic system
Network structure
Cree
Ldyn
Lorenz-63
InvObsOp only
1.95
>1.00e-2
Three operators (our work)
0.751
1.52e-4
Lorenz-96
InvObsOp only
0.763
> 1100e-2
Three operators (our work)
0.367
2.36e-4
Table 4: Results of the ablation study with the perturbators and the flow operatirs removed. "InvOb- sOp" is short for "inverse observation operator". The first row for each dynamic system corresponds to the network with only the inverse observation operator; the second row corresponds to our entire network, with eight and ten blocks, respectively.
Experiment Based on the Lorenz-96 system, we design two different settings for observations. In the first setting, the observations are sampled every four steps for random half of the compo- nents; in the second setting, the observations are sampled every four steps for the components whose indexes are even. We conduct the experiments by the following steps.
Step 1: Train all three operators the first experimental setting.
Step 2: Only train the inverse observation operator in the second experimental setting.
Step 3: Concatenate the inverse observation operator trained in Step with the perturbator and the
flow operator trained in Step 1 (Without further training tuning).
Step 4: Test the concatenated network the second experimental setting.
The results are reported in Table Meanwhile, we also report the results of another experiment, in which the observation settings are switched, and the same experiment procedure conducted again with the above steps. It can be seen that when the perturbator and the flow operator are transferred to another observation setting, they still work well, despite a slight increase in the reconstruction error (smaller than 10%). It also shows that the function of the observation operator and that of the perturbator and the flow operator can be decomposed well.
E.3 EXPERIMENTS ON A NON-LORENZ SYSTEM
In order to evaluate our framework's applicability to other chaotic systems, we test it Vissio- Lucarini20 (Vissio Lucarini] 2020) system. Lussii-uuc coupled system intended to represent minimal model of the earth's 's atmosphere, which is composed of two types of variables, namely, kinetic variables and thermodynamical variables. Since both of the variables are coupled, this system is complicated than the Lorenz systems. The corresponding ordinary differential
17

---

Under review as a conference paper at ICLR 2023
Observation setting
Observation setting (for training)
Lrea
Lyym
(for test) Random
Random
0.367
2.36e-4
Even
0.402
3.6le-4
Even
Random
0.304
3.6le-4
Even
0.291
2.56e-4
Table 5: Results the ablation study with the refining network (the perturbators and the flow operators) switched, on the Lorenz-96 system. The first column indicates the observation setting corresponding to the test results; the second column indicates which observation setting the pertur- bators and the flow operators are trained in. "Random" corresponds to the first setting, and "Even" corresponds to the second one.
equation is shown in Equation I1 where = ...,....N - alp 7 + F
{
dXa
-K
d0g XR+108+2 X aXx 70s +
(11)
Experimental settings The numbers of both kinetic variables and thermodynamical variables are The observations are sampled every four steps for random
set to 40. a = ? = 1.0, F 10, G half of both kinds of variables. The variance of Gaussian noise for observations is 2.0. The prior dynamics are set to unbiased.
Network design In our proposed method, we adopt three-layer U-Net similar to that in the Lorenz-96 system. The difference that we change the channel number from one two, in order to fit the characteristics of the system with two types of variables. As for the baseline, Fablet et al. (222bb does not implement their neural network for Vissio-Lucarini20 system, and we do our best realize their algorithm.
Experimental results Similar to the results in the Lorenz-96 system, our network outperforms Fablet et al. (2021b) in terms of both the dynamic loss and test time. We report the details in
Table 6
Dynamic System
Model
# of Blocks
Lree
Layn 3.57e-4
AVG (s)
SD (s)
Vissio-Luca
Our work
10
0.352
23.6
1.5
Fablet al. (2021)
0.358
3.25e-2
64.2
3.4
Table Results of comparative experiments. We report the reconstruction and dynamical loss of the proposed framework for the Vissio- Lucariniu0 system. Bold values refer to the best score. Both of the models are tested on an independent test set with 2000 samples. The batch size set to eight, and all the evaluations are performed on one TITAN RTX GPU. Test time may vary across experiments, and we provide the average and standard deviation of results from 10 independent tests.
18


Reviews:
[{"review_id": "oJBsamZaNyl", "reviewer": "Reviewer_6BrX", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "## Summary\n\nThe authors propose a new end-to-end fully data-driven method for data assimilation in which the different components of the system are fully parameterized using neural networks. Their model enjoys good accuracy and across different dynamical system parameterizations, specifically on the classical Lorenz systems.\n\n## Review Summary\n\n Promising work but requires significant improvements.\n \n1. Please clarify the role of prior dynamics in the model / training procedure. \n2. The claims regarding state-of-the-art, especially regarding 4DVar should be weakened (1) there is no explicit comparison to 4DVar, (2) there are no experiments approaching the scale or real-world applications in which 4DVar is used. The first pages of the paper read as if the authors recommend this as a substitute.\n\n## High-level Feedback\n\n- \"Complicated modeling and solving process of 4DVar...\"\n\t- If this is an important model, please remind the read of the basic outlines of how this model works. Ideally, within the mathematical framework that you use to define your model.\n- The definition of prior dynamics need to be defined early so that the paper makes sense. It is very clearly defined, but only in the results section of the paper. Essentially, it is the parameters of the dynamical system in question.\n\t- The experimental setup regarding generalization across prior dynamics needs to be clarified. Either this or the training setup. It's hard to know without answering the question: How do the prior dynamics enter your model / training setup? This key question should be clarified multiple times, e.g. in text as well as in diagrams and figures.\n- Please make this explicit: Your argument is that you have superior accuracy and generalization across different systems (e.g. prior dynamics). I believe you are missing stability, an important property of numerical simulations. You don't have to tackle this, but please mention it. Ideally, include an analysis of the stability of your learned system over long unrolls, even if brief.\n\n## Detailed Feedback\n\n- Equation 2 is sloppy because it omits time. All numerical simulations of chaotic systems diverge for large enough $t$. So what exactly does $C$ measure?\n- \"Our main goal is to construct an inverse map, implemented in the form of a neural network [...]\"\n\t- Yes, great. Please provide citations for this general idea. I was first exposed to it in the work of Max Welling and his collaborators who used this idea for Bayesian inference to map observations to latent variables. You may need to dig deeper to find the right sources.\n\t\t- https://arxiv.org/pdf/1906.02691.pdf\n- Equation 6: What is the utility of this notation if ultimately $\\Phi$ is simply a neural network? \n\t- How exactly do the prior dynamics come into play here? \n- Equation 7: Are you saying here that you want $\\mathcal F(x)$ to match the dynamics? \n\t- Why do you want to do this if you already have an implementation of the dynamics? \n\t- What if the function of the other terms in this equation? \n\t- Are you somehow trying to beat the numerical errors produced by the solver? Or this is where you get a mismatch between prior dynamics and your model (to be minimised during training).\n- Equation 8: This is a neural ODE model. Maybe it is isn't exactly the same as the original paper, but it is in that spirit. Please make mention of this.\n- Equation 9: Where is $\\mathcal L_{dyn-pr}$ defined?\n\t- This entire paragraph is strange. Please compare and contrast with a simpler classical neural network training.\n- Function of the Perturbator:\n\t- What does perturbation have to do with mapping states to observations? Or is it vice-versa? Is this a classical notion in physics? Is there any previous work / grounding of this idea? If the Perturbator is simply a network, how do you get multiple states out of single input (without adding noise somewhere in the process).\n- Numerical Experiments\n\t- Time-to-solution: what is this? Training time? Unroll time?\n\t- Replace GENN aG-Conv with the citation, e.g. Fablet et al. 2021.\n\t- You seem to use a few tricks: stopping condition, pretraining, and perhaps more. Please clarify the improvements gained by each trick using ablation experiments. Knowing what did not work is very useful to the community and provides insight into the decisions that you made for this work.\n\t- If you did 10 independent tests, please provide std for all values.\n\t- Undefined terms: assimilation window, random half of the components, a variance of 2.0 is involved (what does this mean??)\n\t- Please move the basic properties of the dataset in the appendix to the main paper: number of examples, train / test / validation setup, etc.\n\t- Please move all figures after where they are first mentioned in the text. It is very frustrating to see a figure with undefined features only to discover that it is described in some random section below.\n\t- Figure 4:\n\t\t- Don't use \"3 blocks,\" \"4 blocks,\" etc. if what you are actually testing here is the role (e.g. smoothing / unsmoothing) of flow and perturbation operators.\n\t\t- The legend is laughably small.\n\t\t- This blue arrow seems to be guiding the reader about how to read the figure. Would be easier to just make a simpler figure which is more legible. \n\t\t- The MSE plot should be separate. \n\t\t\t- Label the y-axis \"mean-squared error\"\n\t\t\t- Remove the title which has too many words\n\t\t\t- Add legend\n\t- Figure 5\n\t\t- R-score means something else, not reconstruction loss. \n\t\t\t- https://en.wikipedia.org/wiki/R_score\n\t\t- You have already define $\\mathcal L_{rec}$ , just use that.\n- Appendix\n\t- Figure 6\n\t\t- Is this the dataset? I'm confused. A simple description, e.g. 70 / 20 / 10 % train / validation / test split would suffice.\n\t\t\t- Anyway, as mentioned above, this should go into the main paper.\n\t- Figure 7\n\t\t- What is happening when you switch to fine-tuning? Why this blow up?\n\t- Figure 9\n\t\t- From 6 to 10 blocks there is little change in error. What is happening here? Shouldn't you be overfitting as model complexity grows?\n\t- Does not belong in this part of my review but: What do you mean by self-supervised loss? Please clarify this since it probably does not match up with the classical ML usage of this term.", "summary_of_the_review": "See above", "strength_and_weaknesses": "See above", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "iE-JNN0c4d", "reviewer": "Reviewer_oVQ6", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This study proposes a new method for data assimilation, i.e., for learning model-constrained dynamics from empirical data and estimating the state of the system of study. The method is composed of three types of networks (the observation operator, the flow operator and the perturbator), allowing it to flexibly balance the problem of reconstruction of the system latent states and the conforming to the specified model dynamics. On two classical problems in weather sciences, the Lorenz-63 and Lorenz-96, the method is shown to achieve higher performance than the state-of-the-art method 4D-Var (Fablet et al. 2021) in terms of reconstruction, dynamics loss, and computational cost.", "summary_of_the_review": "Overall, the paper is technically sound and has a novel methodological contribution with somewhat convincing empirical results. However, some of the empirical results need to be further expanded and the clarity of the manuscript could be improved.", "strength_and_weaknesses": "The paper is technically sound, previous work is cited and discussed, and claims are for the most part well supported by empirical evaluation. However, I have a few concerns I would appreciate the authors to address:\n\n-Table 1 shows the average results over 10 independent tests. It would be important to also report the respective standard deviations and assess the empirical performance in light of these standard deviations (e.g., performing statistical tests);\n\n-Table 1 reports test times. What is considered test times? On one hand, this should be clarified in the manuscript, and on the other hand, training times should also be reported (if reporting training times does not make sense, then this should be justified);\n\n-Figure 5 quantifies the performance of the new method under different degrees of bias in the dynamics. While this is a useful experiment, it would be equally important to show how the state-of-the-art method performs under these same circumstances.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "WRNeUTqykJ", "reviewer": "Reviewer_BfUY", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "\nThe paper proposes a deep learning framework for data assimilation technique for physics and earth science, such as meteorological dynamical processes. The framework is based on an estimation of a coarse state of the system from noisy and sparse observations, associated to a refinement using dynamical priors. Do do so, different U-net based architectures are developed: an observation operator, a flow operator, and a perturbator. Numerical experiments were performed on Lorenz-63 and Lorenz-96 one-dimentional dynamical systems, having 3 state variables. Both models were also tested with biased versions of the prior dynamics, which reproduces more accurately realistic situations.", "summary_of_the_review": "This paper is interesting and proposes a new end-to-end deep learning framework for a difficult yet important application. In particular, the authors tackled the prior dynamics bias problem. The method seems to be in its early stages, but I think it would be beneficial for the community. I have to add that while I have some knowledge in the field, I am not a data assimilation expert.", "strength_and_weaknesses": "Strenghs:\n\t- the problem tackled is of importance for a large community\n\t- the proposed method is end-to-end and able to provide a solution from both observation and prior dynamics \n\t- the results are promising, reducing considerably the loss dynamics and the computational time from previous study\n\t- the pipeline is efficient and the different U-net blocks are assembled in an intelligent way\n\t- the prior dynamic bias version seems to indeed work, which is one of the main novelties and assets.\n\t\nWeaknesses:\n\t- The technique seems to be in its early stages and it seems as if tuning needs to be performed for every new dynamical system and setting. First, the different training phases are difficult to follow and makes me wondering if the pipeline is robust enough. Also, please explain how was chosen to use a 4-layer U-Net for the Lorenz-63 and only a 3-layer U-Net for the Lorenz-96. \n\t- Network structure: it is not clear which losses are minimized at the different phases, in particular, the perturbator+flow operator is not clear. From eq. 11, it looks as if the dynamics and the data fidelity losses were applied on the same current state x_hat, while my understanding was that L_rec was calculated on the perturbator's output while L_dyn was calculated on the flow operator's output. If the hybrid loss is calculated on the final output of the 2 blocks, I don't se how we can enforce the decoupling of the two goals as stated.\n\t- It is difficult to understand the size of the input data and latent data. In particular, please give axis labels on Figure 1 images, as it first looks like a 2D spatial problem, while later it is explained that the input is of size time and location. Please also clarify what does T represent in Figure 2: I guess the time dimension, and in this case where is the location dimension?\n\t- Please explain best the following sentence: 'For the case in which the prior dynamics are unbiased': in a real setting, how can we know if the prior dynamics are biased or not?\n\t- 'espilon(t) represents the white Gaussian noise.' --> isn't it reductive? \n\t- Why stopping at 8 and 10 blocks for the experiments, while it looks as if the results are always improving with more blocks?\n\nTypos:\nSpecifically, The perturbator\n\nThe perturbator uses the observations and labels it has learned to perturb the reconstructed states to make it deviate from the original flow --> please rephrase\n\nFigure 3: the color do not match Figure 1, as here the colors are also linked to training/no training. Maybe use another sign to indicate training/no training, as a red line surrounding the box.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "74wWf6oNfH", "reviewer": "Reviewer_wimX", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "A network architecture and the corresponding training strategy for solving data assimilation problems are proposed.", "summary_of_the_review": "While the direction of the paper looks nice, the current empirical study does not tell much about the specific benefit of the proposed learning procedure, which makes the paper incomplete.", "strength_and_weaknesses": "### Strengths\n\n- Method is simple.\n- Improvement is clearly shown, compared to at least one baseline.\n\n### Weaknesses\n\n1. The motivation behind the design choice of the proposed method is not overly clear. In the caption of Figure 3, the authors explain that the proposed learning procedure is necessary for making the models perform as they expect. Though I agree with it per se, I do not think it explains the fundamental motivation. In particular, why do you want to iteratively apply the $\\mathcal{F}$ and $\\mathcal{P}$? How is the finetune-based procedure beneficial compared to learning all the networks at once?\n\n2. The experiments lack comparison to an obvious baseline, that is, a network comprising UNet blocks trained altogether. This lack of comparison makes it difficult to assess the particular benefit of the proposed learning procedure.\n\n3. Solving state estimation problems just like data augmentation using neural nets has been an active research area; e.g.:\n\n- Krishnan+, Structured inference networks for nonlinear state space models, AAAI 2017\n- Karl+, Deep variational Bayes filters: Unsupervised learning of state space models from raw data, ICLR 2017\n- Marino+, A general method for amortizing variational filtering, NeurIPS 2018\n\namong many others. However, the current paper does not address the potential relation to these kinds of studies. Maybe the authors' interest does not lie in Bayesian inference, but references from the outside of the data augmentation community would be highly valuable to make the paper more complete.\n\n---\n\n### Minor points\n\n- Naming a network an \"observation operator\" without cautionary statements may be confusing especially for machine learning researchers. A natural interpretation is that such an operator takes the hidden state $x$ as input and gives the observation $y$ as output. But, what could be read from the paper is the opposite, that is, it takes $y$ as input and spits out an estimation of $x$.\n- In the \"Modified U-Net Structure\" part, how you do the dimensionality reduction is unclear. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "6Hrz85P0_Yl", "reviewer": "Reviewer_1MTc", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The work presents an extension to the 4D-Var approach to data assimilation geared towards application on climate data based on neural operator learning. To do so the authors split the neural operator into three distinct operators, the flow-operator, the perturbation-operator, and the observation operator. With a checker board-style training algorithm the individual operators are trained and then tested on two Lorenz-benchmarks.", "summary_of_the_review": "The authors present a compelling extension to existing 4-D Var data assimilation approaches with a neural operator learning framework, which splits its operators into three distinct operators fulfilling different complementary functions. Overall a strong paper with a very clear exposition, and presenting a novel approach to operator learning for data assimilation, and showing strong performance on Lorenz-systems, the paper would benefit from more ablation analyses, examination across a wider spectrum of benchmarks, as well as at times more references to adjacent literature such as iterative refinement, and similar training algorithms. ", "strength_and_weaknesses": "Strengths:\n- Clear algorithm derivation, with a very clear exposition of its individual components which make the contribution of the authors very clear.\n- Strong relation of the algorithm design to the roots in the physical process itself, hence making this a true physical process-driven design.\n- Highly intriguing approach to keep the neural operator's dynamics faithful to the dynamics of the physical process. An approach which should also carry over to similar problems in scientific machine learning\n- Strong performance of the algorithm across the presented benchmarks, clearly outperforming previous results.\n\nWeaknesses:\n- Missing discussion of the training costs of the proposed algorithm, which make it hard to gauge the extensibility of the algorithm to larger, more complex problems.\n- No clear ablations of the influence that each of the 3 operators has on the general performance of the approach. While not immediately clear to the reviewer if the algorithm could be decomposed as such, it would greatly improve the technical depth of the paper if the influences of each operator on the performance could be discerned.\n- Evaluations are restricted to the Lorenz-63, and the Lorenz-96 system. Evaluation on more difficult and/or non-Lorenz systems would be of great help to evaluate its applicability to other chaotic systems.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject