Key-value compression method in self-attention mechanism, large language model and electronic device

Abstract

The present application discloses a key-value compression method, a large language model, and an electronic device in a self-attention mechanism, and relates to the field of computer technology. The compression method includes performing multiple residual decompositions on the key matrix and the value matrix, respectively, to obtain a key residual vector and a value residual vector after each decomposition; clustering and compressing the key residual vector and the value residual vector after each decomposition, and performing attention calculation on the query matrix, the compressed key residual vector, and the value residual vector; and accumulating all attention calculation results. The present application solves the problem that the Linear Transformer cannot use the standard Softmax Transformer parameters and has significant differences from the standard Softmax Transformer.

Description

Technical Field

[0001] The present application relates to the field of computer technology, and in particular to a key-value compression method in a self-attention mechanism, a large language model, and an electronic device. Background Art

[0002] Attention mechanisms have been the cornerstone of recent advances in natural language processing, computer vision, and related fields. However, traditional softmax-based self-attention suffers from quadratic complexity, significantly limiting memory capacity and computational efficiency, particularly for tasks involving long sequences. To address this bottleneck, various linearized attention mechanisms have been proposed. These aim to reduce the complexity of self-attention from quadratic to linear, enabling more efficient inference with fewer time and space resources. This improvement is particularly beneficial for long sequence processing tasks, where the computational overhead of traditional attention methods can be prohibitive.

[0003] Linear Attention is one of the pioneering methods in this field. It replaces the softmax operation with kernel functions such as the Exponential Linear Unit (ELU) activation and rearranges the computation order to achieve linear complexity. Linear Attention is compatible with recurrent neural network structures, bridging the gap between the Transformer model and recurrent neural networks (RNNs). Recently introduced models such as Gated Linear Attention (GLA), Gated Slot Attention (GSA), and Meta Linear Attention (MetaLA) have further improved the efficiency and accuracy of linear attention. For example, GLA employs a dynamic gating mechanism to more efficiently manage and update stored information, further unifying previous linear methods, including structured state-space models (SSMs), into this broader framework, becoming a special form of GLA. GSA also introduces an extended ABC storage mechanism that compresses the key and value matrices into fixed-size slots, making it closer to softmax attention than earlier linearized attention models. In addition to linear Attention and SSM models, many recursive models also use the Neural Turing Mechanism (NTM) method to write and read historical information, thereby improving the recursive model's ability to process historical information. Summary of the Invention

[0004] This application provides a key-value compression method, a large language model, and an electronic device in a self-attention mechanism to achieve a more flexible and robust sequence representation, which can adapt to various attention mechanisms and achieve high-level performance in long sequence tasks.

[0005] In a first aspect, the present application provides a key-value compression method in a self-attention mechanism, which is applied to a large language model. The compression method includes:

[0006] Perform multiple residual decompositions on the key matrix and the value matrix respectively to obtain the key residual vector and the value residual vector after each decomposition; wherein the key residual vector after each decomposition includes the key residual vector after the 0th decomposition, and the value residual vector after each decomposition includes the value residual vector after the 0th decomposition;

[0007] The key residual vector and value residual vector after each decomposition are clustered and compressed respectively, and attention calculation is performed on the query matrix, the compressed key residual vector and the value residual vector;

[0008] Accumulate all attention calculation results.

[0009] Furthermore, the specific formula for the residual decomposition of the bond matrix is:

[0010] K i =K i-1 -gather(C i-1 [0], index1);

[0011] index1=argmax(matmul(K i-1 , C i-1 [0].transpose));

[0012] Among them, K i Represents the key residual vector after the i-th residual decomposition, K i-1 Represents the key residual vector after the i-1th residual decomposition. When i=1, K0 represents the key matrix; gather represents the selection of C according to the index position index1. i-1 [0] and splice; C i-1 Represents the trainable matrix, C i-1 [0] represents the trainable matrix C i-1 argmax means returning the index of the maximum value of the matrix in the specified dimension, matmul means matrix multiplication, and transpose means matrix transposition;

[0013] The specific formula for the residual decomposition of the value matrix is:

[0014] V i =V i-1 -gather(C i-1 [1], index2);

[0015] index2=argmax(matmul(Vi-1 , C i-1 [1].transpose));

