FedCPD:Personalized Federated Learning with Prototype-Enhanced Representation and Memory Distillation

1 Introduction

Federated Learning (FL) [McMahan et al., 2017] is a decentralized machine learning strategy that improves data privacy by enabling local model training on various devices with updates communicated to a central server. However, training a single global model becomes inefficient when the data is heterogeneous or the devices handle distinct tasks, as traditional FL approaches such as FedAvg [McMahan et al., 2017] and FedProx [Li et al., 2020] inadequately address personalized demands. Client data heterogeneity poses a fairness issue, a major challenge for federated learning in diverse settings. Personalized Federated Learning (PFL), [Tan et al., 2022a] addresses these issues by allowing client-specific models that leverage global insights to enhance local outcomes. The main challenge in PFL lies in balancing global knowledge sharing with preserving client-specific information, making the tradeoff an important research area.

∗Corresponding author

In personalized federated learning, historical information forgetting proves challenging as it causes local test accuracy to drop post-global model update. This happens because global model optimization prioritizes overall performance, often overlooking the requirements of individual clients. Consequently, personalized local knowledge diminishes over time. Commonly, personalized federated learning (PFL) employs parameter decoupling to tackle this issue. Techniques like FedPer [Arivazhagan et al., 2019], FedRep [Collins et al., 2021], and FedGH [Yi et al., 2023] separate local private parameters from global ones, allowing clients to tailor task-specific representations using local data. While a shared backbone network (feature extractor) exists, the classifier (head) is customized for local tasks. Despite partially mitigating historical information forgetting, this strategy's efficacy in extracting local features declines when client data exhibits substantial heterogeneity. The core challenge is that task-specific knowledge preservation mechanisms are lacking, making local knowledge loss during global optimization almost unavoidable. Additionally, global optimization objectives might gradually neglect some client-specific tasks, curtailing local model test accuracy improvements.

Personalized federated learning contends with limited generalization performance alongside historical information forgetting. To boost generalization, earlier approaches used prototype methods in representation learning. Prototype learning involves shared class prototypes to curb overfitting and enhance generalization. FedProto [Tan et al., 2022b] enhances generalization by aligning prototypes through measuring distances between similar label representations. FedProc [Mu et al., 2023] and FedCRL [Huang et al., 2024] utilize contrastive loss to connect local features with global prototypes, minimizing representation variations. Yet, as most methods tweak FedAvg's basic averaging for data aggregation, they struggle with diverse data and fall short in representation learning. Thus, current methods inadequately merge prototype alignment with contrastive techniques, limiting their generalization efficacy on varied data.

Both challenges pertain to optimizing personalized federated learning, yet their solutions don't cross-apply. Parameter decoupling protects local knowledge to prevent forgetting but falls short on sharing global insights, thus struggling with generalization. On the other hand, prototype learning curbs overfitting and boosts generalization by sharing class prototypes, yet it misses retaining clients' historical task knowledge.

In this paper, we introduce a Federated Class Prototype and Feature Distillation (FedCPD) framework that tackles the challenges of historical information loss and restricted generalization in personalized federated learning. Key contributions are as follows.

Advanced Cross-Layer Fusion and Prototype Alignment for Optimal Historical Conservation and Improved Generalization. In preliminary experiments, we identified historical information forgetting and limitations in parameter decoupling. To address these, we developed techniques like cross-layer feature fusion, hierarchical feature distillation, and attention mechanisms to reduce information loss from global model aggregation and improve personalized model accuracy. We also explored global prototypes for better model generalization and proposed a prototype contrast and alignment strategy to optimize global-local prototype relationships, mitigate overfitting from local data scarcity, enhance cross-label knowledge sharing, and boost generalization in data-limited environments.

Comprehensive Theoretical Analysis Uncovers Key Mechanisms Enhancing Model Performance. We rigorously prove the upper bound of FedCPD algorithm convergence, highlighting two findings. (1) Feature distillation mitigates historical information forgetting from global model aggregation bias by retaining key local features across training, enhancing stability and speed. (2) Aligning global and local prototypes enhances generalization in data-scarce environments, reducing overfitting, aiding cross-label knowledge sharing, accelerating convergence, and improving adaptability to new data.

