QuPeD: Quantized Personalization via Distillation with Applications to Federated Learning

Abstract: 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 diverse resources. In this work, we introduce a quantized and personalized FL algorithm QuPeD that facilitates collective (personalized model compression) training via knowledge distillation (KD)  among clients who have access to heterogeneous data and resources. For personalization, we allow clients to learn 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 of 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.


Reviews and author responses:

 We hope that we have addressed all the concerns raised earlier, and we are happy to clarify any further questions you may have. We would like to update you on some additional numerical results we have regarding client sampling, as described below.

**Client sampling with 0.1 sampling ratio:** We did additional experiments for training CNN1 on MNIST with a client sampling fraction of 0.1 with 50 clients (i.e., 5 random clients are sampled for training in any round), where each client has data from 3 random classes. We obtained the following results (which are averaged over 2 runs):

FedAvg (FP): $92.87\\% \pm 0.05 $

QuPeD (FP): $98.17\\% \pm 0.32$

QuPeD (2 Bits): $98.01\\% \pm 0.15$

QuPeD (1 Bit): $97.58\\% \pm 0.23$

Per-FedAvg (FP):$95.80\\% \pm 0.29$

pFedMe (FP): $97.79\\% \pm 0.03$ 

We observe that under 0.1 client sampling ratio our method still outperforms others. Note that even QuPeD (1 Bit) is quite competitive compared to other methods. This indicates the robustness of our compression method to client sampling.

We hope that we have addressed all the concerns raised earlier, along with new results on client sampling. We would greatly appreciate it if you could reevaluate your rating based on our responses. Please let us know if you have any more concerns.

 We hope that we have addressed all the concerns raised earlier, and we are happy to clarify any further questions you may have. We would like to update you on some additional numerical results we have regarding client sampling, as described below.

**Client sampling with 0.1 client sampling ratio:** In the first response, we reported additional numerical results pertaining to client sampling with a client sampling fraction of 1/3rd. We further decreased this ratio and did additional experiments with 0.1 client sampling fraction. Specifically, we did additional experiments for training CNN1 on MNIST with 50 clients, where each client has data from 3 random classes. We obtained the following results (which are averaged over 2 runs):

FedAvg (FP): $92.87\\% \pm 0.05 $
	
QuPeD (FP): $98.17\\% \pm 0.32$
	
QuPeD (2 Bits): $98.01\\% \pm 0.15$
	
QuPeD (1 Bit): $97.58\\% \pm 0.23$
	
Per-FedAvg (FP):$95.80\\% \pm 0.29$
	
pFedMe (FP): $97.79\\% \pm 0.03$ 

We observe that under 0.1 client sampling ratio our method still outperforms others. Note that even QuPeD (1 Bit) is quite competitive compared to other methods. This indicates the robustness of our compression method to client sampling.

We hope that we have addressed all the concerns raised, along with new results on client sampling. We would greatly appreciate it if you could reevaluate your rating based on our response. Please let us know if you have any more concerns.

  We thank the reviewer for suggesting acceptance of our work. We appreciate that you went over all reviews and our responses and decided to keep your rating. In the following, we summarize the additions to the appendices in the final version of the paper (if accepted):

* A subsection on soft quantization function, in Appendix A.

* A subsection on knowledge distillation, in Appendix A. 

* A subsection extending our proof of Theorem 2 to include client sampling, at the end of Appendix C.

* A subsection demonstrating experiments with client sampling, at the end of Appendix E. 

We want to note that these sections are ready to be added. We will add further details/discussions on experimental setup as suggested by you and other reviewers. 

 I have read other reviews and find the author's responses reasonable. I still remain in my initial evaluation and recommend the paper for *acceptance*.

In general, I agree with other reviewers that some aspects of the paper is not explained in details and for the non-expert reader it is hard to gauge the contribution. However, given the space limitation and several FL related practical consideration of this work, I think the authors did a good job condensing essential details in the main part of the paper. I would highly recommend the authors to include separate sections in the appendix explaining key components of their method/approach (such as knowledge distillation and soft quantization) to facilitate the reading and discuss further relaxations (such as partial participation and experimental setups). 

 We thank the reviewer for increasing their rating.

