Optimizing Protein Tokenization: Reduced Amino Acid Alphabets for Efficient and Accurate Protein Language Models

该研究通过结合基于理化性质分组简化的氨基酸字母表与字节对编码（BPE）子词分词技术，成功在蛋白质语言模型中实现了显著的训练与推理效率提升，同时保持了甚至改善了在多种下游任务中的预测性能。

要点

这篇论文就像是在教我们如何给蛋白质“减肥”和“提速”，同时还能让它们保持聪明。

想象一下，蛋白质是由 20 种不同的“乐高积木”（氨基酸）搭建而成的复杂结构。以前的科学家在教人工智能（AI）认识这些蛋白质时，习惯把每一种积木都当成一个独立的单词，就像用 20 个不同的字母拼写文章一样。

这就带来了一个大问题：
如果一篇文章有 1000 个字母，AI 就要读 1000 次。这不仅读得慢，还特别费电（计算成本高）。而且，因为字母太多，AI 很难发现那些重复出现的、有规律的“词组”（比如某种特定的结构模式）。

这篇论文的解决方案是：
与其死记硬背 20 种积木，不如把它们归类！

1. 核心创意：把“字母表”变小

作者们想：“如果我们把 20 种积木按性格（比如是亲水的还是疏水的、是酸性的还是碱性的）分成几大类，会发生什么？”

他们尝试了五种不同的“分类法”：

• 20 类（原版）： 每种积木一个名字。
• 12 类、8 类、4 类： 把性格相似的积木归为一组，给它们起个统称。
• 2 类（极简版）： 直接分成“喜欢水”和“讨厌水”两派。

2. 魔法时刻：BPE 分词术

有了这些简化的分类，作者们使用了一种叫BPE（字节对编码）的技术。你可以把 BPE 想象成一个聪明的“缩略语生成器”。

• 在原版（20 类）中： 因为积木种类太多，重复的长模式很少，AI 只能看到很多零散的短词。
• 在简化版（如 2 类或 4 类）中： 因为积木被归类了，原本不重复的长序列现在变成了重复出现的“大词”。
  • 比喻： 就像把“苹果、香蕉、橙子”都叫“水果”。以前你要说“苹果香蕉橙子苹果”，现在可以说“水果水果水果水果”。AI 一下子就能读懂更长的句子，而且读起来快多了！

3. 实验结果：快，而且不笨

作者们训练了不同版本的 AI 模型（叫 ProtBERTa），让它们去解决各种生物难题（比如预测蛋白质能不能溶解、能不能发光、能不能在高温下存活）。

• 速度提升惊人：
  使用简化分类的模型，输入给 AI 的“句子”变短了。

  • 最简化的模型（2 类）比原版模型快了 3 倍，训练时间也缩短了一半以上。
  • 这就像是用快递送包裹，以前是 1000 个小包裹，现在打包成了 300 个大箱子，运输效率自然大大提升。

• 准确度如何？
  这是最精彩的部分：

  • 大部分任务： 简化版模型的表现几乎和原版一样好，甚至有时候更好！
  • 特定任务：
    • 预测最佳工作温度时，最简化的模型（2 类）表现最好。这可能是因为太细节的信息反而成了干扰，简化后的模型抓住了“大局”（比如整体是耐热还是怕热）。
    • 预测蛋白质相互作用（两个蛋白质怎么握手）时，原版（20 类）稍微好一点。因为这时候需要知道具体是哪块积木在握手，太简化的分类会丢失这些细节。

4. 总结与启示

这篇论文告诉我们一个重要的道理：有时候，少即是多（Less is More）。

• 以前： 我们总认为给 AI 越多的细节（20 种氨基酸），它就越聪明。
• 现在： 我们发现，通过聪明的分类（减少字母表），我们可以让 AI 读得更快、更省资源，而且在很多情况下，它依然能保持极高的智慧。

这就好比：
如果你想教一个人识别“森林”，你不需要让他记住每一棵树的学名（20 种积木）。你只需要教他区分“针叶树”和“阔叶树”（简化分类），他就能更快地理解森林的布局，而且往往能更准确地回答关于森林整体特征的问题。

这项研究为未来的生物 AI 模型指明了一条新路：通过简化输入，让 AI 跑得更轻、更快、更聪明。

技术摘要

论文技术总结：优化蛋白质分词——用于高效准确蛋白质语言模型的简化氨基酸字母表

1. 研究背景与问题 (Problem)

蛋白质语言模型（pLMs）通常采用单氨基酸级别的分词（tokenization），即使用包含 20 种氨基酸的标准字母表。这种方法虽然能保留位置特异性信息，但会导致输入序列过长，从而显著增加计算成本（训练和推理时间）。
虽然子词分词方法（如字节对编码，BPE）可以通过学习高频共现模式来缩短序列长度，但在标准的 20 种氨基酸字母表中，长序列模式往往非常稀疏，限制了 BPE 生成更长、信息量更大的 token 的能力。
此前，简化氨基酸字母表（Reduced Amino Acid Alphabets，将具有相似理化性质的氨基酸归为一类）已被用于生物信息学，但在 pLM 中的评估主要集中在单残基分词上，且结果显示性能有所下降。目前，尚未有系统研究探讨将“简化氨基酸字母表”与“子词分词（BPE）”结合使用对蛋白质语言模型效率及性能的影响。

2. 方法论 (Methodology)

本研究提出了一种结合简化字母表与 BPE 分词的新策略，具体步骤如下：

2.1 数据准备

