Method for improving reasoning efficiency of long-context large language model, electronic equipment and storage medium

Abstract

The embodiment of the invention discloses a method for improving the reasoning efficiency of a long-context large-scale language model, electronic equipment and a storage medium, and the method for improving the reasoning efficiency of the long-context large-scale language model comprises a method for improving the reasoning efficiency of the long-context large-scale language model. Comprising the following steps: analyzing the similarity of attention scores among layers in a given large language model, and classifying continuous similar layers into the same block; applying attention sharing in each block, and performing post-training on the large language model; and performing efficient reasoning by using the post-trained large language model.

Description

Technical Field

[0001] The present invention belongs to the technical field of large language models, and in particular to a method, electronic device and storage medium for improving the reasoning efficiency of a large language model with a long context. Background Art

[0002] In related technologies, the context window size of large language models (LLMs) continues to increase, enabling them to excel in handling complex tasks that require in-depth exploration of lengthy texts. However, this also poses challenges to the computation and memory usage of LLMs. Specifically, on the one hand, since most LLMs are based on the Transformer architecture, the computational complexity of the attention module increases quadratically with the increase in the context window size. On the other hand, the size of the KV cache is linearly related to the context window size, and the KV (Key-Value) cache is a commonly used technique designed to prevent redundant computations. Therefore, it is crucial to improve the efficiency of LLMs by using extended context windows.

[0003] In this context, many researchers have proposed methods to improve the reasoning efficiency of LLM by discarding some tokens (units) in the context. Among them, the Window attention method and StreamingLLM identify the "attention sink" phenomenon while retaining the initial token and the most recent token (see Figure 1 (a) in the figure). H2O takes into account the different importance of tokens in the context and selectively retains only the most important tokens in the KV cache based on the attention score. Although this approach improves the efficiency of LLM in processing long contexts, it also brings a major flaw: key words required for later text generation may be irreversibly discarded early in the process. Figure 1 As shown in (a) in Figure 1, when important tokens (evidence in the example) fall outside the window, the prediction fails. In addition, the performance degradation of StreamingLLM and H2O on two real benchmarks further confirms this (see (b) in Figure 1). Summary of the invention

[0004] Embodiments of the present invention provide a method, an electronic device, and a storage medium for improving the reasoning efficiency of a long-context large-scale language model, which are used to solve at least one of the above-mentioned technical problems.

[0005] In a first aspect, an embodiment of the present invention provides a method for improving the reasoning efficiency of a large language model with a long context, comprising: analyzing the similarity of attention scores between layers in a given large language model, and classifying consecutive similar layers into the same block; applying attention sharing in each block, and post-training the large language model; and using the post-trained large language model for efficient reasoning.

[0006] In a second aspect, an electronic device is provided, comprising: at least one processor, and a memory communicatively connected to the at least one processor, wherein the memory stores instructions executable by the at least one processor, and the instructions are executed by the at least one processor so that the at least one processor can perform the steps of the method for improving the efficiency of reasoning of a long-context large language model according to any embodiment of the present invention.

[0007] In a third aspect, an embodiment of the present invention further provides a computer program product, comprising a computer program stored on a non-volatile computer-readable storage medium, wherein the computer program comprises program instructions, and when the program instructions are executed by a computer, the computer executes the steps of the method for improving the reasoning efficiency of a long-context large language model according to any embodiment of the present invention.

[0008] In the method of the embodiment of the present application, by first analyzing the similarity of the attention scores between layers, and then classifying the continuous similar layers into the same block according to the similarity, and then applying attention sharing in each block, the large language model is post-trained, and the post-trained large language model is used for efficient reasoning, thereby improving the reasoning efficiency of the large language model with a long context. BRIEF DESCRIPTION OF THE DRAWINGS

[0009] In order to more clearly illustrate the technical solutions of the embodiments of the present invention, the accompanying drawings required for use in the description of the embodiments will be briefly introduced below. Obviously, the accompanying drawings described below are some embodiments of the present invention. For ordinary technicians in this field, other accompanying drawings can be obtained based on these accompanying drawings without paying any creative work.

