[LLM 学习笔记] Transformer 基础

Transformer 基础

Transformer 模型架构

主要组成: Encoder, Decoder, Generator.

Encoder (编码器)

由 $N$ 层结构相同(参数不同)的 EncoderLayer 网络组成.
$\text{In}:[batch\_sz,seq\_len,d_{model}]$, $\text{Out}:[batch\_sz,seq\_len,d_{model}]$

EncoderLayer: 由一层自注意力 Multi-Head Attention (多头注意力) 子网络, 一层 Position-wise Feed-Forward (基于位置的前馈) 子网络, 以及用于连接子网络的 Residual Connection (残差连接) 和 Layer Normalization (层标准化) 组成.
$\text{In}:[batch\_sz,seq\_len,d_{model}],\text{Out}:[batch\_sz,seq\_len,d_{model}]$

• 自注意力 Multi-Head Attention 网络: Q, K, V 均来自上一层(Input Embedding/EncoderLayer)网络.
  $\text{In}:[batch\_sz,seq\_len,d_{model}],\text{Out}:[batch\_sz,seq\_len,d_{model}]$

Decoder (解码器)

由 $N$ 层结构相同(参数不同)的 DecoderLayer 网络组成.
$\text{In}:[batch\_sz,seq\_len,d_{model}],\text{Out}:[batch\_sz,seq\_len,d_{model}]$

DecoderLayer: 由一层自注意力 Masked Multi-Head Attention 子网络, 一层(Encoder-Decoder)注意力 Multi-Head Attention 子网络, 一层 Position-wise Feed-Forward (基于位置的前馈) 子网络, 以及用于连接子网络的 Residual Connection (残差连接) 和 Layer Normalization (层标准化) 组成.
$\text{In}:[batch\_sz,seq\_len,d_{model}],\text{Out}:[batch\_sz,seq\_len,d_{model}]$

• 自注意力 Masked Multi-Head Attention 网络: Q, K, V 均来自上一层(Output Embedding/DecoderLayer)网络. “Masked” 是通过掩码($[1,seq\_len,seq\_len]$)将后续位置屏蔽, 仅关注需要预测的下一个位置.
  $\text{In}:[batch\_sz,seq\_len,d_{model}],\text{Out}:[batch\_sz,seq\_len,d_{model}]$
• (Encoder-Decoder)注意力 Multi-Head Attention 网络: Q 来自上一层(Masked Multi-Head Attention)网络; K,V 来自 Encoder 的输出 memory.
  $\text{In}:[batch\_sz,seq\_len,d_{model}],\text{Out}:[batch\_sz,seq\_len,d_{model}]$

Generator (生成器)

由 $[\text{In}:d_{model},\text{Out}:vocab\_sz]$ 的线性网络和 Softmax 操作组成.
$$y=\mathrm{softmax}(\mathrm{Linear}(x))=\mathrm{softmax}(x{A}^T+b)$$
生成器是按序列顺序一次只输出下一个位置的预测概率.
$\text{In}:[batch\_sz,d_{model}],\text{Out}:[batch\_sz,vocab\_sz]$

※ Multi-Head Attention

Scaled Dot-Product Attention (缩放点积注意力):
$$\mathrm{Attention}(Q,K,V)=\mathrm{softmax}(\frac{Q{K}^{\top }}{\sqrt{d_k}})V$$
维度变化:

1. 输入:
   • $Q[batch\_sz,h,seq\_len,d_k]$
   • $K[batch\_sz,h,seq\_len,d_k]$, $K^{\top }[batch\_sz,h,d_k,seq\_len]$
   • $V[batch\_sz,h,seq\_len,d_k]$
2. $Q{K}^{\top }[batch\_sz,h,seq\_len,seq\_len]$
3. $\frac{Q{K}^{\top }}{\sqrt{d_k}}$ 与 Mask 操作: 不改变形状 $[batch\_sz,h,seq\_len,seq\_len]$
4. $\mathrm{softmax}(\frac{Q{K}^{\top }}{\sqrt{d_k}})$: 最后一维进行 Softmax 操作, 不改变形状 $[batch\_sz,h,seq\_len,seq\_len]$
5. $\mathrm{softmax}(\frac{Q{K}^{\top }}{\sqrt{d_k}})V$: $[batch\_sz,h,seq\_len,d_k]$
   