This paper proposes a new FL algorithm QuPeD for quantization and personalization in FL. 

QuPeD combines soft quantization surrogate and knowledge distillation to learn compressed personalized models with different quantization and model structures and uses local updates to improve communication efficiency. The idea is novel. 

The paper well structured. It first explores a centralized training scheme (as well as its convergence property) and then extends it to FL setting where the mentioned three techniques are fused.

Their experiments show its effectiveness.  Originality:

QuPeD combines soft quantization surrogate and knowledge distillation to learn compressed personalized models with different quantization and model structures and uses local updates to improve communication efficiency. The idea is novel. 

Quality: 

The method is technically sound. The proof seems correct.

Clarity:
The paper is well structured and clear. It first explores a centralized training scheme (as well as its convergence property) and then extends it to FL setting where the mentioned three techniques are fused. 

Significance:

Their experiments show its effectiveness. It shows that QuPeD has better test accuracy on image classification than previous work.

 The author leaves the relaxation of assumption as future work.

However, I have some other concerns:

1. Why should we have equal weights for $f(\mathbf{x})$ and $f(\widetilde{Q}_{\mathbf{c}}(\mathbf{x}))$ in (6)? Intuitively, different weights might force the algorithm to focus on different goals, for better accuracy (in optimization) or better quantization (for reference). After all, (6) can be viewed as a way towards multi-tasks.
2. Theorem 1 shows the gradient norm decays without any constant error term, while Theorem 2 didn’t. The author didn’t explain the reason and only say that we should choose the clients whose $\kappa_i$ is small in Remark 4. From the previous analysis in FL, the data heterogeneity should not introduce additional error (e.g., see [1-2]). It will only slow convergence rather than introduce error. So why here we can’t obtain a similar result?

[1] Khaled, Ahmed, Konstantin Mishchenko, and Peter Richtárik. "Tighter theory for local SGD on identical and heterogeneous data." International Conference on Artificial Intelligence and Statistics. PMLR, 2020.

[2] Haddadpour F, Mahdavi M. On the convergence of local descent methods in federated learning[J]. arXiv preprint arXiv:1910.14425, 2019.

3. Typically, if no local update is used (i.e., $\tau=1$), we should recover the result of Theorem 1 by setting $\tau=1$ in Theorem 2. However,  if we do that for line 314, I can’t recover the $O(1/T)$ convergence rate. What’s more, I can’t even remove the constant error term $\bar{\kappa}$, since it doesn’t depend on $\tau$. With Point 1 and this, I guess the analysis of Theorem 2 might not be tight. 
4. The work didn’t consider partial participation, so it is unclear how can we select the clients whose $\kappa_i$ is small in practice, as suggested in Remark 4. Besides, the work also didn’t consider the stochastic gradient variant. But, I think it is fine to defer them as future work since the main idea is well delivered already. 
5. In experiments, I think it is better to stress the importance of updating the centers (line 989) in the main text rather than in the appendix. Besides, the increased data heterogeneity will decrease the performance of (almost all) algorithms (from Table 3 and Table 5). It actually supports the theory and worth mentioning in the main text. 



===================
I have read the author's response. The author almost handles all my concerns. I am willing to increase my score to 7.


 We thank the reviewer for their comments and time. We will address their concerns one by one:

* **Different weights to $f(\mathbf{x})$ and $f_{Q_\mathbf{c}}(\mathbf{x})$:** Thank you for pointing this out; this is indeed a good point. However, since our algorithm comprises of many components, it already has 5 hyper-parameters (3 learning rates, $\lambda$, and $\lambda_p$). Adding other hyper-parameters, which may be less critical for our method, will add too much complication to the parameter tuning task; and that is why we gave equal weights to both the loss functions. We agree that giving different weights can only increase the performance of our algorithm, which, with our current formulation, is already the best performing among the competing methods. We can certainly try your suggestion and incorporate in a future revision if it is not too difficult and gives more improvements.

