超越均匀查询分布：基于键驱动的分组查询注意力
计算机视觉与模式识别
2024-08-29 v2 机器学习
摘要
Transformer 架构通过自注意力机制革新了深度学习，该机制能有效捕获上下文信息。然而，自注意力的内存占用对长序列任务构成了重大挑战。分组查询注意力（GQA）通过将查询分组并对相应的键值头进行均值池化来解决此问题——以灵活的方式减少了总体参数和内存需求，而不会不利地损害模型精度。在本工作中，我们引入了对 GQA 的增强，聚焦于两种偏离静态分组性质的新方法：基于键分布的分组查询注意力（KDGQA）和动态键分布分组查询注意力（DGQA），它们利用键值头的范数信息来指导查询分配。具体而言，KDGQA 在每次前向传播时查看键值头范数的比值，而 DGQA 考察范数在训练过程中的演化比值。此外，我们提出扰动 GQA（PGQA）作为案例研究，通过从注意力图中减去噪声来引入分组形成中的变异性。我们在经过微调的 Vision Transformer 上进行的实验，针对 CIFAR-10、CIFAR-100、Food101 和 Tiny ImageNet 等数据集的图像分类，证明了这些变体通过更信息丰富和自适应的分组机制改进原始 GQA 的潜力：具体而言，ViT-L 在使用 DGQA 时相比 GQA 和其他变体实现了高达 8% 的准确率提升。我们进一步分析了键值头数量对性能的影响，强调了利用查询-键亲和力的重要性。代码可在 GitHub 上获取。
引用
@article{arxiv.2408.08454,
title = {Beyond Uniform Query Distribution: Key-Driven Grouped Query Attention},
author = {Zohaib Khan and Muhammad Khaquan and Omer Tafveez and Burhanuddin Samiwala and Agha Ali Raza},
journal= {arXiv preprint arXiv:2408.08454},
year = {2024}
}
备注
11 pages, 9 figures