A Hardware-Aware Compression Method for Large-Scale Pre-Trained Language Models

Abstract

The present invention discloses a method for compressing large-scale pre-trained language models based on hardware awareness, aiming at the problems existing in the model compression methods in the prior art. The invention comprises the following steps: loading the pre-trained language model that has been trained and fine-tuning it on downstream tasks; selecting the pruning granularity; constructing a proxy dataset according to the selected pruning granularity; designing a latency prediction network and training the latency prediction network using the proxy dataset; setting the desired latency, and embedding the latency prediction network into a multi-mixed granularity pruning framework for pruning, and at the same time adopting a knowledge distillation method with dynamic mapping to transfer the knowledge of the teacher model to the changing pruned model; using the best pruning mask obtained by pruning to perform model pruning and performing fine-tuning to achieve accuracy recovery, so as to obtain a compressed model. The present invention solves the problem of efficient and simple deployment of a large number of multi-head attention mechanism-based architecture network models and has good application prospects.

Description

Technical Field

[0001] The present invention relates to the field of deep learning technology, and in particular to a large-scale pre-trained language model compression method based on hardware perception. Background Art

[0002] In recent years, deep learning has made great achievements in many fields, including object detection, super-resolution image generation, natural language processing, etc. The performance of various network models has been continuously improved, but it is also accompanied by an explosive growth in the number of model parameters. For example, the GPT4 model recently released by OpenAI has 17.5 billion parameters, which is ten times the size of its previous generation GPT3 model. The surge in model size makes it a major challenge to deploy it on devices with limited computing power and memory. Therefore, how to effectively compress the model and maintain its performance as much as possible has become a hot topic in academia and industry.

[0003] At present, common model compression methods include pruning, knowledge distillation, quantization, network structure search, etc. Pruning selects the importance of each submodule in the network, and then sorts them according to the importance and prunes the submodules with lower importance until the model parameters meet the preset pruning rate, so as to achieve the purpose of reducing the number of parameters and speeding up the reasoning speed. The main idea of ​​knowledge distillation is to transfer knowledge from the teacher model with more parameters to the student model with smaller parameters, so as to achieve the purpose of making the performance of the student model close to that of the teacher model, but greatly improving the reasoning speed. Model quantization converts the floating-point type parameters in the model into integers of different bit lengths through preset mapping rules without causing a large loss of accuracy. Network structure search starts from network design and automatically searches in a defined network structure space to find the best subnetwork that achieves a trade-off between performance and parameter quantity, thereby achieving the purpose of designing a lightweight network.

[0004] Model pruning has attracted widespread attention due to its advantage of being able to flexibly prune sub-networks of different scales according to the pruning rate. However, existing model pruning methods are based on parameter quantity constraints, that is, pruning models with a predetermined parameter quantity. However, the model parameter quantity and inference speed are not strictly linearly related, which means that if you want to prune a model with a specific inference delay, you may need to repeatedly adjust the pruning rate based on expert experience. In addition, due to the differences in hardware architectures, the pruning rate that is carefully fine-tuned on a certain hardware to meet the delay requirements may not be migrated to another hardware, which further increases the cost of model compression and reduces the degree of automation. At the same time, since the structure of the intermediate layer of the model changes dynamically during the pruning process, it is difficult to use the knowledge distillation method for the intermediate layer of the model to restore the accuracy of the model during the pruning process. Summary of the invention

[0005] Aiming at the deficiencies of existing model compression methods in the prior art, the present invention provides a hardware-aware large-scale pre-trained language model compression method with good generality, high flexibility and effectiveness.

[0006] The technical solution of the present invention is to provide a hardware-aware large-scale pre-trained language model compression method with the following steps: including the following steps,

[0007] Step 1: Load the pre-trained language model after training and fine-tune it on downstream tasks; Step 2: Select the pruning granularity; Step 3: Construct a proxy dataset according to the selected pruning granularity; Step 4: Design a latency prediction network and train the latency prediction network using the proxy dataset; Step 5: Set the expected latency, embed the latency prediction network into a multi-mixed granularity pruning framework for pruning, and use the dynamic mapping knowledge distillation method to transfer the knowledge of the teacher model to the changing pruned model; Step 6: Use the best pruning mask obtained by pruning to prune the model, restore the accuracy after fine-tuning, and obtain the compressed model.

[0008] Preferably, the pre-trained model loaded in Step 1 is the BERT model, and it is fine-tuned on the SST-2 dataset for 4 rounds with a learning rate of 3e -5 。

[0009] Preferably, the pruning granularity in Step 2 includes: the middle layer of the network, the feed-forward layer, the attention layer, the attention head, the hidden layer dimension and the middle layer dimension, which act as the effect of parameter pruning in the form of a mask. The pruning masks corresponding to the pruning grains involved are:

[0010] Feed-forward layer: The corresponding pruning mask is represented by Z FLL The mask dimension is 12, and Z FLL (i) represents whether the i-th middle layer of the network is pruned or not;

[0011] Attention layer: The corresponding pruning mask is represented by Z MHA The mask dimension is 12, and Z MHA (i) represents whether the i-th middle layer of the network is pruned or not;

[0012] Attention head: The corresponding pruning mask is represented by Z HEAD The mask dimension is 12x12, and Z HEAD (i, j) represents whether the j-th attention head in the i-th attention layer is pruned or not;

