新知一下
海量新知
6 5 2 3 1 9 7

比标准Attention快197倍!Meta推出多头注意力机制“九头蛇”

OpenCV中文网 | 2022/09/23 17:46

尽管Transformer已经开始在诸多视觉任务上“大展身手”,但还有一个问题。

那就是在处理 大图像 上计算比较费劲。

比如面对一个1080p的图时,它会有 超过60%的计算量 都耗费在了创建和应用注意矩阵上。


究其原因,主要是因为自注意力头的数量是token的平方,而token的数量又与图形大小呈二次方的关系。

那能怎么办呢?

好消息是——

现在Meta捣鼓出了一种 多头注意力 操作方法,可以做到足足 比标准注意力快197倍

而且在提高计算速度的同时,它也不会牺牲准确率,有时甚至还能将准确率提高1-2个点。

新知达人, 比标准Attention快197倍!Meta推出多头注意力机制“九头蛇”

具体怎么回事儿?

思路来源一个“矛盾点”

这个方法名叫 Hydra Attention ,主要针对Vision Transformer。

(“Hydra”有“九头蛇”之义,来自希腊神话。)

Hydra Attention的思路源于线性注意力中的一种 有点矛盾 的点:

使用标准的多头自注意力,再向模型中添加更多头可以保持计算量不变。

而在线性注意力中改变操作顺序后,增加更多的头实际上还会降低层的计算成本。

于是,作者通过将模型中的 注意力头数量设置成特征 (feature) ,创建出一个token和feature的计算都是线性的注意力模块,从而把上述特性发挥到极致。

具体来说:

新知达人, 比标准Attention快197倍!Meta推出多头注意力机制“九头蛇”

当标准自注意力头是token数的平方 (O(T2D)) 时,通过使用可分解核 (decomposable kernel) ,我们重新安排操作顺序,让注意力头的数量变为特征D的平方。

然后再使用Hydra Trick,最大化注意力头H的数量,让H=D,最终就可以化为一个在空间和时间上的O (TD) 简化操作了。

其中,Hydra Trick的依据见下图:

新知达人, 比标准Attention快197倍!Meta推出多头注意力机制“九头蛇”

作者在ImageNet-1k上训练了具有不同头数的DeiT-B模型,包括使用标准自注意力 (蓝线,基于softmax) 和使用多头线性注意 (红线,基于余弦相似性) 的。

前者在H>96、后者在H<3时出现内存不足的情况。

当他们往模型中添加更多的头时,Softmax注意力模型的准确性似乎会崩溃,而多头线性注意力仍可以保持一致,因此就有了上述操作。

(需要注意的是,H必须除以D=768。)

速度快197倍,准确率还能更上层楼

来看看Hydra Attention交出的成绩单。

可以看到,Hydra注意力比标准注意力快197倍 (T=197)

随着图像大小的增加,它显著提高了模型的FLOP数,在创建和应用注意力矩阵所占的计算量上也 始终只有0.02%

新知达人, 比标准Attention快197倍!Meta推出多头注意力机制“九头蛇”

除此之外,作者发现,使用不同的内核,Hydra Attention还能让模型的准确率提高大约两个百分点。

新知达人, 比标准Attention快197倍!Meta推出多头注意力机制“九头蛇”

或者用Hydra Attention替换特定的注意力层,也能将模型的精度在ImageNet上提高1%或者与基线维持不变。

新知达人, 比标准Attention快197倍!Meta推出多头注意力机制“九头蛇”

新知达人, 比标准Attention快197倍!Meta推出多头注意力机制“九头蛇”

当然,最多可替换8层。

另外,作者表示,这个方法应该可以扩展到NLP领域,不过他们还没试。

新知达人, 比标准Attention快197倍!Meta推出多头注意力机制“九头蛇”



更多“三个”相关内容

更多“三个”相关内容

新知精选

更多新知精选