Token fusion based inference optimization method for large-scale language models

Abstract

The application discloses a large-scale language model inference optimization method based on token fusion, wherein a text sequence after being preprocessed by a word segmenter enters an embedding layer for coding to generate a word vector and a position vector, the word vector and the position vector are added to obtain a hidden state as an input matrix of a Transform module; the large-scale language model is subjected to layer-by-layer calculation of an MHA module, token grouping, token fusion, calculation of an MLP module and token restoration operation; the above steps are repeated until the last layer of the large-scale language model; and the hidden state generated by the last layer enters a tail linear layer for calculation to output the probability of each word in a word table. The application uses a pruning method to compress the LLM, and mainly focuses on how to reduce the inference delay of the large language model. Reducing the inference delay requires using a model compression technology to compress the large language model, reducing the calculation amount of model inference to reduce the inference delay, accelerating the inference process, and at the same time, the precision loss can be guaranteed to be as small as possible.

Description

TECHNICAL FIELD

[0001] The present application belongs to the field of inference optimization of large language models, and particularly relates to a large language model (LLM) inference optimization method based on token fusion. BACKGROUND

[0002] The current large language model (LLM) performs outstandingly on various complex language tasks, but its parameter quantity is usually as high as ten billion or even one hundred billion, thereby bringing about expensive computing and memory costs, while the existing GPU has relatively low memory storage and memory IO, resulting in high inference delay of such models, so the LLM is very expensive for applications with high real-time requirements or delay sensitivity.

[0003] In the existing pruning methods for LLM, most of them are directed to pruning the weight parameters of the model to reduce the parameter quantity and computing quantity of the model. It is found that there is universal similarity in tokens in the LLM, that is, the similarity of many tokens is very high in the input and output of each layer of the transformer decoder, so the present application adopts a method based on token fusion to compress the model, and the compressed matrix after the token fusion is involved in the operation, thereby reducing the computing quantity and delay of the model inference.

[0004] LLM based on Transformer structure has shown excellent performance on many complex language processing tasks in recent years, especially GPT-style generative models, which have shown outstanding performance and excellent context learning ability. However, due to its huge size and computational cost, it is difficult to deploy such models on resource-constrained devices, and large language models are very expensive during inference, making it still a big challenge to use such models in delay-sensitive applications. Generative large language models are mainly based on decoder-only models based on the Transformer structure, which are stacked by multiple decoder blocks, and the parameter quantity is usually up to the order of ten billion. Current inference optimization for LLM mainly focuses on throughput and latency. Throughput focuses on the amount of data processed per unit time, and latency focuses on user experience, i.e. the time required for the model to generate answer text after the user inputs the question. Current research on LLM compression methods mainly focuses on pruning (sparsification) and quantization. Quantization maps the parameters and activation values of the model to lower-precision numerical values, reducing the precision of numerical representation in the model, thereby reducing the memory overhead and computational cost of the model. Although quantization can greatly reduce the storage requirements of the model, it requires specific hardware to support low-precision numerical calculations, such as 8-bit integer matrix calculations. In fact, pruning as an early emerging model compression method has been successful in many neural network models, but its application in LLM is very limited, and natural language understanding tasks are still mainly applied to small language models such as Bert models. Pruning removes redundant or unimportant components in the model, from individual weights (unstructured pruning) to higher granularity structures such as rows or columns of weight matrices (structured pruning), thereby reducing the size and computational complexity, but the excellent pruning method usually requires additional retraining or fine-tuning to restore accuracy, but it has not been widely applied to LLM. On the one hand, pruning or sparsification methods often require retraining or iterative fine-tuning to restore the accuracy of the model, but this is not suitable for models with ten billion parameters; on the other hand, the unique context learning ability of LLM is difficult to maintain after pruning.

[0005] The current mainstream pruning applied to the LLM is weight-based pruning. This method first defines a standard for evaluating the importance or redundancy of weights, then prunes the weights that meet the standard one by one according to the standard, and then fine-tunes the pruned model to recover the accuracy loss caused by pruning. This scheme belongs to static pruning, that is, the model is fixed after the entire pruning scheme is implemented, and the same parameters are used for inference. There are some static pruning schemes for LLMs at present. Although it is a one-shot method that does not need to be retrained, some methods still require a small amount of weight update process, and have not achieved true inference delay acceleration. Another weight-based pruning is dynamic pruning, which prunes different contents according to different inputs of the model, that is, the input determines which weight elements of the model will be pruned.

