Online learning in idealized ocean gyres

Ocean turbulence parameterization has principally been based on processed-based approaches, seeking to embed physical principles so that coarser resolution calculations can capture the net influence of smaller scale unresolved processes. More recently there has been an increasing focus on the application of data-driven approaches to this problem. Here we consider the application of online learning to data-driven eddy parameterization, constructing an end-to-end differentiable dynamical solver forced by a neural network, and training the neural network based on the dynamics of the combined hybrid system. This approach is applied to the classic barotropic Stommel-Munk gyre problem – a highly idealized configuration which nevertheless includes multiple flow regimes, boundary dynamics, and a separating jet, and therefore presents a challenging test case for the online learning approach. It is found that a neural network which is suitably trained can lead to a coarse resolution neural network parameterized model which is stable, and has both a reasonable mean state and intrinsic variability. This suggests that online learning is a powerful tool for studying the problem of ocean turbulence parameterization.

💡 Research Summary

This paper introduces an online‑learning framework for ocean‑turbulence parameterisation that tightly couples a differentiable numerical solver with a neural‑network (NN) forcing term. The authors focus on the classic barotropic Stommel‑Munk double‑gyre configuration, a wind‑driven basin that, despite its idealised nature, exhibits multiple flow regimes, boundary dynamics, anisotropic currents, and a separating jet—features that make it a stringent test case for any sub‑grid‑scale parameterisation.

The governing equations are expressed in stream‑function/vorticity form on a β‑plane with linear drag, Laplacian viscosity, and wind‑stress curl forcing. Spatial discretisation uses the Arakawa Jacobian for advection, second‑order centred differences for diffusion and drag, and Fast Fourier Transforms for the Poisson solve. Time integration employs a Crank‑Nicolson/Adams‑Bashforth 2 (CNAB2) scheme: linear terms are treated implicitly, while the nonlinear advection and NN forcing are explicit.

A key methodological advance is the insertion of a NN term Fθ(·) into the vorticity tendency equation. The NN receives a non‑dimensionalised vorticity field as input (scaled by 1/|β|L) and outputs a non‑dimensionalised wind‑stress‑like forcing (scaled back by |τ0|πρ0DL). Because both the solver and the NN are built with JAX and Keras (JAX backend), the entire hybrid system is end‑to‑end differentiable and runs efficiently on GPUs. This enables the definition of a loss function as the L2 discrepancy between the coarse‑resolution model trajectory (with NN forcing) and a high‑resolution reference trajectory, and the computation of exact gradients through the solver via automatic differentiation.

Reference data are generated by a high‑resolution (2049 × 2049, ≈2 km spacing) simulation with ν = 10 m² s⁻¹, r = 10⁻⁷ s⁻¹, and a cosine wind stress. After a 12‑year spin‑up, the model is integrated for another 12 years; every hour the vorticity field is Gaussian‑filtered (σ = 62.5 km) and down‑sampled to a 65 × 65 grid (≈62.5 km spacing). This filtered, coarse‑grained dataset serves as the training target.

Training proceeds in single precision on a single GPU; the loss is accumulated over a sequence of time steps, and back‑propagation updates the NN weights. The resulting NN‑augmented coarse model remains numerically stable, reproduces the mean stream‑function and jet structure of the high‑resolution run, and retains a realistic level of intrinsic variability. Notably, isolated eddies emerge spontaneously in the coarse simulation, suggesting that the NN has learned an implicit backscatter mechanism that injects energy at large scales while dissipating at small scales.

Generalisation tests, where wind amplitude or viscosity are altered from the training values, show that the mean flow is still captured reasonably, but variability and eddy statistics degrade, indicating that the current NN architecture is tuned to a specific parameter regime. The authors acknowledge that deeper or physics‑constrained networks may improve robustness.

Overall, the study makes three principal contributions: (1) a fully differentiable ocean model implementation, (2) an online‑learning scheme that directly penalises dynamical instability, and (3) a demonstration that a NN‑based eddy parameterisation can restore both mean state and variability at coarse resolution. The work highlights the promise of online learning for ocean modelling, especially for multi‑scale or hybrid high‑/low‑resolution frameworks, while also pointing to future challenges in generalisability, computational cost, and incorporation of physical constraints.