Understanding the Stability-based Generalization of Personalized Federated Learning

We greatly appreciate your suggestions and comments, which will make our work better than this version. Thanks for your positive comments, and we address the weaknesses and minor comments below. All the suggestions have been modified in the revision paper in purple from you and in blue for all reviewers.

W1: We add more experiments on CIFAR-100 with VGG11 in the appendix D in the revision. The additional results can further strongly support our theoretical findings.

• Empirical setup: We conduct experiments on CIFAR-100 in the Pathological distribution (Non-IID c = 20) with VGG-11 for both C-PFL and D-PFL. To verify the impacts of the key hyperparameters, we study the parameter distance when disturbing only one data, the generalization gap of the difference between training and testing error, and testing performance when personalized training. We explore the impact of the four factors: 1) Local Learning Epochs, 2) Local Learning Rates, 3) Client Fraction / Communication topology, 4) Total Client Number. The detailed setup can be seen in Appendix D.1.
• Empirical analysis: From the empirical results, we draw the similar conclusions as CIFAR-10: 1)Less local learning epochs and lower learning rates lead to better generalization performance, but they affect the convergence speed more seriously; 2) More client participation and denser network connection enlarge the generalization gap, but they speed up the convergence to the same extent; 3) A larger total participation clients and a smaller number of local training samples increase the generalization error and reduce the convergence speed simultaneously; 4) C-PFL outperforms D-PFL in both generalization performance and convergence performance when their upper communication bandwidths are at the same level. Compared to the easier dataset CIFAR-10, the generalization error on CIFAR-100 is larger.

W2: We have polished it in the revision. $U= sup_{u,v_i,z}f(u,v_i;z)$ is the upper bound of the personalzied loss function.

Thanks for your comments and your suggestions are very meaningful to us. All the suggestions have been modified in the revision paper in purple for you and in blue for all reviewers. lf you have any further questions, please do not hesitate to contact us. Thanks again!

Dear reviewer,

We hope this message finds you well. We greatly appreciate your time and effort in reviewing our paper. We have carefully answered all the questions you raised about our paper and look forward to further discussions with you.

We greatly appreciate your suggestions and comments, which will make our work better than this version. Thanks for your positive comments, and we address the weaknesses and minor comments below. All the suggestions have been modified in the revision paper in violet for you and in blue for all reviewers.

Weaknesses:

W1: Thank you for your kind advice and we have corrected the typos in the revision.

• $u^{t+1}=\frac{1}{n} \sum u_i^{t+1}$ in Algorithm 2 should be $u^{t+1}=\frac{1}{n} \sum_{i\in S^t} u_i^{t+1}$. $\frac{1}{n} \sum_{i \in \mathcal{N}} u_{i, K_u}^t$. Line 1002 should be $\frac{1}{n} \sum_{i \in S^t } u_{i}^{t+1} = \frac{1}{n} \sum_{i \in S^t } u_{i, K_u}^t$ in the revision.
• $\mathcal{N}$ represents the client set in D-PFL defined in Line 207 in the revision. We check all the use of $\mathcal{N}$ in the revision.
• $I$ is the index of the first time to sample the perturbation sample $z_{i^*,j^*}$. When $t_0K+k_0 < I$, $\Delta_{k_0}^{t_0} = 0$. We add this explanation in Line 931 in the revision.

W2:

• Thank you for your kind advice and we added the perturbations experiments on CIFAR10/100 in Figure 1(a) and Figure 3(a) in the revision. We measure the model parameter difference between the perturbations datasets as [1-2] to further establish the nature of generalization stability. The empirical results also strongly support our theoretical insights that: 1) Fewer local learning epochs and lower learning rates lead to better generalization performance; 2) More client participation and denser network connection enlarge the generalization gap; 3) Larger total participation clients and a smaller number of local training samples increase the generalization error; 4) C-PFL outperforms D-PFL in generalization performance when their upper communication bandwidths are at the same level.
• We present the gap between training loss and test loss followed [2-5] rather than the gap between training accuracy and test accuracy in the last version. These results are reserved in Figure 1(b) and Figure 3(b) in the revision. In realistic experiments, we want to improve the accuracy in the test dataset with a well-trained model in the training dataset with the same data distribution. Therefore, we also reserve the curves of test accuracy in Figure 1(c) and Figure 3(c), which is the final aim in PFL.

