跳转至

LEMON-ZEST: Evolution-Informed Tokenization for Efficient Protein Language Modeling

会议: NeurIPS2026(任务清单归属;所读版本为 arXiv 预印本)
arXiv: 2609.37675v1
领域: 计算生物学
关键词: 蛋白质语言模型、进化信息分词、远程同源检索、层级对比学习、随机切分

一句话总结

LEMON-ZEST 将进化保守片段聚类为共享词元,以随机切分和双头编码器同时学习全局检索表征与残基信息,使 200M 参数模型在多个折叠层级检索指标上超过大型基线,但并非所有分类层级或下游任务都最优。

研究背景与动机

蛋白质语言模型通常把每个氨基酸当成一个词元,再通过掩码语言建模学习序列规律。这种做法保留了细粒度信息,却把长序列直接交给注意力机制,也要求模型从海量数据中重新发现哪些局部片段具有稳定的生物学意义。更重要的是,远程同源检索关注的不是字面相似:序列已经明显不同的蛋白质仍可能保留相近的折叠。单纯擅长恢复被遮蔽的残基,并不意味着全局嵌入的余弦相似度已经适合检索。

普通 BPE 可以缩短输入,但其合并依据是相邻字符串频率,不是进化保守性。把统计上常见的片段作为一个词元,未必能将发生替换却仍保留结构意义的片段视为同一单位。反过来,直接使用结构词元虽然引入了更强的信息,却可能在评测配置中要求预测或实验结构作为输入。本文选择在词表构建阶段吸收保守性信息,使模型推理时仍只需要序列;这不同于声称整个训练过程没有结构来源的信息。

因此,本文没有只扩大模型,而是改变模型看到的基本单位:先把保守片段及其相近变体压缩到共享词元,再用对比学习塑造全局检索空间,同时保留恢复残基信息的约束。核心 idea:让进化保守性决定输入词表的等价类,以多粒度切分和双头训练把压缩后的序列表征对齐到远程同源关系。

方法详解

整体框架

ZEST 是词表与分词机制,LEMON 是使用该词表的编码模型。离线阶段从 TEDLH 的 HMM 共识序列提取保守区段并聚类;在线阶段将输入序列切成可变长度词元,并记录每个词元覆盖的残基数。

共享 Transformer 在较短的词元序列上计算上下文信息。其全局分支通过注意力池化输出用于余弦检索的向量,局部分支把词元表示展开到残基位置并预测氨基酸;另有与输入嵌入绑定的词元级 MLM 输出。

下图实线表示数据流,虚线表示仅在训练中使用的监督。推理时的同源检索不需要为查询序列提供结构、MSA 或 CATH 标签,但训练用到了结构分类层级。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400, 'subGraphTitleMargin': {'top': 8, 'bottom': 16}}}%%
flowchart TD
    P["TEDLH 共识序列"] --> V["进化等价词表"]
    V --> T["前缀树随机切分"]
    X["输入序列"] --> T
    subgraph Y["双头表征"]
        E["共享词元编码器"] --> R["注意力池化<br/>全局检索向量"]
        E --> S["残基广播<br/>全局与局部位置"]
        S --> A["20 类残基预测"]
        E --> M["绑定嵌入的<br/>词元预测"]
    end
    T -->|词元与跨度| E
    C["训练:CATH 层级"] -.->|对比监督| R
    G["训练:掩码目标"] -.->|MLM 监督| A
    G -.->|MLM 监督| M
    R --> O["推理:余弦检索"]

关键设计

1. 进化等价词表:把保守片段的近似变体映射为同一输入单位

作者使用 HHsuite 为 TEDLH 的 765,248 个 HMM profile 生成的共识序列,提取长度至少 3 个残基、且不是同聚物的保守区段。这里的 HMM profile 来自多序列比对,区段表达的是跨相关序列保留下来的模式,而不是直接从预训练语料挑出高频字符串。随后通过 MMseqs2 linclust 以 70% 序列一致性聚类,并用兼顾跨 profile 频率与区段长度的分数排序。正文没有给出这一排名分数的完整公式,因此不能把它改写成一个确定的数学目标。

词表保留排名最高的 31,975 个簇,再加入 20 个标准氨基酸回退词元和 5 个特殊词元,共 32,000 个 ID。关键不只是“多个残基合并”:同一簇内的不同区段成员共享同一 ID,使多种字面形式在输入层就被视为一个进化近似单位。这是多对一映射,而不是为每个区段分别学习一个独立词元。它有利于忽略部分序列变异,但也会压掉某些真实残基差别,后续必须有机制约束这种信息损失。