[0016] Among them, V i Represents the residual vector after the i-th residual decomposition, V i-1 Represents the residual vector after the i-1th residual decomposition. When i=1, V0 represents the value matrix; gather represents the selection of C according to the index position index2. i-1 [1] and splice; C i-1 [1] represents the trainable matrix C i-1 The first dimension of .

[0017] Furthermore, clustering and compression are performed on the key residual vector and value residual vector after each decomposition, including:

[0018] Calculating a correlation matrix between the key residual vector and the trainable matrix;

[0019] Performing a mapping transformation on the correlation matrix to obtain a weight matrix;

[0020] Clustering the key residual vector and the weight matrix to obtain a compressed key residual vector;

[0021] Clustering is performed on the value residual vector and the weight matrix to obtain a compressed value residual vector.

[0022] Furthermore, performing a mapping transformation on the correlation matrix includes:

[0023] A feedforward neural network or a sofimax function is used to perform mapping transformation on the correlation matrix.

[0024] In a second aspect, the present application provides a large language model, comprising an embedding layer, a position encoding layer, an output layer, and a plurality of Transformer Blocks; each Transformer Block comprises an attention layer and an MLP layer;

[0025] The embedding layer is used to convert each word in the text data into a feature vector of fixed length; the position encoding layer is used to map the position of each word in the text data to the feature space and superimpose it on the feature vector corresponding to the word to obtain an input embedding matrix;

[0026] The attention layer includes a first regularization module, a self-attention transformation module and a first superposition module; the first regularization module is used to regularize the input embedding matrix; the self-attention transformation module is used to perform QKV transformation on the regularized matrix to obtain a query matrix, a key matrix and a value matrix, and perform multiple residual decompositions on the key matrix and the value matrix respectively to obtain a key residual vector and a value residual vector after each decomposition; the key residual vector and the value residual vector after each decomposition are clustered and compressed respectively, and attention calculation is performed on the query matrix, the compressed key residual vector and the value residual vector; all attention calculation results are accumulated to obtain an attention output matrix; the first superposition module is used to superimpose the input embedding matrix and the attention output matrix to obtain an attention layer output matrix;

[0027] The MLP layer includes a second regularization module, an MLP transformation module and a second superposition module; the second regularization module is used to regularize the attention layer output matrix; the MLP transformation module is used to extract features from the regularized matrix to obtain a feature matrix; the second superposition module is used to superimpose the attention layer output matrix and the feature matrix.

[0028] In a third aspect, the present application also provides an electronic device comprising a memory, a processor, and a computer program / instructions stored on the memory, wherein the processor executes the computer program / instructions to implement the steps of the key-value compression method in the self-attention mechanism as described above in the present application.

[0029] In a fourth aspect, the present application also provides a computer-readable storage medium having a computer program / instruction stored thereon, which, when executed by a processor, implements the steps of the key-value compression method in the self-attention mechanism as described above in the present application.

[0030] In a fifth aspect, the present application also provides a computer program product, including a computer program / instruction, which, when executed by a processor, implements the steps of the key-value compression method in the self-attention mechanism as described above in the present application.

[0031] The beneficial effects of this application are:

[0032] This application performs multiple residual decompositions on the key matrix and the value matrix, and then clusters and compresses the key matrix, the value matrix, the key residual vector and the value residual vector after each decomposition, respectively, and performs attention calculation and accumulation to obtain approximate attention output, thereby improving the efficiency of the Transformer in processing long sequences; at the same time, it solves the problem that the Linear Transformer cannot use the standard Softmax Transformer parameters and is quite different from the standard Softmax Transformer.

[0033] This application implements limited space key-value storage for long sequences, reducing the training storage overhead and computational complexity from N 2 Reduced to N, the inference storage overhead is reduced from N to 1, and the inference calculation amount is reduced from N 2 Compared to LinearAttention, this method has the advantage of not requiring re-pretraining. Furthermore, the model weights fine-tuned using this method can be used for both the large language model (i.e., the approximate Transformer model) and the standard Transformer model inference, thus offering greater model application flexibility compared to LinearAttention. Furthermore, because this method still uses softmax for weight rearrangement, the model performance is closer to the standard Transformer model than various LinearAttention models. BRIEF DESCRIPTION OF THE DRAWINGS

[0034] In order to more clearly illustrate the technical solution of the present application, the following is a brief introduction to the drawings required for use in the description of the embodiments. Obviously, the drawings described below are only one embodiment of the present application. For ordinary technicians in this field, other drawings can be obtained based on these drawings without any creative work.

