theo]Corollary theo]Lemma

# Layer-wise and Dimension-wise Locally Adaptive Federated Learning

Belhal Karimi, Ping Li, Xiaoyun Li

Cognitive Computing Lab
Baidu Research
10900 NE 8th St. Bellevue, WA 98004

{belhal.karimi, pingli98<EMAIL_ADDRESS>
###### Abstract

In the emerging paradigm of Federated Learning (FL), large amount of clients
such as mobile devices are used to train possibly high-dimensional models on
their respective data. Combining (dimension-wise) adaptive gradient methods
(e.g. Adam, AMSGrad) with FL has been an active direction, which is shown to
outperform traditional SGD based FL in many cases. In this paper, we focus on
the problem of training federated deep neural networks, and propose a novel FL
framework which further introduces _layer-wise_ adaptivity to the local model
updates. Our framework can be applied to locally adaptive FL methods including
two recent algorithms, Mime (Karimireddy et al., 2020) and Fed-AMS (Chen et
al., 2020). Theoretically, we provide a convergence analysis of our layer-wise
FL methods, coined Fed-LAMB and Mime-LAMB, which matches the convergence rate
of state-of-the-art results in FL and exhibits linear speedup in terms of the
number of workers. Experimental results on various datasets and models, under
both IID and non-IID local data settings, show that both Fed-LAMB and Mime-
LAMB achieve faster convergence speed and better generalization performance,
compared to the various recent adaptive FL methods.

##  1 Introduction

A growing and important task while learning models on observed data, is the
ability to train over a large number of clients which could either be personal
devices or distinct entities. In the paradigm of Federated Learning (FL)
(Konečnỳ et al., 2016; McMahan et al., 2017), a central server orchestrates
the optimization over those clients under the constraint that the data can
neither be gathered nor shared among the clients. This is computationally more
efficient, since more distributed computing resources are used; also, this is
a very practical scenario which allows individual data holders (e.g., mobile
devices) to train a model jointly without leaking private data. In this paper,
we consider the following optimization problem:

$\min_{\theta}f(\theta)\mathrel{\mathop{:}}=\frac{1}{n}\sum_{i=1}^{n}f_{i}(\theta)=\frac{1}{n}\sum_{i=1}^{n}\mathbb{E}_{\xi\sim\mathcal{X}_{i}}[F_{i}(\theta;\xi)],$
(1)

where the nonconvex function (e.g., deep networks) $f_{i}$ represents the
average loss over the local data samples for worker $i\in\llbracket
n\rrbracket$, and $\theta\in\mathbb{R}^{d}$ the global model parameter.
$\mathcal{X}_{i}$ is the data distribution on each client $i$. While (1)
reminds that of standard distributed optimization, the principle and setting
of FL are different from the classical distributed paradigm: (i) Local
updates: FL allows clients to perform multiple updates on the local models
before the global aggregation; (ii) Data heterogeneity: in FL, the local data
distributions $\mathcal{X}_{i}$ are usually different across workers,
hindering the convergence of the global model. FL aims at finding a solution
of (1) in fewest number of communication rounds.

One of the most popular framework for FL is called Fed-SGD (McMahan et al.,
2017): we adopt multiple local Stochastic Gradient Descent (SGD) steps in each
device, send those local models to the server that computes the average over
the received local model parameters, and broadcasts it back to the devices.
Moreover, momentum can be added to local SGD training for faster convergence
and better learning performance (Yu et al., 2019). On the other hand, adaptive
gradient methods, e.g., Adam (Kingma and Ba, 2015), AMSGrad (Reddi et al.,
2018), have shown great success in many deep learning tasks. For instance, the
update rule of Adam reads as

$\displaystyle\theta_{t}=\theta_{t-1}-\frac{\alpha m_{t}}{\sqrt{v_{t}}},$
$\displaystyle m_{t}=\beta_{1}m_{t-1}+(1-\beta_{1})g_{t},$ (2) $\displaystyle
v_{t}=\beta_{2}v_{t-1}+(1-\beta_{2})g_{t}^{2},$

where $\alpha$ is the learning rate and $g_{t}$ is the gradient at time $t$.
We note that the effective learning rate of Adam is $\alpha/\sqrt{v}$, which
is different across dimensions, i.e., dimension-wise adaptive. Recently, we
have seen growing research efforts in the design of FL frameworks that adopt
adaptive gradient methods as the protocols for local model training instead of
SGD. Examples include federated AMSGrad (Fed-AMS) (Chen et al., 2020) and Mime
(Karimireddy et al., 2020) with Adam updates. Specifically, in both methods,
in each round the global server not only aggregates the local models, but also
broadcasts to the workers a “global” second moment estimation to reconcile the
dimension-wise adaptive learning rates across the clients. Therefore, this
step can be regarded as a natural mitigation to data heterogeneity, which is a
common and important practical scenario that affects the performance of FL
algorithms (Li et al., 2020a; Liang et al., 2019; Karimireddy et al., 2019).

In this paper, we focus on improving adaptive FL algorithms. For (single-
machine) training of deep neural networks using Adam, You et al. (2020)
proposed a layer-wise adjusted learning rate scheme called LAMB, where in each
update, the ratio $m_{t}/\sqrt{v_{t}}$ is further normalized by the weight of
the deep network, respectively for each layer. LAMB allows large-batch
training which could in particular speed up training large datasets and models
like ImageNet (Deng et al., 2009) and BERT (Devlin et al., 2019). Inspired by
the acceleration effect of LAMB, we propose an improved framework for locally
adaptive FL algorithms, integrating both _dimension-wise_ and _layer-wise_
adaptive learning rates in each device’s local update. More specifically, our
contributions are summarized as follows:

  * •

We develop Fed-LAMB and Mime-LAMB, two instances of our layer-wise adaptive
optimization framework for federated learning, following a principled layer-
wise adaptive strategy to accelerate the training of deep neural networks.

  * •

We show that our algorithm converges at the rate of
$\mathcal{O}\left(\frac{1}{\sqrt{n\mathsf{h}R}}\right)$ to a stationary point,
where $\mathsf{h}$ is the number of layers of the network, $n$ is the number
of clients and $R$ is the number of communication rounds. This matches the
convergence rate of LAMB, AMSGrad, as well as the state-of-the-art results in
federated learning. The theoretical communication efficiency matches that of
Fed-AMS (Chen et al., 2020).

  * •

We empirically compare several recent adaptive FL methods under both
homogeneous and heterogeneous data setting on various benchmark datasets. Our
results confirm the accelerated empirical convergence of Fed-LAMB and Mime-
LAMB over the baseline methods, including Fed-AMS and Mime. In addition, Fed-
LAMB and Mine-LAMB can also reach similar, or better, test accuracy than their
corresponding baselines.

##  2 Background and Related Work

Firstly, we summarize some relevant work on adaptive optimization, layer-wise
adaptivity and federated learning, which compose the key ingredients of our
algorithm.

Adaptive gradient methods. Adaptive methods have proven to be the spearhead
for many nonconvex optimization tasks. Gradient based optimization algorithms
alleviate the possibly high nonconvexity of the objective function by
adaptively updating each coordinate of their learning rate using past
gradients. Common used examples include RMSprop (Tieleman and Hinton, 2012),
Adadelta (Zeiler, 2012), Adam (Kingma and Ba, 2015), Nadam (Dozat, 2016) and
AMSGrad (Reddi et al., 2018). Their popularity owes to their great performance
in training deep neural networks. They generally combine the idea of
adaptivity from AdaGrad (Duchi et al., 2011; McMahan and Streeter, 2010), as
explained above, and the idea of momentum from Nesterov’s Method (Nesterov,
2004) or Heavy ball method (Polyak, 1964) using past gradients. AdaGrad
displays superiority when the gradient is sparse compared to other classical
methods (Duchi et al., 2011). Yet, when applying AdaGrad to train deep neural
networks, it is observed that the learning rate might decay too fast.
Consequently, (Kingma and Ba, 2015) developed Adam whose updating rule is
presented in (2). A variant, called AMSGrad described in Reddi et al. (2018),
forces $v$ to be monotone to fix the convergence issue. The convergence and
generalization of adaptive methods are studied in, e.g., (Zhou et al., 2018b;
Chen et al., 2019; Zhou et al., 2020). Li et al. (2022) extended adaptive
gradient method to distributed setting with communication-efficient gradient
compression.

Layer-wise Adaptivity. When training deep networks, in many cases the scale of
gradients differs a lot across the network layers. When we use the same
learning rate for the whole network, the update might be too preservative for
some specific layers (with large weights) which may slow down the convergence.
Based on this observation, You et al. (2018) proposed LARS, an extension of
SGD with layer-wise adjusted scaling, whose performance, however, is not
consistent across tasks. Later, You et al. (2020) proposed LAMB, an analogue
layer-wise adaptive variant of Adam. The update rule of LAMB for the $\ell$-th
layer of the network can be expressed as

$\displaystyle\theta_{t}^{\ell}=\theta_{t-1}^{\ell}-\frac{\alpha\|\theta_{t-1}^{\ell}\|}{\|\psi_{t}^{\ell}\|}\psi_{t}^{\ell},\
\text{with}\ \psi_{t}^{\ell}=m_{t}^{\ell}/\sqrt{v_{t}^{\ell}},$

where $m_{t}$ and $v_{t}$ are defined in (2). Intuitively, for the $\ell$-th
layer, when the gradient magnitude is too small compared to the scale of the
model parameter, we increase the effective learning rate to make the model
move sufficiently far. Theoretically, (You et al., 2020) showed that LAMB
achieves the same convergence rate as Adam; empirically, LAMB can
significantly accelerate the convergence of Adam, allowing the use of large
mini-batch size with fewer training iterations for large datasets.

Federated learning. An extension of the well known parameter server framework,
where a model is being trained on several servers in a distributed manner, is
called federated learning (FL) (Konečnỳ et al., 2016; McMahan et al., 2017)
which has seen many applications in various fields (Yang et al., 2019; Leroy
et al., 2019; Bonawitz et al., 2019; Niknam et al., 2020; Xu et al., 2021).
For Fed-SGD (where clients perform SGD-based updates), recent variants and
theoretical analysis on the convergence can be found in Yu et al. (2019);
Karimireddy et al. (2019); Khaled et al. (2020); Li et al. (2020b); Woodworth
et al. (2020); Wang et al. (2020).

Recently, several works have considered integrating adaptive gradient methods
with FL. Reddi et al. (2021) proposed Adp-Fed where the central server applies
Adam-type updates. However, the local clients still perform SGD updates. Chen
et al. (2020); Karimireddy et al. (2020) proposed Fed-AMS and Mime
respectively, to adopt Adam/AMSGrad at the client level. Both works mitigate
the influence of data heterogeneity by “sharing” the second moment $v$ (which
controls the effective learning rate): in Fed-AMS, a global $v$ is computed
and synchronized in each round by averaging the $v_{i}$’s, $i=1,...,n$ from
the local clients; in Mime, a global $v$ is directly calculated using full-
batches (averaged over all clients) and maintained at the central server.
Hence, Mime requires at least twice computation as Fed-AMS. On many tasks,
these methods outperform Fed-SGD and other popular methods like SCAFFOLD
(Karimireddy et al., 2019) and FedProx (Sahu et al., 2018).

##  3 Layer-wise Adaptive Federated Learning

Figure 1: Illustration of Fed-LAMB framework (Algorithm 1), with a three-layer
network and $\phi(x)=x$ as an example. The depth of each network layer
represents the norm of its weights. For device $i$ and each local iteration in
round $r$, the adaptive ratio of $j$-th layer $\psi_{r,i}^{j}$ is normalized
according to $\|\theta_{r,i}^{j}\|$, and then used for updating the local
model. At the end of each round $r$, client $i$ sends
$\theta_{r,i}=[\theta_{r,i}^{\ell}]_{\ell=1}^{\mathsf{h}}$ and $v_{r,i}$ to
the central server, which transmits back aggregated $\theta$ and $\hat{v}$ to
devices to complete a round of training.

In this section, we introduce our proposed federated learning framework,
admitting both dimension-wise adaptivity (of adaptive learning rate) and
layer-wise adaptivity (of layer-wise scaling). For conciseness, we mainly
consider AMSGrad (Reddi et al., 2018) as the prototype method as it enjoys
better theoretical convergence properties. We assume the loss function
$f(\cdot)$ is induced by a multi-layer structured neural network, which
includes a broad class of network architectures such as MLP, CNN, ResNet and
Transformers.

Notations. We denote by $\theta$ the vector of parameters taking values in
$\mathbb{R}^{p}$. Suppose the neural network has $\mathsf{h}$ layers, each
with size $p_{\ell}$ (thus, $p=\sum_{\ell=1}^{\mathsf{h}}p_{\ell}$). For each
layer $\ell\in\llbracket\mathsf{h}\rrbracket$, denote $\theta^{\ell}$ as the
sub-vector corresponding to the $\ell$-th layer. Let $R$ be the number of
communication rounds and $T$ be the number of local iterations per round.
$\theta_{r,i}^{\ell,t}$ is the model parameter of layer $\ell$ at round $r$,
local iteration $t$ and for worker $i$.

Algorithm 1 Fed-LAMB and Mime-LAMB

1:Input: parameter $0<\beta_{1},\beta_{2}<1$; learning rate $\alpha$; weight
decaying rate $\lambda\in[0,1]$.

2:Initialize: $\theta_{0,i}\in\Theta\subseteq\mathbb{R}^{d}$;
$m^{0}_{0,i}=\hat{v}^{0}_{0,i}=v^{0}_{0,i}=0$, $\forall i\in\llbracket
n\rrbracket$; $\bar{\theta}_{0}=\frac{1}{n}\sum_{i=1}^{n}\theta_{0,i}$;
$\hat{v}_{0}=\epsilon$

3:for $r=1$ to $R$ do

4: Sample a set of clients $D^{r}$

5: for parallel for device $i\in D^{r}$ do

6: Set $\theta_{r,i}^{0}=\bar{\theta}_{r-1}$, $m^{0}_{r,i}=m^{T}_{r-1,i}$ ,
$v^{0}_{r,i}=\hat{v}_{r-1}$

7: for $t=1$ to $T$ do

8: Sample a mini-batch from the local data

9: Compute stochastic gradient $g^{t}_{r,i}$ at $\theta_{r,i}^{t-1}$

10: $m^{t}_{r,i}=\beta_{1}m^{t-1}_{r,i}+(1-\beta_{1})g^{t}_{r,i}$

11: $v^{t}_{r,i}=\beta_{2}v^{t}_{r-1,i}+(1-\beta_{2})(g^{t}_{r,i})^{2}$

12: Compute the ratio $\psi_{r,i}^{t}=m^{t}_{r,i}/(\sqrt{\hat{v}_{r-1}})$.

13: Update local model for each layer $\ell\in\llbracket\mathsf{h}\rrbracket$:
$\theta_{r,i}^{\ell,t}=\theta_{r,i}^{\ell,t-1}-\frac{\alpha_{r}\phi(\|\theta_{r,i}^{\ell,t-1}\|)(\psi_{r,i}^{\ell,t}+\lambda\theta_{r,i}^{\ell,t-1})}{\|\psi_{r,i}^{\ell,t}+\lambda\theta_{r,i}^{\ell,t-1}\|}$
(3)

14: end for

15: Communicate
$\theta_{r,i}^{T}=[\theta_{r,i}^{\ell,T}]_{\ell=1}^{\mathsf{h}}$ to server

16: Communicate $v_{r,i}^{T}$ to server

17: Communicate $\nabla f_{i}(\bar{\theta}_{r-1})$ using full local data

18: end for

19: Server compute $\bar{\theta}_{r}=\frac{1}{|D^{r}|}\sum_{i\in
D^{r}}\theta_{r,i}^{T}$

20: Server compute $\hat{v}_{r}=\max(\hat{v}_{r-1},\frac{1}{|D^{r}|}\sum_{i\in
D^{r}}v^{T}_{r,i})$

21: Compute $\nabla f(\bar{\theta}_{r-1})=\frac{1}{|D^{r}|}\sum_{i\in
D_{r}}\nabla f_{i}(\bar{\theta}_{r-1})$

22: Compute $v_{r}=\beta_{2}v_{r-1}+(1-\beta_{2})\nabla
f(\bar{\theta}_{r-1})^{2})$