[0010] Figure 1 The experimental results of the customized LLaMA3-8B-32K model provided in one embodiment of the present invention; wherein, Figure 1 (a) shows an example of StreamingLLM prediction failure; Figure 1 (b) shows some experimental results of LongBench and LEval; Figure 1 (c) in Figure 1 shows that the test data consists of 1000 32K-length text samples sampled from Dolma; for each sample sequence, we compared the last 100 tokens; Figure 1 (d) in Figure 2 shows that we sampled 1,000 sequences of 32K length from Dolma and extracted the attention scores of the last 16 tokens in each sample; Figure 2 A flowchart of a method for improving the inference efficiency of a long-context large language model provided by an embodiment of the present invention Figure 3 Pseudo code of a greedy layer grouping algorithm provided by an embodiment of the present invention, wherein input: head attention similarity between layers, output: head layer block; Figure 4 The three stages of PoD (Proximal tokens over Distant tokens) provided by an embodiment of the present invention are shown; Figure 5 The search results of the needle-in-a-haystack test provided by an embodiment of the present invention are shown; Figure 6 The evaluation results of different methods provided by an embodiment of the present invention on two well-known long-context benchmarks are shown; Figure 7 The theoretical and actual memory usage savings provided by an embodiment of the present invention are shown; Figure 8 The relationship between the computation saving rate and performance loss rate provided by an embodiment of the present invention and the gate threshold τ mentioned in the above embodiment is shown. ; Fig. 9 The relationship between the model performance on LEval provided by an embodiment of the present invention and two factors is shown; Fig.10 A case study showing different methods provided by an embodiment of the present invention; Fig.11 It is a schematic diagram of the structure of an electronic device provided by an embodiment of the present invention. DETAILED DESCRIPTION

[0011] In order to make the purpose, technical solution and advantages of the embodiments of the present invention clearer, the technical solution in the embodiments of the present invention will be clearly and completely described below in conjunction with the drawings in the embodiments of the present invention. Obviously, the described embodiments are part of the embodiments of the present invention, not all of the embodiments. Based on the embodiments of the present invention, all other embodiments obtained by ordinary technicians in this field without creative work are within the scope of protection of the present invention.

[0012] Methods in related technologies improve the inference efficiency of large language models (LLMs) by discarding certain tokens in the context. In particular, the Window Attention method retains a fixed-size window on the KV state of recent tokens. LM-Infinite and StreamingLLM discovered the "attention sink" phenomenon and retained the initial tokens and recent tokens. H2O considers the importance differences of tokens in the context and selectively retains the most important tokens in the KV cache based on the attention score.

[0013] The inventors found that the above scheme may irreversibly discard key tokens required for subsequent text generation in the early stage of the decoding process. This defect is caused by the fact that the above methods only focus on some tokens in the context. To solve the above defects, those skilled in the art usually use KV sharing between layers. This is different from the method in the embodiment of the present application, which uses the similarity of attention scores between layers and shares them between layers.

[0014] The inventors based their work on two observations: 1) Proximal tokens (initial tokens + recent tokens) are more important than distant tokens. 2) The attention scores between consecutive layers have strong similarities.

[0015] Based on these two observations, we save the size of the KV cache (key-value cache) by sharing the attention scores of distant tokens between layers.

[0016] The three stages of PoD (Proximal tokens over Distant tokens). (1): Example of grouping layers based on the similarity of attention scores calculated offline for each head. Consecutive layers with the same color belong to the same block and therefore use the same attention score. (2): Example model structure of the attention module of head 3 in PoD. Each token pays attention to two groups: proximal tokens (the adjacent r tokens and the initial s tokens, in this example, r=2, s=1) and distant tokens (the rest of the tokens). The attention of the two groups is calculated separately and weighted according to their respective attention scores. (3): Example of KV cache changes for the 7th token and avoiding attention calculation for distant tokens.

[0017] In the process of implementing this application, the inventors have also adopted the following alternative solutions and Beta versions. Among them, the alternative solution: each layer shares the KV of the last layer. Advantages: high efficiency, disadvantages: 1) complex training process, 2) large loss of model performance. Beta version: for relatively unimportant distant tokens, consecutive layers do not share attention distribution, but share key values. Advantages: simpler implementation, disadvantages: since the similarity of keys between consecutive layers is not very high, it will cause certain loss to model performance.

[0018] The solution of the embodiment of the present application can save 35% of the KV cache size while maintaining the model performance. In addition, since remote tokens are not necessary in most scenarios, it is possible to use this phenomenon to accelerate the reasoning of large language models in the future.

[0019] Please refer to Figure 2 , which shows a flowchart of an embodiment of a method for improving the reasoning efficiency of a long-context large language model of the present application.

[0020] like Figure 2 As shown, in step 201, the similarity of attention scores between layers is analyzed in a given large language model, and consecutive similar layers are classified into the same block; In step 202, attention sharing is applied in each block and the large language model is post-trained; In step 203, the post-trained large language model is used for efficient reasoning.

[0021] The solution of the embodiment of the present application first analyzes the similarity of attention scores between layers, and then classifies consecutive similar layers into the same block based on the similarity, and then applies attention sharing in each block, post-trains the large language model, and uses the post-trained large language model for efficient reasoning, thereby improving the reasoning efficiency of large language models with long contexts.

