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$，维度为$(m,v)$
输出：根据自注意力机制得到的输出$y$

首先计算$q$和$m$个key的相似度，得到一个$m$维的权重，然后将权重进行sigmoid标准化。这样得到的权重的每一个$W_i$都表示q和第$i$个Key的相似度或者说相关度，其值越大表示越相关。
将这个权重和对应的$value$相乘，也就是赋予对应位置的$V_i$更高的注意力权重。最后的返回值就是这个将带权重的Value叠加作用到最终结果的一个过程。

上面这个过程就是一个简单的attention过程。

二、Multi-head Attention

从上述简单的attention过程，引入Multihead的概念。

• 输入$x$:张量维度为$d$
• $M$为所有$m$时刻的输入，维度为$[m,d]$
• $P\_q$实现的功能为：1. 提供一个线性映射转换，2. 转换后的$x$ reshape成$h$个head，用于表征Multi-head这个概念。
• $P\_k$,$P\_v$的功能为：1. 为$m$个时刻的输入向量$M$提供一个线性映射变换，2. 将转换之后的$key$,$value$ reshape成Multi-head
• $h$表示head的数量，$k$为每个head中张量的维度，一般的实现中，$h\times k=d$，而且$k=v$,也就是key和value的最后一个维度是一致的。

在论文的表示中，大写的$KV$表示张量名称，小写的$k、v$表示长度。需要注意不要搞混了。

可以看到Multi-head Attention和上面的最简单形式的attention相比：

1. 多了$qKV$的计算方式，不再是直接存在$q,K,V$，而是给出了$KV$是通过$m$个时刻的输入和投影矩阵$P\_k,P\_v$计算得到，$q$是通过输入$x$和$P\_q$计算得到的。
2. 多了head的概念，将当前时刻的输入$x$和$key$，$value$的计算结果都reshape成h个集合。
3. 输出也不是简单的通过value的权重和得到，而是需要增加一个输出映射$P_o$

三、训练过程中的Multi-Head-attention

训练过程和上面这个过程的区别在于

1. 训练过程多了batch这个维度，可以同时计算不同序列的输出。提高了显存利用率。
2. 训练过程有GT，知道每个时刻的输入和输出，因此可以通过mask方式来实现并行计算，可以计算同一个序列上不同位置的输出。而不需要向推理那样等待前一个时刻的输出，来作为后一时刻的输入。

符号说明

• $m$表示的整个序列的输入，通过mask来实现遮蔽指定位置之后的其他输入。
• 这里的$n$表示$n$个时刻的输入（因为原因2，所以可以并行前向推理），也就是可以通过一次前向过程计算$n$个时刻的输出。一般实现的时候，$n=m$。这里的$n$一般为序列的长度，也就是会在依次前向过程中预测整个序列的每一个元素。通过$mask[n,m]$来实现mask掉后面的输入。具体来说，假设$m=n=10$,那么对于时刻$i=4$来说，他的$mask=[1,1,1,1,-inf,-inf,-inf,...]$，这样在计算KV计算只使用该时刻之前的输入，而不使用后面时刻的输入。

其余的和上面的基本没有什么变化

分析一下计算量，建立在

1. $b=batchsize$,
2. $n=m=$序列的长度，
3. $h\ast k=d$，
4. $k=v$

这四条实践经验的基础上，我们知道矩阵乘法的计算量为最终输出的元素个数乘以每个元素的计算量。 我们以矩阵一次乘-加运算作为计算基础。

• $QKV$的计算量都是$bn{d}^2$,那么总计算量为$3bn{d}^2$
• logits的计算量为$bh{n}^{2}k=b{n}^{2}d$
• weights的计算量为 （忽略一下）
• O计算量为$bhnv\ast n=bh{n}^{2}v=b{n}^{2}d$
• Y的计算量$bnd\ast d=bn{d}^2$

综合来看，计算复杂度为$\theta (bn{d}^2)$,考虑到$n⩽d$的情况:

计算参数量

