Federal learning framework optimization method and system for non-independent identically distributed data

By adopting layer-by-layer aggregation strategy and zero-order relaxation backpropagation algorithm in federated learning, the problem of model performance bias and training time in zero-order federated learning methods is solved, and faster training speed and better model performance are achieved.

CN120494045APending Publication Date: 2025-08-15NAT UNIV OF DEFENSE TECH
View PDF 0 Cites 0 Cited by

Patent Information

Application Number
CN202510540289.4
Authority / Receiving Office
CN · China
Patent Type
Applications(China)
Current Assignee / Owner
Filing Date
2025-04-27
Publication Date
2025-08-15

AI Technical Summary

Technical Problem

When the existing zero-order federated learning method updates the model of the aggregate edge device, it leads to a large deviation between the global least squares solution and the mathematical optimal solution, resulting in poor model performance and too long training time.

Method used

The layer-by-layer aggregation strategy is adopted, and the client model is initialized through the parameter server. The client calculates and transmits the local least squares solution based on the zero-order relaxation backpropagation algorithm. The parameter server performs addition and aggregation to ensure that the global least squares solution approximates the mathematical optimal solution.

Benefits of technology

It speeds up training speed, improves model training efficiency, ensures global model performance optimization, reduces the computing pressure of edge devices, and achieves better model performance.

✦ Generated by Eureka AI based on patent content.

Smart Images

  • Figure CN120494045A_ABST
    Figure CN120494045A_ABST
Patent Text Reader

Abstract

The invention discloses a federal learning framework optimization method and system for non-independent identically distributed data, and the method comprises the steps: a client receives a random seed transmitted by a parameter server, and carries out the initialization of a local model; the client obtains non-independent identically distributed user tag data in the training data set, performs reverse mapping on the user tag data based on a zero-order relaxation back propagation algorithm, and calculates a least square solution of each layer of neural network of a local model; transmitting the least square solution of the current layer to a parameter server when the least square solution of one layer of neural network is calculated; and the client sequentially receives the global least square solution of each layer of neural network returned by the parameter server, updates the network parameters of the corresponding layer of neural network in the local model, and obtains a global optimal model after completing the updating of all layers of neural networks. According to the method, the overall training speed can be increased, the training efficiency is improved, and a global model with better performance can be obtained after training is completed.
Need to check novelty before this filing date? Find Prior Art

Description

Technical Field

[0001] The present invention relates to the field of machine learning technology, and in particular to a method and system for optimizing a federated learning framework for non-independent and identically distributed data. Background Art

[0002] Federated learning is a distributed machine learning paradigm. Its typical steps are to train the model locally on multiple edge devices and upload the parameters of the calculation results to the server. The server updates the global model based on the uploaded parameters and feeds them back to the edge devices for a new round of local calculations until the loss function converges and reaches a certain model training accuracy.

[0003] Currently, typical federated learning training schemes include first-order methods (such as stochastic gradient descent) and second-order methods. However, these methods are computationally intensive, and edge devices with limited computing power cannot perform frequent gradient calculations. Therefore, some practitioners have proposed using zero-order methods to enable edge devices to perform gradient-free calculations. However, existing zero-order training methods use finite difference techniques to approximate gradients. This gradient estimation is biased and requires a large number of sampling processes to reduce the estimation error, making the training time required by existing zero-order training methods too long.

[0004] In order to reduce the training time, some practitioners have proposed to accelerate the training process of federated learning models by calculating the least squares solution of each edge device model through zero-order relaxation backpropagation (ZORB). However, when aggregating the least squares solutions of different device models, they often simply average the model update parameters from different edge devices. This will cause the final global least squares solution to deviate significantly from the optimal solution obtained by directly calculating the model data of all edge devices through the zero-order relaxation backpropagation algorithm from a mathematical perspective, resulting in poor performance of the final global model. Summary of the Invention

[0005] The technical problem to be solved by the present invention is: In response to the above-mentioned problems in the prior art, a method and system for optimizing a federated learning framework for non-independent and identically distributed data with high training efficiency and good model performance is provided, thereby shortening the training time while making the final model have better performance.

[0006] In order to solve the above technical problems, the technical solution adopted by the present invention is: A method for optimizing a federated learning framework for non-independent and identically distributed data is disclosed. The method is applied to a distributed system consisting of a parameter server and a plurality of clients. The method comprises: Step S1: The client receives a random seed transmitted by a parameter server and uses the random seed to initialize a local model. The random seed is the same initialization model parameter sent by the parameter server to each client connected to the parameter server. The local model is a neural network model. In step S2, the client obtains non-independent and identically distributed user tag data in the training dataset, performs reverse mapping on the user tag data based on a zero-order relaxed backpropagation algorithm, and calculates the least squares solution of each layer of the neural network in the local model. Whenever the least squares solution of a layer of the neural network is calculated, the least squares solution of the current layer is transmitted to the parameter server; the non-independent and identically distributed user tag data indicates that the user tag data of each client is correlated and comes from different distributions; In step S3, the client receives the global least squares solution of each layer of the neural network returned by the parameter server in turn, and uses the received global least squares solution of each layer of the neural network to update the network parameters of the corresponding layer of the neural network in the local model. After completing the update of all layers of the neural network, the global optimal model is obtained. The global least squares solution is obtained by adding and aggregating the least squares solutions of the same layer of the neural network transmitted by each client.

[0007] Furthermore, the training of each layer of the neural network of the global optimal model is modeled as a distributed least squares problem, which is expressed as:

[0008] In the above formula, To be trained Layer neural network weight matrix, c is the client number, C is the total number of clients, Indicates the local data sample on the cth client The first step obtained by inputting the neural network corresponding to the local model The actual input of the layer, Indicates that the local tag on the cth client The inverse mapping The layer expects output labels.

[0009] Furthermore, the calculation expression of the global least squares solution is:

[0010]

[0011]

[0012] In the above formula, To be trained Layer neural network weight matrix, c is the client number, C is the total number of clients, Indicates that the local tag on the cth client The inverse mapping The layer expects the output label, The first The least squares solution of the layer, is the local loss function of the cth client The second-order partial derivative of , where the superscript T indicates the transpose.

