Multi-Head Attention和Multi-Query Attention的计算分析

本文主要是阅读论文《Fast Transformer Decoding: One Write-Head is All
You Need》的学习记录，这是一篇2019年的改善Multi-Head Attention带来的显存占用瓶颈问题的一种解决方案。

本文根据论文的撰写顺序，从Multi-Head Attention的代码开始，分析MHA在训练阶段和推理阶段的计算量和显存占用参数量；然后给出Multi-Query Attention的代码，分析MQA在训练阶段和推理阶段的计算量和显存占用情况。

提前预备：

1. 假设读者了解decoder解码器的结构，包括L层Transformer，每层Transformer由一个self-attention层和两个MLP层构成，self-attention层默认为采用H头的注意力层。
2. 了解在训练过程中，是通过mask来实现并行训练的。如果对这部分不是很熟悉，可以去看一下相关的知乎或者博客。
3. 了解在推理过程中，模型的输入是当前时刻m的word embedding输入，输出是当前时刻的输出。在中间计算过程中需要用到过去所有时刻的输入来计算Key和Value，实现注意力机制，因此max_length越长的模型，推理的时间开销越大，因为模型需要一次次的重新计算key和value张量。因此在实现中往往采用KVcache的方式来进行加速。

关于KVcache的概念和计算可以参考知乎的文章《分析transformer模型的参数量、计算量、中间激活、KV cache》
在这里为了加快模型的推理速度，网络的输入为当前时刻$index=i$的word embedding，模型会在self-attention层保存前面$(i-1)$个$K,V$的值$previous\_K,previous\_V$，然后在得到当前$index=i$的$new\_k,new\_v$之后，
更新$previous\_K=concat(previous\_k,new\_k)$,
更新$previous\_V=concat(previous\_v,new\_v)$;
这样做的好处是加快当前token的推理速度，缺点是增加了模型的显存占用。且随着max_length的增加，这部分的$K{V}_{cache}$会成为推理的显存占用瓶颈。

一、 Attention_function

这部分主要介绍的是self-attention中的q,K,V是怎么实现注意力机制的。
需要注意的是这里的$q$为$m$时刻的hidden表示，而$K,V$为序列中所有时刻的hidden表示，因此维度为$m$.
输入：单个时刻的输入$q$,维度为$(1,k)$,前面m个时刻的$K$，维度为$(m,k)$,前面$m$时刻的$V$，维度为