# Computational Tradeoffs in Image Synthesis: Diffusion, Masked-Token, and
Next-Token Prediction

Maciej Kilianπ Varun Jampaniπ Luke Zettlemoyerμ
πStabilityAI μUniversity of Washington

###### Abstract

Nearly every recent image synthesis approach, including diffusion, masked-
token prediction, and next-token prediction, uses a Transformer network
architecture. Despite this common backbone, there has been no direct, compute
controlled comparison of how these approaches affect performance and
efficiency. We analyze the scalability of each approach through the lens of
compute budget measured in FLOPs. We find that token prediction methods, led
by next-token prediction, significantly outperform diffusion on prompt
following. On image quality, while next-token prediction initially performs
better, scaling trends suggest it is eventually matched by diffusion. We
compare the inference compute efficiency of each approach and find that next
token prediction is by far the most efficient. Based on our findings we
recommend diffusion for applications targeting image quality and low latency;
and next-token prediction when prompt following or throughput is more
important.

##  1 Introduction

Following the work of Peebles and Xie (2023), deep image synthesis, including
diffusion (Sohl-Dickstein et al., 2015; Song and Ermon, 2019; Song et al.,
2020; Ho et al., 2020; Rombach et al., 2022; Esser et al., 2024), masked-token
prediction (Chang et al., 2022, 2023; Villegas et al., 2022; Yu et al., 2024),
and next-token prediction (Gafni et al., 2022; Yu et al., 2022b; Esser et al.,
2021), are all build on a common Transformer architecture (Vaswani et al.,
2023). Although these approaches are all known to scale well with compute and
data, there has been relatively little controlled comparisons of their
relative training and inference efficiency. Comparing these latent image
synthesis approaches is challenging since the objectives they optimize often
have different requirements which limit the set of applicable modules for each
approach and influence their optimal configurations. For example, next-token
prediction requires discrete input data which makes it unfit for continuous
latent space regularization advances. In fact, latent image synthesis will be
strongly influenced by the state of autoencoding research, often in an
unbalanced way. Examples of this can be found in Section 2.

In this paper, we measure the computational tradeoffs between popular
transformer-based latent image synthesis approaches - diffusion, masked-token
prediction, and next-token prediction. We investigate the impact of the
autoencoder, which encodes the latent space, on generative results and train a
large grid of models with the different approaches, model sizes, and dataset
sizes. Samples from some of our most capable models can be found in Figure 1.
Our findings indicate that (i) at smaller compute budgets, next-token
prediction yields the best image quality but scaling trends suggest it is
eventually matched by diffusion. (ii) Token-based approaches achieve superior
controllability. (iii) The quality of the autoencoder impacts the FID more
than the CLIP score of diffusion models trained on its latent space. (iv) We
find preliminary evidence for improved diffusion training practices. Based on
our findings, we recommend diffusion models for applications targeting low
latency and high image quality; and next-token prediction for applications
where prompt following and throughput are priorities.

Figure 1: Images generated using our best models. Top row is from a next-token
prediction model, bottom row is from a diffusion model. Both models are XL
size and trained for 500k steps.

##  2 Related Work

Scaling transformer-based generative models. Scaling compute budgets for
transformer based generative models is a predictable method for improving
performance. Kaplan et al. (2020); Hoffmann et al. (2022); Clark et al. (2022)
showed that for text, final training loss can be accurately predicted as a
power law of training compute which depends on model size and dataset size.
Following those practices many capable text generation models were trained
(Touvron et al., 2023; Brown et al., 2020; Rae et al., 2022). Similar results
have been found for vision (Zhai et al., 2022; Alabdulmohsin et al., 2024;
Esser et al., 2024; Dehghani et al., 2023) and even mixed modal data
(Aghajanyan et al., 2023). We follow these intuitions and analyze image
synthesis performance as a function of compute budget.

Latent generative modeling. Training latent generative vision models has
emerged as an efficient alternative to the computationally intensive modeling
of high-dimensional pixel space. Studies have demonstrated the advantages of
imposing specific structural regularizations within the latent space for
enhancing the performance of various generative models. For instance, Rombach
et al. (2022) observed that latent diffusion models operating in VAE-style
Kingma and Welling (2022) latent spaces, when regularized towards a standard
Gaussian structure, outperform models trained with alternative regularization
techniques. Yu et al. (2024); Mentzer et al. (2023); Yu et al. (2022a) have
shown that simplifying vector quantization methods can mitigate common issues
such as poor codebook utilization and enhancing the transfer between
autoencoder reconstruction quality and downstream generative model performance
for token-based approaches. Tian et al. (2024) demonstrated that employing
hierarchical next-scale latents enables transformers using next token
prediction to leverage their in-context learning capabilities more
effectively, significantly improving performance. Jin et al. (2024) use image
latents dynamically sized based on their information content which allows
generative models to allocate more computational resources to complex samples,
as these will contain more tokens. We minimize potential bias coming from
autoencoding asymmetries by studying the impact of the autoencoder on the
generative model trained on top of it.

##  3 Background

#### Autoencoding

To train latent generative models, we establish an encoder-decoder pair
$(\mathcal{E},\mathcal{D})$. For an image $x\in\mathbb{R}^{H\times W\times
3}$, the encoder maps $x$ to a latent representation $z=\mathcal{E}(x)$, where
$z\in\mathbb{R}^{H/f\times W/f\times c}$ and $f=2^{m\in\mathbb{N}}$ represents
the factor of dimensionality reduction. The decoder then reconstructs
$\hat{x}=\mathcal{D}(z)$, aiming for high perceptual similarity to $x$,
effectively making $z$ a perceptually compressed representation of the input.
To avoid high-variance latent spaces and ensure structured representations, we
employ regularization methods classified into two main types: discrete and
continuous. The regularization function q for a continuous regularizer maps
$\mathbb{R}^{d}\to\mathbb{R}^{d}$, while a discrete regularizer maps
$\textbf{q}:\mathbb{R}^{d}\to\\{0,1,2,\ldots,N\\}$, making the latent space
finite. In the following subsections we use $z\in\mathbb{R}^{s\times d}$ to
denote the flattened representation output by the encoder $\mathcal{E}$.
$p(z)$ is the latent data distribution we are interested in estimating using
our generative models. We recover images by inputting sampled latents into the
corresponding decoder $\mathcal{D}$.

#### Next token prediction

In the context of sequences of discrete tokens represented as
$z\in\\{0,1,2,\ldots,N\\}^{s}$, we employ the chain rule of conditional
probability to decompose the target distribution into a product of conditional
distributions which are tractable since the range of $z_{i}$ is finite. To
model this distribution, we use a neural network $f$, parameterized by weights
$\theta$. The parameters are optimized by minimizing the negative log-
likelihood $\mathcal{L}_{NT}$.

$\displaystyle p(z)$
$\displaystyle=\prod_{i=1}^{n}p(z_{i}|z_{i-1},\ldots,z_{1})$
$\displaystyle\mathcal{L}_{NT}$
$\displaystyle=\mathbb{E}_{i}[-\log{p(z_{i}|z_{<i};\theta})]$ (1)

Sampling from our learned distribution begins with an empty sequence (in
practice, a "start of text" token is sampled with 1.0 probability). We then
sample the first token unconditionally and append it to our sequence. The
process continues by iteratively evaluating the conditionals and sampling from
them, with each step increasing the sequence length by one.

#### Masked token prediction

Masked token prediction is a form of iterative denoising and can be viewed as
a discrete diffusion process. In this process, tokens progressively transition
to an absorbing [MASK] state according to a probability defined by a noise
schedule $\gamma(t)\in(0,1]$ where $t\sim\mathcal{U}(0,1)$. This transition
can also be mathematically expressed as a product of conditionals, except in a
perturbed order $\sigma$, and implemented as a neural network. Here,
$\sigma(i)$ is a surjective function mapping $[0,N]\mapsto[0,N]$. We follow
Chang et al. (2022, 2023) where $\sigma(i)=\sigma(i,t)$ such that
$p(\sigma(i,t)<j)=\gamma(t)$ meaning the likelihood a token can be attended to
is independent of position. In this formulation, we utilize a truncated
$\arccos$ distribution for our noise schedule:
$\gamma(t)=\frac{2}{\pi}(1-t^{2})^{-\frac{1}{2}}$. To apply this method, we
generate a mask tensor $M\in\\{0,1\\}^{s}$ by sampling $t\sim\mathcal{U}(0,1)$
and $m_{i}\sim\text{Bernoulli}(\gamma(t))$. The tensor $M$ is applied
elementwise to the latents, replacing $z_{i}$ with the [MASK] token if
$m_{i}=1$; otherwise, $z_{i}$ remains unchanged. Denote the resultant noised
sequence as $z_{M}$. The network is then trained to minimize the masked token
loss $\mathcal{L}_{MT}$.

$\displaystyle p(z)$
$\displaystyle=\prod_{i,\sigma(i)=j}^{n}p(z_{j}|z_{\sigma(i)<j})$
$\displaystyle\quad\mathcal{L}_{MT}$
$\displaystyle=\mathbb{E}_{i,m_{i}=1}[-\log{p(z_{i}|z_{\overline{M}};\theta})]$
(2)

Sampling from the distribution starts with a fully masked sequence and
iterates through a discretized noise schedule $t_{i}=i/N$ over $N$ desired
steps. At each step, the model estimates $p(z|z_{\overline{M}})$ for sampling,
followed by re-noising using $\gamma(t_{i+1})$. This iterative re-noising and
sampling process is repeated $N$ times to yield the final sample.

#### Diffusion

We adopt the flow matching framework outlined by Lipman et al. (2023),
focusing on models that map samples from a noise distribution $p_{1}$ to a
data distribution $p_{0}$ using continuous trajectories governed by an
ordinary differential equation (ODE). Furthermore, we enforce straight paths
between the terminal distributions (by setting $\alpha_{t}=1-t$ and
$\beta_{t}=t$) since this has been shown to perform well at scale (Esser et
al., 2024).

$\displaystyle d\phi_{t}(x)$ $\displaystyle=v_{t}(\phi_{t}(x))\,dt$
$\displaystyle\phi_{0}(x)$ $\displaystyle=x$ $\displaystyle z_{t}$
$\displaystyle=\alpha_{t}x_{0}+\beta_{t}\epsilon,\quad\epsilon\sim\mathcal{N}(0,1)$
(3)

Here, $v_{t}:[0,1]\times\mathbb{R}^{d}\mapsto\mathbb{R}^{d}$ represents a
time-dependent vector field, which we aim to parameterize using a neural
network $\theta$ and $\phi_{t}:[0,1]\times\mathbb{R}^{d}\mapsto\mathbb{R}^{d}$
is the flow. To optimize our weights, we regress the vector field $u_{t}$,
which generates our paths $z_{t}$ by employing conditional flow matching which
we reformulate as a noise-prediction objective $\mathcal{L}_{CFM}$. Sampling
is performed by using an ODE solver to solve Equation 3 in reverse time,
utilizing our trained neural network $v_{\Theta}(z,t)$.

$\mathcal{L}_{CFM}(x_{0})=\mathbb{E}_{t\sim\mathcal{U}(t),\epsilon\sim\mathcal{N}(0,I)}\left[\frac{1}{(1-t)^{2}}\|\epsilon_{\theta}(z_{t},t)-\epsilon\|^{2}\right]\;$

##  4 Experimental Setup

#### Data

We train both the autoencoders and the generative models on a large web
dataset of image and text pairs at 256x256 resolution. For the conditioning we
use the pooled text embedding of the OpenCLIP bigG/14 model from Cherti et al.
(2022). Once the autoencoders are trained we pre-encode the entire dataset
with them for improved training speed.

#### Evaluation metrics

Since we are in the infinite data regime, we look at the final train loss and
do not compare across objectives since the losses represent different
quantities. We also look at CLIP score (Radford et al., 2021; Hessel et al.,
2022) and FID computed on CLIP features (Sauer et al., 2021) based on the
decoded samples $\hat{x}=\mathcal{D}(z)$.

