A federated machine learning method based on self-supervised learning

Abstract

The application provides a kind of federated machine learning method based on self-supervised learning.The method proposes a kind of general framework of federated learning, and the overall structure of machine learning model is divided into three parts of encoder, mapping head and prediction head.The encoder can be any general network architecture, such as CNN, ResNet, etc.The prediction head is a fully connected layer.The mapping head part is composed of a linear layer, a Batch Normalization (BN) layer, a relu activation function and another linear layer.When the client performs local training, the cosine similarity between the local model and the global model after passing through the mapping head is calculated, and this item is introduced into the loss function of local model training, to narrow the gap between the global model output and the local model output, so as to alleviate the adverse effects of data heterogeneity on federated learning accuracy.

Description

TECHNICAL FIELD

[0001] The present application belongs to the field of machine learning, and particularly relates to a federated machine learning method. BACKGROUND

[0002] In recent years, the number of terminal edge devices has grown rapidly, and these devices generate a large amount of user data every day. Due to various privacy issues and data regulations, it is difficult for us to obtain data on various devices for centralized training. Recently, federated learning (FL) has become an increasingly popular solution that enables multiple parties to collaboratively train machine learning models without exchanging local data. However, the client data involved in federated learning training is usually non-identically distributed, resulting in a data heterogeneity problem. The heterogeneity of data can reduce the training accuracy of the global model. Therefore, how to overcome the influence of data heterogeneity on learning effect in federated learning has attracted widespread attention from relevant researchers.

[0003] As of now, relevant researchers have proposed a large number of methods to alleviate the data heterogeneity problem in federated learning. Among them, MOON (Model-Contrastive Federated Learning) is one of the most advanced algorithms. It draws inspiration from the contrastive loss of NT-Xent loss and introduces a model-level contrastive loss in the local client for training, so that the local model output is as close as possible to the global model output and far away from the last round of local model output. However, this method requires saving all the last round models of the participating training clients to obtain high accuracy, which requires a large amount of memory overhead and is relatively slow. Therefore, how to reduce memory overhead and improve speed in model contrastive federated learning is a valuable research direction. SUMMARY

[0004] The purpose of the present application is that in federated machine learning, due to the different data between different clients, there is a data heterogeneity problem due to different distributions, and the heterogeneity of data can cause the learning effect of the system to deteriorate. The present application provides a federated machine learning method based on self-supervised learning, which reduces the gap between the global model output and the local model output by improving the loss function, and improves the training accuracy of the federated machine learning method.

[0005] SUMMARY: To solve the above technical problems, the present application provides a federated machine learning method based on self-supervised learning, characterized in that: the method divides the machine learning model into three parts: encoder, mapping head and prediction head, introduces a Batch Normalization (BN) layer in the mapping head composed of MLP, and introduces the cosine similarity between the outputs of the global model and the local model mapping head into the loss function of the local training. The method comprises the following steps:

[0006] (1) The server initializes the global model and sets a client participation rate, randomly selects clients for training in this round according to the client participation rate, and then distributes the model parameters to the selected clients;

[0007] (2) The client receives the global model parameters distributed by the server and uses them for local model initialization, updates the model using the sum of the supervised learning loss of the local model and the contrastive loss between the local model and the global model, and then uploads the model parameters to the server;

[0008] (3) After the server receives the model parameters uploaded by the client, the model parameters are aggregated according to the proportion of the client data set to the total data set size, and then the next round of clients participating in training are reselected according to the set client participation rate, and the next round of model parameter training begins.

[0009] Further, in step (1), the server initializes the global model and distributes it to the client, characterized in that: the model distributed by the server consists of an encoder, a mapping head and a prediction head. The encoder can be any general network architecture, such as CNN, ResNet, etc.; the mapping head consists of a linear layer, a BN layer, a relu activation function and another linear layer; the prediction head is a fully connected layer.

