Method for optimizing probability mass distribution in knowledge distillation based on alpha-beta-divergence

Abstract

The invention discloses a probability mass distribution optimization method in knowledge distillation based on alpha-beta-divergence, and the method comprises the steps: setting a distillation frame for distillation to obtain a student model: setting a monitoring quantity model which is used for representing the probability mass change during gradient updating of each step; reducing the upper bound of the monitoring quantity model through a divergence function FKLD; reducing the upper bound of the monitoring quantity model through a divergence function RKLD; unifying the two monitoring quantity models of which the upper bounds are reduced to obtain an adjustable unified upper bound; analyzing an adjustable unified upper bound, and setting a difference model; and setting an optimization target based on the difference model, and obtaining optimal parameters of the student model through an optimization method, thereby completing training of the student model. According to the method disclosed by the invention, it can be ensured that the student model can focus on the category with large errors in the training process, and learning of global knowledge of the teacher model can be kept, so that the student model can show excellent performance in multiple tasks.

Description

Technical Field

[0001] The present invention relates to a probability mass allocation optimization method in knowledge distillation based on α-β-divergence, and belongs to the technical field of image generation. Background Art

[0002] Knowledge distillation (KD) technology is widely used to transfer knowledge from large-scale models to small-scale models. By directly minimizing the output distribution difference between the teacher model and the student model, knowledge distillation enables the student model to learn rich information from the teacher model, often achieving better performance than traditional supervised fine-tuning.

[0003] Existing knowledge distillation methods often use forward Kullback-Leibler divergence (FKLD). However, the asymmetry of FKLD forces the student distribution to cover all support sets of the teacher distribution, resulting in an overly smoothed student distribution.

[0004] Existing techniques also include Reverse Knowledge Distillation (RKLD), which enables the student model to focus on the key patterns of the teacher model. While RKLD is effective in certain situations, existing experimental results show that its performance on various tasks is suboptimal. Furthermore, the lack of systematic methods to identify potential issues hinders the development of more general and efficient knowledge distillation frameworks.

[0005] Therefore, it is necessary to conduct more in-depth research on existing knowledge distillation methods to solve the above problems. Summary of the Invention

[0006] In order to overcome the above problems, an in-depth study was conducted and a probability mass allocation optimization method in knowledge distillation based on α-β-divergence was proposed, which includes the following steps:

[0007] Use the trained teacher model to predict or generate samples in the dataset to obtain sample soft labels;

[0008] Set up a distillation framework and train the student model based on the sample soft labels;

[0009] The trained student model is used to predict or generate data to obtain prediction or generation results.

[0010] In a preferred embodiment, the setting of the distillation framework to train the student model based on the sample soft labels includes the following sub-steps:

[0011] S21. Set a monitoring quantity model to characterize the probability mass change at each step of gradient update;

[0012] S22, lowering the upper bound of the monitoring quantity model by using the divergence function FKLD; lowering the upper bound of the monitoring quantity model by using the divergence function RKLD;

[0013] S23, unifying the two monitoring quantity models after lowering the upper bound to obtain an adjustable unified upper bound;

[0014] S24, analyze the adjustable unified upper bound and set the difference model;

[0015] S25. Based on the difference model, set the optimization goal and obtain the optimal parameters of the student model through the optimization method, thereby completing the training of the student model.

[0016] In a preferred embodiment, in S21, the monitoring amount is expressed as:

[0017]

[0018] in, For monitoring, represents the knowledge distillation framework, is the probability mass of category y obtained by knowledge distillation at step t+1, q t (y) represents the original probability mass of category y at step t of the student model.

[0019] In a preferred embodiment, q t (y) is obtained by approximating the category probability through the softmax function, expressed as:

[0020]

[0021] in, Represents the logit value of category y at step t.

[0022] In a preferred embodiment, Expressed as:

[0023]

[0024] in, represents the logit value of category k at step t, η represents the learning rate, l represents the training loss, Represents the gradient.

[0025] In a preferred embodiment, in S22, the upper bound of the monitoring quantity model is reduced by the divergence function FKLD as follows:

[0026]

[0027] The upper bound of the monitoring model is reduced by the divergence function RKLD as follows:

[0028]

[0029] Among them, (a), (b), (a1), (b1), (a2), (b2) represent different terms in the model, p(y) represents the probability mass of category y in the teacher model, and q t (y) represents the original probability mass of category y at step t of the student model, represents the constant term in the divergence function FKLD that is independent of the category y, Represents the constant term in the divergence function RKLD that is independent of the category y.

