# Posterior Annealing:
Fast Calibrated Uncertainty for Regression

Uddeshya Upadhyay1 Jae Myung Kim1

Cordelia Schmid2,3 Bernhard Schölkopf${}^{\text{ }4}$ Zeynep Akata1,4

1University of Tübingen 2Inria 3Google Research 4Max Planck Institute for
Intelligent Systems, Tübingen

###### Abstract

Bayesian deep learning approaches that allow uncertainty estimation for
regression problems often converge slowly and yield poorly calibrated
uncertainty estimates that can not be effectively used for quantification.
Recently proposed post hoc calibration techniques are seldom applicable to
regression problems and often add overhead to an already slow model training
phase. This work presents a fast calibrated uncertainty estimation method for
regression tasks, called posterior annealing, that consistently improves the
convergence of deep regression models and yields calibrated uncertainty
without any post hoc calibration phase. Unlike previous methods for calibrated
uncertainty in regression that focus only on low-dimensional regression
problems, our method works well on a wide spectrum of regression problems. Our
empirical analysis shows that our approach is generalizable to various network
architectures including, multilayer perceptrons, 1D/2D convolutional networks,
and graph neural networks, on five vastly diverse tasks, i.e., chaotic
particle trajectory denoising, physical property prediction of molecules using
3D atomistic representation, natural image super-resolution, and medical image
translation using MRI images.

##  1 Introduction

Uncertainty estimation is an essential building block to provide
interpretability and secure reliability in modern machine learning systems
(Shafaei et al., 2018; Kläs & Vollmer, 2018; Varshney & Alemzadeh, 2017;
Hüllermeier & Waegeman, 2021) that offer intelligent solutions for numerous
real-world applications, ranging from medical analytics (Leibig et al., 2017;
Gillmann et al., 2021; Upadhyay et al., 2021b) to autonomous driving (Xu et
al., 2014; Shafaei et al., 2018; Besnier et al., 2021). Recent advances have
explored various formulations to provide accurate predictions and uncertainty
estimates for deep neural networks, as represented by Bayesian approaches (Gal
& Ghahramani, 2016; Kendall & Gal, 2017; Maddox et al., 2019), ensembles
(Lakshminarayanan et al., 2016), pseudo-ensembles (Mehrtash et al., 2020;
Franchi et al., 2020), and quantile regression (Romano et al., 2019; Yan et
al., 2018; Feldman et al., 2021) methods. However, these existing methods are
often computationally expensive – e.g., slow convergence rate during training
or inefficient inference cost due to multiple forward passes – while being
poorly calibrated for uncertainty estimates. Moreover, some of these methods
are proposed for low-dimensional regression tasks (Chung et al., 2021; Zhou et
al., 2021; Chen et al., 2021) (i.e., regressing a scalar value) and do not
scale for high-dimensional regression (i.e., regressing large matrices or
tensors). This paper presents a unified formulation to resolve these issues
for estimating fast, well-calibrated uncertainty in deep regression models for
a wide spectrum of regression problems.

We propose to revisit deep regression models trained via maximum likelihood
estimation (MLE), which assumes a Gaussian distribution over the regression
output and optimizes the negative log-likelihood to estimate the target and
uncertainty. Although such models can ensure low regression error (i.e., high
accuracy) and encapsulate the predictive uncertainty, they often converge
slowly at the beginning of training due to a flat gradient landscape. Further,
they may even risk gradient explosion caused by a steep gradient landscape
when reaching the optima (detailed in Section 3.1), leading to poorly
calibrated uncertainty estimates that do not offer credible interpretability
for the model.

To reshape the aforementioned ill-posed gradient landscape that causes slow
convergence and poorly calibrated uncertainty, we propose a novel posterior
annealing (POSTA) scheme for deep regression models, that alters the original
gradients by formulating a temperature-dependent posterior to be optimized
during the learning phase. In contrast to the standard posterior for
regression that enforces a fixed Gaussian distribution on the target, we
introduce a temperature hyperparameter to impose an evolving distribution.

The proposed temperature-dependent posterior brings crucial properties to
regression uncertainty. First, the multimodal distribution on regression
target ensures that at high residuals (between output and ground-truth,
occurring in the initial learning phase), the gradients are much larger than
the standard unimodal Gaussian distribution (explained in detail in Section 3
and Figure 1) leading to faster convergence at the beginning of the learning
phase. Second, we also anneal the learning rate over the course of training
along with the temperature that avoids gradient explosion towards the end of
the learning phase as the posterior distribution evolves to standard unimodal
distribution with sharp gradients at lower errors. Third, we construct the
temperature-dependent posterior such that the predicted uncertainty is
constrained to be calibrated at every step.

The standard unimodal distribution will face slow convergence in the beginning
and potential gradient explosion towards the end of the learning phase and
provides poorly calibrated uncertainty estimates. In contrast, our POSTA
method allows faster convergence and offers well-calibrated uncertainty
estimates for a wide spectrum of regressions. This also differs from
uncertainty regression methods that estimate the full quantile as they are
only shown to be effective on low-dimensional regression.

Contributions. We introduce a temperature-dependent posterior annealing scheme
for deep regression models with uncertainty that leads to faster model
convergence and offers calibrated uncertainty. We conduct comprehensive
evaluation on a wide variety of datasets, including chaotic particle
trajectory denoising, physical property prediction of molecules using 3D
atomistic representation, image super-resolution, and medical image
translation using MRI images.

##  2 Related Work

Deep neural networks (DNNs) typically estimate inaccurate uncertainty due to
their deterministic form that is insufficient for characterizing the accurate
confidence (Gal, 2016; Guo et al., 2017). Bayesian inference has been widely
studied to effectively estimate uncertainty. Directly performing Bayesian
inference on deep nonlinear networks is infeasible due to intractable
computations. Hence, approximate inference has been explored by variational
inference (Graves, 2011; Blundell et al., 2015; Daxberger et al., 2021; Maddox
et al., 2019) or MCMC-based approximation (Welling & Teh, 2011; Chen et al.,
2014). However, due to its approximation, the estimated uncertainty may fail
to follow the true uncertainty quantification (Lakshminarayanan et al., 2016).
Moreover, compared with typical DNNs, approximate Bayesian inference is
computationally more expensive and has slower convergence in practice. Non-
Bayesian methods have been proposed as an alternative. For instance, (Kendall
& Gal, 2017; Lakshminarayanan et al., 2016) modeled two terms, i.e. predictive
mean and variance, as an output of DNN to estimate the uncertainty directly
from the network’s output. Another line of work estimates the uncertainty in
the prediction in a non-parametric manner by estimating different quantiles
for a given input (Lin et al., 2021; Chen et al., 2021; Zhou et al., 2021;
Chung et al., 2021).

Calibrating the inaccurate uncertainty is another way to estimate accurate
uncertainty (Guo et al., 2017). In the regression task, calibration was first
defined in a quantile manner (Kuleshov et al., 2018). That is, the estimated
credible interval with confidence level $\alpha$ (e.g. 95%) is calibrated if
$\alpha$% of the ground-truth target is covered in that interval. There are
post-processing methods for regression calibration (Kuleshov et al., 2018;
Pearce et al., 2018; Tagasovska & Lopez-Paz, 2019). For instance, (Kuleshov et
al., 2018) introduced an auxiliary model to adjust the output of the pre-
trained model based on Platt-scaling, while others use Gaussian process (Song
et al., 2019) or maximum mean discrepancy (Cui et al., 2020). However, an
auxiliary model with enough capacity will always be able to recalibrate, even
if the predicted uncertainty is completely uncorrelated with the real
uncertainty (Laves et al., 2020). Recently, (Levi et al., 2019) extended the
definition of calibration where a regressor is well calibrated if the
predicted error is equal to the difference between the ground truth and the
predicted mean. Using this definition, (Laves et al., 2020) proposed unbiasing
the predicted error by optimizing a scaling factor in the post-processing
step. However, such methods often add overhead to an already slow model
training phase.

##  3 Methodology: Posterior Annealing

Our framework that we name Posterior Annealing (POSTA) belongs to the family
of models that are designed to predict a distribution for the outputs (Kendall
& Gal, 2017; Upadhyay et al., 2021c; a; 2022) and the model is trained via a
loss function derived from maximum likelihood estimation (MLE).

In this section, we first describe the problem formulation, discuss related
methods and their limitations in Section 3.1. We present posterior annealing
method that constructs temperature dependent posterior to learn faster,
better-calibrated regression uncertainty in Section 3.2, and analyze the
effects of temperature annealing in Section 3.3.

###  3.1 Background and Motivation

Let $\mathcal{D}=\\{(\mathbf{x}_{i},\mathbf{y}_{i})\\}_{i=1}^{i=N}$ be the
dataset that comprises of samples from domain $\mathbf{X}$ and $\mathbf{Y}$
(i.e., $\mathbf{x}_{i}\in\mathbf{X},\mathbf{y}_{i}\in\mathbf{Y},\forall i$),
where $\mathbf{X},\mathbf{Y}$ lies in $\mathbb{R}^{m}$ and $\mathbb{R}^{n}$,
respectively. The goal of a regression task is to learn a function
$\mathbf{\Psi}(\cdot;\theta):\mathbb{R}^{m}\rightarrow\mathbb{R}^{n}$
(parameterized by $\theta$) that maps the input $\mathbf{x}$ to the output
$\mathbf{y}$. Let $\mathbf{\hat{y}}_{i}:=\mathbf{\Psi}(\mathbf{x}_{i};\theta)$
be the estimate for the $\mathbf{y}_{i}$ and
$\mathbf{\epsilon}_{i}:=\mathbf{\hat{y}}_{i}-\mathbf{y}_{i}$ be the residual
between the prediction and the ground-truth. The optimal parameters
($\theta^{*}$) is learnt by minimizing the error (e.g., $\ell_{1}$ or
$\ell_{2}$ loss) between the prediction and ground-truth.

Figure 1:  Our temperature dependent loss from Eq4 as a function of residual
and the estimated standard deviation (orange surface). In the beginning of the
learning phase, the error is high and the gradient of our loss is higher
(orange) than the gradient for standard loss (blue). At the end of the
training (with close to zero error) the loss function is guided by the
standard loss with non-zero gradient, whereas the gradients from the
temperature dependent terms become zero.

The $\ell_{1}$/$\ell_{2}$ loss function to train regression models originate
by treating the residuals (i.e., $\epsilon_{i}$) as following the i.i.d
Laplace/Gaussian distribution. However, the i.i.d assumption will not capture
the heteroscedasticity and will not allow uncertainty estimation.

To estimate the uncertainty, the existing works (Kendall & Gal, 2017) relax
the i.i.d assumption and learn to model the heteroscedasticity as well. Such
models are learned by maximizing the likelihood. Assuming that residuals
follow Gaussian distribution, i.e.,
$\mathbf{\epsilon}_{i}\sim\mathcal{N}(0,\hat{\mathbf{\sigma}}_{i})$, the
likelihood, $P(\mathcal{D}|\theta)$, is a factored Gaussian distribution,
$P(\mathcal{D}|\theta)=\bm{\prod}_{i=1}^{i=N}\frac{1}{\sqrt{2\pi\hat{\mathbf{\sigma}}_{i}^{2}}}\exp({-\frac{|\hat{\mathbf{y}}_{i}-\mathbf{y}_{i}|^{2}}{2\hat{\mathbf{\sigma}}_{i}^{2}}}).$
the MLE estimates for the parameters are obtained by minimizing the negative-
log likelihood,

$\displaystyle-\log
P(\mathcal{D}|\theta)=\bm{\sum}_{i=1}^{i=N}\frac{\log{\hat{\mathbf{\sigma}}_{i}^{2}}}{2}+\frac{|\hat{\mathbf{y}}_{i}-\mathbf{y}_{i}|^{2}}{2\hat{\mathbf{\sigma}}_{i}^{2}}.$
(1)

The DNN is modified to output both the prediction (i.e., the mean of Gaussian)
as well as the uncertainty estimate (i.e., the variance of Gaussian) learned
using the above equation, i.e.,
$\mathbf{\Psi}(\mathbf{x}_{i};\theta)=\\{\hat{\mathbf{y}}_{i},\hat{\mathbf{\sigma}}_{i}\\}$.
While this method allows predicting the uncertainty estimates in single
forward pass post training, it has several downsides.

The blue surface in Figure 1 shows the loss from Equation 1 derived by taking
negative log of Gaussian likelihood, which consists of two variables: the
residual $\mathbf{y}_{i}-\hat{\mathbf{y}}_{i}$ (denoted by $u$) and the
standard deviation $\hat{\mathbf{\sigma}}_{i}$ (denoted by $v$). At the
beginning of the training phase, the residual between the prediction and the
ground-truth is large but the corresponding gradient at that point is small
(see Point a), leading to slower convergence towards optima. As the learning
progresses, the residual between prediction and ground-truth reduces
substantially with very high gradients leading to gradient explosion (see
Point b) if the learning rate is not tuned well. Together, this leads to
slower model convergence as gradients in the beginning are too small while the
learning rate would also have to be substantially smaller to avoid gradient
explosion later. Moreover, (Laves et al., 2020; Levi et al., 2019; Phan et
al., 2018) have shown that this method requires an additional post hoc
calibration phase as derived uncertainties are miscalibrated.

###  3.2 Constructing Temperature Dependent Posterior

To tackle the slow convergence issue while providing well-calibrated
uncertainty estimates, we formulate a temperature-dependent posterior that
facilitates faster convergence with temperature annealing. Our formulation
imposes an explicit condition on the uncertainty estimates which are
calibrated throughout the learning phase, leading to calibrated uncertainty
estimates without any post-hoc calibration phase. We formulate a new posterior
distribution on the network output given by,