[0013] Furthermore, in step S2, the least squares solution of a layer of neural network is calculated, which also includes calculating the second-order partial derivative. , transmit the least squares solution and the second-order partial derivatives to the parameter server together, where Indicates the local data sample on the cth client The first step obtained by inputting the neural network corresponding to the local model The actual input of the layer.

[0014] A method for optimizing a federated learning framework for non-independent and identically distributed data, comprising the following steps: The client receives a random seed transmitted by the parameter server and uses the random seed to initialize the local model. The random seed is the same initialization model parameter sent by the parameter server to each client connected to the parameter server. The local model is a neural network model, and the local model is divided into a feature extraction layer and a classification layer. The feature extraction layer is all convolutional network layers, and the classifier layer is all fully connected layers. Feature extraction layer training: The client obtains the non-independent and identically distributed user label data in the training data set, reverse maps the user label data based on the zero-order relaxed back-propagation algorithm, and calculates the least squares solution of each layer of the local model neural network. Whenever the least squares solution of a layer of neural network is calculated, the least squares solution of the current layer is transmitted to the parameter server; the client sequentially receives the global least squares solution of each layer of feature extraction layer returned by the parameter server and updates the network parameters of the neural network of the corresponding layer of the feature extraction layer of the local model until the training of the feature extraction layer is completed. The global least squares solution is obtained by the parameter server by adding and aggregating the least squares solutions of the same layer of feature extraction layer transmitted by each client; during the training of the feature extraction layer, the client transmits the local data distribution information to the parameter server; the non-independent and identically distributed user label data indicates that the user label data of each client is correlated and comes from different distributions; Classifier layer training: The client receives the grouping results from the parameter server and performs classifier training based on the grouping results, where all clients in each group jointly update a classification layer. The grouping results are obtained by the parameter server by grouping each client based on the local data distribution information of each client.

[0015] Furthermore, grouping the clients according to their local data distribution information specifically includes: The client performs ZORB reverse mapping based on the user tag data, sums the sample number dimension of the reverse output to obtain an intermediate output containing local data distribution information, and performs a clustering operation based on the intermediate output containing local data distribution information to obtain several different groups. The calculation expression of the intermediate output containing local data distribution information is:

[0016] In the above formula, is the intermediate output after summation, is the local sample index, is the total number of local samples, is a single sample label.

[0017] Furthermore, all clients in each group jointly update a classification layer, specifically including: The client calculates the least squares solution of each layer in the classification layer of the local model, and transmits the least squares solution of the current layer to the parameter server every time it completes the calculation; The client receives the intra-group aggregation results of each layer in the corresponding classification layer transmitted by the parameter server in turn, and updates the network parameters of the neural network of the classification layer corresponding to the local model according to the intra-group aggregation results until the update of all layers is completed, and the final personalized classifier is trained. The intra-group aggregation result is obtained by the parameter server adding and aggregating the least squares solutions of the same classification layer of all clients in the group where the current client is located.

[0018] A method for optimizing a federated learning framework for non-independent and identically distributed data, comprising: The parameter server sends the same random seed to each client connected to the parameter server, where the random seed is used to initialize the model parameters of the client's local model, where the local model is a neural network model. The parameter server receives the least squares solutions of each layer of the neural network of the local model returned by each client, performs additive aggregation based on the least squares solutions of the same layer of the neural network transmitted by each client to obtain the global least squares solution of the same layer of the neural network, and sends the global least squares solution as the global optimal model parameters to each client; wherein, the least squares solution of each layer of the neural network of the local model of each client is calculated by the client based on the non-independent and identically distributed user label data in the acquired training data set and the model parameters of the local model.

[0019] A method for optimizing a federated learning framework for non-independent and identically distributed data, comprising: The parameter server sends the same random seed to each client connected to the parameter server. The random seed is used to initialize the model parameters of the client's local model. The local model is a neural network model, including a feature extraction layer and a classification layer. The feature extraction layer is all convolutional network layers, and the classifier layer is all fully connected layers. The parameter server receives the least squares solutions of the neural networks of each layer of the local model feature extraction layer returned by each client, performs additive aggregation based on the least squares solutions of the neural networks of the same layer of the feature extraction layer transmitted by each client to obtain a global least squares solution of the neural network of the same layer of the feature extraction layer, and sends the global least squares solution as the global optimal model parameters of the feature extraction layer to each client; wherein, the least squares solutions of the neural networks of each layer of the local model feature extraction layer of each client are calculated by the client based on the non-independent and identically distributed user label data in the acquired training data set and the model parameters of the local model; The parameter server groups the clients according to their local data distribution information and sends the grouping results to each client.

[0020] A federated learning framework optimization system for non-independent and identically distributed data, comprising: a parameter server and multiple clients; The parameter server sends the same random seed to each client connected to the parameter server through the initialization module to initialize the model parameters of the client's local model, receives the least squares solution of each layer of the local model neural network returned by each client through the aggregation module and performs additive aggregation to obtain the global least squares solution of the same layer of the neural network, and sends the global least squares solution as the global optimal model parameter to each client through the data transmission module; Each of the clients receives a random seed transmitted by the parameter server through an initialization module, and uses the random seed to initialize the local model, obtains non-independent and identically distributed user label data in the training data set through a calculation module, reverse maps the user label data based on a zero-order relaxation backpropagation algorithm, and calculates the least squares solution of each layer of the neural network in the local model. Whenever the least squares solution of a layer of neural network is calculated, the least squares solution of the current layer is transmitted to the parameter server through the data transmission module, and the global least squares solution of each layer of the neural network returned by the parameter server is received in turn. The global least squares solution of each layer of the neural network received is used to update the network parameters of the corresponding layer of the neural network in the local model, and the global optimal model is obtained after completing the update of all layers of the neural network.

[0021] Compared with the prior art, the advantages of the present invention are: The client of the present invention is initialized by the same initialization model parameters, which can ensure the consistency of each model parameter and improve the stability of subsequent model training; the local model calculates the least squares solution through the zero-order relaxation back propagation algorithm, which can reduce the computational pressure of edge devices with limited computing power; at the same time, a layer-by-layer aggregation strategy is adopted, and the least squares solution of each layer of neural network calculated by the client is transmitted to the parameter server, which can enable each client to train independently according to its own calculation progress, speed up the overall training speed, and improve training efficiency. The least squares solution of the same layer of neural network transmitted by each client is additively aggregated to obtain the global least squares solution, which can make the obtained global least squares solution better approximate the global optimal solution directly calculated by the zero-order relaxation back propagation algorithm for all local model data from a mathematical perspective, thereby obtaining a global model with better performance when the training is completed, and effectively optimizing the federated learning model training process. BRIEF DESCRIPTION OF THE DRAWINGS

