ARES( Atomic Rotationally Equivariant Scorer )

Geometric deep learning of RNA structure

Year: 2021
Authors: Raphael J. L. Townshend, Stephan Eismann, Andrew M. Watkins, Ramya Rangan, Maria Karelina, Rhiju Das, Ron O. Dror
Journal Name: Science

Method

网络结构如下图所示

Equivariant convolution

等变卷积基于滤波函数来进行节点的特征更新，同时对于输入的平移和旋转需具有等变性。对于给定的原子 $a$ （源原子），等变卷积对每个在其邻域内的原子 $b$ （邻原子），应用滤波函数来更新源原子特征。定义 ${\vec{r}}_{ab}$ 为 $a$ 和 $b$ 的三维距离向量。滤波函数包括 Radial 函数和 Angular 函数两部分。

Radial functions

辐射函数本质上是输入为 $G(r_{ab})$ 的 dense neural network ， $r_{ab}=\Vert {\vec{r}}_{ab}\Vert$
G ( r a b ) = [ G 0 ( r a b ) , G 1 ( r a b ) , . . . , G n ( r a b ) ] G j ( r a b ) = 1 2 π σ e x p { − ( r a b − μ j ) 2 2 σ 2 } G(r_{ab}) = [G_0(r_{ab}), G_1(r_{ab}), ..., G_n(r_{ab})] \\ G_j(r_{ab}) = \frac{1}{\sqrt{2\pi \sigma}} exp\{ -\frac{(r_{ab} - \mu_j)^2}{2\sigma^2} \} G(rab​)=[G0​(rab​),G1​(rab​),...,Gn​(rab​)]Gj​(rab​)=2πσ ​1​exp{−2σ2(rab​−μj​)2​}

其中， $\sigma =1,n=11,{\mu }_j=\frac{12}{11}j$ 。神将网络含有一包含 12 个神经元的隐层，输出 $E+1$ 个神经元
$$[R_0(r_{ab}),R_1(r_{ab}),...,R_E(r_{ab})]=Dense(G(r_{ab}))$$

因为输入为两原子之间的距离，所以对于输入的平移和旋转具有不变性。

Angular functions

使用球谐函数 $Y$ 作为角函数，球谐波按其角分辨率 $l\in {\mathbb{Z}}_0^{+}$ 分组，称之为角阶，每个阶有 $2l+1$ 个谐波。使用角索引 m ∈ { − l , − l + 1 , . . . , l − 1 , l } m \in \{ -l, -l+1, ..., l-1, l \} m∈{−l,−l+1,...,l−1,l} 。角函数的输入为单位向量 ${\hat{r}}_{ab}={\vec{r}}_{ab}/r_{ab}$ 。
定义 $L$ 为使用的最大阶， $M={\sum }_{l=0}^L(2l+1)$ 为角函数标量值总计，文中采用 $L=2$ 。球谐函数形成单位球面上函数的标准正交基，并且对于空间中的旋转有等变性。

Filter functions

定义滤波函数 $F_{cm}^l$ 为辐射函数的标量与角函数的结合
$$F_{cm}^l({\vec{r}}_{ab})=R_c(r_{ab})Y_m^l({\hat{r}}_{ab})$$

其中， c ∈ { 0 , 1 , . . . , E } , l ∈ { 0 , 1 , . . . , L } c \in \{0, 1, ..., E \}, l \in \{ 0, 1, ..., L \} c∈{0,1,...,E},l∈{0,1,...,L} ，等变卷积维度为 $E$ 。

滤波函数与特征结合

对于第一个网络层的输入，我们只有标量特征（角度阶 $l=0$ ）和总共 $E=3$ 的辐射特征。
等变卷积 $L$ 使用张量积，来组合点相关特征和过滤函数来更新特征 $V$
$$L_{ac{m}_o}^{l_o}(V_{ac{m}_i}^{l_i})=\sum _{m_i,m_f}{C}_{(l_f,m_f)(l_i,m_i)}^{(l_o,m_o)}\sum _{b\in neighbors(a)}{F}_{cmf}^{l_f}({\vec{r}}_{ab})V_{bc{m}_i}^{l_i}$$

其中，下角标 $f,i,o$ 分别表示滤波特征，输入特征和输出特征，选取 $K$ 个最近的邻点进行卷积。 $C$ 为 Clebsch-Gordan 系数，只有当 $|l_i-l_f|\le l_o\le l_i+l_f$ 时系数不为零 ，但不同的 $l_i$ 和 $l_f$ 可以输出相同的 $l_o$ （比如当 $l_o$ 为 1 时， $l_i$ 和 $l_f$ 分别为 1 0 和 1 1 都能使 $C$ 不为零），所以将具有相同 $l_o$ 的输出沿着维度 $c$ 拼接。

Pointwise normalization

节点标准化操作 $N$ 为
$$N(V_{acm}^l)=\frac{V_{acm}^l}{\sqrt{{\sum }_{c,m}(V_{acm}^l{)}^2}}$$

如果只对维度 c 进行标准化呢？

Pointwise non-linearity

节点非线性操作 $P$ 为
$$P(V_{acm}^l)=\{\begin{array}{rclc}& \eta (V_{acm}^l)& & ifl=0\\ & V_{acm}^l\cdot \eta (\sqrt{{\sum }_m(V_{acm}^l{)}^2}+b^l)& & otherwise\end{array}$$

其中， $b^l$ 为可学习的标量。
$$\eta (x)=\ln (0.5e^x+0.5)$$

Pointwise self-interaction

节点自交互层 $S$ 为
$$S_{adm}^l(V_{acm}^l)=b_d+\sum _c{V}_{acm}^l{W}_{cd}$$

其中， $W$ 为可学习的矩阵， $b$ 为可学习的偏差项 ， $d$ 为新的辐射索引（ $c\to d$ ）。只有当 $l=0$ 时才使用偏差项，文中未给出原因，求大佬解答。

Atomic embedding

根据每个原子为碳、氮或氧，使用 one-hot 嵌入来表示最初的特征 $V$ （角阶 $l$ 默认为 0 ）。

Per-channel mean

将零阶谐波对应的卷积输出进行平均可以获得分子嵌入，这会导致之后的层具有旋转不变性，因为零阶谐波本身对旋转是不变的。新特征 $E$ 为
$$E_c=\sum _a{V}_{ac0}^0$$

Dense layers

该层为全连接网络
$$E_d=b_d+\sum _c{W}_{cd}{E}_c$$

Training

模型输入为候选分子的三维结构，输出为候选分子结构与实验测得结构之间的均方根误差 $RMSD$
$$RMSD=\sqrt{\frac{1}{N}\sum _a|{\vec{p}}_a-{\vec{p}}_a^{\prime }{|}^2}$$

其中， $N$ 为原子数， ${\vec{p}}_a$ 和 ${\vec{p}}_a^{\prime }$ 分别为原子 $a$ 在候选结构和实验结构中的位置。