MAPE-PPI：通过微环境感知蛋白质嵌入实现有效且高效的蛋白质-蛋白质相互作用预测

https://github.com/LirongWu/MAPE-PPI

PPI是各种生物过程的基础，在生命活动中具有关键作用，现有PPI预测方法很大程度依赖于蛋白质序列，但是结构才是决定相互作用的关键因素。本研究将两种模式都考虑在内。通过氨基酸结构和序列的上下文来定义氨基酸的微环境。其中微环境描述了氨基酸的周围的化学特征和几何特征。以往工作定义的微环境主要基于实验测定的物理化学性质，很难涵盖微环境的多样性和复杂性。本文提出来微环境感知蛋白质嵌入PPI预测，通过大量微环境‘词汇’，将微环境编码成具有化学意义的离散代码。使用一个新的预训练策略（Masked Codebook Modeling）MCM,通过随机掩码本（masking the codebook）和重构输入来捕捉不同微环境之间的依赖关系。

研究背景

蛋白质相互作用是所有生物必不可少的特定生物功能。对于医学、制药和遗传学研究都很重要。过去的几十年，已经提出了一些高通量PPI预测的实验方法（酵母双交杂筛选、质谱学蛋白质复合体鉴定），然而基于湿实验的方法是非常耗时和昂贵的。
人工智能的发展，PPI预测由机器学习向深度学习转变，现有基于深度学习的方法可以分为两类：基于序列和基于结构的方法。其中基于序列的方法通过CNN、RNN、Transformer等从氨基酸序列中提取蛋白质表示，直接根据表示来预测他们的相互作用。基于序列的方法不能具有好的预测效果。然而决定蛋白质功能和与其他蛋白质相互作用的是蛋白质的结构，基于结构的方法通过使用图神经网络来建模蛋白质2D或者3D结构，但是基于结构的方法会产生沉重的计算负担。
自然界中有20种氨基酸类型，但是不同类型的残基分布并不均匀，并且微环境不同，同一类的残基也可能表现出不同的理化性质，不同残基的微环境也可能相似或者重叠。如果两个残基的微环境高度相似，那么在训练时期对他们分别进行编码将是多余的。
如果我们能学到记录那些最常见微环境模式的“词汇表”(即代码簿)，可以把它作为一种现成的工具，将蛋白质编码成微环境感知的嵌入。因此，通过将端到端的PPI预测分解为学习微环境码本和基于码本的蛋白质编码的两个阶段，可以提高训练效率。
以前有一些工作将目标残基周围的序列和结构上下文定义为其微环境，并根据其物理化学性质将其分类为十几种类型，这构成了一个特殊的微环境“词汇”。与语言中丰富的词汇不同，实验微环境的词汇通常是粗粒度的，大小有限。因此，为了学习具有丰富信息的细粒度词汇，我们提出了一种用于微环境感知的蛋白质嵌入的VQ-VAE(Van Den Oord等人，2017)变体，它包括两个方面的设计：(1)微环境发现：学习大而细粒度的微环境词汇(即码本)；(2)微环境嵌入：使用码本将微环境编码成具有化学意义的离散代码。此外，我们还提出了一种新的码本预训练任务，即掩蔽码本建模(MASMASED CODEDBOOK Model，MCM)，通过随机掩蔽码本而不是输入特征或离散代码来捕捉不同微环境之间的依赖关系。有了学习的微环境码本，我们可以把它作为一个现成的工具，将蛋白质编码成微环境感知的蛋白质嵌入。得到的嵌入可以进一步用作大规模PPI图网络上的节点特征，然后由GNN对其进行编码以预测PPI。

文章的贡献：

1、定义了基于序列和结构contexts的残基的微环境，有效编码蛋白质的序列和结构
2、提出VQ-VAE的变体，来发现细粒度的微环境，并学习微环境感知的蛋白质嵌入
3、提出Masked Codebook Modeling，直接屏蔽码本，捕获不同微环境的依赖关系

Methods

异构蛋白质图的构建

