Knowledge distillation optimization method based on sparse mixed experts and low-rank adaptation

Abstract

The present invention relates to the technical field of natural language processing, and discloses a knowledge distillation optimization method based on sparse mixed experts and low-rank adaptation. A student model is designed according to a teacher model, and the student model is trained. The open question-answering task is completed by the trained student model, including: both the teacher model and the student model adopt a large language model, and the parameter scale of the large language model adopted by the teacher model is larger than the parameter scale of the large language model adopted by the student model; the student model is improved as follows: for each multi-layer perceptron in the backbone network of the large language model adopted by the student model, multiple low-rank experts and a route are added to form a low-rank mixed expert module; SRKL divergence is used as the knowledge distillation loss function from the teacher model to the student model, and the total loss is constructed in combination with the supervised fine-tuning loss function. Through the combination of sparse mixed experts and low-rank adaptation, the student model can achieve performance improvement on multiple tasks.

Description

Technical Field

[0001] The present invention relates to the field of natural language processing, and in particular to a knowledge distillation optimization method based on sparse mixed experts and low-rank adaptation. Background Art

[0002] With the widespread application of large language models in the field of natural language processing, the demand for computing resources and storage space is growing. These models usually require a large amount of data for training and have high hardware requirements when deployed. To solve this problem, knowledge distillation technology came into being, whose goal is to compress the knowledge of large teacher models into smaller and more efficient student models.

[0003] The general process of language model knowledge distillation mainly includes the following steps: First, prepare a large, well-performing teacher model (dense model), which is usually trained on a large amount of data to obtain rich knowledge and strong predictive ability. Then, design a small, simple and computationally efficient student model (dense model). Next, the feature distribution output by the teacher model is passed to the student model through some form of training process. This process usually involves minimizing the difference between the output of the student model and the output of the teacher model, for example, by designing a specific loss function. Finally, the distilled student model can reduce the demand for computing resources while maintaining high performance, making it more suitable for deployment in resource-constrained environments.

[0004] Although knowledge distillation has made some progress in reducing model size and computational cost, existing methods mainly focus on optimizing loss functions and training strategies, while rarely involve improving the structure of student models. This leads to a large performance gap between student models and teacher models, limiting the further development of knowledge distillation technology. Existing technologies rarely consider optimization from the perspective of improving the structure of student models, mainly due to the following two problems:

[0005] 1) Difficulty in cross-model architecture knowledge distillation. Different model series have different tokenizers, which results in different vector dimensions for different models. Therefore, it is impossible to measure the gap between the teacher model and the student model during distillation, and distillation cannot be performed.

[0006] 2) Difficulty in distilling after improving the original architecture. This solution can avoid the problem of different word segmenters, but it is difficult to determine how to improve the model structure, and after changing the model structure, the model needs to be re-pre-trained, which requires hundreds of times the computing power of knowledge distillation with the same architecture.

[0007] The above difficulties have resulted in few relevant techniques to optimize the student model architecture. Summary of the invention

[0008] In view of the shortcomings of existing knowledge distillation technology in improving the student model structure, the present invention aims to propose a knowledge distillation method that combines the sparse mixture of experts (MoE) architecture and the low-rank adaptation (LoRA) architecture. This method can use the MoE architecture for knowledge distillation without re-pretraining the student model, thereby improving the performance of the student model and reducing computational overhead.

[0009] In order to solve the above technical problems, the present invention adopts the following technical solutions:

[0010] A knowledge distillation optimization method based on sparse mixed experts and low-rank adaptation, designing a student model according to a teacher model, training the student model, inputting a question into the trained student model, and outputting an answer corresponding to the question, the process of designing a student model according to the teacher model and training the student model specifically includes:

[0011] Both the teacher model and the student model use a large language model, and the parameter scale of the large language model used by the teacher model is larger than that of the student model;

