面向长文本处理的键值缓存压缩技术：智能压缩，无损性能，免微调

面向长文本处理的键值缓存压缩技术：智能压缩，无损性能，免微调

引用

CSDN

1.

https://blog.csdn.net/weixin_44292902/article/details/139424886

随着输入长度的增加，大型语言模型（LLMs）中的键值（KV）缓存需要存储更多的上下文信息以维持性能，这导致内存消耗和计算时间急剧上升。为了解决这一问题，SnapKV应运而生，它是一种创新的KV缓存压缩方法，能够在不牺牲准确性的前提下显著减少KV缓存的大小。

KV缓存面临的挑战

在深入分析大型语言模型（LLMs）的注意力机制时，研究者发现了一些关键模式，这些模式对于优化KV缓存至关重要。某些键在令牌生成期间始终吸引着模型的注意力，无论上下文的长度如何，这些“活跃”的键展现出稳定的高注意力权重。在长摘要和问答任务中，问题的位置（无论是在提示的开头还是结尾）对模型的注意力分配模式影响不大，显示出模型在处理长文本时的鲁棒性。研究者还发现注意力模式高度依赖于上下文，与用户的具体指令密切相关，这意味着不同的指令会引导模型关注不同的信息。

SnapKV的核心原理

SnapKV的创新之处在于它提出了一种无需微调的压缩方法，通过观察模型在生成过程中的注意力分配模式，自动识别并压缩KV缓存中的关键信息。SnapKV通过“投票”机制选出每个注意力头关注的关键KV位置，并通过聚类算法保留这些关键特征周围的信息，从而在不牺牲准确性的前提下显著减少KV缓存的大小。这种方法不仅减少了计算开销，还提高了内存效率，使得模型在处理长文本时更为高效。

SnapKV的实现方法

SnapKV的实现分为两个主要阶段：

1. 投票选择重要特征（Voting for Important Previous Features）:

• 利用定义好的投票过程（如公式1所示），基于观察窗口（observation window）——即提示的最后部分——来选择重要的特征。
• 通过分析发现，这些特征在整个序列生成过程中表现出显著的一致性，表明它们对后续生成至关重要。
• 此外，实施聚类算法以保留选定特征周围的特征，这有助于保留信息的完整性并避免丢失上下文。

2. 更新和存储截断的键和值（Update and Store Truncated Key and Value）:

• 将选定的特征与观察窗口的特征连接起来，这些特征包含了所有提示信息。
• 将连接后的 KV 缓存存储起来，以供后续生成使用，同时节省内存使用。

SnapKV的实验结果

在对SnapKV进行的实验中，研究团队采取了一系列严谨的测试，旨在评估该方法在不同模型和长文本数据集上的性能。实验的目的是验证SnapKV在减少计算和内存负担的同时，是否能够保持或甚至提升模型的生成质量和效率。

实验首先在LWM-Text-Chat-1M模型上进行了压力测试，这是当时最先进的模型之一，能够处理长达一百万个令牌的上下文。测试中，SnapKV展现了其算法效率，特别是在硬件优化方面。通过“Needle-in-a-Haystack”测试，即在长达380K令牌的文档中准确检索特定句子的能力，SnapKV证明了其在极端条件下处理长文本的能力，即便在极高的压缩比下也能保持精确性。

在单 A100-80GB GPU 上，使用 SnapKV 优化的模型与原始实现在 "Needle-in-a-Haystack" 测试中的性能比较。图表显示了 SnapKV 在处理极长输入上下文时准确管理细节的潜力

实验通过不同的批处理大小对LWM-Text-Chat-1M模型进行了解码速度和内存限制的基准测试。结果表明，SnapKV优化的模型在解码速度上保持了稳定，与输入序列长度的增加无关，这与基线实现形成了鲜明对比，后者的解码速度随输入长度的增加而指数级增长。SnapKV显著提高了模型处理长序列的能力，显著减少了内存消耗。

在不同批量大小设置下，基线实现与 SnapKV 优化模型在解码速度和输入序列长度方面的比较。图表显示了 SnapKV 在保持解码速度恒定方面的优势

为了进一步验证SnapKV的有效性，研究团队还对Mistral-7B-Instruct-v0.2模型进行了消融研究，以理解池化技术对模型信息检索性能的影响。消融研究结果表明，通过池化增强了检索准确性，这可能是因为强大的注意力机制倾向于关注令牌序列的初始部分。

池化技术在 LongEval-Lines 测试中对检索准确性的影响。结果表明，使用池化的模型在压缩 KV 缓存时的性能显著优于不使用池化的模型

最后，实验使用了LongBench，这是一个多任务基准测试，旨在全面评估长文本理解能力。SnapKV在多个不同设置下进行了测试，包括压缩KV缓存到1024、2048和4096个令牌，并使用最大池化和观察窗口。测试结果显示，即使在压缩率高达92%的情况下，SnapKV与原始实现相比，在16个不同数据集上的性能下降可以忽略不计，某些情况下甚至超过了基线模型。

SnapKV 和 H2O 在 LongBench 数据集上不同 LLMs 的性能比较。表格显示了 SnapKV 在压缩提示 KV 缓存时的性能与原始实现相比几乎没有下降

实验结果，SnapKV证明了其作为一种有效的KV缓存压缩方法，在保持大型语言模型处理长文本的能力的同时，显著提升了效率和减少了资源消耗。这些发现不仅证实了SnapKV的实用性，还为未来的研究和应用提供了有价值的见解。

SnapKV代码可在https://github.com/FasterDecoding/SnapKV上找到
论文链接：https://arxiv.org/abs/2404.14469