A hardware-friendly pruning method based on Vision Transformer neural network

Abstract

This invention discloses a hardware-friendly pruning method based on the Vision Transformer (ViT) neural network, belonging to the technical field of computation, inference, or counting. This pruning method combines layer-by-layer N:M weight pruning with non-parametric multi-stage adaptive token pruning. It includes: calculating the average Hessian trace of each layer of the network based on the neural network and pre-trained weights; performing layer-by-layer N:M weight pruning based on the distribution of the average Hessian trace; retraining and updating the weights; deploying an adaptive token pruning module; retraining and updating the weights; and finally deploying the neural network using this pruning method on a hardware platform. This invention effectively reduces the number of parameters and computational cost of the ViT network by utilizing a hardware-friendly weight and token pruning method, reducing inference latency and hardware overhead while maintaining high accuracy, providing an optimized method for deploying ViT networks at resource-constrained edge environments.

Description

Technical Field

[0001] This invention relates to neural network compression and acceleration technology, specifically disclosing a hardware-friendly pruning method based on the Vision Transformer neural network, belonging to the technical field of computation, calculation, or counting. Background Technology

[0002] In recent years, the development of deep learning has driven progress in various computer vision tasks such as image classification and object detection. Among current neural networks, the Vision Transformer (ViT) architecture has demonstrated superior performance and flexibility. Compared with traditional convolutional neural networks, ViT can better capture global information in images through its self-attention mechanism, thus achieving state-of-the-art performance in many vision tasks.

[0003] However, the high performance of ViT neural networks comes with high computational complexity and high storage requirements: there are a large number of matrix multiplications in the network, and the computational complexity of the multi-head self-attention mechanism is quadratic relative to the number of input tokens. The number of parameters in the network is also much higher than that of traditional convolutional neural networks. These characteristics of ViT neural networks limit their application in resource-constrained embedded systems, mobile devices and other edge devices.

[0004] Pruning is an effective solution to the aforementioned problems. Removing redundant weights from the network can significantly reduce the number of parameters and computational cost. Removing redundant tokens from images can quadratically reduce the computational cost of the multi-head self-attention module, thereby improving the model's inference speed and reducing hardware overhead. However, existing research mainly includes the following three types of token pruning schemes: The first type directly prunes the multi-head self-attention module, which is an extremely coarse-grained pruning that severely reduces network accuracy; the second type introduces neural networks and nonlinear functions to achieve token recognition and pruning, but this introduces additional parameters, inference latency, and hardware overhead, making it unfriendly to resource-constrained edge devices; the third type performs fixed pruning of tokens, pruning the same number of tokens for images of different complexities. This pruning method will prune the tokens corresponding to objects in complex images, which has a significant impact on network accuracy. In current weight pruning methods, all layers of the network use the same pruning rate, resulting in a significant loss of network accuracy. Furthermore, to address the urgent problem of reducing network computational cost and parameter count while maintaining good network accuracy, few studies currently consider both token pruning and weight pruning simultaneously.

[0005] Therefore, the present invention aims to propose a hardware-friendly pruning method based on the Vision Transformer neural network to overcome the above-mentioned defects. Summary of the Invention

[0006] The purpose of this invention is to address the shortcomings of the aforementioned background technology by providing a hardware-friendly pruning method based on the VisionTransformer neural network. This method combines layer-by-layer N:M weight pruning with adaptive token pruning to solve the technical problems of high computational complexity and high storage requirements of the ViT neural network. The invention aims to reduce the inference latency of the ViT neural network, reduce the hardware overhead of edge deployment, and improve the inference speed and energy efficiency of the network in resource-constrained environments.

[0007] To achieve the above-mentioned objectives, the present invention employs the following technical solution:

[0008] A hardware-friendly pruning method based on Vision Transformer neural networks includes the following steps:

[0009] Step 1: Calculate the average Hessian trace for each layer of the neural network;

[0010] Step 2: Perform layer-by-layer N:M weighted pruning based on the distribution of the average Hessian traces at each layer, where N and M are positive integers and N... <M;

