LLM 学习笔记——Transformer 架构解读
这份笔记关注大语言模型中的 Transformer 架构的解读和理解。
对于我来说,这份笔记写得太晚了,而且很多内容都是过时的,还是在理解最初的那几篇工作中的内容,但是理解 Transformer 本身是一个在 LLM 时代必学必会的内容,因此还是需要整理一下。
LLM 已经飞速发展近十年了,最初只是相关专业的研究人员在鼓捣,但是现在已经彻底普及并且在改变所有行业了。按照目前的趋势,每一个本科生都应该要掌握大语言模型 Transformer 架构的基础知识,这些知识未来甚至会下沉到高中课程中。毫不客气地说,不理解 LLM,不理解 Transformer,就是 AI 时代的文盲。
准备
Transformer 最初用于序列到序列任务,后来逐渐发展出 Encoder-only、Encoder–Decoder 和 Decoder-only 三条路线。
上图整理了 2018—2023 年间代表性模型的架构路线,其中蓝色表示 Decoder-only 模型。
三种结构的特点如下表
| 模型类型 | 输入输出模式 | 典型任务或模型 |
|---|---|---|
| Encoder-only | 输入序列 $\to$ 上下文表示 | 文本理解、分类,例如 BERT |
| Encoder–Decoder | 输入序列 $\to$ 输出序列 | 翻译、摘要,例如原始 Transformer、T5 |
| Decoder-only | 输入序列 $\to$ 自回归续写 | 文本生成,例如 GPT、LLaMA、Qwen |
三种结构使用的基本模块相同,主要差别在于模块的组合方式以及注意力可以读取哪些位置:
- Encoder 使用双向自注意力,每个位置可以读取整个输入序列;
- Decoder 使用因果自注意力,每个位置只能读取自己及之前的位置;
- Encoder–Decoder 的 Decoder 还包含交叉注意力,用于读取 Encoder 的输出。
目前 GPT、LLaMA、Qwen 等生成式大语言模型主要采用 Decoder-only 架构,因此本文会先介绍 Transformer 的通用结构,再重点解释 Decoder-only 模型。
整体结构
大语言模型的输入和输出是文本字符串,但模型内部处理的都是浮点数向量和矩阵。
整体过程可以概括为:
1 | 文本 -> token -> token ID -> 向量表示 -> Transformer Block x L -> 词表上的概率分布 -> token -> 文本 |
下面依次解释各部分。
文本与向量转换
首先考虑从文本字符串到浮点数向量的转换过程。
分词(Tokenization):原始文本经过分词器转换为 token 序列。一个 token 可能是一个字、一个单词、一个子词或字节片段,也可能是特殊标记。分词器通常预先在语料上训练完成,在大模型训练过程中保持不变。
Token 转换为 Token ID:每个 token 在词表中都有唯一的整数 ID,因此 token 序列可以转换为整数序列 $$ (i_1,i_2,\ldots,i_T),\qquad i_t\in{1,\ldots,V}, $$ 其中 $T$ 是序列长度,$V$ 是词表大小(通常很大)。
Token ID 转换为向量表示(Embedding):Embedding 层是一张可学习的参数表 $$ E\in\mathbb{R}^{V\times d}. $$ 根据 token ID 查表,得到 $$ X_{\text{token}} =E[i_1,i_2,\ldots,i_T] \in\mathbb{R}^{T\times d}. $$ 矩阵的每一行对应一个 token,每一列对应一个隐藏特征。相比维度等于词表大小的独热编码,embedding 使用更低维的稠密向量表示 token,而且这些向量会在训练过程中不断更新。
加入位置信息:自注意力本身并不包含 token 的先后顺序,需要额外注入位置信息。原始 Transformer 使用固定的正弦、余弦位置编码: $$ X=X_{\text{token}}+X_{\text{pos}}. $$ 也可以使用可学习的位置 embedding。
现代大语言模型也常使用 RoPE(Rotary Position Embedding),根据 token 的位置旋转 $Q$ 和 $K$,使注意力分数包含相对位置信息。
经过以上步骤,每个 token 都被表示为一个 $d$ 维浮点数向量,完整的长度为 $T$ 的 Token 序列表示为矩阵 $X\in\mathbb{R}^{T\times d}$。
注意力机制
注意力汇聚
假设有一个查询向量 $q$ 和 $m$ 个键值对
$$ {(k_i,v_i)}_{i=1}^{m}. $$
首先使用注意力评分函数计算 query 与第 $i$ 个 key 的匹配分数:
$$ s_i=a(q,k_i). $$
再对全部分数进行 Softmax,得到注意力权重
$$ \alpha_i =\frac{\exp(s_i)}{\sum_{j=1}^{m}\exp(s_j)}, \qquad i=1,2,\ldots,m. $$
最后按照注意力权重对 value 求加权和:
$$ f(q)=\sum_{i=1}^{m}\alpha_i v_i. $$
因为 $\alpha_i\ge 0$ 且 $\sum_{i=1}^{m}\alpha_i=1$,输出可以看成所有 value 的加权平均。因此可以把注意力理解为一次检索:
- Query 表示想查询什么;
- Key 表示每条信息用什么来匹配;
- Value 表示匹配之后实际取回什么。
缩放点积注意力
Transformer 使用缩放点积作为注意力评分函数:
$$ a(q,k)=\frac{q^T k}{\sqrt{d_k}}, $$
其中 $d_k$ 是 query 和 key 的维度。除以 $\sqrt{d_k}$ 是为了避免维度增大后点积的数值过大,导致 Softmax 过于尖锐。
考虑矩阵形式。给定 $n$ 个 query、$m$ 个 key 和 value:
$$ Q\in\mathbb{R}^{n\times d_k},\qquad K\in\mathbb{R}^{m\times d_k},\qquad V\in\mathbb{R}^{m\times d_v}, $$
缩放点积注意力为
$$ \operatorname{Attention}(Q,K,V) =\operatorname{Softmax}\left(\frac{QK^T}{\sqrt{d_k}}+M\right)V \in\mathbb{R}^{n\times d_v}. $$
这里的 Softmax 按行操作。矩阵 $QK^T\in\mathbb{R}^{n\times m}$ 保存每个 query 与每个 key 的匹配分数,经过 Softmax 后再对 $V$ 的各行加权求和。
$M$ 是可选的掩码矩阵:
- 双向注意力不屏蔽序列中的前后位置;
- 因果注意力把未来位置对应的分数设为 $-\infty$;
- Padding mask 用于屏蔽批处理中为了对齐长度而添加的填充位置。
Transformer 的基本模块
Transformer 主要由两种子层反复组合构建:
- 注意力(Attention)子层;
- 前馈网络(Feed-Forward Network, FFN)子层。
每个子层外还包含归一化和残差连接。
自注意力层
对于输入矩阵 $X\in\mathbb{R}^{T\times d}$,首先进行三次线性变换:
$$ Q=XW_Q,\qquad K=XW_K,\qquad V=XW_V, $$
其中 $W_Q$、$W_K$ 和 $W_V$ 是可学习的参数矩阵。因为 $Q$、$K$、$V$ 来自同一个输入,所以称为自注意力(Self-Attention)。注意这并不意味着三者相同:不同的线性变换使它们分别承担查询、匹配和传递信息的作用。
在 Decoder-only 模型中使用因果掩码
$$ M_{ij}= \begin{cases} 0,&j\le i,\ -\infty,&j>i. \end{cases} $$
因此第 $i$ 个位置只能读取第 $1$ 到第 $i$ 个位置的信息。掩码经过 Softmax 后,未来位置对应的注意力权重为 $0$。
多头注意力
为了在不同子空间中学习不同的匹配方式,Transformer 使用多头注意力(Multi-Head Attention):
$$ \operatorname{head}_r =\operatorname{Attention} \left(XW_Q^{(r)},XW_K^{(r)},XW_V^{(r)}\right), $$
$$ \operatorname{MHA}(X) =\operatorname{Concat} (\operatorname{head}_1,\ldots,\operatorname{head}_h)W_O. $$
如果模型隐藏维度为 $d$,头数为 $h$,通常每个头的维度为 $d_h=d/h$。各个头独立计算注意力,拼接后再通过 $W_O$ 映射回 $d$ 维。
不同注意力头的职责不是预先指定的,而是在训练过程中形成。它们可能关注不同位置、不同语义关系或不同类型的上下文信息。
MQA 让所有 query 头共享一组 key/value,GQA 则让一组 query 头共享一组 key/value,可以减少推理时的显存占用。
前馈全连接层
前馈网络对矩阵逐行操作,每个 token 使用相同的参数,但不同 token 在这一子层中不会交换信息。
原始 Transformer 使用包含一个隐藏层的 MLP:
$$ \operatorname{FFN}(X) =\max(0,XW_1+b_1)W_2+b_2, $$
其中输入维度和输出维度都是 $d$,中间维度 $d_{\mathrm{ff}}$ 通常更大:
$$ W_1\in\mathbb{R}^{d\times d_{\mathrm{ff}}},\qquad W_2\in\mathbb{R}^{d_{\mathrm{ff}}\times d}. $$
现代大语言模型常使用 SwiGLU 等门控前馈网络:增加一条门控分支,经过逐元素乘法后再投影回隐藏维度。
注意力负责不同 token 之间的信息交换,前馈网络负责对每个 token 的特征进行非线性变换。Transformer Block 基本就是这两种操作的重复组合。
残差连接与归一化
注意力子层和 FFN 子层的输入、输出形状相同,都是 $T\times d$,因此可以在子层外加入残差连接。设子层为 $F$,残差连接为
$$ Y=X+F(X). $$
子层只需要学习相对于输入 $X$ 的变化量;同时,加法提供了一条从 $Y$ 直接回到 $X$ 的梯度传播路径,使很多个 Transformer Block 可以稳定堆叠。
归一化用于控制隐藏状态的数值尺度。Transformer 中的归一化对每个 token 的隐藏向量分别进行,不会混合不同 token,也不依赖同一批次中的其它样本。
LayerNorm
对于一个 token 的隐藏向量 $x\in\mathbb{R}^d$,首先计算各分量的均值与方差:
$$ \mu=\frac{1}{d}\sum_{j=1}^{d}x_j, \qquad \sigma^2=\frac{1}{d}\sum_{j=1}^{d}(x_j-\mu)^2. $$
LayerNorm 定义为
$$ \operatorname{LayerNorm}(x) =\gamma\odot \frac{x-\mu} {\sqrt{\sigma^2+\varepsilon}} +\beta, $$
其中 $\gamma,\beta\in\mathbb{R}^d$ 是可学习参数,$\varepsilon$ 用于避免分母为零。经过标准化后,向量被重新缩放和平移,因此归一化不会限制模型只能使用固定的均值和方差。
RMSNorm
RMSNorm 不减去均值,只使用均方根进行缩放:
$$ \operatorname{RMSNorm}(x) =\gamma\odot \frac{x} {\sqrt{\frac{1}{d}\sum_{j=1}^{d}x_j^2+\varepsilon}}. $$
与 LayerNorm 相比,RMSNorm 省略了中心化和偏移参数 $\beta$。LLaMA 等现代大语言模型通常采用 RMSNorm。
Post-Norm 与 Pre-Norm
归一化与残差连接存在两种常见的排列方式。原始 Transformer 使用 Post-Norm,先计算子层和残差,再进行归一化:
$$ Y=\operatorname{Norm}(X+F(X)). $$
现代大语言模型通常使用 Pre-Norm,先对输入进行归一化,再进入子层,最后与原输入相加:
$$ Y=X+F(\operatorname{Norm}(X)). $$
Pre-Norm 中的残差分支从 $X$ 到 $Y$ 不经过归一化层,梯度可以沿这条路径直接传播,因此更适合堆叠很深的网络。无论采用哪种排列方式,每个 Transformer Block 中的注意力和 FFN 都各自包含一次归一化与残差连接。
编码器与解码器
原始 Transformer 采用 Encoder–Decoder 架构。
Encoder 由多个相同的 Encoder Layer 级联而成。每层包含:
- 双向自注意力子层;
- 前馈网络子层。
Decoder 也由多个 Decoder Layer 级联而成。每层包含:
- 因果自注意力子层;
- 交叉注意力子层;
- 前馈网络子层。
在交叉注意力中,query 来自 Decoder 当前的隐藏状态,key 和 value 来自 Encoder 的最终输出:
$$ Q=X_{\text{decoder}}W_Q,\qquad K=X_{\text{encoder}}W_K,\qquad V=X_{\text{encoder}}W_V. $$
因此 Decoder 可以一边读取已经生成的前缀,一边读取 Encoder 对输入序列的表示。
Decoder-only 模型去掉整个 Encoder,同时也去掉交叉注意力。输入提示词和模型生成的内容被放在同一条 token 序列中,所有 Transformer Block 都使用因果自注意力:
1 | Embedding |
这里的“Decoder-only”沿用了原始 Transformer 的命名,但其中的 Block 并不等同于原始 Decoder Layer:它不再包含读取 Encoder 输出的交叉注意力。
输出与采样
最后一个 Transformer Block 输出
$$ H\in\mathbb{R}^{T\times d}. $$
对于常见的 Pre-Norm Decoder-only 模型,最后一个 Transformer Block 的输出还会经过一次归一化:
$$ \widetilde{H}=\operatorname{Norm}(H). $$
然后通过线性层映射到词表维度,得到 logits:
$$ Z=\widetilde{H}W_{\text{vocab}}\in\mathbb{R}^{T\times V}. $$
对每一行进行 Softmax,就得到对应位置对下一个 token 的条件概率分布:
$$ p(x_{t+1}=j\mid x_{\le t}) =\frac{\exp(Z_{t,j})} {\sum_{k=1}^{V}\exp(Z_{t,k})}. $$
在输入、输出使用同一个词表且维度匹配时,可以令输出矩阵与 token embedding 矩阵共享参数,这称为权重绑定(Weight Tying)。
模型直接输出的是 logits,而不是已经确定的 token。生成过程中需要采用具体的采样策略选择下一个 token:
- 贪心搜索(Greedy Search):每一步选择 logits 最大的 token;
- Top-K 采样:只保留概率最高的 $K$ 个 token,重新归一化后随机采样;
- Top-P 采样:按概率从高到低选择累计概率达到阈值 $P$ 的最小 token 集合,重新归一化后随机采样。
还可以在 Softmax 中加入温度参数 $\tau>0$:
$$ p_j=\frac{\exp(z_j/\tau)}{\sum_k\exp(z_k/\tau)}. $$
$\tau<1$ 时概率分布更加集中,生成结果更确定;$\tau>1$ 时概率分布更加平缓,随机性更强。温度通常与 Top-K 或 Top-P 配合使用。
Encoder–Decoder 模型
Encoder–Decoder 模型“吃进去一个序列,吐出来另一个序列”。输入序列由 Encoder 编码,输出序列由 Decoder 自回归生成,两条序列通过交叉注意力连接。典型任务包括机器翻译、摘要和文本改写。
推理/生成过程
Encoder–Decoder 模型先编码完整的输入序列,再逐个生成输出 token:
- 输入文本经过分词器转换为 token 序列,再转换为 embedding 并加入位置信息;
- Encoder 一次处理完整的输入序列,得到上下文表示 $Y=\operatorname{Encoder}(X)$;
- Decoder 的输入从起始标记开始,通过因果自注意力读取已经生成的输出前缀;
- 在每个 Decoder Layer 的交叉注意力中,query 来自 Decoder,key 和 value 来自 Encoder 输出 $Y$;
- Decoder 最后一个位置的隐藏状态经过输出层,得到下一个 token 的概率分布;
- 选择一个 token 追加到 Decoder 输入末尾,重复生成,直到产生结束标记或达到长度上限。
可以简写为
1 | 输入序列 -> Encoder -> 上下文表示 Y |
输入序列只需要经过一次 Encoder。自回归生成期间,Encoder 的输出保持不变,并在每一步通过交叉注意力供 Decoder 读取。
训练过程
设输入序列为
$$ (x_1,x_2,\ldots,x_S), $$
目标输出序列为
$$ (y_1,y_2,\ldots,y_U). $$
训练过程如下:
Encoder 一次读取完整的输入序列,得到上下文表示;
将目标序列右移一位,使用 $(\mathrm{BOS},y_1,\ldots,y_{U-1})$ 作为 Decoder 输入,$(y_1,y_2,\ldots,y_U)$ 作为预测目标;
Decoder 通过因果自注意力读取目标前缀,并通过交叉注意力读取 Encoder 的输出;
Decoder 在所有位置并行输出 logits,并计算交叉熵损失:
$$ \mathcal{L} =-\sum_{t=1}^{U}\log p(y_t\mid y_{<t},x_1,\ldots,x_S). $$
训练时直接把真实的目标前缀送入 Decoder,这种方式称为 Teacher Forcing。因为输入序列和完整的目标序列都已经给出,Encoder 可以并行处理所有输入位置,Decoder 也可以在因果掩码下并行计算所有目标位置。推理时没有真实的目标序列,只能逐个生成 token。
Decoder-only 模型
Decoder-only 模型把提示词和生成结果放在同一条序列中,可以理解为“吃进去一段前缀,继续吐出它的后继”。模型只使用因果自注意力,不需要单独的 Encoder 和交叉注意力。
推理/生成过程
Decoder-only 大语言模型按照以下步骤生成文本:
- 输入文本字符串,经过分词器得到 token 序列;
- 将 token ID 转换为 embedding,并按照所采用的位置编码方法注入位置信息;
- 将隐藏状态依次传入所有 Transformer Block;
- 使用 LM Head 得到 logits 矩阵;
- 只取最后一行 logits,得到下一个 token 的概率分布;
- 按照采样策略选择一个 token,将它追加到原序列末尾;
- 重复上述过程,直到生成结束标记或达到长度上限。
可以简写为
1 | 输入 token 序列 |
实际推理通常使用 KV Cache 保存历史 token 在各层产生的 key 和 value,避免每生成一个 token 都重新计算完整前缀。
训练过程
Decoder-only 模型使用完整的 token 序列训练,并将序列错开一个位置作为监督信号。对于序列
$$ (c_1,c_2,\ldots,c_T), $$
训练过程如下:
将 token 序列转换为输入矩阵 $X\in\mathbb{R}^{T\times d}$;
将 $X$ 输入 Decoder-only Transformer,使用因果掩码保证第 $i$ 个位置只能读取前 $i$ 个 token;
输出 logits 矩阵 $Z\in\mathbb{R}^{T\times V}$,其中第 $i$ 行用于预测 $c_{i+1}$;
对所有有效位置计算交叉熵损失:
$$ \mathcal{L} =-\sum_{i=1}^{T-1}\log p(c_{i+1}\mid c_1,\ldots,c_i). $$
一次前向传播同时完成
1 | 基于 {c_1} 预测 c_2 |
因果掩码保证各位置不会读取未来信息,因此这些预测可以通过矩阵运算并行计算。推理时,下一个 token 尚未确定,必须先生成 $c_{t+1}$ 才能继续生成 $c_{t+2}$,所以 token 之间仍然是串行的。
总结
以典型的 Pre-Norm Decoder-only Transformer 为例,完整数据流可以总结为:
1 | token IDs (T,) |
各部分的作用分别是:
- Embedding 把离散 token 转换到连续向量空间;
- 位置编码或 RoPE 提供 token 的顺序信息;
- 因果自注意力让当前位置有选择地读取历史信息;
- FFN 分别对每个位置的特征进行非线性变换;
- 残差连接和归一化使很多层 Block 可以稳定堆叠;
- LM Head 把隐藏状态转换为词表上的预测分数。
Transformer 可以看成两类操作的反复组合:注意力负责不同 token 之间的信息交换,前馈网络负责每个 token 内部的特征变换。现代大语言模型虽然在位置编码、归一化、注意力头、前馈网络和稀疏结构等方面存在很多变体,但主体仍然可以放进这一框架理解。
原来是 LLM 生成文本啊,不知道的还以为是把输入剁成一堆拿大量文本统计出的最常见字符串组合的词元,高频组合各占一个词元而生僻词拆成更小的子词元,查询一张几万行的大对照表把每个词元换成几千维的数字串每一层除以各元素的均方根,下面把每串数字各与三个不同的矩阵相乘变出 QKV 三个值,在 Q 和 K 上按各自的位置乘一个旋转矩阵把位置信息融合进去,所有的位置用 Q、K 两两点积再除以根号 d_k 把数值压回正常量级,算完先把当前位置后面的未来词元全部用负无穷遮住,再用 softmax 把分数归一化按权重把 V 加权混合成新的数字串,并行跑好几个头再横向拼在一起,再乘一个 W_O 输出投影矩阵线性变换,把结果通过残差加回原来的数字串,省得十几层下来原始信息洗没梯度也传不回去。再过一遍 RMSNorm,扔进一个先把维度扩到约八除以三倍再缩回来的两层矩阵变换里,中间把输入拆两路一路过 SiLu 激活一路不过逐元素相乘,再把结果加回去。这个流程叠几十次,乘一个把几千维压成词表大小的矩阵,过一遍 softmax 变成概率分布,拿温度决定分布宽窄,再或者只保留分数前 k 个,把后面直接归零,或者搞核采样从高到低累加,超过阈值 p 就砍掉后面,只留核心。然后,按这个分布随机选择一个词元,选完把它接到输入后面接着跑下一轮,每层每个注意力头把之前所有步骤的 KV 都留在显存里。新词元算自己一行 Q 跟存好的 K 点积,全程一直循环直到选到那个 eos 的特殊词元,把每次选择的词元拼到一起组合成为最终输出文本呢。