23: Update $\hat{v}_{r}=\max(\hat{v}_{r-1},v_{r}$)

24:end for

In general, our proposed algorithm can be viewed as a novel extension of LAMB
to the more complicated federated setting. Based on the two recent works
regarding locally adaptive FL mentioned above, we present the framework with
two instances, Fed-LAMB and Mime-LAMB, as summarized in Algorithm 1 and
depicted in Figure 1. We differentiate the steps of these two methods by
blue(Fed-LAMB)and red(Mime-LAMB)boxes surrounding the text. Both methods use
layer-wise adaptive LAMB for local updates (Line 13). The update in (3) on
local workers can be expressed as

$\displaystyle\theta\leftarrow\theta-\alpha\frac{\phi(\|\theta\|)}{\|\psi+\lambda\theta\|}(\psi+\lambda\theta),$

where $\phi(\cdot):\mathbb{R}_{+}\mapsto\mathbb{R}_{+}$ is a scaling function
(usually chosen to be the identity function in practice) and $\lambda$ is the
weight decay rate. The main difference between Fed-LAMB and Mime-LAMB is the
way the second moment $\hat{v}$ is synchronized, i.e., the dimension-wise
adaptive learning rate. Both methods maintain a global $\hat{v}$ at the
central server:

  * •

Fed-LAMB (Line 20): at the end of each round, clients $i$ communicates the
local $v_{i}$; the server updates the global $\hat{v}$ by max operation with
the averaged $v$, and sends back the $\hat{v}$.

  * •

Mime-LAMB (Line 21-23): in each round $r$, the client computes and transmits
the gradient at the global model $\bar{\theta}_{r}$ using full local data; the
server updates the global $v$ and $\hat{v}$ in the same manner as AMSGrad.

Conceptually, both approaches aim at alleviating the impact of data
heterogeneity by “globally” reconciling the adaptive learning rates. Note that
Mime-LAMB needs to calculate the gradients twice, leading to double the
computational cost as Fed-LAMB.

Data Heterogeneity: Dealing with data heterogeneity is an important topic in
federated learning. Works have been conducted (e.g., Karimireddy et al.
(2019)) on designing specific strategies to alleviate the negative influence
of non-IID data by using techniques like control variate. We note that our
scheme is, in some sense, naturally capable of balancing the heterogeneity in
different local data distributions. This is largely due to the “moment
sharing” steps in Algorithm 1), where the adaptive learning rates guided by
the second moment estimation are aggregated among clients periodically. In
Chen et al. (2020) and Karimireddy et al. (2020), the authors have shown that
Fed-AMS and Mime would perform much worse, or even diverge, without
aggregating and sharing the second moment $\hat{v}$ (please refer to the
papers for details). Intuitively, synchronizing $\hat{v}$ makes all the
clients “on the same pace” which is crucial for the convergence of locally
adaptive FL methods. We will provide more discussion on this moment
synchronization in the next section.

##  4 Theoretical Analysis

For conciseness, we summarize in Table 1 some important notations that will be
used in our analysis.

$R,T$ | $\mathrel{\mathop{:}}=$ |  Number of communications rounds and local iterations (resp.)
---|---|---
$n,D,i$ | $\mathrel{\mathop{:}}=$ |  Total number of clients, portion sampled uniformly and client index
$\mathsf{h},\ell$ | $\mathrel{\mathop{:}}=$ |  Total number of layers in the DNN and its index
$\phi(\cdot)$ | $\mathrel{\mathop{:}}=$ |  Scaling factor in Fed-LAMB update
$\bar{\theta}$ | $\mathrel{\mathop{:}}=$ |  Global model (after periodic averaging)
$\psi_{r,i}^{t}$ | $\mathrel{\mathop{:}}=$ |  ratio computed at round $r$, local iteration $t$ and for device $i$. $\psi_{r,i}^{\ell,t}$ denotes its component at layer $\ell$
Table 1: Summary of notations used in the paper.

We need the following analytical assumptions.

######  Assumption .

(Smoothness) For all $i\in\llbracket n\rrbracket$ and
$\ell\in\llbracket\mathsf{h}\rrbracket$, the local loss function is
$L_{\ell}$-smooth: $\left\|\nabla f_{i}(\theta^{\ell})-\nabla
f_{i}(\vartheta^{\ell})\right\|\leq
L_{\ell}\left\|\theta^{\ell}-\vartheta^{\ell}\right\|$.

######  Assumption .

(Unbiased and bounded gradient) The stochastic gradient is unbiased for
$\forall r,t,i$: $\mathbb{E}[g_{r,i}^{t}]=\nabla f_{i}(\theta_{r}^{t})$ and
bounded by $\left\|g_{r,i}^{t}\right\|\leq M$.

######  Assumption .

(Bounded variance) The stochastic gradient admits (_locally_)
$\mathbb{E}[|g_{r,i}^{j}-\nabla f_{i}(\theta_{r})^{j}|^{2}]<\sigma^{2}$, and
(_globally_) $\frac{1}{n}\sum_{i=1}^{n}||\nabla f_{i}(\theta_{r})-\nabla
f(\theta_{r})||^{2}]<G^{2}$.

Assumption 4 and Assumption 4 are commonly used in the analysis of adaptive
gradients methods (Reddi et al., 2018; Chen et al., 2019; Karimireddy et al.,
2020; Reddi et al., 2021), Assumption 4 characterizes the data heterogeneity
among local devices, and $G=0$ when local data are IID Following (You et al.,
2020), we use the following assumption on the scaling function $\phi$.

######  Assumption .

For any $a\in\mathbb{R}_{+}$, there exist $\phi_{m}>0,\phi_{M}>0$ such that
$\phi_{m}\leq\phi(a)\leq\phi_{M}$.

We now state our main result regarding the convergence rate of the proposed
Algorithm 1.

Under Assumption 4-Assumption 4, consider $\\{\overline{\theta_{r}}\\}_{r>0}$
obtained from Algorithm 1 with a constant learning rate $\alpha$. Let
$\lambda=0$. Then, for any round $R>0$, we have

$\displaystyle\frac{1}{R}\sum_{r=1}^{R}\mathbb{E}\left[\left\|\frac{\nabla
f(\overline{\theta_{r}})}{\hat{v}_{r}^{1/4}}\right\|^{2}\right]$
$\displaystyle\leq\sqrt{\frac{M^{2}p}{n}}\frac{\triangle}{\mathsf{h}\alpha
R}+\frac{4\alpha^{2}\overline{L}M^{2}(T-1)^{2}\phi_{M}^{2}(1-\beta_{2})p}{\sqrt{\epsilon}}$
$\displaystyle+4\alpha\frac{M^{2}}{\sqrt{\epsilon}}+\frac{\phi_{M}\sigma^{2}}{Rn}\sqrt{\frac{1-\beta_{2}}{M^{2}p}}+4\alpha\left[\phi_{M}^{2}\sqrt{M^{2}+p\sigma^{2}}\right]$
$\displaystyle+4\frac{\alpha^{2}\overline{L}}{\sqrt{\epsilon}}M^{2}(T-1)^{2}G^{2}(1-\beta_{2})p+4\alpha\left[\phi_{M}\frac{\mathsf{h}\sigma^{2}}{\sqrt{n}}\right],$
(4)

where
$\triangle=\mathbb{E}[f(\bar{\theta}_{1})]-\min\limits_{\theta\in\Theta}f(\theta)$
and $\overline{L}=\sum_{\ell=1}^{\mathsf{h}}L_{\ell}$.

Note that this result holds for both Fed-LAMB and Mime-LAMB variants. Also,
the manifestation of $p$ in the rate is because the variance bound is assumed
on each dimension in Assumption 4. This dependency on $p$ can be removed when
Assumption 4 is assumed globally, which is also common in optimization
literature.

Using a uniform bound on the moment $\|\hat{v}_{r}\|\leq M^{2}$ and by
choosing a suitable decreasing learning rate, we have the following simplified
statement.

Under the same setting as Theorem 4, with
$\alpha=\mathcal{O}(\frac{1}{\sqrt{\mathsf{h}R}})$, it holds that

$\displaystyle\frac{1}{R}\sum_{r=1}^{R}\mathbb{E}\left[\left\|\nabla
f(\overline{\theta_{r}})\right\|^{2}\right]\leq\mathcal{O}\left(\frac{\sqrt{p}}{\sqrt{n\mathsf{h}R}}+\frac{\sqrt{\mathsf{h}}\sigma^{2}}{\sqrt{nR}}+\frac{G^{2}(T-1)^{2}p}{R\mathsf{h}}\right).$
(5)

The leading two terms display a dependence of the convergence rate of Fed-LAMB
on the initialization and the local variance of the stochastic gradients
(Assumption 4). The last term involves the number of local updates $T$, and
the global variance $G^{2}$ characterizing the data heterogeneity. Next, we
provide detailed discussion and comparison of our bound to related prior
results.

LAMB bound in You et al. (2020):  We start our discussion with the comparison
of convergence rate of Fed-LAMB with that of LAMB, Theorem 3 in You et al.
(2020). In a single-machine setting, the convergence rate of LAMB is
$\mathcal{O}(\sqrt{p}{\sqrt{\mathsf{h}T}})$ where $T$ is the number of
training iterations. Note the convergence rate of Fed-LAMB is different from
that of LAMB in the sense that, the convergence criterion is given at the
averaged parameters (global model) at the end of each round. In Corollary 4,
our rate would match LAMB if we take number of local step $T=1$. This also
holds true for any fixed $T$ and $R$ sufficiently large. In addition, the
$\mathcal{O}(\frac{1}{\sqrt{nR}})$ rate of Fed-LAMB implies a linear speedup
effect: the number of iterations to reach a $\delta$-stationary point of Fed-
LAMB decreases linearly in $n$, which displays the merit of distributed
(federated) learning.

Fed-AMS bound in Chen et al. (2020):  We now compare our method theoretically
with Fed-AMS, the baseline distributed adaptive method developed in Chen et
al. (2020). Their results state that when $T\leq\mathcal{O}(R^{1/3})$, the
convergence rate of Fed-AMS is $\mathcal{O}(\sqrt{\frac{p}{nR}})$. Firstly,
when the number of rounds $R$ is sufficiently large, both our rate (5) and the
rate of Fed-AMS are dominated by $\mathcal{O}(\frac{\sqrt{p}}{\sqrt{nR}})$,
improving the convergence rate of the standard AMSGrad, e.g. (Zhou et al.,
2018a) by $\mathcal{O}(1/\sqrt{n})$. Secondly, in (5), the last term
containing the number of local updates $T$ is small as long as
$T^{4}\leq\mathcal{O}(\frac{Rh}{G^{2}})$. If we further assume $h\simeq T$,
then we get the same rate of convergence as Fed-AMS with
$T\leq\mathcal{O}(R^{1/3})$ local iterations, identical to the condition of
Fed-AMS. Moreover, when $G$ is small (less data heterogeneity), the bound on
$T$ would increase, i.e., we can conduct more local updates. This is
intuitive, for example, when $G=0$ in the IID data setting, $T$ can be very
large.

In general, similar to the result that single machine LAMB matches the
convergence rate of Adam theoretically, we show that Fed-LAMB achieves the
same convergence rate as Fed-AMS, in the federated (distributed) setting,
which is the first such result in literature for locally adaptive FL method.
Also similar to the findings in You et al. (2020), we will show that Fed-LAMB
and its variants provide impressive acceleration empirically, in our
experimental study (Section 5).

Practical Considerations: Under our analytic setting and conditions mentioned
above, the convergence rate of Fed-LAMB also matches many popular federated
learning methods in nonconvex optimization, e.g., Fed-SGD (McMahan et al.,
2017), Mime (Karimireddy et al., 2020) and Adp-Fed (Reddi et al., 2021), at
$\mathcal{O}(\frac{1}{\sqrt{nRT}})$. In practice, when trained with same
number of rounds and $R$ and local iterations $T$, Mime, Fed-LAMB and Fed-Mime
all require communicating two tensors, while Fed-SGD, Adp-Fed only communicate
one tensor (the local model parameter). Next, we discuss a simple
implementation trick of our algorithm that leads to less communication.