$\displaystyle P(\theta|\mathcal{D})$ $\displaystyle\propto
P(\theta)\times\bm{\prod}_{i=1}^{i=N}\frac{e^{\frac{-|\hat{\mathbf{y}}_{i}-\mathbf{y}_{i}|^{2}}{(2\hat{\mathbf{\sigma}}_{i}^{2})}}}{\sqrt{2\pi\hat{\mathbf{\sigma}}_{i}^{2}}}\times\frac{e^{-T_{2}(|\hat{\mathbf{y}}_{i}-\mathbf{y}_{i}|^{2})}}{\sqrt{2\pi\frac{1}{T_{2}}}}\times\frac{e^{-T_{3}\left\\{\begin{array}[]{l}|\hat{\mathbf{y}}_{i}-(\mathbf{y}_{i}+\hat{\mathbf{\sigma}}_{i})|^{2},\hat{\mathbf{y}}_{i}\geq\mathbf{y}_{i}\\\
|\hat{\mathbf{y}}_{i}-(\mathbf{y}_{i}-\hat{\mathbf{\sigma}}_{i})|^{2},\hat{\mathbf{y}}_{i}<\mathbf{y}_{i}\end{array}\right\\}}}{\sqrt{2\pi\frac{1}{T_{3}}}}$
(4)

$T_{2},T_{3}$ are hyper-parameters that we refer to as temperature and
$P(\theta)$ is the prior. The maximum a posterior (MAP) of Equation 4 leads to
the following loss ($-\log P(\theta|\mathcal{D})$, ignoring constants):

$\displaystyle-\log
P(\theta)+\bm{\sum}_{i=1}^{i=N}\frac{\log{\hat{\mathbf{\sigma}}_{i}^{2}}}{2}+\frac{|\hat{\mathbf{y}}_{i}-\mathbf{y}_{i}|^{2}}{2\hat{\mathbf{\sigma}}_{i}^{2}}+T_{2}(|\hat{\mathbf{y}}_{i}-\mathbf{y}_{i}|^{2})+T_{3}\left\\{\begin{array}[]{l}|\hat{\mathbf{y}}_{i}-(\mathbf{y}_{i}+\hat{\mathbf{\sigma}}_{i})|^{2},\hat{\mathbf{y}}_{i}\geq\mathbf{y}_{i}\\\
|\hat{\mathbf{y}}_{i}-(\mathbf{y}_{i}-\hat{\mathbf{\sigma}}_{i})|^{2},\hat{\mathbf{y}}_{i}<\mathbf{y}_{i}\end{array}\right\\}.$
(7)

With a uniform prior, Equation 7 can be re-written as,

$\displaystyle-\log P(\theta|\mathcal{D})$
$\displaystyle=\bm{\sum}_{i=1}^{i=N}\frac{\log{\hat{\mathbf{\sigma}}_{i}^{2}}}{2}+\frac{|\hat{\mathbf{y}}_{i}-\mathbf{y}_{i}|^{2}}{2\hat{\mathbf{\sigma}}_{i}^{2}}+T_{2}(|\hat{\mathbf{y}}_{i}-\mathbf{y}_{i}|^{2})+T_{3}(|\hat{\mathbf{\sigma}}_{i}-|\hat{\mathbf{y}}_{i}-\mathbf{y}_{i}||^{2}).$
(8)

Figure 2:  Schematic of the temperature dependent posterior for a single data
point characterized by
$\\{\mathbf{y}_{i},\hat{\mathbf{y}}_{i},\hat{\mathbf{\sigma}}_{i}\\}$. This
enforces the prediction to be close to ground-truth and the uncertainty
estimate to be close the error, i.e., calibrated.

Figure 2 shows the schematic of our proposed temperature dependent posterior
distribution (Equation 4) for a single data point (characterized by
$\\{\mathbf{y}_{i},\hat{\mathbf{y}}_{i},\hat{\mathbf{\sigma}}_{i}\\}$). The
first temperature dependent factor in Equation 4:

$\frac{1}{\sqrt{2\pi\frac{1}{T_{2}}}}e^{-T_{2}(|\hat{\mathbf{y}}_{i}-\mathbf{y}_{i}|^{2})}$

enforces the central peak in Figure 2, which pushes $\hat{\mathbf{y}}_{i}$
close to $\mathbf{y}_{i}$. The second temperature dependent factor in Equation
4:

$\frac{1}{\sqrt{2\pi\frac{1}{T_{3}}}}e^{-T_{3}\left\\{\begin{array}[]{l}|\hat{\mathbf{y}}_{i}-(\mathbf{y}_{i}+\hat{\mathbf{\sigma}}_{i})|^{2},\hat{\mathbf{y}}_{i}\geq\mathbf{y}_{i}\\\
|\hat{\mathbf{y}}_{i}-(\mathbf{y}_{i}-\hat{\mathbf{\sigma}}_{i})|^{2},\hat{\mathbf{y}}_{i}<\mathbf{y}_{i}\end{array}\right\\}}$

ensures that if the prediction $\hat{\mathbf{y}}_{i}$ deviates from the
ground-truth $\mathbf{y}_{i}$, then the predicted uncertainty
$\hat{\mathbf{\sigma}}_{i}$ is close to the error (i.e.,
$\hat{\mathbf{y}}_{i}=\mathbf{y}_{i}\pm\hat{\mathbf{\sigma}}_{i}$), indicating
calibration, which correspond to the left/right peaks.

Figure 3:  Effects of temperature annealing. As we anneal the temperature in
Equation 4, the proposed loss surface (in orange) gradually changes from (a),
(b), (c) to (d), which provides faster convergence at the beginning of
training while ensuring convergence to the same optima as the standard loss
function as described in Equation 1 (in blue).

###  3.3 Effects of Temperature Annealing

The temperature dependent posterior in Equation 4 allows us to control the
contribution of individual loss terms in the final loss Equation 8 by changing
the temperature hyper-parameters $T_{2},T_{3}$. As described in Section 3,
controlling the temperature hyperparameters (i.e., decreasing the values)
allows faster convergence of the uncertainty-aware regression with better
calibrated uncertainties. We start by initializing $T_{2},T_{3}$ with a high
value of $100$ and progressively reduce them according to the training epochs
using exponential annealing – referred to as temperature annealing.

At higher temperature, the overall loss is dominated by the temperature-
dependent loss terms which lead to higher gradients (see Figure 1-(point a on
orange curve) and Figure 3(a)). As the temperature decreases, the overall loss
is close to the standard loss function given that the gradients from the
temperature-dependent terms are close to zero (see Figure 1-(point b on orange
surface) and Figure 3 (d)). This dynamic contribution from different loss
terms allows the network to converge faster in the beginning (as gradients
from the temperature-dependent loss terms are higher than the standard loss
term), and ensure stable convergence to the same optima as the standard loss,
thus leading to faster, better-calibrated uncertainty.

##  4 Experiments

We first describe our experimental setup (i.e., datasets, evaluation metrics
and implementation details) in Section 4.1. We compare our model to a wide
variety of state-of-the-art methods quantitatively and qualitatively in
Section 4.2. Finally, we also provide an ablation analysis in Section 4.2 to
study the rationale of our model formulation.

###  4.1 Experimental Setup

Datasets and Tasks. We conduct experiments on five datasets (three small scale
problems, two large scale problems) to solve the regression task and provide
uncertainty estimation.

We choose the following three low-dimensional regression problems. They
highlight the different complexities and network architectures that are
required to solve them. In Chaotic System using Lorenz Attractor (Lorenz
Attractor), the Lorenz equations describe non-linear chaotic systems given by,
$\frac{\partial z_{1}}{\partial t}=10(z_{2}-z_{1})$, $\frac{\partial
z_{2}}{\partial t}=z_{1}(28-z_{3})-z_{2}$, $\frac{\partial z_{3}}{\partial
t}=z_{1}z_{2}-8z_{3}/3$. Similar to (Garcia Satorras et al., 2019), to
generate a trajectory we run the Lorenz equations with a $\partial t=10^{-5}$
from which we sample with a time step of $t=0.05$. Each point is then
perturbed with Gaussian noise of standard deviation $0.5$ to produce pairs of
noisy and clean trajectories split into non-overlapping train/validation/test
sets. We use a 1D CNN to map the noisy input to clean output. The Physical
Properties of Molecules (Atom3D) (Townshend et al., 2020) is a 3D molecular
structure dataset aiming to predict the physical property such as the dipole
moment given the 3D atomistic representation. We use the standard Graph Neural
Network (GNN) for this task. The House Price Prediction (Boston-housing)
(Harrison Jr & Rubinfeld, 1978; Belsley et al., 2005) dataset is used to
predict the house prices using various attributes using Multi Layer
Perceptrons (MLPs).

To show the generalization of our method to high-dimensional regression
problems, we use the following two datasets. In Super-resolution of Natural
Images (Super-resolution), we learn mapping from low-resolution to high-
resolution images using CNNs, using DIV2K dataset (Timofte et al., 2018;
Ignatov et al., 2019). We do 4x downsampling to create the corresponding low-
resolution images. The dataset is split into 800/100/100 images for
training/val/test sets. In Medical Image Translation (MRI Translation), as T1
and T2 MRI from the same patient in the same orientation are often not
available and T2 takes longer to acquire, learning a mapping from T1 to T2 is
desirable. As in (Upadhyay et al., 2021a), we use T1 and T2 MRI of 500
patients from IXI dataset (Robinson et al., 2010) (200/100/200 for
training/val/test) in a 2D CNN based on U-Net (Ronneberger et al., 2015).

Evaluation Metrics. To measure the quality of regression output, we adopt the
standard metrics: mean absolute error (MAE) and mean square error (MSE). In
addition, for the super-resolution and medical image translation tasks, we use
PSNR and SSIM to measure the structural similarity between two images (Wang et
al., 2004). To measure the quality of uncertainty estimates
($\hat{\mathbf{\sigma}}^{2}$), we compute (i) the correlation coefficient
(Corr. Coeff.) between uncertainty estimates ($\hat{\mathbf{\sigma}}^{2}$) and
the error ($|\hat{\mathbf{y}}-\mathbf{y}|^{2}$). (ii) Uncertainty calibration
error (UCE) for regression tasks (Laves et al., 2020; Levi et al., 2019).
Following (Guo et al., 2017), the uncertainty output
$\hat{\mathbf{\sigma}}^{2}$ of a deep model is partitioned into $M$ bins with
equal width (each represented by $B_{m}$ for $\forall m\in\\{1,2...M\\}$). A
weighted average of the difference between the predictive error and
uncertainty is used,
$\text{UCE}=\sum_{m=1}^{M}\frac{|B_{m}|}{N}|\text{err}(B_{m})-\text{uncer}(B_{m})|.\text{
Where, }\texttt{err}(B_{m}):=\frac{1}{|B_{m}|}\sum_{i\in
B_{m}}||\hat{\mathbf{y}}_{i}-\mathbf{y}_{i}||^{2}\text{ and
}\texttt{uncer}(B_{m}):=\frac{1}{|B_{m}|}\sum_{i\in
B_{m}}\hat{\mathbf{\sigma}}_{i}^{2}\,.$ (iii) UCE for the re-calibrated
uncertainty estimates (R.UCE). We use post-hoc calibration technique
introduced in (Laves et al., 2020), called $\sigma$-scaling, that optimizes
for the scaling factor ($s$), post training to produce uncertainty estimates
($\hat{\mathbf{\sigma}}^{2}$) and predictions ($\hat{\mathbf{y}}$) using,
$s^{*}=\underset{s}{\text{argmin}}\left[N\log(s)+\frac{1}{2s^{2}}\sum_{i=1}^{N}\frac{|\hat{\mathbf{y}}_{i}-\mathbf{y}_{i}|^{2}}{\hat{\mathbf{\sigma}}_{i}^{2}}\right]$.
In addition, we present the (iv) expected calibration error (ECE) and (v)
sharpness (Sharpness). While ECE is another metric to quantify the calibration
of the uncertainty estimates, one must note that it may be possible to have an
uninformative, yet average calibrated model (Chung et al., 2021; Zhou et al.,
2021). Therefore it is necessary to also present the Sharpness metric that
encourages more-concentrated distributions. Finally, we present the (vi)
predictive log-likelihood that assesses how well the predicted conditional
distribution fits the data.

Implementation Details. Our POSTA method is generalizable across different
types of architectures. Here we perform experiments with MLPs, 1D/2D CNNs, and
GNNs. We take the well-established networks for the respective problems and
modify them to produce the uncertainty estimates as described in (Kendall &
Gal, 2017; Sudarshan et al., 2021). All the networks were trained using Adam
optimizer (Kingma & Ba, 2014). The initial learning rate was set to $2e^{-4}$
and cosine annealing was used to decay the learning rate over the course of
the learning phase. The hyper-parameters, $(T_{2},T_{3})$ (Equation 8) were
set to $(100,100)$ and scheduled to exponentially decay over the course of the
training. We provide the code in the supplementary.

###  4.2 Comparing to Uncertainty Estimation Methods

Figure 4:  Plots comparing the required convergence time (number of epochs to
converge) for different methods and corresponding ECE during the training on
(i) Super-resolution, (ii) MRI translation, (iii) Atom3D. More results in
supplementary.

Compared methods. For each of the regression tasks, we compare our model
(POSTA) to eight representative state-of-the-art methods for uncertainty
estimation using DNNs for regression tasks, belonging to a diverse class of
methods, i.e. Bayesian, ensemble, test-time data augmentation, maximum
likelihood, and finally quantile regression methods.

Bayesian methods: In (DO) (Gal & Ghahramani, 2016) the weights of the neural
network are randomly dropped at training and inference time. Multiple forward
passes for the same input at inference time allow us to estimate the
uncertainty. In Concrete Dropout (Conc. DO.) (Gal et al., 2017) the optimal
dropout probability for the weights of the neural network are learnt at
training.

