Adaptive KV cache management method based on deep reinforcement learning

Abstract

The invention discloses a self-adaptive KV cache management method based on deep reinforcement learning, belongs to the technical field of artificial intelligence, deep learning and reinforcement learning, and particularly relates to a method for modeling a KV cache management problem in big language model reasoning into a Markov decision process, and dynamically determining a cache expelling strategy by training a deep reinforcement learning agent. According to the method, state representation fusing multi-dimensional information is designed, an intelligent agent comprising a strategy network and a value network is constructed, a multi-target reward function for balancing reasoning quality, memory efficiency and calculation delay is defined, a near-end strategy optimization algorithm is used for training, a meta-learning mechanism is introduced to achieve rapid task adaptation, and the method is suitable for large-scale popularization and application. And an intelligent and adaptive solution is provided for efficient deployment of the large language model.

Description

Technical Field

[0001] This invention belongs to the fields of artificial intelligence, deep learning and reinforcement learning, and specifically relates to an adaptive key-value cache management method based on deep reinforcement learning. Background Technology

[0002] Large Language Models (LLMs) demonstrate superior performance in applications such as text generation and dialogue systems, but their deployment faces a significant memory bottleneck. During the generation process, to avoid redundant computation, the model requires a large amount of key-value (KV) cache, the size of which grows linearly with the sequence length and batch size. For example, a model with 30 billion parameters, a batch size of 128, and a sequence length of 1024 would require 180 GB of memory for its KV cache.

[0003] Existing KV cache management methods mainly include: (1) static sparse attention mechanisms, such as SparseTransformer, which predefines fixed attention patterns; (2) heuristic eviction strategies, such as Least Recently Used (LRU) or Least Frequently Used (LFU); and (3) observation-based methods, such as the Heavy Hitter strategy proposed by H2O, which retains tokens with high cumulative attention scores.

[0004] However, these methods have the following limitations:

[0005] 1. Lack of adaptability: These methods use fixed strategies and cannot dynamically adjust according to different input content, task types and model characteristics;

[0006] 2. Ignoring long-term benefits: Greedy local decision-making may lead to suboptimal global performance, especially in long sequence generation tasks;

[0007] 3. Difficulty in parameter tuning: For example, the proportion of heavy hitters to be retained in H2O needs to be set manually, making it difficult to find the optimal parameters for different scenarios;

[0008] 4. Lack of learning ability: Unable to learn from historical experience to improve strategies, starting from scratch every time they reason.

[0009] Reinforcement learning, as a powerful tool for solving sequential decision-making problems, can learn optimal policies through interaction with the environment. Modeling key-value cache management as a reinforcement learning problem can achieve automatic policy learning and dynamic optimization, overcoming the limitations of existing methods. Summary of the Invention

[0010] To address the technical problems existing in the above-mentioned key-value (KV) cache management methods, this invention provides an adaptive KV cache management method based on deep reinforcement learning. By modeling the KV cache eviction decision as a Markov decision process (MDP), a deep reinforcement learning agent is trained to dynamically determine which KVs should be retained or evicted, thereby achieving an optimal balance between memory efficiency and inference quality.

[0011] The technical solution adopted in this invention is as follows:

[0012] An adaptive key-value cache management method based on deep reinforcement learning includes the following steps:

[0013] S1. Environment modeling and state construction, specifically:

[0014] Collect a dataset related to the text generation task; use a target large language model to infer the meaning of each data point in the dataset, and complete environmental modeling based on the information recorded by the inference, including the state at different decoding times; define the state at time t. It includes five modal features: the attention score matrix at time t, token semantic features, location encoding, historical access frequency, and cache utilization rate;

[0015] S2, agent construction, including policy network and value network;

[0016] The policy network is state-based. As input, after feature encoding and feature fusion, the output is the probability distribution of buffer eviction actions. ;in, Indicates the state The following actions were taken; Representation strategy;