[1] Miaoxi Zhu, Li Shen, Bo Du, and Dacheng Tao. Stability and generalization of the decentralized stochastic gradient descent ascent algorithm. Advances in Neural Information Processing Systems, 36, 2024.

[2] Moritz Hardt, Ben Recht, and Yoram Singer. Train faster, generalize better: Stability of stochastic gradient descent. In International conference on machine learning, pp. 1225–1234. PMLR, 2016.

[3] Zhenyu Sun, Xiaochun Niu, and Ermin Wei. Understanding generalization of federated learning via stability: Heterogeneity matters. In International Conference on Artificial Intelligence and Statistics, pp. 676–684. PMLR, 2024b.

[4] Tongtian Zhu, Fengxiang He, Lan Zhang, Zhengyang Niu, Mingli Song, and Dacheng Tao. Topology aware generalization of decentralized sgd. In International Conference on Machine Learning, ICML, pp. 27479–27503. PMLR, 2022.

[5] Tao Sun, Dongsheng Li, and Bao Wang. Stability and generalization of decentralized stochastic gradient descent. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp.9756–9764, 2021.

Questions:

Q1: We define $U=\sup_{w_i,z}f(w_i;z)$ so $\|f(w_i; z)-f(\tilde{w}_i ; z)\| \leq \sup_{w_i,z}f(w_i;z)$ in Line 286 and 927 in the revision.

Q2: Thank you and we have polished it in the revision.

Q3: $I$ in Line 812 (last version) is the index of the first time to sample the perturbation sample $z_{i^*,j^*}$. When $t_0K+k_0 < I$, $\Delta_{k_0}^{t_0} = 0$. And $P\left(\xi^c\right)=P\left(\Delta_{k_0}^{t_0}>0\right) \leq P\left(I \leq t_0 K+k_0\right)$. We add this explanation in Line 931 in the revision.

Q4: We polish the comparisons with generalization stability of central and decentralized learning in Table 1 in the revision. And the comparisons with generalization bound with PFL methods are removed to Appendix B.

Thanks for your comments and your suggestions are very meaningful to us. All the suggestions have been modified in the revision paper in violet for you and in blue for all reviewers.. lf you have any further questions, please do not hesitate to contact us. Thanks again!

Dear reviewer,

We hope this message finds you well. We greatly appreciate your time and effort in reviewing our paper. We have carefully answered all the questions you raised about our paper and look forward to further discussions with you.

Thanks for your comments and your suggestions are very meaningful to us. We reply to them below and all the suggestions have been modified in the revision paper in cyan for you and in blue for all reviewers. We sincerely hope that you could reconsider the contributions of our work.

Technical issues:

W1:

• Thank you for your kind comments. In the non-convex conditions, we can only obtain the upper bound of excess risk without the lower bound, so it is unsuitable to call it "optimal" number of communication rounds. We polish our analysis about the recommended stop point $T$ and check all the "optimal" in the revision.
• In our analysis, the excess risk is based on a specific algorithm and setting, so we assume that the problem-specific parameters of smoothness, stochastic gradient variance are fixed to discuss the training hyper-parameters like $T$ and $K$.

W2: Our technical contributions are as follows:

(1) We analyze the effect of the multiple local iterations and partial partition, which bridges centralized SGD and the partial partition FL. The technical difficulty is that local updates fail to be an unbiased gradient estimation after multiple local iterates, thereby merging the multiple local iterations is non-trivial.

(2) We analyze the effect of the personalized training and federated aggregation, which is the first generalization analysis to bridges FL and PFL. The technical difficulty is that the alignment errors exist in local training of the personalized variables after combining the two parts, thereby understanding the relationship between FL and PFL is non-trivial.

(3) We analyze the effect of communication topologies in Decentralized PFL, which is the first generalization analysis to bridges C-PFL and D-PFL. Lack of the central server, we need more complicated analysis of the aggregation error introduced by peer-to-peer communication with different gossip matrix $W$. Therefore, understanding the relationship between C-PFL and D-PFL is non-trivial.

