Alphafold2
Alphafold2
Title: Highly accurate protein structure prediction with AlphaFold
Authors: Jumper, John
DOI: https://doi.org/10.1038/s41586-021-03819-2
Date: October 11, 2022
Finish time: 2022/10/15
Future: 应用于别的很多领域,药物靶点啥的
Meaning: 解决了生物学中存在了50年的难题,二维氨基酸序列预测3D蛋白质结构
Theoretical/Conceptual Framework: Transformer
Year: 2021
关键词: AI,Bioinformatics
期刊杂志: nature2021
Reference: 【AlphaFold 2 论文精读【论文精读】https://www.bilibili.com/video/BV1oR4y1K7Xr?vd_source=5ec85dfc5468a21a485b1b1d4d271219
标题&作者
Highly accurate protein structure prediction with AlphaFold
使用AlphaFold进行非常精确的蛋白质结构预测
与RoseTTAFold(Accurate prediction of protein structures and interactions using a three-track neural network)同一天发表
前面四行作者后面都写了个4,4的意思是说这些作者的贡献是一样的,第一个和最后一个是通讯作者
摘要
前面大篇幅介绍问题和解决这个问题的重要性,介绍实验的结果,最后很简单的提了一下他们的算法
-
蛋白质结构的重要性决定功能
蛋白质通常是指比较长的一串氨基酸序列,比如100个氨基酸,那么长一串不稳定,就会卷在一起,卷完之后比较稳定,因为这些氨基酸之间相互吸引,然后形成了一个独特的3D结构,结构的形状决定了蛋白质的功能
氨基酸序列可以认为是一段代码,要让他干活的话就要把它编译成一个可执行的文件,序列卷起来的过程就可以认为是把代码编译成可执行文件,你的代码就决定了你最后的可执行文件的样子
有时候你会发生错乱,比如说一个序列可以卷成两个不同的形状,这时候通常会带来疾病,所以正常情况下都是唯一决定的
这里的工作就是,这个代码是怎么变成这个可执行文件的,这个叫做蛋白质结构的预测。我给你一串氨基酸序列,你去预测3D结构长什么样子
-
“protein folding problem”——蛋白质结构预测的困难
目前为止我们大概知道10万左右的蛋白质结构,但是我们已知的蛋白质有将近10亿种,所以我们只对很少的一部分的蛋白质了解它的功能,目前我们可能需要数月或者数年的时间才能了解一个蛋白质的结构
具体来说就是,把这个蛋白质动起来,从不同的角度用显微镜去看他的投影,再还原出他的3D结构,这个费时又费力,所以我们需要更精确的基于计算的办法,能够更便宜更方便的得到一个蛋白质的结构,这个问题叫做蛋白质折叠问题,这个问题在生物学里面存在了50年了
-
现在方法的问题
alphafold1精度不够,不在原子级别的精度,也就是说,你对一个氨基酸位置的预测和实验真正测到的那一个结果,偏差不在一个原子大小的级别
特别是你不知道跟这个蛋白质功能比较,预测的其它蛋白质的结构的时候,这个精度会更差
-
介绍了alphafold2
- 我们终于提供了第一个能达到一个原子精度预测的模型,即你预测的位置和你真实测到的位置之间的偏差在一个原子大小的级别
- 我们用这个模型参加了CASP14竞赛(这是两年一届的比赛,科学家把最近两年里面测出来结构的蛋白质但是没公布的拿出来做竞赛,看一下你跟实验室测得了精度的比较),得到了大家的认可
- alphafold2是一种比较新奇的机器学习的算法,使用了物理和生物学的一些知识,同样也用了深度学习算
经验:这是一类应用型的文章,你主要的目的是用机器学习这个工具解决这个领域里面一个重要的问题,所以你关心的是两点:1、这是一个什么样的问题,这个问题重不重要,对这个领域来说有什么意义,2、你结果的好坏,你是不是解决了这个问题,具体来说你用的模型的好坏并不重要。比如说alphafold2就是一个很简单的mlp,只要你也能达到原子级别的精度,那你一样的能上nature封面。这个和你做一个特别复杂的模型,得到同样的结果,可能共识上来说是差不多的。反过来说,如果你能用很简单的算法解决这个问题的话,很有可能已经轮不到你了,别人就已经解决了。所以你就有两个办法,要么你找一个新的问题,别人都没用机器学习来做过,你把它变成一个机器学习的问题,收集数据可能用简单的模型也就行了,要么你的问题数据已经摆在那里了,之前而且有举办过竞赛,大家都做了很多不一样的方法,这时候你就需要用一个很不一样的很强大的算法,在算法上创新才行呢
导论
摘要里故事的详细版本,讲了一下问题,讲了一下实验结果,讲了一下具体的数字,没有讲太多的算法
看一下图
-
这次比赛的结果
x轴:参赛队伍
y轴:在某一个特定置信区间里面,你的平均预测的位置跟真实的位置的一个平均的区别。
单位:Å($10^{-10}$)
alphafold平均误差在1Å左右,别人基本上只在3Å,一个碳原子的大小大概是1.5Å的样子,所以可以达到一个原子精度的误差
经验:表达自己的成果的重要性时,把你的结果换算成一个大家能理解的一个概念,而不是一个纯粹数字的概念,大家对数字的敏感性远远不如对一个常见事物的比较的敏感度来的高
- “我比别人的结果都要好三倍” “极大的改善了别人的结果” ➡️ “我达到了原子精度的结果”
- “我的图片识别比别人好了10%” ➡️ “我的图片识别别人类识别的精度还要高”
-
实验结果
绿色表示实验室做出来结果,蓝色表示预测
在PDB数据集(目前为止所有我们了解的蛋白质的结构都放在这个数据集里)上的预测结果
a: x轴:预测的误差
y轴:有多少比例的蛋白质你的误差在这个里面
可以看到一半以上的预测在原子精度里面
✨模型介绍和训练
整个模型和训练的细节还是比较复杂的,但是就那么一点点篇幅,不可能在这个篇幅里面把整个东西给你写清楚,它采用的策略是说,大概给你讲一下这个模型是干什么,还是基于你的整个神经网络不那么了解的前提下,然后再高亮了一下他们的贡献,他们做的东西跟别人东西有什么不一样的地方,以至于你知道它的每个词是什么意思,但是读完下来也不知道这个模型是什么样子
-
算法总览图
输入:一个蛋白质的氨基酸序列,画了个人,表示是人的,每一个点表示一个氨基酸
输出:对每个氨基酸序列预测它在三维空间中的位置
概论:算法结构就分为了三个部分,拿到一个我们想要预测结构的蛋白质序列之后
-
先抽取特征主要是两大类特征,一个是氨基酸序列特征,一个是氨基酸对之间特征
-
进去完之后有一个编码器和解码器,编码器是一个transformer编码器的一个改进版本,因为他要处理两个不一样的输入,而且每个输入里面在行和列上是都有关系的,还要在里面不断的对两个模块进行数据的交互
-
做完编码之后再进入解码器,解码器就是根据编码器的信息去还原出每个氨基酸里面在3D中的一个位置,它仍然是基于注意力的机制,但是他在里面显式的对位置信息进行了建模
这个地方更像是一个RNN的架构,这也多少解释了为什么他的编码器叫做evoformer,听上去就是一个Transformer的变种,但是他的解码器就是简简单单叫做结构模块
-
然后还有一个回收机制,能把最后3D的结构的信息和编码器的输出都可以加回到编码器的输入里面,然后重复4次,就可以做成一个更长一点的网络,但是因为不进行误差反传,只是在计算上贵了4倍,但是在内存上并没有太多的瓶颈。
细节很多,附录里面有大概30个算法
-
-
算法结构
-
feature——抽特征
输入序列进来
-
直接导入到最后的神经网络
-
去一个基因数据库里面搜一搜,看一看跟这个蛋白质的序列有没有相似的蛋白质的序列。
比如说这个地方他发现了在鱼、兔子、鸡身上有相似片段的蛋白质,把它拿出来形成一个MSA(Multi sequence alignment 多序列比对)
具体来说,拿出来相似蛋白质,一一比对氨基酸,这个地方是相似而不是一样,就意味着说有一些片段是一样的,有些片段不一样,那么需要这些相似的片段或相似的氨基酸跟你一一的对起来(字符串匹配的过程)
-
另一个重要的特征是氨基酸之间的关系。蛋白质卷起来是因为氨基酸之间原子的相互作用,那么我们这也可以表示每两个氨基酸之间的一些关系,如果蛋白质的长度是2的话,就会整成一个2*2的表,每一项表示的是一对氨基酸之间的关系
当然最好的是能表示成在卷完之后这两个氨基酸在3D空间之间的距离
-
另外额外的特征是,把这个序列放进结构数据库里面去搜。我们知道很多蛋白质的结构信息,那我们就知道真实的氨基酸对之间在空间中它的距离什么样子,这样子就得到了很多模板
抽完特征之后,我们得到两大类特征,第一个是不同序列的特征,第二个是氨基酸之间的一类特征
这两大类特征拼上我们后面来的一些东西(后讲),就进入我们的编码器
-
-
encoder —— 编码器,将输入的特征编码成神经网络想要的形状
-
输入:两个三维的向量
-
有s行,s表示的是你一共有多少个蛋白质,第一个是要预测的人类的蛋白质,后面s-1个是我们在数据库里面匹配而来的s-1个蛋白质
长度是r,表示这个蛋白质里面一共有r个氨基酸
长度为c,表示我把每一个氨基酸表示成一个长为c的一个向量。对于图片来讲,就是每一个像素的通道数,对于文本来讲,对应的概念就是每一个词对应的那一个嵌入的长度
-
一个氨基酸对的表示
一共有r个氨基酸,所以就是一个r乘以r
每一个氨基酸对也是用一个长为c的向量来表示他的特征
-
-
Evoformer(48 blocks)——Transformer变种
-
不同点
-
不再是一个序列里面关系,现在是一个二维之间的关系
-
MSA
原来,不管是对文字还是图片来说输入都是序列;
现在,不再是一个序列,而是一个二维矩阵,在每一行上面可以认为就是这一个氨基酸的序列,列里面也是有关系的,因为这是同一个氨基酸在不同的蛋白质里面的一个表现,所以你的关系不再是一个一维的序列,而是一个二维的矩阵了
-
Pair
每一行表示的是一个氨基酸跟别的氨基酸的关系
在列上面也表示的是一个氨基酸跟别的氨基酸的关系,这个地方有一定的对称性
而且如果你认为这个对表示就应该表示这氨基酸在3D空间中的距离的话,那么这个地方你还有些物理上的限制,比如说你要满足一个三角关系
-
-
现在输入的是两个不同的张量
之前就输入一个,现在输入两个把数据融合起来
-
-
相同点
transformer里面基本上就是一些transformer块,每个块一个输入进去,输出应该是跟输入一样的,但是抽取了他们元素之间的相关信息,然后通过48个块之后,最后得到的输出跟输入之间的大小没有发生变化
但是输出那个地方里面的那些向量就是c那个维度,已经很好的表征了每个氨基酸之间他们的关系,这就完成了编码器的输出
-
-
-
decoder
解码器,将编码器的输入最后变成一个3D的位置。
解码器拿到编码器的输出,具体来说,拿到要预测的那个人类蛋白质里面氨基酸所有的特征的表示以及氨基酸之间的相关信息
根据这些相关信息对每一个氨基酸来预测它的3D位置,最后得到我的输出
-
Recycling(回收机制)
-
他把编码器的输出和解码器的输出都通过回收机制变成了编码器的输入
-
大概意思是说,做完这一次预测的时候,可能结构还不是那么的准确,把这个不是那么准确的结构传回我们的输入,然后在他的基础上再做一次精调,希望精度会更好,有点像循环神经网络的机制,可以认为把这个模型循环了3次,变成了一个四倍更深的一个网络,这样子达到更好的精度
-
但是他比复制三次的好处在于是说,在网络里面是共享权重的,也就是说下次再做这个东西的时候,权重还是基于前面的
-
另外一个是,在做回收的时候,我们的梯度是不反传的,也就是说正常从前面进来之后,损失函数在最后输出3D结构的地方就要往回了,只要经过48+8=56层,就可以计算梯度了
但如果你允许梯度在回收时候传递的话,也就是像循环神经网络一样,在一条长的序列里面传递梯度的话,那么你要经过56*4=224层才能进行梯度反传,内存有可能会吃不消,内存消耗率跟层数是线性关系
-
-
-
架构具体图
-
encoder
通过自注意力模块来对每个元素之间的关系进行建模,然后通过一个全连接层来做一个信息的空间的转换,不同是说输入里面既有行的信息,也有列的信息,所以这些模块都是成双成对的出现的
然后有两个输入,分别表示的是氨基酸在一个序列中的信息和氨基酸对之间的信息,需要做一些信息的传递,使得这两块之间能够做信息的融合
-
对空间信息建模
-
跟transformer的非常相似。
首先进入的是一个多头自注意力的模块,同样道理有残差连接,然后就是一个在元素上作用的MLP
-
不同点
-
氨基酸对之间的信息,会通过绿色这条路加入到对序列的建模里面,序列的信息也会通过橙色这条路来反馈到氨基酸对的信息上面,所以就完成了一次信息的交互,而且是每一块里面都有信息的交互
-
每一个自注意力机制更多了,因为之前他只要在一个序列里面就做了,所以这里的话我们有个二维的关系
紫色框框处理的是按行的序列信息
绿色框框处理的是按列的序列关系
粉色框框是元素对之间的关系,要满足一些之前说到的物理关系,比如说三角不等式之间的关系,所以他也通过这个性质来控制了你在做自注意力的时候,他们是怎么样去设计你的query、key和value
-
-
MSA(row-wise gated self-attention with pair bias)
按行的 带门的 自注意力机制 带按对的偏移
-
按行的
-
在MSA里面每一次拿出一行做成一个序列,但我们知道里面的每一个元素就是一个氨基酸,有一个长为cm长的长度作为他的特征,那么就可以用最原始的多头自注意力了
-
具体来说就是对里面每一个元素做投影,来得到我们的query、key、value,q k v长度都是等于c的,而且是多头的话我们要一共做h次,query和key做点乘,来计算他们的相似度
-
通过一个softmax,计算他的自注意力权重,然后拿注意力分数跟你的value做乘法,得到我的输出
因为是多头的话,所以这个地方一共要进行h次,可以看到这里是有多个矩阵,把这些输出全部跟你并在一起,最后做一次线性的投影,就得到我们的输出了,就是得到这一行他的对应的新的值
-
-
带门的
- 多了一行红色框框的东西,他对每个元素跟之前一样也做一个线性投影,然后用一个sigmoid函数将这个输出变成0和1之间,在每个头里面都做一次这样子的计算,然后跟你这个输出做点乘起来,按元素的乘法,所以这就完成了门这个操作
- 红色框框这一个线性投影让我们学习一个方法,使得我们来控制哪些元素应该输出出去,哪些不应该输出
- 输出出去的意思是说,我这个元素的值比较大,我做完sigmoid的时候变成了1,那么就是我对这个输出不改变,让他放他出去
- 如果对某个元素我不想把它放出去的话,那么学习一个比较小的值,通过sigmoid之后,基本上把它变成0,0乘以这个元素的话就变成0了,就能实现控制,哪些应该输出,哪些不应该输出,这就是门的意思
-
对偏移
-
在算注意力分数的时候,对每个q和每个k做内积得到他的相似度,也就是说作为query的那一个氨基酸和作为key的那一个氨基酸,计算他们的相似度。但是pair representation这个地方,我们还额外的建模了每一个氨基酸对他们之间的关系,所以信息跟dot-product干的事情有一定的相似度
-
假设我要计算第i个氨基酸作为q,然后第j个氨基酸作为key的时候,我可以把它对应的那一个第i行第j列的那一个元素拿出来,这里也表示了他们之间的关系,通过一个线性投影投影到一维,然后作为一个偏移加进去
就是之前做的那个dot-product内积再加上这一个偏移,然后再进入softmax,这样我的注意力分数既来自于我们这个投影算出来的相似度,也来自于我们额外建模的相似度,这个就是pair bias要干的事情
-
所以看到他就是在原始的多头注意力上面做了三个改进,用来建模我的MSA里面的每一行的序列信息,以及能用上我们之前建模的按对的那些信息
-
伪代码
-
把张量表示成向量的一个集合
$m_{si}$表示的是第s个序列中间的第i个氨基酸对应的那一个向量他的长度是cm
-
$z_{ij}$表示的是第i个氨基酸和第j个氨基酸,他们之间的关系的向量表示也是cm
-
c是一个超参数,表示我每个头里面投影的大小是32,所以算法是一个多头注意力,有一个投影步骤,一个注意力步骤和一个输出步骤
-
-
-
MSA(Column-wise gated self-attention)
他跟之前唯一的区别是说,在之前是按行的,也就是只见过的每一行里面,每一个序列里面氨基酸之间的信息
现在要建模每一个列里面的信息,也就是同样一个氨基酸,但是在不同的蛋白质之间,他的表示之间的一个相关性
如果把输入那一列拿出来之后,剩下的部分跟之前是没有区别的,除了没有把对信息作为一个偏移加进去之外,其他都是一样的
做完前面两个MSA块,我既对于每一行里面序列信息做了建模,也对每一列信息里面做了建模,虽然我们不是对整个行列里面所有信息同时建模,通过两个模块,计算复杂度相对来说是比较低的
-
MLP模块(Transition)
- Transformer里面自注意力机制主要的作用是去混合不同元素之间的信息,真正的做信息的提炼还是在MLP模块里面
- 简单说就是对每一个元素,通过一个全连接层,投影到原始的4倍长的大小,4这个常数也是来自于Transformer里面,然后再通过一个ReLU的激活层,再投影回原来大小
- 线性层的权重对每一个元素是共享的,也就是对每个元素,我们用同样的全连接层进行投影,等价于transformer里面的那个MLP模块
你就想象的把每一个氨基酸这个矩阵拉成一条长的向量就行了,这是等价的
这就完成了对MSA它里面空间信息的一个建模
-
-
对对信息进行建模
来自MSA氨基酸序列的表示是怎么样通过这个紫色的模块融入到氨基酸对的表示?
-
氨基酸i在MSA序列里面表达它是一个矩阵,具体来说,氨基酸i在每一个蛋白质里面,它是用一个长为cm的象征来表示的,然后这里一共有s条蛋白质,所以呢就表示成一个s乘以cm的一个矩阵来表示i这个氨基酸,同样道理对于氨基酸j来说,也是一个s乘以cm的矩阵来表示
-
在对表示里面,从氨基酸i连接到氨基酸j的这一个表示是一个长为cm的向量,那么意味着说我们要把两个矩阵转换成一条向量
-
具体来说,将这两个矩阵投影到c维,就得到一个s*c的矩阵,然后做外积得到(s,c,c)的张量,可认为是他先把这个向量加一个维度变成一个(s,c,1)的一个三维的张量,下面变成(s,1,c)一个三维的张量,这两个张量做batch 的dot,得到结果是(s,c,c)的张量,最后两个维度c c表示的是氨基酸i和氨基酸j的一些交互的信息,s表示的是这里一共有s个蛋白质
-
但我们关心的是氨基酸i到氨基酸j的信息,所以下面他取了一个均值,在s的这个维度上就得到一个c * c 的矩阵了,最后把这个矩阵拉直,然后再投影到一个长为cm的维度,就会得到一个长为cm的一个向量,把这个向量加回到最右边那里
这样就把来自MSA里面的序列的信息加入到了对表示里面,就完成了这一个信息的融合的操作
-
-
对氨基酸对之间的关系做一些建模
后面红色的两块跟前面对空间信息建模的红色两块是很像的
-
Triangle self-attention around starting node
跟前面很像,走势很像
-
输入只有$z_{ij}$,氨基酸对的表示,已经不再有前面的序列的表示了,仍然是一个自注意力的一个函数,函数有三个部分,一部分是对输入做投影,然后计算注意力分数和你的注意力输出,最后把每一个头之间全部连接起来,投影得到最终的输出
-
要计算的是$z_{ij}$的输出,那么就是第$i$行里面第$j$个元素,然后我们把它对应的这个query,$q_{ij}$跟所有这一行里面就是第$i$行里面别的那些元素对应的key,然后做内积,就可以得到自注意力的一个分数
同样跟之前一样加上一个偏移,这个偏移是$b_{jk}$,在每一行里面做softmax,得到结果
-
这个公式跟之前是一模一样的 ,就是物理意义不太一样
设计的是三个氨基酸$i,j,k$,$q_{ij}$表示的是氨基酸$i$到氨基酸$j$的这一个信息,$k_{ik}$表示的是氨基酸$i$到氨基酸$k$的这一个信息,他们做内积表示的是这两个边之间的一个关系
从$j$到$k$的这个偏移,加上偏移的意思是说想让你的自注意力机制在计算相似度的时候去看这个三角的关系,就是$i$到$k$的这个距离应该是小于$i$到$j$再$j$到$k$的这个距离,这就是叫做triangular的原因,别的跟之前的就没有太多区别
-
得到注意力分数之后,跟他对应的那一个v相乘,然后在每一行就是遍历k,所以相加得到我的自注意力的输出,然后再乘以我的门得到我最终的输出,最后的话我们就是把它连接在一起投影下去就行了
-
-
Triangle self-attention around ending node
-
之前是按行做自注意力,现在变成了按列做自注意力,跟之前下标不一样
$q_{ij}$是第$j$列里面的第$i$个元素,我们需要把这个第$j$列里面所有的别的元素,所以$k$是从1到$r$里面全部遍历一遍,然后相乘,最后同样道理的话他也是做的是$b_{ki}$,$v_{kj}$,跟之前的相比就是按行做softmax,变成按列做softmax
-
所以经过红色这两个模块之后,在按行层面和按列层面都做上了信息的融合,但是有一个先做和后做的关系,导致他的结果是不对称的,所以导致第$ij$和第$ji$的向量是不相等的
但是没关系,因为我们用的是一个有向图,$ij$就表示氨基酸$i$在氨基酸$j$的前面,$ji$就表示氨基酸$j$在氨基酸$i$的前面,所以两个氨基酸在序列中的位置一调换,导致他们的关系发生变化也是没问题的
-
-
triangle update using outgoing edges
-
跟红色的一样,也是想对氨基酸对信息之间进行一些交互,但是这个模块的设计比后面的自注意力要简单很多,作者本来是想用计算量相对来说比较小的模块来替换掉自注意力,使得整个计算更简单一点
-
干这个事情是因为它使用的一个蛋白质序列里面氨基酸的那个个数大于你有序列的个数的,r>s,导致pair输入的大小比MSA输入大一些
-
主要的计算是来自于自注意力层的计算,所以如果能用便宜的模块来替代他的话,那么整个算法的计算量会降低
但是作者又发现如果只使用黄色这两个模块的话,那么精度不如在后面再加上红色这两个模块来的好
-
首先对所有输入做一个Layer Norm
-
然后对每一个对的向量$z_{ij}$,计算一个a和一个b的版本,具体的计算是先把$z_{ij}$投影到$c$就是128的维度,在ab那做了个门,就是把它做另外一次投影然后加一个sigmoid,控制的是相乘的那个Linear($z_{ij}$)里面哪一些元素我需要,哪些元素就阻止掉,不让他进去了
-
$g_{ij}$也是一个输入的门,$z_{ij}$等于这个门再加上$a_{ik}$和$b_{jk}$做点积,然后在$k$上做求和,做一个LayerNorm再做一个投影
具体看一下是什么意思,$z_{ij}$要计算的是氨基酸$i$到氨基酸$j$的这条边,他的更新是说,我们对所有的氨基酸$k$,把$a_{ik}$拿出来,对$a_{jk}$和$b_{jk}$两条边做一个按元素的点乘,然后用一个LayerNorm把它normaliz一下,然后$z_{ij}$的信息由所有的这些$ik$、$jk$这些对的信息汇总而来
-
-
triangle update using incoming edges
- 跟之前唯一的区别就是从$a_{ik}$变成了$a_{ki}$,$b_{jk}$变成了$b_{kj}$,i与k互换,j与k互换,意味着这个边,箭头从往外伸变成了往回伸,这就完成了一个对称性的一个交换
- 就是我既汇聚了出去的边的信息,也汇聚了回来的边的信息,所以这两个黄色的模块和后面红色的模块有一定的相似度,也就是对氨基酸对之间做一些消息的传递
-
MLP模块
具体来说,就是对于每一个元素里面的这个向量,都用同样一个MLP对他做一次投影,这就完成了对所有氨基酸对的表示的一个变化
-
-
-
decoder
-
蛋白质的3D结构
- 最简单的表达方式就是把它在3D空间中的坐标记录下来,那就是三个值,然后我们对每个原子去预测这三个值就行了
- 一个蛋白质你把它在空间中进行旋转好位移也好,是不会影响它的结构的,但是任何的旋转或者位移,都会导致3D空间的绝对位置发生变化,所以如果你有绝对位置的话,那么你对平移、旋转这样的变化是不友好的,所以他这里用到的方法是相对位置
-
具体来说可以认为一个蛋白质就是n个氨基酸串起来了,中间这条网络叫做主干网,边上就是一些枝
-
所谓的相对位置是说,下一个氨基酸相对于上一个氨基酸的位置而言是怎么样的,只要我们知道所有这种相邻位置是怎么变换,给我任何一个起始点,都能还原出这条主干网络,用的是欧几里得的变换或者叫做刚体变换
-
假设这个点在空间中的位置是x,x是一个长为3的向量,这个点的位置是y,那么怎么从x变换成y呢?
在欧几里得里面可以写成,y可以等于对x做旋转就是一个R,再加上一个偏移,这个地方R是一个33的一个矩阵,t是一个31的一个向量,只要能确定r和t的值的话,那就能从x还原到y
-
这里的好处是说,如果我们对这个蛋白质的整体的结构做旋转或位移,就是把一个全局的r或者t加在整个蛋白质上面的话,他不会影响这些局部的变换,就是导致x变成y的这个r和t不会发生变化,这样这个表示就是对于全局的刚体变换是无关的
-
确定了主干的结构之后,上面的每一个氨基酸的下面这些部分可以认为是在每一个可以变换的原子那个地方,我只要去关心他旋转的角度是什么就行了,只要头是固定的话,只要关注那些能够转的地方,到底转了多少就行了,这些变化还是要满足物理和生物学的定义的
-
比如说两个氨基酸只能通过一些特定的角度连在一起,你可以做的办法是说在训练的时候,我可以限制住我的这些r和t可以要选取的值,或者是我在训练完之后再做一步校验,验证是这个结构在物理上和生物上是可行的
-
decoder结构
-
输入:编码器的输出,包括了之前我们讨论的氨基酸对的那个表示和氨基酸序列的那个表示
这里只画了一条人类的蛋白质,因为我们只要对这一条进行预测就行了,别的那些来自其它动物的相似的序列我们就不用管了
-
上面来自于对的表示(Pair representation)和序列的表示(Single repr.),进入一个叫做IPA的模块之后,它的输出是氨基酸序列(Single repr.),但是其中每一个氨基酸的表示,应该含有了更多的位置的信息
-
接下来用紫色模块(Predict relative rotations and translations)来预测主干里面那些旋转和偏移,就是这个蛋白质的主干结构是串成一条这样子(Backbone frames)的形状
-
拿到主干的形状和之前的氨基酸表示的话,我们再去预测那些枝叶上面的旋转,最后拿到Backbone frames和上面的结构之后,就可以还原出整个蛋白质的3D结构,然后这个更新的序列信息会进入下一个块里面
-
对于主干里面的那些旋转和偏移,也会输入进入下一个块里面。所以在IPA这个模块里面,其实他不仅拿到了氨基酸对的信息和不断在更新的序列信息,以及在不断的去调整的主干的旋转和偏移的信息,再拿到三个信息之后再去上面做一步的调整,再做新的预测,然后跟之前的做叠加
所谓的叠加是说,我对一个刚体做变换的时候,旋转一点位移一点,然后我还可以在上面再旋转一点偏移一点,偏移一点你可以认为是在不断的调整直到我要到最后的形态
-
一开始,所有的氨基酸都认为在原点,那么在解码器里面经过8个这样子的blocks,逐渐的把所有的原点的氨基酸然后慢慢的变形成我们最后要的形状(3D结构),每一次里面做一些的结构的调整不用一步到位,所以这就是解码器块的一个整体的设计
-
这里面的所有解码器块是共享权重的,所有就导致整个解码器块特别像一个RNN的架构,输入进来的序列表示和他的主干的旋转和偏移都是他的隐藏状态,pair representation块可以认为是来自编码器的上下文,最后我们关心的是最后一个时刻他的输出作为我的最终的一个输出。结构跟LSTM太像了
-
IPA模块(Invariant point attention——不动点注意力)
首先是注意力机制,另外一个叫做不动点
-
注意力机制
-
输入:
$s_i$表示的是第$i$个氨基酸在序列里面的那个表示
$z_{ij}$表示的是氨基酸对从$i$到$j$的那个表示
$T_i$表示的是第$i$个氨基酸他的那一个变化就是旋转和位移的那个变换全部写在里面了,因为是个多头注意力,所有这个地方有头的信息,有意思的是这里有3个头
-
接下来跟之前一样对$s_i$做投影,得到他的$q,k,v$,有意思的是下面还有一些$q,k,v$出来
具体来说,$s_i$先算了一个$q$和一个$k$,又算了一个$v$,之所以分开写是因为头的个数不一样,而且看到是说每一个是一个三维的东西,三维就是表示$q,k,v$这个结构在3D中空间中的位置是相关的,然后跟之前一样是根据你的对信息算出一个偏移
-
注意力分数是怎么样计算的?
第7行,跟之前一样,我们还是来自于第$i$个氨基酸的$q$和第$j$个氨基酸的key做内积,再加上来自对表示里面的偏移
这个地方新加了一项,我们可以忽略掉这些常数,可以看到是说我们先把第$i$个氨基酸上的变换作用在$q_i$上面,再减去第$j$个氨基酸的变换作用在$k_j$上面,然后计算他的距离
$q_i$可以认为是表示的是第$i$个氨基酸当前的位置,那么$k_j$也表示的是第$j$个氨基酸他的位置,分别作用变换之后,中间的圈表示的把这个变换作用在这个向量上面,就得到了第$i$个氨基酸后的那一个氨基酸的位置和第$j$个氨基酸后的那个氨基酸的位置
如果他们后面两个氨基酸的位置相隔的比较远的话,因为前面有个负号的关系,就表示他们就不应该那么相似,因为之前是说,如果你们的夹角比较大的话,那表示相似,那么他的后一个氨基酸,如果我把起始点拉到附近的话,长得很不一样的话,那表示这两个氨基酸确实是不那么一样的
所以这个新加的一项跟前面的主要的区别是我们将这个变换了和位置的信息显式的建模了进来,最后我们计算一些输出
-
跟之前不一样的地方,在于我们这里(8)使用的$z_{ij}$就表示了第$i$个氨基酸到第$j$个氨基酸之间的关系,或者是说他的距离就是相对距离,然后(10)也使用了跟位置和变换相关的一个输出,-1表示的是变换回去,就是对每一个氨基酸,我们把它做完变换之后,然后(10)通过权重加起来,然后再反变换回去
因为$o_i^{hp}$这个东西,我们是要在上面做变换的,这样的话(11)把三个输出全部放在一起,还有$o_i^{hp}$的长度全部放在一起来计算$s_i$,那么这个时候你就认为你的$s_i$即你更新后的氨基酸的那一个表示里面含有了我的位置信息,这也是跟之前的编码器不一样的地方,在解码器里面,我们的向量里面显式的加入了位置的信息
-
-
不动点——点不变的
-
具体来说先看一下第7行后面平方的那一项就是计算距离的那一项
假设我们要对所有的氨基酸做一个全局的变化,叫做$T_{global}$的话,我们就是说对$q_i$这个点先做$T_i$的变换再做全局变换,同样对于$k_j$也是做先做$T_j$变换再做全局变换
-
因为变换是可以展开的,那么就等价于是说先对$q_i$做变换和对$k_j$做变换,然后相减,减完出的那个值再做全局变换,就把全局写在外面
-
因为全局变换我们知道在做偏移的时候,因为有个减的关系,所以偏移会被抵消掉,那么剩下来的就是一个旋转,但旋转是不改变那个向量的长度的,$T_{global}$这一项就被写没了
-
意思是说,在计算距离的时候,我们不管做什么样的全局变换,其实他都不影响这个值,因为有个减号的关系
-
此外在第10行要计算输出,算$o_i$,同样道理,如果我们要对整个氨基酸做了一个全局变换的话,那么在第10行后面$T_i$有一个全局变换,因为欧几里得变换是线性可以相加的,所以可以一直往前写出来,然后这里有一个逆变换,然后作为一个全局变换之后,做一个局部的变换,再回过来做一个全局的逆变换,这两个变换会被抵消掉,这是因为矩阵的乘法是可以交换的
所以不管是第7行计算相似度也好,还是第10行计算输出也好,我对整个蛋白质做一个全局的欧几里得变换是不会影响到我整个的结果的,所以这就是IPA里面的IP是来自于什么意思
-
-
-
预测模块(紫色的和绿色的)
-
预测模块比较简单,因为每一个氨基酸的向量表示里面已经含有了位置信息,那么要去预测他的实际的位置的时候,我们就做一个线性的投影层就可以得到我们的输出了
当然这个地方稍微复杂一点,是说我们的预测是有物理意义的,所以在结构上有需求,而不是我们之前预测一个类别,基本上任何输出都是行的
-
以紫色模块(Predict relative rotations and translations)为例,分析一下这个物理信息是怎么样放进来的
-
名字叫做主干更新,具体要干的事情是说,对第$i$个氨基酸的向量表示,我们去预测它对应的那一个变换,变换$T_i$里面有两个东西,一个是他的33旋转矩阵,另一个是一个31的向量,对于偏移来说我们没有任何的要求,怎么偏都可以,但是对一个旋转矩阵来说,是一个正交矩阵,而且他的Norm是等于1的,所以在做预测的时候,他不是去直接预测$R_i$里面的12个值,而是说把$s_i$投影到一个6维里面(1),其中后面的三维$t_i$直接作为我的偏移,然后前面的三维($b_i,c_i,d_i$)用来构造旋转矩阵,虽然旋转矩阵有12个元素,但是其实就有3个值可以决定
具体来说,通过下面(3)这个变换来得到一个合法的单位旋转矩阵,这样子满足我们物理上的要求,这就是他怎么样做主干变换的预测的
-
-
-
-
-
训练中的细节
- 主损失函数:FAPE(Frame aligned point error),大概意思是说,根据我预测出来的那些变换,然后我就能把整个蛋白质的结构还原出来
- 绿色是表示预测出来的结构,灰色表示真实的结构,然后就知道对应的原子在真实中的位置和预测出来的位置,然后把这两个距离一减就得到我的损失,距离越大损失越大,距离越小就损失越小
使用了两个额外的技术
-
使用没有标号的数据
- 有标号了的数据全部在PDB里面,也就是已知结构的所有的蛋白质,但是还是觉得不够
- 用了一个方法叫做噪音的学生自蒸馏方法,核心思想是,现在在有标号的数据上训练一个模型,然后我用这个模型去预测一个大的没有标号的数据集,然后把那些比较置信的标号拿出来,跟之前的有标号的数据集一起拼成一个更大的一个数据,这样子我们可以在上面重新再训练一个模型出来,这样子就等于是说我的数据集变大了,然后可以重复做很多次
- 这里一个核心的关键点是说你要加噪音进去,不然的话如果我之前的一个模型预测一个样本预测错了,而且我又特别置信,那么这个错的标号将会进入到你的训练样本,使得你下一次训练的时候,在这个错上可能会加错,但是你加入噪音比如说大量的数据增强,甚至是把标号做一些改变的话,那么你的模型就比较好的处理这些不正确的标号
- 具体来说,他从一个新的数据集里面找到了35万个不是很一样的训练,他先在PDB上训练一个模型,然后把这些上面预测的置信度比较高的那一些蛋白质序列拿出来跟PDB一起做成一个新的数据集再重新开始训练,这个地方模型架构并没有变化,但是多训练几次
-
来自BERT
在蛋白质序列里面随机遮住一些氨基酸,或者甚至是把一些氨基酸把它做一些变换,然后像BERT去预测这些被遮住的氨基酸,然后他发现在训练的时候同时加入这个任务的话,整个模型对整个序列的建模上更加好一点
训练参数:
序列长度:256
batch size : 128
具体是用128个TPU v3训练,从随机初始训练的话,大概是一个星期,如果你做微调的话还需要额外的4天
最大的问题是内存不够,因为整个模型对内存的占用率是非常非常高的,大概需要几百GB的内存的样子,但是一个TPU v3也就是几十GB的样子
具体的做法也是个常见的做法,你的内存基本来自于误差反传的时候,因为你要算出来中间的那些结果,结果给你存下来,后面算梯度时候要用,存的那些结果跟网络的深度成正比
可以做的是说,你把中间一些算的比较快的结果丢掉,然后再算梯度的时候,发现你要的时候再去重新算一遍,这样就是用计算来换空间,使得计算也是可以往前走的
预测的一些性能,而预测性能跟你的蛋白质的长度是相关的
如果你的蛋白质里面有256个氨基酸的话,那么你在一个单卡V100的卡上面,大概是5分钟的样子,如果有384个氨基酸的话,那么是9分钟,2500个的话,那么是18个小时
结果的分析
- 主要讲的是一些消融实验,主要的结果是放在图四,左边是CASP14就是竞赛中的结果,右边是PDB也就是他用来训练模型的测试数据集上的结果。
- 灰线,0表示的是他的基线,右边比0大的话表示比基线要好,如果在灰线的左边的话表示比0小,那么他就比基线要差。其中每一行表示的是一个方法,基线的结果大概就是在这个灰线的样子
- 使用自蒸馏会怎么样,也就是使用额外的没有标注的训练数据集对他进行训练,在竞赛上的效果好一点,但是在PDB这个数据集上还是高了很多,很有可能是说PDB这个数据集他的多样性更大一点,他的测试数据更大一点
- 去掉一些模块会怎么样,前面的这些都是去掉一些输入的数据,比如说把模板信息去掉,损失那么一点,如果把直方图去掉,也损失一点
- 如果把原始的MSA但是用一个替代来去掉的话,也会损失掉一些
- 接下来是对模型的一些变换,比如说把IPA去掉,意味着是说你在解码器的时候,我不再使用一个注意力机制,把你的位置信息放进你的氨基酸的向量里面,可以看到会损失一些,但是在PDB好像损失的并不大
- 如果没有使用BERT那种带掩码的机制会怎么样,在竞赛数据集上还好,但是在PDB这个数据集上,损失还是比较大的
- 如果你不做回收,不做那四次回收的话,那损失是比较大的
- 如果你使用最简单的注意力机制,就是说我先按行做标准的自注意力,然后再按列做标准的自注意力,而不使用三角更新或者是用对信息来做偏差或者使用门,可以看到的话,这里差的还是比较大的,在PDB的数据上差的更多一些,或者说如果不做端对端的结构梯度,我理解就是你的编码器不参与计算梯度那损失更多了
- 最后是你不做IPA也不做回收,那损失就是相当的明显
- 所以这个图想表达的核心思想是说,我这个网络虽然复杂,但是里面没有一块能去掉,把所有块加起来都是有好处的,但反过来讲这个模型还是相当复杂,虽然这个地方已经做了很多消融实验了,但是你很难从系统的对每一个模块看一下,他到底是不是重要的,因为里面还是有很多模块没有逐一检查到的
-
我们的编码器里面有48块,然后我们做了4次回收,所以等价于一共做了192个块
-
如果我给你一个蛋白质,只使用前面的一些块,效果会怎么样,y轴表示的是你的好坏越往高越好,可以看到如果是两个简单蛋白质的话,基本上到了48就差不多了,也就是说你不用做回收也没关系
如果你的编码器里面有48个块,那其实这个地方精度就差不多了,但是对于比较复杂的话,48个还是不够的,甚至做到192还是有上升的趋势,说不定可以做的更深一点
评论
-
文章的写作
这篇文章算比较短的,虽然页数不少,但是你把图去掉之后,他的文字的部分并不多,话特别经典,每一句话都是在讲件事情,但是这个地方的网络其实比Transformer那个地方可能还要复杂个几倍的样子,而且他的大量的篇幅都在介绍问题和对自己结果的歌颂上面,所以以至于是讲模型的部分是非常非常少,读正文读很多遍你也不知道他模型的细节在什么地方,以至于你要读它提供的补充材料,才能知道里面的细节是什么样子,或者去读源代码
可以回过去看一下他的正文是怎么写的,学习一下对一个比较复杂的算法,你怎么样用相对来说比较简单的篇幅去对他进行介绍,在结果上不需要做太多的评论,因为整篇文章的主要的卖点是我的结果非常好
-
模型
这个模型真的是比较复杂的,细节很多,为什么想出了那么多的模块,可以用来提升精度,是他的写作的方法造成的一个误解,他的正文里面没有讲太多细节,基本上就是一个很简单的一个描述,然后大概讲了一下我的一些贡献,有一个非常简短的相关工作,然后就直接跳到了补充材料,基本上是代码的一个伪代码版本,只讲了具体的实验细节,究极融合怪
你就说这是一个什么的问题,我们提出把其中一个东西改掉之后,就解决了这个问题,一句话就可以讲清楚的结论甚至可以放进标题