[0006] The existing LLM compression methods including model sparsification and quantization require additional hardware support, such as quantization requiring fixed-point computing GPU, and sparsification requiring hardware co-design to achieve acceleration. In terms of pruning, weight-based static pruning mainly focuses on weight sparsification of linear layers, and also requires support for semi-structured GPU hardware to achieve efficient inference deployment. SUMMARY

[0007] In view of the above prior art, the present application provides a large-scale language model inference optimization method based on token fusion. The present application uses pruning method to compress LLM, mainly focusing on how to reduce the inference delay of large language model. Reducing inference delay requires using model compression technology to compress large language model, reducing the calculation amount of model inference to reduce inference delay, accelerating the inference process, and making the application of LLM more popular. The compression of the present application for LLM is a one-shot pruning method without retraining or fine-tuning, which can reduce the inference delay while ensuring the smallest accuracy loss.

[0008] To solve the above technical problems, the present application provides a large-scale language model inference optimization method based on token fusion. The large-scale language model includes an embedding layer, a plurality of Transformer modules connected in sequence, and a linear layer connected at the tail. The Transformer module includes a multi-head attention mechanism module, a multi-layer perception module, a token grouping module, a token fusion module, a token difference predictor, and a token restoration module. The multi-head attention mechanism module includes three linear layers. The method comprises:

[0009] Step 1: The text sequence preprocessed by the tokenizer enters the embedding layer for encoding to generate word vectors and position vectors. The word vectors and position vectors are added to obtain a hidden state as an input matrix of the Transformer module.

[0010] Step 2, for the large-scale language model, the calculation of the multi-head attention mechanism module, token grouping, token fusion, the calculation of the multi-layer perception module and the token restoration operation are performed layer by layer; comprising the following steps:

[0011] Step 2-1) first perform the calculation of the multi-head attention mechanism module, and obtain a preliminary token similarity matrix by using the attention score matrix generated by the multi-head attention mechanism module;

[0012] Step 2-2) using the token grouping module, based on the similarity obtained from the token similarity matrix and the set similarity threshold, the tokens are grouped, and the grouping index of each token is calculated by using the multi-head attention mechanism module;

[0013] Step 2-3) using the token fusion module, based on the obtained grouping index of each token, the input of the multi-layer perception module is token fused, and the tokens in the same group are fused into a token vector, thereby obtaining a compressed input matrix;

[0014] Step 2-4) the obtained compressed input matrix is input into the multi-layer perception module for calculation, and the calculation result is input into the next layer of the multi-head attention mechanism module for calculation, and the query matrix Q, the key matrix K and the value matrix V are generated. At the same time, the token difference predictor generates the difference of each token in the group before and after token fusion;

[0015] Step 2-5) the token restoration module restores the query matrix Q, the key matrix K and the value matrix V output by the three linear layers of the multi-head attention mechanism module to the original number of tokens by copying the group token vectors, and then adds the query matrix Q, the key matrix K and the value matrix V to the output result of the token difference predictor, to obtain the restored query matrix Q, the key matrix K and the value matrix V, which are used for the calculation after the three linear layers of the multi-head attention mechanism module;

[0016] Step 3, repeat the above step 2 until the last layer of the large-scale language model;

[0017] Step 4, the hidden state generated by the last layer of the Transformer module enters the tail linear layer calculation, and outputs the probability of each word in the word table.

[0018] Further, the large-scale language model inference optimization method provided by the application, wherein:

[0019] The specific content of step 2-1) includes:

[0020] The multi-head attention mechanism module adopts a query-key-value mode, and first maps the input of the multi-head attention mechanism module to three different spaces through three linear layers to obtain a query matrix Q, a key matrix K and a value matrix V, the vector groups at corresponding positions in the key matrix K and the value matrix V represent input information; for each query vector q in the query matrix Q, vector matrix multiplication is performed with the key matrix K to generate different attention score matrices, which represent the attention scores of each token to other tokens; the different tokens are weighted and averaged using the attention scores, and finally linear mapping is performed to generate the output of the multi-head attention mechanism module.

[0021] In the multi-head attention mechanism module, the query matrix Q and the key matrix K are scaled and dot-multiplied to generate an attention score matrix, and the values in each row of the attention score matrix represent the attention weight of the token to other tokens, indicating the similarity between the token and other tokens; the attention score matrix is a preliminary token similarity matrix.