[0022] Figure 1 Schematic diagram of the process of optimizing the federated learning framework for non-independent and identically distributed data of the present invention.

[0023] Figure 2 The figure is a flowchart of a method for optimizing a federated learning framework for non-independent and identically distributed data in a specific application embodiment.

[0024] Figure 3 A flowchart of a method for optimizing a federated learning framework for non-independent and identically distributed data in another specific application embodiment is shown.

[0025] Figure 4 A statistical comparison of the total training time for optimizing the federated learning framework using different algorithms. DETAILED DESCRIPTION

[0026] In order to better understand the above technical solution, the above technical solution will be described in detail below with reference to the accompanying drawings and specific implementation methods.

[0027] Federated learning is a new distributed machine learning paradigm that allows models to be trained on multiple edge devices without uploading raw data to a central parameter server (PS). Training a general federated learning model typically involves three steps: (i) a group of edge devices perform local computations based on their own datasets, typically using stochastic gradient descent (SGD), and upload the resulting parameters to the server; (ii) the server aggregates the uploaded parameters and updates the global model; and (iii) the server feeds the new global model back to the edge devices for a new round of local computation. This process is repeated until the loss function converges and a certain model training accuracy is achieved. The entire training process is often referred to as the Federated Averaging (FedAvg) algorithm.

[0028] Typical federated learning training schemes include first-order methods (such as stochastic gradient descent) and second-order methods, which require the computation of gradients and Hessian matrices. While these gradient descent (GD)-based methods are highly efficient, their reliance on gradient computation presents challenges when applied to edge devices. Specifically, edge devices with limited computational power cannot frequently perform gradient computations. For example, mobile phones or displays equipped only with a neural network processor (NPU) are primarily designed for inference, specifically for inference rather than gradient computation. In these cases, only small, pre-trained, or pruned models are computationally suitable for execution on these devices, making first-order and second-order federated learning methods difficult to implement. Furthermore, gradient information is not always available to support gradient computation. For example, in scenarios such as federated hyperparameter tuning and federated black-box attacks, gradients cannot be derived due to the non-differentiability of the loss function. Therefore, gradient descent-based federated learning frameworks are less feasible in these scenarios.

[0029] To address these issues, practitioners have proposed zero-order FL training as an efficient gradient-free FL training method. Existing zero-order FL training methods typically use finite difference techniques to approximate the gradient, effectively replacing automatic differentiation with multiple inference passes in the gradient calculation. However, this gradient estimation suffers from bias, requiring a large number of sampling passes to reduce the estimation error, resulting in significant computational overhead and excessively long training times for these methods. To address this issue, practitioners have further proposed zero-order relaxed backpropagation (ZORB) as an alternative to finite difference techniques. This method uses Moore-Penrose (MP) inversion to sequentially generate a least-squares solution for each layer of the neural network. ZORB is significantly faster (for example, approximately 300 times faster than the Adam optimizer) when training shallow neural networks with comparable accuracy. Therefore, leveraging the properties of ZORB to further accelerate FL model training is an effective approach.

[0030] Currently, when training directly with ZORB and FedAvg, since the weights calculated for each layer by ZORB are least-squares solutions, it is necessary to aggregate the weights calculated on the local datasets of each edge device (hereinafter referred to as the client) using the FedAvg method to obtain a global least-squares solution. A simple aggregation method is to average the least-squares solutions. However, from a mathematical perspective, when decomposing the analytical expression of the global model into the sum of multiple sub-equations, the first-order derivative of the analytical expression of the least-squares problem in ZORB is zero. The average of the least-squares solutions of each sub-equation is obviously not equal to the least-squares solution of the analytical expression of the global model. As a result, the final global least-squares solution is mathematically inconsistent with the least-squares solution calculated by inputting all data into ZORB. Therefore, an aggregation strategy is needed that can make the aggregated least-squares solution as close as possible to the mathematically global least-squares solution.

[0031] In order to solve the above problems, after initialization, the client of the present invention first calculates the least squares solution of each layer of the neural network of the local model based on the zero-order relaxation back propagation algorithm according to the user label data in the training data set, and then sequentially receives the global least squares solution of each layer of the neural network returned by the parameter server to update the network parameters of the neural network in the local model to obtain the global optimal model, wherein the global least squares solution is obtained by additively aggregating the least squares solutions of the same layer of the neural network transmitted by each client through the parameter server, and adopts a layer-by-layer aggregation strategy so that the obtained global least squares solution can better approximate the optimal solution from a mathematical perspective, so as to obtain a global model with better performance when the training is completed.

[0032] The following is a brief description of the basic knowledge of zero-order relaxed backpropagation (ZORB) involved in the present invention.

[0033] ZORB solves the problems of slow convergence and hyperparameter tuning in gradient-based deep learning (DL) by exploiting the least squares property of the Moore-Penrose (MP) generalized inverse in processing linear equations to obtain the least squares solution of the neural network layer by layer. This least squares property can be mathematically expressed as follows:

[0034] in is the explanatory variable, is the response variable, is its least squares solution.

[0035] Next, consider an n-layer neural network (NN): (2) in, X∈ ,Y∈ is the input matrix and label matrix of the neural network with b training samples. The input and output dimensions of the neural network are and ,and ,..., ∈ ,..., ∈ is the weight matrix of each layer. represents the activation function of the kth layer, which operates element-by-element on the intermediate output matrix. The goal of ZORB is to To begin, determine the ,..., The weight update process of any layer can be expressed mathematically as follows:

[0036]

[0037] in represents the MP inverse, Indicates updated layer weights, () represents the inverse mapping. The inverse mapping defined by ZORB is shown in the following table:

[0038] Example 1 The optimization method of this embodiment is specifically applied to a distributed system consisting of a parameter server and several clients, such as Figure 1As shown, the steps of the federated learning framework optimization method for non-independent and identically distributed data in this embodiment include: Step S1: The client receives a random seed transmitted by a parameter server and uses the random seed to initialize a local model. The random seed is the same initialization model parameter sent by the parameter server to each client connected to the parameter server. The local model is a neural network model. In step S2, the client obtains the user label data in the training data set, performs reverse mapping on the user label data based on the zero-order relaxation backpropagation algorithm, and calculates the least squares solution of each layer of the local model neural network. Whenever the least squares solution of a layer of neural network is calculated, the least squares solution of the current layer is transmitted to the parameter server; In step S3, the client receives the global least squares solution of each layer of the neural network returned by the parameter server in turn, and uses the received global least squares solution of each layer of the neural network to update the network parameters of the corresponding layer of the neural network in the local model. After completing the update of all layers of the neural network, the global optimal model is obtained. The global least squares solution is obtained by adding and aggregating the least squares solutions of the same layer of the neural network transmitted by each client.

[0039] It can be understood that the client of this embodiment is initialized by the same initialization model parameters, which can ensure the consistency of each model parameter and improve the stability of subsequent model training; the local model calculates the least squares solution through the zero-order relaxation back propagation algorithm, which can reduce the computational pressure of edge devices with limited computing power; at the same time, a layer-by-layer aggregation strategy is adopted, and each time the least squares solution of a layer of neural network is calculated by the client, it is transmitted to the parameter server, which can enable each client to train independently according to its own calculation progress, speed up the overall training speed, and improve training efficiency. The least squares solution of the same layer of neural network transmitted by each client is additively aggregated to obtain the global least squares solution, which can make the obtained global least squares solution better approximate the global optimal solution directly calculated by the zero-order relaxation back propagation algorithm for all local model data from a mathematical perspective, thereby obtaining a global model with better performance when the training is completed, effectively optimizing the federated learning model training process.

[0040] In this embodiment, the parameter server sends the same random seed to each client connected to the parameter server. The random seed is used to initialize the model parameters of the client's local model, which is a neural network model. The parameter server receives the least squares solutions of each layer of the neural network of the local model returned by each client, performs additive aggregation based on the least squares solutions of the same layer of the neural network transmitted by each client to obtain the global least squares solution of the same layer of the neural network, and sends the global least squares solution as the global optimal model parameters to each client; wherein, the least squares solution of each layer of the neural network of each client's local model is calculated by the client based on the user label data in the acquired training data set and the model parameters of the local model.

[0041] In this embodiment, the training model of each layer of the neural network of the global optimal model is a distributed least squares problem, which is expressed as follows:

[0042] In the above formula, To be trained Layer neural network weight matrix, c is the client number, C is the total number of clients, Indicates the local data sample on the cth client The first step obtained by inputting the neural network corresponding to the local model The actual input of the layer, Indicates that the local tag on the cth client The inverse mapping The layer expects output labels.

[0043] In this embodiment, the calculation expression of the global least squares solution is:

[0044]

[0045]

[0046] In the above formula, To be trained Layer neural network weight matrix, c is the client number, C is the total number of clients, Indicates that the local tag on the cth client The inverse mapping The layer expects the output label, The first The least squares solution of the layer, is the local loss function of the cth client The second-order partial derivative of , where the superscript T indicates the transpose.

[0047] In step S2 of this embodiment, the least squares solution of a layer of neural network is calculated, which also includes calculating the second-order partial derivative. , transmit the least squares solution and the second-order partial derivatives to the parameter server together, where Indicates the local data sample on the cth client The first step obtained by inputting the neural network corresponding to the local model The actual input of the layer.

[0048] In a specific application embodiment, since the weighted least squares estimator has been proven to be successful in combining the least squares solutions obtained by each subsystem in a distributed system, the neural network can be used in the federated learning scenario. Layer training is modeled as a distributed least squares problem as shown in Equation (4). Equation (4) is redefined as follows:

[0049]

[0050] Then make a second-order Taylor expansion of equation (6):

[0051]

[0052]

[0053]

[0054] In formula (8), , For formula (7) The second-order partial derivative of , in the line of the approximate equality, we use the fact that the first-order derivative of the least squares method is equal to 0, that is, =0, and in formula (8), the following redefinition is made:

[0055]

[0056] Finally, by transforming Equation (8) into The partial derivative of is set to 0, and the formula (5) can be obtained ,as follows:

[0057] Equation (9) is the parameter aggregation expression for the least squares solution of the same layer of neural network transmitted by each client calculated through ZORB.

[0058] The following will be combined Figure 2The process of the federated learning framework optimization method for non-independent and identically distributed data in this embodiment is described. Specifically, the method is a federated learning method based on the FedMPI (Federated Learning via Moore-Penrose Pseudoinverse) framework with zero-order relaxed backpropagation.

[0059] Step A1, model initialization: Each client receives the same random seed from the parameter server (PS) and uses this seed to initialize its local model , to ensure that a unified global model is initialized. Then, the input layer (that is, layer 1) of the neural network corresponding to the local model is used as the first round of training layer.