[0035] Figure 1 This is a flow chart of the key-value compression method in the self-attention mechanism in an embodiment of the present application;

[0036] Figure 2 This is a detailed diagram of the key value compression process in the embodiment of this application;

[0037] Figure 3 This is the loss curve in the embodiment of this application;

[0038] Figure 4 This is a diagram of the large language model architecture in the embodiment of this application;

[0039] Figure 5 This is a schematic diagram of the structure of an electronic device in an embodiment of the present application. DETAILED DESCRIPTION

[0040] The following is a clear and complete description of the technical solutions in this application in conjunction with the accompanying drawings in the embodiments of this application. Obviously, the embodiments described are only part of the embodiments of this application, not all of them. Based on the embodiments in this application, all other embodiments obtained by ordinary technicians in this field without making any creative efforts are within the scope of protection of this application.

[0041] The following specific embodiments are used to describe the technical solution of the present application in detail. The following specific embodiments can be combined with each other, and the same or similar concepts or processes may not be described in detail in some embodiments.

[0042] Example 1

[0043] Figure 1 The flowchart of the key value compression method in the self-attention mechanism provided by this application is shown. Figure 1 As shown, the compression method includes the following steps:

[0044] Step S1: Perform multiple residual decompositions on the key matrix and the value matrix respectively to obtain a key residual vector and a value residual vector after each decomposition; wherein the key residual vector after each decomposition includes the key residual vector after the 0th decomposition (i.e., the key matrix), and the value residual vector after each decomposition includes the value residual vector after the 0th decomposition (i.e., the value matrix);

[0045] Step S2: cluster and compress the key residual vector and value residual vector after each decomposition, and perform attention calculation on the query matrix, the compressed key residual vector and the value residual vector;

[0046] Step S3: Accumulate all attention calculation results.

[0047] The query, key, and value in the attention mechanism are represented by Q, K, and V, respectively. Therefore, the query matrix, key matrix, and value matrix are represented by Q0, K0, and V0, respectively. K0 and V0 also represent the key residual vector and value residual vector after the 0th residual decomposition. Performing QKV decomposition on the attention input matrix X can obtain the query matrix, key matrix, and value matrix. The specific QKV decomposition formula is:

[0048] Q0=XW q , K0=XW k , V0=XW v (1)

[0049] Among them, W q 、W k 、W v All are weight matrices. Note that the size of the input matrix X is N×d, where N represents the sequence length and d represents the embedding dimension size. The size of the query matrix, key matrix, and value matrix are all N×d, and the weight matrix W q 、W k 、W v The size of is d×d. The weight matrix W q 、W k 、W vThe value of is gradually adjusted during training as the loss function gradient descends. It should be noted that the embedding dimension d of the query matrix can be further divided into multiple heads, and the number of heads in the embedding dimension of the query matrix, key matrix, and value matrix may not be the same. For ease of explanation, but without affecting generality, it is assumed that the embedding dimension of the query matrix, key matrix, and value matrix is ​​the same, d.

[0050] The idea of ​​key-value compression is to train a fixed-length multi-head codebook token sequence, where each head is used to cluster the token vector residuals of the key matrix and the value matrix, and then compress the key residual vector and value residual vector after clustering, ultimately obtaining the compressed key residual vector and value residual vector. This example uses three-head KV compression as an example to illustrate the specific compression process of the key matrix and value matrix. Figure 2 As shown in the figure, the specific compression process is:

[0051] Step B1: Perform a residual decomposition on the key matrix and the value matrix respectively to obtain the key residual vector and the value residual vector after the decomposition. The specific decomposition formula is:

[0052] K1=K0-gather(C0[0], index1) (2)

[0053] index1=argmax(matmul(K0,C0[0].transpose)) (3)

[0054] V1=V0-gather(C0[1], index2) (4)

[0055] index2=argmax(matmul(V0,C0[1].transpose)) (5)

