Tokenizer

Tokenizer 的作用是:将原始文本转换为模型可以处理的离散 Token ID 序列,并在生成后将 Token ID 重新解码为文本

常见切分方式:

  • Word-level
  • Character-level
  • Subword-level
  • Byte-level

目前大多数大语言模型采用子词切分

Word-level Tokenization

将文本按空格和标点符号分割成单词

英文中可以近似按照空格和标点切分,但中文、日文等语言通常需要额外的分词算法

  • 优点:简单直接;
  • 缺点:词表通常较大,单词形态变化会占用不同词表项,遇到词表外单词时只能映射为 [UNK]
1
2
I love machine learning
["I", "love", "machine", "learning"]

Character-level Tokenization

将每个字符作为一个 Token

1
2
hello
["h", "e", "l", "l", "o"]

优点:

  • 词表很小;
  • 几乎不会出现未知 Token;
  • 对拼写错误、罕见单词较鲁棒。

缺点:

  • Token 序列明显变长;
  • 增加模型计算复杂度和训练时间。

Subword Tokenization

现代 LLM 最常使用的方案

  • 高频词或高频片段保留为一个 Token
  • 低频词拆成多个较小的 Token

常见方法包括 BPE、WordPiece 和 Unigram

BPE

BPE:Byte Pair Encoding

核心思想:从较小的基本单位开始,不断合并训练语料中最常出现的相邻 Token Pair

1
2
3
4
5
6
7
8
9
10
11
Initial Units
↓
统计相邻 Pair Frequency
↓
选择最高频 Pair
↓
Merge
↓
加入 Vocabulary
↓
重复

直到达到指定的 Vocabulary Size

WordPiece

和 BPE 很像,核心区别在于:

  • BPE 主要根据 Pair Frequency 决定合并,
  • WordPiece 更倾向选择能够提高训练语料语言建模效果的 Subword 合并

例如:

1
playing

可能被切成:

1
2
play
##ing

其中 ##表示当前 Subword 位于一个单词内部,而不是新单词开头

Unigram

和 BPE 的思路相反

从较大的候选 Vocabulary 开始,不断删除不重要 Token

Unigram 会根据 Token 概率选择概率较高的切分

方法 核心思想 典型特点
BPE 反复合并高频 Pair 简单、常见
WordPiece 根据训练目标选择 Subword BERT 中经典
Unigram 从大词表逐步删除 Token 可以考虑多种切分方案

现代大模型核心组件

现代大模型通常会在原始 Transformer 上采用一组工程化改进

组件 作用
RMSNorm 简化归一化计算并稳定激活尺度
SwiGLU 使用门控结构增强 FFN 表达能力
RoPE 旋转 Q/K,使注意力分数包含相对位置关系
MoE 通过稀疏激活扩大模型参数容量

RMSNorm

[1910.07467] Root Mean Square Layer Normalization

LayerNorm:

$$ \operatorname{LayerNorm}(x) = \gamma \frac{x-\mu} {\sqrt{\sigma^2+\epsilon}} + \beta $$
RMSNorm 可以看作 LayerNorm 的简化版本:
$$ \operatorname{RMSNorm}(x)=\gamma\odot\frac{x}{\sqrt{\frac{1}{d}\sum_{i=1}^{d}x_i^2+\epsilon}} $$

相比 LayerNorm,RMSNorm:

  • 减少了均值计算和相关运算;
  • 保留输入的整体方向信息;
  • 对输入的正比例缩放具有不变性;
  • 在许多模型中能以更低计算成本取得与 LayerNorm 相近的效果

RMSNorm 已被 Llama、Mistral、Qwen、Gemma 等许多现代 Decoder-only LLM 广泛采用

SwiGLU

[2002.05202] GLU Variants Improve Transformer

传统 Transformer FFN 通常写成:

$$ \operatorname{FFN}(x)=W_2\phi(W_1x+b_1)+b_2 $$
普通 FFN 只有一条升维分支

SwiGLU加入门控,使用两条分支

1
2
3
4
5
           +--> ---------内容分支 ----------+
| |
输入 x ----+ * --> 降维 --> 输出
| |
+--> 门控分支 --> 激活函数 -------+