[0060] Step A2, local calculation (for any k-th round Each client uses formula (3-2) to perform the ZORB reverse mapping process on all the user's tag data, and calculates the local first The least squares solution of the layer neural network, that is , at the same time, calculate the second-order partial derivatives required for the next aggregation step, that is , and finally upload the least squares solution and the second-order partial derivative to the parameter server.

[0061] Step A3, transmission and aggregation: The parameter server receives the least squares solution and second-order partial derivative of each client, and performs additive aggregation using formula (9) to calculate the first Global least squares solution of a multi-layer neural network , and sends the solution back to each client.

[0062] Step A4, local model update: Each client receives the global least squares solution from the parameter server and uses it to replace the first Layer network parameters, then the FedMPI framework will enter the next layer (that is, layers) until all layers of the neural network are trained.

[0063] It can be understood that any k-th round corresponds to the k-th layer training, and in any federated learning In each round, each client only needs to upload as well as You can complete the The parameters of the neural network layer are aggregated. In other words, FedMPI only needs The number of training rounds matches the number of neural network layers in federated learning. Federated learning only requires the same number of rounds as the number of layers, which can further reduce the total training time and improve training efficiency.

[0064] Example 2 This embodiment is applicable to a distributed system consisting of a parameter server and several clients, and also implements personalized model customization for client training. The optimization method of the federated learning framework for non-independent and identically distributed data in this embodiment includes: The client receives a random seed transmitted by the parameter server and uses the random seed to initialize the local model. The random seed is the same initialization model parameter sent by the parameter server to each client connected to the parameter server. The local model is a neural network model. The local model is divided into a feature extraction layer and a classifier layer. The feature extraction layer is all convolutional network layers, and the classifier layer is all fully connected layers. Feature extraction layer training: The client obtains the user label data in the training data set, performs reverse mapping on the user label data based on the zero-order relaxation backpropagation algorithm, and calculates the least squares solution of each layer of the local model neural network. Whenever the least squares solution of a layer of neural network is calculated, the least squares solution of the current layer is transmitted to the parameter server; the client sequentially receives the global least squares solution of each layer of feature extraction layer returned by the parameter server and updates the network parameters of the neural network of the corresponding layer of the feature extraction layer of the local model until the training of the feature extraction layer is completed. The global least squares solution is obtained by the parameter server by adding and aggregating the least squares solutions of the same layer of feature extraction layer transmitted by each client; during the training of the feature extraction layer, the client transmits information containing local data distribution to the parameter server; Classifier layer training: The client receives the grouping results from the parameter server and performs classifier training based on the grouping results, where all clients in each group jointly update a classification layer. The grouping results are obtained by the parameter server by grouping each client based on the local data distribution information of each client.

[0065] It can be understood that dividing the local model of each client into a feature extraction layer and a classifier layer, and grouping each client according to the local data distribution information of each client so that all clients in each group jointly update a classification layer can alleviate the model performance degradation caused by data heterogeneity between federated learning edge devices (clients). Each client trains a personalized model for the data distribution of its assigned group, so that each client can customize the personalized model according to its local data distribution, thereby improving model performance.

[0066] In this embodiment, the parameter server sends the same random seed to each client connected to the parameter server. The random seed is used to initialize the model parameters of the client's local model. The local model is a neural network model, including a feature extraction layer and a classification layer. The feature extraction layer is all convolutional network layers, and the classifier layer is all fully connected layers. The parameter server receives the least squares solutions of the neural networks of each layer of the local model feature extraction layer returned by each client, performs additive aggregation on the least squares solutions of the neural networks of the same layer of the feature extraction layer transmitted by each client to obtain a global least squares solution of the neural network of the same layer of the feature extraction layer, and sends the global least squares solution as the global optimal model parameters of the feature extraction layer to each client; wherein the least squares solutions of the neural networks of each layer of the local model feature extraction layer of each client are calculated by the client based on the user label data in the acquired training data set and the model parameters of the local model; The parameter server groups the clients according to their local data distribution information and sends the grouping results to each client.

[0067] In this embodiment, grouping the clients according to their local data distribution information specifically includes: The client performs ZORB reverse mapping based on the user tag data, sums the sample number dimension of the reverse output to obtain an intermediate output containing local data distribution information, and performs a clustering operation based on the intermediate output containing local data distribution information to obtain several different groups. The calculation expression of the intermediate output containing local data distribution information is:

[0068] In the above formula, is the intermediate output after summation, is the local sample index, is the total number of local samples, is a single sample label.

[0069] In this embodiment, all clients in each group jointly update a classification layer, specifically including: The client calculates the least squares solution of each layer in the classification layer of the local model, and transmits the least squares solution of the current layer to the parameter server every time it completes the calculation; The client receives the intra-group aggregation results of each layer in the corresponding classification layer transmitted by the parameter server in turn, and updates the network parameters of the neural network of the classification layer corresponding to the local model according to the intra-group aggregation results until the update of all layers is completed, and the final personalized classifier is trained. The intra-group aggregation result is obtained by the parameter server adding and aggregating the least squares solutions of the same classification layer of all clients in the group where the current client is located.

[0070] In a specific application embodiment, combined with Figure 3 The specific process of the federated learning framework optimization method for non-independent and identically distributed data in this embodiment is described. The method is specifically a federated learning method based on a personalized FedMPI (pFedMPI) framework with zero-order relaxed backpropagation.

[0071] Model sharding: Like FedMPI, each client receives the same random seed from the parameter server and uses it to initialize its model. , to ensure a unified global model is initialized. The client then splits the local model into two parts: the feature extraction layer and the classifier layer. The feature extraction layer is all convolutional network layers, and the classifier layer is all fully connected layers.

[0072] Phase 1: Each client uses formula (3-2) to reversely map the local label matrix. After passing through the fully connected layer, the sample number dimension of its reverse output is summed up. The mathematical expression is shown in formula (10).

[0073] Get intermediate output After the matrix is generated (this matrix implies the local data distribution information), the client will upload it to the parameter server. After the parameter server receives the intermediate output of all clients, it executes the clustering algorithm (specifically, -means algorithm), at the same time, the client and parameter server follow the above FedMPI process (i.e., steps A2-A4) to train the feature extraction layer layer by layer. At the end of the first stage, the parameter server obtains G groups according to the K-Means algorithm, namely: , and complete the training of the feature extraction layer.

[0074] Phase 2: After completing the training of the feature extraction layer, the personalized classifier layer is trained. To achieve personalization, clients with similar data distributions are trained together to share a common classifier. In each subsequent round of federated learning, the client performs local computations and uploads the least squares solution and second-order partial derivatives, similar to FedMPI. However, unlike FedMPI, the aggregation step generates G intra-group aggregate results. Each intra-group aggregate result can be mathematically expressed as:

[0075] In the above formula, represents the g-th result within any k-th round. At the end of each round, the parameter server sends the corresponding aggregated results for each client's group. After receiving the aggregated results, each client updates its local model and proceeds to training the next layer. Finally, after completing the training of the personalized classifier layer by layer, the model is personalized for the local private data distribution, thereby improving model performance.

[0076] In summary, pFedMPI splits the local model into a feature extractor and a classifier, and achieves personalized customization by having clients with similar data distributions jointly train a personalized classifier. The pFedMPI workflow consists of two phases. In the first phase, pFedMPI updates the feature extractor layer by layer, just like the original FedMPI. The difference is that pFedMPI also requires a parameter server to group clients based on the similarity of their data distributions. Then, in the second phase, based on the grouping results of the first phase, the clients that were grouped together jointly train a classifier. It can be understood that parallelizing the grouping process with model training can further reduce the time consumption during the training process and achieve better performance levels than FedMPI.

[0077] The following experiments in specific application scenarios verify the effectiveness of the federated learning framework optimization method for non-independent and identically distributed data of the present invention.

[0078] (1) Experimental setup Testbed: A FedMPI prototype was implemented on a CPU cluster-based testbed with 100 nodes, each equipped with an Intel Xeon Platinum 9242 CPU. Each node implemented a single process using one CPU core and 4GB of memory to simulate edge devices with limited computing power. The nodes were connected via an InfiniBand network with a bandwidth of 100Gbps.

[0079] Models and Datasets: Experiments are conducted on three types of neural network models: 1) a three-layer deep neural network; 2) LeNet (6×10^6 parameters); and 3) MobileNetV3-small (2.9 million parameters). The deep neural network and LeNet represent shallow models, while MobileNet demonstrates transfer learning from deep models. The deep neural network is trained on the MNIST dataset, while LeNet is trained on the CIFAR-10 dataset. MobileNet is trained on the OfficeHome dataset, which contains 65 categories from four different domains. Next, the paper explains how to partition the training data under the non-independent and identically distributed (IID) setting. Non-IID here stands for non-independence and non-identical distribution. Non-independence refers to a degree of correlation or association between two or more random variables, meaning the value of one random variable may be affected by the value of one or more other random variables. Non-identical distribution refers to samples that are not drawn from the same distribution.