[0013] Hidden feature dimension: The corresponding pruning mask is represented by Z HID The mask dimension is 768, and Z HID (i) represents whether the i-th hidden layer dimension is pruned or not;

[0014] Intermediate feature dimension: The corresponding pruning mask is represented by Z INT It is represented, the mask dimension is 3072, Z INT (i) Indicates whether it is the i-th intermediate layer dimension;

[0015] The specific formula is as follows:

[0016]

[0017] FFN(X) = z FFN ·gelu(XW U )·diag(z int )·W D

[0018] Among them, MHA(X) is the output of the multi-head attention layer, Nh is the number of multi-heads of this layer, are the necessary query, key, value, and output matrices for constructing the attention mechanism of the multi-head attention layer, FFL(X) is the output of the feed-forward layer, W U is the weight matrix of the intermediate hidden layer in the feed-forward layer, W D is the weight matrix of the output layer in the feed-forward layer.

[0019] Preferably, step 3 includes the following sub-steps:

[0020] Step 3.1: Determine the corresponding pruning mask for the selected pruning granularity, and determine a pruning model through the combination of multiple groups of pruning masks. The pruning mask is composed of 0 and 1, where 0 indicates pruning this group of parameters, and 1 indicates retaining this group of parameters. The pruning granularity and the corresponding pruning mask involved are:

[0021] Feed-forward layer: The corresponding pruning mask is represented by Z FLL It is represented, the mask dimension is 12, Z FLL (i) Indicates whether to prune the i-th network intermediate layer;

[0022] Attention layer: The corresponding pruning mask is represented by Z MHA It is represented, the mask dimension is 12, Z MHA (i) Indicates whether to prune the i-th network intermediate layer;

[0023] Attention head: The corresponding pruning mask is represented by Z HEAD It is represented, the mask dimension is 12x12, Z HEAD (i, j) Indicates whether to prune the j-th attention head in the i-th attention layer;

[0024] Hidden feature dimension: The corresponding pruning mask is represented by Z HID It is represented, the mask dimension is 768, Z HID (i) Indicates whether it is the i-th hidden layer dimension;

[0025] Intermediate feature dimension: The corresponding cropping mask is denoted by Z INT which indicates that the mask dimension is 3072, and Z INT (i) indicates whether it is the i-th intermediate layer dimension;

[0026] Step 3.2: Randomly sample the cropping mask, crop the model according to the adopted cropping mask, test its inference latency, construct a proxy dataset, use the randomly sampled cropping mask as the feature of the proxy dataset, and the measured inference latency as the label.

[0027] Preferably, in step 4, the latency features of different modules in the model are first extracted through each different branch. The aggregation branch is responsible for receiving the latency features output by other branches, and further linearly transforms the received latency features and then outputs the inference latency. The latency prediction network is a multi-branch tree-structured network. The features input to this network are: a one-dimensional binary vector of length 25, and the specific meanings of its each dimension are: bits 1-12 represent the retention ratios of the attention heads in the 12 attention layers in BERT, bits 13-24 represent the retention ratios of the intermediate feature dimensions in the 12 feed-forward layers in BERT, and bit 25 represents the retention ratio of the hidden feature dimension in BERT; this network includes the following branches:

[0028] (1) Attention layer latency feature extraction branch: The input is the 1st - 12th bits of the total input, and it passes through fully connected layers with 12, 24, 12, and 1 neurons in sequence, and finally outputs to the aggregation branch. And leakRelu activation functions are used between different layers to prevent all negative information from being lost;

[0029] (2) Feed-forward layer latency feature extraction branch: The input is the 13th - 24th bits of the total input, and its structure is the same as that of the attention layer latency feature extraction branch, and finally outputs to the aggregation branch;

[0030] (3) Hidden dimension latency feature extraction branch: The input is the 25th bit of the total input, and it is connected to a fully connected layer with 1 neuron, and finally outputs to the aggregation branch;

[0031] (4) Aggregation branch: The input is a one-dimensional vector of length 3. Concatenate the values of the above three branches to form a one-dimensional vector of length 3, and input it into a fully connected network. Finally, this network outputs the latency information on a specific hardware. This fully connected network consists of two linear layers. The first layer contains 3 neurons, and the second layer contains 1 neuron.

[0032] Preferably, the pruning framework with multiple mixed granularities in step 5 includes: an l0 regularization module, a model component library, and a pruning training module; the composition and functions of each module include:

[0033] Regularization module: This module uniformly controls the learnable structure parameters corresponding to the cropping mask and uses a differentiable l0 norm form to model the cropping mask. During the cropping training process, the regularization module is used to calculate the sparse loss based on the l0 coefficient;

[0034] Model component library: This component library contains the fine-tuned sub-modules that make up the BERT model. The sub-modules have an added control mechanism without changing their original functions. The control mechanism means that the representational ability of the module is constrained by the incoming cropping mask;

[0035] Dynamic mapping knowledge distillation module: Based on the outputs of the teacher model and the student model, dynamic mapping is used to calculate the corresponding prediction layer distillation loss and intermediate layer distillation loss. The prediction layer distillation loss is calculated based on the prediction probabilities of the student model and the teacher model, and the intermediate layer distillation loss is calculated based on the similarity between the feature distributions output by the student model and the teacher model. The higher the similarity between different layers of the teacher model and the student model, the lower the knowledge transfer cost of the teacher model, and correspondingly, more knowledge is transferred to this layer, that is, it accounts for a larger proportion in the distillation loss;

