过去的蛋白质语言模型以单个序列为输入,MSA Transformer以多序列比对的形式将一组序列作为输入。该模型将行和列注意力交织在输入序列中,并在许多蛋白质家族中使用mask语言建模目标进行训练。模型的性能远超过了当时最先进的无监督学习方法,其参数效率远高于当时最先进的蛋白质语言模型。
来自:MSA Transformer, ICML2021
ESM-MSA-1b概述
ESM-MSA-1b是一个在UR50上训练的掩码语言模型,ESM-MSA-1b在三种下游任务:无监督氨基酸接触图预测、有监督氨基酸接触图预测和二级结构预测上均达到了当时的SOTA水平。
模型采用随机mask,其中,Transformer做了改进,使内存占用降低。模型与普通Transformer的区别是,其特殊的轴向注意力机制,见图1所示。对于蛋白质通用规律的学习,MSA信息中同源序列不同位置上的氨基酸对于当前的氨基酸的权重信息并不大,当前氨基酸最重要的关注点还是同一序列其他氨基酸(行)和同源不同序列同一位置的其他氨基酸(列),因此注意力限制在横纵这两条轴向范围就可以了,可以大大降低时间复杂度。
- 图1左,注意力稀疏结构。通过将注意力限制在行和列上,计算成本从 O ( ( L M ) 2 ) O((LM)^{2}) O((LM)2)降低到 O ( L M 2 ) + O ( L 2 M ) O(LM^{2})+O(L^{2}M) O(LM2)+O(L