[0022] The specific content of step 2-2) includes:

[0023] The token grouping module adopts a threshold decay grouping method, and a similarity threshold interval and a decay step are defined in advance, and the tokens are grouped in turn according to the step from the maximum threshold value. The similarity matrix is a matrix obtained by normalizing the scaled dot product result of the matrix Q and the matrix K, and one row or one column represents a token. The similarity matrix is a diagonal matrix, and the elements in the matrix represent the similarity between the token represented by the row and the token represented by the column.

[0024] According to the threshold interval and the step, a plurality of thresholds are obtained, and grouping is performed according to each threshold, and each grouping is performed on the basis of the previous step;

[0025] In each grouping, matrix multiplication is used for matrix calculation. First, set the elements greater than the threshold in the similarity matrix to 1 and the other elements to 0 to obtain a preliminary grouping mask matrix. Then, extract the rows with more than one 1 element in the grouping mask as a group, concatenate the mask of each group to generate the grouping mask matrix of this time, and set the similarity of the tokens in the determined group to -1.

[0026] Concatenate the grouping mask matrix generated each time in the row direction, and form a one-dimensional vector by taking all tokens that do not meet the threshold condition as separate groups. Concatenate the one-dimensional vector with the grouping mask matrix to form the final grouping mask matrix, which is used as the token fusion in the next step.

[0027] The grouping mask matrix is in one-hot encoding format, each row represents a group, each column represents a token, and the element at position (i, j) indicates that the ith group contains the jth token; the fusion method is to weight and average the tokens in the group based on similarity, and each group is finally fused into a token; the grouping mask matrix represents the grouping index of each token.

[0028] The specific content of step 2-3) includes:

[0029] First, the grouping mask matrix and the input matrix of the multi-layer perception module are multiplied;

[0030] Then, divide each row by the number of tokens in the group to obtain the weighted average matrix;

[0031] Finally, the input size of the multi-layer perception module changes from (sequence length, hidden dimension) to (group number, hidden dimension), and the fused matrix is input into the multi-layer perception module for calculation.

[0032] The specific content of step 2-4) includes:

[0033] The calculation result of the multi-layer perception module is input into the next layer of the multi-head attention mechanism module to perform three linear layer calculations to generate query matrix Q, key matrix K and value matrix V;

[0034] At the same time, the compressed matrix after token fusion is restored to the shape and vector order of the original multi-layer perception module input matrix by copying the vector of each group, and then the restored input matrix is subtracted from the original input matrix to obtain the difference before and after fusion. The difference value is input into the token difference predictor to predict the difference value;

[0035] The structure of the token difference predictor is a two-layer multi-layer perception structure, which contains a Sigmoid activation function in the middle. The output and input of the token difference predictor are both difference matrices, and the shape is the same as the input matrix of the multi-layer perception module in the Transformer. The output of the token difference predictor is the difference between the output of the compressed matrix after token fusion after multi-layer perception operation and the output of the three linear layers of the multi-head attention mechanism module and the output of the uncompressed matrix after multi-layer perception module and the three linear layers of the multi-head attention mechanism module. Thus, the token is restored through the difference value.

[0036] The specific content of step 2-5) includes:

[0037] The token recovery module replicates the query matrix Q, the key matrix K and the value matrix V output by the three linear layers of the multi-head attention mechanism module into the intra-group token vector, recovers to the original number of tokens, and then adds the output results of the token difference predictor to the query matrix Q, the key matrix K and the value matrix V to obtain the recovered query matrix Q, the key matrix K and the value matrix V, and then continue to calculate after the three linear layers in the multi-head attention mechanism module, that is, the Q matrix and the K matrix scale dot product to generate an attention score matrix, and the attention score matrix is multiplied with the V matrix to enter the linear layer to generate the output of the multi-head attention mechanism module.

[0038] Compared with the prior art, the beneficial effects of the present application are:

[0039] For a generative large language model, grouping and compressing tokens based on token similarity can reduce the computational complexity of LLM inference, especially in the FFN module where parameters account for a large proportion of the model. Using the compressed input for forward propagation can greatly reduce the computational overhead, thereby reducing the inference delay of the LLM, thereby realizing efficient deployment of the model. Only conventional hardware can realize the method of the present application. BRIEF DESCRIPTION OF DRAWINGS

[0040] Figure 1 is the overall step diagram of the token fusion-based large-scale language model inference optimization method.