Empirical Validation Shows Outstanding Performance Across Various Real-World Datasets. We conducted extensive experiments on multiple real-world datasets with varying data heterogeneity. The results showed that, compared to traditional methods, FedCPD significantly improved generalization ability by up to 10.40% and maintained personalization accuracy with an improvement of up to 4.90%, validating the effectiveness and superiority of our approach.

2 RelatedWork

2.1 Personalized Federated Learning

In personalized federated learning, various approaches have been developed to address data heterogeneity between clients.

Parameter decoupling techniques improve model personalization by isolating specific parameters from local modules. For example, FedRep [Collins et al., 2021] alternates between training the global extractor and the local classifier following FedPer's [Arivazhagan et al., 2019] separation strategy. Prototype learning methods mitigate data heterogeneity by computing the average feature representation for each category, efficiently utilizing sparse sample information for classification tasks. These strategies hold significant promise in scenarios with limited learning samples. For example, FedTGP [Zhang et al., 2024a] uses an adaptive margin contrastive learning technique on the server to greatly enhance the representational capabilities of feature vectors. Meanwhile, FedKTL [Zhang et al., 2024b] creates prototype image vector pairs aligned with client tasks through a server-based pretrained generator, thus supporting the learning of the client model and improving performance. These approaches effectively handle heterogeneous data with shared representations, but still offer room for advances through better utilization of these representations. We merge parameter decoupling with prototype learning in the personalized federated learning framework, efficiently tackling the challenges of data heterogeneity between clients and enhancing the model's generalization.

2.2 Feature Distillation

Knowledge Distillation (KD) is a model compression technique that transfers knowledge from well-trained large models to simpler, smaller models, making them suitable for deployment on various devices. Proposed by Hinton et al. [Hinton, 2015], KD reduces model complexity while maintaining performance. FitNets [Adriana et al., 2015] uses intermediate teacher features to guide the student model, ensuring similar predictions. FEED [Park and Kwak, 2019] facilitates knowledge transfer through distillation of feature map-level features through non-linear transformations. Heo et al. [Heo et al., 2019] improve the synergy between teacher and student models with methods such as feature localization and distance functions. DKKR [Chen et al., 2021] uses multilevel feature distillation to guide the student layer by layer. Our method introduces feature distillation into federated learning, combined with attention mechanisms and feature fusion, to facilitate efficient knowledge transfer between teacher and student models.

3 Methodology

In this section, we dive into the details of the implementation of the proposed FedCPD, as illustrated in Figure 1 (global structure diagram) and Algorithm 1.

3.1 Problem Statement

In PFL, an architecture consists of N clients and one server, where each client i has non-IID private data Di. Clients collaboratively train their personalized models W1, ...,WN. Similarly to FedPer and FedRep, the backbone network is split into a feature extractor f and a classifier g, where f : RD → RK maps the input samples into a feature space and g : RK → RC transforms the feature vectors into the label space. For each client i, the local model Wi = [W f i ;W g i ] is trained on the data set Di. For each sample-label pair (xi, yi) ∈ Di, the model predicts ˆ yi = g(f(xi;W f i );W g i ) and the objective is to minimize the empirical risk across all clients as follows.

where Li(W) = E(xi,yi)∼Di L(xi, yi;W f i ,W g i ) represents the empirical loss for client i, given by the expected loss across all samples in Di. We adopt FedRep as our baseline algorithm. In the t-th round of communication, the central server sends the global extractor W f,t to the active client set Ct ⊆ C. Each client Ci ∈ Ct initializes its parameters W f,t i,0 and performs K optimization iterations using its local data, alternating updates between the local extractor and the local classifier in each iteration. The server collects the local extractors W f,t i and computes the global extractor for the next round W f,t by averaging the parameters, repeating the process until convergence.

3.2 Motivation

Global Model Performance Gap and Historical Informa-

