[论文解读] Efficient and Equivariant Graph Networks for Predicting Quantum Hamiltonian

Haiyang Yu, Xu Zhao+3 more|arXiv (Cornell University)|Jun 8, 2023

Machine Learning in Materials ScienceMaterials Science被引用 3

一句话总结

本文提出 QHNet，一种用于预测分子体系中量子哈密顿矩阵的高效 SE(3)-等变图神经网络。通过将张量积运算减少 92%，并固定所有原子对的输出块形状，QHNet 在保持 MD17 数据集上具有竞争力的准确率的同时，实现比最先进模型快 3 倍的训练速度和 50% 的内存占用降低，且可灵活泛化至混合分子数据集。

ABSTRACT

We consider the prediction of the Hamiltonian matrix, which finds use in quantum chemistry and condensed matter physics. Efficiency and equivariance are two important, but conflicting factors. In this work, we propose a SE(3)-equivariant network, named QHNet, that achieves efficiency and equivariance. Our key advance lies at the innovative design of QHNet architecture, which not only obeys the underlying symmetries, but also enables the reduction of number of tensor products by 92\%. In addition, QHNet prevents the exponential growth of channel dimension when more atom types are involved. We perform experiments on MD17 datasets, including four molecular systems. Experimental results show that our QHNet can achieve comparable performance to the state of the art methods at a significantly faster speed. Besides, our QHNet consumes 50\% less memory due to its streamlined architecture. Our code is publicly available as part of the AIRS library (\url{https://github.com/divelab/AIRS}).

研究动机与目标

• 开发一种深度学习模型，能够以高效率和空间变换下的等变性，准确预测量子哈密顿矩阵。
• 解决现有用于量子体系的图网络中模型效率与 SE(3) 等变性之间的权衡问题。
• 设计一种架构，可在不增加模型复杂度或通道维度的前提下，灵活扩展至多种原子类型。
• 实现对多样化分子体系（包括混合数据集场景）的量子张量通用预测。
• 减少张量积运算带来的计算开销，此类运算在等变模型中主导了时间和内存成本。

提出的方法

• QHNet 采用一种新颖架构，与先前的 SE(3)-等变模型相比，将张量积运算数量减少了 92%，显著提升了效率。
• 其使用固定形状的扩展模块，为所有原子对输出中间块，且与原子类型无关，从而实现可扩展且统一的处理。
• 通过设计确保 SE(3) 等变性——即预测结果在旋转、平移和置换下保持一致变换。
• 基于原子轨道配置，通过后处理方式重构如哈密顿矩阵等量子张量。
• 通过预先定义具有固定形状的完整轨道矩阵，避免在处理多种原子类型时通道维度呈指数增长。
• 在 MD17 数据集上使用标准损失函数进行训练，包括哈密顿矩阵、轨道能级和波函数系数的预测。

实验结果

研究问题

• RQ1 SE(3)-等变图网络是否能在预测量子哈密顿矩阵时同时实现高准确率和显著提升的效率？
• RQ2 在不牺牲模型表达能力或等变性的情况下，张量积运算最多可减少到何种程度？
• RQ3 该模型在涉及多种分子类型的混合数据集场景下，泛化性能如何？
• RQ4 固定形状的输出模块是否能实现对多样化分子体系的可扩展且内存高效的预测？
• RQ5 QHNet 的计算成本降低是否带来更快的收敛速度和更低的内存占用，同时不损害预测准确率？

主要发现

• 与 SOTA 基线 PhiSNet 相比，QHNet 实现了超过 3 倍的训练速度提升，并将 GPU 内存消耗降低了 50%。
• 在 MD17 数据集上，QHNet 在哈密顿矩阵预测的平均绝对误差（MAE）方面表现相当，且在水分子数据集上实现了 4.88×10⁻⁵ Eh 的改进。
• 在乙醇、丙二醛和尿素数据集上，QHNet 在轨道能级预测方面始终优于 PhiSNet，且平均绝对误差更低。
• 分子轨道系数的余弦相似度保持在 ≥99.8%，表明波函数重构准确。
• 在包含水、乙醇、丙二醛和尿素的混合数据集训练中，QHNet 实现了有效的泛化，哈密顿矩阵的 MAE 达到 83.12×10⁻⁶ Eh。
• 固定形状输出设计使得模型可无缝扩展至多分子数据集，而 PhiSNet 在此类场景下则面临架构复杂性问题。