#### Autoencoding

We study well-established autoencoder configurations that have proven
effective without special handling for each data type. We adhere to the
training and architectural guidelines provided by Rombach et al. (2022). Each
autoencoder is trained with a downsampling factor $f=8$, reducing $256\times
256$ images to a $32\times 32$ grid of latents. For continuous variants,
$\textbf{q}(z)$ implements a KL penalty aiming towards the standard normal
distribution (Kingma and Welling, 2022; Rombach et al., 2022), while for
discrete variants, we utilize lookup-free quantization (LFQ) (Yu et al.,
2024). Further details on the selection of discrete regularizers are available
in Appendix A. To circumvent potential challenges associated with large
vocabulary sizes, as highlighted by Yu et al. (2024), our LFQ-regularized
autoencoder is trained with a vocabulary size of 16384 (Esser et al., 2021).
Assessing the comparability of autoencoders is difficult since there are many
variables of interest such as the (1) information capacity of the latent
space; (2) compute used to train the autoencoder; (3) reconstruction quality
achieved by the autoencoder. To explore the influence of these factors on the
performance of generative models, we train a set of autoencoders similar to
those in Esser et al. (2024), which exhibit a range of information capacities
and reconstruction qualities. Additionally, we experiment with targeting
specific reconstruction qualities, irrespective of other factors, by training
a KL-regularized autoencoder with early stopping to match the reconstruction
quality of our discrete autoencoder within a certain threshold $\epsilon$
111Discrete autoencoders typically have worse reconstruction qualities since
the information bottleneck is tighter. This can be shown by comparing
$\log(\text{codebook size})$ to $\text{num\\_channels}*\text{sizeof(dtype)}$
for common values of these quantities. In our case we needed to stop at 75k
steps vs. 1M for the discrete autoencoder.. Table 2 provides detailed
information about the autoencoders.

Regularizer | Latent space capacity | rFID ($\downarrow$)
---|---|---
KL | 16 channels | 1.060
KL | 8 channels | 1.560
KLes | 8 channels | 2.856
KL | 4 channels | 2.410
LFQ | 16384 vocabulary | 2.784
Table 1: Autoencoders. Reconstruction metrics for differently regularized and
trained autoencoders (downsampling $f=8$). "es" is early stopping to match the
LFQ autoencoder.

Model size |  Layers $N$ |  Hidden size $d$ | Heads
---|---|---|---
S | 12 | 768 | 12
M | 24 | 1024 | 16
L | 24 | 1536 | 16
XL | 32 | 2304 | 32
Table 2: Transformer configurations. Base transformer hyperparameters for
models we train. Common across all approaches

Autoencoder ablation. We train an L-size diffusion model on top of the latent
space of each continuous autoencoder. We then evaluate the models using the
metrics described in Section 4 and plot them against the number of training
steps. Results are shown in Figure 2. We find that the autoencoder’s
reconstruction quality has a consistently significant impact on the FID score,
while its effect on the CLIP score diminishes with larger dataset sizes, where
the models tend to yield similar results. This trend likely emerges because
improvements in autoencoder quality enhance perceptual reconstruction metrics
similar to FID, rather than affecting language or semantic capabilities. Upon
examining the number of channels in the autoencoders, our findings concur with
those reported by Esser et al. (2024), indicating that leveraging larger and
better latent spaces requires more compute and model capacity. Additionally,
the model trained on our early-stopped autoencoder’s latent space performed
significantly worse than the 4-channel autoencoder, which achieves similar
reconstruction quality. This confirms the importance of latent space structure
for overall performance. Building on these insights, we have chosen to use the
4-channel autoencoder for our main diffusion experiments. This model most
closely matches the latent space capacity and reconstruction quality of our
discrete autoencoder, while also ensuring that the latent structure is
adequately developed to support the diffusion model trained on it. Although
more advanced autoencoders have been developed—such as those featuring
increased channel counts or expanded codebook sizes—our primary focus in this
study is to maintain comparability across objectives.

Figure 2: Impact of autoencoder quality on diffusion models. We train L-size
diffusion models on our set of continuous latent space autoencoders. The
choice of autoencoder has more impact on FID than CLIP score. Effectively
using a larger latent space requires more compute and model capacity.

###  4.1 Network Architecture

Backbone. We opt for the transformer architecture as our primary network
backbone, recognizing its capability to scale effectively with computational
resources and its status as the state-of-the-art (SOTA) across all evaluated
approaches. Configuring a transformer involves many decisions, such as
choosing normalization methods, feed-forward layer configurations, positional
embedding schemes, conditioning methods, and initialization strategies. Given
the prohibitive cost of exploring all possible hyperparameters, we adhere to
established practices in recent studies.

Design differences. For approaches utilizing discrete representations, we
primarily follow the configurations used in the LLaMa model (Touvron et al.,
2023), incorporating SwiGLU feed-forward layers with an expansion ratio of
$\frac{2}{3}4$ and rotary positional embeddings (Su et al., 2023). An
exception is made for masked token prediction, where learned positional
embeddings are preferred to address positional ambiguities that degrade
performance near the center of the image. For continuous representation
approaches, we align with diffusion transformers (Peebles and Xie, 2023),
employing GELU feed-forward layers with an expansion ratio of 4 and learned
positional embeddings. All models use QK-normalization (Dehghani et al., 2023)
for better training stability.

Conditioning ablation. We choose to ablate the conditioning method, as it
significantly impacts the computational cost of model operations. Adaptive
layer normalization (AdaLN) (Perez et al., 2017) has shown promise in latent
image synthesis for both continuous (Peebles and Xie, 2023) and discrete (Tian
et al., 2024) settings. To validate this choice in the discrete context, we
conduct small-scale ablations on S-size models, comparing AdaLNzero (Peebles
and Xie, 2023), with two other common conditioning methods: prepending a
projected embedding in the context of the transformer and cross-attention. The
outcomes of these ablations are presented in Table 4, informing our choice of
conditioning method for subsequent experiments.

Compute cost. To assess the computational cost of each model, we first
standardize a set of hyperparameters across all transformers, detailed in
Table 2. We then calculate the forward pass FLOPs for a single sample (a
sequence of 1024 embeddings) for each approach and model size, and present
them in Table 4. Assuming the backward pass is twice the cost of the forward
pass, we compute the training FLOPs for each model as
$(1+2)\times(\text{forward FLOPs})\times D$, where $D$ represents the total
number of training samples.

Model | n-parameters (%) | Forward TFLOPs
---|---|---
DiT-S | 131.13 M (97.2%) | 0.2133
DiT-M | 459.19 M (98.7%) | 0.7234
DiT-L | 1031.67 M (98.8%) | 1.5485
DiT-XL | 3083.69 M (99.2%) | 4.4901
NT/MT-S | 153.73 M (82.9%) | 0.2261
NT/MT-M | 494.56 M (92.9%) | 0.6631
NT/MT-L | 1072.14 M (95.1%) | 1.3421
NT/MT-XL | 3137.29 M (97.5%) | 3.7166

Table 3: Model forward pass costs. Number of parameters and FLOPs used in each
forward pass for all models we trained. DiT - Diffusion Transformer; NT -
Next-token; MT - Masked-token

Objective | Conditioning | FID | CLIP
---|---|---|---
NT | adaLNzero | 83.052 | 0.2213
NT | in context | 88.176 | 0.2041
NT | cross attention | 92.852 | 0.2062
MT | adaLNzero | 97.021 | 0.2164
MT | in context | 100.646 | 0.1925
MT | cross attention | 103.221 | 0.1960

Table 4: Conditioning method ablation. Results for different objectives and
conditioning methods. adaLNzero conditioning is used for the remainder of
experiments. NT - Next-token; MT - Masked-token.

###  4.2 Training

Each approach also has associated training hyperparameters which past work has
found to work well and for the same reasons as stated in 4.1 we follow them.

Optimization and conditioning. For diffusion experiments we follow Esser et
al. (2024) and use a constant learning rate schedule with a maximum value of
$1^{-4}$. For next and masked token prediction we use a cosine decay learning
rate with a maximum value of $3^{-3}$ which decays down to $3^{-5}$. All
models have a linear learning rate warmup lasting 1000 steps up to the maximum
value. We use the AdamW (Loshchilov and Hutter, 2019) optimizer with
$\beta_{1}=0.9$, $\beta_{2}=0.95$, decay=0.01, and epsilon=1e-15 for improved
transformer training stability (Wortsman et al., 2023). All models are trained
at bf16-mixed precision (Chen et al., 2019). We intend to use classifier free
guidance (CFG) (Ho and Salimans, 2022) during sampling so we randomly drop
conditioning 10% of the time during training. Since its inexpensive and does
not influence training, for all models, we store a copy of the model weights
which gets updated every 100 training batches with an exponential moving
average (EMA) using a decay factor of 0.99 and during evaluation we evaluate
both sets of weights.

Training steps. For each objective and model size we scale to at least 250k
training steps with a batch size of 512. For diffusion we decide to go up to
500k steps since constant learning rate schedule allows more flexibility with
dataset size222With a decaying learning rate, each dataset size we want to
study requires a separate run from scratch whereas for constant learning rate
schedules you can simply continue from a past checkpoint. Occasionally we
train models for longer to attempt to illustrate convergence or crossing
points.

###  4.3 Sampling

Classifier free guidance. Ho and Salimans (2022) introduced it in diffusion
models as an elegant way of trading off diversity for fidelity and has been
demonstrated to improve results for all approaches we consider in this study
(Chang et al., 2023; Gafni et al., 2022; Ho and Salimans, 2022). We use it
here in the form

$x_{g}=(1+w)x_{c}-wx_{u}$ (4)

where $w$ is the guidance scale. For diffusion $x$ will be the position in the
denoising trajectory and for token based methods $x$ is the logit distribution
at a given timestep.

Hyperparameters. For our diffusion models we follow Esser et al. (2024) and
use 50 sampling steps with a CFG scale of 5. Since the conditioning and input
data is slightly different we also perform a small sweep around those
parameters to confirm they are still optimal. For the token based models we
could not find good resources on reasonable sampling hyperparameters so we
perform small sweeps for S-size models to find the best configurations and
verify the robustness of those values for larger models. Common between them,
we use nucleus sampling (Holtzman et al., 2020) with a top-p value of $0.9$
and a temperature of $1.0$. For next token prediction and masked token
prediction we use CFG scales 8 and 5 respectively. For masked token prediction
we perform 10 sampling steps.

##  5 Results

###  5.1 Training tradeoffs

For all models, we measure our evaluation metrics every 50k steps of training
and plot them in log scale against the log of training compute. Figure 3
presents this for FID and CLIP score. There we can see that for FID, next
token prediction starts out more compute efficient but scaling trends suggest
that its eventually matched by diffusion. When looking at CLIP score we see
that token prediction is significantly better than diffusion, implying the
models generate images that follow the input prompt better. This could be a
feature of using more compressed latent spaces which is supported by Figure 2
where the 4 channel continuous autoencoder outperforms both the 8 and 16
channel autoencoder on CLIP score near the end of training. This is also
supported in Figure 7 with interpretable features like human faces emerging
sooner in the token based methods. Extending a finding from Mei et al. (2024),
we observe that, for all approaches studied, smaller models trained for longer
often surpass larger models. In Figure 4 we show the final training loss of
each model against training compute to show that it follow similar scaling
trends to what has been shown in past work on scaling deep neural networks,
briefly described in Section 2. Samples from the most capable XL sized next-
token prediction and diffusion models can be found in Figure 1.

Figure 3: Training compute efficiency on perceptual metrics. Performance on
CLIP and FID scores for various models and dataset sizes across different
image synthesis approaches. On FID, next-token prediction is initially the
most compute-efficient but scaling trends suggest it is eventually matched by
diffusion. Token-based methods significantly outperform diffusion in CLIP
score. Both axes are in log scale.  Figure 4: Training compute efficiency on
final loss. All objectives follow predictable scaling trends. Right plot shows
the difference in loss scale between diffusion models trained on top of
different autoencoders. FLOPs axis is in log scale.

