###### Abstract

We consider estimating the marginal likelihood in settings with independent
and identically distributed (i.i.d.) data. We propose estimating the
predictive distributions in a sequential factorization of the marginal
likelihood in such settings by using stochastic gradient Markov Chain Monte
Carlo techniques. This approach is far more efficient than traditional
marginal likelihood estimation techniques such as nested sampling and annealed
importance sampling due to its use of mini-batches to approximate the
likelihood. Stability of the estimates is provided by an adaptive annealing
schedule. The resulting stochastic gradient annealed importance sampling
(SGAIS) technique, which is the key contribution of our paper, enables us to
estimate the marginal likelihood of a number of models considerably faster
than traditional approaches, with no noticeable loss of accuracy. An important
benefit of our approach is that the marginal likelihood is calculated in an
online fashion as data becomes available, allowing the estimates to be used
for applications such as online weighted model combination.

###### keywords:

marginal likelihood; evidence; nested sampling; annealed importance sampling;
Monte Carlo; stochastic gradients; SGHMC

21 11 1109 Received: 30 September 2019; Accepted: 4 October 2019; Published:
12 November 2019 yes Stochastic Gradient Annealed Importance Sampling for
Efficient Online Marginal Likelihood Estimation † Scott A. Cameron 1,2,*, Hans
C. Eggers 1,2 and Steve Kroon 3Scott A. Cameron, Hans C. Eggers and Steve
Kroon Correspondence<EMAIL_ADDRESS>MaxEnt 2019

##  1 Introduction

Marginal likelihood (ML), sometimes called evidence, is a quantitative measure
of how well a model can describe a particular data set; it is the probability
that the data set occurred within that model. Consider a Bayesian model with
parameters ${\bm{\theta}}$ for a data set
${\mathcal{D}}=\\{{\mathbf{y}}_{n}\\}_{n=1}^{N}$. The ML is the integral

${\mathcal{Z}}:=p({\mathcal{D}})=\int
p({\mathcal{D}}|{\bm{\theta}})p({\bm{\theta}})\,d{\bm{\theta}},$

where $p({\mathcal{D}}|{\bm{\theta}})$ is the likelihood and
$p({\bm{\theta}})$ is the prior. In this paper, we consider the case where the
data are conditionally independent given the parameters
$p({\mathcal{D}}|{\bm{\theta}})=\prod_{n}p({\mathbf{y}}_{n}|{\bm{\theta}})$,
as is common in many parametric models. This restriction is to ensure that a
central limit theorem applies to the stochastic likelihood approximation. This
can be weakened to any factorization that exhibits a central limit theorem,
such as conditionally Markov data and autoregressive models. The posterior
distribution over a set of models is proportional to their MLs, and so
approximations to ML are sometimes used for model comparison and weighted
model averaging Barber (2012) (chapter 12). The above integral is typically
analytically intractable for any but the simplest models, so one must resort
to numerical approximation methods.

Nested sampling (NS) Skilling (2006) and annealed importance sampling (AIS)
Neal (1998) are two algorithms able to produce accurate estimates of the ML.
NS accomplishes this by sampling from the prior under constraints of
increasing likelihood and AIS by sampling from a temperature annealed
distribution
$\propto~{}p({\mathcal{D}}|{\bm{\theta}})^{\lambda}p({\bm{\theta}})$ and
averaging over samples with appropriately calculated importance weights.
Although NS and AIS produce accurate estimates of the ML, they tend to scale
poorly to large data sets due to the fact that they need to repeatedly
calculate the likelihood function: for NS, this is to ensure staying within
the constrained likelihood contour; for AIS, the likelihood must be calculated
both to sample from the annealed distributions using some Markov chain Monte
Carlo (MCMC) method as well as to calculate the importance weights.
Calculation of the likelihood is computationally expensive on large data sets.
To combat this problem, various optimization and sampling algorithms instead
use stochastic approximations of the likelihood by sub-sampling the data set
into mini-batches Welling and Teh (2011).

Unfortunately, NS and AIS cannot trivially use mini-batching in their vanilla
form to improve scalability. Using stochastic likelihood approximations
changes the statistics of the likelihood contours in NS, allowing particles to
occasionally move to lower likelihood instead of higher, violating the basic
assumptions of the algorithm. Vanilla AIS could benefit from using stochastic
likelihood gradients during the MCMC steps, but introducing stochasticity into
the importance weights would bias the results.

The key contributions of this work are as follows:

  * •

We introduce stochastic gradient annealed importance sampling (SGAIS), which
combines stochastic gradient MCMC with annealed importance sampling to
estimate the ML in an online fashion using mini-batch Bayesian updating.

  * •

SGAIS enables efficient ML estimation for streaming data and for large data
sets, which was not previously feasible.

  * •

We illustrate how SGAIS can be used to identify distribution shift in the data
when applied in an online setting.

  * •

We empirically analyze the behavior of SGAIS and its robustness to various
choices of algorithm parameters.

We illustrate our approach by calculating ML estimates on simulated data
generated with three simple models. For these models, we obtain considerable
speedup over nested sampling and annealed importance sampling on data sets
with one million observations, without noticeable loss in accuracy.

##  2 Sequential Marginal Likelihood Estimation

The ML can be factorized, through the product rule, into a product of
predictive distributions
${\mathcal{Z}}=\prod_{n}p({\mathbf{y}}_{n}|{\mathbf{y}}_{<n})$, where

$p({\mathbf{y}}_{n}|{\mathbf{y}}_{<n})=\int
p({\mathbf{y}}_{n}|{\bm{\theta}})p({\bm{\theta}}|{\mathbf{y}}_{<n})\,d{\bm{\theta}}.$
(1)

Throughout this manuscript, we will present our approach as Bayesian updating
based on one observation at a time. We do this for notational simplicity, with
the understanding that the extension to Bayesian updating with multiple
observations is mathematically trivial. In our experiments, Section 6, we
decompose the data into chunks of a fixed size, rather than one data point at
a time. Assuming one is able to produce accurate estimates
$\hat{p}({\mathbf{y}}_{n}|{\mathbf{y}}_{<n})$ of the predictive probabilities,
the log-ML can be approximated by
$\log\hat{\mathcal{Z}}=\sum_{n}\log\hat{p}({\mathbf{y}}_{n}|{\mathbf{y}}_{<n})$.
In this way, the difficult problem of estimating an integral of an extremely
peaked function, $p({\mathcal{D}}|{\bm{\theta}})$, reduces to the easier
problem of estimating many integrals of smoother functions
$p({\mathbf{y}}_{n}|{\bm{\theta}})$.

Many sequential Monte Carlo (SMC) methods use a similar approach and calculate
predictive estimates using a combination of importance resampling and MCMC
mutation steps Naesseth et al. (2019). Generic examples of such algorithms are
the bootstrap particle filter Gordon et al. (1993), which is often used for
posterior inference in hidden Markov models and other latent variable sequence
models Naesseth et al. (2019), and the “left-to-right” algorithm, which is
used in Wallach et al. (2009) to evaluate topic models.

The computational efficiency of using this approach depends on the method of
approximating Equation (1). In previous work Cameron et al. (2019), we
proposed the estimator
$\hat{p}({\mathbf{y}}_{n}|{\mathbf{y}}_{<n})=\frac{1}{M}\sum_{i=1}^{M}p({\mathbf{y}}_{n}|{\bm{\theta}}_{i})$,
where each ${\bm{\theta}}_{i}$ is drawn from the posterior distribution
$p({\bm{\theta}}|{\mathbf{y}}_{<n})$ using MCMC methods. This paper extends
this idea by combining Bayesian updating with annealing, which we describe in
Section 3. Since samples from the previous posterior,
$p({\bm{\theta}}|{\mathbf{y}}_{<n-1})$, would generally be available at each
step, we expect that only a small number of steps will be needed to accurately
sample from the next posterior distribution,
$p({\bm{\theta}}|{\mathbf{y}}_{<n})$. Metropolis–Hastings-based MCMC
algorithms would have to iterate over all previous $n-1$ data points in order
to calculate the acceptance probability for each Markov transition, and so
using them to estimate $\log{\mathcal{Z}}$ in this sequential manner would
scale at least quadratically in $N$. The key computational improvement in our
approach comprises the use of stochastic gradient-based MCMC algorithms such
as the stochastic gradient Hamiltonian Monte Carlo (SGHMC) Chen et al. (2014).
SGHMC utilizes mini-batching, allowing one to efficiently draw samples from
the posterior distribution $p({\bm{\theta}}|{\mathbf{y}}_{<n})$ even when $n$
is large. We call this approach stochastic gradient annealed importance
sampling (SGAIS).

##  3 Bayesian Updating with Annealing

Annealed importance sampling (AIS) Neal (1998) is a Monte Carlo method that
generates samples from a sequence of intermediate distributions, bridging from
a simple distribution that can be sampled easily (usually the prior) to a
desired distribution (usually the posterior). AIS also produces an unbiased
estimate of the normalizing constant (the ML) for the desired distribution as
a byproduct. The sequence of distributions is typically chosen to have the
form
$f_{t}({\bm{\theta}})=p({\mathcal{D}}|{\bm{\theta}})^{\lambda_{t}}p({\bm{\theta}})$,
where 0 = $\lambda_{0}<\lambda_{1}<\ldots<\lambda_{T}=1$. At time step $t=0$,
a number of particles are drawn from the prior, and at each subsequent time
step $t=1,\ldots,T$, an MCMC operator that leaves the distribution $f_{t}$
invariant is applied to each particle. The importance weights are initialized
as the normalizing constant of the prior, which is typically 1, and at each
time step are updated as follows:

$w_{i}^{(t)}\leftarrow
w_{i}^{(t-1)}\frac{f_{t}({\bm{\theta}}_{i}^{(t-1)})}{f_{t-1}({\bm{\theta}}_{i}^{(t-1)})}=w_{i}^{(t-1)}\,p({\mathcal{D}}|{\bm{\theta}}_{i}^{(t-1)})^{\lambda_{t}-\lambda_{t-1}}\qquad
i=1,\ldots,M,$

where ${\bm{\theta}}_{i}^{(t)}$ is the $i$th particle at time step $t$. For
any function $h$, the unnormalized posterior expectation

$\int h({\bm{\theta}})f_{T}({\bm{\theta}})\,d{\bm{\theta}}=\int
h({\bm{\theta}})p({\mathcal{D}}|{\bm{\theta}})p({\bm{\theta}})\,d{\bm{\theta}}$

can be estimated using $M$ particles by
$\frac{1}{M}\sum_{i}h({\bm{\theta}}^{(T-1)}_{i})w^{(T)}_{i}$. This estimator
$\hat{H}$ is unbiased and corresponds to the ML in the case
$h({\bm{\theta}})=1$ and the product of the ML and the posterior predictive
$p({\mathbf{y}}^{\prime}|{\mathcal{D}})$ in the case
$h({\bm{\theta}})=p({\mathbf{y}}^{\prime}|{\bm{\theta}})$. Although it is
possible to use stochastic gradient-based MCMC algorithms to update the
particles in this setting, the importance weight updates would normally still
require iterating over the whole data set at each time step. To circumvent
this, we propose using AIS sequentially to calculate predictive probabilities
in a Bayesian updating setting, as described in Section 2, instead of on the
full data set. This is equivalent to using the following sequence of
intermediate distributions:

$f_{n}^{(t)}({\bm{\theta}})=p({\mathbf{y}}_{n}|{\bm{\theta}})^{\lambda_{t}}\left[\prod_{k<n}p({\mathbf{y}}_{k}|{\bm{\theta}})\right]p({\bm{\theta}}).$
(2)

