Knowledge distillation method based on TinyBERT

Abstract

The invention discloses a knowledge distillation method based on TinyBERT, and the method comprises the following steps: S1, employing BERT as a teacher model, employing TinyBERT as a student model, and carrying out the distillation; s2, defining a distillation loss function, and updating a student attention layer; s3, optimizing hierarchical mapping, and generating optimization output in combination with the attention weight; s4, data enhancement and pseudo target generation are carried out, and the generalization ability is improved; s5, dynamic weight adjustment is carried out, and a loss weight ratio is optimized; according to the invention, by optimizing the knowledge distillation technology, the knowledge of a large-scale model is efficiently transmitted to a small-scale student model, and the light weight, high performance and wide applicability of the model are realized.

Description

Technical Field

[0001] The present invention relates to the field of knowledge machine learning technology, and in particular to a knowledge distillation method based on TinyBERT. Background Art

[0002] In recent years, the rapid development of natural language processing has spawned a large number of models with excellent performance, especially the emergence of pre-trained models based on the Transformer architecture, such as BERT and GPT. These models learn rich language features through large-scale corpus training, significantly improving the performance of tasks such as text classification, question-answering systems, and machine translation, and promoting the development of natural language processing technology. However, as the model size and number of parameters continue to increase, these large models have exposed a series of problems in practical applications.

[0003] First of all, the number of parameters of models such as BERT and GPT is extremely large, which directly leads to high demand for computing resources. During the training and inference stages, these models usually require high-performance computing hardware, such as GPU or TPU support. This high resource demand limits the widespread application of models in the industry and increases the usage threshold for small and medium-sized enterprises and developers. In addition, large models have high storage requirements and operating environment requirements, resulting in high deployment costs on resource-constrained devices (such as mobile devices and embedded systems). This high cost of computing resource consumption and deployment threshold makes it difficult to apply the model in low computing power scenarios.

[0004] Secondly, large models usually need to be fine-tuned in practical applications to adapt to specific tasks. However, this fine-tuning process places extremely high demands on data scale and computing power, further increasing costs. At the same time, when large models are migrated to new domains, they usually require a large amount of domain data to be retrained to achieve domain adaptation. This migration difficulty not only limits the flexibility of the model, but also significantly increases the complexity of development and maintenance.

[0005] The above problems show that large language models are difficult to fully meet the needs of practical applications under the current computing environment and hardware conditions. Therefore, how to reduce the scale and computational complexity of the model while maintaining its performance has become a research hotspot of common concern in academia and industry. In order to solve this problem, model compression technology has emerged. Model compression makes the model more efficient and lightweight by reducing the number of model parameters and computing requirements. Among them, knowledge distillation, as an efficient model compression method, has received widespread attention in recent years.

[0006] The core idea of ​​knowledge distillation is to build a teacher-student model framework, let the student model imitate the output behavior of the teacher model, learn its internal structure and decision rules, and thus achieve model compression and acceleration while maintaining performance. This method introduces distillation loss to minimize the difference between the output of the student model and the teacher model, so that the student model gradually approaches the teacher model. However, existing knowledge distillation methods still have obvious shortcomings in loss function design, attention distribution learning, dynamic weight adjustment, and domain generalization capabilities.

[0007] Take TinyBERT as an example. As a compact version of BERT, this model achieves model compression through knowledge distillation technology. However, the distillation process of TinyBERT adopts a fixed hierarchical mapping method, ignoring the importance of features at different levels in knowledge learning. In addition, its Transformer layer loss mainly calculates the difference in attention distribution through the mean square error (MSE), fails to fully utilize the probability distribution information of the teacher model, and is difficult to maintain robustness in the case of noisy labels. At the same time, when dealing with hidden layer representation loss and attention loss, TinyBERT simply adds the two directly, and fails to dynamically adjust the weight distribution for different task scenarios. More importantly, its data augmentation strategy is relatively limited and fails to significantly improve the generalization ability of the student model in unseen fields.