• 语料库：来自 EBI MGnify 和 NCBI GenBank 的微生物蛋白质序列（去除了真菌、后生动物和植物），经 CD-HIT 去冗余后，划分为训练集和测试集（80:20），确保同源序列不交叉。
• 数据集：构建了涵盖多种下游任务的基准数据集，包括：
  • 分类任务：溶解度、酶识别、转运蛋白识别、双组分系统、蛋白质 - 蛋白质相互作用（PPI）。
  • 回归任务：荧光强度、最适温度、稳定性预测。
  • 零样本评估：同源性检测、信号肽预测。

2.2 简化字母表设计

研究使用了五种不同粒度的简化字母表（基于理化性质分组）：

1. 20 字母（基线）：标准氨基酸。
2. 12 字母：基于 Linclust 聚类算法的字母表。
3. 8 字母：基于功能基团分组。
4. 4 字母：基于极性分组。
5. 2 字母：基于亲水性/疏水性分组。

2.3 模型训练与评估

• 分词器训练：针对每种字母表，在 2830 万蛋白质序列上训练 BPE 分词器（词汇量固定为 5000），以控制变量，仅考察字母表大小对 token 长度的影响。
• 预训练：基于 RoBERTa 架构（12 层，768 隐藏维度）从头预训练蛋白质语言模型（命名为 ProtBERTa_X，X 为字母表大小）。使用掩码语言建模（MLM）目标，训练 5 个 epoch。
• 微调与评估：
  • 分类任务：添加分类头，冻结编码器权重（PPI 任务除外，需全量微调）。
  • 回归任务：使用平均池化输出，添加回归头。
  • 基准测试：在 DGEB（多样化基因组嵌入基准）上评估嵌入质量，并计算同源性检测、信号肽预测等零样本性能。
• 效率评估：测量训练时间、推理时间及序列压缩率。

3. 关键贡献 (Key Contributions)

1. 首次系统结合：首次系统性地研究了简化氨基酸字母表与 BPE 子词分词在蛋白质语言模型中的协同作用。
2. 效率与性能的权衡分析：证明了简化字母表能显著增加长序列模式的频率，使 BPE 能够生成更长的 token，从而大幅压缩输入序列长度。
3. 任务特异性发现：揭示了不同任务对字母表分辨率的敏感性不同。某些任务（如最适温度预测）在低分辨率字母表下表现更好，而精细任务（如 PPI）则依赖高分辨率。
4. 开源资源：提供了训练好的分词器和预训练模型，供社区使用。

4. 主要结果 (Results)

4.1 分词与序列压缩

• 随着字母表尺寸减小（从 20 降至 2），BPE 生成的平均 token 长度显著增加。
• 序列压缩率：ProtBERTa_2（2 字母）相比基线 ProtBERTa_20 实现了约 4 倍 的序列长度压缩（即输入序列长度减少至原来的 25%）。

4.2 嵌入质量与基准测试 (DGEB)

• 综合得分：ProtBERTa_12 在 DGEB 基准测试中取得了最高综合得分（0.35），略优于基线 ProtBERTa_20（0.347）。
• 零样本性能：在同源性检测和信号肽预测任务中，ProtBERTa_20 和 ProtBERTa_12 表现最佳，但简化模型（如 12 字母）在保持高性能的同时提供了显著的压缩优势（1.28 倍）。

4.3 下游任务表现

• 分类任务：
  • 溶解度、酶、转运蛋白：ProtBERTa_20 表现略优或持平，但 ProtBERTa_12 和 ProtBERTa_8 在性能损失极小（甚至无显著差异）的情况下，实现了 1.5 倍以上 的输入压缩。
  • PPI 任务：性能随字母表减小而下降，表明精细的残基身份对相互作用预测至关重要。
• 回归任务：
  • 最适温度：表现随字母表减小而提升，ProtBERTa_2 表现最佳。这可能是因为小数据集下，简化字母表有助于过滤噪声，学习更通用的热力学特征。
  • 稳定性：ProtBERTa_4 表现最佳。
  • 荧光：ProtBERTa_12 表现最佳。

4.4 效率提升

• 训练与推理时间：训练和推理时间大致与序列压缩率成正比。
  • ProtBERTa_4 的训练时间约为基线的 50%。
  • ProtBERTa_2 的训练时间约为基线的 33%。
• 推理扩展性：推理时间随数据集大小呈线性增长，且简化模型在所有规模下均显著快于基线模型。

5. 意义与结论 (Significance & Conclusion)

• 效率革命：该研究证明，通过引入先验的氨基酸理化知识（简化字母表）并结合 BPE 分词，可以在不显著牺牲预测性能（甚至在特定任务上提升性能）的前提下，大幅降低蛋白质语言模型的计算成本。
• 任务适应性：不同生物任务对序列信息的粒度需求不同。对于需要全局特征或受限于小数据集的任务（如温度预测），简化字母表通过减少过拟合风险反而提升了性能；而对于依赖精确残基相互作用的复杂任务，则需保留更多细节。
• 未来方向：建议在实际应用中，针对特定任务评估多种字母表配置。该方法为构建更高效、可扩展的蛋白质语言模型提供了一条切实可行的路径，特别适用于资源受限或需要处理大规模蛋白质组数据的场景。

总结：本文提出了一种“简化字母表 + BPE"的优化策略，成功解决了 pLM 输入序列过长的问题，实现了训练和推理效率的显著提升，同时保持了甚至在部分任务上超越了标准 20 氨基酸模型的性能。