[0017] The value network is in state and cache eviction probability distribution As input, after feature encoding and feature fusion, the output is the expected cumulative value. This indicates that the value network is related to the state. Execute action The state obtained later regarding time t+1 The assessed value;

[0018] S3. Reward function design, specifically:

[0019] S3.1 Define the composite reward function:

[0020]

[0021] in, Indicates the execution of an action The immediate reward afterwards; Indicates a reward for the quality of reasoning; This indicates a memory efficiency bonus; Indicates delayed reward; , and All represent weights;

[0022] S3.2 Define the global evaluation function:

[0023]

[0024] in, Indicates the state Next action The final overall assessment score; Indicates the balance coefficient;

[0025] S4. Train the agent through deep reinforcement learning, specifically:

[0026] Initialize policy network parameters and value network parameters Based on the environment and the overall evaluation function, the proximal policy optimization (PPO) algorithm is used to train the agent and obtain the optimal policy network parameters. and optimal value network parameters This allows them to learn how to make optimal strategies based on the state to maximize the overall evaluation score, which can then be used to deploy actual key-value cache management tasks.

[0027] Furthermore, in S3.1 , and The calculation formulas are as follows:

[0028]

[0029]

[0030]

[0031] in, Indicates the execution of an action The change in confusion was subsequently caused by cache eviction. Indicates the cache usage coefficient; Indicates the time cost coefficient; This indicates the cache occupancy rate at time t; This indicates the time cost of caching operations.

[0032] Furthermore, both the policy network and the value network adopt a lightweight Transformer architecture design.

[0033] Furthermore, for the policy network of S2, during the feature encoding stage, the state is... The attention score matrix, token semantic features, location encoding, and historical access frequency are each input into an independent encoder, and a cache occupancy rate perceptron is used to process the cache occupancy rate to obtain the hidden feature vectors of each modality. In the feature fusion stage, the hidden feature vectors of each modality are concatenated and then input into the multi-head attention module to generate the fused feature representation. In the output stage, the fused feature representation is input into a fully connected layer for linear mapping to obtain the unnormalized scores (Logits) of each action. After processing by the Softmax function, the cache eviction action probability distribution is output. .

[0034] Furthermore, for the value network of S2, in the feature encoding stage, the state is... The attention score matrix, token semantic features, location encoding, and historical access frequency are each input into an independent encoder, and a cache occupancy perceptron processes the cache occupancy rate. This is then combined with the cache eviction probability distribution output by the policy network. The hidden feature vectors of each modality are obtained. In the feature fusion stage, the hidden feature vectors of each modality are concatenated and then input into the multi-head attention mechanism module to generate the fused feature representation. In the output stage, the fused feature representation is input into a fully connected layer for linear mapping to represent the probability distribution of the current cache eviction action. Conduct a value assessment and output the expected cumulative value. .

[0035] Furthermore, in the process of training the agent using the proximal policy optimization algorithm in S4, the loss function formula used is as follows:

[0036]

[0037]

[0038]

[0039] in, Indicates the current text generation task The loss; Indicates the importance sampling ratio; Indicates the current parameter and State-based Get action The probability of; Indicates the parameters used in the previous update. and State-based Get action The probability of; The dominant function represents the action to be performed. Relative state Advantages; It represents a constant used to control the range of values; Represents a numerical clipping function; This indicates taking the minimum value; Represent the expected function; Representation strategy Entropy;

[0040] By calculating the loss function with respect to the policy network parameters, respectively and value network parameters The gradient, and then the learning rate. Update strategy network parameters and value network parameters To obtain the optimal policy network parameters and optimal value network parameters .

[0041] Furthermore, to accommodate more text generation tasks, a dataset containing multiple text generation tasks is collected in S1. A target large language model is used to infer the data for each text generation task in the dataset and to complete environment modeling. Finally, based on deep reinforcement learning, meta-learning is combined to enhance the training of the agent, obtaining the globally optimal policy network parameters. and global optimal value network parameters .

[0042] Furthermore, the specific process of combining meta-learning to enhance the training of the agent is as follows:

[0043] First, the proximal policy optimization algorithm is used to perform inner loop training for each text generation task in the training set to obtain the optimal policy network parameters for each text generation task. and optimal value network parameters Then, for all text generation tasks, the optimal policy network parameters are applied. The calculated losses are summed to calculate their relationship with the optimal policy network parameters. The gradient, and then the learning rate. Update strategy network parameters To obtain the globally optimal policy network parameters For all text generation tasks, the optimal value network parameters are used. Sum the calculated losses and calculate their relationship with the optimal value network parameters. The gradient, and then the learning rate. Update value network parameters To obtain the globally optimal value network parameters .

[0044] Furthermore, the dataset containing various text generation tasks specifically includes data for dialogue tasks (DailyDialog), data for summarizing tasks (CNN / DailyMail), data for translation tasks (WMT), and data for question answering tasks, etc.

[0045] Compared with the prior art, the beneficial effects of the present invention are as follows:

[0046] 1. This invention proposes an adaptive KV cache management method based on deep reinforcement learning. It models the KV cache management problem in large language model inference as a Markov decision process. By training a deep reinforcement learning agent, it dynamically determines the cache eviction strategy, thereby achieving adaptive and efficient KV cache management. While maintaining or even improving the model inference quality, it significantly reduces memory consumption and increases inference throughput.

[0047] 2. This invention specifically designs a state representation that integrates multi-dimensional information, constructs an agent that includes a policy network and a value network, defines a multi-objective reward function that balances inference quality, memory efficiency, and computational latency, uses a proximal policy optimization algorithm for training, and introduces a meta-learning mechanism to achieve rapid task adaptation, providing an intelligent and adaptive solution for the efficient deployment of large language models. Attached Figure Description

[0048] Figure 1 This is a schematic diagram of the principle architecture of the adaptive KV cache management method based on deep reinforcement learning proposed in Example 1;

[0049] Figure 2 This is a schematic diagram of the policy network structure in Example 1;

[0050] Figure 3 This is a schematic diagram of the value network structure in Example 1;

[0051] Figure 4 This is a schematic diagram of the training process of the agent in Example 1. Detailed Implementation

[0052] To make the objectives, technical solutions, and advantages of the present invention clearer, the present invention will be further described in detail below with reference to the embodiments and accompanying drawings. The illustrative embodiments and descriptions of the present invention are only used to explain the present invention and are not intended to limit the present invention.

[0053] Example 1

[0054] This embodiment proposes an adaptive key-value cache management method based on deep reinforcement learning, the principle of which is as follows: Figure 1 As shown, it includes the following steps:

[0055] S1. Environment modeling and state construction, specifically:

[0056] S1.1 Collect datasets containing various text generation tasks, specifically including data for dialogue tasks, summarization tasks, translation tasks, and question-answering tasks;

[0057] S1.2 Use the target large language model to infer each piece of data for each text generation task in the dataset, and record the complete attention pattern and KV cache access trajectory;

[0058] S1.3 Model the LLM inference process as a partially observable Markov decision process (POMDP), specifically by modeling the environment based on the information recorded during inference, including the state at different decoding moments;

[0059] Define the state space: specifically define the state at time t. It includes five modal features: the attention score matrix at time t, token semantic features, position encoding, historical access frequency, and cache utilization rate. Among them, the token semantic features are extracted from the target large language model; the position encoding represents the relative position of each token; the historical access frequency represents the number of times each token is accessed, which is obtained by extracting from the running device; and the cache utilization rate is the ratio of the current cache size to the maximum cache size.

[0060] Define the action space: At each moment, the agent needs to decide whether to evict a cached entry, using... The symbol indicates that 0 means no expulsion and 1 means expulsion.

[0061] S2, agent construction, including policy network and value network, both adopt a lightweight Transformer architecture design;

[0062] The policy network is state-based. As input, after feature encoding and feature fusion, the output is the probability distribution of buffer eviction actions. ;in, Indicates the state The following actions were taken; Representation strategy;