[0012] The student model is improved as follows: for each multi-layer perceptron in the backbone network of the large language model used by the student model, N low-rank experts E and a router are added to form a low-rank mixed expert module; the weight matrix of the multi-layer perceptron is Replaced by low-rank matrix A and low-rank matrix B to form the low-rank expert; weight matrix The matrix shape is ,in is the weight matrix The input dimension is Weight Matrix The output dimensions of the low-rank matrix A and the low-rank matrix B are , ,in is the rank of the low-rank expert, which is a controllable parameter; the routing Is a trainable weight matrix; multiply the input matrix x of the routing and multilayer perceptron to obtain the intermediate matrix , normalize the intermediate matrix to get the probability matrix ; Probability matrix Each row element in the probability matrix forms a row vector, and each row vector of the probability matrix represents the weight of each low-rank expert on the input matrix x; the output of the two low-rank experts with the highest weight is weighted and summed, and then added to the output of the multi-layer perceptron to obtain the output matrix y; the input matrix x is the vector representation of the training sample after the word segmenter and the embedding layer, where the training sample is a text pair consisting of a question and an answer, and a text pair containing S words is recorded as ; The output matrix y is the probability distribution of the next word predicted by the student model;

[0013] When distilling knowledge from the teacher model to the student model, a hybrid loss function is used : ,in represents the SRKL divergence; is the supervised fine-tuning loss function, ; Indicates that the large language model has the first k-1 words Under the condition of The probability of represents the kth word of a text pair, Represents the first k-1 words of the text pair, that is, , Represents the total number of words in the text pair; To control the parameters of the supervised fine-tuning loss function in the hybrid loss function.

[0014] Furthermore, the routing The matrix shape is ,in is the dimension of the hidden state of the student model.

[0015] Furthermore, the routing and multilayer perceptron input matrices x are multiplied to obtain an intermediate matrix , normalize the intermediate matrix to get the probability matrix ; Probability matrix Each row element in the probability matrix forms a row vector. Each row vector of the probability matrix represents the weight of each low-rank expert on the input matrix x. After weighted summing the outputs of the two low-rank experts with the highest weight, the output matrix y is obtained by adding the output of the multilayer perceptron. Specifically, it includes:

[0016] Multiply the routing and input matrices x to obtain the intermediate matrix h(x):

[0017] ;

[0018] The superscript T indicates transpose; the i-th row vector of the intermediate matrix h(x) Represents the i-th low-rank expert The weight influence on the input matrix x;

[0019] Use the softmax function to normalize the intermediate matrix h(x) to obtain the probability matrix :

[0020] ;

[0021] is the probability matrix The i-th row vector of represents the weight proportion of the i-th low-rank expert to the input matrix x;

[0022] Select the two low-rank experts with the largest weight , , and then the low-rank expert , The results obtained by multiplying the input matrix x are used as low-rank experts , Output , Finally, the outputs of the two largest low-rank experts are weighted and summed together with the output of the multilayer perceptron. Sum and get the output matrix y:

[0023] ;

[0024] , They are the two low-rank experts with the largest weights. , The index of .

[0025] Furthermore, the SRKL divergence for:

[0026] ;

[0027] represents KL divergence; represents the probability distribution of the teacher model output, represents the probability distribution of the student model output, represents the parameters of the student model, is a hyperparameter that controls the weight of the probability distribution output by the teacher model in the SRKL divergence calculation.

[0028] The present invention optimizes the knowledge distillation of the language model from the perspective of improving the student model structure, and has the following advantages:

[0029] Performance Improvement: By combining sparse mixture of experts and low-rank adaptation architecture, the student model is able to achieve significant performance improvements on multiple tasks.

[0030] Computational efficiency: The introduction of the low-rank adaptive architecture significantly reduces the computational cost of model fine-tuning, while the dynamic selection mechanism of the sparse hybrid expert architecture also reduces unnecessary calculations.

[0031] Generalization ability: Combining SRKL divergence and supervised fine-tuning loss function, the student model not only performs well on open-ended question answering tasks, but also has good generalization ability.

[0032] Structural flexibility: The method of the present invention does not require retraining of the student model, making the adjustment of the model structure more flexible and efficient. BRIEF DESCRIPTION OF THE DRAWINGS

[0033] Figure 1 A comparison chart of the distillation architecture;

[0034] Figure 2 This is a comparison chart between the dense model and the sparse mixed expert model;

[0035] Figure 3 is a schematic diagram of a low-rank hybrid expert module of the present invention;

[0036] Figure 4a and Figure 4b Graph showing the experimental results of an embodiment of the present invention on an open question answering task; Figure 4a Represents the experimental results for Exact Match. Figure 4b Represents the experimental results for Rouge-L similarity (Rouge-L). DETAILED DESCRIPTION