[0041] Figure 2 is the model structure diagram of the optimized large-scale language model, which shows the relationship and forward inference process of each module, wherein the dashed arrow actually points to the next layer, and the structure of each layer is the same, which is not separately shown in the diagram.

[0042] Figure 3 is a detailed structure diagram of the token fusion-based large-scale language model inference optimization method in the model, that is, the detailed structure of part of the modules in Figure 2 is displayed.

[0043] Figure 4 is a schematic diagram of the key-value pair attention mechanism.

[0044] Figure 5 is a token grouping mask matrix diagram, which shows the generation process of the grouping mask matrix.

[0045] Figure 6 is a token fusion algorithm diagram, which shows the token fusion calculation process of the multi-layer perceptron module input matrix using the grouping mask matrix.

[0046] Figure 7is an input data schematic diagram of a token difference predictor, showing the generation process of input data.

[0047] Figure 8 is a structural schematic diagram of a token difference predictor.

[0048] Figure 9 is a token restoration schematic diagram, showing the process of token restoration using the difference matrix generated by the token difference predictor. DETAILED DESCRIPTION

[0049] The design concept of the token fusion-based large-scale language model inference optimization method proposed in the application mainly focuses on compressing the number of tokens to prune the model to reduce inference delay. In terms of the content of the pruning, the pruning method of the LLM mainly includes two categories: weight-based pruning and token-based pruning.

[0050] There are many fully connected layers in the Transformer structure, and the size of the input data value can indirectly determine the size of the output value, which relates to the importance of the related weights, so the pruning weights are dynamically determined according to the input. This method does not need to be trained and fine-tuned again, and reduces the memory bandwidth and inference delay by reducing part of the calculation process.

[0051] In addition to the weight-based pruning, there is a token-based pruning method on the language model. The input of the language model is usually one or more sentences, each sentence is composed of multiple words, i.e., tokens, and then converted to word vectors of the same dimension through the Embedding layer. Token-based pruning is to prune part of the input word vectors of each layer to reduce the amount of calculation. Token-based pruning is currently mainly applied to vision transformer models and small language models (such as Bert) in the vision field, and has not been applied to LLM. In addition, similar to token pruning, there is a token fusion scheme based on token, i.e., fusing part of the tokens into one token. This method is also only applied to vision models.

[0052] In the LLM, there are many tokens that are very similar. In the application, a pruning method applied to the LLM is designed by combining the token fusion and dynamic pruning scheme. The method does not need to be trained and fine-tuned again, and reduces the amount of calculation in the model inference by fusing similar tokens to reduce the input data amount of each layer, thereby accelerating the inference process.

[0053] The compression technology in the application is aimed at a generative pre-training large language model based on a Transformer structure, such as an OPT or Llama model. In such a model, the basic structure is a Transformer block, and token fusion and restoration operations are performed layer by layer based on the Transformer block.

[0054] The application provides a large-scale language model inference optimization method based on token fusion. The large-scale language model comprises an embedding layer, a plurality of Transformer modules, and a linear layer. The Transformer module comprises a multi-head attention mechanism module, a multi-layer perceptron module, a token grouping module, a token fusion module, a token difference predictor, and a token restoration module. The multi-head attention mechanism module comprises three linear layers. The input and output of the multi-head attention mechanism module and the multi-layer perceptron module are matrices composed of embedding vectors of a plurality of tokens, and each row represents a token in a sentence sequence. The pruning scheme of the application first obtains a token similarity matrix according to the intermediate output of the multi-head attention mechanism module, then groups the tokens based on the similarity and a threshold set in advance to obtain a grouping index of each token. This process is actually performed in parallel with the attention calculation process of the multi-head attention mechanism module. After the multi-head attention mechanism module is calculated, the input of the multi-layer perceptron module is fused based on the grouping index obtained in the foregoing process, that is, tokens in the same group are fused into a token vector to obtain a compressed input matrix, and then the input matrix is input into the multi-layer perceptron module and the first half of the multi-head attention mechanism module in the next layer. At the same time, the token difference predictor generates the difference between the tokens in the group before and after token fusion, so that the token restoration operation is performed according to the generated difference before the second half of the attention module is calculated, the original number of tokens is restored, and then the same operation in the next layer is continued.

[0055] As shown in Figure 1 , the method comprises the following steps:

[0056] Step 1: The text sequence preprocessed by a tokenizer enters an embedding layer for encoding to generate a word vector and a position vector, and the word vector and the position vector are added to obtain a hidden state as an input matrix of the Transformer module.

[0057] Step 2, for the large-scale language model, the calculation of the multi-head attention mechanism module, token grouping, token fusion, the calculation of the multi-layer perception module and the token restoration operation are performed layer by layer; comprising the following steps:

[0058] Step 2-1) first, the calculation of the multi-head attention mechanism module is performed, and a preliminary token similarity matrix is obtained by using the attention score generated by the multi-head attention mechanism module; the specific content includes:

[0059] The multi-head attention mechanism module adopts a query-key-value mode, first maps the input of the multi-head attention mechanism module to three different spaces through three linear layers to obtain a query matrix Q, a key matrix K and a value matrix V, and the vector groups at the corresponding positions in the key matrix K and the value matrix V represent the input information in the form of key-value pairs; for each query vector q in the query matrix Q, vector matrix multiplication is performed with the key matrix K to generate different attention score matrices, which represent the attention score of each token to other tokens; the different tokens are weighted and averaged using the attention score, and finally linear mapping is performed to generate the output of the multi-head attention mechanism module.

[0060] In the multi-head attention mechanism module, the query matrix Q and the key matrix K are scaled and dot multiplied to generate an attention score, and the values in each row of the attention score matrix represent the attention weight of the token to other tokens, indicating the similarity between the token and other tokens; the attention score matrix is a preliminary token similarity matrix.

[0061] Step 2-2) using the token grouping module, based on the similarity obtained from the token similarity matrix and the set similarity threshold, the tokens are grouped, and the grouping index of each token is obtained by using the multi-head attention mechanism module; the specific content includes:

[0062] The token grouping module adopts a threshold decay grouping method, and a similarity threshold interval and a decay step are defined in advance, and the tokens are grouped by sequentially decreasing the threshold from the maximum value by the step, the similarity matrix is a matrix obtained by normalizing the scaled dot product result of the matrix Q and the matrix K, and one row or one column represents a token, and the similarity matrix is a diagonal matrix, wherein the elements represent the similarity between the token represented by the row and the token represented by the column.

[0063] According to the threshold interval and the step, a plurality of thresholds are obtained, and each threshold is grouped once, and each grouping is performed on the basis of the previous step.

[0064] In each grouping, matrix calculation such as matrix multiplication is used. First, 1 is set in the similarity matrix greater than the threshold value, and other elements are set to 0, to obtain a preliminary grouping mask matrix; then, the number of 1 elements in the grouping mask is extracted as a group, and each group of masks is spliced to generate the grouping mask matrix this time, and the similarity of the token that has been determined is set to -1.

[0065] The grouping mask matrix generated each time is spliced in the row direction, all tokens that do not meet the threshold condition are respectively taken as a separate group, a one-dimensional vector is formed, and the one-dimensional vector is spliced with the grouping mask matrix to form the final grouping mask matrix, which is used as the token fusion in the next step.

[0066] The grouping mask matrix is in one-hot encoding format, each row represents a group, each column represents a token, and the element at position (i, j) indicates that the i-th group contains the j-th token; the fusion method is to weight and average the tokens in the group based on similarity, and each group is finally fused into a token; the grouping mask matrix represents the grouping index of each token.

[0067] Taking the Llama2 model as an example, in order to model the long-distance dependence of the model and process the variable-length input sequence, the self-attention mechanism is used in the Transformer model, that is, to “dynamically” generate different connection weights. The self-attention mechanism adopts the “query-key-value” mode, as shown in Figure 3 First, the input of MHA is mapped to three different spaces to obtain the query matrix Q, the key matrix K and the value matrix V, and the vector group at the corresponding position in K and V represents the key-value pair representing the input information. As shown in Figure 4 For each query vector q in Q, it respectively performs vector matrix multiplication with K to generate different attention distributions, which represent the attention scores of each token to other tokens, that is, the Attention Score matrix after matmul in the figure, and then the matrix is softmaxed and multiplied with the V matrix representing the value information, that is, the different tokens are weighted and averaged using the attention score, and finally linear mapping is performed to generate the output of the MHA module.

[0068] In the MHA module, the attention score generated by the scaled dot product of Q and K actually also represents a similarity relationship, and the values in each row represent the attention weight of the token to other tokens, which indirectly reflects the similarity between the token and other tokens, so the grouping of the token is based on the matrix scaled dot product result of Q and K. The grouping process of the token is parallel to the multiplication of the attention score matrix in MHA and V and the subsequent calculation process.

