A model compression method based on adaptive knowledge distillation

Abstract

The application discloses a model compression method based on adaptive knowledge distillation, comprising the following steps: constructing a unified knowledge-aware search space, and designing a distillation knowledge form into a candidate scheme composed of a knowledge type, a knowledge position, a knowledge objective function and a balance coefficient; on the basis, using a knowledge adapter based on BO-TPE to perform adaptive search with the performance of a student model as an optimization target, and combining an early pruning strategy to accelerate the search, so as to obtain an optimal distillation configuration and output a distilled lightweight student model.

Description

Technical Field

[0001] This invention belongs to the technical field of knowledge distillation, specifically relating to a model compression method based on adaptive knowledge distillation. Background Technology

[0002] Pre-trained language models have achieved remarkable results in natural language processing tasks such as text classification, natural language inference, and semantic matching. However, these models typically have large parameter sizes and high inference computation costs, making them difficult to deploy directly in resource-constrained environments such as mobile phones, wearable devices, and edge servers. Therefore, how to compress and accelerate pre-trained language models while maintaining performance as much as possible is a crucial technical challenge in current service computing and intelligent terminal deployment.

[0003] Existing model compression techniques mainly include quantization, pruning, and knowledge distillation. Among these, knowledge distillation, by transferring knowledge learned from a large model to a lightweight student model, maintains deployment friendliness while preserving model capabilities as much as possible, thus becoming an important approach in pre-trained language model compression. For knowledge distillation of pre-trained language models, existing methods not only distill the final output of the teacher model but also attempt to distill various internal knowledge such as intermediate layer features, attention maps, and relational structures.

[0004] However, most existing distillation schemes rely on human experience to predetermine what knowledge to distill, at which layers to distill, and what loss function to use for knowledge matching.

[0005] Existing research has moved beyond simply distilling the final output of teacher models, exploring distillation knowledge from various perspectives such as intermediate layer features, attention distribution, and relational structures. Distillation of pre-trained language models can be summarized into three fundamental questions: first, what knowledge to distill; second, where to distill knowledge; and third, how to distill knowledge. Current methods typically rely on human experience to answer these questions, meaning researchers manually specify the type of knowledge to distill, the distillation layer, and the loss function. In pre-trained language model distillation research, knowledge can be broadly categorized into four types: response-related knowledge, feature-related knowledge, attention-related knowledge, and relational knowledge. BERT-PKD primarily distills the intermediate layers and the final output; TinyBERT distills the embedding layer, attention layer, and feature layer in task-independent and task-related stages; MiniLM emphasizes the query-key-value relationship in attention; and CKD and MGSKD further introduce relational knowledge such as lexical units and samples. Although these methods each have their own characteristics, their distillation schemes largely depend on manual pre-setting.

[0006] Different tasks, teacher-student model combinations, and student layers often exhibit different preferences for distillation knowledge design. This leads to extensive trial and error and parameter tuning required for manual design, resulting in poor generalization and a high risk of missing the optimal distillation configuration. Furthermore, the cost of manual combination and search increases rapidly as the number of candidate knowledge types, candidate distillation positions, and candidate loss functions continues to grow.

[0007] Although existing methods for acquiring student models have made significant progress, current pre-trained language model knowledge distillation methods lack a unified technical solution that can adaptively determine the optimal distillation knowledge design for different tasks and different teacher-student model combinations. As a result, the performance improvement of student models depends on human experience, and the distillation configuration is often suboptimal.

[0008] In addition to the main drawbacks mentioned above, existing technologies also present the following problems:

[0009] (1) Existing solutions usually fix the knowledge type, distillation level and objective function respectively, and cannot uniformly represent and automatically search for "what to distill, where to distill and how to distill".

[0010] (2) As the number of candidate distillation configurations increases, manual parameter adjustment and exhaustive experiments become costly and inefficient.

[0011] (3) There is a lack of a black-box optimization mechanism that dynamically updates the search direction based on historical distillation results.

[0012] (4) The lack of an early elimination mechanism for obviously unpromising candidate distillation configurations leads to excessively long search times.

[0013] (5) Existing distillation methods lack transferability between different teacher-student combinations and are difficult to accommodate different student sizes such as 3, 4, and 6 layers.

