3. 蛋白序列分词与词表构建

蛋白序列语言模型的第一步不是搭 Transformer，而是把一条由 20 种标准氨基酸组成的字符串，转换成模型可以处理的整数 ID 序列。这一步叫分词（tokenization），它决定了模型能看到多大的“语义颗粒”，也直接决定了词表规模、序列长度、显存占用和训练稳定性。与自然语言不同，蛋白序列没有空格、没有标点，字母表天然只有约 20～25 个符号（20 种标准氨基酸加上 X、B、Z、U、O 等模糊或特殊残基），因此分词策略的选择空间和权衡点都与文本模型有明显差异。

四种主流分词策略

氨基酸级分词是最直接的方案：每个氨基酸残基映射为一个 token，词表大小约 25～30。它的优点是词表极小、无 OOV（未登录词）、每条序列的 token 数等于残基数，位置信息与真实残基一一对应，便于做残基级的下游任务（如突变效应预测、二级结构标注）。缺点是序列过长：一条 500 残基的蛋白会产生 500 个 token，而 Transformer 的注意力开销随长度平方增长，长序列训练代价高；同时单个氨基酸携带的语义信息有限，模型需要更多层才能学到 motif 级别的模式。

k-mer 分词把连续的 k 个残基当作一个 token。例如 k=3 时，“ACDEF”会被切成 “ACD”“CDE”“DEF”。理论上 20 种氨基酸的 3-mer 有 8000 种组合，实际词表通常只保留训练集中出现频次较高的 k-mer，其余用 [UNK] 兜底。k-mer 缩短了序列长度（约为原来的 1/k），并让每个 token 自带局部序列模式，对识别短 motif 有帮助。代价是词表膨胀、低频 k-mer 学不充分，且切分边界固定，跨边界模式会被割裂。

BPE（Byte Pair Encoding）从单氨基酸出发，反复统计语料中最高频的相邻 token 对并合并，直到词表达到预设大小。它数据驱动、自适应，能自动发现如 “GXGXXG”（核苷酸结合 motif）这类高频片段，兼顾了常见模式的整体性和低频区域的细粒度。WordPiece 与 BPE 思路相近，但合并准则不是频次而是似然增益，倾向于选择能最大提升语料概率的对。两者在蛋白序列上都能工作，BPE 实现更简单、更常用。

Unigram / SentencePiece则从大词表出发，按概率剪枝，适合多语言或需要概率化切分的场景，在蛋白领域使用相对较少。

词表规模的权衡

词表大小不是一个可以随意设定的超参数，它与序列长度、嵌入矩阵参数量、低频 token 的学习质量三者相互牵制。设词表为 V、平均序列 token 数为 L、隐藏维度为 d，则嵌入层参数量约为 V×d，注意力计算量约为 O(L²d)。增大 V 会缩短 L，从而降低注意力开销，但嵌入层变大，且每个 token 在语料中的出现次数减少，低频 token 的向量难以训练充分。经验上，氨基酸级词表（V≈25）适合残基级任务和短序列；BPE 词表在 1000～8000 之间是蛋白语言模型的常见区间，例如在 UniRef50 上训练时，V=2000～4000 往往能在序列长度和 token 质量之间取得较好平衡。若语料只有几十万条序列，词表超过 8000 后大量 token 出现次数低于 100，通常得不偿失。

特殊 token 设计

蛋白序列模型至少需要以下几类特殊 token，它们不参与氨基酸语义，但承担结构控制功能：

• [CLS]：序列级表示，放在序列开头，其最终隐藏状态用于分类、回归等下游任务。
• [SEP]：分隔符，用于拼接两条序列（如抗原-抗体对）或标记序列结束。
• [PAD]：批内对齐填充，必须配合 attention mask 使用，否则会污染注意力。
• [MASK]：掩码语言建模（MLM）的训练目标，随机替换被预测的残基。
• [UNK]：未登录 token 兜底，k-mer 和 BPE 都需要。
• [BOS]/[EOS]：自回归生成任务中标记序列起止。

一个容易被忽视的细节是：蛋白序列中 X（未知残基）、B（Asx）、Z（Glx）、U（硒代半胱氨酸）、O（吡咯赖氨酸）等非标准符号必须显式纳入词表，否则会被错误地映射为 [UNK]，丢失信息。建议在构建词表前先对语料做字符统计，确认实际出现的符号集合。

工程实现：从 FASTA 到词表

下面给出一个可运行的 BPE 词表构建流程，使用 HuggingFace tokenizers 库。假设已有 uniref50_sample.fasta，约 50 万条序列。

from tokenizers import Tokenizer, models, trainers, pre_tokenizers

# 1. 初始化 BPE 模型，unk_token 必须显式指定
tokenizer = Tokenizer(models.BPE(unk_token="[UNK]"))

# 2. 蛋白序列无空格，用 Whitespace 会失效，改用字符级预分词
tokenizer.pre_tokenizer = pre_tokenizers.Split(
    pattern="", behavior="isolated"
)

# 3. 训练器：词表大小、特殊 token、最小频率
trainer = trainers.BpeTrainer(
    vocab_size=4000,
    min_frequency=50,
    special_tokens=["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]"],
    show_progress=True,
)

# 4. 从 FASTA 读取序列并训练
def read_fasta(path):
    seqs = []
    with open(path) as f:
        for line in f:
            if not line.startswith(">"):
                seqs.append(line.strip())
    return seqs

tokenizer.train_from_iterator(read_fasta("uniref50_sample.fasta"), trainer)

# 5. 保存并验证
tokenizer.save("protein_bpe_4000.json")
enc = tokenizer.encode("MKTAGXGXXGACD")
print(enc.tokens)
# 示例输出: ['[CLS]', 'MK', 'TA', 'GX', 'GX', 'XG', 'AC', 'D', '[SEP]']
print("token 数:", len(enc.ids))

训练完成后需要做三项检查。第一，覆盖率：统计验证集中 [UNK] 的比例，若超过 1%，说明 min_frequency 设得过高或词表过小，应下调 min_frequency 或增大 vocab_size。第二，压缩率：计算平均每条序列的 token 数与原残基数的比值，BPE 通常能压到 0.5～0.7，若接近 1.0 说明合并几乎没生效，需检查预分词是否正确。第三，特殊 token 位置：确认 [CLS] 和 [SEP] 被正确添加，且 [PAD] 的 ID 为 0（便于后续 mask 构造）。

最后给出一个词表规模选择的经验参考：残基级任务（突变预测、二级结构）用氨基酸级或小 BPE（V≤1000）；通用蛋白语言模型用 BPE V=2000～8000；若要做跨物种、跨家族的生成任务，可适当增大到 16000，但需保证语料规模在千万条以上，否则低频 token 会成为训练噪声。分词策略没有绝对最优，关键是让词表规模、序列长度和语料规模三者匹配。