Extension: lazy synchronization of $\hat{v}_{t}$. As mentioned before, the
synchronization of the second moment $\hat{v}$ local workers keep in step with
each other. Yet, may not need to update and broadcast the global $\hat{v}$ in
every round. To reduce the extra communication overhead of transmitting
$\hat{v}$, one trick in practice is to reduce the aggregation frequency of
$\hat{v}$ (e.g., we synchronize $\hat{v}$ every $Z$ rounds). It can be shown
that this “lazy” aggregation does not affect the convergence rate of our Fed-
LAMB. Yet it can effectively reduce the communication of $\hat{v}$ by a factor
of $Z$, which to a great extent alleviates the extra communication cost of
locally adaptive methods. We will also show empirical evidence of this trick
in our experiments.

##  5 Experiments

In this section, we conduct experiments on benchmark datasets with various
network architectures to justify the effectiveness of our proposed method in
practice. Our method empirically confirms its merit in terms of convergence
speed. Basically, Fed-LAMB and Mime-LAMB reduce the number of rounds and thus
the communication cost required to achieve a similar stationary point (or test
accuracy) than the baseline methods. In many cases, Fed-LAMB also brings
notable improvement in generalization over baselines.

Methods. We evaluate the following five FL algorithms, mainly focusing on
recent federated optimization approaches based on adaptive methods:

  1. 1.

Fed-SGD (McMahan et al., 2017), standard federated averaging with local SGD
updates.

  2. 2.

Adp-Fed (_Adaptive Federated Optimization_ , see Appendix A), the federated
adaptive algorithm proposed by (Reddi et al., 2021). Adp-Fed performs local
SGD updates. In each round $r$, the changes in local models,
$\triangle_{i}=w_{r,i}^{T}-w_{r,i}^{0}$, $i=1,...,n$, are sent to the central
server for an aggregated Adam update.

  3. 3.

Fed-AMS (Chen et al., 2020), locally adaptive AMSGrad.

  4. 4.

Mime (Karimireddy et al., 2020) with AMSGrad, which performs adaptive local
updates with central-server-guided global adaptive learning rate.

  5. 5.

Our proposed Fed-LAMB and Mime-LAMB (Algorithm 1).

For all the adaptive gradient methods, we set $\beta_{1}=0.9$,
$\beta_{2}=0.999$ by default (Reddi et al., 2018). We present the results of
$n=50$ clients with $0.5$ participation rate, i.e., we randomly pick half of
the clients to be active for training in each round, and the local mini-batch
size is set as 128. In each round, the training samples are allocated to the
active devices, and one local epoch is completed after all the local devices
run one pass over their received samples via mini-batch training. Example
results with more clients can be found in Figure 6.

We tune the initial learning rate $\alpha$ for each algorithm over a fine
grid. For Adp-Fed, there are two learning rates involved (global and local),
both of which are tuned. More tuning details can be found in Appendix A. For
Fed-LAMB and Mime-LAMB, the weight decay rate $\lambda$ is tuned from
$\\{0,0.01,0.1\\}$, and we use the identity scaling function for
$\phi(\cdot)$. For each run, we report the best test accuracy. The results are
averaged over 3 runs, each with same initialization for every method.

Datasets and models. We experiment with four popular benchmark image
classification datasets: MNIST (LeCun, 1998), Fashion MNIST (FMNIST) (Xiao et
al., 2017), CIFAR-10 (Krizhevsky, 2009) and TinyImageNet (Deng et al., 2009).
For MNIST, we apply 1) a simple multi-layer perceptron (MLP), which has one
hidden layer containing 200 cells; 2) Convolutional Neural Network (CNN),
which has two max-pooled convolutional layers followed by a dropout layer and
two fully-connected layers with 320 and 50 cells respectively. This CNN is
also implemented for FMNIST. For CIFAR-10 and TinyImageNet, we use ResNet-18
(He et al., 2016).

###  5.1 Comparison under IID settings

Figure 2: IID data setting. Test accuracy against the number of communication
rounds.

In Figure 2, we report the test accuracy of MLP trained on MNIST, as well as
CNN trained on MNIST and FMNIST, where the data are IID allocated among the
clients. We test 1 local epoch and 3 local epochs (more local iterations). In
all the figures, we observe a clear advantage of Fed-LAMB over the competing
methods in terms of the convergence speed. In particular, we can see that Fed-
LAMB is able to achieve the same accuracy with fewest number of communication
rounds, thus improving the communication efficiency. For instance, this can be
observed as follows: on MNIST + CNN (1 local epoch), Fed-AMS requires 20
rounds to achieves 90% accuracy, while Fed-LAMB only takes 5 rounds. This
implies a 75% reduction in the communication cost. Moreover, on MNIST, Fed-
LAMB also leads to improved generalisation performance, i.e., test accuracy.
We can draw same conclusions with 3 local epochs. Also, similar comparison
holds for Mime-LAMB vs. Mime. In general, the Mime-LAMB variant matches the
performance of Fed-LAMB closely.

###  5.2 Comparison under non-IID settings

Figure 3: non-IID data setting. Test accuracy against the number of
communication rounds.

In Figure 3, we provide the results on MNIST and FMNIST with non-IID local
data distribution. In particular, in each round of federated training, every
local device only receives samples from one or two class (out of ten). We see
that for experiments with 1 local epoch, in all cases our proposed Fed-LAMB
outperforms all the baseline methods. Similar to the IID data setting, Fed-
LAMB provides faster convergence speed and achieves higher test accuracy than
Fed-SGD and Fed-AMS. The advantage is especially significant for the CNN
model, e.g., it improves the accuracy of Fed-SGD and Fed-AMS by more than 10%
on FMNIST at the 50-th round. The other baseline method, Adp-Fed, performs as
good as our Fed-LAMB on FMNIST, but worse than other methods on MNIST. Mime-
LAMB also considerably improves Mime on all the tasks, see Figure 3. In
general, the two variants Fed-LAMB and Mime-LAMB perform similarly.

The relative comparison is basically the same when the local models are
trained for 3 epochs before model aggregation. The proposed methods converge
faster than the underlying baselines. Adp-Fed performs well on FMNIST but
worse on MNIST. Yet, we notice that the advantage of Fed-LAMB becomes less
significant than what we observed in Figure 2 with IID data. One plausible
reason is that when the local data is highly non-IID, the local learning
objective functions become very different. Intuitively, with more local steps,
learning the local models fast might not always do good to the global model,
as local models target at different loss functions. This is consistent with
extensive empirical observations in literature that in FL, allowing too many
local training steps usually pushes local models too far which may hinder the
convergence of the global model.

Figure 4: non-IID data. Test accuracy of CIFAR-10 and TinyImagenet on
ResNet-18.

In Figure 4, we present the results on CIFAR-10 and TinyImageNet datasets
trained by ResNet-18. When training these two models, we decrease the learning
rate to $1/10$ at the 30-th and 70-th communication round. From Figure 4, we
can draw similar conclusion as before: the proposed Fed-LAMB is the best
method in terms of both convergence speed and generalization accuracy. In
particular, on TinyImageNet, we see that Fed-LAMB has a significant advantage
over all three baselines. Although Adp-Fed performs better than Fed-SGD and
Fed-AMS, it is considerably worse than Fed-LAMB. We report the test accuracy
at the end of training in Table 2. Fed-LAMB achieves the highest accuracy on
both datasets. Mime-LAMB also substantially improves Mime.

| Fed-SGD | Adp-Fed | Fed-AMS | Fed-LAMB | Mime | Mime-LAMB
---|---|---|---|---|---|---
CIFAR-10 | 90.75 $\pm$ 0.48 | 91.57 $\pm$ 0.38 | 90.93 $\pm$ 0.22 | 92.44 $\pm$ 0.53 | 90.94 $\pm$ 0.13 | 92.00 $\pm$ 0.21
TinyImageNet | 67.58 $\pm$ 0.21 | 74.17 $\pm$ 0.43 | 64.86 $\pm$ 0.83 | 76.00 $\pm$ 0.26 | 67.82 $\pm$ 0.24 | 73.46 $\pm$ 0.25

Table 2: Test accuracy with ResNet-18 network after 100 communication rounds.

Figure 5: non-IID data. Fed-LAMB and Mime-Fed with lazy synchronization of
$\hat{v}$.

In addition, in Figure 5 we present the result of our methods with lazy
synchronization of $\hat{v}$, where the server updates and broadcasts
$\hat{v}$ every $Z=3,5$ rounds, instead of very single round. We see that
practically the performance is similar to the standard algorithm with every-
round global $\hat{v}$ updates (sometimes even slightly better).

More Workers: In Figure 6, we provide additional figures with larger number of
workers $n=200$, on MNIST and FMNIST with non-IID data. The conclusions stay
the same: we see that the proposed Fed-LAMB and Mime-LAMB perform much better
than the baseline algorithms, with faster convergence and better accuracy at
the end of 100 FL training rounds.

Figure 6: Test accuracy with $n=200$ workers, full participation, local batch
size 64. Data are non-IID distributed among clients.

###  5.3 Summary of empirical findings

We provide a brief summary of this section. On all datasets, the primary
comparison of most importance appears evident:

Fed-LAMB$\approx$ Mime-LAMB$>$Fed-AMS$\approx$Mime.

The proposed scheme (with two variants) exhibits faster convergence and better
generalisation accuracy than recently proposed adaptive FL algorithms. Our
results suggest that, similar as in the single-machine training, layer-wise
acceleration can also be effective and beneficial in federated learning.
Moreover, in practice we may adopt the lazy aggregation trick to further
reduce the additional communication required for Fed-LAMB.

##  6 Conclusion

We study a doubly adaptive method in the particular framework of federated
learning (FL). Built upon the acceleration effect of layer-wise learning rate
scheduling and of state-of-the-art adaptive gradient methods, we derive a
locally layer-wise FL framework that performs local updates using adaptive
AMSGrad on each worker and periodically averages local models stored on each
device. The core of our Fed-LAMB scheme, is to speedup up local training by
adopting layer-wise adaptive learning rates. To out knowledge, this is the
first FL algorithm in literature that possess both the _dimension-wise_
adaptivity (by AMSGrad) and _layer-wise_ adaptivity (by layer-wise learning
rate). We provide the convergence analysis of Fed-LAMB that matches many
existing methods, with a linear speedup against the number of clients.
Extensive experiments on various datasets and models, under both IID and non-
IID data settings, validate that both Fed-LAMB and Mime-LAMB are able to
provide faster convergence which in turn could lead to reduced communication
cost. In many cases, our framework also improves the overall performance of
federated learning over prior methods.

