Federal learning optimization method for classification tasks

Abstract

The invention belongs to the technical field of federated learning, and particularly relates to a federated learning optimization method for classification tasks. According to the method, several existing ideas for processing the data heterogeneity problem are combined, and then two new regular terms, namely a missing distillation regular term and a dimension maintaining regular term, are provided. In the method, the probability prediction of the missing class output by the local model is limited to be as close as possible to the prediction of the global model, and in addition, the classification vector of the linear classification layer of the historical neural network is stored and the rank of the historical classification vector is regularized, so that the feature learned by the model is prevented from collapsing to a low dimension. The effectiveness of the method is verified through comprehensive experiments.

Description

Technical Field

[0001] The present invention belongs to the technical field of federated learning, and particularly relates to a federated learning optimization method for classification tasks. Background Art

[0002] At present, the remarkable progress made by artificial intelligence technology based on deep learning is largely due to the availability of a large amount of and extensive data. However, with the growth of the data volume, aggregating such data has become very difficult. Therefore, federated learning emerges as a distributed framework for deep learning. Federated learning distributes the training data on numerous edge devices such as smartphones and PCs, and these devices cooperate to train a global model under the orchestration of a central server. Federated learning methods are of great significance in privacy-sensitive deep learning applications. However, federated learning also faces its own challenges, especially in an environment where the data is non-independent and identically distributed (Non-IID). Simply taking the weighted average of the model parameters of each participant under non-independent and identically distributed conditions may lead to a serious decline in performance compared to centralized learning.

[0003] One of the bottlenecks of federated learning is the problem of non-independent and identically distributed data. In an actual federated learning system, each client participating in federated learning often has its own local data distribution. For example, different users have different usage habits when using the same mobile phone APP. In a scenario of non-independent and identically distributed data, there are large differences in data distribution among clients. This not only causes deviations in the gradients of local model updates but also may prevent the global model from quickly achieving coordination and consistency during the aggregation process, thus delaying the convergence speed of the model and reducing its prediction accuracy. At the same time, each client needs to periodically transmit update information to the central server, and in an actual wireless or mobile network, the limited bandwidth often results in low efficiency in transmitting a large number of model parameters or gradient information, increasing the communication burden during the training process.

[0004] Therefore, a key challenge in federated learning is to design strategies to improve the local learning method or global aggregation method of federated learning, overcome the local update bias and global fusion bias caused by data heterogeneity, and improve the performance of the finally obtained global model. At present, in response to the challenges brought by non-independent and identically distributed data, the academic and industrial communities have proposed a variety of solutions from different perspectives, such as the MOON algorithm proposed from the perspective of contrastive learning, the FedProx algorithm proposed from the perspective of parameter consistency, or the FedRS algorithm proposed from the perspective of class missing information. However, these existing technologies usually solve the problem of non-independent and identically distributed data from a single perspective. This limitation makes them perform poorly in complex distribution scenarios and real-world tasks and unable to effectively solve the challenges brought by the combined action of multiple factors. Summary of the Invention

[0005] To solve the above problems, the present invention provides a federated learning optimization method based on a deep neural network for classification tasks.

[0006] The technical solution of the present invention is as follows:

[0007] A federated learning optimization method for classification tasks, which is used in a federated learning system composed of a central server and multiple clients and for classification tasks. The method is characterized in that model training is performed by the clients, the optimization goal is to minimize the loss of the aggregated global model on the entire data set. After the clients finish training the model, the clients upload the model to the central server, and the central server aggregates it to obtain the global model, and then distributes the global model to the clients. A federated learning method for data heterogeneity conditions includes the following steps:

[0008] S1. Establish the optimization goal of the neural network as:

[0009]

[0010] where \(w\in R\) D is the parameter to be optimized, \(R\) represents the real number field, \(D\) represents the number of model parameters, represents the data set, \(N\) is the number of samples, \(l(x\) n ,y n ;w) represents the loss of the model on the sample . Suppose there are \(K\) clients, and each client can access its local data set with a size of Rewrite the optimization goal as:

[0011]

