Heterogeneous federal learning method based on attention guidance aggregation and prototype enhancement

Abstract

The invention discloses a heterogeneous federated learning method based on attention guidance aggregation and prototype enhancement, and relates to a heterogeneous federated learning method. The invention aims to solve the problem that the existing heterogeneous federal learning method cannot improve the granularity of client feature representation learning and cannot solve class imbalance in global head training. The method comprises the following steps: step 1, constructing a feature aggregation mechanism based on attention guidance; 2, designing a contrast learning objective function, and explicitly enhancing inter-class separability and intra-class compactness of feature representation in a local training process; 3, designing a self-adaptive prototype enhancement strategy, and relieving a class imbalance problem; and 4, constructing a collaborative optimization framework of the global model and the local model. The invention belongs to the technical field of distributed collaborative learning.

Description

Technical Field

[0001] The present invention relates to a heterogeneous federated learning method, belonging to the technical field of distributed collaborative learning. Background Art

[0002] With the rapid adoption and development of smart devices, edge computing has become a key computing paradigm for processing large-scale distributed data. Traditional centralized machine learning approaches require aggregating all data onto a central server, which not only raises serious privacy concerns but also requires extensive communication resources. Federated learning, an emerging distributed machine learning framework, effectively addresses these challenges by enabling multiple participants to collaboratively train models while protecting data privacy. In a typical federated learning framework, participants train models on local data and upload only the model parameters to the server. The server aggregates these parameters to form a global model, which is then distributed to participants for the next round of training. However, real-world federated learning faces two key challenges: system heterogeneity and data heterogeneity. System heterogeneity stems from differences in client devices' computing power, memory capacity, and battery life, making it difficult to uniformly deploy large models across all devices. Data heterogeneity manifests as non-independent and identically distributed (IID) data, meaning that data distributions may differ significantly between different clients. Traditional federated learning approaches, such as FedAvg, require all clients to train the same model architecture, and therefore perform poorly in the face of this dual heterogeneity.

[0003] To address the challenges of system and data heterogeneity, heterogeneous federated learning has emerged, allowing different clients to deploy models with different structures. Existing heterogeneous federated learning methods can be mainly divided into two categories: knowledge distillation methods that require public datasets and methods that do not rely on public data. The former, such as FedMD and FedDF, require high-quality public datasets, which limits their practical application. The latter, such as FedRep, LG-FedAvg, and FedGen, avoid the reliance on public data but typically separate the model into a feature extractor and a classification head, sharing only one component. This reduces communication efficiency and makes it difficult to capture global knowledge.

[0004] In particular, the FedGH approach achieved progress by sharing a global prediction head, but its simple class prototype averaging aggregation mechanism ignored intra-class variations, resulting in suboptimal prototype learning. Furthermore, FedProto significantly reduced communication costs and alleviated heterogeneity by uploading class prototypes rather than model parameters, but still used a simple averaging mechanism to generate prototypes, ignoring differences in sample representation quality. FedPCL introduced prototype contrastive learning to improve feature representation quality in federated learning, but failed to effectively address the class imbalance problem.

[0005] Existing heterogeneous federated learning methods generally overlook two key issues: how to improve the granularity of client-side feature representation learning and how to address class imbalance in global head training. In heterogeneous data environments, some classes may only have prototype representations from a small number of clients, making it difficult for existing methods to effectively address this imbalance. These issues severely limit the effectiveness of heterogeneous federated learning in highly non-IID environments, necessitating a novel federated learning framework that can simultaneously address these issues. Summary of the Invention

[0006] In order to solve the problems that existing heterogeneous federated learning methods cannot improve the granularity of client feature representation learning and cannot solve the problem of class imbalance in global head training, the present invention proposes a heterogeneous federated learning method based on attention-guided aggregation and prototype enhancement.

[0007] The technical solution adopted by the present invention to solve the above problems is: the steps of the present invention include: Step 1: Construct a feature aggregation mechanism based on attention guidance; Step 2: Design a contrastive learning objective function to explicitly enhance the inter-class separability and intra-class compactness of feature representation during local training. Step 3: Design an adaptive prototype enhancement strategy to alleviate the class imbalance problem; Step 4: Build a collaborative optimization framework for the global model and the local model.

[0008] Furthermore, step 1 specifically includes: Step 101: The client trains a heterogeneous model on local data and obtains a category feature representation for a given image through a feature extractor; Step 102: Collect all sample representations of category s to form a representation set ,in is the number of samples in this category; Step 103: Perform feature aggregation preparation; Step 104: Establish a prototype center guidance mechanism; Step 105: Implement dynamic weight fusion and stabilization; Step 106: Generate the final class prototype representation.

[0009] Furthermore, step 103 specifically includes: Perform L2 normalization on all feature representations: ; Calculate pairwise similarity (attention score) between samples: ; Calculate the self-attention importance score for each sample: ; Step 104 specifically includes: Computing Prototype Center: ; Perform L2 normalization to obtain the center vector: ; Calculate the similarity weight between each sample and the center: ; Step 105 specifically includes: Design an adaptive adjustment strategy based on the standard deviation of the self-attention weight distribution: ; Integrate dual weight information: ; Applying exponential moving average update mechanism to enhance time stability: ; Step 106 specifically includes: Generate initial class prototypes through weighted aggregation: ; Add regularization terms to enhance feature discrimination: ; Perform normalization to obtain the final class prototype: .