* **Error term on the convergence:** Thank you for this important observation. The reason why the referenced works do not have the constant terms is because they present the convergence guarantee for the global model, whereas, we give the convergence guarantee for personalized models. We believe that the error in our convergence is unavoidable as clients learn different personalized models (tailored to their local data) and so it is likely that the convergence error will depend on the data heterogeneity.

* **Not recovering the $\mathcal{O}(\frac{1}{T})$ rate and the error term:** Thank you for asking this question. First we would like to point out that the underlying settings of Theorem 1 and Theorem 2 not only differ in local iterations, but also in that the former is a result for centralized setting whereas the latter is a result for distributed setting where clients have heterogeneous data. In the following, first we answer both the concerns and then describe how we can also recover $\mathcal{0}(\frac{1}{T})$ convergence rate of Theorem 1:

	**Presence of the error term:** We believe that the error in our convergence is unavoidable as clients learn different personalized models (tailored to their local data) and so it is likely that the convergence error will depend on the data heterogeneity. We also would like to point out that there are in fact other papers on personalized federated learning (albeit in a much simpler setting than ours) on heterogeneous data that get convergence rate and error similar to ours; see, for example, [1] below. 

	**$\mathcal{O}(\frac{1}{T})$ convergence rate:** In fact, it all depends on what error we want to tolerate; for example, we could also get $\mathcal{O}(\frac{1}{T})$ rate, but with a slightly increased error, as explained below. Let us consider the setting with $\tau=1$ (to match the setting in your argument) in our bound in Theorem 2. Note that with our current convergence rate of $\mathcal{O}(\frac{1}{\sqrt{T}})$ in Theorem 2, we get an error that scales purely with data heterogeneity and nothing else (see the last term on line 756 in the proof of Theorem 2). To get a better rate, consider the bound on the gradients in line 755. Note that the learning rate $\eta_3$ has to satisfy $\eta_3\leq\frac{1}{4\lambda_pL_w\sqrt{C_L}}$. If we simply put this in the bound on the RHS of line 755, we would also get a rate of $\mathcal{O}(\frac{1}{T})$ but with an error that now depends not only on the heterogeneity but also on the Lipschitz constants and other parameters, which may be slightly bigger than the error that we get with the $\mathcal{O}(\frac{1}{\sqrt{T}})$ convergence rate, as in Theorem 2. In order to get an error as small as possible, we choose $\eta_3=\frac{1}{4\lambda_pL_w\sqrt{C_L}\sqrt{T}}$, which gives $\mathcal{O}(\frac{1}{\sqrt{T}})$ rate with an error that only depends on the data heterogeneity. 
	
	**Recovering the $\mathcal{O}(\frac{1}{T})$ convergence rate of Theorem 1:** Since Theorem 1 is for the centralized case and Theorem 2 is for the distributed case, in order to recover the convergence rate of Theorem 1 as a special case of Theorem 2, we need to take $\tau=1$ and $n=1$ (note that $n=1$ implies that there is only one client, i.e., the centralized case, which in turn implies that the data heterogeneity is zero, i.e., $\overline{\kappa}=0$ in Assumption A.6). Substituting these and $\eta_3=\frac{1}{4\lambda_pL_w\sqrt{C_L}}$ in the proof of Theorem 2 (on line 755, RHS), we can recover the $\mathcal{O}(\frac{1}{T})$ convergence rate of Theorem 1 without any error term.
