论文阅读 Neural Euler’s Rotation Equation

今天带来一篇 NeurIPS 2023 的论文，作者主要工作是提出了神经网络旋转欧拉方程，并将这种方法用于基于能量的模型训练。

Unsupervised Protein-Ligand Binding Energy Prediction via Neural Euler’s Rotation Equation 阅读笔记

摘要

<未编辑>

1 介绍

药物发现的挑战之一是设计对靶蛋白具有高结合亲和力的配体（小分子或抗体）。近年来，已经为小分子开发了许多蛋白质-配体结合预测模型。这些模型通常是在大量标有实验结合数据的晶体结构上进行训练的。然而，很难将这种监督学习方法应用于结合亲和力数据有限的其他配体类别（如抗体）。例如，最大的抗体结合亲和力数据集只有566个数据点。而且从小分子到抗体的迁移学习也很困难，因为它们的结构非常不同。

鉴于这一数据稀缺挑战，我们的目标是开发一种通用的用于小分子和抗体的无监督结合能预测框架。其基本思想是通过最大化Protein Data Bank(PDB) 中晶体结构的对数似然来学习蛋白质-配体复合物的 energy-based model(EBM)¹。这种生成建模方法的动机是 PLM 的最新进展表明蛋白质序列与蛋白质突变效应相关。不过，PLM 不能用于小分子，并且它建模的是蛋白质序列的似然而非结构的似然。我们设想蛋白质复合物结构的似然能更好地用于结合亲和力预测，这种结构是由蛋白质和配体的相对方位来决定的。 实话说，我们的方法与传统的由专家设计的基于物理能量的方法具有相同的思想，也就是说 晶体结构具有最低的能量 （亦即，最高的似然）。我们的关键出发点是使用神经网络来学习一个更昂贵的数据驱动方法的能量函数。重要的是，我们的模型可以用于大规模虚拟筛选，这是因为它不需要昂贵的分子动态模拟。

特别地，我们使用 SE(3) 去噪得分匹配 (DSM) 来训练 EBM，因为最大似然估计(MLE) 很困难。举个例子，对于每次训练而言，我们通过随机旋转/平移配体来创造一个扰动的蛋白质配体复合物，并让我们的模型恢复旋转/平移的噪声。 在我们的背景下，SE(3) DSM 的主要挑战是如何以等变的方式根据 EBM 的分数预测旋转噪声。受物理学启发，我们开发了一种新的等变旋转预测网络，名为神经网络欧拉旋转方程(NERE)。这里面最关键的发现是一个EBM的得分等于每个原子受到的力，这种力进一步地定义了蛋白质和配体的扭力。使用欧拉旋转方程，我们将这种扭力转换为配体的角速度和它对应的用于计算SE(3)DSM损失的旋转矩阵。 重要的是，NERE兼容任何SE(3)不变的网络并且保证网络结构中无需任何更进一步的需求也能具有等变性。

我们使用蛋白质-配体和抗体-抗原两种结合亲和力基准评估了NERE，分别来自 PDBBind 和 SAbDab。我们将NERE与像MM/GBSA这样的无监督的基于物理的模型、与蛋白质语言模型（ESM-1v和ESM-IF）、与使用实验亲和力数据进行回归的大量有监督模型进行比较。为了模拟真实事件的虚拟筛选场景，我们考虑两种设置：输入的复合物是晶体化的或者由分子拼接软件预测的。NERE在所有的设置中都表现的比无监督基线更好，并在抗体设置中超越了有监督模型，这突出了当亲和力数据受到限制时无监督学习的好处。

2 相关工作

2.1 蛋白质配体结合模型

2.2 抗体抗原结合模型

2.3 去噪得分匹配

2.4 等变旋转预测

3 无监督的结合能预测

一个蛋白质配体（或者一个抗体抗原）复合物是一个包括蛋白质和配体（一个小分子或抗原）的几何体。这可以表示为一个元组 ，其中原子特征 ，原子坐标 （按列拼接）。一个复合物的关键属性是它的结合能 。一个更低的结合能意味着一个配体与一个蛋白质之间有着更强的结合。在本节，我们描述如何将 参数化并设计正确的训练目标以从一列没有结合力标签的晶体结构中推理实际的结合能函数。

3.1 EBM 结构

一个能量函数 有一个蛋白质编码器和一个输出层组成。编码器是一个帧平均神经网络²，它学习每个原子的SE(3)不变的表征 。我们使用这个结构是因为它的简单性（详情在附文A）。输出层 是一个带有一个隐藏层的前馈神经网络，它预测每对原子的结合能 。最后，我们定义 为每对结合能之和：

由于原子相互作用在超过某个距离后会消失，我们只考虑那些小于某个距离 的原子对。至此，我们已经将我们的方法描述为了通用术语。我们现在开始阐述为小分子和抗体定制的输入特征和预处理步骤。

3.1.1 小分子

当配体是小分子是，我们模型的输入是一个蛋白质配体复合物，其中蛋白质被裁剪为它的结合口袋（距离配体以内的残基）。在蛋白质这一边，每个在结合口袋中的原子被表示为它的原子名称的one-hot编码（ 等等）。我们考虑所有骨架和侧链上的原子。在配体这一边，原子特征 由基于配体分子图的消息传递网络（MPN）³学习。MPN和能量函数在训练时联合优化。