## References

  * Bonawitz et al. [2019] Kallista A. Bonawitz, Hubert Eichner, Wolfgang Grieskamp, Dzmitry Huba, Alex Ingerman, Vladimir Ivanov, Chloé Kiddon, Jakub Konečný, Stefano Mazzocchi, Brendan McMahan, Timon Van Overveldt, David Petrou, Daniel Ramage, and Jason Roselander.  Towards federated learning at scale: System design.  In _Proceedings of Machine Learning and Systems (MLSys)_ , Stanford, CA, 2019.
  * Chen et al. [2019] Xiangyi Chen, Sijia Liu, Ruoyu Sun, and Mingyi Hong.  On the convergence of A class of adam-type algorithms for non-convex optimization.  In _Proceedings of the 7th International Conference on Learning Representations (ICLR)_ , New Orleans, LA, 2019.
  * Chen et al. [2020] Xiangyi Chen, Xiaoyun Li, and Ping Li.  Toward communication efficient adaptive gradient method.  In _Proceedings of the ACM-IMS Foundations of Data Science Conference (FODS)_ , pages 119–128, Virtual Event, USA, 2020.
  * Deng et al. [2009] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei.  ImageNet: A large-scale hierarchical image database.  In _Proceedings of the 2009 IEEE Computer Society Conference on Computer Vision and Pattern Recognition (CVPR)_ , pages 248–255, Miami, FL, 2009.
  * Devlin et al. [2019] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.  BERT: pre-training of deep bidirectional transformers for language understanding.  In _Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (NAACL-HLT)_ , pages 4171–4186, Minneapolis, MN, 2019.
  * Dozat [2016] Timothy Dozat.  Incorporating nesterov momentum into Adam.  In _Proceedings of the 4th International Conference on Learning Representations (ICLR Workshop)_ , San Juan, Puerto Rico, 2016.
  * Duchi et al. [2011] John C. Duchi, Elad Hazan, and Yoram Singer.  Adaptive subgradient methods for online learning and stochastic optimization.  _J. Mach. Learn. Res._ , 12:2121–2159, 2011.
  * He et al. [2016] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Deep residual learning for image recognition.  In _Proceedings of the 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR)_ , pages 770–778, Las Vegas, NV, 2016.
  * Karimireddy et al. [2019] Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank J Reddi, Sebastian U Stich, and Ananda Theertha Suresh.  Scaffold: Stochastic controlled averaging for on-device federated learning.  _arXiv preprint arXiv:1910.06378_ , 2019.
  * Karimireddy et al. [2020] Sai Praneeth Karimireddy, Martin Jaggi, Satyen Kale, Mehryar Mohri, Sashank J Reddi, Sebastian U Stich, and Ananda Theertha Suresh.  Mime: Mimicking centralized stochastic algorithms in federated learning.  _arXiv preprint arXiv:2008.03606_ , 2020.
  * Khaled et al. [2020] Ahmed Khaled, Konstantin Mishchenko, and Peter Richtárik.  Tighter theory for local SGD on identical and heterogeneous data.  In _Proceedings of the 23rd International Conference on Artificial Intelligence and Statistics (AISTATS)_ , pages 4519–4529, Online [Palermo, Sicily, Italy], 2020.
  * Kingma and Ba [2015] Diederik P. Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  In _Proceedings of the 3rd International Conference on Learning Representations (ICLR)_ , San Diego, CA, 2015.
  * Konečnỳ et al. [2016] Jakub Konečnỳ, H Brendan McMahan, Felix X Yu, Peter Richtárik, Ananda Theertha Suresh, and Dave Bacon.  Federated learning: Strategies for improving communication efficiency.  _arXiv preprint arXiv:1610.05492_ , 2016.
  * Krizhevsky [2009] Alex Krizhevsky.  Learning multiple layers of features from tiny images.  _Master’s thesis, Department of Computer Science, University of Toronto_ , 2009.
  * LeCun [1998] Yann LeCun.  The mnist database of handwritten digits.  _http://yann. lecun. com/exdb/mnist/_ , 1998.
  * Leroy et al. [2019] David Leroy, Alice Coucke, Thibaut Lavril, Thibault Gisselbrecht, and Joseph Dureau.  Federated learning for keyword spotting.  In _Proceedings of the IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)_ , pages 6341–6345, Brighton, UK, 2019.
  * Li et al. [2020a] Tian Li, Anit Kumar Sahu, Ameet Talwalkar, and Virginia Smith.  Federated learning: Challenges, methods, and future directions.  _IEEE Signal Process. Mag._ , 37(3):50–60, 2020a.
  * Li et al. [2020b] Xiang Li, Kaixuan Huang, Wenhao Yang, Shusen Wang, and Zhihua Zhang.  On the convergence of fedavg on non-iid data.  In _Proceedings of the 8th International Conference on Learning Representations (ICLR)_ , Addis Ababa, Ethiopia, 2020b.
  * Li et al. [2022] Xiaoyun Li, Belhal Karimi, and Ping Li.  On distributed adaptive optimization with gradient compression.  In _Proceedings of the 10th International Conference on Learning Representations (ICLR)_ , Virtual Event, 2022.
  * Liang et al. [2019] Xianfeng Liang, Shuheng Shen, Jingchang Liu, Zhen Pan, Enhong Chen, and Yifei Cheng.  Variance reduced local sgd with lower communication complexity.  _arXiv preprint arXiv:1912.12844_ , 2019.
  * McMahan and Streeter [2010] Brendan McMahan and Matthew J. Streeter.  Adaptive bound optimization for online convex optimization.  In _Proceedings of the 23rd Conference on Learning Theory (COLT)_ , pages 244–256, Haifa, Israel, 2010.
  * McMahan et al. [2017] Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Agüera y Arcas.  Communication-efficient learning of deep networks from decentralized data.  In _Proceedings of the 20th International Conference on Artificial Intelligence and Statistics (AISTATS)_ , pages 1273–1282, Fort Lauderdale, FL, 2017.
  * Nesterov [2004] Yurii Nesterov.  Introductory lectures on convex optimization: A basic course.  _Springer_ , 2004.
  * Niknam et al. [2020] Solmaz Niknam, Harpreet S. Dhillon, and Jeffrey H. Reed.  Federated learning for wireless communications: Motivation, opportunities, and challenges.  _IEEE Commun. Mag._ , 58(6):46–51, 2020.
  * Polyak [1964] B. T. Polyak.  Some methods of speeding up the convergence of iteration methods.  _Mathematics and Mathematical Physics_ , 1964.
  * Reddi et al. [2018] Sashank J. Reddi, Satyen Kale, and Sanjiv Kumar.  On the convergence of adam and beyond.  In _Proceedings of the 6th International Conference on Learning Representations (ICLR)_ , Vancouver, Canada, 2018.
  * Reddi et al. [2021] Sashank J. Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Konečný, Sanjiv Kumar, and Hugh Brendan McMahan.  Adaptive federated optimization.  In _Proceedings of the 9th International Conference on Learning Representations (ICLR)_ , Virtual Event, Austria, 2021.
  * Sahu et al. [2018] Anit Kumar Sahu, Tian Li, Maziar Sanjabi, Manzil Zaheer, Ameet Talwalkar, and Virginia Smith.  On the convergence of federated optimization in heterogeneous networks.  _arXiv preprint arXiv:1812.06127_ , 2018.
  * Tieleman and Hinton [2012] T. Tieleman and G. Hinton.  Rmsprop: Divide the gradient by a running average of its recent magnitude.  _COURSERA: Neural Networks for Machine Learning_ , 2012.
  * Wang et al. [2020] Jianyu Wang, Vinayak Tantia, Nicolas Ballas, and Michael G. Rabbat.  Slowmo: Improving communication-efficient distributed SGD with slow momentum.  In _Proceedings of the 8th International Conference on Learning Representations (ICLR)_ , Addis Ababa, Ethiopia, 2020.
  * Woodworth et al. [2020] Blake E. Woodworth, Kumar Kshitij Patel, Sebastian U. Stich, Zhen Dai, Brian Bullins, H. Brendan McMahan, Ohad Shamir, and Nathan Srebro.  Is local SGD better than minibatch sgd?  In _Proceedings of the 37th International Conference on Machine Learning (ICML)_ , pages 10334–10343, Virtual Event, 2020.
  * Xiao et al. [2017] Han Xiao, Kashif Rasul, and Roland Vollgraf.  Fashion-MNIST: a novel image dataset for benchmarking machine learning algorithms.  _arXiv preprint arXiv:1708.07747_ , 2017.
  * Xu et al. [2021] Jie Xu, Benjamin S. Glicksberg, Chang Su, Peter B. Walker, Jiang Bian, and Fei Wang.  Federated learning for healthcare informatics.  _J. Heal. Informatics Res._ , 5(1):1–19, 2021\.
  * Yang et al. [2019] Qiang Yang, Yang Liu, Tianjian Chen, and Yongxin Tong.  Federated machine learning: Concept and applications.  _ACM Trans. Intell. Syst. Technol._ , 10(2):12:1–12:19, 2019.
  * You et al. [2018] Yang You, Zhao Zhang, Cho-Jui Hsieh, James Demmel, and Kurt Keutzer.  Imagenet training in minutes.  In _Proceedings of the 47th International Conference on Parallel Processing (ICPP)_ , pages 1:1–1:10, Eugene, OR, 2018.
  * You et al. [2020] Yang You, Jing Li, Sashank J. Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsieh.  Large batch optimization for deep learning: Training BERT in 76 minutes.  In _Proceedings of the 8th International Conference on Learning Representations (ICLR)_ , Addis Ababa, Ethiopia, 2020.
  * Yu et al. [2019] Hao Yu, Rong Jin, and Sen Yang.  On the linear speedup analysis of communication efficient momentum SGD for distributed non-convex optimization.  In _Proceedings of the 36th International Conference on Machine Learning (ICML)_ , pages 7184–7193, Long Beach, CA, 2019.
  * Zeiler [2012] Matthew D Zeiler.  Adadelta: an adaptive learning rate method.  _arXiv preprint arXiv:1212.5701_ , 2012.
  * Zhou et al. [2018a] Dongruo Zhou, Jinghui Chen, Yuan Cao, Yiqi Tang, Ziyan Yang, and Quanquan Gu.  On the convergence of adaptive gradient methods for nonconvex optimization.  _arXiv preprint arXiv:1808.05671_ , 2018a.
  * Zhou et al. [2018b] Dongruo Zhou, Jinghui Chen, Yuan Cao, Yiqi Tang, Ziyan Yang, and Quanquan Gu.  On the convergence of adaptive gradient methods for nonconvex optimization.  _arXiv preprint arXiv:1808.05671_ , 2018b.
  * Zhou et al. [2020] Yingxue Zhou, Belhal Karimi, Jinxing Yu, Zhiqiang Xu, and Ping Li.  Towards better generalization of adaptive gradient methods.  In _Advances in Neural Information Processing Systems (NeurIPS)_ , virtual, 2020.

##  Appendix A Hyper-parameter Tuning and Algorithms

###  A.1 The Adp-Fed Algorithm [Reddi et al., 2021]

The Adp-Fed (Adaptive Federated Optimization) is one of the baseline methods
compared with Fed-LAMB in our paper. The algorithm is given in Algorithm 2.
The key difference between Adp-Fed and Fed-AMS [Chen et al., 2020] is that, in
Adp-Fed, each client runs local SGD (Line 8), and an Adam optimizer is
maintained for the global adaptive optimization (Line 15). In the Fed-AMS
framework (as well as our Fed-LAMB), each clients runs local (adaptive)
AMSGrad method, and the global model is simply obtained by averaging the local
models.

Algorithm 2 Adp-Fed: Adaptive Federated Optimization [Reddi et al., 2021]

1:Input: parameter $0<\beta_{1},\beta_{2}<1$, and learning rate $\alpha_{t}$,
weight decaying parameter $\lambda\in[0,1]$.

2:Initialize: $\theta_{0,i}\in\Theta\subseteq\mathbb{R}^{d}$, $m_{0}=0$,
$v_{0}=\epsilon$, $\forall i\in\llbracket n\rrbracket$, and
$\theta_{0}=\frac{1}{n}\sum_{i=1}^{n}\theta_{0,i}$.

3:for $r=1,\ldots,R$ do

4: parallel for device $i$ do:

5: Set $\theta_{r,i}^{0}=\theta_{r-1}$.

6: for $t=1,\ldots,T$ do

7: Compute stochastic gradient $g^{t}_{r,i}$ at $\theta_{r,i}^{0}$.

8: $\theta_{r,i}^{t}=\theta_{r,i}^{t-1}-\eta_{l}g_{r,i}^{t}$

9: end for

10: Devices send $\triangle_{r,i}=\theta_{r,i}^{T}-\theta_{r,i}^{0}$ to
server.

11: end for

12: Server computes
$\bar{\triangle}_{r}=\frac{1}{n}\sum_{i=1}^{n}\triangle_{r,i}$

13: $m_{r}=\beta_{1}m_{r-1}+(1-\beta_{1})\bar{\triangle}_{r}$

14: $v_{r}=\beta_{2}v_{r-1}+(1-\beta_{2})\bar{\triangle}_{r}^{2}$

15: $\theta_{r}=\theta_{r-1}+\eta_{g}\frac{m_{r}}{\sqrt{v_{r}}}$

16:end for

17:Output: Global model parameter $\theta_{R}$.

###  A.2 Hyper-parameter Tuning

In our empirical study, we tune the learning rate of each algorithm carefully
such that the best performance is achieved. The search grids in all our
experiments are provided in Table 3.

Table 3: Search grids of the learning rate. | Learning rate range
---|---
Fed-SGD | $[0.001,0.003,0.005,0.01,0.03,0.05,0.1,0.3,0.5]$
Fed-AMS | $[0.0001,0.0003,0.0005,0.001,0.003,0.005,0.01,0.03,0.05,0.1]$
Fed-LAMB | $[0.001,0.003,0.005,0.01,0.03,0.05,0.1,0.3,0.5]$
Adp-Fed | Local $\eta_{l}$: $[0.0001,0.0003,0.0005,0.001,0.003,0.005,0.01,0.03,0.05,0.1,0.3,0.5]$
Global $\eta_{g}$:
$[0.0001,0.0003,0.0005,0.001,0.003,0.005,0.01,0.03,0.05,0.1]$
Mime | $[0.0001,0.0003,0.0005,0.001,0.003,0.005,0.01,0.03,0.05,0.1]$
Mime-LAMB | $[0.001,0.003,0.005,0.01,0.03,0.05,0.1,0.3,0.5]$

##  Appendix B Theoretical Analysis

###  B.1 Intermediary Lemma

We now develop the proof of the convergence rate of Fed-LAMB. We need a
supporting Lemma B.1 for this.

Consider $\\{\overline{\theta_{r}}\\}_{r>0}$, the sequence of parameters
obtained running Algorithm 1. Then for $i\in\llbracket n\rrbracket$:

$\|\overline{\theta_{r}}-\theta_{r,i}\|^{2}\leq\alpha^{2}M^{2}\phi_{M}^{2}\frac{(1-\beta_{2})p}{\epsilon}\;,$

where $\phi_{M}$ is defined in Assumption 4 and p is the total number of
dimensions $p=\sum_{\ell=1}^{\mathsf{h}}p_{\ell}$.

###### Proof.

Assuming the simplest case when $T=1$, i.e., one local iteration, then by
construction of Algorithm 1, we have for all
$\ell\in\llbracket\mathsf{h}\rrbracket$, $i\in\llbracket n\rrbracket$ and
$r>0$:

$\theta^{\ell}_{r,i}=\overline{\theta_{r}}^{\ell}-\alpha\phi(\|\theta_{r,i}^{\ell,t-1}\|)\psi_{r,i}^{j}/\|\psi_{r,i}^{\ell}\|=\overline{\theta_{r}}^{\ell}-\alpha\phi(\|\theta_{r,i}^{\ell,t-1}\|)\frac{m^{t}_{r,i}}{\sqrt{v^{t}_{r}}}\frac{1}{\|\psi_{r,i}^{\ell}\|}$

leading to

$\begin{split}\|\overline{\theta_{r}}-\theta_{r,i}\|^{2}=\sum_{\ell=1}^{\mathsf{h}}\left\langle\overline{\theta_{r}}^{\ell}-\theta^{\ell}_{r,i}\,|\,\overline{\theta_{r}}^{\ell}-\theta^{\ell}_{r,i}\right\rangle\leq\alpha^{2}M^{2}\phi_{M}^{2}\frac{(1-\beta_{2})p}{\epsilon}\;,\end{split}$

which concludes the proof. ∎

###  B.2 Proof of Theorem 4

######  Theorem.

Suppose Assumption 4-Assumption 4 hold. Consider
$\\{\overline{\theta_{r}}\\}_{r>0}$, the sequence of parameters obtained
running Algorithm 1 with a constant learning rate $\alpha$. Let the number of
local epochs be $T\geq 1$ and $\lambda=0$. Then, for any round $R>0$, we have

$\displaystyle\frac{1}{R}\sum_{r=1}^{R}\mathbb{E}\left[\left\|\frac{\nabla
f(\overline{\theta_{r}})}{\hat{v}_{r}^{1/4}}\right\|^{2}\right]$
$\displaystyle\leq\sqrt{\frac{M^{2}p}{n}}\frac{\triangle}{\mathsf{h}\alpha
R}+\frac{4\alpha\alpha^{2}LM^{2}(T-1)^{2}\phi_{M}^{2}(1-\beta_{2})p}{\sqrt{\epsilon}}$
(6)
$\displaystyle+4\alpha\frac{M^{2}}{\sqrt{\epsilon}}+\frac{\phi_{M}\sigma^{2}}{Rn}\sqrt{\frac{1-\beta_{2}}{M^{2}p}}+4\alpha\left[\phi_{M}\frac{\mathsf{h}\sigma^{2}}{\sqrt{n}}\right]+4\alpha\left[\phi_{M}^{2}\sqrt{M^{2}+p\sigma^{2}}\right],$