[0010] Furthermore, step 2 specifically includes: Step 201: Construct a contrastive learning framework; Step 202: Design a supervised contrast loss function: ; Step 203: Dynamically adjust the balance between contrastive learning and the main task.

[0011] Furthermore, step 201 specifically includes: For batch sample feature representation , perform L2 normalization: ; Tag-based Define the positive sample set: ; Define the negative sample set: ; Step 203 specifically includes: Use linear attenuation strategy to dynamically adjust contrastive learning weights: ; Define the client-side total loss function: .

[0012] Furthermore, step 3 specifically includes: Step 301: Analyze the server-side class distribution; Step 302: Establish a dual-standard rare class identification mechanism; Step 303: Implementing a differentiated prototype generation strategy; Step 304: Control the number and combination of prototype generation.

[0013] Furthermore, step 301 includes: Count the number of prototypes of each category received by the server: ; Calculate the average number of prototypes: ; Step 302 includes: Determine the adaptive threshold: ,in is the balance factor; Preliminary identification of rare class sets based on thresholds: ; If the number of rare classes is insufficient, a ranking-based selection strategy is used to select the class with the least number of prototypes. Class; Step 303 includes: For rare classes with at least 2 prototypes ( ), apply linear interpolation in feature space: , From the interval Random sampling; For rare classes with only one prototype ( ), using the controlled noise perturbation method: ,in is the standard normal distribution noise vector; Step 304 includes: Set a spawn limit for each rarity class: ; The generated synthetic prototypes are combined with the original prototypes to form an enhanced prototype set: .

[0014] Furthermore, step 4 specifically includes: Step 401: The server performs global model training. Step 402: The client interacts with the server. Step 403: Iterate and optimize until convergence.

[0015] Furthermore, in step 401, the enhanced prototype set is used to train the global prediction head: ,in is the global prediction head, is the cross entropy loss function; The interaction process between the client and the server in step 402 is as follows: The server will update the global prediction head parameters Distribute to each client; The client combines the local feature extractor with the global prediction head to form a complete local model: ; The client trains the model on local data: ; The client extracts the class prototype and uploads it to the server; Repeat steps 1 to 4 until the model reaches the preset convergence condition; The optimization goals of heterogeneous federated learning are: .

[0016] The beneficial effects of the present invention are: 1. The FedAPE framework proposed in this invention effectively solves the core challenges in heterogeneous federated learning through a mechanism that combines attention-guided feature aggregation and contrastive learning. The feature aggregation method based on the attention mechanism dynamically evaluates the intra-class similarity between samples and can generate more representative class prototype representations, effectively solving the problem that the traditional prototype averaging method ignores intra-class changes. At the same time, the introduced contrastive learning objective significantly improves the inter-class separability and intra-class compactness of the feature representation, enabling the model to learn more discriminative feature representations. These technological innovations achieved significant accuracy improvements of 2.82%, 5.65% and 2.94% on the CIFAR-10, CIFAR-100 and GTSRB datasets, respectively, achieving significant advantages over the existing state-of-the-art methods.

[0017] 2. The innovative adaptive prototype enhancement strategy of this invention can automatically identify rare categories and synthesize high-quality pseudo-prototypes in the feature space, effectively alleviating the class distribution imbalance problem in non-independent and identically distributed data environments, enabling the model to maintain stable performance in highly heterogeneous environments. In addition, the FedAPE framework supports client-side deployment of models with different structures, adapting to devices with various computing power and resource constraints while maintaining efficient knowledge exchange. This greatly improves the applicability of federated learning systems in real-world heterogeneous environments, enabling them to cope with complex and diverse real-world application scenarios.

[0018] 3. By transmitting prototypes rather than complete model parameters, the present invention significantly reduces the communication burden, while effectively protecting client data privacy and improving model training efficiency. Experimental results show that the FedAPE framework exhibits strong generalization capabilities and robustness in highly non-identically distributed data environments, and can better cope with the challenges of data heterogeneity in practical environments. In summary, the present invention successfully solves the key technical difficulties in heterogeneous federated learning through an innovative federated learning framework design, provides a more efficient, accurate and adaptable solution for collaborative learning in practical application environments, and promotes the application and development of heterogeneous federated learning technology in practical scenarios. BRIEF DESCRIPTION OF THE DRAWINGS

[0019] Figure 1 Schematic diagram comparing the FedAPE method of the present invention with a conventional polymerization method; Figure 2 1 is an overview of the workflow of the FedAPE method of the present invention; Figure 3 It is a visualization diagram of heterogeneity under different data distribution settings; Figure 4 This is a graph comparing the convergence of the FedAPE method of the present invention and the existing federated learning algorithm under heterogeneous data settings; Figure 5 This is a performance comparison chart of the FedAPE method of the present invention under different feature dimensions; Figure 6 3 is a performance comparison chart of the FedAPE method of the present invention under different degrees of model heterogeneity. DETAILED DESCRIPTION

