BERT 一个蛋白质 - 季军 - 英特尔创新大师杯冷冻电镜蛋白质结构建模大赛 -paipai

2024-04-28
浙江
本文字数：3136 字
阅读完需：约 10 分钟

解决方案

团队介绍

paipai 队、取自 PAIN + AI，核心成员如我本人 IvanaXu(IvanaXu · GitHub)，从事于金融科技业，面向银行信用贷款的风控、运营场景。但我们团队先后打过很多比赛，其中跨领域居多，如天文、海洋，也非常有幸参加本次蛋白质结构建模大赛。我们将延续“他山之石，可以攻玉”的基本思想，这也将在后续方案中体现。

主要方案

01 赛题解析

金融场景下其实就是在追求精准定位问题，并以最小成本换取更大收益。同样的，从现有解析蛋白质结构的主流方法来看，能覆盖较多场景（如生物大分子、静动两态、蛋白质折叠等）且精度较高，但为此需要付出大量人力成本，甚至是在专业软件辅助下才能完成的人工操作。另外，由 DeepMind 研发的 AlphaFold2 却能高效几乎 0 人工输出蛋白质结构，当然这个代价是覆盖场景有限（如生物大分子易错）、局部精度高（需要强依赖于已知蛋白质结构/PDB）。据此我们提出的解决方案是，拥抱类似于 AlphaFold2 的 AI 预测技术但充分结合传统方法，或许是最大收益方案。即复杂生物应用场景下，求一种高效、准确的蛋白质结构解析方法。所以我们将解决方案聚焦到以下两点：（1）仅从氨基酸序列出发的 AI 预测技术为主，通过已知蛋白质结构去做一些自衍生处理，尝试去摆脱对已知蛋白质结构高度依赖；（2）当遇到已知蛋白质结构甚至是自衍生蛋白质结构都不能很好识别时，再来考虑结合冷冻电镜数据，作为预测结果的辅助性修正，毕竟获取冷冻电镜数据就已经投入了一定的成本；

02 算法设计

如示例 2145 这个氨基酸序列，我们看到前 4 位的 MSSK、就是一个甲硫氨酸+两个丝氨酸+一个赖氨酸的序列，但我们将它转为以下文本处理：第一步，将氨基酸序列-蛋白质结构作为一组 X-Y 对，即使用氨基酸序列预测蛋白质结构，并计划进行自衍生处理 1：将衍生后的氨基酸序列-蛋白质结构对进行随机混淆打乱，这一步复赛阶段未实现，但后续我们做了一些尝试；第二步，将氨基酸序列转化为文本序列，在这里会做自衍生处理 2：将它切分为不同长度的文本序列，如左侧 V2 算法的第一步，将 MSSKS 每隔 4 个氨基酸就切分为 MSSK、SSKS 等等，同样的还会按照 5、6 不等的长度进行切分，最终使用 4、5、6、8、10、12、15、18，分别代表了 4、5、6 长度段的 1 倍、2 倍、3 倍。从生物意义上就代表了一些短氨基酸序列。这种切分方法的优点是可以是因为把其充分序列化，那么蛋白质中的对称重复架构会在序列中特别突出，有点像 121213 这组数字，我们听一遍后，在脑海中对 12、121、21 这样的组合印象就很深刻；第三步，文本处理：尝试提取序列之间的潜在关系，事实上可考虑引入谷歌提出的基于 Transformers 的双向编码语言模型 BERT 来解决，但 BERT 很重，考虑上述成本效益，在此简化应用 TF-IDF 来处理；第四步，将处理后序列信息输入至预测模型，输出最接近蛋白质结构；

为进一步比较自衍生处理 1、自衍生处理 2、以及不同文本处理算法之间的差异，我们先后设计了 V1、V2、V3 三个算法版本：

V1，无自衍生处理 1、无自衍生处理 2，即 MSSK 序列会当作 M、S、S、K 四个单词组成的一句话，直接使用 BM25 文本模型进行预测；
V2，在 V1 基础上增加自衍生处理 2，即 MSSK 序列会切分为如图 MSSK、SSKS、等等单词组成的一句更长的话；

可以通过比较 V1、V2 来体现自衍生处理 2 带来的预测增益。