[0011] Step 3: Retrain the network after weight pruning in Step 2 and update the weights;

[0012] Step 4: Perform non-parametric multi-stage adaptive token pruning on the network retrained in Step 3.

[0013] Step 5: Retrain the network after token pruning in Step 4 to obtain the Vision Transformer neural network, and deploy the Vision Transformer neural network to the hardware platform.

[0014] As a further optimization of the hardware-friendly pruning method based on the Vision Transformer neural network, step 1 calculates the average Hessian trace of each layer of the neural network according to the pre-trained weights. The average Hessian trace of each layer includes the average Hessian trace of the multi-head self-attention module and the average Hessian trace of the feedforward module. The average Hessian trace of the multi-head self-attention module includes the average Hessian trace of the input mapping and the average Hessian trace of the output mapping. The average Hessian trace of the feedforward module includes the average Hessian trace of the two fully connected layers.

[0015] As a further optimization of the hardware-friendly pruning method based on the Vision Transformer neural network, step 2, which performs layer-by-layer N:M weight pruning based on the distribution of the average Hessian traces of each layer, is as follows: Based on the distribution of the average Hessian traces of the multi-head self-attention modules in each layer, the optimal numerical interval division method for the average Hessian traces of the multi-head self-attention modules in each layer is determined, and the N:M values ​​are configured for each numerical interval of the average Hessian traces of the multi-head self-attention modules in each layer according to the optimal numerical interval division method; the N:M values ​​are configured for each numerical interval of the average Hessian traces of the feedforward modules in each layer in the same way; and the weight matrices of the multi-head self-attention modules and the feedforward modules in each layer are pruned column-by-column according to the configured N:M.

[0016] As a further optimization of the hardware-friendly pruning method based on the Vision Transformer neural network, step 4, which performs non-parametric multi-stage adaptive token pruning on the network retrained in step 3, is as follows: During the forward inference process of the network, a token pruning layer is selected. Based on the output matrix of the Softmax function of all heads of the multi-head self-attention module in the token pruning layer, the token score of each image in the token pruning layer and the pruning threshold are calculated. The token scores of each image in the token pruning layer are sorted in descending order. The token scores of each image after descending order are sequentially subjected to partial sum accumulation operation. When the accumulated partial sum is greater than the pruning threshold of the token pruning layer, the accumulation operation is stopped. A pruning decision instruction is generated based on the image tokens that did not participate in the partial sum accumulation operation. The image tokens input to the layer normalization module before the feedforward module of the token pruning layer are used to execute the pruning decision instruction.

[0017] As a further optimization of the hardware-friendly pruning method based on the Vision Transformer neural network, the weight matrices of the multi-head self-attention modules and the feedforward modules of each layer are pruned column by column according to the configured N:M. Specifically, the weights of each column are grouped according to the value of M, and the N largest weights in each group are kept at their original values ​​while the other weights are assigned to zero.

[0018] As a further optimization of the hardware-friendly pruning method based on the Vision Transformer neural network, the optimal numerical interval partitioning method for the average Hessian trace of the multi-head self-attention modules in each layer and the optimal numerical interval partitioning method for the average Hessian trace of the feedforward modules in each layer are determined based on network accuracy and weight pruning benefits.

[0019] As a further optimization of the hardware-friendly pruning method based on the Vision Transformer neural network, the token score and pruning threshold of each image in the token pruning layer are calculated based on the output matrix of the Softmax function of all heads of the multi-head self-attention module in the token pruning layer. Specifically, the first row of data in the output matrix of the Softmax function of all heads of the multi-head self-attention module in the token pruning layer is summed to obtain a new row of data. The elements remaining after removing the first element in the new row of data correspond to the token scores of each image in the token pruning layer. The token scores of each image in the token pruning layer are summed and multiplied by a scaling factor to obtain the pruning threshold of the token pruning layer.

[0020] As a further optimization of the hardware-friendly pruning method based on the Vision Transformer neural network, the pruning threshold of the token pruning layer is... Where Threshold is the pruning threshold of the token pruning layer, α is the pruning scaling factor, 0 < α < 1, and token score is... i Let be the score of the i-th image token, and N be the number of image tokens input in the token pruning layer.