[0020] Specific implementation method 1: Figures 1 to 6 As shown in the figure, a heterogeneous federated learning method based on attention-guided aggregation and prototype enhancement is proposed. The specific steps include: Step 1: Construct a feature aggregation mechanism based on attention guidance; specifically including: Step 101: The client trains a heterogeneous model on local data and obtains a category feature representation for a given image through a feature extractor; Step 102: Collect all sample representations of category s to form a representation set ,in is the number of samples in this category; Step 103: Perform feature aggregation preparation; Perform L2 normalization on all feature representations: ; Calculate pairwise similarity (attention score) between samples: ; Calculate the self-attention importance score for each sample: ; Step 104: Establish a prototype center guidance mechanism; Computing Prototype Center: ; Perform L2 normalization to obtain the center vector: ; Calculate the similarity weight between each sample and the center: ; Step 105: Implement dynamic weight fusion and stabilization; Design an adaptive adjustment strategy based on the standard deviation of the self-attention weight distribution: ; Integrate dual weight information: ; Applying exponential moving average update mechanism to enhance time stability: ; Step 106: Generate the final class prototype representation; Generate initial class prototypes through weighted aggregation: ; Add regularization terms to enhance feature discrimination: ; Perform normalization to obtain the final class prototype: ; Step 2: Design a contrastive learning objective function to explicitly enhance the inter-class separability and intra-class compactness of feature representations during local training. Specifically, Step 201: Construct a contrastive learning framework; For batch sample feature representation , perform L2 normalization: ; Tag-based Define the positive sample set: ; Define the negative sample set: ; Step 202: Design a supervised contrast loss function: ; Step 203: Dynamically adjust the balance between contrastive learning and the main task; Use linear attenuation strategy to dynamically adjust contrastive learning weights: ; Define the client-side total loss function: ; Step 3: Design an adaptive prototype enhancement strategy to alleviate the class imbalance problem; specifically, Step 301: Analyze the server-side class distribution; Count the number of prototypes of each category received by the server: ; Calculate the average number of prototypes: ; Step 302: Establish a dual-standard rare class identification mechanism; Determine the adaptive threshold: ,in is the balance factor; Preliminary identification of rare class sets based on thresholds: ; If the number of rare classes is insufficient, a ranking-based selection strategy is used to select the class with the least number of prototypes. Class; Step 303: Implementing a differentiated prototype generation strategy; For rare classes with at least 2 prototypes ( ), apply linear interpolation in feature space: , From the interval Random sampling; For rare classes with only one prototype ( ), using the controlled noise perturbation method: ,in is the standard normal distribution noise vector; Step 304: Control the number and combination of prototype generation; Set a spawn limit for each rarity class: ; The generated synthetic prototypes are combined with the original prototypes to form an enhanced prototype set: ; Step 4: Build a collaborative optimization framework for the global model and the local model; specifically, Step 401: The server performs global model training; the global prediction head is trained using the enhanced prototype set: ,in is the global prediction head, is the cross entropy loss function; Step 402: The client interacts with the server. The client-server interaction process is as follows: The server will update the global prediction head parameters Distribute to each client; The client combines the local feature extractor with the global prediction head to form a complete local model: ; The client trains the model on local data: ; The client extracts the class prototype and uploads it to the server; Repeat steps 1 to 4 until the model reaches the preset convergence condition; The optimization goals of heterogeneous federated learning are: ; Step 403: Iterate and optimize until convergence.

[0021] The model of the present invention is mainly compared with the following models: 1) LG-FedAvg: This method allows for different underlying architectures by sharing the same classification head, reducing communication overhead but limiting performance in non-IID environments. While LG-FedAvg supports model heterogeneity, its simple feature aggregation approach fails to fully leverage complementary knowledge between clients, in stark contrast to the attention-guided aggregation mechanism of our invention.

[0022] 2) FML: The Federated Mutual Learning model achieves knowledge transfer through mutual learning between models, but lacks an effective mechanism for evaluating feature quality, resulting in unstable performance in highly heterogeneous environments. In contrast, our FedAPE can more effectively screen and aggregate high-quality feature representations by accurately evaluating sample importance.

[0023] 3) FedGen: As a data-free knowledge distillation method that does not require a public dataset, FedGen can handle heterogeneous environments, but it only transmits a subset of model parameters, making it difficult to fully capture global knowledge. Our unique prototype enhancement strategy addresses FedGen's limitations in addressing class imbalance.

[0024] 4) FedKD: This method achieves heterogeneous federated learning through knowledge distillation, but relies on high-quality communication channels and a large amount of parameter transmission. In contrast, the present invention significantly reduces communication costs through prototype transmission while maintaining higher model performance.