3.1.2 抗体

当配体是抗体时，我们模型的输入是一个抗体抗原结合接口。结合接口由抗体中包括轻重链的互补决定区（CDR）和最接近CDR的50个抗原残基（抗原决定基）组成。 与小分子的情况不同，我们只考虑每个抗体/抗原残基的 原子。每个残基由预训练的ESM-2表示为2560维的嵌入表征。为了计算效率，训练时我们冻结了ESM-2的权重。

3.2 使用去噪得分匹配训练 EBM

鉴于我们的训练集没有包含结合亲和力标签，我们需要设计一个无监督训练目标，这与有监督的回归损失不同。我们的关键猜想是，我们可以通过最大化我们的训练集中晶体结构的似然来推理真正的能量函数（达到仿射等价）。我们猜想的动机是晶体结构有着蛋白质配体复合物最小的能量状态。 最大似然目标是寻找晶体结构的最小能量，这是因为复合物的似然是 。

由于边际化，最大似然估计（MLE）对EBM来说是困难的，而最近的工作成功地使用去噪得分匹配（DSM）训练了EBM，并证明DSM很好地接近MLE。在标准的DSM中，我们通过添加高斯噪声到配体原子坐标来创建了一个扰动的复合物，也就是说，，其中，。DSM目标尝试匹配我们模型的得分 以及噪声分布的得分 :

直觉上，当输入复合物是一个晶体结构时()， 使梯度变为零。结果，一个晶体结构形态将会是在DSM目标下的一个EBM的局部最小值。

3.3 使用刚体变换噪声扰动复合物

尽管如此，添加高斯噪声对蛋白质配体结合不是理想的，因为它可能会造成无意义的违反物理约束的异构体（例如，芳香环必须是平面结构的）。一个更好的解决方案是类似分子拼接一样通过随机配体旋转和平移创造扰动的复合物 。为了构建随机旋转，我们从一个基于 SO(3) 旋转群的以 为方差的各向同性高斯分布 中取样 向量⁴。每个向量 具有形式 ，其中 是一个从单位球均匀采样的向量， 是一个旋转角，其密度函数

类似地，我们从正态分布 取样随机平移向量 。最后，我们将这个刚体变换应用到配体并计算它的扰动坐标 ，其中 是由旋转向量 给出的旋转矩阵。

此处的 表示矩阵幂而 是一个无穷小旋转矩阵。由于 是一个反对称矩阵，它的矩阵幂具有下面的近似形式

更进一步地，我们不需要显式地计算矩阵幂 因为 是叉积的线性映射，也就是说 。因此，使用旋转矩阵是需要包含叉积操作，这十分高效：

3.4 使用SE(3)去噪得分匹配训练EBM

在这种扰动方法下，DSM的目标是匹配模型得分 以及旋转和平移噪声 。我们的SE(3)DSM的目标是两个损失的和：，其中 和 对应DSM的平移和旋转损失。平移部分是直接的，因为 遵循一个正态分布，并且 ：

DSM的旋转部分更复杂一些。由于 是从一个圆上的均匀分布（密度函数是常数）取样的，密度函数和 的得分是

实际上，我们预先计算了 中截断的无穷级数。不过，最主要的挑战是模型得分 是原子坐标上定义的，它无法直接与 进行比较，因为它们具有不同的维度。为了解决这个问题，我们需要将 通过一个旋转预测器 映射到一个旋转向量 并且将得分匹配表示到旋转空间上：

旋转预测器 是基于神经网络欧拉旋转方程的，它将会在下一节中阐述。

4 神经网络欧拉旋转方程

在本节中，我们展示NERE，它可以用于SE(3)DSM以推理来自梯度 的得分。为了理解我们的方法，让我们首先回顾一些欧拉旋转方程相关的基础概念。

4.1 欧拉旋转方程

在经典机械学中，欧拉旋转方程式一个一阶常微分方程，它用于描述刚体旋转。假设一个配体绕着它的中心质量 以角速度 旋转，欧拉旋转方程在惯性参考系中的定义为：

其中 是一个配体的惯性矩阵， 是它收到的扭力，而 是应用到配体原子 的力。惯性矩阵描述了配体的质量分布和对应角加速度所需的扭力，其定义为：

在一个短时间 内，我们可以近似一个新的角速度 ：

因为我们假设系统是一个惯性参考系()，则有 (我们设 )。我们注意到计算矩阵的逆 是廉价的，因为它是 的矩阵。

4.2 应用NERE到SE(3)DSM

在 SE(3)DSM 中，我们的目标是将梯度 映射到一个旋转向量 。在物理学中，梯度 其实是原子 所受的力 。因此，旋转预测器 在式(10)中是一个简单的欧拉旋转方程的应用：

NERE的一个好的性质是在SO(3)旋转群下是等变的，因为它由物理学推导而出。我们正式地描述这个命题如下，并在附文中给出证明。

命题1：

一旦我们建立了 SO(3) 等变性，就可以简单地满足 SE(3) 等变性，方法是首先置旋转中心于原点()，通过NERE预测旋转，然后将 添加回每个原子。

5 实验

6 讨论

附文

A NERE 旋转不变性

B 实验详情