A Communication-Efficient Distributed Gradient Clipping Algorithm for Training Deep Neural Networks

Abstract

In distributed training of deep neural networks or Federated Learning (FL), people usually run Stochastic Gradient Descent (SGD) or its variants on each machine and communicate with other machines periodically. However, SGD might converge slowly in training some deep neural networks (e.g., RNN, LSTM) because of the exploding gradient issue. Gradient clipping is usually employed to address this issue in the single machine setting, but exploring this technique in the FL setting is still in its infancy: it remains mysterious whether the gradient clipping scheme can take advantage of multiple machines to enjoy parallel speedup. The main technical difficulty lies in dealing with nonconvex loss function, non-Lipschitz continuous gradient, and skipping communication rounds simultaneously. In this paper, we explore a relaxed-smoothness assumption of the loss landscape which LSTM was shown to satisfy in previous works and design a communication-efficient gradient clipping algorithm. This algorithm can be run on multiple machines, where each machine employs a gradient clipping scheme and communicate with other machines after multiple steps of gradient-based updates. Our algorithm is proved to have O 1 N 4 iteration complexity for finding an -stationary point, where N is the number of machines. This indicates that our algorithm enjoys linear speedup. We prove this result by introducing novel analysis techniques of estimating truncated random variables, which we believe are of independent interest. Our experiments on several benchmark datasets and various scenarios demonstrate that our algorithm indeed exhibits fast convergence speed in practice and thus validates our theory.

EPISODE: Episodic Gradient Clipping with Periodic Resampled Corrections for Federated Learning with Heterogeneous Data

“…Gradient clipping (Pascanu et al, 2012; is a well-known strategy to improve the training of deep neural networks with the exploding gradient issue such as Recurrent Neural Networks (RNN) (Rumelhart et al, 1986;Elman, 1990;Werbos, 1988) and Long Short-Term Memory (LSTM) (Hochreiter & Schmidhuber, 1997). Although it is a widely-used strategy, formally analyzing gradient clipping in deep neural networks under the framework of nonconvex optimization only happened recently (Zhang et al, 2019a;Cutkosky & Mehta, 2021;Liu et al, 2022). In particular, Zhang et al (2019a) showed empirically that the gradient Lipschitz constant scales linearly in terms of the gradient norm when training certain neural networks such as AWD-LSTM (Merity et al, 2018), introduced the relaxed smoothness condition (i.e., (L 0 , L 1 )-smoothness 1 ), and proved that clipped gradient descent converges faster than any fixed step size gradient descent.…”

“…Although there is a vast literature on FL (see (Kairouz et al, 2019) and references therein), the theoretical and algorithmic understanding of gradient clipping algorithms for training deep neural networks in the FL setting remains nascent. To the best of our knowledge, Liu et al (2022) is the only work that has considered a communication-efficient distributed gradient clipping algorithm under the nonconvex and relaxed smoothness conditions in the FL setting. In particular, Liu et al (2022) proved that their algorithm achieves linear speedup in terms of the number of clients and reduced communication rounds.…”

“…To the best of our knowledge, Liu et al (2022) is the only work that has considered a communication-efficient distributed gradient clipping algorithm under the nonconvex and relaxed smoothness conditions in the FL setting. In particular, Liu et al (2022) proved that their algorithm achieves linear speedup in terms of the number of clients and reduced communication rounds. Nevertheless, their algorithm and analysis are only applicable to the case of homogeneous data.…”

“…Nevertheless, their algorithm and analysis are only applicable to the case of homogeneous data. In addition, the analyses of the stochastic gradient clipping algorithms in both single machine (Zhang et al, 2020a) and multiple-machine setting (Liu et al, 2022) require strong distributional assumptions on the stochastic gradient noise 3 , which may not hold in practice.…”

“…In this work, we introduce a provably computation and communication efficient gradient clipping algorithm for nonconvex and relaxed-smooth functions in the general FL setting (i.e., heterogeneous data, limited communication) and without any distributional assumptions on the stochastic gradient noise. Compared with previous work on gradient clipping (Zhang et al, 2019a;Cutkosky & Mehta, 2020;Liu et al, 2022) and FL with heterogeneous data (Li et al, 2020a;Karimireddy et al, 2020), our algorithm design relies on two novel techniques: episodic gradient clipping and periodic resampled corrections. In a nutshell, at the beginning of each communication round, the algorithm resamples each client's stochastic gradient; this information is used to decide whether to apply clipping in the current round (i.e., episodic gradient clipping), and to perform local corrections to each client's update (i.e., periodic resampled corrections).…”