Model compression federated learning framework based on interpretable optimization heuristic deep network

Abstract

The invention discloses a model compression federated learning framework based on an interpretable optimization heuristic deep network, and the method comprises a pre-training step: S1, collecting multi-round gradient data of a plurality of models on different data sets, carrying out the vectorization and partitioning of gradient parameters, and constructing a training data set; a network training step S2, based on an ISTA principle, designing and training an explainable optimized heuristic deep network as a gradient codec; a federated learning step S3, the server initializes and distributes a global model, and after the client performs local training, a sampling matrix is used to compress a gradient and upload the gradient; s4, the server reconstructs the gradient by using the optimized heuristic deep network and corrects the gradient; and S5, the server aggregates the gradient to update the global model and issues the global model. And further, the server circularly extracts the clients for multiple rounds of training until the model meets a preset condition. The method has the advantages of low communication overhead, good model performance, wide applicability and the like.

Description

Technical Field

[0001] The present invention relates to the fields of federated learning and model compression, and in particular to a method and system for compressing a federated learning model based on an interpretable optimization heuristic deep network. Background Art

[0002] Federated learning, a distributed machine learning method, allows multiple clients to collaboratively train machine learning models by transmitting gradients without uploading datasets to a central server, effectively reducing the risk of privacy leaks. However, in practice, communication overhead is a major challenge for federated learning. Due to the large number of model parameters, transmitting these parameters consumes significant bandwidth resources. This communication overhead becomes a bottleneck limiting the application of federated learning, especially when client devices have limited network connectivity.

[0003] Compressed sensing (CS), a signal processing method that samples and reconstructs signals, can accurately reconstruct the original signal at sampling frequencies far below the Nyquist sampling frequency. This has made CS technology extremely popular in various fields, such as the compression and reconstruction of natural images. This is because images exhibit structured sparsity to some extent. Image reconstruction is achieved by solving a sparsity-regularized optimization problem. This approach offers strong convergence in most cases and facilitates theoretical analysis, but consumes significant computational resources and time in practice. Meanwhile, CS reconstruction methods based on deep learning have made significant progress. By directly learning the inverse mapping from the measured signal to the original signal, they significantly reduce communication overhead and achieve excellent performance. However, these methods suffer from two major drawbacks: a lack of interpretability and limited generalization. To address these issues, researchers proposed an interpretable optimization-inspired deep network, ISTA-Net. ISTA-Net maps the L1 normalization process in ISTA to network layers, with a fixed number of layers similar to one iteration of the algorithm. This approach retains the efficiency of deep learning while inheriting the interpretability of traditional algorithms. This method has achieved significant success in the field of image compression.

[0004] Existing technologies primarily address the communication overhead in federated learning by using two approaches: gradient sparsification, which transmits only the important gradient parameters; and gradient quantization, which reduces the precision of each parameter. Compressed sensing (CS) is being proposed as an orthogonal technique, combining it with these two approaches to further improve compression. The gradient parameters uploaded by federated learning are generally normally distributed, fully meeting the sparsity requirements of CS. However, similar approaches in the field of image compression, combining traditional neural networks, have not been explored. How to handle the larger and more complex feature distributions of model gradients in federated learning compared to images? Whether gradients trained on different model structures and datasets generalize well is also a pressing issue. Finally, because federated learning is an iterative process, errors in each reconstruction round can accumulate, leading to a sharp decline in model performance. Therefore, methods to correct for erroneous reconstructions should be developed accordingly. Summary of the Invention

[0005] The technical problem to be solved by the present invention is: In response to the problems existing in the existing technology, the present invention provides a model compression federated learning framework based on an interpretable optimization-heuristic deep network. This method uses the structure of the optimization algorithm to design the deep network model to achieve efficient gradient compression and accurate reconstruction. At the same time, it adopts a gradient correction strategy to prevent error accumulation, reduce the communication overhead of federated learning, and maintain model performance.

[0006] To solve the above technical problems, the present invention proposes a technical solution: a federated learning model compression method based on an interpretable optimization heuristic deep network, comprising the following steps:

[0007] S1, the pre-training stage, collects multiple rounds of gradient data generated by training multiple models on different datasets, vectorizes and divides the gradient parameters into blocks, and constructs a dataset for training and optimizing heuristic deep networks;

[0008] S2, the network training phase, designs and trains an interpretable optimization heuristic deep network as a gradient encoder-decoder based on the principle of iterative soft thresholding algorithm (ISTA);

[0009] S3, the federated learning initialization phase, where the server initializes the global model and distributes the model to the clients participating in the training;

[0010] S4, client local training phase, the client trains the model based on local data and calculates the gradient;

[0011] S5, gradient compression phase, the client compresses the gradient using a predefined sampling matrix and uploads the compressed gradient to the server;