一条分支产生内容,一条分支决定这些内容应该保留、放大还是抑制

标准 GLU:

$$ \operatorname{GLU}(x)=(W_v x+b_v)\odot\sigma(W_g x+b_g) $$
SwiGLU 将门控分支中的激活函数换成 SiLU:
$$ \operatorname{SwiGLU}(x)=\operatorname{SiLU}(W_g x)\odot(W_u x) $$

混合专家结构 MoE

[2101.03961] Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity

[2401.06066] DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models

在普通稠密 Transformer 中,每个 Token 都经过同一个 FFN

主要问题:

  • 所有 Token 使用相同参数;
  • 增大 FFN 参数量时,每个 Token 的计算量同步增加;
  • 总参数量与单 Token 计算量紧密绑定。

MoE 的目标是:增加模型总参数容量,但每个 Token 只激活其中少量专家,使总参数量大幅增长,而单 Token 计算量不按相同比例增长。

可以把 Transformer 粗略理解为:Attention 做跨 Token 的信息混合,FFN 做 Token 内的通道/特征变换

因此 FFN 很适合进行专家化

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
输入 Token x
↓
Router
|
+--------> Expert 1
+--------> Expert 2
+--------> Expert 3
|
...
|
+--------> Expert N
↓
Router 选择 Top-K Experts
↓
专家输出加权求和

Router 通常是一个较小的线性层:

$$ z=W_rx, \qquad W_r\in\mathbb{R}^{E\times d} $$
其中 $z$ 包含每个专家的路由分数,再根据具体模型使用 Softmax、Sigmoid 或其他方式选择 Top-K

第 $i$ 个专家可以写成:

$$ E_i(x) = W_{\mathrm{down},i} \left[ \operatorname{SiLU}(W_{\mathrm{gate},i}x) \odot (W_{\mathrm{up},i}x) \right] \qquad \text{SiLU} = x\sigma(x) $$
最终输出:
$$ y(x) = \sum_{i\in\operatorname{TopK}(g(x))} p_i(x)E_i(x) $$
一个 Token 可以经过:
  1. 始终启用的 Shared Expert,共享专家(非必须)
  2. Router 选择的 Routed Expert,路由专家
1
2
3
                    +--> Router --> Routed FFN Experts --+
x --> RMSNorm ------+ +--> 相加
+--> Shared FFN Experts -------------+

忽略残差连接时:

$$ y(x) = \sum_{j=1}^{N_s}E_j^{\mathrm{shared}}(x) + \sum_{i\in\operatorname{TopK}(g(x))} p_i(x)E_i^{\mathrm{routed}}(x) $$
DeepSeekMoE 将这种设计称为 Shared Expert Isolation,共享专家隔离

MoE 的主要训练问题是:路由不均衡,热门专家可能过载,而冷门专家得不到充分训练

旋转位置编码 RoPE

[2104.09864] RoFormer: Enhanced Transformer with Rotary Position Embedding

位置编码的目标是让 Attention 感知 Token 顺序和位置关系

RoPE 不把位置向量直接加到 Token 表示上,而是根据 Token 的位置,对 Query 和 Key 的二维子空间进行旋转

假设 Query 中有两个维度,RoPE 根据位置 $m$,将它旋转 $m\theta$

$$ R_{m,\theta}=\begin{bmatrix}\cos(m\theta)&-\sin(m\theta)\\\sin(m\theta)&\cos(m\theta)\end{bmatrix} $$
旋转后的 Query:
$$ \widetilde q_m = R_{m,\theta}q_m $$
位置 $n$ 的 Key:
$$ \widetilde k_n = R_{n,\theta}k_n $$
二维旋转矩阵是正交矩阵,因此 RoPE 改变向量方向和相位,但不改变其欧氏范数

利用:

$$ R_m^\top R_n = R_{n-m} $$
得到:
$$ \boxed{\widetilde q_m^\top \widetilde k_n = q_m^\top R_{n-m} k_n} $$
这意味着 RoPE 之后的 Attention 内积自然包含相对位置 $n-m$

实际 Query / Key 是 $d_h$ 维,RoPE 会把相邻维度两两分组

