在Swin-Transformer一文中，有这样两个公式，分别为：

1. Transformer中提出的Multi-head Self-Attention模块（MSA）
2. Swin-Transformer中提出的Window Multi-head Self-Attention模块（W-MSA）

两者计算量即计算复杂度分别为：

$$\Omega (MSA)=4hw{C}^2+2(hw{)}^{2}C$$
$$\Omega (W-MSA)=4hw{C}^2+2M^{2}hwC$$

一、MSA模块计算量

Transformer中提出的多头自注意力模块运算公式为
$$Attention(Q,K,V)=Softmax(\frac{Q{K}^T}{\sqrt{d_k}})V$$
运算过程以及计算量如下图：
（注：图中黑色字体为矩阵维度，橙色字体为计算量）
1. 首先，一张维度为H * W * C的图片，分别通过乘上一个CC的变换矩阵，分别变成QKV矩阵，维度仍为H * W * C。这样三个变换用到的计算量都为：HW * C^2，加起来为：$$3HW\ast C^2$$
2. 然后做QK的转置，(HW * C) * (C * HW)，得到的矩阵维度为HW * HW，计算量为：$$(HW{)}^2\ast C$$
3. 忽略根号dk和softmax的计算量，2中得到的 A矩阵 × V矩阵 ，结果的维度为HW * C，计算量为：$$(HW{)}^2\ast C$$
4. 因为是多头自注意力，所以在做完矩阵乘法后，需要矩阵拼接融合，做一个Proj映射，得到输出矩阵，维度和输入一样为HW * C，计算量为：$$HW\ast C^2$$
5. 最后计算量加起来，为： $$\Omega (MSA)=4hw{C}^2+2(hw{)}^{2}C$$

二、W-MSA模块计算量

W-MSA与MSA总体的计算过程是一致的，区别在于：W-MSA的长宽不再是H和W，而是$$窗口:M\ast M$$并且有$$\frac{H}{M}\ast \frac{W}{M}个窗口需要计算$$

所以它的计算量为：
$$(\frac{h}{M}\ast \frac{w}{M})\ast (4M^2{C}^2+2M^{4}C)=4hw{C}^2+2M^{2}hwC$$

分析

从他们公式可以看出区别主要在于两个公式的后半部分
带一点数进去就可以看出W-MSA在计算量上比MSA少很多，比如以原文中的一些参数设定为例：HW都为56，C为96，M为7

前者MSA为：$$2\ast (56\ast 56{)}^2\ast 96=1888223232$$
后者为W-MSA为：$$2\ast 7^2\ast 56\ast 56\ast 96=29503488$$
二者相差了64倍。

论文名称：Swin Transformer: Hierarchical Vision Transformer using Shifted Windows
原论文地址： https://arxiv.org/abs/2103.14030