Transformer 相关模型的参数量计算

如何计算Transformer 相关模型的参数量呢？
先回忆一下Transformer模型论文《Attention is all your need》中的两个图。

设Transformer模型的层数为N，每个Transformer层主要由self-attention 和 Feed Forward组成。设self-attention模块的head个数为$n_{head}$，每一个head对应的维度为$d_{head}$，self-attention输出维度为$d_{model}=n_{\text{heads}}\cdot d_{\text{head}}$。我们可以得到一个Transformer层的参数量为$12d_{model}^2+13d_{model}$，具体如下：

• self-attention块的模型参数有Q、K、V的权重矩阵$W_Q、W_K、W_V$和偏置，输出矩阵$W_O$及其偏置。这4个权重矩阵的大小为$[d_{model},d_{model}]$，4个偏置的大小为$[d_{model}]$，所以self-attention块的参数量为$4d_{model}^2+4d_{model}$。

• Feed Forward块一般由2个线性层组成，第一个线性层将维度从$d_{model}$ 映射成 $4d_{model}$， 其权重矩阵$W_1$的大小为$[d_{model},4d_{model}]$ ，其偏置的大小为$[4d_{model}]$。 第二个线性层将维度从$4d_{model}$ 映射成 $d_{model}$，其权重矩阵$W_2$的大小为$[4d_{model},d_{model}]$ ，其偏置的大小为$[d_{model}]$。所以Feed Forward的参数量为$8d_{model}^2+5d_{model}$。

• self-attention 和 Feed Forward都跟随着layer normalization，它有两个可训练模型参数，形状都是$[d_{model}]$。所以2个layer normalization的参数量为$4d_{model}$。

除了Transformer层之外的参数有：

• 词embedding矩阵的参数量，embedding的维度通常等于$d_{model}$，设词表的大小为V，则词embedding的参数量为$V{d}_{model}$。
• 位置向量相关，有些位置向量表示方式需要学习参数。

所以N层Transformer模型的可训练模型参数量为$N(12d_{model}^2+13d_{model})+V{d}_{model}$。当$d_{model}$较大时，可以忽略一次项，模型参数量近似为$12N{d}_{model}^2$。

最后试验一下模型参数估计量与论文是否对的上，下表是GPT3和LLaMA的计算对比，可以发现数量级是可以对的上的，因为我们忽略了一次项，所以具体数据与论文不一致。

参考资料

1. Transformer 论文（模型图来自论文）、GPT3的论文等

2. 整理过程中参考的blog: 1. 知乎用户回旋托马斯x 的文章，除了计算量外，还算了计算量、中间激活等 ， 2 transformer 参数量计算, 3 flops 计算, 4 transformers 参数量计算公式

3. transfomers 库如何得到参数量