编辑 | ScienceAI
蛋白质布局相比于序列往往被认为更加具有信息量，因为其直接决定了蛋白质的功能。而随着AlphaFold2带来的巨大突破，大量的猜测布局被发布出来供人研究应用。如何利用这些蛋白质布局来训练强大且通用的表征模型是一个值得研究的方向。
西湖大学的研究人员利用Foldseek来处理蛋白质布局，将其编码成一维的离散token，并与传统的氨基酸举行结合，形成了布局感知词表（Structure-aware Vocabulary），以此将布局信息嵌入到模型输入中，增强模型的表征才能。
在预训练上，论文应用了目前最多的蛋白质布局（identity过滤后4000万），在64张A100上训练了3个月，最终开源了具备650M参数量的模型SaProt（同时包括了35M的版本）。尝试结果表明SaProt各种蛋白质任意上都要好于之前的序列和布局模型。
研究《SaProt: Protein Language Modeling with Structure-aware Vocabulary》的预印版本，于 2024 年 3 月 21 日发布在 bioRxiv 预印平台。
论文链接：https://www.biorxiv.org/content/10.1101/2023.10.01.560349v4
github链接：https://github.com/westlake-repl/SaProt
步骤
本文利用Foldseek将蛋白质举行编码，生成了一维的3Di布局序列（应用了Foldseek的布局词表，每种3Di token代表不同的局部布局），这样的布局序列与氨基酸序列是等长的。
因此本文应用了一种简单而有效的布局嵌入形式：将布局词表和氨基酸词表计算笛卡尔积（即两两组合），形成新的布局感知词表。这样对于蛋白质的每个位点，其氨基酸类型和对应的局部布局都能组合成新词表中的某个元素，从而让模型同时考虑到蛋白质的序列与布局信息。
本文应用Bert架构举行掩码语言建模（Masked Language Modeling ）预训练（关于训练的更多细节可参考原论文）。
图：布局感知词表
尝试
步骤对比
一个可能令人疑惑的问题就是为什么需要这样编码布局？论文展示了应用不同的布局编码形式举行预训练的结果图：
图：不同布局模型训练的loss曲线图
图左和图中是两种经典的蛋白质布局建模形式，即将布局信息编码成bias后添加到transformer的attention map中（如Evoformer，Uni-Mol），或者应用图神经网络的形式建模蛋白质的空间关系（如MIF，GearNet等）。
然而从loss图中可以发现，当上述两种建模形式在AF2布局上应用MLM的训练目标举行预训练时，模型会非常迅速地过拟合（表现为在AF2猜测布局上猜测loss非常低，但在PDB真实布局上loss停滞甚至上升）。
作者推测这是由于AF2猜测出来的蛋白质布局带有一些隐藏的模式（patterns），由于前两种形式是直接对蛋白质的三维坐标举行建模，这些隐藏的pattern可能很轻易地就被模型识别出来，从而造成了信息泄露的问题，让模型无需真正学习到蛋白质的进化信息就能轻松地完成训练目标。
而布局感知词表通过将蛋白质布局编码成一维的布局序列，在尽可能保留布局模式的情况下忽略了精细的坐标数值，因此模型能够有效地利用布局信息而不受到隐藏pattern的影响。
Zero-shot尝试
作者在蛋白质突变数据集（ProteinGym）上和真实人类临床疾病数据集（ClinVar）上尝试了SaProt的zero-shot才能，结果以下：
图：Zero-shot尝试结果
SaProt在两个数据集上都超越了以往的所有布局和序列模型，证明了其在zero-shot猜测突变上具备优秀的才能。
监督微调尝试
本文还涵盖了各种下游任意来尝试模型表现，结果以下：
图：下游任意fine-tune结果
SaProt在各个下游任意上都超越了以往的序列和布局模型，展示出了其强大且通用的表征才能。
布局信息尝试
SaProt在4000万的蛋白质布局上举行训练，获得了强大的表征才能。一个可能的疑问是如何确定SaProt学到了更多的布局信息而不是模型被训练得更好？
论文对SaProt和ESM-2在残基接触猜测任意（Contact Prediction Task）上举行了尝试。作者冻住了模型的backbone，只训练一个线性分类层。尝试结果以下：
图：Contact Prediction Task的结果
从结果可以看到，由于布局token的嵌入，SaProt的表现大大超越了ESM-2，这表明SaProt蕴含了非常丰富的布局信息，使其能够在布局猜测任意上获得十分优秀的结果。同时，论文在SCOPe数据库上对alpha蛋白质和beta蛋白质举行了可视化，结果以下：
图：在SCOPe数据库上的Embedding可视化
SaProt的可视化结果非常清晰地将alpha蛋白质和beta蛋白质区分开来，而ESM-2的可视化结果却将两种蛋白质混杂在一起，这说明了SaProt对布局的变化有很强的感知才能。
不同布局猜测步骤的比较
除了AF2，目前还存在许多其他的单序列布局猜测步骤（如ESMFold），因此本文额外尝试了其他步骤猜测出来的布局对SaProt性能的作用。结果以下：
图：不同布局猜测步骤的fine-tune结果
从尝试结果可以看出，虽然SaProt在AF2布局上的表现最好（模型本身也是基于AF2布局举行训练的），但其他的布局猜测步骤也能让SaProt与ESM-2等模型性能相当。这意味着考虑到计算与时间成本，单序列布局猜测模型也能作为替代步骤输入到SaProt中。
局限
虽然SaProt经过训练展示出了优秀的性能，但依然还有一些可以改进的地方，例如：
Foldseek默认的布局词表大小只有20，如果有更加精准的布局编码模型，扩大布局表征的词表大小，是不是能进一步提升模型利用布局的才能？
由于计算才能的限制，SaProt只在650M上完成了训练。如果能够继续扩大模型规模，是否可以进一步地提升模型表现？
论文虽然已经尝试了很多的蛋白质任意，但还有一些其他任意可以应用探索，例如蛋白质序列设计（给定backbone猜测氨基酸序列）等。