完整公式(参考 FlashAttention):
$$\begin{array}{rl}& S=\tau Q{K}^{\top }\in {\mathbb{R}}^{N\times N}\\ & S^{\text{masked}}=\text{MASK}(S)\in {\mathbb{R}}^{N\times N}\\ & P=\text{softmax}(S^{\text{masked}})\in {\mathbb{R}}^{N\times N}\\ & P^{\text{dropped}}=\text{dropout}(P,p_{drop})]\\ & \text{Attention}(Q,K,V)=O=P^{\text{dropped}}V\in {\mathbb{R}}^{N\times d}\end{array}$$

Multi-Head Attention (多头注意力) 机制:
$$\begin{array}{rl}MultiHeadAttn(Q,K,V)& =Concat(hea{d}_1,...,hea{d}_h)W^O\\ \mathrm{where}hea{d}_i& =Attention(Q{W}_i^Q,K{W}_i^K,V{W}_i^V)\end{array}$$

其中, $W_i^Q\in {\mathbb{R}}^{d_{model\times d_k}},W_i^K\in {\mathbb{R}}^{d_{model}\times d_k},W_i^V\in {\mathbb{R}}^{d_{model}\times d_v},W^O\in {\mathbb{R}}^{h{d}_v\times d_{model}}$
在实现中, $W^Q=(W_1^Q,...,W_h^Q)$, $W^K=(W_1^K,...,W_h^K)$, $W^V=(W_1^V,...,W_h^V)$, $W^O$, 由 4 个 $[\text{In}:d_{model},\text{Out}:d_{model}]$ 的线性网络组成, $d_k=d_v=d_{model}/h$

维度变化:

1. 输入: $X[batch\_sz,seq\_len,d_{model}]$
2. 多头预处理: $X[batch\_sz,seq\_len,d_{model}]$ → $X[batch\_sz,h,seq\_len,d_k]$
3. 注意力机制: $X[batch\_sz,h,seq\_len,d_k]$ → $Q,K,V[batch\_sz,h,seq\_len,d_k]$ → $\mathrm{Attention}(Q,K,V)[batch\_sz,h,seq\_len,d_k]$
4. 拼接多头结果: $Concat(hea{d}_1,...,hea{d}_h)[batch\_sz,h,seq\_len,d_k]$
5. 输出: $MultiHeadAttn(Q,K,V)[batch\_sz,seq\_len,d_{model}]$

Position-wise Feed-Forward

$$\mathrm{FFN}(x)={\mathrm{Linear}}_2(\mathrm{ReLU}({\mathrm{Linear}}_1(x)))=\max (0,x{W}_1+b_1)W_2+b_2$$

${\mathrm{Linear}}_1(x)$ : $[\text{In}:d_{model},\text{Out}:d_{ff}]$
${\mathrm{Linear}}_2(x)$ : $[\text{In}:d_{ff},\text{Out}:d_{model}]$
$\text{In}:[batch\_sz,seq\_len,d_{model}],\text{Out}:[batch\_sz,seq\_len,d_{model}]$

Add&Norm

论文中: (post-Norm)
$$\mathrm{SublayerConnection}(X)=\mathrm{LayerNorm}(X+\mathrm{Sublayer}(X))$$

AnnotatedTransformer 实现中: (pre-Norm)
$$\mathrm{SublayerConnection}(X)=X+\mathrm{Sublayer}(\mathrm{LayerNorm}(X))$$

$\text{In}:[batch\_sz,seq\_len,d_{model}],\text{Out}:[batch\_sz,seq\_len,d_{model}]$

其中:

• S u b l a y e r ∈ { M u l t i H e a d A t t n , F F N } \mathrm{Sublayer}\in\{\mathrm{MultiHeadAttn},\mathrm{FFN}\} Sublayer∈{MultiHeadAttn,FFN}
• 层标准化 $\mathrm{LayerNorm}(X)$: 对张量 $X$ 的最后一维($d_{model}$ 维, 表示每个样本) $x=X[b,pos,:]\in {\mathbb{R}}^{d_{model}}$ 进行标准化.
  $\mathrm{Norm}(x)=\frac{x-E(x)}{SD(x)+ϵ}\ast \gamma +\beta$. 其中, $E(x)$ 为平均值(期望), $SD(x)$ 为标准差, $\gamma ,\beta \in {\mathbb{R}}^{d_{model}}$ 为可学习的参数, $ϵ$ 是用于数值稳定性(避免除 0)在分母上加的一个极小值标量.
• 残差连接 (Residual Connection): $y=x+\mathcal{F}(x)$
• 注: pre-Norm 与 post-Norm 的区别, 参考: 【重新了解Transformer模型系列_1】PostNorm/PreNorm的差别 - 知乎

Token Embedding

大小为 $vocab\_sz$ 嵌入维度为 $d_{model}$ 的查询表(lookup table).
$\text{In}:[batch\_sz,seq\_len],\text{Out}:[batch\_sz,seq\_len,d_{model}]$
$$\mathrm{Embedding}(\mathrm{x})=\mathrm{lut}(x)\cdot \sqrt{d_{model}}$$

Positional Encoding

用于
$$\begin{array}{rl}& P{E}_{(pos,2i)}=\sin (pos/10000^{2i/d_{\text{model}}})\\ & P{E}_{(pos,2i+1)}=\cos (pos/10000^{2i/d_{\text{model}}})\end{array}$$

$$\mathrm{PE}(X)=X+P,\text{where}(p_{(b,pos,i)})=P,p_{(b,pos,i)}=P{E}_{(pos,i)}$$

其中, $X,P\in {\mathbb{R}}^{batch\_sz\times seq\_len\times d_{model}}$, 即 $X$ 和 $P$ 为 $[batch\_sz,seq\_len,d_{model}]$ 形状的张量; $p_{(b,pos,i)}$ 为 $P$ 对应位置的元素, $pos$ 为 token 在 $seq\_len$ 长度的序列中位置, $i$ 为 $d_{model}$ 中的维度.
$\text{In}:[batch\_sz,seq\_len,d_{model}],\text{Out}:[batch\_sz,seq\_len,d_{model}]$

Subsequent Mask

也称为 “Causal Attention Mask”, 因果注意力掩码("FlashAttention"中的说法). 用于 Decoder 的注意力网络中屏蔽预测位置之后的信息, 即仅根据预测位置及之前的信息进行预测.
掩码应用于矩阵 $Q{K}^T/\sqrt{d_k}$, 是一个包括对角线的下三角矩阵(对应保留 $Q$ 的 $seq\_len$ 索引 $i$ 大于等于 $K^T$ 的 $seq\_len$ 索引 $j$ 的计算结果), 将掩码为 0 部分(上三角部分为 0)对应的矩阵数据替换为极小值(如 -1e9).
$\text{shape}:[1,seq\_len,seq\_len]$

代码实现

• The Annotated Transformer 官方 Colab 代码: AnnotatedTransformer.ipynb
• 带详细中文注释的 Colab 代码: AnnotatedTransformer.ipynb
• The Annotated Transformer 官方 GitHub 仓库: harvardnlp/annotated-transformer
• 带详细中文注释且模型代码分离的 GitHub 仓库: peakcrosser7/annotated-transformer

参考资料

• Vaswani A, Shazeer N, Parmar N, et al. Attention is all you need[J]. Advances in neural information processing systems, 2017, 30. https://dl.acm.org/doi/10.5555/3295222.3295349
• The Annotated Transformer - Harvard University
• Self-Attention v/s Attention: understanding the differences | by Nishant Usapkar | Medium
• Self attention vs attention in transformers | MLearning.ai
• 【重新了解Transformer模型系列_1】PostNorm/PreNorm的差别 - 知乎