In principle, $\lambda$ may now depend on both $n$ and $t$. Since we will
choose the annealing schedule adaptively anyway, this notation is suppressed.
From an SMC perspective, this can be considered a combination of thermal and
data tempering Naesseth et al. (2019). The corresponding importance weights
can then be calculated without iterating over the entire data set;

$w_{i}^{(t)}\leftarrow
w_{i}^{(t-1)}\frac{f_{n}^{(t)}({\bm{\theta}}_{i}^{(t-1)})}{f_{n}^{(t-1)}({\bm{\theta}}_{i}^{(t-1)})}=w_{i}^{(t-1)}p({\mathbf{y}}_{n}|{\bm{\theta}}_{i}^{(t-1)})^{\lambda_{t}-\lambda_{t-1}}.$

This result is of central importance to this paper: only by using AIS within a
Bayesian updating setting can we effectively take advantage of mini-batching
and stochastic gradients.

So far, the choice of each $\lambda_{t}$ is not specified. While any
increasing sequence of values for $\lambda_{t}$ (annealing schedule)
guarantees an unbiased estimator of the ML, a poor choice can result in high
variance of the ML estimator. Grosse et al. Grosse et al. (2015) point out
that a linear sequence typically results in poor performance, and instead,
they recommend a sigmoidal annealing schedule,
$\lambda_{t}=\sigma(\delta(2t/T-1))$ for some $\delta$. We expect that during
early Bayesian updating steps, i.e., when $n$ is small, the relative prior
$p({\bm{\theta}}|{\mathbf{y}}_{<n-1})$ and posterior
$p({\bm{\theta}}|{\mathbf{y}}_{<n})$ may differ substantially, and so many
intermediate distributions may be needed. However, once $n$ becomes large, the
relative prior and posterior will likely be similar and only a few
intermediate distributions should be required. For this reason, it is
preferable that the annealing schedule be chosen adaptively. An adaptive
annealing schedule can be chosen by specifying a target effective sample size
(ESS) Kong et al. (1994) and choosing the next $\lambda$ such that the ESS is
approximately equal to this target Buchholz et al. (2018); Beskos et al.
(2016). For the choice of intermediate distributions given in Equation (2),
the ESS used for the weight update is

$\operatorname*{ESS}(\Delta):=\frac{(\sum_{i}\omega_{i}(\Delta))^{2}}{\sum_{i}\omega_{i}(\Delta)^{2}},$

where $\Delta:=\lambda_{t}-\lambda_{t-1}$ and

$\omega_{i}(\Delta):=\frac{f_{n}^{(t)}({\bm{\theta}}_{i}^{(t-1)})}{f_{n}^{(t-1)}({\bm{\theta}}_{i}^{(t-1)})}=p({\mathbf{y}}_{n}|{\bm{\theta}}_{i}^{(t-1)})^{\Delta}.$

Our proposed method combining sequential Bayesian updating and AIS is outlined
in Algorithm 1. As in other SMC algorithms, one can optionally resample the
particles ${\bm{\theta}}_{i}$ proportionally to their importance weights
$w_{i}$ before performing each MCMC update on line 9, thereafter setting
$w_{i}~{}\leftarrow~{}\frac{1}{M}\sum_{i}w_{i}$. Resampling can help to reduce
variance in the estimator but can cause mode collapse if few particles are
used. See Naesseth et al. (2019); Buchholz et al. (2018) for more details on
resampling. After each annealing sequence is completed (i.e., after the
completion of the inner while-loop in the algorithm), one can estimate the ML
of the data observed so far by $\hat{p}({\mathbf{y}}_{\leq
n}):=\frac{1}{M}\sum_{i}w_{i}$.

Algorithm 1 Stochastic Gradient Annealed Importance Sampling

Input: Data ${\mathcal{D}}=\\{{\mathbf{y}}_{n}\\}_{n=1}^{N}$, number of
particles $M$, pdfs $p({\mathbf{y}}|{\bm{\theta}}),p({\bm{\theta}})$, target
ESS: $\operatorname*{ESS}^{*}$
Output: $\hat{{\mathcal{Z}}}$ estimator of marginal likelihood

1:$\forall_{i}$: sample ${\bm{\theta}}_{i}\sim p({\bm{\theta}})$

2:$\forall_{i}$: $w_{i}\leftarrow 1$

3:for $n=1,\ldots,N$ do

4: $\lambda\leftarrow 0$

5: while $\lambda<1$ do

6:
$\Delta\leftarrow\operatorname*{argmin}_{\Delta}[\operatorname*{ESS}(\Delta)-\operatorname*{ESS}^{*}]$
$\triangleright$ $\Delta\in(0,1-\lambda]$

7: $\lambda\leftarrow\lambda+\Delta$

8: $\forall_{i}$: $w_{i}\leftarrow
w_{i}p({\mathbf{y}}_{n}|{\bm{\theta}}_{i})^{\Delta}$

9: $\forall_{i}$:
${\bm{\theta}}_{i}\leftarrow\textsc{Sghmc}({\bm{\theta}}_{i},\hat{U}_{n}^{(\lambda)})$
$\triangleright$ potential energy $\hat{U}_{n}^{(\lambda)}$ defined in
Equation (4)

10: end while

11:end for

12:return $\hat{{\mathcal{Z}}}=\frac{1}{M}\sum_{i}w_{i}$

##  4 Stochastic Gradient Hamiltonian Monte Carlo

SGHMC Chen et al. (2014) generically simulates a Brownian particle in a
potential $U({\bm{\theta}})$ by numerically integrating the Langevin equation

$\begin{array}[]{ll}d{\bm{\theta}}&={\mathbf{v}}\,dt,\\\
d{\mathbf{v}}&=-\nabla
U({\bm{\theta}})\,dt-\gamma{\mathbf{v}}\,dt+\sqrt{2\gamma}\,dW,\end{array}$

where $\gamma$ is the friction coefficient and $W$ is the standard Wiener
process. See Durrett (1996) (chapter 5) for an introduction to stochastic
differential equations. It can be shown through the use of a Fokker–Planck
equation Gardiner (2004) that the above dynamics converge to the stationary
distribution
$p({\bm{\theta}},{\mathbf{v}})\propto\exp(-U({\bm{\theta}})-\frac{1}{2}{\mathbf{v}}^{2})$.
We can use this to sample from the full data posterior by using a potential
energy equal to the negative log joint,
$U({\bm{\theta}}):=-\log{p}({\mathcal{D}},{\bm{\theta}})$. The numeric
integration is typically discretized Chen et al. (2014); Zhang et al. (2019)
as

$\begin{array}[]{ll}\Delta{\bm{\theta}}&={\mathbf{v}},\\\
\Delta{\mathbf{v}}&=-\eta\nabla\hat{U}({\bm{\theta}})-\alpha{\mathbf{v}}+\bm{\epsilon}\sqrt{2(\alpha-\hat{\beta})\eta},\end{array}$
(3)

where $\eta$ is called the learning rate, $1-\alpha$ is the momentum decay,
$\bm{\epsilon}$ is a standard Gaussian random vector, $\hat{\beta}$ is an
optional parameter to offset the variance of the stochastic gradient term, and
$\hat{U}({\bm{\theta}})$ is an unbiased estimate of $U({\bm{\theta}})$
calculated on independently sampled mini-batches. Since $U({\bm{\theta}})$
grows with the size of the data set, a small learning rate
$\eta\sim\mathcal{O}\left(\frac{1}{|{\mathcal{D}}|}\right)$ is required to
minimize the discretization error. The variance of the stochastic gradient
term is proportional to $\eta^{2}$, while the variance of the injected noise
is proportional to $\eta$, so in the limit $\eta\to 0$, stochasticity in the
gradient estimates becomes negligible and the correct continuum dynamics are
recovered, even if one ignores the errors from the stochastic gradient noise
by using $\hat{\beta}=0$. We refer the reader to Ma et al. (2015);
Springenberg et al. (2016); Chen et al. (2014) for an in-depth analysis of the
algorithm parameters.

For the purposes of Bayesian updating and annealing, we use SGHMC with a
potential energy that leaves the distribution $f_{n}^{(t)}$ in Equation (2)
invariant, $U({\bm{\theta}})=-\log f_{n}^{(t)}({\bm{\theta}})$. We can
approximate this potential energy stochastically by

$\hat{U}_{n}^{(\lambda)}({\bm{\theta}})=-\lambda\log{p}({\mathbf{y}}_{n}|{\bm{\theta}})-\frac{n-1}{|B|}\sum_{{\mathbf{y}}\in
B}\log{p}({\mathbf{y}}|{\bm{\theta}})-\log{p({\bm{\theta}})},$ (4)

where the mini-batch is drawn independently and identically distributed
(i.i.d.) with replacement from the set of all previous data points, i.e.,
$B\subset\\{{\mathbf{y}}_{k}|k<n\\}$. Since
$\hat{U}_{n}^{(\lambda)}({\bm{\theta}})$ is an unbiased estimator of $-\log
f_{n}^{(t)}({\bm{\theta}})$ with finite variance (assuming the model assigns a
non-zero probability to all observations), SGHMC with this potential energy
will leave $f_{n}^{(t)}$ invariant in the limit $\eta\to 0$.

##  5 Online Marginal Likelihood Estimation

Our proposed approach is particularly efficient for ML estimation in
applications in which new data are continually becoming available, since the
marginal cost to update the ML estimate based on new data is independent of
the amount of previously processed data. Some examples of such applications
are: finance modeling, monitoring audio or video or other sensor readings,
applications in security, process control, etc. In these types of
applications, vanilla NS or AIS would have to recalculate the ML from scratch,
without leveraging the previous ML estimates.

One challenge with estimating ML online using SGAIS is the generation of mini-
batches that are uniformly sampled from historical data. This challenge can be
tackled with the help of reservoir sampling. If a large enough reservoir is
kept to be representative of the previously processed data, mini-batches could
be drawn uniformly from the reservoir instead Vitter (1985). If the data set
is too large to fit on disk or is only available during streaming, then this
approach makes it possible to estimate ML, whereas using NS or AIS to estimate
ML would not even be possible.

In online scenarios, the data may also exhibit various degrees of non-
stationarity. This may happen due to a shift in the underlying data-generating
distribution. Such distribution shifts would typically cause a noticeable
change in the ML, which can then be used for change-point detection.

##  6 Experiments

###  6.1 Experimental Setup

To evaluate the accuracy and run-time performance of our proposed approach, we
estimate the ML for three simple models on simulated data sets in an online
fashion. We further evaluate the robustness of SGAIS under various choices of
algorithm parameter values.

Default Parameters. We use mini-batches of size 500, with the following SGHMC
parameters: $\eta=0.1/N$, $\alpha=0.2$, and $\hat{\beta}=0$. Predictive
distributions are approximated using $M=10$ particles and 20 burn-in steps for
each intermediate distribution. We use a target $\operatorname*{ESS}$ of 5 for
adaptive annealing. As mentioned in Section 2, rather than Bayesian updating
by adding a single observation at a time, we add chunks of data at a time that
are the same size as the mini-batches. These parameter choices are not
necessarily the optimal choice for the models below. Instead, we chose the
SGHMC parameters based approximately on those suggested in Chen et al. (2014),
and we chose the number of particles and target $\operatorname*{ESS}$ to
hopefully be sufficient for adaptive annealing but small enough to result in a
short running time.