V3，会更充分考虑实际蛋白质结构预测中会遇到的问题，并根据本次比赛提供的 E-HPC 阿里云弹性高性能计算平台改进算法底层一些计算设计，如图：

（1）在高性能计算平台上应用 CPU 多核并发，能同时计算多个蛋白质，并将氨基酸序列计算转入稀疏矩阵存储 + TF-IDF 限制某些高频率出现的简单序列，减少大量空值存储和无效计算，充分压缩氨基酸序列结果，避免内存错误；（2）在 V2 的基础上调整 BM25 文本模型至 TF-IDF+LR；这就是复赛阶段的解决方案。（3）另外，复赛后我们还在此基础上设计如左虚框部分的自衍生处理 1，对算法拓展性做了一部分探索。引入 BERT 中 MLM/Masked Language Model 即遮蔽语言模型的概念，其灵感来自于我们学英语很熟悉的完形填空，如我们通过线上会议 XX 了答辩，就会填写“参加”这个词。这种概念其实类似于自衍生处理 2，本质上是为了充分提取氨基酸序列的潜在关系，如前面提出的不同长度切分就为了获取一定的上下文信息，即氨基酸序列中 S 丝氨酸前后的氨基酸信息，但这种方式也仍然是单向或一个正向+一个反向这种假双向。反之 MLM 能更充分的解决这个问题，因为这个时候模型学到不是一个向量，而是**“一种学习能力”**。TODO

所以在 AAAI2020 的工作中，已经有论文提出 K-BERT 即知识图谱+BERT 的概念，使用类似 MLM 的方式尝试进行知识发现。

同样的，可以使用这个概念来“知识发现”一些未知的蛋白质结构，进而打破上述对已知蛋白质结构的高度依赖。

以下我们来对照 AlphaFold2，对比一下上述算法：

以 2148 为例，我们也尝试使用了 AlphaFold2：

通过前两栏的 RMSD、TM-score 来看：

无疑是 AlphaFold2 最强，但对比 V2、V1 可以看到，在做了自衍生处理 2 即序列切分后，TM-score 自 0.19 提升到了 0.77，V2 算法也就是初赛核心方案。另外，改进后的 V3 算法从个例来看相对接近于 AlphaFold2。

另外，由于我们在调研阶段从 PDB 中发现了远比训练数据更大的蛋白质，所以这里的 long 评估项即长氨基酸序列的兼容，AlphaFold2/V1/V2 都不同程度地出现在长氨基酸序列计算时内存错误的情况。前面介绍过，V3 在做了一些相关设计所以相对可控。且能在不需要 GPU 的情况下进一步兼容 HPC 高性能计算，尽管 AlphaFold2 也提出支持 HPC，对此我们未深入体验。

接下来我们看一下可视化结果对比：

AlphaFold2 已经接近严丝合缝，V1 差异还比较大，但 V2、V3 已经有了相似的蛋白质结构雏形，未来可能可以进一步优化。

03 总结

1、直接转换为文本问题进行解决事实上我们也发现其实这与 Alphafold2 的序列处理有类似的地方，但不同 AI 场景下 2D 和 3D transformers 的发展，可能还可以进一步借鉴。但氨基酸序列的潜在关系我们还是交由更为简单的文本处理算法（TF-IDF）进行提取，这里考虑降低模型复杂性、提高应用效率。如我们还没有提到的冷冻电镜图片处理，CVPR2022 的工作中已经提出了 Point-BERT 基于掩码建模的 3D 点云自注意力模型，实现将类似于冷冻电镜数据这种 3D 结构表达为一个“词汇”集合。跟我们上述将氨基酸序列转化为文本处理的逻辑类似，我们的解决方案也提出了类似的思想，先用氨基酸序列进行文本预测，再对预测不佳的部分根据冷冻电镜数据也进行文本预测，作为辅助性修正，但这部分还没有来得及完成。2、避免高度依赖已知蛋白质结构为此前面我们设计了自衍生处理 1、自衍生处理 2，并通过 V1/V2 比较证明了自衍生处理 2 的有效性。也正是引入自衍生处理 1，通过 V2/V3 比较看到了突破这种蛋白质结构依赖带来的额外效益。