Weighted Grouped Query Attention in Transformers

注意机制是Transformer语言模型的基础模块。最近的方法表明，模型的规模扩大可以达到人类水平的性能。然而，随着对规模的需求增加和硬件内存的限制，这些模型的推理成本仍然很高。为了减少推理时间，在（Shazeer，2019）和（Ainslie等人，2023）中提出了多查询注意力（MQA）和分组查询注意力（GQA）。在本文中，我们提出了一种分组查询注意力的变体，称为加权分组查询注意力（WGQA）。我们为T5解码器注意力块中的每个键和值头引入了新的可学习参数，使模型在微调期间能够进行加权平均。我们的模型比GQA平均提高了0.53％，且性能在推理期间没有额外的开销时收敛到传统的多头注意力（MHA）。我们通过比较T5-small和T5-base架构之间的结果来评估我们分析中的缩放定律。我们评估了引入这些参数和随后的微调如何在训练期间通知模型关于分组机制，从而提高性能。

本文旨在解决transformer语言模型中推理时间长的问题，提出了一种Weighted Grouped-Query Attention（WGQA）的变体来减少推理时间。

本文提出的WGQA方案引入新的可学习参数，使得模型在微调过程中能够进行加权平均，从而减少推理时间，同时保持性能。

实验结果显示，WGQA相比于Grouped-Query Attention（GQA）平均提高了0.53%的性能，并且在推理过程中没有额外的开销。此外，本文还展示了T5-small和T5-base架构之间的性能比较和分析，为后续的研究提供了参考。

此前的研究提出了Multi-Query Attention（MQA）和GQA等方案来减少推理时间，本文则在GQA的基础上提出了WGQA。