[0069] The grouping process takes the form of threshold decay. The present application defines a similarity threshold interval and a decay step, and sequentially decreases the threshold from the maximum value by the step size to group the tokens. The similarity matrix is a matrix after normalizing the scaled dot product result of Q K, and one row or one column represents a token, and the specific element represents the similarity between the token represented by the row and the token represented by the column, so the similarity matrix is a diagonal matrix. The present application will obtain multiple thresholds according to the threshold interval and the step size, and the present application will group according to each threshold, and each grouping is based on the previous step.

[0070] Specific to each step of grouping, in order to reduce the additional time cost, the present application uses all matrix operations such as matrix multiplication to complete each grouping, avoiding the large time consumption of for loop. As shown in the following formula, first, set 1 in the similarity matrix greater than the threshold, and set 0 in other elements, to obtain a preliminary grouping mask matrix; in order to ensure that tokens will not appear repeatedly in different groups, the grouping mask is processed to ensure that there is only one 1 in each column; then extract the rows with more than one 1 elements in the grouping mask as a group, and concatenate the mask of each group to generate the grouping mask matrix of this step, and also set the similarity of the tokens that have been determined to be grouped to -1, so as to ensure that these tokens are not considered in the next grouping. Figure 5

[0071] Finally, concatenate the grouping mask matrix generated each time in the row direction. In addition, all tokens that do not meet the threshold condition are separately regarded as a single group, that is, each group contains only one token, and the grouping mask matrix is concatenated to form the final grouping mask matrix, which is used as the token fusion of the next step.

[0072] Step 2-3) using the token fusion module to perform token fusion based on the obtained grouping index of each token to fuse the tokens in the same group into a token vector, thereby obtaining a compressed input matrix; the specific content includes:

[0073] First, multiply the grouping mask matrix and the input matrix of the multi-layer perception module;

[0074] Then, divide each row by the number of tokens in the group to obtain a weighted average matrix;

[0075] Finally, the input size of the multi-layer perception module changes from (sequence length, hidden dimension) to (group number, hidden dimension), and the fused matrix is input into the multi-layer perception module for calculation.

[0076] ​Specifically, after obtaining the grouping information of the token, the token needs to be fused. Because there is interaction between tokens in the calculation process of the attention mechanism in the MHA module, fusion cannot be performed in this process, and because the calculation of the MLP is not interactive for tokens, the input to the MLP module is selected for token fusion. As shown in Figure 3 , the output of the MHA is the input of the MLP, and in order to reduce the amount of calculation, the fusion algorithm is applied before the residual connection branch.

[0077] The input of the MLP is a matrix composed of the Embedding vectors of all tokens. The purpose of the present application is to compress the input matrix by reducing the number of tokens and reducing the amount of subsequent calculation. The grouping mask matrix generated by the grouping algorithm is in one-hot encoding format, with each row representing a group and each column representing a token. The position (i, j) of the element with a value of 1 indicates that the i-th group contains the j-th token. The fusion method of the present application is to perform a weighted average of the tokens in each group based on similarity, and each group is finally fused into one token, as shown in Figure 6 . The specific calculation method is to directly multiply the grouping mask matrix and the input matrix of the MLP, and then divide each row by the number of tokens in the group to obtain the weighted average matrix. Finally, the input size of the MLP changes from (sequence length, hidden dimension) to (group number, hidden dimension), and the matrix dimension is greatly reduced. The fused matrix then enters the MLP module for subsequent calculation. The MLP module of the general Transformer structure contains two Linear layers, Figure 3 , as shown in the structure of Llama2, which includes three Linear layers. Therefore, the compressed matrix after token fusion will reduce the operation of many linear layers, thereby speeding up the inference.

[0078] Step 2-4) input the obtained compressed input matrix into the multi-layer perception module for calculation, and input the calculation result into the next layer of the multi-head attention mechanism module for calculation of three linear layers to generate a query matrix Q, a key matrix K and a value matrix V; at the same time, the token difference predictor generates the difference between each token in the group before and after token fusion; the specific content includes:

[0079] The calculation process of the token difference predictor is as follows: the compressed matrix after token fusion is restored to the shape and vector order of the original multi-layer perception module input matrix by copying the vector of each group, and then the restored input matrix is subtracted from the original input matrix to obtain the difference before and after fusion. The difference is input into the token difference predictor for prediction, as shown in Figure 7 .

