Generalizing Denoising to Non-Equilibrium Structures Improves Equivariant Force Fields
Primary Area: applications to physical sciences (physics, chemistry, biology, etc.)
Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.
Keywords: equivariant neural networks, graph neural networks, computational physics
Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2024/AuthorGuide.
TL;DR: We propose to use denoising non-equilibrium structures as an auxiliary task to improve the performance of equivariant networks on OC20, OC22 and MD17 datasets.
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 DeNS, we first corrupt a 3D structure by adding noise to its 3D coordinates and then predict the noise. Different from previous works on pre-training via denoising, which are limited to equilibrium structures, the proposed DeNS generalizes to a much larger set of non-equilibrium structures without relying on another dataset for pre-training. The key enabler is the encoding of input forces. A non-equilibrium structure has non-zero forces and thus many possible atomic positions, making denoising an ill-posed problem. To address the issue, we additionally take the forces of the original structure as inputs 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 demonstrate the effectiveness of training equivariant networks with DeNS on OC20, OC22 and MD17 datasets. For OC20, EquiformerV2 trained with DeNS achieves better S2EF results and comparable IS2RE results compared to EquiformerV2 trained without DeNS. For OC22, EquiformerV2 trained with DeNS establishs new state-of-the-art results. For MD17, Equiformer ($L_{max} = 2$) trained
with DeNS achieves better results than Equiformer ($L_{max} = 3$) without DeNS and saves 3.1$\times$ training time. We also show that DeNS can improve other equivariant networks like eSCN on OC20 and SEGNN-like networks on MD17.
Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors' identity.
No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.
Submission Number: 487
Loading