{'fraction_non_alphanumeric': 0.044012084592145014, 'fraction_numerical': 0.017570996978851963, 'mean_word_length': 4.190816710575837, 'pattern_counts': {'":': 0, '<': 8, '<?xml version=': 0, '>': 1, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 9, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
QuPeD: Quantized Personalization via Distillation with Applications to Federated Learning

Traditionally, federated learning (FL) aims to train a single global model while collaboratively using multiple clients and a server. Two natural challenges that FL algorithms face are heterogeneity in data across clients and collaboration of clients with {\em diverse resources}. In this work, we introduce a \textit{quantized} and \textit{personalized} FL algorithm QuPeD that facilitates collective (personalized model compression) training via \textit{knowledge distillation} (KD) among clients who have access to heterogeneous data and resources. For personalization, we allow clients to learn \textit{compressed personalized models} with different quantization parameters and model dimensions/structures. Towards this, first we propose an algorithm for learning quantized models through a relaxed optimization problem, where quantization values are also optimized over. When each client participating in the (federated) learning process has different requirements for the compressed model (both in model dimension and precision), we formulate a compressed personalization framework by introducing knowledge distillation loss for local client objectives collaborating through a global model. We develop an alternating proximal gradient update for solving this compressed personalization problem, and analyze its convergence properties. Numerically, we validate that QuPeD outperforms competing personalized FL methods, FedAvg, and local training of clients in various heterogeneous settings.


Introduction
Federated Learning (FL) is a learning procedure where the aim is to utilize vast amount of data residing in numerous (in millions) edge devices (clients) to train machine learning models without collecting clients' data [27]. Formally, if there are n clients and f i denotes the local loss function at client i, then traditional FL learns a single global model by minimizing arg min It has been realized lately that a single model may not provide good performance to all the clients in settings where data is distributed heterogeneously. This leads to the need for personalized learning, where each client wants to learn its own model [8,9]. Since a locally learned client model may not generalize well due to insufficient data, in personalized FL process, clients maintain personalized models locally and utilize other clients' data via a global model. Resource diversity among clients, which is inherent to FL as the participating edge devices may vary widely in terms of resources, is often overlooked in personalized FL literature. This resource diversity may necessitate clients to learn personalized models with different as well as different dimension/architecture. Systematically studying both these resource heterogeneity together with data heterogeneity in personalized FL is the primary objective of this paper.
In this work, we propose a model compression framework 1 for personalized FL via knowledge distillation (KD) [16] that addresses both data and resource heterogeneity in a unified manner. Our framework allows collaboration among clients with different resource requirements both in terms of precision as well as model dimension/structure, for learning personalized quantized models (PQMs). Motivated by FL, where edge devices are resource constrained when actively used (e.g. when several applications are actively running on a battery powered smartphone) and available for training when not in use (e.g., while charging and on wi-fi), we do training in full precision for learning compressed models to be deployed for inference time. For efficient model compression, we learn the quantization parameters for each client by including quantization levels in the optimization problem itself. First, we investigate our approach in a centralized setup, by formulating a relaxed optimization problem and minimizing it through alternating proximal gradient steps, inspired by [4]. To extend this to FL for learning PQMs with different dimensions/architectures, we employ our centralized algorithm locally at clients and introduce KD loss for collaboration of personalized and global models. Although there exist empirical works where KD is used in personalized FL [21], we formalize it as an optimization problem, solve it using alternating proximal updates, and analyze its convergence.

Contributions.
Our contributions can be summarized as follows: • In the centralized case, we propose a novel relaxed optimization problem that enables optimization over quantization values (centers) as well as model parameters. We use alternating proximal updates to minimize the objective and analyze its convergence properties.
• More importantly, our work is the first to formulate a personalized FL optimization problem where clients may have different model dimensions and precision requirements for their personalized models. Our proposed scheme combines alternating proximal updates with knowledge distillation.
• For optimizing a non-convex objective, in the centralized setup, we recover the standard convergence rate of O( 1 /T ) (despite optimizing over quantization centers), and for federated setting, we recover the standard convergence rate of O( 1 / √ T ) (despite learning PQMs with different precisions/dimensions). In the federated setting, our convergence bound has an error term that depends on multiplication of two terms averaged over clients: one characterizing client's local model smoothness and the other data heterogeneity with respect to overall data distribution. 2

Problem Formulation
Our goal in this paper is for clients to collaboratively learn personalized quantized models (with potentially different precision and model sizes/types). To this end, below, we first state our final objective function that we will end up optimizing in this paper for learning personalized quantized models, and then in the rest of this section we will describe the genesis of this objective. Recall from (1), in the traditional FL setting, the local loss function at client i is denoted by f i . For personalized compressed model training, we define the following augmented loss function at client i: Here, w ∈ R d denotes the global model, x i ∈ R di denotes the personalized model of dimension d i at client i, c i ∈ R mi denotes the model quantization centers (where m i is the number of centers), Q ci denotes the soft-quantization function with respect to (w.r.t.) the set of centers c i , R(x i , c i ) denotes the distance function, f KD i denotes the knowledge distillation (KD) loss [16] between the two input models on client i's dataset, λ is a design parameter for enforcing quantization (large λ forces weights to be close to respective centers), and λ p controls the weighted average of regular loss and KD loss functions (higher λ p can be used when client data is limited). We will formally define the undefined quantities later in this section. Consequently, our main objective becomes: Thus, our formulation allows different clients to have personalized models x 1 , . . . , x n with different dimensions d 1 , . . . , d n and architectures, different number of quantization levels m 1 , . . . , m n (larger the m i , higher the precision), and different quantization values in those quantization levels. Note that there are two layers of personalization, first is due to data heterogeneity, which is reflected in clients learning different models x 1 , . . . , x n , and second is due to resource diversity, which is reflected in clients learning models with different sizes, both in terms in dimension as well as precision.
In Section 2.1, we motivate how we came up with the first three terms in (2), which are in fact about a centralized setting because the function f i and the parameters involved, i.e., x i , c i , are local to client i; and then, in Section 2.2, we motivate the use of the last two terms containing f KD

Towards Personalized Quantized Federated Learning: Knowledge Distillation
Note that the objective function defined in (6) can be used for learning a quantized model locally at any client. There are multiple ways to extend that objective for learning personalized quantized models (PQMs) via collaboration. For example, when all clients want to learn personalized models with the same dimension (but with different quantization levels), then one natural approach is to add an 2 penalty term in the objective that would prevent local models from drifting away from the global model and from simply fitting to local data. This approach, in fact, has been adopted in [8,14] for learning personalized models and in [23] for heterogeneous FL, though not quantized ones. In our previous work QuPeL [28] we analyzed a quantized approach for learning PQMs (having the same dimension). In Section 5, we demonstrate that QuPeD (for the same task but using KD as opposed to the 2 penalty) outperforms QuPeL.
In this paper, since we allow clients to learn PQMs with potentially different dimensions, the above approach of adding a 2 penalty term in the objective is not feasible. Observe that, the purpose of incorporating a 2 penalty in the objective is to ensure that the personalized models do not have significantly different output class scores compared to the global model which is trained using the data generated at all clients; this does not, however, require the global model to have the same dimension as that of local models and can be satisfied by augmenting the local objective (6) with a certain knowledge distillation (KD) loss. In our setting, since clients' goal is to learn personalized models with different dimensions that may also have different quantization levels, we augment the local objective (6) with two separate KD losses: where the first one ensures that the behavior of x i ∈ R di is not very different from that of w ∈ R d , and the second one ensures the same for the quantized version of x i and w. Formally, we define them using KL divergence as follows: , where s w i and s i denote functions whose inputs are global and personalized models, respectively -and data samples implicitly -and outputs are the softmax classification probabilities of the network.
We need to train x i and w mutually. Identifying the limitations of existing approaches for theoretical analysis (as mentioned in related work in Section 1), we use reverse KL updates (i.e., taking gradient steps w.r.t. the first parameter in D KL (·, ·)) to train the teacher network w from the student network x i . This type of update can be shown to converge and also empirically outperforms [31] (see Section 5). We want to emphasize that though there are works [21,24] that have used KD in FL and studied its performance (only empirically), ours is the first work that carefully formalizes it as an optimization problem (that also incorporate quantization) which is necessary to analyze convergence properties.

Centralized Model Quantization Training
In this section, we propose a centralized training scheme (Algorithm 1) for minimizing (6) by optimizing over x ∈ R d (the model parameters) and c ∈ R m (quantization values/centers). During training, we keep x full precision and learn the optimal quantization parameters c. The learned quantization values are then used to hard-quantize the personalized models to get quantized models for deployment in a memory-constrained setting.
Description of the algorithm. We optimize (6) through alternating proximal gradient descent. The model parameters and the quantization vector are initialized to random vectors x 0 and c 0 . The objective Algorithm 1 Centralized Model Quantization Input: Regularization parameter λ; initialize the full precision model x 0 and quantization centers c 0 ; a penalty function enforcing quantization R(x, c); a soft quantizer Q c (x); and learning rates η 1 , η 2 .
Then, under A.1-A.5 and for L min = min{ 1 η1 , 1 η2 }, L max = max{ 1 η1 , 1 η2 }, we have: Theorem 1 is proved in Section 6. As a result, we would have L min = L and L max = 2L. Note that the ratio L 2 max/L min increases as the quantization becomes more aggressive, and consequently, aggressive quantization has a scaling effect on convergence rate.

Personalized Quantization for FL via Knowledge Distillation
We now consider the FL setting where we aim to learn quantized and personalized models for each client with different precision and model dimensions in heterogeneous data setting. Our proposed method QuPeD (Algorithm 2), utilizes the centralized scheme of Algorithm 1 locally at each client to minimize (3) Here, x i , c i , denote the model parameters and the quantization vector (centers) for client i, and w denotes the global model that facilitates collaboration among clients which is encouraged through the knowledge distillation (KD) loss in the local objectives (2).
Description of the algorithm. Since clients perform local iterations, apart from maintaining x t i , c t i at each client i ∈ [n], it also maintains a model w t i that helps in utilizing other clients' data via collaboration. We call {w t i } n i=1 local copies of the global model at clients at time t. Client i updates w t i in between communication rounds based on its local data and synchronizes that with the server which aggregates them to update the global model. Note that the local objective in (2) can be split into the weighted average of loss functions ( and the term enforcing quantization λR(x i , c i ). At any iteration t that is not a communication round (line 3), client i first computes the gradient g t i of the loss function w.r.t. x t i (line 4) and then takes a gradient step followed by the proximal step using R (line 5) to update from x t i to x t+1 i . Then it computes the gradient h t i of the loss function w.r.t. c t i (line 6) and updates the centers followed by the proximal step (line 7). Finally, it updates w t i to w t+1 i by taking a gradient step of the loss function at w t i (line 8). Thus, the local training of x t i , c t i also incorporates knowledge from other clients' data through w t i . When t is divisible by τ , clients upload {w t i } to the server (line 10) which aggregates them (line 15) and broadcasts the updated global model (line 16). At the end of training, clients learn their person- Assumptions. In addition to assumptions A.
for some constants

Experiments
In this section, we first compare numerical results for our underlying model quantization scheme (Algorithm 1) in a centralized case against related works [2,34]. Here, both [2,34] considers proximal algorithms with 2 and 1 penalty respectively without optimizing the quantization levels; moreover, both methods are restricted to 1 bit quantization. For a major part of the section, we then compare QuPeD (Algorithm 2) against other personalization schemes [8,9,28,31] for data heterogeneous clients and demonstrate its effectiveness in resource heterogeneous environments.

Centralized Training
We compare Algorithm 1 with [2,34] for ResNet-20 and ResNet-32 [15] models trained on CIFAR-10 [19] dataset. While both [2,34] are limited to binary quantization, [2] can be seen as a specific case of our centralized method where the centers do not get updated. Specifically, previous works did not optimize over centers as we did in (6). From Table 1, we see that updating centers (Algorithm 1) significantly improves the performance (0.48% increase in test accuracy). Allowing quantization with 2 bits instead of 1bit for Algorithm 1 further increases the test accuracy. Our algorithm allows us to employ any number of bits for quantization.

Personalized Training:
We consider an image classification task on FEMNIST [5] and CIFAR-10 [19] datasets. We consider two CNN architectures: (i) CNN1 (used in [27]): has 2 convolutional and 3 fully connected layers, (ii) CNN2: this is CNN1 with an additional convolutional layer with 32 filters and 5 × 5 kernel size. For CIFAR-10 we choose a batch size of 25. For FEMNIST, we choose variable batch sizes to have 60 iterations for all clients per epoch. We train each algorithm for 250 epochs on CIFAR-10 and 30  ) 96.10 ± 0.14 73.52 ± 0.51 QuPeL (1 Bits) 94.06 ± 0.28 71.01 ± 0.32 pFedMe (FP) [8] 96.60 ± 0.37 73.66 ± 0.65 Per-FedAvg (FP) [9] 97. 16  epochs on FEMNIST. For quantized training, as standard practice [30], we let the first and last layers of networks to be in full precision. We use last 50 epochs on CIFAR-10, and 5 epochs on FEMNIST for the fine-tuning phase. Data Heterogeneity (DH): We consider n = 50 clients for CIFAR-10 and n = 66 for FEMNIST. To simulate data heterogeneity on CIFAR-10, similar to [27], we allow each client to have access to data samples from only 4 randomly chosen classes. Thus, each client has 1000 training samples and 200 test samples. On FEMNIST, we use a subset of 198 writers from the dataset and distribute the data so that each client has access to data samples written by 3 randomly chosen writers. The number of training samples per client varies between 203-336 and test samples per client varies between 25-40. Test samples are sampled from the same class/writer that training samples are sampled from, in parallel with previous works in heterogeneous FL.  [8] 97.79 ± 0.03 93.70 ± 0.39 Per-FedAvg (FP) [9] 95.80 ± 0.29 92.10 ± 0.22 Federated ML (FP) [31] 98.03 ± 0.31 92.73 ± 0.36

Resource Heterogeneity (RH):
To simulate resource heterogeneity for QuPeD, we consider 4 settings: (i) half of the clients have CNN1 in full precision (FP) and the other half CNN2 in FP, (ii) half of the clients have CNN1 in 2 bits and the other half CNN2 in FP, (iii) half of the clients have CNN1 in 2 bits and the other in FP, (iv) half of the clients have CNN1 in 2 bits and the other half CNN2 in 2 bits.

Results (DH):
We compare QuPeD against FedAvg [27], local training of clients (without any collaboration), and personalized FL methods: pFedMe [8], Per-FedAvg [9], Federated Mutual Learning [31], and QuPeL [28]. For all methods, if applicable, we set τ = 10 local iterations, use learning rate decay 0.99 and use weight decay of 10 −4 ; we fine tune the initial learning rate for each method independently, see Appendix D for details. The results are provided in Table 2 with full client participation (K = n), plotted in Figure 1 for CIFAR-10, and in Table 3 with client sampling where we state average results over 3 runs; all clients train CNN1 (see Appendix D for CNN2) and quantization values are indicated in parenthesis. Thus, we only consider model personalization for data heterogeneity. In Table 2, we observe that full precision QuPeD consistently outperforms all other methods for both datasets. Furthermore, we observe QuPeD with 2-bit quantization is the second best performing method on CIFAR-10 (after QuPeD (FP)) and third best performing method on FEMNIST despite the loss due to quantization. Hence, QuPeD is highly effective for quantized training in data heterogeneous settings. Since QuPeD outperforms QuPeL, we can also (empirically) claim that considering KD loss to encourage collaboration is superior to 2 distance loss. Lastly, we observe from Table 3 that QuPeD continues to outperform other methods under client sampling.

Results (DH+RH):
We now discuss personalized FL setting with both data and resource heterogeneity across clients. Note that since FedAvg, pFedMe, and Per-FedAvg cannot work in settings where clients have different model dimensions, we only provide comparisons of QuPeD with local training (no collaboration) to demonstrate its effectiveness. The results are given in Table 4. We observe that QuPeD (collaborative training) significantly outperforms local training in all cases (about 3.5% or higher on FEMNIST and 2.5% or higher on CIFAR-10) and works remarkably well even in cases where clients have quantized models without any significant loss in performance.

Proof of Theorem 1
This proof consists of two parts. First we show the sufficient decrease property by sequentially using Lipschitz properties for each update step in Algorithm 1. For each variable x and c we find the decrease inequalities and then combine them to obtain an overall sufficient decrease. Then we bound the norm of the gradient using optimality conditions of the proximal updates in Algorithm 1. Using sufficient decrease and bound on the gradient we arrive at the result. We leave some of the derivations and proof of the claims to Appendix B. Alternating updates. Remember that for the Algorithm 1 we have the following alternating updates: These translate to following optimization problems for x and c respectively (see end of the section for derivation):

Sufficient Decrease
This section is divided into two, first we will show sufficient decrease property with respect to x, then we will show sufficient decrease property with respect to c.

Sufficient Decrease Due to x
Using Claim 1 we have,

Claim 2. Let
Then Now we use Claim 2 and get, Using (10) we have, and obtain the decrease property for x:

Sufficient Decrease Due to c
Now we state the counterpart of Claim 2 for c.

Claim 3. Let
Then and using the bound in (12), we obtain the sufficient decrease for c:

Overall Decrease
Summing the bounds in (11) and (13), we have the overall decrease property: Let us define L min = min{L Then from (14): Telescoping the above bound for t = 0, . . . , T − 1, and dividing by T :

Bound on the Gradient
We now find the first order stationarity guarantee. Taking the derivative of (7) with respect to x at x = x t+1 and setting it to 0 gives us the first order optimality condition: Combining the above equality and Claim 1: . First order optimality condition in (8) for c t+1 gives: Combining the above equality and Claim 11: . Then: Summing over all time points and dividing by T : where in the last inequality we use (15). This concludes the proof of Theorem 1.

Proof of Theorem 2
In this part, different than Section 3, we have an additional update due to local iterations. The key is to integrate the local iterations into our alternating update scheme. To do this, we utilize Assumptions A. 6 and A.7. This proof consists of two parts. First, we show the sufficient decrease property by sequentially using and combining Lipschitz properties for each update step in Algorithm 2. Then, we bound the norm of the gradient using optimality conditions of the proximal updates in Algorithm 2. Then, by combining the sufficient decrease results and bounds on partial gradients we will derive our result. We defer proofs of the claims and some derivation details to Appendix C. In this analysis we take w t = 1 n n i=1 w t i , so that w t is defined for every time point. Alternating updates. Let us first restate the alternating updates for x i and c i : The alternating updates are equivalent to solving the following two optimization problems.
Note that the update on w t from Algorithm 2 can be written as: In the convergence analysis we require smoothness of the local functions F i w.r.t. the global parameter w. Recall the definition of (2). It follows that from Assumption A.7 that F i is (λ p (L D2 + L DQ3 ))-smooth with respect to w: Now let us move on with the proof.

Sufficient Decrease
We will divide this part into three and obtain sufficient decrease properties for each variable: x i , c i , w.

Sufficient Decrease Due to x i
We begin with a useful claim.
From Claim 4 ,after some algebra (deferred to Appendix C), we have: .
. Using the inequality from Claim 5 in (19) gives To obtain (a), we substituted the value of A(x t+1 i ) from Claim 5 into (19).

Sufficient Decrease Due to c i
In parallel with Claim 4, we have following smoothness result for c: From Claim 6 we have: gives us:

Sufficient Decrease Due to w
Now, we use (λ p (L D2 + L DQ3 ))-smoothness of F i (x, c, w) with respect to w: After some algebraic manipulations (see Appendix C) we have:

Overall Decrease
Let L (i) x , L (i) c for any i ∈ [n] and L w are defined as follows: Summing (21), (23), (24) we get the overall decrease property: Let L (i) min for any i ∈ [n] and L min are defined as follows: 18 Then, We have obtained the sufficient decrease property for the alternating steps; now, we need to arrive at the first order stationarity of the gradient of general loss function. To do this we move on with bounding the gradients with respect to each type of variables.

Bound on the Gradient
Now, we will use the first order optimality conditions due to proximal updates and bound the partial gradients with respect to variables x and c. After obtaining bounds for partial gradients we will bound the overall gradient and use our results from Section 7.1 to arrive at the final bound.

Bound on the Gradient w.r.t. x i
Taking the derivative inside the minimization problem (17) with respect to x at x = x t+1 i and setting it to 0 gives the following optimality condition: Then we have, is due to Claim 10 and A.1, and (c) is due to A.7. This implies: we have:

Bound on the Gradient w.r.t. c i
Similarly, taking the derivative inside the minimization problem (18) with respect to c at c = c t+1 i and setting it to 0 gives the following optimality condition: Then we have (34) and the last inequality is due to Claim 11 and Assumption A.7. As a result we have, we have:

21
Then, where in (a) we use the bound from (31), and L min is defined in (30). Now we state a useful lemma that bounds the average deviation between the local versions of the global model at all clients, and the global model itself. See Appendix C for a proof. (36)), then we have,

Lemma 1. Let η 3 be chosen such that
As a corollary: . Then, we have: max is defined in (36), and L max , L (min) max are defined in Lemma 1.
max }) 2 , using Lemma 1 and Corollary 1, summing the bound in (37) over time and clients, dividing by T and n: . Choice of η 3 . Note that in Lemma 1 we chose η 3 such that η 3 ≤ 1 6τ 2 λ 2 p L 2 w C L . Now, we further introduce upper bounds on η 3 .
• We can choose η 3 small enough so that L min = η 3 − 2λ p L w η 2 3 ; see the definition of L min in (30). • We can choose η 3 small enough so that η 3 −2λ p L w η 2 3 ≥ η3 2 . This is equivalent to choosing η 3 ≤ 1 4λpLw . These two choices imply L min ≥ η3 2 . In the end, we have 2 critical constraints on η 3 , {η 3 : T this choice clearly satisfies the above constraints. From (38) we have, In (a) we used L min ≥ η3 2 . Now, we plug in This concludes the proof.

Proof Outline with Client Sampling
Incorporating partial client participation and analyzing the resulting algorithm is fairly simple. Essentially only changes are in Lemma 1 and Corollary 1, as everything before that is for local updates only. Now we give a summary of what changes: Let K t denote the set of clients that participates at time t, where |K t | = K, i.e., K clients participate in the training process at any time. In this case, we define the average parameter w t and the gradient g t as the average over the respective parameters of only the active clients at time t; we also define γ t similarly.
• Change in the proof of Lemma 1: In the proof of Lemma 1, the second term on the RHS of the second inequality, with the above modification will be equal to 1 Earlier, the average was over all clients from 1 to n and this term was bounded by κ i using Assumption A.6. Now, we can use the Jensen's inequality (iteratively) and Assumption A.6 and bound this by 2κ i + 2 K j∈Kt κ j . This change will propagate over until the end.
• Change in the proof of Corollary 1: Since this is a corollary to Lemma 1, this will also see a similar change. • Remaining convergence proof: Now, continuing the exact same convergence proof and using the modified bounds of Lemma 1 and Corollary 1 will give the bound of our algorithm with partial client participation. This is the modification in the entire proof.

A.1 Notation
• Given a composite function g(x, y) we will denote ∇g(x, y) or ∇ (x,y) g(x, y) as the gradient; ∇ x g(x, y) and ∇ y g(x, y) as the partial gradients with respect to x and y.
• For a vector u, u denotes the 2 -norm u 2 . For a matrix A, A F denotes the Frobenius norm.
• Unless otherwise stated, for a given vector x, x i denotes the i th element in vector x; and x i denotes that the vector belongs to client i. Furthermore, x t i denotes a vector that belongs to client i at time t.

A.2 Equivalence of Assumption A.6 to Assumptions in Related Work
In particular, diversity assumption (Assumption 5) in [9] is as follows: where f i is local function, B is a constant and ∇ w f (w) = 1 n n j=1 ∇ w f i (w). Now we will show the equivalence to our stated assumption A.6. Let us define, Then we can define, as a result, we can further define g i (w t ) := ψ i (x i (w t ), c i (w t ), w t ). Therefore, our assumption A.6 is equivalent to stating the following assumption: At any t ∈ [T ] and any client i ∈ [n], the variance of the local gradient (at client i) w.r.t. the global gradient is bounded, i.e., there exists κ i < ∞, such that for every w t ∈ R d , we have: And we also define κ := 1 n n i=1 κ i and then, Hence, our assumption is equivalent to assumptions that are found in aforementioned works.

A.3 Alternating Proximal Steps
We define the following functions: f : Q(x, c). Throughout our paper we will use Q c (x) to denote Q(x, c), in other words both c and x are inputs to the function Q c (x). We propose an alternating proximal gradient algorithm. Our updates are as follows: For simplicity we assume the functions in the objective function are differentiable, however, our analysis could also be done using subdifferentials. Our method is inspired by [4] where the authors introduce an alternating proximal minimization algorithm to solve a broad class of non-convex problems as an alternative to coordinate descent methods. In this work we construct another optimization problem that can be used as a surrogate in learning quantized networks where both model parameters and quantization levels are subject to optimization. In particular, [4] considers a general objective function of the form F (x, y) = f (x) + g(y) + λH(x, y), whereas, our objective function is tailored for learning quantized networks: , c). Furthermore, they consider updates where the proximal mappings are with respect to functions f, g, whereas in our case the proximal mappings are with respect to the distance function R(x, c) to capture the soft projection.

