↓ Skip to main content
111qqz’s blog
  1. Posts/

从 Token ID 到 Dense Vector：Embedding 到底在做什么

·3886 words·8 mins
Deep-Learning Embedding Tokenizer Cs336 6.s191 Ml-Infra
Table of Contents
Note: This article is available in Chinese only. 本文暂无英文版本。 View original

起因
#

最近在继续 CS336 和 MIT 6.S191。

之前那篇写了 BPE tokenizer——从训练到编解码，关注的是 raw text 怎么变成 token，token 怎么映射成 token ID。做完之后，对这条链路的前半段算是搞清楚了：

1raw text → tokenizer / BPE → token → token ID

在 6.S191 重新看到 one-hot 和 embedding 的时候，发现后半段也能接上：

1token ID → one-hot → embedding → dense vector → neural network

这两段以前是分开学的。BPE 的文章讲完 token ID 就结束了，embedding 的文章从 nn.Embedding 开始讲。中间差一个连接：BPE 做完以后，我们手里只有一串整数 ID。这些 ID 本身没有数值意义，神经网络到底如何从 token ID 得到真正可以参与计算的 representation？

把整条链路画出来：

flowchart LR
    A["Raw Text"] --> B["Tokenizer / BPE"]
    B --> C["Token"]
    C --> D["Token ID"]
    D --> E["One-hot"]
    E --> F["Embedding Lookup"]
    F --> G["Dense Vector"]
    G --> H["RNN / Transformer"]

    style D stroke-dasharray: 5 5
    style E stroke-dasharray: 5 5

这篇笔记把中间这段补完。

BPE 之后，我们手里有什么
#

先回顾一下 BPE 做了什么（算法细节见上一篇，这里不重复）。

BPE 解决的核心问题是：如何把连续文本划分成有限 vocabulary 中的离散 symbol。

比如：

1"unbelievable" → ["un", "believ", "able"]

这些 subword token 再通过 vocabulary 映射为整数 ID：

1["un", "believ", "able"] → [381, 9281, 521]

这里有三个不同的东西，容易混在一起：

  • token：文本被切分后的基本单位（“un”、“believ”、“able”）
  • vocabulary：所有合法 token 的集合，大小为 \(V\)
  • token ID：每个 token 在 vocabulary 中的编号

Token ID 只是一个编号。381 不比 12 “更大”，也不意味着任何语义距离。

换一个更明显的例子。假设：

1cat → 2
2dog → 3
3car → 100

如果把这些 ID 直接当数值输入网络：

  • 3 - 2 = 1 并不说明 dog 比 cat “更接近”
  • 100 也不意味着 car 比 dog “大很多”

这些数字只有 identity，没有 magnitude。就像数据库的 primary key 或者 enum 的 ID——user_id = 42 和 user_id = 43 之间不存在"距离为 1"的关系。

Token ID 是 categorical variable，不是 numerical variable。直接拿来做算术运算没有意义。

那怎么办？

One-hot：消除伪数值关系
#

一个自然的想法：不用单个数字表示 token，改用向量，并且让不同 token 之间的表示不存在任何隐含的数值偏序。

One-hot encoding 就是这个思路。假设 vocabulary size \(V = 5\)，给每个 token 分配一个 \(V\) 维的标准基向量：

$$e_i \in \mathbb{R}^V$$

具体展开：

1a    → e_0 = [1, 0, 0, 0, 0]
2cat  → e_1 = [0, 1, 0, 0, 0]
3dog  → e_2 = [0, 0, 1, 0, 0]
4car  → e_3 = [0, 0, 0, 1, 0]
5the  → e_4 = [0, 0, 0, 0, 1]

每个向量只有一个位置是 1，其余全是 0。

这种表示有一个重要的数学性质：所有不同 token 的 one-hot 向量两两正交。

$$e_{\text{cat}} \cdot e_{\text{dog}} = 0$$$$e_{\text{cat}} \cdot e_{\text{car}} = 0$$

在 one-hot 空间里，cat 和 dog 的距离与 cat 和 car 的距离完全一样——不存在"有的更近，有的更远"。

这恰恰是我们想要的。One-hot 只表达"这是哪个离散 symbol"，不引入任何语义距离的假设。

但也要注意：one-hot 没有学到任何东西。它只是对 categorical variable 的一种数学编码方式。有些材料（包括 6.S191 的 slide）会出现 “one-hot embedding” 的说法——更严格的叫法是 one-hot encoding。Encoding 是固定的规则，embedding 通常指可学习的向量表示，两者不是一回事。

