Inference model key value cache compression method based on supervised training and reinforcement learning
By using supervised training and reinforcement learning methods to dynamically adjust the key-value cache eviction strategy, the problem of capturing the importance of key-value pairs in existing technologies is solved, resulting in better model performance and lower computational and storage overhead in long inference tasks.
Patent Information
- Authority / Receiving Office
- CN · China
- Patent Type
- Applications(China)
- Current Assignee / Owner
- Filing Date
- 2025-11-04
- Publication Date
- 2026-03-10
AI Technical Summary
Existing key-value caching compression methods struggle to dynamically capture the importance of key-value pairs while covering various complex features of the attention head, leading to reduced model performance.
We employ a supervised training and reinforcement learning approach, establishing a dynamic key-value cache eviction mechanism. We use an MLP scorer to evaluate the importance of key-value pairs and optimize the eviction strategy through a reinforcement learning algorithm. We model key-value cache eviction as a Markov decision process, dynamically adjusting the cache size to maintain key-value pairs.
It achieves better model performance in long inference tasks, reduces computational and storage overhead, and maintains the model's sample capacity with low training overhead.
Smart Images
Figure CN121636153A_ABST
Abstract
Description
Technical Field
[0002] This invention relates to the fields of artificial intelligence and natural language processing, and in particular to a method for compressing key-value caches for inference models based on supervised training and reinforcement learning. Background Technology
[0004] When solving complex reasoning tasks involving generation of long key-value pairs, large language models often exhibit high memory usage and computational overhead due to the existence of long key-value caches. To address this issue, existing methods reduce this overhead by evicting unimportant key-value pairs. These methods are mainly divided into training-free methods and training methods. Training-free methods typically rely on predefined heuristics, using attention scores, positional information, and information about the key-value pairs themselves to score each pair and determine which pairs need to be evicted based on the scores. Training methods, on the other hand, use a lightweight scorer to score each key-value pair and decide whether to retain or evict it.
[0005] However, existing key-value caching compression methods still have some limitations. Training-free methods dynamically score key-value pairs based on predefined heuristic rules; however, the importance of each key-value pair often exhibits diverse and complex characteristics across different attention heads and pruning pairs. Heuristic rules alone cannot cover all cases, leading to reduced model performance. Training-based methods, primarily designed for long input scenarios, typically only determine the importance of each key-value pair and immediately discard them, failing to capture the importance of each key-value pair at different positions in the generation of long sequences. Therefore, a method is needed that can capture the dynamic importance of key-value pairs while covering the various complex features of attention heads.
[0006] The information disclosed in this background section is intended only to enhance the understanding of the general background of the invention and should not be construed as an admission or in any way implying that the information constitutes prior art known to those skilled in the art. Summary of the Invention
[0008] The purpose of this invention is to provide a key-value cache compression method for inference models based on supervised training and reinforcement learning, so as to solve the technical problems existing in the prior art.
[0009] To achieve the above objectives, the present invention adopts the following technical solution:
[0010] A method for compressing key-value caches inference models based on supervised training and reinforcement learning includes the following steps:
[0011] S1. Establish a dynamic key-value cache eviction mechanism in long text generation: use a key-value pair scorer to determine the importance of each key-value pair in the reasoning process, and select important key-value pairs to retain at each step;
[0012] S2. Supervised training: Enables the scoring model to identify key key-value pairs;
[0013] S3. Reinforcement Learning: The key-value cache eviction process in the inference process is modeled as a Markov decision process, and a reinforcement learning algorithm is used to optimize the eviction strategy.
[0014] Furthermore, the specific implementation process of step S1 is as follows:
[0015] Within the framework, a key-value cache dynamic compression process is defined, controlled by two hyperparameters:
[0016] Cache budget B: Specifies the maximum number of key-value pairs that need to be retained;
[0017] Eviction length L: determines the frequency at which cache eviction is triggered;
[0018] During the generation process, whenever L new tokens are generated, the cache size will grow to B+L. At this point, the latest L key-value pairs are retained, and the most important BL key-value pairs are selected from the remaining cache through an eviction algorithm, thereby restoring the total cache size to the budget B.
[0019] Specifically, in each elimination step, a scoring model is used to evaluate the importance of each key-value pair; to strike a balance between efficiency and performance, an MLP scorer is employed. To predict the importance score of the nth key-value pair Input features Construct a vector consisting of the following parts: key ,value and their corresponding attention features The attention feature is a fixed-length vector derived from the attention score.
[0020] ;
[0021] Based on importance scores, the top 2L key-value pairs with the highest scores are first selected through Top-K sampling; then, L candidates are selected from these 2L candidates through multinomial sampling; finally, the remaining L relatively unimportant key-value pairs are discarded, and the rest are retained for subsequent attention calculations.
[0022] Furthermore, the specific implementation process of step S2 is as follows:
[0023] In the two-stage framework, the first stage is supervised training, which aims to enable the scoring model to identify key key-value pairs.
[0024] To predict future importance, a golden elimination algorithm is first proposed to construct target labels and elimination trajectories. Attention dimensions are aggregated based on the elimination steps and group sizes in grouped attention. Specifically, for a complete inference trajectory, the corresponding attention score matrix is calculated on the original model. The attention matrix is then divided into blocks along the query dimension with a step size L; each block starts at the first elimination position B+L, and the last block is padded to match the maximum length; within each block, pooling is performed along the query dimension to obtain a block score for each key-value pair. :
[0025]
[0026] in, , indicating the total number of elimination steps; in grouped query attention, each g attention head shares the same key-value cache, and these scores are further pooled within each attention group to achieve unified elimination at the group level; to simplify calculation, average pooling is used in both pooling operations:
[0027]
[0028] In each elimination step t, the goal is to minimize the impact on future attention computation; it is desired that the eliminated key-value pairs have low attention scores in all subsequent blocks; therefore, the maximum block score for each key-value pair across all future blocks is calculated as its future score. :
[0029]
[0030] After obtaining the key-value cache elimination trajectory generated by the golden elimination algorithm, a scoring model is used to calculate a score for each key-value pair in each elimination step. Next, the elimination process is modeled as a ranking task, and a pairwise ranking loss is used to train the scoring model, ensuring that the model's predicted score ranking is the reverse of the actual future score order. Formally, the loss function is defined as follows:
[0031]
[0032] Where m is the hyperparameter of the loss.
[0033] Furthermore, the specific implementation process of step S3 is as follows:
[0034] Given a complete inference sequence and a large language model with frozen parameters, the goal is to optimize the set of scoring models. To maximize the overall reward; below, we define the components of the reinforcement learning framework:
[0035] state : Indicates the remaining key-value cache content at step t;
[0036] action In state Next, action This indicates that a subset of key-value pairs' indices will be selected from the current cache to be retained, in order to meet the preset cache budget;
[0037] Strategy For each attention group, the scoring model is treated as a policy model, which is used to assign a score to each key-value pair to determine whether it should be retained or discarded.
[0038] award To evaluate the quality of each elimination action, a sequence-level reward function is defined; a subset E of tokens is defined, satisfying the following two conditions: 1) The original entropy of the token is located in the 80% of the lowest entropy values in the sequence, denoted as . ;2) Its incremental loss Exceeding the threshold Thus, we obtain the set: The reward is defined as the negative average of the squared loss increments of all tokens in set E, i.e.:
[0039]
[0040] After defining the Markov decision process for key-value cache eviction, the GRPO algorithm is used to train the scoring model; policy model Initialize the model obtained during the self-supervised training phase; this model also serves as the reference policy. In each training step, for a given sequence, the old strategy is used. Top-K multinomial sampling is used to generate key-value cache elimination trajectories. Since different trajectories retain different key-value caches, the same token will generate different hidden states in different trajectories, resulting in different rewards for the entire sequence. Group relative normalization is then used to calculate the estimated advantage values of these trajectories.
[0041]
[0042] Next, these advantage values are broadcast to all layers and the head elimination process, and all scoring models are jointly optimized; the training objective function is defined as follows:
[0043] in, To truncate the threshold, and To control the weight of the KL regularization term.
[0044] By adopting the above technical solution, the present invention has the following beneficial effects:
[0045] Existing key-value cache eviction algorithms are generally divided into two categories: heuristic methods that do not require training and static eviction methods based on training. While the former can capture the importance of key values as the sequence length changes, it often struggles to cover all key-value characteristics. The latter learns to evaluate the importance of key-value pairs, but because it can only perform a static estimation and eviction once, it struggles to capture the changes in importance over time within the sequence, and also incurs significant training overhead. To address these shortcomings, this patent proposes a training method that can dynamically capture the importance of key-value caches at different positions, requiring only fine-tuning of the key-value pair scorer, resulting in lower training overhead. Furthermore, this patent is the first to model the key-value cache eviction problem as a Markov decision process and introduces reinforcement learning into the field of key-value cache compression. Specifically, this patent treats the currently retained key-value cache as a state, discarding the key-value cache as an action, and the model's loss on the sequence as a reward. Optimization is performed using the GRPO algorithm to maximize the reward, enabling the large model and the key-value cache eviction strategy to adapt to each other, thereby achieving better performance on long inference tasks. Attached Figure Description
[0047] To more clearly illustrate the specific embodiments of the present invention or the technical solutions in the prior art, the drawings used in the description of the specific embodiments or the prior art will be briefly introduced below. Obviously, the drawings described below are some embodiments of the present invention. For those skilled in the art, other drawings can be obtained from these drawings without creative effort.
[0048] Figure 1 This is a system architecture diagram of the inference model key-value caching compression method based on supervised training and reinforcement learning provided in an embodiment of the present invention. Detailed Implementation
[0050] The technical solution of the present invention will now be clearly and completely described with reference to the accompanying drawings. Obviously, the described embodiments are only some, not all, of the embodiments of the present invention. Based on the embodiments of the present invention, all other embodiments obtained by those skilled in the art without creative effort are within the scope of protection of the present invention.
[0051] The specific embodiments of the present invention will be described in detail below with reference to the accompanying drawings. It should be understood that the specific embodiments described herein are for illustration and explanation only and are not intended to limit the present invention.
[0052] like Figure 1 As shown, this invention proposes a key-value cache compression method for inference models based on supervised training and reinforcement learning. This patent aims to reduce the computational and storage overhead of the model during long sequence generation by utilizing a key-value cache eviction method, while better preserving the model's sample capabilities. To this end, this application proposes a training-based key-value cache eviction framework, which uses a key-value pair scorer to determine the importance of each key-value pair in the inference process and selects important key-value pairs for retention at each step. Furthermore, this application, based on supervised training and reinforcement learning, can efficiently balance model performance and efficiency simply by training the key-value pair scorer.
[0053] The key-value cache compression method for inference models based on supervised training and reinforcement learning in this application specifically includes:
[0054] 1. Dynamic key-value caching and eviction mechanism in long text generation:
[0055] Within this framework, this application defines a key-value cache dynamic compression process, controlled by two hyperparameters:
[0056] Cache budget B: Specifies the maximum number of key-value pairs that need to be retained;
[0057] Eviction length L: Determines the frequency at which cache eviction is triggered.
[0058] During the generation process, the cache size grows to B+L whenever L new tokens are generated. At this point, this application retains the latest L key-value pairs and uses an eviction algorithm to select the most important BL key-value pairs from the remaining cache, thereby restoring the total cache size to the budgeted B.
[0059] Specifically, in each elimination step, this application uses a scoring model to evaluate the importance of each key-value pair. To strike a balance between efficiency and performance, this application employs an MLP scorer. To predict the importance score of the nth key-value pair This application will input features. Construct a vector consisting of the following parts: key ,value and their corresponding attention features The attention feature is a fixed-length vector derived from the attention score.
[0060]
[0061] Based on importance scores, this application first selects the top 2L key-value pairs with the highest scores using Top-K sampling; then, it selects L candidates from these 2L candidates using multinomial sampling. Finally, this application discards the remaining L relatively unimportant key-value pairs, retaining the rest for subsequent attention calculations.
[0062] 2. Supervised training:
[0063] In the two-stage framework of this application, a supervised training phase is first conducted to enable the scoring model to identify key key-value pairs. Furthermore, reinforcement learning is used for optimization.
[0064] To predict future importance, this application first proposes a golden elimination algorithm for constructing target labels and elimination trajectories. This application aggregates the attention dimension based on the elimination steps and the group size in grouped attention. Specifically, for a complete inference trajectory, this application calculates the corresponding attention score matrix on the original model. The attention matrix is then divided into blocks along the query dimension with a step size L. Each block starts at the first elimination position B+L, and the last block is padded to match the maximum length. Within each block, pooling is performed along the query dimension to obtain a block score for each key-value pair. :
[0065]
[0066] in, , representing the total number of elimination steps. In grouped query attention, each g attention head shares the same key-value cache. This application further pools these scores within each attention group to achieve unified elimination at the group level. To simplify computation, this application uses average pooling in both pooling operations:
[0067]
[0068] In each elimination step t, the objective of this application is to minimize the impact on future attention computation. In other words, this application aims to ensure that the eliminated key-value pairs have low attention scores in all subsequent blocks. Therefore, this application calculates the maximum block score for each key-value pair across all future blocks as its future score. :
[0069]
[0070] After obtaining the key-value cache elimination trajectory generated by the golden elimination algorithm, this application utilizes a scoring model to calculate a score for each key-value pair in each elimination step. Next, this application models the elimination process as a ranking task and uses pairwise ranking loss to train the scoring model, ensuring that the model's predicted score ranking is the reverse of the actual future score order. Formally, the loss function is defined as follows (where m is a hyperparameter of the loss):
[0071]
[0072] 3. Enhance learning:
[0073] In long context decoding, retained key-value caches not only affect the model's current generation results but also subsequent cache eviction decisions. Therefore, this application models key-value cache eviction during inference as a Markov decision process and employs reinforcement learning algorithms to optimize the eviction strategy.
[0074] Given a complete inference sequence and a large language model with frozen parameters, the goal of this application is to optimize the scoring model set. To maximize overall reward, the main components of the reinforcement learning framework are defined below:
[0075] state : Indicates the remaining key-value cache content at step t.
[0076] action In state Next, action This indicates that a subset of key-value pairs' indices will be selected from the current cache to be retained, in order to meet the preset cache budget.
[0077] Strategy For each attention group, this application treats the scoring model as a policy model to assign a score to each key-value pair, thereby determining whether it should be retained or discarded.
[0078] award To evaluate the quality of each elimination action, this application defines a sequence-level reward function. This application defines a subset E of tokens that satisfies the following two conditions: (1) The original entropy of the token is located in the 80% of the sequence with the lowest entropy values (denoted as E). (2) Its loss increment Exceeding the threshold Therefore, this application yields the following set: This application defines the reward as the negative average of the squared loss increments of all tokens in set E, i.e.:
[0079]
[0080] After defining the Markov decision process for key-value cache eviction, this application employs the GRPO algorithm to train the scoring model. Policy Model Initialize the model obtained during the self-supervised training phase; this model also serves as the reference policy. In each training step, for a given sequence, this application uses the old strategy. G key-value cache elimination trajectories are generated through Top-K multinomial sampling. Since different trajectories retain different key-value caches, the same token will generate different hidden states in different trajectories, resulting in different rewards for the entire sequence. This application then uses group relative normalization to calculate the estimated advantage values of these trajectories:
[0081]
[0082] Next, this application broadcasts these advantage values to all layers and the head elimination process, and jointly optimizes all scoring models. The training objective function is defined as follows:
[0083] in, To truncate the threshold, and Control the weights of the KL regularization terms.
[0084] Finally, it should be noted that the above embodiments are only used to illustrate the technical solutions of the present invention, and not to limit them; although the present invention has been described in detail with reference to the foregoing embodiments, those skilled in the art should understand that modifications can still be made to the technical solutions described in the foregoing embodiments, or equivalent substitutions can be made to some or all of the technical features; and these modifications or substitutions do not cause the essence of the corresponding technical solutions to deviate from the scope of the technical solutions of the embodiments of the present invention.
Claims
1. An inference model key-value cache compression method based on supervised training and reinforcement learning, characterized in that, The method comprises the following steps: S1, establishing a dynamic key-value cache eviction mechanism in long text generation: judging the importance of each key-value pair in the reasoning process through a key-value pair scorer, and selecting important key-value pairs for retention at each step; S2, supervised training: enabling the scoring model to identify key key-value pairs; S3, reinforcement learning: modeling the key-value cache eviction in the reasoning process as a Markov decision process, and using a reinforcement learning algorithm to optimize the eviction strategy.
2. The inference model key-value cache compression method based on supervised training and reinforcement learning according to claim 1, characterized in that, The specific implementation process of step S1 is as follows: In the framework, a key-value cache dynamic compression process is defined, which is controlled by two hyperparameters: Cache budget B: specifies the upper limit of the number of key-value pairs to be retained; Eviction length L: determines the frequency of triggering cache eviction; During the generation process, whenever L new tokens are generated, the cache size will grow to B+L; at this time, the latest L key-value pairs are retained, and the most important B-L key-value pairs are selected from the remaining cache through the eviction algorithm, so that the total cache size is restored to the budget B; Specifically, in each elimination step, a scoring model is used to evaluate the importance of each key-value pair; to strike a balance between efficiency and performance, an MLP scorer is employed. To predict the importance score of the nth key-value pair Input features Construct a vector consisting of the following parts: key ,value and their corresponding attention features ; Where, The attention feature is a fixed-length vector converted from the attention score: ； According to the importance score, first select the top 2L key-value pairs with the highest score through Top-K sampling; Then, select L from the 2L candidates through polynomial sampling; finally, discard the remaining L relatively unimportant key-value pairs, and retain the rest for subsequent attention calculation. 3.The inference model key-value cache compression method based on supervised training and reinforcement learning of claim 1, wherein, The specific implementation process of step S2 is as follows: In the two-stage framework, first, the supervised training stage is carried out, aiming to enable the scoring model to identify key key-value pairs; To predict the future importance, a gold-knockout algorithm is first proposed to construct target labels and knockout trajectories; the attention dimension is aggregated according to the knockout step and the group size in the group attention; specifically, for a complete reasoning trajectory, the corresponding attention score matrix is calculated on the original model , and the attention matrix is blocked in the query dimension with a step length L; where the blocking starts from the first knockout position B+L, and the last block is padded to match the maximum length; in each block, a pooling operation is performed in the query dimension to obtain the block score of each key-value pair : ； where, , denotes the total number of elimination steps; in grouped query attention, each g attention heads share the same key-value cache, and the scores are further pooled within each attention group to achieve group-level uniform elimination; for simplicity, average pooling is used in both pooling operations: ； At each elimination step t, the goal is to minimize the impact on future attention computations; key-value pairs that are expected to be eliminated have low attention scores in all subsequent splits; thus, the maximum split score of each key-value pair in all future splits is computed as its future score : ； After obtaining the key-value cache eviction trajectory generated by the golden eviction algorithm, the scoring model is used to calculate the score of each key-value pair at each eviction step; then, the eviction process is modeled as a ranking task, and a pairwise ranking loss is used to train the scoring model, so that the model's predicted score ranking should be opposite to the real future score order; formally, the loss function is defined as follows: ； Where m is the hyperparameter of the loss. 4.The inference model key-value cache compression method based on supervised training and reinforcement learning of claim 1, wherein, The specific implementation process of step S3 is as follows: Given a complete sequence of reasoning and a frozen parameter large language model, the goal is to maximize the overall reward by optimizing a set of scoring models Below, the components of the reinforcement learning framework are defined: State : represents the current remaining key-value cache content at step t; Actions : In the state , the action indicates that a portion of the indexes that select a portion of the key-value pairs from the current cache are reserved to meet the preset cache budget; Strategy : For each attention group, treat the scoring model as a strategy model for assigning scores to each key-value pair, deciding whether it should be kept or culled; reward : To evaluate the quality of each pruning action, define a sequence-level reward function; define a token subset E that satisfies the following two conditions: 1) the original entropy of the token is located in the lowest 80% of the entropy values in the sequence, denoted as ; 2) the loss increment of the token exceeds the threshold ; thus obtaining the set: ; define the reward as the negative average of the squared loss increments of all tokens in the set E, that is: ； After defining the Markov decision process of key-value cache eviction, the GRPO algorithm is used to train the scoring model; the strategy model The model obtained in the self-supervised training phase is initialized, which is used as a reference strategy In each training step, for a given sequence, the old strategy Top-K polynomial sampling is used to generate key-value cache eviction trajectories; since the retained key-value caches in different trajectories are different, the same token will produce different hidden states in different trajectories, resulting in different rewards for the entire sequence; then, group relative normalization is used to calculate the estimated advantage value of these trajectories: ； Then, these advantage values are broadcast to the eviction process of all layers and heads, and all scoring models are jointly optimized; the training objective function is defined as follows: wherein, is a truncation threshold, and is a weight controlling the KL regularizer.