[0008] Therefore, how to provide a knowledge distillation method based on TinyBERT is an urgent problem that technicians in this field need to solve. Summary of the invention

[0009] One object of the present invention is to propose a knowledge distillation method based on TinyBERT. The present invention significantly reduces the number of parameters and computational complexity of the student model by optimizing the key technologies of knowledge distillation, so that it can run efficiently in a resource-constrained environment while maintaining performance close to that of the teacher model, thereby achieving a breakthrough in model compression technology in practical applications.

[0010] A knowledge distillation method based on TinyBERT according to an embodiment of the present invention includes the following steps:

[0011] S1. Use the network architecture of teacher model and student model, use BERT model as teacher model and TinyBERT model as student model to perform knowledge distillation.

[0012] S2. During the distillation process, define the distillation loss, including word embedding loss, Transformer layer loss and prediction loss, and update the attention layer output of the student model;

[0013] S3. Based on the updated attention layer output of the student model, perform hierarchical mapping optimization. By automatically learning the nonlinear mapping function, the hierarchical mapping relationship of the student model is matched with the hierarchical mapping of the teacher model. The hierarchical mapping relationship is optimized in combination with the attention distribution weight of the teacher model to generate the optimized student model output.

[0014] S4. Using data augmentation and pseudo target generation techniques, based on the optimized student model output, further improve the performance and generalization ability of the student model on unseen domain data, and generate enhanced student model output;

[0015] S5. Introduce a dynamic weight adjustment mechanism to optimize the weight ratio of hidden layer representation loss and attention loss according to the enhanced student model output, so as to improve the training effect of the final student model.

[0016] Optionally, the S2 specifically includes:

[0017] S21, Transformer layer loss measures the difference in attention distribution between the student model and the teacher model through the improved KL divergence loss function;

[0018] S22, use temperature adjustment method to optimize KL divergence loss and enhance robustness to noisy labels;

[0019]

[0020] Among them, KL LOSS () is the KL divergence loss function, τ is the temperature adjustment parameter, is the probability of the jth category of the teacher model, is the probability of the jth category of the student model, max(τ,τ 2 ) is the dynamic coefficient, is the distribution difference between the teacher model and the student model on the jth category.

[0021] Optionally, the S3 specifically includes:

[0022] S31, dividing the layers of the teacher model into multiple groups, and corresponding to the layers of the student model respectively;

[0023] S32, dynamically adjust the learning focus of the student model by automatically learning the nonlinear mapping weights and combining the attention distribution of the teacher model;

[0024]

[0025] Among them, V k Describes the corresponding relationship between the teacher model and the student model, softmax() is the normalization function, M represents the number of student layers, θ k is the corresponding weight parameter;

[0026] S33. By optimizing the attention weight matrix, the mapping relationship is made more consistent with the knowledge distribution of different layers in the teacher model:

[0027] Vfinal=V k softmax(α attn );

[0028] Among them, Vfinal represents the mapping weight of each layer, α attn is the attention distribution of each layer.

[0029] Optionally, the S4 specifically includes:

[0030] S41, Generate out-of-domain sample data through a generative pre-training model to expand the training set;

[0031] S42, improve the learning ability of the student model for minority class samples through random oversampling and SMOTE oversampling technology;

[0032] S43. Use model ensembles to identify minority class samples and enhance the weights of these samples during the distillation process.

[0033] S44, generating multiple pseudo targets through beam search, and increasing the diversity of pseudo targets through random sampling;

[0034] S45, the pseudo targets generated by the joint teacher model and the student model are used alternately during the training process to reduce the exposure bias of the model;

[0035] S46. Improve the generalization performance of the student model by optimizing the conditional distribution of pseudo targets.

[0036] Optionally, the dynamic weight adjustment mechanism specifically includes dynamically controlling the weight ratio of hidden layer representation loss and attention loss by introducing a variable weight parameter:

[0037] l all =αl hidn +(1-α)l attn ;

[0038] Among them, l all is the total loss function value, l hidn is the hidden layer representation loss, and α is a variable weight parameter.