[0080] The structure of the token difference predictor is a two-layer multi-layer perception structure, which contains a Sigmoid activation function in the middle. The output is the difference matrix, which is the input of the token difference predictor, and the shape is the same as the input matrix of the multi-layer perception module in the Transformer, as shown in Figure 8 .

[0081] The output of the token difference predictor is the difference between the output of the compressed matrix after token fusion after multi-layer perception operation and the output of the three linear layers of the multi-head attention mechanism module, and the output of the uncompressed matrix after multi-layer perception module and the three linear layers of the multi-head attention mechanism module. Thus, the token is restored through the difference.

[0082] In addition, the token difference predictor needs to be trained in advance by collecting data. The entire model process is run according to the positive operation steps and the steps containing the token fusion operation. The token difference predictor input of each layer is collected and the label is made according to the output, so as to train the token difference predictor of each layer respectively, so as to learn the token restoration ability.

[0083] Step 2-5) The token restoration module replicates the token vectors within the three linear layers of the multi-head attention mechanism module to restore the original number of tokens, and then adds the query matrix Q, the key matrix K and the value matrix V to the output result of the token difference predictor to obtain the restored query matrix Q, the key matrix K and the value matrix V, which are used for calculation after the three linear layers of the multi-head attention mechanism module.

[0084] After token fusion, the shape needs to be restored after several module calculations. As mentioned earlier, the Q KV matrix after the three linear layers in the MHA module is restored to the original number of tokens. Specifically, the token difference predictor generates the difference between the Q KV matrix before and after fusion, and adds the Q KV matrix generated by the three linear layers to the corresponding output result of the token difference predictor, as shown in Figure 9 , and then continues the attention calculation as the restored Q KV matrix, that is, the Q and K perform scaling dot product to generate attention scores, and the attention scores are multiplied by the V matrix to realize the aggregation of token information, as shown in Figure 3 .

[0085] Step 3, repeat step 2 above until the last layer of the large-scale language model;

[0086] Step 4: The hidden state generated by the last layer of the Transformer module enters the tail linear layer to calculate and output the probability of each word in the word table.

[0087] Although the present application is described above with reference to the accompanying drawings, the present application is not limited to the above-described specific embodiments, and the above-described specific embodiments are merely illustrative, not restrictive. Those skilled in the art can make many improvements and changes to the present application without departing from the spirit of the present application, and these are all within the protection scope of the present application.

Claims

1. A token fusion-based large-scale language model inference optimization method, characterized in that, The large-scale language model comprises sequentially connected embedding layers, multi-layer Transformer modules, and a tail linear layer, and the Transformer modules comprise multi-head attention mechanism modules, multi-layer perception modules, token grouping modules, token fusion modules, token difference predictors, and token restoration modules; The multi-head attention mechanism module comprises three linear layers; and the method comprises: Step 1, a text sequence preprocessed by a tokenizer enters an embedding layer for coding to generate a word vector and a position vector, and the word vector and the position vector are added to obtain a hidden state as an input matrix of the Transformer module; Step 2, for the large-scale language model, the calculation of the multi-head attention mechanism module, token grouping, token fusion, the calculation of the multi-layer perception module, and the token restoration operation are performed layer by layer; and the step 2 comprises the following steps: Step 2-1), the calculation of the multi-head attention mechanism module is first performed, and a similarity matrix of tokens is obtained by using an attention score matrix generated by the multi-head attention mechanism module; Step 2-2), the tokens are grouped based on the similarity obtained from the similarity matrix of the tokens and a set threshold of the similarity by using the token grouping module, and a grouping index of each token is obtained by using the multi-head attention mechanism module; Step 2-3), the token fusion module is used to perform token fusion on the input of the multi-layer perception module based on the obtained grouping index of each token, and tokens in the same group are fused into a token vector, so as to obtain a compressed input matrix; Step 2-4), the obtained compressed input matrix is input into the multi-layer perception module for calculation, and the calculation result is input into the multi-head attention mechanism module of the next layer for calculation of the three linear layers to generate a query matrix Q, a key matrix K, and a value matrix V; at the same time, the token difference predictor generates a difference between each token in the group before and after the token fusion; Step 2-5), the token restoration module restores the number of original tokens by duplicating the query matrix Q, the key matrix K, and the value matrix V output by the three linear layers of the multi-head attention mechanism module into group token vectors, and then adds the query matrix Q, the key matrix K, and the value matrix V and the output result of the token difference predictor to obtain the restored query matrix Q, the key matrix K, and the value matrix V, which are used for calculation after the three linear layers of the multi-head attention mechanism module; Step 3, the step 2 is repeated until the last layer of the large-scale language model; Step 4, the hidden state generated by the last layer of the Transformer module enters the tail linear layer for calculation, and the probabilities of each word in the word table are output.