###  5.2 Inference tradeoffs

Inference cost. We evaluated all models trained for 250k steps to understand
the impact of inference FLOPs on perceptual metrics. To adjust the number of
inference FLOPs for a single model, we varied the number of sampling steps,
applicable only to iterative denoising methods like masked token prediction
and diffusion. As shown in Figure 5, next-token prediction demonstrates far
greater inference compute efficiency compared to other objectives. This
efficiency arises because when using key-value caching, sampling N tokens
autoregressively uses the same amount of FLOPs as forwarding those N tokens in
parallel once. However, for iterative denoising methods, this value is
multiplied by the number of sampling steps. Interestingly, despite being
trained for iterative denoising, the number of steps in masked token
prediction appears to have minimal impact on sample quality.

Figure 5: Inference compute efficiency on perceptual metrics. Diffusion and
masked token prediction evaluated at 4, 10, 20, 50, and 100 sampling steps.
Next token prediction is 1 forward pass factorized over each token
individually. Masked token prediction isn’t influenced by the number of
sampling steps very much. Next token prediction is the most compute efficient.
Both axes are in log scale.

Sampling latency and throughput. While next-token prediction requires much
less compute per sample, the autoregressive dependency of each token causes it
to be data bound when few queries are being processed in parallel which
results in high latency. Conversely, bidirectional denoising approaches
utilize a more parallel sampling process which, despite its high cost,
facilitates low latency especially in low-volume settings with models that fit
on local devices (Chang et al., 2022). For high-volume sampling, where
throughput becomes more important, such as serving many users via an API, next
token prediction could use a batching algorithm to maximize GPU utilization by
choosing batch sizes inversely proportional to sequence lengths. The
effectiveness of this method is ensured by the fact that, for next-token
prediction image synthesis, all responses are the same length so you can
easily plan your batches ahead. This way, for high-volume sampling, next-token
prediction would enjoy the same benefits over the other approaches as
presented in the cost section above but for sample throughput.

###  5.3 EMA ablations

Among the various training practices distinguishing these methods, the use of
an exponential moving average (EMA) on the model weights stands out. In the
diffusion literature (Karras et al., 2024, 2022; Peebles and Xie, 2023; Esser
et al., 2024) EMA is an essential component of the training pipeline. In
contrast, this practice has not received equivalent attention in other
approaches. The differential impact of EMA is evident in Figure 6. For token-
based approaches, the influence of EMA is either negligible or, in some cases,
harmful, whereas for diffusion models, it is beneficial almost universally. We
hypothesize that the impact of EMA may be linked to the learning rate
schedule, where decaying schedules similarly minimize weight variation towards
the end of training. To test this hypothesis, we conducted an ablation study
on an M-sized next-token prediction and diffusion model trained over 250k
steps. Our findings verify our hypothesis that EMA enhances performance under
a constant learning rate schedule; however, it does not exceed the
improvements seen with a cosine decay learning rate schedule. This implies
that future diffusion models should consider substituting the EMA for a cosine
decay learning rate schedule if they are willing to pay the cost of decreased
training length flexibility. Results from this ablation study are presented in
Table 5.

Figure 6: Impact of EMA. EMA significantly improves FID for diffusion models
but hurts token based approaches. On CLIP score the effect on diffusion models
stays consistent however for token based methods the influence is negligible.

Objective | LR schedule | EMA | FID | CLIP
---|---|---|---|---
Next-token | constant | ✗ | 81.976 | 0.2208
Next-token | constant | ✓ | 79.571 | 0.2230
Next-token | cosine | ✗ | 75.715 | 0.2256
Next-token | cosine | ✓ | 76.404 | 0.2257
Diffusion | constant | ✗ | 74.087 | 0.2153
Diffusion | constant | ✓ | 71.789 | 0.2166
Diffusion | cosine | ✗ | 69.284 | 0.2195
Diffusion | cosine | ✓ | 69.468 | 0.2192

Table 5: EMA and learning rate schedules. EMA on model weights improves
results under a constant learning rate schedule but does not exceed the gains
from using a cosine decay schedule.

###  5.4 Limitations

Our analysis has several limitations which result from resource limitations
and project scope. We only investigate pretraining whereas most production
systems utilize a progression of pretraining, finetuning, and distillation
stages. We do not investigate high resolution images. We only measure loss and
perceptual metrics and leave out an analysis of utility for potential
downstream tasks. There are many others approaches that we leave out such as
other discrete diffusion approaches (Austin et al., 2023; Pernias et al.,
2023), causally masked token prediction (Aghajanyan et al., 2022), and many
more. We choose most hyperparameters by following past work instead of
exhaustively sweeping to find the best configurations. And finally, we do not
compare approaches using the best possible autoencoders.

Figure 7: Increasing training compute improves sample quality for all
approaches. For each approach and prompt we sample an image with all
combinations of S, M, L model sizes and 50k, 150k, 250k dataset sizes. Going
down or right in the 3x3 increases dataset and model size respectively.

##  6 Conclusion

We conduct a compute-controlled analysis comparing transformer-based
diffusion, next-token prediction, and masked-token prediction latent image
synthesis models. Our findings indicate that token based methods, led by next-
token prediction, achieve superior CLIP scores, indicating greater
controllability. In terms of FID, and therefore image quality, while next-
token prediction is much better at low training compute scales, scaling trends
suggest it is eventually matched by diffusion. We find that next token
prediction has, by far, the best inference compute efficiency but this comes
at the cost of high latency in low data intensity settings. Based on our
findings recommend diffusion models when image quality and low latency is
important; and next-token prediction for better prompt following and
throughput.