A.4 A Soft Quantization Function
In this section we give an example of the soft quantization function that can be used in previous sections. In particular, we can define the following soft quantization function: )) + c 1 where σ denotes the sigmoid function and P is a parameter controlling how closely Q c (x) approximates Q c (x). Note that as P → ∞, Q c (x) → Q c (x). This function can be seen as a simplification of the function that was used in [33]. Assumption. For all j ∈ [m], c j is in a compact set. In other words, there exists a finite c max such that |c j | ≤ c max for all j ∈ [m].
In addition, we assume that the centers are sorted, i.e., c 1 < · · · < c m . Now, we state several useful facts.

Fact 4. Product of bounded and Lipschitz continuous functions is also Lipschitz continuous.
Fact 5. Let g : R n → R m . Then, the coordinate-wise Lipschitz continuity implies overall Lipschitz continuity. In other words, let g i be the i'th output then if g i is Lipschitz continuous for all i, then g is also Lipschitz continuous.
In our convergence analysis, we require that Q c (x) is Lipschitz continuous as well as smooth with respect to both x and c.
Claim 8. Q c (x) is l Q1 -Lipschitz continuous and L Q1 -smooth with respect to x.
Proof. First we prove Lipschitz continuity. Note, x − x ; using Fact 5 and the fact that i was arbitrary, there exists l Q1 such that 1:d,: . Now we focus on an arbitrary term of ∇ x Q c (x) j,i . From (2) we know that this term is 0 if i = j, and a weighted sum of product of sigmoid functions if i = j. Then, using the Facts 1-4 the function ∇ x Q c (x) j,i is Lipschitz continuous. Since i, j were arbitrarily Proof. For Lipschitz continuity we have, As a result, ∂ Qc(x)i ∂cj ≤ 2 + c max P 2 . Similar to Claim 8 using the facts that i is arbitrary and the Fact 5, we find there exists l Q2 such that In other words, Q c (x) is Lipschitz continuous. And for the smoothness, following the same idea from the proof of Claim 9 we find Q c (x) is L Q2 -smooth with respect to c.
The example we gave in this section is simple yet provides technical necessities we require in the analysis. Other examples can also be used as long as they provide the smoothness properties that we utilize in the next sections.

