Beyond Self-attention: External Attention using Two Linear Layers for Visual Tasks

URL

https://arxiv.org/pdf/2105.02358.pdf

TL;DR

• 本文提出一种 Multi-Head External Attention 使用两层全连接层和两个 normalization 层替代 transformer 的 attention 层，以降低计算复杂度
• 提出一种 DoubleNorm 标准化层替代 attention 中的 softmax 层
• 在一些任务中没有超越 transformer 达到 SOTA，只是提出一种关于 attention 的思考

Algorithm

网络结构

• external attention 对比 self attention

• multi-head external attention 对比 multi-head self attention

数学表示与伪代码表示

• self attention

  $A = (\alpha)_{i,j} = softmax(QK^T),\ \ \ (\cdot)\ means\ matrix$

  $F_{out} = AV$

  其中:

  $F\in\mathbb{R}^{N\times d}$ 表示输入，$Q=FW_1,K=FW_1,V=F,W_1\in\mathbb{R}^{d\times d'}$

  $Q\in\mathbb{R}^{N\times d'},K\in\mathbb{R}^{N\times d'},A\in\mathbb{R}^{N\times N}$

  $F_{out}\in\mathbb{R}^{N\times d}$ 表示输出

• 简化 self attention

  $A = (\alpha)_{i,j} = softmax(FF^T)$

  $F_{out} = AV$

• external attention

  $A = (\alpha)_{i,j} = Norm(FM^T)$

  $F_{out} = AM$

  以上两步可以共享 MLP 权重，$M\in\mathbb{R}^{S\times d}$，也可以不共享，分成 $M_k$、$M_v$

  计算复杂度 O(dSN)

  伪代码表示：

• multi-head external attention

$h_i = ExternalAttention(F_i,M_k,M_v)$

$F_{out} = MultiHead(F,M_k,M_v) = Concat(h_1,...,h_H)W_o$

伪代码表示：

• DoubleNorm

本文提出一种 double-normalization 结构，先在第一维做 softmax，再在第二维做 average

$(\tilde{\alpha})_{i,j} = FM_k^T$

$\hat{\alpha}_{i,j}=exp(\tilde{\alpha}_{i,j})/\sum_k exp(\tilde{\alpha}_{k,j})$

$\alpha_{i,j}=\hat{\alpha}_{i,j}/\sum_k \hat{\alpha}_{i,k}$

EA 结构用法示例

External Attention 和 DoubleNorm 的消融实验

算法表现

• 本文中对 EANet 与其他网络结构在不同任务上的对比实验做的非常详细，建议去看原文

Thought

• 在一些任务上超越了 self-attention 并不能说明 external attention 结构优于 self-attention 结构
• 在参数量较小的情况下，external attention 结构可带来较多增益
• 本文更像是对 《Attention is all you need》 的嘲讽之作