## References

  * Aghajanyan et al. [2022] Armen Aghajanyan, Bernie Huang, Candace Ross, Vladimir Karpukhin, Hu Xu, Naman Goyal, Dmytro Okhonko, Mandar Joshi, Gargi Ghosh, Mike Lewis, and Luke Zettlemoyer.  Cm3: A causal masked multimodal model of the internet, 2022.
  * Aghajanyan et al. [2023] Armen Aghajanyan, Lili Yu, Alexis Conneau, Wei-Ning Hsu, Karen Hambardzumyan, Susan Zhang, Stephen Roller, Naman Goyal, Omer Levy, and Luke Zettlemoyer.  Scaling laws for generative mixed-modal language models, 2023.
  * Alabdulmohsin et al. [2024] Ibrahim Alabdulmohsin, Xiaohua Zhai, Alexander Kolesnikov, and Lucas Beyer.  Getting vit in shape: Scaling laws for compute-optimal model design, 2024.
  * Austin et al. [2023] Jacob Austin, Daniel D. Johnson, Jonathan Ho, Daniel Tarlow, and Rianne van den Berg.  Structured denoising diffusion models in discrete state-spaces, 2023.
  * Brown et al. [2020] Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei.  Language models are few-shot learners, 2020.
  * Chang et al. [2022] Huiwen Chang, Han Zhang, Lu Jiang, Ce Liu, and William T. Freeman.  Maskgit: Masked generative image transformer, 2022.
  * Chang et al. [2023] Huiwen Chang, Han Zhang, Jarred Barber, AJ Maschinot, Jose Lezama, Lu Jiang, Ming-Hsuan Yang, Kevin Murphy, William T. Freeman, Michael Rubinstein, Yuanzhen Li, and Dilip Krishnan.  Muse: Text-to-image generation via masked generative transformers, 2023.
  * Chen et al. [2019] Dehao Chen, Chiachen Chou, Yuanzhong Xu, and Jonathan Hseu.  Bfloat16: The secret to high performance on cloud tpus, 2019.  URL https://cloud.google.com/blog/products/ai-machine-learning/bfloat16-the-secret-to-high-performance-on-cloud-tpus?hl=en.
  * Cherti et al. [2022] Mehdi Cherti, Romain Beaumont, Ross Wightman, Mitchell Wortsman, Gabriel Ilharco, Cade Gordon, Christoph Schuhmann, Ludwig Schmidt, and Jenia Jitsev.  Reproducible scaling laws for contrastive language-image learning, 2022.
  * Clark et al. [2022] Aidan Clark, Diego de las Casas, Aurelia Guy, Arthur Mensch, Michela Paganini, Jordan Hoffmann, Bogdan Damoc, Blake Hechtman, Trevor Cai, Sebastian Borgeaud, George van den Driessche, Eliza Rutherford, Tom Hennigan, Matthew Johnson, Katie Millican, Albin Cassirer, Chris Jones, Elena Buchatskaya, David Budden, Laurent Sifre, Simon Osindero, Oriol Vinyals, Jack Rae, Erich Elsen, Koray Kavukcuoglu, and Karen Simonyan.  Unified scaling laws for routed language models, 2022.
  * Dehghani et al. [2023] Mostafa Dehghani, Josip Djolonga, Basil Mustafa, Piotr Padlewski, Jonathan Heek, Justin Gilmer, Andreas Steiner, Mathilde Caron, Robert Geirhos, Ibrahim Alabdulmohsin, Rodolphe Jenatton, Lucas Beyer, Michael Tschannen, Anurag Arnab, Xiao Wang, Carlos Riquelme, Matthias Minderer, Joan Puigcerver, Utku Evci, Manoj Kumar, Sjoerd van Steenkiste, Gamaleldin F. Elsayed, Aravindh Mahendran, Fisher Yu, Avital Oliver, Fantine Huot, Jasmijn Bastings, Mark Patrick Collier, Alexey Gritsenko, Vighnesh Birodkar, Cristina Vasconcelos, Yi Tay, Thomas Mensink, Alexander Kolesnikov, Filip Pavetić, Dustin Tran, Thomas Kipf, Mario Lučić, Xiaohua Zhai, Daniel Keysers, Jeremiah Harmsen, and Neil Houlsby.  Scaling vision transformers to 22 billion parameters, 2023.
  * Esser et al. [2021] Patrick Esser, Robin Rombach, and Björn Ommer.  Taming transformers for high-resolution image synthesis, 2021.
  * Esser et al. [2024] Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, Dustin Podell, Tim Dockhorn, Zion English, Kyle Lacey, Alex Goodwin, Yannik Marek, and Robin Rombach.  Scaling rectified flow transformers for high-resolution image synthesis, 2024.
  * Gafni et al. [2022] Oran Gafni, Adam Polyak, Oron Ashual, Shelly Sheynin, Devi Parikh, and Yaniv Taigman.  Make-a-scene: Scene-based text-to-image generation with human priors, 2022.
  * Hessel et al. [2022] Jack Hessel, Ari Holtzman, Maxwell Forbes, Ronan Le Bras, and Yejin Choi.  Clipscore: A reference-free evaluation metric for image captioning, 2022.
  * Ho and Salimans [2022] Jonathan Ho and Tim Salimans.  Classifier-free diffusion guidance, 2022.
  * Ho et al. [2020] Jonathan Ho, Ajay Jain, and Pieter Abbeel.  Denoising diffusion probabilistic models.  _Advances in neural information processing systems_ , 33:6840–6851, 2020.
  * Hoffmann et al. [2022] Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hennigan, Eric Noland, Katie Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero, Karen Simonyan, Erich Elsen, Jack W. Rae, Oriol Vinyals, and Laurent Sifre.  Training compute-optimal large language models, 2022.
  * Holtzman et al. [2020] Ari Holtzman, Jan Buys, Li Du, Maxwell Forbes, and Yejin Choi.  The curious case of neural text degeneration, 2020.
  * Jin et al. [2024] Yang Jin, Kun Xu, Kun Xu, Liwei Chen, Chao Liao, Jianchao Tan, Quzhe Huang, Bin Chen, Chenyi Lei, An Liu, Chengru Song, Xiaoqiang Lei, Di Zhang, Wenwu Ou, Kun Gai, and Yadong Mu.  Unified language-vision pretraining in llm with dynamic discrete visual tokenization, 2024.
  * Kaplan et al. [2020] Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei.  Scaling laws for neural language models, 2020.
  * Karras et al. [2022] Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine.  Elucidating the design space of diffusion-based generative models, 2022.
  * Karras et al. [2024] Tero Karras, Miika Aittala, Jaakko Lehtinen, Janne Hellsten, Timo Aila, and Samuli Laine.  Analyzing and improving the training dynamics of diffusion models, 2024.
  * Kingma and Welling [2022] Diederik P Kingma and Max Welling.  Auto-encoding variational bayes, 2022.
  * Lipman et al. [2023] Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le.  Flow matching for generative modeling, 2023.
  * Loshchilov and Hutter [2019] Ilya Loshchilov and Frank Hutter.  Decoupled weight decay regularization, 2019.
  * Mei et al. [2024] Kangfu Mei, Zhengzhong Tu, Mauricio Delbracio, Hossein Talebi, Vishal M. Patel, and Peyman Milanfar.  Bigger is not always better: Scaling properties of latent diffusion models, 2024.
  * Mentzer et al. [2023] Fabian Mentzer, David Minnen, Eirikur Agustsson, and Michael Tschannen.  Finite scalar quantization: Vq-vae made simple, 2023.
  * Peebles and Xie [2023] William Peebles and Saining Xie.  Scalable diffusion models with transformers, 2023.
  * Perez et al. [2017] Ethan Perez, Florian Strub, Harm de Vries, Vincent Dumoulin, and Aaron Courville.  Film: Visual reasoning with a general conditioning layer, 2017.
  * Pernias et al. [2023] Pablo Pernias, Dominic Rampas, Mats L. Richter, Christopher J. Pal, and Marc Aubreville.  Wuerstchen: An efficient architecture for large-scale text-to-image diffusion models, 2023.
  * Radford et al. [2021] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and Ilya Sutskever.  Learning transferable visual models from natural language supervision, 2021.
  * Rae et al. [2022] Jack W. Rae, Sebastian Borgeaud, Trevor Cai, Katie Millican, Jordan Hoffmann, Francis Song, John Aslanides, Sarah Henderson, Roman Ring, Susannah Young, Eliza Rutherford, Tom Hennigan, Jacob Menick, Albin Cassirer, Richard Powell, George van den Driessche, Lisa Anne Hendricks, Maribeth Rauh, Po-Sen Huang, Amelia Glaese, Johannes Welbl, Sumanth Dathathri, Saffron Huang, Jonathan Uesato, John Mellor, Irina Higgins, Antonia Creswell, Nat McAleese, Amy Wu, Erich Elsen, Siddhant Jayakumar, Elena Buchatskaya, David Budden, Esme Sutherland, Karen Simonyan, Michela Paganini, Laurent Sifre, Lena Martens, Xiang Lorraine Li, Adhiguna Kuncoro, Aida Nematzadeh, Elena Gribovskaya, Domenic Donato, Angeliki Lazaridou, Arthur Mensch, Jean-Baptiste Lespiau, Maria Tsimpoukelli, Nikolai Grigorev, Doug Fritz, Thibault Sottiaux, Mantas Pajarskas, Toby Pohlen, Zhitao Gong, Daniel Toyama, Cyprien de Masson d’Autume, Yujia Li, Tayfun Terzi, Vladimir Mikulik, Igor Babuschkin, Aidan Clark, Diego de Las Casas, Aurelia Guy, Chris Jones, James Bradbury, Matthew Johnson, Blake Hechtman, Laura Weidinger, Iason Gabriel, William Isaac, Ed Lockhart, Simon Osindero, Laura Rimell, Chris Dyer, Oriol Vinyals, Kareem Ayoub, Jeff Stanway, Lorrayne Bennett, Demis Hassabis, Koray Kavukcuoglu, and Geoffrey Irving.  Scaling language models: Methods, analysis and insights from training gopher, 2022.
  * Rombach et al. [2022] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer.  High-resolution image synthesis with latent diffusion models, 2022.
  * Sauer et al. [2021] Axel Sauer, Kashyap Chitta, Jens Müller, and Andreas Geiger.  Projected gans converge faster, 2021.
  * Sohl-Dickstein et al. [2015] Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli.  Deep unsupervised learning using nonequilibrium thermodynamics.  In _International conference on machine learning_ , pages 2256–2265. PMLR, 2015.
  * Song and Ermon [2019] Yang Song and Stefano Ermon.  Generative modeling by estimating gradients of the data distribution.  _Advances in neural information processing systems_ , 32, 2019.
  * Song et al. [2020] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole.  Score-based generative modeling through stochastic differential equations.  _arXiv preprint arXiv:2011.13456_ , 2020.
  * Su et al. [2023] Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu.  Roformer: Enhanced transformer with rotary position embedding, 2023.
  * Tian et al. [2024] Keyu Tian, Yi Jiang, Zehuan Yuan, Bingyue Peng, and Liwei Wang.  Visual autoregressive modeling: Scalable image generation via next-scale prediction, 2024.
  * Touvron et al. [2023] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurelien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample.  Llama: Open and efficient foundation language models, 2023.
  * van den Oord et al. [2018] Aaron van den Oord, Oriol Vinyals, and Koray Kavukcuoglu.  Neural discrete representation learning, 2018.
  * Vaswani et al. [2023] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin.  Attention is all you need, 2023.
  * Villegas et al. [2022] Ruben Villegas, Mohammad Babaeizadeh, Pieter-Jan Kindermans, Hernan Moraldo, Han Zhang, Mohammad Taghi Saffar, Santiago Castro, Julius Kunze, and Dumitru Erhan.  Phenaki: Variable length video generation from open domain textual description, 2022.
  * Wortsman et al. [2023] Mitchell Wortsman, Peter J. Liu, Lechao Xiao, Katie Everett, Alex Alemi, Ben Adlam, John D. Co-Reyes, Izzeddin Gur, Abhishek Kumar, Roman Novak, Jeffrey Pennington, Jascha Sohl-dickstein, Kelvin Xu, Jaehoon Lee, Justin Gilmer, and Simon Kornblith.  Small-scale proxies for large-scale transformer training instabilities, 2023.
  * Yu et al. [2022a] Jiahui Yu, Xin Li, Jing Yu Koh, Han Zhang, Ruoming Pang, James Qin, Alexander Ku, Yuanzhong Xu, Jason Baldridge, and Yonghui Wu.  Vector-quantized image modeling with improved vqgan, 2022a.
  * Yu et al. [2022b] Jiahui Yu, Yuanzhong Xu, Jing Yu Koh, Thang Luong, Gunjan Baid, Zirui Wang, Vijay Vasudevan, Alexander Ku, Yinfei Yang, Burcu Karagol Ayan, Ben Hutchinson, Wei Han, Zarana Parekh, Xin Li, Han Zhang, Jason Baldridge, and Yonghui Wu.  Scaling autoregressive models for content-rich text-to-image generation, 2022b.
  * Yu et al. [2024] Lijun Yu, José Lezama, Nitesh B. Gundavarapu, Luca Versari, Kihyuk Sohn, David Minnen, Yong Cheng, Vighnesh Birodkar, Agrim Gupta, Xiuye Gu, Alexander G. Hauptmann, Boqing Gong, Ming-Hsuan Yang, Irfan Essa, David A. Ross, and Lu Jiang.  Language model beats diffusion – tokenizer is key to visual generation, 2024.
  * Zhai et al. [2022] Xiaohua Zhai, Alexander Kolesnikov, Neil Houlsby, and Lucas Beyer.  Scaling vision transformers, 2022.

Supplementary

##  Appendix A Discrete regularizers

To select a simple but performant vector quantization method for our discrete
latent space, we compare the classic vector quantization (VQ) van den Oord et
al. [2018] without additional complexity like codebook reinitialization,
lookup free quantization (LFQ) Yu et al. [2024], and finite scalar
quantization (FSQ) Mentzer et al. [2023]. While training these autoencoders we
observe interesting differences in training dynamics with multiple crossing
points between FSQ and LFQ for certain metrics. We present those in Figure 8
where we can see that FSQ often takes the lead in the beginning phases of
training but eventually gives it up to LFQ. We can also see that both of these
methods outperform classic VQ which struggles without additional aids.

Figure 8: Perceptual reconstruction metrics for various discrete
regularization methods. Classic vector quantization (VQ) struggles without
tricks like codebook reinitialization. LFQ and FSQ have different training
dynamics, often trading the lead in the beginning phases of training which is
highlighted by the red X’s.


# Quasi-geostrophic projection of rotating shallow-water equations

Louis Thiry1,2  Long Li1  louis[dot]thiry[at]inria[dot]fr Etienne Mémin1
Guillaume Roullet3  1Equipe ODYSSEY, INRIA Rennes, IRMAR Rennes, Université
Rennes 1, Rennes, France 2Equipe ANGE, INRIA Paris, Laboratoire Jacques-Louis
Lions, Sorbonne Universite, Rennes, France 3Université de Bretagne Occidentale
, CNRS, IRD, Ifremer, Laboratoire d’Océanographie Physique et Spatiale (LOPS),
IUEM, Brest, France

###### Abstract

With a new unifying model for layered rotating shallow-water (RSW) and quasi-
geostrophic (QG) equations, this paper sheds light on the relation between
these two sets of equations. We propose here a new formulation of the quasi-
geostrophic equations as a projection of the rotating shallow-water equations.
This QG formulation uses the same prognostic variables as RSW, namely velocity
and layer thickness, restoring the proximity of these two sets of equations.
It allows direct access to the ageostrophic velocities hidden in geostrophic
velocities resolved by the QG equations. It opens the path of studying the
difference between QG and RSW using the same underlying numerical
discretization. We illustrate this new possibility on a vortex shear
instability and a double-gyre configuration.

###### keywords:

Quasi-geostrophic, Rotating shallow-water.

##  1 Introduction

Large-scale ocean models offers a natural trade-off between complexity and
realism. On the one side, the primitive equations (PE) are realistic enough to
be used for climate simulations and real-world data assimilation. However,
they are rather complex as they describe the coupling between equation of
motions (mass and momentum) and conservation of tracers via the equation of
state following thermodynamics laws. On the other side, the barotropic
planetary geostrophic equations can be efficiently solved with a single
prognostic variable (the thickness), yet they ignore thermodynamics and
vertical variations.

In between, there exists some approximate models of the PE such as the multi-
layer rotating shallow-water (RSW) and the multi-layer quasi-geostrophic (QG)
models. The latter can be either derived from the former using an asymptotic
approach (Pedlosky, 2013; Vallis, 2017) or considered as a vertical
discretization of the continuously stratified QG model derived from the PE.
These layered models describe the dynamics of vertically stratified flow in
isentropic (or isopycnal) coordinates and only require solving the horizontal
momentum and mass equations. For instance, the multi-layer RSW model can
fairly reproduce the ocean dynamics in the Gulf-Stream region with solely five
vertical levels (Hurlburt & Hogan, 2000). The multi-layer QG equations are
widely adopted for the development of determistic mesoscale eddies
parameterizations (Marshall et al., 2012; Jansen & Held, 2014; Fox-Kemper et
al., 2014; Ryzhov et al., 2020; Uchida et al., 2022) as well as stochastic
ones (Grooms et al., 2015; Zanna et al., 2017; Bauer et al., 2020). These two
sets of equations can be used as research tools as their numerical integration
is light enough to run on laptop computers.