[0036] Cropping training module: This module provides a complete cropping pipeline by integrating the l0 regularization module described in step 5.1, the model component library described in step 5.2, and the delay prediction network described in step 4. The specific steps are as follows:

[0037] 5.1. Initialize the baseline model, that is, the BERT model, and fine-tune it on the downstream task dataset. Use the fine-tuned model as the teacher model during the cropping process;

[0038] 5.2. Replace the original modules in BERT with the modules in the model component library as the initialized cropping model;

[0039] 5.3. Initialize the l0 regularization module, including the initialization of the target sparsity and target delay, the initialization of the cropping granularity, the embedding of the delay prediction network, and the initialization of the structure parameters;

[0040] 5.4. Preprocess the downstream task data;

[0041] 5.5. Use the l0 regularization module to output the cropping soft mask in the single-training mode according to the structure parameters and use it as part of the input data in a batch training;

[0042] 5.6. Input a batch of data into the teacher model, and at the same time input the same batch of data into the cropping model after adding the mask parameter to complete a forward propagation;

[0043] 5.7. Calculate the distillation loss based on the output of the teacher model and the output of the pruning model in 5.6, and calculate the sparsity loss according to the l0 regularization module. The sum of the two losses is the total loss;

[0044] 5.8. Use the adaptive moment estimation optimizer with weight decay to update the model parameters and structural parameters;

[0045] 5.9. Repeat steps 5.6 - 5.8 until the preset number of iterations;

[0046] 5.10. Input the converged structural parameters into the l0 regularization module to obtain the binary pruning mask under the inference model, prune the model through this mask, and input the pruned small model.

[0047] Compared with the prior art, the hardware-aware large-scale pre-trained language model compression method of the present invention has the following advantages:

[0048] First, better compression performance: Since the present invention adopts a hybrid-grained pruning method, the fine-grained pruning unit provides the accuracy of pruning, and the coarse-grained pruning unit provides more acceleration effects, making the pruning have both pruning accuracy and acceleration practicality. At the same time, the present invention introduces a dynamic knowledge distillation mapping strategy for the structure-changing network, better transferring the knowledge of the teacher model to the pruned student model, and maximizing the restoration of accuracy, thus achieving better compression performance.

[0049] Second, hardware awareness: In the present invention, a delay prediction network with good hardware generalization is designed. By sampling a small amount of data on the target device and training it, the delay can be accurately predicted. Embedding it into the pruning framework, it is easy to prune a sub-model with a specific inference speed on the target device, so as to achieve the purpose of perceiving the best delay-friendly sub-model on a specific hardware.

[0050] Third, strong versatility: The model compression method of the present invention is not limited to the BERT model in terms of application scope. It can compress any network with a multi-head attention layer and a feed-forward layer in an end-to-end manner, and only requires a small amount of code modification, having the advantages of good automatic compression.

[0051] IV. Based on the hardware-aware pre-training language model compression, the present invention designs a small model extracted from a large model with hardware generalization and a compression method that maximally retains the performance of the original large model, which well solves the deployment problems of a large number of existing network models with multi-head attention mechanism architectures. Specifically, a multi-mixed granularity pruning method is adopted to combine the accuracy of pruning and the practicality of inference acceleration. In addition, the present invention designs a latency prediction network with hardware awareness, which can accurately predict the latency of the pruned network on the target hardware by collecting a small amount of data for training on the target hardware. Then, by embedding the pruning framework, the time for repeatedly fine-tuning the pruning rate to achieve a specific inference acceleration effect can be avoided, and the latency information for the target device can be sensed to search for the best hardware-friendly small model. Furthermore, by using a dynamic mapping knowledge distillation strategy during the pruning process, reasonable knowledge transfer from the teacher model to the student model is realized, and the accuracy recovery during the pruning process is achieved. Finally, the large model is pruned into a high-performance small model that is easy to deploy and has a specified inference speed on the specified hardware. BRIEF DESCRIPTION OF THE DRAWINGS

[0052] Figure 1 is a schematic diagram of the implementation process of the present invention;

[0053] Figure 2 is a schematic diagram of the implementation of mixed granularity pruning in the present invention;

[0054] Figure 3 is a schematic diagram of the embedding of the latency network into the pruning training in the present invention;

[0055] Figure 4 is a schematic diagram of knowledge distillation for accuracy recovery in the present invention. DETAILED DESCRIPTION OF THE EMBODIMENTS

[0056] In order to enable those skilled in the art to better understand the solution of the present invention, the technical solutions in the embodiments of the present invention will be clearly and completely described below with reference to the accompanying drawings in the embodiments of the present invention. Obviously, the described embodiments are only a part of the embodiments of the present invention, rather than all the embodiments. Based on the embodiments of the present invention, all other embodiments obtained by those of ordinary skill in the art without creative efforts shall fall within the protection scope of the present invention.

