# Gradient-free online learning of subgrid-scale dynamics with neural
emulators

Hugo Frezat https://hrkz.github.io<EMAIL_ADDRESS>Université Grenoble-
Alpes, LEGI, Grenoble, France  Université Grenoble-Alpes, IGE, Grenoble,
France  Inria team Odyssey, Brest, France  Ronan Fablet  Inria team Odyssey,
Brest, France  Guillaume Balarac  Université Grenoble-Alpes, LEGI, Grenoble,
France  Institut Universitaire de France (IUF), Paris, France  Julien Le
Sommer  Université Grenoble-Alpes, IGE, Grenoble, France

(October 2023)

###### Abstract

In this paper, we propose a generic algorithm to train machine learning-based
subgrid parametrizations online, i.e., with a posteriori loss functions for
non-differentiable numerical solvers. The proposed approach leverage neural
emulators to train an approximation of the reduced state-space solver, which
is then used to allows gradient propagation through temporal integration
steps. The algorithm is able to recover most of the benefit of online
strategies without having to compute the gradient of the original solver. It
is demonstrated that training the neural emulator and parametrization
components separately with respective loss quantities is necessary in order to
minimize the propagation of some approximation bias.

##  I Introduction

Numerical models used for simulating the evolution of fluid flows generally
require subgrid scales (SGS) and physical parameterizations. These serve the
purpose of representing the impact of scales that remain unresolved by the
discretized solver as for instance subgrid turbulent processes [1, 2]. They
may also account for other physical phenomena that are not explicitly
accounted for by the partial differential equations. SGS and physical
parameterizations are generally considered to be key components of numerical
models in Earth System science [3, 4]. In these numerical models, they are
often rate-controlling components which largely govern the long term behavior
of model solutions [5, 6]. This is why significant efforts are concentrated on
the development of SGS and physical parameterizations in these communities.
Machine learning (ML) techniques have now been widely applied to SGS and
physical parametrizations for numerical simulations, as already discussed in
reviews [7, 8]. These ML-based components have been demonstrated in different
applications, including atmospheric circulation [9], oceanic circulation [10],
and magnetohydrodynamics models [11]. The ambition of this emerging field is
to design hybrid numerical models that improves upon existing ones by
combining physics-based PDE solvers and ML-based trainable components. In this
framework, well-constrained aspects are incorporated into the resolved
equations, while components that are less rigorously constrained or follow a
statistical behavior are represented by ML-based algorithms. When possible,
physical knowledge can also be incorporated in the ML-based component [12, 13,
14]. Hybrid numerical models are currently designed through an optimization
procedure that attempts to improve its ability to replicate the results of a
higher resolution model. The hybrid numerical models are typically run at
coarser resolution and the parameters of their ML components are optimized in
order to fit a given training objective. This objective is usually formulated
from pre-existing high-resolution, process resolving numerical models (that is
direct numerical simulation (DNS), or high-resolution version of the same
code) but also in principle be formulated from observations or lab
experiments. By essence, this optimization is very much related to the tuning
of model parameters, or to the estimation of state-dependent correction of
models.

Several recent results suggest that ML-based SGS parametrization give better
results when the ML components are trained with online strategies [15, 16, 17,
18, 19]. The training of ML components for hybrid numerical models is said to
use an online strategy if the evaluation of the training objective involves
integrating the hybrid numerical model for several time steps. Note that the
same idea has also been referred to as a posteriori [20], solver-in-the-loop
[21], differentiable [22, 23] or end-to-end in the ML community. In practice,
the strategy has been shown to lead to more stable simulations, with better
performance in a posteriori tests. However, the online strategy typically
requires to compute the derivative of the hybrid numerical model with respect
to the parameters of the ML components. These results suggest that, as in
other fields of physical science, the emerging paradign of differentiable
programming shows great potential for the design of hybrid numerical models
applied to the simulations of turbulent systems. Differentiable programming
(DP) is often combined with automatic differentiation (AD) which allows to
automatically obtain the derivatives of a chain of operations in order to be
used in an optimization procedure based on a gradient descent variant. AD
frameworks are typically implemented by specific libraries or programming
languages, and have seen many new development applications in physical
sciences [24, 25, 26, 27]. There has been a particularly active community
around the problem of turbulence modeling [28, 29, 30] and climate-
geoscientific modeling [31, 32, 33].

However, most numerical models are usually not easily differentiable. This
major practical issue makes the optimization of ML components with online
strategies difficult. For performance and historic reasons, legacy codes used
in production are written in Fortran or C programming languages, with
relatively low abstraction levels. This is the case in many physics community
including for Earth System models used for climate projections. These hand-
optimized numerical models are also relatively large codebases ( hundreds of
thousands of lines of codes) and would thus require lots of effort to be re-
implemented in a modern AD language. Generating adjoints automatically using
code analysis libraries is possible, but practically complicated to deploy and
lead to suboptimal codes (in particular in terms of memory footprint). In
practice, online learning strategies are thus not easily amenable with
production scale codes. Different approaches that includes the temporal aspect
provided by online learning have also been proposed and rely on gradient-free
optimization algorithms. While gradient-free optimization is a large subfield,
we identify two large families that have recently seen many applications to
training large ML models in the context of hybrid physical models. First,
methods adapted from data assimilation [34, 35] and inverse problems as for
instance the ensemble kalman inversion (EKI) [36, 37, 38]. More recently, many
studies applying gradient-free policy search in reinforcement learning (RL)
have been explored in particular for turbulence modeling [39, 40, 41, 42].
Overall, it is not clear whether which strategy performs best or should be
preferred in practical situations.

Alternatively, we here propose to leverage neural emulation as a strategy to
perform online supervised learning while avoiding the need for a
differentiable solver. The use of neural emulators also connects with some
algorithms in the field of simulation based inference which use neural
networks to solve inverse problems [43] and the field of data assimilation for
model error correction [44, 45]. Here, the neural emulator aims at
approximating the full solver in the (reduced) model state space. We can
rephrase the neural emulation as an approximation of the temporal evolution of
some PDE. We describe in this paper a generic algorithm which allows to train
jointly a ML-based SGS parametrization for a non-differentiable solver with an
online strategy and a neural emulator of the same solver. To overcome bias
separation between the two components, we introduce a two-step procedure and a
compensated loss function that isolates the different training objectives. The
algorithm is implemented and demonstrated to produce stable models on a
challenging subgrid parametrization problem. We also show that we are able to
recover most of the benefits of online strategies without having to compute
the gradient of the reduced resolution solver.

The paper is structured as follows: in Section 2, we outline the general issue
being addressed and its particular application to SGS modeling. Section 3
describes the two-step algorithm that allows training a SGS parametrization
with online strategy using a differentiable neural emulator of the target
solver. In Section 4, we describe our evaluation benchmark, consisting of a
two-dimensional quasi-geostrophic system subject to kinetic energy
backscatter. The performance of our approach is then illustrated in Section 5
and discussed in Section 6.

##  II Problem statement

As explained in the introduction, training a model online requires the ability
to compute the gradient of the loss function w.r.t. model parameters, which
involves solver calls and consequently also requires its gradient. Let us
define a typical regression task where one would seek to minimize a cost
between some prediction
$\mathring{\mathbf{z}}=\mathcal{M}(\mathbf{y}\,|\,\theta)$ and ground truth of
the same input space $\mathbf{z}$, i.e.

$\operatorname*{arg\,min}_{\theta}\mathcal{L}(\mathbf{z},\mathcal{M}(\mathbf{y}\,|\,\theta)).$
(1)

During the training process, the parameters of $\mathcal{M}$, $\theta$ are
optimized in order to minimize $\mathcal{L}$. The minimization algorithm
requires the computation of the gradient of the loss function w.r.t. the
trainable parameters $\theta$,

$\frac{\partial\mathcal{L}}{\partial\theta}(\mathbf{z},\mathcal{M}(\mathbf{y}\,|\,\theta))=\frac{\partial\mathcal{M}}{\partial\theta}(\mathbf{y}\,|\,\theta)\frac{\partial\mathcal{L}}{\partial\mathcal{M}}.$
(2)

In practice, $\mathcal{M}$ could be a simple operator from which the
analytical gradient would be computed beforehand and provided to the
optimization process. However, in differentiable solvers, we are interested in
the temporal evolution of a system
$E(\mathbf{y}):\mathbb{R}^{n}\rightarrow\mathbb{R}^{n}$ of a vector-valued
quantity of interest $\mathbf{y}(t)$. In a discretized formulation, the
operator $E$ is typically defined by a sequence of operations,

$\mathbf{y}(t+\Delta t)=E_{m}\circ\cdots\circ E_{1}(\mathbf{y}(t)).$ (3)

Now, if $\mathcal{M}\equiv E$, the required partial derivative of the system
is defined as a Jacobian in the minimization formulation (2),

