MindSpore AI科学计算系列 | 基于深度学习模型来替代传统DFT模型以及DeephE3模型的分析综述

最新推荐文章于 2024-05-22 10:31:15 发布

昇思MindSpore

最新推荐文章于 2024-05-22 10:31:15 发布

阅读量254

点赞数

分类专栏：技术博客 AI工程与实践文章标签：人工智能深度学习

本文链接：https://blog.csdn.net/Kenji_Shinji/article/details/134034136

版权

技术博客同时被 2 个专栏收录

416 篇文章 55 订阅

订阅专栏

AI工程与实践

29 篇文章 7 订阅

订阅专栏

背景

近年来深度学习在计算量子化学领域取得了重大进展，现有的深度学习方法已经证明了其在解决各种具有挑战性的量子力学模拟任务方面的效率和表现力，随着当今算力的不断增加，深度学习和第一性原理方面结合越来越深，从基本的物理学定律出发，在原子和电子的尺度结合等变神经网络来预测量子领域的值，尤其是在DFT领域里，也诞生了一些深度学习模型，例如DeephE3，QhNet, 这些网络都是基于等变网络来预测DFT哈密顿量。这篇文章会介绍密度泛函理论(DFT)的由来，等变神经网络的原理以及一种较为通用的等变神经网络E3nn, 最后会概述DeepHE3模型，一个由清华大学团队提出的一个基于E3等变网络来预测DFT哈密顿量的的模型 [1]。

1、密度泛函理论

密度泛函理论 (density functional theory，简称DFT) 是一种研究多电子体系电子结构的量子力学方法, 将多电子系统通过电子的密度泛函来进行表示。密度泛函理论在物理和化学上都有广泛的应用，特别是用来研究分子和凝聚态的性质。

在没有密度泛函理论之前，通过求解薛定谔方程来计算系统波函数。薛定谔方程作为量子力学的基本方程，形式如下：

而密度泛函理论用电子密度取代波函数，既通过电子的密度研究多电子体系电子结构的方法。其中密度是一个三位坐标的函数来描述电子的密度, 泛函是指把密度映射成能量E的函数。DFT把电子密度作为基本量，将能量表示为电子密度的泛函。由于电子密度只是空间坐标的函数，让多电子体系的维度直接降到3，简化了薛定谔方程的求解过程。

1965年，美国加利福尼亚大学圣迭戈分校的沃尔特·科恩与沈吕九提出了Kohn–Sham equation。作为密度泛函理论里最普遍的通用的代表，KS方程把把相互作用的多粒子体系转换成一个无相互作用的单粒子体系，把电子之间的相互作用归结到未知的交换关联势中：

式4. KS方程[4]

其中交换关联势项是指有相互作用的多粒子系统和无互作用的多粒子系统的能量差。且该能项的准确函数形式未知，只能表示成电子密度的近似泛函，比如局域密度近似（LDA）等。电子密度由上述单电子波函数方程的解决定，因此该方程的具体形式依赖于其解本身，需要通过自洽迭代进行求解。

图1. 近似计算过程[4]

其计算复杂度为O(N^3)，N为电子个数，对大体系仍然求解难度大。

2、等变网络

在用神经网络计算一些量子特性的时候，通常需要考虑到粒子的旋转对这些特性的变换。对于一些标量值，例如能量值，粒子之间的距离等，不会受粒子的旋转而影响。而对于一些多维向量的特性，例如力，哈密顿量等的值，需要根据粒子的旋转进行相应的改变，而且这种改变需要从网络开始到结束都保持一致性。因此对于大多数第一性原理有关的模型都会使用等变网络。

2.1 什么是等变性

以一个函数来举例子，如果你对其输入施加的变换同样也会反应在输出上，那么这个函数就具有等变性。f(g(x)) = g(f(x))。

2.2 什么是等边网络

（1）对网络输入的变换，需要对称的映射到内部和输出结果上。

（2）例如有一个三维原子的结构，我们要去用神经网络去预测他的各种性质，例如势能，电子数量，受力方向。如果我们旋转了该原子结构，它的势能和电子数量应该保持不变，因为它们是标量；而它们的受力方向结果，应该相应的做出进行变换，因为它们是多维向量。这种对称映射需要反应在网络中间和结果上。因此保证这种映射关系就需要一个等变的网络。

2.3 为什么要实现等变？

