Guo M H, Liu Z N, Mu T J, et al. Beyond self-attention: External attention using two linear layers for visual tasks[J]. arXiv preprint arXiv:2105.02358, 2021.
清华

自注意力机制在同一个样本内, 任意一个部位的特征都可以聚合所有位置的特征进行加权输出。但是自注意力拥有二次复杂度, 并且不能计算多个样本之间的潜在联系。

External-Attention(EAT) 希望在学习某个数据集时, 能够找到多个样本之间的潜在联系。其通过保持一定的key memory, 以找到跨越所有样本的最具有辨识性的特征。这种思想类似于sparse coding 和 dictionary learning。并且由于key memory设计的很小, 因此EAT计算上具有O(n)的复杂度, 比起自注意力高效很多。

External Attention (EAT)

\[ \begin{aligned} A&=(\alpha_{i,j})=Norm(FM_k^T)\\ F_{out}&=AM_v \end{aligned} \]

\(F \in \mathbb{R}^{N \times d}\), 其中\(M \in \mathbb{R}^{S \times d}\)是一个独立于所有的输入特征进行学习的单元,即贯穿整个数据集的一个记忆模块, 实践中使用了两个不同的单元 \(M_k\) 和 \(M_v\) 来增强效果。\(\alpha_{i,j}\) 用于衡量第i个词元与\(M\)中第j行记忆特征的关系。

Normalization: 在自注意力中attention map是由余弦相似度的组成的一维向量, 因此可以使用Softmax来进行权重归一化。EAT中假如把图像的每个像素看作一个词元, 那么这个图像的attention map就是一个二维矩阵, 因此使用了二维的softmax来进行归一化, 即double-normalization(2021, PCT):

\[ \begin{aligned} (\tilde{\alpha})_{i, j} &=F M_{k}^{T} \\ \hat{\alpha}_{i, j} &=\exp \left(\tilde{\alpha}_{i, j}\right) / \sum_{k} \exp \left(\tilde{\alpha}_{k, j}\right) \\ \alpha_{i, j} &=\hat{\alpha}_{i, j} / \sum_{k} \hat{\alpha}_{i, k} \end{aligned} \]

Multi-head: 自注意力中常用多头机制来增强多视角的学习能力, 因此EAT中改造如下:

\[ \begin{aligned} h_i&=ExternalAttention(F_i,M_k,M_v) \\ F_out&=MultiHead(h,M_k,M_v) \\ &=Concat(h_1,...h_H)W_o \end{aligned} \]

\(h_i\) 为每个头计算的注意力输出, \(W_o\)是用于使维度一致的线性层。\(M\)为在多个头之间共享权重的memory单元。