Learning to Assimilate in Chaotic Dynamical Systems
Authors: Michael McCabe
Jed Brown
Venue: NeurIPS 2021 Poster (2021)
Keywords: Data assimilation
self-supervised denoising
dynamical systems
uncertainty quantification

Abstract: The accuracy of simulation-based forecasting in chaotic systems is heavily dependent on high-quality estimates of the system state at the beginning of the forecast. Data assimilation methods are used to infer these initial conditions by systematically combining noisy, incomplete observations and numerical models of system dynamics to produce highly effective estimation schemes. We introduce a self-supervised framework, which we call \textit{amortized assimilation}, for learning to assimilate in dynamical systems. Amortized assimilation combines deep learning-based denoising with differentiable simulation, using independent neural networks to assimilate specific observation types while connecting the gradient flow between these sub-tasks with differentiable simulation and shared recurrent memory. This hybrid architecture admits a self-supervised training objective which is minimized by an unbiased estimator of the true system state even in the presence of only noisy training data. Numerical experiments across several chaotic benchmark systems highlight the improved effectiveness of our approach compared to widely-used data assimilation methods.

Learning to Assimilate in Chaotic Dynamical Systems
Michael McCabe
Jed Brown Department of Computer Science University of Colorado Boulder jed@jedbrown. org
Department Computer Science University of Colorado Boulder mi michael, mccabeccolooado..o....... edu
Abstract
The accuracy of dimulation-based forecasting in chaotic systems is heavily de- pendent on high-quality estimates of the system state at the time the forecast is initialized. Data assimilation methods are used infer these initial conditions by systematically combining noisy, incomplete observations and numerical models of system dynamics produce effective estimation schemes. We introduce amorlized 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. Experimen- tal results across several benchmark systems highlight the improved effectiveness of our approach over widely-used data assimilation methods.
Introduction
Forecasting in the geosciences is 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/7), forecasts produced by estimating the current state of the dynamical system of interest and integrating that state forward 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 trivilivill. Discretization introduces numerical errors and it is often too computationally expensive to directly simulate the system of interest 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 [ This can be problematic when the initial condition for 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 firee-dimensional discretization of the Earth's 's atmosphere ( To perfectly initialize a simulation, one would need to know the true value all of these state variables simultaneously. This data is simply not available. In reality, what is available are noisy, partial measurements scattered non-uniformly time and space.
The inverse problem of estimating the true state of a dynamical system from imperfect observations is the target of work and of the broader field known 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 sparse M.]. Thus for accurate estimates of the full system state, it is important to utilize both the model and observations.
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

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 [ but have also been used in the context of
numerical simulation [3], and meta-learning 14 Our work extends this approach to variational
data assimilation [1S
In this work, we introduce 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 neural network that then behaves as 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 the dynamical systems setting through the use of differentiable simulation. We then support this theory through a set of numerical experiments) 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 In this work, we consider a dy- namical system with state variable x(1) C Rdd where C' is compact subset of Id and system dy- namics defined by the differential equation:
yk-
yk
yk++1
dx g(x(t))
(1)
dt
with g as time-invariant, Lipschitz continuous map from Rd + Rd. The state is observed at a sequence
Tk
Kk
Øk+1
of K discrete time points {To: FI..... TK generating the time-series Yo.K (No. YI.... ................................ .x} where 0 < K is an index corresponding to the time point Tk- These observations imperfect representations of the system states Xo:K {xxo X1, xx}
Figure 1: Observations assumed to be generated from a Hidden Markov Model.
generated from observation operators of the form:
Ya Hx(( + 77
where Hy is an arbitrary potentially nonlinear function and nhs ~ (N(00000) (1). This noise model is often referred to Hidden Markov Model (Figure I In the classical assimilation setting, the observation operators are assumed to be known priori though our proposed method is capable of learning subset of observation operators. For problem like numerical weather prediction, these operators may represent measurements taken of certain atmospheric quantities at particular measurement location.
Data assimilation is then the inverse problem of estimating the true trajectory No:K from our noisy observations Yo:k and our model of the system evolution:
x(Tx11) MMRR+1 =x(Tk) + f 9(x(t)) dt.
Sequential Filtering In the absence of ground truth, is practical to present data assimilation from statistical perspective. In particular, the quantity of interest that we are concerned with in this paper is the filtering distribution p(xx Yo.4). Efficient algorithms for estimating the filtering distribution
Code available htttpp://gicchb..co//lleeeeeeeee http:////ghhooooooa.acccaea//a/a

---

Analysis
x(t)
v(T++1
(Tk)
xa (Tk+1)
(TK+1)
7k
Forecast
Tk+1
Forecast
Tk+2
Figure 2: Ensemble filters model the time-evolution of uncertainty under nonlinear dynamics by maintaining set of samples from the state distribution and simulating their trajectories forward in time. Observations are then used refine these forecast estimates in what called the analysis step.
are derived from exploiting dependence relationships in the generative model to produce the two step cycle:
p(xe Yo.k-1) = / p(xx Xic-1) dx
Forecast:
(4a)
Analysis:
p(x: Yo.k) x p(y& xx)p((x Yock-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, will use x and xf to refer to forecast and analysis estimates of X₃ respectively.
In the case of linear dynamics with Gaussian uncertainties, the Kalman filter [8 provides both the forecast and analysis distributions exactly in closed form. However, the general case significantly more complicated. Under nonlinear dynamics, computing the forecast distribution requires the solution of the Fokker-Plance partial differential equation [19]. While finding an exact solution to these equations is computationally infeasible, one produce what is called 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 the form of 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 particle filter approaches while enforcing Gaussian assumptions leads to the Ensemble Kalman filter (EnKF) [16]. The former is rarely used data assimilation settings as particle filters can become degenerate in high dimensions (1 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 often reasonable, but the former is 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 system by finding the state No which minimizes the negative log posterior density through nonlinear optimization. This 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

---

Discretizatin
Ye
Obss indicator
Observation
( Am
Ma
Forecast
Residual Blocks
Linear Conv
fou
Dynamics
(x()))=
Covariance
Sigmoid Conv
Memory
Trainable
-Set-valued Out
Static Operation -Snngle Out
Numerical Simulation
Figure 3: The AmEnF replaces the analysis step of conventional ensemble filter with neural network trained using historical data. The model uses an augmented state which includes both standard inputs as well as recurrent memory.
methods be augmented by auxiliary objectives promoting concordance with physical properties like conservation laws.
3 Amortized Assimilation
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 act as an analysis update mechanism. Before diving into the theory, we present the general amortized assimilation framework from the perspective of specific architecture which we call the Amortized Ensemble Filter (AmEnF). The AmEnF (Figure 3 replaces the EnKF analysis equations with parameterized function in the form of neural network whose weights are optimized for specific dynamical system during training. By unrolling fixed number of assimilation cycles as first proposed in Haarnoja et al. [25), the AmEnF can be trained efficiently using backpropagation through time 226 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 xX_ for each ensemble member i E [m]] with denoting the number of ensemble members, the ensemble covariance P and the observations Yę- However, we augment these features with the dynamics g(x,,)) and coordinate-wist recurrent memory c resulting in the following update equations:
x -M-1: (x,,
pf Cov(( [Cou((x{A)
2 Zo; = [x Pl. Nx. s
Ae, A = SN P Fr, ((x (, (5) C, = Ao, ① C +2c
x = Ae, 2 x Zv.
where Sr. and JN are neural networks with linear and sigmoid final activations respectively. We implement these as single network which 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 additional feature channels where each channel represents the covariance with the state value at a fixed relative spatial position. For instance a ID system, the relative covariance channels may contain the variance at the given point and the covariances with the points 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 memory allows the model to learn from the past without explicitly including previous states and observations. This is significant benefit in practice. The other non-standard inclusion, the local dynamics, led a small performance boost as well. our method 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- self-supe 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 J-invariance framework which exploits independence assumptions between noise different partitions of the state vector to derive valuable decomposition of the denoising loss. Techniques developed using this approach have exhibited comparable performance to explicitly supervised denoising across multiple applications [28-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 J -invarriance framework to derive results, we begin by restricting Proposition from Batson and Royer [17] to our setting,
Lemma (Noise2Self- Restricted). Suppose concatenated noisy observation vector y = D'gi is an unbiased estimator of concatenated state vector = [x+i and that the noise in yk is independent from the noise in YR+1- Now let z = f(y) = [kx:Z++1] 1. I is a function such that Zk+1 does not depend on the value of Ms+1 then:
ff - Mx+1 2= Ey f)) - -xx+1l/ + xx+1
(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 sule-supe 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 fo denote a neural network parameterized by weights 0, we can formalize the objective that we'd d actually like to minimize as the supervised analysis loss:
(²
fo NE. ((x(x))) (.
K
This supervised loss reflects that our goal is to produce the best initial condition estimates the start of a forecast window. Unfortunately, X₁ 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:
K Σ kml
H4+1 folx( Ys+1 +
-0)
( K-
(8)
The self-supervised approach (Figure has the obvious advantage that it uses readily available noisy observations the training target. However, we can also show that it 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 S 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' (t) = g(x(t)) is deterministic and L--Lipschitz continuous in x thus admitting a unique solution the initial value problem under the Lindel- theorem (32).
Assumption (Unbiased Observation Operators). For all observation operators Ha. the observation yk is an unbiased estimator of some subset of Ik. ie. x is the restriction of x some subset of features S then Eyyy 2 x
5

---

Analysis
Analysis Step
Analysis Step
Step
M1.
MK-18 1:
(Learndd)
H11
H2
H2
HK
H K
Forward Pass
L
Gradient Flow
xtt
Figure 4: The self- self-suuereised training process is unrolled across multiple filtering steps. Loss is evaluated between "pseudo--osservations and noisy observations obtained from the system.
We also define an additional loss, the supervised forecast loss L the mean squared error between the forecast state and the true state. For the purposes of our analysis, we'll ll assume that S is the full state such that Elys æ&] = zs 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 fe denote family of functions parameterized by 0 for which min 8 L (+)) = O. For any O which achieves this minimum, is also true that L (0) O and that 0 is in set of minimizers of Foo. (C (0))
The proof can be found in the supplementary materials (C.2). This case provides important motivation by showing that a perfect denoiser exists and is contained within our search space then 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, only consider the supervised losses as per the decomposition in Lemma 1, minimizing the expected self-supervised forecast loss 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 foreeaas loss as:
K
c(( T.) K-1
R- Σ k=1
( J0(1)) Xx S
max im(m]
M So(x x
(9)
K
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 number of common test systems designed to simulate real-world phenomena.
3.3 Implementation Details
Avoiding Ensemble Collapse Ensemble basee uncertainty estimates are an elegant solution to the problem of computing the evolution of uncertainty under nonlincar dynamics, but without assumptions on the noise model, evaluating the analysis uncertainty is non non-trivial. A well-known problem in deep learning-bast based uncertainty quantification is feature collapse B33 34] a phenomenon in which the hidden representations resulting from regions of the input space are pushed arbitrarily closely together by the neural network. our ensemble- based uncertainty quantification scheme, this poses an issue ensemble members can quickly converge ensemble collapse.
small number of assimilation steps resulting in
We address this by combining ensemble estimation with MC Dropout 35 which interprets pass through rropaut-regullrrzzed neural network as sample from a variational distribution over network weights. Thus, for standard usage of MC Dropout can estimate uncertainty in the predictive

---

MC Dmppout
No Dropout
10
2.6-
0.4-
0.2
Att1
Otsereec
-10
100
150
200
250
300
350
50
80
300
20
40
60
100
Cycle
Cyde
Cycle
(a) Mean ensemble spread over as- (b) Ensemble trajectories with ob- (c) Ensemble trajectories with obser- similation eycles
servation noise 1.0
vation noise 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.
distribution by using multiple dropout passes to integrate out weight uncertainty. our case, we further have to account for uncertainty over the input distribution which is represented by forecast ensemble. We employ a relatively simple Monte Carlo scheme in which the forecast ensemble members are assimilated by single pass through rtropout-reuularized network with independent dropout masks giving us m samples from the predictive distribution as our analysis ensemble. While this is crude approximation, found 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 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 be used with any observation type, relationships to the system state. These unknown observation operators can be trained by evaluating the loss points with future, known observation operators and passing the gradient signal back in time through the differentiable simulation connecting the time points.
ones with unknown
The second, which 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.
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 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 43)
The work most comparable our own is that which uses machine learning methods learn component of the assimilation mechanism directly in continuous nonlinear systems 44 45 Our method differs from Grönquist 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 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-54).
7

---

Experiments
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 training 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 further 10,000 steps. AmEnF models are developed in Py (55] using the torchdiffeq [56) library for ODE integration. Models are trained on single GTX 1070 GPU for 500 epochs using the Adam [22 optimizer with initial learning rate 8e 4 with a warm-up 50 iterations followed by halving the learning rate every 200 iterations. All experiments are repeated over ten independent noise samples and error bars indicate single standard deviation.
5.1 Qualitative Evaluation
Experiments in this section are performed using the Lorenz 96 system 577 Lorenz 96 is a system of coupled differential equations intended emulate the evolution of single atmospheric quantity across a latitude circle.
-
F
(10)
The system defined to have periodic boundary conditions ID+1 = I1, 20 JD, and x- DD where D is the number of system dimensions. We set the number of dimensions to 40 and the forcing value 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 o=1 and m = 10.
Ensemble Behavior Figure demonstrates the benefit of incorporating MC Dropout. In that the ensemble variance
the left figure, quickly collapses to zero using deterministic net- works while it varies along 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 single spatial coordinate
Training Loss Comparison
Vallda validetion Analysis RMSE
0.9
0.8
Training Loss
0.7
Sule-Super Forecast
0.6
Supervised Analysis
across several noise settings. We can see that the uncertainty patterns in our ensemble behave
Self-supervised Analysis
0.5
largely as one would hope. Less certainty in the system results larger ensemble spread.
0.4
Of particular note the behavior near outlier ob- servations. We can that while the ensemble trajectories occasionally approach such obser- vations, when the ensemble spread is tight the learned assimilator trusts the current forecast trajectory and the ensemble members do move onto worse path to match the observation.
75
200
Training Epoch
Figure 6: Analysis RMSE evaluated on the val- idation set across models trained using different losses for 1.
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 a = 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- sele-vupe forecast training approaches result in learning legitimate denoising function.

---

L96
KS
VL20
Test Analyds RMSE
1.25
1FTXT
IFTX
d 1.0
4D%ar
lEncS
40
17085
ETK5
0.75
2
0.25
Thst Analysis AMSE
AmEEd
a = 2.5
1.25
40
4D
0.25
2
1.0
20
25
30
0.00
15
20
25
30
10
20
30
Ensemble Size
Ensemble Size
Ensemble Size
Figure 7: Results from numerical comparisons across multiple ensemble sizes and observation noises.
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 definitive as the performance gap 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
We evaluate performance on three benchmark systems. Loss is reported as the time-averaged RMSE between the mean analysis state estimate and the state on the state. Each system evaluated across a variety of ensemble sizes (denoted by m) and isotropic Gaussian observation noise levels (whose standard deviation is denoted by r). All results are reported in Figure
We compare performance against a set of widely used filtering methods for data assimilation imple- mented in the Python DAPPER library
(58). These methods include 4D-Var [24), the Local Ensemble Transform Kalman Filter (LETKFF) [59], and the Iterative Ensemble Kalman Smoother (iEnKS) [60] 4D-Var and the iEnKF used in 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.
For method comparisons, we examine the more challenging case a partially observed system. At every assimilation step, observe only rotating subset consisting of every fourth spatial dimension so that only 25% the system is observed at any time. These experiments use an RK-4 61 integrator with step sizes of .05 and partial observations every integration steps.
Lorenz 96
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
Garamoto-Shivashinse
U1 + tyg + Myazz + uu
(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 to represent a minimal model of the earth's atmosphere. It augments the Lorenz system with "temperature" variables allowing for more complicated behavior:
coupled system intended
dX, dt
X)-1(X5+1 10:11xx-22) X,-2) all - -xx,
(12)
d0;
dt - - -00 +G

---

We set = 10, G -,, 1, 1, and use 36 spatial points. Experiments using this system are also partially observed with every fourth spatial coordinate being observed 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 time component and the KS experiments occur at the largest time scales, Nonetheless, despite these potential theoretical snags, the AmFnF 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 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 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 concern as 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 [T] 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 of a self-supervised approach for learning to assimilate in chaotic dynamical systems. This approach uses hybrid deep learning- numerical simulation architecture which outperforms 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
[1] Peter Bauer, Alan Thorpe, and Gilbert Brunet. The quiet revolution of numerical weather prediction. Nature, 525(7567):47-55, 2015.
[2] David A Randall, Richard A Wood, Sandrine Bony, Robert Colman, Thierry Fichefet, John Fyfe, Vladimir Kattsov, Andrew Pitman, Jagadish Shukla, Jayaraman Srinivasan, et 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.
[3] M. Bocquet, H. Elbern, H. Eskes, M. Hirtl, R. 7abkar, 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. Atnospheric Chemistry and Physics, 15(10):535-5335 2015. doi: 10.51994/ep-1155555220555 URL https://acp.. htttpss//acp.coeemii.cc cus. htgg/aaticlee///52222222
[4] Andreas Fichtner. Full seismic waveform modelling and inversion. Springer Science & Business Media, 2010.
10

---

[5] Didier Bresch, Thierry Colin, Emmanuel Grenier, Benjamin Ribba, and Olivier Saut. Com- putational modeling of solid tumor growth: the avascular stage. SIAM Journal on Scientific Computing, 32(4))2321-234 2010.
[6] Cristóbal Bertoglio, Philippe Moireau, and Jean-Frederic Gerbeau. Sequential parameter estimation for fluid-structure problems: application hemodynamics. International Journal for Numerical Methods in Biomedical Engineering, 28(4):434 455, 2012.
[7] Heinz W Engl, Christoph Flamm, Philipp Kügler,Jamas Lu, Stefan Müller, and Peter Schuster. Inverse problems in systems biology. Inverse Problems, 25( 25(12)1130144 2009.
[8] Stephen Wiggins, Stephen Wiggins, and Martin Golubitsky. Introduction to applied nonlinear dynamical systems and chaos, volume Springer, 1990.
[9] Alberto Carrassi, Mare 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 /tt/ttts::////// htttps:/
wwwooonnnnnaarrrlo.ooeeeeeaaaaal wwwey. com/ 1o0/abb//0..00///vcec 1002/wcc
535
[10] 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 172, 2017. doi: 10.1
[11] Eugenia Kalnay. Atmospheric modeling, data assimilation and predictabiliiy. Cambridge university press, 2003.
[12] Diederik P. Kingma and Max Welling, Auto-Encoding Variational Bayes. In 2nd International Conference on Learning Representations, ICLR 2014, Bansf, AB, Canada, April 14-16, 2014, Conference Track Proceedings, 2014.
[13] Tianju Xue, Alex Beatson, Sigrid Adriaenssens, and Ryan Adams. Amortized Finite Element Analysis for Fast Constronsr Optimization. In Hal Daumé and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Pro- ceedings of Machine Learning Research, pages 10638-10647. PMLR, 13-18 Jul 2020. URL http://prooceedings mlr. wwwpppressseee v1 http////11xxxxxaaeeaaa xue20a. htt.m20a.e.mmlllllla html
[14] Sachin Ravi and Alex Beatson. Amortized Bayesian Meta-Learming. In International Con- ference on Learning Representations, 2019. URL BuRLitttps///ppeeeeeeee..tttt wttp://viiiieeeeee//uuuuiaaaii Tkgpy3C5tX
[15] Mark Asch, Marc Bocquet, and Maëlle Nodet. Data assimilation: methods, algorithms, and applications, volume 11. SIAM, 2016.
[16] G. Evensen. Sequential data assimilation with nonlinear quasi-geostrophic model using Monte Carlo methods to forecast error statistics. Journal of Geophysical Research, 1994. ISSN 01480227. doi: 10. 1099944c000/7..
[17] Joshua Batson and Loic Royer. Noise2Self: Blind Denoising by Self- Supervisrri... In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 524-533, PMLR, 09-15 Jun 2019. URL https: http::///ooceeeeee/eeeee mlr press /9eees//////// ///0///batoo//t 19a. html
[18] R. Kalman, A new approach linear filtering and prediction problems. Journal of Fluids Engineering, Transactions of the ASME, 1960. ISSN 1528901X. doi: 10.1115/1.362552..
[19] A H Jazwinski. Stochastic processes and filtering theory. IEEE Transactions on Automatic Control, 1972.
[20] Arnaud Doucet, Nando Freitas, and Neil Gordon. An Introduction to Sequential Monte Carlo Methods. In Sequential Monte Carlo Methods Practice. 2001. doi: 10.1007/ 978-1-4757-3437.1
[21] Chris Snyder, Thomas Bengtsson, Peter Bickel, and Jeff Anderson. Obstacles to high- dimensional particle filtering. Monthly Weather Review, 2008. ISSN 00270644. doi: 10.11752008MWRR2299.11

---

[22] J. L. Anderson and S. L. Anderson. A Monte Carlo implementation of the nonlinear filtering problem produce ensemble assimilations and forecasts. Mon. Wea. Rev. 127:2741-2758, 1999.
[23] Thomas M. Hamill, Jeffrey S. Whitaker, and Chris Snyder. Distance-Dependent Filter- ing of Background Error Covariance Estimates in an Ensemble Kalman Filter. Monthly Weather Review, 129(1 (29)))277 2790, 2001. URL BtLhhtts::// https://journall. journals. ametsoc. wrg.ooogooo....... view/ fournals/nwrre/ 11////////11/
1498_2001__299 2_6_2_276____2011 21w_222766666e6______________e12 ddfddddfoeedddeed __e_obeee2000 O. co_2.0_222.2.... xm1
[24] F Rabier, H Järvinen, E Klinker, J.-F. Mahfouf, and A Simmons. The ECMWF operational implementation of four-dimensional variational assimilation. Experimental results with simplified physics. Quarterly Journal of the Royal Meteorological Society, 126(564): 1443- 1170, 2000. URL htttss///mmet::e onl wollinelibee wwleeellbrrrrrr. wi wwwey. www 49712656415
[25] Tuomas Haarnoja, Anurag Ajay, Sergey Levine, and Pieter Abbeel. Backprop KF: Learning discriminative deterministic state estimators. Advances Neural Information Processing Systems, 2016.
[26] P.J. Werbos. Backpropagation through time: what it does and how to do it. Proceedings of the IEEE, 15510115500-1500 1990. doi: 10.1109/5.58377.
[27] Jaakko Lehtinen, Jacob Munkberg, Jon Hasselgren, Samuli Laine, Tero Karras, Miik Ait- tala, and Timo Aila. Noise2Noise: Learning image restoration without clean data. In 35th International Conference on Machine Learning, ICML 2018, 2018. ISBN 17865108.7963.
[28] Samuli Laine, Tero Karras, Jaakko Lehtinen, and Timo Aila. High-quality self-supervised deep image denoising. In Advances Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL /ttpssttss///proceeddi (tt.://9988000aacccc5ae0eeea555.
www.gss.euripp..cmm hww..ip..c/ppppr/2000.. pdf
[29] Shreyas Fadnavis, Joshua Batson, and Fleftherios 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: h/wp:/eedigss......o.. httpse/////
wwwpeediinnnnnnuns..no.. nwurips..cccoo cc/ Paper .p.ff
[30] Alexander Krull, Tim-Oliver Buchholz, and Florian Jug. Noise2Void Learning Denoising From Single Noisy Images. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition (CVPR), June 2019.
[31] Kiwon Um. Robert Brand, Yun (Raymond) Fei, Philipp Holl, and Nils Thuerey. Solver- in-the-Loop: Learning from Differentiable Physics Interact with Iterative PDE-Solvers. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Ad- vances in Neural Information Processing Systems, volume 33, pages 6111-6122. Curran Associates, Inc., 2020. URL http:::s//rocceedingg neurips cc/ /tt/p////r////1lle// Paper. pdf
[32] Earl A Coddington and Norman Levinson. Theory of ordinary differential equations. Tata McGraw-Hill Education, 1955.
[33] Joost Amersfoort, Lewis Smith, Andrew Jesson, Oscar Key, and Yarin Gal. Improving Deterministic Uncertainty Estimation in Deep Learning for Classification and Regression. CoRR, 11s/2102.11409, 2021. URL httpps://aaateeeeee orgg www.:/ogggg/./eeeeeee wtwp//112.... 11409
[34] Vardan Papyan, X. 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):4652-2446 2020. ISSN 0027-8424. doi: 00.1073/paa.2211550111 URL https //www.pnas. wwt////g///eeeeooo /ttp://on/nn/1/////455/003
[35] Yarin Gal and Zoubin Ghahramani. Dropout as Bayesian approximation: Representing model uncertainty in deep learning. 33rd International Conference on Machine Learning, ICMI.
2016, 2016. ISBN 97815 10829008.
12

---

[36] 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 Pro- ceedings of Machine Learning Research, pages 9690-9700 PMLR, 13-18 Jul 2020. URL https: https://proceedinng mlr wwwpp//es//ep.. wttpsss/vvvvv/vee
html.
[37] 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 NeurlPS Workshops, volume 137 of Proceedings of Machine Learning Research, pages 21-32. PMLR, 12 Dec 2020.
URL URLhttps: ptattps://ppoceediiee wwwroceedding.m.r..... mlr. www..pprsssvvivv 1aosss//13//rosca2062 html
[38] Fabrício Härter and Haroldo Fraga de Campos Velho. New approach applying neural network in nonlinear dynamic model. Applied Mathematical Modelling, 32(12):2621-2633
dec 2008. ISSN 0307-904X. doi: N0.1016JJ.PMM.0007 .9.0000......6......
[39] Rosângela Cintra, Haroldo Campos Velho, and Steven Cocke. Tracking the model: Data assimilation by artificial neural network. pages 403-410, 2016. doi: 10.1 2016 7727227.
[40] 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: case study with the Lorenz 96 model. Geoscientific Model Development Discussions, 2019:1-21, 2019. doi: 10. 5 5194/mmd-2019 136. URL BtRpsttps:/// : http:////ggg.oopocccc...ooo////
gmdgggd--000 - 136/
[41] 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/
[42] Mare Bocquet, Julien Brajard, Alberto Carrassi, and Laurent Bertino. Data assimilation as a learning tool to infer ordinary differential equation representations dynamical models. Nonlinear Processes in Geophysics, 143-112,, 2019.
[43] 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 a case study with ERA5, Geoscientific Model Development Discussions, 2021: 2021. doi: 00.5194/mmd-221-11 httppprrinnt/mmmm2211----11 2021-717
wttp:///mmcccococccccccomm cus
[44] Peter Grönquist, Chengyuan Yao, Tal Ben-Nun, Nikoli Dryden, Peter Dueben, Shigang Li, and Torsten Hoefler. Deep learning for post- 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. 1..18//sta.2222000000
[45] Thomas Frerix, Dmitrii Kochkov, Jamie Smith, Daniel Cremers, Michael Brenner, and Stephan Hoyer. Variational data assimilation with 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 httpt:tst/pppeeeeeeeeeo
mlr wwwpppeessvvvie. hwt.3///freixx11a... htmm
[46] 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. 109110CC.20020..889.
[47] 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.
[48] Rahul G. Krishnan, Uri Shalit, and David Sontag. Deep Kalman Filters, 2015.
13

---

[49] Syama Sundar Rangapuram, Matthias W Seeger, Jan Gasthaus, Lorenzo Stella, Yuyang Wang, and Tim Januschowski. Deep State Space Models for Time Series Forecast- ing. In S Bengio, H Wallach, H Larochelle, K Grauman, N Cesa-Bianchi, and Garnett, editors, Advances in Neural Information Processing Systems 31, pages 7785-7794. Curran Associates, Inc., 2018. URL s004-deep-state- space- spaceeeees-----eees----ee models for-time- seriee-eiie-- foreeasst......inngg pdf
nips cctcccc.pppeee cc/paper/
[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.
[51] Shixiang Gu, Zoubin Ghahramani, and Richard E. Turner. Neural adaptive Sequential Monte Carlo. In Advances in Neural Information Processing Systems, 2015.
[52] Tuan Anh Le, Maximilian Igl, Tom Rainforth, Tom Jin, and Frank Wood. Auto-encoding Sequential Monte Carlo. 6th International Conference on Learning Representations, ICLR 2018 Conference Track Proceedings, 2018.
[53] Hongyuan Mei, Guanghui Qin, and Jason Eisner. Imputing missing events in continuous-time event streams. 36th International Conference on Machine Learning, ICML 2019, 2019. ISBN 97815 7781510886988.
[54] Philipp Becker, Harit Pandya, Gregor Gebhardt, Cheng Zhao, C. James Taylor, and Gerhard Neumann. Recurrent Kalman Networks: Factorized Inference High-Dimensional Deep Feature Spaces. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 544-552. PMLR, 09-15 Jun 2019, URL hwt.h:tp///rrtee https:://proceedings
[55] Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang. Zachary Devito Facebook, 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.
[56] Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary differential equations. Advances in Neural Information Processing Systems, 2018.
[57] Edward N. Lorenz. Predictability--a problem partly solved. In Predictability of Weather and Climate. 2006. ISBN 968051 doi: 10. 10.1017/CB0988851 1617552.044
[58] Patrick N Raanes and Others. nansencenter/DAPPER.. Version 0.8, 2018. URL https://doi
org/10. 500.222/eeeoooo..... 2029296
[59] M. Bocquet. Ensemble Kalman filtering without the intrinsic need for inflation. Nonlinear Processes in Geophysics, 2011. ISSN 10235809. doi: 10.5194/npg- 18-735-201
[60] Pavel Sakov, Dean S. Oliver, and Laurent Bertino. An iterative EnKF for strongly nonlinear systems, Monthly Weather Review, 2012. ISSN 00270644. doi: 10.175//WRR--11100177.,
[61] C. Runge. Ueber die numerische Auflösung von Differentialgleichnngn. Mathematische Annalen, 1895. ISSN 00255831. doi: 10. 10.1007/FF01446007
[62] Yoshiki Kuramoto. Diffusion-Inducel Chaos in Reaction Systems. Progress of Theoretical Physies Supplement, 1978. ISSN 0375-9687. doi: 10. 10.1144./pp...4.346..
[63] 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.1 1
1-
14


Reviews:
[{"review_id": "M9MJ-IH0OL", "reviewer": "Reviewer_kG8H", "summary": "The paper presents, amortized assimilation, way to perform data assimilation through deep learning utilizing self-supervised training to obtain an unbiased estimate of the true state. ", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "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": "The abstract of the paper is overwhelmed with technical jargons. While the need for technical jargons maybe be essential at most points in a paper, I would expect the authors to try and write a more easy-to-understand abstract so that the work may appeal to both the deep learning and the DA community.\n\n\"Line 44\" what is oft-repeated in numerical simulations of Earth system ? I do not think that there are any oft-repeated processes, else the authors should cite something that explains this. \n\nLines 51-60 are extremely hard to read. In fact, throughout the paper, the authors might want to keep in mind that their application is useful for both ML and DA audiences and can rephrase the end of their introduction to read a little more easily to a vaster range of audience, say from geoscience, DA, dynamical systems, and physics along with ML. \n\n\"Lines 102-106\" At small ensemble sizes DA has been successfully used for even the most complex NWPs through localization and inflation.  Generating large ensembles has also been successfully performed with deep learning methods as well. The statements in between these lines are not enough to showcase a need for the authors' method. A clear motivation is either  absent in this introduction or might get hidden owing to the overly technical writing . \n\n\"Lines 112-113\" Sequential filters do not respect conservation laws ? When making claims like these, the authors should cite papers or show evidence. \n\n\"that we view sequential filtering as a particular recurrent prediction model which can be unrolled into a multi-step computation graph allowing us the train using backpropogation through time\" -- What is this computation graph? While the authors present a very well thought out method, with enough novelty, the writing makes it very hard to follow. This is not helping the authors make their case for the importance or relevance of this method\n\nThe results are interesting but does not show the advantage of this method over other traditional DA methods. What happens with small ensemble sizes and localization and inflation? In the absence of fair comparisons , the results are not complete.\n\nReferences: No journal names in the references\n\n[39] Julien Brajard, Alberto Carassi, Marc Bocquet, and Laurent Bertino. Combining data assimilation and machine learning to emulate a dynamical model from sparse and noisy observations: a case study with the lorenz 96 model, 2020. -- incomplete reference, where's the journal name? https://gmd.copernicus.org/preprints/gmd-2019-136/\n\n [42] Ashesh Chattopadhyay, Mustafa Mustafa, Pedram Hassanzadeh, Eviatar Bach, and Karthik\n Kashinath. Towards physically consistent data-driven weather forecasting: Integrating data  assimilation with equivariance-preserving deep spatial transformers, 2021. --- same problem. Here's the paper: https://gmd.copernicus.org/preprints/gmd-2021-71/\n\nWith an improvement in the writing, and more experiments with comparisons with DA+localization and inflation, this will be a very good contribution.\n\n\n\n\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes", "review_text": "", "extra_scores": {"ethical_concerns": "N/A", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "IQaS5ZTqoVu", "reviewer": "Reviewer_NmPU", "summary": "The overarching goal of this paper is to improve assimilation of (chaotic) dynamical systems. The authors introduced a self-supervised framework, amortized assimilation, which uses methods from Ensemble filters and supervised denoising. \n In Section 3, they develop their own method, first focusing on amortized ensemble filters and then deriving their self-supervised assimilation method. In Section 4, their method is compared with standard data assimilation methods on chaotic benchmarks (such as the Lorenz 96). ", "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": "1) I cannot really comment on the Originality of this paper. The methods they use seem common in the dynamical system literature (which they do not review), but maybe there is a crucial difference when moving to assimilation I overlook. Thus I don’t feel confident enough to make a statement.\n\n2) I found the paper to be hard to understand. Some points are definitely lack of knowledge on my part, but other parts could probably be better explained\n\n- The state that the problem they want to solve is a high-quality estimate of the initial state (see abstract). They actually never explain how they estimate the initial state with their model, why their model should be better suited to estimate the initial value nor do they compare a measure of quality for the initial state in their empirical section. \n\n- In the abstract they claim that they use \"independent neural networks to assimilate specific observation types”(l. 8f), but later they claim that they didn’t do that, but instead handled them as follows: \n(Line 226 )“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\" \n\nEither I misunderstand the authors, or their abstract and later explanations do not match …\n\n- They claim:” \n(Line 251)  All methods tested 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”\nI am not sure that is true ... \n\n- Their core section about self-supervised assimilation could be a lot clearer. At one point the authors write \"letting us use the framework from [16] to analyze the problem” without further explaining the framework … \nThis makes it very arduous for the reader to understand what they do \n\n- Figure 6 should have a legend or a more elaborate caption. The plot is not understandable without reading the text and even in the text the colors of the lines plotted is not explained …\n\n3) I should point out some comments that need to be addressed (please see below):\n\n- First, there are some grammatical mistakes as well as missing words and typos …\n\n- Figure labels especially in Figure 7 should be a lot larger. I had to zoom a lot to see the figures\n\n- They should explain what they mean by “independent neural networks to assimilate specific observation types” in the abstract and how that compares to line 226ff (see above)\n\n- If possible they should further explain, why they focus so much on the error in the estimate of the initial state of the system (abstract & introduction) as opposed to other model error and why they do not compare the quality of initial state estimation in their experimental section, if that estimation is so essential \n\n- Lastly, in their subsection Training objective, the authors should consider more training-runs. They find an “intriguing” result but because they only consider one training-run (why???) it’s entirely unclear whether this is to chance or an significant effect. ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5 hours"}, "extra_text": {}}, {"review_id": "6HEj7PUKGYT", "reviewer": "Reviewer_UkND", "summary": "The authors introduce amortized assimilation or AmEnF, a novel method for inferring dynamics from sampled states, and predicting initial and future states from it.", "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": "The key idea of AmEnF is that we can run the model from several different timepoint and compare the different predictions we get from the model to one another, results in a self-supervised loss. They prove that the self supervised loss is bounded by the standard supervised loss. I think this is a really neat idea, and an interesting new direction for self-supervised learning for dynamical system. \n\nThe authors show that their model performs well on standard data assimilation tasks such as Lorenz 96, Vissio-Lucarini 20 and Kuramoto-Shivashinsk, especially when observed data is sparse and noise is high, which are usually the instances where self-supervised are especially helpful. \n\nOverall, the paper does a good job at explaining the background, previous works, and their innovation and contributions. \n\nMinor remarks: All figures are to small and a quite hard to read, especially figure Figure 1.\n\nI recommend acceptance.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "I see no negative societal impact from this work.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2 hours"}, "extra_text": {}}]

Author rebuttal:

Meta-review: From the SAC. This is an instance where the rebuttal and the discussion worked. While the original decision for this paper was to not accept, it is being raised to a recommended accept. The primary reason is the quality of the rebuttal, and the useful technical discussion between authors and reviewers that ensued that seems to have been revealing (in particular, reviewer kG8H). To the authors: I trust that you will take all reviewer feedback into account and most importantly, that all of the things in your rebuttal and discussion that were promised will be done in the next version of the paper.
Decision: Accept (Poster)