A.5 Lipschitz Relations
In this section we will use the assumptions A.1-5 and show useful relations for partial gradients derived from the assumptions. We have the following gradient for the composite function: where dim (∇h(x, c) that the soft quantization functions of our interest are elementwise which implies ∂ Qc(x)i ∂xj = 0 if i = j. In particular, for the gradient of the quantization function we have, Moreover for the composite function we have, In (5) and (4) we use x i , c j to denote (x) i and (c) j (i'th, j'th element respectively) for notational simplicity. Now, we prove two claims that will be useful in the main analysis.
To obtain (a) we have used the fact Ax 2 ≤ A F x 2 .
We can follow similar steps,

A.6 Assumption A.7 is a Corollary of other Assumptions
In this section we discuss how Assumption A.7 can be inferred from A.1-A.5. Here we drop client indices i for notational simplicity. Let us define f w (w) as the neural network loss function with model w. First we will argue that f KD (w, x) is smooth, given that f w (w) and f (x) are two smooth neural network loss functions with Cross Entropy as the loss function, and that f w (w) and f (x) have bounded gradients. These two standard assumptions imply the smoothness of f KD (w, x) individually with respect to both input parameters.
Proposition 1. f KD (w, x) is L D1 -smooth with respect to x and L D2 -smooth with respect to w for some positive constants L D1 , L D2 .
where i denotes the index of data sample, y i is the one hot encoding label vector, N is the total number of data samples, log denotes elementwise logarithm and 1 s(x;ξi) denotes elementwise inverse; softmax function s(x) : R d → R K , where K denotes the number of classes (similarly s w (w) : R d → R K is the function whose output is a vector of softmax probabilities and input is global model), is defined in Section 2, here we explicitly state that data samples ξ i is a parameterization of s. Assuming f (x) is smooth for any possible pair of (y i , ξ i ) implies log( 1 s(x)j ) is C x -smooth for some constant C x , here we used s(x) j to denote j'th output of s(x) and we omitted ξ i since log( 1 We have, where the operations are elementwise as before. In this expression only the last term depends on x and for each i, since s w (w; ξ i ) j ≤ 1 with j s w (w; ξ i ) j = 1, the expression is a weighted average of smooth functions log( 1 s(x)j ); as a result, f KD (w, x) is L D1 -smooth with respect to x for some constant L D1 . Now we investigate smoothness with respect to w. First, note that we can assume for all j and w, s w (w) j is lower bounded by a positive constant M > 0 and upper bounded by a positive constant P < 1 since, by definition, output vector of the softmax function contains values between 0 and 1 (we ignore the limiting case when a logit is infinitely large). Then, note that assuming a gradient bound on f w (w) = 1 (w;ξi) ) implies that for all j ∇ log( 1 s w (w)j ) = ∇s w (w)j s w (w)j ≤ G w for some constant G w (again the division operation is elementwise); since 0 < s w (w) j < 1 we have ∇s w (w) j ≤ G w . Moreover, similar to the first part, by assuming f w (w) is smooth we obtain that log( 1 s w (w)j ) is smooth for all j. This implies having a bounded Hessian: where (a) is due to reverse triangular inequality and 0 < s w (w) j < 1. As a result we obtain C + Proof. The proof is exactly the same as in Proposition 1, instead of using smoothness of f (x), using smoothness of f ( Q c (x)) with respect to x, c from Claims 10 and 11 gives the result.

B Omitted Details in Proof of Theorem 1
First we derive the optimization problems that the alternating updates correspond to. Remember we had the following alternating updates: For x t+1 , from the definition of proximal mapping we have: Note, in the third equality we remove the terms that do not depend on x. Similarly, for c t+1 we have: Minimization problems in (6) and (7) are the main problems to characterize the update rules and we use them in multiple places throughout the section.
Proof. From our assumptions, we have f is L-smooth. And from Claim 10 we have f ( Q c (x)) is (GL Q1 + G Q1 LL Q1 )-smooth. Using the fact that if two functions g 1 and g 2 are L 1 and L 2 smooth respectively, then g 1 + g 2 is (L 1 + L 2 )-smooth concludes the proof.

Claim (Restating Claim 2). Let
Then Proof. Let A(x) denote the expression inside the arg min in (6) and we know that (6) is minimized when . This proves the claim.

Claim (Restating Claim 3). Let
Then Proof. Let B(c) denote the expression inside the arg min in (7) and we know that (7) is minimized when c = c t+1 . So we have B(c t+1 ) ≤ B(c t ). This proves the claim.

C Omitted Details in Proof of Theorem 2
Claim (Restating Claim 4.).
Proof. From our assumptions, we have w) is L DQ1 -smooth with respect to x. And applying the Claim 10 to each client separately gives that f i ( Q c (x)) is (G (i) L Q1 + G (i) Q1 Ll Q1 )-smooth. Using the fact that if two functions g 1 and g 2 (defined over the same space) are L 1 and L 2 -smooth respectively, then g 1 + g 2 is (L 1 + L 2 )-smooth, and the fact that αg 1 is αL 1 -smooth for a given constant α concludes the proof. .