Ensemble Methods: In Deep Ensemble (Ens) (Lakshminarayanan et al., 2016)
multiple deterministic networks are trained to make the final prediction with
uncertainty estimates. In Anchored Ensemble (Anch. Ens.) (Pearce et al., 2020)
the weights of the neural networks in the ensemble are regularized about
values drawn from a prior distribution, allowing approximate Bayesian
inference.

Test-time Data Augmentation Methods: In Test Time Data Augmentation (TTDA)
(Wang et al., 2019; Ayhan & Berens, 2018; Gawlikowski et al., 2021) multiple
perturbed copies of the input are passed through a deterministic network to
estimate the predictive uncertainty at the inference stage.

Maximum likelihood methods: In this method (NLL) (Kendall & Gal, 2017;
Sudarshan et al., 2021) the network is modified to predict the mean and
variance and then trained by optimizing negative log-likelihood. The variance
head then provides uncertainty estimates for the prediction at the inference
time.

Quantile Regression Methods: In Calibrated Quantile Regression Method (BPLoss)
(Chung et al., 2021) proposes a model that specifies the full quantile
function for the predictions and achieves a balance between calibration and
sharpness. In Collaborating Networks for estimating uncertainty intervals (CN)
(Zhou et al., 2021) two networks are trained simultaneously, one to estimate
the cumulative distribution function, and the other approximates its inverse.
We note that some baseline methods (i.e., BPLoss and CN) have only been
proposed for low-dimensional regression settings (where the output of a model
is single scalar) and it is non-trivial and inefficient to scale it to high-
dimensional regression settings (e.g., image translation, where the output for
an input is a high-dimensional matrix/tensor). Therefore such models are
compared only on low-dimensional regression tasks where they are applicable.

Quantitative results on convergence. In this experiment, we train different
models to perform the different kinds of regression task and keep track of the
training and validation loss to identify if the model has converged. For all
the models we used the same optimizer (i.e., Adam (Kingma & Ba, 2014)) with
the same initial learning rate (i.e., lr$=2e^{-4}$) and identical decaying
schedule (i.e., cosine annealing for lr).

We observe in Figure 4 that the baseline methods consistently take longer time
to converge while our proposed method (POSTA) consistently has faster
convergence. For instance, on the super-resolution task, our method takes
about 4,000 epochs to converge while the other baseline methods consistently
take longer than 8000 epochs to converge. In particular, the NLL baseline
takes the longest to converge. We also note that in the early phase of
training, our POSTA has much higher loss, this is due to the additional
temperature dependent loss terms (in Equation 8) that contribute to the
overall loss. However, the higher values of the temperature $T_{2}$ and
$T_{3}$ in the beginning of the training phase also allow faster convergence,
as explained in Section 3. Moreover, towards the end of the training phase,
the temperature parameters are annealed to a low value (close to zero) and the
over all loss function reduces to a low value.

Figure 4 (second row) shows the evolution of ECE for the derived uncertainty
using various methods during the training. Again we see that our POSTA
achieves the lowest ECE much faster than the other methods. A similar trend is
observed for the other datasets. For example, on Atom3D dataset, the proposed
method converges at about 2000 epochs, much faster than other baselines,
similarly it achieves lowest ECE much faster than other methods. These results
show that our method converges much faster than the other methods, which is in
line with our motivation to ensure a faster convergence for regression
uncertainty model along with better calibrated uncertainty as described in
Section 3.3.

T | Methods | Metrics
---|---|---
MAE $\downarrow$ | MSE $\downarrow$ | SSIM $\uparrow$ | PSNR $\uparrow$ | C.Coeff. $\uparrow$ | UCE$\downarrow$ | R.UCE.$\downarrow$ | Log-likeli. $\uparrow$ | ECE $\downarrow$ | Sharp. $\downarrow$
Boston-housing | DO (Gal & Ghahramani, 2016) | 2.851 | 13.26 | - | - | 0.014 | 10.76 | 10.18 | -2.46 | 10.2 | 8.66
Conc. DO (Gal et al., 2017) | 2.413 | 10.18 | - | - | 0.135 | 9.882 | 9.126 | -2.15 | 9.18 | 9.16
Ens. (Lakshminarayanan et al., 2016) | 2.971 | 13.76 | - | - | 0.011 | 11.26 | 10.78 | -2.41 | 10.3 | 8.87
Anch. Ens. (Pearce et al., 2020) | 2.553 | 10.11 | - | - | 0.154 | 9.547 | 9.135 | -2.32 | 9.92 | 9.82
TTDA (Wang et al., 2019; Ayhan & Berens, 2018; Gawlikowski et al., 2021) | 2.584 | 10.30 | - | - | 0.007 | 14.32 | 13.85 | -2.24 | 11.8 | 9.28
NLL (Kendall & Gal, 2017) | 2.663 | 10.75 | - | - | 0.107 | 12.67 | 12.23 | -2.42 | 11.5 | 8.21
BPLoss (Chung et al., 2021) | 2.684 | 11.49 | - | - | 0.237 | 9.216 | 8.837 | -2.11 | 9.72 | 9.01
CN (Zhou et al., 2021) | 2.594 | 11.13 | - | - | 0.213 | 10.84 | 9.722 | -2.23 | 9.65 | 9.67
POSTA (ours) | 2.593 | 10.51 | - | - | 0.348 | 0.756 | 6.374 | -2.06 | 6.37 | 8.22
Atom3D | DO (Gal & Ghahramani, 2016) | 1.950 | 5.828 | - | - | 0.085 | 5.380 | 5.054 | -0.24 | 2.12 | 4.32
Conc. DO (Gal et al., 2017) | 1.834 | 5.212 | - | - | 0.136 | 4.879 | 4.122 | -0.21 | 1.81 | 4.18
Ens. (Lakshminarayanan et al., 2016) | 1.215 | 2.388 | - | - | 0.138 | 4.623 | 4.376 | -0.23 | 1.69 | 4.17
Anch. Ens. (Pearce et al., 2020) | 1.087 | 1.743 | - | - | 0.182 | 4.124 | 3.763 | -0.26 | 1.42 | 3.95
TTDA (Wang et al., 2019; Ayhan & Berens, 2018; Gawlikowski et al., 2021) | 0.903 | 1.301 | - | - | 0.157 | 4.167 | 3.988 | -0.38 | 1.94 | 4.78
NLL (Kendall & Gal, 2017) | 0.498 | 0.463 | - | - | 0.164 | 3.358 | 3.335 | -0.22 | 1.38 | 3.32
BPLoss (Chung et al., 2021) | 0.527 | 0.873 | - | - | 0.189 | 3.527 | 3.166 | -0.21 | 1.55 | 3.12
CN (Zhou et al., 2021) | 0.521 | 0.845 | - | - | 0.087 | 4.311 | 2.971 | -0.16 | 1.77 | 3.18
POSTA (ours) | 0.513 | 0.495 | - | - | 0.567 | 0.296 | 0.277 | -0.18 | 1.37 | 3.17
Lorenz Attractor | DO (Gal & Ghahramani, 2016) | 1.373 | 3.463 | - | 29.85 | 0.281 | 2.864 | 2.134 | -0.16 | 4.34 | 5.67
Conc. DO (Gal et al., 2017) | 1.247 | 3.198 | - | 30.34 | 0.311 | 2.379 | 2.136 | -0.14 | 4.13 | 5.22
Ens. (Lakshminarayanan et al., 2016) | 2.544 | 11.65 | - | 24.32 | 0.778 | 6.726 | 6.294 | -0.22 | 10.4 | 8.43
Anch. Ens. (Pearce et al., 2020) | 2.122 | 10.12 | - | 25.64 | 0.432 | 8.756 | 8.154 | -0.29 | 10.7 | 9.43
TTDA (Wang et al., 2019; Ayhan & Berens, 2018; Gawlikowski et al., 2021) | 1.391 | 3.764 | - | 29.16 | 0.438 | 3.325 | 3.077 | -0.17 | 5.96 | 8.91
NLL (Kendall & Gal, 2017) | 0.172 | 0.048 | - | 31.28 | 0.588 | 2.368 | 1.933 | -0.13 | 4.33 | 7.87
POSTA (ours) | 0.153 | 0.029 | - | 32.33 | 0.821 | 0.779 | 0.356 | -0.11 | 4.36 | 9.12
Super-resolution | DO (Gal & Ghahramani, 2016) | 0.832 | 0.548 | 0.947 | 35.64 | 0.033 | 0.748 | 0.519 | -0.38 | 4.67 | 6.32
Conc. DO (Gal et al., 2017) | 0.801 | 0.423 | 0.951 | 35.71 | 0.134 | 0.711 | 0.494 | -0.36 | 4.43 | 8.82
Ens. (Lakshminarayanan et al., 2016) | 0.793 | 0.462 | 0.953 | 36.61 | 0.029 | 0.941 | 0.733 | -0.36 | 8.76 | 10.2
Anch. Ens. (Pearce et al., 2020) | 0.755 | 0.441 | 0.957 | 36.63 | 0.178 | 0.883 | 0.713 | -0.41 | 8.11 | 9.21
TTDA (Wang et al., 2019; Ayhan & Berens, 2018; Gawlikowski et al., 2021) | 0.883 | 0.691 | 0.939 | 34.94 | 0.047 | 1.175 | 0.994 | -0.39 | 11.3 | 10.3
NLL (Kendall & Gal, 2017) | 0.693 | 0.414 | 0.955 | 37.15 | 0.189 | 0.581 | 0.512 | -0.36 | 1.45 | 2.73
POSTA (ours) | 0.618 | 0.351 | 0.962 | 37.87 | 0.518 | 0.104 | 0.053 | -0.16 | 0.74 | 0.83
MRI Translation | DO (Gal & Ghahramani, 2016) | 0.732 | 0.683 | 0.912 | 32.45 | 0.159 | 0.864 | 0.771 | -0.33 | 4.48 | 6.23
Conc. DO (Gal et al., 2017) | 0.715 | 0.612 | 0.917 | 32.97 | 0.189 | 1.125 | 0.932 | -0.31 | 4.12 | 7.89
Ens. (Lakshminarayanan et al., 2016) | 0.681 | 0.611 | 0.927 | 33.76 | 0.110 | 1.143 | 0.974 | -0.36 | 4.86 | 7.21
Anch. Ens. (Pearce et al., 2020) | 0.655 | 0.532 | 0.933 | 33.84 | 0.166 | 1.122 | 0.913 | -0.34 | 5.88 | 7.32
TTDA (Wang et al., 2019; Ayhan & Berens, 2018; Gawlikowski et al., 2021) | 0.755 | 0.729 | 0.904 | 32.18 | 0.128 | 1.483 | 1.153 | -0.37 | 7.21 | 9.74
NLL (Kendall & Gal, 2017) | 0.632 | 0.582 | 0.938 | 34.34 | 0.134 | 1.673 | 1.448 | -0.28 | 4.03 | 5.12
POSTA (ours) | 0.615 | 0.537 | 0.946 | 35.27 | 0.432 | 0.098 | 0.062 | -0.30 | 3.26 | 5.78

Table 1:  Evaluating different methods on five datasets using MAE, MSE, PSNR,
SSIM (where applicable, to evaluate regression) and C.Coeff., UCE, R.UCE, Log-
Likeli., ECE, Sharp. (to measure quality of uncertainty estimates).
$\uparrow$/$\downarrow$ indicates higher/lower is better. “T”: tasks. Best
results are in bold.

Quantitative results on regression and uncertainty. Uncertainty-aware
regression models must be evaluated on two fronts which are (i) the regression
performance, i.e., the quality of the target predictions and (ii) the quality
of estimated uncertainty (the uncertainty should be sharp and well
calibrated). We evaluate the model performance based on two set of metrics:
(1) task-specific metrics that evaluate the regression results using MAE, MSE,
PSNR, SSIM, and (2) calibration-specific metrics that evaluate the quality of
the uncertainty estimates using C.Coeff., UCE, R.UCE, ECE, Sharp., and Log-
likeli. Table 1 shows the quantitative results that evaluate regression and
the quality of uncertainty estimates for different methods on multiple
regression tasks. Our POSTA method also obtains high quality regression
outputs. In two tasks (including super-resolution, and MRI translation), our
POSTA achieves the best or competitive performance compared to the other
methods. We note that while no single metric can indicate the “goodness” of
uncertainty estimates (as there is no groundtruth for uncertainty values), the
collective set of metrics such as C.Coeff., UCE, R.UCE, ECE, Log-likeli,
Sharp. provide a holisitic indication of “goodness” of uncertainty metric. The
proposed method, POSTA, consistently performs well in terms of the above
metrics. Overall, these quantitative results show that our method performs
well in providing both satisfactory regression results and uncertainty
estimates.

Figure 5:  Qualitative results: input, predictions, groundtruth, and the
error.

Qualitative results on regression and uncertainty. Figure 5 shows the
regression output on different datasets. Figure 5-(i) & (ii) visualizes the
generated images for image super-resolution and MRI translation tasks. While
the other methods often generate relatively blurry images with artifacts in
colours, our model produces better output visually more similar to the ground-
truth. Moreover, Figure 5-(i) & (ii) also shows the uncertainty maps, along
with the prediction and error for super-resolution and MRI translation. We
observe that for compared methods, uncertainty maps do not always agree with
error maps at pixel level (i.e., higher/lower uncertainty than the
corresponding error), whereas our uncertainty maps are in agreement with the
errors. This suggests that our model provides better-calibrated uncertainty.
Figure 5-(iii) shows the plots for predictions vs ground-truth on the Atom3D
dataset. We can see that compared to other methods, our method yields
predictions much closer to the ground-truth e.g., on the Atom3D dataset, our
method produces regression output more highly correlated with the ground-
truth. Figure 5-(iv) shows the input noisy trajectory, denoised output and the
corresponding ground-truth for the Lorentz attractor dataset. We can see that
compared to other methods, our method yields smoother trajectories.

