A federated learning method

Abstract

The application belongs to the technical field of edge artificial intelligence, and particularly relates to a federated learning method. In the method, each client performs sparse training on a network model downloaded from a server. In each iteration of the training process, if the accuracy of the network model is greater than an accuracy threshold and the pruning rate is less than a target pruning rate, pruning is performed. For the weights set to 0, the corresponding positions in the corresponding mask are also set to 0. The iteration termination condition of the training is reached until the local model and the mask of each client are obtained and uploaded to the server. The server aggregates the local models uploaded by each client to obtain a global model. Then, according to the global model and the masks of each client, the subnetwork of the global model is obtained and is issued. In the application, the network model is pruned after sparse training by the client, so that the network model is more lightweight. Compared with the heuristic method of direct pruning, the convergence speed is accelerated, and the problem of limited client resources is better solved.

Description

Technical Field

[0001] This invention belongs to the field of edge artificial intelligence technology, specifically relating to a federated learning method. Background Technology

[0002] According to the Global Mobile Market Report, the number of smartphone users worldwide is projected to exceed four billion by 2023, generating massive amounts of data every second. Uploading this data to the cloud for processing introduces significant latency. Even for applications with low real-time requirements, the cloud faces increasing pressure as the number of devices continues to grow. Offloading computing to the edge is one way to alleviate this pressure. Edge computing provides services closer to the user or data input, and performing some tasks at the edge can better protect user data privacy.

[0003] Many studies have achieved significant breakthroughs by applying machine learning to other fields. Edge devices, being closer to users, can collect vast amounts of data, and machine learning can extract useful information from that data. Using machine learning theories and techniques in edge computing can provide users with better services. Federated learning can jointly train neural network models on both the client and server sides while protecting data security.

[0004] However, training models on the client side with the assistance of the server presents certain challenges. Neural network models are large in scale and require multiple iterations during training, while edge devices have limited resources and computing power, potentially making it difficult to bear the training and communication overhead. To reduce costs, many studies have attempted to lightweight models by reducing the number of parameters, the amount of communication data, or the number of communications, but this ignores the fact that the data among clients is not independently and identically distributed. Furthermore, models trained from heterogeneous data across different clients often exhibit heterogeneity, which may lead to performance degradation after aggregation on the server. Personalized federated learning is a good solution for data heterogeneity, and common methods for constructing personalized federated learning systems include meta-learning, transfer learning, and adaptive adjustment, but none of these methods consider the limitations of client-side computing and communication capabilities.

[0005] The methods mentioned above for reducing model overhead will inevitably compromise model performance. Methods for heterogeneous data require additional parameters to represent the data, further increasing model overhead. While pruning strategies in federated learning can remove redundant parameters and train personalized models for clients, the commonly used pruning strategies are heuristic and do not consider the dynamic nature of weights during training, leading to significant accuracy loss. This requires additional retraining to recover, making this approach still relatively costly. Summary of the Invention

[0006] The purpose of this invention is to provide a federated learning method to solve the problem of model accuracy loss caused by using heuristic pruning strategies.

[0007] To address the aforementioned technical problems, this invention provides a federated learning method, comprising the following steps:

[0008] 1) Each client performs sparse training on the network model downloaded from the server. During each sparse training iteration, the accuracy and pruning rate of the network model are judged. If the accuracy of the network model is greater than the accuracy threshold and the pruning rate is less than the target pruning rate, some weights in the network model that are close to 0 are reset to 0 to prune. For the weights that are set to 0, the corresponding positions in the corresponding mask are also set to 0 until the iteration termination condition of sparse training is reached, thereby obtaining the local model and mask of each client and uploading them to the server.

[0009] 2) The server aggregates the models uploaded to the local machine by each client to obtain the global model; based on the global model and the mask of each client, it obtains the subnet of the global model and sends it to each client.

[0010] The beneficial effects of the above technical solution are as follows: In this invention, the client performs sparse training on the network model before pruning, making the network model more lightweight. Compared with the heuristic approach of direct pruning, this accelerates the convergence speed, better solves the problem of limited client resources, and ultimately enables the network model to run better on the client. Moreover, compared with methods without sparse training, pruning has less impact on model accuracy. Furthermore, the network model accuracy and pruning rate are judged before pruning, and pruning is only performed when both meet the requirements, preventing permanent damage to the model's accuracy.

[0011] Furthermore, when performing local model aggregation, only the unpruned parts of the local model are aggregated.

