Federal learning model aggregation optimization method capable of adaptively adjusting weight and learning rate

Abstract

The invention discloses a federal learning model aggregation optimization method for adaptive adjustment of weight and learning rate, and the core of the method lies in collaborative optimization of aggregation weight and local learning rate: on one hand, through introducing an adaptive weight mechanism, the contribution degree of a low-performance client is improved, and the low-performance client is prevented from being marginalized in long-term training; on the other hand, an appropriate learning rate is configured for the client according to the computing power, the learning process of the vulnerable party is accelerated, and global convergence is balanced. According to the method, the problems of slow convergence and generalization reduction caused by calculation performance difference of clients in federated learning are solved, the method has remarkable advantages in convergence speed, communication round and final precision, and the training efficiency and accuracy are improved.

Description

Technical Field

[0001] This invention belongs to the field of machine learning technology, and in particular, relates to a method for aggregation optimization of federated learning models with adaptive adjustment of weights and learning rates. Background Technology

[0002] Federated learning, as an innovative distributed machine learning paradigm, boasts the core advantage of allowing participants to collaboratively train a global model without sharing their local raw data, demonstrating significant application potential in privacy-sensitive scenarios such as finance and healthcare. However, in practical deployments, federated learning faces challenges due to varying client-side computing performance. Significant differences in computing power between different terminal devices lead to inconsistent client-side training progress and unbalanced model update steps, resulting in slow or even biased convergence of the global model. Therefore, designing optimization algorithms that effectively address these performance differences is crucial for improving the practicality and fairness of federated learning and remains a critical scientific problem that urgently needs to be solved.

[0003] Currently, model aggregation strategies to address computational performance differences can be summarized in three main aspects. First, there are communication-based strategies, such as asynchronous learning. This aggregation method is suitable for resource-constrained devices, improving the overall efficiency and scalability of federated learning. However, this immediate aggregation approach can lead to inconsistencies between the updates of slow nodes and the global model, potentially reducing the training efficiency and performance of the federated learning model. Second, there are client-selection-based model aggregation strategies. These accelerate training by selecting efficient devices, which may speed up the process but deprives slower clients of the opportunity to participate, potentially missing some data and affecting the model's training effectiveness. Finally, there are regularization-based model operations. These methods reduce local biases in client models through different regularization operations or normalization constraints on parameter update directions, thereby improving the generalization ability of the global model. FedProx introduces regularization terms to mitigate biases based on FedAvg, but its suppression of model biases caused by computational performance differences is insufficient. FedNova normalizes client-local updates before uploading them for weighted summation, reducing model bias caused by unbalanced local iterations and partially alleviating the slow model convergence problem caused by computational performance differences. However, none of the above methods have solved the problem of computational performance differences at the root. They either passively adapt or impose one-sided constraints, but they all ignore the collaborative optimization of the two core parameters, learning rate and aggregate weight.

[0004] A patent application with publication number CN119647546A discloses a federated learning algorithm based on adaptive aggregation and dynamic training, belonging to the field of deep learning. The specific steps of this learning algorithm are as follows: S1, each client receives the initial global model parameters and trains the model using its local data; S2, stochastic gradient descent is used to update the model weights locally; S3, after local training is completed, the client calculates the model loss value for the current round and retains this loss as a reference for subsequent adaptive aggregation; S4, on the client side, each client uses a dynamic training strategy based on its historical training time to predict the time requirement for the next round of training. While the general framework of this patent provides a foundation for solving the problem of computing power differences, the extreme performance optimization for specific hardware architectures (such as from low-power MCUs to high-end GPUs), and the new security challenges arising from this (such as clients with weaker computing power being more vulnerable to attacks), have not yet been explored in depth in the application. Summary of the Invention

[0005] This invention addresses the shortcomings of existing technologies, such as FedProx, which introduces regularization terms to mitigate model bias by building upon FedAvg, but which does not adequately suppress model bias caused by differences in computational performance. FedNova, by normalizing client-side local updates before uploading them for weighted summation, reduces model bias caused by unbalanced local iterations and partially alleviates the slow model convergence problem caused by differences in computational performance. This invention proposes a federated learning model aggregation optimization method with adaptive adjustment of weights and learning rates.