[0014] Therefore, how to optimize the distillation configuration and output a lightweight student model after distillation is an urgent problem to be solved. Summary of the Invention

[0015] The main objective of this invention is to overcome the shortcomings and deficiencies of the prior art and provide a model compression method based on adaptive knowledge distillation. By constructing a unified knowledge-aware search space and adaptive search based on a BO-TPE knowledge adapter, and combining an early pruning strategy to accelerate the search, the optimal distillation configuration is obtained and a lightweight student model after distillation is output.

[0016] To achieve the above objectives, the present invention adopts the following technical solution:

[0017] In a first aspect, the present invention provides a model compression method based on adaptive knowledge distillation, comprising the following steps:

[0018] S1. Design and construct a unified knowledge-aware search space based on candidate distillation knowledge;

[0019] S2. Initialize the BO-TPE-based knowledge adapter, including the surrogate model, observation set, and acquisition function; the observation set is used to store historical distillation designs and their corresponding student model performance.

[0020] S3. The knowledge adapter selects the next candidate distillation knowledge design in the knowledge perception search space, and performs a knowledge distillation training from teacher model to student model according to the candidate distillation knowledge design. According to the knowledge distillation constraints, the student model performance under the current candidate distillation knowledge design is obtained.

[0021] S4. Update the observation set and surrogate model based on the performance of the student model designed according to the current candidate distillation knowledge;

[0022] S5. Repeat steps S3 to S4 until the preset number of iterations is reached or the convergence condition is met, and output the optimal distillation knowledge design and the corresponding student model after distillation.

[0023] As a preferred technical solution, the candidate distillation knowledge design includes knowledge type, distillation candidate position, distillation candidate objective function, and candidate balance coefficient;

[0024] The knowledge types include response-based knowledge, feature-based knowledge, attention-based knowledge, and relationship-based knowledge;

[0025] The distillation candidate positions include a prediction layer, an intermediate layer, and an embedding layer;

[0026] The distillation candidate objective functions include KL, L1, L2, SmoothL1, cosine similarity, Pearson distance, and cross-entropy;

[0027] The candidate balance coefficients include task loss weights and weights for different knowledge items.

[0028] As a preferred technical solution, the BO-TPE-based knowledge adapter sets a maximization objective function. The optimal distillation knowledge design was used to search for the following: ,

[0029] in, This represents the optimal design of distillation knowledge. This represents a specific design of distillation knowledge. This represents the knowledge-aware search space, and argmax(·) represents the parameter function that takes the maximum value.

[0030] As a preferred technical solution, step S3 specifically includes:

[0031] The BO-TPE algorithm is used to continuously update the surrogate model based on the original observation set to estimate... ,in Let be the performance of the student model under the t-th DKD obtained in a given search process; in each iteration t, the next The choice is determined by the acquisition function;

[0032] In the t-th iteration, select Calculate the acquisition function, in acquiring Then, the new distillation knowledge will be used to design observations. Add to the original observation set to update the surrogate model. The new DKD observations are added to the original observation set, as shown in the following formula: ;

[0033] In each iteration, the BO-TPE algorithm selects an optimal acquisition function to calculate the result. Ultimately, this leads to the optimal performance of DKD and student models. .

[0034] As a preferred technical solution, the proxy model is constructed using conditional probability. and marginal probability To achieve this, the knowledge adapter divides the target values ​​of the probability distribution into two groups: good and bad, and the conditional probability... The model is as follows: ,

[0035] in, This represents the good density distribution of student model performance when selected. This represents the bad density distribution of student model performance at the time of selection. This represents the threshold used to judge the performance of the student model. This represents the optimal design of distillation knowledge. This represents a specific distillation knowledge design;

[0036] The marginal probability As shown in the following formula: ,

[0037] in, Indicates quantiles.

[0038] As a preferred technical solution, the threshold Quantiles of student model performance observations Sure.

[0039] As a preferred technical solution, the acquisition function is the one that is expected to be improved. As shown in the following formula: .

[0040] As a preferred technical solution, the knowledge distillation constraint includes task loss constraint, response-type knowledge distillation constraint, and at least one internal knowledge distillation constraint.