tion Forgetting. In federated learning, clients receive and train on a global model each round. However, this global model, being an aggregate of multiple client models, may not perform well on clients with significantly different data distributions. This often results in a gap in personalized testing accuracy between the global model and previous local models of clients, as shown in Figure 2a, indicating a potential decrease in personalized performance after updating the global model, especially in cases of high data heterogeneity.

Methods such as FedRep [Collins et al., 2021] tackle the challenge of forgetting historical information [Jin et al., 2022] through techniques such as parameter decoupling.

However, they still face limitations in fully leveraging historical local knowledge, which can introduce new challenges. For example, while FedRep preserves the local personalized head with each update, the shared feature extractor may still be susceptible to influence from the global model. This can result in an inadequate representation of local data nuances,

Algorithm 1 FedCPD

Require: Total communication rounds T, {Di}Ni=1 , global feature extractor parameters W f,0 , learning rate η.

Ensure: {W f,∗

i }Ni=1 .

• 1: for iteration t = 1, . . . , T do
• 2: Server aggregates {W f,t i }Ni=1 and {P t i }Ni=1 to obtain W f,t and P t ;
• 3: Server sends W f,t and P t to all clients;
• 4: for the i-th client in parallel do
• 5: Fix W f,t i , update W g,t i ; ▷ Classifier Update
• 6: for layer l = 1, . . . , L do
• 7: Calculate Lt+1 fd,i = ‖Al,∗,t+1 i -Al,t i ‖22+
• 8: ‖F l,∗,t+1 i -F l,t i ‖22 ;
• 9: end for

exp(s(pt+1,+

c∈C\c exp(s(pt+1,-

i,j )/τ)+∑ˆ

i,j )/τ) ;

• 15: Fix W g,t i ,update W f,t i ;
• 16: Same as Classifier Update;
• 17: Train W f,t+1 i ←W f,t+1 - η∇W f,t+1Lt+1 i ;
• 18: Collect and uploadW f,t+1 i and P t+1 i ;
• 19: end for
• 20: end for
• 21: return {W f,∗ i }Ni=1 .

ultimately restricting further enhancements in personalized model performance. As illustrated in Figure 2b, the performance of the model after training is inferior to that of FedAvg. By introducing the hierarchical feature distillation module in FedRep, we verify that the upper limit of local model accuracy can be improved to be comparable with that of local model in FedAvg after making full use of historical knowledge.

Prototype Alignment and Contrastive Learning for Improved Generalization. In the context of federated learning, the data between clients is typically non-IID. However, the success of centralized deep learning indicates that data share global feature representations, with statistical heterogeneity reflected primarily in labels. Due to the limited amount of data available to each client, locally trained features are prone to overfitting, leading to insufficient generalization. Therefore, sharing common prototypes to leverage data from other clients is an effective strategy. Inspired by this, we introduced prototype alignment in our FedAvg and FedRep experiments to ensure global prototype consistency across classes, improving generalization. We also applied prototype contrastive learning to enable cross-label collaboration. In FedAvg, the prototype alignment significantly increases accuracy, outperforming both the combination of alignment and contrastive losses, and contrastive loss alone

(see Figure 2c). This suggests that prototype alignment enhances representation learning and reduces overfitting due to limited data.

In FedRep, combining prototype contrast and alignment is even more effective (see Figure 2d), improving class boundary understanding and classification accuracy. The difference arises from the methods used: FedAvg relies on prototype consistency and a simple averaging aggregation, where alignment reduces parameter diversity and stabilizes the global model. In contrast, FedRep focuses on representation learning and uses contrastive learning to capture subtle sample differences, allowing the combination of both methods to better leverage class structure, significantly boosting performance.

3.3 Attention-Guided Hierarchical Feature Distillation

