Transformer 中 Self-attention 的计算复杂度

在 Transformer 中，Multi-head attention 的计算过程是： $\text{MultiHeadAttn}(z_q,x)={\sum }_{m=1}^M{W}_m[{\sum }_{k\in {\Omega }_k}{A}_{mqk}\cdot W_m^{\prime }{x}_k]$.

其中$m$是 attention head 的索引，${W_m}^{\prime }\in {\mathbb{R}}^{C_v\times C}$ 是输入的映射矩阵，$W_m\in {\mathbb{R}}^{C\times C_v}$ 是输出的映射矩阵，二者都是可学习的权重（$C_v=C/M$）。Attention 权重 $A_{mqk}\propto \exp \{\frac{z_q^T{U}_m^T{V}_m{x}_k}{\sqrt{C_v}}\}$，并且 ${\sum }_{k\in {\Omega }_k}{A}_{mqk}=1$，其中$U_m,V_m\in {\mathbb{R}}^{C_v\times C}$分别是 query 的映射矩阵和 key 的映射矩阵，也都是可学习权重。设 query 和 key 元素的个数分别是$N_q$和$N_k$. $\text{MultiHeadAttn}(z_q,x)$的计算复杂度是$O(N_q{C}^2+N_k{C}^2+N_q{N}_{k}C)$。

1. 输入是 $X\in {\mathbb{R}}^{N\times C}$，用$U_m,V_m\in {\mathbb{R}}^{C_v\times C}$分别对 query 和 key 做线性变换，计算得到$Q,K\in {\mathbb{R}}^{N\times C}$矩阵。这样，计算$Q$和$K$的复杂度就是$O(N_q\times C^2)$和$O(N_k\times C^2)$.
2. 然后计算$A_{mqk}\propto \exp \{\frac{z_q^T{U}_m^T{V}_m{x}_k}{\sqrt{C_v}}\}$，复杂度是$O(N_q\times N_k\times C)$.
3. $A_{mqk}$与$x_k$相乘，计算复杂度是$O(N_q\times N_k\times C)$.
4. 总体的计算复杂度就是$O(N_q\times C^2+N_k\times C^2+N_q{N}_{k}C)$.

在 DETR 中，Transformer encoder 的 query 和 key 元素就是特征图上的像素点，假设输入特征图的宽度和高度分别是$W$和$H$。

1. Encoder 中的 self-attention 的计算复杂度就是 $O(H^2{W}^{2}C)$.
2. Decoder 包括了 self attention 和 cross attention，输入包括来自于 encoder 的特征图、$N$个 object queries。
3. 在 decoder 的 cross attention 中，query 元素来自于 object queries，key 元素来自于 encoder 特征图，从 encoder 提供的特征图上提取 key 元素，$N_q=N,N_k=H\times W$，计算复杂度是$O(N_k{C}^2+N{N}_{k}C)=O(HW{C}^2+NHW{C}^2)$.
4. 在 decoder 的 self attention 中，object queries 相互作用，query 和 key 元素都来自于 object queries。$N_q=N_k=N$，复杂度就是$O(2N{C}^2+N^{2}C)$.

引用

• https://stackoverflow.com/questions/65703260/computational-complexity-of-self-attention-in-the-transformer-model
• Deformable DETR