[0041] The task loss constraint is used for cross-entropy loss of a specific task;

[0042] The response-type knowledge distillation constraint originates from the final output representation of the teacher model and has rich predictive information;

[0043] The internal knowledge distillation constraints include feature-based knowledge, attention-based knowledge, and relation-based knowledge.

[0044] As a preferred technical solution, step S4 further includes determining whether the current candidate distillation knowledge design meets the early pruning condition. If it does, the training is terminated early and the candidate distillation knowledge design is eliminated.

[0045] As a preferred technical solution, the step of determining whether the current candidate distillation knowledge design meets the early pruning conditions specifically includes:

[0046] The performance of the student model corresponding to each DKD is evaluated during the search process. If the performance of the student model of a certain DKD fails to surpass that of the student model of the previous DKD, it is considered that the condition for early pruning is met.

[0047] Compared with the prior art, the present invention has the following advantages and beneficial effects:

[0048] (1) By modeling distillation knowledge design as a searchable object, this invention can reduce the reliance on manual distillation experience and improve the automation of distillation configuration design.

[0049] (2) By simultaneously searching for knowledge type, distillation location, objective function and weight configuration, this invention can adaptively find a better distillation scheme for different tasks and different teacher-student model combinations, thereby improving the performance of the student model.

[0050] (3) By introducing the BO-TPE knowledge adapter, the present invention can use historical distillation results to dynamically guide the subsequent search direction, thereby improving search efficiency and effectiveness.

[0051] (4) The present invention can eliminate unpromising candidate distillation designs as early as possible through the early pruning mechanism, significantly reducing the overall search time. Experiments show that the search efficiency can be improved by 6 times.

[0052] (5) Under multiple GLUE tasks and different teacher-student combinations, the corresponding framework of this invention has achieved results superior to many existing distillation methods, indicating that it has good generalization and practical value. Attached Figure Description

[0053] To more clearly illustrate the technical solutions in the embodiments of this application, the accompanying drawings used in the description of the embodiments will be briefly introduced below. Obviously, the accompanying drawings described below are only some embodiments of this application. For those skilled in the art, other drawings can be obtained based on these drawings without creative effort.

[0054] Figure 1 This is an embodiment of the present invention: an adaptive knowledge distillation graph based on knowledge awareness.

[0055] Figure 2 This is a flowchart illustrating the process of searching for the optimal DKD based on BO-TPE in an embodiment of the present invention. Detailed Implementation

[0056] To enable those skilled in the art to better understand the present application, the technical solutions in the embodiments of the present application will be clearly and completely described below with reference to the accompanying drawings. Obviously, the described embodiments are merely some embodiments of the present application, and not all embodiments. All other embodiments obtained by those skilled in the art based on the embodiments of the present application without creative effort are within the scope of protection of the present application.

[0057] In this application, the reference to "embodiment" means that a specific feature, structure, or characteristic described in connection with an embodiment may be included in at least one embodiment of this application. The appearance of this phrase in various places throughout the specification does not necessarily refer to the same embodiment, nor is it a mutually exclusive, independent, or alternative embodiment. It will be explicitly and implicitly understood by those skilled in the art that the embodiments described in this application can be combined with other embodiments.

[0058] Distillation Knowledge Design (DKD).

[0059] Please see Figure 1 This embodiment provides a model compression method based on adaptive knowledge distillation, including the following steps:

[0060] S1. Based on candidate distillation knowledge, design and construct a unified knowledge-aware search space.

[0061] In this embodiment, before constructing the knowledge-aware search space, the pre-trained language model to be compressed (teacher model T), student model S, target task training set, and validation set are taken, and a candidate distillation knowledge design (DKD) is set, including distillation candidate knowledge type, distillation candidate position, distillation candidate objective function, and balance coefficients for each candidate. Each candidate distillation knowledge design (DKD) is a combination of "what to distill, where to distill, how to distill, and the weights of each loss."

[0062] To adaptively determine the optimal DKD to improve the performance of the student model, this embodiment first constructs a knowledge-aware search space. This invention addresses three key issues in KD (Knowledge Distillation) and breaks down distillation knowledge design into four dimensions:

[0063] Knowledge type dimension: including response-based knowledge, feature-based knowledge, attention-based knowledge, and relationship-based knowledge.

[0064] Knowledge location dimension: includes prediction layer, intermediate layer and embedding layer.

[0065] Knowledge objective dimensions include KL, L1, L2, SmoothL1, cosine similarity, Pearson distance, and cross-entropy.

[0066] Balance coefficient dimension: includes task loss weights and the weights of different knowledge items , , .

[0067] Therefore, the knowledge design z for each candidate distillation can be expressed as: z={ , , , , , , },

[0068] in, Indicates the selected combination of knowledge types. This indicates the corresponding combination of distillation positions. This represents the objective function configuration corresponding to each knowledge item. Through this unified expression z, this embodiment can simultaneously cover existing manual distillation schemes and new knowledge combination configurations.

[0069] S2. Initialize the BO-TPE-based knowledge adapter, including the surrogate model, observation set, and acquisition function; the observation set is used to store historical distillation designs and their corresponding student model performance.

[0070] This embodiment utilizes a knowledge-aware search space. In design, maximize the objective function To find the optimal DKD, the following is a detailed analysis: ,

[0071] in, To represent a specific DKD, Indicates that in a given The performance of the student model, for example This represents the t-th DKD obtained during the search process. Indicates that in a given The performance of the student model.

[0072] See again Figure 1 In this embodiment, the knowledge adapter specifically employs a Bayesian optimization with Tree-structured Parzen Estimator (BO-TPE) algorithm. BO-TPE is an efficient black-box optimization algorithm that constructs... proxy model To search for the best DKD to improve student model performance.

[0073] Specifically, BO-TPE is based on the DKD set observed in previous iterations, i.e., the observation set. Continuously update the proxy model To estimate In each iteration t, the next... The choice is determined by a collection function that balances exploration and utilization, as detailed in steps S3 to S4.

[0074] S3. The knowledge adapter selects the next candidate distillation knowledge design in the knowledge perception search space, and performs a knowledge distillation training from teacher model to student model according to the candidate distillation knowledge design. According to the knowledge distillation constraints, the performance of the student model under the current candidate distillation knowledge design is obtained.

[0075] S4. Update the observation set and surrogate model based on the performance of the student model designed under the current candidate distillation knowledge.

[0076] like Figure 2 The diagram illustrates the process of the BO-TPE algorithm searching for the optimal DKD of a student model. To evaluate subsequent selections... right The contribution to the improvement is assessed using the Expected Improvement (EI) as the acquisition function. Given... In the case of exceeding the threshold The improvement probability is as follows: ,

[0077] in, Indicates that in a given Under the conditions The probability distribution.

[0078] In the t-th iteration, select To compute local optima That is, to obtain the maximum expected improvement .

[0079] In calculation Then, the new DKD observations will be... Add to the original observation set to update the surrogate model. The new DKD observations are added to the original observation set, as shown in the following formula: .

[0080] It is worth explaining that BO-TPE does not directly affect... Instead of modeling, through and To model this, BO-TPE constructs conditional probabilities by dividing the target value of the probability distribution into two groups: good and bad. The model is specifically defined as follows: ,

[0081] in, This represents the good density distribution of student model performance when selected. Indicates a bad density distribution. This represents the threshold used to judge the performance of the student model. The threshold represents the optimal distillation knowledge design. Quantiles of student model performance observations Sure.

[0082] Next, based on the above division and Bayes' theorem, the acquisition function can be restated as follows: ,

[0083] The marginal probability As shown in the following formula: ,

[0084] Ultimately expected improvement It can be represented as: ,

[0085] It is evident that improvements are expected. That is, maximizing the expected improvement in EI and maximizing Positive correlation, the ultimate goal is to maximize The probability of this decreases simultaneously. The probability of [the improvement in EI]. In each iteration, BO-TPE selects [a variable] with a higher expected improvement in EI. .

[0086] Finally, based on the knowledge distillation constraint, the best candidate DKD and student model performance for the current round are obtained. .

[0087] It is worth explaining that the task loss constraint is used for the cross-entropy loss of a specific task.