$$ (q_0,q_1), \quad (q_2,q_3), \quad \dots, \quad (q_{d_h-2},q_{d_h-1}) $$

省略了位置下标,其实应该是 $(q_{m,2i},q_{m,2i+1})$

第 $i$ 个二维分组的常见频率:

$$ \theta_i = 10000^{-\frac{2i}{d_h}}, \qquad i=0,1,\dots,\frac{d_h}{2}-1 $$
不同频率承担不同尺度的位置表示:
  • 高频维度:相邻位置相位变化明显,提供较细的位置分辨率;
  • 低频维度:相位变化较慢,更适合表达较长距离关系。

可以类比:

1
2
3
高频 → 秒针
中频 → 分针
低频 → 时针
优点 局限
Attention 分数自然包含相对位置关系 超出训练长度时外推性能下降
不引入额外可学习位置参数 不同位置可能产生相似的组合相位
只需旋转 Q、K,计算简单 长序列下可能出现相位混叠
保持旋转部分的欧氏范数 更偏向相对位置表示
兼容 MHA、MQA、GQA 等标准点积注意力 与 MLA 的低秩矩阵吸收存在结构冲突

注意力机制

不同注意力结构的核心区别之一,是如何保存和访问历史信息

方法 历史信息组织方式 缓存特点 核心目标
MHA 每个 Query 头独立使用 K/V 最大 保留完整多头表达能力
MQA 所有 Query 头共享一组 K/V 最小 最大程度降低 KV Cache
GQA 每组 Query 头共享一组 K/V 中等 折中表达能力与缓存
MLA 每个 token 保存压缩 KV 潜变量 显著小于 MHA 低秩压缩并保留多头表达
线性注意力 全部历史压缩为固定状态 不随序列长度增长 同时降低计算量和缓存量

可以把结构优化大致分成两条路线:

模型结构路线 代表方法 主要目标
Token KV 压缩 MQA、GQA、MLA 保留 Softmax 检索并减少 KV Cache
固定状态压缩 线性注意力、GLA 同时降低序列计算量和缓存量

FlashAttention 属于计算实现优化,主要减少中间张量存储和 HBM IO

标准多头注意力 MHA

[1706.03762] Attention Is All You Need

标准 MHA 对第 $t$ 个 Token 的 Query、Key、Value 拆分为多个 Head:

$$ q_t=[q_{t,1};\ldots;q_{t,n_h}],\qquad k_t=[k_{t,1};\ldots;k_{t,n_h}],\qquad v_t=[v_{t,1};\ldots;v_{t,n_h}] $$
每个 Head 独立执行:
$$ o_{t,i}=\sum_{j=1}^{t}\operatorname{softmax}_j\left(\frac{q_{t,i}^{\top}k_{j,i}}{\sqrt{d_h}}\right)v_{j,i} $$
自回归生成时,每生成一个新 Token,都需要让当前 Query 与全部历史 Key 计算相似度,并读取相应的 Value

因此所有历史 Token 的 $(K,V)$ 都必须保存在 KV Cache 中

如果:

  • Transformer 层数为 $L$
  • Context Length 为 $T$
  • KV Head 数为 $H_{KV}$
  • Head Dimension 为 $d_h$
  • 每个元素占 $b$ Bytes

单条 Sequence 的 KV Cache 约为:

$$ M_{\mathrm{KV}}=2LTH_{KV}d_hb $$
前面的 $2$ 表示 Key 和 Value

因此长上下文、大 Batch 推理时,KV Cache 会占用大量显存和显存带宽

MQA

[1911.02150] Fast Transformer Decoding: One Write-Head is All You Need

保留多个 Query Head,但所有 Query Head 共享同一组 Key 和 Value

1
2
3
4
5
6
7
8
Q1 ─┐
Q2 ─┤
Q3 ─┤
Q4 ─┤
Q5 ─┤──> K1, V1
Q6 ─┤
Q7 ─┤
Q8 ─┘

显著减少自回归推理时 KV Cache 的大小和 HBM 访存,从而提高 Decode 吞吐

