利用跨层注意力减少 Transformer 键值缓存大小
机器学习
2024-05-22 v1 计算与语言
摘要
键值（KV）缓存在加速基于 Transformer 的自回归大型语言模型（LLMs）的解码过程中扮演着至关重要的角色。然而，在长序列长度和大批量大小下，存储 KV 缓存所需的内存容量可能变得不可承受。自 Transformer 发明以来，为减少 KV 缓存大小而发现的最有效的两种干预措施是 Multi-Query Attention（MQA）及其推广形式 Grouped-Query Attention（GQA）。MQA 和 GQA 都修改了注意力块的设计，使得多个查询头可以共享单个键/值头，从而大幅减少了不同键/值头的数量，同时仅极小地降低准确率。在本文中，我们表明可以通过在相邻层之间也共享键和值头，将 Multi-Query Attention 进一步推进，产生一种我们称之为 Cross-Layer Attention（CLA）的新注意力设计。借助 CLA，我们发现可以在保持与未修改 MQA 几乎相同准确率的同时，将 KV 缓存大小再减少 2 倍。在从零开始训练 1B 和 3B 参数模型的实验中，我们证明 CLA 相较于传统 MQA 能够实现的内存/准确率权衡提供了 Pareto 改进，从而支持了比原本可能更长的序列长度和更大的批量大小进行推理。
引用
@article{arxiv.2405.12981,
title = {Reducing Transformer Key-Value Cache Size with Cross-Layer Attention},
author = {William Brandon and Mayank Mishra and Aniruddha Nrusimha and Rameswar Panda and Jonathan Ragan Kelly},
journal= {arXiv preprint arXiv:2405.12981},
year = {2024}
}