[0080] For the MNIST and CIFAR-10 datasets, Dirichlet sampling is used. Generate non-IID partitions of the dataset for 100 clients, where represents the probability of assigning a sample with label k to client j. A smaller β (lower level of non-IID data) means greater diversity in the data distribution among local clients. For OfficeHome, the total number of clients is set to C = 40, and 10 clients are generated from each domain to simulate a non-IID scenario.

[0081] Benchmarks and Metrics. We select three classic and novel algorithms, namely BAFFLE, FedAvg, and FedZO. We then use test accuracy, training time, and memory cost to evaluate the performance and training efficiency of FedMPI, pFedMPI, and their baseline models.

[0082] Table 1 MNIST and CIFAR-10 datasets with non-IID labels

[0083] As can be seen from Table 1, on the MNIST and CIFAR-10 datasets with non-IID labels, FedMPI and pFedMPI rank higher in accuracy.

[0084] (2) Model performance on non-IID clients Following section (1), the model performance of the framework is evaluated in the scenarios where labels are not independent and identically distributed and features are not independent and identically distributed, as shown in Table 1.

[0085] For the non-IID label setting, FedZO and BAFFLE achieved the best test accuracy within 1000 / 1000 and 20 / 40 communication rounds, respectively, following their original work on MNIST and CIFAR-10. BAFFLE was tested with different perturbation K (ranging from 100 to 500 on MNIST and from 200 to 1000 on CIFAR-10). For FedAvg, the best accuracy was achieved within 100 rounds, with 5 local training rounds. Since pFedMPI exhibited the best model performance in this setting, the number of clusters in pFedMPI was set to 25. As shown in Table 1, FedMPI achieved state-of-the-art accuracy when training LeNet on CIFAR-10, compared to FedZO and BAFFLE. BAFFLE's performance degraded by 5% to 6% when K was reduced, as the reduction in the number of inference rounds may result in reduced sampling, which can amplify errors in the gradient approximation. Compared to vanilla FedMPI, pFedMPI improves accuracy by 8% to 30%. Notably, under certain non-IID data settings (β = 0.1 and β = 0.3), pFedMPI outperforms the first-order FedAvg framework, while other zero-order FL frameworks underperform the first-order framework. Furthermore, as β decreases, FedMPI has little impact on test accuracy, while pFedMPI provides a significant improvement in model performance, while the other frameworks experience a 7% to 20% drop in test accuracy.

[0086] As can be understood, the aggregation strategy in FedMPI does not introduce additional error terms, so FedMPI can achieve the same model performance as full batch training. Therefore, FedMPI demonstrates strong robustness to non-IID data levels. For pFedMPI, a smaller β means greater diversity in data distribution across clients, making it easier to group clients and implement personalization. Therefore, pFedMPI can provide greater improvements in model performance when the non-IID data distribution is more diverse.

[0087] For scenarios where features are not independent and identically distributed, we adapted MobileNet-based transfer learning by loading a model pre-trained on ImageNet and fine-tuning the fully connected layers. FedMPI does not require an optimizer, while FedZO, BAFFLE, and FedAvg all use the Adam optimizer with a learning rate of 0.005 and train for 40 epochs. The results are shown in Table 2.

[0088] Table 2: OfficeHome dataset with non-IID participation features

[0089] Table 2 presents tests on the art, clipart, product, and real-world domains. As can be seen from Table 2, under the setting where features are not independent and identically distributed, FedMPI surpasses all zero-order methods in accuracy, falling only 2.64% behind FedAvg. Notably, FedMPI and pFedMPI outperform FedAvg in the art and product domains, respectively, demonstrating their potential for application to downstream tasks in specific domains. pFedMPI improves test accuracy by 3%, achieving nearly the same performance level as FedAvg.

[0090] (3) Training efficiency The training time of the aforementioned methods is compared. Table III lists the average communication time per training round and the peak client memory usage. Each round of training time consists of both computation time and communication time. Mini-batch training plays a key role in the computation time, while communication time refers to the time spent transmitting data across the network. Since pFedMPI does not impose additional storage burden on the client compared to the original FedMPI, only the communication time of pFedMPI is reported in Table III.

[0091] Table 3 Communication time and peak memory usage per training round for FedMPI, FedZO, BAFFLE, and FedAVG

