SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Discussion on the proposed protocol and related work. Our proposed wait-free local training protocol is specially designed for the PFL+ scenario, where each client trains two models locally. For simplicity, we use this protocol in our experiments, under the assumption of synchronized PFL+. For comparison in the asynchronous PFL+ setting [47], see Appendix. Related work that reduce the total training runtime, such as delayed gradient averaging [48] and wait-free decentralized FL training [49], are designed for conventional FL and does not deal with the PFL+ scenario. Furthermore, we also provide a discussion on how our wait-free local training protocol could be adapted to the partial client participation scheme in FL in the Appendix.
Methods	α = 1	α = 0.5	α = 0.1
GM	PM	GM	PM	GM	PM
FedAvg	85.35 ± 0.11	(80.33 ± 0.38)	80.76 ± 0.13	(74.51 ± 0.48)	73.51 ± 0.17	(72.68 ± 0.39)
FedProx	85.61 ± 0.08	(86.28 ± 0.21)	80.54 ± 0.14	(76.88 ± 0.30)	71.96 ± 0.12	(73.77 ± 0.30)
FedDyn	86.03 ± 0.13	(85.33 ± 0.19)	80.88 ± 0.18	(78.93 ± 0.25)	73.62 ± 0.14	(74.25 ± 0.58)
FedGen	86.17 ± 0.32	(85.24 ± 0.47)	79.86 ± 0.34	(77.52 ± 0.43)	71.36 ± 0.28	(71.42 ± 0.63)
FedAvgM	85.44 ± 0.05	(82.85 ± 0.28)	81.04 ± 0.09	(75.71 ± 0.33)	72.87 ± 0.06	(72.96 ± 0.14)
pFedMe	85.58 ± 0.23	88.17 ± 0.17	79.33 ± 0.14	84.66 ± 0.17	72.11 ± 0.23	81.18 ± 0.15
Ditto	85.34 ± 0.10	87.55 ± 0.09	80.70 ± 0.13	83.39 ± 0.12	73.45 ± 0.18	80.08 ± 0.20
FedRep	(85.61 ± 0.19)	87.32 ± 0.11	(80.33 ±0.23)	84.10 ± 0.10	(73.50 ± 0.24)	79.74 ± 0.31
FedRoD	86.02 ± 0.12	91.67 ± 0.16	81.31 ± 0.15	85.91 ± 0.15	74.64 ± 0.07	81.37 ± 0.17
FedBABU	(85.67 ± 0.24)	91.34 ± 0.19	(79.57 ± 0.23)	83.22 ± 0.33	(73.88 ±0.19)	80.58 ± 0.22
Ours	86.37 ± 0.15	92.25 ± 0.21	81.27 ± 0.18	86.59 ± 0.17	75.52 ± 0.11	82.69 ± 0.16
Table 1: Average (3 trials) and standard deviation of the best test accuracies for generic/personalized models of various methods on CIFAR-10 with different non-IID settings. See also Remark 4.1.
Methods	α = 1	α = 0.1
GM	PM	GM	PM
FedAvg	48.37 ± 0.22	(52.64 ± 0.48)	38.61 ± 0.27	(39.27 ± 0.42)
FedProx	47.33 ± 0.15	(53.85 ± 0.33)	39.55 ± 0.18	(41.33 ± 0.38)
FedDyn	49.24 ± 0.27	(57.20 ± 0.35)	40.43 ± 0.14	(40.92 ± 0.26)
FedAvgM	48.55 ± 0.19	(55.60 ± 0.26)	39.03 ± 0.08	(40.85 ± 0.19)
pFedMe	47.29 ± 0.27	61.52 ± 0.25	38.22 ± 0.23	45.88 ± 0.32
Ditto	48.37 ± 0.25	60.47 ± 0.27	39.61 ± 0.19	43.12 ± 0.28
FedRep	(46.32 ± 0.23)	58.76 ± 0.36	(40.11± 0.35)	45.22 ± 0.19
FedRoD	50.07 ± 0.16	62.51 ± 0.15	40.58 ± 0.22	45.99 ± 0.14
FedBABU	(48.52 ± 0.30)	60.33 ± 0.28	(37.35 ± 0.29)	44.72 ± 0.28
Ours	51.39 ± 0.22	63.15 ± 0.16	40.67 ± 0.14	46.82 ± 0.23
Table 2: Average (3 trials) and standard deviation of the best test accuracies for generic/personalized models of various methods on CIFAR-100 with different non-IID settings. See also Remark 4.1.
3.5 Algorithm Summary
Our proposed spectral co-distillation framework combined with our wait-free local training protocol, is given in Algorithm 1. As an overview, we begin every communication round  $t$  with the server broadcasting the global generic model  $w_{\mathrm{G}}^{t-1}$  to each client for local computation. Each client  $i$  would send back the updated generic model  $w_{\mathrm{G},i}^{t}$  after  $E_{\mathrm{G}}$  local computation steps for global model aggregation, then immediately start the personalized model training and continue until the global generic model  $w_{\mathrm{G}}^{t}$  is received, which marks the start of the next communication round  $t+1$ .
Remark on convergence analysis. Note that the global loss function includes a weighted sum of the local loss functions and a regularizer. The regularizer is given in the form of the divergence function  $\mathfrak{D}$ , which is equivalent to KL divergence; cf. Sec. 3.2. As demonstrated in [50], the KL divergence usually exhibits convexity in terms of the model parameters. Consequently, since the model training undergoes (stochastic) gradient descent, it is possible to establish a convergence rate for the training of the global model (under the commonly employed assumption of smoothness of the local loss functions).
Datasets, DNN models, federated settings, and evaluation metrics. We evaluated our proposed PFL+ framework with  $N$  clients on CIFAR-10/100 [51], and iNaturalist-2017, using model archi
Methods	FedProx	FedDyn	Ditto	FedRep	FedRoD	FedBABU	Ours
GM	39.46±0.39	39.35±0.27	39.33±0.33	39.81±0.41	40.16±0.35	39.23±0.53	41.75±0.37
PM	41.58±0.27	40.99±0.35	41.88±0.41	42.07±0.24	44.54±0.29	42.36±0.44	45.87±0.21
tectures ResNet-18/34 [52] and ResNet-50, respectively. For the experiments on CIFAR-10 (resp. CIFAR-100), we used  $N = 100$  (resp.  $N = 50$ ). For experiments on iNaturalist-2017 [53], we used  $N = 20$ . For dataset partition, we used the symmetric Dirichlet distribution to emulate real-world heterogeneous data distributions [9, 11], where the heterogeneity is controlled by the concentration parameter  $\alpha$ . (A smaller  $\alpha$  indicates a higher degree of data heterogeneity.) For evaluation, we used two performance metrics:
- Generic model evaluation: global test accuracy (same metric in conventional FL).
- Personalized model evaluation: weighted average of local test accuracies.
For every client, the PM is evaluated on a local test set, whose underlying distribution is the same as that for the local training set. All the experiments are implemented with a full client participation scheme. Further experiment details, results on partial client participation, and the computation overhead discussion are provided in the Appendix.
Remark 4.1. For generic FL methods, personalized model (PM) accuracies are obtained by evaluating the generic model (GM) on local test sets. For PFL methods without GM training, GM accuracies are obtained by evaluating the averaged PM on the global test set.
Baselines. We compared our proposed method with the following SOTA PFL methods: pFedMe [17], Ditto [15], FedRoD [36], FedRep [24], and FedBABU [30]. Moreover, to have a fair performance evaluation of the generic models, we also include methods designed for conventional FL as baselines: FedAvg [19], FedProx [10], FedDyn [11], FedGen[54], and FedAvgM [9].
4.2 Performance comparison with state-of-the-art methods
We evaluated the generalizability of our proposed spectral co-distillation framework, as well as the communication cost performance of our wait-free training protocol for PFL+.
Generalizability over heterogeneous settings. We compared the best test accuracies with multiple baselines over the different levels of data heterogeneity, using the same system configuration. Tab. 1 and Tab. 2 give the main results on CIFAR-10 and CIFAR-100, respectively. In summary, our proposed framework achieves the best test accuracies across diverse heterogeneous data settings, outperforming all PFL and conventional FL baselines on both PM and GM test accuracies concurrently. We also investigated the performance on the real-world dataset iNaturalist2017 in Tab. 3, where our proposed method also achieves the best GM/PM test accuracies. We attribute such consistent outperformance to the bi-directional co-distillation design.
This demonstrates that: a) the spectral information of the generic model is useful for knowledge distillation during personalized model training; and b) using truncated spectral information of the personalized models could boost the performance of the generic model via careful spectrum truncation. (See Appendix for a sensitivity analysis of the truncation ratio  $\tau$  and other hyper-parameters.)
Communication cost comparison. To demonstrate the superiority of the wait-free training protocol (WF), we evaluated the communica
Table 3: Average (3 trials) and standard deviation of the best test accuracies for generic/personalized models of various methods on iNaturalist-2017 with non-IID setting  $\alpha = 0.1$ . See also Remark 4.1.
Methods	3 epochs	5 epochs
40%	80%	40%	80%
Speedup
Ours (w/ WF)	1.82 ×	1.56 ×	2.21 ×	1.85 ×
Ditto w/ WF	1.97 ×	1.38 ×	2.87 ×	1.93 ×
FedRoD w/ WF	1.75 ×	1.54 ×	2.42 ×	2.19 ×
Table 4: Communication cost comparison of various methods for personalized model accuracies on CIFAR-10 to reach target accuracy (40%/80%) with non-IID setting  $\alpha = 0.1$ . The speedup factors are with respect to the performance of the corresponding methods without WF.
tion cost performance of SOTA methods with/without the protocol on non-IID CIFAR-10 ( $\alpha = 0.1$ ), in terms of the total runtime  $\zeta_{\mathrm{total}}$  for PM to reach the target test accuracy (40%/80%). A smaller  $\zeta_{\mathrm{total}}$  indicates higher communication efficiency. For PFL methods that train generic and personalized models using the compute-and-wait local training protocol, we evaluated Ditto and FedRoD. We conduct experiments with different numbers of epochs for local PM training (3 or 5 epochs). As shown in Tab. 4, our proposed wait-free training protocol could significantly improve the efficiency of convergence time and has the potential to boost the time efficiency of PFL+ methods.
Method	α = 1	α = 0.1
GM	PM	GM	PM
Ours	86.37 ± 0.15	92.25 ± 0.21	75.52 ± 0.11	82.69 ± 0.16
Ours w/o SCD-GM	85.35 ± 0.11	91.86 ± 0.17	73.51 ± 0.17	81.03 ± 0.20
Ours w/o SCD-PM	82.74 ± 0.39	79.65 ± 0.83	68.96 ± 0.47	70.51 ± 1.21
Ours w/o Both	85.35 ± 0.11	79.65 ± 0.83	73.51 ± 0.17	70.51 ± 1.21
Table 5: Ablation study results on non-IID CIFAR-10 (average and standard deviation of 3 trials). SCD-GM (resp. SCD-PM) represents the spectral distillation approaches adopted during the training of generic (resp. personalized) model.
4.3 Ablation results
Ablation study. In our proposed spectral co-distillation framework, we introduce the bi-directional spectrum knowledge distillation to bridge the training of generic and personalized models with the target for training good generic and personalized models simultaneously. To achieve the target, truncated and full model spectrum information are adopted in different training stages. Here, we conduct an ablation study to evaluate the effectiveness of these two components (see Tab. 5 for the effects of each component), in which we apply the distillation approaches in the two training stages separately. In the setup where both SCD-PM and SCD-GM are removed (Case I), the GM training is identical to FedAvg. In the case of removing only SCD-PM while keeping SCD-GM (Case II), each PM would be trained locally without any knowledge distilled from the GM. This is akin to the client training its model by itself, separately from the server. Naturally, the PM performance would be drastically lower. As SCD-GM is kept in Case II, where the GM is the student and the PM is the teacher, since the PM's performance is drastically lower, we would expect a drop in the GM's performance. Informally, the model would be worse off with the distillation of bad knowledge, than without distillation.
As demonstrated in Tab. 5, both the distillation methods can boost the accuracy performance of generic and personalized models, whereas the bi-directional distillation can bridge the training performance of the generic and personalized models. Specifically, we can observe that, the SCD-PM module effectively transfers the knowledge from the generic model to the personalized model and avoids over-fitting during local training.
Generalizability on new joining clients. In a real-world PFL system, dynamic client participation should be regarded as an important factor to consider during algorithm design, in which there would be continually new clients joining the system during training. The PFL system needs to rapidly train a good personalized model that could generalize well on the new client's local data. To evaluate the generalizability of the system, we simulate a dynamic participation system with 80 in-training clients and 20 new clients on CIFAR-10 (partitioned
Figure 3: Performance comparison for generalizability on new clients of various methods.
by the Dirichlet distribution with
$\alpha = 0.1)$ , and deal with new clients with the global model-based fine-tuning approach. Fig. 3 gives the results of the average test accuracies of the new clients. Among all evaluated methods, our method has the best average test accuracies, illustrating the fast adaptive capability of our method.
In this work, we propose a spectral co-distillation framework for PFL to learn better generic and personalized models simultaneously. As far as we know, this is the first work in PFL that represents the (dis-)similarity of models via their Fourier spectra. Even without co-distillation, there have been no other works that explore spectral distillation in PFL (or even in FL). The advantage of this new approach is clear from our experiments: We achieved outperformance in both generic and personalized model training. Our framework also incorporates a simple yet effective wait-free local training protocol to reduce the overall local training time.
Limitations. Our proposed spectral co-distillation framework, as currently formulated, does not deal with stragglers and adversarial attacks. Their influence on performance would require further investigation. Also, our protocol assumes a synchronized network connection, which may not be practical for scenarios with large system/network heterogeneity. Moreover, it would be good to consider a more realistic local training protocol design that takes into account the issues of network/system heterogeneity; we leave the extension as future work.
This work is supported by the National Research Foundation, Singapore under its AI Singapore Program (AISG Award No: AISG-RP-2019-015), and by Ministry of Education, Singapore, under its Tier 2 Research Fund (MOE-T2EP20221-0016), and under its SUTD Kisktarter Initiative (SKI 2021_03_01), in part by the National Research Foundation, Singapore and Infocomm Media Development Authority under its Future Communications Research & Development Programme, in part by the National Natural Science Foundation of China under Grant 62271513, in part by the Zhejiang Lab Open Research Project (No. K2022PD0AB05), and in part by the Zhejiang-Singapore Innovation and AI Joint Research Lab.
[1] E. A. Van Dis, J. Bollen, W. Zuidema, R. van Rooij, and C. L. Bockting, "Chatgpt: five priorities for research," Nature, vol. 614, no. 7947, pp. 224-226, 2023.
[2] E. J. Hu, Y. Shen, P. Wallis, Z. Allen-Zhu, Y. Li, S. Wang, L. Wang, and W. Chen, "Lora: Low-rank adaptation of large language models," arXiv preprint arXiv:2106.09685, 2021.
[3] A. Z. Tan, H. Yu, L. Cui, and Q. Yang, "Towards personalized federated learning," IEEE Transactions on Neural Networks and Learning Systems, 2022.
[4] J. P. Albrecht, “How the gdpr will change the world,” Eur. Data Prot. L. Rev., vol. 2, p. 287, 2016.
[5] Y. Deng, M. M. Kamani, and M. Mahdavi, "Adaptive personalized federated learning," arXiv preprint arXiv:2003.13461, 2020.
[6] J. Wang, Z. Charles, Z. Xu, G. Joshi, H. B. McMahan, M. Al-Shedivat, G. Andrew, S. Avestimehr, K. Daly, D. Data, et al., “A field guide to federated optimization,” arXiv preprint arXiv:2107.06917, 2021.
[7] J. Xu, Z. Chen, T. Q. Quek, and K. F. E. Chong, "Fedcorr: Multi-stage federated learning for label noise correction," in Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 10184-10193, 2022.
[8] Q. Yang, Y. Liu, T. Chen, and Y. Tong, "Federated machine learning: Concept and applications," ACM Transactions on Intelligent Systems and Technology (TIST), vol. 10, no. 2, pp. 1-19, 2019.
In the appendix, we provide further details organized into the following parts:
- Sec. B introduces the implementation details for our proposed method with respect to the federated system settings, evaluation metrics, and the data partitions used in our experiments. In particular, we provide the implementation details for the experiments in "Generalizability on new joining clients".
- Sec. C gives further sensitivity analysis on the key parameters  $\tau$ ,  $\lambda_{\mathrm{G}}$ , and  $\lambda_{\mathrm{p}}$ .
- Sec. D gives a discussion on the computational overhead.
- Sec. E gives a discussion on the differences and relations to asynchronous (personalized) federated learning (FL).
- Sec. F gives a discussion on how our wait-free local training protocol could be adapted for partial client participation.
Federated system settings. All experiments were implemented using Pytorch. For all methods, we used an SGD local optimizer with a momentum of 0.5 and with no weight decay. We train all methods over a total number of  $T = 500$  global communication rounds. Batch size for CIFAR-10/100 [51] and iNaturalist-2017 [53] are 10 and 128, respectively. As our aim is to train a personalized model for each client, we conducted the training using a full participation scheme, i.e., in every communication round, all clients would be selected to participate in training. For our proposed method, we used a learning rate of 0.01 (resp. 0.003) for both  $\eta_{\mathrm{G}}$  and  $\eta_{\mathrm{p}}$  when training on CIFAR-10/100 (resp. iNaturalist-2017).
Details on the data partition. We evaluated the performance of all methods with non-IID data partition. We simulated a non-IID data distribution via the symmetric Dirichlet distribution, in which we used the concentration parameter  $\alpha$  to control data heterogeneity across the FL system [9, 11]. For the entire dataset  $\mathcal{D}$  with  $M$  classes, we sample a stochastic vector  $\mathbf{d}_j$  with length  $N$  (total number of clients) for each  $1 \leq j \leq M$  via the Dirichlet distribution with common parameter  $\alpha$ . We shall treat  $\mathbf{d}_j$  as the vector representing the portion of the data samples in class  $j$  to be allocated to the clients. For example, if the second entry of  $\mathbf{d}_1$  equals 0.05, then client 2 would be randomly allocated with  $5\%$  of the data samples from  $\mathcal{D}$  in class 1. This allocation has no overlaps; two different clients have no common data samples. In particular, a smaller  $\alpha$  indicates a higher degree of data heterogeneity, whereas a larger  $\alpha$  will tend to generate identical-like distributions. (As  $\alpha \to +\infty$ , the data allocations become uniform.)
Performance metrics. The detailed calculation methods for the two metrics used in the evaluation are given as follows:
- We used the global test accuracy (same metric used in conventional FL) to evaluate the collaboratively trained global generic model, where the accuracy refers to the global test set (e.g., for CIFAR-10, the test set has 10,000 images).
- We used the weighted average of local test accuracies for personalized model evaluations, where each local test accuracy is obtained with respect to a local test set, and where the weights of the weighted average are the ratios of local dataset sizes  $\frac{n_1}{n}, \frac{n_2}{n}, \dots, \frac{n_N}{n}$ . Let  $acc_i$  denote the test accuracy for the personalized model at client  $i$ . Then the weighted average of local test accuracies is computed as  $\sum_{i=1}^{N} \frac{n_i}{n} acc_i$ . Note that for each client, the distribution of the local test set for each client is the same as the distribution of the local training set. Specifically, we perform sampling for the local test sets according to the distribution of training data for each respective client, where sampling is performed without replacement on the test set of the corresponding benchmark dataset, which is obtained via the Dirichlet distribution.
Details for the experiments on new joining clients. To evaluate the generalizability of our proposed federated co-distillation framework when there are new joining clients, we simulated a PFL system with 80 existing clients and 20 new clients on the CIFAR-10 dataset with non-IID data partition. In particular, we first partitioned the dataset into 100 local training sets via the Dirichlet distribution. We then performed training for the pre-determined 80 clients over 300 communication rounds for all methods, after which we conducted local fine-tuning [55, 23] on the new clients independently with a
certain number of local epochs based on the well-trained global generic model. (We tried 5 and 10 local epochs in our experiments.) The reported performance is the weighted average of the local test accuracies over the 20 new clients, in which the weights are the ratios of the dataset sizes relative to these 20 clients.
C Sensitivity analysis
Recall from Sec. 3.3 of the main paper that  $\tau$  is the truncation ratio used to compute the truncated spectrum vectors. In our experiment results (as reported in the main paper), we used  $\tau = 0.4$ . To demonstrate the effect of this key hyperparameter, we conduct sensitivity analysis with multiple values of  $0 < \tau \leq 1$  on two non-IID CIFAR-10 settings; see Tab. 6.  $\tau = 0$  represents there is no spectral distillation for generic model training (which is equivalent to the setup "Ours w/o SCD-GM" in Tab. 5 of the ablation studies in the main paper), while  $\tau = 1$  indicates there is no spectrum truncation. In summary,  $\tau = 0.4$  achieves the best performance for both generic model (GM) and personalized models (PM), which shows the effectiveness of the truncated spectrum in generic model training and implicitly demonstrates the characteristics of the low frequency of model Fourier spectrum.
Method	α = 1	α = 0.1
GM	PM	GM	PM
τ = 0	85.35 ± 0.11	91.86 ± 0.17	73.51 ± 0.17	81.03 ± 0.20
τ = 0.1	85.18 ± 0.26	90.85 ± 0.39	72.66 ± 0.11	80.51 ± 0.25
τ = 0.3	85.42 ± 0.22	90.93 ± 0.38	72.59 ± 0.38	82.60 ± 0.19
τ = 0.4	86.37 ± 0.15	92.25 ± 0.21	75.52 ± 0.11	82.69 ± 0.16
τ = 0.5	86.33 ± 0.08	92.04 ± 0.29	74.91 ± 0.31	82.50 ± 0.36
τ = 0.7	86.08 ± 0.43	91.79 ± 0.44	72.85 ± 0.28	81.44 ± 0.43
τ = 0.9	86.25 ± 0.39	92.07 ± 0.28	72.57 ± 0.31	82.09 ± 0.37
τ = 1.0	85.47 ± 0.26	90.33 ± 0.41	73.55 ± 0.17	81.36 ± 0.22
We also provide experiment results for the sensitivity analysis of the coefficients  $\lambda_{\mathrm{G}}$  and  $\lambda_{\mathrm{p}}$  on CIFAR-10 with data setting  $\alpha = 1$  in Tab. 7 and Tab. 8.
Table 6: Sensitivity analysis for hyper-parameter  $\tau$  (i.e., the truncation ratio for spectrum truncation, used in generic model training) on non-IID CIFAR-10 (average and standard deviation of 3 trials).
Setup	0.01	0.05	0.1	0.3
GM	84.76	86.37	86.20	83.55
PM	90.33	92.25	91.35	89.64
Table 7: Sensitivity analysis for hyper-parameter  ${\lambda }_{\mathrm{G}}$  on non-IID CIFAR-10.
Setup	0.01	0.05	0.1	0.3
GM	86.37	86.33	85.61	85.72
PM	92.25	91.89	89.47	89.33
Table 8: Sensitivity analysis for hyper-parameter  ${\lambda }_{\mathrm{p}}$  on non-IID CIFAR-10.
D Computational overhead of the proposed method
In our proposed method, the additional computational overhead is dominated by the computation of spectral vectors, when compared to other model similarity/difference-based PFL methods [15, 17]. In our experiments, we use the famous fast Fourier transform (FFT) algorithm to compute the discrete Fourier transform of an input (complex) vector. Informally, the FFT algorithm is an efficient and
widely used algorithm that converts an input signal from the original domain into a representation in the frequency domain.
In our framework, when given an input weight vector  $w_{\mathrm{p},i}^{t} \in \mathbb{R}^{d}$  (i.e. there are a total of  $d$  model weights), its associated spectral vector  $s(w_{\mathrm{p},i}^{t})$  is computed using the FFT algorithm, with corresponding computational complexity  $\Theta(d \log d)$ ; see, e.g., [56].
E Discussion on the differences/relations to asynchronized FL
Our main goal for proposing the wait-free  $\mathrm{PFL + }$  training protocol is to reduce the total running time overhead and thereby increase the training time efficiency. It should be noted that our wait-free local training protocol is synchronous, meaning that all the clients perform local computing, pushing generic model gradients, as well as generic model updates, in a synchronized manner. Similar to [48], our proposed protocol allows local computing during global communication.