©PaperWeekly 原创 · 作者 | 吴昊一、吴悠
单位 | 上海科技大学
研究方向 | 自然语言处理
为了对这类方法建立更加全面的认识，我们进一步提出了一个统一的框架，在这个框架中对比了 LCKV、YOCO 和 CLA，以及一些之前未被尝试的层间 KV 共享方法，并有幸取得了一些新的发现。我们希望可以借此帮助对层间共享 KV 感兴趣的使用者根据需求选择合适的方法，并为大模型高效推理的研究带来一些有益的启发。
论文链接：
代码链接：
论文链接：
代码链接：
Layer-Condensed KV Cache
Layer-Condensed KV Cache（LCKV）的核心思想非常简单：只计算 Transformer 顶层的 KV，并共享给其他所有层，将所有层的查询（query）与其配对，如图所示。
▲ LCKV 模型结构图。图中每个节点代表一个词元在一个 Transformer 层的计算过程，每个从 a 到 b 的箭头表示将 b 的查询与 a 的 KV 配对。
这个思路来自于我们先前的一份工作 Probabilistic Transformer。在这份工作中，我们将 Transformer 多层的计算过程解释为迭代地改进单词表示的过程。在这样的解释下，顶层的表示应当包含了最多、最完整的信息。因此直觉上看，在生成后文时，与前文顶层配对应当是最为恰当的。
从另一个角度来看，这也类似编码器-解码器（encoder-decoder）架构的 Transformer 中的交叉注意力（cross-attention）机制，其中解码器的所有层都关注编码器的顶层。
有些细节上的问题我们还需要处理一下。注意到在推理时，尽管每个单词能够看到前文的顶层 KV，但它自身顶层的 KV 还没有被计算；要计算它自身顶层的 KV，又需要先完成它自身底层的计算，这就形成了一个循环依赖的问题。
▲ 预热层数量对模型性能和吞吐量的影响
▲ 串行和迭代并行的计算图，两者相互等价
在我们今年 2 月完成这份工作并向 ACL 投稿之后，5 月 arXiv 上出现了多篇基于层间 KV 共享的其他工作。这些方法提出了与 LCKV 不同的共享策略，在不同的场景下取得了不错的效果。我们在这里简单介绍一下其中的两种方法。
▲ YOCO模型结构图
正因如此，该工作重点关注了长文本场景，验证了模型能够在超长文本下保持较好的表现，同时大量减少 KV 缓存的内存消耗和预填充时间，并提高吞吐量。
▲ CLA模型结构图
上述三种方法都是基于层间共享 KV 的思想，只是共享方式各有不同。如果相互借鉴，还可以有更多的可能性。于是，我们整理了一下思路，用一个统一的框架来描述这些方法，以便更好地进行对比。
首先，我们考虑只保留 Transformer 中一层的 KV，而所有层都使用这一层计算得到的 KV。我们将这一层称为所有层的目标层（target layer），即图中箭头所指向的那一层。
由于这一层计算并使用自身的 KV，我们也将其称为一个 KV 层（KV layer）。我们可以取不同位置的层作为目标层，这里我们只考虑顶部（top）、中部（middle）和底部（bottom）三种情况。如果目标层位置为顶部或中部，模型计算时会遇到和 LCKV 相同的循环依赖和顺序依赖问题，我们同样采用屏蔽对角线和迭代计算的方式来解决这两个问题。
实验表明，仅有一层提供 KV 不足以使得模型具有可靠的效果。因此，我们需要考虑增加 KV 层的数量。
一种方式是在模型的底部加入一些 KV 层，待模型隐层表示能够包含足够的信息后再共享 KV，这就类似于披萨（pizza）的饼底，把上层的馅料托住。
或者像 LCKV 中提到的那样，在顶部和底部各加入一些 KV 层，这就类似于三明治（sandwich）的面包，把馅料夹在中间。
再或者，我们可以将所有层均匀地划分为若干组连续的层，每组中的所有层都使用同组中某一层 KV，这就类似于千层面（lasagna），由面皮和馅料交替堆叠而成。
▲ 统一框架的 9 种设定。图中红色的层代表 KV 层，每个箭头从一个非 KV 层指向它的目标层
为了研究不同的层间 KV 共享方式对推理速度和模型性能的影响，我们将 Llama 作为基线模型，与框架中的 9 种设定进行了实验对比。
在推理速度方面，我们对比了 1.1B 的模型在不同上下文长度下的最大吞吐量，如图 a 所示。在模型性能方面，我们首先将 110M 和 1.1B 的模型在小规模的 Minipile 数据集上进行训练，它们的困惑度（perplexity，PPL）如图 b 所示。
我们从图中可以发现，大多数设定都可以在不显著降低模型性能的前提下，大幅提升推理吞吐量，并且随着 KV 层数的减少，吞吐量和模型性能分别呈现上升和下降的趋势。
图中最值得关注的是目标层为顶部和中部的情况：在推理效率方面，因为需要对提示词进行迭代计算，所以较长的提示词（512+1024）会导致吞吐量显著下降；在模型性能方面，在 KV 层数较少时，将目标层位置设为顶部或中部，可以比设为底部更好地维持性能。
我们还注意到，当 KV 层数较少时，三明治-中部设定具有最好的性能，这个发现与 LCKV 的出发点不太一致。LCKV 中，我们认为顶层的 KV 包含了最丰富和重要的信息，但实验表明使用中间层的 KV 效果更好。
总而言之，通过层间共享来减少 KV 缓存的显存消耗是一个极具潜力的方向，现在也正在受到越来越多的关注。我们率先提出了 LCKV，通过只保留顶层的 KV 并共享给其他所有层，实现了显著的显存节省和推理加速。
课题组曾获 ACL 2023 杰出论文奖、SemEval 2022 和 SemEval 2023 两次最佳系统论文奖，以及多个顶会最佳论文提名。课题组招收硕士生（推免学硕）、博士生（推免直博、普博工博）、博士后和研究助理。
更多信息请访问屠可伟老师主页：
参考文献
[2] Brandon, W., Mishra, M., Nrusimha, A., Panda, R., & Kelly, J. R. (2024). Reducing Transformer Key-Value Cache Size with Cross-Layer Attention. arXiv preprint arXiv:2405.12981.
[3] Sun, Q., Pickett, M., Nain, A. K., & Jones, L. (2024). Transformer layers as painters. arXiv preprint arXiv:2407.09298.
更多阅读
#投 稿 通 道#
让你的文字被更多人看到
如何才能让更多的优质内容以更短路径到达读者群体，缩短读者寻找优质内容的成本呢？答案就是：你不认识的人。
总有一些你不认识的人，知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁，促使不同背景、不同方向的学者和学术灵感相互碰撞，迸发出更多的可能性。
PaperWeekly 鼓励高校实验室或个人，在我们的平台上分享各类优质内容，可以是最新论文解读，也可以是学术热点剖析、科研心得或竞赛经验讲解等。我们的目的只有一个，让知识真正流动起来。
📝 稿件基本要求：
• 文章确系个人原创作品，未曾在公开渠道发表，如为其他平台已发表或待发表的文章，请明确标注
• 稿件建议以 markdown 格式撰写，文中配图以附件形式发送，要求图片清晰，无版权问题
• PaperWeekly 尊重原作者署名权，并将为每篇被采纳的原创首发稿件，提供业内具有竞争力稿酬，具体依据文章阅读量和文章质量阶梯制结算
📬 投稿通道：
• 投稿邮箱：hr@paperweekly.site
• 来稿请备注即时联系方式（微信），以便我们在稿件选用的第一时间联系作者
• 您也可以直接添加小编微信（pwbot02）快速投稿，备注：姓名-投稿
△长按添加PaperWeekly小编
🔍
现在，在「知乎」也能找到我们了
进入知乎首页搜索「PaperWeekly」
点击「关注」订阅我们的专栏吧