[0010] Further, in step (2), the client trains the local model, characterized in that: the local training loss consists of the cross-entropy loss in supervised learning and the contrastive loss between the local model and the global model mapping head output. Specifically, for data x in the local data set D i of client i, z x represents the output of x after passing through the local encoder and projection head, z glob,x represents the output of x after passing through the mapping head of the global model, and the contrastive loss l con between the global model and the local model is defined as:

[0011]

[0012] Where <z x ,z glob,x > represents the dot product between z x and z glob,x , ||z x ||2 and ||z glob,x ||2 represent the modulus of z x and z glob,x , respectively.

[0013] Further, the method for the server to aggregate the model parameters in the step (3) is characterized in that: in the tth round, there are N clients participating in the training, the total size of the data sets of the N clients is |D|, the size of the local data set of the client i is |D i The calculation method of the global model parameters w of the t+1th round is as follows: t+1

[0014]

[0015] Beneficial effects:

[0016] The present application considers the problem of data heterogeneity in federated machine learning, divides the model into three parts of encoder, mapping head and prediction head, adds a Batch Normalization (BN) layer in the mapping head usually composed of MLP, and introduces the cosine similarity between the outputs of the local model and the global model through the mapping head in the loss function of the local training, so as to reduce the distance between the outputs of the global model and the local model, and to alleviate the influence of data heterogeneity on the accuracy of federated learning. At the same time, compared with the existing model comparison federated learning method (MOON), it does not need to save the model of the last round on the client, and at the same time of obtaining higher accuracy, it speeds up the training speed. BRIEF DESCRIPTION OF DRAWINGS

[0017] Figure 1 is the flow chart of the federated machine learning method based on self-supervised learning provided by the present application.

[0018] Figure 2 is the overall structure diagram of the machine learning model in the method of the present application. DETAILED DESCRIPTION

[0019] The present application provides a federated machine learning method based on self-supervised learning, characterized in that: the method divides the model into three parts of encoder, mapping head and prediction head, introduces a BN layer in the mapping head composed of MLP, and introduces the cosine similarity between the outputs of the data through the mapping heads of the global model and the local model into the loss function of the local training. Through the combination of the above methods, the training accuracy of federated learning on heterogeneous data sets is improved.

[0020] The present application scheme design is further specifically described below in combination with the drawings and related formulas.

[0021] As shown in the accompanying Figure 1 The method provided by the present application comprises the following steps:

[0022] ​(1) The server randomly initializes a global model and sets a client participation rate, randomly selects part of the clients to participate in this round of training according to the client participation rate, and then distributes the model parameters to the selected clients in this round. The distributed machine learning model is composed of an encoder, a mapping head and a prediction head. The encoder can be any general network architecture, such as CNN, ResNet, etc.; the mapping head is composed of a linear layer, a Batch Normalization (BN) layer, a relu activation function and another linear layer; and the prediction head is a fully connected layer.

[0023] (2) The client receives the global model parameters distributed by the server and loads them into the local model for the initialization of the local model in this round of communication. Then, the model is updated using the sum of the supervised learning loss of the local model and the comparison loss between the local model and the global model, and then the updated model parameters are uploaded to the server.

[0024] (3) After the server receives the model parameters uploaded by the client, the model parameters are aggregated using the following method: suppose that there are N clients participating in the training in the tth round, the total size of the data sets of the N clients is |D|, the size of the local data set of client i is |D i , and the model parameters of client i in the tth round are The calculation method of the global model parameters w t+1 in the t+1th round is as follows:

[0025]

[0026] After obtaining the updated global model, the server reselects the clients according to the set client participation rate and distributes the model parameters, starting the next round of model training.

[0027] As shown in the accompanying drawings, Figure 2 step (2) The local loss used in the local training process of the client is composed of two parts: the first part l sup is the cross-entropy loss in supervised learning, and the second part l con is the model comparison loss term proposed in the application.

[0028] For input pair (x, y), y is the label corresponding to data x, and the loss l(x, y) is calculated as follows:

[0029]

[0030] where w t represents the global model weight in the tth round, is the model weight of client i in the tth round, and μ is a hyperparameter for controlling the weight of the comparison loss.

[0031] For a classification problem of M categories, lsup The calculation is

[0032]