[0012] The beneficial effect of the above technical solution is that it can reduce the impact of data heterogeneity on model performance by aggregating only the unpruned parts.

[0013] Furthermore, the average value of the corresponding positions of the unpruned parts in each local model is taken to achieve aggregation.

[0014] Further, in step 2), the global model is multiplied by the corresponding positions of the masks of each client to obtain the subnet of each client.

[0015] Furthermore, the loss function is transformed into an optimization problem with sparsity requirements, and the ADMM algorithm is used to solve the optimization objective in order to achieve sparsified training of the network model.

[0016] The beneficial effects of the above technical solution are: ADMM can solve non-convex optimization problems under sparse constraints very well, therefore ADMM is used for solving. Attached Figure Description

[0017] Figure 1 This is a diagram of the personalized federated learning process of ADMM weight pruning in this invention;

[0018] Figure 2 This is a flowchart of the ADMM weight pruning process of this invention;

[0019] Figure 3 This is a schematic diagram of the aggregated local model of the present invention;

[0020] Figure 4 This is a schematic diagram of the subnet acquisition method of the present invention;

[0021] Figure 5(a) shows the accuracy obtained when different target pruning rates are set for CIFAR-10;

[0022] Figure 5(b) shows the accuracy obtained when different target pruning rates are set for MNIST;

[0023] Figure 6(a) shows the accuracy diagrams obtained using different algorithms for CIFAR-10;

[0024] Figure 6(b) shows the accuracy graphs obtained for MNIST using different algorithms;

[0025] Figure 7(a) is a comparison of communication costs for CIFAR-10 using different algorithms;

[0026] Figure 7(b) is a comparison of communication costs for MNIST using different algorithms;

[0027] Figure 8(a) is a diagram of the main interface of the system platform;

[0028] Figure 8(b) is a diagram of the flower recognition function interface of the system platform;

[0029] Figure 9 This is the interface diagram for selecting the model;

[0030] Figure 10 It is to select the image interface diagram;

[0031] Figure 11(a) and Figure 11(b) show the flower recognition results of the server and client models for the same image, respectively. Detailed Implementation

[0032] To make the objectives, technical solutions, and advantages of the present invention clearer, the present invention will be further described in detail below with reference to the accompanying drawings and embodiments.

[0033] Method Implementation Examples:

[0034] Assume there is one server and n clients in the network. Client 1, Client 2, and Client 3 are any three clients communicating with the server. The server and these three clients collaboratively train the model. This embodiment uses the server, Client 1, Client 2, and Client 3 as an example to provide a detailed explanation of the model training process. The entire learning process is as follows: Figure 1 As shown, the process mainly includes three steps: ADMM weight pruning, local model aggregation, and subnet acquisition.

[0035] Step one: Each client performs ADMM sparsity training and pruning operations on the network model downloaded from the server, obtaining its local model and mask, and then uploads it to the server. For example... Figure 2 As shown, specifically:

[0036] The network model is trained using sparsity optimization. The model's loss function is transformed into an optimization problem with sparsity requirements. The network model (in this example, a neural network) is then trained using sparsity optimization, as shown below:

[0037]

[0038] The above formula represents the desired loss function L for client i under the constraint of sparse model weights. i Minimum. Among them, W i The neural network weights of client i are represented by W. g ⊙M i Received; W g M represents the parameters of the global model in the server. i S represents the mask for client i; i ={W i |card(W i )≤N i} is a non-convex set, card(W) i Returns the number of non-zero elements in the neural network of client i, N. i This represents the expected number of non-zero elements.

[0039] Then, Equation (1) is solved. Equation (1) is a nonconvex optimization problem under sparse constraints. ADMM is a powerful algorithm for solving this type of problem. Therefore, ADMM is used to solve the optimization objective, i.e., Equation (1). ADMM combines the advantages of dual ascent and augmented Lagrange algorithms, decomposes the original problem into several subproblems, and solves the variables alternately. It is widely used in large-scale distributed learning and can bring good results. In order to minimize the objective function, S is first introduced. i The indicator function g(W) i ):

[0040]

[0041] Add the indicator function to formula (1) to transform the original inequality constraint into the equality constraint in formula (3).

[0042]

[0043] Formula (3) is consistent with the standard form of ADMM, and its augmented Lagrangian function is:

[0044]

[0045] Among them, Y T ρ is a Lagrange multiplier, and ρ is the penalty parameter. Let Then equation (4) can be transformed into the following equation:

[0046]

[0047] The ADMM algorithm iterates multiple times using the following formula to solve:

[0048]

[0049]

[0050]

[0051] Solve for W using the stochastic gradient descent method. i :

[0052]

[0053] The solution to formula (7) can be expressed as:

[0054]

[0055] in, S represents i The Euclidean projection on the surface. (For) Analytical solution, preserving The first N i Large parameters remain unchanged, while the rest are set to zero. Update U using formula (8). i This completes one iteration of ADMM.

[0056] In practice, the loss function can be set as follows: The training is sparse, and W, Z, and U are dynamically updated after each training round. W is updated to reflect the model weights after this iteration. The top N values ​​in W+U are maintained. i The major parameters remain unchanged, while the others are set to zero, and then updated to Z; U is updated to U+WZ.

[0057] In each iteration process, the accuracy acc and pruning rate prune of the network model are calculated. If the accuracy acc of the network model is greater than the accuracy threshold acc0 and the pruning rate prune is less than the target pruning rate prune0, that is, acc > acc0 and prune < prune0, then some weights approaching 0 in the network model are set to 0 for pruning. Repeat until the iteration termination condition for sparse training is reached, determine the corresponding mask, and obtain the local models of each server. Specifically, in each iteration process, it is necessary to judge the accuracy and pruning rate of the network model: if acc > acc0 and prune < prune0, it indicates that pruning can be performed to achieve the purpose of model lightweighting, and continue the next iteration cycle for training after pruning; if acc > acc0 and prune ≥ prune0, it indicates that pruning cannot continue, and directly execute the next iteration cycle for training at this time; if acc ≤ acc0 and prune < prune0, it indicates that the model accuracy has not yet reached the requirement, and directly execute the next iteration cycle for training at this time; if acc ≤ acc0 and prune ≥ prune0, it indicates that the model accuracy has not yet reached the requirement, and directly execute the next iteration cycle for training at this time. That is to say, the entire training process still takes the required number of iterations as the iteration termination condition, but pruning can be performed if the situation of acc > acc0 and prune < prune0 is found during the training process.

[0058] Among them, in the sparse training of the model, a large number of iteration times are required to prune p% of the parameters in the model (p% is the iterative pruning rate), which does not meet the lightweight goal we proposed. Therefore, when the client reaches the predetermined number of iterations, the training stops. At this time, the weights that need to be pruned may not be completely 0, and many are approaching 0. To meet the sparsity requirement, prune p% of the parameters in the model (after sparse training, these p% of the parameters are already close to 0) and prune the weights close to 0. Since these weights are very small themselves, compared with the method without sparse training, pruning has less impact on the model accuracy. Before pruning, it is necessary to judge whether the current accuracy and pruning rate meet the requirements to prevent permanent damage to the model accuracy. If there are parameters set to 0 in the neural network model of the client, the corresponding values in the mask are also set to 0, and the weights with a value of 0 will be frozen and not updated in subsequent training, thereby reducing the calculation and enabling it to adapt to small edge devices.

[0059] Step 2, the server aggregates the local models uploaded by each client to obtain the global model.

[0060] After performing ADMM weight pruning on the client side, the model retaining its personalized features is uploaded to the server. Models trained on datasets with similar labels also exhibit similarity. Since the model is pruned on the client side, the traditional method of averaging all parameters is no longer suitable for this invention. This invention only aggregates the unpruned portions, averaging the partially or fully overlapping parameters across models. If a parameter is pruned in all clients except one, it is excluded from aggregation, reducing the impact of heterogeneous data on accuracy. This aggregation method, by finding "partners" for the client side, preserves the personalized attributes of the embedded models, improving model performance.

[0061] For example, such as Figure 3 As shown, the first parameter of the first layer of the local model of client 1, client 2 and client 3 is 3, 0 and 2 respectively, and the average value obtained after aggregating the non-zero values ​​is 5 / 2.

[0062] Step 3: The server obtains the subnet of the global model based on the global model and the mask of each client, and sends it to each client.

[0063] After local training, the client uses a mask to mark the model pruning status and uploads the mask to the server. The mask embeds information from the local data and is used to distinguish the models of different clients on the server. The global model is multiplied by the corresponding positions of the mask to obtain a subnet, which is then downloaded to the client. Simply put, obtaining a subnet involves the client selecting a portion of the network from the global model based on the characteristics of its local data. After randomly selecting a subset S of clients, each client in set S receives a different subnet from the server, thus achieving personalization of federated learning.