$\frac{\partial E_{i}}{\partial\theta}=\begin{bmatrix}\dfrac{\partial
E_{i,1}}{\partial\theta_{1}}&\cdots&\dfrac{\partial
E_{i,1}}{\partial\theta_{n}}\\\ \vdots&\ddots&\vdots\\\ \dfrac{\partial
E_{i,n}}{\partial\theta_{1}}&\cdots&\dfrac{\partial
E_{i,n}}{\partial\theta_{n}}\end{bmatrix}$ (4)

for a single step of the temporal evolution operator $E$. Composing the
partial derivatives of the sequences of operators in $E$ can be done by
applying the chain rule,

$\frac{\partial E}{\partial\theta}=\frac{\partial(E_{m}\circ\cdots\circ
E_{1})}{\partial\theta}=\frac{\partial E_{m}}{\partial
E_{m-1}}\cdots\frac{\partial E_{2}}{\partial E_{1}}\frac{\partial
E_{1}}{\partial\theta}$ (5)

The gradient of $E$ can quickly become difficult to maintain by hand. Note
that when training NNs, it is required to have a scalar-valued loss function
$\mathcal{L}$, which means that its derivative is a gradient instead of a
Jacobian. The way partial derivatives are computed is called reverse-mode
differentiation [46], which here consists of a series of matrix-vector
multiplications starting from

$\frac{\partial
E_{1}}{\partial\theta}\frac{\partial\mathcal{L}}{\partial\mathcal{M}}:\mathbb{R}^{m\times
n}\rightarrow\mathbb{R}^{m}.$ (6)

Solving (5) can be as trivial as implementing a solver leveraging automatic
differentiation (AD) languages or libraries. However, as pointed out in the
introduction, this is a challenging task for large-scale solvers such as GCMs
that would requires a lot of development work and might lead to the vanishing
gradient problem [47] due to many function compositions.

###  II.1 The specific SGS modeling problem

Here, we are interested in modeling small-scale quantities arising in
differential equations restricting the grid resolution and thus the
performance of the simulations [48]. We assume an underlying differential
equation involving the time evolution of variables $\mathbf{y}(t)$ to be known
and defined by a direct operator $f(\mathbf{y})$. The aim is to solve an
equivalent approximation for reduced variables $\bar{\mathbf{y}}(t)$ such that

$\displaystyle\begin{cases}\dfrac{\partial\mathbf{y}}{\partial
t}=f(\mathbf{y}),\hskip 14.22636pt&\mathbf{y}\in\Omega\\\\[6.45831pt]
\dfrac{\partial\bar{\mathbf{y}}}{\partial
t}=g(\bar{\mathbf{y}})+\mathcal{M}(\bar{\mathbf{y}}),\hskip
14.22636pt&\bar{\mathbf{y}}\in\bar{\Omega}\\\\[6.45831pt]
\mathcal{T}(\mathbf{y})=\bar{\mathbf{y}}\end{cases}$ (7)

where $\bar{\Omega}\subset\Omega$, $g$ a reduced operator, $\mathcal{M}$ a SGS
model and $\mathcal{T}$ a projection that maps direct variables to reduced
ones. The objective is formulated as an inverse problem where operator
$\mathcal{M}$ has to be determined such that the evolution of the reduced
variables match the projection $\mathcal{T}(\mathbf{y})$ of the direct
variables $\mathbf{y}$. In most situations, we have $f=g$ with variables
existing on different spaces or dimensionalities. Note that this can be
applied to any type of partial differential equation without any loss of
generality. Within a learning framework, one states the identification of a
general reduced term

$\tau(\mathbf{y})=\mathcal{T}(f(\mathbf{y}))-g(\mathcal{T}(\mathbf{y}))\approx\mathcal{M}(\bar{\mathbf{y}}\,|\,\theta)$
(8)

where $\theta$ are trainable model parameters. Under the assumption that
$\mathcal{T}$ commutes with partial derivatives, the classical approach comes
to train a model $\mathcal{M}(\bar{\mathbf{y}}\,|\,\theta)$ as a functional
approximation of the missing term $\tau(\mathbf{y})$. The classical strategy
to optimize $\mathcal{M}$ is to generate a dataset
$\mathbb{D}=\\{\bar{\mathbf{y}}\\}\rightarrow\\{\tau(\mathbf{y})\\}$ that maps
input states $\bar{\mathbf{y}}$ to the SGS term $\tau(\mathbf{y})$ beforehand
and then finding optimal parameters $\theta$ offline, i.e., without any
knowledge of the dynamics $\partial\bar{\mathbf{y}}/\partial t$. Here, let us
refer to an offline loss $\mathcal{L}_{<}$, the corresponding minimization
problem (1) writes

$\operatorname*{arg\,min}_{\theta}\mathcal{L}_{<}=\operatorname*{arg\,min}_{\theta}\ell_{<}(\tau(\mathbf{y}),\mathcal{M}(\bar{\mathbf{y}}\,|\,\theta))$
(9)

where $\ell_{<}$ is an instantaneous scalar-valued loss function
$\ell_{<}:(\mathbb{R}^{n},\mathbb{R}^{n})\rightarrow\mathbb{R}$. In this
context, $\mathcal{M}$ is defined as a NN implemented in a standard deep
learning framework such as Pytorch or TensorFlow. This allows for the
computation of the gradient of the loss function, i.e., injecting (9) in (2),

$\frac{\partial\mathcal{L}_{<}}{\partial\theta}=\frac{\partial\ell_{<}}{\partial\mathcal{M}}\frac{\partial\mathcal{M}}{\partial\theta}(\bar{\mathbf{y}}\,|\,\theta).$
(10)

Using this formulation, the problem is simple to setup since both terms on the
r.h.s. are automatically determined using AD capabilities. However, while the
offline approach is efficient to accurately reproduces the statistical
properties of the missing term $\tau(\mathbf{y})$, it has been shown to lead
to unstable simulations in non-linear, chaotic dynamical systems such as those
found in turbulence modeling [49]. Improving the stability of the SGS model
can be done by increasing the size of the dataset, in particular by reducing
the inter-sample correlation [50] or adding physics-based terms in the loss
function [51]. As pointed out in [20], the offline training approach may not
be able to optimize a model based on relevant metrics in the context of a
dynamical system. An equivalent training approach tailored to this task,
called online, has already been explored successfully in turbulence modeling
[24, 25, 19]. Formally, [20] described the minimization problem (1) of the
online loss $\mathcal{L}_{>}$,

$\operatorname*{arg\,min}_{\theta}\mathcal{L}_{>}=\operatorname*{arg\,min}_{\theta}\ell_{>}(\\{\mathbf{y}(t)\\}_{t\in[t_{0},t_{1}]},\\{\Phi_{\theta}^{t}(\bar{\mathbf{y}}(t_{0}))\\}_{t\in[t_{0},t_{1}]})$
(11)

where $\Phi$ is the flow operator that advances the reduced variables in time
using the currently training model $\mathcal{M}$,

$\Phi_{\theta}^{t}(\bar{\mathbf{y}}(t_{0}))=\bar{\mathbf{y}}(t_{0})+\int_{t_{0}}^{t}g(\bar{\mathbf{y}}(t^{\prime}))+\mathcal{M}(\bar{\mathbf{y}}(t^{\prime})\,|\,\theta)\,\mathrm{d}t^{\prime}$
(12)

and $\ell_{>}$ is a temporal scalar-valued loss function
$\ell_{>}:(\mathbb{R}^{n\times M},\mathbb{R}^{n\times
M})\rightarrow\mathbb{R}$ defined on $M$ timesteps that discretizes the
temporal interval $[t_{0},t_{1}]$. Let us expand the gradient of the online
loss $\mathcal{L}_{>}$ from (2) at final time $t_{1}$,

$\frac{\partial\mathcal{L}_{>}}{\partial\theta}=\frac{\partial\ell_{>}}{\partial\Phi_{\theta}^{t}}\frac{\partial\Phi_{\theta}^{t}}{\partial\theta}(\bar{\mathbf{y}}(t_{0}))=\frac{\partial\ell_{>}}{\partial\Phi_{\theta}^{t}}\left(\int_{t_{0}}^{t}\frac{\partial
g(\bar{\mathbf{y}}(t^{\prime}))}{\partial\theta}+\frac{\partial\mathcal{M}(\bar{\mathbf{y}}(t^{\prime}))}{\partial\theta}\,\mathrm{d}t^{\prime}\right).$
(13)

This process requires the ability to compute the partial derivatives of the
reduced solver $g$ w.r.t. the trained model parameters $\theta$, i.e.
$\partial g/\partial\theta$. Here we are interested in determining this
quantity such that $\tau$ can be learnt online without re-implementing $g$
using an AD framework.

##  III Coupled SGS-emulation

Building on [45], we propose an algorithm that is able to jointly train a
differentiable emulator of $g$ and its corresponding SGS term. Let us
introduce a differentiable emulator $\mathcal{E}$ of $g$

$\mathcal{E}(\bar{\mathbf{y}})\approx g(\bar{\mathbf{y}}),\hskip
2.84526pt\text{s.t. }\frac{\partial\mathcal{E}}{\partial\theta}\text{ is
known}.$ (14)

