Di's Notes
LLM Basics · 连载中 · 第一篇

Step by Step 读懂 LLM(一):Embedding 与 BPE

这个系列想一步一步把 LLM 讲通:Transformer 的每个组件,到推理时的 Prefill 和 Decode,再到 Encode。首先从 Transformer 的基本架构开始,一个 LLM 从读入文字到输出下一个词,依次要经过这几部分:

  1. BPE:把文本切成 token
  2. Embedding Layer:把 token 变成向量
  3. Layer Norm / RMSNorm
  4. Self Attention
  5. MLP
  6. Cross Entropy 分类器:输出下一个 token 的概率
一个 LLM 的组成overview
图 1点任意一块看它做什么。蓝色是这一篇讲的两块;3–5 组成一个 Transformer block,真实模型会把它堆叠几十层。

1Embedding Layer

列表里 BPE 排在第一,但我想先讲 Embedding,因为理解了 embedding,才能更好地理解我们为什么需要 BPE。

以英文为例,一段话进来,可能包含词典里的任何一个词。但 LLM 和决策树、LDA 这类传统模型不同,它需要的输入是数字,这样才能做矩阵乘法、激活函数等各种计算。所以 NLP 和语言模型一直以来的一个重要方向,就是如何用数字、用向量来表示输入。

这中间经历过大家耳熟能详的 one-hot,再到后来的 word embedding。word embedding 之后,行业建立了一个共识:用一个固定维度的向量来代表一个词。最早的通用 word embedding(比如 word2vec[1])通过“预测周围的词”这样的任务来训练:意思相近的词,周围出现的词也相近,训练出来的向量也就相近。训练好之后,任何模型都可以直接拿来使用。

word2vec:用一个词预测它周围的词
图 2点任意一个词,或用箭头移动。蓝色是中心词,黄色是要预测的上下文。窗口在整个语料上滑过,产生海量的(中心词 → 上下文)训练样本;这是 word2vec 中 skip-gram 的做法。

但后来大家把它用在各类不同的任务里,发现把 embedding 向量也作为模型参数的一部分,跟着具体任务一起训练,更合理,效果也更好。这就是 Transformer 和 LLM 里 embedding layer 的由来。

假设我们的词表有 n 个词,每个词用一个 b 维的向量来表达它的特征,那 embedding 总共就是一个 n×b 的矩阵 E,词 i 用的就是矩阵中的第 i 行。

这种情况下,一个很自然的、把一段话的每个词都变成 embedding 的方法是:每个词先用 one-hot 表示。一段话有 t 个词,就是一个 t×n 的矩阵,每行只有一个位置是 1,其他全是 0。用这个矩阵乘以 n×b 的 embedding 矩阵,结果就是 t×b 的特征矩阵:

X = onehot(ids) · E, (t×n)(n×b) = t×b
one-hot 矩阵乘法 = 查表
换成真实的大小
图 3小图里 n = 10,b = 4,t = 4,embedding 的数值是示意。下方按一段 t = 1000 个 token、b = 4096 维估算:one-hot 乘法里每一行只有一个 1,其余的乘加全部是“乘 0”。

但这样做,在词库很大的情况下,构建 t 个 n 维的 one-hot 向量本身就是很大的工作量,还要再做一个绝大多数运算都是在浪费时间“乘 0”的矩阵乘法。one-hot 乘 E 的结果,其实就是把 E 的第 id 行原样复制出来。所以现代的做法都是直接用 token id 查表(lookup),从矩阵里取出对应的那一行。两种方法的结果一模一样,查表只是省掉了所有乘 0 的计算。

反向传播时也一样

训练时,梯度也只会流回被查到的那几行:没有在这批数据里出现的词,它那一行 embedding 这一步不会被更新。这一点在下一节会变得很重要。

2为什么需要 BPE

说完 embedding,回来看看为什么我们需要 BPE。如果直接用“词”作为 embedding 的单位,会遇到三个问题。

第一,词太多。单以英文为例,仅牛津英语词典就收录了超过 60 万个词[4],这还仅仅只是一种语言。每个词都要占 embedding 矩阵的一行,矩阵会大得离谱,而且总会遇到词表里没有的新词、拼写错误和人名。