Rearranging the terms gives (24).
Lemma (Restating Lemma 1). Let η 3 be chosen such that (36)), then we have, Proof. Let t c be the latest synchronization time before t. Then: in (a) we use the facts that K i=1 a i 2 ≤ K K i=1 a i 2 , t ≤ τ + t c and that we are summing over non-negative terms. As a result, we have: . Then, we have: max is defined in (36), and L max , L (min) max are defined in Lemma 1.
Proof. From ( 1) ≤ ( 2) in the proof of Lemma 1, we have: Summing over all t c and dividing by T :

D Additional Details for Experiments
In this section, we first discuss the implementation details for the prox steps for Algorithm 1 and Algorithm 2 in Section D.1. Section D.2 discusses implementation details for the algorithms along with hyperparameters which was omitted in Section 5 of the main paper due to space constraints. and implemented ProxQuant using their published code 9 . We use a learning schedule for η 2 , particularly, we start with η 2 = 10 −4 and multiply it with 0.1 at epochs 80 and 140.

D.2.2 Federated Setting
For each of the methods we tuned the local step learning rate separately on the set {0.2, 0.15, 0.125, 0.1, 0.075, 0.05}. We observed that except for the two cases, for all other cases, 0.1 was the best choice for the learning rate in terms of accuracy: The two exceptions are the local training methods on FEMNIST and Per-FedAvg on CIFAR-10, for which, respectively, 0.075 and 0.125 were the best choices for the learning rate.
• Federated Mutual Learning [31]: Since authors do not discuss the hyperparameters in the paper, we used α = β = 0.25 for CNN1 and α = β = 0.15 for CNN2, similar to our use of λ p in QuPeD. Global model has the same learning schedule as the personalized models.
For QuPeD and Federated ML we used CNN1 as the global model in all settings. For the other methods where global and personalized models cannot be different we used the same structure as personalized models.

D.3 Additional Results for Federated Setting
In this section we provide additional experimental results for comparison of QuPeD with other pearsonalized learning schems from literature. Comparison on another CNN architecture (CNN2). We first report experimental results on CIFAR-10 for CNN2 in Table 5 (with the same setting we have for Table 2). This is a deeper architecture than CNN1, as described in Section 5 in the main paper.
For the results in Table 5, it can be seen that the comments made for Table 2 in the main paper directly hold as QuPeD is able to outperform other schemes by a significant margin. This demonstrates that QuPeD also works for a deeper neural network (than CNN1 considered in the main paper).
Another Type of Data Heterogeneity. We report results for another data heterogeneity setting where each client has access to data samples from random 3 classes on CIFAR-10. Sampling data from 3 random classes per client is a more challenging setting compared 4 classes per client considered in Section 5. In Table 6 we see that FedAvg's performance further decreased due to increased heterogeneity. Moreover, most of the other personalized FL methods are outperformed by local training whereas QuPeD still performs better than local training, and other personalized FL methods. We observe that QuPeD with 2 Bits aggressive quantization outperforms all the other competing methods except Federated ML [31] (for which it shows a similar accuracy). Moreover, QuPeD (1Bit) is able to outperform Per-FedAvg.