Although the multi-layer QG equations are derived from the multi-layer RSW
equations, their relationship is still not clear from a numerical point of
view. Indeed, they are usually written with different progonstic variables.
The prognostic variables of the RSW system are the horizontal velocity
$(\mathbf{u},\mathbf{v})$ and the layer thickness $\mathbf{h}$, whereas one
can formulate the QG system with a single prognostic variable, typically the
potential vorticity $\mathbf{q}$ from which one can diagnose the
streamfunction $\boldsymbol{\psi}$ and the pressure $\mathbf{p}$. This
difference breaks the conceptual continuity between these two equation sets in
the ocean model hierarchy. It also brings some undesirable consequences in
practice. For instance, many eddy parameterizations (e.g. Bachman, 2019; Li et
al., 2023) yield different formulations and/or discretizations when applied to
the RSW model, using $(\mathbf{u},\mathbf{v},\mathbf{h})$ as the prognostic
variables, or applied to the QG model, using $\mathbf{q}$ as the prognostic
variable. Moreover, it is not straightforward to use the same discrete schemes
for these two different models in order to compare them under the same
configuration.

In the present work, we propose to reformulate the multi-layer QG model using
$(\mathbf{u},\mathbf{v},\mathbf{h})$ as the prognostic variable. The QG
equations can be written as a projection acting on the RSW equations. This is
the main novelty of our paper. Using this projection approach, one can build a
QG discretization on top of any RSW discretization. Moreover, this formulation
provides direct access to the ageostrophic velocity that is hidden in the
standard QG model formulation, yet contributes to QG dynamics.

For the numerical experiments, we adopt the RSW discretization proposed by
Roullet & Gaillard (2022) which relies on the vector invariant formulation and
uses high-order WENO reconstructions to provide implicit dissipation. The
discretized QG model uses the same dynamical core than the multi-layer RSW
model. The only modification is to add the projection operator. We end up with
a compact, efficient, and CPU/GPU portable Python code using the PyTorch
library (Paszke et al., 2019). We first use a test case of vortex shear
instability to investigate the similarities and differences of the QG and RSW
solutions according to different Rossby numbers. We then study the solutions
produced by QG and RSW on an idealized double-gyre configuration.

This paper is organized as follows. In Section 2, we briefly recall the
derivation of the multi-layer RSW and QG equations. In Section 3, we present
our projected QG formulation.In Section 4, we numerically test our new
formulation on two different configurations. We conclude and evoke further
perspectives in Section 5.

##  2 Multi-layer RSW and QG equations

In this section, we first review the governing equations of the RSW system,
then explain briefly the QG scaling of RSW equation and give the usual multi-
layer QG equations using $\mathbf{q}$ as the prognostic variable.

###  2.1 Multi-layer RSW equations

Figure 1: Vertical cross-section of a three-layer RSW model displaying layer
thicknesses $H_{i}+h_{i}$ and interface heights $\eta_{i}$

The stratification of a multilayer RSW model consists of a stack of $n$
isopycnal layers, as illustrated in Figure 1 with $n=3$ layers. By convention
we index the layers from $i=1$ for the top layer to $i=n$ for the bottom one.
These layers have a uniform reference thickness $H_{i}$ and a density
$\rho_{i}$. As shown in Figure 1, the total thickness of a layer $i$ is the
sum of the reference thickness $H_{i}$ and the thickness anomaly $h_{i}(x,y)$.
The interface vertical displacment $\eta_{i}(x,y)$ and the hydrostatic
pressure $p_{i}(x,y)$ are given by

$\eta_{i}=\sum_{j=i}^{n}h_{j}\ ,\quad
p_{i}=\rho_{1}\sum_{j=1}^{i}g_{j}^{\prime}\eta_{j}$ (1)

with, for $i>1$, the reduced gravities
$g^{\prime}_{i}=g(\rho_{i}-\rho_{i-1})/\rho_{1}$ and for the top layer
$g^{\prime}_{1}=g$.

Using the following vector notation

$\displaystyle\mathbf{H}$ $\displaystyle=\left(H_{1},\ \ldots,\ H_{n}\right),$
(2a) $\displaystyle\mathbf{h}$ $\displaystyle=\left(h_{1}(x,y),\ \ldots,\
h_{n}(x,y)\right),$ (2b) $\displaystyle\boldsymbol{\eta}$
$\displaystyle=\left(\eta_{1}(x,y),\ \ldots,\ \eta_{n}(x,y)\right),$ (2c)
$\displaystyle\mathbf{p}$ $\displaystyle=\left(p_{1}(x,y),\ \ldots,\
p_{n}(x,y)\right),$ (2d)

we can rewrite the linear relations between $\mathbf{h},\boldsymbol{\eta}$ and
$\mathbf{p}$ in a compact form

$\displaystyle\mathbf{h}$
$\displaystyle=\mathop{diag}(\mathbf{H})A\mathbf{p},$ (3a)
$\displaystyle\boldsymbol{\eta}$ $\displaystyle=C\mathbf{h}$ (3b)
$\displaystyle\mathbf{p}$ $\displaystyle=M\mathbf{h},$ (3c) with the matrices
$\displaystyle A$
$\displaystyle=\frac{1}{\rho_{1}}\begin{bmatrix}\frac{1}{H_{1}g_{1}^{\prime}}+\frac{1}{H_{1}g_{2}^{\prime}}&\frac{-1}{H_{1}g_{2}^{\prime}}&.&.\\\
\frac{-1}{H_{2}g_{2}^{\prime}}&\frac{1}{H_{2}}\left(\frac{1}{g_{2}^{\prime}}+\frac{1}{g_{3}^{\prime}}\right)&\frac{-1}{H_{2}g_{3}^{\prime}}&.\\\
.&\ldots&\ldots&\ldots\\\
.&.&\frac{-1}{H_{n}g_{n}^{\prime}}&\frac{1}{H_{n}g_{n}^{\prime}}\\\
\end{bmatrix}\ ,$ (3d) $\displaystyle B$
$\displaystyle=\rho_{1}\begin{bmatrix}g^{\prime}_{1}&0&0&\ldots&0\\\
g^{\prime}_{1}&g^{\prime}_{2}&0&\ldots&0\\\ .&.&.&.&.\\\
g^{\prime}_{1}&g^{\prime}_{2}&g^{\prime}_{3}&\ldots&g^{\prime}_{n}\end{bmatrix},\
\ C=\begin{bmatrix}1&1&\ldots&1&1\\\ 0&1&\ldots&1&1\\\ .&.&.&.&.\\\
0&0&\ldots&0&1\end{bmatrix},$ (3e)

and $M=BC$. We then introduce the velocity $(\mathbf{u},\mathbf{v})$, the
kinetic energy $\mathbf{k}$, and the relative vorticity $\boldsymbol{\omega}$

$\displaystyle\mathbf{u}$
$\displaystyle=\left(u_{1}(x,y),\ldots,u_{n}(x,y)\right),$ (4a)
$\displaystyle\mathbf{v}$
$\displaystyle=\left(v_{1}(x,y),\ldots,v_{n}(x,y)\right),$ (4b)
$\displaystyle\mathbf{k}$ $\displaystyle=(\mathbf{u}^{2}+\mathbf{v}^{2})/2,$
(4c) $\displaystyle\boldsymbol{\omega}$
$\displaystyle=\partial_{x}\mathbf{v}-\partial_{y}\mathbf{u}\ .$ (4d)

With these variables, the multi-layer RSW equations read

$\displaystyle\partial_{t}\,\mathbf{u}=(\boldsymbol{\omega}+f)\mathbf{v}-\partial_{x}(\mathbf{p}+\mathbf{k}),$
(5a)
$\displaystyle\partial_{t}\,\mathbf{v}=-(\boldsymbol{\omega}+f)\mathbf{u}-\partial_{y}(\mathbf{p}+\mathbf{k}),$
(5b)
$\displaystyle\partial_{t}\,\mathbf{h}=-\mathbf{H}(\partial_{x}\mathbf{u}+\partial_{y}\mathbf{v})-\partial_{x}(\mathbf{u}\mathbf{h})-\partial_{y}(\mathbf{v}\mathbf{h})\,,$
(5c)

where $f$ is the Coriolis parameter. They can be formulated in the compact
form

$\partial_{t}\mathbf{X}=F(\mathbf{X})\,,$ (6a) with
$\mathbf{X}=\left(\mathbf{u},\mathbf{v},\mathbf{h}\right)^{T}$, the state
variable, and $F\begin{pmatrix}\mathbf{u}\\\ \mathbf{v}\\\
\mathbf{h}\end{pmatrix}=\begin{pmatrix}(\partial_{x}\mathbf{v}-\partial_{y}\mathbf{u}+f)\mathbf{v}-\partial_{x}\left(M\mathbf{h}+(\mathbf{u}^{2}+\mathbf{v}^{2})/2\right)\\\
-(\partial_{x}\mathbf{v}-\partial_{y}\mathbf{u}+f)\mathbf{u}-\partial_{y}\left(M\mathbf{h}+(\mathbf{u}^{2}+\mathbf{v}^{2})/2\right)\\\
-\mathbf{H}(\partial_{x}\mathbf{u}+\partial_{y}\mathbf{v})-\partial_{x}(\mathbf{u}\mathbf{h})-\partial_{y}(\mathbf{v}\mathbf{h})\end{pmatrix}\,,$
(6b) the RSW model operator.

###  2.2 QG scaling of RSW equations

The QG model is derived from the RSW model under two scaling assumptions:
$Ro\ll 1$ and $Bu\sim 1$ (McWilliams, 2006; Zeitlin, 2018) where $Ro$ is the
Rossby number and $Bu$ the Burger number. The consistency of the QG scaling
imposes a beta-plane approximation for the Coriolis parameter $f=f_{0}+\beta
y$, where $\beta$ is the meridional Coriolis parameter gradient. $Ro\ll 1$
implies that the velocity is close to the geostrophic balance, and
specifically that the ageostrophic correction is $O(Ro)$. We define
$(\mathbf{u}_{\rm g},\mathbf{v}_{\rm g})$, the geostrophic velocity that obeys
the geostrophic balance

$\displaystyle-f_{0}\mathbf{u}_{\rm g}=\partial_{y}\mathbf{p},\ \quad
f_{0}\mathbf{v}_{\rm g}=\partial_{x}\mathbf{p}\,,$ (7a) and decompose the
velocity into the geostrophic and the ageostrophic velocity $(\mathbf{u}_{\rm
a},\mathbf{v}_{\rm a})$ $\displaystyle\mathbf{u}=\mathbf{u}_{\rm
g}+\mathbf{u}_{\rm a},\quad\mathbf{v}=\mathbf{v}_{\rm g}+\mathbf{v}_{\rm
a}\,.$ (7b)

$Bu\sim 1$ and $Ro\ll 1$ imply that $\mathbf{h}/\mathbf{H}\sim O(Ro)$.
Following the approach of Holton (1973); McWilliams (2006); Cushman-Roisin &
Beckers (2011); Zeitlin (2018), we do not decompose
$\mathbf{p},\mathbf{h},\boldsymbol{\eta}$ in a geostrophic and an ageostrophic
part. This would require an extra assumption to decide what is the
ageostrophic part of the stratification. This is unecessary to derive the QG
model. There is therefore a single mass variable and two velocities per layer.

With these assumptions, the multi-layer QG equations can be written in terms
of three equations for the prognostic variables $(\mathbf{u}_{\rm
g},\mathbf{v}_{\rm g},\mathbf{h})$ with the right-hand side terms that depend
on the ageostrophic velocities $(\mathbf{u}_{\rm a},\mathbf{v}_{\rm a})$

$\displaystyle\partial_{t}\,\mathbf{u}_{g}=(\boldsymbol{\omega}_{g}+f_{0}+\beta
y)\mathbf{v}_{g}+f_{0}\mathbf{v}_{a}-\partial_{x}(M\mathbf{h}+\mathbf{k}_{g}),$
(8a)
$\displaystyle\partial_{t}\,\mathbf{v}_{g}=-(\boldsymbol{\omega}_{g}+f_{0}+\beta
y)\mathbf{u}_{g}-f_{0}\mathbf{u}_{a}-\partial_{y}(M\mathbf{h}+\mathbf{k}_{g}),$
(8b)
$\displaystyle\partial_{t}\,\mathbf{h}=-\mathbf{H}(\partial_{x}\mathbf{u}_{g}+\partial_{y}\mathbf{v}_{g})-\partial_{x}(\mathbf{u}_{g}\mathbf{h})-\partial_{y}(\mathbf{v}_{g}\mathbf{h})-\mathbf{H}(\partial_{x}\mathbf{u}_{a}+\partial_{y}\mathbf{v}_{a}).$
(8c)

