Deep Learning

Attention / Conv 大锅烩

Self-Attentions and Convolutions

Renovamen
2021-08-31
15 min

长期记录和实现看过的各种论文里的自注意力和卷积机制,咕咕咕,实现地址在: Github

提前定义,FRN×dF \in \Reals^{N \times d} 是一个特征,在文本任务里,NN 是文本长度,dd 是词向量维度;在视觉任务里,NN 是特征图的 W×HW \times Hdd 是通道数量。

# Attention

# Self-Attenion

Attention Is All You Need. Ashish Vaswani, et al. NIPS 2017. [Paper] [Code]

self-attention

图片来源:论文 Beyond Self-attention: External Attention using Two Linear Layers for Visual Tasks

Q,V,KQ, V, K 由同一个值 FRN×dF \in \Reals^{N \times d} 经过线性变换得到,然后 attention map 和 context 的计算公式为:

A=softmax(QKTdk)Fout=AV\begin{gathered} A = \text{softmax}(\frac{QK^T}{\sqrt{d_k}}) \\ F_{\text{out}} = A V \end{gathered}

其中 dk=d/hd_k = d / h 是一个注意力头的维度,hh 是注意力头的数量。除以 dk\sqrt{d_k} 是因为,dkd_k 越大 QKTQK^T 就会越大,可能就会将 softmax 函数推入梯度极小的区域,导致更新缓慢,所以要用 dk\sqrt{d_k}QKTQK^T 进行缩放。

最后 FoutF_{\text{out}} 会经过线性层、残差连接和 Layer Normalization。

一个简化版本是,把线性变换扔掉,直接用 FF 计算 attention,即 Q=K=V=FQ = K = V = F,从而省掉那三个线性层的计算量:

simplified-self-attention

图片来源:论文 Beyond Self-attention: External Attention using Two Linear Layers for Visual Tasks

A=softmax(FFTdk)Fout=AF\begin{gathered} A = \text{softmax}(\frac{FF^T}{\sqrt{d_k}}) \\ F_{\text{out}} = A F \end{gathered}

Self Attention 是在计算任意两个位置之间的相关性,得到一个 N×NN \times N 的 attention map,所以计算复杂度是 O(dN2)O(d N^2)(“简化版本”也是这样),这是一个很高的计算开销。所以 Bert 的最大长度只有 N=512N = 512。而对于视觉任务来说,N=W×HN = W \times H,所以特征图也不能太大,一般 WWHH 都需要先被降到 10+ 的数量级。

# SAGAN

Self-Attention Generative Adversarial Networks. Han Zhang, et al. ICML 2019. [Paper] [Code]

