[论文解读] TensorNet: Cartesian Tensor Representations for Efficient Learning of Molecular Potentials

Guillem Simeon, Gianni De Fabritiis|arXiv (Cornell University)|Jun 10, 2023

Machine Learning in Materials Science被引用 17

一句话总结

TensorNet 引入一个使用笛卡尔秩-2 张量的 O(3)-等变神经网络，以高效学习分子势能，在比某些球面张量模型少得多的参数和更低的计算量下实现最先进的准确性。它还能够在预测能量和力的同时预测向量和张量分子量。

ABSTRACT

The development of efficient machine learning models for molecular systems representation is becoming crucial in scientific research. We introduce TensorNet, an innovative O(3)-equivariant message-passing neural network architecture that leverages Cartesian tensor representations. By using Cartesian tensor atomic embeddings, feature mixing is simplified through matrix product operations. Furthermore, the cost-effective decomposition of these tensors into rotation group irreducible representations allows for the separate processing of scalars, vectors, and tensors when necessary. Compared to higher-rank spherical tensor models, TensorNet demonstrates state-of-the-art performance with significantly fewer parameters. For small molecule potential energies, this can be achieved even with a single interaction layer. As a result of all these properties, the model's computational cost is substantially decreased. Moreover, the accurate prediction of vector and tensor molecular quantities on top of potential energies and forces is possible. In summary, TensorNet's framework opens up a new space for the design of state-of-the-art equivariant models.

研究动机与目标

• 在分子模拟与药物发现中说明对高效、准确的原子间势能的需求。
• 提出使用笛卡尔秩-2 张量进行原子嵌入的 O(3)-等变架构。
• 展示相对于更高秩的球形模型的计算效率和参数效率。
• 展示模型在预测能量、力和张量量上的能力。
• 在多样数据集上评估 TensorNet，以确立在更少参数的前提下的最先进性能。

提出的方法

• 用完整的 3x3 笛卡尔张量表示原子，并分解为标量、向量和张量分量（I、A、S）。
• 使用基于边的相对位置和原子序数初始化来形成每条边的张量特征。
• 计算不变量范数，并用它们通过保持对称性的矩阵运算产生每个原子的更新。
• 应用矩阵乘积风格的相互作用，在各层中保持 O(3) 等变性和奇偶性。
• 将总能量预测为来自最终张量范数的逐原子输出之和。
• 通过从张量分量中提取（包括斜对称部分）实现向量和张量输出。

Figure 1: Key steps, from top to bottom, in the embedding and interaction modules for some central atom $i$ and neighbors $j$ and $k$ found within the cutoff radius. a) Relative position vectors are used to initialize edge-wise tensor components, modified using edge-wise invariant functions, and sum

实验结果

研究问题

• RQ1 笛卡尔秩-2 张量表示是否能以比更高秩球形模型更少的参数达到分子势能的最先进精度？
• RQ2 是否仅凭借简单的矩阵乘积和不可约分解即可实现对能量、力和张量性质的准确预测的 O(3) 等变张量运算？
• RQ3 在多样的分子数据集（QM9、rMD17、SPICE/ANI1x/COMP6）上，TensorNet 相对于现有基线的表现如何？
• RQ4 等变性、相互作用设计和截断对预测精度有何影响？
• RQ5 TensorNet 是否能同时可靠地预测标量、向量和张量分子性质？

主要发现

• TensorNet 在 QM9 的 U0、U、H 以及 U0 的基准上超越 Allegro 和 MACE，但参数仅为 Allegro 的 23%，实现了最先进的精度。
• 在 rMD17 上，TensorNet 通过 1–2 层相互作用实现了具有竞争力的能量和力的 MAE，同时具备显著更少的参数（低至 0.535–0.770M），优于其他球形模型。
• TensorNet 在 ANI1x/COMP6 上的泛化表现出色，使用轻量的 1L 模型，在多个基准测试中提供较低的能量和力 MAE 相较于若干基线。
• 对于乙醇在真空中的情形，TensorNet 的能量和力 MAE 大约比 FieldSchNet 和 PaiNN 好 2 倍左右，同时还能提供准确的偶极矩和极化率预测。
• 消融研究表明，完整的 O(3) 等变性和相互作用乘积能显著提升精度，0L 与 1L/2L 的权衡取决于截断和分子大小。
• TensorNet 在 GPU 上显示出较好的推理与训练速度，并在多分子达到数百原子时保持高效，与较小的等变变换器在更大系统上竞争。

(b) Interaction and node update