[0006] An aggregation optimization method for federated learning models with adaptive adjustment of weights and learning rates includes the following steps: S1. Initialize the global model through the central server and distribute the global model to N clients participating in federated learning; each client performs local training based on its local dataset and records the actual training rounds of each client during the local training process. and local data volume ; S2. The central server receives local training feedback information uploaded by each client, the feedback information including the actual training rounds of the client. Local data volume and the number of steps to update gradient descent during local training; S3. In the global model aggregation stage of federated learning, the aggregation weights and learning rates are dynamically adjusted according to the computing power of each client by collaboratively optimizing the aggregation weights and learning rates. Specifically, the aggregation weights of clients with low computing power are increased, and clients with strong computing performance are given smaller learning rates, while clients with weak performance are given larger learning rates. S4. The central server performs weighted aggregation on the model update results uploaded by each client according to the adjusted aggregation weight, and generates a new global model. S5. Repeat steps S1 to S4, iterating until the global model meets the convergence condition, and complete the federated learning training.

[0007] Further, in step S1, each client uses a stochastic gradient descent optimizer to perform local model training to minimize the average loss of the local data, expressed as:

[0008] In the above formula, To participate in federal learning Average loss of local data per client For the first Local model parameters for each client. , For input data and labels, The loss function used for each client.

[0009] Furthermore, in step S3, the expression for the dynamically adjusted weights is as follows:

[0010] In the above formula, For the first Each client-side global model updates the aggregate weight amplification factor. The target number of local training rounds for the federated training client. For the first The actual number of update rounds for each client. The aggregation coefficient represents the proportion of data. For the first Data volume per client This represents the total amount of data.

[0011] Furthermore, in step S3, the expression for dynamically adjusting the learning rate is as follows:

[0012] In the above formula, Indicates the first Each client updates the gradient descent steps in a training round. The initial uniform learning rate.

[0013] Furthermore, the computing power updates the number of steps through local training gradient descent within one round. The expression is as follows:

[0014] In the above formula, This refers to the size of the mini-batch samples used for local training on the client side. The number of clients participating in federated learning.

[0015] Furthermore, the local dataset is a multi-classification task dataset for privacy-sensitive scenarios, selected from the rail defect vibration signal dataset or the CIFAR10 image classification dataset. The training objective of the global model is to minimize the weighted average loss of all clients, and its expression is:

[0016] In the above formula, This represents the weighted average loss for all clients under the current parameters, i.e., the objective function.

[0017] Furthermore, when the local dataset is a rail defect vibration signal dataset, the initial global model adopts a LeNet neural network structure; when the local dataset is a CIFAR10 image classification dataset, the initial global model adopts a convolutional neural network structure.

[0018] Furthermore, in step S5, the convergence condition includes: the accuracy of the global model on the test set reaches a preset threshold or the fluctuation of the global model's test accuracy is less than 0.005 in 10 consecutive rounds of communication iterations.

[0019] Furthermore, when the local dataset is for rail defect identification, the preset threshold for the convergence condition is a test accuracy of not less than 0.9; when the local dataset is for CIFAR10 image classification dataset, the preset threshold for the convergence condition is a test accuracy of not less than 0.56.

[0020] Furthermore, for multi-class classification tasks, an accuracy evaluation metric and a loss value are used for comprehensive evaluation. The expression for the accuracy is as follows:

[0021] In the above formula, , , and These represent the true negative class, the false negative class, the false positive class, and the true negative class, respectively.

[0022] Compared with the prior art, the beneficial effects of the present invention are as follows: This invention collaboratively optimizes aggregate weights and local learning rates: on the one hand, by introducing an adaptive weight mechanism, it enhances the contribution of low-performance clients, preventing them from being marginalized during long-term training; on the other hand, it configures appropriate learning rates for clients based on their computing power, accelerating the learning process of weaker clients and balancing global convergence. Experimental results show that compared with mainstream algorithms such as FedAvg, FedProx, and FedNova, FedAWR exhibits significant advantages in convergence speed, communication rounds, and final accuracy, achieving higher training efficiency and accuracy. Furthermore, generalization tests on the CIFAR-10 dataset further confirm the superiority of FedAWR, demonstrating faster accuracy improvement and lower loss convergence, fully validating the effectiveness of this method for federated learning optimization in heterogeneous computing environments. Attached Figure Description

