4.1 Tokens, Vocabulary, and Embeddings

Transformer 的输入不是“自然语言本身”，而是离散 token 序列。LLM 的许多性质，包括上下文长度、罕见词处理、代码能力、多语言能力和推理成本，都首先受 tokenizer 与 embedding 设计影响。

Running Example: Why Tokenization Matters

考虑一句中文夹英文和代码的 prompt：

用 PyTorch 写一个 self_attention 函数

模型真正看到的不是这句话，而是一串 token ids。一个 tokenizer 可能把它切成：

用 | Py | Torch | 写 | 一个 | self | _ | attention | 函数

另一个 tokenizer 可能切成：

用 | PyTorch | 写 | 一个 | self_attention | 函数

第二种切法让 PyTorch 和 self_attention 作为更完整的语义单元出现，sequence length 更短；第一种切法更通用，但模型要靠上下文重新组合碎片。于是 tokenizer 并不是预处理小工具，而是在定义模型的离散世界。

Tokenization

NoteDefinition: Tokenization

Tokenization maps raw text \(s\) into a discrete sequence \[ x_{1:T}=\operatorname{Tok}(s), \qquad x_t\in\{1,\ldots,V\}, \] where \(V\) is the vocabulary size.

常见 tokenizer 包括 BPE、WordPiece、Unigram LM 和 byte-level BPE。它们都在做一个折中：词表越大，序列越短，但 embedding/head 参数越大；词表越小，序列越长，attention 计算更贵。

Choice Benefit Cost
character open vocabulary very long sequences
word short sequences OOV and huge vocab
subword robust compromise segmentation artifacts
byte-level no OOV less semantic units

Byte Pair Encoding Step by Step

BPE 的思想是从小单位开始，不断合并语料里最常见的相邻 pair。假设语料是：

low lower lowest

初始可以看成 character sequence：

l o w
l o w e r
l o w e s t

最常见 pair 是 l o，合并成 lo：

lo w
lo w e r
lo w e s t

下一步常见 pair 是 lo w，合并成 low：

low
low e r
low e s t

继续合并可能得到 er、est、lower、lowest。训练结束后，tokenizer 保存 merge table。推理时遇到新词 lowland，也可以切成已知 subwords：

low | land

这就是 subword tokenizer 的优势：它既能表示常见词，也能把未知词拆成可处理的片段。

Vocabulary Size as a System Parameter

Vocabulary size \(V\) 不是纯 NLP 细节，而是模型结构参数。它直接决定：

  1. input embedding 参数量；
  2. output LM head 参数量；
  3. logits tensor 大小；
  4. sampling / top-\(p\) / top-\(k\) 后处理成本；
  5. rare token 的训练频率。

若 hidden size 为 \(d\)，embedding table 是：

\[ E\in\mathbb{R}^{V\times d}. \]

仅输入 embedding 就有 \(Vd\) 个参数。若 output head 不 tied，还要再来一个：

\[ W_U\in\mathbb{R}^{V\times d}. \]

所以 vocabulary 变大时，参数增加：

\[ \Delta N_{\text{embed+head}} = \begin{cases} \Delta V\cdot d, & \text{tied embeddings},\\ 2\Delta V\cdot d, & \text{untied embeddings}. \end{cases} \]

例如 \(d=4096\)，词表增加 \(10000\) 个 token，tied 情况下也增加约 \(41\)M 参数；BF16 权重约 \(82\) MB，AdamW 训练状态按 12 bytes/param 则约 \(492\) MB。

WarningPitfall: Larger Vocabulary Is Not Free

A larger vocabulary may shorten sequences, but it increases embedding/head parameters and logits memory. Tokenizer design is a compute-memory trade-off, not only a linguistic choice.

Vocabulary Resize as a Checkpoint Operation

给 tokenizer 增加 token 后，模型结构也变了。设旧词表大小为 \(V\)，新增 \(\Delta V\) 个 special tokens 或 domain tokens，新 embedding 变成：