设一个具有$M$个氨基酸的蛋白质$P_i\in P$,可以表示为序列$S_i=(v_1,v_2,.....,v_M)$,其中$v_M\in S_i$表示氨基酸类型。其次氨基酸序列又可以折叠成稳定的3D结构。文中将蛋白质$P_i$的结构和序列表示胃残基水平的蛋白质异构图$G_P^{(i)}=(V_P^{(i)},E_P^{(i)},R)$,其中$V_P^{(i)}$表示$M$个氨基酸的集合，$E_P^{(i)}$表示连接残基边的集合，文中采用三种不同的序列和结构边类型：
（1） 序列边：两个残基按照序列连接。即：序列中相邻氨基酸进行连接
（2） 半径边：设置阈值，空间中两个残基的欧氏距离小于阈值$d_r$则相连
（3） k近邻边：两个残基在空间中是邻居，则相连，邻居数目设置为$K$

残基级微环境定义

残基微环境基于序列和结构的背景描述了他周围的化学和几何特征，每个残基$v_m$的微环境$E_m$定义为蛋白质图$G_P^{(i)}$的子图。其微环境中的节点集$V_{E_m}\subseteq V_P^i$:
$V_{E_m}=\left\{v_n||m-n\mid \le d_s,\Vert C{\alpha }_m-C{\alpha }_n\Vert \le d_r,v_n\in {\mathcal{N}}_m^{(K)}\right\}$
文中主要学习一个微环境码本$A$和映射$\tau$：$E_m⟶A$,实现端到端的微环境发现和蛋白质嵌入。

微环境感知蛋白质嵌入

异构蛋白质编码

再给定不同边类型$R$的蛋白质异构图$G_P=(V_P,E_P)$,使用异构图神经网络座位蛋白质的编码器，每个边类型$r\in R$对应的权重矩阵$W_r$。异构蛋白质编码器有每个编码层处的每个节点$v_m$的三个关键计算组成。
（1）聚集节点$v_m$的对应不同边$r\in R$的来自领居的信息 N r ( m ) = { v n ∣ e m , n = r } N_r(m)=\left \{v_n | e_{m,n}=r\right \} Nr​(m)={vn​∣em,n​=r}，（这里的邻居是指来自于不同边的邻居，对邻居信息进行求和，参考CNN卷积操作）
（2）对每个边的类型$r$进行线性变换$W_r$
（3）通过线性变换$W_h$对来自不同边类型的邻域的更新消息求和
对于L层的异构蛋白编码器，第$l$的公式定义如下：
${\mathbf{h}}_m^{(l)}=\mathrm{BN}\left(\mathrm{ReLU}\left({\mathbf{W}}_h^{(l)}\cdot {\sum }_{r\in \mathcal{R}}{\mathbf{W}}_r^{(l)}{\sum }_{v_n\in {\mathcal{N}}_r(m)}{\mathbf{h}}_n^{(l-1)}\right)\right)\text{, for }1\le l\le L$

微环境码本的发现与嵌入

异构蛋白质编码器通过聚集来自不同邻域的消息来模拟残基的微环境，即，顺序和结构上下文。得到$h_m^{(i)}$包含关于微环境$E_m$的所有信息。
给定一个微环境码本$A=\left\{e_1,e_2,...,e_{|A|}\right\}\in R^{|A|\times F}$,蛋白质的微环境 { E 1 , E 2 , . . . , E M } \left\{E_1,E_2,...,E_M \right\} {E1​,E2​,...,EM​}通过在码本$A$中查找每个嵌入$h_m^{(L)}$的最近邻居，通过矢量量化被标记为离散代码 { z 1 , . . . . , z M } \left\{z_1,....,z_M \right\} {z1​,....,zM​},
$z_m={\mathrm{argmin}}_n{\Vert {\mathbf{h}}_m^{(L)}-{\mathbf{e}}_n\Vert }_2，\text{ where }1\le m\le M\text{.}$