[0033] where p x represents the probability of x outputting M categories after passing through the local encoder, the mapping head and the prediction head, p x,c represents the probability of the prediction output x corresponding to the category c, y x,c is a Boolean variable with a value of 0 or 1, and is defined as follows:

[0034]

[0035] For data x in the local data set D i of the client i, z x represents the output of x after passing through the local encoder and the projection head, z glob,x represents the output of x after passing through the mapping head of the global model. The contrastive loss l con between the global model and the local model is defined as:

[0036]

[0037] where <z x ,z glob,x > represents the dot product between z x and z glob,x , ||z x ||2 and ||z glob,x ||2 represent the modulus of z x and z glob,x , respectively.

[0038] In summary, the optimization objective of the local model is

[0039]

[0040] The above only describes the preferred embodiments of the present application, and the protection scope of the present application is not limited to the above-mentioned embodiments, but any equivalent modifications or changes made by those skilled in the art according to the disclosed content of the present application shall be included in the protection scope recorded in the claims.

Claims

1. A federated machine learning method based on self-supervised learning, characterized in that: This method divides the machine learning model into three parts: encoder, mapping head, and prediction head. In the mapping head part composed of MLP, a Batch Normalization (BN) layer is introduced, and the data is processed by the cosine similarity between the outputs of the global model and the local model mapping head, and then introduced into the loss function of local training. The method includes the following steps: (1) The server initializes the global model and sets a client participation rate. Based on the client participation rate, it randomly selects the client to be used for training in this round and then sends the model parameters to the selected client. (2) The client receives the global model parameters sent by the server and uses them to initialize the local model. It updates the model by using the sum of the supervised learning loss of the local model and the contrast loss between the local model and the global model, and then uploads the model parameters to the server. (3) After receiving the model parameters uploaded by the client, the server aggregates the model parameters according to the proportion of the client's dataset to the total dataset size, and then selects the next round of clients to participate in the training according to the set client participation rate, and starts the next round of model parameter training. In step (2), the client trains a local model. The local training loss consists of two parts: the cross-entropy loss from supervised learning and the contrast loss between the output of the local model and the global model mapping head. Part One It is the cross-entropy loss in supervised learning, Part 2. It is the contrastive loss of the model; For input pairs , For data The corresponding label, its loss The calculation is as follows: (2) in, Indicates the first The global model weight parameters of the wheel, For the client In the Wheel model weight parameters, To control the hyperparameters of the contrast loss weights; for A classification problem involving several categories. The calculation is as follows (3) in, express The corresponding output after local encoding, mapping header, and prediction header The probability of each category, Indicates the predicted output The corresponding category is The probability, Let it be a Boolean variable with a value of 0 or 1, defined as follows: (4) For the client Local dataset Data x in This represents the output of x after passing through the local encoder and mapping head. This represents the output of x through the global model's mapping head; the contrast loss between the global and local models. Defined as: (5) in, represent and dot product between and Represent and The model; The optimization objective of the local model is （6） Indicates the first The global model weight parameters of the wheel, For the client In the Wheel model weight parameters, client Local dataset Part 1: Cross-entropy Loss in Supervised Learning The second part compares the loss terms of the model. .

2. The federated machine learning method based on self-supervised learning according to claim 1, characterized in that: In step (1), the server initializes the global model and sends it to the client. The model sent by the server consists of three parts: an encoder, a mapping head, and a prediction head. The encoder is an arbitrary general network architecture; the mapping head consists of a linear layer, a BN layer, a ReLU activation function, and another linear layer; the prediction head is a fully connected layer.

3. The federated machine learning method based on self-supervised learning according to claim 1, characterized in that: The method for determining the server aggregation model parameters in step (3) is as follows: Wheel One client participates in the training, this The total size of the client datasets is Client The local dataset size is Client No. The model weight parameters of the wheel are denoted as , No. Global model weight parameters of the wheel The calculation method is as follows: 。 4. The federated machine learning method based on self-supervised learning according to claim 1, characterized in that: The encoder can be any general network architecture, such as CNN or ResNet.

Images