Quantization Robust Federated Learning for Efficient Inference on Heterogeneous Devices
Authors: Kartik Gupta
Marios Fournarakis
Matthias Reisser
Christos Louizos
Markus Nagel
Venue: Accepted by TMLR ()

Abstract: 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.

Published in Transactions on Machine Learning Research (09/2023)
Quantization Robust Federated Learning for Efficient Infer- ence on Heterogeneous Devices
Kartik Guptal; Marios Fournarakis', Matthias Reisser!, Christos Louizos!, Markus Nagel
Australian National University, & Qualcomm AI Research
Kkartik gww.ikkkuuptaaau..mmm au
[[ffurrnr, www.meerre clouizos, hat.a//s..maaaa/oo.....ula......
Reviewed on OpenRevicw: http:://oppaoorrre..../////ooood
Abstract
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 learn- ing is to enable trained models to be quantized to various bit-widths based on the energy needs and heterogencous hardware designs across the federation. In this work, we introduce multiple variants of federated averaging algorithm that train neural networks robust. to quan- tization. Such networks can be quantized to various bit-widths with only limited reduction in full precision model accuracy. We perform extensive experiments standard FL bench- marks to evaluate our proposed FEDAVG variants for quantization robustness and provide a convergence analysis for our Quavtization-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.
1 Introduction
Federated Learning (FL) is 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 an important requirement in modern machine learning and acts as 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. Quantization robust.ness aims lo achieve this objective where a single model trained with the constraint of robustness to various quantization bit-widths on-the-go at inference time without re-training finetuning.
Recent works (Shkoliik et al. (2020); Alizadeh et al. (2020 Kim et al. 2020 introduced novel ways to train quantization robust models in the centralized training setting, but the application of such quantization robustness mechanisms has not been a focus in FL yet. In this work, we address the problem of learning quantization robust models trained using the standard federated learning algorithm, known as Federated
* Work done during internship Qualcomm AI Research. Qualcomm Research is initiative of Qualcomm Technologies, Inc. and/or its subsidiaries.

---

Published in Transactions on Machine Learning Research (09/2023)
Averaging (FEDAvG) (McMahan et al. (017 To this end, we introduce multiple variants of FEDAVG algorithms to incorporate quantization robustness in FL; this is done either via regularisation-based meth- ods for quantization robustness or modified quantization-awars training methods. Firstly, we propose the integration of a standard quantization robustness approach known as Kurtosis Regularization (KURE), which involves a regularization lerm the local clients' loss function, enforce uniform distribution on the weights and activations. Secondly, we present a quantization robustness approach that involves adding random pseudo-quantization noise during the training procedure. The adopted mechanism is also inspired by the introduction of additive pseuudo-uantizztionn noise (Défossez et al. (2021)) for QAT that discards the Straight Through Estimator (STE) approximation (Courbariaux et al. (2015)) for non-differentiable uniform quantization. Straight Through Estimator (STE) approximation is de-facto method in QAT, that
allows the backpropagation through a non-differentiable quantization operalor by assuming it as an identity function.
Quantization-Awar Training (QAT) methods (Zhou al. (2016): Krishnamoorthi (2018) ; Esser al. 20020; Nagel et al.] (2021)) 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 (Kim et al. 2020). To address this limitation of conventional QAT methods, we further introduce Multi-bit Quantization AiaantizationAAware 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 cousidered quantiza- tion 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.
Below we summarize the contributions in this paper:
We introduce multiple quantization robustness methods such as Kurtosis Regularization (KURE), Addi- tive
Noise (APQN) into the federated learning setup to achieve quantization robust models that can be used for efficient inference at multiple bil-widths.
As the standard form of QAT integration into federated learning fails to generalise across multiple bit- width, we propose Multi-bit Quantization-Awar Training (MQAT), a new variant to achieve quantization robust models in FL.
We study the thcoretical convergence properties for our QAT variants in federated learning and show that the convergence rate for the server-side weights is similar to traditional FL, albeit with QAT- method-specific error floor.
We perform extensive experimental evaluations of baselines and our FEDAVG variants CIFAR-10, CIFAR-100, FEMNIST and TinyTmaged with different network architectures. We empirically demon- strate 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.
2
Preliminaries
We first provide some brief background on FEDAVG and neural network quantization robustness.
2.1
Federated Averaging
Federated learning is 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 an optimization objective
min F(w) = E.- (F(W),
(1)
2

---

Published in Transactions on Machine Learning Research (09/2023)
where F((w, D;) BgwD, (ww, is the the local objective function at client i, RR represents the parameters for the global model, and P denotes distribution over the population of clients I. The local data distribution D, often varies between clients, resulting in data heterogeneity.
Federated Averaging (FEDAvG) [McMahan et al. (2017) ) 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, 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 Te 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 wi, During the server update phase, the updates of the sampled clients are averaged to build the server-side update A, The server then applies that update with learning rate T to receive the next round's parameters W++1- Reddi et al. (2020) describe a generalization of the server-side update rule include more advanced adaptive optimizers.
2.2 Quantization Robustness
The objective for robust quantization is learn single model that can be quantizod to different bit-widths without significant degradation in the full precision performance. Given neural network parameterized by that is optimized using standard loss function F, such as the cross entropy, quantization robustness aims at mininizing the following loss:
min
Σ P(Q(w,b)) D). BEB
(2)
Here, quantizer Q(-) with bit-width b and quantization step size As, we have that
w A 20-1 1-2--120-1-11
111
Q(w, b) As - clip
(3)
where 1-1 denotes the rounding-t -eaarissstt integer operation, and clip(-) clamps its imput such that it lies in the rangej-22 2 1]. The quantization step size can be estimated either post-training or learnt using QAT Krishnamoorthi (2018); Esser et al. (2020); Nagel et al. (2021)). The above objective intends to learn neural network that robust various sit-widths in the quantization bit set B. B could also include the high precision 32-bit floating-point 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 (Alizadeh et al. (2020 Shkolnik et al. (2020)) have explored alternate ways of satisfying the above objective by adding a regularization term in the standard training procodure instcad of directly solving the aforementioned optimization problem.
3 Federated Learning Framework with Quantization Robustness
In quantization robust federated learning we aim to solve an optimization problem of the following form:
min F(w) = [. (w)),
(4)
where (w, D₂) BEND. If is modified local loss to encourage quantization robustness at client i and B is the set of larget quantization bit-widths. Instead of directly optimizing this loss, which involves multiple forward- backwadd 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 Kurtosis regularization Shholnik et al. (220)) which enforce uniform distribution on the weight tensors, can be incorporated in the FEDAVG framework by modifying the loss function F; for each client as follows:
F (w, D.) EE-D, (w,, {)] + A LKURE(w).
(5)
3

---

Published in Transactions on Machine Learning Research (09/2023)
The Kurtosis regularization term for an M- layered neural network can be expressed
M = (k,,)) w ( )
LKURE
M Σ (((w)) KC,
(6)
Here, μ and a are the mean and standard deviation of tensor w and Kr denotes a scalar value that defines the distribution enforced on the tensors. Similar to Shkolnik et al. (2020), we employ fr = 1.8 to enforce uniform distribution. We provide the pseudo-code for FEDAVG with Kurtosis regularization, FED
Algorithm
Algorithm FEDAVG KURE
1: Require (Wo: No, fla A) do sample a subset S of clients for all i E S in parallel do (broadcast server state to c
for t = 0......,T
5: W --
6: for k = O....,K
do
7:
Likk Vfi(W
8:
J W ki Si = f.(W; St.k + LKUUR( w.
fi.k
V (w
10:
W 1 w Te Di.k {client update}
11:
end for
12: end for
13:
A |S| E,es(wi,r wh.o)
14: Wi+1 t Wr + no D. (server update)
15: end for
Additive Reeudo-quantzztionn Noise (APQN). The quantization robustness problem has similaritics 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 quantization robustness approach that involves adding random pseudo-quantization noise during the training procedure. This motivated by the recent success of Randomized Smoothing (Cohen et al. (2019)) 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 introducod Differentiable Quantizer (Défossez al. (2021) a replacement of the commonly used Straight Through Estimator (STE) based quantizer. Their proposed quantizer has only been used to achieve quantized models for single target bit-width or mixed-precisson 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 learn models that are robust to varying levels of quantization noise and thus can be quantized different bit-widths. The local loss function this case can be reformulated
(w, D.) EE (Q(. b).
(7)
Here, Q(-) is greudo-quantize with bit-width that adds noise sampled from the uniform distribution ul-A,/2, A,, and can be defined as
A A
Q(w, - 4
(8)
2
Provided that A sufficiently large, the sampled pscudo-quantization noise can have support for vari- ous target bit-widths, thus encouraging quantization robustness. The pseudo-code for FEDAVG-APQN is provided in Algorithm 2

---

Published in Transactions on Machine Learning Research (09/2023)
QAT and MQAT. A standard way of learning network with low bit-widths is to constrain the parame- ters /nr/or activations of the model to fixed quantization levels. This procedure of training neural network with standard Projected Gradient Descent (PGD) algorithm with quantization function as projection, is of- ten termed Quantization-Aware Training (QAT). To perform "quantization-awan"
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 bithwi In this "quantization-aware FL, the client-level loss function can be reformulated as:
17 (w, D.) EE-D, (QW,, b).
The quantization step-size As can be either learnt as parameter (Esser et al.) (2020) or be estimated before the start of training and kept fixed thereafter. Note that the backward pass of the network involves gradient estimate through the non-differentiable rounding operation of Q(-). To this end, similar to prior
QAT literature we use the standard STE approximation (Bengio et al. (2013) ) which approximates the gradient of the rounding operator as 1:
i - Ōx
(10)
By using this STE approximation for QAT, we 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:
2
A
2a 1 As.
(11)
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 perfor- mance degradation when quantizing to other bit-widths (Kim et al. (2020)). For this reason, QAT alone is not suitable for tackling the heterogeneous hardware requirements that one can encounter in 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 bitwidth b E B either randomly during local training by fixing it based on each client's hardware requirements. The sampled bit-width is then kept same for all the layers.
Similar to QAT, the quantization step-size As for different bit-widths can be either learnt or estimated before the start training and kept fixed thereafter. The quantization step-size for different bit-widths is then shared along the model parameters with all clients. We provide the pseudo-code for FEDAVG-MQAT in Algorithm 2
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. We do not discuss the convergence bohaviour of KURE, as KURE uses standard unbiased gradients of regularized objective and therefore the standard FEDAVG guarantees apply for this objective. What follows now convergence analysis in the non-convex setting with the help of the following assumptions:
Assumption (Lipschitz Gradient). Each local loss function F, is L-smooth Vs E S, i.e. [VV (F.())) < L||= vll, Vx, y E IRD.
Assumption 2 (Bounded variance). Each F, has bounded local variance, i.e., V -F of, where fa is stochastic estimate of the local loss based on w IRD and is random mini-batch Furthermore, the global variance is also bounded,, i.e., 3 E. VV,
FI < oz. Vw E IR
Assumption 3 (Bounded quantization noisc). Let w IR,, j be any of its dimensions and Q(-) the quantization operation with a step size A E IR. The quantization noise ry E IR added to wy, i.c,, Tj = (ww,)) Wi; is bounded by half the step size of Q(-), ie., rj $ 4.
5

---

Published in Transactions on Machine Learning Research (09/2023)
Algorithm FEDAVG-APON, FEDAVG-QAT, and FEDAVG-MQAT
Require (Wo. Ne. no. 6, b) for 1 0,..., T-....T-11 do sample a subset 8 of clients
4: for all i E S in parallel do
5: wio <- W, (broadcast server state to c
6:
+ [
7:
for = O,....,
do Li t Vw.f(QQ(V W/, b): (7
P..k VwJi(Q/ W (.kt b); St.k
10:
G.k VwJe(Qi W i,k: Si,te
11:
w Wi--1 F wit Pi,k {client update]
12:
end for
13:
end for
14: A. [S| Eies(wirr - wi.))
15:
W++1 <- + Tls D. {server update)
16:
end for
The first two assumptions are common in the non-convex optimization literature for FL (Reddi al. (2020)) and the third is automatically satisfied in quantization schemes, provided the ranges are set up appropriately. Based on these assumptions, we can then make the following statement.
Theorem 1. Let K be the local iterations of each client and w E TRD be the global model parameter vector. Under assumptions
and 3 if the client (Nc) and server (1s) learning rates are chosen such that
?e 10LK
Ne 8LK
(12)
we have that the FEDAVO-/APQN.QATT.OOO
server updates satisfy
F F(w1) - F(w")
min S8ST
/VF(w))]
+ Te n.A
(Bof IKo HI.DRR)
(13)
where we define
A K 2LmmmoK3
(14)
B 41 K2 1.2 + Lnia 2K² + K ))
(15)
r 24 K2 221 2 Lni,
(16)
4m K K²
H
(17)
3n
and R An for APQN, R= A for QAT and R= EB A2 for MQAT.
2
The proof is delegated to appendix due to space constraints. It follows Reddi et al. (0200) while handling quantization noise, such Li et al. (2017). We server side weights is similar to traditional FL (Reddi et al. (2020)), albcit with an additional error floor due to the quanlization noise. Some of this quantization error can be reduced by decreasing the learning rales, up to an irreducible factor of OLL'DR) that depends on the bit-width. Both the convergence term, í.e., F(w,)- (Fw"), and the the error term scale with the number of local itcrations K, with the first decaying with but for the second some of the error terms increase with K, due each step contributing additional quantization noise and drift between the local and the server weights.
that the convergence rate for the (non-quantized)

---

Published in Transactions on Machine Learning Research (09/2023)
4 Related Work
Post-training quantization. PTQ is fast and efficient way of achieving neural network quantization by using little or no data at inference time. Recent literature in PTQ focused on post-training quantization of LLMs and transformers. Frantar & Alistarh (2022) extend the Optimal Brain Surgeon (OBS) frame- work to efficiently quantize and prune NNs in a unified setting. Their method is time and space-efficient while achieving high accuracy in vision and language models. Frantar et al. (2023) propose an extension of OBC Frantar & Alistarh (2022) that is optimized for efficient and accurate quantization of generative pretrained models. This one-shot weight quantization method can quantize GPT models with 175 billion parameters in approximately four GPU hours, reducing the bitwidth down to 3 or 4 bits per weight, with negligible accuracy degradation relative to the wncompessse baseline. Liu et al (2023a) address the issue of quantizing heavy-tailed activations in vision transformers. They discover that for a given quantizer adding a fixed uniform noisy bias to the values being quantized can significantly reduce the quantization error. By adding a noisy bias to each layer they are able to actively alter the activations distribution and make it more quantization-frien frienddly. Yao et al. (2022) proposes a method for efficient and accurate PTQ of large- scale transformers. It comprises of a fine-grained hardware-friendfy quantization scheme for both weight and activations employing layer-by-layer knowledge distillation algorithm (LKD) without the access to the original training data. ZeroQuant can achieve 8-bit weight/activate quantization of GPT--3-stle models with minimal accuracy impact.
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. (2020) first introduced the idea of learning the quantization step-size jointly with the weight achieving near peating-point accuracy in NesNets even with 3-bit quantization. Since then, further advances in quantization-awal training (Han et al. 2021
J. Lee 2021 Gong et al. 2019: Bhalgal et al. 2020 Park & Yoo 2020) have pushed the envelop and enabled ultra liw-bit quantization (2-4 bits) for a wide range of networks and tasks. Recently, [agel et al. (2022)
observe that oscillating latent weights can prevent neural networks from converging to optimal solutions during QAT. They propose freezing oscillating weights or oscillation dampening through regularization and thus improve quantized accuracy in efficient ConvNets. Shin et al. (023) propose training with pseudo-noise noise quantization to prevent unstable convergence induced by the NIPQ formulation allows for naturally learning the bitwidth and quantization parameters lcading to more accurale and efficient mixed-preeision quantized NNs. Liu et. al. (2023b) investigate QAT for LLMs and propose data-free distillation method that leverages generations produced by the pre-trained model, which better preserves the original output distribution and allows quantizing any generative model independent of its training data, similar to PTQ methods. They experiment with LLaMA models of sizes 7B, 13B, and 30B, at quantization levels down to 4 bits.
(STE) in QAT. The
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. [lliadeeh] let al. (0200) model quantization noise as an additive perturbation and show that they can improve quan- tization robustness by regularizing the Lı-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 (Shkolnik et al., 2020). Shkolnik et al. (2020) trains the network with Kurtosis Regularization (KURE) the weights to improve robustness and use the LAPQ NNahshan et al. 2020) algorithm to find the optimal quantization parameters post-training It has been shown in Shkol- nik et al. (2020) 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, Délos-
sez et al. (202I) introduce additive uniform noise to simulate quantization during training for the purpose of QAT: this method can be casily extended for the purpose of quantization robustness as we show in this paper. Kim et al. (2020) introduces a training method that leads to more iuaantization-rienddy weights, by scaling the gradient depending on the distance of the weights from the quantization grid.
7

---

Published in Transactions on Machine Learning Research (09/2023)
Quantization in Federated Learning (FL). Quantization in the context of FL has been studied mostly in the context of compressing communication, e.g., (Amiri et al. 2020 Reisizadeh et al. 2020 Triasteyn et al. 2021). Equally important to reducing the communication overhead is the reduction of the computational cost of training and inference on-dovice. Existing work in this direction can roughly be divided into designing
more efficient models through e.g. sparsification (Caldas et al. 2018b) Jiang et al.] 2019 Louizos et al. 2021 and more effective algorithms (Reddi et al. 2020 He et al. 2020] that reduce the number of training
rounds. Structurally sparse models can reduce training and inference costs (Horvath et al. (2021) targets
heterogeneous computational resources) and may additionally reduce communication (Louizos et al. 2021
Diao et al. ) proposes to assign to each client subset of the global model depending on their resources. A weaker client receives only subset of hidden layers and as such, it is suited for training as well as inference across a helerogenuous compute landscape. Such an approach is orthogonal to MQAT, similar lo how sparsification and quantization are being successfully combined in federated training and centralized settings aliko. Quantization, the focus of this work, is orthogonal to these and as of yet understudied in the literature. QuPeL (Ozkara et al. 2021) is very close in scope to our work. A core proposition of their work however is the flexibility of the quantization mechanism to be non-uniform and the set of quantized values that the model weights can occupy to be learnable. Such a method is highly performant in theory but is entirely unsuitable to low-bit accelerators as can be found in today's hardware. Furthermore, QuPeL addresses hardware helerogeneity through personalization ion, meaning that each client needs to have access a local dataset and perform a finetuning operation in order to select the appropriate centroids for their budget. This in contrast to our work where the quantization happens "zero-shot", i.e., the client does not need any data (i.e., it could be a new client in the federation, only interested in inference) and can just quantize the server model to specific, hardware friendly, bit-width. 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.
5 Experiments
We evaluate the quantization robustness of different proposed variants of FEDAVG using standard bench- marks 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 (Kizhovsky et al. 2009), TinyImageNed and FEMNIST (Caldas et al. 2018a) datasets. We split the data into 100 (O) 500 (CIFAR-100), 500 (TinylmageNett) 3500 (FEMNIST) clients in a non-iid way following Hsu et al. (2019), where in each round only 10 clients participate for all datasets except TinylmageNet dataset, where 100 clients participate. We train different models for 5000 (CIFAR-10 using ResNet-20), 2000 (CIFAR-10 using LeNet-5), 10000 (CIFAR-100), 4500 (TinylmageNet), and 6000 (FEMNIST) rounds. We use small local batch sizes for all clients in our experiments for different datasets: 64 (CIFAR- 10), 20 (CIFAR-100, TinyImageNet, FEMNIST). For all our experiments, we use ADAM optimizer for server training phase and SGD optimizer for client training phase. We use single epoch of local client training for each client participating in a round for all our experiments. For data augmentation, we normalize CIFAR- 10 and CIFAR-100 and TinyImageNet to per-channel zero-mean and standard deviation of one. CIFAR-100 further undergoes random cropping to 28 pixels height and width with zero-padding, followed by random horizontal flipping with 50% probability. FEMNIST requires no preprocessing.
In order to simulate a non-i.i.d. data split that we would expect in the federated sccnario, we artificially split CIFAR-10, CIFAR-100 and TinylmageNet by their label. For CIFAR-10 and TinylmageNet, the label proportions on each client are computed by sampling from a Dirichlet distribution with a - 1.0 (Hsu et al. (2019)). For CIFAR-100 we use the coarse labels provided with the dataset and follow Reddi et al. (2020). For our FEMNIST experiments, the federated split is naturally determined by the writer-id for each client.
ltap::///tiny--magnee hww.://ininnntteeouuu..oommm herokuapp.com/
8

---

Published in Transactions on Machine Learning Research (09/2023)
CIFAR-10 ResNet-20
CIFAR-10 LeNet-5
CIFAR-10 CCT
Global GiobalVallisa Validation Accuracy
Validat Accuracy
Glabal Validation Accuracy
60
860
50
50
HcttriFF
40-
30
20
20
32
Quantization BR-Width
Quantization Bit-Witth
Quantization Bit-iitth
(a)
(b)
(c)
CIFAR-100 ResNet-20
TinylmageNet ResNet-18
Global Validation Accuracy
-50
40
40
Bareline
Feokug SURE
Global Validation
Pookug
dagg (9-3)
20
Fediun DAT (W-4)
20
Feskug (W-2)
10-
0-
Quantization Bit-Width
Quantization Bit-Width
(d)
(e)
Figure 1: Global validation accuracy of the proposed FEDAVG variants at different bil-widths of weight quantization for models trained on (a) 100 using ResNet-20, (b) CIFAR-10 using LeNet-5, (c) CIFAR- 100 using ResNet-20, and (d) TimylmageNet using the R8 architecture. Here, an abbreviation of "W" indicates that we performed Brrigtt-qantizaate only, whereas "W-#" refers to quantization at # bits.
We use LeNet-5, ResNet-20 and Compact Convolutional Transformers (CCT) (Hassani et al. 2021 archi- tecture for evaluation on CIFAR-10 and LeNet-5 on FEMNIST. For experiments with CIFAR-100 and Tiny- ImageNet, we 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 (Reddi et al. 2020). We report the accuracy of the quantized server-side model whenever comparing the models at different quanti- zation bit-widths. To compare quantization robustness, we use bit-widths from the set {32, 8, 6, 4, 3,3,2). We provide the details regarding training hyperparameters 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 (Nagel al. (2021)). We found that the "Scaler" without "static Batch Normalization" lead to unstable training in LeNet-5. For HeteroFI. (Diao al. 2021 implemen- tation, we therefore report results without Scaler Batch Normalization. Batch Normalization negates the scaler's effect, and the absence of Batch Normalization causes the instability. 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 experiments are performed using NVIDIA Tesla V100 GPUs and code in Porch...
5.1 Weight Quantization
Firstly, compare all weight-only quantization-robust FEDAVO 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.
CIFAR-10. Fig. la 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 low bit-widths (2 and 3). Both, FEDAVG-KURE and FEDAVG-APQN perform similarly 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 the FEDAVG-QAT
9

---

Published in Transactions on Machine Learning Research (09/2023)
models are quantized to bit-widths other than the target bit-width they have been trained for. This is known issue with QAT trained models in the context of quantization robustness. Our proposed FED MQAT variant directly targets this issue. The AVG-MQAT model outperforms all other variants at different level of quantizations consistently. Furthermore, it improves validation accuracy at full procision as well. Further investigation revealed that for an over-parameterived network such as ResNet-20, 20, the AEDAVG baseline overfits the training set. Our proposed AVG-AAQT 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 relatively small network, no overfitting is observed for the FEDAvG baseline. It is important to note that the LeNet-ō (achieves 48-55 % accuracy at bits) is more robust to weight quantization compared to ResNet-20 (achieves 28-35% accuracy 2 bits) for the baseline FEDAVG trained model on CIFAR-10. Similar previous comparisons, KERE and AEDAVG-APQN achieve marginal gains low bit-widths. The FEDAVG-MQAT produces better validation accuracy (improvement of Rs 12% al ultra low bit-width of 2 bits but with considerable loss (+= 3%) in accuracy at full precision. We believe this is because small network, such 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. Fig. Ic shows performance on CIFAR-10 10 with the transformer architecture namely CCT (Hassani et al.] 2021). Similar to convolutional neural networks, FEDAVG-MQAT consistently outperforms all other other variants at different level of quantizations.
In order to compare our quantization robust mechanism in FL against existing literature, we implemented HeteroFL. (Diao et al. 2021 HeteroFT. propose to achieve efficient inference by pruning subset of hidden layers based on computer budget. To compare HeteroFL to our proposed methods, we compute the bit operations per second (BOPs) (Van Baalen et al. 2020) for the full-precision activation LeNet-5 at different bit-widths for the weight tensors and chose pruning ratios for HeteroFL that result in the same BOPs. Instead of bits, the HeteroFL curve in Fig. 1b thus presents the evaluation accuracy of the pruncd model at different ratios while being close BOPs to the corresponding quantized LeNet-5 model. We notice that pruning as an alternative approach and dataset. Since these approaches are orthogonal, expect future work to explore the combination of quantization and pruning to ensure efficient inference in FL.
quantization generally underperforms quantization for this model
CIFAR-10 ResNet-20
CIFAR-100 ResNet-20
80
Global Validation Accuracy
Global Validation Accuracy
50
60-
40
30
40
Baseline
20
FedAug-KURE
APON (A-4)
20
Fedkug-APOH APON (A-2)
FedAug QAT (A-4)
FedAug QAT (A 2:
Fedtug MQAT A
Quantization Bit-Width
Quantization Bit-Width
(a)
(b)
Figure 2: Global validation accuracy of the proposed FEDAVG variants at different bit-widths of activation quantization for models trained on (a) CIFAR-10 using ResNet-20, and (b) CIFAR-100 using ResNet-20 ar- chitecture. Here, an abbreviation of "A" indicates that we performed activation-quantiaatiio only, whereas " "A-#" " refers to quanlization al # bils.
CIFAR-100. Fig. [d shows the performance our FEDAVG variants with weight-only quantization on CIFAR-100 using the ResNet-20 architecture. Similar our CIFAR-10 sctup, we observe a large drop the baseline model accuracy at low bits. Our FEDAvG-MQAT variant achieves significant gains, especially at low bit-widths; v 35% at 2-bits and 10% at 3-bits, compared to the baseline FEDAVG model. It is
10

---

Published in Transactions on Machine Learning Research (09/2023)
CIFAR-10 10 ResNet-22 20
Global Validation Accuracy
70-
60-
50-
30-
Baseline IDropoutt
20
Fed/vg QAT IWA
Fed/vg OAT (WA
10
FedAug WA
32
Quantization Bit-Width
Figure 3: Global validation accuracy of the proposed FEDAYG variants at different bit-widths of weight and activation quantization for models trained on CIFAR-10 10 using ResNet-20 architecture. An abbreviation of " "WA" " indicates that we performed both weight and activation guantization.
important to note that these improvements on low bit-widths come at the cost of a small loss of accuracy (fs 2%) at full precision comparison the baseline.
TinylmageNet. We also performed an experimental evaluation on more challenging FL setup; clas- sification on TinyImageNet with the ResNet-18 architecture. The results for all of our proposed FEDAVG variants can be seen in Fig. Iel 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 low 2-bits without any loss in full-preeision accuracy. In comparison to the baseline, it improves the 2-bit, 3-bit quan- tization accuracy by significant margin (20- 35%). We would like to point out that despite being trained for different set of bit-width, our AVG- 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 of the standard datasets used to evaluate FL algo- rithuns. For the sake of completeness, 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, relatively easier and more robust to quantization.
5.2 Activation Quantization
To further demonstrate the effectiveness of our FEDAVG variants, analyze the task of quantizing acti- vations CIFAR-10 and CIFAR-100 with the ResNet-20 architecture. For AEURVG-KRRE, we impose the regularization term on the activations and, in a similar manner, for FEDAVC-(APQN, QAT and MQAT), the noise/ / quantizer is on the activations.
Fig. 2a and Fig. 2b show the validation accuracy different bit-widths for CIFAR-10 and CIFAR-100 re- spectively. For both CIFAR-10 and CIFAR-100, we observe considcrable decline in the full precision model accuracy after Kurtosis regularization on activations compared to the baseline. It sould be noted that the original work KURE (Shkolnik et al. (2020)) considered weight-quantizatinn only. FEDAvG-MQAT achieves significant gains at 2-bit (<e 55% on CIFAR-10 and ~ 24% on CIFAR-100) and 3-bit (~ 26% on CIFAR-10 and s 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 datasot, we observe considerable decrease in model accuracy for AVQ- variant higher bit-widths (and full precision).
11

---

Published in Transactions on Machine Learning Research (09/2023)
Federated Averaging (FedAvg)
Bit Config
Baseline
MQAT-W
MQAT--W
W-32
70.16
74.46
70.44
W-8
69.86
74.54
70.88
W-6
70.02
74.60
68.24
W-4
68.44
74.58
70.02
W-3
64.14
75.64
67.58
W-2
28.02
72.58
68.20
Table 1: Global validation accuracy after quantization at various bit-widths for different FEDAVG variants trained on CIFAR- R-10 10 dalasel using ResNet-20 architeclure. Here, indicntes the client-specffee bit-widthh is chosen at the begining of training and then kept fized throughout.
5.3 Weight and Activation Quantization
The combination 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 quantized before each matrix multiplication. We see that FEDAVG-MAAT outperforms the FEDAvG baseline 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.
5.4 Per-client fixed bit-width.
Performing a forward-pass low bit-width during Quantization-Awarc Training (QAT) does have the ad- ditional benefit of reduced computational requirements also during training. A good assumption to make is that each client implements efficient hardware acceleration for specific bit-width that remains constant during its participation in the federated lcarning process. While in the main experimental setting we con- sidered 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 [ ex-ante fixed bit-widths lead to some degradation in performance, albeit at the ionee benefit. spead-up training.
6 Discussion
Real-world deployments of FL, necessarily require catering to heterogencous device landscape. Quantization-robust server models, f.e., models that have robust performance arbitrary target bit-widths, are step towards effectively navigating such landscape. In this work, we introduced several variants of the AvD algorithm that encourage quantization-robustness for the server model. Experimentally, we demonstrated that our FEDAVG variants can achiove good performance on several target bit-widths, with- out significant. accuracy degradation for the full precision model. No method clearly outperforms all settings, although we see MQAT performing well most situations, especially for lower bit-widths. Theo- retically, we showed that quantization-awaa local training on the clients provides a convergence rate that is similar to traditional FL, albeit with 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 considered, such as non-uniform quantization, different quantization stratogies (symmetric vs. asymmetric) and more subtle differences in inference engines across devices. We believe that the scenario where client's 's hardware characteristies considered fixed throughout training opens up the possibility for advanced client
12

---

Published in Transactions on Machine Learning Research (09/2023)
sub-sampling strategies and necessitates 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 's dataset's influence on the learned function will be different.
References
Milad Alizadeh, Arash Behboodi, Mart van Baalen, Christos Louizos, Tijmen Blankevoort, and Max Welling. Gradient 11 regularization for quantization robustness. ICLR, 2020.
Mohammad Mohammadi Amiri, Deniz Gunduz, Sanjeev R. Kulkarni, and H Vincent Poor. Federated learning with quantized global model updates. arXic preprint arX6::2006.10672, 2020.
Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. Xivi preprint 1rXi::1308.343 2013.
Yash Bhalgat, Jinwon Lee, Markus Nagel, Tijmen Blankevoort, and Nojun Kwak. Lsq+: Improving low- bit quantization through learnable offsets and better initialization. In Proceedings of the IEEE/CVP Conference on Computer Vision and Pattern Recognition (CVPR) Workshops, June 2020.
Sebastian Caldas, Sai Meher Karthik Duddu, Peter Wu, Tian Li, Jakub Konečný, Brendan McMa- han, Virginia Smith, and Amect Talwalkar. Leaf: A benchmark for federated scttings. arXiv preprint arXiv:1812.01099 2018a.
Sebastian Caldas, Jakub Konečny, H Brendan McMahan, and Ameet Talwalkar. Expanding the reach of federated learning by reducing client resource requirements. arXiv preprint arXio:1812.072: 012100, 2018b.
Jeremy Cohen, Elan Rosenfeld, and Zico Kolter. Certified adversarial robustness via randomized smoothing. International Conference on Machine Learning, pp. 1310-1320. PMLR, 2019.
Matthieu Courbariaux, Yoshua Bengio, and Jean-Pierre David. Binaryconnedt: Training deep neural net- works with binary weights during propagations. Advances neural information processing systems, 28, 2015.
Alexandre Défossez, Yossi Adi, and Gabriel Synnaeve. Differentiable model compression via pscudo quanti- zation noise. orXiv preprint arXiv:2104. 099887, 2021.
Eamao Diao, Jie Ding, and Vahid Tarokh. (f]e Computation and communication efficient federated learning for heterogeneous clients. In International Conference on Learning Representations, 2021. URL
https://openrevieet rer.
Steven K. Esser, Jeffrey L. McKinstry, Deepika Bablani, Rathinakumar Appuswamy, and Dharmendra Modha. Learned step size quantization. International Conference on Learning Representations (TCLR), 2020.
Eliss Frantar and Dan Alistarh. Optimal brain compression: A framework for accurate post-training quan- tization and pruming. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.). Advances in Neural Information Processing Systems, 2022. URL httpp:://opaarrriee.eee//// net.oorundddoo ksVGCD10Ebal
Elias Frantar, Saleh Ashkboos, Torsten Hoefler, and Dan Alistarh. OPTQ: Accurate quantization for gen- erative pre-trained transformers. In The Eleventh International Conference on Learning Representations, 2023. URL http:://paaareeeeeeeeeeeeoraaaaa
Ruihao Gong, Xianglong Liu, Shenghu Jiang, Tianxiang Li, Peng Hu, Jiazhen Lin, Fengwei Yu, and Junjic Yan. Differentiable soft quantization: Bridging frel-precision and low-bit neural networks. 2019 IEEE/CVE International Conference on Computer Vision (ICCV), pp. 4851 1851-4860, 2019.
13

---

Published in Transactions on Machine Learning Research (09/2023)
Tiantian Han, Dong Li, Ji Liu, Lu Tian, and Yi Shan. Improving low-precision network quantization via bin regularization. In Proceedings of the IEEE/CVF International Conference on Compuler Vision (ICCV), pp. 1261-52700 October 2021.
Ali Hassani, Steven Walton, Nikhil Shah, Abulikemu Abuduweili, Jiachen Li, and Humphrey Shi. Escaping the big data paradigm with compact transformers. arXiv preprint crXiv:2104.055 05704,, 2021.
Chaoyang He, Murali Annavaram, and Salman Avestimehr. Group knowledge transfer: Federated learning of large cans at the edge. arXiv preprinl arXio:2007. 14513, 2020.
Samuel Horvath, Stefanos Laskaridis, Mario Almeida, Ilias Leontiadis, Stylianos Venieris, and Nicholas D Lane. Fjord: Fair and accurate federated learning under heterogeneous targets with ordered dropout.
arXiv preprint ar Xic: 13451, 2021.
Tzu-Ming Harry Hsu, Hang Qi, and Matthew Brown. Mcasuring the effects of non-identical data distribution for federated visual classification. arXiv preprint arxie: 1909::999..000 1999006335 2019.
B. Ham J. Lee, D. Kim. Network quantization with element-wise gradient scaling. In Proceedings of the IEEE/CVF Conference on Compuler Vision and Pallern Recognition, 2021.
Yuang Jiang, Shiqiang Wang, Victor Valls, Bong Jun Ko, Wei-Han Lee, Kin K Leung, and Leandros Tassiulas. Model pruning enables efficient federated learning on edge devices. ar preprint Arx:v1190.. 12326, 2019.
Jangho Kim, KiYoon Yoo, and Nojun Kwak. Position- based scaled gradient for model quantization and pruning. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin (eds.). Advances in Newral Curran Associates, Inc., 2020. URL https: pdf
Information Processing Systems, volume 33, pp.
Raghuraman Krishnamoorthi.. Quantizing deep convolutional networks for efficient inference: A whitepaper. arXiv preprint arXix:1866.00 08342, 2018.
Alex Krizhevsky et al. Learning multiple layers of features from tiny images. 2009.
Hao Li, Soham De, Zheng Xu, Christoph Studer, Hanan Samet, and Tom Goldstein. Training quantized nets: A deeper understanding. In Proceedings of the 31st International Conference on Neural Information Processing Systems, pp. 5813-5823, 2017.
Yijiang Liu, Huanrui Yang, Zhen Dong, Kurt Keutzer, Li Du, and Shanghang Zhang. Noisyquant: Noisy bias-enhanced post-training activation quantization for vision transformers, 2023a.
Zechun Liu, Barlas Oguz, Changsheng Zhao, Ernie Chang, Pierre Stock, Yashar Mehdad, Yangyang Shi, Raghuraman Krishnamoorthi, and Vikas Chandra. Llm-qat: Data-free quantization aware training for large language models, 2023b.
Christos Louizos, Matthias Reisser, Joseph Soriaga, and Max Welling. An expectation-mmxiiizationn per- spective on federated learning. arXie preprint orXiv:2111.10192 2021.
Brendan McMahan, Fider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Commumication-efiiiiee learning of deep networks from decentralized data. In Artificial intelligence and statistics, pp. 1273-1282. PMLR, 2017.
Markus Nagel, Marios Fournarakis, Rana Ali Amjad, Yelysei Bondarenko, Mart van Baalen, and Tijmen Blankevoort. A white paper neural network quantization. arXiv preprint ar arXiv:2106.088 08295 2021.
Markus Nagel, Marios Fournarakis, Yelysei Bondarenko, and Tijmen Blankevoort. Overcoming oscillations in quantization-awan training. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 11318.16330. PMLR, 17-23 Jul 2022. URL https://proceedings. mlr. wwtppp/eeeeepeaeeeeeeee www.e///agel/// 122a. html
14

---

Published in Transactions on Machine Learning Research (09/2023)
Yury Nahshan, Brian Chmiel, Chaim Baskin, Evgenii Zheltonozhskii, Ron Banner, Alex M. Bronstein, and Avi Mendelson. Loss aware post-training quanlization, 2020.
Kaan Ozkara, Navjot Singh, Deepesh Data, and Suhas Diggavi. Qupel: Quantized personalization with applications to federated learning. orXiv preprinl Xrxiv:2102 0.te2102.117000 11786, 2021.
Eunhycok Park and Sungjoo Yoo. PROFIT: A novel training method for sub-4-bit mobilenct models. In Andrea Vedaldi, Horst Bischof, Thomas Brox, and Jan-Michael Frahm (eds.), Compuler Vision - ECCV 2020 16th European Conference, Glasgow, UK, August 23-28, 2020, Proceedings, Part VI, volume 12351 of Lecture Notes in Computer Science, pp. 430-446. Springer, 2020. doi: 10.1007 00100779978-33000--55553 URL https://dol. http:://ddoo.o/11../1 wrt.:/..11.0//711--- 3- -803958539-- 26
26.
Sashank Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Konečný, Sanjiv Kumar, and H Brendan McMahan. Adaptive federated optimization. arXiv preprint arXiv:2003.00295
2020.
Amirhosscin Reisizadch, Aryan Mokhtari, Hamed Hassani, Ali Jadbabaic, and Ramtin Pedarsani. Fed- paq: A communication-eeliceeut federated learning method with periodic averaging and quantization. International Conference on Artificial Intelligence and Statistics, pp. 2021-2031 PMLR, 2020.
Juncheol Shin, Junhyuk So, Sein Park, Seungyeop Kang, Sungjoo Yoo, and Eunhyeok Park. Nipq: Noise
proxy-based integrated psenuo-quaanizzttonn In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognation (CVPR), pp.
June 2023.
Moran Shkolnik, Brian Chmiel, Ron Banner, Gil Shomron, Yury Nahshan, Alex Bronstein, and Uri Weiser. Robust quantization: One model to rule them all. NewrIPS, 2020.
Alcksci Triastcyn, Matthias Reisser, and Christos Louizos. Dp-rec: Private & commmicationn-fffcientt fed- erated learning. arXiv preprint arXiv:2111.05454 2021.
Mart Van Baalen, Christos Louizos, Markus Nagel, Rana Ali Amjad, Ying Wang, Tijmen Blankevoort, and Max Welling, Bayesian bits: Unifying quantization and pruning. Advances in neural information processing systems, 33:5741-577 5752, 2020.
Zhewei Yao, Reza Yazdani Aminabadi, Minjia Zhang, Xiaoxia Wu, Conglong Li, and Yuxiong Hc. Zero- quant: Efficient and alfordable post-training quantization for large-scale transformers. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. URL https:://ppenewee
net/forum/IG==ffvcEII
Shuchang Zhou, Zekun Ni, Xinyu Zhou, He Wen, Yuxin Wu, and Yuheng Zou. Dorefa-net: Training low bitwidth convolutional neural networks with low bitwidth gradients. arXiv preprint arXiv:1606. 06166, 2016.
A Appendix
Appendix
B 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 Reddi et al. (2020), while handling the quantization noise, such Li et al. 2017). This appendix tries to be verbose 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. in the paper), V F(w.) for different perspectives on L. 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 incqualitics 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.
to be easy to follow along, Theorem
15

---

Published in Transactions on Machine Learning Research (09/2023)
B.1 Notation and definitions
E IRD
...is the flattened vector of model paramcters.
w IR model parameters the k'th (out of K) iteration of local updates
client in communication round
S is the set of all clients ("shards'), of size S = |S|
Σ Fa(w) as shorthand. F(w)= Σ Fo(w), We use 5
(w( - De
D. S. (w,, where D, = [Da] the size of the local dataset D, at client s.
fs( ... is a loss function evaluated on a random mini-batch 6, on client with model parameters w.
[.(F(W)+1) the loss evaluated on parameters al round ++ 1, in expectation over the randomness at round t that influences the transition of parameters from round to round t ++1. 1.
B[F(w) is the loss evaluated on parameters w, averaged over all possible sources of randomness. Precise definition given in context below.
F(w) is the gradient of F with respect w evaluated w. We the shorthand FF(w), similarly for (F,(w) and (.,(w).
A
is the step size of the quantizer / magnitude of the quantization noise for given bit-width b.
B.2 Assumptions
Assumption (Lipschitz Gradient). Each local loss function is L-smooth Vs 8, i.e, [VFF(x)) FF,yy)l S Ll|a 11l=-yl), yll, Vx, y IRD.
Assumption 2 (Bounded variance). Each F, has bounded local variance, i.e., - < of, where f. is stochastic estimate of the loss based on Furthermore, the global variance is also bounded, i.c., $/ E, (F)(w) FF(w)] I oz Vw E IRD
IDD and is random mini-batch
Assumption (Bounded quantization noisc). Let quantization operation with a step size A E IR. The quantization noise rj IR added to Wj, i.c, Tj Q(w,)) Wi; is bounded by half the step size of Q(),, Le., rj S
IR,, j be any of its dimensions and Q(-) the
B.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 2. For any y have that +2a8 S ya² + 3
Proof.
1 3)2 1 +0 502 + B2 + 2a3 +2a3 s 7a² + 1
< (V7a
7
Corollary 3. For any 7 0 we have that (a + 8)2 (1 77 + + 3
16

---

Published in Transactions on Machine Learning Research (09/2023)
Proof.
(a a² + 32 2a,
(18)
From Lemma 2
1 B2 = = (1+ (1 - 1,22
a² + 8 7a²
(19)
7
Lemma 4. For random variables 2p...., 2r we have that
E||| 2 z. 12] S rBll| 11? + ||
(20)
Proof. The proof follows from expanding the square and applying Lemma to each 2z,z term with = 1.
Es [[[]]]]]]]
Σ B.
EB
A
DI, R.
[
(21)
Lemma 5. Let r the quantization noise added w satisfyinn assumption Eq. 3
When performing
QAT, MQAT, APQN we have that
E |||| & DR
(22)
where R = As for QAT, R = Es
for MQAT and R = An for APQN.
v12
Proof. For QAT we have that
D
A 2
E
|||||||]
E
E
(23)
d=1
d=1
due to the bounded quantization noise assumption. For MQAT where consider random bitwidths, we have that
A
A
E |,[E||||||]]
Σ
E, Er &
Eb
DE,
(24)
4
For APQN we have that
D E
A2
2
B
|||r||
D
(25)
12
12
d=1
d=1
d=1
due Fa 4 2 Ao and Er
Var[ra] + EErra]
A
12
Lemma 6. For any local learning rate ne < 1
IoLK: we can bound the difference between the local shadow
weights and the server weights for any lo E (O...., K l} and K 2 at given federated training iteration via
Ellwes -wll 5 4K +6Kog) 322 24K ((
(26)
5
17

---

Published in Transactions on Machine Learning Research (09/2023)
Proof. We begin by noting that
E Fl|w'R
w = E|| w s, We 7/c V/s(W 9, + rxx VF, s + r k-
k + t V (w:))) + V...w.) - VF(w,) + VW( DI
(27)
where we introduced several shorthand notations for easier manipulation of the terms. By expanding the norm using the multinomial theorem have that
+211 +2 + r
2E[a" (ne +++ c)) + 2n (c+ d e)]
B.e c)] Bdd"e.
(28)
We can now use the fact that the expectation of b is zero, since fe s + is unbiased estimate of FF,W - + + In this way, we can simplify the above to
|||||| 3| + / B 6| r n + +f _
21 (ne(e + e)) Bc (d e)] + Edd e].
(29)
We can now use Lemma with = 2K 1 in order to "split" the 2E (nele + d e))l term
< [B[[||]
|||b| 13 [|]]|||]]]]]0]]]
2K |||||||1 + (2K 1) > d +
E (d e)] Bdd e).
(30)
Following that, we can see that several terms cancel, due to + 2B[c (d e)] 2
d + el|?] [ + B[[[a[a] + [[[]e[[]]
+
1+ 2K Bi||a|11 + 11 2K 2 d e|l
(31)
Finally, we will apply Lemma
4 in order to split Eiilc + + e| and thus end up with
+ 2K
[[[aa| P] n. [|| 6K
+ 6K 6Kn.BBll [2] + 6Kn) B(1
(32)
1
2K - [I|| n -
ABll f. (w". .- + -T - (F (w + -.r.,-11)1 1)
6Kn VF (w-. w. 11 V (W.(W))113)
+
- VF(w:)I1 + &
(33)
where we replaced the shorthand notations with their original terms. To proceed, we will make use of assumptions and to arrive at
2K 1 - + no
+ 6Kn, 1.2 w wi -...---WAl w
6K 6K VF.
VF(w:)l11 + 6K //W
(34)
18

---

Published in Transactions on Machine Learning Research (09/2023)
We can now make use of corollary g with y in order to separate the quantization error Fa,k-1 from the difference between the local (shadow) weight and the server weight
E Blll s.-- well o
2K
w/ S w. 8k 24K 1/c T.2
+
BK - VVP(w:)11 + 6K /
(35)
and then use Lemma 5
order to bound the squared norm rs,k-1
2K Blll - wo well
24 1/ L2 w 1a w. + DR
+ 6K Fl - VF((.))l + 6k GKn V/(:))l)
(36)
To proceed and make further use of our assumptions, we will average the aformentioned incquality over the clients and thus have that
Σ Blll B(llw's w.ll
24
1
Σ E [[[[ s.k- - + no.0
5
2K
S
1. + 6K
(37)
s
+ 6K ((w))))))))
(38)
and then make use of assumption 2 in order to bound the "global" variance
1
1 Σ
1
2K 24
Blll 8 woll
- ooo?
$
8Kg DQ
6Kniogg
6K E
(39)
Finally, given our assumption that 7e S TOLK we have that (1 + 2 + 24KTEL2) S (1+ 1 we can simplify the upper bound even further
K-1 and thus
Σ Ell Blllsss w - w. 1 (1
Σ E [lll sk-1 - well +
$
K-1 S
+ + DDQ + 6K
(40)
We have now arrived at point where the average difference between the local shadow weight at ileration l 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. 40 on each weight difference, up until we end up at the server weight wt (since local optimization started from that point) where the difference is zero. Notice that each application of this bound "adds" additional non-ncgative terms, so we have that the upper bound of K iterations would upper bound the bound K K-1 1 iterations. Therefore, the "worst-case" upper bound is the one where k K. In this case, we unroll the recursion and have that
K-1
(+*-T)+
; (T2 + + DD
Σ Blllw's w.
1
6K 6K
(41)
To simplify even further, we can use the fact that (1 + Rー ) is monotonic in j and that there are terms in the sum, thus get
K 1+ K - (v2( + 6Koz) + IZ + 6K o
(42)
19

---

Published in Transactions on Machine Learning Research (09/2023)
and since (1+ n- j S 4 for any K > 1
S (of + 6Koz) + 323 32K22 12 DQ² + 24K [[]
(43)
which proves our claim.
B.4 Proof of Theorem 3.1
We begin by noting that the server-side update rule of the model the case when the clients perform Quan- tization Aware (QA) SGD with a learning rate of Tle and the server does SGD with a learning rate 1/s- The extension of the proof to more involved server-side update rules in Reddi al. (2020) is straightforward.
Wo+1 Wr = GT.
Te EZ G= = - 8
V.s + trr))
(44)
where 8 indexes the clients and S is the total number of clients. indexes the local client iteration number, and we assume that there are K local iterations in total. w'' E IR corresponds shadow weight at iteration and r's E IR corresponds to the quantization noise that is added to it in cach iteration of the local optimization (as the weights are rounded noised before the forward pass).
real valued local
Using the L-smoothness of the global loss function,
L In, Gil
F(w++1) S F(w.) + n, VF(w)) G.
-
(45)
2
= = + VF (
|G.
(46)
We then take an expectation over all randomness at time step
Lafa
[F((ee)))) & F(w.) + )
+
E.
(47)
2
T1s
T21
and work towards upper bounding the Tıe T2, terms separately.
Bounding Tıt
VF(W,)" G. = VFW,, +,
(48)
-nckk VF. VF(W:))" + no K [VW::)
(49)
We will now work towards upper bounding Ts:
Tse = VF(W.)
Toe 5
Σ Σ
(w: + + + KVF(((
(50)
= VF((:)) Σ
17c 5
VF + +.)++n-VF(m))
(51)
8
k
VV.F() Eel VVvr(ws +w/) 1 ΣΣ
VF, (w.)).
(52)
Now by using Lemma. 2 with
K we have that
Tyx
K
VF((w))
No Edll
1
ΣΣ V 1 SE
VF, (w:)l1)
(53)
2
2K S
S
K.K
VF
FF(()
Ex[1 Σ
(VF.(W +r(a))))))-)) 130-VA VF((w:))
(54)
2
2K
k
20

---

Published in Transactions on Machine Learning Research (09/2023)
By then using Lemma 4 to push the squared norm inside the sum
% K
m1Sk 2KS
VF(()))l +
E1
FF W sk + rdke) - VF&(w:))1
(55)
2
k
and from the Lipschitz gradient assumption
$
%
VF(()))l
+
Te 25
E.
|[(WWAk + + -well
(56)
2
k
noL2
(VF((:))l"
B Bell + r -wall]]
(57)
2
2
S
We will now once more use Lemma
to separate the shadow weight difference from the quantization noise
K
/VF(w))]
The
Σ
Esllw's w. ne L
E, x
(58)
2
S
S
and from Lemma 5
< 2
K
VF n
L2
2
1
E Blww -
(59)
S
We see that in order to proceed, we need to upper bound the difference between the local shadow weight at any iteration & and the server weight. This is were we will use our Lemma in order to proceed
$ MaK [FF(w)) Kn,
+ 6k (32
+ L2 ne [E(
(60)
= 2
+ 24K n1
L2 (oî (K
4K +
(Ko + (ncK +
(61)
Finally, we make use of our assumption Te ToLK which leads to 24 & 1/ K and 32 122 K Je.. In this way, can arrive at our final bound for T3t
3rck [VF(W)] + 2 +KKoz)
4
To:
(62)
Now we can apply this bound to Tı, in order to get
Tit S ne - )2 + 3ncK
|VF((,))) + KI. "
+4K 12 + +Koz))
(63)
4
ne K
[F():)) + 3 DR' 12 6Ko
(64)
Which is the final upper bound on Tıc
Bounding Tx
We begin by noting that
L
[1 = Lof R. [| no WcKVF(W k2
(65)
2
and then we can split the squared norm via Corollary with 7 ==
& (E, + 7cKFF( +n (
(66)
T4
To continuc, we will move towards upper bounding Tat by expanding the terms inside the squared norm
T4p = Eg[l
T7c
Σ
( + r +
ne S
Σ Σ (F,,w,
(67)
S
k
222
E|ll
Σ
(W +r(e) - VFs (w + r + (.. + +(L)
V.P(ww)) 17..
(68)
5
k
21

---

Published in Transactions on Machine Learning Research (09/2023)
We will then apply Lemma 4 in order to move the squared norm inside the sums
SK
E. f. (w.k + - VF,(W's + +(k + VF. + ras) (F.(w)))
(69)
8
and will apply Corollary 3
with 7 in order to split the norm
2n K
ΣΣ (f + r9s) VFs(W's + E.IVFF +r'ak)
-VF((
(70)
To proceed, we will make use of our assumptions order to get
S 2n +
2n K Σ
Σ
E Lell + rse - -w))l
(71)
and we will apply Corollary and the quantization error
3 with 7 one more time in order to split the norm of the weight difference
K of KL
Σ
E.llw's well 4n Σ
Σ
Eelll
(72)
that we can apply Lemma in order to bound the latter
1 Σ E,ll w.ll + 4 L 2 on KL Σs
(73)
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 in order to proceed further
s Kf + LZ
+ (AKK + - (
(74)
=2 + T2 16K (o 1 +KK 128 T, K DR
+ 96n 1 F(:))
(75)
In order simplify the aforementioned inequality we will make a use of our assumption on Tey namely that Te ToLK In this way, we will have that 16K o L & Knz along with 128 3 K2 these into account, we have that
Taking
3 s (2 Kn + (4ng + "
(76)
and due 4+ <6
< K2 Kn K3 K + =
(77)
which constitutes our final upper bound on Ta With this bound at hand, we can move back to bounding T2e and thus get
Te: < k7 + Kn + K2 LR + [
+
(78)
= -n K1 = + + 6
+ Ln 1 + _ K /F(w)))1
(79)
22

---

Published in Transactions on Machine Learning Research (09/2023)
Having bounded Tú and T2c. can now apply these bounds to the inequality at Eq. 47 and thus get
LT
(((F(we))) F(w)) + n, (F((w))'
G
(80)
2
Tit
The
no
4 3 + 4n. Ln {
S F(w))
|F((W]11
6Koz)
4
+ ((2 K2 + Kf +
o2 + 6n K2 L2DR
+ (96 K2 + ne K2 |F(w))]1.
(81)
To simplify the aforementioned bound, we can once again make use of our condition S 10LK which leads to 96 K* I k² In this way, we get that
F. & F(w.) - nal McK
4 2L.
+ (4 22 + L2
KnL)lo
4
+ (2An, 2 L
(
+ 6Lm ) DR DR?
(82)
K 2n VF(w:)]1
F(w:) nale
K + ne (4m,K² LL + L2K2 +
5 )f
Anls
+ n (24 K2 I. ne + Lnn K) Kop + noa
K
(83)
3πc
where introduced several shorthand notations for casier manipulation of the incqualities. We can thus re-arrange the inequality to
[.(F(w-))]] - F(w.) - ( + HL
(84)
In order to consider the entire training trajectory, will use telescoping sum, iee, we will sum this inequality over all rounds and take the expectation at each time-step
Σ (E,[F((++1] (F < A [V + l'oo))
HT,
(85)
In doing that, most of the terms on the left-hand-side will cancel across subsequent time-steps and thus will end up with
T
[r([(WT+))] F(w:) A.
(B +
(86)
We can now re-arrange the terms to get
A
FF(w))l F(w1) - Ey (F(((())))] (Bo [Koo) + H
(87)
1=1
23

---

Published in Transactions on Machine Learning Research (09/2023)
Dataset
Network
f/c
Rounds (T)
Batch Size
CIFAR-10
LeNet-5
le 3
5e 2
le 8
2000
64
CIFAR-10
ResNet-20
le 3
5e - 2
le 7
5000
64
C1FAR-100
ResNet-20
le -
5e - 2
le 7
10000
20
FEMNIST
LeNet-5
1e -3 3
5e -2 2
le 8
6000
20
TinylmageNet
ResNet-18
1e 1e-2
le - 2
1e 3
4500
20
Table 2: Hyperpprrmeteers used for the experimental evaluations in the paper. Here 1s Te denote the server and client learning rate and Es refers to the correction term in value in ADAM optimizer of server.
and by considering w to be the parameters of lowest loss, BT([[(T+1)] 2 F(w"), we have that
T
MeN Σ > - F(w") + The + +KKog) + HLL
(88)
order to proceed, we have to impose a condition on A, namely that it has to positive (otherwise, dividing by A reverses the inequality). For this to happen we need that
K 2L - 2Lm + no 8LK1s
1
1
(89)
4
Assuming that this condition is satisfied, we have that
VF((:)) F(w1) F(w") TnoNeA
+ Be (Boî TKo HI.DRR)
(90)
n,A
(91)
Finally, in order to obtain the result of Theorem 3.1, we make use of the fact that min [VF()))]
T Em [VF(w))l and thus arrive at
1
min [VF(w)] F(W1) F(w")
AAA (Bo] + +IKo2 DRL
Te
(92)
1StT
TooliAA
which completes the proof.
C Final Hyperparameters
We provide the final hyperparameters used for all the experiments in Table Since the grid-search for FL is expensive, we tune the hyperparameters such as client learning rate (n.), server learning rate (n,) and epsilon term (e,) used in ADAM optimizer only for the baselines and then use the same set of hyperparameters for all our proposed variants. For AVE-VR-K found À Te 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 3
tuned À from the grid [1e+1, 1e+0, 1e-1, 1c-2, 1e-3, 3, 1e 1e-4].
D Additional Results
Implicit Regularization of MQAT. As discussed in the main paper, ResNet-20 CIFAR-10 exhibits overfitting, which we further investigate here. It was observed in the main text that AVT 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 AwD) and drop-out before the last full-connected layer of the network in
24

---

Published in Transactions on Machine Learning Research (09/2023)
Federated Averaging (FedAvg)
Dataset
Network
MQAT-W
MQAT-A
MQAT-WA
CIFAR-10
LeNet-5
[4,6,8]
CIFAR-10
ResNet-20
[2,,,46,8,322 [2,3,4,6,8,22]
2,3,4,6,8,32 2,3,4,6,8,22)
[2,3,46688,22]
CIFAR-100
ResNet-20
TinyImagNNet
ResNet-18
[2,3,4,6,8]
Table 3: Quantization Bit-Set used to perform FEDAVG-MQAT on different experimental setup in the paper.
Table 4 Despite not being its primary objective, our FEDAVG-MQAT variant achieves considerable gains in full- pileisiciiionn 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 the right. For MQAT, we observe no overfitting according the validation loss.
CIFAR-10 ResNet-20
CIFAR-10 ResNet-20
2.50
Baseline
Global Validation Accuracy
70
Baseline
FedAvg-MOAT-W
Global GlobalValidaa Validation Loss
2.25
FedAvg-MQATH
FedAvg-MOAT-WA
2.00
1.75
50
1.50
40
1.25
1.00
30
0.75
1000
2000
3000
4000
5000
0
1000
2000
3000
4000
5000
Rounds
Rounds
(a)
(b)
Figure 4: Global validation (a) loss and (b) accuracy curves for Baseline and different FEDAVG-MQAT variants trained CIFAR-100 using ResNet-20 architecture. For FEDAVG-MQAT voriants the volidation loss refers to the loss after quantization at lowest bit-width in the bit-set B. While Baseline model clearly suffers from the overfitting issue, our AV--AQQQ variants clearly manage to avoid it. Further, the validation accuracy curces on the baselines reveals that overfilling cannot be prevented even by early stopping.
Detailed results from the main text In this section, provide the exact values used for plotting various figures in the main text in Table 5-11.
25

---

Published in Transactions on Machine Learning Research (09/2023)
FedAvg
Accuracy
Baseline
70.16
Baseline (Dropout)
72.02
AWD = -le- 2
44.94
AWD -ee
71.62
AWD
70.76
AWD
70.06
AWD le
70.10
MQAT-W
74.46
MQAT-A
76.28
MQAT-WA
74.90
Table 4: Global validation accuracy for full-precision models learnt in federation using FEDAVG variants trained CIFAR-10 dataset with ResNet-20 architecture. Here W indicales the client-specific bit-width is chosen at the begining of training and then kept fixed throughout. Here, on abbreviation of "W", "A and "WA" " indicate weight quantization, activation quantization and both weight-acttvation.
Federated Averaging (FedAvg)
Bit Config
Baseline
KURE APQN (W-4) APQN (W-2) QAT (W-4)
QAT (W-2)
MQAT-W
W-32
70.16
69.38
70.5
70.04
72.38
29.06
74.46
W-8
69.86
69.38
70.72
70.19
71.08
11.68
74.54
W-6
70.02
69.08
70.62
70.16
71.6
12.68
74.60
W-4
68.44
68.12
69.44
69.36
72.58
18.5
74.58
W-3
64.14
64.28
63.46
65.66
71.22
39.02
74.64
W-2
28.02
31.28
29.64
33.36
52.32
72.64
72.58
Table 5: Global validation accuracy after weight quantization of various quantization robustness variants for Federated Averoging (FEDAvG) on federated version of CIFAR-10 dataset using ResNet-20 architecture.
Federated Averaging (Fed
Bit Config
Baseline
KURE APQN (W-4)
APQN (W-2) QAT (W-4)
QAT (W-2)
MQAT-W
W-32
49.17
49.57
50.77
50.8
49.19
13.62
48.47
W-8
49.24
49.52
50.86
50.88
46.46
1.7
48.21
W-6
48.99
49.66
50.27
50.81
17.34
1.92
48.28
W-4
45.37
47.75
45.76
46.34
49.7
3.05
47.18
W-3
36.42
41.04
34.87
37.04
44.44
11.68
45.88
W-2
6.62
7.62
24.77
42.92
42.27
Table 6: Global volidation accuracy after weight quantization of various quantization robustness variants for Federated Averoging (FEDAVG) on federated version of CIFAR-100 dataset using ResNet-20 architecture.
Federated Averaging (FedAvg)
Bit Config
Baseline
KURE APQN (W-4)
APQN (W-2) QAT (W-4)
QAT (W-2)
MQAT-W
W-32
37.21
37.51
36.48
37.85
36,83
4.39
37.89
W-8
37.29
37.61
36.41
37.89
36.73
2.3
37.56
W-6
36.61
37.17
35.9
37.35
36.98
2.73
37.23
W-4
31.09
33.29
31.4
33.37
37.43
5.99
37.12
W-3
17.64
17.61
20.08
19.25
31.14
15.66
36.84
W-2
0.86
0.9
1.02
0.89
7.28
34.53
35.45
Table 7: Global validation accuracy after weight quantization of various quantization robustness variants for Federated Averoging (FEDAVG) on federated version of TinglmogeNet dataset using ResNet-18 architecture.
26

---

Published in Transactions on Machine Learning Research (09/2023)
Federated Averaging (FedAvg) APQN (W-2)
Bit Config
Baseline
KURE
APQN (W-4)
QAT (W-4)
QAT (W-2)
MQAT-W
W-32
69
69.4
69.66
68.92
65.68
27.74
65.68
W-8
69.02
69.4
69.66
68.86
55.1
11.52
65.66
W-6
68.72
69.38
69.54
68.54
59.42
11.74
66.26
W-4
68.24
68.8
68.66
68.26
66.74
12.62
66.02
W-3
66.82
67.12
67.06
66.94
58.4
15.36
62.51
W-2
48.68
51.94
51.3
53.14
61.44
61.5
59.32
Table 8: Global validation accuracy after weight quantization of various quantization robustness variants for Federated Averoging (FEDAvG) on federated version of CIFAR-10 dataset using LeNet-5 architecture.
Federated Averaging (Fed
Bit Config
Baseline
KURE
APQN (A-4)
APQN (A-2)
QAT (A-4)
QAT (A-2)
MQAT-A
A-32
70.16
66.82
70.42
50.56
63.74
41.04
76.28
A-8
70.06
66.88
70.3
50.55
73.98
53.24
76.38
A-6
70
66.78
70.1
49.98
73.72
53.22
76.42
A-4
64.5
63.38
65.66
41.49
71.26
53.1
76.14
A-3
47.92
50.03
51.48
22.72
51.52
53.1
74.68
A-2
12.66
26.2
15.2
3.8
11.06
58.2
67.78
Table Global validation accuracy after activation quantization of vorious quantization robustness variants for Federated Averaging (FEDAVG) on federated version of CIFAR-10 10 datasel using ResNet-20 architecture.
Federated Averaging (FedAvg)
Bit Config
Baseline
KURE
APQN (A-4)
APQN (A-2)
QAT (A-4)
QAT (A-2)
MQAT-A
A-32
49.17
37.36
50.94
50.56
18.13
4.74
40.71
A-8
49.37
37.22
50.86
50.55
45.62
23.91
43.79
A-6
48.73
37.25
50
49.98
45.57
23.82
44.23
A-4
41.36
32.5
40.6
41.4
43.06
23.13
42,59
A-3
24.33
20.16
22.22
22.72
25.31
24.14
30.5
A-2
1.59
3.4
4.13
3.8
1.24
7.81
28.79
Table Global validation accuracy after activation quantization of various quantization robustness variants for Federated Averaging (FEDAvG) on federated version of CIFAR-100 dataset using ResNet-P2 20 architecture.
Federated Averaging (FedAvg)
Bit Config
Baseline Baseline (Dropout)
QAT (WA-4)
QAT (WA-2)
MQAT-WA
WA-32/32
70.16
72.02
60.64
26.61
74.9
WA-8/8
69.86
71.62
71
28.08
77.3
WA-6/6
69.92
70.58
71.38
29.42
76.8
WA-4/4
63.12
64.08
70.04
36.44
75.54
WA-3/3
42.98
41.8
42.54
50.12
72.9
WA-2/2
11.04
14.5
9.58
40.18
66.08
Table 11: Global volidation accuracy after activation quantization of various quantization robustness variants for Federated Averaging (FEDAVG) on federaled version of CIFAR- 10 dataset using ResNet-20 architecture.
27


Reviews:
[{"review_id": "y3z1Jy3ITA", "reviewer": "Reviewer_8f8P", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. Well-motivated problem, solution makes sense and is well described. To my knowledge this paper is probably the first to study this problem in-depth through experiments and theoretical analysis.\n2. Experiments are conducted on quite a few FL benchmarks\n3. Theoretical analysis where appropriate.\n\nWeaknesses:\n1. I am not the most familiar with this specific problem, but I was curious as to what is a strong baseline for this problem. Since it is a relatively niche problem, what are the baselines people expect to work? Is the proposed approach a 'strong baseline' (ie taking QAT and extending it to MQAT, and applying to FL), or is it something that we would not have expected to work, and it is cool that it did? If it is the former, phrase it as such and explain why it works with theory and experiments. If it is the second, I would have liked to see what the 'strong baselines' are, and how well they work. Since it's a new problem, I almost don't know how to contextualize how good these results are.\n2. There should be some analysis of additional computational or storage complexity given these modifications to the FL methods.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces variants of federated averaging with the underlying neural network quantized to various bit-widths with only limited reduction in full precision model accuracy so that models deployed on different devices continue to work despite different computational constraints (i.e., bit-width handling requirements) on those devices. The authors conduct experiments on several benchmarks and provide convergence analysis for their approach, and show promising empirical and theoretical results.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "see weakness points above regarding framing of the method, strong baselines, and analysis of computational or storage complexity.", "broader_impact_concerns": "none", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "zcR7cN1SZT", "reviewer": "Reviewer_4rVv", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n1. It considers multiple useful techniques to improve quantization-aware training in federated settings. They are not new, but to the best of my knowledge, they are first formally studied in FL in this submission.\n2. It gives convergence guarantees for almost all the algorithms.\n3. The color coding of the algorithms makes the paper easier to read.\n\nWeaknesses:\n\n1. Writing and motivation need to be improved. The notion of 'robust quantization' is first introduced in the third paragraph of the Introduction, but more context is needed to explain this specific technique or objective. In the introduction, the submission mentions that robust quantization is motivated by randomized smoothing, which would make readers think that the goal of robust quantization is to be robust to adversarial test samples, which is not what this paper means exactly (despite being related). The introduction is also hard to sparse---there are too many technical terms without proper explanation. For instance, what do Kurtosis Regularization, Quantization-Aware Training, Straight Through Estimator offer at a high level. The flow is a bit unnatural. The second paragraph is talking about the need of multi-bit training. But the third paragraph jumps into previous works on robust quantization without mentioning why robust quantization is important, and how it is connected with the previous paragraph.\n\n2. The convergence theorem doesn't cover the convergence of Kurtosis Regularization based objective. It would be interesting to also show Algorithm 1 can converge to the original FedAvg objective. The existing analysis would go through naturally (if the local quantization gradients are unbiased). It would be useful to highlight what is new in the current analysis. Assumption 3 should be written in a more formal form, in consistent with Assumptions 1-2. How does the convergence (along with the error term) scale with the number of local iterations K? How about convex cases? These questions are useful to discuss/answer as well.\n\n\n\n\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper studies model quantization  (quantization-aware training and several variants) in federated learning. The goal is to output a global model that has good average inference accuracy after quantization on each client. It presents multiple algorithms by adapting existing quantization techniques to federated settings, including leveraging Kurtosis Regularization, adding random noise to the weights, and optimizing multi-bit quantization.\n\nThe submission presents convergence guarantees of the quantization algorithms, and empirically compares the proposed methods with the baseline of directly quantizing FedAvg models, under different quantization levels.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Fixing 1. and 2. in the 'weaknesses' part.", "broader_impact_concerns": "None.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "poaXeq12VH", "reviewer": "Reviewer_mxjj", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strengths\n1. The paper addresses a timely topic and is generally well-written.\n2. The paper provides convergence theory, providing a principled foundation for the proposed approaches.\n3. The experiments illustrate the benefits of the proposed approaches, and especially MQAT, across several workloads and ablating several aspects of the formulation.\n\n## Weaknesses\n1. It is not clear what (if anything) is particularly complicated or a major contribution in the convergence theory, beyond combining convergence techniques for FedAvg with those typically used to analyze QAT and APQN in the non-federated setting. The analysis for MQAT is effectively the same as that of QAT.\n2. Some aspects of the experimental results could be investigated and discussed in a bit more depth, to more clearly support the claims  (see below).\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper considers the problem of Quantization-Aware Training (QAT) in the context of federated learning, to train models which can be easily quantized for on-device inference. This is important since many recent devices require that weights be quantized in order to be able to take advantage of on-device accelerator chips. In the context of federated learning, where the trained model is likely to be deployed on a wide range of devices, this paper specifically considers the problem of training a model that can be robustly quantized to a variety of different bit-widths and quantization step-sizes.\n\nTo this end, the paper considers incorporating standard techniques from the model quantization literature (APQN and QAT), while also proposing a straightforward variant of QAT termed MQAT (M for multiple), where at each training iteration a different quantization setting is sampled and used to determine the update. Convergence theory is provided for these three flavors of FedAvg, and experiments illustrate the promise of the proposed MQAT approach on multiple workloads.\n", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "In the revision I would appreciate if the authors can address the following three points:\n\n1. Can the QAT techniques studied in this paper be combined with gradient/update quantization typically used in FL to reduce communication overhead?\n\n2. In Figure 1, what is the intuition for why FedAvg-QAT (W-4) performs well at bit-widths higher than 4, even though it was only trained for 4?\n\n3. In Theorem 1, the term in the upper bound related to QAT depends on the parameter $R$, and for MQAT this depends specifically on the worst-case (largest) quantization step-size, so in theory the performance is no better than QAT for that particular setting. However, the MQAT algorithm uses a randomly sampled bit-width at each iteration. Might one hope to improve the upper bound for MQAT so that it depends on the average quantization step-size (or more generally, something better than the worst-case)?\n", "broader_impact_concerns": "None noted", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "fwGrkDuyX7", "reviewer": "Reviewer_n6mB", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strengths\n* The paper is well-structured;\n* The studied problem is timely and important.\n* The manuscript considers experiments on CIFAR-10/100, TinyImageNet, and FEMNIST. \n\n## Weaknesses\n* The manuscript did not discuss the related work well and ignored a significant fraction of prior work on improving federated/distributed learning on heterogeneous hardware systems. See e.g., [1, 2, 3, 4] and their follow-up works (as well as the discussed papers therein). \n    * The discussion on quantization-aware training is limited (only consider papers before the year 2021).\n    * There is no discussion/comparison of the latest post-training quantization methods, e.g., [5, 6].\n* The evaluation did not consider other standard FL methods (w/o considering the hardware heterogeneity), leaving the compatibility of the proposed method unknown.\n* The evaluation did not compare with other federated learning methods proposed for heterogeneous hardware.\n\n## Reference\n* [1] Heterofl: Computation and communication efficient federated learning for heterogeneous clients, ICLR 2021, https://arxiv.org/abs/2010.01264 (it has more than 200 Google Scholar citations)\n* [2] FedLite: A Scalable Approach for Federated Learning on Resource-constrained Clients, http://arxiv.org/abs/2201.11865\n* [3] QuPeL: Quantized Personalization with Applications to Federated Learning, http://arxiv.org/abs/2102.11786\n* [4] Quantized Distributed Training of Large Models with Convergence Guarantees, http://arxiv.org/abs/2302.02390\n* [5] NoisyQuant: Noisy Bias-Enhanced Post-Training Activation Quantization for Vision Transformers, http://arxiv.org/abs/2211.16056\n* [6] Optimal Brain Compression: A Framework for Accurate Post-Training Quantization and Pruning, http://arxiv.org/abs/2208.11580", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The manuscript considers the need for heterogeneous hardware-aware on-device learning and proposes quantization methods for federated learning. The claimed key contributions are\n* a multi-bit quantization-aware training (MQAT) over the existing QAT methods (up to the year 2021);\n* an additive pseudo-quantization noise\n* some theoretical analysis.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "* The manuscript needs to carefully discuss the related work in the field.\n* The manuscript needs to justify the superiority and compatibility of the proposed method.", "broader_impact_concerns": "NA", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper aims to improve robustness of FL models to quantization. It presents multiple algorithms by adapting existing quantization techniques to federated settings along with convergence guarantees of the quantization algorithms. Empirical evaluation of the proposed methods leaves more to be desired. We thank the authors and reviewers to actively engage in discussion for making the paper better. During discussion many of the reviewer concerns/questions were resolved like the one about convergence and storage requirements. However some concerns remain about baselines and related works. Nevertheless, the proposed method and technique are correct, experimental result demonstrate robustness to quantization in one way, and the approach will be of interest to the community, hence I propose to accept the paper with some minor modifications.

1. Add evaluation comparing to standard FL methods (w/o considering the hardware heterogeneity), which will allow us to understand the compatibility of the proposed method.
2. Add evaluation comparing with other federated learning methods proposed for heterogeneous hardware.

Finally not required but it would be great if the manuscript could also evaluate the proposed method on the transformer architectures and compare results with that of the corresponding quantization methods.
Decision: Accept (poster)