[0012] where \(f(w)\) represents the global optimization goal, and \(l\) n (w) represents the loss of the sample on the model \(w\).

[0013] S2. Calculate the feature representations of the samples under various parameters, including:

[0014] S21. Calculate the feature representation of the current parameters of the client:

[0015]

[0016] where \(R\) represents the feature extraction part of the neural network, represents the parameters of the \(k\)-th client at round \(t\), and \(x\) represents the training sample.

[0017] S22. Calculate the feature representation of the global parameters:

[0018]

[0019] where w t are the parameters of the global model in the current round.

[0020] S23. Calculate the feature representation under the parameters of the previous round of clients

[0021]

[0022] S3. Calculate the classification loss after adding the missing item correction factor:

[0023] S31. Calculate the predicted probability after adding the missing item correction factor:

[0024]

[0025] where C is the set of categories for the entire classification task, is the classification vector corresponding to the m-th category in the linear classification layer of the neural network of the k-th client, the correction factor corresponding to the m-th category of the k-th client. When the client is missing samples of this category, takes a constant value less than 1. This constant is a hyperparameter of the algorithm and is determined through debugging. When the client has samples of this category, takes the value of 1.

[0026] S32. Calculate the classification loss:

[0027] l sup = CrossEntropyLoss(p, y)

[0028] where y is the label corresponding to the sample, and CrossEntropyLoss represents the cross-entropy loss function.

[0029] S4. Calculate the contrastive loss and the distillation loss:

[0030]

[0031] where sim is the cosine similarity function and τ is the temperature parameter (hyperparameter).

[0032]

[0033] where:

[0034]

[0035]

[0036] represents the neural network of the k-th client in the t-th round, Denote the global neural network at the $t$-th round. Let $P$ be the probability prediction vector of sample $x$ calculated by the client model, and $Q$ be the probability prediction vector calculated by the global model. $P(i)$ is the $i$-th element of vector $P$, and $Q(i)$ is the $i$-th element of vector $Q$. Denote the set consisting of the classes missing from the $k$-th client.

[0037] S5. Calculate the dimension-preserving regularization term, including:

[0038] S51. Put the classification vectors of the linear classification layer into a queue:

[0039]

[0040] where is the classification vector corresponding to the $m$-th class of the linear classification layer of the neural network of the $k$-th client in the $t$-th round, arranged in column vector form. $h$ is the length of the queue. Arrange the $h$ $r$ vectors in the history corresponding to the $m$-th class together to obtain matrix $A$ m .

[0041] S52. Normalize the features:

[0042] $A$ m = zscore($A$ m )

[0043] where zscore means performing z-score normalization on each row of the matrix.

[0044] S53. Calculate the regularization term:

[0045]

[0046] where $d$ is the dimension of the sample feature vector, and $\|\cdot\|$ F denotes the Frobenius norm of the matrix.

[0047] S6. Combine the loss functions, perform backpropagation, and update the local model parameters:

[0048] $l = l$ sup $+\mu l$ con $+\beta l$ distill $+\gamma l$ dim

[0049] where $\mu$, $\beta$, $\gamma$ are the weighting coefficients of the three regularization terms, and $l$ is the final loss function.

[0050] After calculating the loss function, update the local model parameters through backpropagation and the gradient descent algorithm.

[0051] S7. Combine the local model parameters to obtain the global model parameter $w$ t, the calculation formula is If the specified number of training rounds is not reached, the global model is redistributed to each client for continued training.

[0052] S8. Construct a federated learning model for classification tasks and conduct training:

[0053] Select a training model based on the requirements of the classification task, collect training data from publicly available datasets, distribute the training data to each client, and after setting the model parameters, train the training model based on the methods of S1 - S7 to finally obtain a trained federated classification model;

[0054] S9. Use the trained federated classification model to classify the target samples.