[0063] Specifically, such as Figure 2 As shown, in the feature encoding stage, the state is... The attention score matrix, token semantic features, location encoding, and historical access frequency are input into independent encoders (the corresponding attention encoder, semantic encoder, location encoder, and historical information encoder, respectively), and a cache occupancy rate perceptron is used to process the cache occupancy rate to obtain the hidden feature vectors of each modality. In the feature fusion stage, the hidden feature vectors of each modality are concatenated and then input into the multi-head attention mechanism module to generate the fused feature representation. In the output stage, the fused feature representation is input into a fully connected layer for linear mapping to obtain the unnormalized score of each action. After processing by the Softmax function, the cache eviction action probability distribution is output. ;

[0064] The value network is in state and cache eviction probability distribution As input, after feature encoding and feature fusion, the output is the expected cumulative value. This indicates that the value network is related to the state. Execute action The state obtained later regarding time t+1 The assessed value;

[0065] Specifically, such as Figure 3 As shown, in the feature encoding stage, the state is... The attention score matrix, token semantic features, location encoding, and historical access frequency are input into independent encoders (the corresponding attention encoder, semantic encoder, location encoder, and historical information encoder, respectively), and a cache occupancy rate perceptron is used to process the cache occupancy rate. This is combined with the cache eviction action probability distribution output by the policy network. The hidden feature vectors of each modality are obtained. In the feature fusion stage, the hidden feature vectors of each modality are concatenated and then input into the multi-head attention mechanism module to generate the fused feature representation. In the output stage, the fused feature representation is input into a fully connected layer for linear mapping to represent the probability distribution of the current cache eviction action. Conduct a value assessment and output the expected cumulative value. .

[0066] S3. Reward function design, specifically:

[0067] S3.1 Define the composite reward function:

[0068]

[0069]

[0070]

[0071]

[0072] in, Indicates the execution of an action The immediate reward afterwards; Indicates a reward for the quality of reasoning; This indicates a memory efficiency bonus; Indicates delayed reward; , and Each represents a weight, which is adaptively adjusted during subsequent training. Indicates the execution of an action The change in confusion was subsequently caused by cache eviction. Indicates cache usage coefficient; Indicates the time cost coefficient; This indicates the cache occupancy rate at time t; Indicates the time cost of cache operations;

[0073] S3.2 Define the global evaluation function:

[0074]

[0075] in, Indicates the state Next action The final overall assessment score; Indicates the balance coefficient;

[0076] This design allows the optimization objective of the policy network to simultaneously consider immediate rewards. and expected cumulative value The two are weighted and summed to form a unified decision-making basis; the value network learns from historical experience to provide a state assessment reference for current decisions, along with immediate rewards. They jointly guide the cache management strategy of the intelligent agents.

[0077] S4. Enhance the training agent through deep reinforcement learning and meta-learning, as follows: Figure 4 As shown, specifically:

[0078] S4.1 Initialize policy network parameters and value network parameters ,

[0079] S4.2. Based on the environment and the overall evaluation function, the proximal policy optimization algorithm is used to perform inner loop training on each text generation task in the training set to obtain the optimal policy network parameters for each text generation task. and optimal value network parameters ;

[0080] Specifically, assuming the current text generation task is The loss function formula used for its training is:

[0081]

[0082]

[0083]

[0084] in, Indicates the current text generation task The loss; Indicates the importance sampling ratio; Indicates the current parameter and State-based Get action The probability of; Indicates the parameters used in the previous update. and State-based Get action The probability of; The dominant function represents the action to be performed. Relative state Advantages; This represents a constant used to control the range of values, with a value of 0.5; Represents a numerical clipping function; This indicates taking the minimum value; Represent the expected function; Representation strategy Entropy;

[0085] By calculating the loss function with respect to the policy network parameters gradient And regarding value network parameters gradient Then use the learning rate Update strategy network parameters and value network parameters To obtain the optimal policy network parameters and optimal value network parameters These are respectively used as temporary policy network parameters. and temporary value network parameters ;

