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

URL

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