vision transformer的计算复杂度

Vision transformer

假设每个图像有 $h\ast w$ 个patch，维度是 $C$

输入的图像 $X$ ( 大小为 $hw\ast C$ )，和三个系数矩阵相乘 ( 大小为 $C\ast C$ )，得到 $qkv$ 三个向量 ( $hw\ast C$ )，复杂度为：
$$3hw{C}^2$$

$q$ ( $hw\ast C$ ) 和 $k^T$ ( $C\ast hw$ ) 相乘得到矩阵 $A$ ( $hw\ast hw$ )，复杂度为：$$(hw{)}^{2}C$$

$A$ ( $hw\ast hw$ ) 和 $v$ ( $hw\ast C$ )相乘，得到多头注意力的结果 ( $hw\ast C$ )，复杂度为：$$(hw{)}^{2}C$$

经过MLP投影层 ( $C\ast C$ )，得到 ( $hw\ast C$ )，复杂度为：
$$hw{C}^2$$

所以复杂度之和为：$$4hw{C}^2+2(hw{)}^{2}C$$

Swin transformer

基于滑动窗口的多头注意力，是在每个窗口内计算注意力

假设每个窗口有 $M\times M$ 个patch

在一个窗口内的复杂度为：

$$4M^2\ast C+2M^{4}C$$

共有 $hw/M^2$ 个窗口，所以复杂度之和为：

$$4hwC+2M^{2}hwC$$

Convolutional vision Transformer

使用 $s\times s$ 卷积进行卷积投影，有 $hw$ 个patch，通道维度为 $C$

输入的图像 $X$ ( 大小为 $hw\ast C$ )，使用三个标准卷积进行投影 ( 大小为 $s\ast s\ast C$ )，得到 $qkv$ 三个向量 ( $hw\ast C$ )，投影的复杂度为：

$$3hw{s}^2{C}^2$$

使用深度可分离卷积，投影的复杂度为：

$$3hw{s}^{2}C$$

使用步长大于1的卷积进行多头注意力的投影，减小后面注意力的计算花销。

key和value的步长为2，query的步长为1，key和value的token数量减小了4倍，所以后续的多头注意力计算花销也减小了4倍。

Cross Attention Transformer

交叉注意力包括IPSA和CPSA，IPSA在单个patch内使用卷积进行投影，CPSA在单个通道计算patch间的注意力

IPSA的复杂度：

patch大小为 $N$，通道数为 $C$

输入的图像 $X$ ( 大小为 $N^2\ast C$ )，使用卷积进行投影 ( 大小为 $1\ast 1\ast C$ )，得到 $qkv$ 三个向量 ( $N^2\ast C$ )，复杂度为：
$$3N^2{C}^2$$

$q$ ( $N^2\ast C$ ) 和 $k$ ( $C\ast N^2$ ) 相乘得到矩阵 $A$ ( $N^2\ast N^2$ )，复杂度为：$$N^4{C}^2$$

$A$ ( $N^2\ast N^2$ ) 和 $v$ ( $N^2\ast C$ )相乘，得到多头注意力的结果 ( $N^2\ast C$ )，复杂度为：$$N^4{C}^2$$

经过MLP投影层 ( $C\ast C$ )，得到 ( $N^2\ast C$ )，复杂度为：
$$N^2{C}^2$$

单个patch内的复杂度为：

$$4N^2{C}^2+2N^4{C}^2$$

共有 $HW/N^2$ 个patch，所以IPSA总复杂度为：
$$4HW{C}^2+2N^{2}HW{C}^2$$

CPSA的复杂度：

patch数目为 $HW/N^2$，patch大小为 $N^2$

输入的图像 $X$ ( 大小为 $HW/N^2\ast N^2$ )，和三个系数矩阵相乘 ( 大小为 $N^2\ast N^2$ )，得到 $qkv$ 三个向量 ( $HW/N^2\ast N^2$ )，复杂度为：
$$3HW{N}^2$$

$q$ ( $HW/N^2\ast N^2$ ) 和 $k$ ( $N^2\ast HW/N^2$ ) 相乘得到矩阵 $A$ ( $HW/N^2\ast HW/N^2$ )，复杂度为：$$(HW{)}^2/N^2$$

$A$ ( $HW/N^2\ast HW/N^2$ ) 和 $v$ ( $HW/N^2\ast N^2$ )相乘，得到多头注意力的结果 ( $HW/N^2\ast N^2$ )，复杂度为：$$(HW{)}^2/N^2$$

经过MLP投影层 ( $N^2\ast N^2$ )，得到 ( $HW/N^2\ast N^2$ )，复杂度为：
$$HW{N}^2$$

单个通道内的复杂度为：

$$4N^{2}HW+2(HW/N{)}^2$$

共有 $C$ 个通道，所以CPSA总复杂度为：
$$4N^{2}HWC+2(HW/N{)}^{2}C$$