上一篇《Mamba学习(一):总体架构》提到,Transformer模型的主要缺点是:自注意力机制的计算量会随着上下文长度的增加呈平方级增长。所以,许多次二次时间架构(指一个函数或算法的增长速度小于二次函数,但大于线性函数),如线性注意力、门控卷积和循环模型,以及结构化状态空间模型(SSM)被开发出来,以解决Transformer在长序列上的计算效率低下问题。
那先来了解一下什么是线性注意力。
关于线性注意力,已经有大量的文章,这里直接引用。
一、Softmax注意力与线性注意力
参见这篇文章《Softmax注意力与线性注意力的优雅融合》
(文章链接:https://mp.weixin.qq.com/s/ndQTSmf3RCCLU0xjd-nrkQ)
由于 Softmax 注意力是平方复杂度,直接进行全局自注意力的计算往往会带来过高的计算量。
针对这一问题,先前的工作通常通过减少参与自注意力计算的特征数量的方法来降低计算量。例如,设计稀疏注意力机制(如 PVT)或将注意力的计算限制在局部窗口中(如 Swin Transformer)。尽管有效,这样的自注意力方法很容易受到计算模式的影响,同时也不可避免地牺牲了自注意力的全局建模能力。
与 Softmax 注意力不同,线性注意力将 Softmax 解耦为两个独立的函数,从而能够将注意力的计算顺序从 (query・key)・value 调整为 query・(key・value),使得总体的计算复杂度降低为线性。
下面的图可以看的很清楚。
二、Softmax注意力与线性注意力的公式表达
参见这篇论文《Demystify Mamba in Vision: A Linear Attention Perspective (Arxiv 2024.05)》(论文地址:https://arxiv.org/pdf/2405.16605)
从上面的公式可以看出,Softmax 注意力的问题是,Softmax是非线性的函数,如果没有 Softmax,那么就是三个矩阵 (query・key)・value 连乘 ,而矩阵乘法满足结合率,可以调整为调整为 query・(key・value)。
(query・key)得到是n✖n的矩阵,(key・value)得到的是d✖d的矩阵。在d远小于n的情况下,复杂度就大致从O(n^2)降低到O(n)量级。
线性注意力将 Softmax 解耦为两个独立的函数,从而能够将注意力的计算顺序从 (query・key)・value 调整为 query・(key・value),使得总体的计算复杂度降低为线性。
三、苏神的几篇文章
《线性Attention的探索:Attention必须有个Softmax吗?》
文章链接:https://mp.weixin.qq.com/s/ziY7y9Lcc4jXf08UtJvKsQ
文章探讨了线性注意力的几种形式。
文章链接:https://mp.weixin.qq.com/s/MhLXO_VE_VYb7iIJCgN_lQ
这篇文章分析指出:当输入序列长度较小时,是线性复杂度的 FFN 占主导,而不是注意力机制部分的计算。按照BERT的base版计算,当序列长度接近 5000 时,Transformer 的复杂度才真正体现出二次性!
《Performer:用随机投影将Attention的复杂度线性化》
文章链接:https://mp.weixin.qq.com/s/NKhp5TWPHEhLl76X3aECpg
《Transformer升级之路:从Performer到线性Attention》