[0037] A preferred embodiment of the present invention is described in detail below with reference to the accompanying drawings.

[0038] The present invention is applied in open question-answering tasks, and the model training samples in the development and answering tasks are text pairs consisting of questions and answers; during reasoning, the questions are input into the model and the corresponding answers are output.

[0039] The technical solution adopted by the present invention includes the following parts.

[0040] 1. Low-rank mixed expert architecture:

[0041] If the hybrid expert structure is directly applied to the improvement of the model structure, the model needs to be re-pretrained. Therefore, the present invention innovatively proposes a new low-rank hybrid expert module suitable for knowledge distillation. The principle of the low-rank hybrid expert module is shown in Figure 3 .

[0042] 1.1 Low-rank experts:

[0043] In the transformer model (an autoregressive model proposed by Google in 2017), an expert usually refers to a multi-layer perceptron (MLP) module. The present invention does not change the structure of the backbone network of the large language model, but adds some low-rank experts consisting of low-rank matrices next to the backbone network. These low-rank experts have the function of improving training effects and performance.

[0044] For each weight matrix in the original multi-layer perceptron , the present invention uses two low-rank matrices, low-rank matrix A and low-rank matrix B, to replace it. The matrix shape is ,in is the weight matrix The input dimension is is the weight matrix The output dimension of , then the shapes of the low-rank matrix A and the low-rank matrix B are , Where r is the rank of the matrix in the low-rank expert, which is a controllable parameter. It is usually relatively small and is set to 16 in this embodiment. Therefore, the number of parameters of the low-rank expert in the present invention is very small, which can reduce the computational burden of the large language model.

[0045] 1.2 Low-rank hybrid expert module routing:

[0046] For a multi-layer perceptron in a backbone network, the present invention sets 2 to 8 low-rank experts to match it. The low-rank experts are separated from the backbone network, so the router only controls the input and output of the low-rank experts. The router is a trainable weight matrix, and its matrix shape is ,in is the dimension of the hidden layer state of the student model, and N is the number of low-rank experts.

[0047] The routing and input matrix x are multiplied to obtain an intermediate matrix h(x), where the vector of the intermediate matrix Represents a low-rank expert The weight influence on the input matrix x. Then use the softmax function (a normalization function) to process h(x) to obtain the probability matrix p(x), Represents a low-rank expert Normalized weight ratio for input matrix x:

[0048] ;

[0049] .

[0050] Then the two low-rank experts with the largest weight will be selected , The expert set Γ is constructed, and then the result of multiplying the two low-rank experts by the input matrix x is calculated. The remaining experts are not calculated. Such a sparse setting can save a lot of calculations. Finally, the outputs of the two largest low-rank experts are weighted and summed to obtain the output matrix y.

[0051] .

[0052] 2. Knowledge distillation method:

[0053] The optimization method of the present invention can be naturally combined with the distillation method. The present invention adopts the Skew Reverse KL (SRKL) divergence proposed in 2024 (see the paper DistiLLM: Towards StreamlinedDistillation for Large Language Models) as the knowledge distillation loss function , this function can better guide the student model to learn the output probability distribution of the teacher model and prevent the probability averaging phenomenon of the output of the student model, thereby improving the distillation effect.

[0054] ;

[0055] It is a hyperparameter that controls the weight of the probability distribution output by the teacher model in the SRKL divergence calculation. In this embodiment, it is set to 0.1;

[0056] In addition, the present invention also combines the loss function of supervised fine-tuning (SFT), which further improves the performance of the student model on a variety of downstream tasks. In the following formula, S represents the text pair The total number of words, It means that the large language model has the first k-1 words. Under the condition of , the output of the kth word (token) is The probability of represents the kth word of a text pair, It represents the first k-1 words of the text pair , which is a simplified way of writing.

[0057] ;

[0058] .

[0059] in, In order to control the parameters of the supervised fine-tuning loss function in the hybrid loss function, it is set to 0.1 in this embodiment.

[0060] 3. Optimization of low-rank expert configuration:

[0061] The present invention experiments on the configuration of low-rank experts, including the setting of the rank r of low-rank experts, the number of low-rank experts in each layer, and the hierarchical distribution of low-rank experts, and selects the optimal configuration. For higher layers, more low-rank experts will be assigned. In a preferred embodiment, r is set to 16, which reduces the number of model parameters and the amount of calculation while ensuring performance. For the hierarchical distribution, the present invention chooses to place low-rank experts in the last half of the layers of the model. Preferably, the low-rank experts are placed in the last six layers of the large language model, that is, 2, 2, 4, 4, 8, and 8 low-rank experts are placed from the 7th layer to the 12th layer.

[0062] Depend on Figure 1 It can be seen that the knowledge distillation optimization method of the present invention is different from the previous distillation method with the same architecture. The previous method generally optimizes the loss function of distillation Or training methods, but without improving the model structure, are mostly distillations of a dense teacher model to a dense student model of the same architecture. The present invention proposes to use a technology combining sparse mixed experts and low-rank adaptation to enhance the capabilities of the student model.

[0063] like Figure 2 As shown in the figure, in the dense model, each input must pass through all the parameters of the model, so the amount of calculation is large. The sparse mixed expert model controls which experts and parameters the input passes through by routing, so that only some parameters are calculated. Under the same amount of calculation, the performance of the sparse mixed expert model is usually better.

[0064] Figure 3 The principle of the low-rank mixed expert module of the present invention is demonstrated. For an input matrix x, it will pass through the multi-layer perceptron of the backbone network. This process does not make any changes, because once the structure of the backbone network is changed, it needs to be pre-trained again. The present invention only sets up multiple low-rank experts and a router next to it to form a low-rank mixed expert module. For an input matrix x, the route is multiplied by x to calculate which low-rank experts process the input matrix, and then the output of the selected low-rank experts and the output of the backbone network are added to obtain the final output matrix y.

[0065] Figure 4a and Figure 4bThe experimental results of using the method of the present invention to distill the GPT2-large model to the GPT2-small model on the open question answering task (Natural Question) are shown, where MoE-KD represents the method of the present invention. SFT, MoE-SFT, and KD refer to the supervised fine-tuning method, the mixed expert supervised fine-tuning method, and the traditional knowledge distillation method, respectively. It can be seen that for the open question answering task, the method of the present invention performs better than other methods in both exact match (Exact Match) and Rouge-L similarity (Rouge-L).

[0066] Example

[0067] The method of the present invention comprises the following steps:

[0068] Step 1, supervised fine-tuning model:

[0069] First, select a dense large language model to be optimized, which is generally a larger teacher model and a student model with the same architecture, such as the GPT2-large model and the GPT2-small model. Perform supervised fine-tuning on the student model on the training data of the open-ended question-answering task. The fine-tuning loss function is the negative log-likelihood function, which can maximize the probability of occurrence of the training data. The experience of supervised fine-tuning training is shown in Table 1.

[0070] Table 1

[0071]

[0072] Step 2, insert the low-rank mixed expert module:

[0073] Improve the structure of the student model that has been trained through supervised fine-tuning. Note that the teacher model does not need to improve its structure. The present invention traverses the structure of the student model, and for each layer of the multilayer perceptron, adds 2 to 8 low-rank experts next to it, and sets a learnable route to form a low-rank mixed expert module. After that, all low-rank expert parameters are initialized by Kaiming (the initialization method proposed by scholar He Kaiming). The original parameters of the student model will not be changed in this process. The configuration of the inserted low-rank mixed expert module is shown in Table 2.

[0074] Table 2

[0075]

[0076] Step 3, knowledge distillation:

[0077] After inserting the low-rank mixed expert module into the student model, the present invention performs knowledge distillation on the student model on the training dataset of the open question answering task. The total loss of knowledge distillation uses a combination of SRKL divergence and supervised fine-tuning loss function. The parameters of the distillation process are shown in Table 3.

[0078] Table 3

[0079]

[0080] It is obvious to those skilled in the art that the present invention is not limited to the details of the exemplary embodiments described above, and that the present invention can be implemented in other specific forms without departing from the spirit or essential features of the present invention. Therefore, from any point of view, the embodiments should be regarded as exemplary and non-limiting, and the scope of the present invention is defined by the appended claims rather than the above description, and it is intended that all changes falling within the meaning and scope of the equivalent elements of the claims are included in the present invention, and any reference numerals in the claims should not be regarded as limiting the claims involved.

