Heterogeneous Personalized Federated Learning by Local-Global Updates Mixing via Convergence Rate

• Is it correctly understood that you use a fixed learning rate for all experiments, which is not optimized?

For classification tasks, we follow the FedBN [1] to set the learning rate as $1e-2$. For the segmentation task, we use the learning rate of $1e-3$. We did a grid search to determine an appropriate range for the learning rate.

[1] Li, Xiaoxiao, Meirui Jiang, Xiaofei Zhang, Michael Kamp, and Qi Dou. "Fedbn: Federated learning on non-iid features via local batch normalization." International Conference on Learning Representations (2021).

• What are the practical implications of the convergence analysis in sec. 4? How is the theorem and proof different from existing convergence results?

The main point is it governs the convergence of our algorithm. The proof adapted the same method used in FL-NTK [1], whereas FL-NTK only shows the convergence of vanilla FedAvg, and we further extended the result to show the convergence for PFL algorithms with combined local and global updates. To be more specific, Prop.2 and Prop.3 in Appendix B.5 Proof of Theorem 1 are directly adapted from FL-NTK, which illustrates the effect of global and local updates on global and local models respectively. Our work mainly focuses on bounding the fluctuation caused by the cross effect, as shown in Prop.4 and Prop.5. We have added more discussion on our convergence analysis in the updated version.

[1] Huang, Baihe, Xiaoxiao Li, Zhao Song, and Xin Yang. "Fl-ntk: A neural tangent kernel-based framework for federated learning analysis." In International Conference on Machine Learning, pp. 4423-4434. PMLR, 2021.

• I do not understand the arguments in the paragraph titled "Feature value distribution by our personalized model". What results is the statement "If a model learns confident representations, then the related neurons should be highly activated..." based on?

Our statement is supported by Bau and Zhou et al.'s studies [1, 2]. They have demonstrated the strong correlation between neuron activation values and representation learning. Specifically, their studies show that neurons that effectively learn important representations exhibit high activation values. Consequently, forcibly deactivating these neurons results in a substantial drop in performance. We have included this reference in our manuscript.

[1] Bau, David, Jun-Yan Zhu, Hendrik Strobelt, Agata Lapedriza, Bolei Zhou, and Antonio Torralba. "Understanding the role of individual units in a deep neural network." Proceedings of the National Academy of Sciences 117, no. 48 (2020): 30071-30078.

[2] ​​Zhou, Bolei, David Bau, Aude Oliva, and Antonio Torralba. "Interpreting deep visual representations via network dissection." IEEE transactions on pattern analysis and machine intelligence 41, no. 9 (2018): 2131-2145.

• Are the methods described only applicable in the context of stochastic gradient descent? It would be beneficial to have a clear definition of the problem setup in which the proposed methods apply.

Yes, our analysis is based on the NTK, which is only applicable in the context of gradient descent.

• In sec. 3.1 both covariate shift and concept drift are mentioned. Is this method applicable to both scenarios? This should preferably be clear from the introduction.

Yes, our method is applicable to both covariate shift and concept drift. We have clarified this in Section 3.1. Our experimental evaluation also includes both covariate shift and concept drift.

• Are there any previous papers that consider the update strategy in eq. 3, or is this an idea that has not yet been explored?

The combination of local and global updates in PFL is not an uncommon practice, we have discussed the one related work APFL in the fourth paragraph in the introduction. However, our aim is to find a smart way to dynamically adjust the combination weight. To determine the weight, we draw inspiration from NTK. Our idea is to take the trace of the Gram matrix as a hint to determine the mixing ratio, which has not been explored yet.

• It is not always clear to me if w_c and u denote models or model weights?

We acknowledge that the notation we initially employed needed to be clarified. We have rectified this issue in the updated manuscript, by using $w$ and $u$ to denote the model weights, and $y(t)=f(w(t),x)$ denotes the model prediction with weights $w(t)$ and $f: \mathbb{R}^d \rightarrow \mathbb{R}$ is the model function.

• Where does eq. 4 come from? There is no details regarding its derivation or any clear reference.

