self-attention原理和改进方向

1. self-attention原理介绍

self-attention的概念起源自seq2seq机器翻译模型中的attention，基本思想都是计算不同部分的权重后进行加权求和，最后的效果是对attend的不同部分予以不同程度的关注，不同是seq2seq的attention是将decoder的一个状态去对比encoder的所有状态，而self-attention可以在encoder和decoder中均可使用，其计算的是同一集合内元素的相互关联。Sequence to Sequence (seq2seq) and Attention很好对二者都有非常好的介绍。

关于self-attention的介绍这里就不详细展开了，重点部分：

可以看到self-attention的基本计算基本都是矩阵计算，其最大的优点是不包含任何RNN、CNN结构，可以解决序列的长程依赖问题，同时矩阵计算可以因为并行化而非常快速。
其缺点也很明显：
1、由于没有任何部分处理序列信息，因此需要额外的positional encoding模块；
2、attention矩阵维度与输入序列长度$t$是$O(t^2)$关系，如果序列较长非常消耗内存，目前的transformer最多能处理512长度的序列；
3、虽然可以对长程依赖有效处理，但是对短程的上下文信息反而处理的不好，这会引出后续很多改进工作

推荐几篇非常好的介绍：

形象化解释：The Illustrated Transformer
哈佛代码介绍：The Annotated Transformer
论文解读：Attention Is All You Need - The Transformer

看到知乎上一个有意思的讨论：self-attention和全连接有什么区别？
全连接：可以看做特征映射，其权重表示的是特征的重要度
self-attention：并不是一种映射，其权重表示的是序列每个实体的重要度

2. self-attention改进方向

1）增强局部上下文信息

Modeling Localness for Self-Attention Networks

构建局部上下文信息：当“Bush”对齐到“held”时，我们希望自注意力模型能同时将更多注意力放到其邻近的词“a talk”上。这样，模型能捕获短语“held a talk”。本文将局部上下文建模设计为一种可学习的高斯偏差矩阵$G$，其中心位置和宽度可以被学习，最终该偏差矩阵与原attention矩阵相加得到修正attention矩阵：$\widehat{Att}=softmax(Att+G)$

Convolutional Self-Attention Networks

本文将局部上下文的拟合构建为类卷积形式的计算，其中1D卷积将原序列长度的attention限制在一个1D卷积核内，而2D卷积则是扩展了head维度后将attention限制在一个2D卷积核内，这种方法可以不增加模型参数

可以参考腾讯AI LAB论文解析

2）降低与序列长度的平方依赖处理更长序列

Longformer: The Long-Document Transformer

1、使用固定窗口的稀疏卷积，堆叠多层同样可以达到全注意力的感受野。同时，增加空洞可以扩大感受野
2、在不同的head中使用不同的窗口和空洞可以让不同的head学习不同大小的局部信息。在底层使用小窗口，而在顶层使用大窗口
3、在特定的token上使用全注意力，这些token可以根据不同任务选择出来的

Linformer: Self-Attention with Linear Complexity

1、论文核心贡献是发现self attention矩阵$P$是低秩矩阵，从而可以通过SVD构建另一个维度小的低秩矩阵$\bar{P}$来近似，这个构建的矩阵仅选取前$k$大的奇异值对应的奇异向量
2、对矩阵进行SVD仍然是耗时操作，因此可以进一步对$K$和$V$引入两个线性变换$E,F:k\times T$来计算逼近$\bar{P}$
$$Attention(Q{W}_Q,EK{W}_K,FV{W}_V)=\underset{\bar{P}:T\times k}{\underbrace{softmax(\frac{Q{W}_Q(EK{W}_K{)}^T}{\sqrt{d}})}}\cdot \underset{k\times d}{\underbrace{FV{W}_V}}$$
这里时间和空间复杂度都变成了线性$O(Tk)$，如果选择$k\ll T$的话，时间和空间要求就会非常低
3、一些关于$E,F$的空间映射技巧可以使用，比如在head里、layer间共享权重，或者在不同的head、layer使用不同的$k$，甚至可以将简单的线性映射替换为pooling、conv等操作