Embedding Lookup = One-hot × Matrix
#

One-hot vector 的维度等于 vocabulary size。当 \(V = 100000\) 时，每个 token 需要一个 100000 维的稀疏向量——维度高、信息密度极低（只有一个 1）。而且所有 token 两两正交，没有任何结构可以利用。

自然的下一步：把这个高维稀疏向量映射到一个低维稠密向量。

定义 embedding matrix：

$$E \in \mathbb{R}^{V \times d}$$

\(V\) 行 \(d\) 列。\(V\) 是 vocabulary size，\(d\) 是 embedding dimension。每一行对应一个 token 的 embedding vector。

Token \(i\) 的 one-hot vector \(e_i \in \mathbb{R}^V\)（行向量）乘以 embedding matrix：

$$e_i^T E = E[i]$$

也就是说：one-hot vector 左乘 embedding matrix，结果就是取出矩阵的第 \(i\) 行。

用一个 \(V=5, d=3\) 的具体例子看。假设 token ID = 1（对应 cat），one-hot 为 \([0, 1, 0, 0, 0]\)：

$$\begin{bmatrix} 0 & 1 & 0 & 0 & 0 \end{bmatrix} \begin{bmatrix} 0.1 & 0.3 & -0.2 \\ 0.7 & -0.1 & 0.4 \\ 0.2 & 0.5 & 0.8 \\ -0.3 & 0.6 & 0.1 \\ 0.9 & -0.4 & 0.3 \end{bmatrix} = \begin{bmatrix} 0.7 & -0.1 & 0.4 \end{bmatrix}$$

结果就是 \(E\) 的第 1 行（0-indexed）。乘法过程中，其他四行全被 0 消掉了——只有 index 为 1 的位置是 1，所以只有第 1 行"活下来"。

Embedding lookup 在数学上做的事：one-hot vector 乘以 embedding matrix，选出对应的行。

马上转到工程视角。

现实系统不会真的去构造 one-hot vector。\(V = 100000\) 时，\([0, 0, \ldots, 1, \ldots, 0]\) 有 100000 个元素，绝大多数是 0。构造这个向量浪费内存，做一次 dense matrix multiplication 更浪费——99999 次乘零加法全是无用功，只有 1 次有效。

实际实现直接用 token ID 做 index：

1embedding_vector = embedding_matrix[token_id]

PyTorch 中对应的就是：

1embed = nn.Embedding(num_embeddings=vocab_size, embedding_dim=d)
2x = embed(token_ids)  # token_ids: [batch, seq_len] → x: [batch, seq_len, d]

底层操作是 indexed gather——根据 index 从 table 中取出对应的行。没有矩阵乘法，没有浮点运算，只有内存读取。

数学上是 one-hot × matrix，工程上是 indexed gather / lookup。同一件事的两种视角。

这也解释了 embedding lookup 的性能特征。之前讨论 Arithmetic Intensity 时提到过：embedding lookup 是典型的 memory-bound 操作——random access，几乎没有计算，和 dense GEMM（数据复用率高、compute-bound）完全不同。

Embedding 增加了什么
#

经过 embedding lookup，一个 token 的表示从整数 ID 变成了 dense vector：

1cat → 314 → E[314] → [0.13, -0.42, 0.71, ...]

比 one-hot 多了什么？

两样东西。

第一，维度从 \(V\) 降到了 \(d\)。\(V\) 通常是几万到十几万，\(d\) 通常是几百到几千。表示变得紧凑。

第二，embedding matrix \(E\) 是可学习参数。它和网络的其他权重一起，通过训练被更新。每个 token 对应的 \(d\) 维向量不是固定的，而是模型在训练过程中逐渐调整出来的。

一个常见但需要谨慎对待的说法：

“embedding 把语义相似的词映射到向量空间中相近的位置。”

这句话作为直觉没问题，但不应该当成严格定义。实际训练出来的 embedding 结构取决于 training objective、数据分布、模型架构等很多因素。“相似的词更近"只是可能出现的现象之一。

还有一点容易忽略。Token embedding 是 static 的——同一个 token 不管出现在什么 context 中，从 embedding matrix 取出的向量都一样。

但同一个 token 在不同 context 中的含义可能完全不同。比如 “bank” 在 “river bank” 和 “bank account” 中：初始的 token embedding 是同一个向量。真正区分这两种含义的，是经过 Transformer 多层处理后得到的 hidden state。

