Model performance optimization method based on heterogeneous federation

Abstract

The invention provides a model performance optimization method based on heterogeneous federation, which comprises the following steps: a server sends a global model, a global prototype and global knowledge to selected clients, and each client initializes a local model by using the global model; on each client, calculating a loss function by using the global prototype, the global knowledge and a local sample, performing gradient descent, executing a local model obtained after epoch rounds, and calculating a local prototype and local knowledge; after training is completed, the client uploads the local model, the local prototype and the local knowledge to the server. And the server aggregates the local models, the local prototypes and the local knowledge to form a global model, a global prototype and global knowledge, and judges whether a convergence condition or a preset training round number is reached, if yes, training is ended, and if not, the next round of training is executed. The method can effectively solve the problems of inter-class feature stacking, prototype deviation and intra-class feature stacking in a prototype learning technology.

Description

Technical Field

[0001] The present invention relates to the field of federated learning in machine learning, and in particular to a model performance optimization method based on heterogeneous federation. Background Art

[0002] With the rapid development of computer applications and Internet technology, massive amounts of data are generated every second, which has brought us into the era of big data. In this context, federated learning has emerged as an emerging machine learning method that aims to solve the problems of data privacy and data dispersion.

[0003] Traditional machine learning methods usually require all data to be concentrated in one place for training, which means that sensitive data in user devices or data centers needs to be transferred to a central server. However, this practice poses privacy risks and the risk of data leakage. In addition, due to various reasons (such as privacy, regulations, or network bandwidth limitations), data may be scattered across different devices or data centers and cannot be stored in one place. In this case, federated learning is proposed as a solution. Federated learning allows model training on distributed devices without transmitting the original data to a central server. In federated learning, the model is trained on the local device and the updated model parameters are sent to a central server for aggregation. This decentralized learning method helps protect the user's data privacy and reduces the need for data transmission. However, due to the heterogeneity of user data, federated learning faces some challenges. There is inhomogeneity between different devices and data, which may lead to the emergence of data bias problems. Therefore, federated learning needs to overcome these problems to ensure that the training results of the model have good generalization performance and are not affected by data heterogeneity.

[0004] Contrastive learning originally compares the similarities and differences between different samples to learn effective feature representations. Prototype contrastive learning is different from this. In prototype contrastive learning, not only are samples from the same category closer to their category prototypes than samples from other categories, but features are also required to separate prototypes of other categories as much as possible.

[0005] At present, how to use prototype learning to solve the heterogeneity of user data in federated learning has become a research hotspot, and some related research results have emerged: Yue Tan et al. (Fedproto: Federated prototype learning across heterogeneous clients. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pages 8432–8440, 2022..) was the first to introduce prototype learning into federated learning. However, since FedProto simply aligns the norm distance between features and corresponding class prototypes, it does not guide feature learning to learn the feature distance between classes, and does not limit the degree of feature alignment, which will make the features too similar to the class prototype and damage the diversity of sample features; Huancheng Chen et al. (The best of bothworlds: Accurate global and personalized models through federated learning with data-free hyper-knowledge distillation.) aggregated soft predictions based on FedProto and proposed the concept of hyper-knowledge, but only simply calculated the norm values ​​of predictions and aggregated soft predictions. Since it is a numerical calculation, such enhancement is limited.Xutong Mu et al. (Fedproc: Prototypical contrastive federated learning on non-iid data.) considered the issue of data privacy and proposed to use a lightweight generator to generate feature representations instead of generating raw data; Yue Tan et al. (Fedproto: FFederated prototype learning across heterogeneous clients. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pages 8432–8440) realized the importance of learning the distance between classes while learning the distance between intra-class features, and thus introduced contrastive learning into prototype learning. When aligning features, the local model not only aligns the prototypes of the corresponding category, but also stays away from the prototypes of other categories. However, there are still problems with prototypes that are too similar and biased; Zexi Li et al. (No fear of classifier biases: Neural collapse inspired federated learning with synthetic and fixed classifier. In Proceedings of the IEEE / CVF International Conference on Computer Vision, pages 5319–5329, 2023.) A fixed ETF classifier is used to replace the original classifier of each client to solve the problem of intra-class feature stacking.

[0006] Based on the above analysis of the current mainstream algorithms using prototype learning technology, the present invention summarizes the following three problems existing in prototype learning technology: First, since only the feature vector is close to its corresponding prototype, rather than guiding it away from other prototypes, this leads to the intersection of classification boundaries and thus the formation of inter-class feature stacking; Second, due to data heterogeneity, the prototype extracted by the offset model is biased, which will seriously affect the quality of obtaining the global prototype on the server; Third, since the sample features are directly aligned with the corresponding class prototype, and the degree of feature alignment of the model is not restricted, this will make the features of samples of the same category too similar, eliminating the diversity of sample features and causing intra-class feature stacking.