[0022] In some optional embodiments, analyzing the similarity of attention scores between layers in a given large language model and classifying consecutive similar layers into blocks includes: calculating the attention scores of each layer of each attention head; calculating the attention similarity between layers based on the attention scores; if there is a similarity between two layers greater than or equal to a preset threshold, determining that the two layers are similar layers; and grouping consecutive similar layers into head blocks.

[0023] In a further optional embodiment, grouping the consecutive similar layers into a header block includes: iteratively merging the consecutive similar layers into the header block using a bottom-up greedy algorithm.

[0024] In some optional embodiments, the applying of attention sharing in each block and post-training of the large language model includes: for each unit, calculating the attention outputs of the close units and the distant units, wherein the attention scores of the distant units are represented by the attention scores of the bottom layer of the corresponding head blocks after grouping, the close units include the adjacent units and the initial units, and the distant units include other units except the adjacent units and the initial units; aggregating the attention outputs of the close units and the distant units.

[0025] In a further optional embodiment, the aggregating the attention outputs of the proximate unit and the distant unit includes: integrating the attention of the proximate unit and the distant unit using a parameter-free gating mechanism.

[0026] In some optional embodiments, the use of a post-trained large language model for efficient reasoning includes: optimizing the memory usage of the KV cache by using a shared attention score in an attention mechanism and optimizing the computation of the remote unit.

[0027] It should be noted that the above method steps are not used to limit the execution order of each step. In fact, some steps may be executed simultaneously or in the opposite order to the specified step order, and this application has no limitation on this.

[0028] The solution of the present application is described below through a specific embodiment so that those skilled in the art can better understand the solution of the present application.

[0029] In the embodiments of the present application, the inventors attempt to improve the efficiency of LLM while minimizing performance degradation. The inventors' core motivation is that less important tokens should be allocated fewer resources rather than being discarded completely. This raises two challenges: 1) For a token, where are the important tokens it should pay attention to distributed; and 2) how to optimize the memory and computation of less important tokens. The inventors attempt to address these two challenges through two key observations: Observation 1: Proximal tokens (initial tokens + recent tokens) are more important than distant tokens (see Figure 1 We conducted experiments to study the proportion of cases where the next token is predicted to be the same when each token only attends to a fixed number of similar tokens in the context, rather than attending to all tokens. Figure 1(c) in Figure 2 shows that for the same input sequence, even if we only focus on 256 similar tokens, in 80% of cases, the model’s prediction of the next token is the same as the model focusing on all tokens. This phenomenon confirms Observation 1.

[0030] Observation 2: Attention scores are similar between consecutive layers. Figure 1 In (d), we can find that the attention scores between the layers in the red box show strong similarities.

[0031] To this end, we propose PoD (Proximal tokens over Distant tokens) to optimize the inference efficiency of the decoding stage. Specifically, it only shares the inter-layer attention scores for distant tokens, while keeping the nearby tokens unchanged.

[0032] This approach consists of three main stages: 1) Exploration of Offline Inter-Layer Attention Sharing: Determine which layers can share attention scores; 2) Lightweight Training Adaptation: Post-train dense models with limited data based on the determined inter-layer attention sharing pattern; 3) Efficient Inference: Sharing the attention scores of distant tokens between layers allows us to retain the key state of a single layer in the KV cache. In addition, we can pre-identify situations where only nearby tokens are needed to predict the next token, thus eliminating the need to calculate the attention of distant tokens.

[0033] We evaluate the performance of PoD on the NeedleInAHaystack test and two real-world long context benchmarks, analyze its efficiency and study the impact of key hyperparameters. We also conduct case studies. Extensive experiments demonstrate that PoD can save 35% of KV cache without affecting model performance. In summary, our contributions are 1) We propose the idea of ​​evaluating the importance of tokens based on their position in the context and improve inference efficiency by reducing the resources allocated to less important tokens. 2) We introduce PoD, a new model that adapts to the intra-layer shared attention distribution of distant tokens. 3) We conduct extensive experiments to demonstrate how PoD works, and plan to open source our code and models in the future.

[0034] Figure 1 The experimental results of our customized LLaMA3-8B-32K model are shown. (a) An example of StreamingLLM prediction failure. (b) Some experimental results of LongBench and LEval. (c) The test data includes 1000 32K length text samples sampled from Dolma. For each sample sequence, we compared the last 100 tokens. (d) We sampled 1000 32K length sequences from Dolma and extracted the attention scores of the last 16 tokens in each sample.

