如何计算Transformer 相关模型的参数量呢?
先回忆一下Transformer模型论文《Attention is all your need》中的两个图。
设Transformer模型的层数为N,每个Transformer层主要由self-attention 和 Feed Forward组成。设self-attention模块的head个数为 n h e a d n_{head} nhead,每一个head对应的维度为 d h e a d d_{head} dhead,self-attention输出维度为 d m o d e l = n heads ⋅ d head d_{model}= n_\text{heads}\cdot d_\text{head} dmodel=nheads⋅dhead。我们可以得到一个Transformer层的参数量为 12 d m o d e l 2 + 13 d m o d e l 12 d_{model}^2 + 13 d_{model} 12dmodel2+13dmodel,具体如下:
-
self-attention块的模型参数有Q、K、V的权重矩阵 W Q 、 W K 、 W V W_Q、W_K 、W_V WQ、WK、WV和偏置,输出矩阵 W O W_O WO及其偏置。这4个权重矩阵的大小为 [ d m o d e l , d m o d e l ] [d_{model}, d_{model}] [dmodel,dmodel],4个偏置的大小为 [ d m o d e l ] [d_{model}] [dmodel],所以self-attention块的参数量为 4 d m o d e l 2 + 4 d m o d e l 4 d_{model}^2 + 4 d_{model} 4dmodel2+4dmodel。
-
Feed Forward块一般由2个线性层组成,第一个线性层将维度从 d m o d e l d_{model} dmodel 映射成 4 d m o d e l 4d_{model} 4dmodel, 其权重矩阵 W 1 W_1 W1的大小为 [ d m o d e l , 4 d m o d e l ] [d_{model}, 4d_{model}] [dmodel,4dmodel] ,其偏置的大小为 [ 4 d m o d e l ] [4d_{model}] [4dmodel]。 第二个线性层将维度从 4 d m o d e l 4d_{model} 4dmodel 映射成 d m o d e l d_{model} dmodel,其权重矩阵 W 2 W_2 W2的大小为 [ 4 d m o d e l , d m o d e l ] [4d_{model}, d_{model}] [4dmodel,dmodel] ,其偏置的大小为 [ d m o d e l ] [d_{model}] [dmodel]。所以Feed Forward的参数量为 8 d m o d e l 2 + 5 d m o d e l 8 d_{model}^2 + 5 d_{model} 8dmodel2+5dmodel。
-
self-attention 和 Feed Forward都跟随着layer normalization,它有两个可训练模型参数,形状都是 [ d m o d e l ] [d_{model}] [dmodel]。所以2个layer normalization的参数量为 4 d m o d e l 4 d_{model} 4dmodel。
除了Transformer层之外的参数有:
- 词embedding矩阵的参数量,embedding的维度通常等于 d m o d e l d_{model} dmodel,设词表的大小为V,则词embedding的参数量为 V d m o d e l Vd_{model} Vdmodel。
- 位置向量相关,有些位置向量表示方式需要学习参数。
所以N层Transformer模型的可训练模型参数量为 N ( 12 d m o d e l 2 + 13 d m o d e l ) + V d m o d e l N(12 d_{model}^2 + 13 d_{model}) + Vd_{model} N(12dmodel2+13dmodel)+Vdmodel。当 d m o d e l d_{model} dmodel较大时,可以忽略一次项,模型参数量近似为 12 N d m o d e l 2 12 N d_{model}^2 12Ndmodel2。
最后试验一下模型参数估计量与论文是否对的上,下表是GPT3和LLaMA的计算对比,可以发现数量级是可以对的上的,因为我们忽略了一次项,所以具体数据与论文不一致。
模型名 | 实际参数量 | n l a y e r n_{layer} nlayer | d m o d e l d_{model} dmodel | n h e a d n_{head} nhead | d h e a d d_{head} dhead | 估计参数量 |
---|---|---|---|---|---|---|
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 |
参考资料
-
Transformer 论文(模型图来自论文)、GPT3的论文等
-
整理过程中参考的blog: 1. 知乎用户回旋托马斯x 的文章,除了计算量外,还算了计算量、中间激活等 , 2 transformer 参数量计算, 3 flops 计算, 4 transformers 参数量计算公式
-
transfomers 库如何得到参数量