[0039] The beneficial effects of the present invention are:

[0040] The present invention optimizes the knowledge distillation technology to efficiently transfer the knowledge of large pre-trained models to smaller and lighter student models, effectively solving the high consumption of storage requirements and computing resources of traditional large models. This improvement enables the student model to significantly reduce the complexity and operating costs of the model while maintaining high performance, and is particularly suitable for resource-constrained environments such as mobile devices, embedded devices, or edge computing devices. Through the design of the knowledge distillation process, the present invention enables the student model to inherit the robustness of the teacher model and show stronger stability to perturbations of the input data, thereby having higher resistance to adversarial attack scenarios. This feature significantly enhances the reliability and security of the student model in practical applications.

[0041] In addition, the present invention realizes efficient transfer of knowledge between different architecture models through an improved hierarchical mapping mechanism, so that the student model can fully learn the core knowledge in the teacher model while maintaining similar interfaces and performance as the teacher model. This design not only simplifies the integration process between different models, but also improves the flexibility of the system, providing researchers and engineers with a more convenient development and deployment experience. By optimizing the generalization ability of the student model, the present invention significantly improves the performance of the student model on test data, especially when the training data is limited, the student model can better adapt to the distribution of unseen data, thereby showing stronger applicability and robustness.

[0042] In summary, the present invention not only improves the traditional knowledge distillation technology, but also further improves the robustness, generalization ability and operating efficiency of the student model, providing an efficient, flexible and reliable solution for natural language processing tasks in resource-constrained scenarios. By migrating the deep knowledge of the teacher model to the student model, the present invention broadens the application scope of small models in diverse scenarios, while lowering the deployment threshold of the model, providing a new development direction for model compression technology. BRIEF DESCRIPTION OF THE DRAWINGS

[0043] The accompanying drawings are used to provide a further understanding of the present invention and constitute a part of the specification. Together with the embodiments of the present invention, they are used to explain the present invention and do not constitute a limitation of the present invention. In the accompanying drawings:

[0044] Figure 1 A flow chart of a knowledge distillation method based on TinyBERT proposed in the present invention;

[0045] Figure 2 This is a schematic diagram of the optimization effect of the knowledge distillation method based on TinyBERT proposed in the present invention. DETAILED DESCRIPTION

[0046] The present invention will now be described in further detail with reference to the accompanying drawings. These drawings are simplified schematic diagrams, which only illustrate the basic structure of the present invention in a schematic manner, and therefore only show the components related to the present invention.

[0047] refer to Figure 1-2 , a knowledge distillation method based on TinyBERT, including the following steps:

[0048] S1. Use the network architecture of teacher model and student model, use BERT model as teacher model and TinyBERT model as student model to perform knowledge distillation.

[0049] S2. During the distillation process, define the distillation loss, including word embedding loss, Transformer layer loss and prediction loss, and update the attention layer output of the student model;

[0050] S3, based on the updated student model’s attention layer output, perform hierarchical mapping optimization, automatically learn the nonlinear mapping function, correspond the student model’s hierarchy to the teacher model’s hierarchy, and optimize the hierarchical mapping relationship in combination with the teacher model’s attention distribution weights to generate the optimized student model output;

[0051] S4. Using data augmentation and pseudo target generation techniques, based on the optimized student model output, further improve the performance and generalization ability of the student model on unseen domain data, and generate enhanced student model output;

[0052] S5. Introduce a dynamic weight adjustment mechanism to optimize the weight ratio of hidden layer representation loss and attention loss according to the enhanced student model output, so as to improve the training effect of the final student model.

[0053] In this implementation, S2 specifically includes:

[0054] S21, Transformer layer loss measures the difference in attention distribution between the student model and the teacher model through the improved KL divergence loss function;

[0055] S22, use temperature adjustment method to optimize KL divergence loss and enhance robustness to noisy labels;

[0056]