[0035] 2 Methods Our method consists of three key steps ( Figure 4 ): 1) Analyze the similarity of attention scores between layers in a given long-context LLM and group consecutive similar layers into blocks; 2) Apply attention sharing in each block and post-train the LLM; 3) Use the post-trained LLM for efficient inference.

[0036] 2.1 Exploration of Offline Inter-LayerAttention Sharing To evaluate the similarity of attention scores between layers, we input several tokens to LLM and collect a series of attention scores. Subsequently, we calculate the attention similarity between layers and group consecutive similar layers into blocks to prepare for improving the inference efficiency of LLM.

[0037] Attention scores calculation Assume that we Samples Input Model , we will collect the final The attention score of a token on the previous position token. Mathematically, we can get: in, Respectively represent the number of layers and attention heads in the model, represents the hth attention head The attention scores collected by the layer.

[0038] Attention similarity evaluation Next, we need to evaluate the similarity between layers based on the collected attention scores. For any two different layers and , the attention similarity between them at the hth head is defined as the average Jensen-Shannon (JS) divergence of the last q tokens in all N samples. Formally, express The jth row of .

[0039] Layer grouping After calculating the head attention similarity between layers, we group consecutive similar layers into head blocks for future use. Our grouping strategy is based on the idea that any two layers in the same block should be similar enough. Specifically, when hour, and are considered similar, where is a hyperparameter. On this basis, we adopt a bottom-up greedy algorithm to iteratively merge consecutive similar layers into blocks, see Figure 3 As shown (Greedy Layer Grouping Algorithm: Greedy Layer Grouping Algorithm). Figure 4 ① in the figure shows a simple example of dividing the heads of each layer based on the similarity of attention scores between layers.

[0040] Figure 4 The three stages of PoD are shown in Figure 1: An example of grouping layers based on the similarity of the attention scores calculated offline for each head. Consecutive layers with the same color belong to the same block and therefore use the same attention score. 2: Example model structure of the attention module of head 3 in PoD. Each token attends to two groups: adjacent tokens (adjacent n r token and initial n s token, in this case n r = 2, n s = 1) and distant tokens (the rest of the tokens). The attention scores of the two groups are calculated separately and weighted according to their respective attention scores. ③: Example of KV cache changes for the 7th token and avoiding attention to distant tokens.

[0041] 2.2 Lightweight Training Adaptation To enhance the model’s ability to share attention across layers within the same block, we first introduce the post-training process required to adapt the model to this mechanism.

[0042] Attention sharing within each block We represent long input sequences as In the LLM based on autoregressive transformer, for the Any token of the layer , it will focus on all the previous tokens As mentioned before, our goal is to optimize the memory and computational cost of distant tokens by sharing the attention layer. Therefore, we first divide the previous tokens into two categories: proximal tokens and distant tokens. According to previous research, considering the "attention pooling" phenomenon, we classify several initial tokens as proximal tokens. Then, token Both groups of tokens will be attended to simultaneously, but the attention scores of distant tokens will utilize the attention scores of the lowest level of the corresponding blocks grouped in the aforementioned embodiment.

[0043] Mathematically, for any attention head, let Respectively represent The query, key, and value status of a layer. Belongs to a block consisting of consecutive layers. For token , the attention output of close tokens and distant tokens is calculated as follows: in express The i-th row of Indicates from The rows from row i to row j, including the boundaries. In addition, (starting size) and (Nearest Size) represents the number of initial tokens and adjacent tokens in the similar tokens. is the attention inner product of similar tokens, It is for similar tokens The attention output for distant tokens is similar.

[0044] Aggregation of attention output to proximal and distant tokens A parameter-free gating mechanism can integrate the attention of proximal and distant tokens in the following way: Figure 4 ② in the figure is an example of parallel training using attention masks.

[0045] 2.3 Efficient Inference Next, we discuss strategies for optimizing memory usage and reducing distant token computation by using layer sharing in the attention mechanism during inference of large language models with long context.

[0046] KV cache memory footprint optimization As shown in Equation 3, for distant tokens, query and key states are shared across layers in the same block. During inference, there is no need to cache query states because they are not reused. Therefore, our approach will reduce the memory consumption of key states in the KV cache. Figure 4 ③ in Figure 3 shows an example of KV cache changes. In this case, the layers that share attention scores during decoding only keep distant tokens once; for example, only layer 1 keeps the key states of x2 and x3, while layers 2 and 3 do not.

[0047] Computation optimization for distant tokens Empirical evidence shows that in many cases, predicting the next token can be done efficiently without paying attention to distant tokens. This is reflected in Equation 4, where is close to 1 in many cases. ≥ τ (0 ≤ τ ≤ 1 is a hyperparameter), the attention calculation for distant tokens can be omitted, thereby reducing the calculation of distant tokens.

