大型语言模型 （LLM） 在今年可谓是风光无限。不过惊艳的效果背后是一个巨大的模型以及夸张的硬件资源。
LLM在现实中部署时通常会面临两个难题：昂贵的KV缓存成本，以及对长序列的泛化能力差。
近日，田渊栋团队发表了一篇论文，成功解决以上两个难题，并将推理系统的吞吐量提高了近30倍！
论文地址：https://arxiv.org/pdf/2306.14048.pdf
代码地址：https://github.com/FMInference/H2O
这个成果也将在NeurIPS'23上展示。
下面，我们来看一下这两个难题的具体情况，以及论文提供的解决方案。
首先是缓存，KV缓存用于存储生成过程中的中间注意力键和值，以避免重新计算。
通常，除了模型参数外，还会将大量瞬态信息（KV缓存）存储在GPU内存中，这部分的内存占用，与序列长度和批处理大小线性相关。
例如，一个输入批次大小为128、序列长度为1024的300亿参数模型需要180GB的KV缓存。
其次，由于硬件限制，LLM会以固定的序列长度进行预训练（例如Llama-2使用固定长度4K的序列）。
然而，这其实也对推理过程中的注意力窗口施加了限制，使得模型在面对更长输入序列时无法发挥作用，阻碍了更广泛的应用。
对此，论文提出了一种实现KV缓存的新方法，显著减少了内存占用，且在长输入序列的任务中表现良好。
方法基于这样一个事实：在计算注意力分数时，一小部分tokens贡献了大部分的价值，——这里称这些tokens为Heavy Hitters (H2)。
通过综合调查，作者发现H2的出现是自然的，且与文本中词组的频繁共现密切相关，而去除它们会导致显著的性能下降。
基于此，作者提出了Heavy Hitter Oracle（ H2O ），一种KV缓存逐出策略，可动态保持最近的tokens和H2 tokens的平衡。
另外，作者将KV缓存驱逐表述为一个动态的子模块问题，为提出的驱逐算法提供了理论保证。
最后，作者使用OPT、LLaMA和GPT-NeoX在各种任务中验证算法的准确性。
其中，在OPT-6.7B和OPT-30B上实现的H2O，将DeepSpeed Zero-Inference、Hugging Face Accelerate和FlexGen这三个推理系统的吞吐量分别提高了29倍、29倍 和3倍，且在相同的批量大小下，H2O最多可以减少1.9倍 的延迟。
论文细节
上图为在 LLM 生成中部署不同 KV 缓存策略的符号图；左下为H2O的框架概述；右下为不同策略下的准确性与内存消耗的对比。
我们可以看出，将前几种方法应用于预训练的LLM ，会导致高未命中率并降低精度。
解决KV缓存问题，面临着三个技术挑战。
首先，目前尚不清楚是否可以限制KV缓存的大小——原则上，每个解码步骤可能需要访问所有先前的注意力键和值。
其次，确定保持生成准确性的最佳逐出策略是一个组合问题。
最后，即使可以暴力解开最佳策略，在实际应用程序上部署也是不可行的。
幸运的是，作者通过研究发现了一些有趣的结果。
小缓存大小的稀疏性：即使在密集训练时，LLM的注意力矩阵在推理时也有超过95% 的稀疏率（图a）。这适用于各种预训练的LLM。
因此，在每个生成步骤中，只需要5% 的KV缓存就足以解码相同的输出tokens，这表明KV缓存大小最多可以减少20倍，而不会降低精度。
Heavy Hitters（ H2 ）：注意力区块中所有tokens的累积注意力分数都遵循幂律分布（图b）。这表明存在一小群有影响力的tokens，这些tokens在生成过程中至关重要，是重量级tokens （ H2 ）。这使我们可以摆脱组合搜索问题，并确定保持准确性的逐出策略。
低成本策略的贪婪算法：在每个解码步骤中保留基于局部统计数据的H2（仅将前面tokens的注意力分数相加）与考虑未来tokens的注意力一样有效（图d）。
基于上述内容，作者定义了在大小受限的KV缓存中， LLM的生成过程，并提出了Heavy-Hitter Oracle （ H2O ），该框架利用了上面提到的性质，并使用简单、低成本的驱逐策略。
方法与分析
LLM的生成过程包括两个不同的阶段：
提示阶段：使用输入序列来生成KV缓存（由键和值嵌入组成），类似于LLM训练期间采用的前向传递；
tokens生成阶段：利用和更新KV缓存以增量方式生成新tokens 。每个生成步骤都依赖于先前生成的tokens。
本文的重点是在tokens生成阶段提高KV缓存的注意力效率，从而加速LLM推理。
作者定义了具有有限KV缓存大小的生成过程，包括注意力查询矩阵Q和键矩阵K。
驱逐策略：
以及采用了驱逐策略的生成过程：
接下来讨论在不降低精度的情况下，减少KV缓存大小的可能性。
上图中，（a）表示预训练OPT模型中的注意力稀疏性；（b）表示累积注意力分数相对于相应单词的分布（红色散点）和数据中单词的共现次数（灰色曲线），x轴表示词汇表中的单词索引；（c）表示具有完整KV缓存的基线模型与本文模型（H2O）的性能比较；（d）表示具有完整KV缓存的基线模型、具有局部统计量的H2O、具有全局统计量的H2O和仅具有最新KV（局部）的模型之间的比较。
给定由查询矩阵Q和键矩阵K计算的归一化注意力得分Softmax矩阵，将阈值设置为每行最大值的百分之一，并计算相应的稀疏度。
然后在Wiki-Text-103的验证集上使用预训练的OPT模型进行零样本推理，绘制注意力块内的逐层稀疏性，并可视化了归一化的注意力得分矩阵。
结果如下图所示，尽管LLM是密集训练的，但由此产生的注意力得分矩阵是高度稀疏的，几乎所有层的稀疏度都超过95%。
注意力块的稀疏性表明，生成下一个tokens时，不需要访问所有先前的键和值嵌入，所以可以逐出不必要的KV嵌入，也就减少了生成过程中对KV缓存的需求。
不过，逐出的策略需要谨慎，因为一旦驱逐了重要的KV，由于LLM生成的顺序依赖性，可能会破坏LLM的性能。
作者研究发现，注意力区块内所有tokens的累积注意力分数都遵循幂律分布，如下图所示。
这表明存在一小部分在生成过程中至关重要的tokens，也就是前文谈到的Heavy-Hitters （H2）。
此外，每个单词的累积注意力分数（红点）与它们在数据中的共现（灰色曲线）具有高度相关性。
作者基于以上现象设计了一种贪婪驱逐策略：
在生成过程中，当令tokens数量超过分配的KV缓存预算时，根据其累积的注意力分数统计数据，以及缓存中的本地tokens来保留重量级tokens。
一般而言，需要使用整个生成过程中的统计数据，才能得到最理想的结果，但这在实际部署中显然是不可行的，因为无法访问未来生成的tokens。
于是，作者进行了下图的实验，发现在每个解码步骤中使用局部统计数据计算的局部H2 ，与考虑未来tokens的情况效果差不多（红线和蓝线）。
随后，作者将这种动态注意力分数计算（有空间限制）定义为一种新的动态的子模块问题（dynamic submodular type problem）：
利用上面的形式定义KV缓存驱逐策略：
上图展示了驱逐算法，以及说明性示例。这里假设KV缓存的预算大小为3 ，完成第四个解码步骤后，根据累积的注意力分数逐出与第三个token关联的KV嵌入，被逐出的KV嵌入在后续解码步骤中将不可访问。
另外，作者还提到了实际实现中的细节。比如，为了保证I/O效率，我们在驱逐存储的KV时不会交换内存，而是直接填充新添加的KV。
实验结果
论文的实验选用了三个具有代表性的LLM模型系列，包括OPT，LLaMA和GPT-NeoX-20B 。
选取了8个评估任务：COPA ， MathQA ， OpenBookQA ， PiQA ， RTE ， Winogrande ， XSUM ， CNN/Daily Mail 。
实验的硬件采用NVIDIA A100 80GB GPU。
考虑到H2O所采用的缓存策略，这里除了完整的KV缓存（Full），还将本地缓存策略（Local）也作为基线方法。
由上面的图和表可知：在不同的KV缓存预算下，本文提出的方法（H2O）在各种不同条件的测试中都优于Local策略。
同时，在低于20%的KV缓存预算之下，H2O实现了与全KV嵌入模型（Full）相当的性能，且在更具挑战性的长序列生成任务、XSUM和CNN/Daily Mail中表现良好。
