Transformer

RNN 的主要问题：

• 梯度消失/爆炸：长距离依赖难以学习
• 顺序计算：无法并行处理序列
• 信息瓶颈：最后时刻隐藏状态需承载全部信息

Transformer 的改进：

• 并行计算：同时处理整个序列
• 自注意力机制：直接建立任意位置间的联系
• 位置编码：显式注入位置信息

结构

图：Transformer 单元

图：Transformer 的详细结构

• 输入
  • 编码器输入
  • 解码器输入
• 输出
  • 线性层
  • Softmax 层
• 编码器
  • 由 N 个编码器层堆叠而成
  • 每个编码器层由两个子层连接结构组成
  • 第一个子层连接结构包括一个多头自注意力子层和规范化层以及一个残差连接
  • 第二个子层连接结构包括一个前馈全连接子层和规范化层以及一个残差连接
• 解码器
  • 由 N 个解码器层堆叠而成
  • 每个解码器层由三个子层连接结构组成
  • 第一个子层连接结构包括一个带掩码的-多头自注意力子层和规范化层以及一个残差连接
  • 第二个子层连接结构包括一个多头注意力子层（编码器到解码器）和规范化层以及一个残差连接
  • 第三个子层连接结构包括一个前馈全连接子层和规范化层以及一个残差连接

核心组件

自注意力机制（Self-Attention）

在传统的神经网络处理序列时，模型只能一步步按顺序处理，难以捕捉长距离依赖关系。自注意力机制就是为了让序列中的每个元素都能直接与序列中所有其他元素进行交互，无论它们直接的距离多远。

定义：

符号	维度	含义
$X$	$n\times d$	输入矩阵（n=序列长度，d=特征维度）
$Q$	$n\times d_k$	Query 矩阵（查询向量）
$K$	$n\times d_k$	Key 矩阵（键向量）
$V$	$n\times d_v$	Value 矩阵（值向量）
$W^Q,W^K,W^V$	$d\times d_k/d_v$	可学习参数矩阵

$$\text{Attention}(Q,K,V)=\text{Softmax}(\frac{Q{K}^T}{\sqrt{d_k}})V$$

TIP

• $Q$ 表示当前需要关注的信息或问题，用于确定输入序列中哪些部分与当前任务相关
• $K$ 用于匹配查询，通过计算相似度判断输入序列中哪些元素与查询匹配
• $V$ 存储实际信息

推导过程

将输入转换为 Query、Key、Value：

$$Q=X{W}^Q,K=X{W}^K,V=X{W}^V$$

计算注意力分数：

$$\text{Scores}=\frac{Q{K}^T}{\sqrt{d_k}}$$

生成注意力权重矩阵：

$$A=\text{Softmax}(\text{Scores})$$

得到最终注意力输出：

$$\text{Output}=AV$$

带掩码自注意力层（Masked Multi-head attention）

编码时，对于 $t$ 时刻的预测，我们知道 $x_1,x_2,\cdots ,x_t,x_{t+1},\cdots ,x_T$ 全部的信息。

解码时，对于 $t$ 时刻的预测，我们仅知道 $x_1,x_2,\cdots ,x_{t-1}$ 的信息。看不到后续的信息，因此需要将后续的信息遮掩起来。

$$\text{Attention}(Q,K,V)=\text{Softmax}(\frac{Q{K}^T\odot M}{\sqrt{d_k}})V$$

多头注意力（Multi-Head Attention）

$$\text{MultiHead}(Q,K,V)=\text{Concat}({\text{head}}_1,\cdots ,{\text{head}}_h)W^O$$

其中：

$${\text{head}}_i=\text{Attention}(Q{W}_i^Q,K{W}_i^K,V{W}_i^V)$$

位置编码（Positional Encoding）

与 RNN 和 LSTM 等顺序算法不同，Transformer 没有内置机制来捕获句子中单词的相对位置，所以在 Transformer 的 encoder 和 decoder 的输入层中，使用了 Positional Encoding，使得最终的输入满足：

$$input=input\mathrm{\_}embedding+positional\mathrm{\_}encoding$$

原始正弦编码公式：

$$P{E}_{(pos,2i)}=\sin \left(\frac{pos}{{10000}^{2i/d}}\right)$$$$P{E}_{(pos,2i+1)}=\cos \left(\frac{pos}{{10000}^{2i/d}}\right)$$

前馈网络（Feed Forward Network）

包括两个线性变换+ReLU 激活：

$$\text{FFN}(x)=\text{ReLU}(x{W}_1+b_1)W_2+b_2$$

计算复杂度

当输入批次大小为 $b$，序列长度为 $N$，词向量的维度（隐藏层的维度）为 $d$ 时，$l$ 层 transformer 的计算复杂度：

Self-Attention 层

$$\text{FLOPs(Self-Attention)}=8bN{d}^2+4b{N}^{2}d$$