第二,分布极不均匀。词表中大量的词是生僻词,在我们能收集到的全部数据中出现次数极少。词频大致服从 Zipf 定律:排第 r 的词,出现次数约和 1/r 成正比。少数高频词占了绝大部分文本,长长的尾巴上是几十万个几乎见不到的词。它们那一行 embedding 几乎得不到训练(还记得上面说的,只有出现过的行才会被更新)。

词频的长尾zipf
图 4按 Zipf 定律(频率 ∝ 1/排名,总共 60 万个词)画的理想化示意,横纵轴都是对数刻度。蓝色是词表保留的词,红色是被扔掉的长尾。

第三,词本身有构成规律。还是以英文为例,单词有很多构词的规律,比如 -tion 是名词后缀,un- 表示否定,-er 表示“做这件事的人或东西”。学习这些基本单元,比学习一个个整词通用得多:模型见过 decode 和 -er,就能理解从没见过的 decoder。

这就是 token 的由来:比字母大、比单词小的单元。常见的词自己就是一个 token,生僻词被拆成几个常见的片段,任何文本都能被表示,词表大小也可以控制(例如 Llama 3 的词表是 128,256 个 token[3])。

3BPE 怎么分词

但如何分词呢?这时候就用到了 BPE(Byte Pair Encoding)算法[2]。它的逻辑是:

  1. 首先人为决定一个目标词表大小,也就是你想要最后分出多少个 token。
  2. 从最小的单元开始,也就是英文字母。词表一开始只有这些字母。
  3. 统计文本数据中所有相邻的 pair(前后连续的两个单元)的出现次数。比如 decode 这个词:de 2 次,ec 1 次,co 1 次,od 1 次。
  4. 把出现次数最多的 pair 合并成一个新 token,加进词表。还是以 decode 为例,de 合并成一个 token 之后,这个词变成 de · c · o · de,新的 pair 是 (de, c)、(c, o)、(o, de)。
  5. 重新统计,再合并。一直这样做,直到词表数量达到目标为止。

下面是一个小语料上的完整过程。每一步都在统计 pair、合并最多的那个;合并的顺序本身也会被记下来,它就是训练好的 tokenizer。

BPE 一步步合并live
14

语料(次数 · 当前切分)

相邻 pair 的出现次数

词表
用学到的规则切分新词
新词会按合并规则出现的顺序依次合并。红色是初始字母表里没有的字符:只从字母开始的 BPE 处理不了它们,下面会讲怎么解决。
图 5语料是 6 个词和它们出现的次数。pair 的次数要乘上词出现的次数:decode 出现 6 次,它里面的 de 就贡献 12 次。黄色是下一步要合并的 pair,蓝色是已经合并出来的 token。

可以看到几件事:

从字节开始:byte-level BPE

上图从“字母”开始,遇到字母表以外的字符(红色)就没办法了。中文、emoji、各种符号都会有这个问题。GPT-2 以来的 LLM 一般从 256 个字节开始做 BPE[5]:任何文本都能写成 UTF-8 字节序列,所以初始词表永远覆盖所有输入,不会出现词表外的字符。一个汉字在 UTF-8 里是 3 个字节,常见的字和词会在合并过程中被合成一个 token,生僻字最坏也能退回到 3 个字节 token。

到这里,一段文字就变成了一串 token id,再通过 embedding 查表,变成一个 t×b 的矩阵。这个矩阵就是接下来 Transformer block 的输入。

未完待续

下一篇继续沿着图 1 往右走:

  1. Layer Norm / RMSNorm
  2. Self Attention
  3. MLP
  4. Cross Entropy 分类器

再往后是推理时的 Prefill 与 Decode,以及 Encode。

参考

  1. Mikolov, Chen, Corrado, Dean. Efficient Estimation of Word Representations in Vector Space. 2013.
  2. Sennrich, Haddow, Birch. Neural Machine Translation of Rare Words with Subword Units. ACL 2016.
  3. Llama Team, AI @ Meta. The Llama 3 Herd of Models. 2024.
  4. Oxford English Dictionary. About the OED.
  5. Radford et al. Language Models are Unsupervised Multitask Learners(GPT-2). 2019.