[0025] 5) FedProto: This method reduces communication costs by uploading class prototypes instead of model parameters, but its simple averaging mechanism ignores the differences in sample representation quality. Our FedAPE significantly improves prototype quality and model generalization by introducing an attention mechanism and contrastive learning.

[0026] 6) FedTGP: This method combines trainable global prototypes and enhanced contrastive learning, but does not fully address the class imbalance problem. The adaptive prototype enhancement strategy of this invention specifically addresses this problem, achieving a more balanced and comprehensive class representation.

[0027] 7) FedGH: This approach has made progress by sharing a global prediction head, but its simple prototype averaging mechanism cannot fully distinguish the quality differences between different client features. Our FedAPE dynamically assigns weights using a dual similarity metric, which can more accurately capture intra-class variations and generate more representative prototypes.

[0028] Compared with the aforementioned methods, the FedAPE framework of this invention demonstrates significant advantages in handling model and data heterogeneity through its innovative combination of attention-guided feature aggregation, feature representation enhanced by contrastive learning, and an adaptive prototype enhancement mechanism. Experimental results demonstrate that FedAPE outperforms existing methods across various benchmarks, demonstrating its superior performance and broad applicability in heterogeneous federated learning environments.

[0029] FedAPE, our heterogeneous federated learning framework, was thoroughly evaluated and validated on three widely used image classification benchmark datasets: CIFAR-10, CIFAR-100, and the German Traffic Sign Recognition Benchmark (GTSRB). These datasets, with varying complexity and characteristics, fully validated our approach's performance in various scenarios.

[0030] CIFAR-10 Dataset: CIFAR-10 is a dataset of 60,000 32×32 color images, evenly divided into 10 categories, with 6,000 images per category. The dataset is divided into 50,000 training images and 10,000 test images. CIFAR-10's categories include common objects such as airplanes, cars, birds, cats, deer, dogs, frogs, horses, ships, and trucks, resulting in a balanced distribution of categories, making it ideal for validating a model's basic classification capabilities. In our federated learning environment, we use CIFAR-10 to evaluate the model's performance on simpler classification tasks.

[0031] CIFAR-100 dataset: CIFAR-100 also contains 60,000 32×32 color images, but is divided into 100 categories, each containing 600 images. This dataset has the same total number of images as CIFAR-10, but the number of categories has increased tenfold, with a corresponding decrease in the number of samples per category, making CIFAR-100 a more challenging classification task. These categories are organized into 20 superclasses, each containing five related fine-grained categories. In this paper, CIFAR-100 is used to evaluate the model's ability to handle more complex multi-class classification problems, particularly in heterogeneous federated learning settings.

[0032] GTSRB (German Traffic Sign Recognition Benchmark): GTSRB is a dataset containing over 50,000 real-world traffic sign images, covering 43 different categories of German traffic signs. These images exhibit significant variation in size, lighting conditions, and viewing angles, presenting models with a challenge closer to real-world application scenarios. The dataset is characterized by an uneven distribution of classes, with some traffic signs having significantly more samples than others, consistent with the data distribution characteristics found in the real world. In this paper, GTSRB is used to evaluate the robustness and accuracy of models in handling the uneven data distribution found in real-world scenarios.

[0033] To simulate the data heterogeneity in real-world federated learning environments, this paper adopts two non-independent and identically distributed (non-IID) data partitioning methods: Pathologically non-IID: In this setting, we assign 2 categories for CIFAR-10, 10 categories for CIFAR-100, and 4 categories for GTSRB to each client, so that each client only has data of a specific category. This extreme data distribution simulates the situation in reality where clients may focus on a specific task.

[0034] Practical Non-IID: Applying the Dirichlet distribution (α=0.1) creates an unbalanced data partition that is closer to reality. In this setting, each client may have data for all categories, but the category distribution is severely skewed, with some categories having significantly more samples than others. This distribution is more closely aligned with the data distribution characteristics in real-world application scenarios.

[0035] For each dataset, we used 75% of the samples as the training set and the remaining 25% as the test set to evaluate the model's generalization ability. This comprehensive experimental setup ensures that the invention can be fully validated on datasets of various complexity and characteristics, providing strong support for the superior performance of the FedAPE framework in heterogeneous federated learning environments.

[0036] We design four heterogeneous model configurations to evaluate the performance of the framework under conditions of different degrees of model heterogeneity: HtFE8 basic configuration: Contains a variety of convolutional neural networks, including FedAvgCNN, GoogleNet, MobileNet_v2, and multiple ResNet variants (ResNet18, ResNet34, ResNet50, ResNet101, and ResNet152).

[0037] HtFE3 configuration: A low-heterogeneity environment containing similar lightweight models.

[0038] HtFE4 configuration: A moderately heterogeneous environment containing a mix of models with some architectural differences.

[0039] HtFE9 configuration: A highly heterogeneous environment containing a combination of models with significantly different structures.

[0040] This paper mainly uses classification accuracy as the core evaluation indicator to calculate the proportion of samples correctly classified by the model on the test dataset: ,This indicator intuitively reflects the overall performance of the model, and facilitates horizontal ,comparison between different methods.

