Hydra Attention 真的快197倍吗

上篇blog说最近会提高点更新频率，这不马上就来了

1、前言

15号Meta发了篇新paper取名hydra attention，各大媒体宣称hydra比vanilla快了197倍，事实真是如此吗，答案是：No，实际上hydra attention的快依赖于你处理的序列长度，该paper的作者主要是在cv上做的测试，对于384px的图片，如果是12层的attention那么FLOPs能减少11%，224px的图片，12层attention的FLOPs减少4%。

如果你想找到一个速度提升足够多的attention，那hydra可能确实不适合你，如果你是好奇其实现，那就继续看下去吧~

2、何如提升attention的计算效率

自从GPT3出来后，大家都朝着大模型方面去卷，attention速度方面的优化工作感觉就越来越少了，19年、20年有很多attention的魔改用于提升attention速度，我把它们简单分为了两块：

• 稀疏化
• 低秩化

稀疏化的代表BigBird，当时取得了摘要的sota，稀疏化其思路就是把attention matrix变得稀疏一些，如下图，看似计算量好像减少了，但代码实现起来很难达到理论的优化效率。

低秩化代表作Performer，其思路是把attention的时间复杂度修改为线性的，这里简单回顾下attention的复杂度，假设序列长度是T，维度是D，那么计算attention的时间复杂度就是$T\ast T\ast D$，即$O(N^2)$，低秩化 attention做的是把希望把时间复杂度尽可能降到$O(N)$，如下图，既然attnetion matrix计算比较耗时，那么我们另辟蹊径先计算$K\ast V$，只是说这样计算是否合理？怎么样才合理？就是我们需要思考的了，而hydra其实就是走的这个路线。

3、Hydra Attention

接下来我们来看下hydra的全貌，上文我们遗留了一个问题，先计算$K\ast V$是否合理，先看下vanilla self attention
$$A(Q,K,V)=\mathrm{softmax}\left(\frac{Q{K}^T}{\sqrt{D}}\right)V$$
公式相信大家都很熟悉了，不过多赘述，softmax部分实际是就是在计算一个相似度矩阵，或者说相关性矩阵，我们把公式做个简化，定义一个sim函数来代替softmax部分
$$A(Q,K,V)=sim(Q,K)V$$
如果这个核函数是个可分解的函数，那么我们可以表示为
$$A(Q,K,V;\varphi )=\left(\varphi (Q)\varphi (K{)}^T\right)V$$
根据乘法的结合律我们有
$$A(Q,K,V;\varphi )=\varphi (Q)(\varphi (K{)}^{T}V)$$
咋们先不考虑$\varphi$函数，先再看下时间复杂度，这个时候时间复杂度就变成了$O(T{D}^2)$，而vanilla版本的时间复杂度是$O(T^{2}D)$，所以到底哪种方法快实际上取决你是你的序列长度长还是hidden size大，所以Performer的测试数据长度都是512,1024这样较长的文本，对于短文本根本没有提速的效果。

那么Hydra难道也只能提速长文本吗？self attention还有一个关键点multi head，多个head实际上对应了多个attention matrix，如果有H个头，那么vanilla时间复杂度是
$$O(H{T}^{2}D/H)=O(T^{2}D)$$
而对于Hydra而言时间复杂度是
$$O(HT(D/H{)}^2)=O(T{D}^2/H)$$
如果head的个数H等于D，那么时间复杂度就变成了
$$O(T{D}^2/H)=O(TD)$$
真正做到了线性级别的attention，头是真的多啊，这也是取名为Hydra的原因。

说完相似度，我们回过来再聊一下核函数，目前比较主流的核函数有以下几种

余弦相似度应该是最好理解的，两种激活函数可以理解为门机制，均值的方式，嗯…我也没理解，有想法的小伙伴可以留言讨论下。

4、提速到底有多少

看完理论部分来看看作者实验结果，作者主要是在图片上做的实验，其中224px的图片提速如下

384px的图片提速如下

2层的情况甚至还有精度的提升，有点出乎意料了。

5、实战环节

ok，接下来我们来复现看看效果到底咋样，我简单写了一个hydra attention，并对比了self attention的效果，场景是一个短文本4分类，每层之间加了layer_norm和残差防止梯度消失，跑5个epochs，初始学习率是1e-3，我们主要对比acc与耗时

	hydra	vanilla
1层	79.24/15	79.31/24
2层	78.89/17	77.99/34
3层	78.83/21	78.15/46
6层	77.99/33	77.64/85
8层	78.19/45	77.80/112
12层	78.44/57	76.36/161

在达到8层之后，vanilla版本模型收敛的并不好，我调小了学习略才稍微提升了些，这里更合理的方法应该是采用warm up，大家可以自行尝试。整体来看hydra效率确实有所提升，并且也比vanilla更容易收敛，但是我这里实验做的比较少，还需要更多的一些验证，所以单纯对比acc说服力不够，条件有限期待后续作者会带来nlp方面的实验结果。

最后简单说下hydra的代码，github地址：https://github.com/terrifyzhao/attention

class HydraAttention(nn.Module):

    def __init__(self, hidden_size, all_head_size):
        super().__init__()
        self.query = nn.Linear(hidden_size, all_head_size)
        self.key = nn.Linear(hidden_size, all_head_size)
        self.value = nn.Linear(hidden_size, all_head_size)

    def forward(self, x):
        q = self.query(x)
        k = self.key(x)
        v = self.value(x)

        q = q / q.norm(dim=-1, keepdim=True)
        k = k / k.norm(dim=-1, keepdim=True)
        kv = (k * v).sum(dim=-2, keepdim=True)
        out = q * kv
        return out

代码很简单，熟悉self attention源码的读者看到这个代码可能会有疑问，为啥没有做head维度上的转置，self attention会把维度变为[batch,head_num,seq_len,head_size]再做matrix的计算，我们按照bert base版本的参数为标注，那么对于hydra来说head_num就是768，head_size就是1，此时维度就是[batch,seq_len,768,1]，转置后是[batch,768,seq_len,1]，$K^T\ast V$后就是[batch,768,1,1]，这个其实就是代码中的$(k\ast v).sum$

6、总结

简单测试下来hydra确实有点小用，但是在nlp方面的表现还是需要有算力的团队来做预训练模型进行对比。20年google出过一篇paper synthesizer ，大意是并不是所有的head都是有用的，如果只保留有价值的头，那么head size数还能进一步减少，也能进一步提升模型速度，感兴趣的读者也可以自己下来实验看看~

7、References

Hydra Attention: Efficient Attention with Many Heads

Big Bird: Transformers for Longer Sequences

RETHINKING ATTENTION WITH PERFORMERS

Poly-NL: Linear Complexity Non-local Layers with Polynomials

Synthesizer: Rethinking Self-Attention for Transformer Models