[0007] However, existing methods cannot effectively solve these three key problems at the same time; therefore, combining the characteristics of contrastive prototype learning and knowledge distillation, designing a method that can effectively solve inter-class feature stacking, prototype offset, and intra-class feature stacking has become a technical problem that needs to be urgently solved by technical personnel in this field. Summary of the invention

[0008] The purpose of the present invention is to propose a model performance optimization method based on heterogeneous federation, which can solve the problems existing in the above-mentioned prototype learning technology to a certain extent.

[0009] To achieve the above object, the technical solution of the present invention is: a model performance optimization method based on heterogeneous federation, comprising the following steps:

[0010] S1. The server sends the global model, global prototype and global knowledge to the selected clients; each client initializes the local model using the global model;

[0011] S2. On each client, use the global prototype, global knowledge and local samples to train the local model, calculate the loss function, and perform gradient descent; execute the local model obtained after epoch rounds, and calculate the local prototype and local knowledge;

[0012] S3. After the training is completed, the client uploads the local model, local prototype and local knowledge to the server; the server aggregates the local models, local prototypes and local knowledge of all selected clients to form a global model, global prototype and global knowledge; it determines whether the convergence condition or the predetermined number of training rounds is reached, if so, the training is terminated, otherwise it returns to S1 to loop and execute the next round of training.

[0013] Preferably, each client uses the global model to initialize the local model specifically as follows:

[0014] Each client initializes the local model using the global model parameters; the global model and the local model are divided into two components: a feature extractor and a classifier. The extractor is responsible for extracting features from the data, and the classifier classifies the output features from the feature extractor. The parameterization of the model is represented as ω = [φ, ψ], where φ is the extractor, ψ is the classifier, and the embedding functions of the extractor and the classifier are F φ and R ψ .

[0015] Preferably, client local model training is based on four indicators: cross entropy loss, prototype contrastive learning loss, penalty loss and global knowledge distillation loss.

[0016] Preferably, in step S2, the client uses the local model, the global prototype, the global knowledge and the local sample to calculate the loss function, which specifically includes the following steps:

[0017] S21. According to formula (1), the cross entropy loss function of client i is obtained:

[0018]

[0019] Where N is the number of samples in an epoch, M is the number of sample categories; y m is the label of sample x, which takes 1 when the category belongs to m, otherwise it takes 0; p m is the probability that sample x is predicted to belong to category m; is the sample data set of client i;

[0020] S22. Calculate the local prototype contrast training loss according to formula (2):

[0021]

[0022] Where i represents the client number, m represents the sample category, and t represents the federated learning round; x i,m is the sample of category m of client i, m' is used to traverse all sample categories except category m, and are the global prototypes of categories m and m', is the cosine similarity, is the input x i,m To the features extracted by the feature extraction network, τ is the scaling factor;

[0023] S23. Calculate the sample penalty loss of the model according to formula (3):

[0024]

[0025] in, is an indicator function used to express whether the sample belongs to the sample set that needs to be punished; is the set of samples of category m that need to be punished;

[0026] S24. Calculate the knowledge distillation loss according to formula (4):

[0027]

[0028] Where W is Sliced ​​Wasserstein Distance, is the global knowledge of category m; R ψ (·) represents the function operation for classifying features;

[0029] S25. The overall loss function of the model calculated based on the above four loss functions is as follows:

[0030]

[0031] Among them, γ and λ are adjustable hyperparameters, and the overall loss function is composed of the cross entropy loss Local prototype contrast training loss Sample Penalization Loss and knowledge distillation loss The model parameters ω are optimized by minimizing the loss function.

[0032] Preferably, in step S3, the server aggregates the local models, local prototypes and local knowledge of all selected clients to form a global model, a global prototype and global knowledge, which specifically includes the following steps:

[0033] S31. Calculate the global model according to formula (6):

[0034]

[0035] Among them, |D i | represents the dataset D of client i i Size, |S t | is the number of selected clients, A dataset representing the selected client The size of represents the local model of client i;

[0036] S32. Calculate the global prototype according to formula (7):

[0037]

[0038] in, represents the local prototype of client i;

[0039] S33. Calculate the global logit according to formula (8):

[0040]

[0041] in, represents the local knowledge distillation of client i.