[0057] Among them, KL LOSS () is the KL divergence loss function, τ is the temperature adjustment parameter, is the probability of the jth category of the teacher model, is the probability of the jth category of the student model, max(τ,τ 2 ) is the dynamic coefficient, is the distribution difference between the teacher model and the student model on the jth category.

[0058] In this implementation, S3 specifically includes:

[0059] S31, dividing the layers of the teacher model into multiple groups, and corresponding to the layers of the student model respectively;

[0060] S32, dynamically adjust the learning focus of the student model by automatically learning the nonlinear mapping weights and combining the attention distribution of the teacher model;

[0061]

[0062] Among them, V k Describes the corresponding relationship between the teacher model and the student model, softmax() is the normalization function, M represents the number of student layers, θ k is the corresponding weight parameter;

[0063] S33. By optimizing the attention weight matrix, the mapping relationship is made more consistent with the knowledge distribution of different layers in the teacher model:

[0064] Vfinal=V k softmax(α attn );

[0065] Among them, Vfinal represents the mapping weight of each layer, α attn is the attention distribution of each layer.

[0066] In this implementation, S4 specifically includes:

[0067] S41, Generate out-of-domain sample data through a generative pre-training model to expand the training set;

[0068] S42, improve the learning ability of the student model for minority class samples through random oversampling and SMOTE oversampling technology;

[0069] S43. Use model ensembles to identify minority class samples and enhance the weights of these samples during the distillation process.

[0070] S44, generating multiple pseudo targets through beam search, and increasing the diversity of pseudo targets through random sampling;

[0071] S45, the pseudo targets generated by the joint teacher model and the student model are used alternately during the training process to reduce the exposure bias of the model;

[0072] S46. Improve the generalization performance of the student model by optimizing the conditional distribution of pseudo targets.

[0073] In this implementation, the dynamic weight adjustment mechanism specifically includes dynamically controlling the weight ratio of hidden layer representation loss and attention loss by introducing a variable weight parameter:

[0074] l all =αl hidn +(1-α)l attn ;

[0075] Among them, l all is the total loss function value, l hidn is the hidden layer representation loss, and α is a variable weight parameter.

[0076] Embodiment 1:

[0077] In order to verify the feasibility of the present invention in implementation, the present invention is applied to the personalized recommendation system of the online education platform. The recommendation system needs to process a large amount of real-time user behavior data, such as course browsing records, test answering status, and learning video viewing preferences, in order to recommend learning resources that better meet personal needs for users. However, although the traditional system using large models (such as BERT) has a high recommendation accuracy, it is difficult to achieve the goals of real-time and low-resource deployment in mobile devices and embedded scenarios due to high computing resources and storage requirements.

[0078] In order to meet this challenge, the present invention proposes a knowledge distillation method based on TinyBERT, which transfers knowledge from the teacher model BERT to the student model TinyBERT to achieve substantial model compression and performance optimization. In the experiment, BERT was used as a teacher model and was pre-trained using user behavior data collected by the platform. The amount of data exceeds 1 million and is distributed in four fields: education, e-commerce, social media, and medical care. Among them, 50% of the data is used for training, 30% for verification, and 20% for testing. After the training is completed, the knowledge distillation method of the present invention is used to migrate BERT's knowledge to the TinyBERT model.

[0079] In the server scenario, TinyBERT is deployed to process large-scale user behavior data of online education platforms, and achieves real-time recommendations through its efficient reasoning capabilities. At the same time, in order to verify its applicability in low-resource environments, the TinyBERT model is also deployed in mobile devices for personalized recommendations in resource-constrained environments. This scenario simulates the real-time response capability of the system when students learn through mobile terminals.

[0080] The TinyBERT model uses the knowledge distillation method to learn the hidden layer representation features and attention distribution from the BERT model. During the training process, the improved dynamic weight adjustment mechanism effectively optimizes the ratio distribution of hidden layer representation loss and attention loss, enhancing the performance of TinyBERT in recommendation tasks. In actual deployment, TinyBERT significantly improves the inference speed on the server side, and achieves lightweight operation on mobile devices to meet real-time requirements.