[0092] from Figure 4 As can be seen, FedMPI achieves a significant 9× to 197× speedup compared to the BAFFLE benchmark, demonstrating that MP inverse is able to speed up training compared to finite difference and gradient descent based methods. p The training efficiency of FedZO, BAFFLE and FedAvg is analyzed from the perspective of convergence rate. The convergence rate of FedZO is √O(d / NHTb1b2), where d is p model size, it is d times slower than FedAvg (for the same batch size, O(1 / NHTb1b2)). Therefore, as the model size increases, FedZO requires more training iterations than FedAvg, although it p n The processing time on the batch may be shorter. The convergence rate of BAFFLE is O K, which means that a larger K can achieve faster convergence. However, in BAFFLE, there is a trade-off between convergence speed and computational cost, as a larger K also means more inference processes. Therefore, it is difficult to determine K to optimize this trade-off. The training efficiency level of pFedMPI is close to that of original FedMPI because the designed parallelization strategy can overlap the time of the grouping process with the training time of the feature extractor.

[0093] In addition to the communication time mentioned above, Table 3 shows the peak memory usage of each method and its ratio to FedAvg. Since K cannot change the communication complexity and memory cost (MB), in BAFFLE, communication time and memory usage are recorded only once for each model. Numerical results show that pFedMPI can reduce communication time by 27%-60% compared to standard FedMPI. This is because the aggregation domain shifts from the global level to the group size level, reducing the network load. Since the completion time of local training in each group varies, each group can fully utilize network resources during communication.

[0094] It can be understood that the FedMPI of the present invention is a computationally efficient zero-order framework that performs federated learning through MP inversion. The design of the hierarchical aggregation strategy makes the number of training rounds match the number of layers in the global model, greatly reducing the total training time. In addition, the pFedMPI of the present invention is a personalized version of FedMPI, which is used to improve model performance when dealing with practical scenarios involving non-independent and identically distributed data in federated learning. In order to evaluate its effectiveness, the framework was implemented on a CPU cluster-based test platform. Experiments show that compared with the zero-order benchmark BAFFLE, FedMPI can achieve up to 9 times to 197 times acceleration in training time, and pFedMPI has an average improvement of 12% in test accuracy when processing non-independent and identically distributed data compared to ordinary FedMPI.

[0095] The present invention further provides a federated learning framework optimization system for non-independent and identically distributed data, comprising: a parameter server and a plurality of clients; The parameter server sends the same random seed to each client connected to the parameter server through the initialization module to initialize the model parameters of the client's local model, receives the least squares solution of each layer of the local model neural network returned by each client through the aggregation module and performs additive aggregation to obtain the global least squares solution of the same layer of the neural network, and sends the global least squares solution as the global optimal model parameter to each client through the data transmission module; Each of the clients receives a random seed transmitted by the parameter server through an initialization module, and uses the random seed to initialize the local model, obtains non-independent and identically distributed user label data in the training data set through a calculation module, reverse maps the user label data based on a zero-order relaxation backpropagation algorithm, and calculates the least squares solution of each layer of the neural network in the local model. Whenever the least squares solution of a layer of neural network is calculated, the least squares solution of the current layer is transmitted to the parameter server through the data transmission module, and the global least squares solution of each layer of the neural network returned by the parameter server is received in turn. The global least squares solution of each layer of the neural network received is used to update the network parameters of the corresponding layer of the neural network in the local model, and the global optimal model is obtained after completing the update of all layers of the neural network.

[0096] The system of the present invention corresponds to the above method and also has the advantages described in the above method.

[0097] The present invention can implement all or part of the process steps in the above-described method embodiments by instructing related hardware through a computer program. The computer program can be stored in a computer-readable storage medium. When executed by a processor, the computer program can implement the steps of the above-described method embodiments. The computer program includes computer program code, which can be in source code form, object code form, executable file, or some intermediate form. Computer-readable media include any entity or device capable of carrying computer program code, recording media, USB flash drives, removable hard drives, magnetic disks, optical disks, computer memory, read-only memory (ROM), random access memory (RAM), electrical carrier signals, telecommunication signals, and software distribution media. Memory is used to store computer programs and / or modules. The processor implements various functions by running or executing the computer programs and / or modules stored in the memory and accessing data stored in the memory. The memory may include a high-speed random access memory and may also include a non-volatile memory, such as a hard disk, a memory, a plug-in hard disk, a Smart Media Card (SMC), a Secure Digital (SD) card, a Flash Card, at least one disk storage device, a flash memory device, or other volatile solid-state storage devices.

[0098] The above description is merely a preferred embodiment of the present invention. The scope of protection of the present invention is not limited to the above embodiment. All technical solutions based on the concept of the present invention are within the scope of protection of the present invention. It should be noted that for those skilled in the art, various improvements and modifications that do not depart from the principles of the present invention should also be considered within the scope of protection of the present invention.

Claims

1. A method for optimizing a federated learning framework for non-IID data, the method being applied to a distributed system consisting of a parameter server and several clients, characterized in that: The method comprises: Step S1: The client receives a random seed transmitted by a parameter server and uses the random seed to initialize a local model. The random seed is the same initialization model parameter sent by the parameter server to each client connected to the parameter server. The local model is a neural network model. In step S2, the client obtains non-independent and identically distributed user tag data in the training dataset, performs reverse mapping on the user tag data based on a zero-order relaxed backpropagation algorithm, and calculates the least squares solution of each layer of the neural network in the local model. Whenever the least squares solution of a layer of the neural network is calculated, the least squares solution of the current layer is transmitted to the parameter server; the non-independent and identically distributed user tag data indicates that the user tag data of each client is correlated and comes from different distributions; In step S3, the client receives the global least squares solution of each layer of the neural network returned by the parameter server in turn, and uses the received global least squares solution of each layer of the neural network to update the network parameters of the corresponding layer of the neural network in the local model. After completing the update of all layers of the neural network, the global optimal model is obtained. The global least squares solution is obtained by adding and aggregating the least squares solutions of the same layer of the neural network transmitted by each client.

2. The method for optimizing a federated learning framework for non-independent and identically distributed data according to claim 1, characterized in that: The training model of each layer of the global optimal model neural network is a distributed least squares problem, which is expressed as: In the above formula, To be trained Layer neural network weight matrix, c is the client number, C is the total number of clients, Indicates the local data sample on the cth client The first step obtained by inputting the neural network corresponding to the local model The actual input of the layer, Indicates that the local tag on the cth client The inverse mapping The layer expects output labels.

