代码
通用分子结构建模网络ViSNet

通用分子结构建模网络ViSNet

通用分子结构建模网络ViSNet

**作者:**于璠 来源:知乎

背景

尽管几何深度学习已经彻底颠覆了分子建模领域,在理解生物活性机制、化学性质预测、药物设计和蛋白质工程方面发挥着关键作用。但在实际应用中依旧存在一些有待解决的问题和局限性。

1. 分子可解释性不足:深层神经网络尽管可以进行预测,但缺乏对分子的深入洞察;

2. 随着分子尺寸的增加,计算成本迅速增加:一些目前最先进的方法中采用的高阶 Clebsch-Gordan 系数计算是计算密集型的,因此阻碍了其在大分子中的应用;

3. 需要实际应用中的盲目测试和评估:模型总是在基准测试上进行测试,同时也需要仔细评估在实际应用中的有效性。

为了解决这些难题,微软研究院科学智能中心的研究员们创新性地提出了通用分子结构建模网络 ViSNet (Vector-Scalar interactive graph neural Network)。ViSNet 能够以低成本高效的方式提取几何特征并建模分子结构,并在多个分子动力学基准测试中均表现优异。

01

方法

微软提出了“方向单元”(direction unit)从中心节点到其任何第一个相邻节点的所有归一化向量的总和,作为中心节点的矢量化表示。再以此将键长、键角和二面角计算扩展到二体、三体和四体相互作用。然后,通过设计运行时几何计算(runtime geometry calculation, RGC)模块来描述模型操作等多体交互。此外,微软进一步提出了向量标量交互式消息传递机制(ViS-MP),其中方向单元会通过构建块由节点和边的标量表示迭代更新,反过来,标量表示由方向单元同时更新 RGC 模块。RGC 和 ViS-MP 的独特设计显著增强了几何编码能力并加速了分子图神经网络中的消息传递过程。

关键创新点如下:

1. 运行时几何计算(RGC)策略:可隐式提取角度、二面角扭转角和非规则角等几何特征。方向单元的概念使用,让ViSNet在保持线性时间复杂度的同时,加速了模型训练和推理过程,降低了内存消耗。

2. 扩展向量表示:引入球谐函数,用内积简化了计算成本昂贵的Clebsch–Gordan乘积。

3. 向量标量交互式消息传递机制(ViS-MP):通过向量隐藏表示和标量表示的交互,充分利用几何特征。

例如,Graph2Edits[1]利用图神经网(GNN)和图编辑操作来表示和预测化学反应,通过对图结构的编辑模拟化学反应过程。该模型采用了自回归的预测方式,从产物出发一步步预测可能的化学改变,直到得到合适的反应物。因为其能够直接在图结构上进行操作,在处理单步反应和较为简单的化学反应中表现良好。但不足是在处理多步反应和需要复杂反应条件的优化时,表现较为有限,难以有效捕捉多种反应条件的复杂关系。

图1 ViSNet整体架构

图2 运行时几何计算(RGC)模块示意图

02

结果

VisNet的效果主要体现在几何特征和分子结构之间合理的模型可解释性上。该网络在预测分子能量、力、HOMO-LUMO间隙和其他量子特性的六个基准测试中均能实现最先进的性能。

图3 在MD17上VisNet与其他算法的MAE对比

图4 在MD22上VisNet与其他算法的MAE对比

此外,在小分子和166原子的Chignolin上,VisNet能够以高保真度驱动分子动力学模拟。

图5 Chignolin构象空间评估与分子动力学模拟

03

感想

ViSNet通过引入精心设计的运行时几何计算(RGC)策略和向量-标量交互等变消息传递(ViS-MP)机制,实现了其加速分子动力学模拟的目标,使大型分子系统的模拟精度接近从头算法。这是一个在分子建模领域具有突破性的模型,它通过结合先进的机器学习技术和精心设计的网络架构,提高了对分子系统模拟的准确性和效率。这启发我们,在应对生物活性机制、化学性质预测、药物设计和蛋白质工程等分子建模领域,不应仅关注在基准测试上的有效性,也需要仔细评估在实际应用中的有效性。同时,ViSNet的实现也提醒我们,图神经网络随着模型变得越来越大、越来越深,提高可解释性和低计算成本是未来提升AI在科学领域的实际应用的重要方向之一。

参考文献

[1] Wang, Y., Wang, T., Li, S. et al. Enhancing geometric representations for molecules with equivariant vector-scalar interactive message passing. Nat Commun 15, 313 (2024). https://doi.org/10.1038/s41467-023-43720-2

[2] Microsoft. MSRA. (n.d.). https://www.msra.cn/zh-cn/news/features/visnet