[0030] In S23, the adjustable unified upper bound is expressed as:

[0031]

[0032] in, is an adjustable unified upper bound, and α and β are adjustable hyperparameters.

[0033] In a preferred embodiment, the difference model in S24 is set to:

[0034]

[0035] in, represents the difference model, C represents the total number of categories, p(k) represents the probability mass of category k of the teacher model, and q(k) represents the probability mass of category k of the student model.

[0036] In a preferred embodiment, in S25, the optimization target is set to:

[0037]

[0038] Among them, l represents the loss function, l CE represents the cross entropy loss and λ is a hyperparameter.

[0039] The present invention also provides an electronic device, comprising:

[0040] At least one processor; and a memory communicatively connected to the at least one processor; wherein the memory stores instructions executable by the at least one processor, and the instructions are executed by the at least one processor to enable the at least one processor to perform any one of the above methods.

[0041] The present invention also provides a computer-readable storage medium storing computer instructions, wherein the computer instructions are used to enable the computer to execute any one of the above methods.

[0042] The beneficial effects of the present invention include:

[0043] (1) Setting up a monitoring model, analyzing the log-mass ratio to understand how different divergence functions affect the probability mass distribution of the student model, especially the balance between difficult classes (classes with larger errors) and classes with high confidence;

[0044] (2) The proposed distillation framework can smoothly interpolate between FKLD and RKLD, thereby precisely regulating the attention effects of the two modes of difficulty concentration and confidence concentration, ensuring that the student model can focus on categories with larger errors during training while maintaining the learning of the global knowledge of the teacher model, so that the student model can perform well in a variety of tasks. BRIEF DESCRIPTION OF THE DRAWINGS

[0045] Figure 1 2. A flow chart of a method for optimizing probability mass allocation in knowledge distillation based on α-β-divergence according to a preferred embodiment of the present invention;

[0046] Figure 2 Schematic diagram of the knowledge distillation process in the probability mass allocation optimization method in knowledge distillation based on α-β-divergence according to a preferred embodiment of the present invention. DETAILED DESCRIPTION

[0047] The present invention will be described in further detail below with reference to the accompanying drawings and examples, through which the features and advantages of the present invention will become more clearly understood.

[0048] The word "exemplary" is used exclusively herein to mean "serving as an example, example, or illustration." Any embodiment described herein as "exemplary" is not necessarily to be construed as preferred or advantageous over other embodiments. Although various aspects of the embodiments are shown in the drawings, the drawings are not necessarily drawn to scale unless otherwise indicated.

[0049] According to the present invention, a probability mass allocation optimization method in knowledge distillation based on α-β-divergence is provided, such as Figure 1 As shown, the following steps are included:

[0050] Use the trained teacher model to predict or generate samples in the dataset to obtain sample soft labels;

[0051] Set up a distillation framework and train the student model based on the sample soft labels;

[0052] The trained student model is used to predict or generate data to obtain prediction or generation results.

[0053] In the present invention, there is no limitation on the specific model structures of the teacher model and the student model, and those skilled in the art may adopt any known model according to actual needs.

[0054] In the present invention, there is no limitation on the specific data type in the dataset, and those skilled in the art can freely set it according to actual needs, such as using a text dataset, an image dataset, an audio dataset, etc.

[0055] Traditional distillation methods generally use one of the divergence functions FKLD or RKLD. However, these divergence functions will cause the training of the student model to focus on the wrong categories or ignore the diverse knowledge of the teacher model.

[0056] To solve the above problems and ensure distribution matching during knowledge distillation, in the present invention, the distillation framework is set up to train the student model based on the sample soft labels, such as Figure 2 As shown, it includes the following sub-steps:

[0057] S21. Set a monitoring quantity model to characterize the probability mass change at each step of gradient update;

[0058] S22, lowering the upper bound of the monitoring quantity model by using the divergence function FKLD; lowering the upper bound of the monitoring quantity model by using the divergence function RKLD;

[0059] S23, unifying the two monitoring quantity models after lowering the upper bound to obtain an adjustable unified upper bound;

[0060] S24, analyze the adjustable unified upper bound and set the difference model;

[0061] S25. Based on the difference model, set the optimization goal and obtain the optimal parameters of the student model through the optimization method, thereby completing the training of the student model.

[0062] Since the probability mass of student distribution must be continuously adjusted during the knowledge distillation process, how to find a suitable probability matching solution is a difficulty in the present invention.

