Replace self attention with MLP

use MLP Instead of Self-Attention

This is a job of Tsinghua University “Beyond Self-attention: External Attention using Two Linear Layers for Visual Tasks”

Replace... With two linear layers Self-Attention Mechanism , In the end, it can improve the speed while maintaining the accuracy .

What's surprising about this job is , We can use MLP Instead of Attention Mechanism , This makes it necessary for us to reconsider Attention The nature of the performance improvement .

Transformer Medium Self-Attention Mechanism

First , As shown in the figure below ：

We give its formal result ：
$$\begin{gathered} A=\text{softmax}(\frac{Q{K}^T}{\sqrt{d_k}}) \\ {F}_{out}=AV \end{gathered}$$
among ,$Q,K\in {\mathbb{R}}^{N\times d^{\prime }}$ meanwhile $V\in {\mathbb{R}}^{N\times d}$

here , We give a simplified version , As shown in the figure below ：

Also is to $Q,K,V$ All based on input features $F$ Instead of , It is formalized as ：
$$\begin{gathered} A=\text{softmax}(F{F}^T) \\ {F}_{out}=AF \end{gathered}$$

However , The computational complexity is $O(d{N}^2)$, This is a Attention A big drawback of the mechanism .

External attention (External Attention)

As shown in the figure below :

Two matrices are introduced $M_k\in {\mathbb{R}}^{S\times d}$ as well as $M_v \in\mathbb{R}^{S\times d} $, Instead of the original $K,V$

Here we give its formalization directly ：
$$\begin{gathered} A=\text{Norm}(F{M}_k^T) \\ {F}_{out}=A{M}_v \end{gathered}$$
This design , Reduce the complexity to $O(dSN)$, The work found that , When $S\ll N$ When , Still able to maintain enough accuracy .

Among them $\text{Norm}(\cdot )$ The operation is to start with the column Softmax, Then normalize the rows .

experimental analysis

First , The article will Transformer Medium Attention The mechanism replaced , And then test on all kinds of tasks , Include ：

• Image classification
• Semantic segmentation
• Image generation
• Point cloud classification
• Point cloud segmentation

Only partial results are given here , Briefly explain the accuracy loss after replacement .

Image classification

Semantic segmentation

Image generation

You can see , On different tasks , There's basically no loss of accuracy .