A Federated Learning Communication Method and System Based on Knowledge Distillation

Abstract

This invention discloses a federated learning communication method and system based on knowledge distillation. The method includes the following steps: Step 1, initializing student model parameters and teacher model parameters, and setting relevant parameters; Step 2, each client performs local training, training both the teacher model and the student model separately; Step 3, calculating the average logits of the student model and uploading it to the server; Step 4, the server calculates the global average logits and the global distillation loss, and distributes the global average logits; Step 5, updating the student model parameters and teacher model parameters. This invention can reduce the amount of communication data while ensuring the adaptability and generalization ability of the global model when processing non-IID data.

Description

Technical Field

[0001] This invention belongs to the field of machine learning technology, specifically relating to a federated learning communication method and system based on knowledge distillation. Background Technology

[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 (e.g., mobile phone, IoT device) to independently train a model on local data while ensuring that the data does not leave the local device, and then upload the updated model to a central server for merging. Compared to traditional centralized deep learning methods, Federated Learning has two significant advantages:

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

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

[0005] Furthermore, federated learning can leverage the computing power of distributed devices to alleviate the burden on the central server and improve overall computational efficiency. However, the implementation of federated learning is accompanied by significant communication overhead, especially as model complexity and the number of clients increase, making communication costs a major bottleneck to system efficiency. The communication overhead problem in federated learning is primarily manifested in the fact that after each round of model training, the client needs to upload its local model parameters or gradient updates (the gradient is the partial derivative of the model error with respect to the parameters, used to optimize model parameters and bring the loss function to convergence) to the server. For complex deep neural networks (a model structure with multiple layers of neurons, widely used to solve complex pattern recognition problems), which may contain millions or even hundreds of millions of parameters, uploading this data in each training round will place a significant burden on network bandwidth and device resources. The communication overhead problem is particularly prominent, specifically manifested in the following aspects:

[0006] (1) Bandwidth limitation: In real-world applications, 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, which leads to a significant increase in the time required for each round of communication, thereby prolonging 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 great deal of resources. Frequent communication not only leads to a surge in energy consumption but may also affect the normal operation of the devices.

[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 uploaded frequently still poses a risk of attack or leakage. Malicious attackers may infer certain characteristics of local data by analyzing the transmitted model parameters, thereby threatening user privacy.

[0009] (4) Challenges posed by Non-Identical Data: In practical applications, the data distribution of different clients is often non-identical (i.e., the data distribution of clients is different, and the data of each client may be biased, leading to an adaptation problem 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 will further increase the communication overhead.

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

[0011] Researchers have explored various methods to reduce communication overhead in federated learning, including gradient compression, quantization, and sparsification. However, these methods have limitations in high-frequency, large-scale, and multi-client applications. Especially when dealing with complex environments containing non-IID data, the trade-off between communication overhead and global model performance becomes a critical challenge. Therefore, designing an efficient federated learning communication mechanism that can significantly reduce communication volume while maintaining model performance is key to solving this problem. Summary of the Invention

[0012] To address the aforementioned issues, this invention proposes a federated learning communication method and system based on knowledge distillation. This invention improves the federated learning framework (by training a student model and using a teacher model to guide the student model's learning, enabling the student model to approach the teacher model's performance with fewer parameters). It replaces the traditional model parameter uploading method with a logits (the raw scores output by the model, which can be converted into a probability distribution for prediction after applying the softmax activation function) transmission mechanism, significantly reducing communication overhead. By transmitting the model's output logits, rather than the complete model parameters, between the client and server, it reduces the amount of communication data while ensuring the global model's adaptability and generalization ability when processing non-IID data, making it particularly suitable for applications with limited bandwidth and complex data distributions.

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

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

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

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

[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 the global distillation loss, and then distributes 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 epochs 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 and outputs logits: in, Let x represent the teacher model function, and x represent the input data. Indicates the teacher model parameters;

[0027] Step 2.2: Student model calculates and outputs logits: in, Represents the student model function. Indicates the parameters of the student model;

[0028] Step 2.3: Calculate task loss:

[0029]

[0030] Where CE represents the Cross Entropy Loss function, and y represents the label;

[0031] Step 2.4, Calculate distillation losses:

[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. This represents the hidden layer loss function of the teacher model;

[0037] Step 2.6: Calculate the total loss:

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

[0039] Step 2.7: Update teacher model parameters:

[0040]

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

[0042] Step 2.8: Update student model parameters:

[0043]

[0044] in, This 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, This represents the total number of samples for each category;

[0048] After training is complete, the client uploads the student model's logits to the server.

[0049] Preferably, step 4 is as follows:

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

[0051] Where N represents the number of clients;

[0052] Calculate global distillation losses:

[0053] Where C represents the number of categories.

[0054] Preferably, step 5 is as follows:

[0055] Update student model parameters:

[0056] Update teacher model parameters:

[0057] This invention also discloses a federated learning communication system based on knowledge distillation for performing the above method, which includes the following modules:

[0058] Initialization module: Initializes student model parameters and teacher model parameters, and sets relevant parameters;

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

[0060] The 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 distributes the global average logits.

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

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

[0064] (1) Logits transmission mechanism to reduce communication overhead: By transmitting the model's logits instead of the complete parameters or gradients, this invention significantly reduces communication overhead. Complete model parameters typically contain a large number of weights, resulting in high transmission costs and bandwidth requirements. Especially in bandwidth-constrained environments, this can prolong training time and cause network congestion. Logits, on the other hand, only contain the model's final prediction scores, with a much smaller data volume than complete parameters. Therefore, it reduces the time and bandwidth consumption per round of communication while maintaining model performance. This mechanism makes federated learning more suitable for resource-constrained applications.

[0065] (2) Combining Knowledge Distillation with Adaptability to Non-IID Data: This invention improves adaptability in non-IID data environments by using knowledge distillation to enable the teacher model to guide the student model. In traditional federated learning, inconsistent client data distributions (non-IID) often lead to unstable global model performance because the large differences in client data make it difficult for the aggregated model to adapt to various data distributions. This invention utilizes knowledge distillation, where the teacher model provides unified guidance to the student models on each client. Even with significant differences in data distribution, the student models can maintain stable performance and faster convergence speed under the guidance of the teacher model, thereby improving the application effect in non-IID environments.

[0066] (3) Flexibility in supporting multiple model architectures: This invention designs a federated learning framework that can support multiple model architectures, exhibiting strong flexibility. Traditional federated learning frameworks are usually limited to a single model, making it difficult to switch between different tasks, especially performing poorly between image tasks (CNN) and text tasks (DNN). In contrast, this 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 framework's versatility.

[0067] (4) Optimizing the global aggregation strategy to improve model generalization ability: This invention significantly improves the generalization ability of the global model by optimizing and aggregating the logits uploaded by clients. In traditional federated learning, the server usually aggregates the model parameters of each client by simple averaging. This method can cause model instability when the data distribution varies greatly, making it difficult for the global model to adapt to the data of all clients. This invention performs refined aggregation of logits for each category, enabling the global model to better capture the data characteristics of different clients, thereby improving the generalization performance under diverse data.

[0068] (5) Dynamic Parameter Adjustment and Adaptive Learning Mechanism: This invention employs an adaptive learning mechanism (a mechanism that dynamically adjusts model parameters according to the real-time environment, enabling the model to adapt to different data distributions and communication environments). The learning rate and hyperparameters of the model are dynamically adjusted based on the client's real-time data distribution and communication conditions, thereby ensuring optimal system performance in dynamically changing environments. Traditional methods typically use fixed hyperparameters, which are difficult to cope with changes in communication environments or data distributions, potentially leading to performance degradation. The adaptive mechanism of this invention enables the model to flexibly respond to different data conditions and communication bandwidth limitations, improving the system's robustness and practicality while maintaining high performance. Attached Figure Description

[0069] Figure 1 This is a flowchart of a federated learning communication method based on knowledge distillation, according to a preferred embodiment of the present invention.

[0070] Figure 2 This is a graph comparing the communication volume of FedAvg and FedRD.

[0071] Figure 3 This is a block diagram of a federated learning communication system based on knowledge distillation, according to a preferred embodiment of the present invention. Detailed Implementation

[0072] The preferred embodiments of the present invention will now be described in detail with reference to the accompanying drawings.

[0073] This embodiment of a federated learning communication method based on knowledge distillation involves a system including:

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

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

[0076] 3. Communication Mechanism: The server is used to centrally aggregate and update the global model. Clients transmit their local learning results to the server by uploading the logits of their student models, and the server updates its local model by globally aggregating the logits.

[0077] This embodiment presents a federated learning communication method based on knowledge distillation, as follows: Figure 1 As shown:

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

[0079] Table 1

[0080]

[0081]

[0082] The specific steps of the 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, random initialization (such as Xavier or Kaiming initialization) or loading pre-trained weights can be used. Random initialization helps prevent the model from getting stuck in symmetry problems, while pre-trained weights can improve training convergence speed. Starting training with a student model in a better initial state can improve local optimization efficiency.

[0085] (2) Initialize teacher model parameters Typically, this comes from a high-performance pre-trained model (such as a CNN model with two fully connected layers or a ResNet18 network model). Pre-trained models have been trained on large-scale data and possess good feature extraction capabilities. Teacher models provide higher-quality knowledge, which is beneficial for student model learning.

[0086] (3) Set the teacher learning rate η t Student learning rate η s Batch size (B) and training epochs (T) are key parameters. Different teacher and student learning rates can affect the model's convergence speed and final performance. The temperature parameter controls the smoothness of the softmax output distribution, amplifying the contribution of lower probability classes and increasing the model's sensitivity to long-tailed distributions. Batch size determines the amount of data used in each training iteration, affecting convergence efficiency. The number of training epochs needs to be set based on the data scale and model complexity, aiming to achieve convergence with as few epochs as possible. After optimizing the hyperparameters, the model becomes more robust in both local and global tasks.

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

[0088] Step 2: Local training:

[0089] Local training is performed on each client, training both the teacher and student models separately. For the local data D... i Each batch B in:

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

[0091] (2) Student model calculates output logits: The student model learns from the teacher model's knowledge through logits, mitigating the impact of label noise on training. The student model gradually improves performance by imitating the teacher model's output. Task loss is calculated as follows: Cross-entropy loss is used to measure the error between the predicted value and the true label. It is used to evaluate the accuracy of the model's predictions in supervised learning and can ensure the student model's ability to predict the true label.

[0092] (3) Calculate distillation loss Distillation loss is used to measure the difference between the student model and the teacher model. KL divergence measures the similarity between the two distributions, and temperature adjustment boosts 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. Hidden layers use feature contrast to enhance the model's representation learning ability, allowing the student model to learn more comprehensive features.

[0094] (5) Total loss: Taking into account the above loss terms, the total loss is calculated. Multi-loss optimization can effectively balance different objectives and ensure that the model develops in a balanced manner across multiple tasks.

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

[0096] (7) Update student model parameters: The student model is updated using gradient descent, allowing it to gradually improve its generalization performance by optimizing the total loss.

[0097] Step 3: Upload logits:

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

[0099]

[0100] After training, the client uploads the student model's logits to the server. Uploading logits instead of model parameters reduces communication costs. The server can effectively aggregate 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 clients, calculates the global average logits, and distributes the global average logits. This weighted average aggregation method preserves the client's knowledge contribution. The global logits demonstrate better model generalization ability.

[0103] (2) Calculate the global distillation loss: By comparing the distribution similarity of client-side logits and global logits using KL divergence, the student model is gradually aligned with the global target. This enhances the global consistency of the student model and improves its generalization ability across 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 backpropagation method optimizes the loss function through gradient descent, gradually adjusting the student model parameters. The student model progressively adapts to global knowledge, improving performance across different tasks and data distributions.

[0106] (2) Update teacher model parameters: The teacher model's participation in the optimization of the global distillation loss can further enhance its ability to represent the global data distribution, thereby strengthening its ability to guide the student model.

[0107] The preferred embodiments of the present invention have been described in detail above. The following is a description in conjunction with... Figure 2 A comparative analysis of communication volume was conducted to compare the communication scheme of the present invention with existing federated learning schemes in order to verify the significant advantages of the present invention in terms of communication volume.

[0108] In the FedAvg algorithm (a traditional federated learning algorithm that aggregates model parameter updates from clients after each training round), the client uploads complete model parameters or gradient updates after each training round, and its communication overhead is linearly related to the size of the model parameters. Assuming the sum 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) per training round is 2 × TNθ (TNθ is the one-way transmission volume (upload or download)). This invention replaces uploading model parameters with uploading logits, whose size is much smaller than the model parameters. Assuming the size of logits is L, and the number of classification categories per client is C, the communication overhead per training round is 2 × CTNL (CTNL is the one-way transmission volume (upload or download)).

[0109] Since L is typically much smaller than θ, the communication overhead of the FedRD method in this invention is significantly less than that of the traditional FedAvg algorithm. Theoretically, the communication overhead ratio is... (CL represents the size of the uploaded logits data, and θ represents the size of the model parameters). In typical deep learning models, assuming the number of model parameters is in the tens of thousands to millions, and the size of logits is only related to the number of classification categories (e.g., in the hundreds), the communication overhead can be significantly reduced to one percent or even lower than that of FedAvg. Assuming the classification task has a small number of categories, the communication volume per training round can be reduced to one-thousandth or even lower than that of the original scheme. In summary, this invention has significant technical advantages compared to existing technologies.

[0110] like Figure 3 As shown, this embodiment discloses a federated learning communication system based on knowledge distillation for executing the above-described method, characterized by comprising the following modules:

[0111] Initialization module: Initializes student model parameters and teacher model parameters, and sets relevant parameters;

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

[0113] The 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 distributes the global average logits.

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

[0116] Other aspects of this embodiment can be found in the above method embodiments.

[0117] The above description is only a preferred embodiment of the present invention. It should be noted that for those skilled in the art, several improvements and modifications can be made without departing from the technical principles of the present invention, and these improvements and modifications should also be considered within the scope of protection of the present invention.

Claims

1. A federated learning communication method based on knowledge distillation, characterized in that, Includes the following steps: Step 1: Initialize student model parameters and teacher model parameters, and set relevant parameters; Step 2: Each client performs local training, training both the teacher model and the student model separately; 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 the global distillation loss, and then distributes the global average logits. Step 5: Update the student model parameters and teacher model parameters; 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 Student learning rate Batch size B, training epochs T; Step 1.4: Set loss weights ; Step 2 is as follows: For local data D i Each batch B in; Step 2.1: The teacher model calculates and outputs logits: in, Let x represent the teacher model function, and x represent the input data. Indicates the teacher model parameters; Step 2.2: Student model calculates and outputs logits: in, Represents the student model function. Indicates the parameters of the student model; Step 2.3: Calculate task loss: Where CE represents the cross-entropy loss function, and y represents the label; Step 2.4, Calculate distillation losses: Where KL represents the divergence function, This represents the softmax function. Indicates the temperature coefficient; Step 2.5: Calculate the hidden layer loss: in, This represents the hidden layer loss function of the student model. This represents the hidden layer loss function of the teacher model; Step 2.6: Calculate the total loss: in, This represents the local knowledge distillation loss value; Step 2.7: Update teacher model parameters: in, The gradient represents the parameters of the teacher model; Step 2.8: Update student model parameters: in, Represents the gradient of the student model parameters; Step 3 is as follows: For each category Calculate the average logits for each category: in, This represents the total number of samples for each category; After training, the client uploads the student model's logits to the server; Step 4 is as follows: The server calculates the global average logits: Distribute the global average logits; Where N represents the number of clients; Calculate global distillation losses: Where C represents the number of categories; Step 5 is as follows: Update student model parameters: Update teacher model parameters: 。 2. A federated learning communication system based on knowledge distillation, for performing the method as described in claim 1, characterized in that, Includes the following modules: Initialization module: Initializes student model parameters and teacher model parameters, and sets relevant parameters; Training module: Each client performs local training, training the teacher model and the student model separately; The 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 distributes the global average logits. Update module: Update student model parameters and teacher model parameters.

Images