NS and AIS. As our reference standards of accuracy, we implemented NS and AIS.
Both NS and AIS implementations use SGHMC as their MCMC kernel. For NS, this
is to sample from the prior; for AIS, this is to sample from the intermediate
annealed distributions. For AIS, we use the same parameters as for our
sequential sampler, except that each MCMC step uses the whole data set instead
of mini-batches. We implement NS with 20 SGHMC steps to sample from the
constrained prior. NS still requires the full data set to check the
constraints and so cannot take advantage of mini-batching. For SGHMC used with
NS, we used parameters $\eta=10^{-3}$, $\alpha=0.1$, and $\hat{\beta}=0$
because there is no gradient noise when sampling from the prior. Results
reported are for two particles; more behave similarly but are slower. We allow
NS to run until the additive terms are less than 1% of the current
$\hat{\mathcal{Z}}$ estimate. This is a popular stopping criterion and is also
used in Grosse et al. (2015). More information about our implementations is
given in Appendix A.

Sensitivity Analysis. We investigate the sensitivity of our approach to the
following parameters: number of particles $M$, the target
$\operatorname*{ESS}$, the number of burn-in steps for each intermediate
distribution, the learning rate $\eta$, and the mini-batch size. Each test is
done by varying one parameter while keeping the others fixed at their default
values.

Run-time Environment. Our experiments were executed on a laptop with an Intel
i7 CPU and 8GB of ram, running Linux. For fair comparison, all code was
single-threaded. Multithreading gives a considerable speedup when calculating
the likelihood on large data sets but can introduce subtle complexities that
are difficult to control for in tests of run-time performance.

Models. We evaluated our approach on 3 models:

  * •

Bayesian linear regression with 6 parameters.

  * •

Bayesian logistic regression with 10-dimensional observations and four
classes. This model has 44 parameters.

  * •

Bayesian Gaussian mixture model for two-dimensional data with 5 mixture
components. This model has 25 parameters.

For each of the models, we generated data sets by sampling i.i.d. from the
model’s conditional distribution with the parameters fixed. See Appendix B for
a more detailed description of the models.

###  6.2 Distribution Shift

As described in Section 5, changes in the data-generating distribution should
be detectable in the ML estimates. To investigate this, we generate simulated
data with varying numbers of clusters. Histograms of the simulated data can be
seen in Figure 1. The first 1000 observations were generated from 3 Gaussian
distributions, the next 9000 observations were generated from 5 Gaussian
distributions, including the 3 used to generate the first observations, and
the remaining 90,000 observations were generated from 7 Gaussian
distributions, including the previous 5. Some of the clusters overlap, so it
is not immediately obvious from the histograms how many clusters there
actually are.

We evaluate the effect of this type of distribution shift on SGAIS by
estimating the ML online for Gaussian mixture models with 3, 5, and 7 mixture
components. We then shuffle the data to enforce stationarity and estimate the
ML for these three models on the shuffled data. If the final ML estimates for
the in-order and shuffled data differ significantly then this may indicate
that the particles are getting trapped in local modes before the change-points
occur.

a

b

c

d

Figure 1:  Histograms of non-stationary simulated data. (a) shows the first
1000 observations, (b) shows the next 9000 observations, (c) shows the last
$90,000$ observations, and (d) shows the total data set. In each of the three
time phases, the data-generating distribution produces data with more clusters
than before.

##  7 Results and Discussion

The log-ML typically grows linearly in the number of data points. For this
reason, it is natural to measure errors in $\log{\mathcal{Z}}/N$ rather than
$\log{\mathcal{Z}}$. For each model, we measured the runtime performance of
NS, AIS, and SGAIS for various data set sizes up to one million observations.
Each of these data sets is taken as the first $N$ observations of the largest
data set. This allows SGAIS to leverage the previous ML estimates for online
ML estimation. Due to computational constraints, we run NS and AIS only for
data set sizes at logarithmically increasing intervals, while SGAIS naturally
produces many more intermediate results in a single run. Since we found AIS to
be much slower than NS, we only run AIS on data sets small enough to finish
within 4000 s. We consider errors that are comparable to the discrepancy
between NS and AIS acceptable.

###  7.1 Accuracy and Speed

####  7.1.1 Linear Regression

For the linear regression model, the exact ML is available analytically and is
shown in Figure 2a for comparison. Each algorithm is able to produce accurate
results for this model for all data set sizes. The final error of SGAIS on one
million data points was only about 0.1%. For this model, our method achieved a
speedup over NS by about a factor of 3.3, and a speedup over AIS by a factor
of 24.9 on one million observations.

a

b

Figure 2: Linear regression model. (a) shows the accuracy of our marginal
likelihood (ML) estimator compared to nested sampling (NS), annealed
importance sampling (AIS), and the exact ML. (b) shows the run-time of each
method. ns is nested sampling, ais is annealed importance sampling, and sgais
is our stochastic gradient annealed importance sampling approach.

####  7.1.2 Logistic Regression

Figure 3a,b shows the log-ML estimates and run-time of each algorithm for the
logistic regression model. For the largest data set, NS and SGAIS produced
estimates that differed by roughly 0.6%, which is negligible. SGAIS was a
factor 10.4 faster than the nested sampler on one million observations for
this model.

a

b

Figure 3: Logistic regression model. (a) shows our ML estimator compared to NS
and AIS. (b) shows the run-time of each method. AIS was not run for larger
data set sizes because each subsequent run would take more than 4000 s.

####  7.1.3 Gaussian Mixture Model

Figure 4a,b shows the log-ML estimates and run-time of each algorithm for the
Gaussian mixture model. The posterior distribution for this model is
multimodal. Some modes are due to permutation symmetries; these modes do not
have to be explored since each one contains the same information. There are
also some local modes that do not necessarily capture meaningful information
about the data; for example, fitting a single Gaussian to the whole data set
may be a poor local optimum of the likelihood function. If an MCMC walker
finds one of these modes, it can get trapped. However, we find that by
Bayesian updating and annealing, the MCMC walkers tend to leave the poor local
modes early on, before they become extremely peaked. The estimates produced by
NS and SGAIS differed on the largest data set by roughly 0.1%. For this model,
SGAIS was about a factor of 4.9 faster than the nested sampler for one million
observations.

a

b

Figure 4: Gaussian mixture model. (a) shows our ML estimator compared to NS
and AIS. (b) shows the run-time of each method.

In all the above experiments, SGAIS converges to the same result as NS with a
negligible error for large $N$.

###  7.2 Distribution Shift

The log-ML estimates shown in Figure 5a display sharp changes at 1000 and
10,000 observations for the non-shuffled data. We are very clearly able to
identify the position of the change-points just by looking at the resulting
plot, without a priori assuming the existence or number of change-points. The
numbers of annealing steps shown in Figure 5b exhibit spikes at the change-
points and remain high once more clusters are added to the data than the model
can describe.

The agreement of the final ML estimates between the shuffled and non-shuffled
data suggests that these estimates can be trusted. The difference between the
online and shuffled estimates is small enough to be able to distinguish
between the three models. The 5- and 7-component models seem to describe the
total data set better than the 3-component model, but the 5 and 7 models have
similar values for their log-ML, presumably due to the overlapping clusters in
the data set.

a

b

Figure 5:  ML estimation under distribution shift. (a) shows the ML estimates
for Gaussian mixture models with different numbers of mixture components. The
solid lines are for the in-order data and the dashed lines are for the
shuffled and therefore stationary data. (b) shows the number of annealing
steps for the online ML estimates for the in-order data.

###  7.3 Sensitivity to Algorithm Parameters

The following results agreed closely with our expectations:

  * •

Increasing the number of particles, $M$, results in higher accuracy and a
longer running time without much effect on the number of annealing steps.

  * •

A smaller number of burn-in SGHMC steps per intermediate distribution
typically resulted in lower accuracy and a shorter run-time. A smaller number
of burn-in steps also resulted in more annealing steps due to the slower
equilibration.

  * •

Larger mini-batch sizes typically result in higher accuracy but more
computation per SGHMC step. Larger mini-batch sizes result in fewer Bayesian
updating steps but require more annealing steps per new chunk of data. Mini-
batch size would typically be chosen based on the hardware capabilities of the
platform and the type of data under consideration.

We report on our more interesting findings below; plots of the results for
each parameter investigated are given in Appendix C.

####  7.3.1 Target $\operatorname*{ESS}$

As expected, a higher target $\operatorname*{ESS}$ tends to result in more
annealing steps—see Figure 6a. Most of this work is done in the early stages
of Bayesian updating. Note that since we used 10 particles, a target
$\operatorname*{ESS}$ of $0.1M=1$ requires no annealing steps because
$\operatorname*{ESS}$ is bounded below by 1. No annealing results in high
variance during the early stages of Bayesian updating, and adaptively
annealing helps to reduce that variance, with only a small impact on the run-
time. This illustrates the importance of the adaptive annealing schedule in
our approach. Figure 6a,b indicates that our approach converges to the log-ML
within acceptable accuracy within a reasonable time for a target
$\operatorname*{ESS}$ larger than 1. Even a small target $\operatorname*{ESS}$
was good enough to match vanilla AIS, on average.

a

b

c

Figure 6: Sensitivity to the target effective sample size
($\operatorname*{ESS}$). (a) shows the log-ML estimates, (b) shows the run-
time, and (c) shows the number of annealing steps for each chunk of data
against the data set size until that chunk.

####  7.3.2 Learning Rate

Interestingly, smaller values of the learning rate tend to result in less
accurate log-ML estimates over a longer time—see Figure 7a,b. We suspect this
to be because a smaller learning rate does not allow the particle to move as
far each step, resulting in a slower equilibration and requiring more
annealing steps per observation. This effect can be seen in Figure 7c; the
smaller learning rates appear to result in a larger number of annealing steps
per observation. To further verify this, we investigated the interaction
between the learning rate and the number of burn-in steps.

a

b

c

Figure 7: Sensitivity to the learning rate. Reported learning rates are per-
observation learning rates, that is, $\eta=\mathrm{lr}/N$. (a) shows the log-
ML estimates, (b) shows the run-time, and (c) shows the number of annealing
steps. For a learning rate of 0.01, the run-time shown in (b) displays a
change in gradient near $10^{5}$ observations. This is the result of a reduced
number of annealing steps but is not visible in (c) since we only show the
number of annealing steps for up to $5\times 10^{4}$ observations.

####  7.3.3 Learning Rate and Burn-in

We investigate the interaction between the number of SGHMC steps taken per
intermediate distribution and the learning rate by varying the learning rate,
while keeping the product of the learning rate and the number of SGHMC steps
constant. Fewer burn-in steps (larger learning rate) tends to make the
algorithm faster, but a smaller learning rate results in higher accuracy in
the log-ML estimates, as seen in Figure 8a. The decrease in accuracy with a
larger learning rate is presumably due to the discretization error in Equation
(3). This is supported by Figure 8c: a larger learning rate requires a larger
number of annealing steps to reach the target $\operatorname*{ESS}$. Figure 8a
indicates that a per-observation learning rate of 1.0 can be used and still
result in estimates of acceptable accuracy on data sets of one million
observations. For a learning rate of 1.0, SGAIS achieved a speedup over NS by
a factor of 83.

a

b

c

Figure 8: Sensitivity to the learning rate while keeping the product of the
learning rate and the number of stochastic gradient Hamiltonian Monte Carlo
(SGHMC) steps constant. Reported learning rates are per-observation learning
rates, that is, $\eta=\mathrm{lr}/N$. (a) shows the log-ML estimates, (b)
shows the run time, and (c) shows the number of annealing steps.

##  8 Materials and Methods

Code for this work was implemented using pytorch Paszke et al. (2017). The
code for our experiments is available at
https://gitlab.com/pleased/sequential-evidence.

