AlphaFold 1:用距离关系锚定空间坐标

蛋白质预测任务

我们有一个链,链条上有不同的小零件。把链条放进水里一泡,它便会自己组装成一个三维结构——这条链叫肽链,而链条的组装就叫蛋白质的折叠。小零件的顺序几乎完全决定了三维结构,所以根据零件顺序预测三维结构便成了避不开的话题。

三维结构怎么表示呢?一般使用分子坐标,即,一个预测要给出所有非氢原子的三维笛卡尔坐标和对该位置的置信度估计。

核心思想

AlphaFold 1 认为,如果把链上的这一个个氨基酸残基当成刚性的小零件,那么对于每一个氨基酸而言,只要给定它的坐标,那么三维结构就解析了出来,至于侧链原子的预测就简单不少。那我们就先只考虑怎么把一条肽链堆起来组装成一个蛋白质。

想象一个三维空间,这个空间里有两个小球,其中一号小球被固定在原点位置。那么,这两个小球的自由度是多少呢?是三,因为一号小球被钉死了位置,而二号小球可以上下左右前后移动,对吧?那如果一号小球和二号小球之间连了一根长为一的杆子,自由度又是多少呢?这时就从三变成了二,因为两个小球之间的距离被固定了下来,所以自由度减少了一。

肽链本身就给了类似的空间信号。理所当然的,甚至不用理解任何化学或生物学知识,三号位的氨基酸残基和四号位的残基肯定离得很近,而二号位的残基和一百号位的残基,恐怕就不太可能离得这么近。如果有 n+1n+1 个氨基酸,其中一号氨基酸被固定在原点,那么整个系统的自由度就是 3n3n;但如果它们被组装成一条肽链,那么就会因为形成的 nn 个肽键下降 nn 个自由度。

然而,3nn=2n3n-n=2n 个自由度对于预测任务来说还是太大了。有没有方法能让自由度下降得更多呢?那么,一个自然的想法就出现了:对于任意一个氨基酸残基而言,如果它得知了自己与其他氨基酸之间的距离,那么它就锚定了自己在整体三维结构中的位置。也就是说,如果我们得知了任意两个氨基酸之间的距离关系,我们就有很大可能把三维结构还原出来。换而言之,n(n1)2\frac{n(n-1)}{2} 根“键”(当然不是真实存在的那种)对上 3n3n 个自由度,简直是手到擒来。

或许你要问,当 nn 比较大的时候,n(n1)2\frac{n(n-1)}{2} 个位置关系从数量级上不是可以完全碾压 3n3n 吗?但是其一,对氨基酸两两距离的预测本来就是“大概率如何”“小概率如何”,它们之间并没有一根正儿八经的实体存在长度不变的“键”。其二,即使是长度固定的键,也不能说一个键就能消掉一个自由度。比如说正方形的四个端点两两连线,那连接四边和一条对角线就足以确定形状,再连一条对角线也没有任何作用了。

具体实现

那么,应该怎么得知氨基酸残基之间的距离情况呢?

输出

先想输出。现在我们需要知道氨基酸与氨基酸之间的距离。假设链长度为 LL,那么两两氨基酸之间的距离就可以用一张 L×LL \times L 的矩阵表示。考虑到我们最后还要做置信度估计,所以最好也给出模型对这个距离的置信度。

然而,如果只给出两个矩阵,其中一个填距离,另一个填置信度,就免不了损失诸如「其他的距离置信度有多高?」「用别的距离会不会更好?」之类的许多信息。AlphaFold 1 则通过把 222A˚2-22 \AA 划分为 64 个等宽区间,输出 P(dijBink),k1,2,,64P(d_{ij} \in \text{Bin}_k),\quad k \in 1, 2, \cdots, 64 来解决了这个问题。因此,AlphaFold 1 的输出就是 64 个 L×LL \times L 的矩阵。

输入

L×LL \times L 个像素,64 个通道——这实在太像 CNN 了,所以 AlphaFold 1 就用了 CNN。CNN 的输入也是多通道图片,所以我们不难构建出 AlphaFold 1 的输入:许多个 L×LL \times L 的矩阵。每个矩阵对应一个特征工程。