[0088] The response-based knowledge distillation constraint refers to knowledge derived from the final output representation of the teacher model, which possesses rich predictive information. Response-based knowledge transfer can be represented as: ,

[0089] Where KL[·|·] is the KL divergence function, It is the softmax function, and τ is the temperature hyperparameter. and These represent the final output representations of the teacher model and the student model, respectively.

[0090] The internal knowledge distillation constraints include feature-based knowledge, attention-based knowledge, and relation-based knowledge.

[0091] The feature-based knowledge includes [CLS] lexical knowledge and sequence lexical knowledge. This type of knowledge enables the student model to learn feature information from the intermediate hidden states of the teacher model. Compared with response-based knowledge distillation constraints, it covers more information from the teacher model's computation process. Feature-based knowledge transfer can generally be represented as: ,

[0092] in, and These represent the intermediate layer feature knowledge of the teacher model and the student model, respectively. This represents the objective function used to convey feature-based knowledge, where X represents the knowledge feature. This represents the hidden state of the intermediate layer of the teacher model. This represents the hidden state of the intermediate layer in the student model.

[0093] The attention-based knowledge includes attention graphs and query-key-value graph knowledge. Multi-head self-attention layers play a crucial role in pre-trained language models, responsible for encoding the contextual dependencies and interactions of the input text, and effectively capturing complex linguistic knowledge. The attention-based knowledge transfer can be represented as: ,

[0094] in, and These represent the attention knowledge of the teacher model and the student model, respectively. This represents the objective function used to convey attention-based knowledge.

[0095] The relation-based knowledge includes lexical, sample, and layer relational knowledge. This type of knowledge deeply analyzes the interrelationships between layers, data samples, and lexical units. Unlike traditional single-sample knowledge transfer methods, relation-based knowledge focuses on capturing the structural information of data in the embedding space. This knowledge transfer can be represented as: ,

[0096] in, and These represent the relationship between the teacher model and the student model, respectively. This represents the objective function used to convey relation-based knowledge.

[0097] In addition, in step S4, in order to improve the efficiency of the optimal DKD search, this embodiment also adopts a pruning strategy, that is, it is determined whether the current candidate distillation knowledge design meets the early pruning condition. If it does, the training of the candidate distillation knowledge design is terminated in advance and the candidate distillation knowledge design is eliminated.

[0098] A pruning strategy is employed to improve the efficiency of finding the optimal DKD by quickly identifying and eliminating poorly performing DKDs during the search process. Specifically, the strategy continuously evaluates the student model performance for each DKD during the search. If a DKD's student model performance fails to outperform previous DKDs—for example, if its performance is lower than the median of previous student models in the same iteration—the strategy terminates further training of that DKD early and removes it from the search, thus reducing overall search time. This pruning strategy, combined with the search algorithm, not only improves search efficiency but also ensures the performance of the selected DKDs.

[0099] S5. Repeat steps S3 to S4 until the preset number of iterations is reached or the convergence condition is met, and output the optimal distillation knowledge design and the corresponding student model after distillation.

[0100] It should be noted that, for the sake of simplicity, the aforementioned method embodiments are all described as a series of actions. However, those skilled in the art should understand that the present invention is not limited to the described order of actions, because according to the present invention, some steps can be performed in other orders or simultaneously.

[0101] The technical features of the above embodiments can be combined in any way. For the sake of brevity, not all possible combinations of the technical features in the above embodiments are described. However, as long as there is no contradiction in the combination of these technical features, they should be considered to be within the scope of this specification.

[0102] The above embodiments are preferred embodiments of the present invention, but the embodiments of the present invention are not limited to the above embodiments. Any changes, modifications, substitutions, combinations, or simplifications made without departing from the spirit and principle of the present invention shall be considered equivalent substitutions and shall be included within the protection scope of the present invention.

Claims

