Federal learning communication method and system based on knowledge distillation

Abstract

The invention discloses a federal learning communication method and system based on knowledge distillation, and the method comprises the following steps: 1, initializing student model parameters and teacher model parameters, and setting related parameters; 2, performing local training on each client, and training the teacher model and the student model respectively; 3, calculating the average logits of the student model and uploading the average logits to a server; step 4, the server calculates the global average logits and the global distillation loss, and issues the global average logits; and 5, updating student model parameters and teacher model parameters. According to the method, the communication data volume can be reduced, and the adaptability and generalization ability of the global model during non-IID data processing can be ensured.

Description

Technical Field

[0001] The present invention belongs to the technical field of machine learning, and in particular relates to a federated learning communication method and system based on knowledge distillation. Background Art

[0002] With the rapid development of deep learning technology, Federated Learning (FL), as a distributed and decentralized machine learning method, is gradually becoming a mainstream technology for data privacy protection and distributed model training. The core principle of federated learning is to allow each client (such as mobile phones, IoT devices, etc.) to independently train models on local data while ensuring that the data does not leave the local device, and then upload the model updates to the central server for merging. Compared with traditional centralized deep learning methods, federated learning has two significant advantages:

[0003] 1. Data privacy protection: Because the data will not be transmitted to the server, the user's privacy is better protected.

[0004] 2. Reduce network bandwidth consumption: Federated learning reduces dependence on network bandwidth. Federated learning only needs to transmit model updates (such as model parameters or gradients) instead of raw data, thus reducing dependence on network bandwidth.

[0005] In addition, federated learning can also use the computing power of distributed devices to share the burden of central servers and improve overall computing efficiency. However, the implementation of federated learning is accompanied by huge communication overhead, especially when the model complexity and the number of clients continue to increase, the communication cost becomes the main bottleneck of system efficiency. The communication overhead problem of federated learning is mainly reflected in the fact that after each round of model training, the client needs to upload local model parameters or gradient updates (gradients are the partial derivatives of the calculation model error with respect to the parameters, which are used to optimize the model parameters and make the loss function converge) to the server. For complex deep neural networks (a model structure with multiple layers of neurons, widely used to solve complex pattern recognition problems) that may contain millions or even hundreds of millions of parameters, uploading this data for each round of training will bring a heavy burden to network bandwidth and device resources. The communication overhead problem is particularly prominent, which is specifically reflected in the following aspects:

[0006] (1) Bandwidth limitation: In real-world application scenarios, federated learning is often deployed on distributed devices in different geographical locations, and the network environments of different devices vary greatly. In scenarios with slow networks (such as remote areas or mobile devices), the transmission speed of model parameters is extremely slow, resulting in a significant increase in the time required for each round of communication, thereby extending the cycle of the entire training process.

[0007] (2) Limited client device resources: Many client devices participating in federated learning (such as smartphones, IoT devices, and edge computing nodes) are limited by their own computing power and battery life. When faced with the huge demand for model parameter transmission, these devices consume a lot of resources. Frequent communication not only leads to a sharp increase in energy consumption, but also may affect the normal operation of the device.

[0008] (3) Data privacy and security risks: Although federated learning avoids direct data transmission by training on local data, the large number of model parameters that are frequently uploaded still poses the risk of being attacked or leaked. Malicious attackers may infer certain features of local data by analyzing the transmitted model parameters, thereby posing a threat to user privacy.

[0009] (4) Challenges brought by non-IID data: In practical applications, the data distribution of different clients is often non-IID (i.e., the data distribution of clients is different, and the data of each client may be biased, resulting in adaptability issues when the model is merged). This will cause the model trained by each client to perform inconsistently when the global model is aggregated. In order to ensure the effectiveness of the global model, more rounds of communication and training are required, which further increases the communication overhead.

