Large language models (LLMs) utilize key-value (KV) cache to store historical information during sequence processing. The size of KV cache grows linearly as the length of the sequence extends, which seriously affects memory usage and decoding efficiency. Current methods for KV cache eviction typically utilize the last window from the pre-filling phase as queries to compute the KV importance scores for eviction. Although this scheme is simple to implement, it tends to overly focus on local information, potentially leading to the neglect or omission of crucial global information. To mitigate this issue, we propose Judge Q, a novel training method which incorporates a soft token list. This method only tunes the model's embedding layer at a low training cost. By concatenating the soft token list at the end of the input sequence, we train these tokens' attention map to the original input sequence to align with that of the actual decoded tokens. In this way, the queries corresponding to the soft tokens can effectively capture global information and better evaluate the importance of the keys and values within the KV cache, thus maintaining decoding quality when KV cache is evicted. Under the same eviction budget, our method exhibits less performance degradation compared to existing eviction approaches. We validate our approach through experiments conducted on models such as Llama-3.1-8B-Instruct and Mistral-7B-Instruct-v0.3, using benchmarks including LongBench, RULER, and Needle-in-a-Haystack. Results indicate an improvement of approximately 1 point on the LongBench and over 3 points on RULER. This proposed methodology can be seamlessly integrated into existing open-source models with minimal training overhead, thereby enhancing performance in KV cache eviction scenarios.

翻译：大型语言模型（LLM）在序列处理过程中利用键值（KV）缓存存储历史信息。随着序列长度的增加，KV缓存的规模呈线性增长，严重影响内存使用和解码效率。现有的KV缓存逐出方法通常采用预填充阶段的最后窗口作为查询来计算KV重要性分数以执行逐出。尽管该方案实现简单，但容易过度关注局部信息，可能导致关键全局信息被忽略或遗漏。为缓解此问题，我们提出Judge Q——一种结合软令牌列表的新型训练方法。该方法仅以较低训练成本微调模型的嵌入层。通过在输入序列末尾拼接软令牌列表，我们训练这些令牌对原始输入序列的注意力图，使其与实际解码令牌的注意力图对齐。通过这种方式，软令牌对应的查询能够有效捕获全局信息，并更好地评估KV缓存中键值的重要性，从而在KV缓存被逐出时保持解码质量。在相同逐出预算下，与现有逐出方法相比，我们的方法表现出更低的性能下降。我们在Llama-3.1-8B-Instruct和Mistral-7B-Instruct-v0.3等模型上进行了实验验证，使用的基准测试包括LongBench、RULER和Needle-in-a-Haystack。实验结果表明，在LongBench上性能提升约1分，在RULER上提升超过3分。该方法能够以极低的训练开销无缝集成到现有开源模型中，从而提升KV缓存逐出场景下的性能表现。