[0081] The experiments further verified the performance of the present invention in different noise label ratio scenarios and cross-domain data. In the noise label scenario, TinyBERT showed strong robustness through the improved KL divergence loss function. In the cross-domain generalization ability test, TinyBERT showed high recommendation accuracy in education, e-commerce, social media, and medical fields, close to the performance of the original BERT.

[0082] Table 1 Model performance comparison

[0083]

[0084]

[0085] Table 2 Recommendation accuracy of the model under different noise label ratios

[0086]

[0087] Table 3 Comparison of model generalization ability

[0088]

[0089] Table 1 shows the inference performance, model parameter amount and recommendation accuracy of the original BERT model, TinyBERT model and compressed GPT model in server and mobile device scenarios. The original BERT model has a parameter amount of up to 340 million and a model size of 420MB. Although it can run at a speed of 250 items / second on the server side, it cannot run on mobile devices. In addition, its recommendation accuracy is 93.5%, which is high but not friendly to resource requirements. In contrast, TinyBERT achieves efficient model compression through the knowledge distillation method of the present invention. Its parameter amount is only 66 million, the model size is only 80MB, the inference speed reaches 800 items / second on the server side, and it can run at a speed of 15 milliseconds / item on mobile devices. At the same time, the recommendation accuracy is only 1.1% lower than BERT, reaching 92.4%. In contrast, the parameter amount and model size of the compressed GPT model are between the two, with an inference speed of 600 items / second on the server side and 20 milliseconds / item on the mobile device side, and a recommendation accuracy of 91.7%. As can be seen from the table, TinyBERT shows superiority in inference speed, recommendation accuracy, and resource requirements, especially its ability to run in mobile devices, indicating that it is more suitable for resource-constrained environments.

[0090] Table 2 compares the recommendation accuracy of TinyBERT, original BERT, compressed GPT and traditional compression models under different noise label ratios. When the noise ratio is 0%, the recommendation accuracy of all models shows performance close to the maximum value. Among them, the recommendation accuracy of TinyBERT is 92.4%, which is only 1.1% lower than the 93.5% of the original BERT. As the noise ratio increases, the recommendation accuracy of all models decreases to varying degrees, but there are significant differences in the degree of decrease. When the noise ratio is 15%, the recommendation accuracy of TinyBERT is 90.3%, which is 2.1 percentage points lower than the noise-free case; while the traditional compression model drops sharply from 91.2% to 84.6%, showing its strong sensitivity to noise labels. The decline of the original BERT and compressed GPT is relatively small, 1.7% and 2.8% respectively. It can be seen from the data that the present invention enhances the robustness of TinyBERT to noise data through the improved KL divergence loss function, so that it can still maintain a high recommendation performance in complex data scenarios.

[0091] Table 3 shows the recommendation accuracy performance of TinyBERT, original BERT, compressed GPT and unoptimized TinyBERT in cross-domain data. The original BERT showed high accuracy in all fields, up to 93.5%, and after TinyBERT was optimized by the knowledge distillation method of the present invention, the recommendation accuracy in the field of online education reached 92.4%, which was only 1.1% lower than the original BERT. In data from other fields, TinyBERT also showed superior generalization ability. For example, in the fields of e-commerce and social media, its recommendation accuracy was 89.8% and 88.7% respectively, which was less than 2 percentage points from the original BERT's 91.5% and 90.4%. The performance of the unoptimized TinyBERT in the same scenario was significantly lower than that of the optimized model, such as only 81.0% in the medical field data, while the optimized TinyBERT was 87.3%. From these data, it can be seen that the automatic hierarchical mapping method and pseudo-target generation technology proposed in the present invention significantly improve the cross-domain generalization ability of TinyBERT, enabling it to maintain a high recommendation performance in a variety of scenarios. Combining the analysis results of the above three tables, it can be concluded that the knowledge distillation method based on TinyBERT proposed in this invention shows significant advantages in performance optimization, noise label processing, and cross-domain generalization ability improvement. These data fully reflect the technical innovation and practical application value of this invention, and provide an effective solution for model optimization and deployment in resource-constrained scenarios.