##  9 Conclusions

This paper introduced SGAIS, a novel algorithm for efficient large-scale ML
estimation, by combining the essential ingredients of Bayesian updating,
thermal annealing and data sub-sampling.We found that our SGAIS implementation
was able to produce accurate ML estimates with a speedup over NS and AIS on
simple models. Furthermore, since the marginal cost of updating the ML
estimates when new data arrives does not depend on the number of previous data
points, this approach may be effective for weighted model averaging in a
setting where one periodically gets access to new data, such as in streaming
applications.

We evaluated the sensitivity of our approach to the parameters and found that
the log-ML estimates were robust to many different parameter choices while
maintaining a speedup over NS and AIS for large data sets.

Potential future work may include further exploring the effects of stochastic
gradients for ML calculation in the full SMC setting, including latent
variable models, and with more general stochastic gradient MCMC algorithms
such as Riemannian manifold SGHMC Ma et al. (2015).

Lastly, we would like to point out that SGAIS naturally inherits some of the
pitfalls of AIS and SGHMC such as sensitivity to phase changes, local modes,
and the curse of dimensionality. Bayesian updating and adaptive scheduling
does help to reduce some of the challenges that come with annealing, just as
annealing helps to reduce problems arising from multi-modality, but to our
knowledge, there is no way to completely remove these difficulties in general.

Conceptualization, methodology, software, formal analysis, investigation,
visualization, writing—original draft preparation: S.A.C.; validation,
writing—review and editing: S.A.C., H.C.E., and S.K.; resources: S.K.; project
administration, supervision: H.C.E. and S.K.; funding acquisition: H.C.E.

S. Cameron received bursary support from the South African National Institute
of Theoretical Physics.

###### Acknowledgements.

We wish to thank the South African National Institute of Theoretical Physics
for financial support and the organizers of MaxEnt 2019 for sponsoring the
publication fee for this paper. The authors declare no conflict of interest.
The following abbreviations are used in this manuscript:
AIS | Annealed importance sampling
---|---
ESS | Effective sample size
i.i.d | Independently and identically distributed
MCMC | Markov chain Monte Carlo
ML | Marginal likelihood
NS | Nested sampling
SGAIS | Stochastic gradient annealed importance sampling
SGHMC | Stochastic gradient Hamiltonian Monte Carlo
SMC | Sequential Monte Carlo
yes

##  Appendix A Correctness of Nested Sampler and Annealed Importance Sampler

###  A.1 Nested Sampler Implementation

Nested sampling requires one to sample from the prior under increasing
likelihood constraints. Sampling under constraints is, in general, a difficult
problem. This sampling can be accomplished by repeatedly sampling from the
prior until the constraint is satisfied. This method of rejection sampling
scales extremely poorly with both the size of the data set and the dimension
of the parameter space and is not feasible on any but the smallest problems.
Instead, the constrained sampling is typically implemented by duplicating a
live particle, which is already within the constrained likelihood contour, and
applying an MCMC kernel to the particle that leaves the prior invariant. While
moving the particle with MCMC, care must be taken to avoid leaving the
constrained region.

Our implementation of NS is based on SGHMC, simply because the code for SGHMC
was already written. In order to sample under constraints, we use a similar
strategy to Galilean Monte Carlo as described in Skilling (2012), where the
particle reflects off of the boundary of the likelihood contour by updating
the momentum as follows:

$\Delta{\mathbf{v}}=-2({\mathbf{v}}\cdot{\mathbf{n}}){\mathbf{n}}.$

Here, ${\mathbf{n}}$ is a unit vector parallel to the likelihood gradient
$\nabla_{{\bm{\theta}}}p({\mathcal{D}}|{\bm{\theta}})$. While it is not the
focus of our paper, we note that we have not previously encountered the idea
of applying SGHMC to sampling under constraints; our implementation allows a
fairly direct approach to implementing NS in a variety of contexts, and the
technique may also potentially be of value in other constrained sampling
contexts.

We found that due to discretization error, the constrained sampler tends to
undersample slightly at the boundaries of the constrained region; however, the
undersampled volume is of the order of the learning rate $\eta$ and so can be
neglected for a sufficiently small learning rate.

###  A.2 Test Results

If either of our NS and AIS implementations were incorrect, we expect that
their ML estimates would disagree. If both implementations were incorrect, the
probability that their errors would cancel is close to zero since NS and AIS
work in completely different ways. Figure 9 shows the result of 5 independent
runs of NS and AIS for data sets of varying sizes from 20 to 1000, in
intervals of 20. The shaded regions are the minimum and maximum log-ML
estimates for each run; these regions are the typical values produced by NS
and AIS on these data sets. Since the typical results of NS and AIS tend to
agree, it is safe to assume that they have been implemented correctly.

Figure 9: Comparison of NS and AIS with SGHMC kernels for a 1-dimensional
Gaussian mixture model with 9 parameters (8 degrees of freedom). The shaded
regions are the minimum and maximum of 5 independent runs of each algorithm.

##  Appendix B Model Details

###  B.1 Linear Regression

The data set consists of pairs $({\mathbf{x}},y)$ related by

$y={\mathbf{w}}^{T}{\mathbf{x}}+b+\epsilon,$

where $\epsilon$ c distributed with known variance $\sigma^{2}$. We do not
assume any distribution over ${\mathbf{x}}$ as it always appears on the right-
hand side of the conditional. The single-observation likelihood is

$p(y|{\mathbf{x}},{\mathbf{w}},b)=\frac{1}{\sqrt{2\pi\sigma^{2}}}\exp\left(-\frac{(y-{\mathbf{w}}^{T}{\mathbf{x}}-b)^{2}}{2\sigma^{2}}\right).$

Parameters are ${\mathbf{w}}$ and $b$, with standard Gaussian priors. ML can
be calculated analytically for this model. We used 5 dimensional vectors
${\mathbf{x}}$; this model therefore has 6 parameters.

###  B.2 Logistic Regression

The data set consists of pairs $({\mathbf{x}},y)$, where ${\mathbf{x}}$ is an
observation vector that is assigned a class label $y\in\\{1,\ldots,K\\}$. The
labels have a discrete distribution with probabilities given by the softmax
function of an affine transform of the observations

$p(y|{\mathbf{x}},{\bm{\theta}})=\frac{\exp({\mathbf{w}}_{y}^{T}{\mathbf{x}}+b_{y})}{\sum_{k}\exp({\mathbf{w}}_{k}^{T}{\mathbf{x}}+b_{k})}.$

Parameters are ${\bm{\theta}}=({\mathbf{w}}_{1:K},b_{1:K})$, with standard
Gaussian priors. Again, we do not assume any distribution over ${\mathbf{x}}$
as it always appears on the right-hand side of the conditional. We used 10
dimensional vectors ${\mathbf{x}}$ with 4 classes; this model has 44
parameters.

###  B.3 Gaussian Mixture Model

The data are modeled by a mixture of $d$-dimensional multivariate Gaussian
distributions with diagonal covariance matrices. Mixture weights, means, and
variances are treated as parameters. This type of model is often treated as a
latent variable model, where the mixture component assignments of each data
point are the latent variables. Here, we marginalize out the latent variables
to obtain the following conditional distribution:

$p({\mathbf{y}}|{\bm{\theta}})=\sum_{k=1}^{K}\beta_{k}\prod_{j=1}^{d}\frac{1}{\sqrt{2\pi\sigma_{k,j}^{2}}}\exp\left(-\frac{(y_{j}-\mu_{k,j})^{2}}{2\sigma_{k,j}^{2}}\right),$
${\bm{\theta}}=(\beta_{1:K},\mu_{1:K,1:d},\sigma_{1:K,1:d}^{2}).$

Mixture weights $\beta_{1:K}$ are modeled by a Dirichlet prior with
$\bm{\alpha}=1$; means $\mu_{k,j}$ are modeled conditionally given the
variances by Gaussian priors, centered around zero with variance
$4\sigma_{k,j}^{2}$; variances $\sigma_{k,j}^{2}$ are modeled by inverse gamma
priors with shape and scale parameters equal to 1. We used 5 Gaussian
components, and observations were 2-dimensional; this model has 25 parameters
with 24 degrees of freedom. We use this model for our sensitivity tests, since
it has the most complex structure.

##  Appendix C Supplementary Figures For Parameter Sensitivity Tests

a Dependence on $M$

b Dependence on $\operatorname*{ESS}$

c Dependence on number of SGHMC steps

d Dependence on learning rate

e Dependence on learning rate with the product of learning rate and SGHMC
steps constant

f Dependence on mini-batch size.

Figure 10: Log-ML sensitivity.

a Dependence on $M$

b Dependence on $\operatorname*{ESS}$

c Dependence on number of SGHMC steps

d Dependence on learning rate

e Dependence on learning rate with the product of learning rate and SGHMC
steps constant

f Dependence on mini-batch size

Figure 11: Running time sensitivity.

a Dependence on $M$

b Dependence on $\operatorname*{ESS}$

c Dependence on number of SGHMC steps

d Dependence on learning rate

e Dependence on learning rate with the product of learning rate and SGHMC
steps constant

f Dependence on mini-batch size

Figure 12: Sensitivity of number of annealing steps.

References