###  4.3 Ablation Analysis of Posterior Annealing

Table 2 shows the ablation study of two temperature hyperparemeters in our
formulated temperature dependent posterior (Equation 4) along with different
choices of priors for the super-resolution task.

We test the baseline that removes both temperature dependent terms (i.e.
$T_{2}=T_{3}=0$) with a uniform prior, this is equivalent to the NLL method
and is shown in the first row (MAE of 0.693). We then study the effect of
fixing one of the temperatures at non-zero value while setting the other
temperature to 0. With $T_{2}=100,T_{3}=0$, we see slight improvement in
regression performance (MAE of 0.614 vs. 0.693) and much poorer performance
with respect to uncertainty calibration (UCE of 1.169 vs. 0.581), this is due
to more weighting of fidelity term between the prediction and the ground-truth
along with suppression of the default calibration effect of NLL. On the other
hand, $T_{2}=0,T_{3}=100$ suppresses the default fidelity term for NLL,
therefore the output is of significantly worse quality (poor regression
scores, MAE of 1.395 vs 0.693) this further degrades the quality of the
uncertainty estimates (UCE 3.733 vs 0.581). We notice that if the model does
not perform good regression, the quality of uncertainty estimate is also
adversely effected.

We then study the effects of decaying one of the temperatures while setting
other to 0. With $T_{2}$ decaying (i.e., $T_{2}=\downarrow,T_{3}=0$) we see
slightly better performance than $T_{2}=100,T_{3}=0$ (MAE of 0.612 vs. 0.614
and UCE of 0.983 vs. 1.169), whereas with $T_{2}$ decaying (i.e.,
$T_{2}=0,T_{3}=\downarrow$) we see good regression performance but also an
improved calibration performance (UCE of 0.152 vs. 0.581). With both the
parameters decaying (i.e., $T_{2}=\downarrow,T_{3}=\downarrow$) we achieve
improved regression and calibration results concluding that annealing works
the best. In addition to uniform prior setup (i.e.,
$P(\theta)=\mathcal{U}(\theta)$), we evaluate two other priors (i) Gaussian
prior on the parameters of the network, i.e., $P(\theta)=\mathcal{N}(\theta)$
that is equivalent to $\ell_{2}$ regularization of weights and (ii) Laplace
prior, i.e., $P(\theta)=\mathcal{E}(\theta)$ that is equivalent to $\ell_{1}$
regularization of weights. With Gaussian/Laplace prior we achieve MAE of
0.625/0.612 showing that carefully crafted priors may further boost the
performance, designing such priors will be explored in future works.

Methods | Metrics
---|---
MAE $\downarrow$ | MSE $\downarrow$ | PSNR $\uparrow$ | SSIM $\uparrow$ | C.Coeff. $\uparrow$ | UCE$\downarrow$ | R.UCE.$\downarrow$ | Log-likeli. $\uparrow$ | ECE $\downarrow$ | Sharp. $\downarrow$
$T_{2}=0,T_{3}=0$ | 0.693 | 0.414 | 37.15 | 0.955 | 0.189 | 0.581 | 0.512 | -0.36 | 1.45 | 2.73
$T_{2}=100,T_{3}=0$ | 0.614 | 0.384 | 37.72 | 0.961 | 0.062 | 1.169 | 0.833 | -0.41 | 1.77 | 2.82
$T_{2}=0,T_{3}=100$ | 1.395 | 7.274 | 20.19 | 0.793 | 0.219 | 3.733 | 2.442 | -0.44 | 2.12 | 3.11
$T_{2}=\downarrow,T_{3}=0$ | 0.612 | 0.344 | 37.76 | 0.961 | 0.077 | 0.983 | 0.797 | -0.27 | 1.03 | 1.35
$T_{2}=0,T_{3}=\downarrow$ | 0.632 | 0.388 | 37.71 | 0.960 | 0.442 | 0.152 | 0.116 | -0.20 | 0.85 | 0.98
$T_{2}=\downarrow,T_{3}=\downarrow$ | 0.618 | 0.351 | 37.87 | 0.962 | 0.518 | 0.104 | 0.083 | -0.16 | 0.74 | 0.83
$T_{2}=\downarrow,T_{3}=\downarrow$ with $P(\theta)=\mathcal{N}(\theta)$ | 0.625 | 0.358 | 36.98 | 0.952 | 0.488 | 0.168 | 0.133 | -0.24 | 1.12 | 1.47
$T_{2}=\downarrow,T_{3}=\downarrow$ with $P(\theta)=\mathcal{E}(\theta)$ | 0.612 | 0.353 | 37.92 | 0.966 | 0.503 | 0.118 | 0.102 | -0.15 | 0.83 | 1.01

Table 2:  Ablation study of temperature hyperparemeters of the temperature
dependent posterior used in the proposed posterior annealing (POSTA) method on
image super-resolution task.  Figure 6:  Evaluation of different methods using
out-of-distribution input samples for MRI translation.

###  4.4 Evaluation on Out-of-Distribution Data

Previous works have studied the performance of various uncertainty-aware
methods in the presence of out-of-distribution (OOD) samples at the inference
time (Ovadia et al., 2019; Hendrycks et al., 2019; Nandy et al., 2020; Mundt
et al., 2019). To evaluate if better quality of uncertainty estimates lead to
better OOD performance, we evaluate all the uncertainty trained model for MRI
Translation on OOD samples. MRI image acquisition is a noisy process that
leads to noisy/corrupted images (Macovski, 1996; Parrish et al., 2000; Wiest-
Daesslé et al., 2008; Aja-Fernández & Vegas-Sánchez-Ferrero, 2016). Similar to
(Upadhyay et al., 2021a; c; Sudarshan et al., 2021), we study the performance
of various uncertainty-aware models in the presence of noisy input samples
(corrupted with varying degrees of noise) at test time. Figure 6-(left) shows
the example of in-distribution (noise-level 0, NL0) and out-of-distribution
samples (NL1 and NL2). The severity of corruption gradually increases from NL0
to NL2. From Figure 6-(middle and right), that shows the regression and
quality of uncertainty estimates in the presence of OOD samples, we observe
that the performance of various models degrades as severity of corruption
increases from NL0 to NL2, however our POSTA method performs much better than
the compared methods even at higher severity of corruption both in terms of
regression and uncertainty calibration metric.

##  5 Conclusion

In this paper, we present a temperature dependent posterior to learn faster,
better-calibrated uncertainty estimates for regression tasks. Our formulation
speeds up model convergence and improves calibration of uncertainty without
any post-hoc calibration phase. By annealing the temperature during training,
we show that our method achieves 1.5 to 6 times faster convergence compared to
other methods. Moreover, we demonstrate that our method achieves strong
regression results while providing better calibrated uncertainty estimates
compared to five uncertainty estimation methods on multiple regression
datasets. Lastly, we present insights of our method with ablation study and
show its potential to generalize well on out-of-distribution data, which
further suggests its robustness.