• 1.$X,M,Q,K,V,O,Y$的参数都是bnd,那么总参数量为$6bnd$,
• 2.logits和weights总参数量为$2bh{n}^2$,
• 3.参数$P\_v,P\_q,P\_k,P\_o$的参数总量为$4d^2$.

考虑到1的部分元素是中间激活值，不会保存，2的全部元素是中间激活值，不会保存。那么空间复杂度可以表示为O(bnd+d^2)

四、推理过程的MHA

这里着重讨论的是生成式模型，因为生成式模型的在某个时刻$i$的输出依赖于从0时刻到当前时刻的所有输入。且在当前时刻的输出得到之前，无法计算后面时刻的输出，因此不像训练阶段那样可以并行执行前向过程。

但是前面时刻的输入主要是为了计算$KV$，$q$只和当前时刻的输入有关。因此想要加速推理，可以将$KV$在推理过程中进行缓存保留。这也就是$KVcache$，这种方式存在的问题是$KVcache$是需要占用显存的。

这个公式和上面训练阶段的公式的区别在于$KV$的计算.

$KV$是通过之前时刻的$prev\_K$和当前时刻根据当前输入$M$（我觉得这里的$M$和$X$表达的是一个意思，都是当前时刻的输入，为了和前面一章节的公式保持一致所以才这么写的。但是不确定我的这种理解对不对，欢迎讨论或者告知这里的$M$是什么）计算得到的$new\_K$,$new\_V$.

类似上一章节的计算复杂度推理，这里的考虑到不能并行计算的情况，那么对于长度为$n$的序列，需要计算$n$次，那么计算复杂度为$\theta (bn{d}^2)$；

空间复杂度也是同样的，在一次计算中，不考虑中间激活值，主要保存是是$KVcache$，以及矩阵$P\_q,P\_k,P\_v,P\_o$.前者的参数量为$O(b{n}^{2}d)$,后者的参数量为$O(d^2)$。可以看到随着序列长度$n$的增加或者隐藏层深度$d$的增加，这部分的显存占用将会称为瓶颈。

五、Multi-Query Attention

介绍完MHA之后，就可以介绍MQA了。

左边是MHA，右边是MQA，中间是没有提到过的GQA（Group-Query Attention，可以理解为不那么激进的MQA）。

可以看到MHA中每个head独立维护自己的$KVcache$。在MQA中则是同一层Transformer中的所有head维护一个$KVcache$，那么$KVcache$的显存可以明显的减少$h$倍。

另外KV的投影映射矩阵$P\_k,p\_v$也会同样的减少$h$倍。

这里好像就没有什么好讲的啦~~和MHA的第三章节唯一的不同就是算法前面提到的两点。其他都是一样一样的

六、 推理阶段的MQA

$x$是某个时刻的输入,维度为$[b,d]$
$KVcache$的更新不再独立于每个head，而是所有head共享，因此$P_k,P_v,prev\_K,prev\_V$的维度都少了一维。

七、实验对比

作者在英文转德文的数据集上进行了实验，采用的模型是具有2.1亿参数、6层layers，d_model=1024,head=8，且位置编码可学习的模型。
在对比MHA和MQA的时候，为了保持总参数量不变，会将MQA中的MLP层的hidden layer从1024*4增大为5440.

训练的效果评估指标如下所示，可以看到相比于MHA有轻微的性能衰减。PPL指标越小越好，BLEU是用于衡量翻译质量的数据集，越大越好。

从训练时间和推理时间上看，训练时有轻微的减少。对于单个training step从MHA的433ms减少到MQA的425ms。这里的训练上并没有很明显的提升。

在推理上的提升是很明显的。实验是在一个英文转德文的数据集上进行的，encoder的长度为128，decoder的长度也是128。encoder模块推理需要222ms,在MHA中，decoder需要47ms。在MQA上，encoder的时间为195ms，decoder只需要3.9ms
下表中时间是每个token上的时间。所以和上面我们提到的时间看起来不同，实际是一样的。