[0023] Figure 1 This is a flowchart of the present invention; Figure 2 This is a schematic diagram of the FedAWR parameter update in this invention; Figure 3 This is a schematic diagram illustrating the average number of training steps for each client in this invention; Figure 4 Convergence curves of test accuracy for different federated learning methods on the rail dataset; Figure 5 The test loss convergence curves for different federated learning methods on the rail dataset; Figure 6 Convergence curves of test accuracy for different federated learning methods on the CIFAR-10 dataset; Figure 7 The test loss convergence curves for different federated learning methods on the CIFAR-10 dataset. Detailed Implementation

[0024] To clearly illustrate the technical features of the present invention, the invention will be described in detail below through specific embodiments and in conjunction with the accompanying drawings. Many specific details are set forth in the following description to provide a thorough understanding of the invention; however, the invention may be implemented in other ways different from those described herein, and therefore, the scope of protection of the invention is not limited to the specific embodiments disclosed below. In the present invention, unless otherwise expressly specified and limited, the first feature "on" or "below" the second feature may mean that the first and second features are in direct contact, or that the first and second features are in indirect contact through an intermediate medium. In the description of this specification, references to terms such as "an embodiment," "some embodiments," "example," "specific example," or "some examples," etc., indicate that the specific feature, structure, material, or characteristic described in connection with that embodiment or example is included in at least one embodiment or example of the invention. In this specification, the illustrative expressions of the above terms do not necessarily refer to the same embodiment or example. Moreover, the specific features, structures, materials, or characteristics described may be combined in any suitable manner in one or more embodiments or examples.

[0025] Example 1 like Figure 1 and Figure 2 As shown, a federated learning model aggregation optimization method with adaptive adjustment of weights and learning rate includes the following steps: S1. Initialize the global model through the central server and distribute the global model to N clients participating in federated learning; each client performs local training based on its local dataset and records the actual training rounds of each client during the local training process. and local data volume ; S2. The central server receives local training feedback information uploaded by each client, the feedback information including the actual training rounds of the client. Local data volume and the number of steps to update gradient descent during local training; S3. In the global model aggregation phase of federated learning, the aggregation weights and learning rates are dynamically adjusted based on the computing power of each client. Specifically, the aggregation weights of clients with low computing power are increased, and clients with high computing performance are given smaller learning rates, while clients with low performance are given larger learning rates. Figure 2 As shown, to mitigate model bias caused by differences in computing performance, a weight adjustment strategy based on computing performance is introduced during the aggregation phase. By increasing the aggregation weight of clients with low computing power, the updated global model parameters can be optimized. Towards global optimum Convergence, avoiding local optima biased towards high-computing-power clients. The weight adjustment formula is as follows:

[0026] In the above formula, For the first Each client-side global model updates the aggregate weight amplification factor. The target number of local training rounds for the federated training client. For the first The actual number of update rounds for each client. The aggregation coefficient represents the proportion of data. For the first Data volume per client Total data volume; S4. The central server performs weighted aggregation on the model update results uploaded by each client according to the adjusted aggregation weight, and generates a new global model. S5. Repeat steps S1 to S4, iterating until the global model meets the convergence condition, and complete the federated learning training.

[0027] In this embodiment, by enhancing the contribution weights of key clients during the aggregation phase and adaptively adjusting the learning rate based on computing power, the model update direction is dynamically optimized and global convergence efficiency is improved. By comprehensively considering the computing power and training contribution of each client, the global aggregation weights and local learning rates are dynamically allocated to balance the impact of high and low computing power devices on model updates, thereby improving the efficiency, stability, and generalization performance of global model training. Figure 2 As shown. By increasing the aggregation weight of the low-computing-power clients on the left, the adjusted aggregation weight is the same as the original weight allocated based on data. This is 5 / 2 times the size of the model, which can prevent it from being marginalized during long-term training, thereby correcting the update direction of the global model and enhancing its generalization performance.

[0028] To mitigate the issues of inappropriate learning rate configuration and decreased convergence efficiency caused by varying computational performance on different clients, the learning rate is dynamically adjusted based on the client's computational capabilities: clients with high computational performance use a smaller learning rate to prevent over-updates, while clients with lower performance use a larger learning rate to accelerate convergence. Computational capabilities are assessed by the number of steps updated via local training gradient descent within one round. The expressions for the training steps and learning rate adjustment process are as follows.