假设:

  • Batch Size:$B$
  • Sequence Length:$L$
  • Attention Heads:$H$
  • Head Dimension:$d_h$

MHA 中 KV Cache大小为

$$ 2BLHd_h $$
MQA 只有 1 个 KV Head
$$ 2BLd_h $$
但共享程度较高,也可能限制不同注意力头表达不同类型的信息

GQA

[2305.13245] GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints

GQA不改变注意力本身公式,但 Query Head 数量和 Key / Value Head 数量不再要求相同

多个 Query Head 被划分成若干组,每组共享一个 Key Head 和一个 Value Head。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
MHA
Q1 → K1,V1
Q2 → K2,V2
Q3 → K3,V3
Q4 → K4,V4

MQA
Q1 ─┐
Q2 ─┤
Q3 ─┤→ K1,V1
Q4 ─┘

GQA
Q1 ─┐
Q2 ─┤→ K1,V1
Q3 ─┐
Q4 ─┤→ K2,V2

每个 KV Head 对应的 Query Head 数量:

$$ G=\frac{H_Q}{H_{KV}} $$
GQA 可以理解成 MHA 与 MQA 的折中:
  • MHA:KV Head 最多,表达能力完整,但 KV Cache 最大;
  • MQA:只保留一组 KV,缓存最小,但共享程度最高;
  • GQA:多个 Query Head 分组共享 KV,在表达能力与缓存开销之间折中。

多头潜在注意力 MLA

[2405.04434] DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model

核心思路:不直接缓存完整的 Multi-Head Key / Value,先把隐藏状态压缩到一个低维的共享 KV Latent,再在计算 Attention 时恢复需要的 Key 和 Value

类似 LDM 的思路,将高维表示压缩到低维潜在空间,在保留任务所需信息的前提下降低计算或存储成本

KV 低秩压缩

对于第 $t$ 个 Token 的隐藏状态

$$ h_t\in\mathbb R^d $$
会直接通过两个大的线性投影得到 Key 和 Value
$$ k_t=W_Kh_t,\qquad v_t=W_Vh_t $$
其中
$$ W_K,W_V\in\mathbb R^{n_hd_h\times d} $$
MLA 先通过一个 Down Projection
$$ c_t^{KV} = W^{DKV}h_t \in\mathbb R^{d_c}\qquad W^{DKV}\in\mathbb R^{d_c\times d} $$
得到的 $c_t^{KV}$ 称为 KV Latent

然后分别通过两个 Up Projection 恢复 Content Key 和 Content Value

$$ k_t^C=W^{UK}c_t^{KV}, \qquad v_t^C=W^{UV}c_t^{KV} \qquad W^{UK},W^{UV} \in \mathbb R^{n_hd_h\times d_c} $$
整体过程
1
2
3
4
5
6
7
8
9
10
11
          h_t
│
W_DKV
│
▼
c_t^KV
KV Latent
/ \
W_UK W_UV
↓ ↓
Content Key Content Value

可以将它理解为对原来的 $K/V$ 投影进行低秩分解:

$$ W_K=W^{UK}W^{DKV}, \qquad W_V=W^{UV}W^{DKV} $$
$$ \operatorname{rank} \left( W^{UK}W^{DKV} \right) \le d_c $$

利用一个低维 Bottleneck代替原来的直接高维映射

这与 LoRA 在数学思路上非常类似,都利用低维 Bottleneck,将一个大的矩阵映射表示成两个小矩阵的乘积

Query 低秩压缩

KV 压缩是 MLA 降低 KV Cache 的核心

Query 并不必须压缩,因为 Query 只在当前计算步骤使用,不需要作为历史缓存长期保存

DeepSeek-V2 进一步对 Query 投影进行低秩分解

