本文是一篇名为《h2o: heavy-hitter oracle for efficient generative inference of large language models》的学术研究论文，发表于第37届神经信息处理系统大会（NeurIPS 2023）。该研究由来自德克萨斯大学奥斯汀分校、斯坦福大学、加州大学圣地亚哥分校、加州大学伯克利分校、Adobe研究院、Meta AI以及卡内基梅隆大学等机构的研究人员共同完成。本研究聚焦于解决大型语言模型（Large Language Models, LLMs）在部署，特别是生成长文本内容时，面临的一个关键性能瓶颈：键值缓存（Key-Value cache， KV cache）的巨大内存占用问题。KV cache存储了生成过程中所有先前token的键（Key）和值（Value）嵌入，用于自回归解码，其大小与序列长度和批次大小成线性增长，在长序列、大批次场景下（例如，一个300亿参数的模型，批次大小为128，序列长度为1024时，KV cache可达180GB）成为限制模型推理效率和可扩展性的主要障碍。
研究的学术背景源于大型语言模型在内容创作、摘要、对话系统等自然语言处理任务上取得的显著成功，但其高昂的部署成本限制了广泛应用。除了模型参数量巨大和注意力层的二次计算复杂度外，KV cache的大小问题日益凸显。尽管已有大量关于训练时稀疏注意力近似的研究，但它们并未被广泛用于缓解KV cache瓶颈。现有方法如Reformer和Flash Attention旨在解决长序列建模时注意力的二次内存问题，但仍需要大的缓存；而稀疏Transformer、低秩Transformer或多查询注意力等方法虽能减少缓存大小，但直接应用于预训练LLM进行生成会导致高丢失率和精度下降。因此，研究团队旨在设计一种理想的KV cache机制，需满足三个目标：缓存尺寸小以降低内存占用、丢失率低以保持LLM性能和长内容生成能力、以及淘汰策略开销低以减少生成时的实际耗时。
本研究的工作流程始于对LLM注意力机制的深入观察与分析。研究团队首先发现，即使在密集训练后，LLM在推理时的注意力矩阵也具有超过95%的稀疏性，这意味着在每个生成步骤中，仅有约5%的KV缓存对生成下一个token是充分且必要的。这为在不损失精度的情况下大幅减少KV缓存大小（理论上可达20倍）提供了可能性。第二个关键观察是，注意力块中所有token的累积注意力分数遵循幂律分布，即存在一小部分对生成过程至关重要的“重要token”，研究团队将其命名为“Heavy Hitters”（H2）。实验表明，移除这些H2 token会导致模型性能显著下降，而它们的出现与文本中token的频繁共现高度相关。基于这两个观察，研究团队提出了Heavy-Hitter Oracle（H2O）框架，其核心是一个动态的KV缓存淘汰策略。
H2O算法的详细工作流程如下：首先，将KV缓存的大小限制在一个固定的预算K（例如，总序列长度的20%）。在生成过程的每个解码步骤i，算法维护一个当前缓存中的token索引集合Si。当需要生成第i个token时，算法将当前新生成的token i的KV对加入候选集合（即Si-1 ∪ {i}）。如果此时集合大小未超过预算K，则直接保留。如果超过预算K，则需要根据一个评分函数f_score淘汰一个KV对。H2O采用的评分函数是基于局部统计的累积注意力分数。具体而言，对于候选集合中的每个元素v，算法计算如果将其从集合中移除后，剩余集合的累积注意力分数总和。然后，选择移除后能使剩余集合分数最高的那个元素v进行淘汰。这个过程在算法1中被形式化描述。其核心思想是，在每一步都动态地保留一个兼顾“最近token”（因其与当前token相关性更强）和“重要token”（H2）的KV缓存组合。研究团队将这一淘汰过程形式化为一个动态子模最大化问题，并在一定假设下为所提出的贪心算法提供了理论保证（定理4.4），表明该算法能提供接近最优解的近似解。
为了验证H2O的有效性，研究团队进行了广泛的实验评估。实验对象包括OPT、LLaMA和GPT-NeoX等多种主流LLM架构，模型规模从67亿到300亿参数不等。评估任务涵盖了来自LM-eval-harness和HELM的多个下游任务，如COPA、MathQA、OpenBookQA、PIQA、RTE、Winogrande、XSum和CNN/Daily Mail等，以及AlpacaEval和MT-Bench等生成基准测试。基线方法包括使用完整KV缓存的模型、仅保留最近token的“局部”策略、以及稀疏Transformer的两种变体（步进式和固定式）。实验的主要指标是在不同KV缓存预算（如4%到100%）下的任务精度、生成吞吐量和延迟。
研究的主要结果体现在多个方面。在精度方面，H2O在仅使用20% KV缓存预算（即内存占用减少5倍）的情况下，在大多数任务上达到了与使用完整KV缓存模型相当甚至偶尔更优的性能。相比之下，“局部”策略和其他稀疏化基线在低缓存预算下性能严重下降。例如，在OPT-30B模型上，使用20%缓存预算的H2O在COPA任务上达到85%的准确率，与完整缓存持平，而“局部”策略仅为48%。研究还发现，H2O策略可以增强现有稀疏化技术，当与“步进式”稀疏注意力结合时，能使其在20%预算下恢复接近完整缓存的性能。在系统性能方面，研究团队在FlexGen推理引擎上实现了H2O，并在NVIDIA T4和A100 GPU上进行了测试。结果表明，H2O（20%预算）相比FlexGen、DeepSpeed Zero-Inference和Hugging Face Accelerate这三个领先的推理系统，吞吐量最高可分别提升3倍、29倍和29倍。在相同批次大小下，H2O能将延迟降低高达1.9倍。此外，研究还展示了H2O能够与量化技术兼容，结合使用时甚至能获得比单独使用任一技术更好的精度；并且H2O能够赋能LLM处理极长序列（实验展示了处理长达四百万token的输入），其性能优于StreamingLLM等方法。一个有趣的额外发现是，H2O还能增加生成文本的多样性，减少重复用词。
本研究的结论是，H2O作为一种简单而有效的KV缓存淘汰策略，通过识别并动态保留注意力计算中的“重要token”（Heavy Hitters），能够显著减少大型语言模型推理时的内存足迹，从而大幅提升吞吐量并降低延迟，且不损害模型在各种任务上的生成质量。研究将KV缓存淘汰问题形式化为动态子模优化，并提供了理论分析，为后续工作提供了指导。
本研究的亮点在于：第一，发现了LLM注意力机制中固有的稀疏性和“重要token”现象，并将其转化为实用的系统优化洞见。第二，提出了H2O这一新颖、低开销、无需重新训练的动态缓存管理算法，该算法仅依赖局部注意力统计信息，易于实现和部署。第三，通过严格的理论分析和广泛的实证评估，全面验证了方法的有效性、高效性和通用性，涵盖了从精度、吞吐量、延迟到长序列处理等多个维度。第四，展示了该方法与现有优化技术（如稀疏注意力、量化）的兼容性和协同增强效果。第五，开源了代码，促进了该领域的研究和应用。这项研究为解决LLM部署中的关键瓶颈提供了重要的思路和工具，对于推动LLM在资源受限环境下的高效应用具有显著的学术价值和实际意义。