[0041] The following key parameter settings were used in the experiment: a federated learning environment with 20 clients was constructed, and the client participation rate was set to 1.0. The global communication rounds were set to 200 to ensure sufficient model convergence. In each communication round, the client trained locally for one epoch, with a batch size of 10 and a learning rate of 0.01. The initial weight of the contrastive loss was 0. The probability of the error is set to 0.3 and decays linearly with the number of training rounds. To reduce experimental randomness, we conducted five independent experiments for each setting and used the average performance as the final test accuracy. All experiments were conducted on a server equipped with an NVIDIA GTX 3090 GPU (24GB of video memory) and implemented using the PyTorch 1.8.0 framework.

[0042] This paper comprehensively compares the proposed method with multiple existing heterogeneous federated learning methods, including LG-FedAvg, FML, FedGen, FedKD, FedProto, FedTGP, and FedGH. Experimental results demonstrate that the FedAPE framework outperforms these existing methods across various datasets and heterogeneous settings, fully validating the effectiveness and advancement of the proposed method.

[0043] Table 1 shows the performance comparison of our proposed FedAPE with existing heterogeneous federated learning methods on the CIFAR-10, CIFAR-100, and GTSRB datasets. To comprehensively evaluate the effectiveness of our proposed method, experiments were conducted under two typical non-IID (non-IID) data settings: pathological non-IID and practical non-IID.

[0044] Under the pathological non-IID setting, the FedAPE of the present invention achieved the best performance on all three datasets. Specifically, FedAPE achieved a classification accuracy of 85.06% on the CIFAR-10 dataset, an improvement of 0.77 percentage points over the baseline method FedGH (84.29%). On the CIFAR-100 dataset, the present invention achieved an accuracy of 58.23%, a significant improvement of 2.53 percentage points over the suboptimal method FedGH (55.70%). This improvement fully demonstrates the excellent ability of the present invention in handling multi-category classification tasks. On the GTSRB dataset, FedAPE achieved an accuracy of 99.19%, an improvement of 1.60 percentage points over FedGH (97.59%).

[0045] These results demonstrate that the attention-guided feature aggregation mechanism proposed in this paper can effectively identify intra-class variation patterns and generate more representative class prototypes, thereby improving the model's classification performance under strictly restricted data distribution environments. In complex multi-class scenarios such as CIFAR-100, the advantages of FedAPE are particularly significant, demonstrating the superiority of this paper in handling inter-class differences in high-dimensional feature spaces.

[0046] In the practical non-independent and identically distributed setting (Dir=0.1) that is closer to actual application scenarios, the performance advantage of the FedAPE of the present invention is more obvious. On the CIFAR-10 dataset, FedAPE achieved an accuracy of 86.41%, an increase of 0.33 percentage points over the second-best method FedKD (86.08%) and a significant increase of 2.82 percentage points over FedGH (83.59%). On the CIFAR-100 dataset, FedAPE achieved an accuracy of 42.14%, an increase of 0.79 percentage points over the second-best method FedTGP (41.35%) and a significant increase of 5.65 percentage points over FedGH (36.49%). On the GTSRB dataset, the present invention achieved an accuracy of 92.86%, an increase of 2.94 percentage points over the second-best method FedGH (89.92%).

[0047] These results demonstrate that FedAPE's performance advantage is even more pronounced under more complex data distribution conditions. Particularly noteworthy is the significant 5.65% improvement on CIFAR-100, demonstrating that the adaptive prototype augmentation strategy proposed in this paper can effectively alleviate class imbalance and improve the model's ability to identify marginal classes. FedAPE demonstrates strong adaptability in handling data heterogeneity, a common occurrence in real-world scenarios.

[0048] Table 1: Performance comparison under different data distribution settings (%)

[0049] To systematically verify the effectiveness of each component of the FedAPE and its contribution to overall performance, this section designed a series of ablation experiments. Table 2 details the incremental performance improvement of each core component under different data distribution settings.

[0050] This paper uses FedGH as a baseline model and gradually adds three core innovative components: a feature aggregation (FA) mechanism, a contrastive learning (CL) strategy, and a prototype enhancement (PE) technique. Experiments are conducted on three datasets, CIFAR-10, CIFAR-100, and GTSRB, examining performance changes under both practical and pathological non-IID settings. This incremental addition of components allows for clear quantification of each component's contribution to the final performance.

[0051] Contribution of the Feature Aggregation (FA) component: After adding the attention-guided feature aggregation mechanism to the baseline model, under the practical non-IID setting, the classification accuracy of CIFAR-10 increased from 83.59% to 84.94%, an increase of 1.35%; the accuracy of CIFAR-100 increased from 36.49% to 40.37%, a significant increase of 3.88%; and the accuracy of GTSRB increased from 89.92% to 90.66%, an increase of 0.74%. These results demonstrate that the proposed attention-guided feature aggregation mechanism can effectively capture the class similarity patterns within samples and generate more representative class prototypes, thereby significantly improving the model's classification performance. In particular, this component exhibits a more significant performance improvement on the more complex CIFAR-100 dataset.