[0042] Preferably, the step S2 of calculating the local prototype is specifically as follows:

[0043]

[0044] in, represents the local prototype of class m for client i, represents the dataset of category m for client i, Where F is the feature extraction embedding function, φ i t is the feature extractor of client i in round t, x is the sample, It means that the sample is fed into the feature extractor of the client in the tth round, and the features can be obtained for the next classification. The local prototype is represented by the mean of the local sample features.

[0045] Preferably, the set of samples to be penalized is defined as the samples whose cosine similarity between the extracted features and the prototype is greater than a threshold β, which is calculated as follows:

[0046]

[0047] in is the set of samples of category m that need to be punished, x m is a sample of category m, and the threshold β is a hyperparameter.

[0048] Preferably, the step S2 of calculating the local knowledge is specifically as follows:

[0049]

[0050] in, represents the local knowledge of class m of client i, Where R is the classification embedding function, is the classifier parameter of client i in round t, It means that the extracted features are fed into the classifier for classification; local knowledge distillation is represented by the mean of local sample feature classification.

[0051] Compared with the prior art, the present invention has the following beneficial effects:

[0052] (1) By utilizing prototype contrastive learning, not only are samples from the same category closer to their category prototypes than samples from other categories, but features are also separated from prototypes of other categories as much as possible, thus solving the problem of inter-category feature stacking.

[0053] (2) Based on the global prototype-guided feature alignment, the present invention introduces a loss penalty term to penalize sample features whose similarity exceeds a threshold, so that the similarity of all features with the corresponding global prototype is within the threshold, solving the problem of intra-class feature stacking.

[0054] (3) The present invention designs a heterogeneous federated learning method based on prototype contrastive learning and knowledge distillation to regularize the classifier of each local model and solve the problem of prototype shift. BRIEF DESCRIPTION OF THE DRAWINGS

[0055] Figure 1 It is an overall flow chart of a heterogeneous federated learning method based on prototype contrast learning and knowledge distillation according to an embodiment of the present invention;

[0056] Figure 2A schematic diagram of a framework of a heterogeneous federated learning method based on prototype contrast learning and knowledge distillation according to an embodiment of the present invention;

[0057] Figure 3 This is a specific flow chart of a heterogeneous federated learning method based on prototype contrast learning and knowledge distillation according to an embodiment of the present invention. DETAILED DESCRIPTION

[0058] The technical solution of the present invention is described in detail below in conjunction with the accompanying drawings.

[0059] Refer to the attached Figure 1-3 The present invention proposes a model performance optimization method based on heterogeneous federation, comprising the following steps:

[0060] S1. The server sends the global model, global prototype and global knowledge to the selected clients. Each client initializes the local model using the global model;

[0061] S2. On each client, use the global prototype, global knowledge and local samples to calculate the loss function and perform gradient descent. Execute the local model obtained after the epoch round and calculate the local prototype and local knowledge;

[0062] S3. After training is completed, the client uploads the local model, local prototype and local knowledge to the server. The server aggregates these local models, local prototypes and local knowledge to form a global model, global prototype and global knowledge. It is determined whether the convergence condition or the predetermined number of training rounds is reached. If so, the training is terminated, otherwise it returns to S1 to loop and execute the next round of training.

[0063] In this embodiment, steps S1, S2, and S3 need to go through multiple cycles to finally achieve convergence. That is, this process needs to go through multiple cycles, continuously iterate and optimize the model, so that it gradually converges to a better state. By iterating these steps multiple times, the federated learning system can improve the performance and generalization ability of the global model while protecting privacy.

[0064] In this embodiment, each client initializes the local model using the global model parameters; the model is divided into two components: feature extractor and classifier (composed of fully connected layers). The former is responsible for extracting features from the data, while the latter classifies the output features from the former. Their parameterization is as follows: the extractor is φ and the classifier is ψ. The present invention represents them as ω = [φ, ψ]. The embedding functions of these two components are F φ and R ψ .

[0065] In this embodiment, the client local model training is based on four indicators: cross entropy loss, prototype contrastive learning loss, penalty loss and global knowledge distillation loss.

[0066] In this embodiment, the client in step S2 uses the local model, global prototype, global knowledge and local samples to calculate the loss function, which specifically includes the following steps:

[0067] S21. According to formula (1), the cross entropy loss function of client i is obtained:

[0068]

[0069] Where N is the number of samples in an epoch, and M is the number of sample categories. m is the label of sample x, which takes 1 when the category belongs to m, otherwise it takes 0. m is the probability that sample x is predicted to belong to category m. is a sample dataset for the client.