[0064] For example, such as Figure 4 As shown, the parameters and mask corresponding to the first neuron in the first layer of the subnet are 3 and 1, respectively, determined by W. g ⊙M1 calculates the corresponding parameter in the subnet to be 3; similarly, the third parameter of the second layer of the subnet is 0.

[0065] This completes the federated learning method of this invention. As can be seen from the above description of model training, it mainly utilizes ADMM weight pruning to remove redundant parameters in the network and builds a proprietary model for the client based on a mask to achieve lightweight and personalized models. This method addresses the problems of high model cost and data heterogeneity in federated learning, which make it difficult to run models on edge devices.

[0066] To better illustrate the effectiveness of the method of this invention, the dynamic pruning algorithm implemented in this invention is compared with several existing algorithms on the CIFAR-10 and MNIST datasets, respectively. Furthermore, a flower recognition system is used to further demonstrate the practicality of this invention. The experiments were run on a Windows 10 64-bit operating system, using PyTorch 1.3.0 and Android Studio 3.4.1 as the experimental platform.

[0067] CIFAR-10 is a color dataset containing 10 image classes, totaling 60,000 images, with 50,000 for training and 10,000 for testing. MINST is a set of handwritten grayscale images from 0 to 9, including 60,000 training images and 10,000 test images. For the CIFAR-10 dataset, a six-layer neural network architecture was used for training, consisting of five hidden layers and one output layer, specifically: two 5×5 convolutional layers followed by three fully connected layers with 120, 84, and 10 channels respectively, and finally the output layer. For the MNIST dataset, a five-layer neural network architecture was used for training, consisting of four hidden layers and one output layer, specifically: two 5×5 convolutional layers followed by two fully connected layers with 50 and 10 channels respectively, and finally the output layer. The experimental parameter settings for the CIFAR-10 and MNIST datasets are shown in Table 1.

[0068] Table 1 CIFAR-10 and MNIST parameter settings

[0069]

[0070] In the communication between the server and client, parameters were iteratively pruned by 10%. The target pruning rates were set to 10%, 20%, 30%, 40%, 50%, 60%, 70%, 80%, and 90%, respectively. The resulting CIFAR-10 and MNIST accuracies are shown in Figures 5(a) and 5(b), respectively. In the CIFAR-10 dataset, the highest accuracy was achieved with a target pruning rate of 30%, while in MNIST, the highest accuracy was achieved with 10%. After reaching the maximum accuracy, the model showed a downward trend due to the pruning of some important parameters. The communication volume of the model also varied under different target pruning rates. In CIFAR-10, a target pruning rate of 30% resulted in higher accuracy than 90%, but correspondingly, more communication was required; the same was true for MNIST.

[0071] Based on the experimental results in Figures 5(a) and 5(b), the target pruning rates for CIFAR-10 were set to 30%, 50%, and 70%, respectively, and for MNIST to 10%, 50%, and 90%, respectively. The method proposed in this invention was compared with Standalone, FedAvg, and LG-FedAvg under different target pruning rates. The following is an introduction to these three algorithms.

[0072] Standalone: ​​The client trains the model using data on its local device without the involvement of a server.

[0073] FedAvg: A classic federated averaging algorithm. The server distributes the global model to each client, which then trains the model locally and uploads it back to the server for aggregation.

[0074] LG-FedAvg: Learns a useful high-dimensional representation from the local data of each client, and the global model operates on this high-dimensional representation, thereby reducing communication and making the model lightweight.

[0075] Figures 6(a) and 6(b) show the accuracy comparison of six sets of experiments conducted on two datasets, and Figures 7(a) and 7(b) show the communication cost comparison (since Standalone does not involve a server, there is no communication cost). From Figures 6(a), 6(b), 7(a), and 7(b), it can be observed that the lightweight personalized federated learning proposed in this invention outperforms Standalone, FedAvg, and LG-FedAvg in both accuracy and communication volume. For the CIFAR-10 dataset, the highest accuracy (90.14%) is achieved when the target pruning rate is 30%, at which point the communication volume is 775.95 MB. LG-FedAvg has an accuracy of 80.78% and a communication volume of 1127.89 MB. The method of this invention improves the accuracy by 9.36% compared to LG-FedAvg, while reducing the communication volume by 1.45 times. For the MNIST dataset, the highest accuracy (99.96%) was achieved when the target pruning rate was 10%, with a communication volume of 212.45 MB. The accuracy of LG-FedAvg was 97.03%, with a communication volume of 244.82 MB, representing an improvement of 2.93% in accuracy and a reduction of 1.15 times in communication volume.

