# Learning by example$\boldsymbol{:}$ fast reliability-aware seismic imaging
with normalizing flows

Ali Siahkoohi and Felix J. Herrmann
School of Computational Science and Engineering,
Georgia Institute of Technology
{alisk<EMAIL_ADDRESS>
###### Abstract

Uncertainty quantification provides quantitative measures on the reliability
of candidate solutions of ill-posed inverse problems. Due to their sequential
nature, Monte Carlo sampling methods require large numbers of sampling steps
for accurate Bayesian inference and are often computationally infeasible for
large-scale inverse problems, such as seismic imaging. Our main contribution
is a data-driven variational inference approach where we train a normalizing
flow (NF), a type of invertible neural net, capable of cheaply sampling the
posterior distribution given previously unseen seismic data from neighboring
surveys. To arrive at this result, we train the NF on pairs of low- and high-
fidelity migrated images. In our numerical example, we obtain high-fidelity
images from the Parihaka dataset and low-fidelity images are derived from
these images through the process of demigration, followed by adding noise and
migration. During inference, given shot records from a new neighboring seismic
survey, we first compute the reverse-time migration image. Next, by feeding
this low-fidelity migrated image to the NF we gain access to samples from the
posterior distribution virtually for free. We use these samples to compute a
high-fidelity image including a first assessment of the image’s reliability.
To our knowledge, this is the first attempt to train a conditional network on
what we know from neighboring images to improve the current image and assess
its reliability.

##  1 Introduction

Seismic imaging is challenged by noise and a non-trivial null-space due to
bandwidth and aperture limitations. These challenges lead to solution non-
uniqueness where different images may fit the data equally well. Solution non-
uniqueness can be characterized via a distribution over the solution space
that assigns probabilities to different estimates—i.e., the posterior
distribution [1]. Due to the high-dimensionality of seismic images, posterior
inference often requires high-dimensional sampling, for instance via Markov
chain Monte Carlo [MCMC, 2] techniques. MCMC methods sample the posterior
distribution via a series of random walks in the probability space where the
posterior probability density function (PDF) needs to be evaluated or
approximated at each step—e.g., via stochastic gradient Langevin dynamics [3].
These sampling methods typically require many steps to traverse the
probability space [4, 5, 6, 7, 8, 9, 10, 11], which fundamentally limits their
applicability to large-scale problems due to costs associated with the forward
operator [12, 13, 14, 15]. Alternatively, variational inference methods [16]
approximate the posterior distribution with a parametric and easy-to-sample
distribution. By virtue of this approximation, sampling is turned into an
optimization problem where the parameters are tuned to minimize the “distance”
between the parametric and posterior distributions. After optimization, this
easy-to-sample distribution is used to conduct Bayesian inference. As
discussed by Blei et al. [17], variational-inference based methods are known
to have computational advantages over MCMC in high-dimensional problems.

In this paper, we take a variational inference approach to seismic imaging and
use a normalizing flow [NF, 18] as a parametric function to approximate the
posterior distribution. NFs are a special type of invertible neural nets [19,
20, 21] capable of approximating complicated functions—e.g., non-Gaussian
PDFs. After training, NFs cheaply sample the distribution of interest and
estimate its PDF. The latter property sets NFs apart from generative
adversarial nets [GANs, 22] since GANs only provide samples from the
distribution of interest and are incapable of explicitly representing
densities by design [22]. Despite this disadvantage, GANs have been used
successfully to carry out stochastic waveform inversions with the Metropolis-
adjusted Langevin algorithm [23]. Contrary to GANs, NFs can be simply trained
via the maximum-likelihood estimation method [18], which is asymptotically
consistent and efficient [24]. On the other hand, training GANs typically
involves a min-max optimization that might lead to unstable training [25].
Finally, the invertibility of NFs allows for memory-efficient training since
the intermediate state variables do not need to be stored [26, 27, 28].

We use a data-driven approach for training the NF to take maximum advantage of
available data in the form of low- and high-fidelity migrated images. One way
of creating training data is to pair high-fidelity least-squares migration
images with low-fidelity reverse-time migration images from nearby surveys. In
this work, we mimic this choice by pairing high-fidelity images from the
Parihaka dataset [29] with low-fidelity images obtained via the process of
demigration, followed by adding noise and migration. This approach circumvents
the challenges of explicitly choosing a prior distribution that captures the
true heterogeneity exhibited by the Earth’s subsurface. Moreover, the NF
trained via this method is not specific to one seismic survey, but thanks to
the NF’s ability to generalize can be used to characterize the posterior
distribution associated with a previously unseen seismic survey, which is in
some sense close—e.g., data from a neighboring survey area. Finally, this
approach, unlike MCMC methods, does not involve repeated evaluations of the
forward operator, however, by providing the reverse-time migrated image to the
NF we implicitly capture the likelihood and prior in our learned posterior
distribution. To our knowledge, this is the first normalizing-flow based
scalable attempt to cheaply sample from the seismic imaging posterior
distribution given previously unseen seismic data.

In the context of variational inference for inverse problems with expensive
forward operators, Stein variational gradient descent [30] is used to
iteratively push samples from the prior to the posterior distribution to
tackle full-waveform inversion [31] and seismic tomography [32]. Rizzuti et
al. [33] and Zhao et al. [34] both independently proposed a physics-based
variational inference approach that approximates the posterior distribution
with a distribution represented by a NF. While this approach does not require
training data, it does need choosing a prior and repeated computationally
expensive evaluations of the forward operator and the adjoint of its Jacobian
(migration). However, there are indications that this formulation can be
faster and easier to scale to large-scale problems compared to MCMC methods
[34]. Motivated by Kruse et al. [21], Siahkoohi et al. [35] take this approach
a step further by proposing a two-stage multifidelity approach where during
the first stage a conditional NF is trained on available pairs of low- and
high-fidelity images. This pretrained NF is subsequently fine tuned during an
optional second physics-based variational inference stage, which is tailored
to the physics of the specific imaging problem at hand. In this work, we focus
on the first stage because it is computational cheap and provides a
rudimentary assessment of the image’s reliability. Moreover, since the first
stage entails training over low- and high-fidelity image pairs of related
imaging problems, it can serve as a warm start for the optional second more
accurate physics-based inference stage [36, 35], greatly reducing its costs.
In related data-driven GAN-based approaches, Adler and Öktem [37] and Kovachki
et al. [38] also train neural nets capable of sampling from the posterior
distribution given new data. Instead of adhering to GANs, we rely on NFs,
which are easier to train and have a much more limited memory footprint during
the training phase, making them applicable to larger-scale problems.

In the following section, we first briefly introduce variational inference and
mathematically state the foundations of NFs. Next, we describe our proposed
normalizing-flow based formulation for seismic imaging and Bayesian inference.
We conclude by showcasing our approach on a “quasi” real example that derives
from 2D subsets of the Parihaka dataset.

##  2 Variational inference

At its core, the method we are proposing characterizes multivariate
distributions given access to samples from these distributions only.
Variational inference provides the statistical framework for this and involves
the problem of approximating a distribution of interest $p$ with a parametric
distribution $q$, represented in our case by a NF. This approximation is
typically achieved by minimizing the Kullback-Leibler (KL) divergence between
$q$ and $p$. The KL divergence offers a distance metric for distributions and
can be explained as the cross-entropy of distribution $p$ relative to
distribution $q$ minus the entropy of distribution $q$. During variational
inference, the KL divergence is minimized with respect to distribution $q$. By
design, $q$ is typically chosen to be easy—i.e., computationally cheap, to
sample from. As a result, after optimization, we can cheaply sample from the
distribution $p$ by sampling $q$ instead. Next, we introduce NFs as a way to
parameterize distribution $q$.

###  2.1 Normalizing flows

A NF $T_{\mathrm{w}}:{X}\rightarrow{Z}$, parameterized by the NF weights
$\mathbf{w}$, is an invertible neural net—i.e., it has a closed-from and exact
inverse (up to numerical precision), with the goal of mapping inputs
$\mathbf{x}\in{X}$ sampled from the distribution of interest, $\mathbf{x}\sim
p(\mathbf{x})$, into samples from the Gaussian distribution $\mathbf{z}\sim
p_{\mathrm{z}}(\mathbf{z})$, where $\mathbf{z}\in{Z}$ and
$p_{\mathrm{z}}(\mathbf{z}):=\mathrm{N}(\mathbf{z}\mid\mathbf{0},\mathbf{I})$.
In other words, given samples $\\{\mathbf{x}^{(i)}\\}_{i=1}^{n}\sim
p(\mathbf{x})$, the parameters $\mathbf{w}$ are optimized so that the KL
divergence between $q_{\mathrm{w}}$ and $p_{z}$ is minimized, where
$q_{\mathrm{w}}$ is the distribution of the NF’s output when $\mathbf{x}\sim
p(\mathbf{x})$ are fed as input. This mapping is useful because after training
we can sample from $p(\mathbf{x})$ by simply inputting Gaussian samples
$\mathbf{z}\sim p_{\mathrm{z}}(\mathbf{z})$ into the inverted network,
$T_{\mathrm{w}}^{-1}$, provided that $q_{\mathrm{w}}$ approximates $p_{z}$
well [18]. Given samples $\\{\mathbf{x}^{(i)}\\}_{i=1}^{n}$ from the unknown
distribution, minimizing the KL divergence between $q_{\mathrm{w}}$ and
$p_{z}$ leads to the following optimization problem for training the NF
$T_{\mathrm{w}}$ [18]:

$\mathop{\rm
min}_{\mathbf{w}}\,\frac{1}{n}\sum_{i=1}^{n}\Big{[}\frac{1}{2}\big{\|}T_{\mathrm{w}}(\mathbf{x}^{(i)})\big{\|}^{2}_{2}-\log\Big{|}\det\nabla_{\mathrm{x}}T_{\mathrm{w}}(\mathbf{x}^{(i)})\Big{|}\Big{]}.$
(1)

In the above objective, the $\ell_{2}$-norm on $T_{w}$ imposes a Gaussian
distribution on the network output, which is what a NF is designed to do—i.e.
to Gaussianize the input, $\mathbf{x}^{(i)}$. The second term involving
$\det\nabla_{\mathrm{x}}T_{\mathrm{w}}(\mathbf{x})$ regularizes the entropy of
$q_{\mathrm{w}}$, which avoids $T_{w}(\mathbf{x})$ from producing a trivial
output—i.e., $T_{w}(\mathbf{x}):=\mathbf{0}$. The extra costs of computing
$\det\nabla_{\mathrm{x}}T_{\mathrm{w}}(\mathbf{x})$ and its gradient are
negligible due to the special design of invertible nets [19]. We solve the
optimization problem in Equation 1 via stochastic optimization where at each
iteration the objective is approximated via randomly selected (without
replacement) training samples $\mathbf{x}^{(i)}$. In addition to sampling from
$p(\mathbf{x})$ after training, we can also use the change-of-variables
formula [39] to cheaply estimate $p(\mathbf{x})$ for a given $\mathbf{x}$ by
$p(\mathbf{x})\approx
p_{\mathrm{z}}(T_{\mathrm{w}}(\mathbf{x}))\,\Big{|}\det\nabla_{\mathrm{x}}T_{\mathrm{w}}(\mathbf{x})\Big{|}$.
In cases where we only have access to samples from the prior distribution,
this relation provides an expression for the prior PDF that can be used in
inverse problems involving estimating $\mathbf{x}$.

###  2.2 Posterior inference with conditional NFs

Our ultimate goal is to capture the posterior distribution associated with
inverse problems involving computationally expensive and possibly nonlinear
forward operators, $F:X\rightarrow Y$, with the forward model,
$\mathbf{y}=F(\mathbf{x})+\mathbf{n}$. Here, $\mathbf{x}\in X$ represents the
unknown model, $\mathbf{y}\in Y$ the observed data, and $\mathbf{n}$ possibly
non-Gaussian noise. Following Kruse et al. [21], NFs can be adapted to sample
from the posterior distribution, $p(\mathbf{x}\mid\mathbf{y})$ given
previously unseen data $\mathbf{y}$. NFs capable of generating samples from
conditional distributions can be defined by imposing a block-triangular
structure on $T_{\mathrm{w}}:Y\times X\rightarrow Z\times Z$, which according
to Kruse et al. [21] leads to the following NF structure:

$\begin{split}T_{\mathrm{w}}(\mathbf{y},\mathbf{x})=\begin{bmatrix}T_{\mathrm{w}_{1}}(\mathbf{y})\\\\[4.0pt]
T_{\mathrm{w}_{2}}(\mathbf{y},\mathbf{x})\end{bmatrix},\
\mathbf{w}=\begin{bmatrix}\mathbf{w}_{1}\\\
\mathbf{w}_{2}\end{bmatrix}.\end{split}$ (2)

In words, Equation 2 describes $T_{\mathrm{w}}$ as a NF that maps the joint
random variable $(\mathbf{y},\mathbf{x})$ into two Gaussian random vectors,
where the first output is only a function of $\mathbf{y}$. As in Equation 1,
this conditional NF can be trained given training pairs,
$\\{\mathbf{y}^{(i)},\mathbf{x}^{(i)}\\}_{i=1}^{n}$, by minimizing the
following objective:

$\mathop{\rm
min}_{\mathbf{w}}\,\frac{1}{n}\sum_{i=1}^{n}\Big{[}\frac{1}{2}\big{\|}T_{\mathrm{w}}(\mathbf{y}^{(i)},\mathbf{x}^{(i)})\big{\|}^{2}_{2}-\log\Big{|}\det\nabla_{(\mathrm{y},\mathrm{x})}T_{\mathrm{w}}(\mathbf{y}^{(i)},\mathbf{x}^{(i)})\Big{|}\Big{]}.$
(3)

This block-diagonal construction lets us sample from the posterior
distribution via [40, 21]

$T_{\mathrm{w}_{2}}^{-1}\big{(}T_{\mathrm{w}_{1}}(\mathbf{y}),\mathbf{z}\big{)}\sim
p(\mathbf{x}\mid\mathbf{y}),\quad\mathbf{z}\sim p_{\mathrm{z}}(\mathbf{z}).$
(4)

According to this expression, samples from the posterior are drawn by first by
feeding the data, $\mathbf{y}$, to $T_{\mathrm{w}_{1}}$. Next, posterior
samples are obtained by feeding $T_{\mathrm{w}_{1}}(\mathbf{y})$ and Gaussian
samples $\mathbf{z}\sim p_{\mathrm{z}}(\mathbf{z})$ into the inverse network,
$T_{\mathrm{w}_{2}}^{-1}$. These Gaussian samples can be considered as
realizations of the posterior in the latent space. Compared to methods such as
MCMC, this sampling procedure does not involve the forward operator and is
therefore extremely fast and computationally cheap. Given samples generated
with Equation 4, we approximate the conditional mean and point-wise standard
deviation point-estimators via

$\boldsymbol{\mu}=\mathbb{E}\big{[}\mathbf{x}\mid\mathbf{y}\big{]},\quad\boldsymbol{\sigma}^{2}=\mathbb{E}\big{[}\left(\mathbf{x}-\boldsymbol{\mu}\right)^{\circ
2}\mid\mathbf{y}\big{]},$ (5)

where ∘2 represents elementwise power of two and the expectations can be
approximated with sample means.

###  2.3 Imaging with conditional NFs

Seismic imaging is the problem of estimating the short-wavelength squared-
slowness perturbation model of the Earth’s subsurface given noisy observed
seismic data and a kinematically accurate background velocity model. We cast
seismic imaging into the deep-learning based framework described in the
previous section to limit the computational burden of high-dimensional
posterior inference. To achieve this, we choose $\mathbf{x}$ to represent the
high-fidelity migrated image and $\mathbf{y}$ the corresponding low-fidelity
reverse-time migrated image obtained by the process of demigration, followed
by adding noise and migration. With this choice, the prior information in this
technique comes from different realizations of high-fidelity migrated images
provided in training pairs $(\mathbf{y}^{(i)},\mathbf{x}^{(i)}),\
i=1,\ldots,n$, obtained from neighboring surveys. We train the NF by solving
the optimization problem in Equation 3 while using the block-triangular
structure (Equation 2). This yields a conditional NF that provides cheap,
fast, and forward-operator free samples from the imaging posterior
distribution via Equation 4. We use these samples to obtain a high-fidelity
image including a first assessment of its uncertainty.

##  3 Numerical experiments

We propose a “quasi” real example in which we generate synthetic data by
applying the linearized Born scattering operator to $4750$ 2D sections with
size $3075\,\mathrm{m}\times 5120\,\mathrm{m}$ extracted from the real
Kirchhoff migrated Parihaka dataset. We consider a $12.5\,\mathrm{m}$ vertical
and $20\,\mathrm{m}$ horizontal grid spacing, and we augment a
$125\,\mathrm{m}$ water column on top of these models to limit the near source
imaging artifacts. We parameterize the linearized Born scattering operator via
a made up smooth background squared-slowness model. To ensure good coverage,
we simulate $102$ shot records with a source spacing of $50\,\mathrm{m}$. Each
shot is recorded for two seconds with $204$ fixed receivers sampled at
$25\,\mathrm{m}$ spread on top of the model. The source is a Ricker wavelet
with a central frequency of $30\,\mathrm{Hz}$. To mimic a more realistic
imaging scenario, we add band-limited noise to the shot records, where the
noise is obtained by filtering white noise with the source wavelet. We create
training pairs $(\mathbf{y}^{(i)},\mathbf{x}^{(i)}),\ i=1,\ldots,4750$, by
first simulating noisy seismic data according to the above mentioned
acquisition design for all $\mathbf{x}^{(i)}$. Next, we obtain
$\mathbf{y}^{(i)}$ by applying reverse-time-migration to the obtained data for
each model $\mathbf{x}^{(i)}$. We train $T_{w}$ according to the objective
function in Equation 3 with the Adam [41] stochastic optimization method with
batch size $16$ for $100$ passes over the training dataset (epochs). We use an
initial stepsize of $10^{-4}$ and multiply it by $0.9$ every two epochs.

After training, we simulate noisy seismic data for a previously unseen
perturbation model extracted from the Parihaka dataset, shown in Figure 1a.
The corresponding low-fidelity reverse-time migrated image is shown in Figure
1b. Given this image, we use Equation 4 to generate $1000$ random samples
$p(\mathbf{x}\mid\mathbf{y})$ at a negligible cost. Two of these samples are
shown in Figures 1c and 1d. With these $1000$ samples, we also estimate the
conditional mean (Figure 1e) and pointwise standard deviation (Figure 1f).
Because of its reported [42] statistical robustness, we use the conditional
mean as an estimate for the high-fidelity image. This estimate includes
pointwise standard deviations that serve as an assessment of the uncertainty.
From Figure 1, we observe that the overall amplitudes are well recovered by
the conditional mean estimate, which includes partially recovered reflectors
in badly illuminated areas close to the boundaries. Although the
reconstructions are not perfect, they significantly improve upon the reverse-
time migrations estimate. As expected, the pointwise standard deviation in
Figure 1f indicates that we have the most uncertainty in areas of complex
geology—e.g., near channels and tortuous reflectors, and in areas with a
relatively poor illumination (deep and close to boundaries). The areas with
large uncertainty align with difficult-to-image parts of the model.

(a)

(b)

(c)

(d)

(e)

(f)

Figure 1: Reliability-aware imaging. (a) High-fidelity image. (b) Reverse-time
migrated image. (c) and (d) Samples from posterior distribution. (e)
Conditional mean and (f) pointwise standard deviation via the proposed method
according to Equation 5.

In this example, we use Devito [43, 44] for the wave-equation based
simulations. For implementation of the network architectures, we rely on
InvertibleNetworks.jl [28], a memory-efficient framework for training
invertible nets in the Julia programming language. Finally, code to reproduce
our results are made available on
https://github.com/slimgroup/Software.SEG2021.

##  4 Conclusions and discussion

Seismic imaging is challenging due to complex and computationally expensive-
to-evaluate forward operators and the highly heterogeneous multiscale
structure of the Earth. These challenges limit the applicability of Markov
Chain sampling methods due to the costs associated with the forward operator.
These difficulties are compounded by difficulties in capturing the
heterogeneity exhibited by the Earth’s subsurface in the form of a prior
distribution. To handle this situation and to assess uncertainty, we propose a
data-driven, conditional normalizing-flow based variational inference
approach. The proposed inference scheme takes advantage of having access to
training pairs—e.g., in form of high-fidelity seismic images and low-fidelity
remigrated images, to train a conditional normalizing flow capable of sampling
from the posterior distribution. Because neural nets are known to generalize,
the trained conditional normalizing flow can be used to quickly generate
samples from the posterior given unseen low-fidelity images. In turn, these
samples provide access to the high-fidelity conditional mean estimate
including a rudimentary assessment of its uncertainty. In that sense, the
presented imaging scheme learns by example, maximally leveraging what is
already known. Since we can not completely rely on generalization of neural
nets, the trained conditional normalizing flow can be used as input to a
second physics-based stage of inference where the normalizing flow undergoes
transfer learning and then samples from the posterior specific to the imaging
problem at hand. By means of a numerical example, we demonstrate that our
proposed scheme indeed provides fast access to posterior samples for problems
that would otherwise be computationally unfeasible to solve with the more
traditional methods. Since the presented method relies on generalization by
conditional normalizing flows, we stop short of claiming that we quantify
uncertainty but rather that we supply a measure of reliability of the
estimated image.

##  5 Acknowledgment

This research was carried out with the support of Georgia Research Alliance
and partners of the ML4Seismic Center.

## References

  * Tarantola [2005] Albert Tarantola.  _Inverse problem theory and methods for model parameter estimation_.  SIAM, 2005.  ISBN 978-0-89871-572-9.  doi: 10.1137/1.9780898717921.
  * Robert and Casella [2004] C. Robert and G. Casella.  _Monte Carlo statistical methods_.  Springer-Verlag, 2004.
  * Welling and Teh [2011] Max Welling and Yee Whye Teh.  Bayesian Learning via Stochastic Gradient Langevin Dynamics.  In _Proceedings of the 28th International Conference on Machine Learning_ , ICML’11, pages 681–688, Madison, WI, USA, 2011. Omnipress.  ISBN 9781450306195.  doi: 10.5555/3104482.3104568.  URL https://dl.acm.org/doi/abs/10.5555/3104482.3104568.
  * Malinverno and Briggs [2004] Alberto Malinverno and Victoria A Briggs.  Expanded uncertainty quantification in inverse problems: Hierarchical bayes and empirical bayes.  _GEOPHYSICS_ , 69(4):1005–1016, 2004.  doi: 10.1190/1.1778243.
  * Malinverno and Parker [2006] Alberto Malinverno and Robert L Parker.  Two ways to quantify uncertainty in geophysical inverse problems.  _GEOPHYSICS_ , 71(3):W15–W27, 2006.  doi: 10.1190/1.2194516.
  * Martin et al. [2012] James Martin, Lucas C. Wilcox, Carsten Burstedde, and OMAR Ghattas.  A Stochastic Newton MCMC Method for Large-scale Statistical Inverse Problems with Application to Seismic Inversion.  _SIAM Journal on Scientific Computing_ , 34(3):A1460–A1487, 2012.  URL http://epubs.siam.org/doi/abs/10.1137/110845598.
  * Ray et al. [2017] Anandaroop Ray, Sam Kaplan, John Washbourne, and Uwe Albertin.  Low frequency full waveform seismic inversion within a tree based Bayesian framework.  _Geophysical Journal International_ , 212(1):522–542, 10 2017.  ISSN 0956-540X.  doi: 10.1093/gji/ggx428.  URL https://doi.org/10.1093/gji/ggx428.
  * Fang et al. [2018] Zhilong Fang, Curt Da Silva, Rachel Kuske, and Felix J. Herrmann.  Uncertainty quantification for inverse problems with weak partial-differential-equation constraints.  _GEOPHYSICS_ , 83(6):R629–R647, 2018.  doi: 10.1190/geo2017-0824.1.
  * Stuart et al. [2019] Georgia K Stuart, Susan E Minkoff, and Felipe Pereira.  A two-stage Markov chain Monte Carlo method for seismic inversion and uncertainty quantification.  _GEOPHYSICS_ , 84(6):R1003–R1020, 11 2019.  doi: 10.1190/geo2018-0893.1.
  * Zhao and Sen [2019] Zeyu Zhao and Mrinal K Sen.  A gradient based MCMC method for FWI and uncertainty analysis.  In _SEG Technical Program Expanded Abstracts 2019_ , pages 1465–1469, 8 2019.  doi: 10.1190/segam2019-3216560.1.
  * Kotsi et al. [2020] M Kotsi, A Malcolm, and G Ely.  Uncertainty quantification in time-lapse seismic imaging: a full-waveform approach.  _Geophysical Journal International_ , 222(2):1245–1263, 05 2020.  ISSN 0956-540X.  doi: 10.1093/gji/ggaa245.
  * Curtis and Lomax [2001] Andrew Curtis and Anthony Lomax.  Prior information, sampling distributions, and the curse of dimensionality.  _Geophysics_ , 66(2):372–378, 2001.
  * Herrmann et al. [2019] Felix J. Herrmann, Ali Siahkoohi, and Gabrio Rizzuti.  Learned imaging with constraints and uncertainty quantification.  In _Neural Information Processing Systems (NeurIPS) 2019 Deep Inverse Workshop_ , 12 2019.  URL https://arxiv.org/pdf/1909.06473.pdf.
  * Siahkoohi et al. [2020a] A. Siahkoohi, G. Rizzuti, and F. J. Herrmann.  A deep-learning based bayesian approach to seismic imaging and uncertainty quantification.  _82nd EAGE Conference and Exhibition 2020_ , 6 2020a.  URL https://slim.gatech.edu/Publications/Public/Submitted/2020/siahkoohi2020EAGEdlb/siahkoohi2020EAGEdlb.html.
  * Siahkoohi et al. [2020b] Ali Siahkoohi, Gabrio Rizzuti, and Felix J. Herrmann.  Uncertainty quantification in imaging and automatic horizon tracking—a Bayesian deep-prior based approach.  In _SEG Technical Program Expanded Abstracts 2020_ , pages 1636–1640, 9 2020b.  doi: 10.1190/segam2020-3417560.1.
  * Jordan et al. [1999] Michael I Jordan, Zoubin Ghahramani, Tommi S Jaakkola, and Lawrence K Saul.  An Introduction to Variational Methods for Graphical Models.  _Machine Learning_ , 37(2):183–233, 1999.  doi: 10.1023/A:1007665907178.
  * Blei et al. [2017] David M Blei, Alp Kucukelbir, and Jon D McAuliffe.  Variational inference: A review for statisticians.  _Journal of the American statistical Association_ , 112(518):859–877, 2017.
  * Rezende and Mohamed [2015] Danilo Rezende and Shakir Mohamed.  Variational Inference with Normalizing Flows.  In _Proceedings of the 32nd International Conference on Machine Learning_ , volume 37 of _Proceedings of Machine Learning Research_ , pages 1530–1538. PMLR, 07–09 Jul 2015.  URL http://proceedings.mlr.press/v37/rezende15.html.
  * Dinh et al. [2016] Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio.  Density estimation using Real NVP.  _arXiv preprint arXiv:1605.08803_ , 2016.
  * Lensink et al. [2019] K. Lensink, E. Haber, and B. Peters.  Fully Hyperbolic Convolutional Neural Networks, 2019.
  * Kruse et al. [2019] Jakob Kruse, Gianluca Detommaso, Robert Scheichl, and Ullrich Köthe.  HINT: Hierarchical Invertible Neural Transport for Density Estimation and Bayesian Inference.  _arXiv preprint arXiv:1905.10687_ , 2019.
  * Goodfellow et al. [2014] Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio.  Generative Adversarial Nets.  In _Proceedings of the 27th International Conference on Neural Information Processing Systems_ , NIPS’14, pages 2672–2680, 2014.  URL http://papers.nips.cc/paper/5423-generative-adversarial-nets.pdf.
  * Mosser et al. [2020] Lukas Mosser, Olivier Dubrule, and Martin J. Blunt.  Stochastic seismic waveform inversion using generative adversarial networks as a geological prior.  _Mathematical Geosciences_ , 52(1):53–79, 2020\.  doi: 10.1007/s11004-019-09832-6.  URL https://doi.org/10.1007/s11004-019-09832-6.
  * Bishop [2006] Christopher M Bishop.  _Pattern Recognition and Machine Learning_.  Springer-Verlag New York, 2006.
  * Arjovsky et al. [2017] Martin Arjovsky, Soumith Chintala, and Léon Bottou.  Wasserstein generative adversarial networks.  In _Proceedings of the 34th International Conference on Machine Learning_ , volume 70 of _Proceedings of Machine Learning Research_ , pages 214–223. PMLR, 06–11 Aug 2017.  URL http://proceedings.mlr.press/v70/arjovsky17a.html.
  * Putzky and Welling [2019] Patrick Putzky and Max Welling.  Invert to learn to invert.  In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, editors, _Advances in Neural Information Processing Systems_ , volume 32. Curran Associates, Inc., 2019.  URL https://proceedings.neurips.cc/paper/2019/file/ac1dd209cbcc5e5d1c6e28598e8cbbe8-Paper.pdf.
  * Peters and Haber [2020] Bas Peters and Eldad Haber.  Fully reversible neural networks for large-scale 3d seismic horizon tracking.  In _82nd EAGE Annual Conference & Exhibition_, pages 1–5. European Association of Geoscientists & Engineers, 2020.
  * Witte et al. [2021] Philipp Witte, Gabrio Rizzuti, Mathias Louboutin, Ali Siahkoohi, Felix Herrmann, and Bas Peters.  InvertibleNetworks.jl: A Julia framework for invertible neural networks, March 2021.  URL https://doi.org/10.5281/zenodo.4610118.
  * WesternGeco. [2012] WesternGeco.  Parihaka 3D PSTM Final Processing Report.  Technical Report New Zealand Petroleum Report 4582, New Zealand Petroleum & Minerals, Wellington, 2012.
  * Liu [2017] Qiang Liu.  Stein variational gradient descent as gradient flow.  In I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, _Advances in Neural Information Processing Systems_ , volume 30. Curran Associates, Inc., 2017.  URL https://proceedings.neurips.cc/paper/2017/file/17ed8abedc255908be746d245e50263a-Paper.pdf.
  * Zhang and Curtis [2020a] Xin Zhang and Andrew Curtis.  Seismic tomography using variational inference methods.  _Journal of Geophysical Research: Solid Earth_ , 125(4):e2019JB018589, 2020a.
  * Zhang and Curtis [2020b] Xin Zhang and Andrew Curtis.  Variational full-waveform inversion.  _Geophysical Journal International_ , 222(1):406–411, 04 2020b.  ISSN 0956-540X.  doi: 10.1093/gji/ggaa170.  URL https://doi.org/10.1093/gji/ggaa170.
  * Rizzuti et al. [2020] Gabrio Rizzuti, Ali Siahkoohi, Philipp A. Witte, and Felix J. Herrmann.  Parameterizing uncertainty by deep invertible networks, an application to reservoir characterization.  In _SEG Technical Program Expanded Abstracts_ , pages 1541–1545, 09 2020.  doi: 10.1190/segam2020-3428150.1.
  * Zhao et al. [2020] Xuebin Zhao, Andrew Curtis, and Xin Zhang.  Bayesian seismic tomography using normalizing flows.  2020\.  doi: 10.31223/X53K6G.  URL https://eartharxiv.org/repository/view/1940/.
  * Siahkoohi et al. [2021] Ali Siahkoohi, Gabrio Rizzuti, Mathias Louboutin, Philipp Witte, and Felix J. Herrmann.  Preconditioned training of normalizing flows for variational inference in inverse problems.  In _3rd Symposium on Advances in Approximate Bayesian Inference_ , 1 2021.  URL https://openreview.net/pdf?id=P9m1sMaNQ8T.
  * Siahkoohi et al. [2020c] Ali Siahkoohi, Gabrio Rizzuti, Philipp A. Witte, and Felix J. Herrmann.  Faster uncertainty quantification for inverse problems with conditional normalizing flows.  Technical Report TR-CSE-2020-2, Georgia Institute of Technology, 07 2020c.  URL https://arxiv.org/abs/2007.07985.
  * Adler and Öktem [2018] Jonas Adler and Ozan Öktem.  Deep Bayesian Inversion.  _arXiv preprint arXiv:1811.05910_ , 2018.
  * Kovachki et al. [2020] Nikola Kovachki, Ricardo Baptista, Bamdad Hosseini, and Youssef Marzouk.  Conditional Sampling With Monotone GANs.  _arXiv preprint arXiv:2006.06755_ , 2020.
  * Villani [2009] Cédric Villani.  _Optimal transport: old and new_.  Springer-Verlag Berlin Heidelberg, 2009.  doi: 10.1007/978-3-540-71050-9.
  * Marzouk et al. [2016] Youssef Marzouk, Tarek Moselhy, Matthew Parno, and Alessio Spantini.  Sampling via measure transport: An introduction.  _Handbook of uncertainty quantification_ , pages 1–41, 2016.
  * Kingma and Ba [2014] Diederik P. Kingma and Jimmy Ba.  Adam: A Method for Stochastic Optimization.  _CoRR_ , abs/1412.6980, 2014.
  * Anderson and Moore [1979] Brian D.O. Anderson and John B. Moore.  _Optimal Filtering_.  Prentice-Hall, Englewood Cliffs, NJ., 1979.
  * Luporini et al. [2018] F. Luporini, M. Lange, M. Louboutin, N. Kukreja, J. Hückelheim, C. Yount, P. Witte, P. H. J. Kelly, F. J. Herrmann, and G. J. Gorman.  Architecture and performance of devito, a system for automated stencil computation.  _CoRR_ , abs/1807.03032, jul 2018.  URL http://arxiv.org/abs/1807.03032.
  * Louboutin et al. [2019] M. Louboutin, M. Lange, F. Luporini, N. Kukreja, P. A. Witte, F. J. Herrmann, P. Velesko, and G. J. Gorman.  Devito (v3.1.0): an embedded domain-specific language for finite differences and geophysical exploration.  _Geoscientific Model Development_ , 12(3):1165–1187, 2019.  doi: 10.5194/gmd-12-1165-2019.  URL https://www.geosci-model-dev.net/12/1165/2019/.


# Algorithms for Solving Optimization Problems Arising from Deep Neural Net
Models: Nonsmooth Problems

Vyacheslav Kungurtsev and Tomas Pevny

###### Abstract.

Machine Learning models incorporating multiple layered learning networks have
been seen to provide effective models for various classification problems. The
resulting optimization problem to solve for the optimal vector minimizing the
empirical risk is, however, highly nonconvex. This alone presents a challenge
to application and development of appropriate optimization algorithms for
solving the problem. However, in addition, there are a number of interesting
problems for which the objective function is nonsmooth and nonseparable. In
this paper, we summarize the primary challenges involved, the state of the
art, and present some numerical results on an interesting and representative
class of problems.

##  1\. Introduction

In recent years there has been a notable increase in the popularity of models
in machine learning incorporating multiple layers of classifiers, referred to
as Deep Neural Nets (DNNs). Appled to text classification, perception and
identification, and a myriad of other settings, deep learning has, after a
prolonged slow start, showed impressive efficacy [12]. The models arising from
deep learning, being the product of a layer of different models, result in a
very nonconvex problem. In the age of big data, the size of the problems that
are desired to be solved is very large, necessitating the use of stochastic or
batch methods.

Consider the general problem,

$\min_{x\in\mathbb{R}^{n}}F(x),$

where $F$ is generally nonconvex. In this paper in particular, we are
interested in nonsmooth problems. Generally, the literature on optimization
algorithms for machine learning, even with deep learning in mind, has focused
on smooth problems, or, sometimes, smooth problems with an additional
separable regularization term (usually the $l1$ norm) [5]. In the case of a
simple regularizer, these are typically separable, that is, decomposable as a
sum of functions of individual variables or blocks (e.g.,
$\|x\|_{1}=\sum_{i=1}^{N}|x_{1}|$), which permits a closed form solution of a
prox operator, as in the iterative shrinkage algorithm. Coordinate descent
type algorithms have proven successful for problems with nonconvex loss
functions and separable regularizers [16, 7]. Although extensions to include
stochastic or batch subproblem/step formulation are necessary to truly apply
them to huge scale deep learning problems, beyond that they can only converge
to a coordinate-wise stationary point, i.e., a point $x^{*}$ such that there
exists $\xi\in\partial_{i}F(x^{*})$ with $(x_{i}-x_{i}^{*})^{T}\xi\geq 0$ for
all $x$, $i$, but this does not necessarily mean stationarity, wherein there
exists $\xi\in\partial F(x^{*})$ with $(x-x^{*})^{T}\xi\geq 0$ for all $x$.

However, in many applications, the actual loss function itself can be
nonsmooth. In particular, some upper layers can involve an activation
function, defined as involving terms of the form $\max(0,a^{T}x+b)$. This
presents a large scale optimization problem that is nonconvex, nonsmooth, and
nonseparable, a particularly challenging class of problems with no algorithm
that stands out as particularly suitable for reliably and efficiently solving
them.

##  2\. Nonsmooth, nonconvex problems, state of the art and prognosis

In general, solving nonconvex nonsmooth problems is a formidable task.
Provably convergent algorithms for solving such problems are uncommon. One of
the most robust such algorithms, gradient sampling [6] requires $2n$ samples
of the gradient to be taken at each iteration, which implies that its
performance scales very poorly with the size of the problem. Yet this
algorithm is regularly used in areas of robust control, for instance. Aside
from this, the provably convergent algorithms include classes of bundle
methods and some other subgradient-approximation type algorithms, namely
discrete gradient and pseudosecant methods. Thorough numerical experiments
given in [1] suggest that for large scale problems, the limited memory bundle
method [10] and the quasisecant method [3] were the most robust and efficient.

However, the quasisecant method requires the computation of quasisecants at
each iteration, which is a very problem dependent and potentially complicated
procedure [2]. Given the convoluted structure of deep neural net models, the
computation of quasisecants at each iteration of an optimization for solving
such a problem would be intractable. With respect to bundle methods, each
iteration involves solving a quadratic program (QP) that summarizes the
function locally from the historical information of (sub)gradients. Often in
applications involving big data, computation of an entire gradient vector is
impossible, given the storage difficulties associated with all of the data
batch functions. Thus either stochastic, or batch averages, or gradient
approximations must be used. The robustness of using such approximations in a
bundle method can be questionable. Although there are bundle methods
incorporating inexact information [8], experience with applying them to this
class of problems is limited, at best. We shall see in the next Section 3 that
it is, at best, not straightforward.

##  3\. Numerical Results

Recall that by Rademacher’s Theorem, it holds that for every point outside a
set of measure zero, a unique gradient vector can be computed [4]. This
implies that gradient methods, or more specifically, stochastic gradient and
batch gradient methods are, in practice, well-defined for nonsmooth problems.
That said, there exists no convergence theory for the basic variation of the
(sub)gradient method, even with a diminishing step-size, for nonsmooth
nonconvex problems. In terms of practical implementation, however, it is
practical to suggest them as a method to try for this class of problems, in
their generic form, as they are the standard bearer for large scale
optimization for problems arising in machine learning in general.

Just as basic SGD methods have been seen to converge slowly for highly
nonconvex problems with a lot of saddle points, however, we expect that their
practical convergence speed for highly nonconvex, and thus ill-conditioned
problems, to be fairly slow. Thus, just as in the smooth case, we expect that
there is a case for higher-order algorithms, but with a number of difficulties
in reliable and efficient implementation. Note that, we at least at least one
application area of strong interest, with the activation function mentioned
above, to be piecewise linear. This implies that the second derivative is zero
almost everywhere (and then undefined on a set of measure zero), and thus
Newton methods (or stochastic variations thereof) are entirely useless.
However, despite a lack of convergence theory, quasi-Newton, specifically BFGS
methods have numerically demonstrated strong efficacy for solving general
nonsmooth, nonconvex problems [13]. Given the scale of the problems of
interest, storing dense Hessians is impractical, and thus the limited memory
variety of BFGS [14] has been tested on optimizations problems arising from
DNN models to good success [9, 15].

We compared the performance of ADAM [11], a contemporary stochastic gradient
algorithm that scales the direction based on accumulated history of apparent
conditioning, sum of functions optimizer (SFO) [17], a second order BFGS-type
algorithm using batches, and the limited memory bundle method (LMBM) [10] on a
model of deep learning of security data from Cisco in Prague, in the Czech
Republic. We studied their performance, indicating the scalability, on a
variety of sizes of batches, each composed of several hundred samples. The
LMBM was run in an entirely sequential and undistributed manner, i.e.,
treating the problem in its entirety, while the other two methods are
stochastic or single-batch based.

We see in Figure 1 the final value of each algorithm after 1000 iterations.
Overall it is clear that the SFO algorithm results in the best objective
value, LMBM is second overall, but more consistent/robust, and finally ADAM is
last.

Figure 1.  Plot comparing the final value after 1000 iterations of ADAM, SFO,
and LMBM.

Taking the final value that ADAM generates after 1000 iterations, we now plot
the time until each algorithm reaches this value in Figure 2. This also
suggests that SFO is the most efficient algorithm, followed by LMBM and ADAM
as a distant second.

Figure 2.  Plot comparing the time to convergence of ADAM, SFO, and LMBM. |
---|---

After 35 batches, the required computation and storage that was needed for the
bundle method was too much. We attempted to modify the LMBM to handle larger
problem sizes. Thus far, none of the attempts we implemented, including,

  1. (1)

computing the entire gradient by computing each batch gradient and averaging,

  2. (2)

choosing a gradient from a random batch each iteration,

  3. (3)

running LMBM to minimize the problem for each batch sequentially,

resulted in an efficiently and quickly converging algorithm. This does not
necessarily suggest that huge-scale bundle methods are impossible, but their
implementation is certainly not straightforward.

##  4\. Discussion

A provably convergent algorithm for huge scale nonconvex, nonsmooth, and
nonseparable optimization problems is, for the most part, a open problem. The
demands from applications in machine learning, with the arrival of the
widespread use of deep neural nets, has shifted the necessity of solving large
convex problems to nonconvex ones. The optimization literature has matched
this demand to some degree, but the added challenge of solving nonsmooth
problems, aside from simple separable regularizers, appears to be far more
difficult. Numerical experience suggests the efficacy of carefully designed
quasi-Newton methods, but convergence is not theoretically guaranteed. The top
competing algorithms for nonconvex, nonsmooth optimizations on a (smaller, but
still) large scale, quasi-secant and bundle methods, are not straightforward
to extend to a larger scope of problem size, in the first case due to the
challenge of computing the step, and in the second due to memory requirements.
We foresee that a proximal type bundle method with stochastic information can
be developed in the future as a provably convergent algorithm for these cases
of problems, and we will focus on this as future work, even if thus far such a
procedure remains elusive.

## References

  * [1] Adil Bagirov, Napsu Karmitsa, and Marko M Mäkelä.  Introduction to Nonsmooth Optimization: theory, practice and software.  Springer, 2014.
  * [2] Adil M Bagirov and Asef Nazari Ganjehlou.  A quasisecant method for minimizing nonsmooth functions.  Optimization Methods & Software, 25(1):3–18, 2010.
  * [3] Adil M Bagirov, L Jin, N Karmitsa, A Al Nuaimat, and Napsu Sultanova.  Subgradient method for nonconvex nonsmooth optimization.  Journal of Optimization Theory and applications, 157(2):416–435, 2013.
  * [4] D Bessis and F Clarke.  Partial subdifferentials, derivates and rademacher’s theorem.  Transactions of the American Mathematical Society, 351(7):2899–2926, 1999.
  * [5] Léon Bottou, Frank E Curtis, and Jorge Nocedal.  Optimization methods for large-scale machine learning.  arXiv preprint arXiv:1606.04838, 2016.
  * [6] James V Burke, Adrian S Lewis, and Michael L Overton.  A robust gradient sampling algorithm for nonsmooth, nonconvex optimization.  SIAM Journal on Optimization, 15(3):751–779, 2005.
  * [7] Amir Daneshmand, Francisco Facchinei, Vyacheslav Kungurtsev, and Gesualdo Scutari.  Hybrid random/deterministic parallel algorithms for convex and nonconvex big data optimization.  IEEE Transactions on Signal Processing, 63(15):3914–3929, 2015\.
  * [8] W de Oliveira, C Sagastizábal, and Claude Lemaréchal.  Convex proximal bundle methods in depth: a unified analysis for inexact oracles.  Mathematical Programming, 148(1-2):241–277, 2014.
  * [9] Jeffrey Dean, Greg Corrado, Rajat Monga, Kai Chen, Matthieu Devin, Mark Mao, Andrew Senior, Paul Tucker, Ke Yang, Quoc V Le, et al.  Large scale distributed deep networks.  In Advances in neural information processing systems, pages 1223–1231, 2012.
  * [10] Marjo Haarala, Kaisa Miettinen†, and Marko M Mäkelä‡.  New limited memory bundle method for large-scale nonsmooth optimization.  Optimization Methods and Software, 19(6):673–692, 2004.
  * [11] Diederik Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  arXiv preprint arXiv:1412.6980, 2014.
  * [12] Yann LeCun, Yoshua Bengio, and Geoffrey Hinton.  Deep learning.  Nature, 521(7553):436–444, 2015.
  * [13] Adrian S Lewis and Michael L Overton.  Nonsmooth optimization via quasi-newton methods.  Mathematical Programming, 141(1-2):135–163, 2013.
  * [14] Dong C Liu and Jorge Nocedal.  On the limited memory bfgs method for large scale optimization.  Mathematical programming, 45(1-3):503–528, 1989.
  * [15] Jiquan Ngiam, Adam Coates, Ahbik Lahiri, Bobby Prochnow, Quoc V Le, and Andrew Y Ng.  On optimization methods for deep learning.  In Proceedings of the 28th International Conference on Machine Learning (ICML-11), pages 265–272, 2011.
  * [16] Peter Richtárik and Martin Takáč.  Parallel coordinate descent methods for big data optimization.  Mathematical Programming, 156(1-2):433–484, 2016.
  * [17] Jascha Sohl-dickstein, Ben Poole, and Surya Ganguli.  Fast large-scale optimization by unifying stochastic gradient and quasi-newton methods.  In Proceedings of the 31st International Conference on Machine Learning (ICML-14), pages 604–612, 2014.


# Unified Keyword Spotting and Audio Tagging on Mobile Devices with
Transformers

###### Abstract

Keyword spotting (KWS) is a core human-machine-interaction front-end task for
most modern intelligent assistants. Recently, a unified (UniKW-AT) framework
has been proposed that adds additional capabilities in the form of audio
tagging (AT) to a KWS model. However, previous work did not consider the real-
world deployment of a UniKW-AT model, where factors such as model size and
inference speed are more important than performance alone. This work
introduces three mobile-device deployable models named Unified Transformers
(UiT). Our best model achieves an mAP of 34.09 on Audioset, and an accuracy of
97.76 on the public Google Speech Commands V1 dataset. Further, we benchmark
our proposed approaches on four mobile platforms, revealing that the proposed
UiT models can achieve a speedup of 2 - 6 times against a competitive
MobileNetV2.

Index Terms—  Keyword spotting, Audio tagging, Vision Transformers, weakly
supervised learning.

##  1 Introduction

Keyword spotting (KWS) is currently a crucial front-end task for most
intelligent voice assistants, which triggers the start of an interaction with
the voice assistant if a user utters a specific keyphrase. Further, Audio
tagging (AT) is a task that aims to label specific audio content into sound
event classes, e.g., the sound of a baby crying. In the previous work [1], the
authors have shown that modelling both tasks via a unified framework (UniKW-
AT) is possible, significantly improving noise robustness without sacrificing
KWS accuracy.

However, if UniKW-AT models were deployed in real-world scenarios, they would
need to fulfil the same requirements as KWS models. First, KWS models are
situated on-device, and their size, i.e., the number of parameters, is
limited. Second, KWS models require a fast inference speed and a small
floating point operations (FLOPs) footprint due to being “always on”. Third,
the delay of a KWS model needs to be as low as possible, such that an
assistant’s wakeup is immediate. While these requirements have already been
researched thoroughly within the KWS community, to the best of our knowledge,
no previous study has focused on lightweight and on-device computation of AT
models. This work aims to bridge the gap and introduce a multitude of
transformer-based models for unified keyword spotting and audio tagging
(UniKW-AT), which can satisfy the requirements mentioned above.

The benefit of such a UniKW-AT model is that outputs from the AT branch can be
passed down further into the automatic speech recognition (ASR) pipeline,
possibly enhancing robustness against noise. At the very least such a UniKW-AT
model can act as a voice activity detector [2, 3, 4].

###  1.1 Previous work

Due to the practical importance of KWS, previous work is focused on decreasing
a model’s parameter size [5], increasing its inference speed [6] and reducing
its false-acceptance rate [7]. In terms of architecture, convolutional neural
networks (CNNs) have been well researched within the community [8, 9, 5],
while more recently transformer-based models [10, 11, 12, 13] and multi-layer
perceptron (MLP) mixers [14, 15] have also been studied. On the contrary,
within the field of AT, most work focuses on improving the state-of-the-art
performance on the well-known Audioset benchmark. Works such as [16]
popularized CNNs, while [17] utilized transformers. However, the majority of
research within AT is solely focused on improving performance without
consideration for the real-world deployment of these models.

##  2 Unified keyword-spotting and audio tagging transformers

This paper contributes a variety of transformer-based networks, further called
unified transformers (UiT), which aim to provide fast inference speed and
reduce the model-parameter size and computational overhead while preserving
KWS and AT performance.

#### Unified Keyword Spotting and Audio Tagging

UniKW-AT has been proposed in [1] and is modelled as follows. Given the target
KWS labelset $\mathbb{L}_{\text{KWS}}$ with $K$ keywords and an AT labelset
$\mathbb{L}_{\text{AT}}$ with $C$ sound events, UniKW-AT merges both labelsets
obtaining $\mathbb{L}=\mathbb{L}_{\text{KWS}}\cup\mathbb{L}_{\text{AT}}$.
Training samples from both KWS and AT datasets are randomly cropped to some
target duration $t$ and the framework is optimized via the binary cross
entropy (BCE) loss. The entire training framework can be seen in Figure 1.

#### Vision Transformers

Transformers were first proposed for machine translation in [18] and quickly
became the state-of-the-art approach within the field of natural language
processing (NLP). Later in [19], the Vision Transformer (ViT) has been
proposed as an adaptation of transformers into the field of computer vision.
Then, ViT-based transformers were used in AT, where images were replaced with
two-dimensional spectrograms [17, 20]. The core idea of the ViT framework is
the “patchification” operation, where an input image (here spectrogram) is
first split into $N$ non-overlapping patches. Each patch is of size
$P=P_{\mathtt{T}}\times P_{\mathtt{F}}$ (time and frequency) and extracted via
a convolution operation. Then these patches are fed into a Transformer model
consisting of $L$ identical blocks of a multi-head attention (MHA) layer
followed by a multi-layer perceptron (MLP) layer. A MHA layer computes:

$\displaystyle\mathbf{A}=\text{softmax}(\frac{\mathbf{QK}^{T}}{\sqrt{D}})\mathbf{V},$
(1)

where $\mathbf{K}\in\mathbb{R}^{N\times D},\mathbf{Q}\in\mathbb{R}^{N\times
D},\mathbf{V}\in\mathbb{R}^{N\times D}$ and are the key, query and value
matrices obtained by a linear transformation
$\mathbf{W}_{j}\in\mathbb{R}^{D\times D},j\in\\{K,Q,V\\}$ of an input
$\mathbf{X}\in\mathbb{R}^{N\times D}$. The complexity of a single block is
$\mathcal{O}(N^{2}D+D^{2}N)$, i.e., is quadratic for the model dimension $D$
and the number of patches $N$.

###  2.1 Proposed Model

Fig. 1: Depiction of the proposed UiT framework for training. We sample from
an AT (Audioset) dataset and a KWS dataset (GSCV1). Then we randomly crop the
Audioset sample to match the KWS target sample length $t$. Afterwards, a pre-
trained model (MobileNetV2) is used to estimate these labels for these
Audioset samples (pseudo-strong labels). During training, each batch is
created using 50% of AT and 50% of KWS data. Training is done by optimizing
the binary cross entropy (BCE) criterion.

We identify the embedding dimension $D$ and the number of patches $N$ as the
primary reason for a large computational overhead and propose patch-reduction
and bottleneck attention to reduce the complexity of our UiT model.

#### Patch-reduction and Subsampling

Common transformer models on Audioset are trained on a 10-s scale, utilize up
to $N=N_{\mathtt{T}}\times N_{\mathtt{F}}=1212$ [20, 17] patches, leading to a
high memory and computational burden for the transformer model. Since time-
frequency information is crucial for AT, we focus on reducing the amount of
available time-patches $N_{\mathtt{T}}$ and fix the number of frequency
patches $N_{\mathtt{F}}$ by feeding the model crops of length $t$ i.e., 1 s.
Limiting the context harms AT performance on Audioset since our model’s
training (1 s), and evaluation (10 s) duration is mismatched. However pseudo
strong labels (PSL) [21] can partially alleviate the performance degradation
by predicting fine-scale labels from a pretrained teacher (MobileNetV2) model.

The first layer within most CNN models, commonly known as stem, maps a
$1\times T\times F$ spectrogram input to
$C\times\nicefrac{{T}}{{2}}\times\nicefrac{{F}}{{2}}$, where $C\geq 16$ [22],
meaning that the overall memory required will be expanded by a factor of $\geq
4$. Our approach uses a subsampling stem, directly reducing the memory
requirement by mapping an input patch of size $P$ to a low-dimensional space
$D$, where $D<P$.

#### Bottleneck attention

Bottleneck attention (BN-A) focuses on reducing the dimension $D$ during the
self-attention stage Equation 1. Our intuition is that each respective patch-
embedding within a spectrogram contains large amounts of redundant
information. Therefore, we propose using a BN-A approach, which reduces the
dimension $D$ during self-attention to a lower space $U$, $U<D$. We set
$U=\frac{D}{4}$ for all our models.

#### Architecture

The proposed model architectures can be seen in Table 1. For all architectures
we use $P=256=16\times 16$ patch-sizes, which sets the delay of our models to
$16$ frames and use two heads for each self-attention operation. This leads to
$N_{\mathtt{F}}=4$ patches along the frequency axis and to $N_{\mathtt{T}}=6$
patches for each input audio second. We use ReLU as the default activation
function due to its faster inference speed and lower computational complexity.
Similar to [11], we use a lower embedding dimension of $3D$ within the MLP,
reducing the memory footprint.

Model | $L$ | $D$ | MLP | #Params | MFLops | $M_{pk}$
---|---|---|---|---|---|---
UiT-XS | 12 | 128 | 384 | 1.5 M | 34 | 7.59
UiT-2XS | 6 | 128 | 384 | 0.8 M | 18 | 4.10
UiT-3XS | 4 | 128 | 384 | 574 k | 13 | 3.15
Table 1: Proposed UiT-based model architectures. The number of MFLOPs and the
peak memory usage $M_{pk}$ (in MB) are calculated over 1 s.

##  3 Experiments

###  3.1 Datasets

This work mainly uses the Google Speech Commands V1 (GSCV1) [23] and Audioset
[24] datasets. We use the common 11 class subset of GSCV1 (V1-11), where the
original 30 classes have been reduced to 10 common keywords: “Yes”, “No”,
“Up”, “Down”, “Left”, “Right”, “On”, “Off”, “Stop”, “Go” while the other 20
keywords are labeled as the AS label “Speech”, where each sample is 1 s long.
We use the official training/validation/testing split containing
51,088/6,798/6,835 utterances, respectively. As for AT training, we use the 60
h long balanced subset of AS containing 21,292 audio clips with a duration of
at most 10 s per clip.

#### Evaluation

Evaluation is split between the KWS and AT subtasks. AT evaluation uses the
common evaluation subset of AS, containing 18,229 audio clips with an overall
duration of 50 h. KWS analysis primarily focuses on the GSCV1 dataset, which
provides 2,567 target keywords and 4,268 non-target samples. Note that for
clips longer than the target length $t$, i.e., 10s, we split the input into
chunks of length $t$ (i.e., 1 s), then feed these chunks into the model and
average all output scores.

###  3.2 Setup

Regarding front-end feature extraction, we use log-Mel spectrograms (LMS) with
64 bins extracted every 10 ms with a window of 32ms and a 16 kHz sampling
rate. Our UiT transformer models use time- and frequency-independent learnable
position embeddings. We use random shifting, volume gain, and polarity
inversion as augmentation methods in the waveform domain. Having obtained a
spectrogram, we augment the data using Specaugment [25]. Training runs with a
batch size of 64 for at most 800 epochs using AdamW optimization [26] with a
linear warmup of 20 epochs to a starting learning rate of 0.001, which is then
gradually decreased using cosine annealing. We use mean average precision
(mAP) during training as our primary validation metric. Computing KWS accuracy
for UniKW-AT requires post-processing since the model can predict multiple
labels simultaneously, i.e., “Keyword + Speech“. Here we use a threshold of
$\gamma=0.2$, indicating the presence of a keyword. The top-4 models achieving
the highest mAP on the joint held-out validation dataset (GSCV1-valid and AS-
balanced) are submitted for evaluation. The neural network back-end is
implemented in Pytorch [27]. To speed up the training procedure, we pretrain a
single UiT-XS on the full Audioset using the masked Autoencoder approach [20]
and initialize all relevant layers for each model from this single checkpoint.
The source code is publicly available111www.github.com/Richermans/UIT_Mobile.

##  4 Results

###  4.1 Main results

The core results of our work on the GSCV1 and Audioset can be seen in Table 2.
To illustrate the difficulty of training a UniKW-AT model, we also ran
baseline experiments using a common TC-ResNet8 [5] model. The results show
that even though TC-ResNet8 can achieve excellent performance on GSCV1 (96.72
Acc), it fails to provide meaningful performance on AT (8.67 mAP). Note that
TC-ResNet8’s performance on GSCV1 improves against the publicly reported
result ($96.10\rightarrow 96.72$) due to UniKW-AT training, where 60h of
“noise” samples from Audioset enhance the model’s robustness to false alarms.
As we can see, our proposed UiT-XS achieves competitive results compared to
the previous MobileNetV2 (MBv2) based method for UniKW-AT, as well as other
works in the literature regarding KWS and AT performance.

Approach | #Params (M) | GSCV1 | AS
---|---|---|---
TC-ResNet8 [5] | 0.06 | 96.10 | -
NAS2 [6] | 0.88 | 97.22 | -
MEGA [28] | 0.3 | 96.92 | -
MatchBoxNet [29] | 0.5 | 96.83 | -
KWT-1 [10] | 0.6 | 97.05 | -
LETR-128 [11] | 0.6 | 97.61 | -
LETR-256 [11] | 1.1 | 97.85 | -
KWT-2 [10] | 2.4 | 97.36 | -
KWT-3 [10] | 5.3 | 97.24 | -
Wav2KWS [30] | 225 | 97.90 | -
MBv2 [31] | 2.9 | - | 26.50
Eff-B0 [31] | 5.3 | - | 33.50
Eff-B2 [31] | 13.6 | - | 34.06
ResNet-50 [31] | 25.6 | - | 31.80
CNN14 [16] | 76 | - | 27.80
AudioMAE [20] | 80 | - | 37.10
TC-ResNet8 | 0.1 | 96.72 | 8.67
MBv2 [1] | 2.9 | 97.53 | 33.42
MBv2${}^{\vardiamondsuit}$ [1] | 2.9 | 97.53 | 32.51
UiT-XS${}^{\vardiamondsuit}$ | 1.5 | 97.76 | 34.09
UiT-2XS${}^{\vardiamondsuit}$ | 0.8 | 97.31 | 32.21
UiT-3XS${}^{\vardiamondsuit}$ | 0.6 | 97.18 | 30.97
Table 2: A comparison between our proposed UiT approaches against other works
in literature. Results for GSCV1 use accuracy, while AS use mAP. Approaches
denoted with ${}^{\vardiamondsuit}$ evaluate on 1 s chunks, influencing
Audioset performance. Results with “-” means not available.

###  4.2 Inference latency on mobile hardware

Here we measure the inference speed of our models using the PyTorch Mobile
Benchmark Tool222https://pytorch.org/tutorials/recipes/mobile_perf.html. In
Table 3, we display measured inference speed on four different mobile devices:
We use two high-end Qualcomm Snapdragon chips, a Snapdragon 865 (SD865) and a
Snapdragon 888 (SD888) and two mid-range chipsets, a MediaTek Helio G90T
(G90T) and a Mediatek Dimensity 700 (MT700). The results are compared to a TC-
ResNet8 and an MBv2. The MBv2 can be viewed as a baseline, representing a
previous UniKW-AT approach, while TC-ResNet8 represents the speed requirement
of a modern KWS system.

Model | SD865 | SD888 | G90T | MT700
---|---|---|---|---
TC-ResNet8 | 0.4 | 0.4 | 1.1 | 1.1
MBv2 | 8.0 | 6.2 | 13.1 | 11.6
Uit-XS | 3.4 | 3.4 | 7.3 | 7.1
UiT-2XS | 1.7 | 1.5 | 2.8 | 3.2
UiT-3XS | 1.2 | 1.1 | 2.2 | 2.2
Table 3: Inference speed comparison on mobile system-on-a-chip (SoC)
platforms, run on the central processing unit (CPU) with float32 precision and
measured in ms. Each speed evaluation is accessed by first warming up the chip
with ten warmup iterations followed by 1000 test trials for an input of length
1 s.

As the latency results demonstrate, our proposed approach can achieve a
speedup of over 2 times (8.0 $\rightarrow$ 3.4 ms) against an MBv2 when using
UiT-XS while achieving a similar performance (see Table 2). Even though
UiT-2XS and UiT-3XS are slower than the baseline TC-ResNet8, they excel at AT
(32.21/30.97 vs. 8.67 mAP, see Table 2). Another important factor worth noting
is that the baseline MBv2 has a delay of 320 ms, while our proposed models
react within 160 ms.

###  4.3 Ablation

Here we present ablation studies focused on the proposed BN-A mechanism and
the choice of a ReLU activation function in favour of a more common Gaussian
Error Linear Unit (GeLU) [32]. For simplicity, we only use the SD865 chipset
for these tests. The results can be seen in Table 4. We observe that the
proposed BN-A mechanism speeds up the inference time by at least 20%
($4.1\rightarrow 3.4$ ms) against standard self-attention without noticeable
performance differences. Moreover, while standard self-attention with GeLU can
provide marginal performance boosts compared to the proposed BN-A + ReLU
approach, they also significantly slow down the inference speed
($3.4\rightarrow 5.7$ ms for UiT-XS) and increase the peak memory usage
($7.59\rightarrow 10.28$ MB for UiT-XS), limiting their potential real-world
use.

Model | Ablation | GSCV1 | AS | Speed | $M_{pk}$
---|---|---|---|---|---
UiT-XS | Proposed | 97.76 | 34.09 | 3.4 | 7.59
w/o BN-A | 97.75 | 33.76 | 4.1 | 10.28
ReLU $\rightarrow$ GeLU | 97.69 | 34.12 | 5.7 | 10.28
UiT-2XS | Proposed | 97.31 | 32.21 | 1.7 | 4.10
w/o BN-A | 96.94 | 32.30 | 2.2 | 5.44
ReLU $\rightarrow$ GeLU | 97.27 | 32.41 | 2.9 | 5.44
UiT-3XS | Proposed | 97.18 | 30.97 | 1.2 | 3.15
w/o BN-A | 96.71 | 30.87 | 1.6 | 3.90
ReLU $\rightarrow$ GeLU | 96.91 | 30.99 | 2.1 | 3.90
Table 4: Ablation studies of the proposed model, where “Speed” represents
measured ms and $M_{pk}$ is the peak memory requirement in MB. ‘w/o BN-A‘
represents the use of a standard self-attention mechanism, whereas ‘ReLU
$\rightarrow$ GeLU‘ utilizes GeLU with standard self-attention. Results are
evaluated on an input of 1 s. Best in bold.

##  5 Conclusion

This paper proposes ViT-based transformer models (UiT) optimized for mobile
device deployment of UniKW-AT. Three lightweight UiT models are introduced,
providing excellent performance for UniKW-AT coupled with fast inference
speed, low peak memory usage and a small parameter footprint. Our best model
(UiT-XS) achieves an accuracy of 97.76 on the GSCV1 dataset and an mAP of
34.09 on Audioset, outperforming a competitive MBv2 baseline while having half
the parameter footprint and a predominantly increased inference speed.

## References

  * [1] Heinrich Dinkel, Yongqing Wang, Zhiyong Yan, Junbo Zhang, and Yujun Wang,  “UniKW-AT: Unified Keyword Spotting and Audio Tagging,”  in Proc. Interspeech 2022, 2022, pp. 3238–3242.
  * [2] Heinrich Dinkel, Shuai Wang, Xuenan Xu, Mengyue Wu, and Kai Yu,  “Voice activity detection in the wild: A data-driven approach using teacher-student training,”  IEEE/ACM Transactions on Audio, Speech, and Language Processing, vol. 29, pp. 1542–1555, 2021.
  * [3] Xuenan Xu, Heinrich Dinkel, Mengyue Wu, and Kai Yu,  “A lightweight framework for online voice activity detection in the wild.,”  in Interspeech, 2021, pp. 371–375.
  * [4] Sashi Novitasari, Takashi Fukuda, and Gakuto Kurata,  “Improving ASR Robustness in Noisy Condition Through VAD Integration,”  in Proc. Interspeech 2022, 2022, pp. 3784–3788.
  * [5] Seungwoo Choi, Seokjun Seo, Beomjun Shin, Hyeongmin Byun, Martin Kersner, Beomsu Kim, Dongyoung Kim, and Sungjoo Ha,  “Temporal Convolution for Real-Time Keyword Spotting on Mobile Devices,”  in Proc. Interspeech 2019, 2019, pp. 3372–3376.
  * [6] Tong Mo, Yakun Yu, Mohammad Salameh, Di Niu, and Shangling Jui,  “Neural Architecture Search for Keyword Spotting,”  in Proc. Interspeech 2020, 2020, pp. 1982–1986.
  * [7] Anirudh Raju, Sankaran Panchapagesan, Xing Liu, Arindam Mandal, and Nikko Strom,  “Data augmentation for robust keyword spotting under playback interference,”  arXiv preprint arXiv:1808.00563, 2018.
  * [8] Tara N. Sainath and Carolina Parada,  “Convolutional neural networks for small-footprint keyword spotting,”  in Proc. Interspeech 2015, 2015, pp. 1478–1482.
  * [9] Byeonggeun Kim, Simyung Chang, Jinkyu Lee, and Dooyong Sung,  “Broadcasted Residual Learning for Efficient Keyword Spotting,”  in Proc. Interspeech 2021, 2021, pp. 4538–4542.
  * [10] Axel Berg, Mark O’Connor, and Miguel Tairum Cruz,  “Keyword Transformer: A Self-Attention Model for Keyword Spotting,”  in Proc. Interspeech 2021, 2021, pp. 4249–4253.
  * [11] Kevin Ding, Martin Zong, Jiakui Li, and Baoxiang Li,  “Letr: A lightweight and efficient transformer for keyword spotting,”  in ICASSP 2022 - 2022 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), 2022, pp. 7987–7991.
  * [12] Yiming Wang, Hang Lv, Daniel Povey, Lei Xie, and Sanjeev Khudanpur,  “Wake word detection with streaming transformers,”  in ICASSP 2021-2021 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP). IEEE, 2021, pp. 5864–5868.
  * [13] Vineet Garg, Wonil Chang, Siddharth Sigtia, Saurabh Adya, Pramod Simha, Pranay Dighe, and Chandra Dhir,  “Streaming Transformer for Hardware Efficient Voice Trigger Detection and False Trigger Mitigation,”  in Proc. Interspeech 2021, 2021, pp. 4209–4213.
  * [14] Mashrur M Morshed and Ahmad Omar Ahsan,  “Attention-free keyword spotting,”  arXiv preprint arXiv:2110.07749, 2021.
  * [15] Dianwen Ng, Yunqi Chen, Biao Tian, Qiang Fu, and Eng Siong Chng,  “Convmixer: Feature interactive convolution with curriculum learning for small footprint and noisy far-field keyword spotting,”  arXiv preprint arXiv:2201.05863, 2022.
  * [16] Qiuqiang Kong, Yin Cao, Turab Iqbal, Yuxuan Wang, Wenwu Wang, and Mark D. Plumbley,  “PANNs: Large-Scale Pretrained Audio Neural Networks for Audio Pattern Recognition,”  IEEE/ACM Transactions on Audio Speech and Language Processing, vol. 28, pp. 2880–2894, dec 2020.
  * [17] Yuan Gong, Yu-An Chung, and James Glass,  “AST: Audio Spectrogram Transformer,”  in Proc. Interspeech 2021, 2021, pp. 571–575.
  * [18] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin,  “Attention is all you need,”  Advances in neural information processing systems, vol. 30, 2017\.
  * [19] Alexander Kolesnikov, Alexey Dosovitskiy, Dirk Weissenborn, Georg Heigold, Jakob Uszkoreit, Lucas Beyer, Matthias Minderer, Mostafa Dehghani, Neil Houlsby, Sylvain Gelly, Thomas Unterthiner, and Xiaohua Zhai,  “An image is worth 16x16 words: Transformers for image recognition at scale,”  2021\.
  * [20] Hu Xu, Juncheng Li, Alexei Baevski, Michael Auli, Wojciech Galuba, Florian Metze, Christoph Feichtenhofer, et al.,  “Masked autoencoders that listen,”  arXiv preprint arXiv:2207.06405, 2022.
  * [21] Heinrich Dinkel, Zhiyong Yan, Yongqing Wang, Junbo Zhang, and Yujun Wang,  “Pseudo strong labels for large scale weakly supervised audio tagging,”  in ICASSP 2022 - 2022 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), 2022, pp. 336–340.
  * [22] Mark Sandler, Andrew Howard, Menglong Zhu, Andrey Zhmoginov, and Liang-Chieh Chen,  “Mobilenetv2: Inverted residuals and linear bottlenecks,”  in Proceedings of the IEEE conference on computer vision and pattern recognition, 2018, pp. 4510–4520.
  * [23] Pete Warden,  “Speech commands: A dataset for limited-vocabulary speech recognition,”  arXiv preprint arXiv:1804.03209, 2018.
  * [24] Jort F Gemmeke, Daniel PW Ellis, Dylan Freedman, Aren Jansen, Wade Lawrence, R Channing Moore, Manoj Plakal, and Marvin Ritter,  “Audio set: An ontology and human-labeled dataset for audio events,”  in 2017 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP). IEEE, 2017, pp. 776–780.
  * [25] Daniel S. Park, William Chan, Yu Zhang, Chung-Cheng Chiu, Barret Zoph, Ekin D. Cubuk, and Quoc V. Le,  “SpecAugment: A Simple Data Augmentation Method for Automatic Speech Recognition,”  in Proc. Interspeech 2019, 2019, pp. 2613–2617.
  * [26] Tim Dettmers, Mike Lewis, Sam Shleifer, and Luke Zettlemoyer,  “8-bit optimizers via block-wise quantization,”  in International Conference on Learning Representations, 2021.
  * [27] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala,  “PyTorch: An Imperative Style, High-Performance Deep Learning Library,”  in Advances in Neural Information Processing Systems 32, H Wallach, H Larochelle, A Beygelzimer, F Alché-Buc, E Fox, and R Garnett, Eds., pp. 8026–8037. Curran Associates, Inc., 2019.
  * [28] Xuezhe Ma, Chunting Zhou, Xiang Kong, Junxian He, Liangke Gui, Graham Neubig, Jonathan May, and Luke Zettlemoyer,  “Mega: Moving average equipped gated attention,”  arXiv preprint arXiv:2209.10655, 2022.
  * [29] Somshubra Majumdar and Boris Ginsburg,  “MatchboxNet: 1D Time-Channel Separable Convolutional Neural Network Architecture for Speech Commands Recognition,”  in Proc. Interspeech 2020, 2020, pp. 3356–3360.
  * [30] Deokjin Seo, Heung-Seon Oh, and Yuchul Jung,  “Wav2kws: Transfer learning from speech representations for keyword spotting,”  IEEE Access, vol. 9, pp. 80682–80691, 2021.
  * [31] Yuan Gong, Yu-An Chung, and James Glass,  “Psla: Improving audio event classification with pretraining, sampling, labeling, and aggregation,”  arXiv preprint arXiv:2102.01243, 2021.
  * [32] Dan Hendrycks and Kevin Gimpel,  “Gaussian error linear units (gelus),”  arXiv preprint arXiv:1606.08415, 2016.