# S2-MLPv2 Improved Spatial-Shift MLP Architecture for Vision

论文地址：

* <https://arxiv.org/abs/2108.01072>

## 整体思路以及计算方式

对S2-MLP的改进，使用了多种spatial-shift，然后split attention进行特征融合：

![](https://2520240655-files.gitbook.io/~/files/v0/b/gitbook-x-prod.appspot.com/o/spaces%2Fb2lKDIgZ0I1cirExLedu%2Fuploads%2Fgit-blob-aa2b0de7f512c562824d48ec58f2c84ef75bbf75%2F2.jpg?alt=media)

这里主要介绍split attention的计算方式：

* 输入：$$\left\[\mathbf{X}*{1}, \mathbf{X}*{2}, \cdots, \mathbf{X}\_{K}\right]\in \mathbb R^{n\times K\times c} , \mathbf X\_k\in \mathbb R^{n\times c}$$
* 特征融合：$$\mathbf a= \sum\_{k=1}^K 1\_n \mathbf X\_k \in \mathbb R^c$$
* $$\mathbf a\_1 = \sigma(\mathbf a \mathbf W\_1) \in \mathbb R^{c\_1}$$
* $$\mathbf a\_2 = \mathbf a\_1 \mathbf W\_2 \in \mathbb R^{K c}$$
* $$\mathbf a\_3 =\mathrm{Softmax}(\mathrm{reshape}(\mathbf a\_2),\mathrm{dim}=1)\in \mathbb R^{1\times K\times c}$$
* $$\mathbf o\_1=\left\[\mathbf {X}*{1}, \mathbf {X}*{2}, \cdots, \mathbf {X}\_{K}\right]\odot \mathbf a\_3 \in \mathbb R^{n\times K \times c }$$
* $$\mathbf o\_2 = \mathrm{Sum}(\mathbf o\_1, \mathrm{dim}=1)\in \mathbb R^{n\times c}$$

该模块主要融合了$$\mathbf X\_k$$的特征，不知道是否可以代替Attention的效果？

## 时间复杂度

split attention模块的时间复杂度为$$O(nKc + cc\_1 + Kc\_1c )$$，其余部分任然为线性复杂度。

## 训练以及loss

不变。

## 代码

* <https://github.com/liuruiyang98/Jittor-MLP/blob/main/models_pytorch/s2_mlp_v2.py>

## 简评

spatial-shift + split attention可以大幅提升性能，可以研究下，然后在nlp中使用。
