MindSpore AI科学计算系列 | 周期性图Transformer提升MindSpore模型对晶体性质的预测
MindSpore AI科学计算系列 | 周期性图Transformer提升MindSpore模型对晶体性质的预测
**作者：**于璠 来源：知乎
背景
晶体材料是一类由3D空间中原子的周期重复形成的大型化合物家族，是许多实际工业应用的基础，如半导体电子学、太阳能电池和化学电池[1]。材料科学通过预测材料性质（例如形成能）和设计具有目标性质的新材料来满足各行各业的巨大需求，但其进展受限于昂贵的实验或耗时的材料模拟。AI方法的成功，使人们尝试在晶体材料上应用这些AI模型。然而与分子不同，晶体材料中的原子排列具有复杂的重复单元和原子的周期排列，开发有效的AI模型需要明确地捕捉这些结构中的对称性。
我们通过晶格矢量L=(L1, L2, L3)，和元胞内原子(Ai, Pi)来表示一种晶体，其中L是3D空间中的晶格重复周期，元胞是晶体的最小重复单元，Ai是其中第i个原子的类型，Pi是其坐标位置。在将晶体结构表示为图结构时，需要具备周期不变性，即学习到的图表示对于最小可重复单元的规模扩大、周期边界变化保持不变。
**1、**CGCNN
较早在晶体图表示中引入周期性的是CGCNN[2]。其构建图的方式如图1所示：
1）图中的节点数等于晶体元胞中的原子个数（元包边界处原子按所属元胞数平分）；
2）周期性重复的原子在图中对应同一个节点；
3）节点间的连接由原子与截断半径内所有原子的连接组成。因为一个原子可以和其周围多个元胞中相同原子相连，而这些原子在晶体图中由同一个节点表示，节点之间可以存在多条不同的边。
图 1. CGCNN构建图表示的方式[2]
这样做的优点是当元胞选取的边界不同时，也能得到同样的晶体图，保证了周期不变性。CGCNN原论文中也存在不足点：
1）没有显式表示周期性信息，微观结构相同，但重复周期不同也可能表示不同晶体；
2）具体实现时对每个原子采用了最近邻n个原子构造连接，当有相同距离的不同原子时，有可能违背周期不变性。
**2、**Matformer
针对这些不足，2022年发表的论文Matformer [3]进行了改进。论文中显式将周期性加入晶体图表示中，并采用attention结构进行图消息传播，在Materials Project、JARVIS等多个数据集上预测晶体属性，得到了不错的结果。
其首先采用和CGCNN类似的方法构建多边图，但将最近邻n的连边方式替换为小于最小距离r则存在边，能够保证周期不变性。图2 a)展示了不使用多边图，当周期边界发生变化时，图表示发生变化的情况。
图 2. a)改变周期边界时满足和不满足周期不变性的图构建示例，b)显式加入周期表示的图构建[3]
然后将晶格矢量显式地以自连边的方式加入图表示中。具体地，如图2 b)所示，将原子与它在3个晶格方向上的邻近复制进行连接和编码。因为使用角度表示空间信息会带来更多的复杂度，额外使用晶格表面的3个对角线隐式的表示晶格矢量的角度关系。
模型的结构如图3所示。原子序号代表的节点信息和原子间距离代表的边信息经过编码后进入多层Matformer Layer。每层Matformer Layer中基于attention架构，针对每个节点将与之相连的节点及边的信息进行汇聚，并以此对该节点进行更新。对于多边图，每条边对应到一个attention head，图中h表示边序号。在生成attention中的K和V时，将边的特征和节点特征进行了拼接。并且在Q和K点乘后，用sigmoid替代了常用的softmax来计算权重，这样可以考虑到节点不同度（连接数）的影响，并且减少了相邻节点间消息的汇聚及广播，提升了计算效率。消融实验也表明sigmoid能得到更好的测试精度。
图 3. a)Matformer的整体网络架构，b)Matformer Layer的网络结构，c)Sigmoid替换为Softmax的消融实验[3]
模型在Materials Project、JARVIS数据集上，对多种属性预测的表现均达到领先。这得益于其显式周期性的引入，以及针对多边图的attention结构的设计。在计算效率上，Matformer比精度次优的ALIGNN也高出很多。预测精度和计算效率对比如图4所示。
图4. Matformer与ALIGNN预测精度和计算效率对比
**3、**感想
通过本文的分析可以看出，加入先验物理知识，例如周期不变性及周期重复特征，对提升模型对材料性质的预测能起到很大的帮助。同时，针对晶体多边图的attention结构的设计在较好地捕捉晶体结构特征的同时，也兼顾了计算效率的优化。当前昇思MindSpore AI框架结合自身优势，基于已有能力开展周期性图Transformer相关工作，从而提升晶体性质预测和计算效率，预计不久将与大家见面。
参考文献
[1]Keith T Butler, Daniel W Davies, Hugh Cartwright, Olexandr Isayev, and Aron Walsh. 2018. Machine learning for molecular and materials science. Nature 559, 7715 (2018), 547–555.
[2]Tian Xie and Jeffrey C Grossman. 2018. Crystal graph convolutional neural networks for an accurate and interpretable prediction of material properties. Physical Review Letters 120, 14 (2018), 145301.
[3]Keqiang Yan, Yi Liu, Yuchao Lin, and Shuiwang Ji. 2022. Periodic Graph Transformers for Crystal Material Property Prediction. In The 36th Annual Conference on Neural Information Processing Systems. 15066–15080.
往期回顾
MindSpore AI科学计算系列 | 初探海洋大模型AI-GOMS，打开大模型在海洋方面的大门
MindSpore AI科学计算系列 | “没有最快，只有更快”，快速傅里叶变换详解