State-dependent preconditioning for the inner-loop in Variational Data Assimilation using Machine Learning

TL;DR

This paper targets the slow convergence of the inner loop in Variational Data Assimilation by learning a state-dependent preconditioner. It constructs a spectral preconditioner through a Deep Neural Network that predicts low-rank eigenpairs of the Gauss–Newton matrix A_x, enabling a preconditioned CG solve with improved conditioning. The loss is defined via a Frobenius-norm objective approximated by Monte Carlo estimates, trained on online-generated linearization states to produce U_θ(x) and Λ_θ(x). Applied to a shallow-water DA setup, the method achieves substantial reductions in CG iterations when a balanced number of eigenpairs are used, while inappropriate rank or eigenpair accuracy can harm performance. The approach is non-intrusive and can complement traditional preconditioning, offering a flexible route to faster, data-driven inner-loop solves in high-dimensional DA problems.

Abstract

Data Assimilation is the process in which we improve the representation of the state of a physical system by combining information coming from a numerical model, real-world observations, and some prior modelling. It is widely used to model and to improve forecast systems in Earth science fields such as meteorology, oceanography and environmental sciences. One key aspect of Data assimilation is the analysis step, where the output of the numerical model is adjusted in order to account for the observational data. In Variational Data Assimilation and under Gaussian assumptions, the analysis step comes down to solving a high-dimensional non-linear least-square problem. In practice, this minimization involves successive inversions of large, and possibly ill-conditioned matrices constructed using linearizations of the forward model. In order to improve the convergence rate of these methods, and thus reduce the computational burden, preconditioning techniques are often used to get better-conditioned matrices, but require either the sparsity pattern of the matrix to inverse, or some spectral information. We propose to use Deep Neural Networks in order to construct a preconditioner. This surrogate is trained using some properties of the singular value decomposition, and is based on a dataset which can be constructed online to reduce the storage requirements.

Figures (7)

• Figure 1: Illustration of minimization using successive quadratic approximations
• Figure 1: Schematic representation of the input/output signature of the Neural Network
• Figure 1: Number of iterations needed to reach the norm threshold. Dotted line indicates the number of iterations for the unprecondition problem.
• Figure 2: Inner and outer loop paradigm for optimization
• Figure 2: Norm of the residuals as a function of the iteration number during CG, depending on the number of retained vectors $r$ and shift parameter $\mu$ as defined in Eq. \ref{['eq:def_split_prec_ml']}