(4) We analyze the effect of model decoupling and alternative training in both FL & DFL & PFL, which is the first generalization analysis to bridges full model personalization and parital model personalization. We analyze the generalization performance of different model parts separately and obtain the alternative errors between them, which provides interesting and useful suggestions for theoretical analysis of model decoupling and alternative training.

W3:

• We add more comparisons between the personalization method and the no personalization method in Table 1 in the revision. Compared to the reference, we are the first stability analysis in both C-PFL and D-PFL with the analysis of multi local updates and the whole training hyperparameters. The contributions in our work can be summarized in the four-holds: (1) extend centralized SGD to the partial partition FL; (2) extend FL to PFL; (3)extend C-PFL to D-PFL; (4)extend full model personalization and parital model personalization.

• Intuitively, whether the personalization generalizes better than no personalization is uncertain，which is related to the data heterogeneity in [1]. Personalization methods generalizes better than no personalization with more heterogeneity, while no personalization is bettern with more homogeneity. But [1] only proposed the analysis for FL and pure local training rather than PFL and only conducted the analysis under strongly convex conditions. Therefore, it is an interesting problem to compare the PFL methods and the FL methods with the analysis of data heterogeneity in the future.

[1] Shuxiao Chen, Qinqing Zheng, Qi Long, and Weijie J Su. A theorem of the alternative for personalized federated learning. arXiv preprint arXiv:2103.01901, 2021.

Presentation issues:

W1:

• "They decay per iteration $\tau = tK+k$" means "The learning rates decay per iteration $\tau = tK+k$". And $\tau_0$ is a specific observation point in the equations below.
• We present eq. (7) and (9) to extablish the changing progress of the stability, which proves that the generalization gaps in eq. (8) and (10) are available and describe the detailed effect of the hyperparameters.
• $\kappa_\lambda=\left(\frac{\alpha}{e}\right)^\alpha\frac{1}{\lambda\left(\ln\frac{1}{\lambda}\right)^\alpha} + \frac{2^\alpha}{\left(1-\alpha\right)e\lambda\ln\frac{1}{\lambda}} + \frac{2^\alpha}{\lambda\ln\frac{1}{\lambda}}$ is a widely used coefficient to measure different communication connections. When $\lambda\rightarrow 1$, it is bounded by $\mathcal{O}(\frac{1}{\lambda ln \lambda})$ and when $\lambda\rightarrow 0$, it is bounded by $\mathcal{O}(\frac{1}{\lambda\left(\ln\frac{1}{\lambda}\right)^\alpha})$. This detailed information can be found in Remark 5 in Line 383. We also add more explanation after Eq.(10) in the revision.
• We refresh the comparisons with the generalization bounds with PFL and add more details of the comparisons methods in Table 2 in Appendix B in the revision.

W2: Thank you and we have polished it in the revision.

• $j$ is the $j$ sample in the local training dataset of client $i$. And the $z_j$ in eq.(4) should be $\xi_j$.
• We have added the bracket in eq. (4) and changed $F_i$ to $F$ in Assumption 2.

W3: Thank you and we have polished it in the revision.

Thanks for your comments and your suggestions are very meaningful to us. All the suggestions above have been modified in cyan and in blue for all reviewers in the revision paper. Please reconsider our contributions to the FL community. If you have any further questions, please do not hesitate to contact us. Thanks again!

Dear reviewer,

We hope this message finds you well. We greatly appreciate your time and effort in reviewing our paper. We have carefully answered all the questions you raised about our paper and look forward to further discussions with you.

A3： We show the innovative results from the perspective of personalized and decentralized analysis.

Innovative Results about Personalized Training: The most basic difference between FL and PFL of the generalization bound comes from the testing dataset. That is to say, the global aggregation does not bring extra generalization errors when the training dataset keeps the same during local training. So the core difference between FL and PFL of generalization bound is the distribution of test dataset. Following the framework of this paper, we can creatively analyze the impact of data heterogeneity on PFL and FL. Here we make a preliminary attempt to analyze the effects of data heterogeneity.