where
$\triangle=\mathbb{E}[f(\bar{\theta}_{1})]-\min\limits_{\theta\in\Theta}f(\theta)$.

###### Proof.

Using Assumption 4, we have

$\displaystyle f(\bar{\vartheta}_{r+1})$ $\displaystyle\leq
f(\bar{\vartheta}_{r})+\left\langle\nabla
f(\bar{\vartheta}_{r})\,|\,\bar{\vartheta}_{r+1}-\bar{\vartheta}_{r}\right\rangle+\sum_{\ell=1}^{L}\frac{L_{\ell}}{2}\|\bar{\vartheta}^{\ell}_{r+1}-\bar{\vartheta}^{\ell}_{r}\|^{2}$
$\displaystyle\leq
f(\bar{\vartheta}_{r})+\sum_{\ell=1}^{\mathsf{h}}\sum_{j=1}^{p_{\ell}}\nabla_{\ell}f(\bar{\vartheta}_{r})^{j}(\bar{\vartheta}^{\ell,j}_{r+1}-\bar{\vartheta}^{\ell,j}_{r})+\sum_{\ell=1}^{L}\frac{L_{\ell}}{2}\|\bar{\vartheta}^{\ell}_{r+1}-\bar{\vartheta}^{\ell}_{r}\|^{2}\;.$

Taking expectations on both sides leads to

$\displaystyle-\mathbb{E}[\left\langle\nabla
f(\bar{\vartheta}_{r})\,|\,\bar{\vartheta}_{r+1}-\bar{\vartheta}_{r}\right\rangle]\leq\mathbb{E}[f(\bar{\vartheta}_{r})-f(\bar{\vartheta}_{r+1})]+\sum_{\ell=1}^{L}\frac{L_{\ell}}{2}\mathbb{E}[\|\bar{\vartheta}^{\ell}_{r+1}-\bar{\vartheta}^{\ell}_{r}\|^{2}]\;.$
(7)

Yet, we observe that, using the classical intermediate quantity used for
proving convergence results of adaptive optimization methods, see for instance
[Reddi et al., 2018], we have

$\bar{\vartheta}_{r}=\bar{\theta}_{r}+\frac{\beta_{1}}{1-\beta_{1}}(\bar{\theta}_{r}-\bar{\theta}_{r-1})\;,$
(8)

where $\bar{\theta_{r}}$ denotes the average of the local models at round $r$.
Then for each layer $\ell$,

$\displaystyle\bar{\vartheta}^{\ell}_{r+1}-\bar{\vartheta}^{\ell}_{r}$
$\displaystyle=\frac{1}{1-\beta_{1}}(\bar{\theta}^{\ell}_{r+1}-\bar{\theta}^{\ell}_{r})-\frac{\beta_{1}}{1-\beta_{1}}(\bar{\theta}^{\ell}_{r}-\bar{\theta}^{\ell}_{r-1})$
$\displaystyle=\frac{\alpha_{r}}{1-\beta_{1}}\frac{1}{n}\sum_{i=1}^{n}\frac{\phi(\|\theta_{r,i}^{\ell}\|)}{\|\psi_{r,i}^{\ell}\|}\psi_{r,i}^{\ell}-\frac{\alpha_{r-1}}{1-\beta_{1}}\frac{1}{n}\sum_{i=1}^{n}\frac{\phi(\|\theta_{r-1,i}^{\ell}\|)}{\|\psi_{r-1,i}^{\ell}\|}\psi_{r-1,i}^{\ell}$
$\displaystyle=\frac{\alpha\beta_{1}}{1-\beta_{1}}\frac{1}{n}\sum_{i=1}^{n}\left(\frac{\phi(\|\theta_{r,i}^{\ell}\|)}{\sqrt{v^{t}_{r}}\|\psi_{r,i}^{\ell}\|}-\frac{\phi(\|\theta_{r-1,i}^{\ell}\|)}{\sqrt{v^{t}_{r-1}}\|\psi_{r-1,i}^{\ell}\|}\right)m^{t}_{r-1}+\frac{\alpha}{n}\sum_{i=1}^{n}\frac{\phi(\|\theta_{r,i}^{\ell}\|)}{\sqrt{v^{t}_{r}}\|\psi_{r,i}^{\ell}\|}g^{t}_{r,i}\;,$
(9)

where we have assumed a constant learning rate $\alpha$.

We note for all $\theta\in\Theta$, the majorant $G>0$ such that
$\phi(\|\theta\|)\leq G$. Then, following (7), we obtain

$\displaystyle-\mathbb{E}[\left\langle\nabla
f(\bar{\vartheta}_{r})\,|\,\bar{\vartheta}_{r+1}-\bar{\vartheta}_{r}\right\rangle]\leq\mathbb{E}[f(\bar{\vartheta}_{r})-f(\bar{\vartheta}_{r+1})]+\sum_{\ell=1}^{L}\frac{L_{\ell}}{2}\mathbb{E}[\|\bar{\vartheta}_{r+1}-\bar{\vartheta}_{r}\|^{2}]\;.$
(10)

Developing the LHS of (10) using (B.2) leads to

$\displaystyle\left\langle\nabla
f(\bar{\vartheta}_{r})\,|\,\bar{\vartheta}_{r+1}-\bar{\vartheta}_{r}\right\rangle$
$\displaystyle=\sum_{\ell=1}^{\mathsf{h}}\sum_{j=1}^{p_{\ell}}\nabla_{\ell}f(\bar{\vartheta}_{r})^{j}(\bar{\vartheta}^{\ell,j}_{r+1}-\bar{\vartheta}^{\ell,j}_{r})$
$\displaystyle=\frac{\alpha\beta_{1}}{1-\beta_{1}}\frac{1}{n}\sum_{\ell=1}^{\mathsf{h}}\sum_{j=1}^{p_{\ell}}\nabla_{\ell}f(\bar{\vartheta}_{r})^{j}\left[\sum_{i=1}^{n}\left(\frac{\phi(\|\theta_{r,i}^{\ell}\|)}{\sqrt{v^{t}_{r}}\|\psi_{r,i}^{\ell}\|}-\frac{\phi(\|\theta_{r-1,i}^{\ell}\|)}{\sqrt{v^{t}_{r-1}}\|\psi_{r-1,i}^{\ell}\|}\right)m^{t}_{r-1}\right]$
$\displaystyle\underbrace{-\frac{\alpha}{n}\sum_{\ell=1}^{\mathsf{h}}\sum_{j=1}^{p_{\ell}}\nabla_{\ell}f(\bar{\vartheta}_{r})^{j}\sum_{i=1}^{n}\frac{\phi(\|\theta_{r,i}^{\ell}\|)}{\sqrt{v^{t}_{r}}\|\psi_{r,i}^{\ell}\|}g_{r,i}^{t,l,j}}_{=A_{1}}\;.$
(11)

Suppose $T$ is the total number of local iterations and $R$ is the number of
rounds. We can write (11) as

$\displaystyle A_{1}=-\alpha\langle\nabla
f(\bar{\vartheta}_{r}),\frac{\bar{g}_{r}}{\sqrt{\hat{v}_{r}}}\rangle,$

where $\bar{g}_{r}=\frac{1}{n}\sum_{i=1}^{n}\bar{g}_{t,i}$, with
$\bar{g}_{t,i}=\Big{[}\frac{\phi(\|\theta_{t,i}^{1}\|)}{\|\psi_{t,i}^{1}\|}g_{t,i}^{1},...,\frac{\phi(\|\theta_{t,i}^{L}\|)}{\|\psi_{t,i}^{L}\|}g_{t,i}^{L}\Big{]}$
representing the normalized gradient (concatenated by layers) of the $i$-th
device. It holds that

$\displaystyle\langle\nabla
f(\bar{\vartheta}_{r}),\frac{\bar{g}_{r}}{\sqrt{\hat{v}_{r}}}\rangle$
$\displaystyle=\frac{1}{2}\|\frac{\nabla
f(\bar{\vartheta}_{r})}{\hat{v}_{r}^{1/4}}\|^{2}+\frac{1}{2}\|\frac{\bar{g}_{r}}{\hat{v}_{r}^{1/4}}\|^{2}-\|\frac{\nabla
f(\bar{\vartheta}_{r})-\bar{g}_{r}}{\hat{v}_{r}^{1/4}}\|^{2}.$ (12)

To bound the last term on the RHS, we have

$\displaystyle\|\frac{\nabla
f(\bar{\vartheta}_{r})-\bar{g}_{r}}{\hat{v}_{r}^{1/4}}\|^{2}=\|\frac{\frac{1}{n}\sum_{i=1}^{n}(\nabla
f(\bar{\vartheta}_{r})-\bar{g}_{t,i})}{\hat{v}_{r}^{1/4}}\|^{2}$
$\displaystyle\leq\frac{1}{n}\sum_{i=1}^{n}\|\frac{\nabla
f(\bar{\vartheta}_{r})-\bar{g}_{t,i}}{\hat{v}_{r}^{1/4}}\|^{2}$
$\displaystyle\leq\frac{2}{n}\sum_{i=1}^{n}\Big{(}\|\frac{\nabla
f(\bar{\vartheta}_{r})-\nabla
f(\bar{\theta}_{r})}{\hat{v}_{r}^{1/4}}\|^{2}+\|\frac{\nabla
f(\bar{\theta}_{r})-\bar{g}_{t,i}}{\hat{v}_{r}^{1/4}}\|^{2}\Big{)}.$

By Lipschitz smoothness of the loss function, the first term admits

$\displaystyle\frac{2}{n}\sum_{i=1}^{n}\|\frac{\nabla
f_{i}(\bar{\vartheta}_{r})-\nabla
f_{i}(\bar{\theta}_{r})}{\hat{v}_{r}^{1/4}}\|^{2}\leq\frac{2}{n\sqrt{\epsilon}}\sum_{i=1}^{n}L_{\ell}\|\bar{\vartheta}_{r}-\bar{\theta}_{r}\|^{2}$
$\displaystyle=\frac{2L_{\ell}}{n\sqrt{\epsilon}}\frac{\beta_{1}^{2}}{(1-\beta_{1})^{2}}\sum_{i=1}^{n}\|\bar{\theta}_{r}-\bar{\theta}_{t-1}\|^{2}$
$\displaystyle\leq\frac{2\alpha^{2}L_{\ell}}{n\sqrt{\epsilon}}\frac{\beta_{1}^{2}}{(1-\beta_{1})^{2}}\sum_{l=1}^{L}\sum_{i=1}^{n}\|\frac{\phi(\|\theta_{t,i}^{l}\|)}{\|\psi_{t,i}^{l}\|}\psi_{t,i}^{l}\|^{2}$
$\displaystyle\leq\frac{2\alpha^{2}L_{\ell}p\phi_{M}^{2}}{\sqrt{\epsilon}}\frac{\beta_{1}^{2}}{(1-\beta_{1})^{2}}.$

For the second term,

$\displaystyle\frac{2}{n}\sum_{i=1}^{n}\|\frac{\nabla
f(\bar{\theta}_{r})-\bar{g}_{t,i}}{\hat{v}_{r}^{1/4}}\|^{2}\leq\frac{4}{n}\Big{(}\underbrace{\sum_{i=1}^{n}\|\frac{\nabla
f(\bar{\theta}_{r})-\nabla
f(\theta_{t,i})}{\hat{v}_{r}^{1/4}}\|^{2}}_{B_{1}}+\underbrace{\sum_{i=1}^{n}\|\frac{\nabla
f(\theta_{t,i})-\bar{g}_{t,i}}{\hat{v}_{r}^{1/4}}\|^{2}}_{B_{2}}\Big{)}.$ (13)

Using the smoothness of $f_{i}$ we can transform $B_{1}$ into consensus error
by

$\displaystyle
B_{1}\leq\frac{L}{\sqrt{\epsilon}}\sum_{i=1}^{n}\|\bar{\theta}_{r}-\theta_{t,i}\|^{2}$
$\displaystyle=\frac{\alpha^{2}L}{\sqrt{\epsilon}}\sum_{i=1}^{n}\sum_{l=1}^{L}\|\sum_{j=\lfloor
t\rfloor_{r}+1}^{t}\Big{(}\frac{\phi(\|\theta_{j,i}^{l}\|)}{\|\psi_{j,i}^{l}\|}\psi_{j,i}^{l}-\frac{1}{n}\sum_{k=1}^{n}\frac{\phi(\|\theta_{j,k}^{l}\|)}{\|\psi_{j,k}^{l}\|}\psi_{j,k}^{l}\Big{)}\|^{2}$
$\displaystyle\leq
n\frac{\alpha^{2}L}{\sqrt{\epsilon}}M^{2}(T-1)^{2}\phi_{M}^{2}(1-\beta_{2})p,$
(14)

where the last inequality stems from Lemma B.1 in the particular case where
$\theta_{t,i}$ are averaged every $ct+1$ local iterations for any integer $c$,
since $(t-1)-(\lfloor t\rfloor_{r}+1)+1\leq T-1$.

We now develop the expectation of $B_{2}$ under the simplification that
$\beta_{1}=0$:

