Layer-Condensed KV——利用跨层注意（CLA）减少 KV 缓存中的内存保持 Transformer 1B 和 3B 参数模型的准确性

导言

论文地址：https://arxiv.org/pdf/2405.12981

近年来，转换器模型在自然语言处理领域取得了长足的进步，在广泛的应用中取得了优异的成绩。然而，要最大限度地提高大规模语言模型的性能，对内存要求较高的键值（KV）缓存是必不可少的。特别是在处理长序列和大批量数据时，其内存消耗量非常大，给实际工作带来了挑战。

为解决这一难题，许多研究人员探索了提高 KV 缓存内存效率的方法。其中，多查询关注（MQA）和分组查询关注（GQA）已被广泛采用，它们允许多个查询头共享一个键/值头，是减少 KV 缓存大小的有效手段。不过，还需要进一步提高内存效率。

在此背景下，麻省理工学院和麻省理工学院-IBM 沃森人工智能实验室的研究人员提出了一种新方法--跨层注意力（CLA）。在保持模型准确性不变的情况下，进一步缩小模型的规模。

相关研究

为了最大限度地提高变压器模型的性能和效率，许多研究人员正在探索不同的方法。本文特别关注提高 KV 缓存器内存效率的相关研究。以下是本文介绍的主要相关研究的摘要。

多重查询注意（MQA）和分组查询注意（GQA）

最相关的研究是多查询注意（MQA）和分组查询注意（GQA），它们改进了转换器模型的注意机制。Ainslie等人（2023年）将其归纳为分组查询注意（GQA），这是一种将查询头分组、每组共享一个键/值头的架构。在这种架构中，每组共享一个键/值头。这在提高内存效率的同时，最大限度地减少了准确性的损失。

KV 缓存压缩

另一种缩小 KV 缓存大小的方法是 KV 缓存压缩：Hooper 等人（2024 年）提出了 KVQuant，它执行键和值量化以实现低精度存储。Zhang 等人（2024 年）还展示了如何通过使用一种称为耦合量化的技术对键和值进行非均匀编码，将 KV 缓存压缩到 1 或 2 比特。

删除不必要的 KV 缓存条目

其他方法包括删除不必要的 KV 缓存项；Zhang 等人（2023 年）提出了 H2O 来删除不重要的 KV 缓存项；Liu 等人（2023 年）使用 Scissorhands 方法展示了如何在生成过程中存储重要标记。他们展示了如何只存储正在生成的重要标记。这些方法可以有效减少 KV 缓存的内存使用量。

通过结构调整减少 KV 缓存大小

本文提出的跨层关注（CLA）是一种通过改变架构来减少 KV 缓存大小的方法。它的独特之处在于 CLA 涉及相邻层之间的键/值共享，而传统的 GQA 和 MQA 涉及单层内的键/值共享。这进一步减少了 KV 缓存的内存占用，同时保持了模型的准确性。

训练记忆的效率

为了提高训练过程中的内存效率，人们也做了很多研究：Shoeybi 等人（2020 年）提出了一种用于高效训练大规模神经网络的模型并行化技术 Megatron-LM；Huang 等人（2019 年）推出了 GPipe，它使用流水线并行化来优化训练内存的使用。CLA也与这些技术兼容，有望进一步提高内存效率。

拟议方法（跨层关注）

为了解决变压器模型中 KV 缓存的内存问题，研究人员提出了一种新方法--跨层注意（CLA），即在相邻层之间共享键和值头，以这种方法在保持模型精度的同时，还能缩小相邻层之间的键和值头。本节将详细介绍 CLA 的设计及其具体操作。

CLA 的基本概念

在传统的转换器架构中，每一层都计算自己的键和值，并将其存储在 KV 缓存中。这种方法需要大量内存来处理长序列和大批量数据。相比之下，CLA 通过与相邻层共享某些层计算的键和值来减少内存使用量。

具体来说，CLA 的工作原理如下。

密钥/值计算和共享：有些层会计算自己的密钥和值，并将其存储在 KV 缓存中。相邻层则重新使用这些计算出的密钥和值（见图 1）。