Note that using the geostrophic balance (7a), one can simplify the above
equations into

$\displaystyle\partial_{t}\,\mathbf{u}_{g}=(\boldsymbol{\omega}_{g}+\beta
y)\mathbf{v}_{g}+f_{0}\mathbf{v}_{a}-\partial_{x}\mathbf{k}_{g}\ ,$ (9a)
$\displaystyle\partial_{t}\,\mathbf{v}_{g}=-(\boldsymbol{\omega}_{g}+\beta
y)\mathbf{u}_{g}-f_{0}\mathbf{u}_{a}-\partial_{y}\mathbf{k}_{g}\ ,$ (9b)
$\displaystyle\partial_{t}\,\mathbf{h}=-\partial_{x}(\mathbf{u}_{g}\mathbf{h})-\partial_{y}(\mathbf{v}_{g}\mathbf{h})-\mathbf{H}(\partial_{x}\mathbf{u}_{a}+\partial_{y}\mathbf{v}_{a}).$
(9c)

In this form the QG equations are still in appearance very close to the RSW
equations, but this is misleading for there is not three degrees of freedom
per layer, but only one. The two components of the geostrophic velocity are
slaved to $\mathbf{h}$ via (7a) and (3a). Another difficulty with this system
is that, though well posed, it is highly implicit. There is no obvious way to
see how to determine the ageostrophic velocity. In practice, and in particular
for numerical models, the model equations are rewritten in a more explicit
form. This form reveals a hidden variable of the system: the potential
vorticity (PV)

$\mathbf{q}=f_{0}+\beta\,y+\boldsymbol{\omega}_{g}-f_{0}\frac{\mathbf{h}}{\mathbf{H}}\,.$
(10)

###  2.3 Multi-layer QG equations

The multi-layer QG equations are obtained upon derivating in time (10) and
using (9a)-(9c)

$\displaystyle\partial_{t}\,\mathbf{q}=-\partial_{x}(\mathbf{u}_{\rm
g}\mathbf{q})-\partial_{y}(\mathbf{v}_{\rm g}\mathbf{q}),$ (11a)
$\displaystyle\Delta\mathbf{p}-f_{0}^{2}A\mathbf{p}=f_{0}\mathbf{q}-f_{0}\beta
y,$ (11b) $\displaystyle-f_{0}\mathbf{u}_{\rm g}=\partial_{y}\mathbf{p},\
\quad f_{0}\mathbf{v}_{\rm g}=\partial_{x}\mathbf{p},$ (11c)

where $\Delta=\partial_{xx}^{2}+\partial_{yy}^{2}$ denotes the horizontal
Laplacian operator and $A$ is the vertical discretization of the stretching
operator introduced in Equation (3d). The PV $\mathbf{q}$ is the sole
prognostic variable (11a) of the model. All the other model variables are
deduced from it via the diagnostic relations (11b) for $\mathbf{p}$ and (11c)
for $\mathbf{u}$ and $\mathbf{v}$. The QG model is written in
$(\mathbf{q},\,\mathbf{p})$ variables and $\mathbf{p}/f_{0}$ is the vector of
stream-function for each layer. Numerical implementations of the QG model
usually rely on this system of equations. However, this system is now very
different from the RSW system of equations. One aspect is the dispareance of
$\mathbf{u}_{\rm a}$ from the equations. This does not mean that
$\mathbf{u}_{\rm a}$ and $\mathbf{v}_{\rm a}$ vanish, rather, it means that
$\mathbf{u}_{\rm a}$ and $\mathbf{v}_{\rm a}$ can be ignored if one is
interested only in the evolution of $\mathbf{p}$ and $\mathbf{q}$. The large
difference between the model equations makes the comparison between them
harder. In particular, and because of that, the numerical versions of QG and
RSW models generally have completely different implementations, with different
numerics on all aspects. We now show how to restore the proximity between the
two.

##  3 Quasi-geostrophic model as projected rotating shallow-water

In this section, we propose a new formulation of the QG equation using
$(\mathbf{u},\mathbf{v},\mathbf{h})$ as state variables, the non-linear RSW
operator $F$ defined in (6b), and a linear projection operator $P$.
Strikingly, to our knowledge, this projection relation has been mostly under
the radar except for one mention in (Charve, 2004).

###  3.1 QG as projected RSW

To derive this projection, we start from the QG equations (8a), (8b), and
(8c), and isolate the part controlled by $F$

$\partial_{t}\left(\begin{matrix}\mathbf{u}_{g}\\\ \mathbf{v}_{g}\\\
\mathbf{h}\end{matrix}\right)=F\left(\begin{matrix}\mathbf{u}_{g}\\\
\mathbf{v}_{g}\\\
\mathbf{h}\end{matrix}\right)+\left(\begin{matrix}f_{0}\mathbf{v}_{a}\\\
-f_{0}\mathbf{u}_{a}\\\
-\mathbf{H}(\partial_{x}\mathbf{u}_{a}+\partial_{y}\mathbf{v}_{a})\end{matrix}\right),$
(12)

in which several terms cancel due to the geostrophic balance. In the form (12)
the QG model can be interpreted as a RSW model forced by an ageostrophic
source term, the second term of the right-hand side. Let us now introduce the
PV linear operator $Q$

$Q\left(\begin{matrix}\mathbf{u}\\\ \mathbf{v}\\\
\mathbf{h}\end{matrix}\right)=\partial_{x}\mathbf{v}-\partial_{y}\mathbf{u}-f_{0}\frac{\mathbf{h}}{\mathbf{H}}\,.$
(13a) This operator is related to the PV $\mathbf{q}$ by
$Q\left(\begin{matrix}\mathbf{u}_{g}\\\ \mathbf{v}_{g}\\\
\mathbf{h}\end{matrix}\right)=\mathbf{q}-\beta y$ (13b) and makes the
contribution of the ageostrophic source term in (12) vanishes
$Q\left(\begin{matrix}f_{0}\mathbf{v}_{a}\\\ -f_{0}\mathbf{u}_{a}\\\
-\mathbf{H}(\partial_{x}\mathbf{u}_{a}+\partial_{y}\mathbf{v}_{a})\end{matrix}\right)=0.$
(13c)

Since $Q$ is independant of time, it commutes with the time derivative,
therefore

$\partial_{t}Q\left(\begin{matrix}\mathbf{u}_{g}\\\ \mathbf{v}_{g}\\\
\mathbf{h}\end{matrix}\right)=Q\circ F\left(\begin{matrix}\mathbf{u}_{g}\\\
\mathbf{v}_{g}\\\ \mathbf{h}\end{matrix}\right).$ (14)

This equation is simply a reformulation of (11a), the conservation of PV.
Notice that because $Q$ decreases the size of the state vector from three to
one variable, we need a reverse operation to retrieve the three variables
$(\mathbf{u}_{g},\mathbf{v}_{g},\mathbf{h})$. This is done by introducing the
geostrophic operator $G$

$G\left(\mathbf{p}\right)=\begin{pmatrix}-\partial_{y}\mathbf{p}\\\
\partial_{x}\mathbf{p}\\\ f_{0}\,\mathbf{H}A\mathbf{p}\end{pmatrix}\,.$ (15a)
$G$ is a linear operator with uniform and time independant coefficients. The
composition of the two operators $Q$ and $G$ results in $(Q\circ
G)\left(\mathbf{p}\right)=\Delta\mathbf{p}-f_{0}^{2}A\mathbf{p}\,$ (15b)

which is the QG elliptic operator (11b) that relates $\mathbf{p}$ to
$\mathbf{q}$. This operator is invertible (Dutton, 1974; Bourgeois & Beale,
1994), therefore $(Q\circ G)^{-1}$ is well defined. Finally, we introduce the
QG projection operator $P$

$\displaystyle P=G\circ(Q\circ G)^{-1}\circ Q\,.$ (16)

$P$ is a projection because $P\circ P=P$, which is easily proven

$\displaystyle P\circ P=G\circ(Q\circ G)^{-1}\circ(Q\circ G)\circ(Q\circ
G)^{-1}\circ Q=P\,.$ (17)

By construction, $P$ preserves the geostrophic state
$\mathbf{X}_{g}=\left(\mathbf{u}_{g},\mathbf{v}_{g},\mathbf{h}\right)^{T}$,
namely $P(\mathbf{X}_{g})=\mathbf{X}_{g}$ . Indeed, since $\mathbf{X}_{g}$ is
geostrophic, we have

$\displaystyle Q(\mathbf{X}_{g})$ $\displaystyle=\mathbf{q}-\beta\,y\,,$ (18)
$\displaystyle(Q\circ G)^{-1}(\mathbf{q}-\beta\,y)$
$\displaystyle=\mathbf{p}/f_{0}\,,$ (19) $\displaystyle G(\mathbf{p}/f_{0})$
$\displaystyle=\mathbf{X}_{g}\,.$ (20)

By applying $P$ on (12), and noting that $P$ commutes with $\partial_{t}$, we
can formulate the multi-layer QG model in the form

$\partial_{t}\mathbf{X}_{g}=P\circ F\left(\mathbf{X}_{g}\right)\,.$ (21)

This form differs from the RSW model (6a) by the additional projection
operator $P$ acting on $F$. By projecting the RSW tendency on the geostrophic
manifold, $P$ ensures that the QG state remains in geostrophic balance.
Another way to express it, is to say that the QG model evolves under the
action of the RSW operator in which the ageostrophic tendency is removed.

###  3.2 Ageostrophic velocities

With this formulation, the ageostrophic velocity
$(\mathbf{u}_{a},\mathbf{v}_{a})$ has a simple closed form. Using equations
(12) and (21), we have

$\left(\begin{matrix}f_{0}\mathbf{v}_{a}\\\ -f_{0}\mathbf{u}_{a}\\\
-\mathbf{H}(\partial_{x}\mathbf{u}_{a}+\partial_{y}\mathbf{v}_{a})\end{matrix}\right)=P\circ
F\left(\begin{matrix}\mathbf{u}_{g}\\\ \mathbf{v}_{g}\\\
\mathbf{h}\end{matrix}\right)-F\left(\begin{matrix}\mathbf{u}_{g}\\\
\mathbf{v}_{g}\\\ \mathbf{h}\end{matrix}\right)\,.$ (22)

In a numerical model this requires no costly computations, only a basic
difference between the QG and the RSW tendencies.

###  3.3 Discussion

With this new projection formulation, one can formulate the QG system using
the same prognostic variables as the RSW system, namely the horizontal
velocity $(\mathbf{u},\mathbf{v})$ and the layer thickness $\mathbf{h}$. This
restores the proximity between these two equation sets in the ocean model
hierarchy. We see four practical implications. First, given a RSW
discretization, one has simply to implement the projection to have the
companion QG model of this particular discretization. Second, this formulation
gives access to the ageostrophic velocity hidden in the QG equations at the
cost of a simple subtraction. This allows an a posteriori diagnostic of the
validity of the QG scaling. For a given geostrophic state
$(\mathbf{u}_{g},\mathbf{v}_{g},\mathbf{h})$ we can compute the corresponding
ageostrophic velocity $(\mathbf{u}_{a},\mathbf{v}_{a})$ and check that

$\frac{\mathbf{k}_{a}}{\mathbf{k}_{g}}\sim
Ro^{2}\,,\qquad\frac{\mathbf{h}}{\mathbf{H}}\sim Ro\ ,$ (23)