SAGAN 是一个用 self-attention 来替代了卷积层的 GAN,它的 self-attention 跟原始 transformer 的 self-attention 的不同在于:

  • 在将 FF 通过线性变换得到 QQKK 时,进行了降维(dd/8d \to d / 8
  • 在残差连接之前,对 FoutF_{\text{out}} 进行了缩放:γFout\gamma F_{\text{out}},其中 γ\gamma 是一个可学习的参数

# External Attention

Beyond Self-attention: External Attention using Two Linear Layers for Visual Tasks. Han Zhang, et al. arXiv 2021. [Paper] [Code]

external-attention

图片来源:论文 Beyond Self-attention: External Attention using Two Linear Layers for Visual Tasks

External Attention 主要考虑了以下问题:

  • 原始 self attention 会计算任意两个位置(pixel)之间的相关性,计算量太大。但实际上,attention map 是很稀疏的,只有很少量的 pixel 之间有关联,所以没有必要用 N×NN \times N 的 attention map
  • 原始 self attention 只考虑了同一个样本中不同 pixel 之间的关系,但没有考虑不同样本间的关联

因此 External Attention 提出用 external memory unit MkM_kMvM_v(都是可学习的参数)来代替 KKVV

A=Norm(QMkT)Fout=AMv\begin{gather} A = \text{Norm}(Q M_k^T) \\ F_{\text{out}} = A M_v \end{gather}

其中 MRS×dM \in \Reals^{S \times d}SS 是一个超参。

对于第一个问题,O(dN2)O(d N^2) 的复杂度来源于 Q,K,VQ, K, V 的矩阵相乘,而这里 K,VK, V 的维度从 N×dN \times d 变成了 S×dS \times d,因此复杂度变为了 O(dSN)O(d S N),与 NN 线性相关。因为 SS 可以取得远远小于 NN(在论文的实验中,S=64S = 64 的时候效果就已经很好了),所以计算复杂度可以降低很多。

对于第二个问题,可以认为 external memory unit 学习到了整个数据集的信息,因此考虑了不同样本间的关联。

(1)(1) 中的 Norm()\text{Norm}(\cdot) 是 double-normalization 操作。Softmax 只会对一个维度进行归一化,但 attention map 对于输入特征的 scale 比较敏感,所以 double-normalization 在 softmax 之后对另一个维度也进行了归一化:

a~=QMkTa^ij=exp(a~ij)kexp(a~kj)Softmaxaij=a^ijka^ik\begin{gathered} \tilde{a} = Q M_k^T \\ \hat{a}_{ij} = \underbrace{\frac{\text{exp} (\tilde{a}_{ij})}{\sum_k \text{exp} (\tilde{a}_{\textcolor{red}{k}j})}}_{\text{Softmax}} \\ a_{ij} = \frac{\hat{a}_{ij}}{\sum_k \hat{a}_{i\textcolor{red}{k}}} \end{gathered}

# Fastformer

Fastformer: Additive Attention Can Be All You Need. Chuhan Wu, et al. arXiv 2021. [Paper] [Code]

又是『XXX is all you need』系列的题目,让人审美疲劳。虽然这个题目包含了它应该包含的信息:“比原始 Transformer 快,因为我们用了 additive attention”,但还是让人审美疲劳。

fastformer

图片来源:论文 Fastformer: Additive Attention Can Be All You Need

同样是想解决 N×NN \times N 的稀疏 attention map 带来的高计算量的问题。它的处理方式是先用 additive attention 把 QQ 融合成了一个 global query vector QglobalRdQ_\text{global} \in \Reals^d

Aq=softmax(QWqTdk)Qglobal=AqQ\begin{gathered} A_q = \text{softmax}(\frac{Q W_q^T}{\sqrt{d_k}}) \\ Q_\text{global} = A_q Q \end{gathered}

其中 WqTRdW_q^T \in \Reals^d(这里直接把 NN 维降到了 1 维)。然后用 KKQglobalQ_\text{global} 矩阵 element-wise 相乘,乘完之后再用和上面一样的 additive attention 得到一个向量 KglobalRdK_\text{global} \in \Reals^d

K=KQglobalAk=softmax((KWkTdk)Kglobal=AqK\begin{gathered} \overline{K} = K \odot Q_\text{global} \\ A_k = \text{softmax}(\frac{(\overline{K} W_k^T}{\sqrt{d_k}}) \\ K_\text{global} = A_q \overline{K} \end{gathered}

最后用 KglobalK_\text{global}VV 矩阵相乘,得到最终的输出:

Fout=KglobalVF_{\text{out}} = K_\text{global} V

因为有 additive attention 这个降维操作,所以上面每一个矩阵乘法的复杂度都是 Q(Nd)Q(Nd),element-wise 乘法的复杂度也是 Q(Nd)Q(Nd),所以总的时间复杂度是线性的 O(Nd)O(Nd)

# HaloNets

Scaling Local Self-Attention for Parameter Efficient Visual Backbones. Ashish Vaswani, et al. CVPR 2019. [Paper]

halo

图片来源:论文 Scaling Local Self-Attention for Parameter Efficient Visual Backbones

用 self-attention 来做类似于卷积的操作,心路历程是:

  • 如果用全局 self-attention,计算量太大;如果用卷积中的滑动窗口,在每个窗口里做 local self-attention,由于每个窗口的结果都要一直存在显存里没法释放,容易造成显存溢出
  • 但卷积中,两个相邻的滑动窗口中的大部分元素都是一样的,为了减少这部分重复的计算量,论文直接将输入的特征图分成了不重复的一些 block(上图中的 blocked image),然后在每个 block 里做 local self-attention,所以论文管这个操作叫 blocked local self-attention
  • 但只在 block 里做 self-attention 也是不合适的,周边共享的元素多少也要考虑一些,因此论文在每个窗口边上都 pad 了一圈(上图中的 neighborhood windows 周边那圈五颜六色的东西),以增加感受野。因为 pad 的这一圈看上去很像光环,所以论文管这个操作叫 haloing

# Linformer

Linformer: Self-Attention with Linear Complexity. Sinong Wang, et al. arXiv 2020. [Paper]

linformer

图片来源:论文 Linformer: Self-Attention with Linear Complexity

从实现上来看,和原始 Transformer 的区别就是,对 K,VK, VE,FRN×kE, F \in \Reals^{N \times k} 进行了降维(上图中的 Projection),即:

Fout=softmax(Q(EK)Tdk)P:N×kFVk×dF_{\text{out}} = \underbrace{\text{softmax} \left ( \frac{Q (E K)^T}{\sqrt{d_k}} \right)}_{\overline{P}: N \times k} \cdot \underbrace{F V}_{k \times d}

kk 是一个超参常数,所以时间复杂度降到 Q(Nkd)Q(Nkd)。同时它还使用了一些参数共享的 trick 来进一步降低计算量。

但论文还证明了为什么可以用 P\overline{P} 来近似原始的 attention map PP(虽然我觉得略显强行)。原文的 Theorem 2 先证明了当 k=5log(nd)/(ϵ2ϵ3)k = 5 \log(nd) / (\epsilon^2 - \epsilon^3) 时,一定存在 E,FE, F 使得 P\overline{P} 可以近似 PP。这里 kk 的取值与 nn 有关,复杂度还不是线性。但如果考虑到 rank(A)=d\text{rank}(A) = d 这一点(我还没搞明白这一点是怎么来的),kk 的选择就可以独立于 nn,从而变为线性复杂度。

# Convolution

# Selective Kernel

Selective Kernel Networks. Xiang Li, et al. CVPR 2019. [Paper] [Code]

希望能够自适应地调整感受野的大小,为了做到这一点,论文采取的方式是用多个大小不同的卷积核生成特征图,然后把这些特征图融合在一起。

selective kernel

图片来源:论文 Selective Kernel Networks

  • Split:用 MM 个大小不同的卷积核生成特征图,上图中举的是 3×33 \times 35×55 \times 5 两个卷积核的例子
  • Fuse:把不同卷积核生成的特征图相加,然后经过平均池化 \to 全连接 \to Batch Norm \to ReLU,得到一个特征向量 zRdz \in \Reals^d,其中 d=max(C/r,L)d = \max (C / r, L)r,Lr, L 都是超参,用于控制通道维度
  • Select:用 zz 来计算一个 attention map,然后把 MM 个特征图加权求和,得到最终的输出

# Squeeze-and-Excitation

Squeeze-and-Excitation Networks. Jie Hu, et al. CVPR 2018. [Paper] [Code]

赋予不同通道不同的权重,这样就可以加强重要的通道特征。

squeeze and excitation

图片来源:论文 Squeeze-and-Excitation Networks

通道权重的计算方式是:

  • Squeeze:先把卷积出的特征图 uRH×W×Cu \in \Reals^{H \times W \times C} 平均池化成一个向量 zRCz \in \Reals^C,这一步抽离了空间上的信息,这样后面就可以只关注通道之间的关系

  • Excitation:zz 只能体现在当前 batch 上计算出来的通道关系,所以论文加了两个全连接层来利用整个数据集的信息:

    s=σ(W2δ(W1z))s = \sigma (W_2 \delta(W_1 z))

    其中 δ\delta 是 ReLU 激活函数,W2W_2W1W_1 是两个全连接层,sRCs \in \Reals^C 是最终各个通道的权重

最后把特征图 uu 跟权重 ss 相乘就行。

# Involution

Involution: Inverting the Inherence of Convolution for Visual Recognition. Duo Li, et al. CVPR 2021. [Paper] [Code]

普通的卷积有两个特点:空间不变(spatial-agnostic)(不同空间位置共享同一个卷积核)和通道特异(channel-specific)(不同通道对应不同的卷积核)。参数数量为 Ci×Co×K×KC_i \times C_o \times K \times K,其中 Ci,Co,KC_i, C_o , K 分别为输入和输出通道数量,以及卷积核大小。

而这篇论文认为:

  • 由于通道数量 Ci,CoC_i, C_o 往往成百上千,所以为了限制参数量和计算量,KK 的取值往往较小,从而限制了感受野大小
  • 空间不变的特征可能会剥夺卷积核适应不同空间位置的能力
  • 卷积核在通道维度往往存在冗余,即很多卷积核是近似线性相关的。也就是说,如果把每个输出通道对应的卷积核铺成一个 Ci×K2C_i \times K^2 大小的矩阵,那么矩阵的秩不会超过 K2K^2。所以缩减卷积核的通道维度可能并不会明显影响效果

因此论文提出了一种跟卷积相反的 involution 操作:空间特异(不同空间位置对应不同的卷积核)和通道不变(不同通道共享同一个卷积核),参数数量为 H×W×K×K×GH \times W \times K \times K \times GGG 表示 CiC_i 个输入通道分成 GG 组(每组 Ci/GC_i / G 个通道),每组通道共享同一个卷积核。

involution

图片来源:论文 Involution: Inverting the Inherence of Convolution for Visual Recognition

和卷积不一样的是,involution 的核是根据输入的特征图自动生成的,一个通用的形式是:

Hij=ϕ(XΨij)\mathcal{H}_{ij} = \phi(X_{\Psi_{ij}})

其中 Ψij\Psi_{ij} 表示坐标 (i,j)(i, j) 的某个邻域上的坐标集合,XΨijX_{\Psi_{ij}} 表示 Ψij\Psi_{ij} 对应的特征。在实例化这个形式时,论文用了比较简单的方式,即令 Ψij={(i,j)}\Psi_{ij} = \{(i, j)\} 这个单点,然后:

Hij=ϕ(X{(i,j)})=W1σ(W0X{(i,j)})\mathcal{H}_{ij} = \phi(X_{\{(i, j)\}}) = W_1 \sigma (W_0 X_{\{(i, j)\}})

论文另一个挺有意思的地方是,它指出 self-attention 公式也可以用这个式子来表示,即:

Ψij=ΩHij=ϕ(XΩ)=(XWQ)(XWK)T\begin{gathered} \Psi_{ij} = \Omega \\ \mathcal{H}_{ij} = \phi(X_\Omega) = (X W^Q)(X W^K)^T \end{gathered}

其中 Ω\Omega 是 query 对应的 key 的范围。