[0010] (5) Increased model complexity: As deep learning models such as ResNet, Transformer, and other large models become more and more complex (ResNet is a deep convolutional neural network suitable for image tasks; Transformer is a model structure based on the attention mechanism, widely used in natural language processing), the number of model parameters has also increased exponentially. In this case, the amount of data required to be transmitted in each round of communication in federated learning has also increased exponentially, resulting in extremely huge bandwidth consumption and increased communication burden.

[0011] At present, researchers have tried various methods to reduce the communication overhead in federated learning, including gradient compression, quantization, sparsification and other technologies. However, these methods have shown certain limitations in high-frequency, large-scale, and multi-client application scenarios. Especially in the complex environment of non-IID data, the trade-off between communication overhead and global model performance has become a difficult problem that needs to be solved urgently. Therefore, designing an efficient federated learning communication mechanism that can significantly reduce the amount of communication while ensuring model performance has become the key to solving the current problem. Summary of the invention

[0012] In order to solve the above problems, the present invention proposes a federated learning communication method and system based on knowledge distillation. The present invention improves the federated learning framework (through training the student model, using the teacher model to guide the student model learning, so that the student model is close to the performance of the teacher model under smaller parameters), and replaces the traditional model parameter uploading method with a transmission mechanism based on logits (the original score output by the model, which can be converted into a probability distribution after applying the softmax activation function for prediction), which greatly reduces the communication overhead. By transmitting the logits output by the model between the client and the server instead of the complete model parameters, it can not only reduce the amount of communication data, but also ensure the adaptability and generalization ability of the global model when processing non-IID data, which is particularly suitable for application scenarios with limited bandwidth and complex data distribution.

[0013] The present invention adopts the following technical scheme:

[0014] A federated learning communication method based on knowledge distillation includes the following steps:

[0015] Step 1: Initialize the student model parameters, teacher model parameters, and set related parameters;

[0016] Step 2: Each client performs local training, training the teacher model and the student model respectively;

[0017] Step 3: Calculate the average logits of the student model and upload it to the server;

[0018] Step 4: The server calculates the global average logits and global distillation loss, and sends the global average logits;

[0019] Step 5: Update the student model parameters and teacher model parameters.

[0020] Preferably, step 1 is as follows:

[0021] Step 1.1: Initialize student model parameters

[0022] Step 1.2: Initialize teacher model parameters

[0023] Step 1.3: Set the teacher learning rate η t , student learning rate η s , batch size B, training rounds T;

[0024] Step 1.4. Set the loss weights α, β, γ, δ.

[0025] Preferably, step 2 is as follows: for local data D i Each batch B in

[0026] Step 2.1, the teacher model calculates the output logits: in, represents the teacher model function, x represents the input data, represents the teacher model parameters;

[0027] Step 2.2, the student model calculates the output logits: in, represents the student model function, represents the student model parameters;

[0028] Step 2.3, calculate the task loss:

[0029]

[0030] Among them, CE represents the cross entropy loss function (Cross Entropy Loss), and y represents the label;

[0031] Step 2.4, calculate the distillation loss:

[0032]

[0033] Where KL represents the KL divergence function, σ represents the softmax function, and τ represents the temperature coefficient;

[0034] Step 2.5, calculate the hidden layer loss:

[0035]

[0036] in, represents the hidden layer loss function of the student model, represents the hidden layer loss function of the teacher model;

[0037] Step 2.6, calculate the total loss:

[0038] in, Represents the local knowledge distillation loss value;

[0039] Step 2.7, Update teacher model parameters:

[0040]

[0041] in, represents the gradient of the teacher model parameters;

[0042] Step 2.8, update the student model parameters:

[0043]

[0044] in, represents the gradient of the student model parameters.

[0045] Preferably, step 3 is as follows: for each category c∈{1,2,...,C}, calculate the average logits of the category:

[0046]

[0047] in, Indicates the total number of samples corresponding to each category;

[0048] After training, the client uploads the logits of the student model to the server.