[0048] 3 Experiments In this section, we explore two key questions: 1) whether PoD leads to performance degradation, and 2) whether PoD can improve the efficiency of long-context reasoning.

[0049] Implementation details In terms of data, we post-trained on a total of 5B tokens extracted from Dolma (Soldaini et al., 2024), ensuring that the total number of tokens in each length interval remained consistent (GLM et al., 2024). For the model, we first initialized with LLaMA3-8B and then post-trained on 5B tokens with a maximum sequence length of 32K, resulting in LLaMA3-8B 32K. Subsequently, we initialized with LLaMA3-8B-32K and continued post-training on 5B tokens with a maximum sequence length of 32K, resulting in n s = 16 and n r = 4080 PoD model.

[0050] The layer similarity threshold δ is set to 0.5, which is equivalent to saving 35% of the KV cache state. During training, the batch size is set to 4 million tokens, the learning rate is 1e-5, and the learning rate scheduler is cosine. In addition, the cardinality in RoPE (Rotation Position Embedding) is increased to 16 million+, the same as the method of Xiong et al. For code execution, we used HuggingFace and DeepSpeed, combined with ZeRO-3 and Ulysses sequence parallelism, and adopted the efficient FlexAttention in PyTorch.

[0051] Baselines We mainly considered the following baselines: 1) Windowed Attention (Beltagy et al., 2020): Each token only attends to a limited number of neighboring tokens. 2) Windowed Attention + CPT (Continuous Pre-training): Continue pre-training LLaMA3-8B-32K on the same 5B tokens while performing windowed attention. 3) StreamingLLM (Xiao et al., 2024): In addition to neighboring tokens, each token also attends to the first few tokens. 4) LM-Infinite (Han et al., 2024): Each token attends to the same tokens as in StreamingLLM, but with different position embeddings. 5) H2O: Each token not only attends to neighboring tokens, but also dynamically adds important tokens and removes unimportant tokens according to the attention scores during decoding.

[0052] 3.1 Performance Evaluation To evaluate the performance of PoD, we conducted experiments in two areas: 1) NeedleInAHaystack and 2) Practical Long ContextBenchmarks.

[0053] Figure 5 The search results for the needle in a haystack test are shown.

[0054] Needle in a haystack test This task places a random sentence in the middle of a long context window and requires the model to retrieve the sentence. Figure 5 The search results of different methods are shown. We find that StreamingLLM and H2O fail when the needle is outside the predefined window. In contrast, our method avoids token loss and performs similarly to dense models, finding almost all needles.

[0055] Long-context benchmarks To ensure that PoD can handle real-world tasks, we evaluate it on two well-known long-context benchmarks: LongBench (English version) (Bai et al., 2024) and LEval (An et al., 2024). We tested it on 14 datasets in LongBench, covering single-document quality assurance, multi-document quality assurance, summarization, few-shot learning, and code completion tasks. LEva consists of 20 subtasks divided into two groups: closed-domain group and open-domain group. The closed-domain group mainly evaluates reasoning and comprehension abilities under longer contexts, while the open-domain group focuses on tasks such as summarization and question answering, which require summarizing information from long documents.

[0056] Figure 6 The evaluation results of different methods on two well-known long-context benchmarks are shown. Figure 6 As shown in the figure, to ensure fairness, all baseline attention mechanisms have the same window size (4K). For PoD, we also ensure that the number of similar tokens that each token pays attention to is consistent with the window size. First, PoD outperforms various baselines that only pay attention to tokens within a fixed window on LongBench and LEval, indicating that our approach of not losing tokens is indeed effective. Second, PoD's performance is comparable to that of standard dense models, indicating that our approach has little impact on the overall performance.

[0057] 3.2 Efficiency Evaluation Figure 7 Theoretical and actual memory footprint savings are shown.

[0058] Memory footprint The savings in memory consumption can be analyzed from both theoretical and empirical perspectives. Theoretically, we can calculate the potential reduction in KV cache size based on the layer sharing results obtained from offline analysis. Empirically, we can perform an end-to-end evaluation to assess the actual memory savings. Following the methods of FlexGen and LCKV, for a prompt of length x, we let the model generate y tokens. The larger b is, the more memory efficient the model is. Figure 7 The memory consumption results are shown in Figure 2. We observe that the maximum batch size of PoD increases by more than 30% with different input text lengths, which is very close to our theoretical KV cache saving rate of 35%, indicating that PoD effectively reduces memory usage.