* **Partial client participation and stochastic gradient variant:** We did not consider partial client participation and stochastic gradient in the main paper because we wanted to focus on our novel problem formulation combining model compression and personalization along-with its convergence analysis. Incorporating partial client participation, stochastic gradient variant, and analyzing the resulting algorithm is fairly simple. We have extended our analysis incorporating partial client participation, and leave the stochastic gradient variant as a future work. Partial client participation essentially changes Lemma 1 and Corollary 1, as everything before that is for local updates only. Now we give a summary of what changes:

	Let $\mathcal{K}_t$ denote the set of clients that participates at time $t$, where $|\mathcal{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 $\gamma_t$ similarly (see line 760). 

	**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 $\||\frac{1}{K} \sum_\{k \in \mathcal{K_t} \} \nabla_\{\mathbf{w}^j}F_k(\mathbf{x_k}^\{j+1\},\mathbf{c_k}^\{j+1\},\mathbf{w}^\{j\}) - \nabla_\{\mathbf{w}^j}F_i(\mathbf{x_i}^\{j+1\},\mathbf{c_i}^\{j+1\},\mathbf{w}^{j})\||^2$. Earlier, the average was over all clients from $1$ to $n$ and this term was bounded by $\kappa_i$ using Assumption A.6. Now, we can use the Jensen's inequality (iteratively) and Assumption A.6 and bound this by $2\kappa_i+\frac{2}{K}\sum_{j\in\mathcal{K}_t}\kappa_j$. This change will propagate over til the end. That's it.

	**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 in lines 743-756 and using the modified bounds of Lemma 1 and Corollary 1 will give the bound of our algorithm with partial client participation. We have essentially given the modification in the entire proof here.

	**New experiments with partial client participation:** We also ran experiments on the FEMNIST dataset with partial client participation. We fixed the ratio of participating clients to be $1/3$'rd in each round and sampled $n/3$ clients accordingly uniformly at random. With the same setup in the main paper we obtained the following results: 

	FedAvg (FP):  $91.30\\% \pm 0.43$
	
	QuPeD (FP):  $94.93 \\%\pm 0.25$
	
	QuPeD (2 Bits): $ 94.56 \\% \pm 0.18$
	
	QuPeD (1 Bit):  $ 92.52 \\% \pm 0.64$
	
	Per-FedAvg (FP):  $92.10 \\% \pm 0.22$
	
	pFedMe (FP):  $ 93.70 \\% \pm 0.39$
	
	It is clear that even with partial client participation, QuPeD outperforms the competing methods even when aggressive compression is employed. We will include these results regarding client sampling in the next revision.
* **Suggestion for experiments.**  Thank you for your suggestion regarding the experiments. We will incorporate that in our revised paper.

We hope to have answered your questions and concerns adequately; if not, we are happy to clarify further.

**References**

 [1] Dinh et al. "Personalized Federated Learning with Moreau Envelopes." Advances in Neural Information Processing Systems 33. 2020. 

  We thank the reviewer for their comments and time. We will address their concerns by one by.

* **Paper Organization:**  Firstly, apologies for missing part in Appendix A. We had prepared a section illustrating an example for soft quantization function and showing its smoothness properties, but commented out in the submission by mistake. In particular, we can define the following soft quantization function: $\widetilde{Q}_\mathbf{c}(\mathbf{x}):\mathbb{R}^{d+m}\rightarrow\mathbb{R}^d$ and $\widetilde{Q}_\mathbf{c}(\mathbf{x})_i := \sum_\{j=2\}^m(c_j-c_\{j-1\})\sigma(P(x_i-\frac{c_j+c_\{j-1\}}{2}))+c_1$ where $\sigma$ denotes the sigmoid function and $P$ is the so called temperature parameter controlling how closely $\widetilde{Q}_\mathbf{c}(\mathbf{x})$ approximates $Q_\mathbf{c}(\mathbf{x})$. This function can be seen as a simplification of the function that was used in [1] and can be shown to provide the theoretical properties that are required for Assumptions A.4 and A.5. Note that as stated in the end of Appendix E.1, in the implementation we let $P\rightarrow\infty$; this simplifies the implementation and provides good experimental results. See Appendix E.1 for the implications. We would like to add the section on soft quantization function in the next revision.
	

	Thank you for your feedback from an FL engineering perspective. We agree, indeed, we would have liked to carry content from Appendix to the main paper, particularly regarding the experiments; however, we followed the literature and left the details to Appendix and provided the main results (both theoretical and experimental) in the main paper. Note that we also could not present a proof sketch in main paper due to space limitations, we sacrificed both from theory and experiments. Nonetheless, we will add more pointers in main text towards implementation details in the revision.

* **Proposed solution:**  Assuming the availability of more resources for training and not for inference time is in fact motivated by FL with edge devices (*e.g.*}, smartphones), where model training is typically done when the phone is charging and is idle (so more power/memory is available), whereas, inference is more frequent and is done when phones are actively used and have other functions. Since relatively less power/memory is available during inference time, the deployed model should respect stricter resource constraints. This is actually, for example, how the Google assistant improves itself; see the following reference: https://support.google.com/assistant/answer/10176224?hl=en#zippy=%2Chow-google-assistant-improves-with-federated-learning for an application where training is done only when the phones are idle. Consequently, we still think our method is practically relevant.