[0049] Preferably, step 4 is as follows:

[0050] The server calculates the global average logits: Send global average logits;

[0051] Where N represents the number of clients;

[0052] Calculate the global distillation loss:

[0053] Where C represents the number of categories.

[0054] Preferably, step 5 is as follows:

[0055] Update the student model parameters:

[0056] Update the teacher model parameters:

[0057] The present invention also discloses a federated learning communication system based on knowledge distillation, which is used to execute the above method, and includes the following modules:

[0058] Initialization module: initialize student model parameters, teacher model parameters, and set related parameters;

[0059] Training module: Each client performs local training, training the teacher model and the student model separately;

[0060] Average logits calculation module: calculates the average logits of the student model and uploads it to the server;

[0061] Global average logits and distillation loss calculation module: The server calculates the global average logits and global distillation loss, and sends the global average logits;

[0062] Update module: Updates the student model parameters and the teacher model parameters.

[0063] Compared with the prior art, the present invention has the following significant technical advances:

[0064] (1) Logits transmission mechanism to reduce communication overhead: By transmitting the logits of the model instead of the complete parameters or gradients, the present invention significantly reduces the communication overhead. The complete model parameters usually contain a large number of weights, which have high transmission costs and large bandwidth requirements. Especially in bandwidth-constrained environments, it will prolong the training time and cause network congestion. Logits only contain the final prediction score of the model, and the amount of data is much smaller than the complete parameters. Therefore, it can reduce the time and bandwidth consumption of each round of communication while maintaining the model performance. This mechanism makes federated learning more suitable for resource-constrained application scenarios.

[0065] (2) Combining knowledge distillation with non-IID data adaptability: The present invention uses knowledge distillation to enable the teacher model to guide the student model, thereby improving the adaptability in non-IID data environments. In traditional federated learning, inconsistent client data distribution (non-IID) often leads to unstable performance of the global model. Since the client data varies greatly, the aggregated model is difficult to adapt to each data distribution. With the help of knowledge distillation, the teacher model of the present invention provides unified guidance for the student models of each client. Even if the data distribution is significantly different, the student model can maintain stable performance and faster convergence speed under the guidance of the teacher model, thereby improving the application effect in a non-IID environment.

[0066] (3) Flexibility in supporting multiple model architectures: The present invention designs a federated learning framework that can support multiple model architectures and has strong flexibility. Traditional federated learning frameworks are usually limited to a single model and are difficult to switch between different tasks, especially poor performance between image tasks (CNN) and text tasks (DNN). The present invention can flexibly support different architectures such as CNN and DNN, and dynamically adjust the model structure and parameters according to specific application scenarios to meet the needs of multiple tasks, avoiding the limitations of a single model architecture and increasing the versatility of the framework.

[0067] (4) Optimize the global aggregation strategy and improve the generalization ability of the model: The present invention significantly improves the generalization ability of the global model by optimizing the aggregation of logits uploaded by the client. In traditional federated learning, the server usually aggregates the model parameters of each client by simple average. This method will cause model instability when the data distribution is very different, making it difficult for the global model to adapt to the data of all clients. The present invention performs refined aggregation of logits of each category, so that the global model can better capture the data characteristics of different clients, thereby improving the generalization performance under diversified data.

[0068] (5) Dynamic parameter adjustment and adaptive learning mechanism: The present invention adopts an adaptive learning mechanism (a mechanism that dynamically adjusts model parameters according to the real-time environment, so that the model can adapt to different data distributions and communication environments) to dynamically adjust the learning rate and hyperparameters of the model according to the real-time data distribution and communication conditions of the client, thereby ensuring the optimal performance of the system in a dynamically changing environment. Traditional methods usually use fixed hyperparameters, which are difficult to cope with changes in the communication environment or data distribution, and may cause performance degradation. The adaptive mechanism of the present invention enables the model to flexibly respond to different data conditions and communication bandwidth limitations, while maintaining high performance, while improving the robustness and practicality of the system. BRIEF DESCRIPTION OF THE DRAWINGS