[0055] The beneficial effects of the present invention are as follows: The present invention integrates various ideas for dealing with data heterogeneity and proposes two new regularization terms on this basis. This method can first ensure that the features learned by the local model do not deviate too much from the global model, while restricting the gradient of the parameters under the condition of class loss. Then, the neural network is regularized through a queue mechanism to prevent the features learned by the network from collapsing to low dimensions, ultimately improving the performance of the global model during testing. The method of the present invention has been comprehensively experimentally verified for its effectiveness. Brief Description of the Drawings

[0056] Figure 1 is the overall flow block diagram of the present invention;

[0057] Figure 2 is the schematic diagram of the sliding window mechanism in the present invention;

[0058] Figure 3 is the schematic diagram of the neural network structure during verification in the present invention.

[0059] Figure 4 is also the schematic diagram of the neural network structure during verification in the present invention. Detailed Embodiments

[0060] The following will describe the specific technical solutions of the present invention with reference to the drawings.

[0061] The present invention combines multiple methods for dealing with data heterogeneity and proposes a mechanism on this basis to slow down the collapse of the representation dimension, as Figure 1 shown is the flow block diagram of the present invention. The present invention specifically includes:

[0062] S1. Establish an optimization objective:

[0063] The training of the neural network can be expressed as the following optimization model:

[0064]