• New assumption: We define the bounded global variance assumption: $\frac{1}{m} \sum_{i=1}^m\left\|| \nabla_u F_i\left(u, v_i\right)-\nabla_u F\left(u, V\right)\right\| |^2 \leq \delta_u^2$, where that bound $\delta$ could be the indicator of the data heterogeneity. The bigger the $\delta$, the more heterogeneous the distribution is. This assumption is commonly used in the convergence analysis for both FL [1-2] and PFL [3-5].
• Analysis technology: We use this assumption to due with $\nabla_u F_i\left(u_i,v_i,z\right)$ and get

$\quad\quad\mathbb{E}||u_{i,k+1}^t$-$\tilde{u}_{i,k+1}^t||$

$\quad\quad\leq(1+\eta_uL_u)\mathbb{E}||u_{i,k}^t-\tilde{u}_{i,k}^t||$

$\quad\quad\quad+\eta_uL_{uv}\mathbb{E}||v_{i,K_v}^t-\tilde{v}_{i,K_v}^t||+2\eta_u(\sigma_u+\delta_u)$

when selecting the different sample in Lemma 7 in Line 1117. Advancing analysis of heterogeneity, we finally deduce a rough generalization bound $\frac{nU\tau_0}{mS}+\left(\frac{TK_u}{\tau_0}\right)^{\mu_u L_u}\frac{2G(\sigma_u+\delta_u)}{mSL_u}$+$\left(\frac{TK_v}{\tau_0}\right)^{\mu_v L_v}$ \left( 1 + \frac{L_{uv}}{L_u} (\frac{TK_u}{\tau_0})^{\mu_u L_u} \right)$$ \frac{2G\sigma_v}{SL_v} for C-PFL (eq.(1)). And $\frac{U\tau_0}{S}+\frac{2G(\sigma_u+\delta_u)}{SL_u}\left(\frac{1+6\sqrt{m}\kappa_\lambda}{m}\right)\left(\frac{T K_u}{\tau_0}\right)^{\mu_uL_u}+\frac{2 \sigma_v G}{S L_v}\left(1+\frac{6\sqrt{m} \kappa_\lambda}{m}\left(\frac{L_{u v}}{L_v}\right)\right)\left(\frac{TK_v}{\tau_0}\right)^{\mu_vL_v}$ for D-PFL (eq.(2)) .

• Remark1: From that bound we can see how data distribution affects the generalization bound of personalization and no personalization. Assuming that $v_i$ attends to the aggregation, the generalization error for $v_i$ will add the heterogeneity impact and eq. (1) changes into $\frac{nU\tau_0}{mS} + \left(\frac{TK_u}{\tau_0}\right)^{\mu_uL_u}\frac{2G(\sigma_u+\delta_u)}{mSL_u} +\left(\frac{TK_v}{\tau_0}\right)^{\mu_vL_v} \left(1+\frac{L_{uv}}{L_u} (\frac{TK_u}{\tau_0})^{\mu_u L_u} \right) \frac{2G(\sigma_v+\delta_v)}{SL_v}$. It can clearly be seen that personalization without the system heterogeneity impact on $v_i$ performs better than no personalization.
• （Delete）Remark 2: As an implication, our analysis with data distribution could obtain a non-iid bound to show that the presumably difficult (infinite dimensional) problem can be reduced to a simple binary decision problem of choosing between classical FedAvg and PureLocalTraining as in [6]. For full participation FedAvg, the generalization bound can be reduced to $\mathcal{O}\left(\frac{U\tau_0}{S} + \left(\frac{TK}{\tau_0}\right)^{\mu L}\frac{2G(\sigma+\delta)}{mSL}\right)$ removing the personalized parameters $v_i$. For PureLocalTraining, the generalization bound can be reduced to $\mathcal{O}\left(\frac{U\tau_0}{S} + \left(\frac{TK}{\tau_0}\right)^{\mu L}\frac{2G\sigma}{SL}\right)$ removing the shared parameters $u$. From the comparisons, when $\delta \rightarrow 0$ with iid distribution, FedAvg generalizes better than PureLocalTraining with m-client data augmentation. With the heterogeneity getting larger, the generalization bounds of these two algorithms will meet at $\delta = (m-1)\sigma$. So it can be an indicator for the binary decision problem.