[0069] Figure 1 It is a flow chart of a federated learning communication method based on knowledge distillation in a preferred embodiment of the present invention.

[0070] Figure 2 This is a graph comparing the traffic between FedAvg and FedRD.

[0071] Figure 3 It is a block diagram of a federated learning communication system based on knowledge distillation according to a preferred embodiment of the present invention. DETAILED DESCRIPTION

[0072] The preferred embodiments of the present invention are described in detail below with reference to the accompanying drawings.

[0073] A system involved in a federated learning communication method based on knowledge distillation in this embodiment includes:

[0074] 1. Server: used to aggregate the logits uploaded by the client and update the global student model.

[0075] 2. Client devices: Client devices are distributed nodes where data is located. Each client device contains a local teacher model and a student model, which are trained on local data.

[0076] 3. Communication mechanism: The server is used to centrally summarize and update the global model. The client transmits the local learning results to the server by uploading the logits of the student model, and the server updates the local model by globally aggregating logits.

[0077] The process of a federated learning communication method based on knowledge distillation in this embodiment is as follows: Figure 1 As shown:

[0078] The meanings of the symbols involved in this embodiment are shown in Table 1.

[0079] Table 1

[0080]

[0081]

[0082] The specific steps of a federated learning communication method based on knowledge distillation in this embodiment are as follows:

[0083] Step 1: Parameter initialization:

[0084] (1) Initialize student model parameters In this step, you can use random initialization (such as Xavier or Kaiming initialization) or load pre-trained weights. Random initialization helps prevent the model from falling into symmetry problems, and pre-trained weights can increase the training convergence speed. The student model starts training with a better initial state, which can improve the efficiency of local optimization.

[0085] (2) Initialize teacher model parameters Usually comes from a pre-trained model with better performance (such as a CNN model with two fully connected layers or a ResNet18 network model). The pre-trained model has been trained on large-scale data and has better feature extraction capabilities. The teacher model can provide higher quality knowledge, which is conducive to the learning of the student model.

[0086] (3) Set the teacher learning rate η t , student learning rate η s , batch size B, training rounds T; different sizes of teacher learning rate and student learning rate can affect the model convergence speed and final performance. The temperature parameter is used to control the smoothness of the softmax output distribution, which can amplify the contribution of smaller probability categories and improve the model's sensitivity to long-tail distributions. The batch size determines the amount of data used in each training and can affect the convergence efficiency. The training rounds need to be set according to the data scale and model complexity to achieve convergence in as few rounds as possible. After tuning the hyperparameters, the model performs more robustly on local and global tasks.

[0087] (4) Set the loss weights α, β, γ, δ to balance the task loss, distillation loss, hidden layer loss, etc. respectively. The specific values ​​can be obtained through hyperparameter tuning. In multi-task learning, the importance of each loss item is adjusted through weights to prevent overfitting of a certain item, so that the model can achieve a balance between the prediction task and the knowledge distillation task.

[0088] Step 2: Local training:

[0089] Perform local training on each client to train the teacher model and the student model separately. i For each batch B:

[0090] (1) The teacher model calculates the output logits: Logits are the unnormalized probability distribution of the output of the last layer of the model in the classification task. The logits of the teacher model provide guiding soft labels for the student model.

[0091] (2) The student model calculates the output logits: The student model learns the knowledge of the teacher model through logits to alleviate the impact of label noise on training. The student model gradually improves its performance by imitating the output of the teacher model. Calculate the task loss: Cross entropy loss is used to measure the error between the predicted value and the true label to evaluate the accuracy of model prediction in supervised learning, which can ensure the predictive ability of the student model on the true label.

[0092] (3) Calculation of distillation loss Distillation loss is used to measure the difference between the student model and the teacher model. KL divergence quantifies the similarity between the two distributions, and temperature adjustment improves the gradient signal, making it easier for the student model to capture the implicit knowledge of the teacher model.