[0056] Among them, K1 represents the key residual vector after the first residual decomposition (i.e., the key residual vector after the first decomposition), K0 represents the key residual vector after the 0th residual decomposition (i.e., the original key matrix); gather represents the selection of C0[0] according to the index position index1 and splicing; C0 represents the trainable matrix and the size is C×d, that is, C0 is the Codebook tokens sequence of length C, C0[0] represents the 0th dimension of the trainable matrix C0, and C0[1] represents the 1st dimension of the trainable matrix C0; argmax represents the index of the maximum value of the matrix in the specified dimension, matmul represents matrix multiplication, and transpose represents matrix transposition; V1 represents the value residual vector after the first residual decomposition (i.e., the value residual vector after the first decomposition), V0 represents the value residual vector after the 0th residual decomposition (i.e., the original value matrix); gather represents the selection of C0[1] according to the index position index2 and splicing.

[0057] Step B2: Perform secondary residual decomposition on the key residual vector and value residual vector after the primary decomposition to obtain the key residual vector and value residual vector after the secondary decomposition. The specific decomposition formula is:

[0058] K2=K1-gather(C1[0], index1) (6)

[0059] index1=argmax(matmul(K1,C1[0].transpose)) (7)

[0060] V2=V1-gather(C1[1], index2) (8)

[0061] index2=argmax(matmul(V1,C1[1].transpose)) (9)

[0062] Among them, K2 represents the key residual vector after the second residual decomposition (i.e., the key residual vector after the secondary decomposition); C1 represents a trainable matrix and its size is C×d, that is, C1 is a codebook tokens sequence of length C, C1[0] represents the 0th dimension of the trainable matrix C1, and C1[1] represents the 1st dimension of the trainable matrix C1; V2 represents the value residual vector after the second residual decomposition (i.e., the value residual vector after the secondary decomposition).

[0063] Step B3: Cluster and compress the key matrix K0, the key residual vector K1 after the first decomposition, the key residual vector K2 after the second decomposition, the value matrix V0, the value residual vector V1 after the first decomposition, and the value residual vector V2 after the second decomposition respectively.

[0064] The key matrix K0 and the value matrix V0 correspond to one compression head, the key residual vector K1 after the first decomposition and the value residual vector V1 after the first decomposition correspond to one compression head, and the key residual vector K2 after the second decomposition and the value residual vector V2 after the second decomposition correspond to one compression head. The clustering compression process of each compression head is the same. This embodiment takes the clustering compression of the key matrix K0 and the value matrix V0 as an example. The specific clustering compression process is:

[0065] Step B3.1: Calculate the correlation matrix between the key matrix K0 and the trainable matrix C0. The specific calculation formula is:

[0066] A0=matmul(C0[0],K0) (10)

[0067] Among them, A0 represents the correlation matrix between the key matrix K0 and the trainable matrix C0, and its size is C×N. Figure 2A1 in represents the correlation matrix between the key residual vector K1 and the trainable matrix C1 after the first decomposition, and A2 represents the correlation matrix between the key residual vector K2 and the trainable matrix C2 after the second decomposition.

[0068] Step B3.2: Perform a mapping transformation on the correlation matrix A0 to obtain the weight matrix F0. The specific formula is:

[0069] F0=F(A0) (11)

[0070] Where F represents a mapping transformation function. In this embodiment, the mapping transformation function is a feedforward neural network or a softmax function. That is, after the correlation matrix A0 is transformed by a feedforward neural network or a softmax function, a weight matrix F0 is obtained. The nonlinear characteristics of the weight matrix F0 are more prominent, and its size remains C×N. Figure 2 Where F1 is the weight matrix corresponding to A1, and F2 is the weight matrix corresponding to A2.

[0071] Step B3.3: Cluster the key matrix K0 and the weight matrix F0 to obtain the compressed key residual vector K 0_z , the specific formula is:

[0072] K 0_z =matmul(F0, K0) (12)

[0073] The compressed key residual vector K 0_z It can be regarded as the cluster compression of the key matrix K0 with the Codebook tokens sequence as the cluster center, and the compressed key residual vector K 0_z The size is C×d. Figure 2 The key residual vector K in 1_z is the clustering compression result of the key residual vector K1 after one decomposition, the key residual vector K 2_z is the clustering compression result of the key residual vector K2 after secondary decomposition.

[0074] Step B3.4: Cluster the value matrix V0 and the weight matrix F0 to obtain the compressed value residual vector V 0_z , the specific formula is:

[0075] V 0_z =matmul(V0,K0) (13)

