蛋白质语言模型工作原理详解

蛋白质语言模型不是一个小号的 AlphaFold。它是一种基于氨基酸序列的掩码 token 模型，输出是每个残基一个向量，而真正有价值的问题是：这些向量是哪些因素的函数。

训练目标

它的框架与 BERT 引入的文本方法一致，只是应用到了一个二十个字母的字母表上。取一条蛋白质序列，将其中一部分残基替换为 mask token，然后训练网络根据剩余的残基恢复原始序列。损失函数中没有标签、结构、注释或比对——唯一的监督信号来自序列对自身的预测。

正是这个目标使得产生的表示变得有趣。要很好地填充一个被掩码的残基，模型必须学习蛋白质中哪些位置约束哪些其他位置，而在真实蛋白质中，这些约束主要来自结构和功能：两个在折叠蛋白质中相邻的残基在进化过程中共同变异，因为其中一个的变化只有在另一个发生补偿性变化时才能被容忍。一个擅长填充任务的模型，作为副产物，吸收了大量的这种共变异信息。

具体的掩码比例和破坏方案（在多大程度上将掩码位置替换为随机残基而非 mask token）在不同代模型之间存在差异，具体数值在每篇论文的方法部分中说明。使用某个检查点时，应该从论文中读取相应的掩码比例，而不是假设文本中的默认值仍然适用。

训练集究竟是什么

这是大多数解释变得模糊的部分，而且是可以查证的。对于 Meta AI 的 ESM 系列，项目自身仓库中发布的模型卡片直接命名了语料库：ESM-1b 在 UniRef50 2018_03 版本上训练，而 ESM-2 检查点在 UniRef50 2021_04 版本上训练，序列从这些 UniRef50 代表序列所在的 UniRef90 簇中采样。检查点名称中包含了这一信息：esm2_t33_650M_UR50D 以 UR50D 结尾就是这个原因。Meta AI 的 ESM 仓库发布了相应的表格。

UniRef50 作为一种选择是有意义的，而不仅仅是一个名字。UniProt 通过在 50% 同一性水平上对序列进行聚类，并每个簇保留一个代表序列来构建它，这是一种有意的去重：原始蛋白质数据库存在大量冗余，对于研究较多的蛋白质有数千个近乎相同的条目，而对于从土壤样本中一次性测序的蛋白质只有一个条目。在原始数据集上训练会使模型偏向于人类测序最多的蛋白质。50% 聚类可以消除这种偏差。从每个 UniRef50 簇内的 UniRef90 成员中采样训练样本，既能恢复簇内的一些变异信息，又不会恢复热门偏见。

对于使用者而言，这意味着一个覆盖范围的说法而非能力的说法。该模型已经学习了沉积在公共数据库中的序列空间，这些空间由细菌和古细菌基因组以及研究较多的真核生物蛋白质组主导。抗体可变区、内在无序蛋白和设计序列相对于人们实际运行模型的频率都存在代表性不足的问题。

嵌入编码了什么

自然输出是每个残基一个向量，来自最终的 transformer 层——对于 esm2_t33_650M_UR50D，那就是 33 层，每层产生 1,280 维的残基表示。包含 300 个残基的蛋白质返回的是一个 300 × 1,280 的矩阵。如果需要整个蛋白质的一个向量，就必须进行聚合，标准选择是对残基位置取平均。

根据已发表评估的探查内容，这些向量包含以下信息：

接触信息。这些模型的注意力图包含与三维空间中接触的残基对对应的模式，可以通过一个小型监督探针来恢复。这是填充目标正在获取结构信息的最直接证据。

家族和同源性。同一家族蛋白质的平均池化嵌入在空间中彼此靠近，这正是基于嵌入的同源性搜索在序列比对失去信号时仍能工作的原因。

位置对替换的耐受性。模型自身对某一位置预测的分布反映了该位置的约束程度。通过突变残基概率与野生型残基概率的对数比来对变体评分——即伪似然——是标准的零样本变体效应方法，完全不需要标注数据。

这些向量中不包含任何序列无法决定的信息。翻译后修饰、表达水平、结合伴侣的存在、缓冲条件：这些都不是输入，所以都不在输出中。

模型规模及其作用

已发布的 ESM-2 检查点从 6 层、320 维、800 万参数的模型扩展到 48 层、5,120 维、150 亿参数的模型，其中 6.5 亿参数、33 层、1,280 维的模型是大多数工作实际使用的选择。这并非保守做法。嵌入维度是构建每个下游特征向量的宽度，而对数十万蛋白质使用 5,120 维表示是一个需要存储和索引的大型矩阵——要认识到维度数量在下游带来的成本。

跨缩放工作的总体发现是，填充损失随规模平稳改善，表示中的结构信息也随之改善，但两者之间的关系并非线性的，在特定下游任务上的收益并不能保证迁移到你的数据集上。诚实的做法是在你自己的标注集上运行最小的检查点和一个更大的检查点并进行比较，这很便宜，因为头部很小；这在小标注集微调中得到了验证。

它没有告诉你什么

最重要的区别在于语言模型和结构预测器。在实验结构上训练出的折叠头可以附加在这些表示之上，这个组合系统预测坐标；而语言模型本身并不能做到这一点。当一篇论文报告这类系统中结构预测的准确性时，报告的是头部加躯干、在结构监督下训练的结果。单独来自躯干的嵌入是不同的产品，用于搜索、分类和变体评分。

其次，伪似然评分描述的是一个序列相对于训练分布的异常程度，而非表型描述。这些评分标记为可容忍的变体中包括真正致病的变体，而标记为破坏性的变体中也包括没有测量到任何影响的变体。这些评分在研究环境中用作先验，并作为优先级排序的输入。它们不是临床诊断，为临床目的解释人类变体是一项受监管的活动，需要由认证实验室根据已发布的标准进行。模型评分不能替代它，也不应该被表现得好像它可以替代。

第三，目标中没有任何东西奖励新颖性。模型评分很高的序列是看起来像训练分布的序列，而这对于设计工作来说往往是你想要的对立面。