[0059] Calculation of distant tokens As mentioned above, the attention output of a particular layer comes from the weighted output of nearby and distant tokens. The weight coefficients can be obtained in advance using the shared attention scores of the lower layer before applying attention to irrelevant tokens. Therefore, when the weight coefficient of nearby tokens is greater than τ, we consider that distant tokens are meaningless for predicting the next token within the current decoding time of this layer, and thus skip their calculation.

[0060] Figure 8 The relationship between the computation saving rate and the performance loss rate and the gate threshold τ mentioned in the above embodiment is shown.

[0061] Figure 8 The figure shows the relationship between the computational savings and performance loss of LEval and the value of τ. We observe that as τ decreases, it becomes easier to ignore the computation of distant tokens, which saves more computation but also brings some performance loss. However, when τ < 0.7, the performance degradation slows down and the computational savings become more obvious. Specifically, when τ = 0.7, the computational cost is reduced by 25%, while the performance is only reduced by 5%.

[0062] Fig. 9 The relationship between model performance on LEval and two factors is shown: the number of similar tokens and the KV cache saving rate.

[0063] 3.3 Additional analysis Starting from LLaMA3-8B-32K initialization, we continue training with 2B data to explore two key hyperparameters in PoD: the number of similar tokens and the degree of sharing of attention scores between layers (reflected in the KV cache saving rate).

[0064] The relationship between model performance and the number of similar tokens is as follows: Fig. 9 As shown on the left side of , the performance of PoD steadily improves with the increase of the number of similar tokens. When the number of tokens reaches 4K, the performance of training with 2B data is acceptable compared to LLaMA3-8B-32K trained with 5B data. Considering the trade-off between performance and efficiency, we finally choose to use 4K similar tokens.

[0065] The relationship between model performance and KV cache saving rate is as follows: Fig. 9 As shown on the right side of , the performance of PoD decreases as the KV cache saving rate increases. Considering the balance between performance and efficiency, we finally choose to compress the KV cache to 35%. This ensures that PoD can achieve comparable performance to LLaMA3-8B-32K using the same training data.

[0066] Case studies in Fig.10 In Figure 2, we provide four different representative cases to compare StreamingLLM, H2O, and our PoD. For case (a), the correct answer is located within a window of adjacent tokens, allowing all methods to attend to it during decoding. Therefore, all three methods can make accurate predictions. For case (b), the answer is located at the beginning, and both StreamingLLM and PoD can notice it during decoding, resulting in the correct prediction. However, for H2O, a long string of irrelevant tokens after the answer causes it to mistakenly abandon the initial answer, resulting in a wrong prediction. For case (c), the answer is located in the middle, and StreamingLLM cannot notice it during decoding, resulting in a wrong prediction. However, the text after the answer and before the question is relevant to the answer, so H2O can keep it in the attention window and make a correct prediction. PoD is able to attend to all tokens and also makes a correct prediction. For case (d), only PoD is able to answer the example in the needle-in-the-haystack test, because the other two methods ignore the answer token.

[0067] Fig.10 A case study of different methods is shown in Figure 1. Here, s∗n represents a string s repeated n times. + represents the concatenation of strings.

[0068] 4 Related Work Long-context LLMs face huge challenges in terms of memory consumption and latency during inference, as they need to handle a large number of parameters and long sequences. From an optimization perspective, we roughly classify these methods into the following categories: Reducing context computation Recent studies have optimized context computation from two directions, namely pre-filling and decoding. In the pre-filling direction, SnapKV (Li et al., 2024), PyramidKV (Zhang et al., 2024), and LazyLLM (Fu et al., 2024) aim to reduce the size of the generated KV cache by selectively caching only important input tokens. At the same time, MInference (Jiang et al., 2024a) and RetrievalAttention (Liu et al., 2024c) minimize pre-filling delays by leveraging the sparse retention mechanism inherent in the transformer. Some methods also improve efficiency by directly compressing the length of input prompts. The scheme of the embodiments of the present application is the opposite of these methods, focusing mainly on optimizing the decoding stage.

[0069] In the decoding stage, H2O (Zhang et al., 2023) removes unimportant tokens from the KV cache based on the attention score, while LM-Infite (Han et al., 2024) and StreamingLLM (Xiao et al., 2024) only retain the most recent token and a few initial tokens at each decoding step. While all of these methods can improve efficiency, they may also discard important tokens required for future text generation. The scheme of the embodiment of the present application aims to improve efficiency while ensuring that no token is ignored.