## References

  * Barber (2012) Barber, D.  Bayesian Reasoning and Machine Learning; Cambridge University Press: New York, NY, USA, 2012.
  * Skilling (2006) Skilling, J.  Nested Sampling for General Bayesian Computation.  Bayesian Anal. 2006, 1, 833–859,  doi:black10.1214/06-BA127.
  * Neal (1998) Neal, R.M.  Annealed Importance Sampling.  arXiv 1998, arXiv:physics/9803008.
  * Welling and Teh (2011) Welling, M.; Teh, Y.W.  Bayesian Learning via Stochastic Gradient Langevin Dynamics.  In Proceedings of the 28th International Conference on Machine Learning, ICML 2011, Bellevue, WA, USA, 28 June–2 July 2011; Getoor, L., Scheffer, T., Eds.; Omnipress: Madison, WI, USA, 2011; pp. 681–688.
  * Naesseth et al. (2019) Naesseth, C.A.; Lindsten, F.; Schön, T.B.  Elements of Sequential Monte Carlo.  arXiv 2019, arXiv:1903.04797.
  * Gordon et al. (1993) Gordon, N.J.; Salmond, D.J.; Smith, A.F.M.  Novel Approach to Nonlinear/Non-Gaussian Bayesian State Estimation.  IEE Proc. F - Radar Signal Process. 1993, 140, 107–113,  doi:black10.1049/ip-f-2.1993.0015.
  * Wallach et al. (2009) Wallach, H.M.; Murray, I.; Salakhutdinov, R.; Mimno, D.  Evaluation Methods for Topic Models.  In Proceedings of the 26th Annual International Conference on Machine Learning, Montreal, QC, Canada, 14–18 June 2009; ACM: New York, NY, USA, 2009; pp. 1105–1112,  doi:black10.1145/1553374.1553515.
  * Cameron et al. (2019) Cameron, S.A.; Eggers, H.C.; Kroon, S.  A Sequential Marginal Likelihood Approximation Using Stochastic Graidients.  In Proceedings of the 39th International Workshop on Bayesian Inference and Maximum Entropy Methods in Science and Engineering, Garching/Munich, Germany, 30 June–5 July 2019.
  * Chen et al. (2014) Chen, T.; Fox, E.; Guestrin, C.  Stochastic Gradient Hamiltonian Monte Carlo.  In Proceedings of the 31st International Conference on Machine Learning, Beijing, China, 21–26 June 2014; Volume 5.
  * Grosse et al. (2015) Grosse, R.B.; Ghahramani, Z.; Adams, R.P.  Sandwiching the marginal likelihood using bidirectional Monte Carlo.  arXiv 2015, arXiv:1511.02543.
  * Kong et al. (1994) Kong, A.; Liu, J.S.; Wong, W.H.  Sequential Imputations and Bayesian Missing Data Problems.  J. Am. Stat. Assoc. 1994, 89, 278–288.
  * Buchholz et al. (2018) Buchholz, A.; Chopin, N.; Jacob, P.E.  Adaptive Tuning Of Hamiltonian Monte Carlo Within Sequential Monte Carlo.  arXiv 2018, arXiv:1808.07730.
  * Beskos et al. (2016) Beskos, A.; Jasra, A.; Kantas, N.; Thiery, A.  On the Convergence of Adaptive Sequential Monte Carlo Methods.  Ann. Appl. Probab. 2016, 26, 1111–1146,  doi:black10.1214/15-AAP1113.
  * Durrett (1996) Durrett, R.  Stochastic Calculus: A Practical Introduction; Probability and Stochastics Series; CRC Press: Boca Raton, FL, USA, 1996; pp. 177–207
  * Gardiner (2004) Gardiner, C.W.  Handbook of Stochastic Methods for Physics, Chemistry and the Natural Sciences, 3rd ed.; Springer Series in Synergetics; Springer: Berlin, Germany, 2004; Volume 13, pp. xviii+415.
  * Zhang et al. (2019) Zhang, R.; Li, C.; Zhang, J.; Chen, C.; Wilson, A.G.  Cyclical Stochastic Gradient MCMC for Bayesian Deep Learning.  arXiv 2019, arXiv:1902.03932.
  * Ma et al. (2015) Ma, Y.A.; Chen, T.; Fox, E.B.  A Complete Recipe for Stochastic Gradient MCMC.  arXiv 2015, arXiv:1506.04696.
  * Springenberg et al. (2016) Springenberg, J.T.; Klein, A.; Falkner, S.; Hutter, F.  Bayesian Optimization with Robust Bayesian Neural Networks. In Advances in Neural Information Processing Systems 29; Lee, D.D., Sugiyama, M., Luxburg, U.V., Guyon, I., Garnett, R., Eds.; Curran Associates, Inc.: Barcelona, Spain 2016; pp. 4134–4142.
  * Vitter (1985) Vitter, J.S.  Random Sampling with a Reservoir.  ACM Trans. Math. Softw. 1985, 11, 37–57,  doi:black10.1145/3147.3165.
  * Paszke et al. (2017) Paszke, A.; Gross, S.; Chintala, S.; Chanan, G.; Yang, E.; DeVito, Z.; Lin, Z.; Desmaison, A.; Antiga, L.; Lerer, A.  Automatic Differentiation in PyTorch. Available online: https://openreview.net/pdJsrmfCZ (accessed on 24 June 2019).
  * Skilling (2012) Skilling, J.  Bayesian Computation in Big Spaces—Nested Sampling and Galilean Monte Carlo.  AIP Conf. Proc. 2012, 1443, 145–156,  doi:black10.1063/1.3703630.


#  DL4DS \- Deep Learning for empirical DownScaling

Carlos Alberto Gomez Gonzalez<EMAIL_ADDRESS>Earth Sciences department,
Barcelona Supercomputing Center, Barcelona, Catalonia, Spain

(28 January 2022)

###### Abstract

A common task in Earth Sciences is to infer climate information at local and
regional scales from global climate models. Dynamical downscaling requires
running expensive numerical models at high resolution which can be prohibitive
due to long model runtimes. On the other hand, statistical downscaling
techniques present an alternative approach for learning links between the
large- and local-scale climate in a more efficient way. A large number of deep
neural network-based approaches for statistical downscaling have been proposed
in recent years, mostly based on convolutional architectures developed for
computer vision and super-resolution tasks. This paper presents DL4DS, Deep
Learning for empirical DownScaling, a python library that implements a wide
variety of state-of-the-art and novel algorithms for downscaling gridded Earth
Science data with deep neural networks. DL4DS has been designed with the goal
of providing a general framework for training convolutional neural networks
with configurable architectures and learning strategies to facilitate the
conduction of comparative and ablation studies in a robust way. We showcase
the capabilities of DL4DS on air quality CAMS data over the western
Mediterranean area. The DL4DS library can be found in this repository:
https://github.com/carlos-gg/dl4ds

###### keywords:

deep learning, downscaling, super-resolution, post-processing

††journal: Environmental Data Science

2022

Gomez Gonzalez

This paper presents DL4DS the first open-source library with state-of-the-art
and novel deep learning algorithms for empirical downscaling.

##  1 Introduction

Downscaling aims to bridge the gap between the large spatial scales
represented by global climate models (GCMs) to the smaller scales required for
assessing regional climate change and its impacts (Maraun and Widmann,, 2017).
This task can be approached via dynamical or statistical downscaling. In the
former, a high-resolution regional climate model is nested into a GCM over the
domain of interest (Rummukainen,, 2010), while the latter aims to learn
empirical links between the large- and local-scale climate that are in turn
applied to low-resolution climate model output. Statistical or empirical
downscaling comes with the benefit of a lower computational cost compared to
its dynamical counterpart.

Machine Learning (ML) can be defined as the study of computer algorithms that
improve automatically by finding statistical structure (building a model) from
training data for automating a given task. The field of ML started to flourish
in the 1990s and has quickly become the most popular and most successful
subfield of Artificial Intelligence (AI) thanks to the availability of faster
hardware and larger training datasets. Deep Learning (DL, Bengio et al.,,
2021) is a specific subfield of ML aimed at learning representations from data
putting emphasis on learning successive layers of increasingly meaningful
representations (Chollet,, 2021). DL has shown potential in a wide variety of
problems in Earth sciences dealing with high-dimensional and complex data and
offers exciting new opportunities for expanding our knowledge about the Earth
system (Reichstein et al.,, 2019; Huntingford et al.,, 2019; Dewitte et al.,,
2021; Irrgang et al.,, 2021).

The layered representations in DL are learned via models called neural
networks, where an input n-dimensional tensor is received and multiplied with
a weights tensor to produce an output. A bias term is usually added to this
output before passing the result through an element-wise nonlinear function
(also called activation). Convolutional Neural Networks (CNNs, LeCun et al.,,
1989) are a type of neural network that rely on the convolution operation, a
linear transformation that takes advantage of the implicit structure of
gridded data. The convolution operation uses a weight tensor that operates in
a sliding window fashion on the data, meaning that only a few input grid
points contribute to a given output and that the weights are reused since they
are applied to multiple locations in the input. CNNs have been used almost
universally for the past few years in computer vision applications, such as
object detection, semantic segmentation and super-resolution. CNNs show
performances on par with more recent models such as visual Transformers (Liu
et al.,, 2022) while being more efficient in terms of memory and computations.

In spite of recent efforts in the computer science community, the field of DL
still lacks solid mathematical foundations. It remains fundamentally an
engineering discipline heavily reliant on experimentation, empirical findings
and software developments. While scientific software tools such as Numpy
(Harris et al.,, 2020), Xarray (Hamman and Hoyer,, 2017) or Jupyter (Kluyver
et al.,, 2016) have an essential role in modern Earth sciences research
workflows, state-of-the-art domain-specific DL-based algorithms are usually
developed as proof-of-concept scripts. For DL to fulfill its potential to
advance Earth sciences, the development of AI- and DL-powered scientific
software must be carried out in a collaborative and robust way following open-
source and modern software development principles.

##  2 CNN-based super-resolution for statistical downscaling

Statistical downscaling of gridded climate variables is a task closely related
to that of super-resolution in computer vision, considering that both aim to
learn a mapping between low-resolution and high-resolution grids (Wang et
al.,, 2021). Unsurprisingly, several DL-based approaches have been proposed
for statistical or empirical downscaling of climate data in recent years
(Vandal et al.,, 2017; Leinonen et al.,, 2020; Stengel et al.,, 2020; Höhlein
et al.,, 2020; Liu et al.,, 2020; Harilal et al.,, 2021). Most of these
methods have in common the use of convolutions for the exploitation of
multivariate spatial or spatio-temporal gridded data, that is 3D
(height/latitude, width/longitude, channel/variable) or 4D (time,
height/latitude, width/longitude, channel/variable) tensors. Therefore, CNNs
are efficient at leveraging the high-resolution signal from heterogeneous
observational datasets, from now on called predictors or auxiliary variables,
while downscaling low-resolution climatological fields.

In our search for efficient architectures for empirical downscaling, we have
developed DL4DS, a library that draws from recent developments in the field of
computer vision for tasks such as image-to-image translation and super-
resolution. DL4DS is implemented in Tensorflow/Keras, a popular DL framework,
and contains a collection of building blocks that abstract and modularize a
few key design strategies for composing and training empirical downscaling DL
models. These strategies are discussed in the following section.

##  3 DL4DS core design principles and building blocks

The core design principles and general architecture of DL4DS are discussed in
the following subsections. An overall architecture of DL4DS is shown in Fig.
1, while some of the building blocks and network architectures are shown in
Figs. 2 and 3.

###  3.1 Type of statistical downscaling

Statistical downscaling methods aim to derive empirical relationships between
an observed high-resolution variable or predictand and low-resolution
predictor variables (Maraun and Widmann,, 2017). Two main types of statistical
downscaling can be defined depending on the origin of the predictors used for
training; Model output statistics (MOS), where the predictors are taken
directly from global climate model outputs, and Perfect Prognosis
(PerfectProg) methods, that relies on observational datasets for both
predictand and predictors (see Fig. 1). Most of the DL-based downscaling
methods proposed to date, work in PerfectProg setup, where a high-resolution
observational dataset is used to create paired training samples via a
downsampling or coarsening operation (Vandal et al.,, 2017; Leinonen et al.,,
2020; Stengel et al.,, 2020). The trained model is then applied to the
desired, unseen during training, low-resolution data in a domain-transfer
fashion. Other approaches such as those proposed by Höhlein et al., (2020) or
Harilal et al., (2021), model a cross-scale transfer function between explicit
low-resolution and high-resolution datasets. DL4DS supports both training
frameworks, with explicit paired samples (in MOS fashion) or with paired
samples simulated from high-resolution fields (in PerfectProg fashion).

Figure 1: The general architecture of DL4DS. A low-resolution gridded dataset
can be downscaled, with the help of auxiliary predictor and static variables,
and a high-resolution reference dataset. The mapping between the low- and
high-resolution data is learned with either a supervised or a conditional
generative adversarial DL model

###  3.2 Multivariate modeling and auxiliary variables

Using high-resolution static predictors fields, such as, topography, distance
from the sea, and land-sea mask, as well as high- or intermediate-resolution
essential climate variables, can help improve the inference capacity of
empirical downscaling methods (Maraun and Widmann,, 2017; Vandal et al.,,
2017; Höhlein et al.,, 2020). The most straightforward way to include these
additional predictors and static variables is to merge (concatenate) them
across the channel dimension with the input low-resolution fields. DL4DS can
handle an arbitrary number of time-varying high- or intermediate-resolution
predictors and high-resolution static variables. Additionally, static high-
resolution fields are passed through a convolutional block in the output
module as shown in panel (E) of Fig. 3 in order to emphasize high-resolution
topographic information.

###  3.3 Pre-processing and normalization