[0012] S6, gradient reconstruction phase, the server uses the trained optimization heuristic deep network to reconstruct the received compressed gradient and corrects the error using the model gradient of the previous round;

[0013] S7, gradient aggregation stage, the server performs weighted averaging on the reconstructed gradients and updates the global model;

[0014] S8, gradient delivery phase, the server delivers the aggregated global model to the client;

[0015] S9, determine whether the preset training rounds are reached, if yes, end the training process, otherwise jump to step S4 for the next round of training.

[0016] Furthermore, in the pre-training stage of step S1, the method for collecting gradient data includes: selecting a variety of typical deep learning models (such as CNN, LeNet, ResNet, Transformer, etc.) for training on different types of data sets (such as MNIST, FashionMNIST, CIFAR10, etc.), and recording the gradient parameters of each round; pulling the collected gradient parameters into a one-dimensional vector, and dividing them into blocks of a fixed size (such as 1089 dimensions) to form training samples; preprocessing the training samples, including normalization and standardization, to improve network training efficiency.

[0017] Furthermore, in the network training phase of step S2, the design of the optimized heuristic deep network includes: expanding the iterative steps of the traditional ISTA algorithm into network layers, where each layer contains a nonlinear transformation and a residual connection; the nonlinear transformation uses a convolutional layer and a learnable activation function to replace the soft threshold function in the traditional ISTA; the training goal is to minimize the reconstruction error, that is, the mean square error between the original gradient and the reconstructed gradient.

[0018] Furthermore, step S2 also includes a method for enhancing network generalization: training different compression ratios during the training process so that the network can adapt to different sampling modes; using data enhancement technology to rotate, scale, and add noise to the training samples to enhance the robustness of the network; designing a network structure suitable for gradient features to capture the spatial features in the gradient.

[0019] Furthermore, in step S5, a random Gaussian matrix satisfying the restricted isometry property (RIP) is used as a sampling matrix for gradient compression, and the compression rate can be dynamically adjusted according to the communication bandwidth and model accuracy requirements.

[0020] Furthermore, in step S6, the gradient reconstruction process includes: inputting the compressed gradient, sampling matrix, initial reconstruction matrix and previous round gradient information into the trained optimized heuristic deep network, initializing it using the initial reconstruction matrix and previous round gradient information, and gradually restoring the original gradient information through multi-layer network conversion; the last layer outputs the reconstructed gradient vector.

[0021] Furthermore, in step S7, the gradient aggregation adopts the FedAvg algorithm to perform weighted averaging on the reconstructed gradients according to the amount of data from each client, and update the global model parameters.

[0022] Compared with the prior art, the advantages of the present invention are:

[0023] 1. This paper proposes a federated learning model compression method based on an interpretable optimization heuristic deep network, which solves the problem of high communication overhead in federated learning while maintaining model performance, providing an effective solution for federated learning in resource-constrained environments.

[0024] 2. The present invention constructs a large-scale pre-training dataset and enhances the generalization of the network, so that the compression reconstruction network can adapt to the gradient characteristics of different models and different datasets, and has wide applicability.

[0025] 3. The present invention incorporates a gradient correction strategy into this solution, which effectively solves the problem of reconstruction error accumulation during the iterative process of federated learning and ensures the stability and convergence of model training. BRIEF DESCRIPTION OF THE DRAWINGS

[0026] Figure 1 Schematic diagram of a specific embodiment of the present invention.

[0027] Figure 2 This is a comparison chart of model accuracy at different compression rates for a specific embodiment of the present invention.

[0028] Figure 3 This is a comparison chart of the original and reconstructed results of some blocks on different neural network structures according to a specific embodiment of the present invention. DETAILED DESCRIPTION

[0029] The present invention will be further described below in conjunction with the accompanying drawings and specific preferred embodiments, but the scope of protection of the present invention is not limited thereby.

[0030] like Figure 1 As shown in Figure 1, the overall process of the federated learning framework for model compression based on an interpretable optimization-heuristic deep network in this embodiment includes an offline phase and an online phase. The offline phase mainly completes gradient data collection and optimization-heuristic deep network training, while the online phase implements gradient compression, reconstruction, and model updates during the federated learning process.

[0031] In this embodiment, the offline phase first collects gradient data generated by training various models on different datasets. Specifically, typical models such as LeNet, ResNet, and CNN, and datasets such as MNIST, FashionMNIST, and CIFAR-10 are selected. Standard stochastic gradient descent training is performed on each model-dataset combination. After each round of training, all gradient parameters of the model are recorded to form the original gradient set G1, G2...G T , where T represents the total number of training rounds; the gradient parameter tensor is then flattened into a one-dimensional vector and divided into blocks of a fixed size (1089 dimensions in this embodiment) to obtain a gradient block set g1, g2...g M , where M represents the total number of blocks. A large-scale dataset containing millions of samples was constructed. To enhance data diversity and improve network generalization, the gradient data was augmented with random rotations, scaling transformations, noise injection, and combined augmentation. This ultimately resulted in a large-scale training dataset containing over one million gradient samples, providing a rich training resource for subsequent deep networks. This large-scale pre-training data is a key step in ensuring good generalization of optimization-inspired deep networks.

