[论文解读] UNiTE: Unitary N-body Tensor Equivariant Network with Applications to Quantum Chemistry.

Zhuoran Qiao, Anders S. Christensen+4 more|arXiv (Cornell University)|May 31, 2021

Machine Learning in Materials Science参考文献 84被引用 23

一句话总结

UNiTE 提出了一种统一的 N 体张量等变神经网络，能够在处理对称张量时实现线性时间复杂度，从而实现可扩展的、对称性保持的量子化学学习。在基准测试中，其性能优于最先进方法超过 110%，并实现了比传统数值方法快三个数量级的稳健零样本泛化能力。

ABSTRACT

Equivariant neural networks have been successful in incorporating various types of symmetries, but are mostly limited to vector representations of geometric objects. Despite the prevalence of higher-order tensors in various application domains, e.g. in quantum chemistry, equivariant neural networks for general tensors remain unexplored. Previous strategies for learning equivariant functions on tensors mostly rely on expensive tensor factorization which is not scalable when the dimensionality of the problem becomes large. In this work, we propose unitary $N$-body tensor equivariant neural network (UNiTE), an architecture for a general class of symmetric tensors called $N$-body tensors. The proposed neural network is equivariant with respect to the actions of a unitary group, such as the group of 3D rotations. Furthermore, it has a linear time complexity with respect to the number of non-zero elements in the tensor. We also introduce a normalization method, viz., Equivariant Normalization, to improve generalization of the neural network while preserving symmetry. When applied to quantum chemistry, UNiTE outperforms all state-of-the-art machine learning methods of that domain with over 110% average improvements on multiple benchmarks. Finally, we show that UNiTE achieves a robust zero-shot generalization performance on diverse down stream chemistry tasks, while being three orders of magnitude faster than conventional numerical methods with competitive accuracy.

研究动机与目标

• 为量子化学等领域的高阶张量缺乏可扩展的等变神经网络提供解决方案。
• 克服现有基于张量分解的方法在等变学习中面临的可扩展性限制。
• 开发一种神经网络架构，对称 N 体张量在酉群作用下（如三维旋转）保持等变性。
• 通过一种新颖的对称性保持归一化技术——等变归一化，提升泛化能力。
• 在多种下游化学任务中实现高速且高精度的零样本泛化。

提出的方法

• 提出一种新颖的神经网络架构 UNiTE，专为在酉群作用（如三维旋转）下保持等变性的对称 N 体张量而设计。
• 采用参数高效的结构设计，实现与非零张量元素数量成线性比例的时间复杂度，从而实现可扩展性。
• 引入等变归一化，一种在保持对称性的同时提升模型泛化能力的归一化方法。
• 以结构化方式利用张量分解原理，保持等变性，同时避免昂贵的分解过程。
• 在 N 体相互作用上采用消息传递框架，并使用可学习的等变层进行特征变换。
• 采用对称张量表示，以编码量子系统中的多体相互作用。

实验结果

研究问题

• RQ1 能否设计一种神经网络架构，使其在一般对称 N 体张量上对酉群作用（如三维旋转）保持等变性？
• RQ2 该网络能否在非零张量元素数量上实现线性时间复杂度，从而实现对高维问题的可扩展性？
• RQ3 等变归一化是否能在保持对称性的同时提升模型泛化能力？
• RQ4 UNiTE 是否能在量子化学基准上实现最先进性能，并显著优于现有方法？
• RQ5 UNiTE 是否能在多样化的化学任务中实现稳健的零样本泛化，并在速度上超越传统数值求解器？

主要发现

• 与最先进机器学习方法相比，UNiTE 在多个量子化学基准上实现了超过 110% 的平均性能提升。
• 该模型在多样化下游化学任务中展现出稳健的零样本泛化能力，表明对称性带来了强大的归纳偏差。
• UNiTE 的运行速度比传统数值方法快三个数量级，同时保持了具有竞争力的精度。
• 该架构在非零张量元素数量上保持线性时间复杂度，从而可扩展至高维系统。
• 等变归一化在不破坏对称性的情况下提升了模型泛化能力，有助于在未见任务上取得更好性能。
• UNiTE 有效学习了高阶张量上的等变函数，这是现有等变神经网络此前未探索的领域。