3. The method for optimizing a federated learning framework for non-independent and identically distributed data according to claim 1, characterized in that: The calculation expression of the global least squares solution is: In the above formula, To be trained Layer neural network weight matrix, c is the client number, C is the total number of clients, Indicates that the local tag on the cth client The inverse mapping The layer expects the output label, The first The least squares solution of the layer, is the local loss function of the cth client The second-order partial derivative of , where the superscript T indicates the transpose.

4. The method for optimizing a federated learning framework for non-independent and identically distributed data according to claim 3, characterized in that: In step S2, the least squares solution of a layer of neural network is calculated, which also includes the calculation of the second-order partial derivative. , transmit the least squares solution and the second-order partial derivatives to the parameter server together, where Indicates the local data sample on the cth client The first step obtained by inputting the neural network corresponding to the local model The actual input of the layer.

5. A method for optimizing a federated learning framework for non-independent and identically distributed data, characterized in that the steps include: The client receives a random seed transmitted by the parameter server and uses the random seed to initialize the local model. The random seed is the same initialization model parameter sent by the parameter server to each client connected to the parameter server. The local model is a neural network model, and the local model is divided into a feature extraction layer and a classification layer. The feature extraction layer is all convolutional network layers, and the classifier layer is all fully connected layers. Feature extraction layer training: The client obtains the non-independent and identically distributed user label data in the training data set, performs reverse mapping on the user label data based on the zero-order relaxed back-propagation algorithm, and calculates the least squares solution of each layer of the local model neural network. Whenever the least squares solution of a layer of neural network is calculated, the least squares solution of the current layer is transmitted to the parameter server; the client sequentially receives the global least squares solution of each layer of feature extraction layer returned by the parameter server and updates the network parameters of the neural network of the corresponding layer of the feature extraction layer of the local model until the training of the feature extraction layer is completed. The global least squares solution is obtained by the parameter server by adding and aggregating the least squares solutions of the same layer of feature extraction layer transmitted by each client; During the training of the feature extraction layer, the client transmits information containing local data distribution to the parameter server; The non-independent and identically distributed user tag data means that the user tag data of each client is correlated and comes from different distributions; Classifier layer training: The client receives the grouping results from the parameter server and performs classifier training based on the grouping results, where all clients in each group jointly update a classification layer. The grouping results are obtained by the parameter server by grouping each client based on the local data distribution information of each client.

6. The method for optimizing a federated learning framework for non-independent and identically distributed data according to claim 5, characterized in that: Grouping each client according to its local data distribution information specifically includes: The client performs ZORB reverse mapping based on the user tag data, sums the sample number dimension of the reverse output to obtain an intermediate output containing local data distribution information, and performs a clustering operation based on the intermediate output containing local data distribution information to obtain several different groups. The calculation expression of the intermediate output containing local data distribution information is: In the above formula, is the intermediate output after summation, is the local sample index, is the total number of local samples, is a single sample label.

7. The method for optimizing a federated learning framework for non-independent and identically distributed data according to claim 5, characterized in that: All clients in each group jointly update a classification layer, specifically including: The client calculates the least squares solution of each layer in the classification layer of the local model, and transmits the least squares solution of the current layer to the parameter server every time it completes the calculation; The client receives the intra-group aggregation results of each layer in the corresponding classification layer transmitted by the parameter server in turn, and updates the network parameters of the neural network of the classification layer corresponding to the local model according to the intra-group aggregation results until the update of all layers is completed, and the final personalized classifier is trained. The intra-group aggregation result is obtained by the parameter server adding and aggregating the least squares solutions of the same classification layer of all clients in the group where the current client is located.

8. A method for optimizing a federated learning framework for non-independent and identically distributed data, characterized in that: include: The parameter server sends the same random seed to each client connected to the parameter server, where the random seed is used to initialize the model parameters of the client's local model, where the local model is a neural network model. The parameter server receives the least squares solutions of each layer of the neural network of the local model returned by each client, performs additive aggregation based on the least squares solutions of the same layer of the neural network transmitted by each client to obtain the global least squares solution of the same layer of the neural network, and sends the global least squares solution as the global optimal model parameters to each client; wherein, the least squares solution of each layer of the neural network of the local model of each client is calculated by the client based on the non-independent and identically distributed user label data in the acquired training data set and the model parameters of the local model.

9. A method for optimizing a federated learning framework for non-independent and identically distributed data, characterized in that: include: The parameter server sends the same random seed to each client connected to the parameter server. The random seed is used to initialize the model parameters of the client's local model. The local model is a neural network model, including a feature extraction layer and a classification layer. The feature extraction layer is all convolutional network layers, and the classifier layer is all fully connected layers. The parameter server receives the least squares solutions of the neural networks of each layer of the local model feature extraction layer returned by each client, performs additive aggregation based on the least squares solutions of the neural networks of the same layer of the feature extraction layer transmitted by each client to obtain a global least squares solution of the neural network of the same layer of the feature extraction layer, and sends the global least squares solution as the global optimal model parameters of the feature extraction layer to each client; wherein, the least squares solutions of the neural networks of each layer of the local model feature extraction layer of each client are calculated by the client based on the non-independent and identically distributed user label data in the acquired training data set and the model parameters of the local model; The parameter server groups the clients according to their local data distribution information and sends the grouping results to each client.

10. A federated learning framework optimization system for non-independent and identically distributed data, characterized in that: include: Parameter server and multiple clients; The parameter server sends the same random seed to each client connected to the parameter server through the initialization module to initialize the model parameters of the client's local model, receives the least squares solution of each layer of the local model neural network returned by each client through the aggregation module and performs additive aggregation to obtain the global least squares solution of the same layer of the neural network, and sends the global least squares solution as the global optimal model parameter to each client through the data transmission module; Each of the clients receives a random seed transmitted by the parameter server through an initialization module, and uses the random seed to initialize the local model, obtains non-independent and identically distributed user label data in the training data set through a calculation module, reverse maps the user label data based on a zero-order relaxation backpropagation algorithm, and calculates the least squares solution of each layer of the neural network in the local model. Whenever the least squares solution of a layer of neural network is calculated, the least squares solution of the current layer is transmitted to the parameter server through the data transmission module, and the global least squares solution of each layer of the neural network returned by the parameter server is received in turn. The global least squares solution of each layer of the neural network received is used to update the network parameters of the corresponding layer of the neural network in the local model, and the global optimal model is obtained after completing the update of all layers of the neural network.