DL4DS implements a couple of pre-processing methods aimed at normalizing or
standardizing multi-dimensional gridded datasets, a required step when
splitting the data into train and validation/test sets used for training and
testing DL models. These procedures, namely mix-max and standardization
scalers, are implemented as classes (dl4ds.MinMaxScaler and
dl4ds.StandardScaler) in the preprocessing module of DL4DS. These classes
extend the functionality of methods from the Scikit-learn package (Pedregosa
et al.,, 2011) aimed at scaling tabular or 2D data. DL4DS scalers can handle
3D and 4D datasets with NaN values in either numpy.ndarray or xarray.Dataarray
formats.

Figure 2: Panel (A) shows the main blocks and layers implemented in DL4DS.
Panel (B) shows the structure of the main spatial convolutional block, a
succession of two convolutional layers with interleaved regularization
operations, such as dropout or normalization. Blocks and operations shown with
dashed lines are optional

###  3.4 Architecture of the super-resolution models in DL4DS

The super-resolution networks implemented in DL4DS aim at learning a mapping
from low or coarse-resolution to high-resolution grids. The super-resolution
networks in DL4DS are composed of two main parts, a backbone section and an
output module, as depicted in Fig. 1.

####  3.4.1 Backbone section

We adopt several strategies for the design of the backbone sections, as shown
in panels (A) through (D) of Fig. 3. The main difference being the particular
arrangement of convolutional layers. In DL4DS, the backbone type can be set
with the backbone parameter of the training classes mentioned subsection 3.5:

  * •

Convnet backbone - Composed of standard convolutional blocks. As shown in
panel (B) of Fig. 2, each spatial convolutional block is composed of two
convolutional layers (using 3x3 kernels) with some interleaved optional
operations, such as dropout (Srivastava et al.,, 2014), batch/layer
normalization or a channel attention block (discussed in subsection 3.4.1).

  * •

Resnet backbone - Composed of residual blocks or convolutional blocks with
skip connections111Skip or shortcut connections are connections that feed the
output of a particular layer to later non-sequential or non-adjacent layers in
the network. Element-wise addition is used in residual blocks, while
concatenation is used in dense blocks or in decoder blocks of a U-Net. that
learn residual functions with reference to the layer inputs. Residual skip
connections have been widely employed in super-resolution models (Wang et
al.,, 2021) in the past. In DL4DS, we implemented cross-layer skip connections
within each residual block and outer skip connections at the backbone section
level (He et al.,, 2016).

  * •

Densenet backbone - Composed of dense blocks (see panel (B) of Fig. 2). As in
the case of the Resnet backbone, we implement cross-layer skip connections
within each dense block and outer skip connections at the backbone section
level (Huang et al.,, 2017).

  * •

Unet backbone - Inspired by the structure of the U-Net DL model (Ronneberger
et al.,, 2015), this backbone is composed of encoder and decoder blocks with
skip connections, as depicted in panel (C) of Fig. 3. The encoder block
consists of a convolutional block plus downsampling via max pooling222Max
pooling is a sample-based discretization process that selects the maximum
element from the region of the feature map covered by the filter., while the
decoder block consists of upsampling, concatenation with skip feature
channels333A feature channel or feature map refers to the output of a
convolutional layer. from the encoder section, followed by a convolutional
block. This backbone is only implemented for pre-upsampled spatial samples.

  * •

Convnext backbone - Based on the recently proposed ConvNext model (Liu et
al.,, 2022), we implemented a backbone with depth-wise convolutional layers
and larger kernel sizes (7x7). To the best of our knowledge, this is the first
time that a ConvNext-style block architecture has been applied to the task of
statistical downscaling.