Eq. 4 comes from Du et al.’s work [1], which uses the Gram matrix to characterize the optimization process of gradient descent for neural networks. This result is originally presented in Section 4 of "Gradient Descent Finds Global Minima of Deep Neural Networks" in the paper.

[1] Du, Simon S., Xiyu Zhai, Barnabas Poczos, and Aarti Singh. "Gradient descent provably optimizes over-parameterized neural networks." ICLR (2019).

• Is proposition 1 a novel result, and if so in which sense?

Regarding the novelty, proposition 1 provides new insight to motivate our idea, the intuition behind this proposition is different from conventional convergence theory. Our algorithm has demonstrated the effectiveness of using trace as an approximation of the convergence rate in practice. However, the accuracy of this approximation requires further investigation, which can be left for future work.

• It is uncler to me if the mixing ratio in eq. 5 is optimal in some sense? Intuitively, the global and local model might both have different error and different convergence properties, so should the model error not also be taken into account?

It would be inaccurate to claim that the ratio presented in Eq. 5 is optimal. Instead, it is a ratio derived from theoretical intuition and has been empirically verified to yield significant improvements.

We agree with you that model error between the local and global models can also affect the convergence property. Further research can be done to investigate whether taking the prediction error into account could speed up the convergence even further.

• Does it make a big difference that you only consider the last layer in the approximation of the Gram matrix? Can the arguments be supported empirically or theoretically?

The approximation should make no big differences in principle. Our approximation is supported by Seleznova et al.’s work [1], where they demonstrate that the empirical NTK with entries of last-layer features highly aligns with the original NTK regarding training dynamics. Their findings are supported both theoretically and empirically.

[1] Seleznova, Mariia, Dana Weitzner, Raja Giryes, Gitta Kutyniok, and Hung-Hsu Chou. "Neural (Tangent Kernel) Collapse." arXiv preprint arXiv:2305.16427 (2023).

• What exactly is meant by this statement: "since the data is distributed and not directly accessible, we propose employing features from local data in conjunction with the global model u"? Do you employ the representation (first layers) of the local model?

To calculate the trace, we employ the representation of both local and global models. Specifically, each client employs its local model to calculate the local feature, while also maintaining a copy of the global model to compute the global feature. The feature representation is extracted from the last hidden layer, that is, parameters up to the penultimate layer.

We sincerely thank you for your positive comments on our idea, method that makes good sense, and the experiments where our method is compared with many other relevant methods. Thank you for providing us with numerous constructive comments and suggestions. After carefully reviewing your comments, we found your concerns regarding the writing can be addressed via minor changes. We uploaded a revised version and highlighted the changes in blue color. In the following, we would like to address your questions regarding further clarification of our method and analysis.

• It is not clear if the results are a fair comparison, i.e. whether or not each algorithm is independently optimized sufficiently well.

Our comparative analysis is performed in a fair manner, as all the methods follow the same experimental settings. Furthermore, the hyperparameters of the compared method are carefully fine-tuned. The reported findings presented in this study are in alignment with recent publications [1, 2], further substantiating the credibility and reliability of our results.

[1] Huang, Wenke, Mang Ye, and Bo Du. "Learn from others and be yourself in heterogeneous federated learning." In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 10143-10153. 2022.

[2] Zhuang, Weiming, and Lingjuan Lyu. "Is Normalization Indispensable for Multi-domain Federated Learning?." arXiv preprint arXiv:2306.05879 (2023).

• The section on the convergence analysis seems a bit disintegrated with the rest of the paper.

We propose Theorem 1 in the convergence analysis to give a bound, while our method is motivated by Proposition 1. Theorem 1 serves as a more general proof of the convergence of the PFL algorithms combining local and global updates, as it poses no constraint on $\lambda$. We admit that more discussion can be made on the convergence analysis, and we have added it to the manuscript.

• Simple experiments that highlight the applicability of the method are not included. I would have appreciated a clear demonstration, e.g. in comparison with fixed local-global ratios.

