# Adjoint-Based Online Learning of Two-Layer Quasi-Geostrophic Baroclinic Turbulence

F. Yan, H. Frezat, Julien Le Sommer, J. Mak, K. Otness

<figure><img src="image_1.png" /></figure>

## To cite this version:

F. Yan, H. Frezat, Julien Le Sommer, J. Mak, K. Otness. Adjoint-Based Online Learning of Two-Layer Quasi-Geostrophic Baroclinic Turbulence. Journal of Advances in Modeling Earth Systems, 2025, 17 (7), 10.1029/2024MS004857. hal-05412230

HAL Id: hal-05412230

https://hal.science/hal-05412230v1

Submitted on 14 Dec 2025

HAL is a multi-disciplinary open access archive for the deposit and dissemination of scientific research documents, whether they are published or not. The documents may come from teaching and research institutions in France or abroad, or from public or private research centers.

L'archive ouverte pluridisciplinaire HAL, est destinée au dépôt et à la diffusion de documents scientifiques de niveau recherche, publiés ou non, émanant des établissements d'enseignement et de recherche français ou étrangers, des laboratoires publics ou privés.

<figure><img src="image_2.png" /></figure>

Distributed under a Creative Commons Attribution 4.0 International License
RESEARCH ARTICLE

10.1029/2024MS004857

**Key Points:**

* Explores an adjoint-based online learning approach for baroclinic turbulence, motivated by ocean eddy parameterization applications

* Considers both an online and an approximate online approach, the latter not requiring a differentiable model

* The full online learning approach, without additional stabilizers, leads to more skillful model integrations

**Supporting Information:**

Supporting Information may be found in the online version of this article.

**Correspondence to:**

F. E. Yan and J. Mak,
[feyan@connect.ust.hk](mailto:feyan@connect.ust.hk);
julian.c.l.mak@googlemail.com

**Citation:**

Yan, F. E., Frezat, H., Le Sommer, J., Mak, J., & Otness, K. (2025). Adjoint-based online learning of two-layer quasi-geostrophic baroclinic turbulence. *Journal of Advances in Modeling Earth Systems, 17*, e2024MS004857. https://doi.org/10.1029/2024MS004857

Received 21 NOV 2024
Accepted 20 JUN 2025

**Author Contributions:**

**Conceptualization:** H. Frezat, J. Le Sommer

**Data curation:** F. E. Yan, J. Mak

**Formal analysis:** F. E. Yan, H. Frezat, J. Le Sommer, J. Mak

**Funding acquisition:** J. Le Sommer, J. Mak

**Investigation:** F. E. Yan, H. Frezat, J. Le Sommer, J. Mak, K. Otness

**Methodology:** F. E. Yan, H. Frezat, J. Le Sommer, J. Mak

**Project administration:** H. Frezat, J. Le Sommer, J. Mak

**Resources:** J. Le Sommer, J. Mak

© 2025 The Author(s). Journal of Advances in Modeling Earth Systems published by Wiley Periodicals LLC on behalf of American Geophysical Union. This is an open access article under the terms of the Creative Commons Attribution License, which permits use, distribution and reproduction in any medium, provided the original work is properly cited.

<figure><img src="image_3.png" /></figure>

# Adjoint-Based Online Learning of Two-Layer Quasi-Geostrophic Baroclinic Turbulence

F. E. Yan¹ ID, H. Frezat² ID, J. Le Sommer³ ID, J. Mak¹,⁴ ID, and K. Otness⁵ ID

¹Department of Ocean Science, Hong Kong University of Science and Technology, Hong Kong, Hong Kong, ²Institut de Physique du Globe de Paris, Paris, France, ³Université Grenoble-Alpes, CNRS, Grenoble, France, ⁴National Oceanography Centre, Southampton, UK, ⁵Courant Institute of Mathematical Sciences, New York University, New York, NY, USA

**Abstract** For reasons of computational constraint, most global ocean circulation models used for Earth System Modeling still rely on parameterizations of sub-grid processes, and limitations in these parameterizations affect the modeled ocean circulation and impact on predictive skill. An increasingly popular approach is to leverage machine learning approaches for parameterizations, regressing for a map between the resolved state and missing feedbacks in a fluid system as a supervised learning task. However, the learning is often performed in an “offline” fashion, without involving the underlying fluid dynamical model during the training stage. Here, we explore “online” approaches that involve the fluid dynamical model during the training stage for the learning of baroclinic turbulence, with reference to ocean eddy parameterization. Two online approaches are considered: A full adjoint-based online approach, related to traditional adjoint optimization approaches that require a “differentiable” dynamical model, and an approximately online approach that approximates the adjoint calculation and does not require a differentiable dynamical model. The full online approach, without the need of additional constraints, is found to result in models that are generally more skillful. Other details relating to online training, such as window size, machine learning model set up and designs of the loss functions are detailed to aid in further explorations of the online training methodology for Earth System Modeling.

**Plain Language Summary** Most global ocean circulation models used in Earth System Modeling are still deficient in representing small-scale processes, affecting the modeled ocean circulation and its trustworthiness for making predictions. An increasingly popular approach is to use machine learning to represent the small-scale feedbacks. However, the learning is often performed in an “offline” fashion, without explicitly involving the underlying fluid dynamical model that may provide additional constraints and information to aid in model stability and skill. Here, we explore two “online” approaches, both of which involves the fluid dynamical model during the training stage for the problem of ocean turbulence. The full online approach without additional modifications is found to be generally more skillful and numerically stable. An approximately online approach, requiring substantially less initial investment in the model development, is found to inherit some but not all desirable properties of the full online approach.

## 1. Introduction

Earth System Models and their component models are invaluable tools that enable us to perform numerical experiments, to probe how the Earth system might respond under various scenarios, to aid in our fundamental understanding of the Earth system itself, and to provide us with information that guide various mitigation strategies as the system evolves (e.g., IPCC, 2021). Yet the various components of Earth System Models possess deficiencies, impacting the fidelity of the models in known and unknown ways. The present work focuses on the representation of physical processes in fluid systems, specifically on baroclinic turbulence. While the principal motivation for the present work is on ocean modeling, baroclinic turbulence is a generic feature in rotating stratified fluid systems, so the work here will have some physical relevance for atmospheric modeling. The numerical methodology should however be of general interest to any component of the Earth System model that possesses a mathematical description of the associated processes, particularly if they are described by partial differential equations.

Fluid systems such as the ocean or the atmosphere are multi-scale in nature, and the explicit representation of their mechanical or thermodynamical processes depend on the choice of computational grid resolution. Ultimately
**Software:** F. E. Yan, H. Frezat, J. Le Sommer, K. Otness

**Supervision:** H. Frezat, J. Le Sommer, J. Mak

**Validation:** F. E. Yan, J. Mak

**Visualization:** F. E. Yan, J. Mak

**Writing – original draft:** F. E. Yan, J. Mak

**Writing – review & editing:** F. E. Yan, H. Frezat, J. Le Sommer, J. Mak, K. Otness

there is a constraint on computational resources, so generally there will be sub-grid processes that lack an explicit representation in numerical models. The representation of these sub-grid processes and their feedbacks onto the resolved state, that is, *parameterizations*, critically impact the performance of the models. There have been increasing use of machine learning techniques to supplement and/or replace existing parameterizations, in the atmospheric context (e.g., Brenowitz & Bretherton, 2019; Connolly et al., 2023; Kashinath et al., 2021; Kochkov et al., 2024; Lopez-Gomez et al., 2022; Mooers et al., 2021; Sun et al., 2023; Yuval & O’Gorman, 2020) as well as the oceanic context (e.g., Bolton & Zanna, 2019; Frezat et al., 2022; Guillaumin & Zanna, 2021; Ouala et al., 2024; Perezhogin et al., 2023; Ross et al., 2023; Sane et al., 2023; Srinivasan et al., 2024; Yan et al., 2023; Zanna & Bolton, 2021). A standard approach in most of these works is to take a high resolution model and/or observational data (e.g., satellite data) as the “truth,” which is then filtered or coarse-grained in some way. The goal is then to obtain a mapping from a machine learning algorithm, where the input would normally be the filtered data, and the output or predicted quantity would be some residual term that arises from the choice of filtering (e.g., the $S^q$ quantity in the present work for a spatial average, or the eddy fluxes for a Reynolds average). The skill of the resulting mapping might then be judged on its ability to predict the withheld data that was not exposed to the model during the training stage (e.g., Bolton & Zanna, 2019; Guillaumin & Zanna, 2021; Yan et al., 2023), referred to as a priori *testing*. The aim is then to implement that mapping into a coarse resolution model with reduced computational demands, and the coarse resolution model is evolved in time with a state-dependent forcing as predicted by the machine learning mapping. There is of course no guarantee the resulting dynamical model has any skill or is even stable in this arguably more relevant a posteriori *test* (e.g., Chattopadhyay & Hassanzadeh, 2023; Frezat et al., 2022; Guillaumin & Zanna, 2021; Ross et al., 2023; Srinivasan et al., 2024; Zanna & Bolton, 2021). Some ways to stabilize the model integrations are often required or desirable, such as improving on the numerics associated with the inherent issues present machine learning algorithms (e.g., Chattopadhyay & Hassanzadeh, 2023), or forcing conservation properties (e.g., Bolton & Zanna, 2019; Ross et al., 2023).