[0057] The following further describes the method for compressing a large-scale pre-trained language model based on hardware perception of the present invention in conjunction with the accompanying drawings and specific embodiments: As shown in Figure (1), in this embodiment, the current mainstream BERT (Bidirectional Encoder Representations from Transformers) pre-trained model is compressed to improve the inference speed on hardware and reduce the number of model parameters, so as to facilitate deployment on the edge side. It should be noted that: The present invention mainly focuses on compressing the encoder unit; The model obtained by fine-tuning BERT on specific downstream tasks in the present invention is called the teacher model, and it is also used as the baseline model for comparing compression performance. The pruned model, that is, the student model, is initialized as the original BERT model without fine-tuning. By inputting pruning masks during the training process of the student model on downstream tasks to simulate the effect of restricting the model size, and at the same time using a delay loss to make the student model finally converge to near the given delay. At the same time, to address the performance loss problem during the pruning process, the present invention uses a dynamic mapping knowledge distillation strategy to help the teacher model better transfer knowledge to the dynamically changing student model.

[0058] As Figure 1 shown, the compression method of the present invention can be divided into six steps:

[0059] Step 1: Load the pre-trained language model that has been trained and fine-tune it on downstream tasks. In this embodiment, the pre-trained model loaded is the BERT model, and it is fine-tuned on the SST-2 dataset (a binary classification dataset for sentiment classification of movie reviews), and the number of fine-tuning epochs is 4 rounds, and the learning rate is 3e -5 . The purpose of this step is to use the fine-tuned BERT model as the teacher model during the pruning training process to help restore the accuracy of the pruned model.

[0060] Step 2: Select the pruning granularity. As Figure 2 shown, in the present invention, multiple pruning granularities can be used to prune the model, including: the middle layer of the network (Transformer Block); the feed-forward layer (Feed-forward Layer); the attention layer (Attention Layer); the attention head (Attention Head); the hidden layer dimension (hidden dimension); the intermediate layer dimension (intermediate dimension). Specifically, coarser-grained pruning units can bring higher acceleration effects, but the disadvantage is that it is difficult to achieve precise pruning, resulting in a performance decline. For example Figure 2Intermediate-level pruning. Fine-grained pruning can achieve precise pruning of redundant parameters, but the resulting acceleration effect is not obvious enough, such as pruning the intermediate layer dimension parameters in the figure. In this example, all the above granularities are adopted to ensure both the accuracy of pruning and the practicality of acceleration. Further, the present invention uses a mask to act as the effect of parameter pruning, and the pruning mask corresponding to the pruning grain involved is:

[0061] Feed-forward Layer: The corresponding pruning mask is represented by Z FLL , and the mask dimension is (12), and Z FLL (i) indicates whether the i-th network intermediate layer is pruned or not;

[0062] Attention Layer: The corresponding pruning mask is represented by Z MHA , and the mask dimension is (12), and Z MHA (i) indicates whether the i-th network intermediate layer is pruned or not;

[0063] Attention Head: The corresponding pruning mask is represented by Z HEAD , and the mask dimension is (12, 12), and Z HEAD (i, j) indicates whether the j-th attention head in the i-th attention layer is pruned or not;

[0064] Hidden Feature Dimension: The corresponding pruning mask is represented by Z HID , and the mask dimension is (768), and Z HID (i) indicates whether the i-th hidden layer dimension is pruned or not;

[0065] Intermediate Feature Dimension: The corresponding pruning mask is represented by Z INT , and the mask dimension is (3072), and Z INT (i) indicates whether the i-th intermediate layer dimension is pruned or not;

[0066] The specific formula is as follows:

[0067]

[0068] FFN(X)Z FFN ·gelu(XW U )·diag(z int )·W D

[0069] Among them, MHA(X) is the output of the multi-head attention layer, and Nh is the number of multi-heads in this layer, The query, key, value, and output matrices necessary for constructing the attention mechanism of the multi-head attention layer. FFN(X) is the output of the feed-forward layer, and W U is the weight matrix of the intermediate hidden layer in the feed-forward layer, and W D is the weight matrix of the output layer in the feed-forward layer.

[0070] Step 3: Construct a (cropping mask - latency) proxy dataset according to the selected cropping granularity. Randomly sample the cropping mask, crop the model according to the adopted cropping mask, and test its inference latency on the specified hardware to form a set of (cropping mask - inference latency) data pairs. Repeat the above operations multiple times to construct a proxy dataset, where the cropping mask is the feature and the inference latency is used as the label. In this example, the present invention uses 2000 sets of data as the training data for the latency prediction network in the environment of Ubuntu 20.04.2, Pytorch 1.10.0, memory 64G, Intel Core i9-10900KF CPU, and Nvidia RTX 3090 GPU.

[0071] Step 4: Design a latency prediction network and train the latency prediction network using the proxy dataset. Since different modules of different models contribute differently to the latency, the present invention designs a multi-branch tree-structured network. First, extract the latency features of different modules of the model through each different branch, and finally perform summary inference and output. Specifically, the design of the network model is as follows: The input is a one-dimensional binary vector of length 25. The specific meanings of the dimensions of this vector are as follows: Bits 1-12 represent the retention ratio of the attention heads in the 12 attention layers in BERT, bits 13-24 represent the retention ratio of the intermediate feature dimensions in the 12 feed-forward layers in BERT, and bit 25 represents the retention ratio of the hidden feature dimensions in BERT. The network structure includes the following branches:

[0072] Attention layer latency feature extraction branch: The input is bits 1-12 of the total input, and it passes through fully connected layers with 12, 24, 12, and 1 neurons in sequence, and finally outputs to the summary branch. And the leakRelu activation function is used between different layers to prevent all negative information from being lost.

[0073] Feed-forward layer latency feature extraction branch: The input is bits 13-24 of the total input, and its structure is the same as that of the attention layer latency feature extraction branch, and finally outputs to the summary branch.