ZEST 的平均词元跨度约为 4 个残基,意味着相同词元窗口通常可覆盖更长的原始序列。原文还报告 4–6 倍压缩,并用 1024 词元窗口容纳约 4,000 残基来说明典型收益;这些是平均或经验描述,不保证任意 4,000 残基序列都能装入窗口。缺少可匹配片段的序列会更多地使用单残基回退,随机切分也可能增加词元数量。

2. 前缀树随机切分:让长片段内部的合法短片段有机会被学习

贪心最长匹配总会优先消耗最长区段,因此某些短区段即使存在于词表,也可能一直被更长的前缀匹配遮住。原文的覆盖分析指出,访问 5 百万条序列后仍有约 139 个词元未被使用。作者把词表组织成前缀树,每个起始位置可以返回一组从长到短的合法匹配,再以给定概率放弃最长匹配,从其余较短匹配中均匀选择一个。剩余后缀随后继续分词,因而同一序列在不同迭代中能够呈现不同边界。

下面将原文叙述整理成选择概率,描述的是当前位置的匹配采样规则,不是论文另外提出的损失函数。

\[ \Pr(k_j)=\begin{cases}1-p_{\text{drop}},&j=1,\\p_{\text{drop}}/(m-1),&j>1,\end{cases}\qquad k_1>k_2>\cdots>k_m,\quad m>1. \]

只有一个合法匹配时没有较短替代项,应保留可用匹配;单残基词元为无法匹配的区段提供回退。这里的 dropout 改变切分边界,不是掩掉词元嵌入,也不是从整个词表均匀抽取词元。它增加短片段的可见性,不能保证有限训练过程中所有词元出现次数相等,甚至不能保证每个词元都出现。

原文称最高 dropout 会回到逐字符切分,但上述规则在存在中间长度匹配时仍可能选中多残基片段。因而更稳妥的结论是:提高 dropout 倾向于更细的切分;“全部变成字符”需要额外的选择条件或实现证据。不能把这一极限描述当成由正文采样公式必然推出的性质。

3. 双头表征:在压缩计算的同时分别约束全局关系与局部内容

共享编码器采用 pre-norm Transformer、RoPE、SwiGLU 和通过 PyTorch scaled_dot_product_attention 使用的 Flash Attention。其主要注意力计算发生在压缩后的词元序列上,而不是先恢复完整残基序列再做同样规模的注意力。输入词元嵌入与词元 MLM 的输出投影绑定,让词元恢复任务直接约束输入表示;这与残基预测分支是不同的输出粒度。

全局表征头用一个学习到的查询向量,以 4 头注意力对全部词元表示做池化。池化后经过 LayerNorm、初始化为恒等映射的对齐投影、瓶颈投影模块以及最终对比嵌入投影,并做 L2 归一化。注意力池化允许模型选择与整体结构关系更相关的区段,而不是默认所有位置对检索同样重要。层级对比学习随后使同源关系反映到向量距离中;单靠池化本身并不会自动产生这种几何结构。

局部分支按照记录的跨度,把每个词元向量广播到其覆盖的所有残基。若只做广播,同一词元内部的位置会收到相同表示,无法区分各个残基。作者因此加入两种位置信号:链上的全局位置,以及词元内部的局部偏移;正文把后者描述为 ALiBi-style。两层瓶颈 MLP 最后输出 20 类氨基酸 logits,使训练梯度能够惩罚粗粒度编码丢掉的局部信息。这不是可逆解码的数学保证,而是通过预测任务促使共享编码器保留残基相关线索。

词表的多对一映射意味着不同片段可能从同一个 ID 开始,因此恢复分支仍要依赖上下文、跨度和位置。它缓解压缩造成的歧义,但并不证明模型能够无损重建任意序列。全局检索头与局部恢复头的职责不同,也解释了为什么远程同源效果好不能直接推出所有残基任务都好。

一个完整示例

考虑一个仅用于说明切分规则的抽象片段:某起始位置有跨度为 6、4、1 的合法匹配。关闭随机切分时会选跨度 6;触发随机切分时,两个较短候选各有一半条件概率被选中,而不是直接强制选跨度 1。若选中跨度 4,剩余两个位置再由后续匹配处理。这里没有给出真实蛋白质序列,也不代表论文的实验样本。

这些词元进入共享编码器时,每个词元只占一个注意力位置。全局头汇总所有上下文向量,得到用于检索的序列嵌入;局部头则把跨度 4 的词元向量广播到 4 个残基位置,再用局部偏移区分这些位置。训练时比较被遮蔽位置的目标残基,同时用 CATH 关系监督全局向量。推理检索时只走序列到全局嵌入的路径,不需要给查询附上结构标签。

同一序列的多次随机切分还可用于测试时增强。以下是原文给出的嵌入平均机制,每次 Tokenize 都应理解为一次新的随机切分。