$$ c_t^Q=W^{DQ}h_t \qquad W^{DQ} \in \mathbb R^{d_c'\times d} $$
从 Query 潜变量生成 Content Query:
$$ q_t^C=W^{UQ}c_t^Q \qquad W^{UQ} \in \mathbb R^{n_hd_h\times d_c'} $$
投影完成后切分多个 Head:
$$ q_t^C = \left[ q_{t,1}^C; \ldots; q_{t,n_h}^C \right] $$

解耦 RoPE

Content Key 的 Up Projection 可以通过矩阵结合律吸收到 Query 一侧

$$ (q_{t,i}^C)^\top W_i^{UK}c_j^{KV} = \left((W_i^{UK})^\top q_{t,i}^C\right)^\top c_j^{KV} $$
因此推理时可以直接让 Query 与缓存的低维 $c_j^{KV}$ 计算,而不需要为所有历史 Token 显式恢复完整 Key

但如果直接对恢复后的 Content Key 应用 RoPE

那么 Attention 中会出现

$$ q_t^\top R_jW^{UK}c_j^{KV} $$
$R_j$ 是与第 $j$ 个 Token 的位置有关的 RoPE 旋转矩阵

问题在于由于 $R_j$ 随位置变化,无法简单将 $W^{UK}$ 吸收到 Query 一侧

这样就会破坏 MLA 只缓存低维 KV Latent 的优势

MLA 将 Query 和 Key 拆分为:

1
2
3
4
5
6
7
Query
├── Content Query
└── RoPE Query

Key
├── Content Key
└── RoPE Key

Content 部分负责语义匹配,RoPE 部分单独负责位置信息

Content 部分

$$ q_t^C=W^{UQ}c_t^Q\qquad k_t^C=W^{UK}c_t^{KV} $$

并进一步按照 Attention Head 切分

$$ [q_{t,1}^C;\ldots;q_{t,n_h}^C]\qquad [k_{t,1}^C;\ldots;k_{t,n_h}^C] $$
Content 部分不直接应用 RoPE,因此仍然可以进行矩阵吸收

RoPE 部分

位置 Query 单独生成

$$ q_t^R = \operatorname{RoPE}_t \left( W^{QR}c_t^Q \right) $$
并按照多个 Attention Head 得到对应的位置 Query
$$ [q_{t,1}^R;\ldots;q_{t,n_h}^R] $$
位置 Key 则直接从隐藏状态生成
$$ k_t^R = \operatorname{RoPE}_t \left( W^{KR}h_t \right) $$
其中 $k_t^R$ 可以在多个 Attention Head 之间共享
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
                    h_t
/ \
/ \
W_DQ W_DKV
↓ ↓
c_t^Q c_t^KV
/ \ / \
W_UQ W_QR W_UK W_UV
↓ ↓ ↓ ↓
Q^C RoPE K^C V^C
↓
Q^R

h_t
↓
W_KR
↓
RoPE
↓
K^R

第 $i$ 个注意力头实际使用:

$$ q_{t,i} = \left[ q_{t,i}^C; q_{t,i}^R \right], \qquad k_{j,i} = \left[ k_{j,i}^C; k_j^R \right] $$
内积可以自然拆成两项
$$ \boxed{ q_{t,i}^{\top}k_{j,i} = \left(q_{t,i}^C\right)^\top k_{j,i}^C + \left(q_{t,i}^R\right)^\top k_j^R } $$
两项分别承担:
  • Content Matching
  • Position Matching

线性注意力

[2006.16236] Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention

线性注意力和 MQA / GQA / MLA 不属于同一种压缩思路

  • MQA / GQA / MLA:减少每个 Token 需要保存的 KV 信息量,但历史 Token 仍然分别存在,因此缓存大小仍然随序列长度增长
  • Linear Attention:不再保存每个历史 Token 的独立 (K,V),而是不断把历史信息累积进一个固定大小的状态(state)

改变计算顺序

线性注意力要求相似度函数可以写成

$$ \operatorname{sim}(q,k) = \phi(q)^\top\phi(k) $$
其中
$$ \phi:\mathbb R^{d_k}\rightarrow\mathbb R^m $$
先分别对 Query 和 Key 做一个特征映射,然后相似度直接写成二者的内积

代入 Attention

$$ y_t = \frac{ \sum_{i=1}^{t} \phi(q_t)^\top\phi(k_i)v_i }{ \sum_{i=1}^{t} \phi(q_t)^\top\phi(k_i) } $$
因为对于固定的 $t$,$\phi(q_t)$ 与求和下标 $i$ 无关,所以可以提到求和外面
$$ y_t = \frac{ \phi(q_t)^\top \left( \sum_{i=1}^{t} \phi(k_i)v_i^\top \right) }{ \phi(q_t)^\top \left( \sum_{i=1}^{t} \phi(k_i) \right) } $$
**利用矩阵乘法结合律,把“Query 分别和所有历史 Key 计算改成先把历史 Key-Value 汇总起来**

定义:

$$ S_t = \sum_{i=1}^{t} \phi(k_i)v_i^\top\qquad z_t = \sum_{i=1}^{t}\phi(k_i) $$
那么 Attention 输出就变成
$$ y_t = \frac{ \phi(q_t)^\top S_t }{ \phi(q_t)^\top z_t } $$
$S_t$ 和 $z_t$ 的大小不随序列长度 $t$ 增长

历史最终都被压缩进

$$ S_t\in\mathbb R^{m\times d_v} \qquad z_t\in\mathbb R^m $$
复杂度约为
$$ O(Nd^2) $$

递归更新

$$ q'_t=\phi(q_t), \qquad k'_t=\phi(k_t) $$

状态更新:

$$ S_t=S_{t-1}+k'_tv_t^\top\\ z_t=z_{t-1}+k'_t $$
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
当前 Token x_t
│
├──> q_t
├──> k_t
└──> v_t
│
▼
历史状态 S_{t-1}, z_{t-1}
│
▼
S_t = S_{t-1} + φ(k_t)v_t^T
z_t = z_{t-1} + φ(k_t)
│
▼
用 q_t 读取状态
│
▼
y_t

和 RNN 非常像

普通 RNN

$$ h_t=F(h_{t-1},x_t)\qquad y_t=G(h_t,x_t) $$
因此,可以把线性注意力的隐藏状态定义为:
$$ h_t^{\mathrm{LA}}=(S_t,z_t) $$
其状态更新函数为:
$$ h_t^{\mathrm{LA}}=F(h_{t-1}^{\mathrm{LA}},x_t)=\left(S_{t-1}+k'_tv_t^\top,z_{t-1}+k'_t\right) $$
从这个角度看,线性注意力本质上是一个具有矩阵隐藏状态的 RNN
1
2
3
4
5
RNN:
历史 → 压进一个向量

Linear Attention:
历史 → 压进一个 Key-Value 统计矩阵

Linear Attention 在自回归推理时可以使用 RNN 式递归更新;训练时则可以利用矩阵计算、前缀扫描或分块方法进行并行计算

特征映射

标准 Softmax Attention 使用

$$ \operatorname{sim}(q,k) = \exp(q^Tk) $$
Linear Attention 将相似度函数替换成另一个可以显式分解的 Kernel,并选择:
$$ \phi(x) = \operatorname{ELU}(x)+1 = \begin{cases} x+1,&x>0,\\ e^x,&x\leq0 \end{cases} $$
可以保证 $\phi(x)>0$,作为一种非负的 Attention 相似度

从 Kernel 的角度,指数点积核理论上可以写成某个特征空间中的内积,但对应的特征空间通常是无限维的

因此 Linear Attention 与标准 Softmax Attention 并不完全等价,而是定义了一种新的 Attention 机制

GLA

[2312.06635] Gated Linear Attention Transformers with Hardware-Efficient Training

基础线性注意力:

$$ S_t = S_{t-1} + k_tv_t^\top $$

旧状态只能持续累积,不能根据当前输入选择性遗忘

一种简单方式是加入固定衰减:

$$ S_t = \gamma S_{t-1} + k_tv_t^\top, \qquad 0<\gamma<1 $$
更一般地:
$$ S_t = G_t\odot S_{t-1} + k_tv_t^\top $$
GLA 根据当前输入计算遗忘门:
$$ \alpha_t = f_\alpha(x_t), \qquad \alpha_t\in(0,1)^{d_k} $$
状态更新:
$$ S_t = \operatorname{Diag}(\alpha_t) S_{t-1} + k_tv_t^\top $$
因此:GLA 仍然是一个具有矩阵隐藏状态的 RNN,只是其状态转移由当前输入动态控制