AlphaFold 1 安排的特征工程既有单残基特征,又有残基对特征。对于位置 ii,有当前氨基酸类型、同源序列中该位置出现氨基酸的频率、插入缺失等等信息;对于位置对 (i,j)(i, j),有两个位置是否协同变化、共变参数等等信息。而在 (i,j)(i, j) 位置,则把特征拼起来:

Xi,j=[位置i的特征, 位置j的特征, 残基对(i,j)的特征]X_{i, j}=[\text{位置} i \text{的特征},\ \text{位置} j \text{的特征},\ \text{残基对} (i, j) \text{的特征}]

于是形成 XRL×L×CX \in \mathbb{R}^{L \times L \times C},其中 CC 是特征通道数。

CNN

我们知道,CNN 中最重要的设计就是卷积。比较常见的卷积手段有只提取特征而保持提取后像素不变的,也有通过拉大步长使像素变少的。无论如何,它们通常都由几个紧挨着的方格组成。然而,这个问题中,CNN 不能只捕捉局部特征。(45,101)(45, 101) 的氨基酸残基对距离固然绝对与 44 号或 100 号氨基酸有关系,但它们同样有可能与 30 号或 120 号氨基酸有关——当我们试图用一维序列推知三维的情况时,一维中的“邻近”便没有那么单一地可靠了。

所以,AlphaFold 1 在传统 3×33 \times 3 卷积核以外,还设计了 (i±2,j±2)(i \pm 2, j \pm 2)(i±4,j±4)(i \pm 4, j \pm 4)(i±8,j±8)(i \pm 8, j \pm 8) 的“\boxdot”状卷积核。这样,不需要特别大的卷积核,也能快速扩大感受野。

简化后的网络可以表示为:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
 L × L × C 输入特征


二维残差卷积块,膨胀率 1


二维残差卷积块,膨胀率 2


二维残差卷积块,膨胀率 4


二维残差卷积块,膨胀率 8


重复很多次


1 × 1 卷积


L × L × 64 距离 logits

又因为 AlphaFold 1 中的 CNN 并不需要进行扩大或缩小图像的操作,于是每一层都可以使用残差连接,保证即使是深层网络也可以有效训练。

Loss

loss=logpi,j,yi,j\text{loss} = - \log p_{i, j, y_{i, j}}

例如,真实距离为 7.4A˚7.4 \AA,对应第十七个分箱,那就只惩罚网络没有给第 17 类足够高的概率:loss=logpi,j,17\text{loss} = - \log p_{i, j, 17}

另外,考虑到距离矩阵应该满足 Pi,j=Pj,iP_{i, j} = P_{j, i},通常用只计算上三角,预测时再镜像的方法解决这个问题。

从距离到坐标

如果我们得到的是精确距离,那么通常会使用多维尺度分析的方式重建三维结构。然而,AlphaFold 1 给出的是分箱后的距离概率分布,所以实际任务应该是在此基础上使衡量三维结构合理性的“能量”函数最小。

完整的能量为

Etotal=Edist+λtorsionEtorsion+λclashEclash+λbondEbond+λangleEangleE_\text{total} = E_\text{dist} + \lambda_\text{torsion} E_\text{torsion} + \lambda_\text{clash} E_\text{clash} + \lambda_\text{bond} E_\text{bond} + \lambda_\text{angle} E_\text{angle}

分别是距离约束、扭转角约束、排斥碰撞项、共价几何项、键角约束项。

距离约束让整体折叠尽可能贴近预测的残基间距离,扭转角约束通过网络预测的 ϕ\phiψ\psi 让局部主链构象合理,排斥碰撞项避免原子重叠,共价几何项约束键长、键角和平面性,键角约束项通过控制键角避免极端尖角或接近直线。

ϕ\phiψ\psi

等等,这两个东西刚刚可没提到——从哪儿蹦出来的?