下面聚焦如何构建和学习一个合适的微环境码本$A$,一个自然的想法是使用实验测定的微环境的分类来建设码本，但是实验测定的成本是昂贵的，仅仅具有十几种可能类别的码本太小并且具有粗粒度。
本研究根据VQVAE的思想，将码本$A$直接参数化为可学习的变量，依托数据重构的训练任务，同时学习码本$A$和训练异构蛋白质编码器。蛋白质训练目标为：
${\mathcal{L}}_{\mathrm{VQ}}=\frac{1}{M}{\sum }_{m=1}^M\left({\Vert {\mathbf{x}}_m-{\hat{\mathbf{x}}}_m\Vert }_2^2+{\Vert \mathrm{sg}\left[{\mathbf{h}}_m^{(L)}\right]-{\mathbf{e}}_{z_m}\Vert }_2^2+\beta {\Vert {\mathbf{h}}_m^{(L)}-\mathrm{sg}\left[{\mathbf{e}}_{z_m}\right]\Vert }_2^2\right),$
包括重构损失，训练码本（embedding）、训练编码器。sg[]为停止梯度操作

掩码码本建模

文中直接对码本$A$进行掩码，不是对输入特征 { x m } m = 1 M \left\{ x_m \right\}_{m=1}^M {xm​}m=1M​或者hidden code { z m } m = 1 M \left\{z_m \right\}_{m=1}^M {zm​}m=1M​.背后的动机是，在序列或结构上相邻的两个残基的微环境中通常存在重叠（依赖性）。为了在码本中捕获这种微环境依赖性，我们提出了Masked Codebook Modeling（MCM）来掩蔽码本A，然后基于未掩蔽的代码重建输入，旨在学习更具表达力的码本A，而不是更强大的蛋白质编码器。
对码本$A$进行采样，并使用掩码向量Mask进行掩码，然后会得到新的微环境${\tilde{e}}_m$被定义为：
${\tilde{\mathbf{e}}}_m=\{\begin{array}{ll}{\mathbf{e}}_{[\mathcal{M}]}& {\mathbf{e}}_m\in \mathcal{M}\\ {\mathbf{e}}_m& {\mathbf{e}}_m\notin \mathcal{M}\end{array}.\text{ where }1\le m\le M$
我们基于掩蔽码本$\tilde{A}$而不是码本A将每个离散码${\tilde{z}}_m$映射回相应的嵌入${\tilde{e}}_{zm}$，并训练蛋白质解码器将嵌入${\tilde{e}}_{zm}$重构为${\tilde{x}}_m$。则掩码码本建模的损失函数：
L M C M = 1 ∣ V [ M ] ∣ ∑ v m ∈ V [ M ] ( 1 − x m ⊤ x ~ m ∥ x m ∥ ⋅ ∥ x ~ m ∥ ) γ ,  where  V [ M ] = { v m ∣ e ~ z m ∈ M } \mathcal{L}_{\mathrm{MCM}}=\frac{1}{\left|\mathcal{V}_{[\mathcal{M}]}\right|} \sum_{v_{m} \in \mathcal{V}_{[\mathcal{M}]}}\left(1-\frac{\mathbf{x}_{m}^{\top} \widetilde{\mathbf{x}}_{m}}{\left\|\mathbf{x}_{m}\right\| \cdot\left\|\widetilde{\mathbf{x}}_{m}\right\|}\right)^{\gamma}, \text { where } \mathcal{V}_{[\mathcal{M}]}=\left\{v_{m} \mid \widetilde{\mathbf{e}}_{z_{m}} \in \mathcal{M}\right\} LMCM​=∣V[M]​∣1​∑vm​∈V[M]​​(1−∥xm​∥⋅∥x m​∥xm⊤​x m​​)γ, where V[M]​={vm​∣e zm​​∈M}
损失函数是为了计算重构特征与原特征是否相似，通过将余弦误差以$\gamma \ge 1$ 的幂数进行缩放，来降低容易学习的样本(easy samples)在训练中的贡献。对于高置信度的预测，其相应的余弦误差通常小于1。这样一来，当缩放因子(scaling factor) 时，余弦误差会更快地衰减为零。