[0076] Compressed value residual vector V 0_z It can be regarded as the cluster compression of the value matrix V0 with the Codebook tokens sequence as the cluster center, and the compressed value residual vector V 0_z The size is C×d. Figure 2 The residual vector V 1_zis the clustering compression result of the residual vector V1 after a decomposition, the residual vector V 2_z is the clustering compression result of the residual vector V2 after secondary decomposition.

[0077] The clustering compression process of the three compression heads is exactly the same, so in the implementation, the compression head can be used as the batch dimension for parallelization, and the output of multiple heads can be merged into K _z ={K 0_z , K 1_z , K 2_z}、V _z ={V 0_z , V 1_z , V 2_z}.

[0078] Step B4: Perform attention calculation on the query matrix, the compressed key residual vector, and the value residual vector.

[0079] The clustering compression result of each compression head is calculated with the query matrix for attention. The specific calculation formula is:

[0080] P[i]=softmax(matmul(Q0,K _z [i].transpose)) (14)

[0081] Out[i]=matmul(P[i],V _z [i]) (15)

[0082] Where P[i] represents the intermediate matrix, i = 0, 1, 2, and the size of P[i] is N×C; K _z [i] represents K _z The i-th element in K i_z ; Out[i] represents the clustering compression result of the i-th compression head and the attention calculation result of the query matrix. The size of Out[i] is N×d.

[0083] Step B5: Accumulate all attention calculation results. The specific accumulation formula is:

[0084] Out=Out[0]+Out[1]+Out[2] (16)

[0085] Among them, Out represents the attention output matrix, that is, the final approximate attention output.

[0086] To illustrate the effectiveness of the method of this application, the weight parameters of the Mistral-7B model are inherited, and then its standard Attention is replaced with the key-value compression method of this application. The loss curve of the fine-tuned Mistral-7B model is as follows: Figure 3As shown. During fine-tuning, only the weight matrix W is trained q 、W k 、W v As well as the parameters of the embedding layer and regularization module, other parameters are frozen to improve training speed. Training uses two A100 computing cards, sequence length 2K, gradient accumulation 64. Figure 3 It can be seen that in less than 200 steps, the linear attention of this application converges to the loss of the original Mistral-7B model.

[0087] After the Mistral-7B model is trained for 1000 steps, the perplexity results tested on the PG-19 dataset and the Proof-Pile dataset are shown in Table 1, with test lengths of 512, 1K, and 2K, respectively.

[0088] Table 1. Perplexity results of Mistral-7B model on different datasets

[0089] Test length PG-19 dataset Proof-Pile Dataset 512 7.39 3.58 1024 7.76 3.69 2048 8.17 3.83

[0090] As can be seen from Table 1, this application achieves considerable perplexity on multiple data sets.

[0091] The effectiveness of multi-head KV compression comes from the residual decomposition and cluster compression of the key matrix and value matrix. Because when compressing KV using the codebook tokens sequence, the feature-related KV is compressed into cluster centers and stored in K i_z 、V i_z . However, due to the limitation of storage size, the shared features stored in a single cluster center are very rough, so this application adopts a multi-level residual decomposition and cluster compression method to solve this problem, that is, a multi-head compression method. For example, the first-level KV maintains the coarsest but largest information storage feature as the cluster center, and the second level performs a finer decomposition on the rough cluster center, that is, decomposing the large categories of the first level into small categories. The second-level KV stores cluster features relative to the first-level category. Its effectiveness lies in the shared feature differences between the first and second levels. The more shared, the better the performance. The third level is similar to the above process, and performs finer clustering within the second-level category. The optimal situation is that the feature differences can be fully shared, thereby achieving a storage space with a power exponential of the parameter size. The upper-level application can adjust the compression rate and restoration degree to a suitable ratio by adjusting the number of compression heads.

[0092] This application uses a multi-level codebook to perform effective clustering compression on the KV matrix. The compressed approximate KV matrix (i.e., K i_z 、V i_z) is of fixed length, making the attention computation between the query matrix and the approximate KV matrix linear in length with the original key and value matrices and requiring only constant-level storage, with computational and storage complexity comparable to linear attention. Compared to traditional linear attention, this application still supports the use of softmax to reweight token similarities, resulting in results closer to standard softmax attention.