[0021] As a further optimization of the hardware-friendly pruning method based on the Vision Transformer neural network, the Vision Transformer neural network obtained in step 5 includes at least one network structure. Each network structure includes a multi-head self-attention module and a feedforward module that execute computation tasks sequentially. The token pruning layer also includes an adaptive token pruning module that executes computation tasks in parallel with the multi-head self-attention module of that layer. The adaptive token pruning module receives the softmax function output matrix of all heads of the multi-head self-attention module in the token pruning layer. The adaptive token pruning module outputs pruning decision instructions to the input of the layer normalization module before the feedforward module in the token pruning layer.

[0022] As a further optimization of the hardware-friendly pruning method based on Vision Transformer neural network, the adaptive token pruning module includes: a token score and pruning threshold calculation unit, a token score descending sorting unit, and a token score partial calculation unit;

[0023] The token score and pruning threshold calculation unit receives the output matrix of the Softmax function of all heads of the multi-head self-attention module in the token pruning layer at its input end, and calculates the token score and pruning threshold of each image in the token pruning layer based on the output matrix of the Softmax function of all heads of the multi-head self-attention module in the token pruning layer.

[0024] The token score descending sorting unit receives the token score of each image in the token pruning layer at its input and sorts the token scores of each image in the token pruning layer in descending order.

[0025] The token score section and calculation unit receive the pruning threshold of the token pruning layer and the token scores of each image after descending order at their input. They perform partial sum accumulation operations on the token scores of each image after descending order in sequence. The accumulation operation stops when the accumulated partial sum is greater than the pruning threshold of the token pruning layer. They generate pruning decision instructions based on the image tokens that did not participate in the partial sum accumulation operation and output the pruning decision instructions to the input of the layer normalization module before the feedforward module in the token pruning layer.

[0026] The present invention, by adopting the above technical solution, has the following beneficial effects:

[0027] (1) This invention provides a hardware-friendly pruning method based on Vision Transformer neural network, which combines weight pruning and token pruning. This method can reduce the computational cost and parameter count of the ViT network, reduce inference latency and hardware overhead while maintaining high accuracy, and provide an optimization method for network deployment at resource-constrained edge devices.

[0028] (2) Considering that the weights of different layers of the ViT network are sensitive to errors differently, this invention adopts the layer-by-layer N:M weight pruning method to prune the weights and applies Hessian traces to the weight pruning. The pruning rate of the weights of different layers and different modules in the same layer is different, which effectively reduces the number of network parameters and storage overhead while reducing the accuracy loss caused by pruning. In addition, the N:M pruning method adopted is very friendly to hardware acceleration.

[0029] (3) The token pruning method used in this invention is an adaptive token pruning scheme. The number of tokens pruned for images of different complexities is different, which is more in line with real images and application scenarios. Moreover, the token pruning method uses the intermediate value of ViT network inference to prune tokens without introducing additional networks, parameters and nonlinear functions, which is very friendly to hardware acceleration.

[0030] (4) The pruning method proposed in this invention has a simple implementation process and clear functional distinction between each step, and can be easily applied to various ViT networks and related variants. Attached Figure Description

[0031] Figure 1 This is an overall flowchart of the pruning method based on the ViT neural network used in this invention.

[0032] Figure 2 This describes the specific structure of a single-layer encoder in a ViT neural network that does not include the adaptive token pruning module involved in this invention.

[0033] Figure 3 This describes the specific structure of a layer encoder in a ViT neural network that includes the adaptive token pruning module involved in this invention.

[0034] Figure 4 This is the average Hessian trace of the multi-head self-attention modules in all layers, calculated based on the network structure and pre-trained weights in an example of the present invention.

[0035] Figure 5 This is a visualization of the pruning results of multi-stage adaptive token pruning in an example of the present invention. Detailed Implementation

[0036] The technical solution of the invention will be described in detail below with reference to the accompanying drawings. It should be understood that these embodiments are only used to illustrate the invention and not to limit the scope of the invention. After reading this invention, any modifications of the invention in various equivalent forms made by those skilled in the art fall within the scope defined by the appended claims.