Figure 3: DL4DS supervised DL models, as well as generators, are composed of a
backbone section (examples in panels (A), (B), (C) and (D)) and an output
module (panel (E). Panel (A) shows the backbone of models for downscaling pre-
upsampled spatial samples using either residual or dense blocks. Panel (B)
presents the backbone of a model for downscaling spatial samples using
ConvNext-like blocks and one of the post-upsampling blocks described in
subsection 3.4.1. Panel (C) shows the backbone of a model for downscaling pre-
upsampled spatial samples using an encoder-decoder structure. Panel (D) shows
the backbone of a model for downscaling spatio-temporal samples using
recurrent-convolutional blocks and a post-upsampling block. These backbones
are followed by the output module (see subsection 3.4.2) shown in panel (E).
The color legend for the blocks used here is shown in panel (A) of Fig. 2

##### Spatial and spatio-temporal modelling

While CNNs excel at modelling spatial or gridded data, hybrid convolutional
and recurrent architectures are designed to exploit spatio-temporal sequences.
Convolutional Long Short-Term Memory (ConvLSTM, Shi et al.,, 2015) and
convolutional Gated-Recurrent-Units (Ballas et al.,, 2015) are examples of
models that have been used for downscaling time-evolving gridded data
(Leinonen et al.,, 2020; Harilal et al.,, 2021). DL4DS can model either
spatial or spatio-temporal data, by either using standard convolutional blocks
or recurrent convolutional blocks. The spatio-temporal network in panel (D) of
Fig. 3 contains recurrent convolutional blocks to handle 4D training samples
(with time dimension). The structure of these recurrent convolutional blocks
is similar to that of the main convolutional block shown in panel (B) of Fig.
2 but using convolutional LSTM layers instead of standard convolutional ones.

##### Channel attention

In DL4DS, we implement a channel attention mechanism based on those of the
Squeeze-and-Excitation networks (Hu et al.,, 2018) and the Convolutional Block
Attention Module (Woo et al.,, 2018). This attention mechanism exploits the
inter-channel relationship of features by providing a weight for each channel
in order to enhance those that contribute the most to the optimization and
learning process. First, it aggregates spatial information of a feature map by
using average pooling444Average pooling is a sample-based discretization
process that computes the average value from the region of the feature map
covered by the filter.. The resulting vector is then passed through two 1x1
convolutional blocks and a sigmoid activation to create the channel weights
(attention maps) which are multiplied element-wise to each corresponding
feature map. The channel attention mechanism is integrated as an optional step
in DL4DS convolutional blocks to get refined feature maps throughout the
network.

##### Upsampling techniques

The upsampling operation is a crucial step for DL-based empirical downscaling
methods. In computer vision, upsampling refers to increasing the number of
rows and columns, and therefore the number of pixels, of an image. Increasing
the size of an image is also often called upscaling, which oddly enough
carries the opposite meaning in the weather and climate jargon. When it comes
to downscaling or super-resolving gridded data, it is of utmost importance to
preserve fine-scale information, or to transfer it from the high-resolution to
the low-resolution fields, while increasing the number or grid points in both
horizontal and vertical directions. In DL4DS, we implement several well-
established upsampling methods proposed in the super-resolution literature
(Wang et al.,, 2021). These methods belong to two main upsampling strategies,
depending on whether the upsampling happens before the network or inside of
it: pre-upsampling and post-upsampling. In the former case, the models are fed
with a pre-upsampled (via interpolation) low-resolution input and a reference
high-resolution predictand. In the latter case, the low-resolution input is
directly fed together with the high-resolution predictand. In this scenario,
we learn a mapping from low-resolution to high-resolution fields in low-
dimensional space with end-to-end learnable layers integrated at the end of
the backbone section (see panels (B) and (D) of Fig. 3). This results in an
increased number of learnable parameters, but an increased computational
efficiency thanks to working on smaller grids. When working with DL4DS, the
upsampling method can be set with the upsampling parameter of the training
classes discussed in the subsection 3.5:

  * •

PIN pre-upsampling - Models trained with interpolation-based pre-upsampled
input.

  * •

RC post-upsampling - Models including a resize convolution block (bilinear
interpolation followed by a convolutional layer).

  * •

DC post-upsampling - Models including a deconvolution (also called transposed
convolution) block (Dong et al.,, 2016). With a transposed convolution we
perform a transformation opposite to that of a normal convolution. In
practice, the image is expanded by inserting zeros and performing a
convolution. This upsampling method can sometimes produce checkerboard-like
patterns that hurt the quality of the downscaled or super-resolved products.

  * •

SPC post-upsampling - Models including a subpixel convolution block (Lim et
al.,, 2017). A subpixel convolution consists of a regular convolution with
$s^{2}$ filters (where $s$ is the scaling factor) followed by an image
reshaping operation called a phase shift. Instead of putting zeros in between
grid points as it is done in a deconvolution layer, we calculate more
convolutional filters in lower resolution and reshape them into an upsampled
grid.

In the context of empirical downscaling, pre-upsampling has been used by
Vandal et al., (2017); Harilal et al., (2021), resize convolution by Leinonen
et al., (2020), transposed convolution by Höhlein et al., (2020) and sub-pixel
convolution by Liu et al., (2020). Choosing an upsampling method depends on
the problem at hand (MOS or PerfectProg training) and the computational
resources available.

####  3.4.2 Output module

The second part of every DL4DS network is the output module, as shown in panel
(E) of Fig. 3. Here we concatenate the intermediate outputs of three branches:
the backbone section feature channels, the backbone section output passed
through a localized convolutional block, and a separate convolutional block
over the input high-resolution static variables. These concatenated feature
channels are in turn passed through a transition block, a 1x1 convolution used
to control the number of channels, and two final convolutional blocks. The
first of this two convolutional blocks in the output module applies a channel
attention block by default.

##### Localized convolutional block

CNNs bring many positive qualities for the exploitation of spatial or gridded
data. However, the property of translation invariance of CNNs is usually sub-
optimal in geospatial applications where location-specific information and
dynamics need to be taken into account. This has been addressed by Uselis et
al., (2020) for the task of weather forecasting. Within DL4DS, we implement a
Localized Convolutional Block (LCB) located in the output module of our
networks. This LCB consists of a bottleneck transition layer, via 1x1
convolutions to compresse the number of incoming feature channels, followed by
a locally connected layer with biases (similar to a convolutional one, except
that the weights here are not shared). To the best of our knowledge, this is
the first time the idea of LCBs has been applied to the task of statistical
downscaling.

###  3.5 Loss functions

Loss functions are used to measure reconstruction error while guiding the
optimization of neural networks during the training phase.

####  3.5.1 Supervised losses

DL4DS implements pixel-wise losses, such as the Mean Squared Error (MSE or L2
loss) or the Mean Absolute Error (MAE or L1 loss), as well as structural
dissimilarity (DSSIM) and multi-scale structural dissimilarity losses (MS-
DSSIM), which are derived from the Structural Similarity Index Measure (SSIM,
Wang et al.,, 2004). The SSIM is a computer vision metric used for measuring
the similarity between two images, based on three comparison measurements:
luminance, contrast and structure. The MS-DSSIM loss has been used by
Chaudhuri and Robertson, (2020) for the task of statistical downscaling.

####  3.5.2 Adversarial losses

Generative adversarial networks (GANs, Goodfellow et al.,, 2014) are a class
of ML models in which two neural networks contest with each other in a zero-
sum game, where one agent’s gain is another agent’s loss. This adversarial
training was extended for its use with images for image-to-image translation
problems by Isola et al., (2017). Conditional GANs map input to output images
while learning a loss function to train this mapping, usually allowing a
better reconstruction of high-frequency fine details. DL4DS implements a
conditional adversarial loss (Isola et al.,, 2017) by training the two
networks depicted in Fig. 4. The role of the generator is to learn to generate
high-resolution fields from their low-resolution counterparts while the
discriminator learns to distinguish synthetic high-resolution fields from
reference high-resolution ones. Through iterative adversarial training, the
resulting generator can produce outputs consistent with the distribution of
real data, while the discriminator cannot distinguish between the generated
high-resolution data and the ground truth. Adversarial training has been used
in DL-based downscaling approaches proposed by Leinonen et al., (2020),
Stengel et al., (2020) and Chaudhuri and Robertson, (2020). In out tests, the
generated high-resolution fields produced by the CGAN generators exhibit
moderate variance, even when using the Monte Carlo dropout technique (Gal and
Ghahramani,, 2016) which amounts to applying dropout at inference time.

Figure 4: Example of a conditional generative adversarial model for spatio-
temporal samples in post-upsampling mode (see subsection 3.4.1). Two networks,
the generator shown in panel (A), and discriminator shown in panel (B), are
trained together optimizing an adversarial loss (see subsection 3.5). The
color legend for the blocks used here is shown in panel (A) of Fig. 2

When working with DL4DS, the training strategy is controlled by choosing one
of the training classes: dl4ds.SupervisedTrainer or dl4ds.CGANTrainer, for
using supervised or adversarial losses accordingly.

##  4 Experimental results

Downscaling applications are highly dependant on the target variable/dataset
at hand, and therefore the optimal model architecture should be found in a
case-by-case basis. In this section, we give a taste of the capabilities of
DL4DS without conducting a rigorous comparison of model architectures and
learning configurations. For our tests, we use data from the Copernicus
Atmosphere Monitoring Service (CAMS) reanalysis (Inness et al.,, 2019), the
latest global reanalysis dataset of atmospheric composition produced by the
European Centre for Medium-Range Weather Forecasts (ECMWF), which provides
estimates of Nitrogen dioxide (NO2) surface concentration. We select NO2 data
from the period between 2014 and 2018 at a 3-hourly temporal resolution which
results in $\sim$14.6k temporal samples. For our low-resolution dataset, we
use the CAMS global reanalysis (CAMSRA) with an horizontal resolution of
$\sim$80 km. Our high-resolution target is the CAMS regional reanalysis
produced for the European domain with a spatial resolution of $\sim$10km
(0.1º). We also include predictor atmospheric variables from the ECMWF ERA5
reanalysis (Hersbach et al.,, 2020), namely 2 meter temperature and 10 meter
wind speed, with an intermediate resolution of $\sim$25 km (0.25º). Finally,
we include as static variables: topography from the Global Land One-km Base
Elevation Project and a land-ocean mask derived from it, and a layer with the
urban area fraction derived from the land cover dataset produced by the
European Space Agency. In order to make our test less heavy on memory
requirements, we spatially subset the data and focus on the western
Mediterranean region, as shown in Fig. 5.

Figure 5: A reference NO2 surface concentration field from the low-resolution CAMS global reanalysis is shown in panel (A). In panel (B), we present a resampled version, via bicubic interpolation, of the low-resolution reference field. This interpolated field looks overly smoothed and showcases the inefficiency of simple resampling methods at restoring fine scale information. Panel (C): corresponding high-resolution field from the CAMS regional reanalysis. Both low- and high-resolution grids were taken from the holdout set for the same time step. The maximum value shown corresponds to the maximum value in the high-resolution grid Table 1: DL4DS models showcased in section 4. The column named ‘panel’ points the corresponding plot in Fig. 6, Fig. 7 and Fig. 8 . Panel | Downscaling | Learning | Sample type | Backbone | Upsampling | LCB | loss
---|---|---|---|---|---|---|---
A | PerfectProg | supervised | spatial | unet | PIN | no | mae
B | MOS | supervised | spatial | unet | PIN | no | mae
C | PerfectProg | supervised | spatial | resnet | SPC | no | dssim+mae
D | MOS | supervised | spatial | resnet | SPC | yes | dssim+mae
E | MOS | supervised | spatial | resnet | SPC | yes | mae
F | MOS | adversarial | spatial | resnet | SPC | yes | mae
G | MOS | supervised | spatiotemp | resnet | SPC | yes | dssim+mae
H | MOS | supervised | spatial | convnext | SPC | yes | mae

We showcase eight models trained with DL4DS, without the intention of a full
exploration of possible architectures and learning strategies. Different loss
functions, backbones, learning strategies and other parameters are combined in
the model architectures detailed in Table 1. The training is carried out using
a single cluster node with four NVIDIA V100 GPUs. The data is split into train
and test sets by reserving the last year of data for test and validation. We
fit the dl4ds.StandardScaler on the training set (getting the global mean and
standard deviation) and applied it to both training and test sets (subtracting
the global mean and dividing by the global standard deviation). All models are
trained with the Adam optimizer (Kingma and Ba,, 2015), for 100 epochs in the
case of supervised CNNs and 18 epochs in the case of the conditional
adversarial models.

In Fig. 6, we show examples of predicted high-resolution fields produced by
the models detailed in Table 1 and corresponding to the reference fields of
Fig. 5. These give a visual impression of the reconstruction ability of each
model. In Fig. 7 and Fig. 8, we show the Pearson correlation and Root Mean
Square Error (RMSE) for each one of the aforementioned models. The metrics in
these maps are computed for each grid point independently and for the whole
year of 2018 (containing 2920 holdout grids). In Table 2, we give a more
complete list of metrics (including the MAE, SSIM and peak signal-to-noise
ratio (PSNR)) computed for each grid pair separately (time step) and then
summarized by taking the mean and standard deviation over the samples of the
holdout year.

Figure 6: Examples of downscaled products obtained with DL4DS, corresponding to the reference grid shown in panel (A) of Fig. 5. The models corresponding to each panel are detailed in Table 1 Figure 7: Pixel-wise Pearson correlation for each model, computed for the whole year of 2018 Figure 8: Pixel-wise RMSE for each model, computed for the whole year of 2018. The dynamic range is shared for all the panels, with a fixed maximum value to facilitate the visual comparison Table 2: Metrics computed for each time step, downscaled product with respect to the reference grid, of the holdout year for the models showcased this section . Panel | MAE | RMSE | PearCorr | SSIM | PSNR
---|---|---|---|---|---
A | 2.58 $\pm$ 0.92 | 4.37 $\pm$ 1.50 | 0.64 $\pm$ 0.11 | 0.81 $\pm$ 0.05 | 32.95 $\pm$ 2.97
B | 1.56 $\pm$ 0.54 | 2.70 $\pm$ 0.87 | 0.85 $\pm$ 0.04 | 0.89 $\pm$ 0.03 | 34.61 $\pm$ 2.78
C | 2.58 $\pm$ 0.93 | 4.40 $\pm$ 1.47 | 0.64 $\pm$ 0.11 | 0.88 $\pm$ 0.04 | 38.85 $\pm$ 3.00
D | 1.60 $\pm$ 0.60 | 4.75 $\pm$ 5.03 | 0.76 $\pm$ 0.20 | 0.99 $\pm$ 0.01 | 64.87 $\pm$ 6.28
E | 1.49 $\pm$ 0.51 | 2.56 $\pm$ 0.84 | 0.88 $\pm$ 0.03 | 0.90 $\pm$ 0.03 | 35.23 $\pm$ 2.83
F | 1.70 $\pm$ 0.58 | 2.87 $\pm$ 0.90 | 0.84 $\pm$ 0.04 | 0.87 $\pm$ 0.03 | 34.60 $\pm$ 2.81
G | 1.53 $\pm$ 0.56 | 2.68 $\pm$ 0.96 | 0.86 $\pm$ 0.04 | 0.89 $\pm$ 0.03 | 34.97 $\pm$ 3.03
H | 1.51 $\pm$ 0.55 | 2.64 $\pm$ 0.90 | 0.87 $\pm$ 0.03 | 0.90 $\pm$ 0.03 | 35.03 $\pm$ 2.97

For this particular task and datasets, and according to the metrics of Table
2, we find that models trained in MOS fashion perform better than those in
PerfectProg and that a supervised model with residual blocks and subpixel
convolution upsampling provides the best results. Also, we note that models
trained with a LCB perform better that those without it, thanks to the fact
that the LCB learns grid point- or location-specific weights. As expected,
when the same network is optimized with respect to a DSSIM+MAE loss, it
reaches higher scores in the SSIM and PSNR metrics. Training the CGAN model
for more epochs could improve the results but would require longer runtimes
and computational resources.

##  5 Summary and future directions

In this paper, we have presented DL4DS, a Python library implementing state-
of-the-art and novel DL algorithms for empirical downscaling of gridded data.
It is built on top of Tensorflow/Keras (Chollet et al.,, 2015), a modern and
versatile DL framework. Among the key strengths of DL4DS are its simple API
consisting of a few user-facing classes controlling the training procedure,
its support for distributed GPU training (via data parallelism) powered by
Horovod (Sergeev and Balso,, 2018), the generous number of customizable DL-
based architectures included, the possibility of composing new architectures
based on the general purpose building blocks offered, the option of saving and
retraining models, and its support of different downscaling frameworks that
learn either with explicit low-resolution and high-resolution data (MOS-like)
or only with a high-resolution dataset (PerfectProg-like). DL4DS is a powerful
tool for reproducible empirical downscaling experiments and for network
architecture optimization. A thorough ablation study varying the depth and
width of the each backbone, in combination with other design choices, such as
the upsampling method, and tested over several benchmark datasets, shall allow
a complete comparison of CNN-based DL models for empirical downscaling.

DL has shown promise in post-proceesing and bias correction tasks (Rasp and
Lerch,, 2018; Grönquist et al.,, 2020), so a natural research direction is to
explore the application of DL4DS to related post-processing problems. Other
interesting research directions are the implementation of uncertainty
estimation techniques besides Monte Carlo dropout, the extension of DL4DS
models to the case of non-gridded reference data (observational station data),
the implementation of alternative adversarial losses such as the Earth mover’s
distance (Wasserstein GAN), and the inclusion of alternative DL algorithms
such as normalizing flows (Groenke et al.,, 2020).

##### Acknowledgments

We wish to acknowledge F. J. Doblas-Reyes, K. Serradell, H. Petetin, Ll. Lledó
and Ll. Palma for helpful discussions and comments, and P-A. Bretonnière and
M. Samsó for assistance with data management.

##### Funding Statement

C.A.G.G. acknowledges funding from the European Union’s Horizon 2020 research
and innovation programme under the Marie Skłodowska-Curie grant agreement
H2020-MSCA-COFUND-2016-754433..

##### Code Availability

DL4DS code repository: https://github.com/carlos-gg/dl4ds.

## References

  * Ballas et al., (2015) Ballas, N., Yao, L., Pal, C., and Courville, A. (2015).  Delving deeper into convolutional networks for learning video representations.  arXiv preprint arXiv: Arxiv-1511.06432.
  * Bengio et al., (2021) Bengio, Y., Lecun, Y., and Hinton, G. (2021).  Deep learning for ai.  Commun. ACM, 64(7):58–65.
  * Chaudhuri and Robertson, (2020) Chaudhuri, C. and Robertson, C. (2020).  Cligan: A structurally sensitive convolutional neural network model for statistical downscaling of precipitation from multi-model ensembles.  Water, 12(12).
  * Chollet, (2021) Chollet, F. (2021).  Deep Learning with Python.  Manning.
  * Chollet et al., (2015) Chollet, F. et al. (2015).  Keras.  https://keras.io.
  * Dewitte et al., (2021) Dewitte, S., Cornelis, J. P., Müller, R., and Munteanu, A. (2021).  Artificial intelligence revolutionises weather forecast, climate monitoring and decadal prediction.  Remote Sensing, 13(16).
  * Dong et al., (2016) Dong, C., Loy, C. C., and Tang, X. (2016).  Accelerating the super-resolution convolutional neural network.  In Leibe, B., Matas, J., Sebe, N., and Welling, M., editors, Computer Vision – ECCV 2016, pages 391–407, Cham. Springer International Publishing.
  * Gal and Ghahramani, (2016) Gal, Y. and Ghahramani, Z. (2016).  Dropout as a bayesian approximation: Representing model uncertainty in deep learning.  In Proceedings of the 33rd International Conference on International Conference on Machine Learning - Volume 48, ICML’16, page 1050–1059. JMLR.org.
  * Goodfellow et al., (2014) Goodfellow, I., Pouget-Abadie, J., Mirza, M., Xu, B., Warde-Farley, D., Ozair, S., Courville, A., and Bengio, Y. (2014).  Generative adversarial nets.  In Ghahramani, Z., Welling, M., Cortes, C., Lawrence, N., and Weinberger, K. Q., editors, Advances in Neural Information Processing Systems, volume 27. Curran Associates, Inc.
  * Groenke et al., (2020) Groenke, B., Madaus, L., and Monteleoni, C. (2020).  Climalign: Unsupervised statistical downscaling of climate variables via normalizing flows.  arXiv preprint arXiv: Arxiv-2008.04679.
  * Grönquist et al., (2020) Grönquist, P., Yao, C., Ben-Nun, T., Dryden, N., Dueben, P., Li, S., and Hoefler, T. (2020).  Deep learning for post-processing ensemble weather forecasts.
  * Hamman and Hoyer, (2017) Hamman, J. and Hoyer, S. (2017).  Xarray: N-d labeled arrays and datasets in python.  Journal of Open Research Software.
  * Harilal et al., (2021) Harilal, N., Singh, M., and Bhatia, U. (2021).  Augmented convolutional lstms for generation of high-resolution climate change projections.  IEEE Access, 9:25208–25218.
  * Harris et al., (2020) Harris, C. R., Millman, K. J., van der Walt, S. J., Gommers, R., Virtanen, P., Cournapeau, D., Wieser, E., Taylor, J., Berg, S., Smith, N. J., Kern, R., Picus, M., Hoyer, S., van Kerkwijk, M. H., Brett, M., Haldane, A., del Río, J. F., Wiebe, M., Peterson, P., Gérard-Marchant, P., Sheppard, K., Reddy, T., Weckesser, W., Abbasi, H., Gohlke, C., and Oliphant, T. E. (2020).  Array programming with NumPy.  Nature, 585(7825):357–362.
  * He et al., (2016) He, K., Zhang, X., Ren, S., and Sun, J. (2016).  Deep residual learning for image recognition.  In 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pages 770–778.
  * Hersbach et al., (2020) Hersbach, H., Bell, B., Berrisford, P., Hirahara, S., Horányi, A., Muñoz-Sabater, J., Nicolas, J., Peubey, C., Radu, R., Schepers, D., Simmons, A., Soci, C., Abdalla, S., Abellan, X., Balsamo, G., Bechtold, P., Biavati, G., Bidlot, J., Bonavita, M., De Chiara, G., Dahlgren, P., Dee, D., Diamantakis, M., Dragani, R., Flemming, J., Forbes, R., Fuentes, M., Geer, A., Haimberger, L., Healy, S., Hogan, R. J., Hólm, E., Janisková, M., Keeley, S., Laloyaux, P., Lopez, P., Lupu, C., Radnoti, G., de Rosnay, P., Rozum, I., Vamborg, F., Villaume, S., and Thépaut, J.-N. (2020).  The era5 global reanalysis.  Quarterly Journal of the Royal Meteorological Society, 146(730):1999–2049.
  * Hu et al., (2018) Hu, J., Shen, L., and Sun, G. (2018).  Squeeze-and-excitation networks.  In 2018 IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 7132–7141.
  * Huang et al., (2017) Huang, G., Liu, Z., Maaten, L. V. D., and Weinberger, K. Q. (2017).  Densely connected convolutional networks.  In 2017 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pages 2261–2269, Los Alamitos, CA, USA. IEEE Computer Society.
  * Huntingford et al., (2019) Huntingford, C., Jeffers, E. S., Bonsall, M. B., Christensen, H. M., Lees, T., and Yang, H. (2019).  Machine learning and artificial intelligence to aid climate change research and preparedness.  Environmental Research Letters, 14(12):124007.
  * Höhlein et al., (2020) Höhlein, K., Kern, M., Hewson, T., and Westermann, R. (2020).  A comparative study of convolutional neural network models for wind field downscaling.  Meteorological Applications, 27(6):e1961.
  * Inness et al., (2019) Inness, A., Ades, M., Agustí-Panareda, A., Barré, J., Benedictow, A., Blechschmidt, A.-M., Dominguez, J. J., Engelen, R., Eskes, H., Flemming, J., Huijnen, V., Jones, L., Kipling, Z., Massart, S., Parrington, M., Peuch, V.-H., Razinger, M., Remy, S., Schulz, M., and Suttie, M. (2019).  The cams reanalysis of atmospheric composition.  Atmospheric Chemistry and Physics, 19(6):3515–3556.
  * Irrgang et al., (2021) Irrgang, C., Boers, N., Sonnewald, M., Barnes, E. A., Kadow, C., Staneva, J., and Saynisch-Wagner, J. (2021).  Towards neural Earth system modelling by integrating artificial intelligence in Earth system science.  Nature Machine Intelligence, 3(8):667–674.
  * Isola et al., (2017) Isola, P., Zhu, J.-Y., Zhou, T., and Efros, A. A. (2017).  Image-to-image translation with conditional adversarial networks.  CVPR.
  * Kingma and Ba, (2015) Kingma, D. P. and Ba, J. (2015).  Adam: A method for stochastic optimization.  In Bengio, Y. and LeCun, Y., editors, 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings.
  * Kluyver et al., (2016) Kluyver, T., Ragan-Kelley, B., Pérez, F., Granger, B., Bussonnier, M., Frederic, J., Kelley, K., Hamrick, J., Grout, J., Corlay, S., Ivanov, P., Avila, D., Abdalla, S., and Willing, C. (2016).  Jupyter notebooks – a publishing format for reproducible computational workflows.  In Loizides, F. and Schmidt, B., editors, Positioning and Power in Academic Publishing: Players, Agents and Agendas, pages 87 – 90. IOS Press.
  * LeCun et al., (1989) LeCun, Y., Boser, B., Denker, J. S., Henderson, D., Howard, R. E., Hubbard, W., and Jackel, L. D. (1989).  Backpropagation applied to handwritten zip code recognition.  Neural Computation, 1(4):541–551.
  * Leinonen et al., (2020) Leinonen, J., Nerini, D., and Berne, A. (2020).  Stochastic super-resolution for downscaling time-evolving atmospheric fields with a generative adversarial network.  IEEE Transactions on Geoscience and Remote Sensing, pages 1–13.
  * Lim et al., (2017) Lim, B., Son, S., Kim, H., Nah, S., and Lee, K. M. (2017).  Enhanced deep residual networks for single image super-resolution.  2017 IEEE Conference on Computer Vision and Pattern Recognition Workshops (CVPRW).
  * Liu et al., (2020) Liu, Y., Ganguly, A. R., and Dy, J. (2020).  Climate Downscaling Using YNet: A Deep Convolutional Network with Skip Connections and Fusion, page 3145–3153.  Association for Computing Machinery, New York, NY, USA.
  * Liu et al., (2022) Liu, Z., Mao, H., Wu, C.-Y., Feichtenhofer, C., Darrell, T., and Xie, S. (2022).  A convnet for the 2020s.  arXiv preprint arXiv: Arxiv-2201.03545.
  * Maraun and Widmann, (2017) Maraun, D. and Widmann, M. (2017).  Statistical Downscaling and Bias Correction for Climate Research.  Cambridge University Press.
  * Pedregosa et al., (2011) Pedregosa, F., Varoquaux, G., Gramfort, A., Michel, V., Thirion, B., Grisel, O., Blondel, M., Prettenhofer, P., Weiss, R., Dubourg, V., Vanderplas, J., Passos, A., Cournapeau, D., Brucher, M., Perrot, M., and Duchesnay, E. (2011).  Scikit-learn: Machine learning in Python.  Journal of Machine Learning Research, 12:2825–2830.
  * Rasp and Lerch, (2018) Rasp, S. and Lerch, S. (2018).  Neural networks for postprocessing ensemble weather forecasts.  Monthly Weather Review, 146(11):3885 – 3900.
  * Reichstein et al., (2019) Reichstein, M., Camps-Valls, G., Stevens, B., Jung, M., Denzler, J., Carvalhais, N., and Prabhat (2019).  Deep learning and process understanding for data-driven Earth system science.  Nature, 566(7743):195–204.
  * Ronneberger et al., (2015) Ronneberger, O., Fischer, P., and Brox, T. (2015).  U-net: Convolutional networks for biomedical image segmentation.  In Navab, N., Hornegger, J., Wells, W. M., and Frangi, A. F., editors, Medical Image Computing and Computer-Assisted Intervention – MICCAI 2015, pages 234–241, Cham. Springer International Publishing.
  * Rummukainen, (2010) Rummukainen, M. (2010).  State-of-the-art with regional climate models.  WIREs Climate Change, 1(1):82–96.
  * Sergeev and Balso, (2018) Sergeev, A. and Balso, M. D. (2018).  Horovod: fast and easy distributed deep learning in tensorflow.  arXiv preprint arXiv: Arxiv-1802.05799.
  * Shi et al., (2015) Shi, X., Chen, Z., Wang, H., Yeung, D.-Y., Wong, W.-k., and Woo, W.-c. (2015).  Convolutional lstm network: A machine learning approach for precipitation nowcasting.  In Proceedings of the 28th International Conference on Neural Information Processing Systems - Volume 1, NIPS’15, page 802–810, Cambridge, MA, USA. MIT Press.
  * Srivastava et al., (2014) Srivastava, N., Hinton, G., Krizhevsky, A., Sutskever, I., and Salakhutdinov, R. (2014).  Dropout: A simple way to prevent neural networks from overfitting.  J. Mach. Learn. Res., 15(1):1929–1958.
  * Stengel et al., (2020) Stengel, K., Glaws, A., Hettinger, D., and King, R. N. (2020).  Adversarial super-resolution of climatological wind and solar data.  Proceedings of the National Academy of Sciences, 117(29):16805–16815.
  * Uselis et al., (2020) Uselis, A., Lukoševičius, M., and Stasytis, L. (2020).  Localized convolutional neural networks for geospatial wind forecasting.  Energies, 13(13).
  * Vandal et al., (2017) Vandal, T., Kodra, E., Ganguly, S., Michaelis, A., Nemani, R., and Ganguly, A. R. (2017).  Deepsd: Generating high resolution climate change projections through single image super-resolution.  In Proceedings of the 23rd ACM SIGKDD International Conference, page 1663–1672, New York, NY, USA. Association for Computing Machinery.
  * Wang et al., (2004) Wang, Z., Bovik, A., Sheikh, H., and Simoncelli, E. (2004).  Image quality assessment: from error visibility to structural similarity.  IEEE Transactions on Image Processing, 13(4):600–612.
  * Wang et al., (2021) Wang, Z., Chen, J., and Hoi, S. H. (2021).  Deep learning for image super-resolution: A survey.  IEEE Transactions on Pattern Analysis and Machine Intelligence, 43(10):3365–3387.
  * Woo et al., (2018) Woo, S., Park, J., Lee, J.-Y., and Kweon, I. S. (2018).  Cbam: Convolutional block attention module.  In Proceedings of the European Conference on Computer Vision (ECCV).