-共享因子：键和值共享的层数称为 "共享因子"。例如，如果共享系数为 2，则每对图层都使用相同的键和值（见图 2）。

这种方法通过共享因子减少了 KV 缓存的内存使用量。

图 1：跨层关注（CLA）概念图

CLA 架构

CLA 设计可与传统的多重/查询注意（MQA）和分组/查询注意（GQA）相结合。传统的 MQA 和 GQA 在同一层中共享键和值，而 CLA 则在多层中共享。这样可以进一步节省内存。

协调和联络局的具体结构如下。

键/值投影 ：有些层执行自己的键和值投影，并将结果存储在 KV 缓存中。其他层则重复使用这些投影。
组合灵活性：CLA 可以与 MQA 和 GQA 结合使用，将各自的优势结合起来，以达到最佳的内存效率。

图 2：具有不同共享因素的 CLA 的构成

试验

本研究使用 1B 和 3B 参数模型进行了一系列实验，以测试所提出的跨层注意力（CLA）方法的有效性。

在所有实验中，模型都是在 SlimPajama 数据集上训练的。GPT-NeoX 标记符号化器被用作模型标记符号化器，并使用字节对编码（BPE）进行标记化。此外，还采用了基于 Llama 架构的预规范化、SwiGLU 激活函数和旋转位置嵌入。训练在 PyTorch 框架中使用英伟达 H100 GPU 进行。

1B 参数模型的实验结果

在 1B 参数模型中测试了各种 CLA 配置。其中，MQA-CLA2 配置表现出色（见图 3）。

图 3：1B 参数模型的实验结果。

∙MQA 和 CLA2 模型：与传统的 MQA 模型相比，头部尺寸从 64 到 512 不等的 MQA 和 CLA2 模型提高了精确度，同时减少了 KV 缓存内存。尤其是头部尺寸为 128 的模型，与传统 MQA 模型相比，内存使用量减少了一半，但精度几乎相同。
GQA/CLA2 模型：还测试了 GQA 和 CLA2 的组合模型，但最有效的是 GQA2/CLA2 配置，它比其他配置显示出更高的精度。

3B 参数模型的实验结果

我们还在 3B 参数模型中进行了实验，以验证 CLA 的有效性。结果再次证明，MQA-CLA2 配置最为有效。

128 字头维度的 MQA 模型：调整学习率后，MQA-CLA2 模型的准确性优于 128 字头维度的传统 MQA 模型。特别是，在维基文本数据集上观察到的困惑度差异很大（见表 5）。

表 5：3B 参数模型的实验结果

审议

1. 提高内存效率：据观察，CLA 可以有效减少 KV 缓存的内存使用量，同时保持接近的精确度，尤其是在共享系数为 2 的情况下。因此，与以前的架构相比，内存效率有了显著提高。

2. 保持准确性：使用 CLA 可以减少内存使用量，同时将准确性损失降到最低，这在序列较长、批量较大的情况下尤为有用。

3. 学习率的重要性：研究表明，调整学习率对模型的性能有显著影响，在 CLA 模型中，较高的学习率尤其有效。这表明 CLA 不仅能提高记忆效率，还能提高训练过程本身的效率。

这些结果表明，CLA 是一种可以成为变压器模型设计新标准的方法，在实用性和效率方面都具有显著优势。

结论

在本文中，我们提出了跨层关注（Cross-Layer Attention，CLA）作为一种新方法，以减少变压器模型中 KV 缓存的内存使用。并保持几乎相同的精度。实验结果，尤其是 1B 和 3B 参数模型的实验结果表明，CLA 在内存效率和精确度方面都表现出色。

未来前景包括 CLA 的进一步优化和扩展。例如，应用于不同的模型架构和更大的模型，以及验证 CLA 在长期序列中的有效性。同样重要的是，通过评估 CLA 在实际应用中的有效性，进一步确认 CLA 的实用性和有效性；CLA 将是推动变压器模型发展的重要一步。