We note that the aforementioned training procedure does not explicitly involve the dynamical model at the training stage, and might be termed an “offline” learning approach: the training is regarded as a regression between inputs and outputs that have no temporal correlation/causality explicitly invoked. Involving the dynamical model during the training stage over some time window in an “online” learning approach would be expected to lead to multiple benefits in the resulting model stability and/or skill, by forcing the training procedure to recognize that the provided data has a causality/correlation as dictated by the dynamical model, sample some of the statistics of the solution as encoded by the integration of the dynamical model, exclude the set of possible mappings that lead to unstable model integrations, and possibly other desirable features. There are many possible methodologies available, with the crucial limitation being computational cost, usually related to the number of model parameters (e.g., number of degrees of freedom in a Convolutional Neural Network) and number of model evaluations involved (e.g., cost of the dynamical model run). Derivative-free methodologies tend to require more model evaluations and has a larger complexity scaling with the number of model parameters: A notable example is the Ensemble Kalman Inversion approach (e.g., Iglesias et al., 2013; Kovachki & Stuart, 2019; Lopez-Gomez et al., 2022; Schneider et al., 2022), although the use of emulations can substantially reduce the cost per model evaluation, and the methodology has some other desirable properties (e.g., Iglesias et al., 2013). On the other hand, derivative-based methods such as adjoint optimization has a lower complexity scaling with number of model parameters (linear complexity scaling for adjoint-based methodologies; e.g., Gunzburger, 2003), making them particularly attractive for problems with a large number of parameters (e.g., for Convolutional Neural Networks being considered in this work). The adjoint-based online approach is the one we consider in this work. The more classical adjoint approach has had important applications in data assimilation (e.g., Kalnay, 2002) and in oceanography (e.g., Forget et al., 2015; Heimbach et al., 2019). The adjoint-based online learning approach has been shown to lead to significant improvements in fluid simulations (called “temporal unrolling” in the works of List et al. (2022, 2025)). In studies of idealized but ocean-relevant studies, the works of Frezat et al. (2022) and Ouala et al. (2024) demonstrate that the adjoint-based online learning approach can lead to substantial improvements in the simulation of rotating turbulence in a one-layer system. A particularly impressive demonstration of improvements to model skill in atmospheric modeling with adjoint-based online learning is reported in the work of Kochkov et al. (2024): The resulting neural general circulation model keeps the underlying atmospheric dynamical core, but achieves substantial speed up from replacing certain model processes by machine learned mappings, with notable improvements in weather and some aspects of climate forecasts over existing operational models.
A notable restriction with deploying adjoint-based online learning is the need for a “differentiable” model (e.g., Frezat et al., 2022; Gelbrecht et al., 2023; Sapienza et al., 2024). The problem is analogous to that of deriving an adjoint model in the applications of adjoint optimization (e.g., Forget et al., 2015; Heimbach et al., 2019): The adjoint associated with the dynamical model is required for performing the back-propagation calculations to compute the derivative of the loss function with respect to the parameters. For this, we either derive the adjoint associated with the dynamical model by hand and then discretize, or we discretize the dynamical model and then derive the numerical adjoint (e.g., Maddison et al., 2019). The former effectively requires maintaining two numerical models at the same time, and presumably any changes to one needs to be consistently changed in the other. The latter is in principle possible with algorithmic differentiation capabilities (e.g., Gelbrecht et al., 2023; Sapienza et al., 2024). However, neither of these choices are normally taken at the initial stage of dynamical model development (with notable exceptions, e.g., MITgcm for adjoint optimization), and adjoint-based online learning as intended is in general not possible because the dynamical model does not by default have an associated adjoint. The previous works on online learning essentially wrote custom codes for idealized models (e.g., Frezat et al., 2022; List et al., 2022, 2025; Ouala et al., 2024), or rewriting an operational dynamical core (Kochkov et al., 2024) with algorithmic differentiation capabilities, enabling the adjoint-based online learning approach as intended to be performed. On that note, however, there has also been recent proposals to alleviate the need of a differentiable dynamical model by approximating some of the calculations accordingly (List et al., 2025; Ouala et al., 2024). These “approximately” online approaches are of particular interest and applicability, since it may inherit some of the benefits of the full online learning and alleviate the need of a differentiable model, although the ideas and proposals require further testing and exploration.

Guided by the recent results and proposals for online learning, a principal aim of the present work is to provide an evaluation of the possible benefits and limitations of online learning over offline learning approaches for *baroclinic turbulence*, that is, turbulence within a rotating stratified fluid systems, such as the Earth's atmosphere or ocean. Specifically, baroclinic turbulence is usually not well-resolved in numerical ocean models covering a sufficiently large geographical area because of the associated computational constraints. Parameterization of the missing or incomplete eddy-mean feedbacks are desirable or sometimes essential, and critically impact the skill of the resulting ocean models (e.g., Fox-Kemper et al., 2019). Deriving parameterizations from machine learning methodologies, given the often black-box nature of the resulting models (but see, e.g., Zanna & Bolton, 2021) and how one enforces underlying physical principles such as conservation laws in the resulting parameterizations, is especially challenging. While one would like the machine learning approach to “learn” those physical principles simply from exposure to data, in reality there is no guarantee that this actually happens unless additional information or constraints are provided. Violation of the underlying physical principles can additionally have impacts on the practical usability of the machine learning derived parameterizations via consequences on numerical model stability (e.g., Frezat et al., 2022; Guillaumin & Zanna, 2021; Ross et al., 2023; Srinivasan et al., 2024; Zanna & Bolton, 2021). In the present work we will evaluate the performance of different training methodologies and the derived parameterizations using the established benchmark procedure documented by Ross et al. (2023) for baroclinic turbulence in a quasi-geostrophic system.

In Section 2 we provide an overview of the methodology, detailing the differences between offline and online learning, the latter in the full form with an adjoint, and an approximate online form proposed by List et al. (2025). The benchmark set up of Ross et al. (2023) as well as details of the differentiable dynamical model with algorithmic differentiation capabilities are detailed as part of the methodology. Section 3 provides a comprehensive evaluation of the results from the different methodologies. Section 4 provides additional details procedures that are more specific to online learning. Section 5 summarizes the results of the work and provides an outlook to further usage of online learning in the Earth System Modeling setting.

## 2. Methodology

In this section we first detail what we mean by offline and online learning, highlighting similarities and links in the principles and nomenclature between the Earth System Science, machine learning, and classical optimization. We will opt for using the terminology in classical optimization (e.g., Gunzburger, 2003) in our expository part. The detailed benchmark setting and the numerical model largely based on Ross et al. (2023) are then documented.
<figure><img src="image_4.png" /><figcaption>**Figure 1.** Pictorial schematic of the three training approaches considered in this work. (a) The offline approach, where there is no dynamical model *M* and no forward integration in time, and only the CNN *N* is present as a map between inputs and outputs. (b) The full online approach, where the hybrid model (a dynamical model *M* supplemented with a CNN *N* for the sub-grid tendency) is used for the forward integration and evaluation of the loss function, and the adjoint of the dynamical model allows the back-propagation calculation in evaluating the sensitivities for updating the CNN parameters. (c) An approximately online approach, which is like the full online approach except there is no adjoint of the dynamical model *M*, and an approximation needs to be made for the back-propagation calculation.</figcaption></figure>

## 2.1. Offline Versus Online Learning

At its core, the problem we are interested in is one of regression: There is a high resolution model truth, and a low resolution model possibly with some extra terms added in, and the aim is to choose the extra terms so that the low resolution model is closer to the high resolution in some metric(s) to be specified. The problem can then be phrased as one of optimization: The objective is to choose the magnitude, distribution and/or form of the extra terms (the control variables), such that differences in the target data from the high resolution model and the predictions from the modified low resolution model in a relevant metric, encoded in a loss function (or sometimes cost, objective or mismatch function) is minimized. The extra terms and the loss function are defined in terms of a state vector, consisting of the variables provided by the model directly, or some functions of those variables (e.g., the velocity from a fluid model, or the transport computed from an integral of the velocity). For concreteness, we will be referring to Convolutional Neural Networks (CNNs) as the machine learning algorithm, although the exposition we provide is valid for other choices.

In this work we consider three training approaches for the CNNs: The more standard offline approach, a full online approach, and a case that we term the approximately online approach. A schematic of the three approaches we consider in this work is pictorially summarized in Figure 1.

### 2.1.1. Offline Learning

We start with the more standard approach of offline training (e.g., Ross et al., 2023), or sometimes one-step learning in the machine learning literature (e.g., List et al., 2025). Training data is obtained for example, by filtering data from a high resolution model onto a low resolution grid, forming our state vector y. A CNN denoted N(x|θ) is our mapping function from some input x to the extra terms that we want to include into a low resolution model; the overbar denotes a filtering procedure, and the CNN weights θ are the control variables. We aim to find the θ that minimizes a loss function, usually defined as the mean squared error between the output of the CNN and the target data, that is, we seek θ such that

$$ \mathcal{L}_{\text{off}} = |\bar{y} - \hat{\bar{y}}|^2, \quad \text{with} \quad \hat{\bar{y}} = N(\bar{x}|\theta), \qquad (1) $$
is minimized. The crucial observation here is that the construction of the loss function is defined without any reference to the dynamical model itself: It is essentially seeking the optimal map that minimizes the mismatch between the target data $\bar{y}$ and the prediction $\hat{\bar{y}} = N(\bar{x}|\theta)$. When a time-series of spatially varying data is provided, one should regard the data as multi-dimensional arrays that are appropriately reshaped or flattened into a collection of data points, and the returned object is the optimal mapping $N(\cdot|\theta)$ over the set of training data.

In that sense, however, the temporal information is absent in offline learning, which is potentially problematic for data with a temporal dimension: Causality in time provides additional constraints that may not be respected by the resulting offline models. While such offline models are likely to have skill in the mapping of the data at fix instances in time (i.e., a priori testing), there is usually no guarantee that the hybrid dynamical models with the offline model implemented has skill or be numerically stable over time (i.e., a posteriori testing). The numerical stability aspect is particularly problematic, and various stabilizers may be required, ideally physically informed. The proposal of Ross et al. (2023) imposes a zero net potential vorticity tendency as a hard constraint in the CNNs trained offline, since eddy parameterization should redistribute but not create addition potential vorticity in the system. There are also other methodologies to enforce conservation, some of which will be discussed in Section 4.3.

### 2.1.2. Full (Adjoint-Based) Online Learning

One might suspect including the dynamical model into the definition of the cost function and/or the training procedure might be beneficial. Instead of looking for optimal choices of control variables $\theta$ over isolated instances in time, we might be more interested in the optimal choice over solution trajectories defined over some time windows. One reason is presumably that the inclusion of the dynamical model creates a mapping that forces the optimization to take into account that there is inherent temporal correlation/causality in the data. Further, the resulting optimization for the control variables is expected to remove the choices that lead to large deviations (e.g., associated with unstable model integrations in time), since that leads to growth in the loss function defined over a time window. The inclusion of the underlying dynamical model in the training and the training procedure over some time window are what we consider to be the key ingredients to online learning (e.g., Frezat et al., 2022; Kochkov et al., 2024; Ouala et al., 2024); this is called temporal unrolling in the machine learning literature (e.g., List et al., 2022, 2025).

We make precise what we mean by an adjoint-based online learning in its “full” form first. Using the same terminology as above, online learning considers a loss function defined over a solution trajectory, that is,

$$ \mathcal{L}_{\text{on}} = \sum_{i=0}^{K} |\bar{y}_{t_0+i\Delta t} - \hat{\bar{y}}_{t_0+i\Delta t}|^2, \quad \text{with} \quad \hat{\bar{y}}_{t_0+i\Delta t} = M^i(\bar{x}_{t_0}|\theta), \qquad (2) $$