Rethinking Attention with Performers（Masked Language Modeling for Proteins via Linearly Scalable Long-Context Transformers）

由于self attention矩阵的维度是$O(t^2)$的，因此如果不显式的计算就可以避免这个限制。这个想法很大胆但是实现起来确实不容易。文章通过正交随机投影和核方法将原最终输出的计算（1）变换为（2）
$$\begin{gathered} Y=D^{-1}AV,A=exp(\frac{Q{K}^T}{\sqrt{d}}),D=diag(A{1}_t)(1) \\ Y={\hat{D}}^{-1}(Q^{\prime }((K^{\prime }{)}^{T}V)),\hat{D}=diag(Q^{\prime }((K^{\prime }{)}^T{1}_t)),Q^{\prime }=D_Q\varphi (Q^T{)}^T,K^{\prime }=D_K\varphi (K^T{)}^T(2) \end{gathered}$$
其中$D_Q=g(Q_i^T),D_K=g(K_i^T)$分别是将行向量映射为标量的函数，$\varphi (X)$通常为正交随机映射，是核函数的一部分。由于$(K^{\prime }{)}^{T}V\in {\mathbb{R}}^{d\times d}$，其中$d$为特征维度常量，因此整体计算的复杂度为$O(t{d}^2)=O(t)$。关于这篇论文的证明部分，可以参考Performer：用随机投影将Attention的复杂度线性化

最近google的一篇论文对此总结Efficient Transformers: A Survey
相关论文参考github论文集

3）attention的kernel技术

Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention

将原始的Scaled-Dot Attention$(1)$输出结果的每一个row改写为$(2)$
$$\begin{gathered} Y=Attention(Q,K)V=softmax(\frac{Q{K}^T}{\sqrt{d}})\mathbf{V}(1) \\ {Y}_i=Attention(Q_i,K)V=\frac{{\sum }_{j=1}^{T}exp(\frac{Q_i\cdot K_j}{\sqrt{d}})V_j}{{\sum }_{j=1}^{T}exp(\frac{Q_i\cdot K_j}{\sqrt{d}})}=\frac{{\sum }_{j=1}^{T}sim(Q_i,K_j)V_j}{{\sum }_{j=1}^{T}sim(Q_i,K_j)}(2) \end{gathered}$$
只要定义相似度函数（核函数）$sim(Q_i,K_j)\ge 0$就可以定义各种类型的attention，原始的attention可以看做指数型，当然还可以有线性、多项式或RBF型的。论文为了实现线性化，将相似度函数通过特征映射函数表示为$sim(Q_i,K_j)=\varphi (Q_i)\cdot \varphi (K_j)$，通过乘法的交换律实现了$\left(\varphi (Q)\varphi (K{)}^T\right)V=\varphi (Q)\left(\varphi (K{)}^{T}V\right)$，其中选取$\varphi (x)=elu(x)+1$
关于这篇论文还可以参考线性Attention的探索：Attention必须有个Softmax吗？

Transformer dissection: An unified understanding for transformer’s attention via the lens of kernel

这篇论文从核函数的角度统一了多篇论文的改进，并把位置编码也融合进了核函数中。具体来说，原始Scaled-Dot Attention的分子部分可以表示成$k_{exp}(f_q+t_q,f_k+t_k)$，这里$f_{\cdot }$表示输入的embedding，$t_{\cdot }$表示绝对位置编码，这表示位置编码是直接加入embedding后再计算$k_{exp}$的，而这篇论文提出一种将位置编码独立计算kernel的计算方式：
$$\begin{gathered} k_F(f_q,f_k)\cdot k_T(t_q,t_k)， \\ where{k}_F(f_q,f_k)=exp(\frac{f_q{W}_F(f_k{W}_F{)}^T}{\sqrt{d}}),k_T(t_q,t_k)=exp(\frac{t_q{W}_T(t_k{W}_T{)}^T}{\sqrt{d}}) \end{gathered}$$
这里位置编码仍然是sin/cos表示的绝对位置，而共享$Q,K$的映射矩阵，这样使得attention矩阵最终是半正定的同时减少了参数（但据我实验效果会下降）。最后，论文表示只在$Q,K$中使用位置编码，$V$不使用任何形式的位置编码效果都会好一些。