[0029]

[0030]

[0031] In the above formula, Indicates the first Each client updates the gradient descent steps in a training round. This refers to the size of the mini-batch samples used for local training on the client side. The number of clients participating in federated learning. The initial learning rate is uniform. The server performs the number of steps of local gradient descent through the client. This allows us to assess its computational capabilities and dynamically configure an appropriate learning rate accordingly to improve overall convergence efficiency.

[0032] The aggregation strategy in this embodiment effectively mitigates the model update direction shift caused by differences in client computing performance by adaptively compensating clients with weaker computing capabilities. Simultaneously, we dynamically configure the learning rate of each client based on its computing performance to accelerate the convergence process of the global model.

[0033] Example 2 In this embodiment, the optimization effect of the FedAWR algorithm of the present invention on the heterogeneous computing performance of the client in an industrial privacy-sensitive scenario (rail defect identification) is verified, achieving fast convergence and high accuracy in defect classification.

[0034] Specifically, the federated learning experimental system consists of: one central server equipped with an NVIDIA GeForce RTX 4060 GPU and 32GB of memory, and 10 clients configured (simulating computing performance differences through VMware Workstation virtual machines).

[0035] Key parameter settings: The same LeNet neural network architecture is used to ensure consistent model configuration when compared with the baseline method; the number of clients N=10; the local dataset consists of rail defect vibration signals, with each client corresponding to data from a different railway section; the communication iteration rounds are 600 rounds, and the target number of local training rounds is [number missing]. Initial uniform learning rate Batch size .

[0036] Convergence criteria: The global model achieves an accuracy of ≥0.90 on the test set, and the accuracy improvement over 10 consecutive rounds is less than 0.005.

[0037] Dataset Preprocessing: The experimental dataset originates from vibration signals collected from different sections of a railway locomotive depot in China. Each section's data serves as a client, and the dataset was acquired using accelerometers integrated into a rail flaw detection vehicle, aiming to realistically reflect the vibration characteristics of rails under operational conditions. Considering the uneven energy distribution, noise interference, and non-stationarity of the raw vibration data, we implemented a standardized preprocessing procedure. This procedure includes: random shuffling to ensure sample randomness; peak value normalization to unify data scale; truncating the first 2500 time steps to meet model input requirements; and finally, dividing the data into training and test sets in a 7:3 ratio and generating corresponding labels. This procedure effectively overcomes the inherent defects of the raw data, improves data quality, and provides a reliable data foundation for subsequent model training.

[0038] Specifically, the verification was performed according to the steps in Example 1, and the experimental results were recorded and analyzed. To quantify the differences in computational performance between clients, the average number of update steps for each client in each training round was statistically analyzed. The average training steps for each client are shown below. Figure 3 As shown.

[0039] Significant heterogeneity exists in the average update steps among different clients. Clients 7, 8, and 9, with higher performance, have average update steps of 42, 49, and 46, respectively; while clients 4 and 6, with lower performance, have average update steps of only 2 and 3, respectively. This large difference in step count directly reflects the widespread imbalance in client computing power within federated learning environments, providing an experimental basis for validating the effectiveness of the proposed method in heterogeneous environments.

[0040] For multi-classification tasks, a comprehensive evaluation is conducted using accuracy metrics and loss values. The specific expression for accuracy is shown below.

[0041]

[0042] In the formula, , , and These represent the true negative class, the false negative class, the false positive class, and the true negative class, respectively.

[0043] This embodiment uses accuracy, loss value, and the number of communication rounds required to reach the target accuracy to measure the model's performance. For example... Figure 4As shown, the accuracy of all algorithms increases with the number of communication rounds, but there are significant differences in convergence speed and final performance: the FedAWR algorithm rapidly improves accuracy in the early stages, reaching close to 0.90 at around 188 rounds, and then consistently leads and stabilizes above 0.93; FedAvg converges the slowest, only gradually approaching 0.90 after 300 rounds, and eventually stabilizing at that level. The performance of FedProx and FedNova falls between the two, both reaching 0.90 around 300 rounds and then stabilizing.

[0044] In this embodiment, cross-entropy is used as the loss function to evaluate the generalization performance of each algorithm on multi-classification tasks. The specific expression is shown below.