To address the limitations of existing methods in overcoming historical information forgetting and its impact on personalized performance (as discussed in Section 3.2), we propose a feature distillation method based on FedRep to reduce the loss of personalized knowledge caused by updates to the feature extractor, thus improving the performance of personalized models. Specifically, we retain the local feature extractor from the previous round to fully leverage historical personalized knowledge, where W f,t i serves as the teacher model and W f,t+1 serves as the student model. Through hierarchical feature distillation, we transfer knowledge from previous models to the current model and guide the transfer of knowledge by integrating an attention mechanism and a feature fusion mechanism, as shown in Figure 3.

AttentionModule. To enhance the effect of feature distillation, we introduce the Convolutional Block Attention Module (CBAM) [Woo et al., 2018], aimed at implementing attention guidance. The goal of feature distillation is to improve the student model's performance through knowledge transfer from the teacher model. With attention mechanisms, we can precisely guide features, enhancing the effectiveness of the distillation process.

The attention map after the CBAM can be represented as:

where l denotes the index of layer, F l is the feature map of the l-th layer, Al is the attention map of the l-th layer.

Hierarchical Feature Distillation: In feature distillation, we distill feature maps and attention maps as the primary targets. We use the global backbone W f,t+1 of the current round as the student model and the final local backbone W f,t i of the previous round as the teacher model. Feature distillation typically relies on same-layer feature alignment techniques, which directly compare and align features from the same layers of the student and teacher models.

However, traditional intra-layer feature alignment methods fail to effectively leverage the advantages of shallow features in small object detection and fully integrate the semantic information captured by high-level features due to the lack of interlayer feature reuse. This limitation results in poor distillation performance, especially when dealing with visual tasks that involve complex semantic structures.

Thus, motivated by the cross-layer feature alignment methods [Chen et al., 2021], we constructed cross-layer connection pathways amongst student networks, depicted in Figure 3. Employing a cross-layer feature fusion approach allows us to effectively merge student features from various layers. This strategy not only boosts the target detection abilities of shallow features but also amplifies the semantic expression of higher-level features, thereby facilitating a more thorough transfer of the teacher model's knowledge.

Firstly, we design three types of convolutional layers Cv1,Cv2,Cv3 for each input feature layer. Cv1 is a transverse convolution that unifies the channel counts of input feature maps of different scales to a fixed output channel count, Cv2 is used for secondary smoothing after feature fusion, and Cv3 is used to enhance low-level features. The feature fusion steps can be represented as follows.

where upsamp represents upscaling using bilinear interpolation, and Al,∗ represents the final fusion result. We used the highest-layer feature map for the first transverse convolution, performed feature fusion from the bottom up, and before fusion, further processed each layer's original input feature map through an enhanced low-level convolution layer to enhance low-level features. Combining the high-level upsampled features, transverse convolution features, and enhanced low-level features, we then apply Cv2 to smooth. The final feature alignment process can be represented as follows.

where s denotes the student, t denotes the teacher, and we use MSE loss for feature alignment.

3.4 Prototype Contrast and Alignment

To effectively improve the generalizability of the model, as mentioned in Section 3.2, we adopted the method of prototype contrast and alignment. The core objective of this approach is to enhance the model's generalization ability while reducing the decline in personalized performance caused by data heterogeneity. Specifically, we dynamically adjust the relationship between global prototypes and local prototypes, ensuring the retention of personalized features while improving the model's generalization ability.

Prototype Alignment. Initially, we generate feature embeddings on the local datasets Di = (x, y)Ni of client i using the backbone network W f,t+1 distributed by the server as follows.

where Ni indicates the number of images at client i. Next, based on the global prototypes P t received from the server, we align the output embeddings zi according to

Here, zi,j is labeled c, j serves as an index for the local data, and P t c represents the global prototype for class c. The objective Lalign aims to minimize the squared Euclidean distance between features and class prototypes. The loss of prototype alignment acts as a harmonizing mechanism, reducing feature representation discrepancies between different models or domains, making feature representations more consistent and stable when learning class boundaries.

Prototype Contrast. With the help of prototype alignment loss, contrastive loss optimizes the relative relations in the feature space without conflicting with supervised loss. It enhances the model's learning of robust feature representations, intra-class compactness, and inter-class separability. Specifically, we learn representations through contrasting positive and negative pairs [Huang et al., 2024], which helps in personalization by leveraging positives and negatives between local and global representations.