[0037] This invention relates to a hardware-friendly pruning method based on the Vision Transformer neural network. It employs a layer-by-layer N:M weight pruning method to remove redundant weights from the network and a non-parametric multi-stage adaptive token pruning method to dynamically remove redundant tokens from the image during inference. By combining these two pruning methods, hardware-friendly weight-token pruning is achieved. During weight-token pruning, different weight and token pruning parameters are selected for each layer based on the network's accuracy and pruning benefits. After determining the pruning parameters, the model is retrained to restore accuracy. Specifically, it includes the following five steps.

[0038] Step 1: Calculate the average Hessian trace of each layer in the neural network.

[0039] The average Hessian trace of each layer includes the average Hessian trace of the multi-head self-attention module and the average Hessian trace of the feedforward module. The average Hessian trace of the multi-head self-attention module includes the average Hessian trace of the input mapping and the average Hessian trace of the output mapping. The average Hessian trace of the feedforward module includes the average Hessian trace of the two fully connected layers. The structure of a single encoder in a ViT network is as follows: Figure 2 As shown, the multi-head self-attention module has two linear mappings: the input mapping and the output mapping. The feedforward module has two fully connected layers. The pre-trained weights of the network are used to calculate the average Hessian trace of the input mapping, the average Hessian trace of the output mapping, and the average Hessian trace of each of the two fully connected layers in each encoder layer.

[0040] Step 2: Perform layer-by-layer N:M weighted pruning based on the distribution of the average Hessian traces at each layer.

[0041] N:M weight pruning means that among every M weights, the largest N weights retain their original values, while the remaining weights are pruned to zero. The value of N:M represents the weight pruning rate, and each layer can have different values ​​for N:M.

[0042] Hessian traces are highly effective in reflecting weight quantization errors. Since both weight quantization and weight pruning are related to error sensitivity, the application of Hessian traces is extended to weight pruning. The N:M value of each layer is determined by the magnitude of the average Hessian trace of that layer. For a network with several layers, the higher the average Hessian trace of a layer, the more sensitive the weights are to error, and the higher the chosen N:M, i.e., the lower the pruning rate. Conversely, the lower the average N:M, the higher the chosen pruning rate.

[0043] When performing layer-by-layer N:M weight pruning on the multi-head self-attention module, the average Hessian trace of the multi-head self-attention module is divided into several numerical intervals based on the distribution of the calculated average Hessian trace. Each numerical interval corresponds to an N:M value, and the higher the average Hessian trace value of the numerical interval, the smaller the weight pruning rate. Subsequently, the same method as for the multi-head self-attention module is used to select the N:M value for weight pruning for each layer of feedforward modules. The weight matrix of the corresponding module is pruned column-wise according to the N:M value configured based on the average Hessian trace of the multi-head self-attention module, and the weight matrix of the feedforward module is pruned column-wise according to the N:M value configured based on the average Hessian trace of the feedforward module. The largest N weights out of every M weights retain their original values, while the others are assigned zero. Then, the network accuracy and the amount of weight pruning are statistically analyzed using the pruned weights. Based on the accuracy and weight pruning benefits, the optimal method for dividing the average Hessian trace numerical intervals and the value of N:M are determined.

[0044] After obtaining the optimal average Hessian trace value interval division method and the N:M values ​​of each module in each layer, the weight matrix of each module in each layer is pruned column by column according to the determined N:M values.

[0045] Step 3: Retrain the model to restore accuracy and update the weights, keeping the pruned weights at zero during training.

[0046] Step 4: Perform token pruning using a non-parametric, multi-stage adaptive token pruning method.

[0047] Based on the weights after retraining with layer-by-layer N:M weight pruning, non-parametric multi-stage adaptive token pruning is performed. Token pruning is carried out in multiple stages during the network's forward inference process, using intermediate values ​​from the network inference without introducing additional networks and parameters. That is, token pruning is performed in several layers during the network's inference process, while other layers are not pruned. Adaptive pruning with different numbers of pruning tokens is performed for images of different complexities.