* **Empirical Results:**  For an empirical convergence plot, please refer to Figure 1 in main paper where we see the convergence behaviour in terms of test accuracy; we can provide additional plots if required. Indeed, $T$ in theoretical results mean the number of epochs. While the convergence rates do not directly translate to experiments, they are an indispensable part of not only FL literature but also theoretical deep learning and optimization literatures. Translating convergence rates to experiments is an active research area, however, it is not developed enough to directly describe empirical behaviour. Convergence bounds are still of great value as they provide insights on algorithms and reveal how different parameters affect the speed of the algorithm (see, for example, Remark 4). As a result, we think that the theoretical results we present are not a detour but an integral part of our work.

We hope to have answered your questions and concerns adequately; if not, we are happy to clarify further.

[1] Yang et al. "Quantization networks." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2019.

 We thank the reviewer for their comments and time. We will address the concerns one by one:

 * **Partial client participation.** We did not consider partial client participation in the main paper because we wanted to focus on our novel problem formulation combining model compression and personalization along-with its convergence analysis. Incorporating partial client participation and analyzing the resulting algorithm is fairly simple. We have extended our analysis incorporating this, which essentially only changes Lemma 1 and Corollary 1, as everything before that is for local updates only. Now we give a summary of what changes:

	Let $\mathcal{K}_t$ denote the set of clients that participates at time $t$, where $|\mathcal{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 $\gamma_t$ similarly (see line 760). 

	**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 $\||\frac{1}{K} \sum_\{k \in \mathcal{K_t} \} \nabla_\{\mathbf{w}^j}F_k(\mathbf{x_k}^\{j+1\},\mathbf{c_k}^\{j+1\},\mathbf{w}^\{j\}) - \nabla_\{\mathbf{w}^j}F_i(\mathbf{x_i}^\{j+1\},\mathbf{c_i}^\{j+1\},\mathbf{w}^{j})\||^2$. Earlier, the average was over all clients from $1$ to $n$ and this term was bounded by $\kappa_i$ using Assumption A.6. Now, we can use the Jensen's inequality (iteratively) and Assumption A.6 and bound this by $2\kappa_i+\frac{2}{K}\sum_{j\in\mathcal{K}_t}\kappa_j$. This change will propagate over til the end. That's it.

	**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 in lines 743-756 and using the modified bounds of Lemma 1 and Corollary 1 will give the bound of our algorithm with partial client participation. We have essentially given the modification in the entire proof here.

	**New experiments with partial client participation:** We also ran experiments on the FEMNIST dataset with partial client participation. We fixed the ratio of participating clients to be $1/3$'rd in each round and sampled $n/3$ clients accordingly uniformly at random. With the same setup in the main paper we obtained the following results: 

	FedAvg (FP):  $91.30\\% \pm 0.43$
	
	QuPeD (FP):  $94.93 \\%\pm 0.25$
	
	QuPeD (2 Bits): $ 94.56 \\% \pm 0.18$
	
	QuPeD (1 Bit):  $ 92.52 \\% \pm 0.64$
	
	Per-FedAvg (FP):  $92.10 \\% \pm 0.22$
	
	pFedMe (FP):  $ 93.70 \\% \pm 0.39$
	
	It is clear that even with partial client participation, QuPeD outperforms the competing methods even when aggressive compression is employed. We will include these results regarding client sampling in the next revision.
* **Client sub-sampling on comparable methods:** Indeed, for a fair comparison, we considered full participation of clients for all comparable methods in the experiments. However, as mentioned above, we re-ran the experiments with client sub-sampling on all the methods and the resulting performance is stated above.

* **Data distribution for experiments:** In the FedAvg paper, authors use heterogeneous distribution only for MNIST (see Section 3 paragraph 2), where each client gets samples from a subset of classes. This is called pathological non-IID distribution, which was used in many subsequent works such as [1,2] and many more. What we meant in lines 355-356 is that we use a similar pathological non-IID distribution for our experiments on CIFAR-10 to distribute the samples such that each client has access to samples from 4 random classes out of 10. We referred to FedAvg because this kind of distribution originated in FedAvg for MNIST. 	

* **Effect of hyper-parameters:** While the individual parameters depend mainly on the structure and the data distribution, there are several empirical observation that we made. Now we will describe our observations for $\lambda$ (the parameter that controls quantization) and $\lambda_p$ (the parameter that controls the personalization). 
	
	As $\lambda$ gets smaller, the effect of proximal steps fades towards regular gradient steps. A too small $\lambda$ results in accuracy loss when hard quantizing the network, as the weights do not get updated in a 'quantization aware' fashion. As $\lambda$ gets larger, the centers' updates are affected more from the location of weights rather than the gradients; similarly, weight updates are affected more from the location of the centers. A too large $\lambda$ can result in network converging too early to a hard quantized network, which would also incur accuracy loss.
	
	For $\lambda_p$, a large choice makes the local updates to be dependent on the global model; this is problematic in early iterations where global model does not perform well. Small value of $\lambda_p$ means the local updates do not rely on the information from global model; this is disadvantageous in settings where global model performs well, such as when the data is distributed rather homogeneously or when global model is pretrained on some data.  
* **New experiments with text data:** Thank you for your suggestion. We tested our algorithm on a text classification problem using AG News dataset (available at https://pytorch.org/text/stable/datasets.html). We used half of the dataset to make the training procedure more challenging. We used EmbeddingBag structure available at https://pytorch.org/tutorials/beginner/text_sentiment_ngrams_tutorial.html and distributed the data such that each of the 42 clients has access to samples from 3 out of 4 classes. We obtained the following results:

	FedAvg (FP):  $83.04\\% \pm 0.60$
	
	Local training (FP):  $84.20\\% \pm 1.53$
	
	Local training (2 Bits):  $82.68\\% \pm 0.34$
	
	Local training (1 Bit):  $82.12\\%\pm 2.17$
	
	QuPeD (FP):  $85.06\\% \pm 1.07$
	
	QuPeD (2 Bits):  $83.62\\% \pm 0.50$
	
	QuPeD (1 Bit):  $82.72\\%\pm 1.20$ 

	The above results clearly demonstrate the effectiveness of QuPeD on text data. We will put these results in the next revision. If required, we can perform additional experiments, however, we expect a similar trend with that as well.
* **Motivation for quantization:** In the personalized FL setting, each client keeps a personalized model to be used for inference on the test data that it possesses/samples. Our method outputs individual models along-with a set of quantization values for each client. Clients then quantize their models in a personalized way with different codebooks according to their training data. In FL, one of the main challenges is resource constraints (such as memory) on clients; these constraints can differ from device to device. Our method provides both different values and precision for quantization to resolve the two challenges: data and resource heterogeneity. 

We hope to have answered your questions and concerns adequately; if not, we are happy to clarify further. 

**References**

[1] Dinh et al. "Personalized Federated Learning with Moreau Envelopes." Advances in Neural Information Processing Systems 33. 2020.

[2] Zhang, Michael, et al. "Personalized federated learning with first order model optimization." ICLR 21.

  We thank the reviewer for their encouraging comments and time. 

**Knowledge distillation loss:** In the following, we omit client index for simplicity. Let $\mathbf{x} \in \mathbb{R}^{d'}$ and $\mathbf{w} \in \mathbb{R}^{d}$ (where $d$ and $d'$ may not be equal) be two models possibly with two different architectures. They can be members of any hypothesis space as long as they work on the same type of data and they output a probability distribution for classification; so they do not have to be neural networks. Then using data samples $\{\xi_i\}$, the KD loss is generally defined as $f^{KD}(\mathbf{x},\mathbf{w}; \xi):= \frac{1}{n}\sum_{i=1}^{n} D_{KL}(s^{w}(\mathbf{w};\xi_i)\||s(\mathbf{x};\xi_i)) = \frac{1}{n} \sum_{i=1}^{n} \sum_{j=1}^{M} \log(s^{w}(\mathbf{w};\xi_i)_j) \frac{(s^{w}(\mathbf{w};\xi_i))_j}{(s(\mathbf{x};\xi_i))_j}$. Note that we explicitly provided data samples as parameterization of the functions for clarity; however, in general, in literature (and in our paper too), data samples are omitted for simplicity from the definition of loss functions. In the definition of KD loss, $M$ is the number of classes, $n$ is the number of samples, $s(\mathbf{x};\xi_i)$ is the output class probabilities with model $\mathbf{x}$ and the data sample $\xi_i$. The KD loss can be defined for any type of classifier such as neural networks, perceptrons, SVMs, or linear classifiers. In neural networks, this would commonly correspond to output of the softmax operation on logit activations. In the definition, $s(\mathbf{x};\xi_i)_j$ for $j\in[M]$ denotes the $j$'th element of the output vector $s(\mathbf{x};\xi_i)$.

We would like to emphasize that KD loss is crucial for the algorithm as a whole, as it allows interaction of global model with local models of different structures. 

The paper provides an algorithm (QuPeD) to learn personalized models in a federated learning (FL) setting where there is heterogeneity in data and computing resources. The method utilizes quantization and knowledge distillation (KD). The paper also formalizes KD in personalized FL as an optimization problem and proposes alternative proximal updates as a solution (with convergence analysis). Empirical evaluation is conducted on CIFAR-10 and FEMNIST, and compared against local training (no collaboration) + two quantization regimes and comparable models: Per-FedAvg, pFedMe, and FedAvg.
  In the basic FL simulation setting, a given $M$ number of client devices (subset of total client devices) are sampled for training. However, I don't see this is the case for QuPeD. Algorithm 2 states "On Clients $i = 1$ to $n$, where from Eq. (1) $n$ is total number of clients? Also, this leads to the questions of how the comparable methods: Per-FedAvg, pFedMe, and FedAvg were set up (since they use random sub-sampling of clients) to make the empirical evaluation comparable.

Line 355,356: "To simulate data heterogeneity on CIFAR-10, similar to [FedAvg],...". But, the FedAvg paper states: "since there isn’t a natural user partitioning of this data, we considered the balanced and IID setting" under CIFAR experiments.

There are several hyper-parameters to be tuned for QuPeD (algorithm 2) and the theoretical analysis provides some insights. However, it would be more insightful to see the effects of hyper-parameters empirically.

Empirical analysis focuses on vision models and datasets, but what are the implication in applying the method to language based models and datasets?

Although the paper states "the main goal of our work is personalized quantization for inference." (line 70), it was not clear (at least not convincing) to me from the supporting motivation w.r.t. the paper on the use of quantization.
 The authors addressed limitations and potential negative societal impact of their work.


This paper investigates Federated Learning setup with several practical consideration including (i) personalization where each client aims to train a local model suited to its local dataset while benefiting from the overall data, (ii) compressed local models to facilitate memory-constraint applications, (iii) deferent model dimension and precision for different clients to address data and resource heterogeneity. To handle this aspects of FL, this work designs an alternative optimization problem, which, besides the standard loss of the models, also minimizes the loss of softly quantized models and knowledge distillation loss between local models and the global model, encourages the models to be close to quantization centers. To solve that optimization problem, a new method called QuPeD has been proposed and theoretical analysis for non-convex and smooth objectives has done. Convergence rates are complemented with extensive numerical evaluations on deep nets.  Overall, the presentation of the paper is clear and clean with necessary motivations. Personalization is an important aspect in FL and the relaxed optimization problem together with the proposed method seems interesting contribution in this direction (especially under data and resource heterogeneity). The paper models several practical aspects into a new optimization problem, provides solid convergence guarantees and extensive numerical experiments.


I have one concern about knowledge distillation, which I am least familiar with.

How KD loss is defined exactly and how important is it for the theory ? It seems that the analysis only cares about smoothness of the loss and not the particular structure of it. As a reader not familiar with knowledge distillation much, description given in lines 193-195 does not explains what is the loss actually (e.g., functions $s_i^w$ and $s_i$) and whether it is only for neural networks or can be defined for other models too. I expected more details in the appendix, but could not find. Yes

This paper tackles the problem of Federated Learning personalization, with both resource and data heterogeneity. The proposed method optimizes three objectives: (a) personalized FL loss, (b) staircase-quantization loss and explicit distance loss, and (c) Knowledge Distillation (KD) loss. This paper makes a lot of assumptions and tricks to make the above quantization-related objectives differentiable.
The major difference between this paper and other FL personalization and FL model compression works is that the proposed algorithm (QuPeD), aims at learning with full accuracy in the Federated Training phase (with the objective that it will be quantized in the adaptation phase), while heterogeneous quantization and knowledge distillation are only conducted in inference time. QuPed conducts alternative FL optimization, on full accuracy model, quantization centers, and local models one-by-one, which shows improved quantized personalized performance, on FEMNIST and Federated CIFAR-10 datasets.  However, I do have some comments, which requires the authors to shed light on:

(1) Paper organization: the main context ignores a lot of key materials due to the limitation of pages, which requires searching for the appendix to understand. For example, the "prox" step, which is very frequently used in the main context, is not shown in the main text. The notation in the main context actually makes it harder and more confusing to understand. Moreover, the soft quantization function is not illustrated clearly. I cannot understand how to implement the soft quantization function, by only reading the main context. The presentation should improve, so that FL engineers like me could understand how to implement the algorithms, without getting lost in details that I don't understand.

(2) Proposed solution: The setup assumes that each client can do full accuracy training, but need different resources to conduct different levels of inference. The best usage in my opinion for the proposed algorithm, is FL on algorithms that will be deployed to IoT devices: Phone conduct full accuracy training, while the trained algorithm is deployed to IoT devices which require different dimension models and different levels of quantization. The proposed algorithm is not a direct fit, to current FL systems, where both training and inference are conducted on the same phone. 

(3) Empirical results: The empirical results support the claim that QuPed improves against its benchmarks. However, I find it hard to understand the theoretical convergence rate and the empirical results. Does the theoretical convergence rate mean the number of training epochs? Without showing the convergence rate empirically, the theoretical results seem to be a relative detour to the main context and could leave to the appendix.

Unfortunately, I am not knowledgeable enough to evaluate the theoretical part of this work, other than understanding enabling differentiability and convergence speed. Due to the above reasons, I leave a borderline rating, and I am willing to read other reviewers’ comments and authors’ answers, to change my rating accordingly. See above.

Meta-review: The reviewers had a few concerns that were addressed in the authors' response.  The majority of the reviewers agreed that the paper provides an interesting contribution to FL. The paper is missing a lot of references, on the quantization side as well as other topics, and I recommend the authors to carefully review the literature and cite the related work properly.  

Decision: accept