Py学习  »  机器学习算法

崔庆华团队提出通过自由设计Clebsch-Gordan变换实现高效机器学习分子力场的方法

BioArtMED • 4 月前 • 107 次点击  

传统的药物设计是基于分子模拟的,该策略也为解释药物作用机制有重要帮助。在分子模拟领域中,分子力场预测对于分子动力学研究,及其下游应用分子对接,化学反应预测有重要作用。从头计算分子力场精度较高,然而量子力学薛定谔方程的求解速度上却不尽如人意;经验力场能够以较快速度完成力场建模,却由于庞大的经验性估计,在精度上难以取得令人满意的效果。机器学习分子力场凭借神经网络的万能拟合能力以及GPU加速的优势,能够在拟合密度泛函理论产出的高精度数据的同时,达到经验力场的预测速度,成为了分子力场研究的前沿领域。然而,为了满足物理规律,欧几里得群与置换群等变性为机器学习分子力场的设计套上了枷锁。

近日,崔庆华团队在机器学习的ICLR (International Conference on Learning Representation,国际学习表征会议) 2025发表文章FreeCG: Free the Design Space of Clebsch–Gordan Transform for Machine Learning Force Fields通过不变性的传递性以及cross attention模块的置换不变性,成功构造置换不变的抽象边(Abstract edges)进而释放Clebsch-Gordan变换的设计空间,设计出了高效且精确的机器学习分子力场模型,FreeCG,在多个标准分子力场评测数据上取得了世界最先进的水平

C:\Users\lenovo\Documents\WeChat Files\wxid_inihevkjiux122\FileStorage\Temp\1741051738267.png

机器学习力场(Machine Learning Force Fields,MLFFs)在药物研发、材料科学、化学反应动力学、纳米技术等领域具有重要意义。它在准确性和效率之间取得了令人满意的平衡,预计其性能可与密度泛函理论(DFT)或其他高精度方法相当,却能在计算速度上快几个数量级。

图神经网络(Graph Neural Networks,GNNs)在多个机器学习力场基准上表现优异。群论和群表示理论在机器学习力场的GNN设计中起到了重要作用。比如,这些工作通常要求旋转不变性,因为我们自然期望势能在分子旋转时保持不变。最近的设计趋势是确保网络对旋转、反射和平移都具有等变性,希望网络的内部特征能随输入分子的几何变化而变化,从而提升模型的表达能力。具备这一特性的GNN被称为等变图神经网络(Equivariant Graph Neural Networks,EGNNs)

为了更好地建模多体相互作用,许多方法会采用不可约表示(irreps)来表示高阶几何对象;在此背景下,Clebsch-Gordan(CG)变换用于在不同不可约表示间进行转换。一些工作利用了高阶不可约表示或张量,确实在性能上带来了显著提升。然而,这种性能提升也带来了较高的计算开销。原因在于,张量是标量和向量的扩展,对它进行CG变换本质上扩展了点积运算,因此不可约表示的阶数越高,所需的计算量就越大。

要满足置换等变性的要求会进一步加剧这一负担。与旋转或平移等变性不同,置换等变性通常在EGNN中被隐式地保证,这意味着网络的内部特征在原子顺序变化时也应当随之变化。为维持置换等变性,EGNN需要让每个节点接收来自其相邻节点以及相连边的信息,而CG变换的高负载运算就在这些邻域原子和边上展开。也就是说,模型设计不能简单地减少邻居之间的计算,否则就会破坏置换等变性。此外,设计空间的狭窄也限制了我们自由构造CG变换层,进而限制了模型的表达能力。比如,我们往往需要对每个邻居原子采用同样的操作方式(例如,通常会使用相同的MLP来对标量边特征进行处理,从而产生中心原子与各邻居原子之间计算的权重)


该研究提出了FreeCG模型,该模型从每个原子周围的边中生成并细化几何特征。作者将不同聚合方式得到的几何特征称为“抽象边”(abstract edges);当关注某一原子与其相邻原子和边所维护的内部特征时,这些抽象边具有置换不变性。通过不变性的传递性,作者证明对这些抽象边进行CG变换同样是置换不变的,不会影响网络层的置换等变性,因此摆脱了此前的计算负担。此外,这些抽象边由不同的真实边共同构成,能同时包含它们经过精炼后的信息,从而提升模型的表达能力。抽象边的这种不变性使模型能够为不同边分配不同权重,而不需要用同一个MLP来计算所有边的权重。本文作者凭借自由的设计空间进一步做出更多优化。他们将抽象边划分为若干组,分别在各组内进行运算,以进一步降低计算开销。此前保持E(3)等变性的模型往往开销更大,因为它们需要一个额外的宇称参数(取1或-1),从而使不可约表示的数量加倍。而作者选择了一种有效的路径来完成CG变换,仅保留E(3)等变性所需的部分计算,相比维持SE(3)等变性的方式更加高效。研究者还借鉴了ShuffleNet的思路,提出一种被称作抽象边洗牌的做法,用于将不可约表示的特征进行重新组合。最后,再将这些抽象边引回到交叉注意力的计算中,以提升注意力得分的质量。以上这些操作都是基于抽象边的置换不变性而得以实现的。

作者在分子力场小分子数据集MD17、rMD17、大分子数据集MD22以及分子性质预测数据集QM9上都进行了实验,证明了所提出模型FreeCG在分子力场与分子性质预测任务上都取得了出色的表现,此外,本研究提出模型在water-1k、LiPS、Chignolin这些带有周期性边界条件的现实分子体系上也取得了出色的效果。该研究所提出的CG变换层设计范式还被证明能很好地提升现有模型的表现,作者在先前领先结构QuinNet上的实验证明了这一点。本研究为机器学习分子力场的建模提出了一种新范式,大大降低了Clebsch-Gordan变换在每层的计算负担,为药物研发、材料科学以及化学反应动力学等相关领域提供了高效的新方法和新技术。

论文链接: https://openreview.net/forum?id=sfi2j1Ot6j

制版人:十一


BioArt

Med

Plants

人才招聘

会议资讯



学术合作组织

(*排名不分先后)

战略合作伙伴
(*排名不分先后)
·
·








转载须知


【非原创文章】本文著作权归文章作者所有,欢迎个人转发分享,未经作者的允许禁止转载,作者拥有所有法定权利,违者必究。


Python社区是高质量的Python/Django开发社区
本文地址:http://www.python88.com/topic/179936
 
107 次点击