Sparse attention mechanism method based on CK-means + +

Abstract

The invention provides a sparse attention mechanism method based on CK-means + +. The sparse attention mechanism method is applied to the situation that clustering needs to be carried out on a sparse attention matrix when a large language model is trained. According to the method, a Canopy clustering and mean value calculation method is adopted to process a Query vector and a Key vector of a transform attention layer so as to determine the number of clustering clusters, a roulette method is adopted to carry out iterative calculation according to the number of the clusters so as to determine an initial center point of each cluster, and then K-means clustering convergence calculation is carried out according to the number of the clustering clusters and the initial center point of each cluster. And finally, a key (Key) vector with high similarity with each Query vector is obtained, and an attention matrix is calculated according to the Query-Key pair. According to the method, the number of clusters and the initial clustering center point can be automatically calculated, the problem that clustering is sensitive to manual selection of the two values is solved, the execution speed is high, high complexity of a large language model caused by various combinations of spatial position codes can be avoided, and the method has the advantages of being better in flexibility, small in calculation amount and high in robustness.

Description

Technical Field

[0001] This invention belongs to the field of large language model training, and specifically relates to a sparse attention mechanism method based on CK-means++. Background Technology

[0002] With the advent of the Transformer self-attention mechanism, large language models have made groundbreaking progress, and training models with billions of parameters based on this mechanism has continuously enhanced their capabilities. However, as the amount and length of training texts increase, it has been found that in the Transformer architecture, the time and storage complexity of the self-attention mechanism are quadratically related to the sequence length, consuming a large amount of computing memory and computational resources. Since researchers generally have limited computing resources, optimizing the time and space complexity of the self-attention mechanism and improving computational efficiency remains a crucial challenge in training large language models.

[0003] Research has revealed that the attention matrices in many well-trained Transformer structures are sparse. Therefore, during the training of large language models, computational complexity can be reduced by limiting the number of query-key pairs. Two existing methods exist: one is a position-based sparse attention mechanism (designing multiple combinations of local positional attention to reduce the computation of global attention), and the other is a content-based sparse attention mechanism (finding and selecting keys with high similarity to a given query, so that self-attention computation only occurs in highly similar query-key pairs). Position-based sparse attention mechanisms typically require computation of spatial location encoding, increasing the complexity of large language models. Furthermore, the combination of selected positional mechanisms relies on manual selection, making it difficult to achieve optimal results. Content-based sparse attention mechanisms, on the other hand, can better capture important information in sequences and handle complex sequences. However, current methods for finding highly similar query-key pairs use K-means clustering, which requires manually setting the number of clusters and the initial centroids of each cluster. This introduces significant uncertainty into the attention mechanism computation during large language model training, potentially affecting the effectiveness of the final large language training computation. Summary of the Invention

[0004] To address the problems existing in content-based sparse attention mechanisms for clustering during the training of large language models, this invention provides a robust, computationally inexpensive, and highly flexible sparse attention method based on CK-means++. It eliminates the need for manually setting the number of clusters and initial center vectors, and can adaptively and dynamically calculate and complete clustering. The technical solution adopted in this invention includes the following steps:

[0005] Determine the query layer that needs to be computed. and bond layer Given the self-attention mechanism, it's known that the matrices at two different levels have the same row dimension d. Therefore, merging the two matrices into a new vector matrix R = [Q, K] is easily understood. The next step is to perform clustering on all column vectors in the vector space R.

[0006] First, the Canopy method is used to automatically determine the number of clusters required for the final K-means clustering. In R, the input is then processed sequentially using R... i Calculate the Euclidean distance between the center vector and all other vectors. Determine if it is less than the distance threshold Where D is the sum of the Euclidean distances between the vectors. If it is less than 1, then this vector R is... i Remove the vectors from R and assign them to the set A of centroids. Repeat the above steps until R is empty. At this point, you can obtain the number of clusters, which is the number of vectors m in set A.