[0048] In layers where token pruning is not required, directly in, such as Figure 2 The encoder shown performs the network's forward reasoning process. In layers where token pruning is required, the network... Figure 3 Forward inference is performed in the ViT encoder with the adaptive token pruning module shown. The adaptive token pruning module and the multi-head self-attention module of this layer are computed in parallel. The adaptive token pruning module includes: a token score and pruning threshold calculation unit, a token score descending sorting unit, and a token score part and calculation unit. First, the output matrix of the Softmax function of all heads of the multi-head self-attention module of this layer is input into the adaptive token pruning module. The token score and pruning threshold calculation unit extracts the first row of data from the output matrix of the Softmax function of all heads, and sums the first row of data from the output matrix of the Softmax function of all heads to obtain a new row of data. After removing the first element of the new row of data, the remaining elements correspond to the score of each image token in this layer, as shown in the following formula. The scores of all image tokens in this layer are summed and multiplied by a scaling factor α (0<α<1) to obtain the pruning threshold:

[0049]

[0050] Where Threshold is the pruning threshold, α is the pruning scaling factor, and token score is... i Let N be the score of the i-th image token, and N be the number of image tokens input into the pruned layer. The token score descending sorting unit then sorts the image token scores in descending order. The token score partial sum calculation unit adds the sorted token scores one by one from high to low to obtain a partial sum. After adding a new token score to the partial sum, the partial sum is compared with a pruning threshold. If the partial sum is less than the threshold, the next token score is added to the partial sum, and the process is repeated. If the partial sum is greater than the threshold, the addition stops, and the image tokens corresponding to the token scores already added to the partial sum are retained, while the image tokens corresponding to the token scores not added are pruned. After completing the partial sum calculation and threshold comparison, the token score partial sum calculation unit generates a pruning decision instruction based on the image tokens corresponding to the token scores not added. The image tokens input to the layer normalization module before the feedforward module in this layer execute the pruning decision instruction. Image tokens corresponding to the token scores not added no longer participate in the layer normalization before the feedforward module, the calculation of the feedforward module, and the remaining layers of the network.

[0051] In the layers where token pruning is required, the adaptive token pruning module begins calculation after the network computes the Softmax function. The pruning process described above is similar to... Figure 3The second matrix multiplication and output mapping in the multi-head self-attention module are computed in parallel, without introducing additional inference latency. Regarding the image token scores, the distribution of token scores differs for images of different complexities. For images with smaller objects to be classified, token scores are concentrated in a few tokens, with large differences between the scores of different tokens. Fewer image tokens are needed to exceed the pruning threshold, resulting in a high number of pruned tokens. For images with larger objects to be classified, token scores are more dispersed, with smaller scores for different tokens and less difference between them. More image tokens are needed to exceed the pruning threshold, resulting in a low number of pruned tokens. In this case, the number of pruned tokens varies across different images, achieving adaptive pruning.

[0052] Step 5: Retrain the token-pruned network and update the weights.

[0053] While keeping the zero-value weights fixed in the N:M weight pruning, the model is retrained to restore accuracy and update the weights to obtain the Vision Transformer neural network, which is then deployed to the hardware platform.

[0054] The layer-by-layer N:M weight pruning method is characterized by the following: it is a static structural pruning scheme in which the weights are pruned before the network inference is performed, which is very friendly to hardware deployment. Furthermore, it takes into account the different sensitivity of each layer of the network to error and adopts different weight pruning rates for each layer, which can minimize the loss of network accuracy.

[0055] The multi-stage adaptive token pruning method is characterized by the following features: it is a dynamic structural pruning scheme where token pruning is related to the input image and is completed during the network's inference process. Pruning uses intermediate values ​​from the network's inference process and the pruning process runs in parallel with the network's inference process, without introducing additional parameters or inference latency. It can effectively accelerate network inference during hardware deployment, and the adaptive scheme can achieve different token pruning rates according to the complexity of the image, closely aligning with real-world application scenarios.