[0070] Hidden state dimension reduction and quantization We also proposed another approach that involves reducing hidden size and quantization, although these techniques are orthogonal to the focus of our work. On the one hand, MQA (Shazeer, 2019) and GQA (Ainslie et al. MLA (Liu et al., 2024a) compress a pair of key-value states into a low-rank latent vector. On the other hand, AWQ (Lin et al., 2024) and QLLM (Liu et al., 2024d) convert model weights and activations into a low-bitwidth format, thereby reducing memory usage and computational overhead.

[0071] Reducing redundancy between layers Another related research aims to improve efficiency by reducing redundancy between layers in transformers. LCKV (Wu&Tu, 2024), CLA (Brandon et al., 2024) and MiniCache (Liu et al., 2024b) exploit the similarity of key values ​​between layers by sharing key-value states across layers. Compared with these methods, our method has two significant features: 1) we exploit the similarity of attention scores between transformer layers and extend the phenomenon to LLMs; 2) unlike their sharing strategies involving adjacent layers or the final layer, our sharing strategy starts from scratch and is derived from the search process.

[0072] 5 Conclusion In the scheme of the embodiment of the present application, we intend to improve the efficiency of LLM. Previous window-based work has suffered from performance degradation due to token loss. Based on the observation that close tokens are more important than distant tokens, we proposed PoD, which allocates fewer resources by sharing attention to distant tokens between similar layers. Evaluation shows that PoD can save 35% of KV cache without affecting model performance. This approach not only optimizes resource allocation, but also provides a way to improve LLM efficiency in the future.

[0073] In some other embodiments, the embodiments of the present invention further provide a non-volatile computer storage medium, the computer storage medium stores computer executable instructions, the computer executable instructions can execute the method for improving the inference efficiency of a long-context large language model in any of the above method embodiments, and a system for improving the inference efficiency of a long-context large language model, wherein the system for improving the inference efficiency of a long-context large language model includes an encoder module, a dual-path module, a decoder module and a merging module, and the decoder module includes an excitation decoder, a channel decoder and a complex spectrum decoder; As an implementation mode, the non-volatile computer storage medium of the present invention stores computer executable instructions, and the computer executable instructions are configured as follows: Analyze the similarity of attention scores between layers in a given large language model and group consecutive similar layers into the same block; Applying shared attention in each block and post-training the large language model; Efficient inference using large post-trained language models.

[0074] The non-volatile computer-readable storage medium may include a program storage area and a data storage area, wherein the program storage area may store an operating system, an application required by at least one function; the data storage area may store data created according to the use of the method and system for improving the efficiency of reasoning of a long-context large language model, etc. In addition, the non-volatile computer-readable storage medium may include a high-speed random access memory, and may also include a non-volatile memory, such as at least one disk storage device, a flash memory device, or other non-volatile solid-state storage device. In some embodiments, the non-volatile computer-readable storage medium may optionally include a memory remotely arranged relative to the processor, and these remote memories may be connected to the method for improving the efficiency of reasoning of a long-context large language model via a network. Examples of the above-mentioned network include, but are not limited to, the Internet, an intranet, a local area network, a mobile communication network, and combinations thereof.

[0075] An embodiment of the present invention also provides a computer program product, which includes a computer program stored on a non-volatile computer-readable storage medium, and the computer program includes program instructions. When the program instructions are executed by a computer, the computer executes any of the above methods for improving the reasoning efficiency of long-context large language models.

[0076] Fig.11 is a schematic diagram of the structure of an electronic device provided by an embodiment of the present invention, such as Fig.11 As shown, the device includes: one or more processors 1110 and a memory 1120, Fig.11 A processor 1110 is taken as an example. The method and system for improving the efficiency of long-context large language model reasoning may also include: an input device 1130 and an output device 1140. The processor 1110, the memory 1120, the input device 1130 and the output device 1140 may be connected via a bus or other means. Fig.11 In the example, the connection via a bus is taken. The memory 1120 is the non-volatile computer-readable storage medium mentioned above. The processor 1110 executes various functional applications and data processing of the server by running the non-volatile software programs, instructions and modules stored in the memory 1120, that is, the method and system for improving the reasoning efficiency of long-context large-scale language models in the above-mentioned method embodiment are implemented. The input device 1130 can receive input digital or character information, and generate key signal input related to user settings and function controls of the device for improving the reasoning efficiency of long-context large-scale language models. The output device 1140 may include display devices such as display screens.

[0077] The above product can execute the method provided by the embodiment of the present invention, and has the functional modules and beneficial effects corresponding to the execution method. For technical details not described in detail in this embodiment, please refer to the method provided by the embodiment of the present invention.

[0078] As an implementation mode, the electronic device is applied to a device for improving the inference efficiency of a long-context large language model, comprising: at least one processor; and a memory connected to the at least one processor in communication; wherein the memory stores instructions executable by the at least one processor, and the instructions are executed by the at least one processor so that the at least one processor can: Analyze the similarity of attention scores between layers in a given large language model and group consecutive similar layers into the same block; Applying shared attention in each block and post-training the large language model; Efficient inference using large post-trained language models.