Innovative Results about Decentralized PFL: Compared to the results of Hardt et al, the analysis of D-PFL demonstrates the effect of communication topologies for generalization error. We introduce the topological spectral $(1-\lambda)$ and $\kappa_\lambda$ to show how communication topologies affect the generalization analysis.

• We conclude that 1) C-PFL generalized better than D-PFL;2)Denser communication topology leads to better generalization performance, which means that the fully connected topology achieves the best generalization performance of shared variables and is compatible with the central ones.
• Compared to the existing decentralized analysis, we obtain the tightest boundaries without the assumptions of bounded gradient [7] or the strong assumption that the weight difference obeys Gaussian distribution [8]. And [9] contains an extra constant term concerned with $\lambda$, which means the upper bound suffers from a nonvanishing influence of the communication network.

Thank you for your kind suggestions. We have answered your questions in the comment below.

A1: Thank you for your constructive suggestions! We have not discussed the tightness of the generalization bound in our paper due to the fact that we are the first work to analyze the stability and generalization of PFL and we can not make comparisons. And We have to admit that we do not analyze the lower bound of our generalization gap, which is still a valuable and open problem for generalization analysis. We are thinking about this problem and will continue on this as future work inspired by [1]. Moreover, we will delete the discussions of $T^*$ in our work.

A2： Thank you for your constructive suggestions! We illustrate the technical contributions from the perspective of personalized and decentralized analysis as below.

New techniques about Personalized FL:

• Compared to the analysis of FL, the analysis of Partial Model Personalization is more difficult to construct the alignment error between the shared aggregation variables and the local update. Therefore, we decompose the generalization errors into aggregation errors from shared variables and fine-tuning errors from both shared and personalized variables, then we build the bridge between $\Delta_v$ and $\Delta_u$ according to the gradient estimation process of the personalized training. It is the first work to analyze the generalization impact from personalized variables to shared variables, which uncovers the interaction mechanism between these two updating processes and provides valuable guidance for alternating personalized optimization.
• Also, it will be an interesting and open problem to propose the unique stability for PFL. We are thinking about this problem and will continue on this as future work.

New techniques about Decentralized PFL:

• We achieve the tight bound without the strong assumption of bounded gradient compared with D-SGD [2]. We transform the bounded gradient assumption into a functional representation related to T, which constructed a more flexible matrix operation in Lemma 5 and Lemma 10 to establish the relationship between the communication topology and the aggregation error. This technique is non-trivial.

[1] Zhang, Yikai, Wenjia Zhang, Sammy Bald, Vamsi Pingali, Chao Chen, and Mayank Goswami. Stability of sgd: Tightness analysis and improved bounds. In Uncertainty in artificial intelligence, pp. 2364-2373. PMLR, 2022.

[2] Tao Sun, Dongsheng Li, and Bao Wang. Stability and generalization of decentralized stochastic gradient descent. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 9756–9764, 2021.

W3:

• Thank you for your kind advise and it is our aim to bridge the generalization bound and data heterogeneity analysis for PFL in the future work. We also point out this meaningful idea in the CONCLUSION in Line 539. With the analysis of data heterogeneity, people can get more ideas to choose the suitable personalized strategies under different data distribution, which will be another significant work for PFL & FL community.

• That analysis on data heterogeneity in the reference work is interesting and inspires us a lot. Carefully considered the heterogeneity analysis, we think the key to consider the data distribution is how to measure the local gradient $g_i(\theta)$ ($\nabla_u F_i\left(u_i, v_i,z \right)$ in our paper), which can be dueled with $\| \nabla R_i(\theta) - \nabla R(\theta) \|$ contained the heterogeneity information $D_i$ in Lemma 1 in (Sun, Niu, Wei, 2024). The similar idea can be extended to our analysis in Lemma 7 and Lemma 9. Thank you and we add the discussion about it in Table 1 and Related Work in our revision.

• Another intuition to consider the data distribution in generalization analysis is to bridge the local gradient $\nabla_u F_i\left(u_i, v_i,z \right)$ and global variance $\frac{1}{m} \sum_{i=1}^m\left\|\nabla_u F_i\left(u_i, v_i\right)-\nabla_u F\left(u_i, V\right)\right\|^2 \leq \delta^2$, which is a commenly used assumption in the convergence analysis in PFL [1-3]. The bigger $\delta$ means the further distance between the global distribution and any local distribution.