[0081] In addition, it should be understood that although this specification is described in terms of implementation methods, not every implementation method contains only one independent technical solution. This narrative method of the specification is only for the sake of clarity. Those skilled in the art should regard the specification as a whole. The technical solutions in each embodiment may also be appropriately combined to form other implementation methods that can be understood by those skilled in the art.

Claims

1. A knowledge distillation optimization method based on sparse mixed experts and low-rank adaptation, characterized in that: A student model is designed according to the teacher model, and the student model is trained. A question is input into the trained student model, and an answer corresponding to the question is output. The process of designing the student model according to the teacher model and training the student model specifically includes: Both the teacher model and the student model use a large language model, and the parameter scale of the large language model used by the teacher model is larger than that of the student model; The student model is improved as follows: for each multi-layer perceptron in the backbone network of the large language model used by the student model, N low-rank experts E and a router are added to form a low-rank mixed expert module; the weight matrix of the multi-layer perceptron is Replaced by low-rank matrix A and low-rank matrix B to form the low-rank expert; weight matrix The matrix shape is ,in is the weight matrix The input dimension is Weight Matrix The output dimensions of the low-rank matrix A and the low-rank matrix B are , ,in is the rank of the low-rank expert, which is a controllable parameter; the routing Is a trainable weight matrix; multiply the input matrix x of the routing and multilayer perceptron to obtain the intermediate matrix , normalize the intermediate matrix to get the probability matrix ; Probability matrix Each row element in the probability matrix forms a row vector. Each row vector of the probability matrix represents the weight of each low-rank expert on the input matrix x. The outputs of the two low-rank experts with the highest weight are weighted and summed, and then added to the output of the multi-layer perceptron to obtain the output matrix y. The input matrix x is the vector representation of the training sample after the word segmenter and the embedding layer. The training sample is a text pair consisting of a question and an answer. A text pair containing S words is recorded as ; The output matrix y is the probability distribution of the next word predicted by the student model; When distilling knowledge from the teacher model to the student model, a hybrid loss function is used : ,in represents the SRKL divergence; is the supervised fine-tuning loss function, ; Indicates that the large language model has the first k-1 words Under the condition of The probability of Represents the kth word of a text pair; Represents the first k-1 words of the text pair, that is, ; Represents the total number of words in the text pair; To control the parameters of the supervised fine-tuning loss function in the hybrid loss function.

2. The knowledge distillation optimization method based on sparse mixed experts and low-rank adaptation according to claim 1, characterized in that: The route The matrix shape is ,in is the dimension of the hidden state of the student model.

3. The knowledge distillation optimization method based on sparse mixed experts and low-rank adaptation according to claim 1, characterized in that: The routing and multilayer perceptron input matrix x are multiplied to obtain the intermediate matrix , normalize the intermediate matrix to get the probability matrix ; Probability matrix Each row element in the probability matrix forms a row vector. Each row vector of the probability matrix represents the weight of each low-rank expert on the input matrix x. After weighted summing the outputs of the two low-rank experts with the highest weight, the output matrix y is obtained by adding the output of the multilayer perceptron. Specifically, it includes: Multiply the routing and input matrices x to obtain the intermediate matrix h(x): ； The superscript T indicates transpose; the i-th row vector of the intermediate matrix h(x) Represents the i-th low-rank expert The weight influence on the input matrix x; Use the softmax function to normalize the intermediate matrix h(x) to obtain the probability matrix : ； is the probability matrix The i-th row vector of represents the weight ratio of the i-th low-rank expert to the input matrix x; Select the two low-rank experts with the largest weight , , and then the low-rank expert , The results obtained by multiplying the input matrix x are used as low-rank experts , Output , Finally, the outputs of the two largest low-rank experts are weighted and summed together with the output of the multilayer perceptron. Sum and get the output matrix y: ； , They are the two low-rank experts with the largest weights. , The index of .

4. The knowledge distillation optimization method based on sparse mixed experts and low-rank adaptation according to claim 1, characterized in that: SRKL divergence for: ； represents KL divergence; represents the probability distribution of the teacher model output, represents the probability distribution of the student model output, represents the parameters of the student model, is a hyperparameter that controls the weight of the probability distribution output by the teacher model in the SRKL divergence calculation.

Images