变权重分组查询注意力机制于 Transformer

注意力机制是 Transformer 语言模型的基础模块。最近的研究表明，通过扩大模型规模可实现接近人类水平的性能。然而，随着规模化需求的增加和硬件内存限制的提升，这些模型的推理成本仍然较高。为降低推理时间，提出了多查询注意力 (MQA) 和分组查询注意力 (GQA)。本文提出一种分组查询注意力的变体，称为加权分组查询注意力 (WGQA)。我们为 T5 解码器注意力模块中的每个 key 和 value head引入新的可学习参数，使模型在微调期间能够进行加权平均。我们的模型在 GQA 上实现平均 0.53% 的性能提升，推理期间无需额外开销即可收敛至传统多头注意力 (MHA)。我们评估了这些参数的引入及随后的微调，表明训练期间模型能了解分组机制，从而提升性能。此外，我们通过比较 T5-small 与 T5-base 架构的结果，展示了我们方法中的比例规律。