[0092] The above description is only a preferred specific implementation manner of the present invention, but the protection scope of the present invention is not limited thereto. Any technician familiar with the technical field can make equivalent replacements or changes according to the technical scheme and inventive concept of the present invention within the technical scope disclosed by the present invention, which should be covered by the protection scope of the present invention.

Claims

1. A knowledge distillation method based on TinyBERT, characterized in that: The steps include: S1. Use the network architecture of teacher model and student model, use BERT model as teacher model and TinyBERT model as student model to perform knowledge distillation. S2. During the distillation process, define the distillation loss, including word embedding loss, Transformer layer loss and prediction loss, and update the attention layer output of the student model; S3. Based on the updated attention layer output of the student model, perform hierarchical mapping optimization. By automatically learning the nonlinear mapping function, the hierarchical mapping relationship of the student model is matched with the hierarchical mapping of the teacher model. The hierarchical mapping relationship is optimized in combination with the attention distribution weight of the teacher model to generate the optimized student model output. S4. Using data augmentation and pseudo target generation techniques, based on the optimized student model output, further improve the performance and generalization ability of the student model on unseen domain data, and generate enhanced student model output; S5. Introduce a dynamic weight adjustment mechanism to optimize the weight ratio of hidden layer representation loss and attention loss according to the enhanced student model output, so as to improve the training effect of the final student model.

2. A knowledge distillation method based on TinyBERT according to claim 1, characterized in that: The S2 specifically includes: S21, Transformer layer loss measures the difference in attention distribution between the student model and the teacher model through the improved KL divergence loss function; S22, use temperature adjustment method to optimize KL divergence loss and enhance robustness to noisy labels; Among them, KL LOSS () is the KL divergence loss function, τ is the temperature adjustment parameter, is the probability of the jth category of the teacher model, is the probability of the jth category of the student model, max(τ,τ 2 ) is the dynamic coefficient, is the distribution difference between the teacher model and the student model on the jth category.

3. A knowledge distillation method based on TinyBERT according to claim 1, characterized in that: The S3 specifically includes: S31, dividing the layers of the teacher model into multiple groups, and corresponding to the layers of the student model respectively; S32, dynamically adjust the learning focus of the student model by automatically learning the nonlinear mapping weights and combining the attention distribution of the teacher model; Among them, V k Describes the corresponding relationship between the teacher model and the student model, softmax() is the normalization function, M represents the number of student layers, θ k is the corresponding weight parameter; S33. By optimizing the attention weight matrix, the mapping relationship is made more consistent with the knowledge distribution of different layers in the teacher model: Vfinal=V k softmax(α attn ); Among them, Vfinal represents the mapping weight of each layer, α attn is the attention distribution of each layer.

4. A knowledge distillation method based on TinyBERT according to claim 1, characterized in that: The S4 specifically includes: S41, Generate out-of-domain sample data through a generative pre-training model to expand the training set; S42, improve the learning ability of the student model for minority class samples through random oversampling and SMOTE oversampling technology; S43. Use model ensembles to identify minority class samples and enhance the weights of these samples during the distillation process. S44, generating multiple pseudo targets through beam search, and increasing the diversity of pseudo targets through random sampling; S45, the pseudo targets generated by the joint teacher model and the student model are used alternately during the training process to reduce the exposure bias of the model; S46. Improve the generalization performance of the student model by optimizing the conditional distribution of pseudo targets.

5. A knowledge distillation method based on TinyBERT according to claim 1, characterized in that: The dynamic weight adjustment mechanism specifically includes dynamically controlling the weight ratio of hidden layer representation loss and attention loss by introducing variable weight parameters: l all ＝αl hidn +(1-α)l attn ； Among them, l all is the total loss function value, l hidn is the hidden layer representation loss, and α is a variable weight parameter.

Images