[0052] Contribution of the contrastive learning (CL) component: After further integrating the contrastive learning strategy based on feature aggregation, the accuracy of CIFAR-10 reached 85.26%, an improvement of 1.67% over the baseline model; the accuracy of CIFAR-100 reached 41.49%, an improvement of 5.00% over the baseline model; and the accuracy of GTSRB reached 92.30%, an improvement of 2.38% over the baseline model. These results demonstrate that the contrastive learning strategy designed in this paper further improves the discriminative ability of the model by enhancing the inter-class separability and intra-class compactness of the feature representation. Notably, in the pathological non-IID setting, contrastive learning significantly improves the performance of CIFAR-100, reaching 3.50%, indicating that contrastive learning can play a greater role in cases of extremely unbalanced data distribution.

[0053] Contribution of the Prototype Enhancement (PE) component: The complete FedAPE framework achieves an accuracy of 86.41% on CIFAR-10 under the practical non-IID setting, an improvement of 2.82% over the baseline model; an accuracy of 42.14% on CIFAR-100, an improvement of 5.65% over the baseline model; and an accuracy of 92.86% on GTSRB, an improvement of 2.94% over the baseline model. Under the pathological non-IID setting, the performance of the three datasets is improved by 1.15%, 4.65%, and 1.64%, respectively. These results fully verify that the adaptive prototype enhancement strategy proposed in this paper can effectively alleviate the problem of class imbalance. By synthesizing representative pseudo-prototypes in underrepresented regions of the feature space, it significantly improves the global model's recognition ability for each category.

[0054] Table 2: Contribution analysis of each component of FedAPE (%)

[0055] To further evaluate the impact of different feature aggregation strategies on model performance, we conducted a series of comparative experiments targeting feature aggregation mechanisms. Table 3 details the performance comparison results of four different feature aggregation methods under a Dir = 0.1 distribution setting, clearly demonstrating the superiority of our proposed dual-weighted comprehensive aggregation strategy.

[0056] The traditional average aggregation strategy, the most basic aggregation method, achieved only 84.07% and 38.19% accuracy on the CIFAR-10 and CIFAR-100 datasets, respectively. This approach assigns equal weight to all sample representations and fails to distinguish between the quality differences between different sample representations. As a result, the generated prototypes may be affected by noise samples or fail to fully reflect the core characteristics of the class. In a heterogeneous federated learning environment, due to significant differences in model structure and data distribution between clients, this simple averaging method struggles to generate class prototypes with strong discriminative capabilities, thus limiting the model's classification performance.

[0057] When using only the self-attention mechanism, model performance improved to 85.74% on CIFAR-10 and 40.16% on CIFAR-100, representing improvements of 1.67% and 1.97% respectively compared to average aggregation. By calculating similarities between samples, the self-attention mechanism automatically identifies highly representative samples and assigns them higher weights, effectively mitigating the impact of noisy samples.

[0058] When using only prototype center similarity weights, the model achieves 85.92% accuracy on CIFAR-10, slightly higher than the method using self-attention alone; however, the accuracy on CIFAR-100 is 39.85%, lower than the performance of self-attention alone. The prototype center similarity strategy assigns weights by calculating the similarity between samples and the center of the class prototype. This effectively maintains the consistency of sample representation with the overall class distribution, but may not fully capture the complex relationships between samples in high-dimensional feature spaces.

[0059] Experimental results clearly demonstrate that FedAPE's dual-weighted comprehensive aggregation approach achieves significant performance improvements over single aggregation strategies on datasets of varying complexity. Specifically, on the complex CIFAR-100 dataset, FedAPE achieves a 1.98% improvement in accuracy over the best single strategy (self-attention) and a 3.95% improvement over average aggregation. This significant difference demonstrates that the proposed dynamic weighted integration strategy is able to more effectively identify high-quality sample representations while balancing local similarity and global consistency, thereby generating more discriminative class prototypes in non-IID environments.

[0060] These experimental results fully demonstrate the innovation and effectiveness of this invention in designing a feature aggregation mechanism, providing important technical support for high-quality feature representation learning in heterogeneous federated learning. FedAPE's dual-weighted aggregation strategy not only improves classification accuracy but also enhances the model's robustness and generalization capabilities in heterogeneous environments.

[0061] Table 3: Performance comparison of different feature aggregation strategies under Dir=0.1 distribution (%)

[0062] To determine the best application of contrastive learning in heterogeneous federated learning, this paper systematically studies the contrastive learning weight parameters ( ) on the performance of the model. Table 4 shows the different The results of the impact of the value on FedAPE performance.

[0063] This experiment focuses on adjusting the weight coefficient of the contrastive learning loss function while keeping other hyperparameters unchanged. , with a value range of 0.1 to 0.5. Through this series of experiments, we aim to determine the optimal balance between contrastive learning and the main classification task, thereby optimizing the discriminative ability of feature representation and improving the overall classification performance.