[0079] The electronic device of the embodiment of the present application exists in various forms, including but not limited to: (1) Mobile communication equipment: This type of equipment is characterized by having mobile communication functions and its main purpose is to provide voice and data communications. This type of terminal includes: smart phones, multimedia phones, functional phones, and low-end phones.

[0080] (2) Ultra-mobile personal computer devices: These devices fall into the category of personal computers, have computing and processing functions, and generally also have mobile Internet access features. These terminals include: PDA, MID and UMPC devices, etc.

[0081] (3) Portable entertainment devices: These devices can display and play multimedia content. They include audio and video players, handheld game consoles, e-books, smart toys, and portable car navigation devices.

[0082] (4) Server: A device that provides computing services. The server consists of a processor, hard disk, memory, system bus, etc. The server is similar to a general computer architecture, but because it needs to provide highly reliable services, it has higher requirements in terms of processing power, stability, reliability, security, scalability, and manageability.

[0083] (5) Other electronic devices with data interaction functions.

[0084] The device embodiments described above are merely illustrative, wherein the units described as separate components may or may not be physically separated, and the components shown as units may or may not be physical units, i.e., they may be located in one place, or they may be distributed on multiple network units. Some or all of the modules may be selected according to actual needs to achieve the purpose of the scheme of this embodiment. Those of ordinary skill in the art may understand and implement it without creative effort.

[0085] Through the description of the above implementation modes, those skilled in the art can clearly understand that each implementation mode can be implemented by means of software plus a necessary general hardware platform, or of course by hardware. Based on such an understanding, the above technical solution can essentially or in other words be embodied in the form of a software product that contributes to the prior art. The computer software product can be stored in a computer-readable storage medium, such as ROM / RAM, a magnetic disk, an optical disk, etc., and includes a number of instructions for a computer device (which can be a personal computer, a server, or a network device, etc.) to execute the methods of each embodiment or some parts of the embodiment.

[0086] Finally, it should be noted that the above embodiments are only used to illustrate the technical solutions of the present invention, rather than to limit it. Although the present invention has been described in detail with reference to the aforementioned embodiments, those skilled in the art should understand that they can still modify the technical solutions described in the aforementioned embodiments, or make equivalent replacements for some of the technical features therein. However, these modifications or replacements do not deviate the essence of the corresponding technical solutions from the spirit and scope of the technical solutions of the embodiments of the present invention.

Claims

1. A method for improving the inference efficiency of a large language model with a long context, comprising: Analyze the similarity of attention scores between layers in a given large language model and group consecutive similar layers into the same block; Applying shared attention in each block and post-training the large language model; Efficient inference using large post-trained language models.

2. The method according to claim 1, wherein: The method of analyzing the similarity of attention scores between layers in a given large language model and classifying consecutive similar layers into blocks includes: Calculate the attention scores of each layer of each attention head; Calculating attention similarity between layers based on the attention scores; If the similarity between two layers is greater than or equal to a preset threshold, the two layers are determined to be similar layers; Consecutive similar layers are grouped into a head block.

3. The method according to claim 2, wherein: The grouping of consecutive similar layers into a header block comprises: A bottom-up greedy algorithm is used to iteratively merge consecutive similar layers into a head block.

4. The method according to claim 3, wherein: The applying attention sharing in each block and post-training the large language model includes: For each unit, calculate the attention outputs of the close units and the distant units, wherein the attention score of the distant unit is represented by the attention score of the bottom layer of the corresponding head block after grouping, the close units include the adjacent units and the initial units, and the distant units include other units except the adjacent units and the initial units; Aggregate the attention outputs of the near unit and the far unit.

5. The method according to claim 4, wherein: The aggregating the attention output of the near unit and the far unit includes: A parameter-free gating mechanism is used to integrate the attention of the nearby units and the distant units.

6. The method according to claim 1, wherein: The method of using a post-trained large language model for efficient reasoning includes: The memory usage of the KV cache and the computation of the distant units are optimized by using the shared attention score in the attention mechanism.

7. An electronic device, comprising: At least one processor, and a memory communicatively connected to the at least one processor, wherein the memory stores instructions executable by the at least one processor, and the instructions are executed by the at least one processor to enable the at least one processor to perform the steps of the method described in any one of claims 1 to 6.

8. A storage medium having a computer program stored thereon, characterized in that: When the computer program is executed by a processor, the steps of the method according to any one of claims 1 to 6 are implemented.

Images