where $\mathbf{k}_{a}=(|\mathbf{u}_{a}|^{2}+|\mathbf{v}_{a}|^{2})/2$ and
$\mathbf{k}_{g}=(|\mathbf{u}_{g}|^{2}+|\mathbf{v}_{g}|^{2})/2$ are
respectively the ageostrophic and the geostrophic kinetic energies. When this
is is not the case, the QG approximation is not valid and the solution
produced by the QG model becomes questionable. Third, using the same variables
$(\mathbf{u},\mathbf{v},\mathbf{h})$ opens the route to formulate eddy
parameterizations, e.g Bachman (2019); Li et al. (2023), that work similarly
on the RSW and the QG models, which is not the case when different
formulations and/or discretizations are used. Four, this formulation might
give new insights and approaches for mathematical analysis of the QG
equations. One could look at the extension of recent results on the well-
posedness properties for a stochastic RSW model (Crisan & Lang, 2023) by
studying the effect of the QG projector on these properties.

Finally, this formulation presents an interesting analogy with the Leray
projector formulation of the Navier-Stokes equations. The Leray projector
enforces the incompressible constraint by filtering out the compressible sound
waves. It involves solving an elliptic Poisson equation for the pressure. The
QG projector presented here enforces the QG balance by filtering out the
gravity waves. It involves solving an elliptic Helmholtz equation for the PV.
It worth emphasizing that the QG balance is not a constraint added to the RSW
model. Indeed, the Hamiltonian structure of the QG model (Holm & Zeitlin,
1998) is very different from the RSW one. In particular, and contrary to the
pressure, the PV is not a Lagrange multiplier enforcing a constraint.

##  4 Numerical experiments

###  4.1 Numerical implementation

To implement our method, we need a RSW solver and a QG projector.

####  4.1.1 RSW discretization

We adopt the RSW solver111Available at https://github.com/pvthinker/pyRSW .
developed by Roullet & Gaillard (2022), that we re-implemented in PyTorch for
seamless GPU acceleration. We have verified that our implementation enables
reproducing their original results up to numerical precision. The key element
of the discretization is a fifth-order WENO upwinded reconstructions on both
the the mass flux and the nonlinear vortex-force term. This provides enough
mixing and dissipation, removing the need of an ad-hoc hyperviscous diffusion,
while ensuring good material conservation of PV (Roullet & Gaillard, 2022).

####  4.1.2 QG discretization

The QG projector requires an elliptic solver. Thiry et al. (2024) released an
efficient Python-PyTorch implementation of the QG equations. They solve the QG
elliptic equation with discrete sine transform implemented with PyTorch FFT,
which uses highly optimized MKL FFT on Intel CPUs and cuFFT on Nvidia GPUs. We
use their elliptic solver to solve the QG elliptic equation (11b). Once the
projector is available, the QG solver is implemented by simply adding the
projection step to the RSW solver. Numerically the two solvers differ by one
line of code, the projection step. All the other ingredients: variable
staggering, time discretization, core RSW equations are identitical.

As a result, we obtain a PyTorch implementation222Will be released at
publication time. which is concise ($\sim$ 800 lines of code), as close as
possible to the equations and which implements both the multi-layer QG and the
RSW equations with the same state variables
$(\mathbf{u},\mathbf{v},\mathbf{h})$.

###  4.2 Vortex shear instability

Figure 2:  Vortex shear instability solved with QG and SW, $Bu=1$. (Top to
bottom) Initial relative vorticity, $Ro=0.01$ and $Ro=0.5$. $Ro=0.01$, QG and
SW final relative vorticities and difference. $Ro=0.5$, QG and SW final
relative vorticities and difference. Unit $s^{-1}$.

To validate our new formulation, we first study a vortex shear instability and
compare its evolution in both the QG and the RSW models as $Ro$ increases for
$Bu=1$. The initial state is a perfectly shielded vortex composed of a core of
uniform vorticity $\omega_{1}$ surrounded with a ring of opposite sign uniform
vorticity $\omega_{2}$. The system involves two lengths: the core radius
$r_{0}$ and the vortex outer radius $r_{1}$. The ratio $\omega_{2}/\omega_{1}$
is such that the total circulation vanishes. This system is unstable and leads
to multipoles formation (Morel & Carton, 1994). The number of poles depends on
the ratio of the vortex radius to the core radius. We focus here on a tripole
formation case, viz. $r_{1}/r_{0}=1.4$. To promote the growth of the most
unstable mode we add a small mode 3 azymuthal perturbation. The experiments
are set up in dimensional form with a square domain of size $L_{x}\times
L_{y}=100\,\text{km}\times 100\,\text{km}$ on a f-plane. There is a single
layer of fluid with thickness $H=1\,\text{km}$. We assume no-flow and free-
slip boundary conditions. The acceleration of gravity is set to
$g=10\,\text{m}\,\text{s}^{-2}$. The vortex core has a radius
$r_{0}=10\,\text{km}$ and a positive vorticity, the vortex radius is
$r_{1}=14\,\text{km}$. We run the simulations on a $512\times 512$ grid, i.e.
200m spatial resolution. The Coriolis parameter is chosen such that $Bu=1$,
viz. $f_{0}=\sqrt{gH}/r_{0}$. The vorticity $\omega_{1}$ is set indirectly via
$u_{\rm max}$, the maximum velocity of the initial condition. We define the
Rossby number as $Ro=u_{\rm max}/(f_{0}r_{0})$. We impose $Ro$ and deduce
$u_{\rm max}$. We have tested four cases:
$Ro\in\\{0.01,\,0.05,\,0.1,\,0.5\\}$. In all cases we apply the QG projector
to the initial state prior starting the time integration. The consequences are
that the initial states are in geostrophic balance and differ only by a
scaling factor. As $Ro$ increases the instability develops faster, in
dimensional time. Therefore we present the results in the rescaled eddy-
turnover time $\tau$, defined as the inverse of the $\ell^{2}$-norm of the
initial vorticity, i.e. $\tau=1/\|\omega\|$. We integrate the simulations for
the time $T=10\,\tau$.

Figure 3: Time evolution of the normalized differences of relative vorticity
between the QG model and the RSW model with $Bu=1$ and different $Ro$.

Figure 2 shows how the initial state has evolved for the two extreme cases
$Ro=0.01$ and $Ro=0.5$. In the $Ro=0.01$ case the differences between the QG
and the RSW solutions are not perceptible to the naked eyes. This is
consistent with the QG scaling and the fact that the QG model is an asymptotic
limit of the RSW model. In the $Ro=0.5$ case the differences are of order one.
In the RSW model the vortex spins faster and the filaments of negative
vorticity are stabilized. The smoothness of the RSW solution at $t=10\,\tau$
might be surprising as we expect a certain amount of gravity waves generated
during the fast cyclo-geostrophic adjustment of the initial state. These
gravity waves are well present but at $t=10\,\tau$ they have bounced back and
forth several times along the boundary and are strongly scattered.
Interestingly the QG solution is exactly the same than in the $Ro=0.01$ case.
This is expected since the QG equations are scaling invariant, viz. the
evolution is invariant under the multiplication by a constant. However this is
also quite remarkable to recover this property in the numerical solutions
because the QG solver relies on the full RSW right-hand side. Another symmetry
of the QG model is the parity invariance. The solution should be invariant
under a sign change of the vorticity. In other words, cyclones and
anticyclones follow the same evolution, up to a change of rotation. This is
not at all the case for the RSW model. By flipping the sign of the initial
vorticity we have checked that the solutions do satisfy this property (not
shown). This means that the QG projector behaves well as it restores two
invariances, scaling and parity, that are absent in the RSW solver.

To assess the influence of $Ro$ on the time evolution, we define the
normalized difference $\delta=2||\omega_{\rm qg}-\omega_{\rm
sw}||/(||\omega_{\rm qg}||+||\omega_{\rm sw}||)$.

Figure 3 shows $\delta(t)$ for the four $Ro$ cases. The oscillations during
the $[0,\,4\,\tau]$ period are due to the gravity waves in the RSW case that
result from the imperfect initial balance. Interestingly, the shortening of
the oscillation period as $Ro$ decreases is due to the time rescaling by
$\tau$. In dimensional time, the periods are the same. A pratical consequence
is that the RSW experiment requires more time steps to reach $t=10\,\tau$ as
$Ro$ decreases, whereas for the QG experiment the number of time steps is
constant. After $t=4\,\tau$ the differences are dominated by the shear
instability developing on the vortex. The small difference seen in the
snapshots in the $Ro=0.01$ case have actually reached a plateau, meaning the
QG solution closely follows the RSW one. This confirms that the QG model is a
good simplified model when the scaling assumption holds. On a longer time
scale and with a chaotic vortex dynamics, the solutions would diverge but on
such a simple setup, the solutions remain close. The time evolution of
$\delta$ is similar for all $Ro$ except $Ro=0.5$, because $\delta$ saturates
at one, the maximum value of this metric. In this case, this metric suggests
that two models predict a completely different solution. Looking back at the
snapshots, this seems exagerated. The order one difference is mostly due to
the difference in the timing, not to the difference in pattern. Depending on
the purpose, the QG solution might still be of interest as it still captures
the main phenomenology. Note that the QG solution could be made closer to the
RSW solution if the two models were started with a different initial state: a
cyclo-geostrophic balance for the RSW model and the associated projected state
for the QG model. The initial state would thus depend on $Ro$. For this
illustrative experiment we have prefered to stick to the same initial state,
up to a scaling constant.

###  4.3 Double-gyre configuration

To explore a richer phenomenology we have tested our new formulation on a
classical oceanic test case, the idealized double-gyre. The domain is a non-
periodic square ocean basin with $N=3$ layers. We assume free-slip boundary
conditions on each boundary. We apply a stationary and symmetric wind stress
$(\tau_{x},\tau_{y})$ with $\tau_{x}=-(\tau_{0}/\rho_{0})\cos(2\pi y/L_{y})$
and $\tau_{y}=0$ in the top layer and a linear drag in the bottom with drag
coefficient $\gamma$. The parameter values are given in Table 1 in Appendix A.

We study this configuration in an eddy-permitting resolution of 20 km, meaning
that the spatial resolution (20 km) is half of the largest baroclinic Rossby
radius (41 km). We expect QG and RSW simulations to produce strong western
boundary currents converging to the middle of the western boundary and an
eastward jet departing from the middle of the western boundary. Starting from
the rest state, this configuration requires about 100 years to spin-up and
have converged statistics (Hogg et al., 2005; Simonnet, 2005).

Figure 4:  QG upper-layer snapshot after 100 years of spin-up from rest state.
KE, geos., rel. stand respectively for kinetic energy, geostrophic and
relative.

In Figure 4 we present snapshots of key quantities of the upper layer after
100 years of integration. The solution exhibits the expected properties of
this kind of setup: a double-gyre system (Fig. 4e) separated by a meandering
eastward jet (Fig. 4a,d) emanating from the western boundary, two strong and
narrow western boundary currents (Fig. 4b) feeding the eastward jet, mesoscale
turbulence everywhere and intensified near the jet (Fig. 4c,d), Rossby waves
propagation (Fig. 4e,f). The local Rossby number defined as $\omega/f_{0}$
peaks at 0.2 which is not small but not large enough to dismiss the solution
either. Interestingly the layer thickness perturbation $h_{1}$ is not small at
all compared to the reference depth $H_{1}=400\,\text{m}$.

Figure 5: QG upper-layer statistics over 100 years after 100 years spin-up
from rest state.

To compliment this one snapshot, we present statistics estimated over the
years [100-200] after the initial state and using one snapshot every 10 days.
We decompose the geostrophic kinetic energy into its mean and its eddy part.
We also compute the $k_{a}/k_{g}$ ratio. The results are shown in Figure 5 for
the upper layer. The statistics confirm the presence of the strong western
boundary currents (Fig. 5a) as well as the strong eastward jet reaching the
middle of the domain (Fig. 5a,d). A remarkable feature of the solution is the
symmetry of all the quantities with respect to the central latitude. This is
due to both the symmetry of the forcing and the fact that the QG model has a
cyclone/anticyclone symmetry which prevents the model to break the symmetry of
the forcing. From the ratio $k_{a}/k_{g}$ and $h/H$ (Fig. 5c,f) we can test
the validity of the QG scaling assumptions (Eq. 23). The colorscale has been
adjusted so that the white intermediate color corresponds to $Ro=0.1$. The red
areas indicate where the QG scaling assumptions are not respected. Counter-
intuively the central jet region, where the kinetic energy is the largest, is
where the QG scaling is the best. The worst regions are the gyres centers
because of too large thickness deviations, and along the boundaries because of
too large ageostrophic velocities.