## References

  * Aja-Fernández & Vegas-Sánchez-Ferrero (2016) Santiago Aja-Fernández and Gonzalo Vegas-Sánchez-Ferrero.  Statistical analysis of noise in mri.  _Switzerland: Springer International Publishing_ , 2016.
  * Ayhan & Berens (2018) Murat Seckin Ayhan and Philipp Berens.  Test-time data augmentation for estimation of heteroscedastic aleatoric uncertainty in deep neural networks.  2018\.
  * Belsley et al. (2005) David A Belsley, Edwin Kuh, and Roy E Welsch.  _Regression diagnostics: Identifying influential data and sources of collinearity_ , volume 571.  John Wiley & Sons, 2005.
  * Besnier et al. (2021) Victor Besnier, Andrei Bursuc, David Picard, and Alexandre Briot.  Triggering failures: Out-of-distribution detection by learning from local adversarial attacks in semantic segmentation.  In _Proceedings of the IEEE/CVF International Conference on Computer Vision_ , 2021.
  * Blundell et al. (2015) Charles Blundell, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra.  Weight uncertainty in neural network.  In _International Conference on Machine Learning_ , pp. 1613–1622. PMLR, 2015.
  * Chen et al. (2021) Haoxian Chen, Ziyi Huang, Henry Lam, Huajie Qian, and Haofeng Zhang.  Learning prediction intervals for regression: Generalization and calibration.  In _International Conference on Artificial Intelligence and Statistics_. PMLR, 2021.
  * Chen et al. (2014) Tianqi Chen, Emily Fox, and Carlos Guestrin.  Stochastic gradient hamiltonian monte carlo.  In _International conference on machine learning_ , pp. 1683–1691. PMLR, 2014.
  * Chung et al. (2021) Youngseog Chung, Willie Neiswanger, Ian Char, and Jeff Schneider.  Beyond pinball loss: Quantile methods for calibrated uncertainty quantification.  _Advances in Neural Information Processing Systems_ , 2021.
  * Cui et al. (2020) Peng Cui, Wenbo Hu, and Jun Zhu.  Calibrated reliable regression using maximum mean discrepancy.  _Advances in Neural Information Processing Systems_ , 33, 2020.
  * Daxberger et al. (2021) Erik Daxberger, Agustinus Kristiadi, Alexander Immer, Runa Eschenhagen, Matthias Bauer, and Philipp Hennig.  Laplace redux-effortless bayesian deep learning.  _Advances in Neural Information Processing Systems_ , 34, 2021.
  * Feldman et al. (2021) Shai Feldman, Stephen Bates, and Yaniv Romano.  Improving conditional coverage via orthogonal quantile regression.  _Advances in Neural Information Processing Systems_ , 2021.
  * Franchi et al. (2020) Gianni Franchi, Andrei Bursuc, Emanuel Aldea, Séverine Dubuisson, and Isabelle Bloch.  Tradi: Tracking deep neural network weight distributions.  In _European Computer Vision Conference_ , 2020.
  * Gal (2016) Yarin Gal.  Uncertainty in deep learning.  2016\.
  * Gal & Ghahramani (2016) Yarin Gal and Zoubin Ghahramani.  Dropout as a bayesian approximation: Representing model uncertainty in deep learning.  In _international conference on machine learning_ , 2016.
  * Gal et al. (2017) Yarin Gal, Jiri Hron, and Alex Kendall.  Concrete dropout.  In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), _Advances in Neural Information Processing Systems_ , 2017.
  * Garcia Satorras et al. (2019) Victor Garcia Satorras, Zeynep Akata, and Max Welling.  Combining generative and discriminative models for hybrid inference.  In _33rd Conference on Neural Information Processing Systems_ , pp. 1–11. Curran Associates, Inc., 2019.
  * Gawlikowski et al. (2021) Jakob Gawlikowski, Cedrique Rovile Njieutcheu Tassi, Mohsin Ali, Jongseok Lee, Matthias Humt, Jianxiang Feng, Anna Kruspe, Rudolph Triebel, Peter Jung, Ribana Roscher, et al.  A survey of uncertainty in deep neural networks.  _arXiv preprint arXiv:2107.03342_ , 2021.
  * Gillmann et al. (2021) Christina Gillmann, Dorothee Saur, Thomas Wischgoll, and Gerik Scheuermann.  Uncertainty-aware visualization in medical imaging-a survey.  In _Computer Graphics Forum_. Wiley Online Library, 2021.
  * Graves (2011) Alex Graves.  Practical variational inference for neural networks.  _Advances in neural information processing systems_ , 24, 2011.
  * Guo et al. (2017) Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q Weinberger.  On calibration of modern neural networks.  In _International Conference on Machine Learning_ , pp. 1321–1330. PMLR, 2017.
  * Harrison Jr & Rubinfeld (1978) David Harrison Jr and Daniel L Rubinfeld.  Hedonic housing prices and the demand for clean air.  _Journal of environmental economics and management_ , 5(1):81–102, 1978.
  * Hendrycks et al. (2019) Dan Hendrycks, Kimin Lee, and Mantas Mazeika.  Using pre-training can improve model robustness and uncertainty.  In _International Conference on Machine Learning_ , 2019.
  * Hüllermeier & Waegeman (2021) Eyke Hüllermeier and Willem Waegeman.  Aleatoric and epistemic uncertainty in machine learning: An introduction to concepts and methods.  _Machine Learning_ , 2021.
  * Ignatov et al. (2019) Andrey Ignatov, Radu Timofte, et al.  Pirm challenge on perceptual image enhancement on smartphones: report.  In _European Conference on Computer Vision (ECCV) Workshops_ , January 2019.
  * Kendall & Gal (2017) Alex Kendall and Yarin Gal.  What uncertainties do we need in bayesian deep learning for computer vision?  _Advances in Neural Information Processing Systems_ , 30:5574–5584, 2017.
  * Kingma & Ba (2014) Diederik P Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  _arXiv preprint arXiv:1412.6980_ , 2014.
  * Kläs & Vollmer (2018) Michael Kläs and Anna Maria Vollmer.  Uncertainty in machine learning applications: A practice-driven classification of uncertainty.  In _International conference on computer safety, reliability, and security_ , 2018.
  * Kuleshov et al. (2018) Volodymyr Kuleshov, Nathan Fenner, and Stefano Ermon.  Accurate uncertainties for deep learning using calibrated regression.  In _International Conference on Machine Learning_ , pp. 2796–2804. PMLR, 2018.
  * Lakshminarayanan et al. (2016) Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell.  Simple and scalable predictive uncertainty estimation using deep ensembles.  _arXiv preprint arXiv:1612.01474_ , 2016.
  * Laves et al. (2020) Max-Heinrich Laves, Sontje Ihler, Jacob F Fast, Lüder A Kahrs, and Tobias Ortmaier.  Well-calibrated regression uncertainty in medical imaging with deep learning.  In _Medical Imaging with Deep Learning_ , pp. 393–412. PMLR, 2020\.
  * Leibig et al. (2017) Christian Leibig, Vaneeda Allken, Murat Seçkin Ayhan, Philipp Berens, and Siegfried Wahl.  Leveraging uncertainty information from deep neural networks for disease detection.  _Scientific reports_ , 2017.
  * Levi et al. (2019) Dan Levi, Liran Gispan, Niv Giladi, and Ethan Fetaya.  Evaluating and calibrating uncertainty prediction in regression tasks.  _arXiv preprint arXiv:1905.11659_ , 2019.
  * Lin et al. (2021) Zhen Lin, Shubhendu Trivedi, and Jimeng Sun.  Locally valid and discriminative prediction intervals for deep learning models.  _Advances in Neural Information Processing Systems_ , 2021.
  * Macovski (1996) Albert Macovski.  Noise in mri.  _Magnetic resonance in medicine_ , 1996.
  * Maddox et al. (2019) Wesley J Maddox, Pavel Izmailov, Timur Garipov, Dmitry P Vetrov, and Andrew Gordon Wilson.  A simple baseline for bayesian uncertainty in deep learning.  _Advances in Neural Information Processing Systems_ , 2019.
  * Mehrtash et al. (2020) Alireza Mehrtash, Purang Abolmaesumi, Polina Golland, Tina Kapur, Demian Wassermann, and William M Wells III.  Pep: Parameter ensembling by perturbation.  _Advances in Neural Information Processing Systems_ , 2020.
  * Mundt et al. (2019) Martin Mundt, Iuliia Pliushch, Sagnik Majumder, and Visvanathan Ramesh.  Open set recognition through deep neural network uncertainty: Does out-of-distribution detection require generative classifiers?  In _Proceedings of the IEEE/CVF International Conference on Computer Vision Workshops_ , 2019.
  * Nandy et al. (2020) Jay Nandy, Wynne Hsu, and Mong Li Lee.  Towards maximizing the representation gap between in-domain & out-of-distribution examples.  _Advances in Neural Information Processing Systems_ , 2020.
  * Ovadia et al. (2019) Yaniv Ovadia, Emily Fertig, Jie Ren, Zachary Nado, David Sculley, Sebastian Nowozin, Joshua Dillon, Balaji Lakshminarayanan, and Jasper Snoek.  Can you trust your model’s uncertainty? evaluating predictive uncertainty under dataset shift.  _Advances in neural information processing systems_ , 2019.
  * Parrish et al. (2000) Todd B Parrish, Darren R Gitelman, Kevin S LaBar, and M-Marsel Mesulam.  Impact of signal-to-noise on functional mri.  _Magnetic Resonance in Medicine_ , 2000.
  * Pearce et al. (2018) Tim Pearce, Alexandra Brintrup, Mohamed Zaki, and Andy Neely.  High-quality prediction intervals for deep learning: A distribution-free, ensembled approach.  In _International Conference on Machine Learning_ , pp. 4075–4084. PMLR, 2018.
  * Pearce et al. (2020) Tim Pearce, Felix Leibfried, and Alexandra Brintrup.  Uncertainty in neural networks: Approximately bayesian ensembling.  In _AISTATS_ , 2020.
  * Phan et al. (2018) Buu Phan, Rick Salay, Krzysztof Czarnecki, Vahdat Abdelzad, Taylor Denouden, and Sachin Vernekar.  Calibrating uncertainties in object localization task.  _arXiv preprint arXiv:1811.11210_ , 2018.
  * Robinson et al. (2010) Emma C Robinson, Alexander Hammers, Anders Ericsson, A David Edwards, and Daniel Rueckert.  Identifying population differences in whole-brain structural networks: a machine learning approach.  _NeuroImage_ , 50(3):910–919, 2010.
  * Romano et al. (2019) Yaniv Romano, Evan Patterson, and Emmanuel Candes.  Conformalized quantile regression.  _Advances in neural information processing systems_ , 2019.
  * Ronneberger et al. (2015) Olaf Ronneberger, Philipp Fischer, and Thomas Brox.  U-net: Convolutional networks for biomedical image segmentation.  In _International Conference on Medical image computing and computer-assisted intervention_ , pp. 234–241. Springer, 2015.
  * Shafaei et al. (2018) Sina Shafaei, Stefan Kugele, Mohd Hafeez Osman, and Alois Knoll.  Uncertainty in machine learning: A safety perspective on autonomous driving.  In _International Conference on Computer Safety, Reliability, and Security_ , 2018.
  * Song et al. (2019) Hao Song, Tom Diethe, Meelis Kull, and Peter Flach.  Distribution calibration for regression.  In _International Conference on Machine Learning_ , pp. 5897–5906. PMLR, 2019.
  * Sudarshan et al. (2021) Viswanath P Sudarshan, Uddeshya Upadhyay, Gary F Egan, Zhaolin Chen, and Suyash P Awate.  Towards lower-dose pet using physics-based uncertainty-aware multimodal learning with robustness to out-of-distribution data.  _Medical Image Analysis_ , 73:102187, 2021.
  * Tagasovska & Lopez-Paz (2019) Natasa Tagasovska and David Lopez-Paz.  Single-model uncertainties for deep learning.  _Advances in Neural Information Processing Systems_ , 32:6417–6428, 2019.
  * Timofte et al. (2018) Radu Timofte, Shuhang Gu, Jiqing Wu, Luc Van Gool, Lei Zhang, Ming-Hsuan Yang, Muhammad Haris, et al.  Ntire 2018 challenge on single image super-resolution: Methods and results.  In _The IEEE Conference on Computer Vision and Pattern Recognition (CVPR) Workshops_ , June 2018.
  * Townshend et al. (2020) Raphael JL Townshend, Martin Vögele, Patricia Suriana, Alexander Derry, Alexander Powers, Yianni Laloudakis, Sidhika Balachandar, Bowen Jing, Brandon Anderson, Stephan Eismann, et al.  Atom3d: Tasks on molecules in three dimensions.  _arXiv preprint arXiv:2012.04035_ , 2020.
  * Upadhyay et al. (2021a) Uddeshya Upadhyay, Yanbei Chen, and Zeynep Akata.  Robustness via uncertainty-aware cycle consistency.  _Advances in Neural Information Processing Systems_ , 2021a.
  * Upadhyay et al. (2021b) Uddeshya Upadhyay, Yanbei Chen, Tobias Hepp, Sergios Gatidis, and Zeynep Akata.  Uncertainty-guided progressive gans for medical image translation.  In _International Conference on Medical Image Computing and Computer-Assisted Intervention_ , 2021b.
  * Upadhyay et al. (2021c) Uddeshya Upadhyay, Viswanath P Sudarshan, and Suyash P Awate.  Uncertainty-aware gan with adaptive loss for robust mri image enhancement.  In _Proceedings of the IEEE/CVF International Conference on Computer Vision_ , pp. 3255–3264, 2021c.
  * Upadhyay et al. (2022) Uddeshya Upadhyay, Shyamgopal Karthik, Yanbei Chen, Massimiliano Mancini, and Zeynep Akata.  Bayescap: Bayesian identity cap for calibrated uncertainty in frozen neural networks.  In _Computer Vision–ECCV 2022: 17th European Conference, Tel Aviv, Israel, October 23–27, 2022, Proceedings, Part XII_ , pp. 299–317. Springer, 2022.
  * Varshney & Alemzadeh (2017) Kush R Varshney and Homa Alemzadeh.  On the safety of machine learning: Cyber-physical systems, decision sciences, and data products.  _Big data_ , 2017.
  * Wang et al. (2019) Guotai Wang, Wenqi Li, Michael Aertsen, Jan Deprest, Sébastien Ourselin, and Tom Vercauteren.  Aleatoric uncertainty estimation with test-time augmentation for medical image segmentation with convolutional neural networks.  _Neurocomputing_ , 2019.
  * Wang et al. (2004) Zhou Wang, Alan C Bovik, Hamid R Sheikh, and Eero P Simoncelli.  Image quality assessment: from error visibility to structural similarity.  _IEEE transactions on image processing_ , 13(4):600–612, 2004.
  * Welling & Teh (2011) Max Welling and Yee W Teh.  Bayesian learning via stochastic gradient langevin dynamics.  In _Proceedings of the 28th international conference on machine learning (ICML-11)_ , pp. 681–688. Citeseer, 2011.
  * Wiest-Daesslé et al. (2008) Nicolas Wiest-Daesslé, Sylvain Prima, Pierrick Coupé, Sean Patrick Morrissey, and Christian Barillot.  Rician noise removal by non-local means filtering for low signal-to-noise ratio mri: applications to dt-mri.  In _International Conference on Medical Image Computing and Computer-assisted Intervention_ , 2008.
  * Xu et al. (2014) Wenda Xu, Jia Pan, Junqing Wei, and John M Dolan.  Motion planning under uncertainty for on-road autonomous driving.  In _2014 IEEE International Conference on Robotics and Automation (ICRA)_ , 2014.
  * Yan et al. (2018) Xing Yan, Weizhong Zhang, Lin Ma, Wei Liu, and Qi Wu.  Parsimonious quantile regression of financial asset tail dynamics via sequential learning.  _Advances in neural information processing systems_ , 2018.
  * Zhou et al. (2021) Tianhui Zhou, Yitong Li, Yuan Wu, and David Carlson.  Estimating uncertainty intervals from collaborating networks.  _Journal of Machine Learning Research_ , 2021.


# The exponential Lagrangian mean

Abhijeet Minz1<EMAIL_ADDRESS>Lois E. Baker1  Hossein A. Kafiabad2
Jacques Vanneste1  1School of Mathematics and Maxwell Institute for
Mathematical Sciences, University of Edinburgh, Edinburgh EH9 3FD, UK
2Department of Mathematical Sciences, Durham University, Durham DH1 3LE, UK

###### Abstract

Lagrangian averaging is a valuable tool for the analysis and modelling of
multiscale processes in fluid dynamics. The numerical computation of
Lagrangian averages from simulation data is considered challenging, however.
In response, we develop a straightforward form of Lagrangian (time) averaging
– exponential averaging – and derive simple partial differential equations
that govern the evolution of the Lagrangian mean fields. These equations can
be solved at minimum cost as part of the numerical simulation of fluid models,
using the same time and space discretisation as for the dynamical equations.

We implement exponential averaging in the rotating shallow-water model and
demonstrate its effectiveness at filtering out large-amplitude Poincaré waves
while retaining the salient features of an underlying turbulent flow. We
generalise the computation of exponential Lagrangian means from scalar to
tensor fields and apply the generalisation to the computation of the
Lagrangian mean momentum and of the related pseudomomentum.

##  1 Introduction

Many fluid dynamical phenomena involve temporal and spatial scales that cannot
be fully resolved by measuring instruments or numerical simulations. Their
modelling requires coarse graining, typically in the form of temporal or
spatial averaging of the equations of motion, and the parameterisation of the
impact of the unresolved scales. To this end, Lagrangian averaging, whereby
averages are computed along fluid trajectories, has advantages over the more
straightforward Eulerian averaging. These advantages stem from the
preservation under Lagrangian averaging of the advective structure of the
equations of motion, leading to valuable properties including the conservation
of Lagrangian mean vorticity and circulation in the absence of dissipation and
forcing.

A conceptual framework for Lagrangian averaging is provided by the generalised
Lagrangian mean (GLM) theory of Andrews & McIntyre (1978) (see Bühler, 2014,
for a comprehensive introduction). Its practical use, with temporal averaging
substituted for the abstract averaging of GLM, has however been hampered by
the challenges posed by the numerical computation of Lagrangian means from
simulation data. Most implementations to date (Nagai et al., 2015; Shakespeare
& Hogg, 2017, 2018, 2019; Shakespeare et al., 2021; Bachman et al., 2020;
Jones et al., 2023) rely on tracking a large number of particles. This is
computationally costly and often does not deliver Lagrangian mean fields in
full agreement with the GLM definitions. Recent papers by Kafiabad (2022) and
Kafiabad & Vanneste (2023, hereafter KV23) develop an approach based on the
formulation of partial differential equations (PDEs) satisfied by the
Lagrangian mean fields. These PDEs can be solved on the fly, together with the
dynamical equations and using the same spatial discretisation. These papers
employ a straightforward ‘top-hat’ time averaging but the approach generalises
to more sophisticated frequency filters (Baker et al., 2024).

This approach remains costly in that independent evolutionary PDEs are
integrated over (fast) time to derive the mean field at a single (slow) time.
The process then needs to be repeated for each slow time at which averaged
fields are desired. There is, however, one particular definition of the time
mean that avoids this complication. This is the exponential mean, obtained by
convolving the signal with a truncated exponential, which enjoys the unique
property of satisfying a closed (first order) differential equation (§2). Its
benefit for Lagrangian averaging were pointed out to the authors by O. Bühler.

The first aim of this paper is to adapt the approach of KV23 to the
exponential mean. This leads to PDEs for the Lagrangian mean fields that share
a single time variable with the dynamical equations and can therefore be
solved in tandem with them, with little overheads. We derive these PDEs for
the Lagrangian averaging of scalar fields (§3). We then demonstrate their
value with an application to a rotating shallow-water flow that combines a
slowly evolving turbulent flow with a large-amplitude fast Poincaré wave (§4).