[0074] Hidden dimension latency feature extraction branch: The input is bit 25 of the total input, and it is connected to a fully connected layer with one neuron, and finally outputs to the summary branch.

[0075] Aggregate branch: The input is a one-dimensional vector of length 3. The values of the above three branches 1, 2, and 3 are concatenated to form a one-dimensional vector of length 3, and then passed through fully connected layers with 3 and 1 neurons in sequence. Finally, the delay information on a specific hardware is inputted.

[0076] In this example, the training batch size is 2, the number of training epochs is 300, the optimizer selected is the Stochastic Gradient Descent optimizer (SGD optimizer), and the learning rate is 1e-4. The present invention uses 80% of the data as the training set and 20% of the data as the test set. Finally, the error between the delay prediction network and the true value in the experimental results is only 5%, which is sufficient to be used as a reliable delay information providing component.

[0077] Step five: Set the desired delay. The target delay in this example is set to 10% of the delay of the complete model. And embed the delay prediction network into the pruning framework of multi-mixed granularity for pruning. The embedding process is shown in Figure 3 as shown. In addition, during pruning, the knowledge distillation method of dynamic mapping is used to transfer the knowledge of the teacher model to the changing pruned model, as shown in Figure 4 as shown. Specifically, the pruning framework of multi-mixed granularity includes 10 regularization modules, a model component library, and a pruning training module. Among them, the I0 regularization module is used as the component providing the model structure for each forward inference and the component providing the sparse loss in the pruning training module; the model component library includes an improved version of the original model, that is, BERT. The improvement includes passing the mask as an input item to control the structure of each module; the pruning training module integrates the above two modules and is trained in the fine-tuning stage of the downstream task, and finally outputs the best small model that meets the preset constraints. Further, the functions of each module and the implementation steps of the corresponding functions are as follows:

[0078] (1) 10 regularization module: This module uniformly controls the learnable structure parameters (score) corresponding to the pruning mask, and models the pruning mask based on the differentiable form of the 10 norm, thereby calculating the sparse loss based on the 10 coefficient during the pruning training process. Here, the hard concrete distribution is used as the differentiable 10 norm, and the specific formula is as follows:

[0079] u~U(0, 1)

[0080] s = Sigmoid((log u - log(1 - u) + logα) / β)

[0081]

[0082]

[0083] U(0, 1) follows the uniform distribution within the interval [0, 1], where l < 0 and r > 0. In this example, l = -0.01 and r = 1.1. Stretch the output of the Sigmoid function to the range (l, r). β is a hyperparameter controlling the smoothness, i.e., the temperature parameter, which is set to 0.66 in this example. logα is a learnable structure parameter, namely the score mentioned above. Control the mask by updating the structure parameter and sample the mask from it during the forward pass, that is, sample the model structure for the current forward inference for pruning training. z represents the sampled structure parameter, namely the z mentioned in step two. Further, when calculating the sparsity loss, multiply the probability that s is greater than 0 by the number of parameters corresponding to the pruning granularity to obtain the penalty term, that is, the expected number of parameters of the model obtained under the current learnable parameter sampling. Take the mean square error between this expectation and the target number of parameters as the sparsity loss regarding the number of parameters. At the same time, the sum of the expectations that the s parameters corresponding to each granularity are greater than 0 represents the retention ratio of this granularity. A corresponding hard mask for the structure of the input delay prediction network can be obtained according to this retention ratio to obtain the corresponding delay loss. That is, the final sparsity loss = delay loss + sparsity loss of the number of parameters.

[0084] The specific functions included in this model are as follows:

[0085] (1) Realize the initialization and reset of the learnable structure parameter (score) corresponding to the pruning mask.

[0086] (2) Realize the initialization of the target sparsity and target delay and the corresponding warm-up process. The warm-up process gradually increases the target sparsity and target delay from 0 to the preset value, ensuring the stability of the pruning process.

[0087] (3) Calculate the sparsity loss of the model mask, including the loss of the number of parameters and the delay loss. Among them, the hard concrete distribution is used to model the pruning mask, so that a differentiable L1 regularization term can be used as an optimization objective during the training process. At the same time, due to different learning rates and pruning granularities, the pruning training may converge to models with very different sizes. Therefore, the present invention adds a Lagrangian term to ensure the convergence of the sparsity loss. The formula is as follows:

[0088]

[0089] Where is the sparsity obtained for the structure parameter z, and t is the target sparsity.

[0090] (4) The mapping rules from structural parameters to cropping masks are implemented, which is the forward propagation method in this module. Specifically, during the training process, the learnable parameters are mapped to soft masks to ensure the differentiability of the parameters during the optimization process; during the inference test, the learnable parameters are mapped to binary hard masks to test the performance of the cropping model.

[0091] (5) The structural parameters are aligned to the input of the delay prediction network described in step (4) of this section, thereby ensuring that the delay prediction network can be coupled to the 10 module and provide the 10 module with the function of predicting the inference delay and serving as the delay loss.

[0092] (II) Model component library:

[0093] This module fine-tunes the original network modules in BERT, and passes the cropping mask into the network as an additional input to control the model structure. This non-invasive model structure control method realizes the decoupling of the original network parameters and the structural parameters (score), making the cropping framework more universal. Any network based on the multi-head attention mechanism structure can be connected to the cropping framework with only a small amount of modification. In the specific implementation, the six cropping masks described in the input step 2 are fine-tuned. The specific implementation is: the output of the corresponding column of the network parameter controlled by the cropping granularity is multiplied by the cropping mask. For example, a mask parameter x controls the importance of the yth row in the weight matrix of a linear layer, then the yth dimension of the output after passing through the linear layer is multiplied by x, because the yth dimension of the output multiplied by x is equivalent to the yth row of the weight matrix multiplied by x, that is, the yth column of the linear layer is controlled by the mask parameter x, and the remaining cropping mask controls the network parameters in a similar way. In this way, the control of the network structure is achieved by inputting the cropping mask parameter.

[0094] 3. Dynamically mapped knowledge distillation module

[0095] like Figure 4 As shown in the figure, this module uses dynamic mapping to calculate the corresponding prediction layer distillation loss and intermediate layer distillation loss based on the output of the teacher model and the output of the student model. The specific steps are:

[0096] (1) The distillation loss of the prediction layer is predicted based on the logarithmic prediction probability of the student model and the teacher model, which is as follows:

[0097]

[0098] where z T and z S They represent the logarithm of the predicted probability of the teacher model and the student model respectively, and t represents the configurable distillation temperature.

[0099] (2) Use the Earth Mover’s Distance, a metric for calculating the distance between two distributions, to measure the similarity between different layers of the student model and the teacher model, that is, the knowledge transfer cost. The higher the similarity between different layers of the teacher and student models, the lower the knowledge transfer cost. Specifically, the similarity includes the similarity of the attention map and the similarity of the feature map. Specifically, the Earth Mover’s Distance between the feature maps of the student and teacher is defined by the following formula:

[0100]

[0101]

[0102]

[0103]

[0104]

[0105]

[0106] where \(H\) T , \(H\) S represent the hidden layer feature map matrices of the teacher model and the student model, represents the mapping flow, that is, the percentage of the knowledge of the \(i\)-th layer of the teacher model allocated to the \(j\)-th layer of the student model, represents the similarity between the \(j\)-th layer of the teacher model and the \(j\)-th layer of the student model, that is, the knowledge transfer cost. and represent the knowledge weight coefficients of different layers of the teacher model and the student model. Initially, they are \(1 / N\) (\(N\) is the number of model layers). This weight characterizes the percentage of the knowledge of the \(i\)-th layer of the teacher model regarding the feature map in the entire model, and the percentage of the knowledge that the \(j\)-th layer of the student model can receive in the entire model. In addition, the Earth Mover’s Distance between the attention maps of the student and teacher is similar to the above definition, and the feature map can be replaced with the attention map.

[0107] (3) Update the knowledge weights of the teacher model layer according to the unit knowledge transfer cost after calculating the Earth Mover’s Distance in step 2 of this part each time, that is and use the learned new weights for the constraint term of the next batch of training. Specifically, the unit knowledge transfer cost is defined as follows:

[0108]

[0109]

[0110] wherein respectively represent the unit knowledge transfer cost of the i-th layer of the teacher model. Further, according to the obtained unit knowledge transfer cost, update the knowledge weights of the teacher model layers, and the specific update formula is as follows:

[0111]

[0112]

[0113] (IV) Pruning Training Module:

[0114] This module provides a complete pruning pipeline by integrating the L0 regularization module and the model component library described in step five, and the delay prediction network described in step four. The specific steps include:

[0115] (1) Initialize the base model, i.e., the BERT model. And fine-tune it on the dataset of the downstream task, and use the fine-tuned model as the teacher model in the pruning process.

[0116] (2) Use the modules in the model component library to replace the original modules in BERT as the initialized pruning model.

[0117] (3) Initialize the L0 regularization module, including the initialization of the target sparsity and the target delay, the initialization of the pruning granularity, the embedding of the delay prediction network, and the initialization of the structural parameters. Initialize the learnable structural parameters to the fully retained state, which ensures that the model is stably pruned from the complete initial structure to the ideal optimal structure. The pruning type is set to include all pruning types in this example. The delay sparsity is set to 10% of the delay of the complete model, and the parameter quantity sparsity is 5% of the complete model. The temperature parameter in the hard concrete distribution is set to 0.66.

[0118] (4) Preprocess the downstream task data. Specifically, use a tokenizer to divide the sentence into small token chunks, generate a vocabulary, and learn better representations through the model.

[0119] (5) Use the L0 regularization module to output a pruning soft mask in the training mode according to the structural parameters and use it as part of the input in a training batch.

[0120] (6) Input a batch of data into the teacher model, and at the same time input the data of the same batch into the pruning model after adding the mask parameter to complete a forward propagation. In this example, the batch size is set to 32.

[0121] (7) Calculate the distillation loss based on the output of the teacher model in (6) and the output of the pruning model. The l0 regularization module calculates the sparsity loss, and the sum of the two losses is the total loss. That is, the goal of pruning training is to find an optimal network substructure to maximize the model performance while achieving the preset values of sparsity and latency.

[0122] (8) Use the AdamW optimizer to update the model parameters and the structure parameters. In this example, the learning rate of the model parameters is 3e -5 , and the learning rate of the learnable parameters is 0.01.

[0123] (9) Repeat steps (5)-(8) of this part until the preset number of rounds. In this example, it is set to 20.

