加权分组查询注意力机制在Transformer中的应用

研究者 : Sai Sena Chinnakonduru†, Astarag Mohapatra†
所属机构 : 印第安纳大学布卢明顿分校
通信邮箱 : saischin@iu.edu, astmohap@iu.edu

摘要

作为Transformer语言模型的核心构成单元，注意力机制在近期研究中展现出重要价值。相关研究表明，当模型容量持续提升时，其表现水平可接近人类认知能力。不过，在模型体量不断扩展以及受限于硬件存储条件的双重影响下，推理阶段的资源消耗也呈现出明显上升趋势。为有效缓解推理耗时问题，多查询注意力（MQA）与分组查询注意力（GQA）等新型机制相继被提出。本研究在此基础上设计了一种改进型分组查询注意力结构，即加权分组查询注意力（WGQA）。该方法通过在T5解码器的注意力模块中为每组键和值头配置可训练参数，在模型微调过程中实现加权平均操作。实验结果显示，相较于GQA基准方案，所提方法在性能指标上实现了0.53%的平均提升，并且在推理阶段未引入额外计算负担。通过对T5-small与T5-base两种架构进行对比测试，进一步验证了该技术方案的可扩展性特征。

引言

语言模型的基础架构依赖于自回归Transformer模型，其运作方式是依据输入序列以及此前生成的输出序列，依次生成新的标记。这一过程具有顺序性，其运算负担主要受