[0056] The following example illustrates the pruning method proposed in this invention in detail.

[0057] We take the pruning of the DeiT-S network, a variant of the ViT neural network, as an example. The layer structure of this network is exactly the same as ViT, but the difference lies in the use of a knowledge distillation training strategy. The DeiT-S neural network has 12 layers, and the multi-head self-attention module in each layer consists of 6 heads. Without pruning, the number of image tokens is 196.

[0058] First, a DeiT-S neural network is constructed, and the average Hessian trace of all layers in the network is calculated. The original, unpruned DeiT-S neural network is constructed first. The average Hessian trace of all layers in the network is then calculated using the DeiT-S neural network and its pre-trained weights. The average Hessian trace of each layer includes the average Hessian trace of the multi-head self-attention module and the average Hessian trace of the feedforward module. The average Hessian trace of the multi-head self-attention module includes the average Hessian trace of the input mapping and the average Hessian trace of the output mapping. The average Hessian trace of the feedforward module includes the average Hessian trace of the two fully connected layers.

[0059] Determine the N:M values ​​for weight pruning at each layer. Calculate the average Hessian trace of all layers in the DeiT-S neural network and obtain the following... Figure 4 The average Hessian trace distribution of the multi-head self-attention module is shown. In this example, the two sets of average Hessian traces of the multi-head self-attention module are divided into three numerical intervals: [0, Th1], [Th1, Th2], and [Th2, 0.001]. The N:M values ​​of the three numerical intervals are 1:8, 1:4, and 1:2, respectively. Similarly, in this example, the two sets of average Hessian traces of the feedforward module are divided into three numerical intervals: [0, Th3], [Th3, Th4], and [Th4, 0.001]. The N:M values ​​of the three numerical intervals are 1:8, 1:4, and 1:2, respectively. Different combinations of numerical intervals are selected for the average Hessian traces of the DeiT-S neural network, and the N:M values ​​are determined for each module in each layer. The accuracy and weight pruning of the network are statistically analyzed, and the results are shown in Table 1.

[0060] 0.1 0.2 0.1 0.2 68.9％ 79.5％ 0.1 0.3 0.1 0.3 71.7％ 79.7％ 0.2 0.4 0.1 0.3 74.4％ 79.0％ 0.2 0.4 0.2 0.4 76.1％ 78.8％ 0.1 0.3 0.2 0.4 73.4％ 79.3％ 0.3 0.5 0.3 0.5 79.3％ 78.2％

[0061] Table 1

[0062] Based on the network precision and weight pruning amounts in Table 1, the parameter values ​​in this example are: Th1 = 0.1, Th2 = 0.3, Th3 = 0.1, and Th4 = 0.3. The N:M values ​​corresponding to the four weight sets in each layer are determined according to the numerical range. It should be noted that the above N:M values ​​and combinations are not limited to 1:8, 1:4, and 1:2; 2:4, 3:4, and other values ​​and combinations of N:M can also be selected.

[0063] Perform weight pruning and retraining. Based on the determined N:M, perform column-wise weight pruning: for each column of weights, group the M weights together, keep the N largest weights of these M weights, and prune the other weights and assign them to zero. Retrain the model to restore accuracy and update the weights, while keeping the pruned weights at zero.

[0064] An adaptive token pruning module is deployed in a DeiT-S neural network to implement adaptive pruning of image tokens. The image token score and pruning threshold are calculated. In this example, the token pruning positions in the DeiT-S neural network are selected as layers 4, 7, and 10. First, it is determined whether the current network inference is in a token pruning layer. If not, token pruning is not performed; otherwise, the adaptive token pruning process is performed: the output matrix of the Softmax function of all heads in the multi-head self-attention module of this layer is fed into the adaptive token pruning module. The first row of data from the output matrix of the Softmax function of all heads is extracted from the token score and pruning threshold calculation unit and summed to obtain a new row of data. After removing the first element from this new row, the remaining elements correspond to the scores of all image tokens in this layer. The scores of all image tokens in this layer are summed and multiplied by a scaling factor α (0 < α < 1) to obtain the pruning threshold. In this example, the scaling factor α can be 0.95, 0.93, 0.85, 0.8, etc. Different scaling factors α can achieve different benefits in terms of network accuracy and computational reduction.