$\displaystyle\mathbb{E}[B_{2}]$
$\displaystyle=\mathbb{E}[\sum_{i=1}^{n}\|\frac{\nabla
f(\theta_{t,i})-\bar{g}_{t,i}}{\hat{v}_{r}^{1/4}}\|^{2}]$
$\displaystyle\leq\frac{nM^{2}}{\sqrt{\epsilon}}+n\phi_{M}^{2}\sqrt{M^{2}+p\sigma^{2}}-2\sum_{i=1}^{n}\mathbb{E}[\langle\nabla
f(\theta_{t,i}),\bar{g}_{t,i}\rangle/\sqrt{\hat{v}_{r}}]$
$\displaystyle=\frac{nM^{2}}{\sqrt{\epsilon}}+n\phi_{M}^{2}\sqrt{M^{2}+p\sigma^{2}}-2\sum_{i=1}^{n}\sum_{\ell=1}^{L}\mathbb{E}[\langle\nabla_{\ell}f(\theta_{t,i}),\frac{\phi(\|\theta_{t,i}^{l}\|)}{\|\psi_{t,i}^{l}\|}g_{t,i}^{l}\rangle/\sqrt{\hat{v}_{r}^{l}}]$
$\displaystyle=\frac{nM^{2}}{\sqrt{\epsilon}}+n\phi_{M}^{2}\sqrt{M^{2}+p\sigma^{2}}-2\sum_{i=1}^{n}\sum_{l=1}^{L}\sum_{i=1}^{p_{l}}\mathbb{E}[\nabla_{l}f(\theta_{t,i})^{j}\frac{\phi(\|\theta_{t,i}^{l,j}\|)}{\sqrt{\hat{v}_{r}^{l,j}}\|\psi_{t,i}^{l,j}\|}g_{t,i}^{l,j}]$
$\displaystyle\leq\frac{nM^{2}}{\sqrt{\epsilon}}+n\phi_{M}^{2}\sqrt{M^{2}+p\sigma^{2}}-2\sum_{i=1}^{n}\sum_{l=1}^{L}\sum_{i=1}^{p_{l}}\mathbb{E}\left[\sqrt{\frac{1-\beta_{2}}{M^{2}p_{\ell}}}\phi(\|\theta_{r,i}^{l,j}\|)\nabla_{l}f(\theta_{t,i})^{j}g_{t,i}^{l,j}\right]$
$\displaystyle\hskip
28.90755pt-2\sum_{i=1}^{n}\sum_{l=1}^{L}\sum_{j=1}^{p_{l}}E\left[\left(\phi(\|\theta_{r,i}^{l,j}\|)\nabla_{l}f(\theta_{t,i})^{j}\frac{g_{r,i}^{t,l,j}}{\|\psi_{r,i}^{l,j}\|}\right)\mathsf{1}\left(\text{sign}(\nabla_{l}f(\theta_{t,i})^{j}\neq\text{sign}(g_{r,i}^{t,l,j})\right)\right],$

where we use assumption Assumption 4, Assumption 4 and Assumption 4. Yet,

$\displaystyle-\mathbb{E}\Bigg{[}\left(\phi(\|\theta_{r,i}^{l,j}\|)\nabla_{l}f(\theta_{t,i})^{j}\frac{g_{r,i}^{t,l,j}}{\|\psi_{r,i}^{l,j}\|}\right)\mathsf{1}\left(\text{sign}(\nabla_{l}f(\theta_{t,i})^{j}\neq\text{sign}(g_{r,i}^{t,l,j})\right)\Bigg{]}$
$\displaystyle\hskip
144.54pt\leq\phi_{M}\nabla_{l}f(\theta_{t,i})^{j}\mathbb{P}\left[\text{sign}(\nabla_{l}f(\theta_{t,i})^{j}\neq\text{sign}(g_{r,i}^{t,l,j})\right].$

Then we have

$\displaystyle\mathbb{E}[B_{2}]\leq\frac{nM^{2}}{\sqrt{\epsilon}}+n\phi_{M}^{2}\sqrt{M^{2}+p\sigma^{2}}-2\phi_{m}\sqrt{\frac{1-\beta_{2}}{M^{2}p}}\sum_{i=1}^{n}\mathbb{E}[\|[\nabla
f(\theta_{t,i})\|^{2}]+\phi_{M}\frac{\mathsf{h}\sigma^{2}}{\sqrt{n}}$

Thus, (13) becomes

$\displaystyle\frac{2}{n}\sum_{i=1}^{n}\|\frac{\nabla
f_{i}(\bar{\theta}_{r})-\bar{g}_{t,i}}{\hat{v}_{r}^{1/4}}\|^{2}\leq
4\left[\frac{\alpha^{2}L_{\ell}}{\sqrt{\epsilon}}\alpha^{2}M^{2}(T-1)^{2}\phi_{M}^{2}(1-\beta_{2})p+\frac{\alpha
M^{2}}{\sqrt{\epsilon}}+\phi_{M}^{2}\sqrt{M^{2}+p\sigma^{2}}+\alpha\phi_{M}\frac{\mathsf{h}\sigma^{2}}{\sqrt{n}}\right]$

Substituting all ingredients into (12), we obtain

$\displaystyle-\alpha\mathbb{E}[\langle\nabla
f(\bar{\vartheta}_{r}),\frac{\bar{g}_{r}}{\sqrt{\hat{v}_{r}}}\rangle]$
$\displaystyle\leq-\frac{\alpha}{2}\mathbb{E}\big{[}\|\frac{\nabla
f(\bar{\vartheta}_{r})}{\hat{v}_{r}^{1/4}}\|^{2}\big{]}-\frac{\alpha}{2}\mathbb{E}\big{[}\|\frac{\bar{g}_{r}}{\hat{v}_{r}^{1/4}}\|^{2}\big{]}+\frac{2\alpha^{3}L_{\ell}p\phi_{M}^{2}}{\sqrt{\epsilon}}\frac{\beta_{1}^{2}}{(1-\beta_{1})^{2}}$
$\displaystyle\hskip
7.22743pt+4\alpha\left[\frac{\alpha^{2}L}{\sqrt{\epsilon}}M^{2}(T-1)^{2}\phi_{M}^{2}(1-\beta_{2})p+\frac{\alpha
M^{2}}{\sqrt{\epsilon}}+\phi_{M}^{2}\sqrt{M^{2}+p\sigma^{2}}+\alpha\phi_{M}\frac{\mathsf{h}\sigma^{2}}{\sqrt{n}}\right].$

To bound the second term on the RHS in above, we notice that

$\displaystyle\mathbb{E}\big{[}\|\frac{\bar{g}_{r}}{\hat{v}_{r}^{1/4}}\|^{2}\big{]}=\frac{1}{n^{2}}\mathbb{E}\big{[}\|\frac{\sum_{i=1}^{n}\bar{g}_{r,i}}{\hat{v}_{r}^{1/4}}\|^{2}\big{]}$
$\displaystyle=\frac{1}{n^{2}}\mathbb{E}\big{[}\sum_{l=1}^{L}\sum_{i=1}^{n}\|\frac{\phi(\|\theta_{r,i}^{l}\|)}{\hat{v}^{1/4}\|\psi_{r,i}^{l}\|}g_{r,i}^{l}\|^{2}\big{]}$
$\displaystyle\geq\phi_{m}^{2}(1-\beta_{2})\mathbb{E}\left[\|\frac{1}{n}\sum_{i=1}^{n}\frac{\nabla
f(\theta_{r,i})}{\hat{v}^{1/4}_{r}}\|^{2}\right]$ (15)
$\displaystyle=\phi_{m}^{2}(1-\beta_{2})\mathbb{E}\left[\|\frac{\overline{\nabla}f(\theta_{r})}{\hat{v}^{1/4}_{r}}\|^{2}\right].$

Regarding
$\left\|\frac{\overline{\nabla}f(\theta_{r})}{\hat{v}_{r}^{1/4}}\right\|^{2}$,
we have

$\displaystyle\left\|\frac{\overline{\nabla}f(\theta_{r})}{\hat{v}_{r}^{1/4}}\right\|^{2}$
$\displaystyle\geq\frac{1}{2}\left\|\frac{\nabla
f(\overline{\theta_{r}})}{\hat{v}_{r}^{1/4}}\right\|^{2}-\left\|\frac{\overline{\nabla}f(\theta_{r})-\nabla
f(\overline{\theta_{r}})}{\hat{v}_{r}^{1/4}}\right\|^{2}$
$\displaystyle\geq\frac{1}{2}\left\|\frac{\nabla
f(\overline{\theta_{r}})}{\hat{v}_{r}^{1/4}}\right\|^{2}-\left\|\frac{\frac{1}{n}\sum_{i=1}^{n}(\nabla
f_{i}(\theta_{r})-\nabla f(\bar{\theta}_{r}))}{\hat{v}_{r}^{1/4}}\right\|^{2}$
$\displaystyle\geq\frac{1}{2}\left\|\frac{\nabla
f(\overline{\theta_{r}})}{\hat{v}_{r}^{1/4}}\right\|^{2}-\frac{\alpha^{2}L_{\ell}}{\sqrt{\epsilon}}M^{2}(T-1)^{2}(\sigma^{2}+G^{2})(1-\beta_{2})p,$

where the last line is due to (14) and Assumption 4. Therefore, we have
obtained

$\displaystyle A_{1}$
$\displaystyle\leq-\frac{\alpha\phi_{m}^{2}(1-\beta_{2})}{4}\left\|\frac{\nabla
f(\overline{\theta_{r}})}{\hat{v}_{r}^{1/4}}\right\|^{2}+\frac{\alpha^{3}L_{\ell}}{\sqrt{\epsilon}}M^{2}(T-1)^{2}\phi_{m}^{2}\phi_{M}^{2}(1-\beta_{2})^{2}p+\frac{2\alpha^{3}L_{\ell}p\phi_{M}^{2}}{\sqrt{\epsilon}}\frac{\beta_{1}^{2}}{(1-\beta_{1})^{2}}$
$\displaystyle\hskip
28.90755pt+4\alpha\left[\frac{\alpha^{2}L}{\sqrt{\epsilon}}M^{2}(T-1)^{2}(\sigma^{2}+G^{2})(1-\beta_{2})p+\frac{M^{2}\alpha}{\sqrt{\epsilon}}+\alpha\phi_{M}^{2}\sqrt{M^{2}+p\sigma^{2}}+\phi_{M}\alpha\frac{\mathsf{h}\sigma^{2}}{\sqrt{n}}\right],$
$\displaystyle\leq-\frac{\alpha\phi_{m}^{2}(1-\beta_{2})}{4}\left\|\frac{\nabla
f(\overline{\theta_{r}})}{\hat{v}_{r}^{1/4}}\right\|^{2}+\frac{\alpha^{3}L_{\ell}}{\sqrt{\epsilon}}M^{2}(T-1)^{2}\phi_{m}^{2}\phi_{M}^{2}(1-\beta_{2})^{2}p+\frac{2\alpha^{3}L_{\ell}p\phi_{M}^{2}}{\sqrt{\epsilon}}\frac{\beta_{1}^{2}}{(1-\beta_{1})^{2}}$
$\displaystyle\hskip
28.90755pt+4\alpha\Big{[}\frac{\alpha^{2}L}{\sqrt{\epsilon}}M^{2}(T-1)^{2}G^{2}(1-\beta_{2})p+\frac{M^{2}\alpha}{\sqrt{\epsilon}}+\alpha\phi_{M}^{2}\sqrt{M^{2}+p\sigma^{2}}$
$\displaystyle\hskip
122.85876pt+\sigma^{2}\left(\frac{\alpha^{2}L}{\sqrt{\epsilon}}M^{2}(T-1)^{2}(1-\beta_{2})p+\phi_{M}\alpha\frac{\mathsf{h}}{\sqrt{n}}\right)\Big{]}.$

Substitute back into (11), assuming $M\leq 1$, we have the following by taking
the telescope sum

$\displaystyle\frac{1}{R}\sum_{t=1}^{R}\mathbb{E}\left[\left\|\frac{\nabla
f(\overline{\theta_{r}})}{\hat{v}_{r}^{1/4}}\right\|^{2}\right]$
$\displaystyle\lesssim\sqrt{\frac{M^{2}p}{n}}\frac{f(\bar{\vartheta}_{1})-\mathbb{E}[f(\bar{\vartheta}_{R+1})]}{\mathsf{h}\alpha
R}+\frac{\alpha}{n^{2}}\sum_{r=1}^{R}\sum_{i=1}^{n}\sigma_{i}^{2}\mathbb{E}\left[\left\|\frac{\phi(\|\theta_{r,i}^{\ell}\|)}{\sqrt{v_{r}}\|\psi_{r,i}^{\ell}\|}\right\|^{2}\right]+\frac{2\alpha^{3}\overline{L}p\phi_{M}^{2}}{\sqrt{\epsilon}}\frac{\beta_{1}^{2}}{(1-\beta_{1})^{2}}$
$\displaystyle+4\Big{[}\frac{\alpha^{2}\overline{L}}{\sqrt{\epsilon}}M^{2}(T-1)^{2}G^{2}(1-\beta_{2})p+\frac{\alpha
M^{2}}{\sqrt{\epsilon}}+\alpha\phi_{M}^{2}\sqrt{M^{2}+p\sigma^{2}}$
$\displaystyle+\sigma^{2}\left(\frac{\alpha^{2}\overline{L}}{\sqrt{\epsilon}}M^{2}(T-1)^{2}(1-\beta_{2})p+\phi_{M}\alpha\frac{\mathsf{h}}{\sqrt{n}}\right)\Big{]}+\frac{\alpha\beta_{1}}{1-\beta_{1}}\sqrt{(1-\beta_{2})p}\frac{\mathsf{h}M^{2}}{\sqrt{\epsilon}}+\overline{L}\alpha^{2}M^{2}\phi_{M}^{2}\frac{(1-\beta_{2})p}{T\epsilon}$
$\displaystyle\leq\sqrt{\frac{M^{2}p}{n}}\frac{\mathbb{E}[f(\bar{\theta}_{1})]-\min\limits_{\theta\in\Theta}f(\theta)}{\mathsf{h}\alpha
R}+\frac{\phi_{M}\sigma^{2}}{Rn}\sqrt{\frac{1-\beta_{2}}{M^{2}p}}$
$\displaystyle+4\Big{[}\frac{\alpha^{2}\overline{L}}{\sqrt{\epsilon}}M^{2}(T-1)^{2}G^{2}(1-\beta_{2})p+\frac{M^{2}\alpha}{\sqrt{\epsilon}}+\phi_{M}^{2}\alpha\sqrt{M^{2}+p\sigma^{2}}$
$\displaystyle+\sigma^{2}\Big{(}\frac{\alpha^{2}\overline{L}}{\sqrt{\epsilon}}M^{2}(T-1)^{2}(1-\beta_{2})p+\phi_{M}\frac{\mathsf{h}}{\sqrt{n}}\Big{)}\Big{]}+\frac{\alpha\beta_{1}}{1-\beta_{1}}\sqrt{(1-\beta_{2})p}\frac{\mathsf{h}M^{2}}{\sqrt{\epsilon}}$
$\displaystyle\hskip
101.17755pt+\overline{L}\alpha^{2}M^{2}\phi_{M}^{2}\frac{(1-\beta_{2})p}{T\epsilon}+\frac{2\alpha^{3}\overline{L}p\phi_{M}^{2}}{\sqrt{\epsilon}}\frac{\beta_{1}^{2}}{(1-\beta_{1})^{2}}.$

Organizing terms, we conclude the proof. ∎


# Voice Quality and Pitch Features in Transformer-Based Speech Recognition

###### Abstract

Jitter and shimmer measurements have shown to be carriers of voice quality and
prosodic information which enhance the performance of tasks like speaker
recognition, diarization or automatic speech recognition (ASR). However, such
features have been seldom used in the context of neural-based ASR, where
spectral features often prevail. In this work, we study the effects of
incorporating voice quality and pitch features altogether and separately to a
Transformer-based ASR model, with the intuition that the attention mechanisms
might exploit latent prosodic traits. For doing so, we propose separated
convolutional front-ends for prosodic and spectral features, showing that this
architectural choice yields better results than simple concatenation of such
pitch and voice quality features to mel-spectrogram filterbanks. Furthermore,
we find mean Word Error Rate relative reductions of up to 5.6% with the
LibriSpeech benchmark. Such findings motivate further research on the
application of prosody knowledge for increasing the robustness of Transformer-
based ASR.

Index Terms: speech recognition, voice quality, prosody, transformers

##  1 Introduction

Research during the last years has shown success of neural network models
applied to automatic speech recognition (ASR), consistently achieving state-
of-the-art results in the field. The architectural choices vary within a wide
range, from hybrid DNN-HMM models [1] to fully neural end-to-end approaches,
like convolutional models benefiting from gated linear units (GLUs) [2, 3, 4],
ResNet blocks [5, 6] or time-depth separable (TDS) convolutions [6, 7].

The mentioned convolutional architectures are efficient approaches in the
sense that they achieve low word error rate (WER) scores, while being very
suitable for online streaming applications [8]. This effectiveness is due to
their ability to extract relevant local features, although other recently
proposed models seem to have an easier time capturing information from larger
contexts. These are sequence-to-sequence models based on the attention
mechanism [9], which is able to focus on the relevant parts of a whole
sequence in order to generate the correct output, like the LAS model [10].
Actually, many state-of-the-art proposals mainly consist of Transformer blocks
[11], which solely rely on the attention mechanism. Transformer models might
be complemented with convolutional operations at their front-ends [6, 12], or
even within the Transformer blocks themselves, a structure known as Conformer
[13]. This way, local feature extraction is enhanced, an important factor
knowing that attention acts by relating time frames in the signal, each one of
them represented by an acoustic feature vector, in a similar way as words are
encoded in embeddings in the NLP field.

Currently, most of the modern architectures work only on cepstral and mel-
frequency spectral coefficients inputs, or even directly with the raw
waveform, tending towards the increasing of granularity at input level by
usually augmenting the number of spectral parameters. Whilst it seems evident
that current end-to-end deep architectures are able to automatically perform
relevant feature extraction for speech tasks, psychical or functional
properties, related to the underlying speech production system, become fuzzy
or difficult to connect with the speech recognition performances. In addition,
it is still unclear how the great quantity and different speech hand-crafted
voice features, carefully developed along past years and based on our
linguistic knowledge, might help and in which degree to the current ASR neural
architectures. To name a few, pitch –or fundamental perceived frequency–, is
an important prosodic information carrier, or jitter and shimmer, which
represent the cycle-to-cycle variations of the fundamental frequency and
amplitude, respectively, and are related to the quality of speech.

The main contribution of this work is assessing the value of two ways of
adding voice quality (VQ) and pitch features to the spectral coefficients,
commonly employed in most neural ASR systems. Firstly, by simple concatenation
before the network’s forward pass. Secondly, by using two separated
convolutional front-ends, one for spectral features and another for voice
quality plus pitch features, concatenating the outcome features from the
front-ends. Experiments are carried out with the LibriSpeech dataset [14],
using the Transformer-based model [11] from fairseq’s toolkit speech-to-text
(S2T) recipe [15, 16]. Up to our knowledge, this is the first attempt to use
jitter and shimmer features within a modern deep neural-based speech
recognition system while keeping easy to identify psychical/functional
properties of the voice and linking them to the ASR performance. Furthermore,
the recipe employed in this research is made available for the community in a
GitHub111https://github.com/gcambara/speechbook/tree/master/recipes/vq_pitch
repository. We hypothesize that prosodic and voice quality features may boost
robustness in state-of-the-art neural ASR systems like the Transformer-based
model used in this paper. Acoustic spectral feature vectors would be
enrichened with such information, in our case specifically using jitter and
shimmer local variants, the best performing ones in [17]. These might
complement pitch features or even be beneficial on themselves.

The structure of the paper is as follows. Section 2 briefly overviews the
related work on pitch and voice quality features, Sections 3 and 4 describe
the model and methodology used in the current work, respectively, Section 5
shows the results obtained, and Section 6 illustrates the final conclusions.

##  2 Pitch and VQ Features in ASR

Some well-known speech recognition frameworks, like Kaldi [18], have
incorporated the use of additional prosodic features, such as the pitch or the
probability of voicing (POV). These are stacked into an input vector together
with those cepstral/spectral ones, and then forwarded to classifiers like HMM
or DNN ensembles. Such features have proved to increase performance in ASR
systems, significantly for tonal languages, like Punjabi [19], as well as for
non-tonal ones like English [20]. Nevertheless, prosodic features are seldom
used in the context of Transformer-based or convolutional speech recognition,
often relying on the sole use of spectral features.

Regarding jitter and shimmer measurements, they have been considered relevant
and often applied to detect voice pathologies for a long time, as in [21, 22].
Although voice quality features differ intrinsically from suprasegmental
prosodic ones like pitch itself, they have shown to be related to prosody. In
[23], the authors showed that voice quality features are relevant markers
signaling paralinguistic information, and that they should even be considered
as prosodic carriers along with pitch and duration, for instance. Furthermore,
in the last decades, jitter and shimmer have shown to be useful in a wide
range of applications; e.g. detection of different speaking styles [24], age
and gender classification [25], emotion detection [26], speaker recognition
[17], speaker diarization [27], and Alzheimer’s and Parkinson Diseases
detection [28, 29], among others.

In ASR literature, some works have reported that prosodic information can
raise the performance of speech recognizer systems. For instance, in [30] the
authors built an ASR for dysarthric speech and [31] reports benefit on the use
of jitter and shimmer for noisy speech recognition, both systems based on
classical HMM acoustic modelling. Actually, our choice of processing prosodic
and spectral features separately before fusion has been inspired by other
works in speaker verification [32] and speaker recognition [33]. In our
previous work [34], we employed voice quality features in a convolutional-
based ASR, showing only slight improvements in WER. Besides such work, and up
to our knowledge, voice quality features have not been used yet for neural-
based ASR. Since Transformer models use attention to find correlations in
larger contexts, we think that prosodic information from pitch and voice
quality features could be better captured compared to convolutional
approaches, due to its suprasegmental nature.

Figure 1: Convolutional front-end for mel-spectrogram filterbanks and pitch
related features in S2T VQ Transformer architecture.

##  3 Model Description

The addition of voice quality and pitch features to spectral coefficients is
tackled early on, in the front-end of a Transformer-based model. The main
structure of such model, named as S2T Transformer from now on, consists of a
convolutional block as a front-end, followed by sinusoidal positional
embedding of features and an encoder-decoder constituted by Transformer
blocks, in the likes of [11]. The convolutional block is formed by two 1-dim
convolutional layers with kernel size 5, stride 2 and padding 2, each followed
by GLU activation functions [4]. The first layer takes $T$ time frames per
$N=40$ mel-spectrograms as input, upsampling them to $1024$ features that are
halved by the GLU activation, resulting in $P=512$ hidden features. These are
passed to the second layer, which outputs the same feature dimension of $512$,
yielding $O=256$ features after GLU activation, which is the embedding size of
all the following attention layers.

Incorporating any combination of voice quality and pitch features with size
$M\in[1,5]$ to the mel-spectrograms causes the input feature vector to be of
size $N+M$. Thus, the implicit pitch-related information is finally contained
within the output feature vector of size $O=256$. However, we wonder whether
such information could be too diluted after convolving with the mel-
spectrograms, being $N>>M$.

Figure 2: LibriSpeech test-clean and test-other WER (%) for several feature
configurations tested, combining mel-spectrogram filterbanks (FB), 3 random
features (Rand), F0+POV+$\Delta$F0 (Pitch), jitter (J) and shimmer (S), using
S2T and S2T VQ models. 40 FB baseline is marked by the dashline.

For such reason, we propose an architectural variant called S2T VQ
Transformer, shown in Figure 1, which consists of two convolutional blocks,
$A$ and $B$, receiving mel-spectrograms and voice quality plus pitch features,
respectively. These are convolved independently and concatenated afterwards,
so the resulting feature vector contains a representation of $K$ spectral
features and $L$ pitch-related features, yielding a total output size of
$O^{\prime}=K+L$. By controlling the proportion between $K$ and $L$ we decide
how much weight we assign to the new features, so the attention layers may
take these more into account. Each block $A$ and $B$ is identical to the block
in the plain S2T Transformer architecture, except for the variations showed in
Table 1.

Table 1: Convolutional front-end blocks for mel-spectrograms and voice quality + pitch features. Block | Input Dim | Hidden Dim | Output Dim
---|---|---|---
A | $N=40$ | $p_{A}=512$ | $K=192$
B | $M\in[1,5]$ | $p_{B}=256$ | $L=64$

Note that the hidden dimension $p_{A}$ is kept for the spectral block $A$, but
for block $B$ is reduced to the half, since the amount of pitch and voice
quality features is smaller and probably we do not need as much up-sampling.
By setting output dimensions for each front-end at $K=192$ and $L=64$,
concatenation of these yields the same total output size as plain S2T
Transformer model: $O^{\prime}=O=256$. This way, we ensure fairness in the
comparison between plain S2T and VQ S2T models, since the simple fact of
increasing $O^{\prime}$ might yield better results due to the increased number
of parameters in the feature vectors, maybe better exploited by attention.
Actually, the VQ variant has slightly less parameters, 29.2M against 29.4M in
the plain model. Furthermore, this choice keeps a proportion of 1 pitch-
related feature for every 3 spectral features, giving a larger weight to the
first ones compared to the plain S2T Transformer model.

##  4 Methodology

###  4.1 Dataset

The proposal of this work is assessed with the LibriSpeech dataset [14], which
contains up to 1000 hours of audio book recordings in English, sampled at
16kHz. LibriSpeech was deployed with a ready-to-go train split subdivided into
three subsets with approximately 100, 360 and 500 hours of speech. For
development and test sets, two subsets were created as well. Hereinafter, the
aforementioned partitions shall be referred as train-clean-100, train-
clean-360, train-other-500, dev-clean, dev-other, test-clean and test-other,
respectively.

###  4.2 Experimental setup

Aiming to explore the information conveyed by pitch and voice quality features
in the speech recognition system, we mainly train and test independent
acoustic models using different feature combinations: mel-spectrograms,
fundamental frequency (F0), POV, $\Delta$F0, jitter and shimmer. To ensure
that additional voice quality and pitch features do not cause a positive
impact simply because they increase the number of parameters in the model, we
also try a configuration of 3 features generated with random numbers sampled
from a uniform distribution between 0 and 10, similar to mel-spectrogram
filterbanks. Performance for every model is evaluated in terms of WER in test-
clean and test-other sets. The whole train and test procedure is repeated for
every feature configuration across 6 different initialization seeds,
calculating the average WER plus the standard deviation of the mean to ensure
statistical significance of the results. We decide to train with the train-
clean-100 set only for the seed extensive experiments, due to computational
constraints. Afterwards, we compare the performance of the filterbank-only
baseline against the fully-packed S2T VQ model with pitch and voice quality
features. For this, we explore different training set hours: 50, 100, 200, 500
and 960.

###  4.3 Acoustic modeling

Training and testing of acoustic models is done with the fairseq toolkit [15],
which contains functionality and examples for the S2T task [16]. Specifically,
we use their LibriSpeech ASR example as a starting point for our experiments,
training the S2T small Transformer model (31M parameters) with mel-spectrogram
features computed on-the-fly by fairseq.

For experiments requiring pitch and voice quality features, we precompute
these with Praat-Parselmouth [35], a Python wrapper for Praat [36], a well
tested library for speech processing and analysis. We compute all features
with a window size of 25 ms and a stride of 10 ms. The pitch, jitter and
shimmer extraction algorithms yield empty values for unvoiced frames, so
following Kaldi’s pitch extraction algorithm recommendation [37], we
interpolate these with the adjacent non-empty values. We keep the empty
indexes to construct a POV vector where voiced and unvoiced segments are
denoted by $1$ and $-1$, respectively. The logarithm is applied to the pitch
vector and pitch, jitter and shimmer vectors are smoothed out by computing the
mean with a centered window across 151 frames, in the style of Kaldi. Finally,
all the combined features are normalized by cepstral mean and variance
normalization.

From this point, features may be handled in two different ways, depending on
the chosen front-end, as described in section 3: simply stacked with mel-
spectrograms, or processed in a dedicated convolutional block and then stacked
with the convolved mel-spectrograms. We use 40 filterbanks, instead of the 80
filterbanks in the original fairseq example, in order to increase the
proportion of pitch and voice quality features in the total amount of
features. This is done to yield a higher numerical importance to the new set
of features, specially for the simple stacking experiments, while keeping a
commonly used number of filterbanks. Since simple stacking configuration is
not used for the training set size experiments, we then switch back to 80
filterbanks.

Models are trained using the cross entropy criterion with label smoothing and
Adam optimization, clipping gradient values above 10.0 to prevent gradient
explosions. The learning rate peaks at 0.002 after warming up for 10k batch
updates, being decreased afterwards by an inverse root scheduler. For the
experiments regarding the different training sizes, we use the exact
hyperparameters of the fairseq recipe, to be closer to its baseline. We leave
them training for many iterations after validation loss plateaus, to ensure
convergence: 300k updates for the 960 hours model, 150k updates for the 500,
200 and 100 hours ones, and 55k for the 50 hours model. However, we restrict
the number of iterations to 20k batch updates for the feature configuration
scan across seeds, since this is computationally costly and no significant
gains are observed for longer trainings.

Table 2: Mean test WER scores for acoustic models trained with different subset sizes of LibriSpeech training. Two architectures are used: the baseline with filterbank features (80 FB) and the S2T VQ model with filterbank, pitch and voice quality features (+VQ). Train set | WER (%)
---|---
Hours | test-clean | test-other
| 80 FB | \+ VQ | 80 FB | \+ VQ
50 | $29.26$ | $\mathbf{28.30}$ | $47.70$ | $\mathbf{46.91}$
100 | $19.33$ | $\mathbf{18.29}$ | $37.04$ | $\mathbf{35.31}$
200 | $10.95$ | $\mathbf{10.80}$ | $\mathbf{24.94}$ | $25.32$
500 | $7.78$ | $\mathbf{6.96}$ | $18.34$ | $\mathbf{17.56}$
960 | $4.62$ | $\mathbf{4.27}$ | $10.64$ | $\mathbf{10.01}$

###  4.4 Decoding

We found better WER scores by averaging the weights from the last 10
checkpoints in training, which is the method in the fairseq recipe, rather
than selecting the checkpoint with best development WER. Decoding is done by
means of the beam search algorithm with beam size 5, using a 10k unigram
lexicon extracted with sentencepiece [38] from the LibriSpeech corpus. Since
the aim of this work is the assessment of voice quality and pitch features for
acoustic modeling, no additional language model is used.

##  5 Results

Figures 2 and 2 depict the WER score distributions for a selection of the
feature combinations across the 6 different seeds employed, for the
LibriSpeech test sets. Overall, mean WER scores for the S2T VQ model tend to
be smaller than plain S2T Transformer using simple concatenation of features.
Furthermore, there is a trend in WER reduction as the number of added features
increases, specially for the S2T VQ architecture. This is particularly notable
when compared to the baseline of 40 filterbanks, or the model using 3 random
features.

Pitch and VQ features seem to benefit from an architecture using specific
convolutions for them, as hypothesized. Jitter and shimmer only slightly
improve the performance of the system, achieving reductions in mean WER of
$1.3\%$ and $0.2\%$ for test-clean and test-other, respectively. Pitch
achieves a better performance, improving $1.5\%$ and $1.0\%$ mean WER for the
same sets. However, the biggest improvement comes from joining both pitch plus
jitter and shimmer features. Mean WER reduction at test-clean is around
$2.3\%$, whereas for test-other is $1.0\%$. Two tailed p-values in the pitch
experiments are $0.134$ and $0.069$, showing a fair significance, whereas the
jitter and shimmer experiments yield low significance, with p-values of
$0.129$ and $0.624$. Nevertheless, the experiments with pitch and VQ features
altogether show a relevant significance, with p-values of $0.027$ and $0.052$
for test-clean and test-other. This suggests that jitter and shimmer might be
a good complement to improve pitch features, although they alone do not cause
a significant change.

Table 2 shows that the fully-packed S2T VQ model surpasses the filterbank-only
baseline in the majority of training size setups. The relative WER reduction
averaged across training setups is $5.6\%$ for test-clean and $3.0\%$ for
test-other, suggesting that the benefits of using pitch and VQ features hold
as the number of train hours increases. The influence of using pitch and VQ
features as training hours grow is better understood looking at the error type
distributions in Figure 3. The evolution of deletions (D) is very similar in
both cases, but the S2T VQ model is significantly better at reducing the
amount of insertions (I), leaving a higher representation of substitutions
(S). This suggests that, with enough data, the S2T VQ model might learn to
better leverage prosody information for delimiting beginning and ending of
words.

Figure 3: Error distributions across training hours. Error types:
substitutions (S), deletions (D) and insertions (I).

##  6 Conclusions

The incorporation of pitch and voice quality features to spectral features has
been studied for a Transformer-based state-of-the-art acoustic ASR, in two
ways: the first being simple concatenation, and the second one exploring
concatenation after separately convolving such features in two separated
front-ends. The latter proposal clearly outperforms the former, showing WER
improvements by using pitch and voice quality features separately,
significantly yielding the best performances when used together, with up to
mean $5.6\%$ relative WER reductions, specially taming insertion errors. Thus,
results show the potential of voice quality features for complementing
prosodic information carried in pitch features. Using separated convolutional
filters for these features increases the representation of such prosodic
information in the acoustic features, which suprasegmental nature might be
better leveraged by the attention layers. Besides further explorations with
bigger and more diverse data sets, we think that the findings of this work
motivate further research in conversational speech recognition or speech
recognition with punctuation signs, where prosody plays an important role.

##  7 Acknowledgements

This work is a part of the INGENIOUS project, funded by the European Union’s
Horizon 2020 Research and Innovation Programme and the Korean Government under
Grant Agreement No 833435. The third author has been funded by the Agencia
Estatal de Investigación (AEI), Ministerio de Ciencia, Innovación y
Universidades and the Fondo Social Europeo (FSE) under grant RYC-2015-17239
(AEI/FSE, UE).

## References

  * [1] K. Veselỳ, A. Ghoshal, L. Burget, and D. Povey, “Sequence-discriminative training of deep neural networks.” in _Interspeech_ , vol. 2013, 2013, pp. 2345–2349.
  * [2] N. Zeghidour, Q. Xu, V. Liptchinsky, N. Usunier, G. Synnaeve, and R. Collobert, “Fully convolutional speech recognition,” 2018.
  * [3] V. Liptchinsky, G. Synnaeve, and R. Collobert, “Letter-based speech recognition with gated convnets,” _ArXiv_ , vol. abs/1712.09444, 2017.
  * [4] Y. N. Dauphin, A. Fan, M. Auli, and D. Grangier, “Language modeling with gated convolutional networks,” _CoRR_ , vol. abs/1612.08083, 2016. [Online]. Available: http://arxiv.org/abs/1612.08083
  * [5] Y. Wang, X. Deng, S. Pu, and Z. Huang, “Residual convolutional ctc networks for automatic speech recognition,” _arXiv preprint arXiv:1702.07793_ , 2017\.
  * [6] G. Synnaeve, Q. Xu, J. Kahn, T. Likhomanenko, E. Grave, V. Pratap, A. Sriram, V. Liptchinsky, and R. Collobert, “End-to-end ASR: from supervised to semi-supervised learning with modern architectures,” 2019.
  * [7] A. Hannun, A. Lee, Q. Xu, and R. Collobert, “Sequence-to-sequence speech recognition with time-depth separable convolutions,” _arXiv preprint arXiv:1904.02619_ , 2019.
  * [8] V. Pratap, Q. Xu, J. Kahn, G. Avidov, T. Likhomanenko, A. Hannun, V. Liptchinsky, G. Synnaeve, and R. Collobert, “Scaling up online speech recognition using convnets,” _arXiv preprint arXiv:2001.09727_ , 2020.
  * [9] D. Bahdanau, K. Cho, and Y. Bengio, “Neural machine translation by jointly learning to align and translate,” _arXiv preprint arXiv:1409.0473_ , 2014\.
  * [10] W. Chan, N. Jaitly, Q. V. Le, and O. Vinyals, “Listen, attend and spell,” _CoRR_ , vol. abs/1508.01211, 2015. [Online]. Available: http://arxiv.org/abs/1508.01211
  * [11] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin, “Attention is all you need,” _arXiv preprint arXiv:1706.03762_ , 2017.
  * [12] A. Mohamed, D. Okhonko, and L. Zettlemoyer, “Transformers with convolutional context for ASR,” _arXiv preprint arXiv:1904.11660_ , 2019.
  * [13] A. Gulati, J. Qin, C.-C. Chiu, N. Parmar, Y. Zhang, J. Yu, W. Han, S. Wang, Z. Zhang, Y. Wu _et al._ , “Conformer: Convolution-augmented transformer for speech recognition,” _arXiv preprint arXiv:2005.08100_ , 2020.
  * [14] V. Panayotov, G. Chen, D. Povey, and S. Khudanpur, “Librispeech: An ASR corpus based on public domain audio books,” 04 2015, pp. 5206–5210.
  * [15] M. Ott, S. Edunov, A. Baevski, A. Fan, S. Gross, N. Ng, D. Grangier, and M. Auli, “fairseq: A fast, extensible toolkit for sequence modeling,” _arXiv preprint arXiv:1904.01038_ , 2019.
  * [16] C. Wang, Y. Tang, X. Ma, A. Wu, D. Okhonko, and J. Pino, “fairseq s2t: Fast speech-to-text modeling with fairseq,” _arXiv preprint arXiv:2010.05171_ , 2020.
  * [17] M. Farrús, J. Hernando, and P. Ejarque, “Jitter and shimmer measurements for speaker recognition,” in _Proceedings of the Interspeech_ , Antwerp, Belgium, 2007.
  * [18] D. Povey, A. Ghoshal, G. Boulianne, L. Burget, O. Glembek, N. Goel, M. Hannemann, P. Motlíček, Y. Qian, P. Schwarz, J. Silovský, G. Stemmer, and K. Vesel, “The Kaldi speech recognition toolkit,” _IEEE Workshop on Automatic Speech Recognition and Understanding_ , 2011.
  * [19] J. Guglani and A. Mishra, “Automatic speech recognition system with pitch dependent features for Punjabi language on Kaldi toolkit,” _Applied Acoustics_ , vol. 167, p. 107386, 2020.
  * [20] M. Magimai-Doss, T. A. Stephenson, and H. Bourlard, “Using pitch frequency information in speech recognition,” in _Eighth European Conference on Speech Communication and Technology_ , 2003.
  * [21] J. Kreiman and B. R. Gerratt, “Perception of aperiodicity in pathological voice,” _The Journal of the Acoustical Society of America_ , vol. 117, no. 4, pp. 2201–2211, 2005.
  * [22] D. Michaelis, M. Fröhlich, H. W. Strube, E. Kruse, B. Story, and I. R. Titze, “Some simulations concerning jitter and shimmer measurement,” in _3rd International Workshop on Advances in Quantitative Laryngoscopy, Aachen, Germany_ , 1998, pp. 744–754.
  * [23] N. Campbell and P. Mokhtari, “Voice quality: the 4th prosodic dimension,” pp. 2417–2420, 2003.
  * [24] R. E. Slyh, W. T. Nelson, and E. G. Hansen, “Analysis of mrate, shimmer, jitter, and f0 contour features across stress and speaking style in the SUSAS database,” in _Proceedings of the ICASSP_ , vol. 4. IEEE, 1999, pp. 2091–2094.
  * [25] F. Wittig and C. Müller, “Implicit feedback for user-adaptive systems by analyzing the users’ speech,” in _Proceedings of the Workshop on Adaptivität und Benutzermodellierung in interaktiven Softwaresystemen (ABIS)_ , Karlsruhe, Germany, 2003.
  * [26] X. Li, J. Tao, M. T. Johnson, J. Soltis, A. Savage, K. M. Leong, and J. D. Newman, “Stress and emotion classification using jitter and shimmer features,” in _Proceedings of the ICASSP_ , vol. 4. IEEE, 2007, pp. IV–1081.
  * [27] A. W. Zewoudie, J. Luque, and J. Hernando, “Jitter and shimmer measurements for speaker diarization,” in _VII Jornadas en Tecnología del Habla and III Iberian SLTech Workshop_ , Las Palmas de Gran Canaria, Spain, 2014, pp. 21–30.
  * [28] S. Mirzaei, M. El Yacoubi, S. Garcia-Salicetti, J. Boudy, C. Kahindo, V. Cristancho-Lacroix, H. Kerhervé, and A.-S. Rigaud, “Two-stage feature selection of voice parameters for early Alzheimer’s disease prediction,” _IRBM_ , vol. 39, no. 6, pp. 430–435, 2018.
  * [29] A. Benba, A. Jilbab, and A. Hammouch, “Hybridization of best acoustic cues for detecting persons with Parkinson’s disease,” in _2014 Second World Conference on Complex Systems (WCCS)_. IEEE, 2014, pp. 622–625.
  * [30] B.-F. Zaidi, M. Boudraa, S.-A. Selouani, D. Addou, and M. S. Yakoub, “Automatic recognition system for dysarthric speech based on MFCC’s, PNCC’s, jitter and shimmer coefficients,” in _Science and Information Conference_ , 2019, pp. 500–510.
  * [31] H. Rahali, Z. Hajaiej, and N. Ellouze, “Robust features for noisy speech recognition using jitter and shimmer,” _International Journal of Innovative Computing, Information and Control_ , vol. 11, pp. 955–963, 01 2015\.
  * [32] S. Soleymani, A. Dabouei, S. M. Iranmanesh, H. Kazemi, J. Dawson, and N. M. Nasrabadi, “Prosodic-enhanced siamese convolutional neural networks for cross-device text-independent speaker verification,” in _9th International Conference on Biometrics Theory, Applications and Systems_. IEEE, 2018, pp. 1–7.
  * [33] M. Farrús, A. Garde, P. Ejarque, J. Luque, and J. Hernando, “On the fusion of prosody, voice spectrum and face features for multimodal person verification,” in _9th International Conference on Spoken Language Processing_ , 2006.
  * [34] G. Cámbara, J. Luque, and M. Farrús, “Convolutional speech recognition with pitch and voice quality features,” _arXiv preprint arXiv:2009.01309_ , 2020.
  * [35] Y. Jadoul, B. Thompson, and B. De Boer, “Introducing parselmouth: A python interface to praat,” _Journal of Phonetics_ , vol. 71, pp. 1–15, 2018.
  * [36] P. Boersma and V. Van Heuven, “Speak and unspeak with Praat,” _Glot International_ , vol. 5, no. 9/10, pp. 341–347, 2001.
  * [37] P. Ghahremani, B. BabaAli, D. Povey, K. Riedhammer, J. Trmal, and S. Khudanpur, “A pitch extraction algorithm tuned for automatic speech recognition,” in _Proceedings of the ICASSP_ , 2014, pp. 2494–2498.
  * [38] T. Kudo and J. Richardson, “Sentencepiece: A simple and language independent subword tokenizer and detokenizer for neural text processing,” _arXiv preprint arXiv:1808.06226_ , 2018.