An implementation method of personalized federated learning based on client self-knowledge distillation

Abstract

The application discloses a kind of based on client self-knowledge distillation's personalized federated learning implementation method.The method includes: server initializes model and is issued to all clients;Server issues global model to client when a new communication round starts;After receiving global model, client updates its local model by local aggregation module;When client local model parameter is trained, the knowledge of historical personalized model is migrated to current local model by knowledge distillation, and the proportion of distillation loss is dynamically adjusted;After training is completed, client updates historical personalized model and sends the parameter of local model to server;Server obtains new model by weighted average according to the model parameter sent by client, for the training of next communication round;The application improves the performance of personalized federated learning under the background of data heterogeneity without additional communication overhead.

Description

TECHNICAL FIELD

[0001] The present application relates to the technical field of federated learning, and particularly relates to a personalized federated learning implementation method based on client self-knowledge distillation. BACKGROUND

[0002] With the development of artificial intelligence technology, especially in the context of the big data era, the rapid growth of data and the popularity of distributed storage make these massive amounts of information usually distributed in various devices and places, including mobile devices, edge devices, and cloud servers. In traditional centralized machine learning, data is usually stored centrally on a centralized server. However, the transmission of large-scale data sets puts a huge demand on network bandwidth. Especially in edge computing and Internet of Things scenarios, the communication cost between devices is more significant. This situation not only affects the efficiency of data transmission, but also can cause serious communication delay, thereby limiting the real-time and overall efficiency of model training. In application scenarios that require rapid response to data changes and maintain the timeliness of the model, the limitations of this traditional method become particularly prominent. Moreover, in sensitive areas involving user privacy, such as the medical and financial industries, there are great privacy protection problems. The traditional centralized machine learning method of transmitting raw data to a central server is vulnerable to sensitive information leakage due to intermediate nodes or network attacks during information transmission.

[0003] In this context, the shortcomings of traditional machine learning methods are apparent, and in most industries, data exists in the form of an island. Due to industry competition, privacy security, complex administrative procedures, and other issues, even between different departments of the same company, data integration faces many obstacles, and in reality, it is almost impossible or extremely costly to integrate data scattered across various locations and institutions.

[0004] To solve these problems, federated learning has emerged. Federated learning is a distributed machine learning model that essentially trains a global model representing all user devices through multiple user devices, without the exchange of user data. Compared to more common distributed machine learning, federated learning emphasizes privacy. Federated learning distributes an initial model to all edge clients through a central server, each client trains the model locally on their own data, and then updates the model information to the central server. That is, without leaving the local data, federated learning performs joint data training, greatly reducing the risk of sensitive information leakage. Secondly, compared to traditional centralized learning that requires the transmission of large amounts of raw data, federated learning only needs to transmit model parameter updates in each iteration, effectively reducing communication burden. Federated learning overcomes the limitations of traditional centralized learning, provides a more private and secure learning environment for users, improves communication efficiency, and promotes data collaboration and sharing.

[0005] Several fundamental challenges exist in general federated learning methods: 1. Poor convergence on highly heterogeneous data. When learning on non-independent and identically distributed (non-IID) data, client drift significantly reduces the accuracy of the global model. 2. Lack of personalized solutions. In the original federated learning framework, a single global shared model is trained to adapt to "averaged clients." When the data distributions of different clients vary significantly, a single global model struggles to handle local distributions that differ drastically from the global distribution. Therefore, we need to address these challenges by providing personalized solutions for each client. These two challenges are precisely what personalized federated learning attempts to solve.

[0006] In recent years, knowledge distillation has attracted increasing attention from the research community, with large-scale deep neural networks achieving remarkable success, especially in real-world scenarios with massive datasets. However, deploying DNN models on mobile devices remains a significant challenge due to their limited computing power and storage. To address this issue, Bucilua et al. (2006) first proposed model compression, transferring information from large or ensemble models to smaller training models without significantly reducing accuracy. The method of learning from large models to smaller models was later formally popularized as knowledge distillation. Self-distillation is one form of knowledge distillation. Its main idea is that the network itself acts as both a teacher and a student model, and the knowledge generated in previous training steps contributes to the current step of model training.

[0007] In summary, while traditional federated learning can reduce the risk of privacy breaches and lower communication overhead, it suffers from low model accuracy when learning from non-independent, identically distributed data. Therefore, introducing knowledge distillation techniques and personalized federated learning allows for the creation of customized models for each client and the transfer of knowledge to local models, thereby improving accuracy. Summary of the Invention