1. A model compression method based on adaptive knowledge distillation, characterized in that, Includes the following steps: S1. Design and construct a unified knowledge-aware search space based on candidate distillation knowledge; S2. Initialize the BO-TPE-based knowledge adapter, including the surrogate model, observation set, and acquisition function; The observation set is used to store historical distillation designs and their corresponding student model performance. S3. The knowledge adapter selects the next candidate distillation knowledge design in the knowledge perception search space, and performs a knowledge distillation training from teacher model to student model according to the candidate distillation knowledge design. According to the knowledge distillation constraints, the student model performance under the current candidate distillation knowledge design is obtained. S4. Update the observation set and surrogate model based on the performance of the student model designed according to the current candidate distillation knowledge; S5. Repeat steps S3 to S4 until the preset number of iterations is reached or the convergence condition is met, and output the optimal distillation knowledge design and the corresponding student model after distillation.

2. The model compression method based on adaptive knowledge distillation according to claim 1, characterized in that, The candidate distillation knowledge design includes knowledge type, distillation candidate position, distillation candidate objective function, and candidate balance coefficient; The knowledge types include response-based knowledge, feature-based knowledge, attention-based knowledge, and relationship-based knowledge; The distillation candidate positions include a prediction layer, an intermediate layer, and an embedding layer; The distillation candidate objective functions include KL, L1, L2, Smooth L1, cosine similarity, Pearson distance, and cross-entropy; The candidate balance coefficients include task loss weights and weights for different knowledge items.

3. The model compression method based on adaptive knowledge distillation according to claim 1, characterized in that, The BO-TPE-based knowledge adapter sets a maximization objective function to search for the optimal distillation knowledge design. As shown in the following formula: ， in, This represents a specific design of distillation knowledge. This represents the knowledge-aware search space, and argmax(·) represents the function that takes the maximum value parameter. Indicates that in a given Performance of the student model.

4. The model compression method based on adaptive knowledge distillation according to claim 3, characterized in that, Step S3 specifically includes: The BO-TPE algorithm is used to continuously update the surrogate model based on the original observation set to estimate... ,in Let be the performance of the student model under the t-th DKD obtained in a given search process; in each iteration t, the next The choice is determined by the acquisition function; In the t-th iteration, select Calculate the acquisition function, in the acquisition Then, the new distillation knowledge will be used to design observations. Add to the original observation set to update the surrogate model. The new DKD observations are added to the original observation set, as shown in the following formula: ; In each iteration, the BO-TPE algorithm selects an optimal acquisition function to calculate the result. Ultimately, this leads to the optimal performance of DKD and student models. .

5. The model compression method based on adaptive knowledge distillation according to claim 1, characterized in that, The proxy model is constructed through conditional probability. and marginal probability To achieve this, the knowledge adapter divides the target values ​​of the probability distribution into two groups: good and bad, and the conditional probability... The model is as follows: ， in, This represents the good density distribution of student model performance when selected. This represents the bad density distribution of student model performance at the time of selection. This represents the threshold used to judge the performance of the student model. This represents the optimal design of distillation knowledge. This represents a specific distillation knowledge design; The marginal probability As shown in the following formula: ， in, Indicates quantiles.

6. The model compression method based on adaptive knowledge distillation according to claim 5, characterized in that, The threshold Quantiles of student model performance observations Sure.

7. The model compression method based on adaptive knowledge distillation according to claim 6, characterized in that, The acquisition function is the one that is expected to be improved. As shown in the following formula: 。 8. The model compression method based on adaptive knowledge distillation according to claim 1, characterized in that, The knowledge distillation constraints include task loss constraints, response-type knowledge distillation constraints, and at least one internal knowledge distillation constraint. The task loss constraint is used for cross-entropy loss of a specific task; The response-type knowledge distillation constraint originates from the final output representation of the teacher model and has rich predictive information; The internal knowledge distillation constraints include feature-based knowledge, attention-based knowledge, and relation-based knowledge.

9. The model compression method based on adaptive knowledge distillation according to claim 1, characterized in that, Step S4 further includes determining whether the current candidate distillation knowledge design meets the early pruning condition. If it does, the training will be terminated early and the candidate distillation knowledge design will be eliminated.

10. The model compression method based on adaptive knowledge distillation according to claim 1, characterized in that, The determination of whether the current candidate distillation knowledge design meets the early pruning condition is as follows: The performance of the student model corresponding to each candidate distillation knowledge design is evaluated during the search process. If the performance of the student model of a candidate distillation knowledge design is lower than that of the student model of the previous distillation knowledge design, it is considered that the condition for early pruning is met.

Images