[0076] Furthermore, the effectiveness of this invention was further demonstrated by training the model using a flower dataset created by Oxford University and MobileNetV2, and then deploying it on an Android device. The dataset includes 17 flower categories, with 80 images per category, totaling 1360 images. MobileNetV2 is a lightweight convolutional neural network based on an inverted residual structure, proposed by the Google team in 2018. The experimental parameters for the flower dataset are shown in Table 2. The final model accuracy was 92.47%, and the communication volume per client was 0.8783 GB.

[0077] Table 2. Parameter settings for the flower dataset

[0078]

[0079] The trained model is deployed to the local recognition system through the following three steps. Figure 8(a) shows the main interface of the system. Clicking "PYTORCH Flower Recognition" will display the flower recognition function interface as shown in Figure 8(b). (1) Convert the .pth model file obtained in PyCharm into a .pt file supported by Android Studio. (2) Add the PyTorch dependency to the build.gradle file of Android Studio, and put the converted model file and label information into Android Studio. (3) Use the model for inference.

[0080] This invention deploys a global model on the server and a personalized model on a client, both of which include the flowers Daisy and Bluebell. For example... Figure 9 As shown, both server and client models are available. Click "Album" to select the image to be recognized from the mobile device, such as... Figure 10 As shown in Figure 11(a) and Figure 11(b), these figures compare the flower recognition results of the server and client models for the same image. It is clear that the personalized model of the client is more accurate than the global model of the server.

[0081] In summary, the algorithm proposed in this invention can better adapt to edge devices, training personalized models for clients while reducing model overhead. This invention uses ADMM-based weight pruning, performing sparse training on the model before pruning, making it lightweight. Compared to heuristic methods of direct pruning, it accelerates convergence and better addresses the problem of limited client resources. A personalized model is built for each client using a mask; the server downloads a subnet of the global model to the client and only aggregates the undipped portions, reducing the impact of data heterogeneity on model performance. This invention addresses the problems of high model cost and data heterogeneity by training lightweight personalized models for clients, enabling them to run better on client devices.

Claims

1. A federated learning method, characterized in that, include: 1) Each client performs sparse training on the network model downloaded from the server. During each sparse training iteration, the accuracy and pruning rate of the network model are judged. If the accuracy of the network model is greater than the accuracy threshold and the pruning rate is less than the target pruning rate, some weights in the network model that are close to 0 are reset to 0 to prune. For the weights that are set to 0, the corresponding positions in the corresponding mask are also set to 0 until the iteration termination condition of sparse training is reached, thereby obtaining the local model and mask of each client and uploading them to the server. During sparsity training, the model's loss function is... This can be transformed into the following optimization problem with sparsity requirements: (1) For the client Network weights, by get; These are the parameters of the global model in the server; For the client The mask; , Return to client The number of non-zero elements in the network, The desired number of non-zero elements; To solve equation (1) using ADMM, first introduce... Indicator function : (2) Will Adding this to equation (2) transforms it into the following equality constraint: (3) The augmented Lagrangian function of equation (3) is: (4) It is a Lagrange multiplier. It is a penalty parameter; let Then equation (4) becomes: (5) ADMM iterates multiple times using the following formula: (6) (7) (8) Solve using stochastic gradient descent. : (9) The solution to equation (7) is expressed as: (10) for Euclidean projection on; for Analytical solution, preserving forward The major parameters remain unchanged, while the rest are set to zero; update using equation (8). This completes one iteration of ADMM; In practice, the loss function is as follows: Sparsity training, after each round of training , and Dynamically updated in China Updated to the model weights trained in this iteration; maintain Center front The major parameters remain unchanged, while the others are set to zero, and then updated to... ; Updated to ;2) The server aggregates the models uploaded to the local machine by each client to obtain the global model; based on the global model and the mask of each client, the subnet of the global model is obtained and sent to each client.

2. The federated learning method according to claim 1, characterized in that, When performing local model aggregation, only the unpruned parts of the local model are aggregated.

3. The federated learning method according to claim 2, characterized in that, Aggregation is achieved by averaging the corresponding positions of the unpruned parts in each local model.

4. The federated learning method according to claim 1, characterized in that, In step 2), the global model is multiplied by the corresponding positions of the masks of each client to obtain the subnet of each client.

5. The federated learning method according to any one of claims 1 to 4, characterized in that, In step 1), training stops when the client reaches the predetermined number of iterations, and p% of the parameters in the model are pruned, where p% is the iteration pruning rate.

Images