[0032] In this embodiment, the encoder selects a Gaussian random matrix of corresponding dimensions according to the degree of compression to be performed. The Gaussian random matrix satisfies the RIP condition required for compressed sensing and is universal. The decoder is designed and trained to optimize the heuristic deep network and uses the constructed model gradient dataset for training as a decoder for gradient reconstruction. Figure 2 As shown in the figure, this network structure, inspired by the ISTA algorithm, expands the sparse transformation process in the iterative steps into network layers, replacing φx with F(x). Each layer can be represented as a forward and a backward process, where the backward process is the inverse transformation from F(x) to φx. Each of these parts consists of two main components: a convolutional module and a residual connection. The convolutional module uses a learnable convolutional layer and a parameterized activation function to replace the traditional linear transformation; the residual connection ensures effective information transmission. Considering the unique distribution characteristics of gradient data, the network depth is designed to be 15 layers, deeper than the traditional ISTA-Net, to capture more complex features in the gradient. The neural network was then trained using the dataset using the Adam optimizer with an initial learning rate of 0.001 for a total of 100 epochs. The mean squared error (MSE) loss function evaluates the difference between the reconstructed gradient and the original gradient.

[0033] In this embodiment, during the online phase of federated learning, the server first initializes the global model and distributes it to all participating clients. This initialization uses standard weight initialization methods, with He initialization used for convolutional layers and Xavier initialization used for fully connected layers, to promote rapid convergence in the early stages of training. After distribution to the client, the client trains the model using local data and obtains gradients.

[0034] In this embodiment, before uploading the model, the client selects the appropriate compression level based on its own computing resources and communication resources, and compresses the trained gradient using a predefined sampling matrix Φ of size m×n, where n is the dimension of the original gradient and m is the dimension after compression. The compression ratio is defined as The default setting is 0.5, which reduces communication traffic to 50%. The compression process is the same as the dataset acquisition method: the model gradients are divided into blocks and normalized, and then the matrix multiplication y = Φx is performed. The compressed gradients are then transmitted to the server via wireless transmission.

[0035] In this embodiment, after receiving the compressed gradients uploaded by the client, the server uses the trained optimized heuristic deep network to reconstruct them. The reconstruction process is a forward propagation operation with high computational efficiency and much lower reconstruction time than the traditional iterative algorithm. i , select the corresponding network according to its size and input it into the neural network together with the predefined compression matrix Φ and the initial reconstruction Q. For the first round of federated training, the network output is directly used as the reconstruction gradient; starting from the second round, a gradient correction strategy is introduced to use the gradient information of the previous round to assist reconstruction: in, is the gradient block reconstructed in round t, f network is the neural network output, α is the weight coefficient, which is initially set to 0.1 and decays as the number of training rounds increases.

[0036] In this embodiment, the server aggregates the reconstructed gradients and uses the FedAvg algorithm to perform weighted averaging on the reconstructed gradients based on the amount of data from each client to calculate the global gradient:

[0037]

[0038] Among them D i is the data volume of client i, N is the total number of all clients in this round, is the reconstructed gradient of client i. The server uses the aggregated gradient to update the global model. After aggregation, the server sends the updated global model to each client to start the next round of training.

[0039] In this embodiment, the technical solution is verified through simulation experiments, and PyTorch is used for model training. In the dataset acquisition stage, the LeNet network is used to train the MNIST and FashionMNIST datasets. By adjusting the seed to obtain different model gradients, a total of 200 model training result datasets for 20 rounds are obtained. In order to evaluate the performance of this method, experiments are conducted on multiple datasets. Figure 3 As shown in Figure 2, the model accuracy of this method is compared with that of the traditional baseline method under different compression rates. The results show that under different compression rates, the model accuracy of this method is only slightly lower than that of the baseline, especially at high compression rates (such as 0.5). In addition, since the training dataset is not obtained on FashionMNIST, but on Figure 3 This ensures that the method has the generalizability to different data sets and neural networks with different structures.

[0040] In this embodiment, if Figure 3 As shown in the figure, the reconstruction effects of some blocks are clearly visualized, including the reconstruction results of some blocks of MNIST, FashionMNIST, and CIFAR10 on LeNet and ResNet9. It can be seen that for the ResNet9 model, the more complex gradient changes of the model gradient can still ensure generalization.