For each feature representation zi,j with class label c, we consider the global prototype P t of class c as the positive sample, while prototypes from other classes are considered negative samples. Therefore, for each feature representation, one positive sample pair pt+1,+ i and |C| - 1 negative sample pairs pt+1,- i can be constructed.

where j serves as an index for the local data, and P t ˆ c represents the remaining prototypes that differ from the category of the output embedding zi,j.

To rectify the knowledge from local training in each client using global prototypes, we introduce Global Prototype Contrastive Loss Lpcl,i. This loss encourages each client's sample to approach its class's global prototype while distancing itself from the other class's prototypes. We define the Global Prototype Contrastive Loss as follows.

where τ is a temperature hyperparameter that adjusts the focus on positive and negative samples, and s(·) is the cosine similarity. This multi-constraint learning mechanism combines supervised loss, alignment loss, and contrastive loss, enhancing the model's ability to distinguish different categories, ensuring consistency in feature representations, and optimizing the structure of the feature space, thereby improving generalization and performance.

4 Convergence Bounds Analysis

In this subsection, we present the convergence analysis for FedCPD and outline the assumptions necessary to prove its convergence, following a framework similar to FedProto [Tan et al., 2022b], FedHKD [Chen et al., 2023] and FedCRL [Huang et al., 2024]. Detailed assumptions and proofs are provided in Appendix B.

Theorem 1. (One-Round Deviation). Let Assumptions 1 to 3 hold (see Appendix B). For an arbitrary client, between the iteration t and the iteration t + 1, we have

Theorem 1 shows that for any client, adjusting the appropriate hyperparameters can bound the deviation in expected loss for the i-th client from iteration t to t+1, thereby ensuring convergence.

Theorem 2. (Non-Convex FedCPD Convergence). λη0 < ηe < η0 , e ∈ { 1 2 , 1, 2, . . . , E}, where λ represents the decay factor for the learning rate. If the learning rate for each epoch satisfies the following condition, the loss function decreases monotonically, leading to convergence

Theorem 2 indicates that, under the specified learning rate conditions, the loss function for any client decreases consistently between successive communication rounds, guaranteeing algorithm convergence.

5 Experiments and Discussion

5.1 Experiment Setup

We evaluated FedCPD using three popular image classification datasets (CIFAR-10, FMNIST, and CIFAR-100). We compare it with isolated local training and eight popular federated learning methods, including these state-of-the-art algorithms FedProto [Tan et al., 2022b], FedGH [Yi et al., 2023], FedALA [Zhang et al., 2023], and FedPA [Jiang et al., 2024]. The experiments run for 200 communication rounds, each round involving 1 epoch of local training and a batch size of 10. All models are trained using the SGD optimizer with a learning rate of 0.01. The model consists of two convolutional layers and two fully connected layers. To simulate data heterogeneity, we use two settings: the "practical setting" based on the Dirichlet distribution and the "pathological setting" with varied class samples from the datasets.

5.2 Performance Comparisons

By default, local clients are configured with β = 0.1 and N = 20.

Test Accuracy. To evaluate FedCPD, we benchmarked it against leading federated learning methods on three datasets, each subjected to two highly heterogeneous partitions (see Table 1). Accuracy is reported for every client's own test set, so the metric directly reflects the personalization performance. Figure 4 trace the evolution of the test accuracy during training. Across both realistic and pathological splits, FedCPD consistently delivers the highest accuracy, outperforming all baselines, and confirming its superior effectiveness.

Results Analysis. FedCPD converges significantly faster than other algorithms, especially in the early stages (e.g., the first 50 rounds), where the improvement in test accuracy is most notable.Moreover, as the number of label categories increases in both heterogeneous settings, FedCPD shows more significant improvements over the second-best method: CIFAR-100 (4.90%/4.58%), CIFAR-10 (0.76%/0.10%), and FMNIST (0.04%/0.15%). This shows that FedCPD can capture category differences more effectively, leading to a significant improvement in classification accuracy.