[0007] Next, the roulette wheel method is used to adaptively determine the initial center vector input required for the final K-means clustering. A vector is randomly selected from the vector space R as the cluster center, based on the existing cluster center vector C. i Calculate the shortest Euclidean distance to each of the other vectors. Calculate these vectors R i Cumulative selection probability in Then, based on a random number λ within the 0-1 interval, determine which vector's cumulative probability interval λ falls into, and select that vector as the next cluster center. Repeat the above steps until m center vectors have been selected, at which point the initial set of center points C = {C i} i＝1,...,m .

[0008] Finally, based on the number of categories m and the initial set of center vectors C obtained from the two methods mentioned above, K-means clustering is performed. Specifically, the K-means clustering calculation is performed on each vector R in R. x and the center points C in set C i Euclidean distance Vector R x Assign the data to the cluster with the shortest centroid vector. After the assignment, update the centroid vector of the new cluster based on the new cluster vector. Repeat the above steps iteratively until the difference between the sum of squared errors (SSE) of all clusters before the update and the sum of squared errors after the update is less than or equal to 10. -6 This proves that the clustering state has approached stability, and the entire method can be stopped.

[0009] Following the steps above, Q and K vectors with high similarity can be clustered into a single cluster (Q, K ∈ C). Subsequent attention calculations can then be performed based on the fact that the selected vector Q only occurs in vector K within the same cluster, greatly reducing the number of query-key pairs and significantly improving computational efficiency.

[0010] The advantages of this invention are:

[0011] (1) This method performs a preprocessing of center vector selection before starting the clustering iteration, which can speed up the subsequent clustering convergence execution speed;

[0012] (2) This method does not require manually determining the number of clusters to be clustered, thus eliminating the influence of manual selection;

[0013] (3) This method uses adaptive calculation of the initial center vector, which alleviates the problem that K-means clustering is sensitive to the initial cluster center vector. Attached Figure Description

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

[0015] Figure 1 This is a flowchart of the computation of the sparse attention mechanism based on CK-means++ of the present invention. Detailed Implementation

[0016] To make the objectives, technical solutions, and advantages of this application clearer, the technical solutions of this application will be clearly and completely described below in conjunction with specific embodiments and corresponding drawings. Obviously, the described embodiments are only a part of the embodiments of this application, and not all of them. Based on the embodiments in this application, all other embodiments obtained by those skilled in the art without creative effort are within the scope of protection of this application.

[0017] The specific steps in this embodiment are as follows: Figure 1 As shown, the technical solution of the present invention will be further described in detail below with examples.

[0018] Taking a query layer Q and a key layer K in the Transformer's self-attention layer as an example, let the weight matrix of the query layer Q be denoted as... The weight matrix of the bond layer K is Since the two matrices have the same row dimension, they are merged into a new vector matrix R = [Q, K].

[0019] The calculation steps for each iteration in determining the number of clusters are as follows:

[0020] (1) Calculate the sum of the Euclidean distances D between the vectors.

[0021] (2) In sequence, vector R i Let ∈R be the center vector, and calculate its relationship with other vectors R. j Euclidean distance ∈R Will The vectors are removed from R and assigned to the set A of center points;

[0022] (3) Determine if R is empty. If it is empty, exit the loop iteration and output the number of vectors m of the center point set A as the number of clusters. If it is not satisfied, continue to use a new vector as the center vector to solve the problem.

[0023] The calculation steps for each iteration in solving for the initial center vector are as follows:

[0024] (1) Randomly select a vector from R as the cluster center, denoted as C. i＝1 Calculate the shortest distance between each vector and the existing cluster centers.

[0025] (2) Based on each vector Calculate the probability that this vector will be selected as the next cluster center. Cumulative probability of each vector Generate a random number λ between 0 and 1, and determine which vector's cumulative probability interval it falls into. Then select this vector as the next cluster center;

[0026] (3) Determine if a total of m cluster centers have been selected. If so, exit the loop iteration and output the set of center points C = {C i} i＝1,...,m If the conditions are not met, the cumulative probability of each vector is calculated and random numbers are generated for judgment and solution.