[0065] The token scores are sorted in descending order, and the partial sum of the token scores is calculated. In the token pruning layer, the image token scores are sorted in descending order in the token score sorting unit. Then, in the token score partial sum calculation unit, the sorted token scores are added one by one from high to low to obtain the partial sum of the token scores. After each new token score is added to the partial sum of the token scores, the partial sum is compared with the pruning threshold. If the partial sum of the token scores is less than the threshold, the next token score is added to the partial sum, and the above process is repeated. If the partial sum of the token scores is greater than the threshold, the addition stops, and the image tokens corresponding to the token scores that have been added to the partial sum are retained. Image tokens whose token scores have not been added are directly discarded after the multi-head self-attention module is completed and no longer participate in the calculation of the feedforward module of this layer and the remaining layers of the network. When using α=0.9 and performing token pruning at layers 4, 7, and 10, the visualization results of multi-stage adaptive token pruning are as follows. Figure 5 As shown, the shaded areas represent tokens that were pruned at that layer.

[0066] Adaptive token pruning is performed for retraining to restore accuracy and update weights. An adaptive token pruning module is added to the DeiT-S neural network. The model is retrained to restore accuracy and update weights while maintaining the retrained weights of the N:M weight pruning. During training, the pruned weights are kept at zero.

[0067] The DeiT-S neural network employing the pruning method proposed in this invention was deployed on a hardware platform, and the network's accuracy and pruning benefits were evaluated. The aforementioned layer-by-layer N:M weight pruning and non-parametric multi-stage adaptive token pruning are both hardware-friendly pruning schemes that can be easily deployed in hardware. With Th1 = 0.1, Th2 = 0.3, Th3 = 0.1, Th4 = 0.3, and α = 0.9, the parameters of the DeiT-S neural network in this example were reduced by 71.7%, and the network inference speed was improved by 1.6 times.

[0068] The foregoing has shown and described the basic principles, main features, and advantages of the present invention. Those skilled in the art should understand that the present invention is not limited to the specific embodiments described above. The specific embodiments and descriptions in the specification are merely for further illustrating the principles and preparation effects of the present invention. Various changes and modifications can be made to the present invention without departing from its spirit and scope, and all such changes and modifications fall within the scope of protection claimed by the present invention. The scope of protection of the present invention is defined by the claims and their equivalents.

Claims

1. A hardware-friendly pruning method based on Vision Transformer neural networks, characterized in that, Includes the following steps: Step 1: Calculate the average Hessian trace for each layer of the neural network; Step 2: Perform layer-by-layer N:M weight pruning based on the distribution of average Hessian traces at each layer: Determine the optimal numerical interval division method for the average Hessian traces of each layer's multi-head self-attention modules based on the distribution of average Hessian traces at each layer. Configure N:M values ​​for each numerical interval of the average Hessian traces of each layer's multi-head self-attention modules according to the optimal numerical interval division method. Similarly, configure N:M values ​​for each numerical interval of the average Hessian traces of each layer's feedforward modules. Prune the weight matrices of each layer's multi-head self-attention modules and feedforward modules column-wise according to the configured N:M, where N and M are positive integers and N is a positive integer. <M; Step 3: Retrain the network after weight pruning in Step 2 and update the weights; Step 4: Perform non-parametric multi-stage adaptive token pruning on the network retrained in Step 3: Select a token pruning layer during the network's forward inference process. Calculate the token score and pruning threshold for each image in the token pruning layer based on the Softmax function output matrix of all heads in the multi-head self-attention module of the token pruning layer. Sort the token scores of each image in the token pruning layer in descending order. Perform partial sum accumulation operations on the sorted token scores in sequence. Stop the accumulation operation when the accumulated partial sum is greater than the pruning threshold of the token pruning layer. Generate a pruning decision instruction based on the image tokens that did not participate in the partial sum accumulation operation. Input the image tokens into the layer normalization module before the feedforward module of the token pruning layer to execute the pruning decision instruction. Step 5: Retrain the network after token pruning in Step 4 to obtain the Vision Transformer neural network, and deploy the Vision Transformer neural network to the hardware platform.

