bert模型的参数量和时间复杂度分析

背景

这里主要比较了bert-base和albert-tiny。前者窄而深，后者宽而浅。需要注意的一点是，albert是对bert瘦了身，但只省了空间，没有省时间。因为虽然不同层共享了参数，但模型还是要经过这么多层算一遍。

假设词表大小为V，句子长度为S，词向量维度为H，层数为L，自注意力头数为A，对于ALBert，其嵌入分解维度为E。模型的下游任务为多分类，标签数为K。

参数设置

$$\begin{array}{cccccc}\text{ 不同模型参数 }& \mathbf{V}& \mathbf{S}& \mathbf{H}& \mathbf{L}& \mathbf{E}\\ \text{ bert-base }& 30000& 512& 768& 12& -\\ \text{ albert-tiny }& 30000& 512& 1024& 4& 128\end{array}$$

参数量比较

$$\begin{array}{cccccc}\text{ 不同模型参数量 }& \text{ 嵌入初始化 }& \text{ 自注意力机制 }& \text{ 前向网络传播 }& \text{ 嵌入正则 }& \text{ 下游任务 }\\ \text{ bert-base }& \mathrm{V}\times \mathrm{H}+\mathrm{S}\times \mathrm{H}+2\mathrm{H}& 3\mathrm{L}\times {\mathrm{H}}^2& 9\mathrm{L}\times {\mathrm{H}}^2& (4\mathrm{L}+2)\times \mathrm{H}& \mathrm{H}\times \mathrm{K}\\ \text{ albert-tiny }& \mathrm{V}\times \mathrm{E}+\mathrm{E}\times \mathrm{H}+\mathrm{S}\times \mathrm{H}+2\mathrm{H}& 3{\mathrm{H}}^2& 9{\mathrm{H}}^2& 6\mathrm{H}& \mathrm{H}\times \mathrm{K}\end{array}$$

运行时间比较

$$\begin{array}{cccccc}\text{ 不同模型时间复杂度 }& \text{ 嵌入初始化 }& \text{ 自注意力机制 }& \text{ 前向网络传播 }& \text{ 嵌入正则 }& \text{ 下游任务 }\\ \text{ bert-base }& \mathrm{O}((\mathrm{V}+\mathrm{S})\mathrm{H})& \mathrm{O}\left({\mathrm{LH}}^2\right)& \mathrm{O}\left({\mathrm{LH}}^2\right)& \mathrm{O}(\mathrm{LH})& \mathrm{O}(\mathrm{HK})\\ \text{ albert-tiny }& \mathrm{O}(\mathrm{VEH})+\mathrm{O}(\mathrm{VE})+\mathrm{O}(\mathrm{SH})& \mathrm{O}\left({\mathrm{LH}}^2\right)& \mathrm{O}\left({\mathrm{LH}}^2\right)& \mathrm{O}(\mathrm{LH})& \mathrm{O}(\mathrm{HK})\end{array}$$

分析

• 词嵌入所需的参数为Token embedding，Position embedding，Segment embedding【对应数据集级别、句子级别、样例级别】：$\mathrm{V}\times \mathrm{H}+\mathrm{S}\times \mathrm{H}+2\mathrm{H}$
• 自注意力机制所需的参数为【3对应QKV】：$L\times 3\times A\times (H\times \frac{H}{A})=3L\times H^2$
• Forward网络所需参数为：$L\times (H\times H+H\times 4H+4H\times H)=9L\times H^2$
• LN层所需参数为（在embedding和block都有用到，有weight和bias两个参数，每个块包含两个LN）：$2H+L\times 2\times 2H$