[0093] The method of this application relies on a set of fine-tuned Codebook token sequences to map long sequence tokens of size N×d to a fixed size C×d token sequence; the method of this application improves the standard SoftmaxAttention method, calculates the approximate attention output, and achieves a better approximation to the existing Softmax Attention. This application inherits the weights of the current large language model based on Softmax Attention and fine-tunes it, converting SoftmaxAttention into a linear recursive method: a sliding window combined with the multi-head KV compression storage mechanism of this application. Experiments show that this application can achieve convergence with very few training resources and achieve considerable perplexity on multiple datasets.

[0094] The proposed method realizes the limited space key-value storage of long sequences, reducing the training storage overhead and computational complexity from N 2 Reduced to N, the inference storage overhead is reduced from N to 1, and the inference calculation amount is reduced from N 2 Compared to Linear Attention, this method has the advantage of not requiring re-pre-training. Furthermore, the model weights fine-tuned using this method can be used for both the large language model (i.e., the approximate Transformer model) and the standard Transformer model inference, thus offering greater model application flexibility compared to Linear Attention. Furthermore, because this method still uses softmax for weight rearrangement, the model performance is closer to the standard Transformer model than various Linear Attention models.

[0095] Example 2

[0096] Figure 4 The large language model architecture diagram provided by this application is shown. The square brackets represent the parameters that need to be trained for the model, mainly including the weight parameters of the embedding layer, the weight matrix W q 、W k 、W v , trainable matrices C0~C2, MLP weight parameters. Figure 4As shown, the large language model (i.e., the approximate Transformer model) includes an embedding layer, a position encoding layer, an output layer, and multiple Transformer Blocks; each Transformer Block includes an attention layer and an MLP layer, the attention layer includes a first regularization module, a self-attention transformation module, and a first superposition module connected in sequence, and the MLP layer includes a second regularization module, an MLP transformation module, and a second superposition module connected in sequence.

[0097] The embedding layer is used to convert each word in the text data into a fixed-length feature vector: the position encoding layer is used to map the position of each word in the text data to the feature space and superimpose it on the feature vector corresponding to the word to obtain the input embedding matrix.

[0098] The first regularization module is used to regularize the input embedding matrix; the self-attention transformation module is used to perform QKV transformation on the regularized matrix to obtain a query matrix, a key matrix and a value matrix, and perform multiple residual decompositions on the key matrix and the value matrix respectively to obtain a key residual vector and a value residual vector after each decomposition; the key residual vector and the value residual vector after each decomposition are clustered and compressed respectively, and attention calculation is performed on the query matrix, the compressed key residual vector and the value residual vector; all attention calculation results are accumulated to obtain the attention output matrix Out. For the specific process, please refer to the method in Example 1 of the present application; the first superposition module is used to superimpose the input embedding matrix and the attention output matrix Out to obtain the attention layer output matrix.

[0099] The second regularization module is used to regularize the output matrix of the attention layer; the MLP transformation module is used to extract features from the regularized matrix to obtain a feature matrix; the second superposition module is used to superimpose the output matrix of the attention layer and the feature matrix to obtain the output matrix of the MLP layer.

[0100] The entire large language model can be implemented by looping through multiple Transformer Blocks. The first call of a Transformer Block takes as input the input embedding matrix obtained by superimposing the input layer and the positional encoding layer. Subsequent calls take as input the output of the previous call, that is, the output matrix of the previous MLP layer.

[0101] The trainable matrices C0 to C2 of the large language model are introduced in the fine-tuning phase. The other parameters that need to be trained have the same parameter structure as the model pre-training phase. The parameter results obtained in the pre-training phase can be reused in the model inference phase.

[0102] The compression mechanism provided by this application theoretically encompasses and expands the most advanced key-value (KV) compression techniques used in current linear attention models, including the gating mechanism of GLA, the ABC gating architecture of GSA, and the Neural Turing Machine (NTM) approach. By separating storage from the current hidden state, this application's method achieves a more flexible and robust sequence representation that can adapt to various attention mechanisms, including standard Softmax Attention and its linearized counterpart. The approximate Transformer model constructed using this compression mechanism has better adaptability and is suitable for achieving high levels of performance in a wide range of long sequence tasks.

[0103] Example 3

[0104] An embodiment of the present invention further provides an electronic device, such as Figure 5 As shown, the electronic device includes: a memory, a processor, and a computer program / instruction stored in the memory, and the processor executes the computer program / instruction to implement the key-value compression method in the self-attention mechanism in Example 1 of the present application.