The second aim of this paper is to extend the computation of Lagrangian means,
in particular of the exponential mean, to tensor fields. This is motivated by
the need to average the dynamical variables of the fluid dynamical equations
which are typically not just scalars. This is vital if GLM theory is to be
used as a practical modelling tool. We formulate the averaging of tensors and
apply it to two fields extracted from the shallow-water flow simulation:
momentum and mass, interpreted as 1- and 2-forms, respectively (§5). This
interpretation is rooted in the geometric view of GLM developed by Gilbert &
Vanneste (2018) (see also Holm, 2002a, b; Gilbert & Vanneste, 2024). We also
compute the pseudomomentum and the mean flow Jacobian, important dynamical
fields that naturally emerge in this view (Andrews & McIntyre, 1978; Bühler &
McIntyre, 1998; Bühler, 2014).

Overall, this paper demonstrates the effectiveness and simplicity of
implementation of the exponential mean which could make it the go-to method
for the computation of Lagrangian means.

##  2 Lagrangian and exponential means

The GLM time average is defined as follows. Given the flow map $\bm{\varphi}$,
such that $\bm{\varphi}(\bm{a},t)$ is the position at time $t$ of a fluid
particle labelled by $\bm{a}$, we first define a mean flow map by

$\bar{\bar{\varphi}}^{i}(\bm{a},t)=\overline{\varphi^{i}(\bm{a},\cdot)}(t),$
(1)

where the superscript $i$ indicates the coordinate and the overbar on the
right-hand side denotes any moving time average evaluated at time $t$. We
emphasise that this construction depends on choosing a coordinate
representation of the flow map, since the average of the flow map itself (as
opposed to its coordinate representation) is ill defined. We use the
unconventional double-bar notation of Gilbert & Vanneste (2024) as a reminder
of this. The Lagrangian mean of scalar field $g(\bm{x},t)$ is defined by

$\overline{g}^{\mathrm{L}}(\bar{\bar{\bm{\varphi}}}(\bm{a},t),t)=\overline{g(\bm{\varphi}(\bm{a},\cdot),\cdot)}(t).$
(2)

In words, the Lagrangian mean $\overline{g}^{\mathrm{L}}(\bm{x},t)$ is the
time average of $g$ following the trajectory of the particle whose mean
position at $t$ is $\bm{x}=\bar{\bar{\bm{\varphi}}}(\bm{a},t)$.

KV23 develop an approach for the computation of Lagrangian means of the form
(2) based on the solution of PDEs. They implement their approach for the ‘top-
hat’ mean, that is, straightforward unweighted integration over a time
interval $T$. This is extended by Baker et al. (2024) to the case of a general
weighted mean. In this paper, we examine the special case of the ‘exponential
mean’. This is defined for functions of time $h(t)$ as

$\overline{h}(t)=\alpha\int_{-\infty}^{t}\mathrm{e}^{-\alpha(t-s)}h(s)\,\mathrm{d}s,$
(3)

where the parameter $\alpha>0$ is interpreted as the inverse of an averaging
time scale. An advantage of the exponential average is that, in the class of
linear averages of the form

$\overline{h}(t)=\int_{-\infty}^{\infty}k(t-s)h(s)\,\mathrm{d}s$ (4)

for some kernel $k(t-s)$, it is the only one which leads to a first-order
differential equation for the average, namely

$\frac{\mathrm{d}\overline{h}}{\mathrm{d}t}=\alpha(h-\overline{h}).$ (5)

This makes it possible to compute the average on-the-fly as new values of
$h(t)$ come in. For the computation of Lagrangian means following KV23, the
exponential mean avoids the need to treat the averaging time as a
(discretised) time variable distinct from the time variable used for the
dynamical equations, as is required for the top-hat and other means. The
exponential Lagrangian average is simply computed by solving additional PDEs
at the same time steps as the dynamical equations. We derive these PDEs next.

##  3 Computing the exponential Lagrangian mean

Particularising (1) and (2) to the exponential mean (3) gives the explicit
forms

$\displaystyle\bar{\bar{\bm{\varphi}}}(\bm{a},t)$
$\displaystyle=\alpha\int_{-\infty}^{t}\mathrm{e}^{-\alpha(t-s)}\bm{\varphi}(\bm{a},s)\,\mathrm{d}s$
(6)
$\displaystyle\textrm{and}\qquad\overline{g}^{\mathrm{L}}(\bar{\bar{\bm{\varphi}}}(\bm{a},t),t)$
$\displaystyle=\alpha\int_{-\infty}^{t}\mathrm{e}^{-\alpha(t-s)}g(\bm{\varphi}(\bm{a},s),s)\,\mathrm{d}s$
(7)

for the mean map and Lagrangian mean of a scalar. (We abuse notation in (6) by
using the same symbol $\bar{\bar{\bm{\varphi}}}$ for for both the mean map and
the list of its components $\bar{\bar{\varphi}}^{i}$. We repeat this abuse
with other maps in what follows.) The Lagrangian velocity $\bar{\bar{\bm{u}}}$
is defined naturally as the time derivative of the mean map:

$\partial_{t}\bar{\bar{\bm{\varphi}}}(\bm{a},t)=\bar{\bar{\bm{u}}}(\bar{\bar{\bm{\varphi}}}(\bm{a},t),t).$
(8)

Time differentiation of (6) gives

$\bar{\bar{\bm{u}}}(\bm{x},t)=\alpha(\bm{\Xi}(\bm{x},t)-\bm{x})$ (9)

on substituting $\bm{x}=\bar{\bar{\bm{\varphi}}}(\bm{a},t)$. Here we introduce
the so-called lifting map

$\bm{\Xi}=\bm{\varphi}\circ\bar{\bar{\bm{\varphi}}}^{-1}$ (10)

such that $\bm{\Xi}(\bm{x},t)$ is the actual position at time $t$ of the fluid
particle with mean position $\bm{x}$. We note that $\bar{\bar{\bm{u}}}$ in (9)
also results from applying the exponential Lagrangian mean (7) to each
component of the velocity field $\bm{u}$. This follows from the computation

$\displaystyle\alpha\int_{-\infty}^{t}\mathrm{e}^{-\alpha(t-s)}\bm{u}(\bm{\varphi}(\bm{a},s),s)\,\mathrm{d}s$
$\displaystyle=\alpha\int_{-\infty}^{t}\mathrm{e}^{-\alpha(t-s)}\partial_{s}\bm{\varphi}(\bm{a},s)\,\mathrm{d}s$
$\displaystyle=\alpha\left[\mathrm{e}^{-\alpha(t-s)}\bm{\varphi}(\bm{a},s)\right]_{-\infty}^{t}-\alpha^{2}\int_{-\infty}^{t}\mathrm{e}^{-\alpha(t-s)}\bm{\varphi}(\bm{a},s)\,\mathrm{d}s$
$\displaystyle=\alpha\left(\bm{\varphi}(\bm{a},t)-\bar{\bar{\bm{\varphi}}}(\bm{a},t)\right)$
$\displaystyle=\alpha\left(\bm{\Xi}(\bm{x},t)-\bm{x}\right)=\bar{\bar{\bm{u}}}(\bm{x},t),$
(11)

where $\bm{x}=\bar{\bar{\bm{\varphi}}}(\bm{a},t)$ and the last equality uses
(9). We nonetheless refrain from identifying $\bar{\bar{\bm{u}}}$ with
$\overline{\bm{u}}^{\mathrm{L}}$ for reasons that become clear in §5.

The lifting map $\bm{\Xi}$ is required for computation of Lagrangian averages.
It can be computed by solving the PDE obtained by differentiating the identity

$\bm{\Xi}(\bar{\bar{\bm{\varphi}}}(\bm{a},t),t)=\bm{\varphi}(\bm{a},t)$ (12)

with respect to $t$ and substituting
$\bm{x}=\bar{\bar{\bm{\varphi}}}(\bm{a},t)$ to find

$\partial_{t}\bm{\Xi}+\bar{\bar{\bm{u}}}\bcdot\bnabla\bm{\Xi}=\bm{u}\circ\bm{\Xi}.$
(13)

Taking (9) into account, this is a closed PDE for $\bm{\Xi}$ that can be
solved alongside the dynamical equations.

Time differentiation of (7) gives a PDE for the Lagrangian mean of the scalar
field $g(\bm{x},t)$, namely

$\partial_{t}\overline{g}^{\mathrm{L}}+\bar{\bar{\bm{u}}}\bcdot\bnabla\overline{g}^{\mathrm{L}}=\alpha(g\circ\bm{\Xi}-\overline{g}^{\mathrm{L}}).$
(14)

We solve this together with (13) since $\bm{\Xi}$ is required both for
$\bar{\bar{\bm{u}}}$ on the left-hand side and for the composition on the
right-hand side. As initial conditions we use $\bm{\Xi}(\bm{x},0)=\bm{x}$ and
$\overline{g}^{\mathrm{L}}(\bm{x},0)=0$. This leads to
$\overline{g}^{\mathrm{L}}(\bm{x},t)$, at least for $t$ large enough that the
integration limit $t\to-\infty$ in (7) can be replaced by $t=0$.

Eqs. (9), (13) and (14) are key results of this paper. They provide a route
for an efficient, easy-to-implement computation of Lagrangian averages
leveraging the unique properties of the exponential mean. Eq. (14) in
particular can be thought of as a Lagrangian version of the ODE (5) for the
exponential mean.

##  4 Shallow-water example

###  4.1 Implementation

We compute Lagrangian means in a simulation of a turbulent flow interacting
with a Poincaré wave in a rotating shallow-water model (e.g. Vallis, 2017;
Zeitlin, 2018). We use the non-dimensional equations for rotating shallow
water with characteristic length $L$, characteristic velocity $U$ and
characteristic time $T=L/U$, that is,

$\displaystyle\partial_{t}{\bm{u}}+\bm{u}\bcdot\bnabla\bm{u}+Ro^{-1}\hat{\bm{z}}\times\bm{u}$
$\displaystyle=-Fr^{-2}\bnabla h,$ (15a)
$\displaystyle\partial_{t}{h}+\bnabla\bcdot(h\bm{u})$ $\displaystyle=0,$ (15b)

which introduces the Rossby and Froude numbers

$Ro=\frac{U}{fL}\quad\textrm{and}\quad Fr=\frac{U}{\sqrt{gH}}.$ (16)

Here $g$ is the gravitational acceleration, $f$ the Coriolis parameter, $H$
the mean depth and $\hat{\bm{z}}$ the vertical unit vector.

We solve the dynamical equations (15) together with the Lagrangian mean
equations (13) and (14) using a pseudospectral discretisation. To solve
exclusively for doubly periodic fields, we express the lifting map $\bm{\Xi}$
in terms of the (doubly periodic) displacement map
$\bm{\xi}(\bm{x},t)=\bm{\Xi}(\bm{x},t)-\bm{x}$ replacing (13) by

$\partial_{t}\bm{\xi}+\bar{\bar{\bm{u}}}\bcdot\bnabla\bm{\xi}=\bm{u}\circ(\mathrm{id}+\bm{\xi})-\bar{\bar{\bm{u}}},$
(17)

where $\mathrm{id}$ denotes the identity map and $\bm{\xi}(\bm{x},0)=0$, and
(14) by

$\partial_{t}\overline{g}^{\mathrm{L}}+\bar{\bar{\bm{u}}}\bcdot\bnabla\overline{g}^{\mathrm{L}}=\alpha(g\circ(\mathrm{id}+\bm{\xi})-\overline{g}^{\mathrm{L}}).$
(18)

The Lagrangian mean velocity is deduced from $\bm{\xi}$ as
$\bar{\bar{\bm{u}}}=\alpha\bm{\xi}$ obtained from (9).

We use of modified version of the code developed by KV23. We discretise (15),
(17) and (18) using $256^{2}$ grid points/Fourier modes, use an RK4 integrator
with time step $\Delta t=0.005$ and integrate over the time range $0\leq t\leq
T=50$. For numerical stability, we employ 2/3 dealiasing and add hyperviscous
dissipation to the momentum shallow-water equation (15a) which amounts to
multiplying the Fourier variables $\hat{\bm{u}}(\bm{k},t)$ by
$\exp(-\kappa|\bm{k}|^{8}\Delta t)$ at each time step. We take
$\kappa=2.6\times 10^{-14}$. We achieve a stable numerical solution of the
Lagrangian mean equations (17) and (18) without any dissipation. However, a
form of dissipation may be required to ensure numerical instability in other
flow configurations. We use bilinear interpolation to evaluate $\bm{u}$ and
$g$ at the position $\bm{x}+\bm{\xi}(\bm{x},t)$ as required for (17) and (18).

We initialise the model with the same initial condition as in KV23
superimposing a geostrophic turbulent flow, obtained by prior solution of a
quasi-geostrophic model, and a right travelling mode-$1$ Poincaré wave. We
take the root-mean square velocity of the geostrophic flow as characteristic
velocity $U$ and the length scale of the first Fourier mode as characteristic
length $L$ so that the doubly periodic domain is $[0,2\pi]^{2}$. The Rossby
and Froude number of the geostrophic flow are $Ro=0.1$ and $Fr=0.5$. With
these parameters, the frequency $\omega=(Ro^{-2}+Fr^{-2})^{\frac{1}{2}}$ of
the mode-$1$ Poincaré wave is $\omega=10.2$ corresponding to a period of
$0.62$ time units. The wave field in the absence of flow are given by

$u^{\prime}=a\hskip 2.84526pt\cos(x-\omega t),\hskip
14.22636ptv^{\prime}=\frac{a}{\omega Ro}\hskip 2.84526pt\sin(x-\omega
t),\hskip 14.22636pth^{\prime}=\frac{a}{\omega}\hskip 2.84526pt\cos(x-\omega
t).$ (19)

At the initial time $t=0$, we add these to the geostrophic flow field. We take
the amplitude $a=-1$ so that the maximum wave velocity is as large as the
geostrophic flow root-mean-square velocity.