[0045]

[0046] In the above formula, For the number of samples, For the number of categories, For the first The sample at the th Real labels in each category The model predicts the first Class probability.

[0047] The test loss convergence curves of different federated learning methods on the rail dataset are as follows: Figure 5 As shown in Table 1, the four algorithms exhibit significant differences in the rate of loss reduction across communication rounds. Specifically, FedAvg demonstrates a slower convergence speed, with a slower decrease in loss value, requiring more communication rounds to approach convergence. In contrast, FedProx and FedNova show slightly improved rates of loss reduction, but still lower than FedAWR. The proposed algorithm achieves a significant decrease in loss value within fewer communication rounds, demonstrating a significantly faster rate of loss reduction. The experimental results for the final accuracy, loss value, and number of rounds required to achieve the target precision are shown in Table 1.

[0048] Table 1 shows the final accuracy, loss value, and number of rounds required to achieve the target accuracy for each method on the rail dataset.

[0049] Compared to the baseline methods, the FedAWR method in this embodiment exhibits superior convergence performance. Specifically, FedAWR achieves a convergence speed 1.72 times faster than FedAvg, and also surpasses FedProx and FedNova, with a final model accuracy of 95.1%, outperforming the other three methods.

[0050] FedAWR's superiority stems primarily from its adaptive mechanism. By dynamically adjusting the learning rate of each client and the aggregation weights of the server, FedAWR guides the global model towards a more optimal direction, thereby effectively improving training efficiency. This advantage is particularly evident on the rail damage dataset: due to significant differences in computational performance among clients, FedAWR achieves faster convergence and lower loss values ​​by linking the aggregation weights and learning rate to the device's computing power.

[0051] Example 3 To further verify the adaptability and generality of this application, generalization verification experiments were conducted on the publicly available image classification dataset CIFAR-10. In the comparative experiments, four algorithms—FedAvg, FedProx, FedNova, and the proposed FedAWR—were selected for performance evaluation. The experimental model adopted a standard convolutional neural network (CNN) structure. The hyperparameters were set as follows: 10 clients, 300 communication rounds, 10 local training rounds, stochastic gradient descent as the local optimizer, a learning rate of 0.001, and a batch size of 24.

[0052] The performance metrics for the generalization experiment in this embodiment are the same as those in the previous section. Under this setting, the test accuracy convergence curves of different federated learning methods on the CIFAR-10 dataset are as follows: Figure 6 As shown in the figure. Experimental results demonstrate that the FedAWR algorithm proposed in this application significantly outperforms the baseline methods in terms of generalization ability. Specifically, although the test accuracy of all algorithms increases with the number of training epochs, FedAWR exhibits a faster convergence speed and higher final accuracy. In the first 100 epochs of training, the accuracy of FedAWR rapidly climbs to over 0.58, significantly outperforming FedAvg and FedProx. At the end of training, the final performance of FedAWR also surpasses FedNova, achieving global optimum.

[0053] like Figure 7 As shown, the test losses of different federated learning methods converged on the CIFAR-10 dataset. In terms of loss descent rate, FedAWR showed a significant advantage in the first 150 rounds of communication, with its loss value decreasing significantly faster than FedAvg and FedProx. In the later stages of convergence, the final loss value of FedAWR reached a similar level to FedNova. The generalization experiment results are shown in Table 2.

[0054] Table 2. Final accuracy, loss value, and number of rounds required to achieve target accuracy for each method on the CIFAR-10 dataset.

[0055] Analysis of accuracy and loss values ​​shows that the FedAWR method outperforms other algorithms in both convergence speed and loss optimization. In particular, FedAWR achieves the target accuracy after 80 rounds of communication, faster than other methods, and its final accuracy is higher than the other three methods. Generalization experiments demonstrate that this method converges more efficiently. Furthermore, the method in this application also shows superiority in loss value reduction, optimizing the loss in a shorter time, further proving its better stability during training.

[0056] Obviously, the embodiments described above are merely examples for clearly illustrating the present invention and are not intended to limit the implementation of the present invention. Those skilled in the art can make other variations or modifications based on the above description. It is neither necessary nor possible to exhaustively list all possible implementations. Any modifications, equivalent substitutions, and improvements made within the spirit and principles of the present invention should be included within the scope of protection of the claims of the present invention.

