Learning to Assimilate in Chaotic Dynamical Systems

1 Introduction . Forecasting in the geosciences is an initial value problem of enormous practical significance . In high value domains like numerical weather prediction [ 1 ] , climate modeling [ 2 ] , atmospheric chemistry [ 3 ] , seismology [ 4 ] , and others [ 5–7 ] , forecasts are produced by estimating the current state of the dynamical system of interest and integrating that state forward in time using numerical models based on the discretization of differential equations . These numerical models are derived from physical principles and possess desirable extrapolation and convergence properties . Yet despite the efficacy of these models and the vast amount of compute power used in generating these forecasts , obtaining highly accurate predictions is non-trivial . Discretization introduces numerical errors and it is often too computationally expensive to directly simulate the system of interest at the resolution necessary to capture all relevant features . Further complicating matters in geoscience applications is that many systems of interest are chaotic , meaning that small errors in the initial condition estimates can lead to significant forecasting errors over relatively small time frame [ 8 ] . This can be problematic when the initial condition for a forecast is current state of the Earth , a large-scale actively evolving system that can not be directly controlled . Modern numerical weather prediction ( NWP ) models utilize hundreds of millions of state variables scattered over a three-dimensional discretization of the Earth ’ s atmosphere [ 9 ] . To perfectly initialize a simulation , one would need to know the true value of all of these state variables simultaneously . This data is simply not available . In reality , what is available are noisy , partial measurements scattered non-uniformly in time and space . The inverse problem of estimating the true state of a dynamical system from imperfect observations is the target of our work and of the broader field known as data assimilation . Data assimilation techniques produce state estimates by systematically combining noisy observations with the numerical model [ 10 , 9 ] . The data assimilation problem differs from other state estimation problems in that the numerical model acts both as an evolution equation and as a mechanism for transporting information from regions where observations are dense to regions where observations are sparse [ 11 ] . Thus for accurate estimates of the full system state , it is important to utilize both the model and observations . 35th Conference on Neural Information Processing Systems ( NeurIPS 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 [ 12 ] but have also been used in the context of numerical simulation [ 13 ] , and meta-learning [ 14 ] . Our work extends this approach to variational data assimilation [ 15 ] . In this work , we introduce a self-supervised framework for learning to assimilate which we call amortized assimilation . Here , we use the term self-supervised to indicate that our method learns to assimilate entirely from trajectories of noisy observations without the use of ground truth data during training . Our design incorporates the objective flexibility of variational assimilation but amortizes the expense of solving the nonlinear optimization problem inherent to variational methods into the training of a neural network that then behaves as a sequential filter . Our contributions are both theoretical and empirical . In Section 3.1 , we introduce an amortized assimilation architecture based on the Ensemble Kalman Filter [ 16 ] . We then develop the theory of amortized assimilation in Section 3.2 by extending the powerful self-supervised denoising results of Batson and Royer [ 17 ] to the dynamical systems setting through the use of differentiable simulation . We then support this theory through a set of numerical experiments1 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 dynamical system with state variable x ( t ) ∈ C ⊂ Rd where C is a compact subset of Rd and system dynamics defined by the differential equation : dx dt = g ( x ( t ) ) ( 1 ) with g as a time-invariant , Lipschitz continuous map from Rd → Rd . The state is observed at a sequence ofK discrete time points { τ0 , τ1 , . . . , τK } generating the time-series y0 : K = { y0 , y1 , . . . , yK } where 0 ≤ k ≤ K is an index corresponding to the time point τk . These observations are imperfect representations of the system states x0 : K = { x0 , x1 , . . . xK } generated from observation operators of the form : yk = Hk ( xk ) + ηk ( 2 ) where Hk is an arbitrary potentially nonlinear function and ηk ∼ N ( 0 , σ2kI ) . This noise model is often referred to as a Hidden Markov Model ( Figure 1 ) . In the classical assimilation setting , the observation operators are assumed to be known a priori though our proposed method is capable of learning a subset of observation operators . For a problem like numerical weather prediction , these operators may represent measurements taken of certain atmospheric quantities at a particular measurement location . Data assimilation is then the inverse problem of estimating the true trajectory x0 : K from our noisy observations y0 : K and our model of the system evolution : x ( τk+1 ) =Mk : k+1 ( x ( τk ) ) = x ( τk ) + ∫ τk+1 τk g ( x ( t ) ) dt . ( 3 ) Sequential Filtering In the absence of ground truth , it is practical to present data assimilation from a statistical perspective . In particular , the quantity of interest that we are concerned with in this paper is the filtering distribution p ( xk | y0 : k ) . Efficient algorithms for estimating the filtering distribution 1Code available at https : //github.com/mikemccabe210/amortizedassimilation . Analysis are derived from exploiting dependence relationships in the generative model to produce the two step cycle : Forecast : p ( xk | y0 : k−1 ) = ∫ Rd p ( xk | xk−1 ) dxk−1 ( 4a ) Analysis : p ( xk | y0 : k ) ∝ p ( yk | xk ) p ( xk | y0 : k−1 ) ( 4b ) When evaluating data assimilation methods , the objective function is often evaluated with respect to the “ analysis ” estimates as these act as these are the initial conditions used for the next forecast phase . In subsequent sections , we will use xfk and x a k to refer to forecast and analysis estimates of xk respectively . In the case of linear dynamics with Gaussian uncertainties , the Kalman filter [ 18 ] provides both the forecast and analysis distributions exactly in closed form . However , the general case is significantly more complicated . Under nonlinear dynamics , computing the forecast distribution requires the solution of the Fokker-Planck partial differential equation [ 19 ] . While finding an exact solution to these equations is computationally infeasible , one can produce what is called an ensemble estimate of the forecast distribution using Sequential Monte Carlo methods [ 20 ] . Ensemble filters ( Figure 2 ) replace an exact representation of the uncertainty over states with an empirical approximation in the form of a small number of samples called particles or ensemble members . The forecast distribution can then be estimated by integrating the dynamics forward in time for each ensemble member independently . Allowing for full generality leads to particle filter approaches while enforcing Gaussian assumptions leads to the Ensemble Kalman filter ( EnKF ) [ 16 ] . The former is rarely used in data assimilation settings as particle filters can become degenerate in high dimensions [ 21 ] while the latter has been quite successful in practice . Like the classical Kalman filter , the EnKF assumes that both the state distribution and observation likelihoods are Gaussian . The latter is often reasonable , but the former is a rather strong assumption under nonlinear dynamics . Ensemble methods possess an inherent trade-off in that larger ensembles more accurately model uncertainty , but each ensemble member requires independently integrating the dynamics forward in time which can be significantly more expensive than the assimilation step itself . Bayesian filters rely on robust covariance estimates to maintain stability which can be difficult to obtain when using a small ensemble . Heuristics such as covariance inflation [ 22 ] and localization [ 23 ] have been developed to improve stability at smaller ensemble sizes , but there remains value in developing methods that improve upon the accuracy of these approaches . Variational Assimilation The alternative to sequential filters is variational assimilation . Variational assimilation methods , like 4D-Var [ 24 ] , directly solve for the initial conditions of a system by finding the state x0 which minimizes the negative log posterior density through nonlinear optimization . This is expensive , but gradient-based optimization utilizing the adjoint of the numerical model avoids the persistent Gaussian state assumptions necessary for tractable Bayesian filters . Furthermore , these methods can be augmented by auxiliary objectives promoting concordance with physical properties like conservation laws . 3 Amortized Assimilation . 3.1 Amortized Ensemble Filters . Our proposed approach attempts to combine the flexibility of variational methods with the efficiency of sequential filters by training a neural network to act as an analysis update mechanism . Before diving into the theory , we present the general amortized assimilation framework from the perspective of a specific architecture which we call the Amortized Ensemble Filter ( AmEnF ) . The AmEnF ( Figure 3 ) replaces the EnKF analysis equations with a parameterized function in the form of a neural network whose weights are optimized for a specific dynamical system during training . By unrolling a fixed number of assimilation cycles as first proposed in Haarnoja et al . [ 25 ] , the AmEnF can be trained efficiently using backpropagation through time [ 26 ] . The loss used for this training will be discussed in more detail in Section 3.2 . One advantage of decoupling from the statistical model in this way is that it allows us to incorporate more information into our analysis process . Like the EnKF , the analysis ensemble members are computed using the forecast values xfi , k for each ensemble member i ∈ [ m ] with m denoting the number of ensemble members , the ensemble covariance P fk , and the observations yk . However , we augment these features with the dynamics g ( xfi , k ) and coordinate-wise recurrent memory ci , k−1 resulting in the following update equations : xfi , k =Mk−1 : k ( x a i , k−1 ) P f k = Cov ( { x f i , k } m k=1 ) zxi , zci = fL ( x f i , k , P f k , yk , g ( x f i , k ) , ci , k−1 ) λxi , λci = fN ( x f i , k , P f k , yk , g ( x f i , k ) , ci , k−1 ) xai , k = λxi x f i , k + zxi ci , k = λci ci , k−1 + zci ( 5 ) where fL and fN are neural networks with linear and sigmoid final activations respectively . We implement these as a single network which is split only at the final layer . We observed notable gains to training stability through the inclusion of this sigmoid-gated final layer . For scalability , we do not use the full covariance matrix as an input to the model . Rather , we include local covariance entries as additional feature channels where each channel represents the covariance with the state value at a fixed relative spatial position . For instance in a 1D system , the relative covariance channels may contain the variance at the given point and the covariances with the points to the left and right . The inclusion of recurrent memory may seem out of place when the dynamics are assumed to be Markovian . However , while the dynamics themselves are Markovian , the data assimilation process is not . The inclusion of memory allows the model to learn from the past without explicitly including previous states and observations . This is a significant benefit in practice . The other non-standard inclusion , the local dynamics , led to a small performance boost as well . As our method is not reliant on generative models of the inputs , additional features could be added trivially .

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.   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).