Understanding the interactions of atoms such as forces in 3D atomistic systems is fundamental to many applications like molecular dynamics and catalyst design. However, simulating these interactions requires compute-intensive ab initio calculations and thus results in limited data for training neural networks. In this paper, we propose to use denoising non-equilibrium structures (DeNS) as an auxiliary task to better leverage training data and improve performance. For training with DeNS, we first corrupt a 3D structure by adding noise to its 3D coordinates and then predict the noise. Different from previous works on denoising, which are limited to equilibrium structures, the proposed method generalizes denoising to a much larger set of non-equilibrium structures. The main difference is that a non-equilibrium structure does not correspond to local energy minima and has non-zero forces, and therefore it can have many possible atomic positions compared to an equilibrium structure. This makes denoising non-equilibrium structures an ill-posed problem since the target of denoising is not uniquely defined. Our key insight is to additionally encode the forces of the original non-equilibrium structure to specify which non-equilibrium structure we are denoising. Concretely, given a corrupted non-equilibrium structure and the forces of the original one, we predict the non-equilibrium structure satisfying the input forces instead of any arbitrary structures. Since DeNS requires encoding forces, DeNS favors equivariant networks, which can easily incorporate forces and other higher-order tensors in node embeddings. We study the effectiveness of training equivariant networks with DeNS on OC20, OC22 and MD17 datasets and demonstrate that DeNS can achieve new state-of-the-art results on OC20 and OC22 and significantly improve training efficiency on MD17.
翻译:理解三维原子体系中原子相互作用(如力场)是分子动力学和催化剂设计等众多应用的基础。然而,模拟这些相互作用需要高计算密度的从头算方法,导致可用于训练神经网络的数据有限。本文提出将非平衡结构去噪(DeNS)作为辅助任务,以更充分利用训练数据并提升性能。在DeNS训练中,我们首先通过向三维坐标添加噪声来破坏原始结构,随后预测该噪声。与以往仅限于平衡结构的去噪研究不同,本方法将去噪推广至规模更大的非平衡结构集合。核心差异在于:非平衡结构并非局部能量最小值对应状态,具有非零力场,因此相较于平衡结构可能对应多种原子位置。这使得非平衡结构去噪成为不适定问题——因为去噪目标并非唯一确定。我们的关键洞察在于额外编码原始非平衡结构的力场,以明确指定去噪对象。具体而言,给定被破坏的非平衡结构及其原始力场,我们预测满足输入力场的非平衡结构,而非任意结构。由于DeNS需要编码力场,该方法天然倾向等变网络——这类网络可便捷地在节点嵌入中整合力场及高阶张量。我们在OC20、OC22和MD17数据集上验证了DeNS训练等变网络的有效性,结果表明DeNS可在OC20和OC22数据集上取得新最优结果,并显著提升MD17数据集上的训练效率。