[0086] S4.3, Perform external loop training, specifically:

[0087] Based on temporary policy network parameters In the current text generation task Calculate the loss again After traversing all text generation tasks, sum the losses corresponding to all text generation tasks and calculate their relationship with the temporary policy network parameters. The gradient, and then the learning rate. Update strategy network parameters To obtain the globally optimal policy network parameters ;

[0088] Based on temporary value network parameters In the current text generation task Calculate the loss again After traversing all text generation tasks, sum the losses corresponding to all text generation tasks and calculate their relationship with the temporary value network parameters. The gradient, and then the learning rate. Update value network parameters To obtain the globally optimal value network parameters ;

[0089] Through deep reinforcement learning and meta-learning enhancement, the agent can quickly adapt to new LLM architectures or task domains, learn how to make optimal policies based on the state to maximize the overall evaluation score, and deploy practical key-value cache management tasks.

[0090] The main innovations of the adaptive KV cache management method based on deep reinforcement learning proposed in this embodiment include:

[0091] 1. Reinforcement Learning Framework Design: For the first time, the key-value cache management of LLM is formalized as a reinforcement learning problem, defining the state space, action space, and reward function, enabling the agent to learn the long-term optimal cache management strategy.

[0092] 2. Multimodal state representation: A comprehensive state representation is designed that integrates attention patterns, token semantic features, location information, and historical access records to provide sufficient information for intelligent decision-making.

[0093] 3. Multi-objective reward mechanism: A composite reward function is proposed to balance inference quality, memory usage and computation latency, and adaptive trade-offs between different objectives are achieved through learnable weights.

[0094] 4. Lightweight network architecture: A policy network and value network based on Transformer were designed, and parameter sharing and knowledge distillation techniques were used to ensure low overhead during inference.

[0095] 5. Meta-reinforcement learning enhancement: Introducing a meta-learning mechanism enables the model to quickly adapt to new task domains and model architectures, improving generalization ability.

[0096] 6. Hybrid Policy Optimization: Combining heuristics such as H2O as prior knowledge, and accelerating convergence and improving performance through hybrid training of behavior cloning and policy gradients.

[0097] It should be noted that this is merely a specific embodiment of the present invention, but the scope of protection of the present invention is not limited thereto. Any modifications, equivalent substitutions, and improvements made by those skilled in the art within the scope of the technology disclosed in the present invention, and within the spirit and principles of the present invention, should be covered within the scope of protection of the present invention.

Claims

