词汇量大小选择的权衡

为您的子词分词器选择词汇量大小，通常表示为 $|V|$，是一个主要的超参数选择，直接影响模型表现、内存占用和计算成本。与词汇表可能自然增长的简单分词方法不同，BPE 或 WordPiece 等子词算法要求您预先定义一个目标词汇量大小。这一决定需要平衡多个相互影响的因素。

对序列长度和信息粒度的影响

较小的词汇量大小会使分词器更频繁地将单词拆分为更小的子词单元。

• 优点： 它能通过常见子词组合来妥善处理稀有词和词汇表外词（OOV）。这能提高对拼写错误和新词的鲁棒性。
• 缺点： 这会导致相同输入文本的标记序列更长。由于Transformer中注意力机制的计算成本与序列长度呈平方关系（$O(L^2)$），更长的序列会显著增加训练和推理时间以及内存使用。

反之，更大的词汇量大小允许分词器将更多常见词或频繁出现的子词序列表示为单个标记。

• 优点： 这会产生更短的标记序列，从而减轻注意力机制的计算负担。如果常见词保持完整，单个标记嵌入中可能会包含更多语义信息。
• 缺点： 如果词汇表对于目标领域不够大或不够多样，会增加出现OOV词的风险。它还会增加模型嵌入层的大小。

以“tokenization”（分词）一词为例。

• 对于较小的词汇量（例如，$|V|=10,000$），它可能被分词为 ['tok', 'en', 'ization']（3个标记）。
• 对于更大的词汇量（例如，$|V|=50,000$），如果在训练时包含“tokenization”，它可能变为 ['tokenization']（1个标记）。
• 对于中等词汇量（例如，$|V|=30,000$），可能为 ['token', 'ization']（2个标记）。

更长的序列直接影响训练期间激活所需的内存，特别是注意力得分矩阵（$L \times L$）。

词汇量越大，对于固定语料库，通常会使平均序列长度变短，从而降低注意力计算成本。

对模型大小和内存的影响

词汇量大小直接决定了输入嵌入矩阵和输出投影层（通常是绑定的）的大小。嵌入矩阵的维度为 $|V| \times d_{model}$，其中 $d_{model}$ 是模型的隐藏维度。

$$\text{嵌入矩阵大小} = |V| \times d_{model} \times \text{每参数字节数}$$

更大的 $|V|$ 会导致嵌入矩阵按比例增大。对于$d_{model}$ 为数千的大型模型，将 $|V|$ 从30,000增加到100,000，仅在嵌入层就可能增加数十亿参数和数千兆字节的模型内存占用。

import torch
import torch.nn as nn

# 嵌入层大小示例
d_model = 4096  # 隐藏维度示例

vocab_size_small = 32000
vocab_size_large = 128000

embedding_small = nn.Embedding(vocab_size_small, d_model)
embedding_large = nn.Embedding(vocab_size_large, d_model)

params_small = sum(p.numel() for p in embedding_small.parameters())
params_large = sum(p.numel() for p in embedding_large.parameters())

# 假设使用 float32 (每个参数4字节) 的内存
mem_small_gb = params_small * 4 / (1024**3)
mem_large_gb = params_large * 4 / (1024**3)

print(
    f"Vocab Size: {vocab_size_small}, "
    f"Embedding Params: {params_small:,}, "
    f"Memory: {mem_small_gb:.2f} GB"
)
print(
    f"Vocab Size: {vocab_size_large}, "
    f"Embedding Params: {params_large:,}, "
    f"Memory: {mem_large_gb:.2f} GB"
)

# --- 输出 ---
# Vocab Size: 32000, Embedding Params: 131,072,000, Memory: 0.49 GB
# Vocab Size: 128000, Embedding Params: 524,288,000, Memory: 1.95 GB

此计算仅包含嵌入层。最终输出层将隐藏状态投影回词汇维度以预测下一个标记，其大小通常也相似（$d_{model} \times |V|$），这会进一步放大 $|V|$ 对模型大小的影响。

对计算（Softmax）的影响

在训练和推理期间，模型通常使用应用于输出logits的softmax函数计算整个词汇表的概率。此最终softmax层的计算成本与 $|V|$ 呈线性关系。

$$\text{Softmax 成本} \propto L \times d_{model} \times |V|$$

虽然对于长序列，注意力机制（$O(L^2 \times d_{model})$）通常占据主导地位，但softmax计算（$O(L \times d_{model} \times |V|)$）可能成为一个重要瓶颈，特别是在词汇量非常大或在推理过程中批量大小可能较小且延迟很关键时。更大的 $|V|$ 直接增加了这项成本。

寻求平衡：常见做法

没有单一的“最佳”词汇量大小；它取决于具体的任务、语言、数据集大小、模型架构和可用的计算资源。然而，存在一些常见的观察和做法：

1. 典型范围： 对于许多英语模型，词汇量大小通常在30,000到100,000之间。32,000、50,000和64,000左右的大小很常见，通常是为了计算便利而选择（例如，可被2的幂或硬件向量单元整除）。多语言模型可能需要更大的词汇量（例如，100,000到256,000）以充分覆盖多种字符集和词形。
2. 收益递减： 增加 $|V|$ 在序列长度缩减和困惑度改进方面，在某个点后会产生递减收益。从3万到5万的收益可能很明显，而从10万到15万的收益可能微不足道，但会带来显著的内存成本。
3. 数据特点： 最佳大小受训练数据的多样性和规模影响。更多样的数据（多种语言、领域、代码）可能会从更大的词汇量中获益。
4. 模型大小： 对于较小的模型，由于词汇量增大而导致的模型大小按比例增长更明显。对于非常大的模型（数千亿参数），嵌入层可能只占总参数的一小部分，如果更大的词汇量能带来性能改善，可能使其更易于接受。
5. 计算预算： 最终，硬件限制（GPU内存）和训练时间约束通常对 $|V|$ 设定了实际的上限。

选择较小或较大词汇量大小（$|V|$）所涉及的权衡。

在实际操作中，选择 $|V|$ 通常涉及一些经验性测试，或采用文献中针对类似模型规模和数据集报告的大小。您需要权衡更短序列和潜在更好表示常见词（更大的 $|V|$）带来的益处，与增加的内存使用、较慢的softmax计算以及潜在更差的稀有词处理（较小的 $|V|$）所产生的成本。