###  4.2 Results

Figure 1: Vorticity field in the rotating shallow-water simulation at
$t=12.5,\,25,\,37.5$ and $50$: instantaneous vorticity $\zeta$ (top row),
Lagrangian mean vorticity $\overline{\zeta}^{\mathrm{L}}$ (middle row) and
Eulerian mean vorticity $\bar{\zeta}$ (bottom row). The exponential mean with
inverse averaging time scale $\alpha=0.5$ is used for the mean vorticity
fields in the middle and bottom rows. See also the animation movie 1.

We first consider the Lagrangian mean of the relative vorticity
$\zeta=\partial_{x}v-\partial_{y}u$ and thus set $g=\zeta$ in (18). The top
row of figure 1 shows $\zeta$ at $t=12.5,\,25,\,37.5$ and $50$. The mode-1
wave which we aim to filter out is a dominant feature. It is distorted by the
flow consisting of vortices and filaments familiar in two-dimensional (quasi-
geostrophic) turbulence. The middle and bottom rows compare the Lagrangian and
Eulerian means $\overline{\zeta}^{\mathrm{L}}$ and $\bar{\zeta}$ computed at
$t=50$ with an inverse averaging time scale $\alpha=0.5$. Both means filter
out the wave satisfactorily, but the Eulerian mean blurs the small-scale
vorticity structures of the flow as a result of the rapid advection by the
velocity field associated with the wave. The Lagrangian mean eliminates this
advection by construction.

Figure 2: Vorticity $\zeta$ (top) and its exponential Lagrangian mean
$\overline{\zeta}^{\mathrm{L}}$ with $\alpha=0.5$ (bottom) as functions of $t$
and $y$ for $x=0.24$ in the simulation in figure 1.

A qualitatively similar Lagrangian mean field is obtained using the top-hat
mean of KV23. However, the exponential mean provides a much simpler way of
computing a Lagrangian mean over an entire time interval rather than a single
snapshot. In particular, since the exponential Lagrangian mean computation
shares the same (small) time step as the solution of the dynamical equations,
it resolves the fast wave time scales in the system, making it straightforward
to examine the wave signal that remains after averaging.

We illustrate this with figure 2, which shows a Hovmöller diagram of $\zeta$
and $\overline{\zeta}^{\mathrm{L}}$ at fixed $x=0.24$, and with an animation
of the evolution of $\zeta$, $\overline{\zeta}^{\mathrm{L}}$ and $\bar{\zeta}$
(movie 1). The wave signal in the instantaneous field $\zeta$ is strongly
reduced by the exponential averaging, yielding a mean field
$\overline{\zeta}^{\mathrm{L}}$ that captures the slow dynamics well, but the
wave signal is not eliminated completely. This is not surprising: the
exponential mean is a low-pass filter with a broad frequency response which
reduces rather than eliminates high frequencies (its Fourier-domain transfer
function at frequency $\Omega$ is $(\alpha-\mathrm{i}\Omega)^{-1}$). This
limitation of the exponential mean is also evident in the relation
$\bar{\bar{\bm{u}}}=\alpha\bm{\xi}$ between the (slow) mean velocity
$\bar{\bar{\bm{u}}}$ and the (fast) displacement $\bm{\xi}$. The notional
slowness of $\bar{\bar{\bm{u}}}$ stems only from the property
$|\bar{\bar{\bm{u}}}|\ll|\partial_{t}\bm{\xi}|$, which is obeyed provided that
$\alpha/\omega\ll 1$, whereas one can ask for stricter slowness conditions,
not satisfied by the exponential mean, which demand that
$|\partial^{n}_{t}\bar{\bar{\bm{u}}}|\sim(\alpha/\omega)^{n}|\partial^{n+1}_{t}\bm{\xi}|$
for $n=0,\,1,\,2,\cdots$.

Figure 3: Vorticity at $(x,y)=(1.2,1.2)$ as a function of $t$: the
instantaneous value $\zeta$ is compared with the exponential Lagrangian
average with $\alpha=0.2,\,0.5$ and $1$.

The choice of $\alpha$ determines which range of frequencies are filtered out
effectively by the exponential mean. Figure 3 illustrates the effect of
varying $\alpha$ by showing $\zeta$ and $\overline{\zeta}^{\mathrm{L}}$ at the
fixed position $(x,y)=(1.2,1.2)$ as a function of time for three values of
$\alpha$, all for which such that $\alpha/\omega\ll 1$. Clearly, decreasing
$\alpha$ reduces the magnitude of the fast oscillations induced by the wave.
For the smallest value of $\alpha$, $\alpha=0.2$,
$\overline{\zeta}^{\mathrm{L}}$ differs markedly from the (Eulerian) average
that is estimated by eye from the instantaneous $\zeta$; this is a reminder
that Lagrangian averaging is non-local in space and depends on the field at
positions other than where it is evaluated. In applications, the choice of
$\alpha$ should be guided by the use made of the Lagrangian averaging. In
particular, when the aim is to assess the impact of waves on the mean flow,
there is a trade off between the requirement to filter out fast waves and the
need to retain the dynamically significant time scales of the flow. We return
to such applications in §5. Our experimentations show that, for the flow
considered, $\alpha=0.5$ is a good compromise. This suggests that, more
generally, exponential averaging with inverse averaging time $\alpha$ is
effective in filtering out waves of frequencies $\omega\gtrsim 20\alpha$.

Figure 4: Same as figure 1 but for the potential vorticity anomaly $q$ at
$t=12.5,\,25,\,37.5$ and $50$: $q$ (top row), $\overline{q}^{\mathrm{L}}$
(middle row) and $\bar{q}$ (bottom row). See also the animation movie 2.

A scalar of particular interest is the potential vorticity (PV)

$q=\frac{Ro^{-1}+\zeta}{h}-Ro^{-1}$ (20)

which we give here in non-dimensional form with the constant background PV
subtracted. Because the PV of the Poincaré wave is zero (at a linear level),
it is impacted by the wave in a very different way to the relative vorticity
$\zeta$: advection by the velocity field associated with the wave induces a
fast oscillation of PV structures that are otherwise unaffected. These
oscillations limit the usefulness of the Eulerian mean PV compared with the
Lagrangian mean PV. We illustrate this in figure 4 which shows snapshots of
the PV and of its exponential Lagrangian and Eulerian means. Animations of
$q$, $\overline{q}^{\mathrm{L}}$ and $\bar{q}$ are available in movie 2. The
Eulerian mean PV is blurred whereas the Lagrangian mean retains all the
features of the instantaneous field. In the absence of dissipation, the
material conservation of $q$ implies that $q=q_{0}\circ\bm{\varphi}^{-1}$,
with $q_{0}$ the initial PV field, hence
$\overline{q}^{\mathrm{L}}=\overline{q\circ\bm{\Xi}}$ is the average of
$q_{0}\circ\bar{\bar{\bm{\varphi}}}^{-1}$ and approximately equal to
$q_{0}\circ\bar{\bar{\bm{\varphi}}}^{-1}$ itself insofar as
$\bar{\bar{\bm{\varphi}}}^{-1}$ can be regarded as slow. Therefore, the
Lagrangian mean PV is approximately a rearrangement of the instantaneous PV,
thus sharing the same extrema and topology.

##  5 Exponential mean of tensors

A key benefit of Lagrangian averaging is that, when applied to the dynamical
equations, it leads to particularly simple mean equations that inherit much of
the structure of the original equations and hence preserve properties of
inviscid fluids such as the material transport of vorticity and the
conservation of Kelvin’s circulation (Andrews & McIntyre, 1978; Holm, 2019,
2002a, 2002b; Gilbert & Vanneste, 2018). Taking advantage of this requires
extending the concept of Lagrangian averaging from the scalar fields
considered in §3 to arbitrary tensor fields. We now discuss this extension. We
build on the geometric interpretation of GLM in Gilbert & Vanneste (2018) and
Gilbert & Vanneste (2024) and refer the reader to these papers for an
introduction to the underlying concepts and notation.

###  5.1 Averaging tensors

The extension of the exponential Lagrangian mean to a tensor field
$\tau(\bm{x},t)$ reads

$(\bar{\bar{\bm{\varphi}}}^{*}\overline{\tau}^{\mathrm{L}})(\bm{a},t)=\alpha\int_{-\infty}^{t}\mathrm{e}^{-\alpha(t-s)}(\bm{\varphi}^{*}\tau)(\bm{a},s)\,\mathrm{d}s.$
(21)

Here, $\bar{\bar{\bm{\varphi}}}^{*}$ and $\bm{\varphi}^{*}$ denotes the pull-
backs of the tensor $\tau$ by $\bar{\bar{\bm{\varphi}}}$ and $\bm{\varphi}$
(e.g. Frankel, 2004). These take different forms depending on the nature of
$\tau$: in particular, for scalar fields,

$(\bm{\varphi}^{*}g)(\bm{a},t)=g(\bm{\varphi}(\bm{a},t),t),$ (22)

making (21) consistent with (7), for (contravariant) vector fields
$\bm{v}=v^{i}(\bm{x},t)\bm{e}_{i}$,

$(\bm{\varphi}^{*}\bm{v})^{i}(\bm{a},t)=\mathsfi{K}^{i}_{\
j}(\bm{a},t)v^{j}(\bm{\varphi}(\bm{a},t),t),$ (23)

and for 1-forms fields (i.e. covariant vectors or covectors)
$\bm{\beta}=\beta_{i}(\bm{x},t)\,dx^{i}$,

$(\bm{\varphi}^{*}\bm{\beta})_{i}(\bm{a},t)=\beta_{j}(\bm{\varphi}(\bm{a},t),t)\mathsfi{J}_{\
i}^{j}(\bm{a},t).$ (24)

The matrix $\mathsfi{J}$ in (24) is the Jacobian matrix

$\mathsfi{J}_{\ i}^{j}(\bm{a},t)=\frac{\partial\varphi^{j}(\bm{a},t)}{\partial
a^{i}},$ (25)

and $\mathsfi{K}$ in (23) is its inverse, $\mathsfi{K}=\mathsfi{J}^{-1}$, i.e.
it satisfies $\mathsfi{K}^{i}_{\ k}\,\mathsfi{J}^{k}_{\ j}=\delta^{i}_{\ j}$.
Note that, when applied to the velocity vector field $\bm{u}$, (21) and (23)
give a definition of $\overline{\bm{u}}^{\mathrm{L}}$ that differs from the
expression (11) for $\bar{\bar{\bm{u}}}$. The latter expression corresponds to
averaging the components of $\bm{u}$ as scalar fields rather than averaging
$\bm{u}$ as a vector field. The distinction justifies the notation
$\bar{\bar{\bm{u}}}$ for the mean velocity.

Taking the time derivative of (21), applying the push-forward
$\bar{\bar{\bm{\varphi}}}_{*}$ (the inverse of the pull-back) and noting that
(10) implies that $\bar{\bar{\bm{\varphi}}}_{*}\bm{\varphi}^{*}=\bm{\Xi}^{*}$
leads to

$(\partial_{t}+\mathcal{L}_{\bar{\bar{\bm{u}}}})\overline{\tau}^{\mathrm{L}}=\alpha(\bm{\Xi}^{*}\tau-\overline{\tau}^{\mathrm{L}}).$
(26)

This is the desired extension of (14) from scalar fields to arbitrary tensor
fields. Here we have introduced the Lie derivative
$\mathcal{L}_{\bar{\bar{\bm{u}}}}$ whose expression differs depending on the
nature of the tensor field it applies to. In particular,

$\mathcal{L}_{\bar{\bar{\bm{u}}}}g=\bar{\bar{u}}^{j}\partial_{j}g,\quad(\mathcal{L}_{\bar{\bar{\bm{u}}}}\bm{v})^{i}=\bar{\bar{u}}^{j}\partial_{j}v^{i}-v^{j}\partial_{j}\bar{\bar{u}}^{i}\quad\textrm{and}\quad(\mathcal{L}_{\bar{\bar{\bm{u}}}}\bm{\beta})_{i}=\bar{\bar{u}}^{j}\partial_{j}\beta_{i}+\beta_{j}\partial_{i}\bar{\bar{u}}^{j}$
(27)

for scalar, vector and 1-form fields. Eq. (26) should be solved together with
(13) for $\bm{\Xi}$ and using (9) for $\bar{\bar{\bm{u}}}$. For tensor fields
other than scalars, the solution requires the evaluation of the derivatives of
$\bm{\Xi}$ and $\bar{\bar{\bm{u}}}$ that appear in the pull-back and Lie
derivative.

###  5.2 Momentum, pseudomomentum and mass

A tensor field of particular interest is the momentum 1-form

$\bm{\nu}=\bm{u}\bcdot d\bm{x}.$ (28)

It appears as the integrand in Kelvin’s circulation and is the central
dynamical variable in formulations of the fluid equations best suited for
Lagrangian averaging. In Euclidean geometry and Cartesian coordinates, the
components of $\bm{\nu}$ are identical to those of $\bm{u}$, $\nu_{i}=u^{i}$,
but the 1-form nature of $\bm{\nu}$ matters for the explicit form of (26).
This is given by

$\partial_{t}\overline{\nu}^{\mathrm{L}}_{i}+\bar{\bar{u}}^{j}\partial_{j}\overline{\nu}^{\mathrm{L}}_{i}+\overline{\nu}^{\mathrm{L}}_{j}\partial_{i}\bar{\bar{u}}^{j}=\alpha\left(\partial_{i}\Xi^{j}\,(\nu_{j}\circ\bm{\Xi})-\overline{\nu}^{\mathrm{L}}_{i}\right),$
(29)

on using (24) and (27).

