{'fraction_non_alphanumeric': 0.031998922897405975, 'fraction_numerical': 0.010322233192711605, 'mean_word_length': 4.507414730598121, 'pattern_counts': {'":': 0, '<': 0, '<?xml version=': 0, '>': 2, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 10, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
rTop-k: A Statistical Estimation Approach to Distributed SGD

The large communication cost for exchanging gradients between different nodes significantly limits the scalability of distributed training for large-scale learning models. Motivated by this observation, there has been significant recent interest in techniques that reduce the communication cost of distributed Stochastic Gradient Descent (SGD), with gradient sparsification techniques such as top-k and random-k shown to be particularly effective. The same observation has also motivated a separate line of work in distributed statistical estimation theory focusing on the impact of communication constraints on the estimation efficiency of different statistical models. The primary goal of this paper is to connect these two research lines and demonstrate how statistical estimation models and their analysis can lead to new insights in the design of communication-efficient training techniques. We propose a simple statistical estimation model for the stochastic gradients which captures the sparsity and skewness of their distribution. The statistically optimal communication scheme arising from the analysis of this model leads to a new sparsification technique for SGD, which concatenates random-k and top-k, considered separately in the prior literature. We show through extensive experiments on both image and language domains with CIFAR-10, ImageNet, and Penn Treebank datasets that the concatenated application of these two sparsification methods consistently and significantly outperforms either method applied alone.


I. INTRODUCTION
Stochastic Gradient Descent (SGD) and its variants have become the workhorse for training large machine learning models on ever growing datasets. Such large-scale training can be accelerated by partitioning datasets across multiple nodes and parallelizing the computation of the stochastic gradient; nodes can compute gradients in parallel based on their local datasets, and these local gradients can then be aggregated at a master node. However, for large models with millions of parameters, it is now well-understood that full-precision gradient aggregation is costly and the associated communication overhead can negate the savings in computation time [1]. Communication is an even more significant bottleneck when training is performed over potentially slow, unreliable and expensive wireless links such as in federated learning [2], where locally processing user data offers additional benefits such as privacy and personalization.
Motivated by these observations, there has been significant recent interest in developing communication-efficient SGD methods [1], [3]- [15]. These works show that the communication cost of SGD can be significantly reduced by using a variety of techniques, including thresholding and sparsification of the local gradients (e.g. communicating only the top-k gradients with largest magnitudes or k randomly selected gradients); quantization and compression of gradients (perhaps with randomization) to a small number of bits; and reducing the number of communication rounds by performing multiple iterations on the local dataset of each node (e.g. federated averaging). The general methodology in these works is to first propose a technique for reducing the communication cost of SGD, and then demonstrate its effectiveness through experimental evaluations and/or prove its convergence under standard assumptions. Such convergence guarantees have been proven in [6], [15]- [17] for various communication efficient training techniques.
In this paper, we take a different approach. We ask the following question: can we develop suitable statistical models for the stochastic gradients and use these models to inform the design of more efficient training techniques? This statistical perspective has been the focus of a recent research line in distributed estimation theory [18]- [26]. Motivated by similar observations as above, these works focus on characterizing the impact of communication constraints on the estimation efficiency of common statistical models, e.g. Gaussian mean estimation and its sparse variants, and discrete distribution estimation. However, it is unclear how to leverage these results to inform new practical training methods with SGD; these works focus on a parameter estimation framework rather the training problem, and their emphasis is on canonical statistical models such as Gaussian mean estimation, which may not accurately reflect the distribution of the stochastic gradients.
In this work, we connect these two research lines by casting each communication round of distributed SGD as a communication-constrained parameter estimation problem and propose a statistical model for the stochastic gradients. This model aims to capture the skewed and sparse distribution of the local gradients observed in experiments, which underlies the experimental May 22, 2020 DRAFT success of existing sparsification methods such as top-k and thresholding [1], [10]. These methods communicate only a few local gradients per node which have the largest magnitudes, and can reduce the communication cost to order k log d, where d is the number of parameters.
By using the toolset of distributed estimation theory, we characterize the fundamental estimation efficiency for our proposed statistical model and the optimal communication and estimation schemes that achieve this performance. The study of this statistical model naturally leads to the idea of communicating a randomly chosen subset of a set of large magnitude gradients. Instead of simply choosing the top-k gradients, the optimal communication scheme chooses a random k-subset of the gradients with large magnitudes. Even though random-k and top-k sparsification have both been separately considered in the literature, selecting a random subset of the top magnitude gradients, which corresponds to their concatenated application, is novel. We call this new sparsification strategy rTop-k. It can be observed from [17] that rTop-k readily enjoys the same convergence guarantees as top-k and rand-k since it is also a k-contraction. However, in extensive experiments on both image and language domains with CIFAR-10, ImageNet, and Penn Treebank datasets we observe that rTop-k significantly outperforms either top-k or random-k applied separately.
The contributions of our paper can be summarized as follows: • We show that each communication round of distributed SGD can be cast as a communicationconstrained statistical parameter estimation problem and that the study of such problems can inform the design of communication-efficient training schemes. To the best of our knowledge, our work is the first to connect distributed stastical estimation theory with communication efficient SGD methods.
• We propose and analyze novel statistical estimation models that capture the skew and sparse distribution of the stochastic gradients. We develop new communication and estimation schemes for these models and prove their optimality via information theoretic lower bounds.
• We introduce a new gradient sparsification scheme, rTop-k, and show that it outperforms either of its two previously known constituent schemes in extensive experiments.

A. Distributed Statistical Parameter Estimation
We begin by formulating the distributed parameter estimation problem. Let where θ ∈ Θ ⊂ R d . We are interested in estimating the parameter θ from the samples X 1 , . . . , X n .
Unlike the traditional statistical setting where samples X 1 , · · · , X n are available to the estimator as they are, we consider a distributed setting where each observation X i is available at a different node in a network and has to be communicated to a master node by using k bits. In other words, each node has to encode its sample X i by a possibly randomized strategy Π i to a k-bit string M i independently of the other nodes and send it to the master processor. The goal of the master node is to produce an estimateθ of the underlying parameter θ from the nk-bit transcript M = (M 1 , . . . , M n ) it receives from the nodes so as to minimize the worst case squared 2 risk: whereθ(M ) is an estimator of θ based on the transcript M . Note that the encoding strategies Π i for i = 1, . . . , n and the estimatorθ can be jointly designed to minimize the risk.

B. Gradient Aggregation in Distributed Training
We next formulate the gradient aggregation problem in distributed training with SGD as a distributed parameter estimation problem. Let g t = E[∇ (ω t ; (X, Y ))] denote the gradient of the population risk in R d where denotes the loss function, ω t ∈ R d denotes the weights of the network at the current iteration t, and the expectation is with respect to the true distribution of the samples (X, Y ). Assume there are n distributed nodes and each node i for i = 1, . . . , n computes the (stochastic) gradient of the loss function with respect to a small batch of i.i.d.
Note that when disjoint subsets of the dataset are assigned to different nodes, can be modeled as independent and identically distributed for different i. Hence, the local gradients g t 1 , g t 2 , . . . , g t n ∈ R d are generated independently from the same (unknown) distribution and have mean equal to g t , the true gradient with respect to the population risk. 1 To model the communication bottleneck, assume each node has k bits to communicate its local gradient vector g t i to the master node. The goal of the master node is to generate an estimateĝ t of the true gradient g t under squared 2 risk as per (1) from the nk-bit transcript M it obtains at the end of this communication round. The g t can be regarded as the parameter of the underlying distribution that we wish to estimate. In order to complete the model description, we need to specify the statistical model P θ , with θ = g t , according to which the samples g t 1 , g t 2 , . . . , g t n are generated and the set Θ ⊂ R d in which the parameter lies.

C. Statistical Models
The distributed parameter estimation problem in (1) has been studied for specific classes of statistical models P θ in the literature. One canonical model that has been of significant interest in the literature is the Gaussian location model P θ = N (θ, σ 2 I d ) with Θ = R d which is studied in [18]- [20], [22]. Adopting this Gaussian model for the gradient aggregation problem would correspond to modeling the local gradients g t i computed at each node as i.i.d. observations of the true gradient vector g t under spherically-symmetric additive Gaussian noise. However, this leads to dense observation vectors g t i , while the distribution of the local gradients is often skewed and sparse like in experiments, i.e. often relatively few entries of g t i have large magnitudes and many entries have magnitudes close to zero.
A sparse variation of the Gaussian mean estimation problem has been studied in [19], [20], [22], where Θ = {θ ∈ R d : θ 0 ≤ s ≤ d}. Adopting this statistical model for the gradient aggregation problem would correspond to modeling the true gradient vector g t as sparse, with only s non-zero components, while the local gradients g t i computed at each node would still correspond to i.i.d. observations of the true gradient vector g t under additive Gaussian noise.
Even though the true gradient vector g t is assumed to be sparse, qualitatively this model behaves similarly to the original Gaussian model. In particular the number of bits k needed to achieve the centralized estimation performance (with no communication constraints) scales linearly with the ambient dimension d as in the original Gaussian model, i.e. it is not impacted by the sparsity s.
Similarly the optimal communication scheme independently quantizes each gradient component as in the case of the dense Gaussian model.
In this work, we instead propose to focus on the following sparse Bernoulli model. Suppose While this is a highly simplified model for real stochastic gradient vectors, it allows us to focus our attention on the impact of a highly skewed distribution for the magnitudes of the gradients.
This model and our main results in the next section can be easily extended to more accurately reflect the distribution of the stochastic gradients, and doing so does not change the optimal encoding or estimation schemes. In particular, we could consider the following refinements: (i) Signed values: The mean parameters θ j could be in [−1, 1] subject to d j=1 |θ j | ≤ s, so that θ j 's can be positive or negative, with the jth component of the X i now being Sign(θ j )Bern(|θ j |).
(ii) Scaling: The Bernoulli random variables could be scaled by any M > 0 with the goal of estimating the scaled Bernoulli parameters M θ.
(iii) Continuous Perturbations: We could add a vector of continuous independent zero-mean random variables Z i to each X i , where X i is distributed according to (2) and each component . As a result, each component of the resultant Y i = Z i + X i will now take continuous values.
With none of these refinements changing the fundamental encoding and estimation schemes, we find that the simple sparse Bernoulli model captures the issues that are central to our distributed estimation problem. Note that by including just the sign and scaling changes, the resulting model also makes a good representation of coarsely quantized stochastic gradients such as with ternary quantization.

D. Theoretical Results and Discussion
We now turn to analyzing the distributed estimation problem under this model. Similarly, the constant scaling in (ii) also does not impact the scheme and the result of Theorem 1.
When continuous perturbations are present as in (iii), the scheme is modified to include a We also prove the following lower bound which shows that our proposed scheme is order optimal up to logarithmic factors. for i = 1, . . . , n dô g t ←ĝ t + rTopEncode(g t i ) end for g t ← r k ·ĝ t rTopEncode(g t i ) : I ← (indices of the r largest mag. components of g t i ) I ← (a random subset of k out of the r elements of I) g ← (g t i with indices not in I set to 0) return g alone in experiments. Intuitively, top-r sparsification allows to focus only on the most significant entries of the gradient vector and random-k sparsification allows to reduce the bias introduced by top-r sparsification, hence concatenation combines the best of the two approaches.
Note that using this algorithm, k is the final number of components that must be communicated from each node to the centralized processor. Since the index for each component can be referred to with log d bits, and the value of each component can be encoded with a constant number of bits of precision, this is also up to logarithmic factors the number of bits needed for communication.
For this reason we have used k to refer to both the number of bits allowed in the distributed statistical estimation framework, and the number of gradient components that are communicated in distributed training.

A. Experiment Settings
We validate our approach over a wide range of experiments including both image and language domains with CIFAR-10, ImageNet, and Penn Treebank datasets and using two training methods May 22, 2020 DRAFT where each node communicates the gradient vector after local training: (1) on one batch of the local data (distributed setting) and (2) for one epoch over the local data (i.e. one iteration over the local data adapting the federated setting in [7], [8]). Note that in the federated setting the number of communication rounds (per epoch) is much smaller than that in the distributed setting.
In We compare our results with the baseline setting where there is no compression, the setting proposed in [1], which uses the gradient top-k selection method, and the random-k sparsification strategy where the gradients to be communicated by each node are chosen uniformly at random [9]. We employ the warm-up strategy and exponentially increase the sparsity in the warm-up period as in [1]. We further employ the local gradient accumulation strategy in [1], which provides substantial improvement in the performance for all sparsification methods. We do not compare our strategy to stochastic quantization techniques such as [6] as our main goal is to compare the performance of different sparsification methods. Also, [17] observes that top-k selection outperforms stochastic quantization in experiments and theoretically it can reduce the communication cost to at most order √ d bits, while sparsification methods can achieve order k log d, where d is the number of parameters.

B. Image Domain
For image classification tasks, we trained ResNet-18 [27] on CIFAR-10 [28] and ResNet-34 [27] on ImageNet [29]. CIFAR-10 dataset consists of 50,000 training images, 5000 for each class, and 10000 test images, 1000 for each class. ImageNet dataset contains over 1 million training images and 50,000 validation images in 1000 classes. In all experiments, we use momentum SGD as the optimizer and cross entropy as the loss function. We split the data into batches of size 128 for CIFAR-10 and 32 for ImageNet experiments. We set the warm-up period as 5 epochs.
In our first two experiments, we train ResNet-18 on CIFAR-10 in both distributed and federated settings and show the results in Figure 2 and Figure 3 respectively. Figure 4 shows  Tables I, II, and III.
We note that in all experiments rTop-k strategy has substantially better performance compared to top-k and random-k strategies under the same compression ratio. In the federated settings in Table II and III, we observe that the accuracy of rTop-k with 99.9% compression ratio is better than the performance of top-k 99% compression ratio, corresponding to more than an order of magnitude improvement in compression ratio. Another interesting observation is that under 99% compression ratio rTop-k outperforms the baseline.

C. Language Domain
We use the Penn Treebank corpus (PTB) dataset, which consists of 923,000 training, 73,000 validation and 82,000 test words [30]. We train the 2-layer LSTM language model architecture with 1500 hidden units per layer [31] and tie the input and the output embeddings [32]. We use the same train/validation/test set split and vocabulary as [33]. We use vanilla SGD with gradient    batch size) as in [34]. We set the warm-up period to 5 epochs.
In our first experiment, we study the distributed setting where each node communicates the gradient vector after the forward-backward pass on a single batch of local data. Figure 5 shows the perplexity and training loss of the trained language model in this setting. In the second experiment, we study the federated setting where each node communicates the gradient vector after local training for one epoch over the local data. Figure 6 shows the perplexity and training loss of the trained language model in this setting. We did not use the local gradient accumulation strategy in this setting since we did not observe an improvement in the performance. In our first experiment, we observe from Figure 5 that the validation perplexity of our rTopk strategy matches the baseline with 99.9% compression ratio. The top-k strategy achieves the same performance with 99% compression ratio. We also note that the random-k strategy has substantially worse performance in this experiment. A similar set of results is obtained in the second experiment with slightly less aggressive compression levels compared to the first experiment as shown in Figure 6. In all cases, the rTop-k strategy has substantially better performance compared with the other methods at the same compression ratio. In both experiments, we observe from training losses and the corresponding test results ( Figure 5 and Table IV for the first experiment and Figure 6 and Table V for the second experiment) that the rTop-k strategy obtains a better perplexity compared to baseline while its training loss is slightly larger. This suggests that training with rTop-k strategy results in better generalization.

IV. PROOF OF THEOREM 1
To prove the desired upper bound, we describe an explicit encoding scheme , 1} k and estimatorθ that is a function of the M i = Π i (X i ) that achieves at most this error.
The encoding functions Π i are defined as follows: (i) The first log d bits of M i = Π i (X i ) encode the number of nonzero entries in X i , i.e., X i 1 .
(ii) We form a codebook for the remaining k − log d bits such that each element of {0, 1} d with at most k ones maps to a unique k − log d bit string. Note that we can set k ≥ k−log d log d . (iii) Take X i , and if X i 1 > k , we formX i by uniformly at random keeping only k out of the original X i 1 ones. If X i 1 ≤ k , thenX i = X i . We encodeX i using the codebook from step (ii).
With this encoding scheme, the estimatorθ has access to bothX i and X i 1 for i = 1, . . . , n.
For convenience, define the subsampling fraction S i by otherwise .
Our estimatorθ is now defined byθ = 1 n n i=1X i S i . This estimator is unbiased in that Finally, we compute the error: The step in (5) follows because k ≤ C 1 k log d and we are assuming k log d ≤ s, and (6) uses the second moment formula for a Poisson binomial distribution.

V. PROOF OF THEOREM 2
To prove Theorem 2, we use part of Theorem 3 from [35] (see also [26]), reproduced below, which is proved using a Fisher information argument. Recall that in the context of Fisher information, the score function vector is the gradient of the log-likelihood, i.e S θ (x) = ∇ θ log f (x|θ). Recall also that the Ψ 2 -Orlicz norm of a random variable X is defined as where Ψ 2 (x) = exp(x 2 ) − 1 , and that a random variable with finite Ψ 2 -Orlicz norm is sub-Gaussian [36]. Proof of Theorem 2. The s n lower bound comes from considering the centralized case where there is no communication constraint. We will focus on the other bound, and will restrict our attention to a subset Θ ⊂ Θ and then use the fact that In particular, let Θ = s 2d , s d d . The distribution f (x|θ) is the product of Bernoulli distributions with parameters θ i , so the score function for each component θ i is and thus S θ i (X i ) Ψ 2 ≤ N . By taking sums of scaled independent sub-Gaussian random variables [36] we get u, S θ (X) Ψ 2 ≤ c 1 max