1. 计算 $Q$、$K$、$V$

输入输出

$$[b,N,d]\times [d,d]\to [b,N,d]$$

计算量为：

$$\text{FLOPs}=\underset{QKV}{\underbrace{3}}\ast b\ast (\underset{[N,d]\times [d,d]\text{乘法}}{\underbrace{Ndd}}+\underset{[N,d]\times [d,d]\text{加法}}{\underbrace{Ndd}})=6bN{d}^2$$

TIP

矩阵加法运算考虑偏差 bias计算量就是 $Ndd$，如果不考虑偏差就是 $Nd(d-1)$，但这个 $-1$ 一般忽略不计。

2. 计算 $Q{K}^T$

输入输出

$$[b,h,N,d_k]\times [b,h,d_k,N]\to [b,h,N,N]$$

$h$ 为注意力头数，$d_k$ 为每个头的维度，$h\ast d_k=d$。

$$\text{FLOPs}=b\ast h\ast (N^2{d}_k+N^2{d}_k)=2b{N}^{2}d$$

3. Softmax 与加权求和

Softmax 计算量较小，通常忽略。

输入输出

$$[b,h,N,N]\times [b,h,N,d_k]\to [b,h,N,d_k]$$

$$\text{FLOPs}=b\ast h\ast (N{d}_{k}N+N{d}_{k}N)=2b{N}^{2}d$$

4. 输出投影

线性变换将结果映射回 $n$ 维：

输入输出

$$[b,N,d]\times [d,d]\to [b,N,d]$$

$$\text{FLOPs}=2bN{d}^2$$

MLP 层

$$\text{FLOPs(MLP)}=16bN{d}^2$$

1. 线性层（扩展层）

输入输出

$$[b,N,d]\times [d,4d]\to [b,N,4d]$$

$$\text{FLOPs}=8bN{d}^2$$

2. 线性层（压缩层）

输入输出

$$[b,N,4d]\times [4d,d]\to [b,N,d]$$

$$\text{FLOPs}=8bN{d}^2$$

logits

Logits 层是将最终的 Transformer 隐藏层输出（维度 $d$）映射到词表大小 $V$，即一个线性投影：

输入输出

$$[b,N,d]\times [d,V]\to [b,N,V]$$

$$\text{FLOPs(logits)}=2bNdV$$

总的计算复杂度

$$\text{FLOPs(Transformer)}=l\ast (24bN{d}^2+4b{N}^{2}d)+2bNdV$$

空间复杂度

大模型在训练过程中通常采用混合精度训练，中间激活值一般是 float16 或者 bfloat16 数据类型的。在分析中间激活的显存占用时，假设中间激活值是以 float16 或 bfloat16 数据格式来保存的，每个元素占了 2 个 bytes，dropout 操作的 mask 矩阵，每个元素只占 1 个 bytes。需要保存的中间激活占用显存大小计算如下：

Self-Attention 层

1. $Q$、$K$、$V$ 共享一个输入 $X$，则显存占用为 $2bNd$
2. 对于 $Q{K}^T$，两个张量形状都是 $[b,N,d]$，显存占用为 $4bNd$
3. 对于 $\text{Softmax}$，函数输入 $Q{K}^T$ 形状为 $[b,h,N,N]$，显存占用为 $2b{N}^{2}h$
4. 计算完 $\text{Softmax}$，会进行 dropout，需要保存一个 mask 矩阵，其形状与 $Q{K}^T$ 相同，显存占用为 $b{N}^{2}h$
5. 计算 $\text{Scores}\cdot V$，二者占用显存大小为 $2b{N}^{2}h+2bNd$
6. 计算输出映射和一个 dropout 操作，二者占用显存大小为 $2bNd+bNd$

综上，Self-Attention 层的显存占用为 $11bNd+5b{N}^{2}a$

MLP 层

1. 第一个线性层的输入占用显存 $2bNd$
2. 激活函数的输入占用显存 $8bNd$
3. 第二个线性层的输入占用显存 $8bNd$
4. 最后的 dropout 操作需要保存的 mask 矩阵占用显存 $bNd$

综上，MLP 层的显存占用为 $19bNd$

LN

Self-Attention 层和 MLP 层分别对应了一个 LN，其输入占用显存为 $2bNd+2bNd$

总的空间复杂度

$$l\ast (34bNd+5b{N}^{2}h)$$

问题

Transformer 的计算复杂度为：$l\ast (24bN{d}^2+4b{N}^{2}d)+2bNdV$，需要保存的中间激活占用显存大小为：$l\ast (34bNd+5b{N}^{2}h)$，即 Transformer 模型的计算量和储存复杂度随着序列长度 N 呈二次方增长。

可以注意到，$4b{N}^{2}d$ 和 $5b{N}^{2}h$ 均产生于 Self-Attention 层。