where $y_{t_0+i\Delta t}$ is the state vector at time $t_0 + i\Delta t$, $\Delta t$ is the (fixed) discretized time step, $\hat{\bar{y}}_{t_0+i\Delta t}$ denotes the prediction of the state vector at the same time, and $K\Delta t$ is the extent of the time window (so $K$ is the number of time steps that defines the time window of training). The mapping $M$ is a hybrid dynamical model, that is, a low resolution dynamical model with terms predicted by the CNNs, which generates the predicted state vector $\hat{\bar{y}}_{t_0+i\Delta t}$ given some initial data denoted $\bar{x}_{t_0}$. The objective is again to choose the CNN weights $\theta$ (the control variable) such that $\mathcal{L}_{\text{on}}$ is minimized. What has been mentioned so far would be what we would regard as “online” learning; the choice we make in employing an adjoint-based methodology utilizing the derivative information for the optimization problem makes it adjoint-based online learning. One notable difference for online learning relative to offline approaches is that the latter regards the whole set of data at the same time, while the former does not. Offline learning is akin to finding a curve/surface of best fit given the whole collection of data embedded in some high dimensional data space, while online learning is aiming to find the best fit trajectory over the analogous high dimensional data space, with time parameterizing how far we are along said trajectory. The approach is identical in spirit to that of Four-Dimensional Variational Assimilation (4DVAR; e.g., Kalnay, 2002): The mapping might be the full dynamical model, the target data would be the observations of the atmosphere/ocean state (or related derived observables, e.g., the volume transport), and the control variable could be the initial state or other model parameters (e.g., free parameters in existing parameterizations).
If we assume appropriate differentiability of $\mathcal{L}_{\text{on}}$, we could solve the problem of $\partial \mathcal{L}_{\text{on}}/\partial \theta = 0$ to find the (local) minimum of $\mathcal{L}_{\text{on}}$ (abusing notation somewhat here since $\theta$ is formally a function, so we should be talking about Fréchet or variational derivatives). This may be done with an adjoint associated with the mapping $M$ given the choice of $\mathcal{L}_{\text{on}}$. Adjoint-based methods use $M$ to perform a forward sweep leading to an evaluation of $\mathcal{L}_{\text{on}}$, and the associated adjoint performs a back-propagation calculation that leads to an evaluation of $\partial \mathcal{L}_{\text{on}}/\partial \theta$ via the chain rule (e.g., Gunzburger, 2003; List et al., 2022, 2025; Ouala et al., 2024). A technical requirement for adjoint-based methods is that $M$ needs to be differentiable, in the sense that one can derive the associated adjoint, by hand or algorithmically. The CNN part is not a problem since these are combinations of maps and activation functions implemented in software that enables automatic differentiation capabilities, but most dynamical models are not constructed to have a manual adjoint or have the capabilities to generate an adjoint automatically via algorithmic differentiation (e.g., List et al., 2025; Ouala et al., 2024; Sapienza et al., 2024).