Following your suggestion, we have extended the investigation on the performance by using different fixed ratios and the optimized ratio. For the fixed optimized ratio, we take the average ratio of our method in the last 10 training rounds. The results are shown in the following table. When the ratio equals to 0, our method will degenerate to the FedAvg, and when the ratio goes to 1, it equals to complete client local training. From the results, it can be observed that introducing the mixing of local and global updates helps improve the personalized model performance, and our dynamic adjustment shows the highest accuracy.

• An empirical demonstration of proposition 1 would have been a strong argument for the specific choices made for the mixing ratio.

Thank you for your insightful comment. Our proposition 1 mainly motivates us to approximate the convergence rate by measuring the trace of the Gram matrix. However, the computation of the Gram matrix $H$ in practical networks poses significant challenges and is often computationally infeasible [1,2,3].

[1] Fort, Stanislav, Gintare Karolina Dziugaite, Mansheej Paul, Sepideh Kharaghani, Daniel M. Roy, and Surya Ganguli. "Deep learning versus kernel learning: an empirical study of loss landscape geometry and the time evolution of the neural tangent kernel." Advances in Neural Information Processing Systems 33 (2020): 5850-5861.

[2] Mohamadi, Mohamad Amin, Wonho Bae, and Danica J. Sutherland. "A fast, well-founded approximation to the empirical neural tangent kernel." In International Conference on Machine Learning, pp. 25061-25081. PMLR, 2023.

[3] Holzmüller, David, Viktor Zaverkin, Johannes Kästner, and Ingo Steinwart. "A framework and benchmark for deep batch active learning for regression." Journal of Machine Learning Research 24, no. 164 (2023): 1-81.

We sincerely thank you for your positive comments on our paper writing, method motivation, and the comparative analysis of ablation studies to verify the details of our method. We would like to address each of your questions below.

• How to calculate the h(x_i)? Does it require storing a global model locally? And adding a storage comparison among baselines.

(a) Calculate the $h(x_i)$: The $h(x_i)$ is the latent feature representation of the input $x_i$ extracted from the last hidden layer, where $h$ denotes the parameters up to the penultimate layer.

(b) Store a global model locally: Yes. A global model is required to keep locally to calculate the global feature.

(c) Storage comparison among baselines： Following your suggestion, we monitored the peak GPU memory cost of each method and listed the costs in the table below. Our method takes a reasonable cost compared with other baseline methods. Please note that we do not report FedRep, FedBABU, and FedHKD on the retinal dataset because they are specifically designed for classification tasks.

GPU memory cost (MB)

• The training wall-clock time

For the computational time, we counted the training time of all methods. The GPU type used for training is GeForce RTX 2080 Ti. From the results in the table, it can be observed that our method lies in a reasonable range with a slightly higher cost than FedAvg, which is a trade-off between computational efficiency and performance.

Training time (hour)

• Draw the loss curve to illustrate the model improvements regarding faster convergence

As per your suggestion, we have included the training and validation loss values in the table provided below. The results indicate that PFL methods generally converge faster than the baseline method FedAvg. Among all the compared methods, our method clearly promotes the convergence speed and presents the lowest training and validation loss. Additionally, we have draw the training/validation loss curves in Section D.2, Figure 4, in the Appendix for your reference.

• What are the possible limitations of the proposed method, and how can they be addressed in future work?

Thanks for the insightful question. One potential limitation of our method is that we approximate the trace calculation by only using the last-layer latent feature. While this approximation is reasonable and reflects the convergence rate in principle, it may not be an accurate estimation when using large foundation models as the model backbone. The convergence behavior may differ due to the model's architecture and complexity. There is a risk of obtaining less precise estimations in these scenarios. To address this limitation, one potential approach is to consider the prediction errors of local and global models. By incorporating the model errors into the mixing process, we can obtain more comprehensive measurements for combining local and global updates. However, exploring the use of the foundation model in federated learning is still an open question. Explorations are needed in future work to validate the effectiveness of our method further. We have included a detailed discussion of our method's limitations and potential extensions in Section E in the Appendix.

• How is eq. 4 derived? please give more explanation on it. y is label, w(t) is the local model weights, what's the meaning of y-w(t)?