Generalization. To evaluate the performance of FedCPD under different levels of heterogeneity, we adjusted the value of β of the Dirichlet distribution in the CIFAR-10 data set to control actual heterogeneity and modified the number of label categories held by each client in the CIFAR-100 data set to control pathological heterogeneity (see Table 2). The results show that FedCPD performed best in both heterogeneity settings.

Results Analysis. As shown in Table 2, in moderate and low heterogeneity environments, FedCPD still outperforms other personalized federated learning methods: CIFAR-10 (1.68%/2.11%), CIFAR-100 (7.55%/10.40%), which demonstrates its strong generalizability. Unlike other algorithms that rely on local data, FedCPD effectively integrates both global and local information, significantly improving model performance in scenarios with many label categories and uneven data distributions, thus improving its adaptability across different data distribution scenarios.

Fairness. We evaluated FedCPD's fairness using the standard deviations in Tables 1 and 2. In all experiments, FedCPD was ranked in the top three, showing minimal performance disparities between client groups and achieving high fairness.

Results Analysis. FedCPD achieves excellent fairness by effectively balancing global and local models, reducing performance disparities among clients. This improves the accuracy of classification and ensures consistency between different clients.

5.3 Discussion

Experimental findings and theoretical insights confirm that feature distillation efficiently reduces the forgetting of historical data by conveying knowledge throughout training iterations, leading to smoother optimization and faster earlystage convergence (see Figure 4). Retaining historical data significantly improves the global model's learning performance, raising classification accuracy by 4.66% under typical conditions (refer to Appendix A) and outperforming leading parameter decoupling techniques by approximately 4.90% across different label distributions (refer to Table 1). Moreover, feature distillation enhances client equity, promoting balanced model training and consistently securing a position among the top three for global fairness.

The global prototype enhances the connection between global and local models via prototype alignment and contrastive learning, leading to improved class separation, faster convergence and a 7.23% boost in classification accuracy in default settings (see Appendix A), especially with numerous labels and imbalanced data (refer to Table 1). These strategies help the model capture global insights and encourage crosslabel knowledge sharing. For generalization, the global prototype combines global knowledge, reducing overfitting due to limited data, and delivers up to a 10.40% boost over the best SOTA in varied environments (refer to Table 2). Furthermore, the global prototype improves fairness by diminishing performance variance among different clients, promoting a balanced training process, particularly with uneven data distributions.

6 Conclusion and FutureWork

To address the issue of historical forgetting and enhance generalization, we introduce FedCPD, a personalized federated learning framework. This framework integrates hierarchical feature distillation, cross-layer feature fusion, and attention mechanisms to effectively balance local and global models while maximizing the utility of historical knowledge. Additionally, shared prototypes are employed to address label skew and accelerate convergence, facilitating privacy-aware and equitable adaptation. Experimental results indicate that FedCPD surpasses leading models, achieving an improvement in generalization by up to 10.40% and personalization by up to 4.90%, thereby demonstrating its efficacy and superiority. Future research will investigate its capability to manage other forms of statistical heterogeneity and evaluate the algorithm's impact on representation learning.

A Ablation Study

To validate the effectiveness of the various components of FedCPD, we performed ablation experiments on the CIFAR100 data set with β = 0.1. The experiments aimed to verify the contributions of feature distillation, prototype alignment, and prototype contrast. Accuracy comparisons can be observed in Table 3, which demonstrates the effectiveness of the different components of FedCPD.

Results Analysis. In the absence of these three components, the algorithm is equivalent to FedRep, exhibiting the poorest performance. With the introduction of prototype contrast or prototype alignment, performance improved by approximately 5%, indicating that these strategies effectively improve the class separability of the model and improve classification accuracy. Furthermore, by incorporating local feature distillation, the model can learn more refined local feature representations, further improving classification accuracy. Ultimately, with the combination of all three components, FedCPD achieves optimal performance, conclusively confirming the effectiveness of these components.