也就是说：

  • token embedding = static initial representation
  • Transformer hidden state = contextual representation

Token embedding 只是起点。上下文相关的表示是后续网络的事——这件事属于 representation learning 的范畴，留到后面再展开。

Vocabulary Size 如何改变模型的计算
#

BPE 那篇已经从 tokenizer 的角度讨论过 vocabulary size 的 trade-off——更大的 vocabulary 压缩效率更高但 vocabulary 本身有成本。这里从 embedding 和模型的角度再补充一层。

Tokenizer 的 vocabulary size \(V\) 不只影响 tokenizer 本身——它一路影响后续模型的参数量和计算量。

较大的 \(V\)：

  • 一个 token 覆盖更多字符 / subword，同一段文本产生的 sequence 更短
  • 但 embedding table 更大：参数量 \(V \times d\)
  • Output projection（LM head）更大：参数量 \(d \times V\)
  • Softmax / logits 的维度更大

较小的 \(V\)：

  • Embedding table 和 LM head 更小
  • 但同样的文本被切成更多 token，sequence length \(T\) 增加
  • Transformer 需要处理更长的序列——attention 的计算量、KV cache 的大小、inference latency 都受影响

这不是"大好还是小好"的问题。Vocabulary size 是一个 cross-layer trade-off，贯穿从 tokenizer 到 Transformer 的整条计算链路：

1vocabulary size V
2    ↕
3embedding / LM head 参数量（V × d）
4    ↕
5sequence length T
6    ↕
7Transformer compute / memory / KV cache

具体数量级：假设 \(V = 50000\)，\(d = 4096\)，那么 embedding table 有约 200M 参数（BF16 下约 400 MB）。如果 \(V\) 翻倍到 100000，这一项多出 200M 参数和约 400 MB 显存。对于 1B 参数量级的模型，这已经是总参数量的相当一部分。

输入和输出的对称性
#

把语言模型的整个 forward pass 拉通来看，embedding 参与了两端。

输入侧：

$$\text{token ID} \xrightarrow{E \in \mathbb{R}^{V \times d}} \text{dense vector} \xrightarrow{\text{Transformer}} \text{hidden state}$$

输出侧（下一篇详细讲了这一段）：

$$\text{hidden state} \xrightarrow{W_{\text{out}} \in \mathbb{R}^{d \times V}} \text{logits} \xrightarrow{\text{softmax}} P(\text{next token})$$

从维度上看：

flowchart LR
    V1["V"] -->|"Input Embedding
E ∈ ℝ^(V×d)"| D1["d"]
    D1 -->|"Transformer"| D2["d"]
    D2 -->|"Output Projection
W ∈ ℝ^(d×V)"| V2["V"]

输入侧把离散的 token 映射到 \(d\) 维连续空间。输出侧把 \(d\) 维 hidden state 映射回 \(V\) 维的词表空间，再通过 softmax 变成概率分布。整个结构是对称的。

因为这个对称性，很多语言模型会做 weight tying——让 input embedding \(E\) 和 output projection \(W_{\text{out}}\) 共享参数（\(W_{\text{out}} = E^T\)）。参数量减少 \(d \times V\)，同时输入和输出共享同一个 token representation space。GPT-2 就是这么做的。

结论
#

  1. Tokenizer / BPE 决定了模型世界中的基本离散单位。它不是在"把文本变成数字”，而是在定义模型看到的 symbol 集合。

  2. Token ID 只是 identity——categorical variable，不是 numerical variable。直接拿来做算术运算没有意义。

  3. One-hot encoding 为离散 symbol 提供了数学表示——标准基向量，两两正交，不引入伪数值关系。Embedding matrix 把它映射到低维连续空间。

  4. Embedding lookup 在数学上等价于 one-hot vector 乘以 embedding matrix（\(e_i^T E = E[i]\)）。工程实现是 indexed gather，不构造 one-hot，也不做 dense matmul。

  5. Vocabulary size 不是 tokenizer 自己的事。它一路影响 embedding 参数量、output projection、sequence length，以及 Transformer 的 compute 和 memory。

Embedding 让 token 进入了连续空间。但这些向量为什么最终能承载语义？同一个 token 在不同 context 中的 representation 为什么不同？Next-token prediction 为什么能逼着模型学出有用的 representation？这些问题留到 representation learning 相关文章继续。

参考资料
#

Related