Eq. 4 is derived by following Du et al.’s work [1], which uses the Gram matrix to characterize the optimization process of gradient descent for neural networks. Eq. 4 is the major theorem (Theorem 4.1) in Du et al.’s work to show the convergence rate of gradient descent. It measures the prediction error, which is the difference between the model prediction and the ground truth label.

$y-w(t)$ should be $y-y(t)$ in Eq.(4), which denotes the prediction error. It is a typo, many thanks for pointing this out, and sorry for the confusion. In our paper, $w(t)$ is used to denote the local model weights, and $y(t)=f(w(t),x)$ denotes the model prediction with model weights $w(t)$ and $f: \mathbb{R}^d \rightarrow \mathbb{R}$ is the model function. We have revised our manuscript's Eq.(4) to avoid confusion.

[1] Du, Simon S., Xiyu Zhai, Barnabas Poczos, and Aarti Singh. "Gradient descent provably optimizes over-parameterized neural networks." ICLR (2019).

• What's the cost (e.g., GPU memory cost, computational time) of the proposed method, compared to the baseline methods?

The GPU memory cost and computational time cost of our method come from two aspects. First, we store a copy of the global model during local training. Second, our method calculates latent features and the trace of latent feature matrices.

For GPU memory cost, we monitored the peak GPU memory cost of each method and listed the costs in the table below. Our method takes a reasonable cost compared with other baseline methods. Please note that we do not report FedRep, FedBABU, and FedHKD on the retinal dataset because they are specifically designed for classification tasks.

GPU memory cost (MB)

For the computational time, we counted the training time of all methods. The GPU type used for training is GeForce RTX 2080 Ti. From the results in the table, it can be observed that our method lies in a reasonable range with a slightly higher cost than FedAvg, which is a trade-off between computational efficiency and performance.

Training time (hour)

We sincerely thank you for your supportive comments on our approach novelty, paper clarity and organization, and informative results. Our detailed responses to your comments are as follows.

• The authors could discuss the scalability and robustness of their method to larger and more complex datasets, and explore potential applications in other domains beyond natural and medical images.

We address your comments from three aspects: (a) Clarify our existing large and complex datasets, (b) Extend the client scalability study (incorporating more clients) on a large complex medical dataset, and (c) Explore applications in the new domain of brain-computer interface data.

(a)Clarify our existing large and complex datasets: We included two real-world complex medical datasets in our experiments. The Camelyon17 dataset is a very large dataset that contains 450,000 histology image patches with different stains from five hospitals. The retinal dataset consists of retinal fundus images acquired from 6 different institutions with various imaging devices and protocols. The superior performance of our proposed method has demonstrated the robustness on real-world large and complex datasets.

(b)Extend the client scalability study (incorporating more clients) on a large complex medical dataset: We have further investigated the client scalability (by incorporating more clients during training) of our method on the medical cancer diagnosis dataset (i.e., Camelyon17). Below, we present the average test performance by incorporating more clients into training. Our method shows a stable increasing trend compared with other methods.

(c) Explore applications in the new domain of brain-computer interface data: Following your suggestion, we further explored the application on the brain-computer interface (BCI) data, which consists in classifying the mental imagery EEG (electroencephalography) datasets. The EEG data is a type of data collected from electrodes placed on the scalp to measure and record electrical activity in the brain. This data can provide valuable insights into brain function and can be used in various fields such as neuroscience, psychology, and medicine. Specifically, we use four datasets from the MOABB BCI benchmark [1], each of those EEG datasets have been acquired on different subjects and have different numbers of channels, resulting in heterogeneous features. We compared our method with all other methods in our manuscript and reported five evaluation metrics, including Accuracy, AUC, Sensitivity, Specificity, and F1-score. The results are shown in the following table. Our method achieves the highest performance on all metrics, with significant improvements of 2.73% accuracy by the second-best method. More details can be found in Section D.2 in the Appendix.

BCI dataset

[1] Jayaram, Vinay, and Alexandre Barachant. "MOABB: trustworthy algorithm benchmarking for BCIs." Journal of neural engineering 15, no. 6 (2018): 066011.