论文

Transformer需要三个投影吗? QKV 变体的系统研究

Do Transformers Need Three Projections? Systematic Study of QKV Variants

模型架构Transformer

摘要

Transformer 已成为各种人工智能任务的标准解决方案,其中查询、键和值 (QKV) 注意力公式发挥着核心作用。然而,这三个预测的各自贡献以及省略某些预测的影响仍然知之甚少。我们系统地评估了三个投影共享约束:a)Q-K=V(共享键值),b)Q=K-V(共享查询键)和c)Q=K=V(单个投影)。最后两个变体产生对称注意力图;为了解决这个问题,我们还通过 2D 位置编码探索不对称注意力。通过涵盖合成任务、视觉(MNIST、CIFAR、TinyImageNet、异常)和语言建模(10B 词元上的 300M 和 1.2B 参数模型)的实验,我们发现我们的 Transformer 的性能与 QKV Transformer 相当或偶尔更好。在语言建模中,Q-K=V 投影共享实现了 50% 的 KV 缓存减少,而困惑度仅降低了 3.1%。至关重要的是,投影共享是头部共享 (GQA/MQA) 的补充:将 Q-K=V 与 GQA-4 相结合可减少 87.5% 的缓存,而 Q-K=V + MQA 则可减少 96.9%,从而实现实际的设备上推理。我们证明 Q-K=V 保留了质量,因为键和值可以占据相似的表示空间,并且注意力在低秩状态下运行,而 Q=K-V 打破了注意力的方向性。我们的结果系统地将投影共享描述为注意力权重绑定的一个未经充分探索的实例,具有直接、可量化的推理记忆优势,对于边缘部署特别有价值。该代码可在 https://github.com/Brainchip-Inc/Do-Transformer-Need-3-Projections 上公开获取