2. The hardware-friendly pruning method based on Vision Transformer neural network according to claim 1, characterized in that, Step 1 calculates the average Hessian trace of each layer of the neural network based on the pre-trained weights. The average Hessian trace of each layer includes the average Hessian trace of the multi-head self-attention module and the average Hessian trace of the feedforward module. The average Hessian trace of the multi-head self-attention module includes the average Hessian trace of the input mapping and the average Hessian trace of the output mapping. The average Hessian trace of the feedforward module includes the average Hessian trace of the two fully connected layers.

3. The hardware-friendly pruning method based on Vision Transformer neural network according to claim 1, characterized in that, The weight matrices of the multi-head self-attention modules and the feedforward modules of each layer are pruned column-wise according to the configured N:M. Specifically, the weights in each column are grouped according to the value of M, and the N largest weights in each group are kept at their original values ​​while the other weights are assigned zero.

4. The hardware-friendly pruning method based on Vision Transformer neural network according to claim 1, characterized in that, The optimal numerical interval division method for the average Hessian trace of each layer's multi-head self-attention module and the optimal numerical interval division method for the average Hessian trace of each layer's feedforward module are determined based on network accuracy and weight pruning benefits.

5. The hardware-friendly pruning method based on Vision Transformer neural network according to claim 1, characterized in that, The step of calculating the token score and pruning threshold for each image in the token pruning layer based on the output matrix of the Softmax function of all heads of the multi-head self-attention module in the token pruning layer is as follows: sum the data in the first row of the output matrix of the Softmax function of all heads of the multi-head self-attention module in the token pruning layer to obtain a new row of data. The elements remaining after removing the first element in the new row of data correspond to the token scores of each image in the token pruning layer. After summing the token scores of each image in the token pruning layer, multiply by a scaling factor to obtain the pruning threshold of the token pruning layer.

6. The hardware-friendly pruning method based on the Vision Transformer neural network according to claim 5, characterized in that, The pruning threshold of the token pruning layer is: Where Threshold is the pruning threshold of the token pruning layer, α is the pruning scaling factor, 0 < α < 1, and token score i Let be the score of the i-th image token, and N be the number of image tokens input in the token pruning layer.

7. A hardware-friendly pruning method based on a Vision Transformer neural network according to any one of claims 1 to 6, characterized in that, The Vision Transformer neural network obtained in step 5 includes at least one network structure. Each network structure includes a multi-head self-attention module and a feedforward module that execute computation tasks sequentially. The token pruning layer also includes an adaptive token pruning module that executes computation tasks in parallel with the multi-head self-attention module of that layer. The adaptive token pruning module receives the softmax function output matrix of all heads of the multi-head self-attention module in the token pruning layer. The adaptive token pruning module outputs pruning decision instructions to the input of the layer normalization module before the feedforward module in the token pruning layer.

8. The hardware-friendly pruning method based on the Vision Transformer neural network according to claim 7, characterized in that, The adaptive token pruning module includes: The token score and pruning threshold calculation unit receives the softmax function output matrix of all heads of the multi-head self-attention module in the token pruning layer at its input end, and calculates the token score and pruning threshold of each image in the token pruning layer based on the softmax function output matrix of all heads of the multi-head self-attention module in the token pruning layer. A token score descending sorting unit receives the token score of each image in the token pruning layer as its input, and sorts the token scores of each image in the token pruning layer in descending order; and... The token score section and calculation unit receive the pruning threshold of the token pruning layer and the token scores of each image after descending order at their input terminals. They perform partial sum accumulation operations on each image token score after descending order in sequence. The accumulation operation stops when the accumulated partial sum is greater than the pruning threshold of the token pruning layer. They generate pruning decision instructions based on the image tokens that did not participate in the partial sum accumulation operation and output the pruning decision instructions to the input terminal of the layer normalization module before the feedforward module in the token pruning layer.

Images