[0063] In S21, the monitoring amount is expressed as:

[0064]

[0065] in, For monitoring, represents the knowledge distillation framework, is the probability mass of category y obtained by knowledge distillation at step t+1, q t (y) represents the original probability mass of category y at step t of the student model.

[0066] In a preferred embodiment, qt (y) is obtained by approximating the category probability through the softmax function, expressed as:

[0067]

[0068] in, Represents the logit value of category y at step t.

[0069] In a preferred embodiment, It can be expressed as:

[0070]

[0071] in, represents the logit value of category k at step t, η represents the learning rate, l represents the training loss, Represents the gradient.

[0072] Furthermore, the above formula can be simplified as:

[0073]

[0074] Taking the logarithm of both sides:

[0075]

[0076] in, is a normalization factor independent of y.

[0077] From the above formula, we can see that when the training loss l is 0, disappears to 0, which shows that the monitoring amount is proportional to

[0078] In addition, the study found that the overall gradient and logit gradient There is a relationship between:

[0079]

[0080] Where J represents the Jacobian matrix of the gradient of the model parameters, whose dimension is C×M, where C is the dimension of the logarithm, M is the dimension of the model parameters, and the superscript T represents the transpose.

[0081] Typically, the matrix J is close to full rank C, in which case if That is to say

[0082]

[0083] Since J is a full-rank matrix, this expression can only be used in When it approaches 0, the whole approach is 0, that is, for all categories y

[0084] From the above analysis, it can be seen that in order to achieve a local minimum of training loss during training, the distillation algorithm automatically reduces The size of the module, and size.

[0085] In S22, based on the above analysis, it is found that the upper bound of the monitoring quantity model can be reduced by the divergence functions FKLD and RKLD.

[0086] Specifically, the upper bound of the monitoring model is reduced by the divergence function FKLD as follows:

[0087]

[0088] The upper bound of the monitoring model is reduced by the divergence function RKLD as follows:

[0089]

[0090] Among them, (a), (b), (a1), (b1), (a2), (b2) represent different terms in the model, p(y) represents the probability mass of category y in the teacher model, and q t (y) represents the original probability mass of category y at step t of the student model, represents the constant term in the divergence function FKLD that is independent of the category y, Represents the constant term in the divergence function RKLD that is independent of the category y.

[0091] According to the present invention, (b), (b1), and (b2) have |s(p(k))-s(q t The general form of (k))| can measure the matching loss between the distribution of student model and teacher model. These terms control the impact of difficulty concentration, and a sharper term with a larger rate of change is manifested as focusing on the most difficult categories, thereby achieving good matching performance.

[0092] (a), (a1), and (a2) are related to the confidence of the student model. These terms control the influence of the confidence set, and a sharper weighting function is expressed as only caring about the matching performance on the labels that the student believes to be the true ones.

[0093] FKLD has weak difficulty and confidence concentration effects, forcing students to evenly match losses across all categories, making it difficult to focus on the target class and unable to establish a reliable decision boundary. In contrast, RKLD has excessive difficulty and confidence concentration effects, causing students to focus almost entirely on the target class, weakening the matching of non-target classes. This causes students to focus solely on the target class and ignore the overall matching effect. In S23 of the present invention, FKLD and RKLD are jointly considered, and the two monitoring quantity models with lowered upper bounds are unified to achieve a balance between the above two effects.

[0094] In S23, the adjustable unified upper bound is expressed as:

[0095]

[0096] in, is an adjustable unified upper bound, and α and β are adjustable hyperparameters.

[0097] According to the adjustable uniform upper bound of the present invention, in (a) and (a1), the power form q t (k) β To achieve the confidence concentration effect, a larger β will enhance the confidence concentration effect, so that the matching is mainly focused on the category that the model is most confident in; at the same time, items (b) and (b1) are The difficulty concentration effect is controlled in the form of α. A smaller α will amplify the difficulty concentration effect and optimize the matching more aggressively by imposing a greater penalty on errors in difficult classes. In this sense, by adjusting α and β, we can flexibly balance the impact of these two effects and avoid extreme situations.

[0098] Based on the above analysis, in the present invention, the difference model in S24 is set as:

[0099]

[0100] in, represents the difference model, C represents the total number of categories, p(k) represents the probability mass of category k of the teacher model, and q(k) represents the probability mass of category k of the student model.

