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的计算对比，可以发现数量级是可以对的上的，因为我们忽略了一次项，所以具体数据与论文不一致。

模型名	实际参数量	$n_{layer}$	$d_{model}$	$n_{head}$	$d_{head}$	估计参数量
GPT-3	175B	96	12288	96	128	173946175488
LLaMA 6.7B	6.7B	32	4096	32	128	6442450944
LLaMA 13.0B	13.0B	40	5120	40	128	12582912000
LLaMA 32.5B	32.5B	60	6656	52	128	31897681920
LLaMA 65.2B	65.2B	80	8192	64	128	64424509440

参考资料

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

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

3. transfomers 库如何得到参数量