[0093] (4) Hidden layer loss: Hidden layer loss measures the difference between the hidden layer features of the teacher and student models. The hidden layer uses feature comparison to enhance the model's representation learning ability, allowing the student model to learn more comprehensive features.

[0094] (5)Total loss: Combining the above loss items, the total loss is calculated. Multi-loss optimization can effectively balance different objectives and ensure balanced development of the model in multiple tasks.

[0095] (6) Update teacher model parameters: The teacher model is updated using gradient descent so that it can adapt to the local data distribution.

[0096] (7) Update student model parameters: Use the gradient descent method to update the student model, so that the student model can gradually improve the generalization performance by optimizing the total loss.

[0097] Step 3: Upload logits:

[0098] For each category c∈{1,2,...,C}, compute the mean logits of the category (student model):

[0099]

[0100] After training, the client uploads the logits of the student model to the server. Uploading logits instead of model parameters can reduce communication costs. The server can effectively summarize knowledge and reduce the risk of data leakage.

[0101] Step 4: Server-side aggregation:

[0102] (1) Calculate the global average logits: The server aggregates the logits uploaded by the client, calculates the global average logits, and sends the global average logits. The weighted average aggregation method can retain the knowledge contribution of the client. Global logits can show better model generalization ability.

[0103] (2) Calculate the global distillation loss: By comparing the distribution similarity of client logits and global logits through KL divergence, the student model is gradually aligned to the global target. This enhances the global consistency of the student model and improves the generalization ability of the model on different client data distributions.

[0104] Step 5: Update the local student model:

[0105] (1) Update student model parameters: The student model parameters are optimized using the gradient of the global distillation loss. The back-propagation method optimizes the loss function by gradient descent and gradually adjusts the parameters of the student model. The student model gradually adapts to the global knowledge and its performance is improved on different tasks and data distributions.

[0106] (2) Update the teacher model parameters: The teacher model participates in the optimization of the global distillation loss, which can further improve its ability to represent the global data distribution and enhance its ability to guide the student model.

[0107] The above describes the preferred embodiments of the present invention in detail. Figure 2 A comparative analysis of communication volume is performed to compare the communication scheme of the present invention with the existing federated learning scheme to verify the significant advantage of the present invention in terms of communication volume.

[0108] In the FedAvg algorithm (a traditional federated learning algorithm that updates the model parameters of the client by aggregating them after each round of training), the client uploads complete model parameters or gradient updates after each round of training, and its communication volume is linearly related to the parameter size of the model. Assuming that the sum of the sizes of all model parameters is θ, the number of clients is N, and the number of training rounds is T, the communication overhead (the sum of upload and download) for each round of training is 2×TNθ (TNθ is the one-way transmission volume (upload or download)). The present invention uploads logits instead of uploading model parameters, and the size of logits is much smaller than the model parameters. Assuming that the size of logits is L and the number of classification categories of each client is C, the communication overhead for each round of training is 2×CTNL (CTNL is the one-way transmission volume (upload or download)).

[0109] Since L is usually much smaller than θ, the communication overhead of the FedRD method of the present invention is significantly smaller than that of the traditional FedAvg algorithm. Theoretically, the ratio of communication overhead is (CL is the size of the uploaded logits data, and θ is the size of the model parameters). In a typical deep learning model, assuming that the number of model parameters is in the order of tens of thousands to millions, and the size of logits is only related to the number of classification categories (such as hundreds of orders of magnitude), the communication overhead can be significantly reduced to one percent of FedAvg or even lower. Assuming that the number of categories of the classification task is small, the amount of communication per round of training can be reduced to one thousandth of the original solution or even lower. From the above analysis, it can be seen that the present invention has significant technical effects compared with the prior art.