[0101] Through the analysis of the adjustable unified upper bound, it can be seen that the difference model can achieve a balance between the two effects. By flexibly adjusting α and β, it can effectively weigh the difficulty concentration and confidence concentration effects.

[0102] In S25, the optimization target is set to:

[0103]

[0104] Among them, l represents the loss function, l CErepresents the cross entropy loss and λ is a hyperparameter.

[0105] In the present invention, the optimization method is not limited, and those skilled in the art can adopt any known optimization method according to actual needs.

[0106] In this paper, through the above steps, we can deeply understand how different divergence functions (such as FKLD and RKLD) affect the probability mass distribution of the student model, especially the balance between more difficult classes and more confident classes. This method provides a theoretical basis for optimizing distribution matching in the knowledge distillation process, effectively solving the problem that existing methods focus on a few classes or ignore the diversity of teacher models.

[0107] Various embodiments of the methods described above in the present invention may be implemented in digital electronic circuit systems, integrated circuit systems, field programmable gate arrays (FPGAs), application specific integrated circuits (ASICs), application specific standard products (ASSPs), system-on-chip systems (SOCs), programmable logic devices (CPLDs), computer hardware, firmware, software, and / or combinations thereof. These various embodiments may include being implemented in one or more computer programs that are executable and / or interpreted on a programmable system comprising at least one programmable processor, which may be a special-purpose or general-purpose programmable processor that can receive data and instructions from a storage system, at least one input device, and at least one output device, and transmit data and instructions to the storage system, the at least one input device, and the at least one output device.

[0108] It should be understood that the various forms of the processes shown above can be used to reorder, add, or delete steps. For example, the steps described in this disclosure can be performed in parallel, sequentially, or in a different order, as long as the desired results of the technical solutions disclosed in this disclosure can be achieved. This is not a limitation herein.

[0109] Example

[0110] Example 1

[0111] A knowledge distillation experiment was conducted. In the experiment, the teacher model used GPT-2XL (1.5B), and the student models were set to GPT-2 (0.1B), GPT-2Medium (0.3B), and GPT-2Large (0.8B). The experiments were conducted on five datasets: Dolly Eval, Self-Instruct, Vicuna Eval, SuperNatural, and Unnatural.

[0112] The knowledge distillation process includes the following steps:

[0113] S21. Set a monitoring quantity model to characterize the probability mass change at each step of gradient update;

[0114] S22, lowering the upper bound of the monitoring quantity model by using the divergence function FKLD; lowering the upper bound of the monitoring quantity model by using the divergence function RKLD;

[0115] S23, unifying the two monitoring quantity models after lowering the upper bound to obtain an adjustable unified upper bound;

[0116] S24, analyze the adjustable unified upper bound and set the difference model;

[0117] S25. Based on the difference model, set the optimization goal and obtain the optimal parameters of the student model through the optimization method, thereby completing the training of the student model.

[0118] In S21, the monitoring amount is expressed as:

[0119]

[0120] q t (y) is obtained by approximating the category probability through the softmax function, expressed as:

[0121]

[0122] Expressed as:

[0123]

[0124] In S22, the upper bound of the monitoring model is reduced by the divergence function FKLD as follows:

[0125]

[0126] The upper bound of the monitoring model is reduced by the divergence function RKLD as follows:

[0127]

[0128] In S23, the adjustable unified upper bound is expressed as:

[0129]

[0130] The difference model described in S24 is set as:

[0131]

[0132] In S25, the optimization target is set to:

[0133]

[0134] Comparative Example 1

[0135] The same experiment as Example 1 was performed, except that the distillation framework used the supervised fine-tuning method (SFT) and KD, SeqKD, MiniLLM, GKD, and DISTILLM as distillation models, respectively.

[0136] KD refers to the literature Hinton, G. Distilling the knowledge in a neural network.arXiv2015;

[0137] For SeqKD, see Wen, Y., Li, Z., Du, W., and Mou, Lf-divergence minimization for sequence-level knowledge distillation. arXiv, 2023;

[0138] MiniLLM, see Gu, Y., Dong, L., Wei, F., and Huang, M. Minillm: Knowledgedistillation of large language models. In ICLR, 2024a;

[0139] For GKD, see the literature Agarwal, R., Vieillard, N., Zhou, Y., Stanczyk, P., Garea, SR, Geist, M., and Bachem, O. On-policy distillation of language models: Learning from self-generated mistakes. In ICLR, 2024;