\[ \bar{\mathbf{e}}=\frac{1}{K}\sum_{k=1}^{K}f_{\theta}\!\bigl(\textsc{Tokenize}(\mathbf{x};\,d)\bigr) \]

模型权重保持冻结,平均的是不同切分得到的嵌入,而不是重新训练一个集成模型。实验使用 5 次前向传播;顺序处理时可控制峰值显存,但总计算随次数线性增加,随机切分产生的不同长度还会影响每次前向成本,因此不能称为免费提升。

损失函数 / 训练策略

预训练使用 UniRef90 的 50%,约 90M 条序列,联合词元级 MLM 与序列展开后的残基级 MLM。随后使用 TEDLH 域注释构建 59.7M 个蛋白质对,在 CATH 的 Architecture、Topology、Superfamily 层级上进行对比学习。正文没有充分提供完整损失形式、各项系数、温度、优化器与训练日程,本笔记不补造一个精确联合目标。

LEMON 主模型为 200M 参数,作者报告在单张 H100 上训练一周。训练与验证对通过 CD-HIT 分离,以减少微调集合内部的序列泄漏。这里的“只用序列”指模型输入和推理模态;CATH 是结构分类,TEDLH/TED 的构建也具有结构来源,所以结论中“没有显式结构监督”的宽泛表述与方法中的结构层级标签存在张力。

对 726,922 条微调训练域序列的审计没有发现与三个基准完全相同的字符串,但约 10% 的基准域与训练序列有至少 70% 的全局一致性,约 2.5% 达到至少 90%。CATH S20 中 80.9% 的折叠类型在训练中出现过。这支持讨论“已知折叠中的序列差异泛化”,而非宣称数据完全无近同源重叠,或已经验证对未见折叠类别的泛化。

实验关键数据

主实验

远程同源检索将每条序列映射为一个向量,按余弦相似度排序候选,再用结构分类层级定义正负关系。AUROC 衡量整体排序,mAP 补充列表靠前位置的检索质量。下表摘录原文 Table 1 的 LEMON 与最重要的对比学习基线 ProtTucker,不将单个层级的优势扩展成全表最优。

基准 / 层级 LEMON AUROC ProtTucker AUROC LEMON mAP ProtTucker mAP
CATH S20 / Architecture 0.825 0.783 0.356 0.250
CATH S20 / Topology 0.898 0.874 0.320 0.293
SCOPe / Fold 0.903 0.876 0.347 0.263
SCOPe / Superfamily 0.959 0.971 0.555 0.693
SCOP / Fold 0.904 0.855 0.287 0.212
SCOP / Superfamily 0.948 0.961 0.421 0.541

LEMON 的折叠层级优势清晰,但在 SCOPe 与 SCOP 的 Superfamily 指标上均低于 ProtTucker。Figure 3 的 SCOP 序列一致性阈值分析另报告最严格 th10 的折叠 AUROC 为 0.875,th95 为 0.904;它与 Table 1 的整体列不是可随意混用的实验口径。

正文声称跨三个基准平均 AUROC 为 0.86、mAP 为 0.31,却没有在所读文本中明确平均的层级及权重。若直接对 Table 1 的六列取算术平均,LEMON 约为 0.906 与 0.381,不能复现上述两数。因此保留其为作者报告值,不将其包装为可验证的六列平均。

Table 1 给 LEMON 标注 5m、ProtTucker 标注 9m,caption 只称单张 H100 的总 GPU wall-clock。根据表格的基准评测语境,应把它视为作者报告的嵌入/评测时间,而非训练时长;具体是否包含完整检索评估并不充分明确。它尤其不能与主模型的一周训练合并为同一个成本结论。

消融实验

Table 2 使用同一 50M 编码器与 8M 头,比较不同词表;有效 batch size 为 1024,在 UniRef90 的分层 20M 子集预训练,再用整个 TEDLH 集合微调。指标为 SCOPe 的 10% 一致性设置,采用 3 个随机种子。ROC-AUC 与 fold ROC 是原表分列标签,缓存没有充分解释两者的聚合区别,下面不擅自合并。

分词器 ROC-AUC mAP fold ROC 预训练时间(mins) 微调时间(mins) 总时间(mins,求和)
ZEST (32K) 0.739 ± 0.017 0.068 ± 0.002 0.749 ± 0.014 304.7 985.5 1290.2
BPE (32K) 0.729 ± 0.008 0.036 ± 0.002 0.743 ± 0.008 278.0 1030.0 1308.0
Char (25) 0.708 ± 0.011 0.059 ± 0.002 0.721 ± 0.010 650.7 1043.1 1693.8

ZEST 与 BPE 都有约 16.4M 嵌入参数,而 Char 只有约 0.01M;相同的是编码器与头,不是整个模型参数量。总时间一列为本笔记对两段时间求和,ZEST 的总时间较低,但 BPE 的预训练单段更快。该结果只支持这一受限训练配置,不证明所有数据规模下 ZEST 必然训练最快,也没有单独隔离词表先验、压缩和随机切分各自的全部贡献。