对于一个二维的图像，例如一张动物图，为了让模型有对称性，通常会做数据增强，把这张动物图旋转10个不同的角度，喂到神经网络里训练，来让网络训练不同角度的一张图。但是对于一个三维的模型，例如原子结构，这种增强就不太现实，通常一个简单的三维模型，如果要用数据增强，需要有至少500次旋转的数据增强，才能足够涵盖一个原子结构在不同角度的特性。而如果使用等变网络，只需要传入一个结构就行。

图2. 二维动物图像

图3. 三维模型图[5]

3、E3nn：基于三维欧几里得空间的空间变换神经网络

E3：即三维欧几里得空间的空间变换群，可分解为平移、旋转（ SO(3) 特殊正交群）与反演，平移的等变性在卷积中已经满足，因此我们重点关注旋转与反演 -> SO(3)×Z2=O(3)

E3NN里的主要概念:

1. Group: 空间的变换类型，例如旋转和反演.

2. Representation: 定义向量空间属于哪一种空间变换组(Group)的表示.

3. Irreducible representation(irreps): 不可约表示，既等于一个不可分解的Representation。每个irreps可以用 (l,p) 标记，l=0,1,2,... 为阶数，p=e,o为奇偶性，l 阶不可约表示的维数为 2l+1 。例如一个矢量，阶数为 1 (代表维度为3)，奇偶性为奇，因此可以简记为1o。

图4. irreps介绍

以下图为例，a1 – a9分别代表9个实数，如果把a1 – a3分别看成3个标量，a4-a6看成一个矢量，a7 – a9看成一个另一个矢量，那么这个矩阵的irreps就是 “ 3 × 0e + 2 × 1o” 来表示。让我们需要对这个矩阵做旋转时，就需要根据irreps里相对应的Group进行不同的变换，对于a1-a3这三个标量，旋转并不影响他们的值，因此乘以1. 而对于a4 – a6，a7-a9两个矢量，就需要乘相对应的旋转矩阵来获取相对应的值。

图5. 旋转矩阵的例子[5]

以下讲解如何分解两相乘的Irreps（如何分解一个tensor product）

式5. Tensor product的分解

例如：2 ⊗ 1 = 1 ⊕ 2 ⊕ 3， 2 ⊗ 2 = 0 ⊕ 1 ⊕ 2 ⊕ 3。从该例子可以看出，e3nn之所以可以保持可以保证等变, 因为它提前确定了网络输入，输出和中间结果的irreps。来保证Group的变换都是根据相应的不可约表示来进行变换，从而防止混淆。

4、DeephE3

一个通用的E{3}等变深度学习框架，通过神经网络将有自旋轨道的原子结构{R}去预测DFT哈密顿量。DeephE3可以通过训练小型材料系统的DFT结果学习去预测更大的材料系统的电子预测。该方法适用于各种材料系统，例如一般的魔角扭曲双层石墨烯或扭曲范德华尔斯材料，并且和直接用DFT计算比成本降低了几个数量级。

下图是整个网络的架构。其中，{Zi}表示原子序数，| rij |表示原子间距离，用于构造阶数等于 0 向量。^rij表示原子之间的相对位置，向量用于构造阶数等于1, 2 的向量。{Zi}传入元素嵌入模块(Elemental embedding)来当作初始顶点; | rij |传入高斯扩展（Gaussian Bias）来当作边特征; ^rij作为原子之间的相对位置会传入球谐函数进行映射生成Y(^rij)，球谐函数Y^l将一个3维向量映射成一个2l+1维向量，其代表输入向量分解成2l+1个基球谐函数时的系数。

图6. DeephE3总体结构[1]

生成的顶点和边特征通过vertex update和edge update更新块更新 L 次，更新块通过等变卷积对原子间距离和相对未知信息进行编码，“⋅”符号代表通道乘法, ||代表向量串联。

随后用message-passing方法，通过获取相邻边的信息，去更新边和顶点的的向量。

最终的边向量会被传递到 Wigner-Eckart 层去展示DFT哈密顿量。如果忽略自旋轨道耦合（SOC），神经网络的输出向量将通过 Wigner Eckart 层使用规则 1 ⊕ 2 ⊕ 3 = 1 ⊗ 2 转换为哈密顿量。如果包括 SOC，则输出由两组实数向量组成，它们组合起来形成复值向量。这些向量会用另外一个规则转换为自旋轨道 DFT 哈密顿量(1 ⊕ 2 ⊕ 3) ⊕ (0 ⊕ 1 ⊕ 2) ⊕ (1 ⊕ 2 ⊕ 3) ⊕ (2 ⊕ 3 ⊕ 4) = (1 ⊕ 1/2) ⊕ (2 ⊕ 1/2 ). ⊕是指tensor add， ⊗是指tensor product。