GQA: Training Generalized Multi-Query Transformer Models fromMulti-Head Checkpoints

https://arxiv.org/abs/2305.13245v1

Multi-head 推理太费时间。

多查询注意力（MQA）可以减少这种开销，仅使用单个键值头极大地加快了decoder推理速度，但会降低模型的容量和性能，并目不想为了更快的推理而单独训练一个模型。

于是本文提出了一种方案，可以将现有的Multi-head 多头注意力模型升级为多查询注意力模型MQA，仅使用原始预训练计算量的5%。

本文还介绍了分组查询注意力 (GQA)，这是一种广义的多查询注意力，它使用中间数量的键值头 （大于一个，小于查询头数）。本文表明，经过训练的GQA可以实现与Multi-head 相当的性能和与Multi-query相当的违度。

多头注意力有同样的查询、键和值个数。

多查询注意力在所有查询头之间共享单个键和值头。

分组查询注意力则为每组查询头共享单个键和值头。

多头注意力机制（MHA）转换为多查询注意力机制（MQA）：

将所有头的键（key）和值（value）投影矩阵进行平均池化，合并成单个头部的矩阵。

为什么使用平均池化，这是消融实验。

T5 XXL模型在MQA和GQA下性能随retraining比例变化的情况。GQA在转换后已经实现了合理的性能，而MQA需要进行retraining才能有用。随着retraining比例增加至5%，MQA和GQA均有所提高，但是从10%开始回报递减。

GQA组数量对推理速度的影响，论文组数量使用8。

 经过5%retraining的GQA可以实现与Multi-head相当的性能和与Multi-query相当的速度。