标准的Attention到稀疏Attention

简介

依旧借鉴苏神的博客【为节约而生：从标准Attention到稀疏Attention】本人只为加深理解，加深印象，所以在本博客对苏神所做的博客进行整合，大家如果有兴趣，可以看原文（杜绝抄袭）。苏神，yyds.
如今在NLP领域，甚至CV领域的transformer都大火。而且现在一众论文都指出，transformer中的attention机制起了至关重要的作用。也就是Q,K,V的交互，即Q,K能够给出向量之间的两两交互的分数，最终把这个分数投射至向量V上。

背景介绍

当然，提到Transformer就必须要提到这篇经典之作:Attention Is All You Need, 其阐述了Attention的一种构造方式，也是就乘性Attention,表示如下：
$\textcolor{red}{~~~~~~~~~~~~~~~~~~~~~~Attention(Q, K, V) = softmax(\frac{QK^T}{\sqrt{d^k}})V~~~~~~~~~~~~~~~~~~~~~}$

以上是transformer中的attention.
然而从理论上来说，self-attention的时间复杂度和空间复杂度都是 $O(n^2)$(n是输入的句子长度) 当句子长度变成2n,时间复杂度和空间复杂度都会增加4倍。当然，当显存的核数足够多时，时间复杂度未必会增加4倍，但是空间复杂度肯定会增加4倍的。所以一旦你的序列变长，transformer就很容易来个OOM.

稀疏Attention

我们通常说的标准的attention通常是指的一段序列中任意两个token都需要关联，得到一个相似度分数，最终得到一个 $n^2$的相似度矩阵.

左边即是标准的attention相似度矩阵,右图表示了 $\textcolor{red}{每个token都需要跟其他任意一个token相关联}$.正因为这种关联性，导致了计算复杂度是 $O(n^2)$.于是最自然的减少计算复杂度的想法即是 $\textcolor{red}{减少这种关联性计算}$，即token只跟一部分token进行关联性的计算，这就是稀疏attention的核心思想了

Atrous Self Attention

第一个要讲的稀疏Attention应用，即是Atrous Self Attention, 中文可翻译成：'膨胀自注意力'，'空洞自注意力'。很明显，这是引用自'空洞卷积（Atrous Convolution）',如下图所示，强行让token只关注相对距离为k, 2k, 3k,... 的元素, k>1, 而相对距离不为k或k的倍数，则会将相似度分数置为0.

由于计算注意力分数是跳着来的，所以每个token都只会跟约 $\frac{n}{k}$个其他的token进行相似度计算。于是计算的复杂度和空间复杂度理论上都变成了 $\textcolor{red}{O(n*(\frac{n}{k})) = O(n^2/k)}$, 也就是说复杂度降低为了原来的 $1/k$倍

Local Self Attention

Local Self Attention顾名思义，即是局部自注意力。也就是放弃标准Attention中的全局注意力思想，只关注局部的tokens.具体就是只关注当前token的前后k个token, 当然包括自己.如下图所示

相当于维护了一个大小为 $2k+1$的滑动窗口，所以从理论上来说，每个token都会跟 $2k+1$个token关联，所以这个时候的复杂度就变成了 $\textcolor{red}{O((2k+1)*n) \approx O(kn)}$。最终也是随着长度的增长而线性增长，当然，缺点就是牺牲了长依赖的关联性.

Sparse Self Attention

从上面的Atrous Self Attention和Local Self Attention可以看出，Atrous的注意力是带洞的，而局部自注意力刚好是紧密连接的，于是，把两者合并起来就是个自然而然的想法了。如下图所示

这样一来，Sparse Attention就将局部紧密关联性和远程稀疏关联完美的结合起来了。 以上即是关于Attention一些优化方面的尝试了。