transformer系列之时间复杂度

一、矩阵运算的复杂度

假设我们有两个矩阵，A的维度为$(m\ast n)$，B的维度为$(n\ast p)$，则结果矩阵C的维度为$(m\ast p)$

对于结果矩阵C中的每个元素，我们需要计算矩阵A中的一行和矩阵B中的一列的点积。这个点积涉及到对应元素的乘积和求和。

为了计算C中的每个元素，我们需要执行n次乘法和n-1次加法，如果再加上偏置项，刚好就是2*n

由于结果矩阵C中有m x p个元素，因此所需的总操作数为$(m\ast p)\ast 2\ast n$。但为了表述简洁性，忽略细节，可以写成$m\ast p\ast 2\ast n$

因此，矩阵乘法的总体复杂度为$O(m\ast n\ast p)$。

二、transformer时间复杂度计算

transformer分为自注意力层与MLP层； 假设当前输入序列长度为 $N$​ 时，只一个头且头的形状为$[d,d]$时，计算复杂度如下：

2.1、自注意力复杂度

首先，我们知道transformer模型由$l$个相同的层组成，每个层分为两部分：self-attention块和MLP块，结构如下：

self.attention模块的结构如下：

而self-attention层的模型参数有两部分，一部分是$Q、K、V$​的权重矩阵$W_Q$、$W_K$、$W_V$和偏置，另一部分是输出权重矩阵$W_O$​和偏置，最终为：8$N{d}^2$ + 4$N^2$d
具体计算步骤如下：
计算​$QKV$:
即​$Q=x{W}_Q,K=x{W}_k,V=x{W}_v$
该矩阵乘法的输入和输出形状为$[N,d]\ast [d,d]$ -->$[N,d]$
计算量为：$3\ast 2N{d}^2$
计算​$Q{K}^T$
该部分的输入和输出形状为[N,d]*[d, N]
计算量为：​2$N^2$d
计算score * $V$
该部分计算的核心在于加法，输入和输出形状为$[N,d]\ast [d,N]$，每个token都是N个V相加，计算量是$N\ast 2d$，$N$个token计算量为2$N^2$d
​ 计算量为：2$N^2$d
attention后的线性映射，矩阵乘法的输入和输出形状为​$[N,d]\ast [d,d]$
计算量为​:2$N{d}^2$

2.2、MLP的复杂度

MLP块由2个线性层组成,维度为$[d,4d][4d,d]$
计算量：$N\ast (2\ast d\ast 4d+2\ast d\ast 4d),即16N{d}^2$

将上述所有表粗所示的计算量相加，得到每个transformer层的计算量大约为:24$N{d}^2+4d{N}^2$

三、换个角度研究transformer计算复杂度

假设向量的维度是$1\ast d$，单独看它的计算量：

$QKV$的计算

$QKV$的矩阵维度假设为$d\ast d$，那么计算量为：2$d^2$ * 3，得到三个$1\ast d$的向量

注意力权重的计算

$Q需要与N个K相乘，计算量为：2d\ast N$

加权注意力的得到新向量

$注意力权重\ast V+偏转向量，[1,d]\ast [d,N]+偏转项,计算量为:2d\ast N$

线性映射 与 MLP层

得到新向量后需要线性映射，再通过MLP层，其计算量：$[1,d]\ast [d,d]$ 得到2$d^2$
MLP层的维度$[d,4d],[4d,d]$，其计算量：$2d\ast 4d\ast 2$，得到16$d^2$

汇总所有token在block的计算量

$N\ast (2d^2\ast 3+2d\ast N+2d\ast N+2d^2+16d^2)=24d^{2}N+4d{N}^2$

四、logits的计算量

词嵌入矩阵的参数量为vocab_size × hidden_size，其中vocab_size是词表的大小，hidden_size是隐藏层的维度。词嵌入矩阵用于将输入的词索引映射为对应的词向量。
而输出层的权重矩阵通常与词嵌入矩阵共享参数。这意味着输出层的权重矩阵的维度也是hidden_size × vocab_size，其中hidden_size是隐藏层的维度，vocab_size是词表的大小。通过共享参数，可以减小模型的参数量，提高模型的效率。
计算量为：$2dvN$

五、结论

embedding维度 $[1,d]$，$QKV$的维度$[d,d]$,注意力Head数为1，最终得到的计算复杂度：$l\ast (24d^{2}N+4d{N}^2)+2dvN$