其实这两个也是用 CNN 预测出来的。ϕ\phi 表示着残基绕 NCαN-C_\alpha 旋转的角度,ψ\psi 表示着残基绕 CαCC_\alpha-C 旋转的角度。但是它对我们理解 AlphaFold 1 「距离建模」的核心思想助益不大,反而更像工程上的改进,所以放到这里讲了。

这个 CNN 的输入依然是那个 L×L×CL \times L \times C 的三维张量,而中间层得到的也是一个 L×L×CL \times L \times C' 的三维张量。距离是对两个氨基酸残基而言的,但它们都是对于一个氨基酸残基而言,所以需要对 L×L×CL \times L \times C' 的中间层做处理。考虑把 L×L×CL \times L \times C' 的三维张量压缩成 L×CL \times C',只需要把代表 LL 的某一维池化处理就解决了。

实际操作中,既会进行平均池化又会进行最大池化,既会沿第一维池化又会沿第二维池化,所以得到的矩阵大约是拼接后得到的 L×4CL \times 4C'

AlphaFold 1 原文中将 ϕ\phiψ\psi 各自 1010 ^\circ 一分,每个角度分 36 片,交叉起来就是 1296 个类别。(预测联合分布是因为 ϕ\phiψ\psi 强相关。)然后我们只需要做和之前的 CNN 差不多的操作,即把 L×4CL \times 4C' 通过 1×11 \times 1 卷积转化成 L×1296L \times 1296 的概率矩阵就可以。

工程细节

尽管论文的主体是 CNN 预测距离关系和最优化 EtotalE_\text{total},但是为了达到 CASP13 的高分数,AlphaFold 1 在工程上也做了不少努力。这里将挑选一两个有趣的工程实现讲一讲。

榨取 MSA 信息

CNN 的输入远不只是氨基酸序列。系统使用 HHblits 和 PSI-BLAST 搜索同源序列,提取:

  • 多种位置特异性 profile
  • 缺口和 deletion 统计
  • MSA 深度
  • Potts 模型偏置
  • 每对残基的 484 个 Potts 耦合参数
  • 耦合参数的 Frobenius 范数

当时很多方法只把 Potts 参数压缩成一个 Frobenius 范数,而 AlphaFold 1 则把原始参数也交给 CNN,让网络自行判断哪些氨基酸替换关系有意义。

为适应浅 MSA,训练时还随机抽取 HHblits MSA 中一半的序列,再重新计算 profile 和共变特征。AlphaFold 1 让每个训练结构域先产生 10 组这样的 MSA 样本,既达成了数据增强,也使网络不至于过度依赖“同源序列特别丰富”的输入。

参考态校正

AlphaFold 1 另外训练了一个不依赖序列和 MSA 的背景模型。它只利用蛋白质长度,以及该距离涉及 CαC_\alpha 还是 CβC_\beta 等背景信息,输出的是 pi,j,krefp^\text{ref}_{i, j, k}。在每个距离箱上都构造能量值 ei,j,k=logpi,j,kpi,j,krefe_{i, j, k} = - \log \frac{p_{i, j, k}}{p^\text{ref}_{i, j, k}},把“常见距离”转换成“这个序列特异地支持哪些距离”。

离散分箱可微化

为了把 L×L×64L \times L \times 64 的距离分布图和 L×1296L \times 1296 的角度分布图变成合理的三维结构,我们要最优化一个能量函数。我们希望用梯度下降的方法搞定,所以当然需要把这些离散的分箱变成有导数的概率密度曲线。

距离分布图的可微化使用的是拟合三次样条,即每个分箱都用中心 ckc_k 做代表,得到 64 个点 (ck,ei,j,k),k=1,,64(c_k, e_{i, j, k}),\quad k = 1, \cdots, 64,然后用三次函数曲线拟合。

而角度比距离多出一个特殊问题:周期性。比如 179179^\circ179-179^\circ,实际上只差 22^\circ,不加处理的话却很容易将它们视作相差 358358^\circ。所以,AlphaFold 1 将角度拟合成单峰 von Mises 分布。它类似高斯分布,不过定义了在圆周上。这使得它的数学形式天然满足周期性,正好用来解决这个问题。