• The data heterogeneity analysis exists from the convergence performance in our excess risk analysis. Referred to the heterogeneity $\delta$ in Eq.(7) in [1] for C_PFL and Eq.(3) in [2] for D_PFL, the data heterogeneity exactly effects our excess risk.

[1] Krishna Pillutla, Kshitiz Malik, Abdel-Rahman Mohamed, Mike Rabbat, Maziar Sanjabi, and Lin Xiao. Federated learning with partial model personalization. In International Conference on Machine Learning, pp. 17716–17758. PMLR, 2022.

[2] Yifan Shi, Yingqi Liu, Yan Sun, Zihao Lin, Li Shen, Xueqian Wang, and Dacheng Tao. Towards more suitable personalization in federated learning via decentralized partial model training. arXiv preprint arXiv:2305.15157, 2023b.

[3] Chen, Y., Cao, L., Yuan, K., & Wen, Z. (2023). Sharper convergence guarantees for federated learning with partial model personalization. arXiv preprint arXiv:2309.17409.

Questions:

Q1: See W1.

Q2: See W2.

All the suggestions above have been modified in orange in the revision paper for you and in blue for all reviewers. Please reconsider our contributions to the FL community. lf you have any further questions, please do not hesitate to contact us. Thanks again!

Thanks for your comments and your suggestions are very meaningful to us. We reply to them below and all the suggestions have been modified in the revision paper in orange for you and in blue for all reviewers. We sincerely hope that you could reconsider the contributions of our work.

Weakness:

W1: Thank you for your kind comment. Our bound in Theorem 1 can further recover the analysis of SGD in（Hardt et al., 2016) and we polish the discussion in Remark 2 in the revision.

• In the last version, we performed only one step of regression analysis. In the physical sense, the first step of the degradation process is to degrade from partial personalization to partial participation in local training, but the result is still for the entire partially participating system, rather than the pure local update for each node. Therefore, our stability for PFL reduce to $\mathcal{O}\left( (nK_vT)^{\frac{\mu_v L_v}{1+\mu_v L_v}}/(m S)\right)$ by removing the $K_u$ and $\sigma_u$ in eq. (7) .
• In the revision, we further degrade the stability on the basis of the first step. When $K_v = 1$ and all participation ($\frac{n}{m} = 1$), $\mathcal{O}\left( (nK_vT)^{\frac{\mu_v L_v}{1+\mu_v L_v}}/(m S)\right)$ can reduce to $\mathcal{O}\left( T^{\frac{\mu_u L_u}{1+\mu_v L_v}}/S\right)$, which is align with the analysis in（Hardt et al., 2016). In the physical sense, the second step of degradation is to assume global participation, and the stability analysis degrades to the pure local training for each node, which is the same as that of SGD in（Hardt et al., 2016).

W2: Our technical contributions are as follows:

(1) We analyze the effect of the multiple local iterations and partial partition, which bridges centralized SGD and the partial partition FL. The technical difficulty is that local updates fail to be an unbiased gradient estimation after multiple local iterates, thereby merging the multiple local iterations is non-trivial.

(2) We analyze the effect of the personalized training and federated aggregation, which is the first generalization analysis to bridges FL and PFL. The technical difficulty is that the alignment errors exist in local training of the personalized variables after combining the two parts, thereby understanding the relationship between FL and PFL is non-trivial.

(3) We analyze the effect of communication topologies in Decentralized PFL, which is the first generalization analysis to bridge C-PFL and D-PFL. Lack of the central server, we need more complicated analysis of the aggregation error introduced by peer-to-peer communication with different gossip matrix $W$. Therefore, understanding the relationship between C-PFL and D-PFL is non-trivial.

(4) We analyze the effect of model decoupling and alternative training in both FL & DFL & PFL, which is the first generalization analysis to bridge full model personalization and partial model personalization. We analyze the generalization performance of different model parts separately and obtain the alternative errors between them, which provides interesting and useful suggestions for theoretical analysis of model decoupling and alternative training.

Dear reviewer,

We hope this message finds you well. We greatly appreciate your time and effort in reviewing our paper. We have carefully answered all the questions you raised about our paper and look forward to further discussions with you.