B Convergence Theory Derivation

Additional Notation. The introduction of additional variables here is to better represent the local model update process. The embedding function corresponding to the i-th client is denoted by fi(W f i ) : RD → RK , gi(W g i ) : RK → RC is the decision function for all clients. Here, D, K, and C denote the sizes of the input, feature, and label spaces, respectively. Thus, the labeling function can be written as Fi(W f i ,W g i ) = gi(W g i ) ◦ fi(W f i ), and Wi can be used as (W f i ,W g i ). Therefore, the local loss function for client i is given by

where S1 = s(fi(W f,t+1

i

, xk), P t+1

S1 = s(fi(W f,t+1

c ), ˆ

i

, xk), P t+1

ˆ

c ).

Let (xk, yk) ∼ Di represent the local data points, where Di denotes the private dataset of client i, and Bi is the training batch size. The loss function L remains fixed within each communication round, but varies between rounds, adding complexity to the convergence analysis.

For notation, t indicates the communication round and e ∈ 1/2, 1, 2, . . . ,E refers to the local iterations, where E is the total number of local updates. Thus, tE + e represents the e-th local iteration in the (t + 1)-th round. Furthermore, tE is the time step just before aggregation on the server, while tE +1/2 marks the step between aggregation and the start of the first local update.

Assumption 1. We assume that each loss function L(W) satisfies L1-Lipschitz smoothness, indicating that its gradient is L1-Lipschitz continuous, and that the embedding function of the feature extractor f(·) is L2-Lipschitz continuous.

which implies the following quadratic bound,

Assumption 2. The stochastic gradient gt i = ∇L(W t i , ξ t i ) is an unbiased estimator of the local gradient for each client. We assume that its expectation satisfies

the variance is bounded by σ2 as follows.

Assumption 3. The expected value of Euclidean norm of the stochastic gradient is bounded by G,

Lemma 1. In the t-th communication round, the loss function for any client i, after conducting E local training rounds, is bounded as

Proof.

Then, by telescoping of E steps and setting the learning step at the beginning of local training to η1/2 = η0, we have,

Lemma 2. After completing E local training rounds in the t-th communication round and before starting local training in the (t + 1)-th communication round, the loss function for any client i remains bounded.

Proof.

Applying expectations to both sides of the inequalities

It should be noted that when the model parameters remain unchanged, the first two terms of the loss function are the same. Therefore, we are left only with the terms related to the prototype representations.

Taking expectation of both sides,

The proof of the other term is as follows.

where

where Dt+1,1/2 = exp(-1/τ )+∑ˆ c∈C\c exp(1/τ ), Dt+1 = exp(1/τ ) + ∑ˆ c∈C\c exp(-1/τ). Since log(C -1) is a constant, and when τ is very small, 2 τ is numch greater than log(C - 1),we can omit log(C - 1).

Summarizing the above derivations, we obtain an upper bound for the loss function.

Theorem 1. (One-round deviation). Assume that Assumptions 1 to 3 are satisfied. For any client, after each communication round, we obtain

Theorem 2. (Non-convex FedCPD convergence). λη0 < ηe < η0, e ∈ { 1 2 , 1, 2, . . . ,E}, where λ represents the decay factor for the learning rate. If the learning rate for each epoch satisfies the following condition, the loss function decreases monotonically, leading to convergence:

where S = ∑E-1 e=1/2 ‖∇Lt,e i ‖22 ,B = S - EG2 ,A = ( L1(Eσ2+E2G2+S) 2 + 2αL2E2G2 ), C = 2β τ .

Proof.

Since A > 0, the quadratic function opens upwards. The inequality Aη2 e + Bηe + C ≤ 0 holds between the roots of the quadratic equation. Therefore, the acceptable range for ηe is between the two roots. However, we are interested in the maximum allowable ηe that satisfies the inequality, so we focus on the larger root:

Acknowledgments

This work is supported by National Natural Science Foundation of China under grants 62171132, 62471139, and U1905211, and Natural Science Foundation of Fujian Province under grant 2024J09032.