Transformer参数量和复杂度

在算法岗面试中经常会问到Transformer相关的基础知识。
首先需要清楚Transformer的参数量和复杂度分别在算什么。

定义：

• 参数量：神经网络中有很多参数矩阵，这个矩阵大小的和就是参数量，静态的，摆在那就在那，定量的。
• 复杂度：与输入的数据有关，动态的，跟计算公式有关，定性的。

Transformer的架构主要分成2部分：

• Encoder：6层，每层包括 Multi-Head Self-Attention（MHSA）和FFN
• Decoder：6层，每层包括Masked MHSA、Multi-Head Cross-Attention和FFN

其他部分包括Input Embedding 、Postion Encoding以及最后解码的Linear层。

其中，每个MHSA、Masked MHSA、Multi-Head Cross-Attention和FFN中有含有Add&Norm操作。

参数量计算：

InputEmbedding: 将vocab中的词映射到d维度，所以： vocab*d
Encoder+Decoder：

• MHSA/Multi-Head Cross-Attention/Masked MHSA：Q K V O 四个矩阵 【没有bias】
  $4\ast d\ast d=4d^2$
• FFN ：第一个矩阵先增到4d，第二个矩阵减到d。所以参数量为
  $d\ast (4d)+4\ast d+(4\ast d)\ast d+d=8d^2+5d$
• layerNorm : 参数量就是 $\gamma$ 和 $\beta$ ,所以是 2d

其实可以看到，MHSA的参数量仅有 FFNN 的一半。

$$\begin{gathered} Total=InputEmbedding+Encoder+Decoder \\ =vocab\ast d+6\ast (4d^2+2d+8d^2+5d+2d)+6（4d^2+2d+4d^2+2d+8d^2+5d+2d） \\ =vocab\ast d+6\ast (12d^2+9d)+6\ast (16d^2+11d) \\ =vocab\ast d+168d^2+120d \end{gathered}$$

复杂度分析

复杂度分成时间复杂度和空间复杂度。
神经网络中，最常见的就是线性映射，涉及到矩阵运算。这里用到一个矩阵相乘运算复杂度的前置知识：

矩阵M1=m * n 矩阵M2=n * k，得到矩阵M=m * k，所以时间复杂度为O(mnk)，空间复杂度为O(m*k)

假设输入序列长度为N

InputEmbedding： 将vocab中的词映射到d维度，类似检索哈希表。
时间复杂度：O(N)
空间复杂度：O(N*d)
Encoder+Decoder：

• MHSA/Multi-Head Cross-Attention/Masked MHSA：
  attention计算是复杂度的关键。
  时间复杂度： $O(N^{2}d)$
  self-attention中计算attention score那里就是
  ● $Q\ast K^T$ 是 $O(N\ast d\ast N)=O(N^{2}d)$；
  ● 让$scores=softmax(Q\ast K^T/\sqrt{d})$ 那么 $scores\ast V$ 是 $O(N^{2}d)$；
  空间复杂度： $O(N^2+Nd)$
  ● $Q\ast K^T$ 是 $O(N^2)$；
  ● $scores\ast V$ 是 $O(Nd)$；
• FFN ：第一个矩阵先增到4d，第二个矩阵减到d。
  时间复杂度： $O(N\ast d\ast 4d)=O(N{d}^2)$
  空间复杂度： $O(Nd)$

所以：
时间复杂度： $O(N^{2}d+N{d}^2)$
空间复杂度： $O(N^2+Nd)$

通常把d作为常量，所以：
时间复杂度： $O(N^2)$
空间复杂度： $O(N^2)$

参考：https://zhuanlan.zhihu.com/p/661804092