[0070] S22. Calculate the local prototype contrast training loss according to formula (2):

[0071]

[0072] where x i,m is the sample of category m of client i, is a global prototype of category m, is the cosine similarity, is the input x i To the features extracted by the feature extraction network, τ is the scaling factor.

[0073] S23. Calculate the sample penalty loss of the model according to formula (3):

[0074]

[0075] in It is an indicator function used to express whether the sample belongs to the sample set that needs to be punished.

[0076] S24. Calculate the knowledge distillation loss according to formula (4):

[0077]

[0078] Where W is Sliced ​​Wasserstein Distance, is the global knowledge of category m, R is the embedding function of the classifier network, that is, R ψ (·) represents the function operation for classifying features.

[0079] S25. The overall loss function of the model calculated based on the four loss functions mentioned above is as follows:

[0080]

[0081] Where L ce is the common local classification cross entropy loss function. γ and λ are hyperparameters that can be adjusted. The overall loss function is the cross entropy loss Local prototype contrast training loss Sample Penalization Loss and knowledge distillation loss The model parameters ω are optimized by minimizing the loss function.

[0082] Through the local cross entropy loss Local prototype contrast training loss Sample Penalization Loss and knowledge distillation loss The introduction of can enhance the generalization ability of the model, accelerate convergence, improve robustness and adversarial ability, thereby achieving better performance and effects in federated learning.

[0083] In this embodiment, the server-side aggregation of the global prototype and the global model in step S3 specifically includes the following steps:

[0084] S31. Calculate the global model according to formula (6):

[0085]

[0086] Among them, |D i | represents the dataset D of client i i Size, |S t | is the number of clients selected.

[0087] A dataset representing the selected client The size of Represents the local model of client i.

[0088] S32. Calculate the global prototype according to formula (7):

[0089]

[0090] in, Represents the local prototype of client i.

[0091] S33. Calculate the global logit according to formula (8):

[0092]

[0093] in, represents the local knowledge distillation of client i.

[0094] In this embodiment, the step S23 further includes:

[0095] The aggregation formula for the local prototype of class m for client i is calculated as follows:

[0096]

[0097] Where i represents the client number, m represents the sample category, and t represents the federated learning round. represents the local prototype of class m for client i, represents the dataset of category m for client i, Where F is the feature extraction embedding function, φ i t is the feature extractor of client i in round t, x is the sample, It means that the sample is fed into the feature extractor of the client in the tth round, and the features can be obtained for the next classification. The local prototype is represented by the mean of the local sample features.

[0098] In this embodiment, the set of samples to which the penalty is imposed is calculated as follows

[0099]

[0100] in It is a set of samples of category m that need to be punished. The present invention defines it as samples whose cosine similarity between the extracted features and the prototype is greater than a threshold β, where the threshold β is a hyperparameter.

[0101] In this embodiment, the step S24 further includes:

[0102]

[0103] in, represents the local knowledge of class m of client i, Where R is the classification embedding function, is the classifier parameter of client i in round t, It means that the extracted features are fed into the classifier for classification; local knowledge distillation is represented by the mean of local sample feature classification.

[0104] By aggregating local models, local prototypes, and local knowledge distillation, it is possible to provide global information when initializing the client. Such a training mechanism helps alleviate the model bias caused by the heterogeneity of local models, thereby improving the effect and performance of federated learning.

[0105] From the above description, we can see that each client uploads its own calculated local model, local prototype, and local knowledge distillation. The server aggregates these data to obtain the next round of global model, global prototype, and global knowledge distillation.

[0106] In summary, the heterogeneous federated learning method based on prototype contrast learning and knowledge distillation provided by the present invention has the following beneficial effects:

[0107] Analysis from inter-class feature stacking: Prototype contrast learning improves this by promoting the difference between features and prototypes of other classes. This makes the features separate prototypes of other categories as much as possible, solving the problem of inter-class feature stacking.

[0108] Analysis from intra-class feature stacking: By imposing intra-class sample feature stacking penalties, the optimal sample feature diversity is maintained. This solves the problem of lack of sample feature diversity caused by intra-class feature stacking.

[0109] Analysis from prototype shift: This problem is usually caused by a biased classifier. The algorithm introduces soft prediction to provide more global knowledge, which regularizes the classifier part and effectively solves the problem caused by prototype shift.

[0110] The above are preferred embodiments of the present invention. Any changes made according to the technical solution of the present invention, as long as the resulting functions do not exceed the scope of the technical solution of the present invention, belong to the protection scope of the present invention.

Claims