[0027] When performing clustering method calculations, the calculation steps at each time step are as follows:

[0028] (1) Calculate the vectors R in R. x and the center points C in set C i Euclidean distance Vector R x Classified to The center point C corresponding to the minimum value i In the cluster;

[0029] (2) The original cluster μ i Updated to μ i ', and set the average vector of each vector in the new cluster as the cluster μ. i ' new center vector Where |μ i '| indicates belonging to cluster μ i The number of vectors;

[0030] (3) Calculate the sum of squared errors before the update. For vector R x To the cluster center vector C before the update i Calculate the Euclidean distance; calculate the updated sum of squared errors. For vector R x To the updated cluster center vector C i The Euclidean distance of '.

[0031] (4) Determine whether the convergence condition is met. ε = 10 -6 If the conditions are met, clustering calculations are stopped, and the query Q-vector and key K-vector for each cluster are output. Subsequent self-attention mechanism calculations only compute QK vectors within the same cluster. T Q, K∈C. If this condition is not met, continue calculating the updated vector R. x and the center points C in set C i The Euclidean distance is used to continue updating the cluster and cluster center vector.

[0032] While the specific embodiments of the present invention have been described above in conjunction with the accompanying drawings, this is not intended to limit the scope of protection of the invention. Those skilled in the art should understand that various modifications or variations that can be made by those skilled in the art without creative effort based on the technical solutions of the present invention are still within the scope of protection of the present invention.

Claims

1. A sparse attention mechanism method based on CK-means++, characterized by The following steps are involved: S1. Determine the query layer that transformer self-attention needs to calculate and key layer Merge the two matrices into a new vector matrix R = [Q, K], then S2. Calculate the sum of the Euclidean distances D between each vector and the distance threshold T that the Canopy method needs to judge: S3. Take the vector R in R in turn i is the center vector, and calculates the relationship between it and other vectors R in R j The Euclidean distance Will The vectors are removed from R and divided into the set A of center points; S4. Determine whether R is empty. If it is empty, output the number m of vectors of the center point set A as the input of the number of clusters for subsequent K-means clustering. If it is not empty, return to S3; S5. Randomly select a vector from R of S1 as the first cluster center, denoted as C i＝1 ; S6. Calculate the shortest distance between each vector and the current cluster center S7. According to each vector Calculate the probability of this vector being selected as the next cluster center: S8. Use the roulette method to calculate and solve the vector that will serve as the next cluster center, specifically: Calculate the cumulative probability of each vector Generate a random number λ∈[0,1], determine which vector's cumulative probability interval λ falls in, and select that vector as the next cluster center, that is, satisfy S9. Determine whether a total of m cluster centers in S4 are selected. If so, output the set of center points C = {C i } i＝1,...,m As the initial center point input for subsequent K-means clustering, if it is not empty, return to S6; S10. Calculate the vectors R in R of S1 x and the center points C in the set C i The Euclidean distance S11. Transform the vector R x Divide into The center point C corresponding to the minimum value i In the cluster, the original cluster μ i Update to μ i ', and set the average vector of each vector in the new cluster to cluster μ i 'The new center vector C i ': where |μ i '| means belonging to cluster μ i 'The number of vectors; S12. Calculate whether the difference between the sum of squared errors (SSE) of all clusters after the update and before the update converges, specifically: For each vector R x , calculate its updated cluster center vector C i 'Distance Then for each updated cluster μ i 'Calculate the sum of squared errors for this cluster Finally, the sum of squared errors of all clusters is added up to get the updated sum of squared errors of all clusters. The sum of squared errors before updating SSE μ Calculate similarly; Determine whether the convergence condition is met If satisfied, proceed to the next step, if not satisfied, return to S10; S13. Output the query Q vector and key K vector in each cluster. When performing subsequent self-attention mechanism calculations, only QK in the same cluster is calculated. T ,Q,K∈C.

Images