Figure 5: Mean velocity $\bar{\bar{u}}_{1}$ (top), Lagrangian mean momentum
$\overline{\nu}^{\mathrm{L}}_{1}$ (middle) and pseudomomentum
$\mathsf{p}_{1}=\bar{\bar{u}}_{1}-\overline{\nu}^{\mathrm{L}}_{1}$ (bottom) as
functions of $t$ and $y$ for $x=0.24$ in the simulation in figure 1.

As a proof of concept, we compute the Lagrangian mean momentum
$\overline{\bm{\nu}}^{\mathrm{L}}$ in the shallow-water simulation of §4. We
solve (29) and show in figure 5 a comparison between the mean velocity
$\bar{\bar{\bm{u}}}$ and Lagrangian mean momentum
$\overline{\bm{\nu}}^{\mathrm{L}}$ in the form of Hovmöller diagrams for their
first components $\bar{\bar{u}}_{1}$ and $\overline{\nu}^{\mathrm{L}}_{1}$.
According to GLM theory, and ignoring the effect of the Coriolis force for
simplicity, the primary impact of perturbations filtered out by the Lagrangian
mean on the dynamics of the mean flow is encoded in the difference between
$\overline{\bm{\nu}}^{\mathrm{L}}$ and $\bar{\bar{\bm{u}}}$, or more precisely
in the pseudomomentum, defined in Cartesian coordinates by

$-\bm{\mathsf{p}}=\overline{\bm{\nu}}^{\mathrm{L}}-\bar{\bar{u}}_{1}\,dx_{1}-\bar{\bar{u}}_{2}\,dx_{2}.$
(30)

The bottom panel of figure 5 shows its first component $\mathsf{p}_{1}$. The
magnitude of $\mathsf{p}_{1}$ is significantly smaller than that of both
$\bar{\bar{u}}_{1}$ and $\overline{\nu}^{\mathrm{L}}_{1}$, and the same
applies to the second components (not shown). This suggests that the Poincaré
wave has only a moderate impact on the evolution of the Lagrangian mean flow.
The pseudomomentum displays a wave signal – the result of the limitations of
the expontial mean discussed in §4.2 – but also low frequency patterns. These
capture the impact of the wave on the mean flow and would be the target for
parameterisation in closed models describing solely the slow dynamics of the
mean flow.

Another tensor of interest is the mass, thought of as the 2-form
$h\,dx_{1}\wedge dx_{2}$. Writing

$\overline{h\,dx_{1}\wedge dx_{2}}^{\mathrm{L}}=\tilde{h}\,dx_{1}\wedge
dx_{2},$ (31)

which defines the effective height $\tilde{h}$ (Bühler & McIntyre, 1998), we
apply (26) to find

$\partial_{t}\tilde{h}+\bnabla\bcdot(\tilde{h}\bar{\bar{\bm{u}}})=\alpha\left(\left|\partial_{i}\Xi^{j}\right|\,(h\circ\bm{\Xi})-\tilde{h}\right),$
(32)

where
$\left|\partial_{i}\Xi^{j}\right|=\partial_{1}\Xi^{1}\partial_{2}\Xi^{2}-\partial_{2}\Xi^{1}\partial_{1}\Xi^{2}$
is the Jacobian of the map $\bm{\Xi}$, and we use that
$\mathcal{L}_{\bar{\bar{\bm{u}}}}(dx_{1}\wedge
dx_{2})=(\bnabla\bcdot\bar{\bar{\bm{u}}})\,dx_{1}\wedge dx_{2}$ and
$\bm{\Xi}^{*}(dx_{1}\wedge
dx_{2})=\left|\partial_{i}\Xi^{j}\right|\,dx_{1}\wedge dx_{2}$. The
conservation of the Lagrangian mean mass,
$(\partial_{t}+\mathcal{L}_{\bar{\bar{\bm{u}}}})\overline{h\,dx_{1}\wedge
dx_{2}}^{\mathrm{L}}=0$, implies that the left-hand side of (32) vanishes,
leaving

$\tilde{h}=(h\circ\bm{\Xi})\left|\partial_{i}\Xi^{j}\right|.$ (33)

(This also follows from the fact that $\Xi^{*}(h\,dx_{1}\wedge
dx_{2})=\bar{\bar{\varphi}}_{*}(h_{0}\,dx_{1}\wedge dx_{2})$ with $h_{0}$ the
initial height field (cf. Andrews & McIntyre, 1978; Bühler, 2014).) This shows
that the effective height, which controls the mass distribution of the
Lagrangian mean flow is not a simple rearrangement of the actual height but
depends on the Jacobian $\left|\partial_{i}\Xi^{j}\right|$. This Jacobian is
sensitive to the choice of $\alpha$ and becomes large for small $\alpha$
because $\Xi$ then captures part of the slow motion which is exponentially
stretching, unlike the wave part.

Figure 6: Jacobian $\left|\partial_{i}\Xi^{j}\right|$ at $t=50$ in the
simulation of figure 1 for inverse averaging time $\alpha=0.2$ (left), $0.5$
(middle) and $1$ (right).

We illustrate this in figure 6 which shows the Jacobian
$\left|\partial_{i}\Xi^{j}\right|$ at the end of the shallow-water simulation
for $\alpha=0.2,\,0.5$ and $1$. The variations of the Jacobian away from the
area-preserving value of $1$ increase as $\alpha$ decreases, corresponding to
longer averaging times. For the smallest value $\alpha=0.2$, the averaging
time is so long as to overlap with the dynamical time scales of the slow flow;
as a result, the map $\bm{\Xi}$ differs significantly from the identity. The
associated divergence effect, that is, the compressibility of the mean flow
arising from the GLM component-wise definition of the mean (Andrews &
McIntyre, 1978; McIntyre, 1988), leads to a Jacobian that does not remain
close to $1$. This has implications both for the interpretation of the mean as
faithfully representing the dynamics and for numerical computations. In
comparison, the largest value $\alpha=1$ leads to a Jacobian that remains
close to $1$, but it does not give a sufficiently long averaging time to
filter out the wave effectively.

##  6 Discussion

In this paper, we propose and test a method for the numerical computation of
temporal Lagrangian means from simulation data. We adapt the PDE-based
algorithm of KV23 to a specific choice of time averaging, namely the
exponential mean (3). This mean has the unique advantage that all the fields
required in the computation satisfy evolution equations with the same time
variable as that of the dynamical equations. The equations for the Lagrangian
mean fields can therefore be solved on-the-fly, together with the dynamical
equations, with no overheads beyond those that result from a few additional
equations (for a scalar field $g$ and in $n$ dimensions, there are $n+1$
additional equations for $\overline{g}^{\mathrm{L}}$ and $\bm{\Xi}$). This is
in contrast with other means such as the top-hat mean used by KV23 and the
more sophisticated frequency filters used by Baker et al. (2024) which require
the solution of a separate initial value problem for each of the times at
which the mean fields are desired. A downside of the exponential mean is
however its broad frequency response which leads to a reduction rather than
elimination of high frequencies from the mean fields. The frequency filters of
Baker et al. (2024) perform markedly better in this respect.

The availability of a simple and efficient method for the computation of
Lagrangian mean fields paves the way for the analysis of realistic simulations
of geophysical fluid flows. In particular, it makes it possible to diagnose
the effects that waves, or more broadly the fast components of the flow
filtered out by averaging, have on the mean flow. According to GLM theory,
these are encoded in the pseudomomentum, best thought of as a 1-form (or co-
vector) and defined as minus the difference between Lagrangian mean momentum
and Lagrangian mean velocity. With this in mind, we consider the exponential
averaging of arbitrary tensor quantities and apply the resulting method to the
momentum 1-form and to the mass 2-form.

The computations in this paper rely on the standard GLM definition of the mean
flow map as a component-wise mean. While this has the advantage of simplicity,
it also has drawbacks, primarily the divergence effect associated with the
compressibility of the mean flow even for incompressible fluids. In the case
of the (compressible) shallow-water model we use here, it leads to changes in
the effective height $\tilde{h}$ that are much larger than changes in the
actual height $h$. Alternative, more geometric definitions of the mean flow
have been proposed that guarantee the incompressibility of the mean velocity
for incompressible fluids and, in the shallow-water context, would lead to
changes in $\tilde{h}$ commensurate with those in $h$ (Soward & Roberts, 2010;
Gilbert & Vanneste, 2018). It would be desirable to develop numerical
procedures for the computation of the Lagrangian mean quantities corresponding
to these definitions.

[Supplementary data]Supplementary movies are available as movie 1 and movie 2.

[Acknowledgements]We are grateful to O. Bühler for suggesting the use of the
exponential mean for Lagrangian averaging.

[Funding]L.E.B. and H.A.K. are supported by the UK Engineering and Physical
Sciences Research Council grants EP/X028135/1 and grant EP/Y021479/1,
respectively. J.V. is supported by the UK Natural Environment Research Council
grant NE/W002876/1.

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

[Author ORCIDs]A. Minz, https://orcid.org/0009-0003-4209-3428; L.E. Baker,
https://orcid.org/0000-0003-2678-3691; H. A. Kafiabad,
https://orcid.org/0000-0002-8791-9217; J. Vanneste,
https://orcid.org/0000-0002-0319-589X

## References

  * Andrews & McIntyre (1978) Andrews, D. G. & McIntyre, M. E. 1978 An exact theory of nonlinear waves on a Lagrangian-mean flow. Journal of fluid Mechanics 89 (4), 609–646.
  * Bachman et al. (2020) Bachman, S. D., Shakespeare, C. J., Kleypas, J., Castruccio, F. S. & Curchitser, E. 2020 Particle-based Lagrangian filtering for locating wave-generated thermal refugia for coral reefs. Journal of Geophysical Research: Oceans 125 (7), e2020JC016106.
  * Baker et al. (2024) Baker, L. E., Kafiabad, H. A. & Vanneste, J. 2024 Lagrangian filtering for wave–mean flow decomposition Preprint, arXiv:2406.03477.
  * Bühler (2014) Bühler, O. 2014 Waves and mean flows. Cambridge University Press.
  * Bühler & McIntyre (1998) Bühler, O. & McIntyre, M. E 1998 On non-dissipative wave–mean interactions in the atmosphere or oceans. Journal of Fluid Mechanics 354, 301–343.
  * Frankel (2004) Frankel, T. 2004 The geometry of physics, 2nd edn. Cambridge University Press.
  * Gilbert & Vanneste (2018) Gilbert, A. D. & Vanneste, J. 2018 Geometric generalised Lagrangian-mean theories. Journal of Fluid Mechanics 839, 95–134.
  * Gilbert & Vanneste (2024) Gilbert, A. D. & Vanneste, J. 2024 Geomeric approaches to Lagrangian averaging. Annual Review of Fluid Mechanics In press, arXiv:2405.04394.
  * Holm (2002a) Holm, D. D. 2002a Lagrangian averages, averaged Lagrangians, and the mean effects of fluctuation in fluid dynamics. Chaos 12, 518–530.
  * Holm (2002b) Holm, D. D. 2002b Variational principles for Lagrangian-averaged fluid dynamics. J. Phys. A. Math. Gen. 35, 679–668.
  * Holm (2019) Holm, D. D. 2019 Stochastic closures for wave–current interaction dynamics. J. Nonlinear Sci. 29, 2987–3031.
  * Jones et al. (2023) Jones, C. S., Xiao, Q., Abernathey, R. P. & Smith, K. S. 2023 Using Lagrangian filtering to remove waves from the ocean surface velocity field. Journal of Advances in Modeling Earth Systems 15 (4), e2022MS003220.
  * Kafiabad (2022) Kafiabad, H. A 2022 Grid-based calculation of the Lagrangian mean. Journal of Fluid Mechanics 940, A21.
  * Kafiabad & Vanneste (2023) Kafiabad, H. A & Vanneste, J. 2023 Computing Lagrangian means. Journal of Fluid Mechanics 960, A36.
  * McIntyre (1988) McIntyre, M. E. 1988 A note on the divergence effect and the Lagrangian-mean surface elevation in periodic water waves. Journal of Fluid Mechanics 189, 235–242.
  * Nagai et al. (2015) Nagai, T., Tandon, A., Kunze, E. & Mahadevan, A. 2015 Spontaneous generation of near-inertial waves by the Kuroshio front. Journal of Physical Oceanography 45 (9), 2381–2406.
  * Shakespeare et al. (2021) Shakespeare, C. J., Gibson, A. H., Hogg, A. M., Bachman, S. D., Keating, S. R. & Velzeboer, N. 2021 A new open source implementation of Lagrangian filtering: A method to identify internal waves in high-resolution simulations. Journal of Advances in Modeling Earth Systems 13 (10).
  * Shakespeare & Hogg (2017) Shakespeare, C. J. & Hogg, A. M. 2017 Spontaneous surface generation and interior amplification of internal waves in a regional-scale ocean model. Journal of Physical Oceanography 47 (4), 811–826.
  * Shakespeare & Hogg (2018) Shakespeare, C. J. & Hogg, A. M. 2018 The life cycle of spontaneously generated internal waves. Journal of Physical Oceanography 48 (2), 343–359.
  * Shakespeare & Hogg (2019) Shakespeare, C. J. & Hogg, A. M. 2019 On the momentum flux of internal tides. Journal of Physical Oceanography 49 (4), 993–1013.
  * Soward & Roberts (2010) Soward, A. M. & Roberts, P. H. 2010 The hybrid Euler–Lagrange procedure using an extension of Moffatt’s method 661, 45–72.
  * Vallis (2017) Vallis, G. K. 2017 Atmospheric and oceanic fluid dynamics. Cambridge University Press.
  * Zeitlin (2018) Zeitlin, V. 2018 Geophysical fluid dynamics: understanding (almost) everything with rotating shallow water models. Oxford University Press.