transformer系列4---transformer结构计算量统计

1 术语解释

1. FLOPs（Floating Point Operations）：浮点运算次数，用来衡量模型计算复杂度，常用来做神经网络模型速度的间接衡量标准。（但该指标与实际模型速度并不一定正相关）

2. MACs(Multiply–Accumulate Operations)：乘加累积操作数，1个MACs包含一个乘法操作与一个加法操作，大约包含2FLOPs。通常MACs与FLOPs存在一个2倍的关系。

2 矩阵相乘FLOPs

1. 对于$A^{1\times n},B^{n\times 1}$两个矩阵相乘，计算AB需要进行n次乘法和n次加法，共计2n次浮点数运算 ,即2n的FLOPs，
2. 对于$A^{m\times n},B^{n\times p}$两个矩阵相乘，计算AB需要进行2nmp次浮点数运算。

3 Transformer的FLOPs估计

假设Transformer的输入每个词向量维度d_model(d) ，词表大小为vocab_size(v)，输入句子最大长度为src_max_len(s)，batchsize为 batch(b)，head头数为head(h)。
对于输入部分，将输入句子分词并且词嵌入步骤没有计算量，位置编码也没有计算量，因此，计算量主要集中在MultiHeadAttention、MLP、以及最后的投影计算。

3.1 MultiHeadAttention

3.1.1 Q,K,V计算

1. 1个矩阵计算量：对于输入I，首先计算$Q=I\ast W^q$，$K=I\ast W^k$，$V=I\ast W^v$，假设输入I的形状为 [b, s, d],1个矩阵乘法的输入和输出形状为[b, s, d] × [d, d] = [b, s, d]，计算量为$2bs{d}^2$
2. 3个矩阵计算量：$6bs{d}^2$

3.1.2 attention计算

1. $Q{K}^T$

  矩阵乘法的输入形状[b, h, s, d] × [b, h, s, d]，输出形状为 [b, h, s, s]，h维度是concat，没有计算量，因此该步骤的计算量为$2b{s}^{2}d$ 。

2. score*V加权
   输入形状为[b, h, s, s] × [b, h, s, d]，输出形状为[b, h, s, d]， h维度是concat，没有计算量，因此该步骤的计算量为$2b{s}^{2}d$ 。

3.1.3 MultiHeadAttention输出线性映射

所有head都concat，输入形状为[b, s, d] × [d, d] $(W^O)$，输出形状为[b, s, d]，计算量$2bs{d}^2$

3.1.4 MultiHeadAttention总计算量

MultiHeadAttention总计算量为上面三部分之和，$2b{s}^{2}d$+$2b{s}^{2}d$+$2bs{d}^2$=$4b{s}^{2}d$+$2bs{d}^2$

3.2 MLP

MLP内包含2个线性层：

1. 第一个线性层，矩阵乘法输入形状为[b, s, d] × [d, 4d]，输出形状为[b, s, 4d]，计算量$8bs{d}^2$
   。
2. 第二个线性层，矩阵乘法输入形状为[b, s, 4d] × [4d, d]，输出形状为[b, s, d]，计算量$8bs{d}^2$

MLP总计算量为$8bs{d}^2$+$8bs{d}^2$=$16bs{d}^2$

3.3 projection输出

logits的计算，将隐藏向量映射为词表大小。矩阵乘法输入形状为[b, s, d] × [d, v]，输出形状为[b, s, v]，计算量$2bsdv$。

3.3 计算量累计

1. Transformer的encoder，包含1个MultiHeadAttention，1个MLP
2. Transformer的decoder，包含2个MultiHeadAttention，1个MLP
3. Transformer的输出为1个projection
   将上面3部分累加，计算量为$4b{s}^{2}d$+$2bs{d}^2$+$16bs{d}^2$+2*（$4b{s}^{2}d$+$2bs{d}^2$）+$16bs{d}^2$+$2bsdv$=$12b{s}^{2}d$+$36bs{d}^2$+$2bsdv$