[0065] where \(w\in\mathbb{R}\) D is the parameter to be optimized. represents the dataset, and \(l(x n ,y n ; w)\) represents the loss of the model on the sample . For simplicity, these terms are abbreviated as \(f(w)\) and \(l n (w)\). For the federated classification setting with \(K\) clients, assume that each client can access its local dataset of size \(N k . Then, the objective function can be rewritten as:

[0066]

[0067] S2. Calculate the feature representations of the samples under various parameters:

[0068] Step 1: Calculate the feature representation \(z\) of the sample under the current client parameters , and this process can be expressed as where \(R\) represents the feature extraction part of the neural network, represents the parameters of the \(k\)-th client at round \(t\), and \(x\) represents the training sample.

[0069] Step 2: Calculate the feature representation \(z glob \) under the global parameters, calculate the feature representation of the sample through the global parameters \(w t sent by the server, and this process can be expressed as

[0070] Step 3: Calculate the feature representation \(z \) under the local parameters of the previous round prev , and this process can be expressed as

[0071] S3. Calculate the corrected classification loss:

[0072] Step 1: Calculate the predicted probability after adding the missing item correction factor:

[0073]

[0074] where \(C\) is the set of categories for the entire classification task, is the classification vector corresponding to the \(m\)-th category in the linear classification layer of the neural network of the \(k\)-th client, is the correction factor corresponding to the \(m\)-th category of the \(k\)-th client. When the client lacks samples of this category, takes a constant value, otherwise it is 1.

[0075] Step 2: Calculate the loss between the predicted probability and the label using cross-entropy loss, which can be expressed as: l sup = CrossEntropyLoss(p, y)

[0076] S4. Calculate the contrastive loss and the distillation loss:

[0077] Step 1: Calculate the contrastive loss according to this formula where sim is the cosine similarity function and τ is the temperature parameter (hyperparameter).

[0078] Step 2: Calculate the distillation loss according to this formula where represents the neural network of the k-th client at the t-th round, represents the global neural network at the t-th round. P is the probability prediction vector of the sample x calculated by the client model, and Q is the probability prediction vector calculated by the global model. P(i) is the i-th element of the vector P, and Q(i) is the i-th element of the vector Q. represents the set of missing classes of the k-th client.

[0079] S5. Calculate the dimension-preserving regularization term, including:

[0080] Step 1: Put the classification vectors of the linear classification layer of the neural network into a queue. This operation is denoted in the following form

[0081]

[0082] where is the classification vector corresponding to the m-th class of the linear classification layer of the neural network of the k-th client at the t-th round. h is the length of the queue. When calculating, we arrange the r vectors in the form of column vectors. After putting the latest vector into the queue, we remove the earliest vector, arrange the h r vectors in the history corresponding to the m-th class together to get a complete queue, and splice the vectors in the queue row by row to get the matrix A m .

[0083] Step 2: Perform z-score normalization on each row of the matrix A m .

[0084] Step 3: Calculate the dimension-preserving regularization term l dim according to the following formula:

[0085]

[0086] where d is the dimension of the sample feature vector, and ||·|| F represents the Frobenius norm of the matrix.

[0087] S6. Combine the loss function according to the formula \(l = l sup +\mu l con +\beta l distill +\gamma l dim (where \(\mu\), \(\beta\), and \(\gamma\) are the weighting coefficients of three regularization terms, and \(l\) is the final loss function), then perform backpropagation and update the local model parameters through gradient descent or other optimization algorithms.

[0088] S7. Combine the local model parameters according to the formula to obtain the global model parameter \(w t . If the specified number of training rounds is not reached, redistribute the global model to each client and return to step S2 to continue training.

[0089] S8. Build a federated learning model for image classification tasks and conduct training and testing;

[0090] S81. Dataset selection and determination of hyperparameters

[0091] Dataset. The present invention focuses on image classification tasks. Extensive experiments were conducted on three public datasets: CIFAR-100, Tiny-Imagenet, and CIFAR-10, and standard training and test splits were used. We generated an imbalanced dataset through the Dirichlet sampling method, with the number of clients being 10 and the distribution parameter being 0.5.

[0092] Baseline algorithms. The present invention compares the present invention with several baseline algorithms: FedAvg, FedRS, and MOON.

[0093] Hyperparameters. For all datasets, the training batch size is 64, the learning rate is set to 0.001, the momentum is 0.9, the local training period is 10, and the global training period is 20. For the MOON algorithm, the temperature parameter \(\tau\) is set to 0.5 on all three datasets, the contrast loss weight \(\mu\) is set to 5 on the CIFAR-10 dataset and 1 on the CIFAR-100 and tiny-ImageNet datasets. The logits weighting coefficient of the FedRS algorithm is set to 0.5. The hyperparameters of the algorithm proposed in the present invention are set as follows: \(\tau = 0.5\), \(\alpha = 0.8\) when the corresponding class is a missing class, \(h = 15\), 5 on the CIFAR-10 dataset and 1 on the CIFAR-100 and tiny-ImageNet datasets, \(\beta = 0.01\), \(\gamma = 10\). The network structure used on the CIFAR-10 dataset is a two-layer convolutional neural network plus a feature projection head, as Figure 3As shown. For the CIFAR-100 and tiny-ImageNet datasets, Resnet50 plus a feature projection head is used. As Figure 4 shown, the structure of Resnet will not be elaborated here. The projection heads of the three datasets project the features to 256 dimensions for the contrastive learning of MOON.

[0094] S52. Validation experiments based on the CIFAR-100, Tiny-Imagenet, and CIFAR-10 image classification datasets;

[0095] The present invention focuses on the CIFAR-100, Tiny-Imagenet, and CIFAR-10 datasets and verifies the advantages of the algorithm proposed in the present invention compared with the baseline algorithms.

[0096] Table 1. Test accuracies achieved by different methods

[0097] CIFAR-10 CIFAR-100 Tiny-Imagenet Method Acc Acc Acc FedAvg 0.6784 0.6503 0.2255 MOON 0.7002 0.6695 0.2327 FedRS 0.7007 0.6635 0.2369 This invention 0.7193 0.6881 0.2493

[0098] Specifically, for the CIFAR-10 dataset, the present invention has a nearly 5% accuracy improvement compared with the FedAvg algorithm without any regularization and a about 2% improvement compared with the baseline method with regularization. For the CIFAR-100 dataset, the present invention has a nearly 4% accuracy improvement compared with the FedAvg algorithm without any regularization and a about 2% improvement compared with the baseline method with regularization. For the Tiny-Imagenet dataset, the present invention has a nearly 2% accuracy improvement compared with the FedAvg algorithm without any regularization and a about 1% improvement compared with the baseline method with regularization. This emphasizes that the method of the present invention has significant advantages on datasets of various scales and is significantly better than the baseline methods for comparison.

[0099] In practical applications, this improvement in accuracy can enable the federated learning system to have higher accuracy, making relevant enterprises more willing to train models under the premise of privacy protection.

[0100] Of course, the above description is not a limitation of the present invention, and the present invention is not limited to the above examples. Changes, modifications, additions, or substitutions made by those skilled in the art within the scope of the essence of the present invention should also fall within the protection scope of the present invention.

Claims

1. A federated learning optimization method for classification tasks. The system implementing the federated learning optimization method includes a central server and multiple clients. The clients are used to train a neural network model, with the goal of minimizing the loss of the aggregated global model on the entire dataset. The clients upload the trained models to the central server, and the central server aggregates them to obtain the global model, and then distributes the global model to the clients for classification tasks; characterized in that, It includes the following steps: S1. Establish the optimization objective of the neural network as: where \(w\in\mathbb{R}\) D is the parameter to be optimized, \(\mathbb{R}\) represents the real number field, \(D\) represents the number of model parameters, represents the dataset, \(N\) is the number of samples, \(l(x n ,y n ; w)\) represents the loss of the model on the sample ; There are K clients set, and each client can access its local dataset with a size of N k , The optimization objective is rewritten as: where f(w) represents the global optimization objective, and l n (w) represents the loss of the sample on the model w; S2. Calculate the feature representations of the samples under various parameters, including: S21. Calculate the feature representation of the client under the current parameters: where R represents the feature extraction part of the neural network, represents the parameters of the k-th client at round t, and x represents the training sample; S22. Calculate the feature representation under the global parameters: where w t are the parameters of the global model in the current round; S23. Calculate the feature representation under the parameters of the previous round of the client: S3. Calculate the classification loss after adding the missing item correction factor: S31. Calculate the predicted probability after adding the missing item correction factor: where \(C\) is the set of categories for the entire classification task, is the classification vector corresponding to the \(m\)-th class in the linear classification layer of the neural network of the \(k\)-th client, is the correction factor corresponding to the \(m\)-th class of the \(k\)-th client. When the client lacks samples of this category, it takes a constant value less than 1. When the client has samples of this category, it takes the value of 1; S32. Calculate the classification loss: l sup = CrossEntropyLoss(p, y), where y is the label corresponding to the sample, and CrossEntropyLoss represents the cross-entropy loss function; S4. Calculate the contrastive loss and the distillation loss: where sim is the cosine similarity function and τ is the temperature parameter; where: Denote the neural network of the k-th client at the t-th round. Denote the global neural network at the t-th round. P is the probability prediction vector of the sample x calculated by the client model, and Q is the probability prediction vector calculated by the global model; P(i) is the i-th element of the vector P, and Q(i) is the i-th element of the vector Q. Denote the set composed of the missing categories of the k-th client. S5. Calculate the dimension-preserving regularization term, including: S51. Put the classification vector of the linear classification layer into a queue: Among them, is the classification vector corresponding to the m-th class of the linear classification layer of the neural network in the t-th round of the k-th client, arranged in column vector form; h is the length of the queue, and the r vectors in the h histories corresponding to the m-th class are arranged together to obtain matrix A m ; S52. Normalize the features: A m = zscore(A m ) where zscore represents performing z-score normalization on each row of the matrix; S53. Calculate the regularization term: where d is the dimension of the sample feature vector, and ||·|| F represents the Frobenius norm of the matrix; S6. Combine the loss function, perform backpropagation, and update the local model parameters: l = l sup + μl con + βl distill + γl dim , where μ, β, and γ are the weighting coefficients of the three regularization terms, and l is the final loss function; After calculating the loss function, update the local model parameters through backpropagation and the gradient descent algorithm; S7. Merge the local model parameters to obtain the global model parameter w t , and the calculation formula is If the specified number of training rounds is not reached, redistribute the global model to each client for continued training; S8. Construct a federated learning model based on the classification task and train it: Select a training model according to the requirements of the classification task, collect training data from the publicly available dataset, distribute the training data to each client, and train the training model based on the methods of S1 - S7 after setting the model parameters, and finally obtain a trained federated classification model; S9. Use the trained federated classification model to classify the target samples.

Images