Joint learning of variational data assimilation models and solvers

This paper addresses variational data assimilation from a learning point of view. Data assimilation aims to reconstruct the time evolution of some state given a series of observations, possibly noisy and irregularly-sampled. Using automatic differentiation tools embedded in deep learning frameworks, we introduce end-to-end neural network (NN) architectures for variational data assimilation. It comprises two key components: a variational model and a gradient-based solver both implemented as neural networks. The latter exploits ideas similar to meta-learning and optimizer learning. A key feature of the proposed end-to-end framework is that we may train the NN models using both supervised and unsupervised strategies. Especially, we may evaluate whether the minimization of the classic definition of variational formulations from ODE-based or PDE-based representations of geophysical dynamics leads to the best reconstruction performance.We report numerical experiments on Lorenz-63 and Lorenz-96 systems for a weak constraint 4D-Var setting with noisy and irregularly-sampled/partial observations. The key features of the proposed neural network framework is two-fold: (i) the learning of fast iterative solvers, which can reach the same minimization performance as a fixed-step gradient descent with only a few tens of iterations, (ii) the significant gain in the reconstruction performance (a relative gain greater than 50%) when considering a supervised solver, i.e. a solver trained to optimize the reconstruction error rather than to minimize the considered variational cost. In this supervised setting, we also show that the joint learning of the variational prior and of the solver significantly outperform NN representations. Intriguingly, the trained representations leading to the best reconstruction performance may lead to significantly worse short-term forecast performance. We believe these results may open new research avenues for the specification of assimilation models and solvers in geoscience, including the design of observation settings to implement learning strategies.