FNO-CG Approach in Variational Data Assimilation

• FNO-CG is a meta-learning framework that approximates the inverse Hessian operator in high-dimensional variational data assimilation problems.
• It leverages Fourier Neural Operators to generate effective initializations for conjugate gradient solvers, significantly reducing computation time.
• Numerical results highlight a 62% reduction in error and a 17% decrease in iterations, demonstrating its potential for operational weather and ocean forecasting.

The FNO-CG approach integrates Fourier Neural Operators (FNOs) with the conjugate gradient (CG) method in variational data assimilation (DA), focusing specifically on efficiently approximating the inverse Hessian operator across a family of high-dimensional DA problems. By leveraging a meta-learning framework, FNOs are trained as surrogates for the inverse Hessian, providing effective initializations for CG solvers and thereby accelerating convergence, particularly in ill-conditioned scenarios. This method directly targets the computational bottlenecks of operational DA systems such as those in numerical weather prediction and ocean forecasting (Moazzami et al., 26 Sep 2025).

1. Motivation and Problem Setting

Variational DA formulates the optimization of initial states in PDE-driven systems by minimizing a cost functional that encodes discrepancies between forecast states and sparse, irregular observational data. The Hessian of the cost function, ∇²J, plays a central role in ensuring rapid and stable convergence during minimization. However, explicit inversion or iterative application of the Hessian (or its preconditioned version) is computationally demanding and becomes prohibitive for high-dimensional, ill-conditioned systems. Conventional approaches rely on CG to approximately solve equations of the form

$\nabla^2 J \cdot u_0^{\text{opt}} = f$

where $f$ is a gradient-based right-hand side involving background information and observational data.

The FNO-CG approach addresses these limitations by meta-learning a mapping from $f$ to $u_0^{\text{opt}} \approx [\nabla^2 J]^{-1}(f)$, using a neural operator architecture capable of representing parametric families of operators with high efficiency.

2. FNO-CG Architecture and Methodology

The core technical innovation is the use of a Fourier Neural Operator as a data-driven surrogate for the inverse Hessian operator. The FNO architecture, composed of multiple spectral convolution layers, is optimized to capture the action of $[\nabla^2 J]^{-1}$ on a family of right-hand sides $f$ derived from diverse initial DA problem configurations. The workflow is as follows:

1. For each DA problem instance (e.g., particular realizations of input parameters, observation locations, noise, etc.), generate $f$ using the adjoint model.
2. Solve for the reference $u_0^{\text{opt}}$ (the ground-truth) by direct or iterative inversion.
3. Train the FNO to minimize

$$\mathcal{L}_{\text{FNO}} = \mathbb{E}_{(f, u_0^{\text{opt}})} \| \text{FNO}(f) - u_0^{\text{opt}} \|_2^2$$

so that FNO(f) approximates the Hessian-inverse solution operator across the distribution of $f$.

1. At inference, obtain an approximate $u_0^{(\text{init})} = \text{FNO}(f)$ for a new DA problem and use it as an initializer for classical CG:
   • For $k=0$: $u_0^{(0)} = \text{FNO}(f)$
   • For $k > 0$: $u_0^{(k+1)} = \text{CG-step}(u_0^{(k)}, f)$

This meta-learned initialization reduces the number of CG iterations required for convergence relative to a standard (zero or background state) initial guess.

3. Numerical Results and Performance

Numerical evaluation focused on a linear advection equation with periodic boundary conditions. Key characteristics include:

• True initial states $u_0^T$ generated as perturbed sines with random parameters and spatial noise.
• $f$ constructed using adjoint methods as in standard 4D-Var.
• Family of DA problem configurations for meta-learning covers a broad parameter regime, including cases with variable background and observation noise.

Quantitative improvements reported:

• FNO-CG achieved a 62% reduction in average relative error compared to standard CG initialized at zero.
• The average number of CG iterations was reduced by 17%.
• Improvements were most pronounced for ill-conditioned Hessians, as shown by performance metrics plotted against the Hessian condition number.

The results indicate that while a standalone FNO is capable of producing a reasonable approximation of $u_0^{\text{opt}}$, combining the FNO-provided initializer with classical CG (FNO-CG) yields both lower error and faster optimization trajectories.

4. Spectral Operator Properties and Implementation

The FNO leverages spectral convolution layers, which efficiently encode global spatial dependencies by acting in the frequency domain: $$v_{t+1}(x) = \sigma \left( W v_t(x) + \mathcal{F}^{-1}\left[ \mathcal{R}(\xi)\ \mathcal{F}(v_t(x)) \right] \right)$$ where $\mathcal{R}(\xi)$ is a set of learnable weights applied to low-frequency Fourier modes, $\mathcal{F}$ denotes the Fourier transform, and $W$ comprises learned local linear transformations. This construction naturally supports generalization across spatially varying $f$ and captures the nonlocal structure typical for inverse Hessian operators in PDE-constrained DA.

In the meta-learning context, the FNO is trained on a population of DA problems to internalize the mapping between $f$ (derived from varying initializations, noise realizations, or observation patterns) and the corresponding optimal $u_0$. The resulting model serves as a universal preconditioner or direct surrogate for the inverse Hessian over the support of the training distribution.

5. Broader Implications and Operational Significance

The FNO-CG approach provides a practical solution to two longstanding challenges in variational DA:

• The high computational cost of traditional Hessian-based solvers for high-dimensional inverse problems.
• The slow convergence and inefficiency of CG in ill-conditioned regimes.

By offloading much of the inversion workload to the non-iterative FNO surrogate, FNO-CG enables efficient, robust DA in operational settings such as numerical weather prediction, ocean state estimation, and related PDE-driven forecasting domains where both speed and accuracy are critical.

A key advantage is the method’s adaptability within a meta-learning framework to families of DA configurations, supporting reusability and extensibility. The architecture preserves the robustness of classical iterative solvers, as the FNO only provides an enhanced initialization and subsequent CG steps guarantee convergence to the desired optimality criteria.

6. Limitations and Future Directions

While the current demonstrations are in a linear PDE/DA context, extension to nonlinear and incremental variational schemes is anticipated. In particular, the approach may be generalized to incremental 4D-Var for chaotic systems (e.g., Kuramoto–Sivashinsky), where nonlinear Hessian-action surrogates become more challenging but also more beneficial due to even worse conditioning.

Extending the FNO meta-learning architecture to parameter regimes with limited training data, incorporating model error, hybrid observation operators, and integrating physics-aware loss regularization are all promising avenues to broaden the applicability of FNO-CG.

---

In summary, FNO-CG represents an overview of data-driven meta-learning of inverse operators with robust iterative solvers, targeted at accelerating and stabilizing high-dimensional variational DA (Moazzami et al., 26 Sep 2025). This framework offers substantial performance gains, particularly in ill-conditioned and operationally demanding scenarios, by combining the generalization strength of neural operators with the reliability of Krylov-subspace optimization.