[0064] Experimental results show that contrastive learning weights have a significant impact on model performance. When the weight is set to 0.1, the model achieves an accuracy of 85.37% on CIFAR-10; as the weight increases to 0.2, the accuracy increases slightly to 85.50%; when When the weight reaches 0.3, the model performance improves significantly, with the accuracy reaching a maximum of 86.41%; however, when the weight is further increased to 0.4 and 0.5, the performance shows a downward trend, with the accuracy dropping to 85.72% and 85.55%, respectively.

[0065] This performance change forms a clear "inverted U-shaped" curve, indicating that there is an optimal value for the contrastive learning weight, which is neither too large nor too small. This phenomenon can be explained by the balance between contrastive learning and the main classification task.

[0066] when If the value is too small (e.g., 0.1), the regularization effect of contrastive learning is weak, making it difficult to effectively enhance the discriminative power of features. In this case, the model is primarily guided by the cross-entropy loss and cannot fully utilize the advantages of structured feature representation provided by contrastive learning. Therefore, it has limited contribution to improving the model's generalization ability in non-IID environments.

[0067] when At a moderate value (0.3), the model achieves an optimal balance between the primary classification task and contrastive learning. Here, contrastive learning effectively promotes inter-class separability and intra-class compactness without excessively interfering with the optimization direction of the primary classification task. This balance enables the model to learn more discriminative feature representations, significantly improving classification performance.

[0068] when When the value is too large (e.g., 0.5), the contrastive learning loss accounts for too high a proportion of the total loss, which may cause the model to over-focus on the similarity relationship between features and ignore the precise location of classification boundaries. This imbalance will weaken the model's professional classification ability and lead to performance degradation.

[0069] Based on these experimental results, the present invention determined = 0.3 as the standard configuration of the FedAPE model and use this value in all subsequent experiments. This finding has important practical value and provides clear parameter guidance for applying contrastive learning techniques in heterogeneous federated learning.

[0070] Furthermore, these results reveal the underlying mechanism by which contrastive learning works synergistically with the primary classification task. Appropriately strong contrastive learning can serve as an effective regularization technique, improving the model's generalization capabilities by enhancing the structured organization of the feature space. However, the introduction of contrastive learning requires careful consideration to ensure that it complements rather than interferes with the primary classification objective.

[0071] This series of experiments not only helped optimize the performance of FedAPE, but also provided valuable experience for the design of multi-task learning in heterogeneous federated learning. Appropriate contrastive learning weight setting is one of the key factors in achieving high-performance heterogeneous federated learning systems.

[0072] Table 4: Comparative impact of learning weights on FedAPE performance (CIFAR-10, Dir=0.1, %)

[0073] To comprehensively evaluate the adaptability of our method under different data distribution conditions, we further compared the performance of FedAPE with existing methods in a low data heterogeneity environment (Dir = 0.01). Table 5 details the classification accuracy results of each method under this setting.

[0074] The low data heterogeneity setting (Dir = 0.01) represents a relatively mild non-IID environment, where the data distributions of each client vary but are relatively close. Compared to the high heterogeneity setting (Dir = 0.1), this setting presents fewer challenges for federated learning, while still reflecting the data distributions likely to be encountered in real-world applications. Evaluating under this setting allows us to examine the adaptability of various methods to varying degrees of heterogeneity, further validating the robustness and generalization capabilities of the algorithms.

[0075] On the CIFAR-10 dataset, our FedAPE achieved an accuracy of 97.37%, surpassing all compared methods. This represents a 0.15 percentage point improvement over the next-best method, FedTGP (97.22%), and a 0.53 percentage point improvement over the baseline method, FedGH (96.84%). While the performance improvement in this relatively simple setting is not as significant as in highly heterogeneous environments, FedAPE still maintains its leading position.

[0076] On the CIFAR-100 dataset, the advantages of our method are even more pronounced, achieving an accuracy of 70.38%, an improvement of 3.32 percentage points over the next-best method, FedTGP (67.06%), and 6.57 percentage points over FedGH (63.81%). This significant difference demonstrates that even in low-heterogeneity environments, FedAPE's attention-guided feature aggregation and prototype enhancement mechanisms still have outstanding advantages when handling complex multi-category tasks.

[0077] On the GTSRB dataset, the proposed method achieved an accuracy of 98.68%, an improvement of 1.16 percentage points over the next-best method, LG-FedAvg (97.52%), and 1.32 percentage points over FedGH (97.36%). This result further confirms FedAPE's outstanding performance in various traffic sign recognition tasks.

[0078] By comparing the performance of various methods in low-heterogeneity environments, it can be shown that the present invention is of great significance for practical applications. FedAPE can adapt to the data distribution that may occur in various practical scenarios and provide reliable technical support for the deployment of heterogeneous federated learning in different application fields.

[0079] Table 5: Performance comparison (%) under low data heterogeneity setting (Dir=0.01)

