[论文解读] Generalizing Denoising to Non-Equilibrium Structures Improves Equivariant Force Fields

Yi-Lun Liao, Tess Smidt+2 more|arXiv (Cornell University)|Mar 14, 2024

Computer Graphics and Visualization Techniques被引用 4

一句话总结

本文提出 DeNS（去噪非平衡结构）作为辅助自监督学习任务，以提升原子机器学习中等变力场的性能。通过在非平衡原子结构去噪过程中引入力作为条件信号，DeNS 将原本病态的问题转化为有良好定义的问题，从而在 OC20、OC22 和 MD17 上实现显著性能提升，且训练开销极低。

ABSTRACT

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 可实现与预训练相当的性能增益，但计算成本显著更低。

提出的方法

• 通过添加噪声对非平衡结构的 3D 原子坐标进行破坏，生成损坏输入。
• 通过原始未损坏结构的力作为条件，唯一地定义目标重构结果，从而实现去噪任务的明确定义。
• 训练模型预测一个噪声向量，使得在给定损坏结构和原始力的条件下，能够恢复原始非平衡结构。
• 将 DeNS 作为辅助损失集成到力场模型的训练过程中，与能量和力预测等主任务并行训练。
• 使用等变消息传递网络（如 Equiformer），通过不可约表示（irreps）自然处理力作为向量特征。
• 通过球谐函数对力进行编码，将力投影到网络的不可约表示空间中，确保等变性。

Figure 1 : Illustration of denoising equilibrium and non-equilibrium structures. In this figure, we relax a non-equilibrium structure ( red point) and form a local relaxation trajectory (black dotted arrow). All the points along the trajectory except the blue point are non-equilibrium structures.

实验结果

研究问题

• RQ1 非平衡原子结构的去噪能否被转化为可处理且有效的自监督预训练任务？
• RQ2 将原始力作为条件信号是否能解决非平衡结构去噪的病态性问题？
• RQ3 DeNS 在 OC20、OC22 和 MD17 等标准基准测试中，对等变力场性能的提升程度如何？
• RQ4 与预训练或架构扩展相比，DeNS 在性能增益和训练效率方面表现如何？
• RQ5 DeNS 在不同等变网络架构和超参数设置下是否具备泛化能力？

主要发现

• DeNS 在 OC20 和 OC22 基准测试中均达到最先进性能，优于先前方法，且无需额外数据集。
• 在 MD17 数据集上，DeNS 将训练时间减少 3.1 倍，相比更深的模型（L_max=3），同时性能优于基线模型（L_max=2）。
• 使用 DeNS 训练使 MD17 上的力 MAE 提升 20.8%，超过在更大数据集（PCQM4Mv2）上进行预训练所获得的 17.2% 提升，且训练成本仅增加 10.5%。
• 力编码至关重要：无力编码的 DeNS 表现不一致，部分分子性能未提升，而有力编码的 DeNS 在所有分子上均一致提升性能。
• 该方法具有良好的架构泛化性：DeNS 在 Equiformer 的多种变体上均表现提升，包括无注意力机制或层归一化的变体，表明其在等变模型中具有广泛适用性。
• DeNS 提供与预训练相当的性能增益，但计算成本显著更低，使其成为数据稀缺的原子机器学习任务中极具效率的替代方案。

Figure 2 : Training process when incorporating DeNS as an auxiliary task. The upper blue block corresponds to the original task of energy and force predictions (Equation 1 ), and the lower red block corresponds to training DeNS (Equation 6 ). “Equivariant GNN” and “energy head” are shared across the