This differentiability condition is automatically fulfilled if $\mathcal{E}$
is implemented using AD capabilities. It is possible to define $\mathcal{E}$
as a trainable neural emulator, which involves training its parameters
$\Theta$ simultaneously with the SGS model parameters $\theta$. Substituting
into (11),

$\operatorname*{arg\,min}_{\Theta,\theta}\ell_{>}(\\{\mathbf{y}(t)\\}_{t\in[t_{0},t_{1}]},\\{\varphi_{\Theta,\theta}^{t}(\bar{\mathbf{y}}(t_{0})\\}_{t\in[t_{0},t_{1}]})$
(15)

where the coupled flow operator that jointly learns parameters $\Theta$ and
$\theta$ is given by

$\varphi_{\Theta,\theta}^{t}(\bar{\mathbf{y}}(t_{0}))=\bar{\mathbf{y}}(t_{0})+\int_{t_{0}}^{t}\mathcal{E}(\bar{\mathbf{y}}(t^{\prime})\,|\,\Theta)+\mathcal{M}(\bar{\mathbf{y}}(t^{\prime})\,|\,\theta)\,\mathrm{d}t^{\prime}.$
(16)

We recall here that the final goal of this algorithm is to provide a stable
SGS model $\mathcal{M}$ that will be used for forward simulations with the
reduced solver $g$. However, with the coupled formulation (15), the model
$\mathcal{M}$ is not explicitly restricted to represent SGS dynamics only.
Indeed the algorithm jointly optimizes $\Theta$ and $\theta$ such that
$\mathcal{E}+\mathcal{M}$ approximates the direct solver $f$. There is
therefore no guarantee that $\theta$ will be optimal for $\mathcal{M}$ to
represent the SGS term.

###  III.1 Joint neural SGS-emulation learning

It is possible to penalize both the neural emulator and the SGS model by
composing the loss function such that
$\ell=\ell_{\mathcal{E}}(\mathcal{E}(\bar{\mathbf{y}}\,|\,\Theta))+\ell_{\mathcal{M}}(\mathcal{M}(\bar{\mathbf{y}}\,|\,\theta))$
using filtering operations but their respective definitions may not be simple
to implement and might not include inter-scale transfers. Instead, we propose
a two-step online formulation that trains $\mathcal{E}$ and $\mathcal{M}$
sequentially,

Step 1 (emulator):
$\displaystyle\operatorname*{arg\,min}_{\Theta}\ell_{>}^{\mathcal{E}}(\\{\bar{\mathbf{y}}(t)\\}_{t\in[t_{0},t_{1}]},\\{\varphi_{\Theta}^{t}(\bar{\mathbf{y}}(t_{0}))\\}_{t\in[t_{0},t_{1}]}),$
(17) Step 2 (model):
$\displaystyle\operatorname*{arg\,min}_{\theta}\ell_{>}^{\mathcal{M}}(\\{\mathbf{y}(t)\\}_{t\in[t_{0},t_{1}]},\\{\varphi_{\theta}^{t}(\bar{\mathbf{y}}(t_{0}))\\}_{t\in[t_{0},t_{1}]}).$
(18)

These two steps are formulated using an online learning strategy, but (17)
minimizes a loss for an integration of the reduced solver $g$ while (18) is
equivalent to the classical online SGS formulation (11) with
$\mathcal{E}\approx g$. Using these sequential steps, it is straightforward to
show that a differentiable version of the reduced solver is not required.
Indeed, let us factor the modified flows $\varphi$ appearing on the right-hand
side of (17) and (18) respectively,

$\displaystyle\varphi_{\Theta}^{t}(\bar{\mathbf{y}}(t_{0}))$
$\displaystyle=\bar{\mathbf{y}}(t_{0})+\int_{t_{0}}^{t}\mathcal{E}(\bar{\mathbf{y}}(t^{\prime})\,|\,\Theta)\,\mathrm{d}t^{\prime},$
(19) $\displaystyle\varphi_{\theta}^{t}(\bar{\mathbf{y}}(t_{0}))$
$\displaystyle=\bar{\mathbf{y}}(t_{0})+\int_{t_{0}}^{t}\mathcal{E}(\bar{\mathbf{y}}(t^{\prime}))+\mathcal{M}(\bar{\mathbf{y}}(t^{\prime})\,|\,\theta)\,\mathrm{d}t^{\prime}.$
(20)

Then, expanding again the gradient of the online losses using (2) yields,

$\displaystyle\frac{\partial\ell_{>}^{\mathcal{E}}}{\partial\varphi_{\Theta}^{t}}\frac{\partial\varphi_{\Theta}^{t}}{\partial\Theta}(\bar{\mathbf{y}}(t_{0}))$
$\displaystyle=\frac{\partial\ell_{>}^{\mathcal{E}}}{\partial\varphi_{\Theta}^{t}}\left(\int_{t_{0}}^{t}\frac{\partial\mathcal{E}(\bar{\mathbf{y}}(t^{\prime}))}{\partial\Theta}\,\mathrm{d}t^{\prime}\right),$
(21)
$\displaystyle\frac{\partial\ell_{>}^{\mathcal{M}}}{\partial\varphi_{\theta}^{t}}\frac{\partial\varphi_{\theta}^{t}}{\partial\theta}(\bar{\mathbf{y}}(t_{0}))$
$\displaystyle=\frac{\partial\ell_{>}^{\mathcal{M}}}{\partial\varphi_{\theta}^{t}}\left(\int_{t_{0}}^{t}\frac{\partial\mathcal{E}(\bar{\mathbf{y}}(t^{\prime}))}{\partial\theta}+\frac{\partial\mathcal{M}(\bar{\mathbf{y}}(t^{\prime}))}{\partial\theta}\,\mathrm{d}t^{\prime}\right)$
(22)

with $\mathcal{E}$, $\mathcal{M}$ are arbitrary NNs implemented using AD and
$\ell_{>}^{\mathcal{E}}$, $\ell_{>}^{\mathcal{M}}$ loss functions also defined
within the machine learning framework, supporting AD by construction.

###  III.2 Neural emulator error compensations

We know that in theory, the two-step algorithm would be equivalent to the
classical online SGS learning scheme with a differentiable solver if and only
if $\mathcal{E}\equiv g$. In practice however, the neural emulator will only
represent an approximation of the reduced solver, i.e., $\mathcal{E}\approx
g$. The performance of the SGS training (18) thus strongly relies on the
emulator’s ability to reproduce the reduced dynamics on a given horizon
$[t_{0},t]$. The instantaneous emulation mismatch between the differentiable
solver and the differentiable solver can be expressed as

$\varepsilon(\bar{\mathbf{y}})=g(\bar{\mathbf{y}})-\mathcal{E}(\bar{\mathbf{y}}).$
(23)

Substituting in the two-step SGS training (18) with the reduced solver on the
left-hand side and the emulation mismatch on the right-hand side, we obtain

$\displaystyle\ell_{>}^{\mathcal{M}}$
$\displaystyle(\\{\mathbf{y}(t)\\}_{t\in[t_{0},t_{1}]},\\{\varphi_{\theta}^{t}(\bar{\mathbf{y}}(t_{0})\\}_{t\in[t_{0},t_{1}]})$
$\displaystyle=\ell_{>}^{\mathcal{M}}(\\{\mathbf{y}(t_{0})+\int_{t_{0}}^{t}f(\mathbf{y}(t^{\prime}))\,\mathrm{d}t^{\prime}\\}_{t\in[t_{0},t_{1}]},$
$\displaystyle\hskip
28.45274pt\\{\bar{\mathbf{y}}(t_{0})+\int_{t_{0}}^{t}\mathcal{E}(\bar{\mathbf{y}}(t^{\prime}))+\mathcal{M}(\bar{\mathbf{y}}(t^{\prime})\,|\,\theta)\,\mathrm{d}t^{\prime}\\}_{t\in[t_{0},t_{1}]})$
$\displaystyle=\ell_{>}^{\mathcal{M}}(\\{\bar{\mathbf{y}}(t_{0})+\int_{t_{0}}^{t}g(\bar{\mathbf{y}}(t^{\prime}))+\tau(\mathbf{y}(t^{\prime}))\,\mathrm{d}t^{\prime}\\}_{t\in[t_{0},t_{1}]},$
$\displaystyle\hskip
28.45274pt\\{\bar{\mathbf{y}}(t_{0})+\int_{t_{0}}^{t}g(\bar{\mathbf{y}}(t^{\prime}))+\mathcal{M}(\bar{\mathbf{y}}(t^{\prime})-\varepsilon(\bar{\mathbf{y}}(t^{\prime}))\,|\,\theta)\,\mathrm{d}t^{\prime}\\}_{t\in[t_{0},t_{1}]}).$
(24)

It is clear here that the emulation mismatch $\varepsilon$ will be absorbed by
the SGS model $\mathcal{M}$ during its training phase, since $\mathcal{E}$ is
already fixed. In order to optimize $\mathcal{M}$ against $\tau$ using this
standard online loss function, it is important that the emulation mismatch
remain small. Controlling this mismatch is challenging in particular due to
the temporal integration. Recall that the loss function
$\ell_{>}^{\mathcal{M}}$ is defined on a sequence of states of the system
$\bar{\mathbf{y}}(t)$ but never involves the missing SGS term $\tau$ that
$\mathcal{M}$ is trying to learn. As a conservative solution, we propose to
transform the loss function $\ell_{>}^{\mathcal{M}}$ so that it operates on
the SGS term of the temporal states,

$\displaystyle\ell_{>}^{\mathcal{M}_{c}}$
$\displaystyle(\\{\mathbf{y}(t)\\}_{t\in[t_{0},t_{1}]},\\{\varphi_{\theta}^{t}(\bar{\mathbf{y}}(t_{0})\\}_{t\in[t_{0},t_{1}]})$
$\displaystyle=\ell_{>}^{\mathcal{M}}(\\{\tau(\mathbf{y}(t))\\}_{t\in[t_{0},t_{1}]},\\{\mathcal{M}(\varphi_{\theta}^{t}(\bar{\mathbf{y}}(t_{0}))\\}_{t\in[t_{0},t_{1}]}).$
(25)

Penalizing these terms in the loss function removes any ambiguities w.r.t. the
emulation mismatch $\varepsilon$ while maintaining the online behavior of the
learning strategy, since these quantities are still evaluated over a temporal
horizon $[t_{0},t]$. The discretized version of the two-step algorithm,
suitable for numerical applications (detailed below), is outlined in Algorithm
1.

Algorithm 1 Two-step online SGS learning algorithm

$\\{\bar{\mathbf{y}}(t)\\}_{t\in[t_{0},t_{1}]}$ // precomputed using reduced
solver $g$

$\\{\mathbf{y}(t),\tau(\mathbf{y}(t))\\}_{t\in[t_{0},t_{1}]}$ // precomputed
using direct solver $f$

$\mathcal{E}(\Theta\sim\mathcal{N})$ // defined emulator NN architecture

$\mathcal{M}(\theta\sim\mathcal{N})$ // defined model NN architecture

$\mathcal{M}(\theta)$ // optimized SGS parameters $\theta$

$\mathcal{M}(\theta)\leftarrow\textsc{model\\_step}(\mathcal{E}(\Theta)\leftarrow\textsc{emulator\\_step})$
// learn $\mathcal{M}$ from the trained emulator $\mathcal{E}$

function emulator_step

while Not converged do // run epochs until convergence

for $t\leftarrow t_{0},t_{1}$ do

$\hat{\mathbf{y}}(t)\leftarrow\bar{\mathbf{y}}(t)$ // Initialize from reduced
states $\bar{\mathbf{y}}$

for $i\leftarrow 0,N_{\mathrm{steps}}$ do // Loop over temporal horizon

$\hat{\mathbf{y}}(t+(i+1)\delta
t)\leftarrow\textsc{ode}(\mathcal{E}(\hat{\mathbf{y}}(t+i\delta
t)\,|\,\Theta))$ // Integrate using ODE solver

end for

$\mathcal{L}_{>}^{\mathcal{E}}\leftarrow\ell_{>}^{\mathcal{E}}(\\{\bar{\mathbf{y}}(t+i\delta
t)\\}_{i\in[1,M]},\\{\hat{\mathbf{y}}(t+i\delta t)\\}_{i\in[1,M]}$ //
classical ”online” loss

$\nabla_{\Theta}\mathcal{L}_{>}^{\mathcal{E}}\leftarrow\textsc{backprop}(\mathcal{L}_{>}^{\mathcal{E}})$
// compute gradient of differentiable loss $\ell_{>}^{\mathcal{E}}$

$\Theta\leftarrow\textsc{optimize}(\nabla_{\Theta}\mathcal{L}_{>}^{\mathcal{E}})$
// apply gradient descent

end for

end while

end function

function model_step($\mathcal{E}(\Theta)$)

while Not converged do // run epochs until convergence

for $t\leftarrow t_{0},t_{1}$ do

$\hat{\mathbf{y}}(t)\leftarrow\mathcal{T}(\mathbf{y}(t))$ // Initialize from
projected direct states $\mathbf{y}$

for $i\leftarrow 0,N_{\mathrm{steps}}$ do // Loop over temporal horizon

$\hat{\mathbf{y}}(t+(i+1)\delta
t)\leftarrow\textsc{ode}(\mathcal{E}(\hat{\mathbf{y}}(t+i\delta
t))+\mathcal{M}(\hat{\mathbf{y}}(t+i\delta t)\,|\,\theta))$ // Integrate using
ODE solver

end for

$\mathcal{L}_{>}^{\mathcal{M}_{c}}\leftarrow\ell_{>}^{\mathcal{M}_{c}}(\\{\tau(\mathbf{y}(t+i\delta
t))\\}_{i\in[1,M]},\\{\mathcal{M}(\hat{\mathbf{y}}(t+i\delta
t)\,|\,\theta)\\}_{i\in[1,M]}$ // temporal ”offline” loss

$\nabla_{\theta}\mathcal{L}_{>}^{\mathcal{M}_{c}}\leftarrow\textsc{backprop}(\mathcal{L}_{>}^{\mathcal{M}_{c}})$
// compute gradient of compensated loss $\ell_{>}^{\mathcal{M}_{c}}$

$\theta\leftarrow\textsc{optimize}(\nabla_{\theta}\mathcal{L}_{>}^{\mathcal{M}_{c}})$
// apply gradient descent

end for

end while

end function

##  IV Application: quasi-geostrophic dynamics

In turbulent flows, the energy cascade drives energy from the large scales to
the small scales until molecular viscous dissipation (forward-scatter), but
the inverse transfer called backscatter in which energy is transferred from
the small scales back to the large scales [52] is also in play, particularly
for geophysical flows. This is explained by the relative dominance of the
Coriolis force which creates vortical structures that appear two-dimensional.
Historically, developing SGS models that account for backscatter is a
challenging task [53, 54, 55]. Indeed, an overprediction of backscatter that
can not be compensated by eddy-viscosity will lead to an accumulation of
small-scale energy causing simulations to become numerically unstable. In two-
dimensional flows, we observe a dual cascade composed of “forward” enstrophy
and “inverse” energy, in a statistical sense. As a consequence, a large number
of SGS models have been proposed in particular for geophysical flows (see [56]
for a review) with well-documented configurations and performance metrics
[57]. SGS modeling is also a key issue for the simulation of ocean and
atmosphere dynamics because of the large range of motions involved [58, 59,
60, 61]. As a case study framework, we consider barotropic QG flows. While
providing an approximate yet representative system for rotating stratified
flows found in the atmosphere and ocean dynamics, it involves relatively
complex SGS features that make the learning problem non-trivial. As such, QG
flows are regarded as an ideal playground to explore and assess the relevance
of machine learning strategies for SGS models in geophysical turbulence. The
governing equations of direct vorticity $\omega$ for the QG equations with
bottom topography $\eta$ are

$\displaystyle\dfrac{\partial\omega}{\partial
t}+\mathcal{N}(\omega+\eta,\psi)$
$\displaystyle=\nu\nabla^{2}\omega-\mu\omega-\beta\partial_{x}\psi+F_{\omega},$
(26) $\displaystyle\nabla^{2}\psi$ $\displaystyle=-\omega$ (27)

where
$\mathcal{N}(\omega+\eta,\psi)=\partial_{x}\psi\partial_{y}(\omega+\eta)-\partial_{y}\psi\partial_{x}(\omega+\psi)$
is the non-linear vorticity advection, $\nu$ is the viscosity, $\mu$ is a
large-scale drag coefficient, $\beta$ is the Earth rotation vector
approximation by beta-plane and $F_{\omega}$ is an additional source term.
Now, the derivation of the reduced system for QG dynamics follows the same
procedure that is typically used for fluid dynamics turbulence modeling (also
called large-eddy simulation). The projection operator $\mathcal{T}$ at
spatial coordinate $\mathbf{x}$ is here given as a discretization (or coarse-
graining) $\mathcal{D}:\Omega\rightarrow\bar{\Omega}$ and the convolution of
$\omega$ with a kernel function $G(\mathbf{x})$ [62],

$\bar{\omega}(\mathbf{x})=\mathcal{D}\left(\int
G(\mathbf{x}-\mathbf{x}^{\prime})\omega(\mathbf{x}^{\prime})\,\mathrm{d}\mathbf{x}^{\prime}\right).$
(28)

We can then derive the equations which govern the evolution of reduced
vorticity $\bar{\omega}$,

$\displaystyle\dfrac{\partial\bar{\omega}}{\partial
t}+\mathcal{N}(\bar{\omega}+\bar{\eta},\bar{\psi})$
$\displaystyle=\nu\nabla^{2}\bar{\omega}-\mu\bar{\omega}-\beta\partial_{x}\bar{\psi}+\bar{F}_{\omega}$
$\displaystyle+\underbrace{\mathcal{N}(\bar{\omega}+\bar{\eta},\bar{\psi})-\overline{\mathcal{N}(\omega+\eta,\psi)}}_{\tau_{\omega}}$
(29)

where $\tau_{\omega}$ is the SGS term. In this context, $\tau_{\omega}$ can
not be determined from the reduced variables because of the non-linear
interactions of small-scale dynamics
$\overline{\mathcal{N}(\omega+\eta,\psi)}$ and must thus be statistically
modeled in order to close equation (29).

Figure 1: Direct vorticity fields $\omega$ (left), reduced vorticity fields
$\bar{\omega}$ (middle) and SGS term $\tau_{\omega}$ (right) at the end of a
testing horizon (3000 temporal iterations) from the direct simulation using
$f$.

###  IV.1 Simulation setup

In order to study the learning problem, we solve equations (26)-(27) using a
pseudospectral code with full 3/2 dealiasing [63] and a classical fourth-order
explicit Runge Kutta time advancement. The system is defined in a squared
domain $\Omega,\bar{\Omega}\in[-\pi,\pi]^{2}$, or domain length $L=2\pi$
discretized with a Fourier basis, i.e. double-periodic boundary conditions on
$N$ and $\bar{N}$ grid points with uniform spacing and a grid ratio
$\Delta^{\prime}=N/\bar{N}>1$ between direct and reduced solver resolutions.
We extract the SGS term with spatial filtering (28) using a spectral cut-off
filter $G$ and sample non-residual quantities along the simulations, i.e.
$\\{\omega\rightarrow\tau_{\omega}\\}$ is obtained from direct solver $f$ and
$\\{\bar{\omega}\\}$ is obtained from reduced solver $g$. We energize the
system using a wind-forced configuration representative of mesoscale oceanic
simulation [64]. We produce an equilibrium solution from a slowly varying in
time circular source at large scale $k$ with steady enstrophy rate injection
$\langle F_{\omega}(t)^{2}\rangle=3$ such that

$F_{\omega}(t)=\cos(4y)-\cos(4x+\pi\sin(1.5t)).$ (30)

We also use unaligned sinusoidal bumps topography [65] in order to create
interesting dnamics with significant nonzonal components in jets induced by
the Coriolis force. Stationary turbulent states are obtained from random
Gaussian initialization $\omega(t=0)\sim\mathcal{N}(0,1)$ at moderate
wavenumbers $k\in[10,32]$, followed by a 500000 iterations spin-up on a
smaller grid ($1024^{2}$) and energy propagation to the smallest scales of the
direct grid ($2048^{2}$) in about 25000 iterations. We show examples of
vorticity fields and SGS term in Fig. 1.

###  IV.2 Learning setup

To generate the corresponding datasets suitable for online learning, we need
continuous trajectories and thus subsample at each iteration from $g$ for (17)
and one state every $\Delta^{\prime}$ iteration performed by $f$ for (18) so
that these states directly correspond to one iteration performed by $g$. For
the SGS model, we use a simple standard CNN with $D=10$ convolution layers
with kernels $K$ of size $5$ and $64$ filters each, followed by non-linear
Relu activations. This architecture has been shown to be able to learn
extremely accurate SGS dynamics in [20]. For the large-scale dynamics
represented by the neural emulator $\mathcal{E}$, we explore a state-of-the-
art neural architecture and a simpler CNN in order to assess the impact of the
emulation quality on the SGS learning:

  * •

Dilated Residual Network (DRN). The architecture combines residual networks
with the encode-process-decode paradigm [66] and dilated convolutions,
recently applied to different types of flows [67].

  * •

Convolutional Neural Network (CNN). The architecture is equivalent to the CNN
used for the SGS model except that we use $D=7$ convolution layers.

Note that for each of these architectures, input boundaries are replicated
periodically given the geometry of the domain. More involved architecture
could further improve the overall performance and might be interesting to
explore. However, we may point out that the goal is not to design an optimal
NN-based architecture but rather to evaluate the training strategy at the same
SGS computational cost during evaluation stage. Concerning the functional loss
$\ell$, we restricted ourselves to simple MSE both for the neural emulator and
the SGS model,

$\ell_{<}(\cdot,\hat{\cdot})=\ell_{>}^{\mathcal{E}}(\cdot,\hat{\cdot})=\ell_{>}^{\mathcal{M}_{c}}(\cdot,\hat{\cdot})=\frac{1}{M}\sum_{i=1}^{M}||\cdot_{i}-\,\hat{\cdot}_{i}\,||_{2}^{2}.$
(31)

##  V Results

We generate simulation data and train corresponding models using parameters
summarized in Table. 1.

Table 1: Parameters of the numerical setup for simulation (top) and learning (bottom). Note that reduced systems use the same parameters, except for grid resolution $\bar{N}$ and timestep $\overline{\Delta t}$, obtained from the grid ratio $\Delta^{\prime}$. The quantities are given in numerical (dimensionless) as directly used in the solver for reproducibility. Simulation parameter |  |
---|---|---
Length of the domain | $L$ | $2\pi$
Linear drag | $\mu$ | $2\times 10^{-2}$
Kinematic viscosity | $\nu$ | $1.025\times 10^{-5}$
Rossby parameter | $\beta$ | $2.195\times 10^{2}$
Topography | $\eta$ | $\sin(3y)+\cos(3x)$
Timestep | $\Delta t$ | $10^{-4}$
Number of grid points | $N$ | 2048
Grid ratio | $\Delta^{\prime}$ | 16
Learning parameter |  |
offline batch size / online iterations | $M$ | $25$
Training trajectories |  | $10$
Testing trajectories |  | $5$
Samples per trajectory |  | $3000$

The objective here being to design a stable and efficient SGS model, we
compare the results of our two-step algorithm with several baselines.

Online model (reference). The main objective here is to reproduce the
performance of the online model implemented using a differentiable solver
using our two-step strategy that does not require differentiability.

Offline model. We train a SGS model using the classical offline strategy with
the same CNN architecture, as described above. This model is known to be
unstable and motivates the advantages of the online strategy.

Physical model. As a common baseline, we also use the Smagorinsky model [68],
defined as an eddy coefficient proportional to the resolved strain rate
$\bar{S}$ with the dynamic procedure proposed by [69, 70] to compute $C_{s}$
and apply spatial averaging with positive clipping [71, 50] in order to avoid
locally negative constants $C_{s}(x,y)<0$, i.e. ensuring that the models are
purely diffusive and $\nu_{e}\geq 0$,

$\tau_{\omega}=\nabla\cdot[\underbrace{(C_{s}\Delta)^{2}|\bar{S}|}_{\nu_{e}}\nabla\bar{\omega}].$
(32)

###  V.1 SGS model evaluation

Figure 2: Reduced vorticity fields $\bar{\omega}$ after a short simulation run
of 500 temporal iterations comparing reference filtered DNS (top-left) and
differentiable online strategy (bottom-left) with baselines (top-middle,right)
and emulation-based models (bottom-middle,right).

We first run a new simulation from different random initial conditions with
our references, baselines and SGS models trained using the two-step algorithm
using both the CNN and DRN architectures. Recall that evaluations are now run
with the reduced solver $g$ and a SGS model $\mathcal{M}$, except for the
direct reference $\mathcal{M}_{\mathrm{DNS}}$ that uses direct solver $f$.

We show in Fig. 2 the final vorticity fields $\bar{\omega}$ after 500 reduced
temporal iterations. It is first clear here that the offline learning strategy
does not lead to a stable model in this configuration, and that numerical
blowup characterized by a point-by-point noisy vorticity field happens
quickly. On the opposite, the physical model is too diffusive, completely
removing small-scale features and the associated dynamics.

We can also validate that the online strategy is indeed a stable and accurate
model to represents the SGS dynamics in QG systems, as it is remarkably close
to the DNS.

Concerning the neural emulators, they both remain stable over this temporal
horizon. However, we can already see that model trained from the simpler CNN
is already starting to differ from the reference.

Figure 3: Domain-averaged invariants energy $E(t)$ (left) and enstrophy $Z(t)$
(right) after a short simulation run of 500 temporal iterations.

To understand the time-evolution of the flow statistics, we look at the
invariants of the QG equations in Fig. 3, in the limit of inviscid flows [72]
for energy

$E(t)=\frac{1}{2}\int\mathbf{u}(t)^{2}\,\mathrm{d}r$ (33)

and enstrophy

$Z(t)=\frac{1}{2}\int\omega(t)^{2}\,\mathrm{d}r.$ (34)

We can confirm previous observations concerning accumulation of enstrophy from
the offline model, exponentially growing until blowup. We also clearly observe
the over-dissipation of enstrophy for the physical model. Both neural
emulators show enstrophy evolution relatively close to the references with
DNR-based model being almost equivalent to the reference online model. While
small scales are well represented by the models trained from the neural
emulators, it is apparent that they are not as efficient on the large scales,
as seen by the constant increase of energy for the CNN-based model. The DNN-
based model does not accumulate energy but performs slightly worse than the
physical model.

Figure 4: Time-averaged enstrophy spectrum (left) and final probability
distribution of the enstrophy transfer term (right) for 5 trajectories of
longer simulation runs of 3000 temporal iterations.

Integrating over a longer temporal horizon composed of 3000 iterations, we can
inspect the statistical behavior of the different models in Fig. 4. First, we
verify that the enstrophy spectrum

$Z(k)=\int_{|\mathbf{k}|=k}|\omega(\mathbf{k})|^{2}\,\mathrm{d}S(\mathbf{k}).$
(35)

and integrated modeled transfers

$\int\bar{\omega}\tau_{\omega}\,\mathrm{d}S=\int(\bar{\mathbf{u}}\,\bar{\omega}-\overline{\mathbf{u}\,\omega})\cdot\nabla\bar{\omega}\,\mathrm{d}S=-\int
T_{Z}\,\mathrm{d}S.$ (36)

of the online (differentiable reference) correctly reproduces the dynamics of
the DNS. On the left, we note a good agreement of the physical model at the
smallest wavenumbers but excessive dissipation at the largest wavenumbers. The
DRN-based model is extremely accurate on the largest wavenumbers which
corresponds with the terms used in the loss function during training, both
$\omega$ and $\tau_{\omega}$ acting mostly on the smallest scales. The DRN-
based model is however not as accurate on the smallest wavenumbers. This might
be improved by adding an additional large-scale term to the loss function such
as velocity $\mathbf{u}$ or streamfunction $\psi$. On the right, it is clear
that the modeled transfers distribution produced by the physical model has a
large peak and a poor reproduction of the distribution tails. On the other
hand, the ML-based models are efficient at maintaining a modeled transfers
distribution close to the references.

It is important to mention that experiments with SGS models trained from
neural emulators with a state loss
$\ell_{>}^{\mathcal{M}}\equiv\ell_{>}^{\mathcal{E}}$ have been omitted from
these plots because of their unstable behavior.

###  V.2 Neural emulator performance

We want here to identify the stability difference and better understand the
source of error between models obtained from neural emulators of different
complexities, mostly in terms of NN architecture and corresponding number of
trainable parameters.

We look at the evolution of the root-mean-squared error (rms) of some
reference $\cdot$ and predicted $\hat{\cdot}$ quantities over temporal
integration steps $M$,

$\ell_{\mathrm{rms}}(\cdot,\hat{\cdot})=\frac{1}{\sigma({\cdot})M}\sqrt{\frac{1}{N}\sum_{i=0}^{N}||\cdot_{i}-\,\hat{\cdot}_{i}\,||_{2}^{2}}$
(37)

where $\sigma$ is the standard deviation.

In Fig. 5, we show the $\ell_{\mathrm{rms}}$ of different SGS models compared
to the DNS.

Figure 5: $\ell_{\mathrm{rms}}$ normalized by the number of temporal
integration steps $M$ for the vorticity $\omega$ (left) and velocity
$\mathbf{u}$ (right). The error for each SGS model is computed w.r.t. the DNS.
As for comparison, we also show the reduced simulation without model
($\mathcal{M}_{0}$) that is known to be unstable.

First, vorticity mismatch indicates that each model is able to improve
substantially compared to the reduced simulation without SGS models (or
$\mathcal{M}_{0}$). In particular, NN-based models produce an error that
increases slower than the physical models. The emulator-based models are
performing slightly worse than the reference online model but remain
competitive. We also see that error for the offline approach grows linearly,
while the other models follow a logarithmic shape. It may indicate that errors
and instabilities develops after more than $M=25$ temporal iterations. The
velocity mismatch is quite different, since error for the CNN-based model
grows faster and is already greater than the no model error. The error for the
DRN-based model is still large compared to the other models, but also seems to
grow linearly. It is clear here that there is still some progress to be made
about the emulator-based predictions on the largest scales, potentially
induced by some error bias in the neural emulator.

We isolate the error produced by the neural emulators compared to the solver a
reduced resolution in Fig. 6.

Figure 6: $\ell_{\mathrm{rms}}$ w.r.t. the reduced DNS, normalized by the
number of temporal integration steps $M$ for the reduced vorticity
$\bar{\omega}$ (left) and velocity $\bar{\mathbf{u}}$ (right) for NN-based
emulators and finite difference discretization of different orders.

In addition to the neural emulators, we show the error produced by finite
difference (FD) schemes (or emulators) of different orders from $h^{2}$ to
$h^{6}$ using modified wavenumbers in a spectral method [73]. We can confirm
that the error produced by the neural emulators is mainly affected by large
scales. Indeed, the vorticity error of the DRN-based model is equivalent to a
6th order FD scheme, while the CNN-based model performs slightly worse than a
4th order FD but better than the 2nd order FD. However, both produce velocity
errors larger than a 2nd order FD.

This further confirms that improvements on the large-scale predictions of the
neural emulator are potentially a limitation that must be addressed.

##  VI Discussion

In this article, we have proposed an algorithm for training an online ML-based
SGS parametrization, i.e., with a posteriori criteria for non-differentiable
numerical solvers. The core idea of our approach is to use neural emulators
that represent the reduced dynamics, equivalent to the original equations in a
lower dimensional space. These emulators are by definition differentiable and
allows us to train the parametrization in an online setting. The main benefit
of our two-step algorithm is to avoid error compensation between the emulator
and the trained ML-based parametrization using different loss targets for the
SGS model and the neural emulator, respectively. Indeed, as demonstrated in
section 3.2 above, learning jointly emulator and parametrization will lead to
error compensation, where the parametrization-part accomodates the imperfect
nature of the emulator-part. When used in a simulation, the trained
parametrization is found to add spurious transfers that decreases overall
performance. Our two-step algorithm has been tested on a two-dimensional SGS
modeling problem and shows performance levels close to those obtained with
online learning on differentiable solvers. In a highly turbulent quasi-
geostrophic setting with bottom topography, the model trained with our two-
step algorithm is able to remain stable over many temporal iterations.
Moreover, it has a matching performance with the online model in terms of
small-scale quantities, and overall enstrophy spectrum and transfers
distribution. Nonetheless, additional work will be required in order to
alleviate some systematic biases in the SGS model associated with the
inherently imperfect nature of the neural emulator. This error is mostly
visible at large scale. We think that constraining the loss function of the
neural emulator on related large scale quantities such as velocities instead
of vorticity could further improve its performance, translating to a smaller
bias in the learned SGS model.

This work represents a first step toward more complex neural emulators. While
the experimental system used here has a considerable amount of state
parameters, more work will still be required for deploying our approach with
full complexity physical models. We have also seen that the SGS learning step
is sensitive to the quality of the neural emulator, which could be a potential
issue when emulating larger numerical systems. At this stage, it is also
unclear how our algorithm compares with other existing gradient free training
techniques and a systematic inter-comparison in terms of cost and performance
would be needed. We note that the neural emulator approach is not only
relevant for SGS parametrization but also more generally for differentiable
computing, with a large potential in Earth System Modeling [31]. Our
emulation-based approach could be leveraged for other task as for instance for
calibrating model parameters [32] or learning model error [74]. Our algorithm
could also provide a practical solution for deploying variational data
assimilation (DA) techniques without having to compute the adjoint of the
direct solver [75]. In practice, variational DA frameworks requires an adjoint
operator which is complicated to maintain. Adjoint-free minimization
techniques have been proposed, including ensemble methods [76] but trades-off
performance and computing time. In previous work [45] and recent applications
to different tasks, differentiable emulation opens a possible alternative
approach for variational DA of numerical models.

It is important to mention that building a neural emulator for a specific
problem can be done using different quantities. Here, we used the neural
emulator to predict the model state-space evolution (or right-hand-side) of
the dynamical system, which allows us to learn on different temporal horizons.
It is also possible to use neural emulation to map model states at different
future times in order to train SGS parametrizations [77].

## Acknowledgments

This research was supported by the CNRS through the 80 PRIME project, the ANR
through the Melody, OceaniX. As part of the M2LiNES project, additional
support was also provided by Schmidt Futures, a philanthropic initiative
founded by Eric and Wendy Schmidt, as part of its Virtual Earth System
Research Institute (VESRI). Computations were performed using GPU resources
from GENCI-IDRIS.

## References

  * Rogallo and Moin [1984] R. S. Rogallo and P. Moin, Numerical simulation of turbulent flows, Annual Review of Fluid Mechanics 16, 99 (1984).
  * Meneveau and Katz [2000] C. Meneveau and J. Katz, Scale-invariance and turbulence models for large-eddy simulation, Annual Review of Fluid Mechanics 32, 1 (2000).
  * Fox-Kemper  _et al._ [2019] B. Fox-Kemper, A. Adcroft, C. W. Böning, E. P. Chassignet, E. Curchitser, G. Danabasoglu, C. Eden, M. H. England, R. Gerdes, R. J. Greatbatch, _et al._ , Challenges and prospects in ocean circulation models, Frontiers in Marine Science 6, 65 (2019).
  * Schneider  _et al._ [2017] T. Schneider, S. Lan, A. Stuart, and J. Teixeira, Earth system modeling 2.0: A blueprint for models that learn from observations and targeted high-resolution simulations, Geophysical Research Letters 44, 12 (2017).
  * Maher  _et al._ [2018] P. Maher, G. K. Vallis, S. C. Sherwood, M. J. Webb, and P. G. Sansom, The impact of parameterized convection on climatological precipitation in atmospheric global climate models, Geophysical Research Letters 45, 3728 (2018).
  * Hewitt  _et al._ [2020] H. T. Hewitt, M. Roberts, P. Mathiot, A. Biastoch, E. Blockley, E. P. Chassignet, B. Fox-Kemper, P. Hyder, D. P. Marshall, E. Popova, _et al._ , Resolving and parameterising the ocean mesoscale in earth system models, Current Climate Change Reports 6, 137 (2020).
  * Brunton  _et al._ [2020] S. L. Brunton, B. R. Noack, and P. Koumoutsakos, Machine learning for fluid mechanics, Annual Review of Fluid Mechanics 52, 477 (2020).
  * Vinuesa and Brunton [2022] R. Vinuesa and S. L. Brunton, Enhancing computational fluid dynamics with machine learning, Nature Computational Science 2, 358 (2022).
  * Arcomano  _et al._ [2022] T. Arcomano, I. Szunyogh, A. Wikner, J. Pathak, B. R. Hunt, and E. Ott, A hybrid approach to atmospheric modeling that combines machine learning with a physics-based numerical model, Journal of Advances in Modeling Earth Systems 14, e2021MS002712 (2022).
  * Sonnewald  _et al._ [2021] M. Sonnewald, R. Lguensat, D. C. Jones, P. D. Dueben, J. Brajard, and V. Balaji, Bridging observations, theory and numerical simulation of the ocean using machine learning, Environmental Research Letters 16, 073008 (2021).
  * Karpov  _et al._ [2022] P. I. Karpov, C. Huang, I. Sitdikov, C. L. Fryer, S. Woosley, and G. Pilania, Physics-informed machine learning for modeling turbulence in supernovae, The Astrophysical Journal 940, 26 (2022).
  * Pawar  _et al._ [2023] S. Pawar, O. San, A. Rasheed, and P. Vedula, Frame invariant neural network closures for kraichnan turbulence, Physica A: Statistical Mechanics and its Applications 609, 128327 (2023).
  * Frezat  _et al._ [2021] H. Frezat, G. Balarac, J. Le Sommer, R. Fablet, and R. Lguensat, Physical invariance in neural networks for subgrid-scale scalar flux modeling, Physical Review Fluids 6, 024607 (2021).
  * Beucler  _et al._ [2021] T. Beucler, M. Pritchard, S. Rasp, J. Ott, P. Baldi, and P. Gentine, Enforcing analytic constraints in neural networks emulating physical systems, Physical Review Letters 126, 098302 (2021).
  * Pahlavan  _et al._ [2023] H. A. Pahlavan, P. Hassanzadeh, and M. J. Alexander, Explainable offline-online training of neural networks for parameterizations: A 1d gravity wave-qbo testbed in the small-data regime, arXiv preprint 10.48550/arXiv.2309.09024 (2023).
  * Sirignano and MacArt [2023] J. Sirignano and J. F. MacArt, Dynamic deep learning les closures: Online optimization with embedded dns, arXiv preprint 10.48550/arXiv:2303.02338 (2023).
  * Sirignano  _et al._ [2020] J. Sirignano, J. F. MacArt, and J. B. Freund, DPM: A deep learning pde augmentation method with application to large-eddy simulation, Journal of Computational Physics 423, 109811 (2020).
  * Duraisamy [2021] K. Duraisamy, Perspectives on machine learning-augmented reynolds-averaged and large eddy simulation models of turbulence, Physical Review Fluids 6, 050504 (2021).
  * MacArt  _et al._ [2021] J. F. MacArt, J. Sirignano, and J. B. Freund, Embedded training of neural-network subgrid-scale turbulence models, Physical Review Fluids 6, 050502 (2021).
  * Frezat  _et al._ [2022] H. Frezat, J. Le Sommer, R. Fablet, G. Balarac, and R. Lguensat, A posteriori learning for quasi-geostrophic turbulence parametrization, Journal of Advances in Modeling Earth Systems 14, e2022MS003124 (2022).
  * List  _et al._ [2022] B. List, L.-W. Chen, and N. Thuerey, Learned turbulence modelling with differentiable fluid solvers: physics-based loss functions and optimisation horizons, Journal of Fluid Mechanics 949, A25 (2022).
  * Fan and Wang [2023] X. Fan and J.-X. Wang, Differentiable hybrid neural modeling for fluid-structure interaction, arXiv preprint 10.48550/arXiv.2303.12971 (2023).
  * Shankar  _et al._ [2023] V. Shankar, V. Puri, R. Balakrishnan, R. Maulik, and V. Viswanathan, Differentiable physics-enabled closure modeling for burgers’ turbulence, Machine Learning: Science and Technology 4, 015017 (2023).
  * Um  _et al._ [2020] K. Um, R. Brand, Y. R. Fei, P. Holl, and N. Thuerey, Solver-in-the-loop: Learning from differentiable physics to interact with iterative pde-solvers, in  _Advances in Neural Information Processing Systems_ , Vol. 33 (2020) pp. 6111–6122.
  * Holl  _et al._ [2020] P. Holl, V. Koltun, and N. Thuerey, Learning to control pdes with differentiable physics, arXiv preprint 10.48550/arXiv.2001.07457 (2020).
  * Heiden  _et al._ [2021] E. Heiden, D. Millard, E. Coumans, Y. Sheng, and G. S. Sukhatme, NeuralSim: Augmenting differentiable simulators with neural networks, in  _2021 IEEE International Conference on Robotics and Automation (ICRA)_ (IEEE, 2021) pp. 9474–9481.
  * Négiar  _et al._ [2022] G. Négiar, M. W. Mahoney, and A. S. Krishnapriyan, Learning differentiable solvers for systems with hard constraints, arXiv preprint 10.48550/arXiv.2207.08675 (2022).
  * Dresdner  _et al._ [2022] G. Dresdner, D. Kochkov, P. Norgaard, L. Zepeda-Núñez, J. A. Smith, M. P. Brenner, and S. Hoyer, Learning to correct spectral methods for simulating turbulent flows, arXiv preprint 10.48550/arXiv.2207.00556 (2022).
  * Kochkov  _et al._ [2021] D. Kochkov, J. A. Smith, A. Alieva, Q. Wang, M. P. Brenner, and S. Hoyer, Machine learning–accelerated computational fluid dynamics, Proceedings of the National Academy of Sciences 118, e2101784118 (2021).
  * Takahashi  _et al._ [2021] T. Takahashi, J. Liang, Y.-L. Qiao, and M. C. Lin, Differentiable fluids with solid coupling for learning and control, in  _Proceedings of the AAAI Conference on Artificial Intelligence_, Vol. 35 (2021) pp. 6138–6146.
  * Gelbrecht  _et al._ [2023] M. Gelbrecht, A. White, S. Bathiany, and N. Boers, Differentiable programming for earth system modeling, Geoscientific Model Development 16, 3123 (2023).
  * Wagner  _et al._ [2023] G. L. Wagner, A. Hillier, N. C. Constantinou, S. Silvestri, A. Souza, K. Burns, A. Ramadhan, C. Hill, J.-M. Campin, J. Marshall, _et al._ , CATKE: a turbulent-kinetic-energy-based parameterization for ocean microturbulence with dynamic convective adjustment, arXiv preprint 10.48550/arXiv.2306.13204 (2023).
  * Ramadhan  _et al._ [2020] A. Ramadhan, J. Marshall, A. Souza, G. L. Wagner, M. Ponnapati, and C. Rackauckas, Capturing missing physics in climate model parameterizations using neural differential equations, arXiv preprint 10.48550/arXiv.2010.12559 (2020).
  * Pawar and San [2021] S. Pawar and O. San, Data assimilation empowered neural network parametrizations for subgrid processes in geophysical flows, Physical Review Fluids 6, 050501 (2021).
  * Huang  _et al._ [2022] D. Z. Huang, J. Huang, S. Reich, and A. M. Stuart, Efficient derivative-free bayesian inference for large-scale inverse problems, Inverse Problems 38, 125006 (2022).
  * Iglesias  _et al._ [2013] M. A. Iglesias, K. J. Law, and A. M. Stuart, Ensemble kalman methods for inverse problems, Inverse Problems 29, 045001 (2013).
  * Kovachki and Stuart [2019] N. B. Kovachki and A. M. Stuart, Ensemble kalman inversion: a derivative-free technique for machine learning tasks, Inverse Problems 35, 095005 (2019).
  * Lopez-Gomez  _et al._ [2022] I. Lopez-Gomez, C. Christopoulos, H. L. Langeland Ervik, O. R. Dunbar, Y. Cohen, and T. Schneider, Training physics-based machine-learning parameterizations with gradient-free ensemble kalman methods, Journal of Advances in Modeling Earth Systems 14, e2022MS003105 (2022).
  * Kim  _et al._ [2022] J. Kim, H. Kim, J. Kim, and C. Lee, Deep reinforcement learning for large-eddy simulation modeling in wall-bounded turbulence, Physics of Fluids 34 (2022).
  * Kurz  _et al._ [2023] M. Kurz, P. Offenhäuser, and A. Beck, Deep reinforcement learning for turbulence modeling in large eddy simulations, International Journal of Heat and Fluid Flow 99, 109094 (2023).
  * Novati  _et al._ [2021] G. Novati, H. L. de Laroussilhe, and P. Koumoutsakos, Automating turbulence modelling by multi-agent reinforcement learning, Nature Machine Intelligence 3, 87 (2021).
  * Bae and Koumoutsakos [2022] H. J. Bae and P. Koumoutsakos, Scientific multi-agent reinforcement learning for wall-models of turbulent flows, Nature Communications 13, 1443 (2022).
  * Cranmer  _et al._ [2020] K. Cranmer, J. Brehmer, and G. Louppe, The frontier of simulation-based inference, Proceedings of the National Academy of Sciences 117, 30055 (2020).
  * Bocquet [2023] M. Bocquet, Surrogate modeling for the climate sciences dynamics with machine learning and data assimilation, Frontiers in Applied Mathematics and Statistics 9, 1133226 (2023).
  * Nonnenmacher and Greenberg [2021] M. Nonnenmacher and D. S. Greenberg, Deep emulators for differentiation, forecasting, and parametrization in earth science simulators, Journal of Advances in Modeling Earth Systems 13, e2021MS002554 (2021).
  * Baydin  _et al._ [2018] A. G. Baydin, B. A. Pearlmutter, A. A. Radul, and J. M. Siskind, Automatic differentiation in machine learning: a survey, Journal of Marchine Learning Research 18, 1 (2018).
  * Hochreiter [1998] S. Hochreiter, The vanishing gradient problem during learning recurrent neural nets and problem solutions, International Journal of Uncertainty, Fuzziness and Knowledge-Based Systems 6, 107 (1998).
  * Sagaut [2005] P. Sagaut,  _Large eddy simulation for incompressible flows: an introduction_ (Springer Science & Business Media, 2005).
  * Maulik  _et al._ [2019] R. Maulik, O. San, A. Rasheed, and P. Vedula, Subgrid modelling for two-dimensional turbulence using neural networks, Journal of Fluid Mechanics 858, 122 (2019).
  * Guan  _et al._ [2022] Y. Guan, A. Chattopadhyay, A. Subel, and P. Hassanzadeh, Stable a posteriori les of 2d turbulence using convolutional neural networks: Backscattering analysis and generalization to higher re via transfer learning, Journal of Computational Physics 458, 111090 (2022).
  * Guan  _et al._ [2023] Y. Guan, A. Subel, A. Chattopadhyay, and P. Hassanzadeh, Learning physics-constrained subgrid-scale closures in the small-data regime for stable and accurate les, Physica D: Nonlinear Phenomena 443, 133568 (2023).
  * Lesieur and Metais [1996] M. Lesieur and O. Metais, New trends in large-eddy simulations of turbulence, Annual Review of Fluid Mechanics 28, 45 (1996).
  * Piomelli  _et al._ [1991] U. Piomelli, W. H. Cabot, P. Moin, and S. Lee, Subgrid-scale backscatter in turbulent and transitional flows, Physics of Fluids A: Fluid Dynamics 3, 1766 (1991).
  * Schumann [1995] U. Schumann, Stochastic backscatter of turbulence energy and scalar variance by random subgrid-scale fluxes, Proceedings of the Royal Society of London. Series A: Mathematical and Physical Sciences 451, 293 (1995).
  * Liu  _et al._ [2011] Y. Liu, L. Lu, L. Fang, and F. Gao, Modification of spalart–allmaras model with consideration of turbulence energy backscatter using velocity helicity, Physics Letters A 375, 2377 (2011).
  * Danilov  _et al._ [2019] S. Danilov, S. Juricke, A. Kutsenko, and M. Oliver, Toward consistent subgrid momentum closures in ocean models, Energy Transfers in Atmosphere and Ocean , 145 (2019).
  * Graham and Ringler [2013] J. P. Graham and T. Ringler, A framework for the evaluation of turbulence closures used in mesoscale ocean large-eddy simulations, Ocean Modelling 65, 25 (2013).
  * Jansen  _et al._ [2015] M. F. Jansen, I. M. Held, A. Adcroft, and R. Hallberg, Energy budget-based backscatter in an eddy permitting primitive equation model, Ocean Modelling 94, 15 (2015).
  * Juricke  _et al._ [2019] S. Juricke, S. Danilov, A. Kutsenko, and M. Oliver, Ocean kinetic energy backscatter parametrizations on unstructured grids: Impact on mesoscale turbulence in a channel, Ocean Modelling 138, 51 (2019).
  * Juricke  _et al._ [2020] S. Juricke, S. Danilov, N. Koldunov, M. Oliver, and D. Sidorenko, Ocean kinetic energy backscatter parametrization on unstructured grids: Impact on global eddy-permitting simulations, Journal of Advances in Modeling Earth Systems 12, e2019MS001855 (2020).
  * Frederiksen  _et al._ [2012] J. S. Frederiksen, T. J. O’Kane, and M. J. Zidikheri, Stochastic subgrid parameterizations for atmospheric and oceanic flows, Physica Scripta 85, 068202 (2012).
  * Leonard [1975] A. Leonard, Energy cascade in large-eddy simulations of turbulent fluid flows, in  _Advances in Geophysics_, Vol. 18 (Elsevier, 1975) pp. 237–248.
  * Canuto  _et al._ [2007] C. Canuto, M. Y. Hussaini, A. Quarteroni, and T. A. Zang,  _Spectral methods: fundamentals in single domains_ (Springer Science & Business Media, 2007).
  * Fox-Kemper and Menemenlis [2008] B. Fox-Kemper and D. Menemenlis, Can large eddy simulation techniques improve mesoscale rich ocean models?, Washington DC American Geophysical Union Geophysical Monograph Series 177, 319 (2008).
  * Thompson [2010] A. F. Thompson, Jet formation and evolution in baroclinic turbulence with simple topography, Journal of Physical Oceanography 40, 257 (2010).
  * Sanchez-Gonzalez  _et al._ [2018] A. Sanchez-Gonzalez, N. Heess, J. T. Springenberg, J. Merel, M. Riedmiller, R. Hadsell, and P. Battaglia, Graph networks as learnable physics engines for inference and control, in  _International Conference on Machine Learning_ (PMLR, 2018) pp. 4470–4479.
  * Stachenfeld  _et al._ [2021] K. Stachenfeld, D. B. Fielding, D. Kochkov, M. Cranmer, T. Pfaff, J. Godwin, C. Cui, S. Ho, P. Battaglia, and A. Sanchez-Gonzalez, Learned coarse models for efficient turbulence simulation, arXiv preprint 10.48550/arXiv.2112.15275 (2021).
  * Smagorinsky [1963] J. Smagorinsky, General circulation experiments with the primitive equations: I. the basic experiment, Monthly Weather Review 91, 99 (1963).
  * Germano  _et al._ [1991] M. Germano, U. Piomelli, P. Moin, and W. H. Cabot, A dynamic subgrid-scale eddy viscosity model, Physics of Fluids A: Fluid Dynamics 3, 1760 (1991).
  * Lilly [1992] D. K. Lilly, A proposed modification of the germano subgrid-scale closure method, Physics of Fluids A: Fluid Dynamics 4, 633 (1992).
  * Pawar  _et al._ [2020] S. Pawar, O. San, A. Rasheed, and P. Vedula, A priori analysis on deep learning of subgrid-scale parameterizations for kraichnan turbulence, Theoretical and Computational Fluid Dynamics 34, 429 (2020).
  * Bouchet and Venaille [2012] F. Bouchet and A. Venaille, Statistical mechanics of two-dimensional and geophysical flows, Physics Reports 515, 227 (2012).
  * Kravchenko and Moin [1997] A. Kravchenko and P. Moin, On the effect of numerical errors in large eddy simulations of turbulent flows, Journal of Computational Physics 131, 310 (1997).
  * Farchi  _et al._ [2021] A. Farchi, P. Laloyaux, M. Bonavita, and M. Bocquet, Using machine learning to correct model error in data assimilation and forecast applications, Quarterly Journal of the Royal Meteorological Society 147, 3067 (2021).
  * Carrassi  _et al._ [2018] A. Carrassi, M. Bocquet, L. Bertino, and G. Evensen, Data assimilation in the geosciences: An overview of methods, issues, and perspectives, Wiley Interdisciplinary Reviews: Climate Change 9, e535 (2018).
  * Yang  _et al._ [2015] Y. Yang, C. Robinson, D. Heitz, and E. Mémin, Enhanced ensemble-based 4dvar scheme for data assimilation, Computers & Fluids 115, 201 (2015).
  * Pedersen  _et al._ [2023] C. Pedersen, L. Zanna, J. Bruna, and P. Perezhogin, Reliable coarse-grained turbulent simulations through combined offline learning and neural emulation, arXiv preprint 10.48550/arXiv.2307.13144 (2023).