参数量计算:
分析transformer模型的参数量、计算量、中间激活、KV cache - 知乎 (zhihu.com)
FLOPs计算:
假设我们有一个自注意力模型,输入序列长度为N,隐藏表示维度为D。
对于每个注意力头(attention head),假设我们使用的是点积注意力(dot-product attention)。
在点积注意力中,对于每个查询向量和键向量的维度,需要进行一个乘法操作。假设每个注意力头中的查询向量和键向量维度均为d。
那么,在计算自注意力的过程中,针对每个注意力头的FLOPs可以计算如下:
- 查询线性变换和键线性变换的点积计算:每个查询向量与每个键向量进行点积操作,共有d个元素。因此,该步骤的FLOPs为 N * N * d。
- Softmax归一化:对于每个查询向量,需要将键线性变换的点积结果进行softmax归一化。该步骤的FLOPs为 N * N * d。
- 自注意力输出的计算:将自注意力权重与值线性变换的点积相乘得到自注意力的输出。每个自注意力权重与每个值线性变换进行点积操作,共有D个元素。因此,该步骤的FLOPs为 N * D * D。
综上所述,每个注意力头的总FLOPs为 N * N * d(查询线性变换和键线性变换的点积)+ N * N * d(softmax归一化)+ N * D * D(自注意力输出的计算)。
请注意,上述计算仅考虑了单个注意力头的FLOPs,并未考虑多头注意力和其它线性变换的FLOPs。在实际应用中,计算总体的模型FLOPs需要考虑所有注意力头以及其他层间的运算量。