[0105] Although not shown, the electronic device includes a processor that can perform various appropriate operations and processes based on programs and / or data stored in a read-only memory (ROM) or programs and / or data loaded from a storage portion into a random access memory (RAM). The processor can be a multi-core processor or can include multiple processors. In some embodiments, the processor can include a general-purpose main processor and one or more special coprocessors, such as a central processing unit, a graphics processing unit (GPU), a neural network processor (NPU), a digital signal processor (DSP), etc. Various programs and data required for device operation are also stored in RAM. The processor, ROM, and RAM are connected to each other via a bus. An input / output (I / O) interface is also connected to the bus.

[0106] The processor and memory are used together to execute the program / instructions stored in the memory. When the program / instructions are executed by the computer, the methods, steps or functions described in the above embodiments can be implemented.

[0107] Although not shown, an embodiment of the present invention further provides a computer-readable storage medium having a computer program / instruction stored thereon, which, when executed by a processor, implements the key-value compression method in the self-attention mechanism in the first embodiment of the present application.

[0108] Computer-readable storage media include permanent and non-permanent, removable and non-removable media that can be implemented by any method or technology to store information. The information can be computer-readable instructions, data structures, program modules, or other data. Examples of computer storage media include, but are not limited to, phase change memory (PRAM), static random access memory (SRAM), dynamic random access memory (DRAM), other types of random access memory (RAM), read-only memory (ROM), electrically erasable programmable read-only memory (EEPROM), flash memory or other memory technology, compact disc read-only memory (CD-ROM), digital versatile disc (DVD) or other optical storage, magnetic cassettes, magnetic disk storage or other magnetic storage devices, or any other non-transmission media that can be used to store information that can be accessed by a computing device. As defined herein, computer-readable media does not include transitory media such as modulated data signals and carrier waves.

[0109] Although not shown, an embodiment of the present invention further provides a computer program product, including: a computer program / instruction, which, when executed by a processor, implements the key-value compression method in the self-attention mechanism in the first embodiment of the present application.

[0110] Although preferred embodiments of the present application have been described, those skilled in the art may make additional changes and modifications to these embodiments once they are aware of the basic inventive concepts. Therefore, the appended claims are intended to be interpreted as including the preferred embodiments and all changes and modifications that fall within the scope of this application. Obviously, those skilled in the art may make various changes and modifications to this application without departing from the spirit and scope of this application. Thus, if such changes and modifications fall within the scope of the claims of this application and their equivalents, then this application is intended to include such changes and modifications.

Claims

1. A key-value compression method in a self-attention mechanism, applied to large language models, characterized in that: The compression method comprises: Perform multiple residual decompositions on the key matrix and the value matrix respectively to obtain the key residual vector and the value residual vector after each decomposition; wherein the key residual vector after each decomposition includes the key residual vector after the 0th decomposition, and the value residual vector after each decomposition includes the value residual vector after the 0th decomposition; The key residual vector and value residual vector after each decomposition are clustered and compressed respectively, and attention calculation is performed on the query matrix, the compressed key residual vector and the value residual vector; Accumulate all attention calculation results.

2. The key-value compression method in the self-attention mechanism according to claim 1, characterized in that: The specific formula for the residual decomposition of the bond matrix is: K i =K i-1 -gather(C i-1 [0],index1); index1=argmax(matmul(K i-1 ,C i-1 [0].transpose)); Among them, K i Represents the key residual vector after the i-th residual decomposition, K i-1 Represents the key residual vector after the i-1th residual decomposition. When i=1, K0 represents the key matrix; gather represents the selection of C according to the index position index1. i-1 [0] and splice; C i-1 Represents the trainable matrix, C i-1 [0] represents the trainable matrix C i-1 argmax means returning the index of the maximum value of the matrix in the specified dimension, matmul means matrix multiplication, and transpose means matrix transposition; The specific formula for the residual decomposition of the value matrix is: V i =V i-1 -gather(C i-1 [1],index2); index2=argmax(matmul(V i-1 ,C i-1 [1].transpose)); Among them, V i Represents the residual vector after the i-th residual decomposition, V i-1 Represents the residual vector after the i-1th residual decomposition. When i=1, V0 represents the value matrix; gather represents the selection of C according to index position index2. i-1 [1] and splice; C i-1 [1] represents the trainable matrix C i-1 The first dimension of .

