用MLP代替掉Self-Attention

這次介紹的清華的一個工作 “Beyond Self-attention: External Attention using Two Linear Layers for Visual Tasks”

用兩個線性層代替掉Self-Attention機制，最終實現了在保持精度的同時實現速度的提昇。

這個工作讓人意外的是，我們可以使用MLP代替掉Attention機制，這使我們應該重新好好考慮Attention帶來的性能提昇的本質。

Transformer中的Self-Attention機制

首先，如下圖所示：

我們給出其形式化的結果：
$$\begin{gathered} A=\text{softmax}(\frac{Q{K}^T}{\sqrt{d_k}}) \\ {F}_{out}=AV \end{gathered}$$
其中，$Q,K\in {\mathbb{R}}^{N\times d^{\prime }}$ 同時 $V\in {\mathbb{R}}^{N\times d}$

這裏，我們給出一個簡化版本，如下圖所示：

也就是將 $Q,K,V$ 都以輸入特征 $F$ 代替掉，其形式化為：
$$\begin{gathered} A=\text{softmax}(F{F}^T) \\ {F}_{out}=AF \end{gathered}$$

然而，這裏面的計算複雜度為 $O(d{N}^2)$，這是Attention機制的一個較大的缺點。

外部注意力 (External Attention)

如下圖所示:

引入了兩個矩陣 $M_k\in {\mathbb{R}}^{S\times d}$ 以及 $M_v \in\mathbb{R}^{S\times d} $, 代替掉原來的 $K,V$

這裏直接給出其形式化：
$$\begin{gathered} A=\text{Norm}(F{M}_k^T) \\ {F}_{out}=A{M}_v \end{gathered}$$
這種設計，將複雜度降低到$O(dSN)$, 該工作發現，當 $S\ll N$ 的時候，仍然能够保持足够的精度。

其中的 $\text{Norm}(\cdot )$ 操作是先對列進行Softmax，然後對行進行歸一化。

實驗分析

首先，文章將Transformer中的Attention機制替換掉，然後在各類任務上進行測試，包括：

• 圖像分類
• 語義分割
• 圖像生成
• 點雲分類
• 點雲分割

這裏只給出部分結果，簡單說明一下替換後的精度損失情况。

圖像分類

語義分割

圖像生成

可以看到，在不同的任務上，基本上不會有精度損失。