1. A deep reinforcement learning based adaptive KV cache management method, characterized in that, Comprising the following steps: S1, environment modeling and state construction, specifically: Collect a dataset about a text generation task; use a target large language model to infer each piece of data in the dataset, and complete environment modeling according to the information recorded by the inference, including the state at different decoding times; define the state at time t , including five kinds of modal characteristics, namely: the attention score matrix corresponding to time t, Token semantic features, position encoding, historical access frequency and cache occupancy rate; S2, agent construction, including policy network and value network; Wherein, the policy network takes state as input, and outputs the cache eviction action probability distribution after feature encoding and feature fusion ; wherein, represents the action taken in state ; represents the policy; The value network takes as input the state and the cache eviction action probability distribution and outputs the expected cumulative value after feature encoding and feature fusion. S3, reward function design, specifically: S3.1, defining a composite reward function: ； wherein, represents an immediate reward for performing an action . represents a reasoning quality reward; represents a memory efficiency reward; represents a delay reward; , and each represent a weight; S3.2, defining an overall evaluation function: ； wherein, represents the overall evaluation score after performing the action in the state ; represents a balancing coefficient; S4, training the agent through deep reinforcement learning, specifically: Initialize the policy network parameters and the value network parameters , train the agent based on the environment and the overall evaluation function using the proximal policy optimization algorithm to obtain the optimal policy network parameters and the optimal value network parameters , and then learn how to make the optimal policy according to the state to maximize the overall evaluation score for deployment in actual KV cache management tasks. 2.The adaptive KV cache management method based on deep reinforcement learning according to claim 1, wherein, The calculation formulas of S3.1, , , and are respectively:​​​​​​ ； ； ； wherein, represents performing an action represents a change in confusion level due to a cache eviction; represents a cache occupancy coefficient; represents a time overhead coefficient; represents a cache occupancy rate at time t; represents a time overhead of a cache operation. 3.The adaptive KV cache management method based on deep reinforcement learning according to claim 1, wherein, Both the policy network and the value network adopt a lightweight Transformer architecture design. 4.The adaptive KV cache management method based on deep reinforcement learning according to claim 3, wherein, For the policy network of S2, in the feature encoding stage, the attention score matrix in the state , token semantic features, position encoding and historical access frequency are respectively input into independent encoders, and the cache occupancy rate is processed by the cache occupancy rate perceiver to obtain the hidden layer feature vectors of each modality feature; in the feature fusion stage, the hidden layer feature vectors of each modality feature are spliced and then input into a multi-head attention mechanism module to generate the fused feature representation; in the output stage, the fused feature representation is input into a fully connected layer for linear mapping to obtain the unnormalized scores of each action, which are processed by a Softmax function to output the cache eviction action probability distribution . 5.The adaptive KV cache management method based on deep reinforcement learning according to claim 4, characterized in that, For the value network of S2, in the feature encoding stage, the attention score matrix in the state , token semantic features, position encoding and historical access frequency are respectively input into independent encoders, and the cache occupancy rate is processed by the cache occupancy rate perceiver, combined with the cache eviction action probability distribution output by the policy network , to obtain the hidden layer feature vectors of each modality feature; in the feature fusion stage, the hidden layer feature vectors of each modality feature are spliced and then input into a multi-head attention mechanism module to generate a fused feature representation; in the output stage, the fused feature representation is input into a fully connected layer for linear mapping to evaluate the value of the current cache eviction action probability distribution , and output the expected cumulative value . 6.The adaptive KV cache management method based on deep reinforcement learning according to claim 5, wherein, In the process of training the agent using the proximal policy optimization algorithm in S4, the loss function formula used is: ； ； ； wherein, represents the loss of the current text generation task ; represents the importance sampling ratio; represents the probability of taking action based on state under the current parameters ; ; represents the probability of taking action based on state under the parameters adopted in the last round of update ; ; is the advantage function, representing the advantage of state after taking action ; represents a constant for controlling the range of numerical values; represents a numerical clipping function; represents taking the minimum; represents the expected function; represents the entropy of policy ; By calculating the loss function respectively regarding and The gradient, and then the learning rate. renew and To obtain the optimal policy network parameters and optimal value network parameters . 7.The adaptive KV cache management method based on deep reinforcement learning according to claim 6, wherein, In order to adapt to more text generation tasks, in S1, a dataset containing multiple text generation tasks is collected, a target large language model is used to infer the data of each text generation task in the dataset, environment modeling is completed, and finally a global optimal strategy network parameter is obtained by combining meta learning to enhance the training agent on the basis of deep reinforcement learning and global optimal value network parameters . 8.The adaptive KV cache management method based on deep reinforcement learning according to claim 7, wherein, The specific process of enhancing the training of the agent in combination with meta-learning is: First, the inner loop training is performed for each text generation task in the training set using the proximal policy optimization algorithm to obtain the and ; then, the loss calculated based on for all text generation tasks is summed, the gradient thereof with respect to is calculated, and the learning rate is used to update to obtain the globally optimal policy network parameter ; the loss calculated based on for all text generation tasks is summed, the gradient thereof with respect to is calculated, and the learning rate is used to update to obtain the globally optimal value network parameter . 9.The adaptive KV cache management method based on deep reinforcement learning according to any one of claims 7 or 8, characterized in that, The dataset containing multiple text generation tasks specifically includes dialogue task data, summary task data, translation task data, and question and answer task data.

Images