3. The key-value compression method in the self-attention mechanism according to claim 1 or 2, characterized in that: The key residual vector and value residual vector after each decomposition are clustered and compressed separately, including: Calculating a correlation matrix between the key residual vector and the trainable matrix; Performing a mapping transformation on the correlation matrix to obtain a weight matrix; Clustering the key residual vector and the weight matrix to obtain a compressed key residual vector; Clustering is performed on the value residual vector and the weight matrix to obtain a compressed value residual vector.

4. The key-value compression method in the self-attention mechanism according to claim 3, characterized in that: Performing a mapping transformation on the correlation matrix includes: A feedforward neural network or a softmax function is used to perform mapping transformation on the correlation matrix.

5. A large language model comprising an embedding layer, a positional encoding layer, an output layer, and multiple Transformer Blocks; each Transformer Block comprises an attention layer and an MLP layer; characterized in that: The embedding layer is used to convert each word in the text data into a feature vector of fixed length; the position encoding layer is used to map the position of each word in the text data to the feature space and superimpose it on the feature vector corresponding to the word to obtain an input embedding matrix; The attention layer includes a first regularization module, a self-attention transformation module and a first superposition module; the first regularization module is used to regularize the input embedding matrix; the self-attention transformation module is used to perform QKV transformation on the regularized matrix to obtain a query matrix, a key matrix and a value matrix, and perform multiple residual decompositions on the key matrix and the value matrix respectively to obtain a key residual vector and a value residual vector after each decomposition; the key residual vector and the value residual vector after each decomposition are clustered and compressed respectively, and attention calculation is performed on the query matrix, the compressed key residual vector and the value residual vector; All attention calculation results are accumulated to obtain an attention output matrix; the first superposition module is used to superimpose the input embedding matrix and the attention output matrix to obtain an attention layer output matrix; The MLP layer includes a second regularization module, an MLP transformation module and a second superposition module; the second regularization module is used to regularize the attention layer output matrix; the MLP transformation module is used to extract features from the regularized matrix to obtain a feature matrix; the second superposition module is used to superimpose the attention layer output matrix and the feature matrix.

6. The large language model according to claim 5, characterized in that The specific formula for the self-attention transformation module to perform multiple residual decompositions on the key matrix and the value matrix is: K i =K i-1 -gather(C i-1 [0],index1); index1=argmax(matmul(K i-1 ,C i-1 [0].transpose)); V i =V i-1 -gather(C i-1 [1],index2); index2=argmax(matmul(V i-1 ,C i-1 [1].transpose)); Among them, K i Represents the key residual vector after the i-th residual decomposition, K i-1 Represents the key residual vector after the i-1th residual decomposition. When i=1, K0 represents the key matrix; gather represents the selection of C according to the index position index1. i-1 [0] and splice; Ci -1 Represents the trainable matrix, C i-1 [0] represents the trainable matrix C i-1 argmax means returning the index of the maximum value of the matrix in the specified dimension, matmul means matrix multiplication, and transpose means matrix transposition; V i Represents the residual vector after the i-th residual decomposition, V i-1 Represents the residual vector after the i-1th residual decomposition. When i=1, V0 represents the value matrix; gather represents the selection of C according to index position index2. i-1 [1] and splice; C i-1 [1] represents the trainable matrix C i-1 The first dimension of .

7. The large language model according to claim 5, characterized in that The self-attention transformation module is used to perform clustering compression on the key residual vector and the value residual vector after each decomposition, including: Calculating a correlation matrix between the key residual vector and the trainable matrix; Performing a mapping transformation on the correlation matrix to obtain a weight matrix; Clustering the key residual vector and the weight matrix to obtain a compressed key residual vector; Clustering is performed on the value residual vector and the weight matrix to obtain a compressed value residual vector.

8. An electronic device comprising a memory, a processor, and a computer program / instruction stored on the memory, wherein the processor executes the computer program / instruction to implement the steps in the key-value compression method in the self-attention mechanism as described in any one of claims 1 to 4.

9. A computer-readable storage medium having a computer program / instruction stored thereon, which, when executed by a processor, implements the steps of the key-value compression method in the self-attention mechanism as described in any one of claims 1 to 4.

10. A computer program product comprising a computer program / instructions, characterized in that When the computer program / instructions are executed by a processor, the steps of the key-value compression method in the self-attention mechanism as described in any one of claims 1 to 4 are implemented.

Images