2. The method of claim 1, wherein, The specific content of step 2-1) comprises: The multi-head attention mechanism module adopts a query-key-value mode, and first maps the input of the multi-head attention mechanism module to three different spaces through three linear layers to obtain a query matrix Q, a key matrix K and a value matrix V, wherein the vector groups at corresponding positions in the key matrix K and the value matrix V represent input information; for each query vector q in the query matrix Q, vector matrix multiplication is performed with the key matrix K to generate different attention score matrices, which represent the attention scores of each token to other tokens; different tokens are weighted and averaged using the attention scores, and finally linear mapping is performed to generate the output of the multi-head attention mechanism module; In the multi-head attention mechanism module, the query matrix Q and the key matrix K are subjected to scaled dot product to generate an attention score matrix, wherein the values in each row of the attention score matrix represent the attention weight of the token to other tokens, indicating the similarity between the token and other tokens; the attention score matrix is a preliminary token similarity matrix.

3. The method of claim 2, wherein, The specific content of step 2-2) includes: The token grouping module adopts a threshold decay grouping method, and a similarity threshold interval and a decay step are defined in advance, and tokens are grouped in descending order according to the step from the maximum threshold value; the similarity matrix is a matrix obtained by normalizing the scaled dot product result of the matrix Q and the matrix K, and one row or one column represents a token; the similarity matrix is a diagonal matrix, and the elements therein represent the similarity between the token represented by the row and the token represented by the column; A plurality of thresholds are obtained according to the threshold interval and the step, and grouping is performed according to each threshold, and each grouping is performed on the basis of the previous step; In each grouping, matrix multiplication is used, first, the elements greater than the threshold in the similarity matrix are set to 1, and the other elements are set to 0, to obtain a preliminary grouping mask matrix; then, the rows with more than one 1 element in the grouping mask are extracted as a group, the grouping mask of each group is spliced to generate the grouping mask matrix of this time, and the similarity of the tokens of the groups that have been determined is set to -1; The grouping mask matrices generated each time are spliced in the row direction, all tokens that do not meet the threshold condition are respectively taken as separate groups, a one-dimensional vector is formed, and the one-dimensional vector and the grouping mask matrix are spliced to form the final grouping mask matrix, which is used as the token fusion of the next step; The grouping mask matrix is in one-hot encoding format, each row represents a group, each column represents a token, and the element at position (i, j) represents that the i-th group contains the j-th token; the fusion method is to weight and average the tokens in the group based on similarity, and each group is finally fused into a token; The grouping mask matrix represents the grouping index of each token.

4. The method of claim 3, wherein, The specific content of step 2-3) includes: First, the grouping mask matrix and the input matrix of the multi-layer perception module are subjected to matrix multiplication; Then, each row is divided by the number of tokens in the group to obtain the weighted average matrix; Finally, the input size of the multi-layer perception module changes from (sequence length, hidden dimension) to (group number, hidden dimension), and the fused matrix is input into the multi-layer perception module for calculation.

5. The method of claim 4, wherein, The specific content of step 2-4) includes: The compressed matrix after token fusion is restored to the shape and vector order of the original multi-layer perception module input matrix by copying the vectors of each group, and then the restored shape input matrix is subtracted from the original input matrix to obtain the difference value before and after fusion, which is used as the input of the token difference value predictor; The calculation of the token difference value predictor: the structure of the token difference value predictor is a two-layer multi-layer perception structure, which contains a Sigmoid activation function in the middle, and the output is a difference value matrix with the same shape as the input matrix of the multi-layer perception module in the Transformer. The output of the token difference value predictor is the difference between the output of the compressed matrix after token fusion after multi-layer perception operation and the output of the three linear layers of the multi-head attention mechanism module, and the output of the uncompressed matrix after multi-layer perception module and the three linear layers of the multi-head attention mechanism module, so as to restore the token through the difference value.

Images