Claims

1. A federated learning model aggregation optimization method with adaptive adjustment of weights and learning rate, characterized in that, Includes the following steps: S1. Initialize the global model through the central server and distribute the global model to N clients participating in federated learning; each client performs local training based on its local dataset and records the actual training rounds of each client during the local training process. and local data volume ; S2. The central server receives local training feedback information uploaded by each client, the feedback information including the actual training rounds of the client. Local data volume and the number of steps to update gradient descent during local training; S3. In the global model aggregation stage of federated learning, the aggregation weights and learning rates are dynamically adjusted according to the computing power of each client by collaboratively optimizing the aggregation weights and learning rates. Specifically, the aggregation weights of clients with low computing power are increased, and clients with strong computing performance are given smaller learning rates, while clients with weak performance are given larger learning rates. S4. The central server performs weighted aggregation on the model update results uploaded by each client according to the adjusted aggregation weight, and generates a new global model. S5. Repeat steps S1 to S4, iterating until the global model meets the convergence condition, and complete the federated learning training.

2. The method for aggregation optimization of federated learning models with adaptive adjustment of weights and learning rates according to claim 1, characterized in that, In step S1, each client uses a stochastic gradient descent optimizer to perform local model training, which minimizes the average loss of the local data. The expression is: In the above formula, To participate in federal learning Average loss of local data per client For the first Local model parameters for each client. , For input data and labels, The loss function used for each client.

3. The method for aggregation optimization of federated learning models with adaptive adjustment of weights and learning rates according to claim 1, characterized in that, In step S3, the expression for dynamically adjusting the weights is as follows: In the above formula, For the first Each client-side global model updates the aggregate weight amplification factor. The target number of local training rounds for the federated training client. For the first The actual number of update rounds for each client. The aggregation coefficient represents the proportion of data. For the first Data volume per client This represents the total amount of data.

4. The method for aggregation optimization of federated learning models with adaptive adjustment of weights and learning rates according to claim 1, characterized in that, In step S3, the expression for dynamically adjusting the learning rate is as follows: In the above formula, Indicates the first Each client updates the gradient descent steps during a training round. The initial uniform learning rate.

5. The method for aggregation optimization of federated learning models with adaptive adjustment of weights and learning rates according to claim 1, characterized in that, The computing power is updated by local training gradient descent steps within one round. The expression is as follows: In the above formula, This refers to the size of the mini-batch samples used for local training on the client side. The number of clients participating in federated learning.

6. The method for aggregation optimization of federated learning models with adaptive adjustment of weights and learning rates according to claim 1, characterized in that, The local dataset is a multi-classification task dataset for privacy-sensitive scenarios, selected from either the rail defect vibration signal dataset or the CIFAR10 image classification dataset. The training objective of the global model is to minimize the weighted average loss of all clients, expressed as: In the above formula, This represents the weighted average loss for all clients under the current parameters, i.e., the objective function.

7. The method for aggregation optimization of federated learning models with adaptive adjustment of weights and learning rates according to claim 6, characterized in that, When the local dataset is a rail defect vibration signal dataset, the initial global model adopts a LeNet neural network structure; when the local dataset is a CIFAR10 image classification dataset, the initial global model adopts a convolutional neural network structure.

8. The method for aggregation optimization of federated learning models with adaptive adjustment of weights and learning rates according to claim 1, characterized in that, In step S5, the convergence condition includes: the accuracy of the global model on the test set reaches a preset threshold or the fluctuation of the global model's test accuracy is less than 0.005 in 10 consecutive rounds of communication iterations.

9. The method for aggregation optimization of federated learning models with adaptive adjustment of weights and learning rates according to claim 8, characterized in that, When the local dataset is for rail defect identification, the preset threshold for the convergence condition is a test accuracy of not less than 0.9; when the local dataset is for CIFAR10 image classification dataset, the preset threshold for the convergence condition is a test accuracy of not less than 0.

56.

10. The method for aggregation optimization of federated learning models with adaptive adjustment of weights and learning rates according to claim 1, characterized in that, For multi-class classification tasks, a comprehensive evaluation is conducted using both accuracy and loss values. The expression for accuracy is as follows: In the above formula, , , and These represent the true negative class, the false negative class, the false positive class, and the true negative class, respectively.

Images