关键发现

Table 3 所称 zero-shot 是冻结嵌入、没有任务专用微调;功能任务仍通过余弦 kNN 从带标签的训练邻居转移注释,K=15。因此不能把它解释为完全不用标注数据。域边界任务使用滑动窗口表征聚类,循环置换任务使用成对比较,不能把所有任务简化为同一个分类器。

任务 指标 LEMON (200M) ESM-2 (650M) ProtTucker (3B)
Circular Permutation AUROC 0.779 0.476 0.720
Circular Permutation AUPRC 0.285 0.081 0.238
EC Number F-max 0.821 0.800 0.878
GO Mol. Function F-max 0.639 0.603 0.635
GO Mol. Function AUPR 0.625 0.584 0.616
Domain Boundary NDO 0.758 0.771 0.750
Domain Boundary NRes 0.637 0.658 0.674
  • 循环置换与部分功能指标表现较好,但 EC Number 低于 ProtTucker;域边界的 NRes 也低于两个基线,不支持“所有任务都超过 SOTA”。
  • NDO 表示预测片段与真实域的覆盖重叠,NRes 表示残基分配准确率;域级覆盖较好并不意味着逐残基边界同样准确。
  • TTA 在 SCOP 折叠检索上采用 3 个种子、5 次切分,原文描述中等 dropout 的有效区间约为 0.05–0.5;缓存未提供图中全部数值,不能补造精确增益表。

亮点与洞察

  • 输入等价类本身可以承担归纳偏置。本文不只压缩字符串,还让进化相近的片段共享 ID,改变模型必须区分和可以忽略的信息。
  • 分词边界可以成为表征增强来源。随机选择合法短片段使同一序列有多种视图,冻结模型时也能利用这些视图,但需要如实计入额外推理成本。
  • 检索几何与局部保真应分别监督。注意力池化负责全局比较,残基恢复约束压缩损失,两者合作并不等于其中任何一个指标能代理所有下游质量。

局限与展望

  • 推理仅需序列不等于训练没有结构来源的监督。更公平的比较应分开报告预训练语料、词表来源、CATH 标签和推理模态。
  • 审计排除了完全相同的训练字符串,却保留部分高一致性邻居与大量已见折叠。未来应单独评估严格同源隔离和未见折叠设置,而不是把现有结果称为完全无泄漏。
  • 字符基线的嵌入参数明显更少,受限消融也未充分拆开所有机制。匹配总参数、报告随机切分覆盖分布和多种子置信区间,会帮助判断提升来自哪里。
  • 多对一词元会损失细节,恢复头只提供学习约束,不能保证无损。原文也承认更细粒度的序列属性任务不及更大规模训练的模型。
  • 损失系数、完整超参数、主表时间口径和部分平均值定义不足,限制严格复现。文本没有给出可核验的代码仓库 URL,因此元数据不填写猜测链接。

相关工作与启发

  • vs BPE / BPE-dropout:BPE 按频率合并,BPE-dropout 随机省略合并;ZEST 从保守片段簇建立词表,再从前缀树中的合法匹配随机改选较短项。二者都改变粒度,但等价类来源不同。
  • vs ProtTucker:两者都利用 CATH 对比监督塑造检索空间。LEMON 进一步改变输入单位,并以较小模型取得更好的折叠层级结果;ProtTucker 在更细的部分层级及 EC 指标上仍较强。
  • vs SaProt / ProstT5:本文对比中的结构条件配置使用 3Di 信息,LEMON 将领域先验放到离线词表与训练监督中,使查询推理只输入序列。这个差别不意味着所有这些方法的所有使用方式都必然要求结构。
  • vs HHblits:profile 比对显式利用 MSA 统计,LEMON 则把一部分保守性压缩成离散输入,再由上下文编码和对比学习建立全局距离;两者的预处理成本与监督条件不能忽略。
  • 正文将参考文献 [33] 称为 EvoBPE,而参考文献标题是 PUMA: Discovery of Protein Units via Mutation-Aware Merging。本笔记保留这一名称差异,不自行断言两者完全等同。

评分

  • 新颖性: 4/5。保守区段聚类共享词元与随机合法切分的组合具有明确领域针对性。
  • 实验充分度: 3/5。多基准和多种子消融有价值,但参数匹配、数据隔离及平均口径仍有缺口。
  • 写作质量: 3/5。主线明确,但监督措辞、极限切分描述与时间定义需要更精确。
  • 价值: 4/5。说明输入单位设计能改善小模型的同源检索,不能视为全面替代大型蛋白质模型的证明。