[0080] The above description is merely a preferred embodiment of the present invention and does not constitute any form of limitation to the present invention. Although the present invention has been disclosed as a preferred embodiment as above, it is not intended to limit the present invention. Any technician familiar with the present profession can make some changes or modifications to equivalent embodiments of equivalent changes using the technical content disclosed above without departing from the scope of the technical solution of the present invention. However, any simple modification, equivalent replacement and improvement of the above embodiments made according to the technical essence of the present invention, within the spirit and principles of the present invention, without departing from the content of the technical solution of the present invention, shall still fall within the scope of protection of the technical solution of the present invention.

Claims

1. A heterogeneous federated learning method based on attention-guided aggregation and prototype enhancement, characterized by: The specific steps include: Step 1: Construct a feature aggregation mechanism based on attention guidance; Step 2: Design a contrastive learning objective function to explicitly enhance the inter-class separability and intra-class compactness of feature representation during local training. Step 3: Design an adaptive prototype enhancement strategy to alleviate the class imbalance problem; Step 4: Build a collaborative optimization framework for the global model and the local model.

2. The heterogeneous federated learning method based on attention-guided aggregation and prototype enhancement according to claim 1 is characterized in that Step 1 specifically includes: Step 101: The client trains a heterogeneous model on local data and obtains a category feature representation for a given image through a feature extractor; Step 102: Collect all sample representations of category s to form a representation set ,in is the number of samples in this category; Step 103: Perform feature aggregation preparation; Step 104: Establish a prototype center guidance mechanism; Step 105: Implement dynamic weight fusion and stabilization; Step 106: Generate the final class prototype representation.

3. The heterogeneous federated learning method based on attention-guided aggregation and prototype enhancement according to claim 2 is characterized in that: Step 103 specifically includes: Perform L2 normalization on all feature representations: ; Calculate pairwise similarity (attention score) between samples: ; Calculate the self-attention importance score for each sample: ; Step 104 specifically includes: Computing Prototype Center: ; Perform L2 normalization to obtain the center vector: ; Calculate the similarity weight between each sample and the center: ; Step 105 specifically includes: Design an adaptive adjustment strategy based on the standard deviation of the self-attention weight distribution: ; Integrate dual weight information: ; Applying exponential moving average update mechanism to enhance time stability: ; Step 106 specifically includes: Generate initial class prototypes through weighted aggregation: ; Add regularization terms to enhance feature discrimination: ; Perform normalization to obtain the final class prototype: .

4. The heterogeneous federated learning method based on attention-guided aggregation and prototype enhancement according to claim 1 is characterized in that Step 2 specifically includes: Step 201: Construct a contrastive learning framework; Step 202: Design a supervised contrast loss function: ； Step 203: Dynamically adjust the balance between contrastive learning and the main task.

5. The heterogeneous federated learning method based on attention-guided aggregation and prototype enhancement according to claim 4 is characterized in that: Step 201 specifically includes: For batch sample feature representation , perform L2 normalization: ; Tag-based Define the positive sample set: ; Define the negative sample set: ; Step 203 specifically includes: Use linear attenuation strategy to dynamically adjust contrastive learning weights: ; Define the client-side total loss function: .

6. The heterogeneous federated learning method based on attention-guided aggregation and prototype enhancement according to claim 1 is characterized in that Step 3 specifically includes: Step 301: Analyze the server-side class distribution; Step 302: Establish a dual-standard rare class identification mechanism; Step 303: Implementing a differentiated prototype generation strategy; Step 304: Control the number and combination of prototype generation.

7. The heterogeneous federated learning method based on attention-guided aggregation and prototype enhancement according to claim 6 is characterized in that: Step 301 includes: Count the number of prototypes of each category received by the server: ; Calculate the average number of prototypes: ; Step 302 includes: Determine the adaptive threshold: ,in is the balance factor; Preliminary identification of rare class sets based on thresholds: ; If the number of rare classes is insufficient, a ranking-based selection strategy is used to select the class with the least number of prototypes. Class; Step 303 includes: For rare classes with at least 2 prototypes ( ), apply linear interpolation in feature space: , From the interval Random sampling; For rare classes with only one prototype ( ), using the controlled noise perturbation method: ,in is the standard normal distribution noise vector; Step 304 includes: Set a spawn limit for each rarity class: ; The generated synthetic prototypes are combined with the original prototypes to form an enhanced prototype set: 。 8. The heterogeneous federated learning method based on attention-guided aggregation and prototype enhancement according to claim 1 is characterized in that: Step 4 specifically includes: Step 401: The server performs global model training. Step 402: The client interacts with the server. Step 403: Iterate and optimize until convergence.

9. The heterogeneous federated learning method based on attention-guided aggregation and prototype enhancement according to claim 1 or 8, characterized in that In step 401, the global prediction head is trained using the enhanced prototype set: ,in is the global prediction head, is the cross entropy loss function; The interaction process between the client and the server in step 402 is as follows: The server will update the global prediction head parameters Distribute to each client; The client combines the local feature extractor with the global prediction head to form a complete local model: ； The client trains the model on local data: ; The client extracts the class prototype and uploads it to the server; Repeat steps 1 to 4 until the model reaches the preset convergence condition; The optimization goals of heterogeneous federated learning are: .

Images