If the dynamical model is differentiable, then the back-propagation of the sensitivities for the gradient calculations can be done throughout the whole time window of interest, and this might be called “end-to-end,” in that the information can be propagated from one end to another in both directions (e.g., List et al., 2025; Ouala et al., 2024; Sapienza et al., 2024). Note that just as the dynamical model can have issues with numerical stability in the forward or prediction mode, the associated adjoint model can also have numerical instabilities for the back-propagation calculations, resulting in numerical issues with the gradients $\partial \mathcal{L}_{\text{on}}/\partial \theta$ and the update stage involving the iterative solver for the next forward-backward sweep. Some manual approaches are possible to force the back-propagation stage to be more stable (e.g., approaches detailed in List et al. (2022), although one could argue there is loss of consistency as a result.

### 2.1.3. “Approximate” Online Learning

While the ideal situation would be that we have access to a differentiable dynamical model with no extra stabilizers needed on the forward or back-propagation calculations, the fact is that most dynamical models are not differentiable, nor will the back-propagation step necessarily be numerically stable. One might however still want to perform a training procedure as detailed above to access some of the desirable features that online learning is expected to offer. One possibility is to use a methodology that does not require derivative information, in exchange for possible increases in computational complexity (e.g., Iglesias et al., 2013; Kovachki & Stuart, 2019; Lopez-Gomez et al., 2022; Schneider et al., 2022). Another is to approximate the relevant derivatives in the adjoint-based approach. The approach of Ouala et al. (2024) considers a first-order approximation to $\partial M/\partial \theta$ (which arises from application of the chain rule to get the derivative of $\mathcal{L}_{\text{on}}$), which they term the Euler Gradient Approximation, and they show it converges back to the full adjoint-based approach in the limit of infinitesimal time-step size. The approach of List et al. (2025) considers the following pipeline: proceed as detailed above, but at each time-step, the back-propagation of the information only reaches the CNN stage, and is then set to zero. The resulting approximation is more drastic and heavily truncates the back-propagation calculation, but is particularly easy to implement. In that regard, both the aforementioned approximations target the back-propagation part, but otherwise solves the analogous optimization problem. One might suspect such approximations might be reasonable for short time windows (or perhaps there are better approximation procedures available, as alluded to in Ouala et al. (2024). For the present work, we consider the approach of List et al. (2025) for simplicity. While we might suspect that the fully differentiable online procedure is most theoretically consistent and might be most skillful, it could also be that the approximated online procedures are in fact good enough and inherit some of the desirable properties of the full online procedure (List et al., 2025; Ouala et al., 2024). A secondary aim of the present work is to explore to what extent the approximate online procedure is comparable to the full online procedure.

### 2.2. Experiment and Model Details

Building on the previous works of Frezat et al. (2022) and Ouala et al. (2024) in a one-layer quasi-geostrophic model, we utilize an idealized two-layer quasi-geostrophic turbulence model that allows for baroclinic instability, following the work of Ross et al. (2023), who demonstrated their work on offline training in the *pyqq* model. To enable the full adjoint-based online approach to be performed, we utilize a model called *pyqq-JAX* (Otness, 2024; Otness et al., 2023), which has capabilities to generate the associated adjoint and perform full back-propagation through the dynamical model via the inbuilt algorithmic differentiation capabilities. We utilize
<table><caption>**Table 1**
Summary of Parameters Used for the Two-Layer Quasi-Geostrophic Model</caption><thead><tr><th>Parameter</th><th>Symbol</th><th>Value and units</th></tr></thead><tbody><tr><td>Domain size</td><td>L × L</td><td>1,000 × 1,000 km</td></tr><tr><td>Meridional planetary vorticity gradient</td><td>β</td><td>1.5 × 10<sup>-11</sup> m<sup>-1</sup> s<sup>-1</sup></td></tr><tr><td>Bottom friction coefficient</td><td>r<sub>ek</sub></td><td>5.787 × 10<sup>-7</sup> s<sup>-1</sup></td></tr><tr><td>Layer thickness</td><td>(H<sub>1</sub>, H<sub>2</sub>)</td><td>(500, 2,500) m</td></tr><tr><td>Rossby deformation radii</td><td>L<sub>d</sub></td><td>15 km</td></tr><tr><td>prescribed mean zonal flows</td><td>(U<sub>1</sub>, U<sub>2</sub>)</td><td>(0.025, 0) m/s</td></tr></tbody></table>

the same model settings as in the previous work of Ross et al. (2023) in the “eddy” regime. Briefly, the model is formulated on a doubly periodic β-plane, governed by

$$ \partial_t q_1 + \nabla \cdot (\mathbf{u}_1 q_1) + \beta \partial_x \psi_1 + U_1 \partial_x q_1 = \text{ssd}, \quad (3) $$

$$ \partial_t q_2 + \nabla \cdot (\mathbf{u}_2 q_2) + \beta \partial_x \psi_2 + U_2 \partial_x q_2 = -r_{ek} \nabla^2 \psi_2 + \text{ssd}, \quad (4) $$

where subscripts 1 and 2 denote the upper and lower layer, and $\partial_{x,y}$ denote the derivatives in the x and y (or zonal and meridional) directions. The mean zonal velocities $U_1$ and $U_2$ are set to maintain a prescribed vertical shear that drives the turbulence in the model. Denoting $l$ as the layer index, the perturbation

variables are the potential vorticity $q_l = \nabla^2 \psi_l + \beta y + (-1)^l f_0^2 / (g' H_l) (\psi_1 - \psi_2)$, $f_0$ the planetary Coriolis parameter, $\beta$ is the meridional variation of the planetary rotation, $g'$ the reduced gravity measure relating to the stratification between layer 1 and 2 with depth $H_{1,2}$, the streamfunction $\psi_l$, with the geostrophic velocity given by $\mathbf{u} = (u, v) = (-\partial_x \psi, \partial_x \psi)$, $r_{ek}$ is the Ekman friction acting only on the bottom later dissipating large-scale energy, and ssd denotes the small-scale dissipation, which dissipates small-scale energy and potential entrophy $q^2$. The relevant parameters of the model are summarized in Table 1.

The model is discretized in the pseudo-spectral formalism, with the derivatives computed in Fourier spectral space while the multiplication is performed in real space, utilizing a 2/3 de-aliasing procedure (e.g., Orszag, 1971). The small-scale dissipation ssd is implemented as a quadruple-exponential filter, which acts as a low-pass filter and attenuates wavenumbers above a cutoff threshold $\kappa_c$, which is 2/3 of the model's Nyquist frequency $\kappa_N = \pi$, that is,

$$ q_k = \begin{cases} q_k & \text{for } \kappa < \kappa_c, \\ q_k e^{-23.6(\kappa - \kappa_c)^4} & \text{for } \kappa \ge \kappa_c, \end{cases} \quad (5) $$

where $q_k$ is the Fourier coefficient of $q$ in spectral space. The value of the e-folding scale $-23.6$ is chosen so that the filtering is numerically zero at the largest wavenumber of the system. The model is discretized in time with a third-order Adams-Bashforth scheme.

As discussed in Hallberg (2013), the model grid size $\Delta x (=\Delta y)$ should be smaller than the deformation radius $L_d = 15$ km, and ideally we should have $L_d / \Delta x \ge 2$ to resolve the mesoscale eddies. We use a model grid with size $256^2$, corresponding to $\Delta x_{high} = 3.9$ km (with a time-step size of $\Delta t_{high} = 1$ hr) for the reference high resolution simulations. For the low resolution simulations, we take instead a grid size of $64^2$, corresponding to $\Delta x_{low} = 15.6$ km (with a time-step size of $\Delta t_{low} = 4\Delta t_{high} = 4$ hr). All models start with a randomly generated initial condition, and all the simulations are run for 10 model years to reach the equilibrium. Sample outputs of the upper layer perturbation potential vorticity $q_1$ are shown in Figures 2a and 2b for the high and low resolution calculations respectively. Note that the present choice of the bottom drag time-scale $r_{ek}$ leads to an “eddy” rather than “jet” regime; we discuss the analogous results for the “jet” regime in Section 5 and in the Supplementary Information. All numerical model values stated in this work are exactly the same as those used in Ross et al. (2023).

## 2.3. Training Strategy

For both offline and online approaches, we follow the strategy of large eddy simulation and the processing as in Ross et al. (2023). We first coarse-grain Equations 3 and 4 to get the filtered equations; in this work we coarse-grain by a spectral truncation, since we already utilize a pseudo-spectral formalism. We take the state vector $\mathbf{y}$ to be the sub-grid forcing $S^q$, which are the resulting terms in the coarse-grained equations, given by

$$ S^q = \overline{\nabla \cdot (\mathbf{u}q)} - \nabla \cdot (\overline{\mathbf{u}q}) = (\overline{\mathbf{u}} \cdot \overline{\nabla})q - (\overline{\mathbf{u}} \cdot \nabla)\overline{q}. \quad (6) $$
<figure><img src="image_5.png" /><figcaption>**Figure 2.** Snapshots of upper-layer perturbation potential vorticity $q_1$ at the end of a 10 years simulation, from the (a) high resolution model truth, (b) low resolution model with no CNN active, (c) offline model, (d) full online model, and (e) approximately online model.</figcaption></figure>

The overbar denotes the coarse-graining by the spectral truncation, and $S^q$ is the residual between the coarse-grained advective potential vorticity tendency and the advective potential vorticity tendency formed from the coarse-grained fields. The input for all training strategies is chosen to be the potential vorticity $q$, obtained either by filtering the high resolution model for a priori testing, or the $q$ field as simulated by the low resolution hybrid models for a posteriori testing. The task of the machine learning algorithms is to predict the missing tendency $S^q$, and the models differ only in how the control variables (i.e., the CNN weights) are adjusted during the training phase.

For systems such as the present multi-layer quasi-geostrophic model, there is a choice on how one chooses the loss function. One choice is to construct one single loss function $\mathcal{L}$ based on a weighted sum of the loss function computed from the separate layers $\mathcal{L}_{1,2}$, that is,

$$ \mathcal{L} = \mathcal{L}_1 + \alpha \mathcal{L}_2, \qquad (7) $$

where, for this work, the layer-wise loss function is defined as

$$ \mathcal{L}_{1,2} = \sum_{i=0}^{K} |S_{1,2}^{q}; t_0 + i\Delta t - \hat{S}_{1,2}^{q}; t_0 + i\Delta t|^2, \qquad (8) $$

where $\hat{S}_{1,2}^q$ denotes the predictions from the CNN; for the offline approach, we formally take $K = 1$ but consider the $S^q$ field as a stacked array. The parameter $\alpha$ is introduced to provide a balance so that the data between the layers have comparable magnitudes (e.g., $S^q$ in layer 1 in this model is at least two orders of magnitude larger than that in layer 2), or as our prior bias in which layer we favor. While one could argue the resulting model will be exposed to a wider regime of data by sampling more of the probability distribution, the model could also be argued to be less specialized, and there is subjectivity in how the free parameter $\alpha$ is chosen. Another approach is simply to train up separate CNNs via separate loss function $\mathcal{L}_{1,2}$, such as that following the work of Ross
et al. (2023); two CNNs are thus at play, costing more in terms of training and prediction. We take the latter approach in this work, but discuss the performance of combined or separate choice of loss function(s) in Section 4.3.

Note that it is in principle possible to use the state variables $q$ or $\psi$ directly for online learning, as was done in a previous work of a one-layer model (Frezat et al., 2022; Ouala et al., 2024), targeting the mean-state variables directly instead of targeting the sub-grid tendencies (which are small-scale quantities). For the present work we were interested in providing a like-for-like comparison between the methodologies as well as results from previous works (e.g., Ross et al., 2023), so training with loss functions defined in terms of $q$ or $\psi$ was not pursued in detail; we present some related results in Section 4.3 and comment on the procedure in Section 5.

A technical issue is that `pyqq-JAX` assumes data uses single-precision floats by default, but the diagnosed magnitude of $S^q$ in both layers are small, roughly $10^{-11}$ in layer 1 and $10^{-13}$ in layer 2. In the present work we employ the single-precision version, and the only standardization we perform is multiplying $S^q$ and perturbation potential vorticity $q$ by $10^{13}$ and $10^6$ in both layers, with the relevant scale factor added or removed at the appropriate stages. Empirical testing for this work suggests the choice of data standardization is not crucial to the performance of the models.

For both offline and online learning, we employ CNNs as the mapping between input and outputs. For commutativity with the differentiable solver `pyqq-JAX`, we use the JAX library `Equinox` (Kidger & Garcia, 2021). We follow the work of Ross et al. (2023) on the choice of CNN architecture in `pyqq`: the CNNs have eight hidden convolutional layers with square kernels, first two layers with size 5 and in circular padding mode, and the remaining six layers have size 3 in constant padding mode. The stride for the cross-correlation is 2 for the first two layers, and 1 for other layers. ReLU is chosen as the activation function between each layer. No batch normalization is performed; this is partly for simplicity, but also for model integration stability, where approximately online models were found to be particularly unstable when batch normalization was active, and models often crash during model integration if additional stabilizers were absent (see next paragraph).

One thing we do differently to Ross et al. (2023) relates to the imposition of a hard constraint so that there is no net tendency of potential vorticity added by the predicted $\hat{S}^q$ that is,

$$ \int \hat{S}^q dA = 0. \qquad (9) $$

The hard constraint is motivated by the observation that the nonlinear transfers being parameterized should transfer but not create or destroy potential vorticity. We note there are at least three possible meanings here: (a) the training for the CNN weights results as usual with no regard of the constraint, and the resulting prediction has the mean removed *post-hoc*; (b) the choice for the CNN weights takes into account that the prediction will need to pass through a de-meaning procedure as the final CNN layer; (c) the choice for the CNN weights is restricted such that the resulting prediction has zero mean. The second case is what is meant here and in Ross et al. (2023); the first case is arguably ad hoc, while the third case would be theoretically most satisfactory, but is likely difficult to execute in practice. In the bulk of the results reported here, the zero-mean constraint is not imposed, and the models are allowed to fail during the training or integration stage to explore stability aspects relating to the three training strategies considered here. Further discussions are given in Section 4.4.

For offline learning, we generate the training data by performing 100 high resolution simulations with different randomized initial conditions, and each of these are integrated forward in time for 10 model years (86,400 time steps). We sample the high resolution simulations every 1,000 time-steps as a means to de-correlate data temporally, coarse-grain the data by spectral truncation, and randomize the ordering of the whole set of data frames. There are 86 images per simulation over 100 simulations, leading to 8,600 frames of training images per model layer, which are randomly selected for the training-testing data sets for the CNN by a choice of the random seed. For statistical significance, we train up an ensemble of 10 CNNs each differing in the choice of the random seed, and results are reported for the ensemble members that did not crash during the model integration. The offline CNNs are trained with a batch size of 32 for 800 epochs in a NVIDIA Tesla T4 GPU in Google Colab. These take about 10 min for the pair of CNNs over the two layers, although these CNNs can be trained independently.
For online learning, a total of 5 high resolution simulations were performed where the data output is at every 4 time-steps (since $\Delta t_{low} = 4\Delta t_{high}$), generating sets of data frames in a temporal sequence. The choice of random seed affects the subset of trajectories selected for training-testing, that is, the hybrid models are initialized with the data frame at the beginning of the time-window, and then integrated over the time window for each given trajectory for the loss function evaluation and associated back-propagation calculations. For the training, we gradually increase the window size from 2 time-steps to 10 time-steps in strides of 2 for training stability, that is, starting from a converged model and using that as the initial condition for the longer set of learning (e.g., Frezat et al., 2022); this is referred to as a “curriculum” in List et al. (2022), and is related to continuation approaches in optimization (solving a related but easier and/or more stable problem, which gradually transforms into the desired problem). Note that the time window size is rather small, and further discussion on the choice of window size is given in Section 4.1. The CNNs for both online learning procedures are trained for 10 epochs at each temporal window size; while the final model could be argued to have been trained for 50 epochs in total, we caution somewhat against this interpretation, since a model trained for 50 epochs is not the same as a sequence of five models trained at 10 epochs each. As in the offline approach, we train up an ensemble of 10 CNNs each differing in the choice of the random seed, and results are reported for the ensemble members that did not crash during the model integration. For every random seed, 67 trajectories are used, that is, 67 forward-and-back calculations simultaneously. For the online learning approaches, these are trained on a (mini-)batch size of 1 out of 16 (or the pure form of stochastic gradient descent) on a NVIDIA Tesla T4 GPU in Google Colab, so over 10 time-steps the full online models are exposed to $67 \times 5 \times 10 = 3,350$ images per update of the CNN weights; mini-batch is used for speed reasons. Both strategies lead to a training of about 30 min for the pair of CNNs over the two layers, but these CNNs need to be trained together unlike the offline case, since there is dynamical coupling between the layers during the forward integration.

We have attempted to vary the amount of data that is exposed to the training strategies: We report on the tests with the online case later in Section 4.1, but the results reported in this work are robust to mild changes in the number of data frames, number of epochs and batch sizes during the model as part of the training procedure. As mentioned, for each of the training methodology (offline, online, and approximate online), we train an ensemble with 10 members, each member within the set of 10 differing only in the choice of the random seed (but the choice of seeds is uniform across the ensembles considered in the three training methodologies). The models are allowed to fail, and we state the number of members that are unstable during training or the model integration stage and remove them from the model score calculations.

Samples of the upper layer perturbation potential vorticity $q_1$ from the reference high resolution, low resolution with no CNN and low resolution hybrid models (with the corresponding CNNs implemented) for a prognostic prediction are shown in Figure 2 at the end of a 10 years simulation. Note that because of sensitive dependence on initial conditions of baroclinic turbulence, the displayed snapshots of $q$ are not expected to be comparable since these are individual realizations, and we are primarily interested in the broad statistical properties. The main feature we highlight here is that the resulting $q$ field from the hybrid models (Figures 2a–2c) are somewhat sharper than the low resolution model with no CNN active (Figure 2b). There is thus some suggestion the resulting CNNs are energizing the fluctuating eddy motions in some way relative to the less vigorous and more diffusive low resolution model. The related statistical properties which will be quantified in the following section.

## 3. Results Relating to Model Skill

To examine the model performance between offline and online learning, we perform both a priori and a posteriori testing. Recall from above that the CNNs expects the (perturbation) potential vorticity $q$ as the input field, and outputs a sub-grid forcing $\hat{S}^q$. A priori testing examines the models' skill in predicting snapshots of diagnosed $S^q$ for given filtered high resolution potential vorticity field $q$, while a posteriori testing will be related to model performance in prognostic mode for relevant model metrics to be defined. We note that, for the present set of results with no batch normalization and no de-meaning layer included in the CNN, none of the models crashed during training or model integration for the choice of random seeds we have selected.

### 3.1. Distribution Similarity

In Figure 3 we show quantile-quantile (Q-Q) plots to assess the distribution similarities between the targets $x$ and the predictions $y$: Dissimilarities in the distribution arise as deviations from the $y = x$ line. We show results from
<figure><img src="image_6.png" /><figcaption>**Figure 3.** Quantile-Quantile (Q-Q) plot, with the distribution of diagnosed target (the $S^q$ from high resolution filtered onto the coarse grid) on the x axis, against the predicted $\hat{S}^q$ generated by the CNNs on the y axis (green: offline model; orange: online (approx) model; blue: online model). Results from (a), (b) a priori testing, and (c), (d) a posteriori testing, for layer 1 and 2 respectively. The red line is the identity line $y = x$.</figcaption></figure>

the a priori test for both fluid layers in Figures 3a and 3b, using the same group of input $q$ and diagnosed $S^q$ for the trained CNNs. Here it would appear that both the offline and the full online approach are able to generate predictions that largely agree in the distribution since they lie closer to the $y = x$ line, with data points from the full online learning lying somewhat closer to the identify line. On the other hand, the approximate online approach leads to biases manifesting as the data points located away from the $y = x$ line, particularly in the top layer.

For the a posteriori testing, the same initial state $q$ is provided to all hybrid models, and these models are integrated forward in time for 10 model years, sampled every 20 time-steps. The predicted $\hat{S}^q$ from the hybrid models are compared to the diagnosed and filtered $S^q$ from the high resolution model, and the corresponding Q-Q plots are shown in Figures 3c and 3d. Since there is inherent data drift as the models run in prognostic mode, the deviations away from the $y = x$ line are larger under a posteriori testing for all cases. While all approaches have comparable performance in the upper layer in terms of the distribution (Figure 3c), the online approaches appears to perform better in the lower layer (blue and orange points closer to the $y = x$ line in Figure 3d). We conclude the online approach improves upon the offline cases in this metric.

## 3.2. Energy Budget

The accurate representation of energy levels is crucial in numerical ocean models, since energy measures the variance of the dynamical variables. We show in Figure 4 the time and depth-averaged total kinetic energy (KE) spectra over the last five model years, where KE density is defined as $\mathbf{\mu}\mathbf{u}^2/2$, and the spectra is computed assuming horizontal isotropy. It can be seen that the low resolution model (without a CNN active; gray line) generally lacks KE across all scales, and there is the “kick-back” feature at the high wavenumbers, symptomatic of energy accumulation at the grid-scale and a sign of the model being under-resolved. The hybrid models on the other hand generally resolve the bias that there is a lack of KE across all scales. The offline model possesses a kick-back in the KE spectra; we note this feature is not something that can be “tuned” out by adjusting the training, and is a robust feature of the resulting offline models. On the other hand, neither of the online models show this kick-back

<figure><img src="image_7.png" /><figcaption>**Figure 4.** Depth-averaged total kinetic energy spectra, normalized by $n_x^2 \times n_y^2$. All panels show the data from the target high resolution simulation (black line) and the low resolution simulation with no CNN active (gray line). Results from the (a) offline model (green line), (b) full online model (blue line), and (c) approximately online model (orange line). The colored lines are from a 10 member ensemble averaged performance of the hybrid models, and the shading denotes the standard deviation displayed by the full ensemble.</figcaption></figure>
<figure><img src="image_8.png" /><figcaption>**Figure 5.** Full spectral energy budget for the APE flux (blue line), APE generation from the imposed background shear (orange line), KE flux (green line), bottom friction (red line), dissipation from the spectral filter (purple line), and input from the CNNs (brown line), from the (a) high resolution model truth, (b) low resolution model with no CNNs active, (c) offline hybrid model, (d) full online hybrid model, and (e) the approximately online hybrid model. For (c, d, e), the colored lines are from the 10 member ensemble averaged performance of the hybrid models, and the shading denotes the standard deviation displayed by the full ensemble.</figcaption></figure>

feature for the reported time window size of 10 time-steps; see Section 4.1 for further related results. All models have a KE that is too large at the large-scales, which is consistent with the results reported in Ross et al. (2023) and Perezhogin et al. (2023) for offline models.

While the KE spectra above tells us about the energy content distribution, it is informative to consider the energy fluxes. The full energy budget for the present idealized ocean model can be represented as (e.g., Jansen & Held, 2014)

$$ \partial_t E = KE_{\text{flux}} + APE_{\text{flux}} + APE_{\text{gen}} - F - D. \quad (10) $$

Here, E is the total depth-averaged energy, APE is the available potential energy, the fluxes are the transfers in spectral space, and F and D are the forcing and dissipation terms respectively. Bottom drag and the spectral filter at the small-scales as part of the fluid model formulation from Equations 3 and 4 are grouped into dissipation, while the predicted $S^d$ in relation to Equation 6 from the implemented CNNs can manifest as forcing, dissipation or as a flux. The tendency terms on the right hand side of the equation are diagnosed from the model runs, and a summary of these from the different model calculations over the same 5 year averaging period are shown in Figure 5. Starting first with the high resolution model truth (Figure 5a), positive values denote energy input at the relevant spatial scales. The generation of APE is from the imposed vertical shear and is always leading to an energy input, while bottom drag acts to remove energy across all scales. The flux of APE is from the large to small-scale, while the KE flux is upscale, consistent with theories of two-dimensional quasi-geostrophic turbulence (e.g., Jansen & Held, 2014; Salmon, 1980). For comparison, while the low resolution model with no CNNs active display a qualitatively similar balance (Figure 5b; similar APE generation, but generally weaker fluxes and dissipation by bottom drag), there is a non-negligible balance at the small-scales between the KE flux related to accumulation at the grid-scale and the dissipation from the spectral filter, at odds with that displayed for the high resolution model truth.
With those observations in mind, it can be seen from Figures 5c–5e that the presence of the trained CNNs, regardless of whether these are trained by offline or the online approach, lead to increased explicit energy fluxes that are more in line with the high resolution model truth. All the resulting CNNs seem to be mimicking some sort of backscatter (e.g., Jansen & Held, 2014), as indicated by the brown lines being positive in the intermediate ranges away from the grid scale in Figures 5c–5e; diagnoses indicate a backscatter in both KE and APE in comparable parts. The fact that CNNs appear to be learning some sort of backscatter is alluded to but not stated explicitly in the work Ross et al. (2023) (because we do not see a decomposition of the energy budget that shows the contribution of the associated CNNs). The resulting CNNs also seem to be mimicking effects of the small-scale dissipation from the spectral filter, in that the explicit small-scale dissipation from the spectral filter (purple) line has reduced in magnitude relative to the no CNN case. Upon computing the integral over spectral space, the decrease in explicit dissipation arising from the spectral filter (Equation 5) in the ensemble mean relative to the no CNN low resolution model was diagnosed to be 42.7% for the offline model, 64.4% for the full online model, and 55.3% for the approximately online model. The behavior of the CNNs are consistent with the observation that the low resolution models without a CNN would be more sluggish at the intermediate scales and has a dissipation that is not sufficiently strong at the grid-scale relative to the high resolution model truth, and so the resulting CNNs are trying to compensate for the discrepancy (note however we only claim the observations are consistent, since there is no evidence to claim anything more at present).

### 3.3. Long-Term Similarity Scores

For a more qualitative measure of the skill in the models in terms of spectral similarity, we compute some of the metrics detailed in Ross et al. (2023). For the distributions of the model variables (potential vorticity $q$, both components of velocity $u$ and $v$, streamfunction $\psi$, so four variables over two layers), we compute the 1-Wasserstein distance $W_1$ (e.g., Villani, 2008), which in the present case can be computed from the unsigned area difference between two cumulative distribution functions, that is,

$$ \text{distrib\_diff}(\text{sim1}, \text{sim2}; f) = \int_{-\infty}^{\infty} |P_{\text{sim1}}(f \le x) - P_{\text{sim2}}(f \le x)| \, dx, \quad (11) $$

where $f$ is the quantity of interest, and $P(f \le x)$ is the empirically determined cumulative distributions, obtained from binning the variable data into a histogram based on their values over real space. For differences in distributions in spectral space (KE and enstrophy in both layers, KE flux, APE flux, APE generation and friction), we compute the root-mean-squared or $L^2$ distance of the spectral distributions, that is,

$$ \text{spectral\_diff}(\text{sim1}, \text{sim2}; f) = \sqrt{\frac{1}{|\mathcal{K}|} \sum_{k \in \mathcal{K}} (f_{\text{sim1}}(k) - f_{\text{sim2}}(k))^2}, \quad (12) $$

where $f$ is the spectral distribution of the quantity of interest and is real, $k$ is the (isotropic) wavenumber, and the normalization is by $|\mathcal{K}|$, which is an isotropic wavenumber common to both simulations, up to 2/3 of the Nyquist frequency of the low resolution simulation.

Once the distribution and spectral differences are computed, the similarity score is computed as

$$ \text{similarity}(\text{hybrid}, \text{HR}, \text{diff}) = 1 - \frac{\text{diff}(\text{hybrid}, \text{HR})}{\text{diff}(\text{LR}, \text{HR})}, \quad (13) $$

where HR, LR and hybrid denote diagnostics from the high resolution, low resolution with no CNN, and the corresponding hybrid models. The present relative measure is bounded above by 1, and the similarity score is maximized when the hybrid model matches the high resolution model truth exactly in the relevant metric. On the other hand, the ratio can be zero or even negative, achieved respectively when the hybrid model does no better or does worse than the low resolution model without the CNNs active.

In Figure 6 we show the similarity scores for quantities of interest relating to the model. The model variable distributions have wide spreads, with biases particularly in the top layer (except for the streamfunction). The spectral distributions on the other hand show substantially less variability, and the inclusion of the CNNs
<figure><img src="image_9.png" /><figcaption>**Figure 6.** Distribution similarity scores for potential vorticity *q*, both components of the velocity *u* and *v*, the streamfunction *ψ*. KE and enstrophy, depth-averaged KE flux, APE flux, APE generation, and the bottom drag (labeled as Friction). Subscripts 1 and 2 denote the upper and lower layer, for the (green) offline model (blue) full online model, and (orange) approximately online model. The solid line is for the ensemble average over the 10 members, and the shading denotes the standard deviation and the shading denotes the standard deviation displayed by the full ensemble.</figcaption></figure>

generally lead to improvements, as indicated by a moderately positive similarity score. The approximately online models have particularly poor performance in the bottom layer potential vorticity distribution. It is generally the case that the full online models improve upon the other two cases in the spectral fluxes, with a higher ensemble average similarity scores, although the overlaps in the ensembles' scores are not entirely well-separated.

It is perhaps interesting to note that the lukewarm performance of the models in the variable distributions do not seem to have a strong imprint on the spectral fluxes, suggesting that “skill” is metric dependent, and one should probably not rely on a single metric for gauging model performance. The mismatch in the performance between model and spectral distributions may be due to the choice of loss function, which targets the sub-grid small-scale forcing terms $S^q$, so it may not be a surprise that we are obtaining improvements largely in the quadratic quantities that are strongly affected by the form of $S^q$. The choice of loss function was explored briefly, with some results reported in Section 4.3.

# 4. Further Comparisons of the Online Procedures

In this section we provide further explorations of the online learning methodology, motivated by some of the observations highlighted in the previous section, namely the dependence on the learning window size, the learning of small-scale dissipation, possible flexibility in the choice of the loss function, and imposition of constraints to control the accumulation of potential vorticity in the approximately online approach.

## 4.1. Dependence on Window Size

One free parameter for the online learning procedure is the time window over which the training is performed. While we would suspect that the longer the window the better (because the model is then exposed to more of the model trajectory and data content), the choice also needs to be balanced by the computational cost and possible stability issues (e.g., that reported in List et al. (2022)).

We have considered choices of window size on the performance of the resulting online model performance, and the main notable result we have found relates to the depth-averaged KE spectra at the small-scales, shown in Figure 7. As a reminder, the models here are trained via a curriculum (e.g., Frezat et al., 2022; Kochkov et al., 2024; List et al., 2022; Ouala et al., 2024), in that a mapping is first learned from a window with restricted size, and that mapping is then used as the initial mapping for learning at the next enlarged window size (cf. continuation approach in optimization and numerical methods). In both online approaches, there is the kick-back feature at the grid-scales when window size is small, but this feature gradually diminishes as window size increases.

In the cases tested for the present set up, marginally longer window lengths beyond the default choice of ten time-steps lead to increases in training time, with not overly significant gains in model skill (not shown; largely similar to Figure 6). An observation we make is that there appears to be an overshoot with increased window size in Figure 7 where the KE content is now a little too small at the grid-scales. Note that Figure 7 is shown with a linear
<figure><img src="image_10.png" /><figcaption>**Figure 7.** Zoomed in depth-averaged KE spectra (cf. Figure 4) focusing on the small scales for the (a) online model and (b) the approximately online model, for models trained with different time window sizes (K) The solid line is for the ensemble average over the 10 members, and the shading denotes the standard deviation and the shading denotes the standard deviation displayed by the full ensemble.</figcaption></figure>

rather than logarithmic scale, so the overshoot is mild. The associated damping/backscatter near the grid-scale is not necessarily a negative feature either, since it may be beneficial for model stability.

We further note that we have not had noticeable issues with the gradient calculations associated with the backpropagation calculations, but this may be because the chosen training window size is rather small.

## 4.2. CNN Represented Small-Scale Dissipation

The observation in the spectral energy budgets in Figure 4 suggests that all the CNNs seem to have some representation of the backscatter and small-scale dissipation from the high resolution model. Here, we perform a posteriori testing of the same hybrid models reported in Section 3 but in the absence of explicit small-scale dissipation, to see how the CNNs perform out-of-sample in this sense. The resulting spectral fluxes are shown in Figure 8. We see in Figure 8a that the offline model performance is generally poor, with large biases over the whole bandwidth, and the ensemble spread is substantial (only 2 out of the 10 model integrations were in fact stable). The spectral fluxes from the online models shown in Figures 8b and 8c are is comparable in shape to Figure 5d but with reduced magnitude, indicating a weaker energy cycle, but with some evidence of model skill when model is taken out-of-sample. The similarity scores for the full and approximately online models are slightly lower (omitted for brevity).

For completeness, we also considered whether it was possible to turn off the explicit small-scale dissipation for online hybrid models at the training stage, so that the online approaches would end up learning the grid-scale dissipation simply via exposure to the high resolution model data that resolves what would the grid-scales of

<figure><img src="image_11.png" /><figcaption>**Figure 8.** Spectral fluxes for a posteriori calculations in a hybrid model with the explicit small-scale dissipation switched off, using the CNNs reported in Section 3. (a) The offline hybrid model. (b) The online hybrid model. (c) The approximately online hybrid model. The solid line is for the ensemble average over the 10 members, and the shading denotes the standard deviation and the shading denotes the standard deviation displayed by the ensemble members that did not crash (2 out of the 10 offline models, and the full ensemble for both online approaches).</figcaption></figure>
the low resolution hybrid model. We note that the Orszag 2/3 dealising is still active, which imparts some numerical stability to the integration, but by itself is insufficient, in that the low resolution model without a CNN will crash in the same experimental setting due to grid-scale accumulation of energy and entrophy. Only the full online method will converge and return CNNs in this no dissipation setting, although the resulting model performance is generally poor, with a significant bias in the KE spectrum over a noticeable bandwidth in the small-scales (not shown). Reported observations are however for training based on a loss function defined with $S^q$; a curious result for an analogous procedure for a different loss function is given in the next subsection.

### 4.3. Choice of State Variable in the Loss Function

In the results presented thus far the choice of the loss function targets the sub-grid forcing $S^q$ given in Equation 6, and for the present two-layer quasi-geostrophic system the CNNs are trained with separate loss functions defined per fluid layer. One choice that we could have made was to consider a combined loss function as in Equation 7, choosing a coupling parameter $\alpha$ accordingly (or we could have considered appropriately normalized $\mathcal{L}_{1,2}$). Analogous training on the combined loss function does not lead to noticeable differences in the qualitative conclusions drawn so far, with similarity scores largely like that of Figure 6, possibly with minor degradation in the skill. The minor degradation may simply be that the complexity (i.e., the number of degrees of freedom) of the CNNs trained on the one combined loss function is lower (since we kept the CNN architecture the same); this is however not something that we have investigated in comprehensive detail.

In the optimization literature, depending on the objective, the loss (or objective) function can often be designed or constructed to target different features. For example, instead of the usual root-mean-squared loss, the work of Guillaumin and Zanna (2021) considers the log-likelihood for their machine learning-based stochastic parameterization. The work of Yan et al. (2023) alludes to the use of negative Sobolev semi-norms as possible choices of loss functions to target larger-scale features (cf. the works in the mixing literature such as Lin et al., 2011; Thiffeault, 2012). For the present work, we note that the models display skill in the quantities relating to spectral fluxes, but less so in the model quantity variables (Figure 6). The observations may have arisen from the fact that $S^q$, which is a small-scale variable relating to the nonlinear transfer term, was used in the definition of the loss function. In some sense $S^q$ is the only choice we have for forming the loss function for offline learning approaches, since we have no natural mapping of $S^q$ to the model variables, as the dynamical model $M: S_{q,t_0}(q_{t_0}|\theta) \mapsto q_{t_0} + \Delta t$ is not involved explicitly. This limitation is not present for online learning approaches, and one could consider using the model variables in the definition of the loss function as a way to target improvements in the model variable predictions (e.g., Frezat et al., 2022; Ouala et al., 2024). One may suspect that if we are able to represent the model state well, then the statistics (e.g., the implied $S^q$) may also be improved at the same time. The work of Frezat et al. (2022) finds that this was the case in a one-layer quasi-geostrophic model with full online learning and a loss function based on potential vorticity $q$.

We have attempted to train online models with loss functions that are solely in terms of model state variables (potential vorticity $q$ or streamfunction $\psi$), or as a mix of sub-grid and model state variables, for example,

$$ \mathcal{L}_{1,2} = \sum_{i=0}^{K} |S_{1,2}^{q}; t_0 + i\Delta t - \hat{S}_{1,2}^{q}; t_0 + i\Delta t|^2 + \tilde{\alpha} \sum_{i=0}^{K} |\bar{q}_{1,2}; t_0 + i\Delta t - \hat{q}_{1,2}; t_0 + i\Delta t|^2, \quad (14) $$

for appropriate values of a weight parameter $\tilde{\alpha}$ (where we would expect the $\tilde{\alpha} \sim 0$ limit to correspond to the primary scenario considered in the present work, although it is also possible that the problem here is singular). For the present case, and without comprehensive investigation into modifications of the CNN structure and/or hyper-parameters, online training (with explicit small-scale dissipation active in the hybrid model) of CNNs based on loss functions defined with the potential vorticity variable $q$ will converge, but the model performance is generally poor, with noticeable degradations in the KE spectrum, spectral fluxes and similarity scores (not shown, but see next paragraph for a possibly anomalous case). On the other hand, it is possible to train models with loss functions involving the streamfunction $\psi$, and the resulting models trained from loss functions defined in terms of $\psi$ have comparable similarity scores to that reported here in Figure 6, but with minor improvements to the model distribution similarity scores. One possible reason could be that the streamfunction $\psi$, being a larger-scale field, is more amenable for extracting information from (cf. Yan et al., 2023). However, the predicted $S^q$ is not so statistically comparable to the diagnosed $S^q$, and the resulting spectral fluxes show large grid-scale dissipation in
<figure><img src="image_12.png" /><figcaption>**Figure 9.** Diagnostics from the full online hybrid model, where the associated CNNs are trained with hybrid models *without* any explicit small-scale dissipation active, using a loss function based on the potential vorticity mismatch given in Equation 14. (a) KE spectrum as Figure 4, and (b) spectral energy budget as Figure 5. The colored lines are from the 10 member ensemble averaged performance of the hybrid models, where the shading denotes the standard deviation of the ensemble.</figcaption></figure>

both KE and enstrophy, raising the question of whether one necessarily gets improvement in the spectral flux statistics even if the model state has reduced biases.

We close by noting that while it is true that the a posteriori model skill leaves a lot to be desired for online training with loss functions defined with the potential vorticity variable $q$ in hybrid models with explicit small-scale dissipation, some aspects are in fact improved for the same full online learning procedure but *without* explicit small-scale dissipation in the hybrid model. Figure 9 shows the resulting KE spectra and spectral fluxes. While the resulting models show some biases in the KE spectra in Figure 9a, a slightly reduced magnitude of energy cycle in Figure 9b, and the resulting similarity scores are slightly worse than the offline model ones shown in Figure 6, the model skill based on the considered metrics are otherwise not unreasonable. We do not have a convincing reason on why training without the explicit small-scale dissipation leads to better model skill for loss functions defined with the potential vorticity variable $q$, and we leave this as an observation for now. The same procedure but for approximately online training fails at the training stage; this is not an unreasonable expectation given the properties of models from the approximately online procedure at least for the problem considered here (e.g., poor lower layer potential vorticity distribution in Figure 6).

## 4.4. Imposition of Constraints

On the topic of stabilizers and constraints, one could build those into the mapping during the training stage, by modifying the CNN architecture or similar (e.g., Beucler et al., 2021), or modify the loss function accordingly. From the traditional optimization point of view, a constraint might be imposed by consider a loss function of the form $\mathcal{L} = \mathcal{L}_{\text{on}} + \lambda \mathcal{C}$, where $\mathcal{C}$ is a constraint of interest. If we wanted to force the CNN here to output a prediction with zero mean, we might consider take $\mathcal{C}$ to be Equation 9 and force the constraint by the method of Lagrange multipliers, which involves augmenting the control variable from $\theta$ to $(\theta, \lambda)$, where $\lambda$ would be interpreted as the Lagrange multiplier. On the other hand, one might consider imposing constraints approximately by viewing $\mathcal{C}$ instead as a penalization (e.g., $\mathcal{C} = \int (S^d)^2 dA$ to keep the differentiability of the cost function); $\lambda$ would then be regarded as a penalization weight. The two approaches might be regarded as a hard and soft constraint optimization (e.g., Beucler et al., 2021), the benefit of the latter being that the control variable $\theta$ is exactly as before with the same number of degrees of freedom. The same algorithms can be applied as normal but introducing subjectivity in the choice of $\lambda$, and the constraint is only enforced approximately. For the present case we considered imposing the zero-mean soft constraint as described above. The notable result is that we get similar results and stability properties as those reported thus far in this work, with the observation that the penalization approach appears to work better for the online than offline training (see Supporting Information S1).

On the other hand, implementing the hard-constraint so that the predicted potential vorticity tendency has zero mean (in the sense described in the text around Equation 9) as in the work of Ross et al. (2023) lead to approximately online and offline models that are noticeably more skillful in the model integrations. The notable
<figure><img src="image_13.png" /><figcaption>**Figure 10.** Distribution similarity scores for CNNs trained with an additional de-meaning final layer imposing the hard constraint in Equation 9, for potential vorticity *q*, both components of the velocity *u* and *v*, the streamfunction *ψ*, KE and enstrophy, depth-averaged KE flux, APE flux, APE generation, and the bottom drag (labeled as Friction). Subscripts 1 and 2 denote the upper and lower layer, for the (green) offline model (blue) full online model, and (orange) approximately online model. The solid line is for the ensemble average over the 10 members, and the shading denotes the standard deviation of the ensemble.</figcaption></figure>

result is in the analogous similarity scores shown in Figure 10. Compared to the analogous result in Figure 6, it is clear the skill in terms of the three approaches are much more comparable. The scores for the approximately online and offline approaches are elevated across the board, although it is also the case that the full online approach has the performance degraded as a result of the hard constraint, particularly in the model variable distributions. The analogous kinetic energy spectrum and the spectral fluxes are visually similar to that of Figures 4 and 5 (not shown).

The point we wish to stress, however, is that while the approximately online and offline approaches can be made to work or improved with extra add-ons (in this case the hard constraint strongly controls the accumulation of potential vorticity tendency), the model integrations from the full online approach were already stable and reasonably skillful in the metrics considered in this work, without the extra constraints.

**5. Conclusion and Discussion**

Building on the previous works demonstrating the efficacy of online learning in geophysical systems (e.g., Frezat et al., 2022; Kochkov et al., 2024; Ouala et al., 2024), known as temporal unrolling in other studies (e.g., List et al., 2022, 2025), the present work provides an examination between online learning approaches and the more standard offline approach for eddy parameterization of baroclinic turbulence. The online approach involves the underlying fluid dynamical model during the training stage to provide a mapping from the predicted quantity (the sub-grid forcing $S^q$ in this work) to the model variable (the potential vorticity $q$ in this work) at the next time-step, with the control variables (the collection of neural network parameters, denoted $\theta$ here) adjusted so that the loss function defined over some time window is minimized. The online learning procedure considered in this work is derivative-based, and at its core is linked to methodologies in optimization and data assimilation, namely adjoint optimization (e.g., Gunzburger, 2003) and Four-Dimensional Variational Assimilation (4DVAR; e.g., Kalnay, 2002) respectively. A plausible explanation for online learning is that the presence of the dynamical model forces the training to recognize that the target data provided has some temporal correlation/causality, exposing more information to the training procedure, and providing constraints of the solution space in which the trained model should reside. While the full adjoint-based online learning requires a differentiable dynamical model to enable the back-propagation calculations to obtain gradient information for the associated optimization problem, proposals exist to approximate the back-propagation and gradient calculations accordingly (e.g., List et al., 2025; Ouala et al., 2024); we considered here the approach put forward by List et al. (2025). One could also employ derivative-free methods (e.g., Iglesias et al., 2013; Kovachki & Stuart, 2019; Lopez-Gomez et al., 2022; Schneider et al., 2022).

For the present investigation, we extend on the previous work of Frezat et al. (2022) and Ouala et al. (2024) for online learning in a one-layer quasi-geostrophic system to a two-layer quasi-geostrophic system that supports baroclinic dynamics, following the same testing framework of Ross et al. (2023). We trained Convolutional Neural Networks (CNN) with loss functions defined with the sub-grid forcing $S^q$ as defined in Equation 6, which is the residual between the filtered high resolution model truth and the low resolution hybrid dynamical model.
The CNNs are trained via definitions of appropriate loss functions defined per fluid layer, and differ only in how they are trained, and without a hard constraint that the predicted $S^q$ should have zero mean (Equation 9) to test for model stability associated with the three training methodologies considered. Models are primarily tested in the a posteriori sense, where the model skill is evaluated on their long term statistics in prognostic mode, which is a harder but more relevant test relating to parameterization (e.g., Srinivasan et al., 2024). For statistical significance, the testing was performed over a 10 member ensemble of CNNs per training strategy. No batch normalization was included, partly for simplicity, and particularly because having it active leads to approximately online models that were particularly unstable during model integrations (7 out of the 10 models crashing). We do not impose controls on the back-propagation calculations since we have not encountered issues like those reported in List et al. (2022); this may be a combination of the problem set up and/or the small training window size that we took to be 10 model time-steps.

We find that models trained with a full online procedure lead to multiple benefits over the other two training approaches, at least in the metrics considered. The predicted sub-grid forcing for the full online model are closer in distribution to the high resolution model truth (see Figure 3). The full online models improve on the overall kinetic energy levels (Figure 4) and lead to more vigorous energy fluxes (Figures 5 and 6) in prognostic calculations. All the investigated CNNs seem to be learning some sort of backscatter that energizes the system and some semblance of the grid-scale dissipation (Figure 5), regardless of learning approach. However, the offline models leads to persistent biases in the kinetic energy spectrum at the grid-scale, while the online models are able to remove that bias (Figure 4). There is additional evidence to suggest that the full online approach can lead to model skill when used out-of-sample, in tests where the hybrid models have no explicit small-scale dissipation (Figure 8). All of the aforementioned properties of the online learning approach are achieved for a rather small window length of 10 time-steps (Figure 7), indicating that we are able to obtain convergent training and obtain notable improvements in the model skill, for a modest investment in computational resources. In our initial study comparing the training methodologies we have focused on the broad behavior of the trained CNNs, but did not consider the resulting variability the CNNs confer to the model beyond presenting calculations relating to the model variable probability distributions (the Wasserstein metric used in Equation 11). Variability and its dependence on data exposure and training methodology is also important, but is beyond the present scope of the work.

When a hard constraint (cf. Ross et al., 2023) is included as a means to control the potential vorticity accumulation in the system, the main beneficiaries are the approximately online and offline models. The three methodologies now display comparable model skill in the similarity scores (Figure 10, although it should be noted that the imposition of such a hard constraint appears to degrade the results from the full online approach, particularly in the model variable distributions. The point we want to emphasize is that while the approximately online and offline models can be made to work with additional constraints and/or stabilizers, the full online approach already works well in a stable fashion without the additional constraints and/or stabilizers, at least in the setting considered in the present work.

We have briefly explored whether including the model variables explicitly into the definition of the loss functions will lead to improvements in the model state variables themselves for online learning (e.g., Frezat et al., 2022; Ouala et al., 2024). We find that models trained on the loss functions defined with the potential vorticity $q$, while stable, tends to have notable issues with skill (except for a curious case reported in Figure 9). Models trained with loss functions defined by the streamfunction $\psi$ (related to $q$ via definition in Section 2.2) have reasonable skill in the similarity scores (cf. Figures 6 and 10 here). The issues with training convergence and/or model skill may be related to the instabilities associated with the back-propagation calculations, and could possibly be made more robust with stabilizations of the back-propagation calculations in the manner suggested by List et al. (2022), and further imposition of constraints or conservation laws via appropriate design of loss functions or CNN architecture (e.g., Beucler et al., 2021; Greydanus et al., 2019; Han et al., 2019; Ling et al., 2016). The approximately online approaches could also be further tuned and/or refined (e.g., taking more accurate/stable approximations for the gradient calculations; cf. Chen et al., 2018; Chattopadhyay & Hassanzadeh, 2023; Ouala et al., 2024). Loss functions defined in terms of statistical quantities (e.g., the energy spectrum) is in principle possible with online learning and another avenue to explore, although the computation costs might end up being a strong limiting factor. Loss functions beyond the usual root-mean-squared error is another possibility (e.g., negative Sobolev semi-norms, Lin et al., 2011; Thiffeault, 2012; Yan et al., 2023). Further investigation is required to draw
conclusive evidence on the benefit in utilizing the model variables to define loss function for online learning in the present baroclinic system.

For the present work, we presented results on a baroclinic turbulence model that was set up to be eddy dominated (cf. Figure 2), via the choice of the bottom drag time-scale $r_{ek}$ (see Table 1). It is known that strong zonal jets can form in rotating turbulence (cf. Ross et al., 2023; Srinivasan et al., 2024). Analogous investigations in the jet setting (by choosing a smaller value of $r_{ek}$) seem to suggest that online learning still lead to some benefits, either via learning from scratch in the jet regime, or utilizing the CNNs trained in the eddy regime reported here but applied to hybrid model a posteriori testing in the jet regime (i.e., applying the trained CNNs out-of-sample), both leading to moderately positive similarity scores. There are however additional subtleties with the training procedure such as the size of time window, and a strong dependence on the numerical model resolution (where the hybrid models' resolution at 64² appears to fundamentally limit the model performance, but is noticeably alleviated when employing models at 128² or higher). One key difference between the jet case and the eddy case is that there is the presence of a significant large-scale structure, and the hybrid models seem to have difficulty reducing the biases at the large-scale (as seen already in Figure 4 of this work). There is likely room for improvement with further training with procedures indicated in the previous paragraph, but the qualitative findings are otherwise not overly different to the eddy case presented here. We refer the interested reader to the Supporting Information S1 for sample results from the investigation on the jet case.

To close, the present work primarily aims to compare at a rather coarse level on the benefits and drawbacks of the online versus offline learning approach. Online learning is found to show better skill in a range of metrics, allow for a more flexible in the definition of the loss function, results in CNNs that demonstrate a degree of skill when utilized out-of-sample, and likely results in CNNs that lead to more numerically stable integrations in prognostic calculations over a wider range of conditions. While the related investment in computational resources is likely worthwhile, in some ways the work has opened up more questions than it has answered. Further work into design and tuning choices of the machine learning models, stabilization techniques, choice and design of loss functions, enforcement of physical properties, and further investigations into the approximate online learning procedure (which does not require a differentiable model) or derivative-free methodologies should be considered, but is beyond the scope of the present work.

## Data Availability Statement

This work utilizes *pyqq-JAX* (ver 0.8.1, Otness et al., 2023; Otness, 2024) and *Equinox* (Kidger & Garcia, 2021). The collection of code, model data and scripts used for generating the plots in this article are available through the repository at Yan and Mak (2025).

## References

* Beucler, T., Pritchard, M., Rasp, S., Ott, J., Baldi, P., & Gentine, P. (2021). Enforcing analytic constraints in neural networks emulating physical systems. *Physical Review Letters*, 126(9), 098302. https://doi.org/10.1103/PhysRevLett.126.098302

* Bolton, T., & Zanna, L. (2019). Applications of deep learning to ocean data inference and sub-grid parameterisation. *Journal of Advances in Modeling Earth Systems*, 11(1), 376–399. https://doi.org/10.1029/2018MS001472

* Brenowitz, N. D., & Bretherton, C. S. (2019). Spatially extended tests of a neural network parametrization trained by coarse-graining. *Journal of Advances in Modeling Earth Systems*, 11(8), 2728–2744. https://doi.org/10.1029/2019MS001711

* Chattopadhyay, A., & Hassanzadeh, P. (2023). Long-term instabilities of deep learning-based digital twins of the climate system: The cause and a solution. Retrieved from https://arxiv.org/abs/2304.07029

* Chen, R. T. Q., Rubanova, Y., Bettencourt, J., & Duvenaud, D. K. (2018). Neural ordinary differential equations. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, & R. Garnett (Eds.), *Advances in neural information processing systems* (Vol. 31). Curran Associates, Inc.

* Connolly, C., Barnes, E. A., Hassanzadeh, P., & Pritchard, M. (2023). Using neural networks to learn the Jet Stream forced response from natural variability. *Artificial Intelligence for the Earth Systems*, 2, e220094. https://doi.org/10.1175/AIES-D-22-0094.1

* Forget, G., Campin, J.-M., Heimbach, P., Hill, C. N., Ponte, R. M., & Wunsch, C. (2015). ECCO version 4: An integrated framework for non-linear inverse modeling and global ocean state estimation. *Geoscientific Model Development*, 8(10), 3071–3104. https://doi.org/10.5194/gmd-8-3071-2015

* Fox-Kemper, B., Adcroft, A. J., Böning, C. W., Chassignet, E. P., Curchitser, E. N., Danabasoglu, G., et al. (2019). Challenges and prospects in ocean circulation models. *Frontiers in Marine Science*, 6, 65. https://doi.org/10.3389/fmars.2019.00065

* Frezat, H., Le Sommer, J., Fablet, R., Balarac, G., & Lguensat, R. (2022). A posteriori learning for quasi-geostrophic turbulence parametrization. *Journal of Advances in Modeling Earth Systems*, 14(11), e2022MS003124. https://doi.org/10.1029/2022MS003124

* Gelbrecht, M., White, A., Bathiany, S., & Boers, N. (2023). Differentiable programming for Earth system modeling. *Geoscientific Model Development*, 16(11), 3123–3125. https://doi.org/10.5194/gmd-16-3123-2023

* Greydanus, S., Dzamba, M., & Yosinski, J. (2019). Hamiltonian neural networks. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d' Alché-Buc, E. Fox, & R. Garnett (Eds.), *Advances in neural information processing systems* (Vol. 32). Curran Associates, Inc.

## Acknowledgments

This study is a contribution to AI4PEX (Artificial Intelligence and Machine Learning for Enhanced Representation of Processes and Extremes in Earth System Models), a project supported by the European Union's Horizon Europe research and innovation programme under grant agreement no. 101137682, and also funded by the RGC General Research Fund 16304021. FEY and JM thanks James Maddison for technical comments for an earlier version of the work and article draft. The authors thank the referees for their comments that led to an improvement in the scientific content and presentation of the article.
Guillaumin, A. P., & Zanna, L. (2021). Stochastic-deep learning parameterization of ocean momentum forcing. *Journal of Advances in Modeling Earth Systems, 13*(9), e2021MS002534. https://doi.org/10.1029/2021MS002534

Gunzburger, M. D. (2003). *Perspectives in flow control and optimization*. SIAM.

Hallberg, R. (2013). Using a resolution function to regulate parameterizations of oceanic mesoscale eddy effects. *Ocean Modelling*, 72, 92–103.
https://doi.org/10.1016/j.oocemod.2013.08.007

Han, J., M. C., Ma, Z., & E. W. (2019). Uniformly accurate machine learning-based hydrodynamic models for kinetic equations. *Proceedings of the National Academy of Sciences of the United States of America* (Vol. 116(44), 21983–21991). https://doi.org/10.1073/pnas.1909854116

Heimbach, P., Fukumori, I., Hill, C. N., Ponte, R. M., Stammer, D., Campin, J., et al. (2019). Putting it all together: Adding value to the global ocean and climate observing systems with complete self-consistent ocean state and parameter estimates. *Frontiers in Marine Science*, 6, 55–2427. https://doi.org/10.3389/fmars.2019.00055

Iglesias, M. A., Law, K. J. H., & Stuart, A. M. (2013). Ensemble Kalman methods for inverse problems. *Inverse Problems*, 29(4), 045001. https://doi.org/10.1088/0266-5611/29/4/045001

IPCC. (2021). In V. Masson-Delmotte (Ed.), et al. (Eds.), *Climate change 2021: The physical science basis. Contribution of working group I to the sixth assessment report of the intergovernmental panel on climate change*. Cambridge University Press. https://doi.org/10.1017/9781009157896

Jansen, M. F., & Held, I. M. (2014). Parameterizing subgrid-scale eddy effects using energetically consistent backscatter. *Ocean Modelling*, 80, 36–48. https://doi.org/10.1016/j.oocemod.2014.06.002

Kalnay, E. (2002). *Atmospheric modeling, data assimilation and predictability* (Vol. 142). Cambridge University Press. https://doi.org/10.1017/CBO9780511802270

Kashinath, K., Mustafa, M., Albert, A., Wu, J., Jiang, C., Esmaeilzadeh, S., et al. (2021). Physics-informed machine learning: Case studies for weather and climate modelling. *Philos. Trans. Royal Soc. A*, 379(2194), 20200093. https://doi.org/10.1098/rsta.2020.0093

Kidger, P., & Garcia, C. (2021). Equinox: Neural networks in JAX via callable PyTrees and filtered transformations. In *Differentiable programming workshop at neural information processing systems 2021*. https://doi.org/10.48550/arXiv.2111.00254

Kochkov, D., Yuval, J., Langmore, I., Norgaard, P., Smith, J., Mooers, G., et al. (2024). Neural general circulation models for weather and climate. *Nature*, 632(8027), 1060–1066. https://doi.org/10.1038/s41586-024-07744-y

Kovachki, N. B., & Stuart, A. M. (2019). Ensemble Kalman inversion: A derivative-free technique for machine learning tasks. *Inverse Problems*, 35(9), 095005. https://doi.org/10.1088/1361-6420/ab1c3a

Lin, Z., Thiffeault, J., & Doering, C. R. (2011). Optimal stirring strategies for passive scalar mixing. *Journal of Fluid Mechanics*, 675, 465–476. https://doi.org/10.1017/S0022112011000292

Ling, J., Jones, R., & Templeton, J. (2016). Machine learning strategies for systems with invariance properties. *Journal of Computational Physics*, 318, 22–35. https://doi.org/10.1016/j.jcp.2016.05.003

List, B., Chen, L.-W., Bali, K., & Thuerey, N. (2025). Differentiability in unrolled training of neural physics simulators on transient dynamics. *Computer Methods in Applied Mechanics and Engineering*, 433, 117441. https://doi.org/10.1016/j.cma.2024.117441

List, B., Chen, L.-W., & Thuerey, N. (2022). Learned turbulence modelling with differentiable fluid solvers: Physics-based loss functions and optimisation horizons. *Journal of Fluid Mechanics*, 949, A25. https://doi.org/10.1017/jfm.2022.738

Lopez-Gomez, I., Christopoulos, C., Langeland Ervik, H. L., Dunbar, O. R. A., Cohen, Y., & Schneider, T. (2022). Training physics-based machine-learning parameterizations with gradient-free ensemble Kalman methods. *Journal of Advances in Modeling Earth Systems*, 14(8), e2022MS003105. https://doi.org/10.1029/2022MS003105

Maddison, J. R., Goldberg, D. N., & Goddard, B. D. (2019). Automated calculation of higher order partial differential equation constrained derivative information. *SIAM Journal on Scientific Computing*, 41(5), C417–C445. https://doi.org/10.1137/18M1209465

Mooers, G., Pritchard, M., Beucler, T., Ott, J., Yacalis, G., Baldi, P., & Gentine, P. (2021). Assessing the potential of deep learning for emulating cloud superparameterization in climate models with real-geography boundary conditions. *Journal of Advances in Modeling Earth Systems*, 13, e2020MS002385. https://doi.org/10.1002/2020MS002385

Orszag, S. A. (1971). On the elimination of aliasing in finite-difference schemes by filtering high-wavenumber components. *Journal of the Atmospheric Sciences*, 28(6), 1074. https://doi.org/10.1175/1520-0469(1971)028(1074:OTEOAI)2.0.CO;2

Otness, K. (2024). pyQG-JAX v0.8.1 [Software]. https://doi.org/10.5281/zenodo.10719906

Otness, K., Zanna, L., & Bruna, J. (2023). Data-driven multiscale modeling of subgrid parameterizations in climate models. Retrieved from https://arxiv.org/abs/2303.17496

Ouala, S., Chapron, B., Collard, F., Gaultier, L., & Fablet, R. (2024). Online calibration of deep learning sub-models for hybrid numerical modeling systems. *Communications Physics*, 7(1), 402. https://doi.org/10.1038/s42005-024-01880-7

Perezhogin, P., Zanna, L., & Fernandez-Granda, C. (2023). Generative data-driven approaches for stochastic subgrid parameterizations in an idealized ocean model. *Journal of Advances in Modeling Earth Systems*, 15(10), e2023MS003681. https://doi.org/10.1029/2023MS003681

Ross, A., Li, Z., Perezhogin, P., Fernandez-Granda, C., & Zanna, L. (2023). Benchmarking of machine learning ocean subgrid parameterizations in an idealized model. *Journal of Advances in Modeling Earth Systems*, 15(1), e2022MS003258. https://doi.org/10.1029/2022MS003258

Salmon, R. (1980). Baroclinic instability and geostrophic turbulence. *Geophysical & Astrophysical Fluid Dynamics*, 15(1), 167–211. https://doi.org/10.1080/03091928008241178

Sane, A., Reichl, B. G., Adcroft, A., & Zanna, L. (2023). Parameterizing vertical mixing coefficients in the ocean surface boundary layer using neural networks. *Journal of Advances in Modeling Earth Systems*, 15(10), e2023MS003890. https://doi.org/10.1029/2023MS003890

Sapienza, F., Bolibar, J., Schäfer, F., Groenke, B., Pal, A., Boussange, V., et al. (2024). Differentiable programming for differential equations: A review. Retrieved from https://arxiv.org/abs/2406.09699

Schneider, T., Stuart, A. M., & Wu, J. (2022). Ensemble Kalman inversion for sparse learning of dynamical systems from time-averaged data. *Journal of Computational Physics*, 470, 111559. https://doi.org/10.1016/j.jcp.2022.111559

Srinivasan, K., Chekroun, M. D., & McWilliams, J. C. (2024). Turbulence closure with small, local neural networks: Forced two-dimensional and β-plane flows. *Journal of Advances in Modeling Earth Systems*, 16(4), e2023MS003795. https://doi.org/10.1029/2023MS003795

Sun, Y. Q., Hassanzadeh, P., Alexander, M. J., & Kruse, C. G. (2023). Quantifying 3D gravity wave drag in a library of tropical convection-permitting simulations for data-driven parameterizations. *Journal of Advances in Modeling Earth Systems*, 15(5), e2022MS003585. https://doi.org/10.1029/2022MS003585

Thiffeault, J.-L. (2012). Using multiscale norms to quantify mixing and transport. *Nonlinearity*, 84(2), R1–R44. https://doi.org/10.1088/0951-7715/25/2/R1

Villani, C. (2008). *Optimal transport: Old and new*. Springer.

Yan, F. E., & Mak, J. (2025). Data collection for online learning in two-layer quasi-geostrophic model ver.2 [Dataset]. https://doi.org/10.5281/zenodo.15244460
Yan, F. E., Mak, J., & Wang, Y. (2023). On the choice of training data for machine learning of geostrophic mesoscale turbulence. *Journal of Advances in Modeling Earth Systems, 16*(2), e2023MS003915. https://doi.org/10.1029/2023MS003915

Yuval, J., & O'Gorman, P. A. (2020). Stable machine-learning parameterization of subgrid processes for climate modeling at a range of resolutions. *Nature Communications, 11*(1), 3295. https://doi.org/10.1038/s41467-020-17142-3

Zanna, L., & Bolton, T. (2021). Data-driven equation discovery of ocean mesoscale closures. *Geophysical Research Letters, 47*(17), e2020GL088376. https://doi.org/10.1029/2020GL088376