1. A model performance optimization method based on heterogeneous federation, comprising the following steps: S1, the server sends the global model, global prototype and global knowledge to the selected client; Each client initializes its local model using the global model; S2. On each client, use the global prototype, global knowledge and local samples to train the local model, calculate the loss function, and perform gradient descent; execute the local model obtained after epoch rounds, and calculate the local prototype and local knowledge; S3. After training is completed, the client uploads the local model, local prototype and local knowledge to the server; The server aggregates the local models, local prototypes and local knowledge of all selected clients to form a global model, global prototype and global knowledge; Determine whether the convergence condition or the predetermined number of training rounds is reached. If yes, end the training; otherwise, return to S1 and loop to execute the next round of training.

2. According to the method for optimizing model performance based on heterogeneous federation according to claim 1, it is characterized in that: Each client uses the global model to initialize the local model specifically as follows: Each client initializes the local model using the global model parameters; the global model and the local model are divided into two components: a feature extractor and a classifier. The extractor is responsible for extracting features from the data, and the classifier classifies the output features from the feature extractor. The parameterization of the model is represented as ω = [φ, ψ], where φ is the extractor, ψ is the classifier, and the embedding functions of the extractor and the classifier are F φ and R ψ .

3. The model performance optimization method based on heterogeneous federation according to claim 2 is characterized in that: Client local model training is based on four indicators: cross entropy loss, prototype contrastive learning loss, penalty loss and global knowledge distillation loss.

4. The model performance optimization method based on heterogeneous federation according to claim 3 is characterized in that: In step S2, the client uses the local model, the global prototype, the global knowledge and the local sample to calculate the loss function, which specifically includes the following steps: S21. According to formula (1), the cross entropy loss function of client i is obtained: Where N is the number of samples in an epoch, M is the number of sample categories; y m is the label of sample x, which takes 1 when the category belongs to m, otherwise it takes 0; p m is the probability that sample x is predicted to belong to category m; is the sample data set of client i; S22. Calculate the local prototype contrast training loss according to formula (2): Where i represents the client number, m represents the sample category, and t represents the federated learning round; x i,m is the sample of category m of client i, m' is used to traverse all sample categories except category m, and are the global prototypes of categories m and m', is the cosine similarity, is the input x i,m To the features extracted by the feature extraction network, τ is the scaling factor; S23. Calculate the sample penalty loss of the model according to formula (3): in, is an indicator function used to express whether the sample belongs to the sample set that needs to be punished; is the set of samples of category m that need to be punished; S24. Calculate the knowledge distillation loss according to formula (4): Where W is Sliced ​​Wasserstein Distance, is the global knowledge of category m; R ψ (·) represents the function operation for classifying features; S25. The overall loss function of the model calculated based on the above four loss functions is as follows: Among them, γ and λ are adjustable hyperparameters, and the overall loss function is composed of the cross entropy loss Local prototype contrast training loss Sample Penalization Loss and knowledge distillation loss The model parameters ω are optimized by minimizing the loss function.

5. The method for optimizing model performance based on heterogeneous federation according to claim 4, characterized in that: In step S3, the server aggregates the local models, local prototypes and local knowledge of all selected clients to form a global model, a global prototype and global knowledge, which specifically includes the following steps: S31. Calculate the global model according to formula (6): Among them, |D i | represents the dataset D of client i i Size, |S t | is the number of selected clients, A dataset representing the selected client The size of represents the local model of client i; S32. Calculate the global prototype according to formula (7): in, represents the local prototype of client i; S33. Calculate the global logit according to formula (8): in, represents the local knowledge distillation of client i.

6. The model performance optimization method based on heterogeneous federation according to claim 4 is characterized in that: The step S2 of calculating the local prototype is specifically as follows: in, represents the local prototype of class m for client i, represents the dataset of category m for client i, Where F is the feature extraction embedding function, φ i t is the feature extractor of client i in round t, x is a sample, and the local prototype is represented by the mean of the local sample features.

7. The method for optimizing model performance based on heterogeneous federation according to claim 4, characterized in that: The set of penalized samples is defined as the samples whose cosine similarity between the extracted features and the prototype is greater than the threshold β, which is calculated as follows in is the set of samples of category m that need to be punished, x m is a sample of category m, and the threshold β is a hyperparameter.

8. The method for optimizing model performance based on heterogeneous federation according to claim 4, characterized in that: The step S2 of calculating the local knowledge is specifically as follows: in, represents the local knowledge of class m of client i, Where R is the classification embedding function, is the classifier parameter of client i in round t, and local knowledge distillation is represented by the mean of local sample feature classification.

Images