Figure 6: (Left) QG upper-layer relative vorticity after 100 years spin-up,
(middle) RSW upper-layer relative vorticity and (right) velocity divergence
after 2 years additional spin-up. Unit of $f_{0}$.

Finally, we compare the QG solution with its companion RSW one. But, because
of the free surface, and the presence of fast barotropic gravity waves, the
RSW solver requires a much smaller time step, typically 200 times smaller.
This 200 factor corresponds to the ratio $\sqrt{gH}/\max(u)$. Integrating from
the rest state over 200 years would be a waste of computational ressources. A
compromise would be to use a barotropic-baroclinic time splitting (Higdon & De
Szoeke, 1997) or an implicit stepping of the free-surface (Roullet & Madec,
2000) but this goes beyond the scope of the present study. Instead, we ran the
RSW solver starting from the year 200 of the QG solution, and we integrate it
over 2 years. Figure 6 compares QG and RSW on a snapshot of vorticity in the
upper layer. The central jet has now a southward component. The two gyres are
no longer symmetric. The mesoscale turbulence has intensified in the Northern
gyre and weakened in the Southern gyre. Interestingly if we apply the QG
projector on the RSW state the resulting state remains very close (Fig. 6c)
albeit with damped fluctuations. The QG projector tends to dampen the short
scales but, since it is applied on the RSW model tendency, it does not affect
the QG state, which does not present this dampening as we can check from a
visual comparison between Fig. 6a and Fig. 6b. The closeness suggests that
while the QG solution might be locally tangent to the RSW solution, when
integrated over a long time, the differences build up and yield a quite
different state.

##  5 Conclusion

In this paper we have shown that the QG model can be formulated as a projected
RSW model, namely by applying the QG projector $P$ (defined in Equation 16) to
the RSW tendency. This allows to manipulate a QG model with the same variables
$(\mathbf{u},\,\mathbf{v},\,\mathbf{h})$ than a RSW model, instead of the more
usual $(\mathbf{p},\,\mathbf{q})$ variables. This helps better understand the
proximity and the differences between the two models. It also allows to have a
numerical model capable of integrating the two models with the same numerics.
We have tested on a vortex shear instability that the resulting QG model has
all the desired properties of symmetry: scaling and parity, that are absent in
the RSW model. We have shown that a QG model implemented in this way recovers
the expected features of a double-gyre experiment. A direct consequence of the
approach is the capability to compare the RSW and the QG solutions.

This new formulation opens the path for studies of differences between QG and
RSW equations, for example study the stability of the geostrophic equilibrium
or the study of spontaneous imbalance. It might be very beneficial for data
assimilation applications as it allows to easily switch between RSW and QG
models, or even to combine them. One could for example run a first data
assimilation algorithm, e.g. ensemble Kalman filter (Evensen, 2003), using the
projected QG to capture the geostrophic dynamics, and then to refine the
obtained solution by switching to the RSW model.

[Acknowledgements] The authors acknowledge the support of the ERC EU project
856408-STUOD. The authors would like to thank Anne-Laure Dalibard for pointing
them to Frederic Charve’s Thesis, and Frederic Charve for interesting
discussion.

[Funding]ERC EU project 856408-STUOD. OcéanIA INRIA Challenge project

[Declaration of interests]The authors report no conflict of interest.

[Data availability statement]Code to reproduce results will be released at
publication time.

[Author contributions] Original idea by GR. LT, LL, and GR derived the math.
LT and LL implemented the software supervised by GR. LT ran the numerical
experiments and produced the figures. LT and GR wrote the article. LL, and EM
helped the redaction with suggestions and corrections. EM provided the
funding.

##  Appendix A Double-gyre configuration

We provide the physical and numerical parameters used in the double-gyre
configuration in Table 1.

Parameter | Value | Description
---|---|---
$L_{x},L_{y}$ |  $5120,5120$ km | Domain size
$H_{k}$ |  $400,1100,2600$ m | Mean layer thickness
$g_{k}^{\prime}$ |  $9.81,0.025,0.0125$ m s-2 | Reduced gravity
$\gamma$ |  $3.6\ 10^{-8}$ s-1 | Bottom drag coefficient
$\tau_{0}$ |  $0.08$ N m-2 | Wind stress magnitude
$\rho$ |  $1000$ kg m-3 | Ocean density
$f_{0}$ |  $9.375\ 10^{-5}$ s-1 | Mean Coriolis parameter
$\beta$ |  $1.754\ 10^{-11}$ (m s)-1 | Coriolis parameter gradient
$L_{d}$ |  $41,25$ km | Baroclinic Rossby radii
$n_{x},n_{y}$ | $256,256$ | Grid size
$dt$ | 4000 s | Time-step
Table 1: Parameters of the idealized double-gyre configuration

There no viscosity nor diffusion coefficient as the grid-scale dissipation is
implicitely handled by the upwinded WENO reconstruction of the masse flux and
the vortex-force (see Roullet & Gaillard, 2022).

## References

  * Bachman (2019) Bachman, S.D. 2019 The GM+E closure: a framework for coupling backscatter with the Gent and McWilliams parameterization. Ocean Modelling 136, 85–106.
  * Bauer et al. (2020) Bauer, Werner, Chandramouli, Pranav, Li, Long & Mémin, Etienne 2020 Stochastic representation of mesoscale eddy effects in coarse-resolution barotropic models. Ocean Modelling 151, 101646.
  * Bourgeois & Beale (1994) Bourgeois, A.J. & Beale, J.T. 1994 Validity of the quasigeostrophic model for large-scale flow in the atmosphere and ocean. SIAM Journal on Mathematical Analysis 25 (4), 1023–1068.
  * Charve (2004) Charve, Frédéric 2004 Etude de phénomenes dispersifs en mécanique des fluides géophysiques. PhD thesis, Palaiseau, Ecole polytechnique.
  * Crisan & Lang (2023) Crisan, Dan & Lang, Oana 2023 Well-posedness properties for a stochastic rotating shallow water model. Journal of Dynamics and Differential Equations pp. 1–31.
  * Cushman-Roisin & Beckers (2011) Cushman-Roisin, B. & Beckers, J.M. 2011 Introduction to geophysical fluid dynamics: physical and numerical aspects. Academic Press.
  * Dutton (1974) Dutton, J.A. 1974 The nonlinear quasi-geostrophic equation: Existence and uniqueness of solutions on a bounded domain. Journal of Atmospheric Sciences 31 (2), 422–433.
  * Evensen (2003) Evensen, G. 2003 The ensemble Kalman filter: Theoretical formulation and practical implementation. Ocean Dynamics 53, 343–367.
  * Fox-Kemper et al. (2014) Fox-Kemper, B., Bachman, S., Pearson, B. & Reckinger, S. 2014 Principles and advances in subgrid modelling for eddy-rich simulations. Clivar Exchanges 19 (2), 42–46.
  * Grooms et al. (2015) Grooms, I., Majda, A.J. & Smith, K.S. 2015 Stochastic superparameterization in a quasigeostrophic model of the Antarctic Circumpolar Current. Ocean Modelling 85, 1–15.
  * Higdon & De Szoeke (1997) Higdon, R.L. & De Szoeke, R.A. 1997 Barotropic-baroclinic time splitting for ocean circulation modeling. Journal of Computational Physics 135 (1), 30–53.
  * Hogg et al. (2005) Hogg, A., Killworth, P.D., Blundell, J.R & Dewar, W.K. 2005 Mechanisms of decadal variability of the wind-driven ocean circulation. Journal of Physical Oceanography 35 (4), 512–531.
  * Holm & Zeitlin (1998) Holm, Darryl D & Zeitlin, Vladimir 1998 Hamilton’s principle for quasigeostrophic motion. Physics of fluids 10 (4), 800–806.
  * Holton (1973) Holton, J.R. 1973 An introduction to dynamic meteorology. American Journal of Physics 41 (5), 752–754.
  * Hurlburt & Hogan (2000) Hurlburt, H.E. & Hogan, P.J. 2000 Impact of 1/8 to 1/64 resolution on Gulf Stream model – data comparisons in basin-scale subtropical Atlantic ocean models. Dynamics of Atmospheres and Oceans 32 (3-4), 283–329.
  * Jansen & Held (2014) Jansen, M.F. & Held, I.M. 2014 Parameterizing subgrid-scale eddy effects using energetically consistent backscatter. Ocean Modelling 80, 36–48.
  * Li et al. (2023) Li, L., Deremble, B., Lahaye, N. & Mémin, E. 2023 Stochastic data-driven parameterization of unresolved eddy effects in a baroclinic quasi-geostrophic model. Journal of Advances in Modeling Earth Systems 15 (2), e2022MS003297.
  * Marshall et al. (2012) Marshall, D.P., Maddison, J.R. & Berloff, P.S. 2012 A framework for parameterizing eddy potential vorticity fluxes. Journal of Physical Oceanography 42 (4), 539–557.
  * McWilliams (2006) McWilliams, J.C. 2006 Fundamentals of geophysical fluid dynamics. Cambridge University Press.
  * Morel & Carton (1994) Morel, Y.G. & Carton, X.J. 1994 Multipolar vortices in two-dimensional incompressible flows. Journal of Fluid Mechanics 267, 23–51.
  * Paszke et al. (2019) Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L. & others 2019 Pytorch: An imperative style, high-performance deep learning library. Advances in Neural Information Processing Systems 32.
  * Pedlosky (2013) Pedlosky, Js 2013 Geophysical fluid dynamics. Springer Science & Business Media.
  * Roullet & Gaillard (2022) Roullet, G. & Gaillard, T. 2022 A fast monotone discretization of the rotating shallow water equations. Journal of Advances in Modeling Earth Systems 14 (2), e2021MS002663.
  * Roullet & Madec (2000) Roullet, G. & Madec, G. 2000 Salt conservation, free surface, and varying levels: a new formulation for ocean general circulation models. Journal of Geophysical Research: Oceans 105 (C10), 23927–23942.
  * Ryzhov et al. (2020) Ryzhov, E.A., Kondrashov, D., Agarwal, N., McWilliams, J.C. & Berloff, P. 2020 On data-driven induction of the low-frequency variability in a coarse-resolution ocean model. Ocean Modelling 153, 101664.
  * Simonnet (2005) Simonnet, E. 2005 Quantization of the low-frequency variability of the double-gyre circulation. Journal of Physical Oceanography 35 (11), 2268–2290.
  * Thiry et al. (2024) Thiry, L., Li, L., Roullet, G. & Mémin, E. 2024 Mqgeometry-1.0: a multi-layer quasi-geostrophic solver on non-rectangular geometries. Geoscientific Model Development 17 (4), 1749–1764.
  * Uchida et al. (2022) Uchida, T., Deremble, B. & Popinet, S. 2022 Deterministic model of the eddy dynamics for a midlatitude ocean model. Journal of Physical Oceanography 52 (6), 1133–1154.
  * Vallis (2017) Vallis, G.K. 2017 Atmospheric and oceanic fluid dynamics. Cambridge University Press.
  * Zanna et al. (2017) Zanna, L., Mana, P.P., Anstey, J., David, T. & Bolton, T. 2017 Scale-aware deterministic and stochastic parametrizations of eddy-mean flow interaction. Ocean Modelling 111, 66–80.
  * Zeitlin (2018) Zeitlin, V. 2018 Geophysical fluid dynamics: understanding (almost) everything with rotating shallow water models. Oxford University Press.