[0124] (10) The l0 regularization module outputs the binary pruning mask under the inference model, prunes the model, and inputs the pruned small model.

[0125] Step Six: Use the optimal pruning mask obtained in Step Five to prune the model and perform fine-tuning to restore the accuracy, obtaining a compressed model; specifically, the accuracy restoration has an obvious effect on the pruning effect with a high sparsity. Among them, for small datasets such as RTE (The Recognizing Textual Entailment datasets), MRPC (The Microsoft Research Paraphrase Corpus), etc., the performance restoration effect is higher than that of large datasets such as MNLI (The Multi-Genre Natural Language Inference Corpus), QQP (The Quora Question Pairs), etc.

[0126] The following further details the parameter settings and experimental results in combination with the experiments on the sst dataset (text classification task). The hardware platform of the present invention is: Ubuntu 20.04.2, Pytorch 1.10.0, memory 64G, Intel Core i9-10900KF CPU, Nvidia RTX 3090 GPU, and the software platform is: Linux16.04 operating system, Python3.6, Pytorch0.4.0.

[0127] Table 1 Experimental parameters of the implementation case

[0128] Parameter Name Value Dataset Sst2 Cropping Granularity FLL, MHA, HEAD, HID, INT Learning Rate of Structure Parameter score 0.01 Learning Rate of Model Parameter <![CDATA[3e -5 > Epoch 20 Batchsize 32 Sparsity of Target Parameter Quantity 95％ Target Technology Ratio 10

[0129] Table 2 Performance Comparison between the Original BERT Model and the Pruned Model

[0130] Text Classification Task (sst2) Original Model (BERT) Cropped Model Comparison Model Size 85M 3.92M Compressed by 95.4% Accuracy 0.92 0.89 Decreased by 3% Speed (Number of Inferences per Second) 972 9932 Speed Increased by 10.2 Times

[0131] The above experimental results show that in this aspect, the large model can be well pruned to the preset volume and maintain good performance. In this example, the number of BERT parameters is compressed to about 5% of the original model, the speed is increased by 10 times, and only 3% of the accuracy drops, which can well ensure the deployment of the large model under resource-constrained conditions.

Claims

1. A hardware-aware large-scale pre-trained language model compression method, characterized in that: It includes the following steps: Step 1: Load the pre-trained language model after training completion and fine-tune it on downstream tasks; Step 2: Select the pruning granularity; Step 3: Construct a pruning mask - latency proxy dataset according to the selected pruning granularity, randomly sample the pruning mask, prune the model according to the adopted pruning mask, and test its inference latency on the specified hardware to form a data pair with the pruning mask - inference latency as a set of data; Step 4: Design a hardware-aware latency prediction network and train the latency prediction network using the proxy dataset; Step 5: Set the desired latency, embed the latency prediction network into the pruning framework of multiple hybrid granularities for pruning, and use the dynamic mapping knowledge distillation method to transfer the knowledge of the teacher model to the changing pruned model, so as to perceive the latency information for the target device and search for the best hardware-friendly small model; Step 6: Use the best pruning mask obtained by pruning to prune the model, restore the accuracy after fine-tuning, and obtain the compressed model.

2. The hardware-aware large-scale pre-trained language model compression method according to claim 1, characterized in that: The pre-trained model loaded in step 1 is the BERT model, which is fine-tuned on the SST-2 dataset for 4 rounds with a learning rate of .

3. The hardware-aware large-scale pre-trained language model compression method according to claim 1, characterized in that: The pruning granularity in Step 2 includes: network intermediate layer, feed-forward layer, attention layer, attention head, hidden layer dimension, and intermediate layer dimension, which act as the effect of parameter pruning in the form of a mask. The pruning masks corresponding to the pruning grains involved are: Feedforward layer: The corresponding pruning mask is represented by with a mask dimension of 12, indicating whether to prune the i-th intermediate layer of the network; Attention layer: The corresponding cropping mask is represented by indicating that the mask dimension is 12, indicating whether the i-th intermediate layer of the network is cropped or not; Attention head: The corresponding cropping mask is represented by indicating that the mask dimension is 12x12, indicating whether it is the j-th attention head in the i-th attention layer; Hidden feature dimension: The corresponding cropping mask is represented by indicating that the mask dimension is 768, indicating whether it is the dimension of the i-th hidden layer; Intermediate feature dimension: The corresponding cropping mask is represented by indicating that the mask dimension is 3072, indicating whether it is the dimension of the i-th intermediate layer; The specific formula is as follows: ； Among them, is the output of the multi-head attention layer, is the number of heads of this layer, , , , are the necessary query, key, value, and output matrices for constructing the attention mechanism of the multi-head attention layer, is the output of the feed-forward layer, is the weight matrix of the intermediate hidden layer in the feed-forward layer, is the weight matrix of the output layer in the feed-forward layer.

4. The hardware-aware large-scale pre-trained language model compression method according to claim 1, characterized in that: Step 3 includes the following sub-steps: Step 3.1: Determine the corresponding pruning mask for the selected pruning granularity, and determine a pruned model through the combination of multiple groups of pruning masks. The pruning mask is composed of 0 and 1, where 0 means pruning this group of parameters and 1 means retaining this group of parameters. The pruning granularity and the corresponding pruning mask involved are: Feedforward layer: The corresponding pruning mask is represented by which indicates that the mask dimension is 12, indicating whether to prune the i-th intermediate layer of the network; Attention layer: The corresponding cropping mask is represented by indicating that the mask dimension is 12, indicating whether the i-th intermediate layer of the network is cropped or not; Attention head: The corresponding cropping mask is represented by indicating that the mask dimension is 12x12, indicating whether the j-th attention head in the i-th attention layer is cropped or not; Hidden feature dimension: The corresponding cropping mask is represented by indicating that the mask dimension is 768, indicating whether it is the dimension of the i-th hidden layer; Intermediate feature dimension: The corresponding cropping mask is represented by indicating that the mask dimension is 3072, indicating whether it is the dimension of the i-th intermediate layer; Step 3.2: Randomly sample the pruning mask, prune the model according to the adopted pruning mask, test its inference latency, construct a proxy dataset, use the randomly sampled pruning mask as the feature of the proxy dataset, and the measured inference latency as the label.

5. The hardware-aware large-scale pre-trained language model compression method according to claim 1, characterized in that: In Step 4, first extract the latency features of different modules in the model through different branches. The aggregation branch is responsible for receiving the latency features output by other branches, and further linearly transforms the received latency features and then outputs the inference latency. The latency prediction network is a multi-branch tree-structured network. The features input to this network are: a one-dimensional binary vector with a length of 25. The specific meanings of its each dimension are: bits 1-12 represent the retention ratio of the attention heads in the 12 attention layers in BERT, bits 13-24 represent the retention ratio of the intermediate feature dimensions in the 12 feed-forward layers in BERT, and bit 25 represents the retention ratio of the hidden feature dimension in BERT; This network includes the following branches: (1) Attention layer latency feature extraction branch: The input is the 1st - 12th bits of the total input, and passes through fully connected layers with 12, 24, 12, and 1 neurons in sequence, and finally outputs to the aggregation branch. The leakRelu activation function is used between different layers to prevent all negative information from being lost; (2) Feed-forward layer latency feature extraction branch: The input is the 13th - 24th bits of the total input, and its structure is the same as that of the attention layer latency feature extraction branch, and finally outputs to the aggregation branch; (3)Hidden Dimension Delay Feature Extraction Branch: The input is the 25th bit of the total input, and it is connected to a fully-connected layer of a neuron, and finally output to the summarization branch; (4)Summarization Branch: The input is a one-dimensional vector of length 3. The values of the above three branches are concatenated to form a one-dimensional vector of length 3, and it is input into a fully-connected network. Finally, this network outputs the delay information on a specific hardware. This fully-connected network consists of two linear layers. The first layer contains 3 neurons, and the second layer contains 1 neuron.

6. The hardware-aware large-scale pre-trained language model compression method according to claim 1, wherein: The pruning framework with multiple hybrid granularities in step 5 includes: an l0 regularization module, a model component library, and a pruning training module; the composition and functions of each module include: Regularization Module: This module uniformly controls the learnable structure parameters corresponding to the pruning mask, and uses a differentiable l0 norm form to model the pruning mask. During the pruning training process, the regularization module is used to calculate the sparse loss based on the l0 coefficient; Model Component Library: This component library contains each sub-module that composes the BERT model after fine-tuning. The sub-module adds a control mechanism without changing the original function. The control mechanism means that the representation ability of this module is restricted by the incoming pruning mask; Dynamic Mapping Knowledge Distillation Module: According to the output of the teacher model and the output of the student model, the dynamic mapping is used to calculate the corresponding prediction layer distillation loss and the intermediate layer distillation loss. Among them, the prediction layer distillation loss is calculated based on the prediction probabilities of the student model and the teacher model, and the intermediate layer distillation loss is calculated based on the similarity between the feature distributions output by the student model and the teacher model. The higher the similarity between different layers of the teacher model and the student model, the lower the knowledge transfer cost of the teacher model, and correspondingly, more knowledge is transferred to this layer, that is, it accounts for a larger proportion in the distillation loss; Pruning Training Module: This module provides a complete pruning pipeline by integrating the l0 regularization module described in step 5.1, the model component library described in step 5.2, and the delay prediction network described in step 4. The specific steps include: 5.

1. Initialize the baseline model, that is, the BERT model, and fine-tune it on the dataset of the downstream task. The fine-tuned model is used as the teacher model during the pruning process; 5.

2. Replace the original module in BERT with the module in the model component library as the initialized pruning model; 5.

3. Initialize the l0 regularization module, including the initialization of the target sparsity and the target delay, the initialization of the pruning granularity, and the embedding and structure parameter initialization of the delay prediction network; 5.

4. Preprocess the downstream task data; 5.

5. Use the l0 regularization module to output a pruning soft mask in the training mode according to the structure parameters, and use it as a part of the input data in a batch training; 5.

6. Input a batch of data into the teacher model, and at the same time input the same batch of data into the pruning model after adding the mask parameter to complete a forward propagation; 5.

7. Calculate the distillation loss according to the output of the teacher model and the output of the pruning model in 5.6, and calculate the sparse loss according to the l0 regularization module. The sum of the two losses is the total loss; 5.

8. Update the model parameters and structural parameters using the adaptive moment estimation optimizer with weight decay; 5.

9. Repeat steps 5.6 - 5.8 until the preset number of iterations; 5.

10. Input the converged structural parameters into the l0 regularization module to obtain the binary pruning mask under the inference model, prune the model through this mask, and input the pruned small model.

Images