4）位置编码

位置编码的应用方式有两种：将位置信息结合进输入序列中的方法称为绝对位置编码，而通过调整self-attention矩阵来计算不同位置上的注意力值的方法称为相对位置编码。还可以参考让研究人员绞尽脑汁的Transformer位置编码

方式一：绝对位置编码

Transformer原论文中就使用了固定的正余弦函数来表达序列token的绝对位置，它具有是固定的参数不随模型训练而改变，至于为何这种方式可以表达位置信息，可以参考这篇解释Transformer architecture: positional encoding。总结起来，一条经过embedding的隐层序列输入，设其token序列位置为$t$，隐层维度位置为$i$，那么正余弦位置编码$PE$的每个元素用公式表示如下：
$$P{E}_t^{(i)}=f(t{)}^{(i)}:=\{\begin{array}{ll}\sin ({\omega }_k.t),& \text{if}i=2k\\ \cos ({\omega }_k.t),& \text{if}i=2k+1\end{array},\text{where}{\omega }_k=\frac{1}{10000^{2k/d}}$$
这样的位置编码具有如下性质
1）每一个位置$t$的位置编码具有唯一性和确定性
2）不同长度的输入序列，任意两个位置的距离保持一致
3）用有限的数值可以表达无限长度的序列
其中第2步的证明如下，对于同频率的${\omega }_k$的一对正余弦组合，总有和位置$t$无关的旋转矩阵$M$满足
$$\left[\begin{array}{c}\sin ({\omega }_k.(t+\varphi ))\\ \cos ({\omega }_k.(t+\varphi ))\end{array}\right]=\left[\begin{array}{cc}\cos ({\omega }_k.\varphi )& \sin ({\omega }_k.\varphi )\\ -\sin ({\omega }_k.\varphi )& \cos ({\omega }_k.\varphi )\end{array}\right]\left[\begin{array}{c}\sin ({\omega }_k.t)\\ \cos ({\omega }_k.t)\end{array}\right]=M\left[\begin{array}{c}\sin ({\omega }_k.t)\\ \cos ({\omega }_k.t)\end{array}\right]$$
BERT则采用了一种更简单的可学习的位置编码，可以证明上述正余弦位置编码是可学习位置编码的一个显式解。

方式二：相对位置编码

相对位置编码考虑从绝对位置编码的计算过程来改造self-attention，绝对位置编码的计算如下
$$\{\begin{array}{rl}{\mathbf{q}}_i=& ({\mathbf{x}}_i+{\mathbf{p}}_i){\mathbf{W}}_Q\\ {\mathbf{k}}_j=& ({\mathbf{x}}_j+{\mathbf{p}}_j){\mathbf{W}}_K\\ {\mathbf{v}}_j=& ({\mathbf{x}}_j+{\mathbf{p}}_j){\mathbf{W}}_V\\ a_{i,j}=& softmax\left({\mathbf{q}}_i{\mathbf{k}}_j^{\top }\right)\\ {\mathbf{o}}_i=& \sum _j{a}_{i,j}{\mathbf{v}}_j\end{array}$$

Self-Attention with Relative Position Representations

Google的这篇论文首先展开${\mathbf{q}}_i{\mathbf{k}}_j^{\top }$
$${\mathbf{q}}_i{\mathbf{k}}_j^{\top }=\left({\mathbf{x}}_i+{\mathbf{p}}_i\right){\mathbf{W}}_Q{\mathbf{W}}_K^{\top }{\left({\mathbf{x}}_j+{\mathbf{p}}_j\right)}^{\top }=\left({\mathbf{x}}_i{\mathbf{W}}_Q+{\mathbf{p}}_i{\mathbf{W}}_Q\right){\left({\mathbf{x}}_j{\mathbf{W}}_K+{\mathbf{p}}_j{\mathbf{W}}_K\right)}^{\top }$$
其中${\mathbf{p}}_i{\mathbf{W}}_Q$和${\mathbf{p}}_j{\mathbf{W}}_K$都包含绝对位置信息，因此可以去掉${\mathbf{p}}_i{\mathbf{W}}_Q$，并将${\mathbf{p}}_j{\mathbf{W}}_K$替换成仅依赖$i-j$的相对距离矩阵${\mathbf{R}}_{i,j}^K$，并同理在value计算中引入相对距离${\mathbf{R}}_{i,j}^V$
$$\begin{array}{rl}{\mathbf{q}}_i{\mathbf{k}}_j^{\top }& ={\mathbf{x}}_i{\mathbf{W}}_Q{\left({\mathbf{x}}_j{\mathbf{W}}_K+{\mathbf{R}}_{i,j}^K\right)}^{\top }\\ {\mathbf{v}}_j& ={\mathbf{x}}_j{\mathbf{W}}_V+{\mathbf{R}}_{i,j}^V\end{array}$$
在最后的${\mathbf{R}}_{i,j}^K$和${\mathbf{R}}_{i,j}^V$的计算中进行截断，表示过远的相对距离可以不用考虑了，并且它们的取值既可以如同transformer选择固定的正余弦编码，也可以如同BERT一样选择可学习编码。

Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context

在Google工作的基础上进一步展开${\mathbf{q}}_i{\mathbf{k}}_j^{\top }$，直接将${\mathbf{p}}_j$替换为正余弦相对位置向量${\mathbf{R}}_{i-j}$，而将后两项${\mathbf{p}}_i{\mathbf{W}}_Q$分别替换为两个可学习的位置编码$\mathbf{u}$和$\mathbf{v}$，并将${\mathbf{W}}_K$拆分成${\mathbf{W}}_{K,p}$和${\mathbf{W}}_{K,R}$，原因是这两项分别是位置-key和位置-位置交叉项，对同一个${\mathbf{q}}_i$都有相同的位置，因此只需要分别学习基于位置的key表示和纯位置表示即可。对于${\mathbf{v}}_j$则选择直接去掉位置编码效果更好。
$$\begin{array}{rl}{\mathbf{q}}_i{\mathbf{k}}_j^{\top }& ={\mathbf{x}}_i{\mathbf{W}}_Q{\mathbf{W}}_K^{\top }{\mathbf{x}}_j^{\top }+{\mathbf{x}}_i{\mathbf{W}}_Q{\mathbf{W}}_K^{\top }{\mathbf{p}}_j^{\top }+{\mathbf{p}}_i{\mathbf{W}}_Q{\mathbf{W}}_K^{\top }{\mathbf{x}}_j^{\top }+{\mathbf{p}}_i{\mathbf{W}}_Q{\mathbf{W}}_K^{\top }{\mathbf{p}}_j^{\top }\\ & ={\mathbf{x}}_i{\mathbf{W}}_Q{\mathbf{W}}_K^{\top }{\mathbf{x}}_j^{\top }+{\mathbf{x}}_i{\mathbf{W}}_Q{\mathbf{W}}_K^{\top }{\mathbf{R}}_{i-j}^{\top }+\mathbf{u}{\mathbf{W}}_{K,p}^{\top }{\mathbf{x}}_j^{\top }+\mathbf{v}{\mathbf{W}}_{K,R}^{\top }{\mathbf{R}}_{i-j}^{\top }\end{array}$$

3. self-attention在其他领域尝试

虽然self-attention提出时主要用在序列模型中，但是其中包含的全局依赖思想可以轻松应用在图像、视频中，那么RNN/CNN所能够发挥的领域都有它的一席之地，甚至有替代传统模型的趋势，近期优秀的self-attention跨界表现主要有以下方面：

1）图像视觉领域

详见self-attention在图像视觉领域的发展

2）时序信号领域

近期的AAAI 2021最佳论文：Informer: Beyond Efficient Transformer for Long Sequence Time-Series Forecasting
另外还可参考self-attention序列数据分析attention-for-time-series-classification-and-forecasting