{'fraction_non_alphanumeric': 0.0361705507272686, 'fraction_numerical': 0.012049286346411309, 'mean_word_length': 4.566502463054187, 'pattern_counts': {'":': 0, '<': 1, '<?xml version=': 0, '>': 0, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 3, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
Quantization Robust Federated Learning for Efficient Inference on Heterogeneous Devices

Federated Learning (FL) is a machine learning paradigm to distributively learn machine learning models from decentralized data that remains on-device. Despite the success of standard Federated optimization methods, such as Federated Averaging (FedAvg) in FL, the energy demands and hardware induced constraints for on-device learning have not been considered sufficiently in the literature. Specifically, an essential demand for on-device learning is to enable trained models to be quantized to various bit-widths based on the energy needs and heterogeneous hardware designs across the federation. In this work, we introduce multiple variants of federated averaging algorithm that train neural networks robust to quantization. Such networks can be quantized to various bit-widths with only limited reduction in full precision model accuracy. We perform extensive experiments on standard FL benchmarks to evaluate our proposed FedAvg variants for quantization robustness and provide a convergence analysis for our Quantization-Aware variants in FL. Our results demonstrate that integrating quantization robustness results in FL models that are significantly more robust to different bit-widths during quantized on-device inference.


Introduction
Federated Learning (FL) is a distributed machine learning paradigm, where a large number of clients, such as consumer smartphones, personal computers or smart home devices learn collaboratively. Clients train on their private local data, which is never shared with other participants in the federation, such as other clients or the server. Despite learning happening on-device, FL results in a single global model at the end of training. The privacy of local client data is an important requirement in modern machine learning and acts as a central motivator for FL. Several challenges arise in the FL setting. For example, different clients might have different computational constraints based on their hardware design specifications. One practically relevant heterogeneous hardware characteristic is the supported quantization bit-width of the hardware accelerator. A suitably trained model should exhibit no significant performance degradation after quantization to various bit-widths represented in the heterogeneous device landscape. in FL; this is done either via regularization-based methods for quantization robustness or modified quantization-aware training methods. Firstly, we propose the integration of a standard quantization robust approach known as Kurtosis Regularization (KURE), which involves a regularization term in the local clients' loss function. Secondly, we present a quantization robustness approach that involves adding random pseudo-quantization noise during the training procedure. This is motivated by the recent success of Randomized Smoothing [7] in the adversarial robustness literature that aims to learn models where the input data samples are corrupted with Gaussian noise. The adopted mechanism is also inspired by the introduction of additive pseudo-quantization noise [8] for QAT that discards the Straight Through Estimator (STE) approximation for non-differentiable uniform quantization.
Quantization-Aware Training (QAT) methods [30,18,9,23] have been successful at training quantized models with ultra-low bit-widths. QAT models perform very well for the target bit-width they have been trained on but can lead to significant degradation for other bit-widths, even for full precision [17]. To address this limitation of conventional QAT methods we introduce Multi-bit Quantization-Aware Training (MQAT), a novel QAT framework that achieves quantized models robust to multiple bit-widths without re-training.
In MQAT for FL, a random bit-width is sampled for each client from the set of considered quantization bit-widths before performing a standard QAT procedure during the client training phase. This small modification enables models trained using the federated regime to be robust to multiple bit-widths during quantized inference. Furthermore, since QAT involves certain heuristics (STE) for computing gradients (due to the non-differentiable rounding operation), we theoretically analyse the convergence behaviour of the global model in the non-convex setting when clients perform local QAT. In summary, we make the following contributions in this paper: • We introduce multiple quantization robustness methods such as Kurtosis Regularization (KURE), Additive Pseudo-Quantization Noise (APQN) into the federated learning setup to achieve quantization robust models that can be used for efficient inference at multiple bit-widths. • As the standard form of QAT integration into federated learning fails to generalise across multiple bit-width, we propose Multi-bit Quantization-Aware Training (MQAT), a new variant to achieve quantization robust models in FL. • We study the theoretical convergence properties for our proposed QAT variants in federated learning and show that the convergence rate for the server-side weights is similar to traditional FL, albeit with a QAT-method-specific error floor. • We perform extensive experimental evaluations of baselines and our FEDAVG variants on CIFAR-10, CIFAR-100, FEMNIST and TinyImageNet with different network architectures. We empirically demonstrate that our proposed modifications yield models that are robust to quantization at multiple bit-widths without significant reduction on the model's full-precision accuracy.

Preliminaries
We first provide some brief background on FEDAVG and neural network quantization robustness.

Federated Averaging
Federated learning is a distributed learning paradigm where multiple clients collaboratively learn a shared model. In this machine learning framework, the local client data is not shared with other clients or the server. The problem of federated learning can be formulated as an optimization objective where F i (w, D i ) = E ξ∼Di [f i (w, ξ)] is the the local objective function at client i, w ∈ R D represents the parameters for the global model, and P denotes a distribution over the population of clients I. The local data distribution D i often varies between clients, resulting in data heterogeneity.
Federated Averaging (FEDAVG) [22] is the standard algorithm for federated learning. It operates via a series of rounds where each round is divided into a client update phase and server update phase. We denote the total number of rounds as T . At the beginning of each round t, a subset of clients S t is sampled from the pool of clients. The server model is then shared with the sampled clients. During the client update phase, each sampled client runs local SGD for K steps with learning rate η c using their own private data. We denote the i-th client's parameters at the k-th local step of the t-th round by w i t,k . During the server update phase, the updates of the sampled clients are averaged to build the server-side update ∆ t . The server then applies that update with learning rate η s to receive the next round's parameters w t+1 . Reddi et al. [26] describe a generalization of the server-side update rule to include more advanced adaptive optimizers.

Quantization Robustness
The objective for robust quantization is to learn a single model that can be quantized to different bit-widths without significant degradation in the full precision performance. Given a neural network parameterized by w that is optimized using a standard loss function F , such as the cross entropy, quantization robustness aims at minimizing the following loss: Here,quantizer Q(·) with bit-width b and a quantization step size ∆ b , we have that where · denotes the rounding-to-nearest integer operation, and clip(·) clamps its input such that it lies in the range The quantization step size can be estimated either post-training or learnt using QAT [18,9,23]. The above objective intends to learn a neural network that is robust to various bit-widths in the quantization bit set B. B could also include the high precision 32-bit floatingpoint format (FP32). Note that the above formulation explicitly enforces robustness to different bit-widths for weight quantization only. It is straightforward to enforce quantization robustness for activations in a similar manner. Recent works [1,28] have explored alternate ways of satisfying the above objective by adding a regularization term in the standard training procedure instead of directly solving the aforementioned optimization problem.

A Federated Learning Framework with Quantization Robustness
In quantization robust federated learning we aim to solve an optimization problem of the following form: where ] is a modified local loss to encourage quantization robustness at client i and B is the set of target quantization bit-widths. Instead of directly optimizing this loss, which involves multiple forward-backward passes through the same batch for each bit-width, we incorporate various, more efficient ways for quantization robustness in the FEDAVG framework.
Regularization Method. Regularization methods such as Kurtosis regularization [28], which enforce a uniform distribution on the weight tensors, can be incorporated in the FEDAVG framework by modifying the loss function F i for each client as follows: The Kurtosis regularization term for an M -layered neural network can be expressed as Here, µ and σ are the mean and standard deviation of tensor w and K τ denotes a scalar value that defines the distribution enforced on the tensors. Similar to [28], we employ K τ = 1.8 to enforce uniform distribution. We provide the pseudo-code for FEDAVG with Kurtosis regularization, FEDAVG-KURE, in Algorithm 1.
Additive Pseudo-Quantization Noise (APQN). The quantization robustness problem has similarities with adversarial robustness in the sense that both aim to keep predictions unaltered in the presence of some form of bounded additive noise. Adversarially robust models aim to be robust to noised-up input, whereas with quantization robustness the noise is added to either the weight tensor or the intermediate activations. We present a quantization robustness approach that involves adding random pseudo-quantization noise during the training procedure. This is motivated by the recent success of Randomized Smoothing [7] in the adversarial robustness literature, where a model is learnt with input data samples corrupted with Gaussian noise. An adaptation similar to ours has been presented in the recently introduced Differentiable Quantizer [8] as a replacement of the commonly used Straight Through Estimator (STE) based quantizer. Their proposed quantizer has only been used to achieve quantized models for a single target bit-width or mixed-precision with fixed computational budget. Since APQN involves additive pseudo quantization noise and the noise does not have any learnable parameters, backpropagation for the rest of the parameters is straightforward.
We propose to learn models that are robust to varying levels of quantization noise and thus can be quantized to different bit-widths. The local loss function in this case can be reformulated as Here,Q(·) is a pseudo-quantizer with bit-width b that adds noise sampled from the uniform distribution U[−∆ b /2, ∆ b /2], and can be defined as Provided that ∆ b is sufficiently large, the sampled pseudo-quantization noise can have support for various target bit-widths, thus encouraging quantization robustness. The pseudo-code for FEDAVG-APQN is provided in Algorithm 2.
QAT and MQAT. A standard way of learning a network with low bit-widths is to constrain the parameters and/or activations of the model to fixed quantization levels. This procedure of training a neural network with standard Projected Gradient Descent (PGD) algorithm with quantization function as projection, is often termed Quantization-Aware Training (QAT). To perform "quantization-aware" FL, we can adopt the QAT procedure for the local optimization at each client to learn a global model that can be quantized to a specific bit-width. In this "quantization-aware" FL, the client-level loss function can be reformulated as: The quantization step-size ∆ b can be either learnt as a parameter [9] or be estimated before the start of training and kept fixed thereafter. Note that the backward pass of the network involves a gradient estimate through the non-differentiable rounding operation of Q(·). To this end, similar to previous QAT literature we use the standard STE approximation [3], which approximates the gradient of the rounding operator as 1: By using this STE approximation for QAT, we can train models that can be quantized to specific bit-widths. In order to use the trained model for bit-widths other than b, we can analytically estimate the quantization step-size using the ranges for bit-width b as follows: Here, a refers to any target bit-width during inference stage and b refers to bit-width for which the model is trained on. The pseudo-code for FEDAVG-QAT is provided in Algorithm 2.
Although QAT trains models that perform favourably at specific bit-widths, they often suffer from performance degradation when quantizing to other bit-widths [17]. For this reason, QAT alone is not suitable for tackling the heterogeneous hardware requirements that one can encounter in a cross-device FL deployment. We propose MQAT to realize quantization-robust FL. MQAT aims to learn models that are robust to a set of bit-widths B, by selecting a bitwidth b ∈ B either randomly during local training or by fixing it on the basis of each client's hardware requirements. The sampled bit-width is then kept same for all the layers.
Similar to QAT, the quantization step-size ∆ b for different bit-widths can be either learnt or estimated before the start of training and kept fixed thereafter. The quantization step-size for different bitwidths is then shared along the model parameters with all clients. We provide the pseudo-code for FEDAVG-MQAT in Algorithm 2.
sample a subset S of clients 4: for all i ∈ S in parallel do 5: w i t,0 ← w t {broadcast server state to client} 6: for k = 0, . . . , K − 1 do 7: 10:  for all i ∈ S in parallel do 5: w i t,0 ← w t {broadcast server state to client} 6: b ← U[B] 7: for k = 0, . . . , K − 1 do 8: 10: w t+1 ← w t + η s · ∆ t {server update} 16: end for Convergence analysis. APQN, QAT and MQAT modify the forward pass of the model, either by adding uniform noise in the case of APQN or approximating the gradient of the non-differentiable rounding operation in the case of QAT and MQAT. Therefore, it is important to study how these modifications / heuristics affect the convergence behaviour of the global model in the federated setting. What follows now is a convergence analysis in the non-convex setting with the help of the following assumptions: Assumption 1 (Lipschitz Gradient). Each local loss function F s is L-smooth ∀s ∈ S, i.e., ∇F s (x) − ∇F s (y) ≤ L x − y , ∀x, y ∈ IR D . Assumption 2 (Bounded variance). Each F s has bounded local variance, i.e., E[ ∇f s (w, ) − ∇F s (w) 2 ] ≤ σ 2 l , where f s is a stochastic estimate of the local loss based on a w ∈ IR D and is a random mini-batch. Furthermore, the global variance is also bounded, i.e., 1 S s ∇F s (w) − ∇F (w) 2 ≤ σ 2 g , ∀w ∈ IR D . Assumption 3 (Bounded quantization noise). Let w be a variable to be quantized, j be any of its dimensions and Q(·) the quantization operation. We assume that the noise added in order to quantize w j , i.e., r j = Q(w j ) − w j is bounded by half the step size of the quantizer, i.e., ∆ 2 .
The first two assumptions are common in the non-convex optimization literature for FL [26] and the third is automatically satisfied in our quantization schemes, provided the ranges are set up appropriately. Based on these assumptions, we can then make the following statement.
we have that the FEDAVG-{APQN,QAT,MQAT} server updates satisfy where we define and The proof is delegated to appendix due to space constraints. It follows [26] while handling quantization noise, such as [20]. We can see that the convergence rate for the (non-quantized) server side weights is similar to traditional FL [26], albeit with an additional error floor due to the quantization noise. Some of this quantization error can be reduced by decreasing the learning rates, up to an irreducible factor of O(L 2 DR 2 ).

Related Work
Quantization-aware training. QAT is one of the more effective and widely used methods for achieving low-bit weight and activation quantization. It relies on simulating the quantization operation operation during training and requires access to labelled training data. Esser et al. [9] first introduced the idea of learning the quantization step-size jointly with the weight achieving near floating-point accuracy in ResNets even with 3-bit quantization. Since then, further advances in quantization-aware training [11,15,10,4,25] have pushed the envelop and enabled ultra low-bit quantization (2-4 bits) for a wide range of networks and tasks.
Quantization robustness. A drawback of QAT is that it can make the trained model highly dependent on the chosen bit-width and quantization parameters. To address this, robust quantization aims at training a single set of weights that are robust to a wider range of quantization choices and bit-widths. Alizadeh et al. [1] model quantization noise as an additive perturbation and show that they can improve quantization robustness by regularizing the L 1 -norm of gradients. Since this type of gradient regularization is computationally expensive due to the second-order gradient information, a simpler alternative regularization procedure was proposed in follow-up work [28]. Shkolnik et al. [28] trains the network with Kurtosis Regularization (KURE) on the weights to improve robustness and use the LAPQ [24] algorithm to find the optimal quantization parameters post-training. It has been shown in [28] that the uniform distributions on the weight tensors achieve better quantization robustness, instead of the Gaussian-like distributions attained during the standard training procedure.
Recently, Defossez et al. [8] introduce additive uniform noise to simulate quantization during training for the purpose of QAT; this method can be easily extended for the purpose of quantization robustness as we show in this paper. Jangho et al. [17] introduces a training method that leads to more "quantization-friendly" weights, by scaling the gradient depending on the distance of the weights from the quantization grid.
Quantization in Federated Learning (FL). Quantization in the context of FL has been studied mostly in the context of compressing communication, e.g., [2,27,29]. Equally important to reducing the communication overhead is the reduction of the computational cost of training and inference on-device. Existing work in this direction can roughly be divided into designing more efficient models through e.g. sparsification [6,16,21] and more effective algorithms [26,12] that reduce the number of training rounds. Structurally sparse models can reduce training and inference costs ( [13] targets heterogeneous computational resources) and may additionally reduce communication [21]. Quantization, the focus of this work, is orthogonal to these and as of yet understudied in the literature.
To the best of our knowledge, this work is the first to investigate the robustness of models to different levels of quantization in the FL setting. Here, an abbreviation of "W" indicates that we performed weight-quantization only, whereas "W-#" refers to quantization at # bits.

Experiments
We evaluate the quantization robustness of different proposed variants of FEDAVG using standard benchmarks in FL. In this work, we mainly focus on weight quantization robustness of various trained models but we also provide additional experimental comparisons for activation quantization and quantizing both weight and activations. For all the results in the paper, we present the accuracy plots at different quantization bit-widths and we refer the reader for exact numbers to Appendix.
Experimental Setup. For the experimental comparisons, we use federated versions of the CIFAR-10, CIFAR-100 [19], TinyImageNet 2 and FEMNIST [5] datasets. We use LeNet-5 and ResNet-20 architecture for evaluation on CIFAR-10 and LeNet-5 on FEMNIST. For experiments with CIFAR-100 and TinyImageNet, we use ResNet-20 and ResNet-18 network architectures respectively. We replace batch-norm with group-norm in ResNet-20 since batch-norm induces training instabilities in FL [26]. We report the accuracy of the quantized server-side model whenever comparing the models at different quantization bit-widths. To compare quantization robustness, we use bit-widths from the set {32, 8, 6, 4, 3, 2}. We provide the details regarding training hyperparameters and datasets in Appendix. For the APQN and KURE methods, we employ post-training quantization methods to set the quantization ranges. Specifically, we find ranges that minimize the mean-squared error [23]. For QAT based methods, range estimation is done at the start of the training and then the same ranges are used at inference time. Our code is written in PyTorch and all our experiments are performed using NVIDIA Tesla V100 GPUs.

Weight Quantization
Firstly, we compare all weight-only quantization-robust FEDAVG variants introduced in this work, i.e., we keep the activations in full precision. We report validation accuracy at the different levels of weight quantization for various datasets and network combinations in Fig. 1.
CIFAR-10. Fig. 1a shows the performance of FEDAVG variants on CIFAR-10 with ResNet-20. We see that the classification acurracy for the baseline FEDAVG trained model drops considerably when weights are quantized to low bit-widths (2 and 3). Both, FEDAVG-KURE and FEDAVG-APQN perform similarly to the baseline with only slight improvements at low bit-widths. Our FEDAVG-QAT variants outperforms the other FEDAVG variants. A signficant drop in validation accuracy is observed as the FEDAVG-QAT models are quantized to bit-widths other than the target bit-width they have been trained for. This is a known issue with QAT trained models in the context of quantization robustness. Our proposed FEDAVG-MQAT variant directly targets this issue. The FEDAVG-MQAT model outperforms all other variants at different level of quantizations consistently. Furthermore, it improves validation accuracy at full precision as well. Further investigation revealed that for an over-parameterized network such as ResNet-20, the FEDAVG baseline overfits the training set. Our proposed FEDAVG-MQAT implicitly regularizes the FL model and avoids the issue of overfitting in this experimental setup and thus achieves better full precision accuracy. We further investigate overfitting and the implicit regularization phenomenon of FEDAVG-MQAT in Appendix. Fig. 1b shows performance on CIFAR-10 with the LeNet-5 architecture. Since LeNet-5 is a relatively small network, no overfitting is observed for the FEDAVG baseline. It is important to note that the LeNet-5 (achieves 48-55 % accuracy at 2 bits) is more robust to weight quantization compared to ResNet-20 (achieves 28-35% accuracy at 2 bits) for the baseline FEDAVG trained model on CIFAR-10. Similar to previous comparisons, FEDAVG-KURE and FEDAVG-APQN achieve marginal gains at low bit-widths. The FEDAVG-MQAT produces better validation accuracy (improvement of ≈ 12% at ultra low bit-width of 2 bits but with considerable loss (≈ 3%) in accuracy at full precision. We believe this is because a small network, such as LeNet-5 on CIFAR-10, is inherently hard to compress (quantize) and the gains at low bit-widths come at the cost of considerable degradation in full precision accuracy.
CIFAR-100. Fig. 1c shows the performance of our FEDAVG variants with weight-only quantization on CIFAR-100 using the ResNet-20 architecture. Similar to our CIFAR-10 setup, we observe a large drop in the baseline model accuracy at low bits. Our FEDAVG-MQAT variant achieves significant gains, especially at low bit-widths; ≈ 35% at 2-bits and ≈ 10% at 3-bits, compared to the baseline FEDAVG model. It is important to note that these improvements on low bit-widths come at the cost of a small loss of accuracy (≈ 2%) at full precision in comparison to the baseline.
TinyImageNet. We also performed an experimental evaluation on a more challenging FL setup; classification on TinyImageNet with the ResNet-18 architecture. The results for all of our proposed FEDAVG variants can be seen in Fig. 1d. Despite being the most challenging task with 200 classes and only 500 training samples for each class, our FEDAVG-MQAT variant remains robust to bit-widths as low as 2-bits without any loss in full-precision accuracy. In comparison to the baseline, it improves the 2-bit, 3-bit quantization accuracy by a significant margin (20 − 35%). We would like to point out that despite being trained for different set of bit-width, our FEDAVG-MQAT variant can achieve accuracy of FEDAVG-QAT for their target bit-widths while preserving the full precision model accuracy.
FEMNIST. The FEMNIST dataset has become one of the standard datasets used to evaluate FL algorithms. For the sake of completeness, we performed quantization robustness experiments on FEMNIST using the LeNet-5 architecture. We observe that the baseline FEDAVG trained model is already robust to various quantization levels, even up to 2-bits. Compared to our other tasks, we believe classification on FEMNIST with LeNet-5, is relatively easier and more robust to quantization.

Activation Quantization
To further demonstrate the effectiveness of our FEDAVG variants, we analyze the task of quantizing activations on CIFAR-10 and CIFAR-100 with the ResNet-20 architecture. For FEDAVG-KURE, we impose the regularization term on the activations and, in a similar manner, for FEDAVG-{APQN, QAT and MQAT}, the noise / quantizer is on the activations. Here, an abbreviation of "A" indicates that we performed activationquantization only, whereas "A-#" refers to quantization at # bits. Fig. 2a and Fig. 2b show the validation accuracy at different bit-widths for CIFAR-10 and CIFAR-100 respectively. For both CIFAR-10 and CIFAR-100, we observe considerable decline in the full precision model accuracy after Kurtosis regularization on activations compared to the baseline. It sould be noted that the original work on KURE [28] considered weight-quantization only. FEDAVG-MQAT achieves significant gains at 2-bit (≈ 55% on CIFAR-10 and ≈ 24% on CIFAR-100) and 3-bit (≈ 26% on CIFAR-10 and ≈ 6% on CIFAR-100) quantization for both CIFAR-10 and CIFAR-100. Perhaps surprisingly, we can see that our FEDAVG-MQAT trained model outperforms the respective FE-DAVG-QAT models trained on their own respective bitwidth for 2 and 4-bits. For CIFAR-100 dataset, we observe a considerable decrease in model accuracy for FEDAVG-MQAT variant at higher bit-widths (and full precision).

Weight and Activation Quantization
The combination of activation and weight quantization promises to fully harness the advantages of specialized hardware accelerators. Fig. 3 illustrates the impact of our quantization robustness variants for joint weight and activation quantization on CIFAR-10 with the ResNet-20 architecture. In this setting, both, weights and activations are quantized before each matrix multiplication. We see that FEDAVG-MQAT outperforms the FEDAVG baseline as well as the bit-specific FEDAVG-QAT across all bit-widths, including full-precision. As noted before, ResNet-20 exhibits overfitting on CIFAR-10. Thus, we include another baseline which introduces dropout (50%) before the final fully-connected layer, which marginally improves performance.

Discussion
Real-world deployments of FL, necessarily require catering to a heterogeneous device landscape. Quantization-robust server models, i.e., models that have robust performance on arbitrary target bit-widths, are a step towards effectively navigating such a landscape. In this work, we introduced several variants of the FEDAVG algorithm that encourage quantization-robustness for the server model. Experimentally, we demonstrated that our FEDAVG variants can achieve good performance on several target bit-widths, without significant accuracy degradation for the full precision model. No method clearly outperforms in all settings, although we see MQAT performing well in most situations, especially for lower bit-widths. Theoretically, we showed that quantization-aware local training on the clients provides a convergence rate that is similar to traditional FL, albeit with an extra error floor that depends on the parameters of the quantization procedure and the model characteristics.
In the future, we would focus on end-to-end quantized training which can enable efficiency for both client-server communication as well as on-device training. Further axes of heterogeneity than bit-width should be considered, such as non-uniform quantization, different quantization strategies (symmetric vs. asymmetric) and more subtle differences in inference engines across devices. We believe that the scenario where a client's hardware characteristics are considered fixed throughout training opens up the possibility for advanced client sub-sampling strategies and necessitates a discussion on the impact of client-specific model performance as a trade-off between efficiency and representation. Device capabilities correlate with the heterogeneous socioeconomic landscape of participating devices and each client's dataset's influence on the learned function will be different.

Appendix A Convergence analysis for quantization aware local training in federated learning
In this appendix we provide a convergence analysis for quantization-aware training in combination with federated learning. We follow the analysis presented in [26], while handling the quantization noise, such as [20]. This appendix tries to be verbose so as to be easy to follow along. Theorem 3.1 contains the formal claim for this proof. On a high level, we will aim to upper-bound the gradient magnitude of the FL objective (Eq. 1 in the paper), ∇ w F (w t ) for different perspectives on t. We begin the proof by making explicit assumptions about the objective as well as the nature of quantization. In terms of techniques we rely only on some general inequalities that we detail before the actual proof. We advise the studious reader to keep a second copy of the pdf open to refer back to these inequalities. Furthermore we preface this section by re-iterating notation.

A.1 Notation and definitions
w ∈ IR D . . . is the flattened vector of model parameters.  E t [F (w t+1 )] . . . is the loss evaluated on parameters w at round t + 1, in expectation over the randomness at round t that influences the transition of parameters from round t to round t + 1.

E[F (w)]
. . . is the loss evaluated on parameters w, averaged over all possible sources of randomness. Precise definition given in context below. ∇ w F (w) . . . is the gradient of F with respect to w evaluated at w.
We use the shorthand ∇F (w), similarly for ∇F s (w) and ∇f s (w). ∆ b . . . is the step size of the quantizer / magnitude of the quantization noise for a given bit-width b.

A.2 Assumptions
Assumption 1 (Lipschitz Gradient). Each local loss function F s is L-smooth ∀s ∈ S, i.e., Assumption 2 (Bounded variance). Each F s has bounded local variance, i.e., E[ ∇f s (w, ) − ∇F s (w) 2 ] ≤ σ 2 l , where f s is a stochastic estimate of the loss based on an w ∈ IR D and is a random mini-batch. Furthermore, the global variance is also bounded, i.e., 1 S s ∇F s (w) − ∇F (w) 2 ≤ σ 2 g , ∀w ∈ IR D . Assumption 3 (Bounded quantization noise). Let w be a variable to be quantized, j be any of its dimensions and Q(·) the quantization operation. We assume that the noise added in order to quantize w j , i.e., r j = Q(w j ) − w j is bounded by half the step size of the quantizer, i.e., ∆ 2 .

A.3 Auxiliary lemmata / inequalities
In this section we will provide some inequalities and lemmata that will be useful for the proof of our main theorem.
Lemma A.3. For random variables z r , . . . , z r we have that Proof. The proof follows from expanding the square and applying Lemma A.1 to each 2z i z j term with γ = 1.
Lemma A.4. Let r be the quantization noise added to w satisfying assumption Eq. (3). When performing QAT, MQAT, APQN we have that where Proof. For QAT we have that due to the bounded quantization noise assumption. For MQAT where we consider multiple bitwidths, it suffices to pick the largest ∆ b to satisfy this upper bound. For APQN we have that Lemma A.5. For any local learning rate η c ≤ 1 10LK , we can bound the difference between the local shadow weights and the server weights for any k ∈ {0, . . . , K − 1} and K ≥ 1 at a given federated training iteration t via Proof. We begin by noting that where we introduced several shorthand notations for easier manipulation of the terms. By expanding the norm using the multinomial theorem we have that We can now use the fact that the expectation of b is zero, since ∇f s (w t s,k−1 + r t s,k−1 ) is an unbiased estimate of ∇F s (w t s,k−1 + r t s,k−1 ). In this way, we can simplify the above to We can now use Lemma A.1 with a γ = 2K − 1 in order to "split" the 2E[a T (η c (c + d + e))] term Following that, we can see that several terms cancel, due to Finally, we will apply Lemma A.3 in order to split E[ c + d + e 2 ] and thus end up with where we replaced the shorthand notations with their original terms. To proceed, we will make use of assumptions 2 and 1 to arrive at We can now make use of corollary A.2 with a γ = 3 in order to separate the quantization error r s,k−1 from the difference between the local (shadow) weight and the server weight and then use Lemma A.4 in order to bound the squared norm of r s,k−1 To proceed and make further use of our assumptions, we will average the aformentioned inequality over the clients and thus have that and then make use of assumption 2 in order to bound the "global" variance Finally, given our assumption that η c ≤ 1 10LK , we have that (1 + 1 2K−1 + 24Kη 2 c L 2 ) ≤ (1 + 1 K−1 ) and thus we can simplify the upper bound even further We have now arrived at a point where the average difference between the local shadow weight at iteration k and the server weight is upper bounded by two things; the average difference at iteration k − 1 along with some constant terms that are independent of the actual weights or iteration. We can thus continue further by sequentially applying the bound at Eq. 36 on each weight difference, up until we end up at the server weight w t (since local optimization started from that point) where the difference is zero. Notice that each application of this bound "adds" additional non-negative terms, so we have that the upper bound of K iterations would upper bound the bound on K − 1 iterations. Therefore, the "worst-case" upper bound is the one where k = K. In this case, we can unroll the recursion and have that To simplify even further, we can use the fact that (1 + 1 K−1 ) j is monotonic in j and that there are K terms in the sum, thus get and since (1 which proves our claim.

A.4 Proof of Theorem 3.1
We begin by noting that the server-side update rule of the model in the case when the clients perform Quantization Aware (QA) SGD with a learning rate of η c and the server does SGD with a learning rate η s . The extension of the proof to more involved server-side update rules as in [26] is straightforward.
where s indexes the clients and S is the total number of clients. k indexes the local client iteration number, and we assume that there are K local iterations in total. w t sk ∈ IR D corresponds to a real valued local shadow weight at iteration t and r t sk ∈ IR D corresponds to the quantization noise that is added to it in each iteration of the local optimization (as the weights are rounded / noised before the forward pass).
Using the L-smoothness of the global loss function, We then take an expectation over all randomness at time step t, and work towards upper bounding the T 1t , T 2t terms separately.
We will now work towards upper bounding T 3t Now by using Lemma A.1 with γ = K we have that By then using Lemma A.3 to push the squared norm inside the sum and from the Lipschitz gradient assumption We will now once more use Lemma A.3 to separate the shadow weight difference from the quantization noise and from Lemma A.4 We see that in order to proceed, we need to upper bound the difference between the local shadow weight at any iteration k and the server weight. This is were we will use our Lemma A.5 in order to proceed Finally, we can make use of our assumption η c ≤ 1 10LK which leads to 24η 3 In this way, we can arrive at our final bound for T 3t Now we can apply this bound to T 1t in order to get Which is the final upper bound on T 1t .
Bounding T 2t We begin by noting that and then we can split the squared norm via Corollary A.2 with γ = 1 To continue, we will move towards upper bounding T 4t by expanding the terms inside the squared norm We will then apply Lemma A.3 in order to move the squared norm inside the sums and will apply Corollary A.2 with γ = 1 in order to split the norm (66) To proceed, we will make use of our assumptions 1, 2 in order to get and we will apply Corollary A.2 with γ = 1 one more time in order to split the norm of the weight difference and the quantization error so that we can apply Lemma A.4 in order to bound the latter By observing the above, we see that we again end up with the average difference between the shadow weights at each iteration k and the server weight. As a result, we can apply Lemma A.5 in order to proceed further ≤ 2η 2 (71) In order to simplify the aforementioned inequality we will make a use of our assumption on η c , namely that η c ≤ 1 10LK . In this way, we will have that 16K 3 η 4 Taking these into account, we have that and due to 4 + 3 2 < 6 which constitutes our final upper bound on T 4t . With this bound at hand, we can move back to bounding T 2t and thus get Having bounded T 1t and T 2t , we can now apply these bounds to the inequality at Eq. 43 and thus get To simplify the aforementioned bound, we can once again make use of our condition η c ≤ 1 10LK which leads to 96η 4 c K 4 L 2 ≤ η 2 c K 2 . In this way, we get that where we introduced several shorthand notations for easier manipulation of the inequalities. We can thus now re-arrange the inequality to In order to consider the entire training trajectory, we will use a telescoping sum, i.e., we will sum this inequality over all rounds and take the expectation at each time-step In doing that, most of the terms on the left-hand-side will cancel across subsequent time-steps and thus we will end up with We can now re-arrange the terms to get and by considering w * to be the parameters of lowest loss, In order to proceed, we have to impose a condition on A, namely that it has to be positive (otherwise, dividing by A reverses the inequality). For this to happen we need that Assuming that this condition is satisfied, we have that  Table 1: Hyperparameters used for the experimental evaluations in the paper. Here η s , η c denote the server and client learning rate and s refers to the correction term in value in ADAM optimizer of server.
Finally, in order to obtain the result of Theorem 3.1, we make use of the fact that min 1≤t≤T ∇F (w t ) 2 ≤ 1 T T t=1 ∇F (w t ) 2 and thus arrive at which completes the proof.
In order to simulate a non-i.i.d. data split that we would expect in the federated scenario, we artificially split CIFAR-10, CIFAR-100 and TinyImageNet by their label. For CIFAR-10 and TinyImageNet, the label proportions on each client are computed by sampling from a Dirichlet distribution with α − 1.0 [14]. For CIFAR-100 we use the coarse labels provided with the dataset and follow [26]. For our FEMNIST experiments, the federated split is naturally determined by the writer-id for each client.
We provide the final hyperparameters used for all the experiments in Table 1. Since the grid-search for FL is expensive, we tune the hyperparameters such as client learning rate (η c ), server learning rate (η s ) and epsilon term ( s ) used in ADAM optimizer only for the baselines and then use the same set of hyperparameters for all our proposed variants. For FEDAVG-KURE, we tuned λ from the grid [1e + 1, 1e + 0, 1e − 1, 1e − 2, 1e − 3, 1e − 4], found λ = 1e − 1 to be optimal and use that for all our experiments. We also provide the details about bit-set used to perform FEDAVG-MQAT variants in Table 2.

C Additional Results
Per-client fixed bit-width. Performing a forward-pass in low bit-width during Quantization-Aware Training (QAT) does have the additional benefit of reduced computational requirements also during training. A good assumption to make is that each client implements efficient hardware acceleration for a specific bit-width that remains constant during its participation in the federated learning process. While in the main experimental setting we considered per-round sampling of the bit-width for MQAT, here we sample a client-specific bit-width at the beginning of training and keep it fixed throughout.
As we can see in Table 3, ex-ante fixed bit-widths lead to some degradation in performance, albeit at the aforementioned benefit of spead-up training.   Table 3: Global validation accuracy after quantization at various bit-widths for different FEDAVG variants trained on CIFAR-10 dataset using ResNet-20 architecture. Here, * indicates the client-specific bit-width is chosen at the begining of training and then kept fixed throughout.
Implicit Regularization of MQAT. As discussed in the main paper, ResNet-20 on CIFAR-10 exhibits overfitting, which we further investigate here. It was observed in the maint ext that FEDAVG-MQAT has an implicit regularisation effect and thus achieves better validation accuracy for the full-precision model by avoiding overfitting despite that not being the main objective. We show experimental comparisons of our proposed MQAT to other standard methods of regularization with different strength of weight decay (measured by regularisation term λ WD ) and drop-out before the last full-connected layer of the network in Table 4. Despite not being its primary objective, our FEDAVG-MQAT variant achieves considerable gains in full-precision accuracy in comparison to standard approaches used to avoid overfitting. Fig. 4 shows how the validation loss of the baseline is increasing after 2.5k rounds and the corresponding validation accuracy on the right. For MQAT, we observe no overfitting according to the validation loss.   Table 4: Global validation accuracy for full-precision models learnt in federation using FEDAVG variants trained on CIFAR-10 dataset with ResNet-20 architecture. Here W indicates the client-specific bit-width is chosen at the begining of training and then kept fixed throughout. Here, an abbreviation of "W", "A" and "WA" indicate weight quantization, activation quantization and both weight-activation.  Table 5: Global validation accuracy after weight quantization of various quantization robustness variants for Federated Averaging (FEDAVG) on federated version of CIFAR-10 dataset using ResNet-20 architecture.
Detailed results from the main text In this section, we provide the exact values used for plotting various figures in the main text in Table 5 Table 9: Global validation accuracy after activation quantization of various quantization robustness variants for Federated Averaging (FEDAVG) on federated version of CIFAR-10 dataset using ResNet-20 architecture.  Table 10: Global validation accuracy after activation quantization of various quantization robustness variants for Federated Averaging (FEDAVG) on federated version of CIFAR-100 dataset using ResNet-20 architecture.