[0041] The above description is merely a preferred embodiment of the present invention and does not constitute any form of limitation to the present invention. Although the present invention has been disclosed above based on preferred embodiments, it is not intended to limit the present invention. Therefore, any simple modifications, equivalent variations, and modifications made to the above embodiments based on the technical essence of the present invention without departing from the content of the technical solution of the present invention shall fall within the scope of protection of the technical solution of the present invention.

Claims

1. A model compression federated learning framework based on interpretable optimization heuristic deep networks. Characterized by: include: S1, the pre-training stage, collects multiple rounds of gradient data generated by training multiple models on different datasets, vectorizes and divides the gradient parameters into blocks, and constructs a dataset for training and optimizing heuristic deep networks; S2, the network training phase, designs and trains an interpretable optimization heuristic deep network as a gradient encoder-decoder based on the principle of iterative soft thresholding algorithm (ISTA); S3, the federated learning initialization phase, where the server initializes the global model and distributes the model to the clients participating in the training; S4, client local training phase, the client trains the model based on local data and calculates the gradient; S5, gradient compression phase, the client compresses the gradient using a predefined sampling matrix and uploads the compressed gradient to the server; S6, gradient reconstruction phase, the server uses the trained optimization heuristic deep network to reconstruct the received compressed gradient and corrects the error using the model gradient of the previous round; S7, gradient aggregation stage, the server performs weighted averaging on the reconstructed gradients and updates the global model; S8, gradient delivery phase, the server delivers the aggregated global model to the client; S9, determine whether the preset training rounds are reached, if yes, end the training process, otherwise jump to step S4 for the next round of training.

2. The method for compressing a federated learning model based on an interpretable optimization-heuristic deep network according to claim 1, wherein: In the pre-training stage of step S1, the method for collecting gradient data includes: selecting a variety of typical deep learning models (such as CNN, LeNet, ResNet, etc.) and training them on different types of data sets (such as MNIST, FashionMNIST, CIFAR10, etc.), and recording the gradient parameters of each round; pulling the collected gradient parameters into one-dimensional vectors and dividing them into blocks of a fixed size (such as 1089 dimensions) to form training samples; preprocessing the training samples, including normalization and standardization, to improve network training efficiency.

3. The method for compressing a federated learning model based on an interpretable optimization-heuristic deep network according to claim 1, wherein: In the network training phase of step S2, the design of the optimized heuristic deep network includes: expanding the iterative steps of the traditional ISTA algorithm into network layers, where each layer contains nonlinear transformations and residual connections; nonlinear transformations use convolutional layers and learnable activation functions to replace the soft threshold function in traditional ISTA; the training goal is to minimize the reconstruction error, that is, the mean square error between the original gradient and the reconstructed gradient.

4. The method for compressing a federated learning model based on an interpretable optimization-heuristic deep network according to claim 3, wherein: Step S2 also includes a method for enhancing network generalization: training different compression ratios during the training process to enable the network to adapt to different sampling modes; Data augmentation technology is used to rotate, scale, and add noise to training samples to enhance the robustness of the network; Design a network structure suitable for gradient features to capture the spatial features in the gradient.

5. The method for compressing a federated learning model based on an interpretable optimization-heuristic deep network according to claim 1, wherein: In step S5, gradient compression uses a random Gaussian matrix that satisfies the restricted isometry property (RIP) as the sampling matrix, and the compression rate can be dynamically adjusted according to the communication bandwidth and model accuracy requirements.

6. The method for compressing a federated learning model based on an interpretable optimization-heuristic deep network according to claim 1, wherein: In step S6, the gradient reconstruction process includes: inputting the compressed gradient, sampling matrix, initial reconstruction matrix and the previous round of gradient information into the trained optimization heuristic deep network, initializing it with the initial reconstruction matrix and the previous round of gradient information, and gradually restoring the original gradient information through multi-layer network conversion; the last layer outputs the reconstructed gradient vector.

7. The method for compressing a federated learning model based on an interpretable optimization-heuristic deep network according to claim 1, wherein: In step S7, the gradient aggregation adopts the FedAvg algorithm to perform weighted averaging on the reconstructed gradients according to the amount of data from each client and update the global model parameters.

8. The method for compressing a federated learning model based on an interpretable optimization heuristic deep network according to any one of claims 1 to 7, characterized in that: The gradient correction strategy implemented in step S6 includes: starting from the second round of federated training, using the gradient information of the previous round to assist reconstruction, optimizing the reconstruction effect through weighted combination, and the weight coefficient decays as the number of training rounds increases.

9. The method for compressing a federated learning model based on an interpretable optimization heuristic deep network according to any one of claims 1 to 7, characterized in that: In step S5, the compression ratio can be dynamically adjusted according to the client's network conditions and computing resources to achieve adaptive compression.

Images