[0110] like Figure 3 As shown, this embodiment discloses a federated learning communication system based on knowledge distillation, which is used to execute the above method, and is characterized in that it includes the following modules:

[0111] Initialization module: initialize student model parameters, teacher model parameters, and set related parameters;

[0112] Training module: Each client performs local training, training the teacher model and the student model separately;

[0113] Average logits calculation module: calculates the average logits of the student model and uploads it to the server;

[0114] Global average logits and distillation loss calculation module: The server calculates the global average logits and global distillation loss, and sends the global average logits;

[0115] Update module: Updates the student model parameters and the teacher model parameters.

[0116] For other contents of this embodiment, please refer to the above method embodiment.

[0117] The above is only a preferred embodiment of the present invention. It should be pointed out that for ordinary technicians in this technical field, several improvements and modifications can be made without departing from the technical principles of the present invention. These improvements and modifications should also be regarded as the scope of protection of the present invention.

Claims

1. A federated learning communication method based on knowledge distillation, characterized in that: The steps include: Step 1: Initialize the student model parameters, teacher model parameters, and set related parameters; Step 2: Each client performs local training, training the teacher model and the student model respectively; Step 3: Calculate the average logits of the student model and upload it to the server; Step 4: The server calculates the global average logits and global distillation loss, and sends the global average logits; Step 5: Update the student model parameters and teacher model parameters.

2. A federated learning communication method based on knowledge distillation as claimed in claim 1, characterized in that: Step 1 is as follows: Step 1.1: Initialize student model parameters Step 1.2: Initialize teacher model parameters Step 1.3: Set the teacher learning rate η t , student learning rate η s , batch size B, training rounds T; Step 1.

4. Set the loss weights α, β, γ, δ.

3. A federated learning communication method based on knowledge distillation as claimed in claim 2, characterized in that: Step 2 is as follows: For local data D i Each batch B in Step 2.1, the teacher model calculates the output logits: Among them, f i t represents the teacher model function, x represents the input data, represents the teacher model parameters; Step 2.2, the student model calculates the output logits: Among them, f i s represents the student model function, represents the student model parameters; Step 2.3, calculate the task loss: Among them, CE represents the cross entropy loss function, and y represents the label; Step 2.4, calculate the distillation loss: Among them, KL represents the divergence function, σ represents the softmax function, and τ represents the temperature coefficient; Step 2.5, calculate the hidden layer loss: in, represents the hidden layer loss function of the student model, represents the hidden layer loss function of the teacher model; Step 2.6, calculate the total loss: in, Represents the local knowledge distillation loss value; Step 2.7, Update teacher model parameters: in, represents the gradient of the teacher model parameters; Step 2.8, update the student model parameters: in, represents the gradient of the student model parameters.

4. A federated learning communication method based on knowledge distillation as claimed in claim 3, characterized in that: Step 3 is as follows: For each category c∈{1,2,...,C}, calculate the average logits of the category: in, Indicates the total number of samples corresponding to each category; After training, the client uploads the logits of the student model to the server.

5. A federated learning communication method based on knowledge distillation as claimed in claim 4, characterized in that: Step 4 is as follows: The server calculates the global average logits: Send global average logits; Where N represents the number of clients; Calculate the global distillation loss: Where C represents the number of categories.

6. A federated learning communication method based on knowledge distillation as claimed in claim 5, characterized in that: Step 5 is as follows: Update the student model parameters: Update the teacher model parameters:

7. A federated learning communication system based on knowledge distillation, used to execute the method according to any one of claims 1 to 6, characterized in that: Includes the following modules: Initialization module: initialize student model parameters, teacher model parameters, and set related parameters; Training module: Each client performs local training, training the teacher model and the student model separately; Average logits calculation module: calculates the average logits of the student model and uploads it to the server; Global average logits and distillation loss calculation module: The server calculates the global average logits and global distillation loss, and sends the global average logits; Update module: Updates the student model parameters and the teacher model parameters.

Images