\[ E'\in\mathbb{R}^{(V+\Delta V)\times d}. \]

旧 checkpoint 只包含前 \(V\) 行：

\[ E'_{0:V}=E_{\text{old}}, \qquad E'_{V:V+\Delta V}\sim \text{new init}. \]

HuggingFace 风格代码通常是：

num_added = tokenizer.add_special_tokens(
    {
        "pad_token": "<|pad|>",
        "additional_special_tokens": ["<|tool|>", "</think>"],
    }
)

if num_added > 0:
    model.resize_token_embeddings(len(tokenizer))

这个操作至少有三层含义：

  1. input embedding rows 增加；
  2. LM head rows 增加，若 tied 则共享同一张表；
  3. optimizer state 对新增行不存在，必须重新初始化或让 optimizer 在下一步创建 state。
NoteDefinition: Vocabulary-Model Contract

The vocabulary-model contract requires tokenizer ids, embedding rows, LM-head rows, special-token ids, and checkpoint tensors to agree exactly.

如果只是改 tokenizer 而不 resize，新增 token id 会越界；如果 resize 但没有训练新增行，模型可以读写这些 token，但它们的 embedding 仍接近随机初始化。对 pad token 这类不应预测的 token，随机初始化通常没问题；对 <|tool|>、</think> 这类协议 token，则需要 SFT/post-training 数据给它们稳定语义。

WarningPitfall: Resizing Does Not Teach New Tokens

resize_token_embeddings only changes tensor shapes and initializes new rows. It does not teach the model what the new tokens mean.

Fertility and Effective Context

同一段原始文本被 tokenizer 切成的 token 数不同。定义：

\[ \operatorname{fertility}(s) = \frac{\#\text{tokens}}{\#\text{characters or words}}. \]

若模型 context length 是 \(L\) tokens，那么原始文本可容纳长度近似为：

\[ \text{raw length} \approx \frac{L}{\operatorname{fertility}}. \]

因此 tokenizer 对多语言公平性很关键。一个中文 fertility 很高的 tokenizer，会让中文在同样 \(L\) 下获得更短的有效上下文，也会让中文训练样本贡献更多 token-level loss positions。

Embeddings

token id 通过 embedding table 映射到向量：

\[ e_t = E[x_t], \qquad E\in\mathbb{R}^{V\times d}. \]

输出 head 通常是

\[ \operatorname{logits}_t = h_t W_U^\top. \]

很多语言模型使用 tied embedding：

\[ W_U = E. \]

它减少参数量，也表达了一个直觉：输入 token 的语义空间和输出 token 的分类空间应该共享结构。

NoteDefinition: Embedding Lookup

Embedding lookup maps a discrete token id \(x_t\) to the corresponding row of an embedding matrix: \[ e_t=E_{x_t}. \] It is equivalent to multiplying a one-hot vector by \(E\).

若 \(o_t\in\{0,1\}^{V}\) 是 token \(x_t\) 的 one-hot vector，则：

\[ e_t=o_t^\top E. \]

这个写法说明 embedding lookup 是可微的，但不是对 token id 可微。梯度会更新 \(E\) 的某些行，而不会更新离散 id 本身。

nn.Embedding Semantics

PyTorch 的 nn.Embedding(V, d) 本质上就是一张 lookup table：

import torch
from torch import nn

embed = nn.Embedding(num_embeddings=V, embedding_dim=d, padding_idx=pad_id)
input_ids = torch.tensor([[1, 5, pad_id]])
x = embed(input_ids)  # [B, T, d]

padding_idx 有一个特殊语义：对应 row 在 backward 中不会累积梯度，并且默认初始化为 0。它适合 <pad>，但不适合 <eos>。很多 decoder-only 模型为了方便把 pad_token_id = eos_token_id，这时就不能把 padding_idx 也设成 EOS，否则 EOS embedding 可能学不到或被意外固定。

WarningPitfall: PAD-as-EOS Changes Embedding Semantics

Using EOS as PAD can be practical for batching, but setting padding_idx=eos_id freezes or special-cases the EOS row. Decide separately whether PAD and EOS share an id and whether that row should receive gradients.

Embedding 还有 sparse=True 选项，让 backward 产生 sparse gradient：

embed = nn.Embedding(V, d, sparse=True)

这对很大词表的推荐/检索模型有用，但不是所有 optimizer 都支持 sparse gradients；LLM 里还常有 tied LM head，使 output side 每步涉及大量 vocab rows，因此 sparse input embedding 并不能自动解决 LM head 的 dense gradient 问题。

Setting Gradient shape Optimizer constraint
sparse=False dense [V,d], most rows zero-like AdamW/SGD common
sparse=True sparse rows for input lookup limited optimizer support
tied LM head output gradient dense in vocab sparse input trick no longer enough

Sparse Gradient Structure

对单个输入 token \(x_t\)，如果后续 loss 对 embedding 输出的梯度是：

\[ \frac{\partial \mathcal{L}}{\partial e_t}=g_t, \]

则 embedding table 的梯度只有第 \(x_t\) 行收到贡献：

\[ \frac{\partial \mathcal{L}}{\partial E_i} = \begin{cases} g_t, & i=x_t,\\ 0, & i\neq x_t. \end{cases} \]

若一个 token 在 batch 中出现多次，对应行的梯度会累加。高频 token 的 embedding 被频繁更新，低频 token 更新很少。这是 subword tokenizer 比 word tokenizer 稳定的重要原因：罕见词可以拆成更常见的片段，从而共享统计强度。

WarningPitfall: Token IDs Are Indices, Not Numeric Features

The model does not know that token id 100 and token id 101 are numerically close. Token ids are only row indices into the embedding table.

Embedding Frequency and Optimizer State

如果 token \(i\) 在数据中的出现概率是 \(p_i\)，一个 batch 有 \(N=BT\) 个 token，那么它作为输入被直接更新的期望次数约为：

\[ \mathbb{E}[\text{count}_i]\approx Np_i. \]

低频 token 的 embedding 不只是更新少；AdamW 这类 optimizer 的一阶/二阶矩估计也更稀疏、更噪声。对 token row \(i\)：

\[ m_i\leftarrow \beta_1 m_i+(1-\beta_1)g_i, \qquad v_i\leftarrow \beta_2 v_i+(1-\beta_2)g_i^2. \]

若很多 step 中 \(g_i=0\)，这个 row 的 optimizer state 主要在衰减；偶尔出现时又会被单个 batch 强烈影响。这解释了为什么 rare token embedding 更容易质量差，也解释了为什么 tokenizer 训练要避免大量“看起来漂亮但数据里很少出现”的 token。

NoteDefinition: Token Frequency Imbalance

Token frequency imbalance means embedding rows receive very different numbers of direct updates, causing common tokens to have much better-estimated representations than rare tokens.

实践中可以做两个统计：

counts = torch.bincount(input_ids.flatten(), minlength=len(tokenizer))
rare = (counts == 0).sum()
top = counts.topk(20)

这在继续预训练、领域 tokenizer 扩展、代码/数学数据混合时特别有用。不是所有 token 都需要同样频率，但极端长尾会把词表参数浪费在几乎不训练的 rows 上。

Shape Walkthrough

设 batch size 为 \(B=2\)，sequence length 为 \(T=4\)，vocab size 为 \(V=50000\)，hidden size 为 \(d=768\)。

token id tensor:

\[ X_{\text{id}}\in\mathbb{N}^{2\times4}. \]

embedding lookup 后：

\[ X=E[X_{\text{id}}]\in\mathbb{R}^{2\times4\times768}. \]

Transformer 输出 hidden states：

\[ H\in\mathbb{R}^{2\times4\times768}. \]

输出 logits：

\[ Z=HW_U^\top\in\mathbb{R}^{2\times4\times50000}. \]

所以语言模型每个位置都在做一个 \(50000\) 类分类。大模型推理时 logits 后处理、采样和 vocabulary projection 不是免费操作。

Logits Memory

logits tensor 的形状是：

\[ Z\in\mathbb{R}^{B\times T\times V}. \]

如果 \(B=8,T=2048,V=150000\)，BF16 logits 仅存储就需要：

\[ 8\cdot2048\cdot150000\cdot2 \approx 4.9\text{ GB}. \]

训练时通常不会长期保存完整 logits；loss kernel 会尽量 fused 或及时释放。但这个数量级解释了为什么大词表模型的 output projection 和 CE loss 是系统热点。

NoteDefinition: Unembedding

The unembedding or LM head maps hidden states back to vocabulary logits, usually through a matrix multiplication with shape \(d\rightarrow V\).

在 tied embedding 情况下：

\[ z_t = h_t E^\top. \]

这意味着每个输出 token logit 是 hidden state 和该 token embedding 的 dot product：

\[ z_{t,k}=h_t^\top E_k. \]

于是输出概率不仅取决于 hidden state，也取决于 embedding 空间的几何结构。

Tied vs Untied LM Head

tied embeddings:

logits = hidden @ embed.weight.T

untied LM head:

lm_head = nn.Linear(d, V, bias=False)
logits = lm_head(hidden)

二者的差异：

Design Parameters Gradient coupling Common use
tied \(Vd\) input/output share rows many decoder LMs
untied \(2Vd\) separate input and output geometry some large models / ablations
output bias \(+V\) shifts token prior sometimes omitted

如果使用 output bias \(b\in\mathbb{R}^V\)：

\[ z_k=h^\top E_k+b_k. \]

\(b_k\) 可以学习 unigram frequency prior：常见 token 即使 hidden state 没有强偏向，也可以有较高 base logit。很多现代 LLM 省略 LM-head bias，让 token prior 主要由 hidden state 和 embedding geometry 表达。

WarningPitfall: Tied Setting Is a Checkpoint Contract

Changing tied vs untied LM head changes parameter names, tensor sharing, optimizer state, and sometimes model quality. It is not a cosmetic refactor.

Discrete Likelihood

语言模型训练的核心 loss 是 categorical negative log-likelihood：

\[ \mathcal{L}_{\text{LM}} = -\sum_{t=1}^{T} \log p_\theta(x_t\mid x_{<t}). \]

softmax 定义为

\[ p_\theta(x_t=k\mid x_{<t}) = \frac{\exp(z_{t,k})}{\sum_{j=1}^{V}\exp(z_{t,j})}. \]

这说明 LLM 的输出层本质上是在一个巨大离散词表上做分类。大词表会让 logits 和 sampling 成为重要系统成本。

ImportantTheorem: Softmax Cross-Entropy Gradient

For logits \(z\in\mathbb{R}^{V}\), target one-hot vector \(y\), and \[ L=-\sum_k y_k\log \operatorname{softmax}(z)_k, \] the gradient is \[ \frac{\partial L}{\partial z_k} = p_k-y_k, \qquad p=\operatorname{softmax}(z). \]

写成

\[ L=-z_{y}+\log\sum_j e^{z_j}. \]

对 \(z_k\) 求导：

\[ \frac{\partial L}{\partial z_k} = -\mathbf{1}[k=y] + \frac{e^{z_k}}{\sum_j e^{z_j}} = p_k-y_k. \]

这说明 cross entropy 的梯度形式非常直接：提高正确 token logit，降低模型当前过高估计的其他 token。

Output Embedding Gradient

若使用 tied LM head \(z=hE^\top\)，单位置 loss 的 logits 梯度为：

\[ \frac{\partial L}{\partial z}=p-y. \]

对 hidden state：

\[ \frac{\partial L}{\partial h} = (p-y)E. \]

对第 \(k\) 个 output embedding row：

\[ \frac{\partial L}{\partial E_k} = (p_k-y_k)h. \]

这说明正确 token row 被推向 \(h\)，而模型赋予概率的其他 token rows 会被按 \(p_k\) 推离 \(h\)。但 tied embedding 还要同时承担输入表示功能，所以它既是“词的输入语义表”，又是“输出分类器权重表”。

WarningPitfall: Weight Tying Couples Two Roles

With tied embeddings, changing the embedding matrix changes both how tokens enter the model and how hidden states are scored as output tokens.

Rare Tokens and Gradient Starvation

如果 token \(k\) 很少作为输入出现，它的 input embedding 很少被直接更新；但在 softmax 输出里，它每一步都会得到一个小的 \(p_k h\) 梯度。对大词表来说，大多数 rare tokens 的 \(p_k\) 很小，所以它们学习很慢。

这也是为什么 tokenizer 不能随便塞入大量很少出现的 whole-word tokens。它们会占据词表和 embedding 参数，却缺少足够训练信号。

Sampling from Logits

训练时我们最小化 cross entropy；生成时则要从 logits 变成 token。常见采样策略：

Strategy Formula / rule Effect
greedy \(\arg\max_k z_k\) deterministic, can be repetitive
temperature \(z_k/\tau\) lower \(\tau\) sharper, higher \(\tau\) more random
top-k keep largest \(k\) logits removes long tail
nucleus keep smallest set with prob mass \(p\) adaptive diversity

采样策略不改变模型参数，却会显著改变输出行为。这也是为什么“同一个 LLM”在不同 temperature、top-p 下像不同的人。

Efficient Sampling

若完整 softmax：

\[ p_k=\frac{e^{z_k}}{\sum_{j=1}^{V}e^{z_j}}, \]

需要扫过全部 \(V\) 个 logits。Top-\(k\) sampling 至少需要找出最大 \(k\) 个元素；top-\(p\) sampling 往往需要排序或近似排序概率质量。对 \(V=150000\) 的模型，sampling 是 CPU/GPU runtime 里真实存在的后处理成本。

在训练中，cross entropy 也需要对所有 vocabulary logits 做 log-sum-exp：

\[ \log\sum_{j=1}^{V}e^{z_j}. \]

因此大词表模型常依赖 fused CE kernel、vocabulary parallelism 或分片 LM head 来降低内存和通信压力。

Chunked Cross Entropy

如果一次性 materialize [B,T,V] logits 太大，可以按 token chunk 计算 loss：

import torch.nn.functional as F


def chunked_lm_loss(hidden, labels, lm_head, chunk_size):
    # hidden: [B, T, d], labels: [B, T]
    flat_h = hidden.reshape(-1, hidden.size(-1))
    flat_y = labels.reshape(-1)
    total = flat_h.new_zeros(())
    count = flat_h.new_zeros(())

    for start in range(0, flat_h.size(0), chunk_size):
        h = flat_h[start : start + chunk_size]
        y = flat_y[start : start + chunk_size]
        mask = y.ne(-100)
        if not mask.any():
            continue
        logits = lm_head(h[mask])
        total = total + F.cross_entropy(logits, y[mask], reduction="sum")
        count = count + mask.sum()
    return total / count.clamp_min(1)

这不会减少 LM head 的总 FLOPs，但会减少峰值 logits memory。代价是多次 GEMM/kernel launch，吞吐可能下降。训练系统里常见的 fused CE kernel、vocab-parallel CE、sequence parallel，本质上都是在处理同一个问题：\(V\) 太大时，输出分类层是内存和通信热点。

NoteDefinition: Vocabulary-Parallel LM Head

A vocabulary-parallel LM head shards the vocabulary dimension across devices, computes partial logits or partial log-sum-exp locally, and communicates only the reductions needed for cross entropy or sampling.

分片 CE 的数学关键是 log-sum-exp 可分解。若 vocab 被切成 shards \(\mathcal{V}_r\)：

\[ \log\sum_{k\in\mathcal{V}}e^{z_k} = \log\sum_r \sum_{k\in\mathcal{V}_r}e^{z_k}. \]

数值稳定实现会先 all-reduce 全局最大值 \(m=\max_k z_k\)，再各 shard 计算：

\[ s_r=\sum_{k\in\mathcal{V}_r}e^{z_k-m}, \]

最后 all-reduce \(\sum_r s_r\)。真实 label 所在 shard 负责提供 \(z_y\)。所以 vocab parallel 不是近似 softmax，而是把同一个 exact CE 拆到多设备上。

Vocabulary as Modeling Bias

tokenizer 会改变模型看到的世界。中文、代码、数学、罕见 Unicode、空格缩进都会受到 segmentation 影响。一个糟糕 tokenizer 可能让模型把常见语义单元切成许多碎片，增加长程依赖负担。

WarningPitfall: Tokens Are Not Words

Never assume one token equals one word. Token boundaries are artifacts of the tokenizer. This matters for context length, evaluation, prompt budgeting, and multilingual behavior.

Special Tokens as Protocol

Special tokens 是 vocabulary 中带协议语义的条目：

Token kind Modeling role
BOS/EOS sequence start/stop
PAD batching placeholder
role tokens system/user/assistant boundary
tool tokens structured tool-call boundary
thinking tags hidden reasoning / final-answer boundary

它们必须和训练数据格式、chat template、loss mask 和 inference parser 对齐。若 tokenizer 中的 <|assistant|> 被拆成普通文本片段，模型就不会学到稳定的角色边界。

NoteDefinition: Protocol Token

A protocol token is a special vocabulary item whose meaning comes from the data format or runtime protocol rather than ordinary text frequency.

LLM 的“会聊天”不是模型结构里有一个 chat 模块，而是大量训练样本用 role/protocol tokens 序列化后，模型学会了这些 token 条件下的 next-token distribution。

Special Token ID Semantics

special token 不只是字符串，还必须绑定到固定 id：

ids = {
    "bos": tokenizer.bos_token_id,
    "eos": tokenizer.eos_token_id,
    "pad": tokenizer.pad_token_id,
    "assistant": tokenizer.convert_tokens_to_ids("<|assistant|>"),
}

这些 id 进入不同系统层：

ID Used by Failure mode
bos_token_id data formatting, generation start extra/missing BOS changes prefix distribution
eos_token_id loss target, stopping model never stops or stops too early
pad_token_id batching attention mask pad leaks into context
role ids SFT label masks, chat parser model learns wrong speaker boundary
tool/thinking ids runtime parser tool calls or reasoning blocks corrupt

EOS 特别容易混淆。训练中 EOS 是普通 target：模型应该学会在回答结束时预测 EOS。padding 中的 PAD 不应贡献 loss。如果 pad_id=eos_id，那么同一个 id 在不同位置有两种语义：真实 EOS target 要训练，padding EOS 要 mask 掉。

WarningPitfall: Same ID Can Have Different Batch Semantics

When PAD reuses EOS, the id alone is insufficient to decide loss behavior. The label mask and attention mask define whether that occurrence is real EOS or padding.

Stop Strings vs Stop Tokens

推理服务可以按 token id 停止，也可以按字符串停止：

Stop type Example Risk
token id eos_token_id tokenizer/version mismatch
special token </think> id must be atomic
string "\nUser:" may span multiple tokens

字符串 stop 需要维护 detokenized suffix，因为 stop pattern 可能跨 token 边界。例如：

tokens: ["</", "think", ">"]
string stop: "</think>"

如果 </think> 不是 atomic special token，服务端 parser 必须在文本层面识别它；如果它是 special token，则 token-level parser 更稳。训练数据、tokenizer 和 runtime parser 应该选同一套协议，而不是各自猜。

From Tokens to Context

token embedding 只提供 identity，位置和上下文来自后续模块：

\[ h_0 = E[x] + P, \]

然后经过多层 self-attention 和 MLP。后续几节会解释 Transformer 如何用 attention 在这些离散 token 之间建立可学习的依赖关系。

Implementation Checklist

训练或替换 tokenizer/embedding 时，至少检查：

  1. len(tokenizer) 是否等于 embedding rows 和 LM head rows；
  2. added special tokens 后是否调用 resize_token_embeddings；
  3. pad token 是否在 labels 中 mask 成 -100；
  4. BOS/EOS 是否和训练数据格式一致；
  5. chat template 是否产生训练时见过的 role tokens；
  6. tied/untied embedding 设置是否和 checkpoint 一致；
  7. logits dtype、fused CE、vocab parallel 是否能承受当前 \(V\)；
  8. 多语言 fertility 是否让某些语言有效上下文过短；
  9. decode/encode round-trip 是否保留空格、换行、Unicode 和代码缩进。
  10. padding_idx 是否意外冻结 EOS 或其他真实 token；
  11. 新增 token rows 是否经过训练，而不只是 resize；
  12. chunked/vocab-parallel CE 的 denominator 是否和普通 CE 一致；
  13. stop token/string parser 是否和 tokenizer atomicity 一致。

这一页的核心 mental model 是：

\[ \text{raw text} \rightarrow \text{token ids} \rightarrow \text{embedding rows} \rightarrow \text{contextual hidden states} \rightarrow \text{unembedding logits over vocabulary}. \]

后面所有 attention、GPT-2、Qwen3、post-training 和 serving 细节，都建立在这个离散接口上。