[0140] DISTILLM See the literature Ko, J., Kim, S., Chen, T., and Yun, S.-Y. Distillm: Towards streamlined distillation for large language models. arXiv, 2024.

[0141] Comparing the results of Example 1 with those of Comparative Example 1, ROUGE-L is used as the evaluation index, which reflects the similarity between the generated text and the answer text. The experimental results are shown in Table 1.

[0142] Table 1

[0143]

[0144] As can be seen from Table 1, the method in Example 1 outperforms both standard KD and SFT on various datasets when refining GPT-2XL (1.5B) to a smaller-scale GPT-2 (0.1B to 0.8B). Furthermore, compared to state-of-the-art methods (SeqKD, MiniLLM, GKD, and DISTILLM), the method in Example 1 achieves competitive or superior results, particularly on the SuperNatural and Unnatural datasets, achieving the best performance.

[0145] The present invention has been described above with reference to preferred embodiments, but these embodiments are merely exemplary and serve only as illustrations. On this basis, various replacements and improvements can be made to the present invention, all of which fall within the scope of protection of the present invention.

Claims

1. A probability mass allocation optimization method in knowledge distillation based on α-β-divergence, characterized by: The following steps are involved: Use the trained teacher model to predict or generate samples in the dataset to obtain sample soft labels; Set up a distillation framework and train the student model based on the sample soft labels; The trained student model is used to predict or generate data to obtain prediction or generation results.

2. The method for optimizing probability mass allocation in knowledge distillation based on α-β-divergence according to claim 1, characterized in that: Setting up the distillation framework to train the student model based on sample soft labels includes the following sub-steps: S21. Set a monitoring quantity model to characterize the probability mass change at each step of gradient update; S22, lowering the upper bound of the monitoring quantity model by using the divergence function FKLD; lowering the upper bound of the monitoring quantity model by using the divergence function RKLD; S23, unifying the two monitoring quantity models after lowering the upper bound to obtain an adjustable unified upper bound; S24, analyze the adjustable unified upper bound and set the difference model; S25. Based on the difference model, set the optimization goal and obtain the optimal parameters of the student model through the optimization method, thereby completing the training of the student model.

3. The method for optimizing probability mass allocation in knowledge distillation based on α-β-divergence according to claim 2, characterized in that: In S21, the monitoring amount is expressed as: in, For monitoring, represents the knowledge distillation framework, is the probability mass of category y obtained by knowledge distillation at step t+1, q t (y) represents the original probability mass of category y at step t of the student model.

4. The method for optimizing probability mass allocation in knowledge distillation based on α-β-divergence according to claim 3, characterized in that: q t (y) is obtained by approximating the category probability through the softmax function, expressed as: in, Represents the logit value of category y at step t.

5. The method for optimizing probability mass allocation in knowledge distillation based on α-β-divergence according to claim 2, characterized in that: In S23, the adjustable unified upper bound is expressed as: in, is an adjustable unified upper bound, and α and β are adjustable hyperparameters.

6. The method for optimizing probability mass allocation in knowledge distillation based on α-β-divergence according to claim 2, characterized in that: In S22, the upper bound of the monitoring model is reduced by the divergence function FKLD as follows: The upper bound of the monitoring model is reduced by the divergence function RKLD as follows: Among them, (a), (b), (a1), (b1), (a2), (b2) represent different terms in the model, p(y) represents the probability mass of category y in the teacher model, and q t (y) represents the original probability mass of category y at step t of the student model, represents the constant term in the divergence function FKLD that is independent of the category y, Represents the constant term in the divergence function RKLD that is independent of the category y.

7. The method for optimizing probability mass allocation in knowledge distillation based on α-β-divergence according to claim 2, characterized in that: The difference model described in S24 is set as: in, represents the difference model, C represents the total number of categories, p(k) represents the probability mass of category k of the teacher model, and q(k) represents the probability mass of category k of the student model.

8. The method for optimizing probability mass allocation in knowledge distillation based on α-β-divergence according to claim 2, characterized in that: In S25, the optimization target is set to: Among them, l represents the loss function, l CE represents the cross entropy loss and λ is a hyperparameter.

9. An electronic device comprising: at least one processor; and a memory communicatively connected to the at least one processor; wherein the memory stores instructions executable by the at least one processor, and the instructions are executed by the at least one processor to enable the at least one processor to execute the method according to any one of claims 1 to 8.

10. A computer-readable storage medium storing computer instructions, wherein: The computer instructions are used to cause the computer to execute the method according to any one of claims 1-8.

Images