[论文解读] DeepH-2: Enhancing deep-learning electronic structure via an equivariant local-coordinate transformer

Yuxiang Wang, He Li+7 more|arXiv (Cornell University)|Jan 30, 2024

Parallel Computing and Optimization Techniques被引用 4

一句话总结

DeepH-2 提出了一种等变局部坐标变换器（ELCT），将 E(3)-等变神经网络与局部坐标变换及注意力机制相结合，实现了深度学习基 DFT 哈密顿量预测的最先进精度与效率。通过将 SO(3) 对称性通过边对齐的局部坐标缩减为 SO(2)，实现了 O(L³) 的计算量缩放，并在单层石墨烯上实现了亚毫电子伏特精度（0.12 meV MAE），相比之前模型在精度上提升了 10 倍以上。

ABSTRACT

Deep-learning electronic structure calculations show great potential for revolutionizing the landscape of computational materials research. However, current neural-network architectures are not deemed suitable for widespread general-purpose application. Here we introduce a framework of equivariant local-coordinate transformer, designed to enhance the deep-learning density functional theory Hamiltonian referred to as DeepH-2. Unlike previous models such as DeepH and DeepH-E3, DeepH-2 seamlessly integrates the simplicity of local-coordinate transformations and the mathematical elegance of equivariant neural networks, effectively overcoming their respective disadvantages. Based on our comprehensive experiments, DeepH-2 demonstrates superiority over its predecessors in both efficiency and accuracy, showcasing state-of-the-art performance. This advancement opens up opportunities for exploring universal neural network models or even large materials models.

研究动机与目标

• 为克服现有基于深度学习的 DFT 哈密顿量预测模型计算成本高且精度有限的问题。
• 利用神经网络有效建模电子结构计算中的高角动量轨道。
• 将 E(3)-等变性与局部坐标变换及注意力机制相结合，以提升数据效率与泛化能力。
• 开发一种可扩展、高精度的框架，适用于大规模材料模拟与数据库构建。
• 在基于深度学习的电子结构预测中，实现精度与训练效率的最先进性能。

提出的方法

• 该框架采用等变局部坐标变换器（ELCT），利用边对齐的局部 z 轴将 SO(3) 对称性简化为 SO(2)，从而实现高效的 O(L³) 计算。
• 每个边应用局部坐标变换，简化角动量耦合的表示，同时保持 E(3) 等变性。
• 模型整合基于变换器的架构与原子边上的消息传递机制，利用注意力机制加权邻近原子的贡献。
• 通过张量积层实现等变消息传递，保持 SO(2) 等变性，实现角动量通道的直接混合。
• 该架构支持高维特征空间，并通过减少对称性约束，优化了 GPU 并行计算。
• 模型在单层与双层石墨烯及 MoS₂ 的 DFT 计算哈密顿矩阵上进行端到端训练，采用标准的训练/验证/测试划分。

实验结果

研究问题

• RQ1 深度学习模型能否在复杂材料的 DFT 哈密顿矩阵元预测中实现亚毫电子伏特精度？
• RQ2 能否在保持等变性与精度的前提下，将 E(3)-等变网络的计算成本从 O(L⁶) 降低至 O(L³)？
• RQ3 将局部坐标变换与注意力机制结合，能否提升电子结构预测中的泛化能力与数据效率？
• RQ4 该模型能否在小到大超胞结构（如双层石墨烯与 MoS₂）之间实现有效泛化？
• RQ5 在不牺牲训练效率的前提下，能否充分利用更高的模型容量（例如参数量增加 60 倍）？

主要发现

• DeepH-2 在单层石墨烯上实现了 0.12 meV 的平均绝对误差（MAE），达到与 DFT 数值误差同量级的精度。
• 在所有测试材料中，包括双层石墨烯与 MoS₂，DeepH-2 的 MAE 值均低于 0.2 meV，表明在所有轨道上均实现亚毫电子伏特精度。
• 尽管参数量增加了 60 倍（从 0.07M 增至 4.37M），DeepH-2 将双层石墨烯的每轮训练时间从 200 秒减少至 100 秒。
• 与 DeepH-E3 相比，DeepH-2 在所有数据集上的精度均提升两倍以上，双层石墨烯的 MAE 为 0.17 meV，双层 MoS₂ 的 MAE 为 0.19 meV。
• 模型的 O(L³) 缩放特性使其能够高效处理高角动量通道，这对精确的电子结构建模至关重要。
• 注意力机制的引入使模型能够动态加权边的贡献，从而在复杂成键环境中提升性能。