[0008] This invention overcomes the difficulty of existing federated learning methods in adapting well to the different data distributions of each client under heterogeneous data conditions, where a single global model cannot adequately adapt. It implements a personalized federated learning method that performs local aggregation and dynamic self-distillation on the client side. This method utilizes the general knowledge in the global model and the historical personalized model from the previous communication round to extract the inherent general and personalized knowledge to improve the local model on the client side. This invention improves model performance without introducing additional communication overhead.

[0009] A personalized federated learning implementation method based on client-side self-knowledge distillation includes: at the start of a new communication round, the server sends a global model to each client. The global model is a weighted average of the model parameters submitted to the server by all clients in the previously selected client subset in the previous communication round, with the weights determined by the proportion of each client's dataset size to the total dataset size.

[0010] After each client receives the global model parameter set sent by the server, it uses the local aggregation module to extract the information needed by the current client from the global model to initialize the local model. The local aggregation module is a weight matrix used to extract low-level general information in the deep neural network. The local model is a model stored on the client's local machine.

[0011] After the client initializes the local model, it begins training the local model parameters for the current communication round. During local model parameter training, knowledge distillation is used to transfer the knowledge of the historical personalized local model from the previous communication round to the local model for the current communication round. At the same time, the proportion of distillation loss is dynamically adjusted according to the increase of the communication round.

[0012] After the client completes the training for the current communication round, it sends the parameter set of the local model to the server.

[0013] After all clients in the waiting subset send back the new set of model parameters, the server generates a new set by weighted averaging the new set of model parameters, which is used for training in the next communication round, until the upper limit of communication rounds is reached.

[0014] At the start of a new communication round, the server sends a global set of model parameters to each client, including:

[0015] Before the start of the new communication round t+1, the server sends data to the client set S according to the following formula. t Each client in the process sends a global model Θ t .

[0016]

[0017] Where, Θ t S is the set of model parameters sent to the server by all clients in the client subset after the previous communication round t. t It is the set of clients selected in the previous communication round. Where D i It is the size of the local dataset owned by the i-th client.

[0018] When each client receives the global model parameter set sent by the server, the local model initialization for each client includes:

[0019] The client i selects s% local data to train the local aggregation weight matrix W according to the following formula.

[0020]

[0021] When the client i receives the global model parameter, the client updates the local model according to the following formula

[0022]

[0023] wherein, is the local model of the last communication round, is the update of the model, and is the Hadamard product. W i is the local aggregation weight matrix of the i-th client, and all ω ∈ [0, 1]. |Θ i | is the number of layers of the global model, and the low layers of the local model have the same structure and are constant 1. W i p and the p high layers of the local model have the same structure. p is a hyperparameter to control the range of the local aggregation module.

[0024] When training the local model parameters, the knowledge of the historical personalized local model is migrated to the current local model, including:

[0025] At any client k in the communication round t, the knowledge distillation is used on the local model of the t-th round using the local model of the t-1-th round on the client k, and the proportion of the distillation loss is dynamically adjusted. The loss function is defined as:

[0026]

[0027] wherein, the hyperparameter λ controls the contribution degree of the knowledge distillation. F k (·) represents the cross-entropy loss of the client k, represents the KL divergence between the input structure of the historical personalized model and the output result of the current local model. The soft label q(·) is the result after the output z of the last fully connected layer is normalized by the exponential (Softmax) function, that is, is a temperature hyperparameter used to control the Softmax function. The client updates the local model parameters using the stochastic gradient descent according to the loss function of the following formula.

[0028]

[0029] in Let η represent the gradient of the loss function, and η be the learning rate.

[0030] When training local model parameters, dynamically adjusting the proportion of distillation loss includes:

[0031] The loss function on client k is as follows:

[0032]

[0033] in, r and R represent the current training rounds and the total number of communication rounds, respectively. This represents the threshold that fixes the ratio of the two loss functions.

[0034] When the client completes training for the current communication round, it sends the following parameters of the local model to the server:

[0035] At any client k in communication round t, the updated parameters will be sent to the server. At the same time, client k updates the historical personalized model to

[0036] After waiting for all clients to send back new model parameters, the server generates a new set based on these new model parameters for training in the next communication round, including:

[0037] The server will use a subset S of the clients. t The set of parameters sent by all clients in {Θ 1 ,Θ 2 ,…,Θ N Update Θ according to Formula 1 t-1 The parameters are then used to begin training for the next communication round, where... This represents the parameters that client i sends to the server in the t-th communication round.

[0038] The core components of this invention are the client and the server. The server includes the following steps:

[0039] Step 1: The server first initializes model Θ 0 The server will use the global model Θ 0 Send to all clients. Set the current communication round variable to t, with the initial communication round t = 0.

[0040] Step 2: The server selects a subset of clients to connect. Connected clients upload the parameters of their local model to the server. Given the model parameters for the i-th client in round t, the server updates the model according to Formula 1 to obtain Θ. t+1 .

[0041] Step 3: The server randomly selects a subset of clients S from the client set C according to the participation rate p. t The server sends the parameter set Θ t to each client in S t+1 ; the current communication round t is incremented by 1.

[0042] Each client runs in parallel after receiving the global model parameter set Θ t+1 from the server, and the steps are as follows:

[0043] S1: After receiving the global model parameter set Θ from the server, the client k extracts the knowledge needed for the client k in the global model according to the local aggregation module according to formula 3.

[0044] S2: After obtaining the local model i of the tthcommunication round, the client k obtains the value of the local model output

[0045] S3: The client calculates the task-related loss. For classification tasks, the task loss is composed of the cross-entropy loss function and the distillation loss function. The distillation loss is the KL divergence of the historical personalized local model and the current local model output. The sum of the distillation loss is obtained by formula 6.

[0046] The KL divergence helps to migrate the knowledge in the historical personalized model to the current client, and the distillation loss proportion increases with the increase of the communication round. Wherein, r and R represent the current model training round and the total number of communication rounds, respectively, represents the threshold for fixing the proportion of the two loss functions.

[0047] S4: The client reduces the loss function to update the parameters of the local model

[0048] S5: The client sends the updated model parameters back to the server.

[0049] After all the clients in the S t set have sent the model parameters, the server will update the parameter set Θ t+1 .

[0050] Step 4: If the maximum communication round has been reached, the training is complete, otherwise it will continue to step 3.

[0051] The beneficial effects of the present application are as follows: the improved client-based self-knowledge distillation personalized federated learning implementation method of the present application utilizes the general information of the global model and the knowledge in the historical personalized model, and utilizes it during client initialization and local training, effectively improving the performance of the client local model without additional federated learning communication overhead; the local aggregation module is used to migrate the general information of the above global model to the local model, and the general information of the model from different clients is extracted; during local training, the historical personalized model is regarded as a teacher model, and the client local model is regarded as a student model, and the information migration process is performed. Compared with the traditional federated learning which improves the performance of the global model on each client by aggregating the model on the server, the client utilizes the general knowledge of the global model and performs self-distillation of the local model to directly improve the performance of the local model. BRIEF DESCRIPTION OF DRAWINGS

[0052] In order to more clearly illustrate the technical solutions of the present application, the drawings required to be used in the examples will be briefly introduced as follows: obviously, for those skilled in the art, other drawings can also be obtained without creative labor on the basis of these drawings.

[0053] Figure 1 The working process of the classical federated learning method;

[0054] Figure 2 The federated learning implementation method of the present application;

[0055] Figure 3 The schematic diagram of the client performing local aggregation. DETAILED DESCRIPTION

[0056] In order to make the purpose, technical solutions and advantages of the present application clearer, the technical solutions of the present application will be described clearly and completely in combination with specific embodiments of the present application and corresponding drawings. Obviously, the described embodiments are only some of the embodiments of the present application, not all the embodiments. Based on the embodiments in the present application, all other embodiments obtained by those skilled in the art without creative labor are within the scope of protection of the present application. The technical solutions provided by each embodiment of the present application will be described in detail below in combination with the drawings.

[0057] Figure 1This diagram illustrates the workflow of federated learning. One goal of federated learning is to build a general global model that performs well across various clients. However, due to the different data distributions on each client, a single global model is unlikely to perform well on every client. Therefore, each client's local model needs to extract the knowledge it needs from the global model. The local aggregation module is used to utilize the knowledge it requires from the global model. In deep neural networks, lower layers learn more general knowledge, while higher layers learn more personalized knowledge. When updating the client's local model, a weight matrix is ​​used to control the model update. The lower layers of the local model retain the parameters of the global model, while the higher layers retain more of the parameters of the local model to achieve personalization. The transfer of knowledge from one model to another has been extensively studied, one approach being knowledge distillation. Knowledge distillation involves two types of models: a teacher model with a large number of parameters and a student model with relatively few parameters. The goal is to transfer the knowledge information contained in the teacher model, which has strong representational capabilities, to the student model, enabling the latter to achieve performance similar to the former with fewer model parameters. In this invention, each client needs to use a local aggregation module to extract general knowledge of the global model during model initialization, and during local training on the client, the local model of the previous communication round is used as the teacher model. Knowledge distillation technology is used to transfer the knowledge contained in the historical personalized local model to the local model of the current communication round, thereby improving the performance of the local model.

[0058] This invention provides a novel implementation of personalized federated learning to address the low accuracy problem of traditional federated learning in heterogeneous data environments. In this invention, the client obtains the weighted average models sent by other clients to the server and initializes its local model using a local aggregation module. The client treats the historical personalized model from the previous communication round as the teacher model and then uses knowledge distillation to transfer the information it contains into its local model. To help users better understand the purpose and workflow of this invention, the technical details will be described below with reference to the accompanying drawings. Figure 2 In this scenario, there are k clients. Initially, the server generates an initialization model and sends it to all clients. Assume the current task is image classification, and client k has image data X in its local dataset. i and the corresponding data label Y i Save as D k When the client trains, it will proceed sequentially from D. k Reading data from the middle - tag pairs {(X i ,Y i). The model used by the clients are all the same structure, each client has an independent local model whose parameters are Θ k .

[0059] Step 1: The server randomly generates an initial model Θ 0 and sends it to all the clients.

[0060] Step 2: t = t + 1, the client trains or reuses the local aggregation module according to t, the following is the step of updating the local aggregation module of any client k:

[0061] S1: The client k selects the local data set s% according to the local data.

[0062] S2: When t = 2, the local aggregation module is trained according to formula 2 in a loop until it converges. Then the weight matrix W is pruned to prevent the weight value from being too large or too small.

[0063] S3: When t > 2, the client k only needs to train the W matrix according to formula 2 once and prune it once.

[0064] Step 3: The client updates its local model according to formula 3 through the local aggregation module.

[0065] Step 4: After the client initializes the local model, all clients will run the training in parallel, the following is the step of updating the local model Θ k of any client k:

[0066] S1: Get the data-label pair (X k , Y i ) based on the current client k's data set D i , where i represents the index of the data

[0067] S2: When the client initializes the local model, it will calculate two loss functions. The steps are as follows:

[0068] A: Get the predicted label output f(X i , Θ i ) of the current client local model based on X k ;

[0069] B: Get the cross-entropy loss between the model output and the hard label based on X i and using formula ;

[0070] C: Calculate the distillation loss between the model output and the historical personalized model using formula ;

[0071] D: Define the total loss function as:

[0072] L = pL hard + (1-p)L soft wherein r and R represent the current model training round and the total number of communication rounds respectively, represents the threshold value of fixing the proportion of two loss functions;

[0073] E: update the model Θ locally on the client according to formula 5 k .

[0074] S3: run steps A to E in a loop until i = |D k |.

[0075] S4: when the client k completes the training, send the new Θ k back to the server.

[0076] Step 5: after the server receives the new parameters Θ k from all clients, a new global model Θ t is formed by weighted average according to the size of the data of each client, and jump to step 2.

[0077] Unlike the traditional federated learning algorithm which constructs a general global model by performing weighted average of model parameters on the server, the method described in the application improves the performance of the local model by performing local aggregation and knowledge distillation locally on the client to achieve individualization. Through the feedback of the example, using the method described in the application, the learning efficiency of the model and the performance of the model are effectively improved without increasing the additional communication overhead.

[0078] The embodiment of the application also provides a storage medium, wherein the storage medium stores a computer program, and the computer program is executed by a processor to realize part or all steps in each embodiment of the federated learning implementation method based on client knowledge distillation provided by the application. The storage medium can be a disk, an optical disc, a read-only memory (English: Read-Only Memory, abbreviated: ROM) or a random access memory (English: Random Access Memory, abbreviated: RAM) and the like.

[0079] Those skilled in the art can clearly understand that the technical solution in the embodiments of the present application can be realized by means of software plus a necessary general hardware platform. Based on such an understanding, the technical solution in the embodiments of the present application can be embodied in the form of a software product, which can be stored in a storage medium, such as a ROM / RAM, a magnetic disk, an optical disk, and the like, and includes a number of instructions to make a computer device (which can be a personal computer, a server, or a network device, etc.) execute the methods described in the various embodiments or some parts of the embodiments of the present application.

[0080] The above-described embodiments of the present application do not constitute a limitation on the protection scope of the present application.

Claims

1. A personalized federated learning implementation method based on client-side self-knowledge distillation, characterized in that, include: At the start of a new communication round, the server sends a global model to each client. The global model is a weighted average of the model parameters submitted to the server by all clients in the previously selected client subset in the previous communication round. The weights are determined by the proportion of each client's dataset size to the total dataset size. After each client receives the global model parameter set sent by the server, it uses the local aggregation module to extract the information needed by the current client from the global model to initialize the local model. The local aggregation module is a weight matrix used to extract low-level general information in the deep neural network. The local model is a model stored on the client's local machine. After the client initializes the local model, it begins training the local model parameters for the current communication round. During local model parameter training, knowledge distillation is used to transfer knowledge from the historical personalized local model of the previous communication round to the local model of the current communication round. Simultaneously, the weight of the distillation loss is dynamically adjusted according to the number of communication rounds. The specific steps are as follows: The client performs an image classification task, and image data X exists in the local data of client k. i and the corresponding data label Y i Save as D k； When the client performs training, it will proceed sequentially from D... k Reading data from the middle - tag pairs {(X i ,Y i )}; S1: Dataset D based on the current client k k Obtain data – label pairs (X i ,Y i ), where i represents the index of the data. S2: After the client initializes the local model, two loss functions will be calculated; the steps are as follows: A: Based on X i Obtain the predicted label output f(X) of the current client-side local model. i ,Θ k ); B: Based on X i And using the formula We obtain the cross-entropy loss between the model output and the hard labels, where F k (·) represents the cross-entropy loss of client k; C: Using the formula The calculation model output is compared with the distillation loss of the historical personalized model, where This represents the KL divergence between the input structure of the historical personalized model and the output of the current local model. D: Define the overall loss function as: L=ρL hard +(1-ρ)L soft ,in r and R represent the current training rounds and the total number of communication rounds, respectively. The threshold representing the ratio of two fixed loss functions; E: Update the local model using stochastic gradient descent. k ; S3: Repeat steps A to E until i = |D k |; After the client completes the training for the current communication round, it sends the parameter set of the local model to the server. After all clients in the waiting subset send back the new set of model parameters, the server generates a new set by weighted averaging the new set of model parameters, which is used for training in the next communication round, until the upper limit of communication rounds is reached.

2. The method as described in claim 1, characterized in that, At the start of a new communication round, the server sends a global set of model parameters to each client, including: Before the start of the new communication round t+1, the server sends data to the client set S according to the following formula. t Each client in the process sends a global model Θ t ; Where, Θ t S is the set of model parameters sent to the server by all clients in the client subset after the previous communication round t; t It is the set of clients selected in the previous communication round; Where D i It is the size of the local dataset owned by the i-th client.

3. The method as described in claim 1, characterized in that, After each client receives the global model parameter set sent by the server, local model initialization is performed for each client, including: Client i selects s% of local data and trains the local aggregated weight matrix W according to the following formula; After receiving the global model parameters, the client updates its local model according to the following formula. in, This is the local model from the previous communication round. For model updates, ⊙ is the Hadamard product; W i Let be the local aggregation weight matrix for the i-th client. All have ω∈[0,1]; |Θ i | represents the number of layers in the global model. It has the same structure as the lower layers of the local model, and all of them are constants of 1; W i p It has the same structure as the p high-level modules of the local model; p is a hyperparameter that controls the scope of the local aggregation module.

4. The method as described in claim 1, characterized in that, When the client completes training for the current communication round, it sends the following parameters of the local model to the server: At any client k in communication round t, the updated parameters will be sent to the server. At the same time, client k updates the historical personalized model to 5. The method as described in claim 2, characterized in that, After waiting for all clients to send back new model parameters, the server generates a new set based on these new model parameters for training in the next communication round, including: The server will use a subset S of the clients. t The set of parameters sent by all clients in {Θ 1 ,Θ 2 ,…,Θ N Update Θ according to Formula 1 t-1 The parameters are then used to begin training for the next communication round, where... This represents the parameters that client i sends to the server in the t-th communication round.

Images