[논문 리뷰] Generalizing Denoising to Non-Equilibrium Structures Improves Equivariant Force Fields

Yi-Lun Liao, Tess Smidt외 2명|arXiv (Cornell University)|2024. 03. 14.

Computer Graphics and Visualization Techniques인용 수 4

한 줄 요약

이 논문은 원자적 기계학습에서 등변성 힘장의 성능을 향상시키기 위해 보조적 자기지도 학습 과제인 DeNS(Denoising Non-Equilibrium Structures)를 제안한다. 비평형 원자 구조의 노이즈 제거 과정에서 힘을 조건부 신호로 활용함으로써, 기존에 잘 정의되지 않은 문제를 잘 정의된 문제로 전환시켜 OC20, OC22, MD17에서 매우 낮은 학습 오 overhead 를 유지하면서도 뚜렷한 성능 향상을 이룬다.

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가 사전학습과 유사한 성능 향상을 이끌어내지만, 훨씬 낮은 계산 비용을 요구한다는 것을 입증한다.

제안 방법

• 비평형 원자 구조의 3차원 좌표에 노이즈를 추가하여 손상된 입력을 생성한다.
• 원래의 비손상된 구조의 힘을 조건부 신호로 사용하여 목표 재구성의 유일성을 보장한다.
• 손상된 구조와 원래의 힘을 입력으로 받아 원래의 비평형 구조로 복원하는 데 필요한 노이즈 벡터를 예측하는 모델을 훈련한다.
• 힘장 모델의 훈련 과정에서 주된 과제(에너지 및 힘 예측)와 함께 DeNS를 보조 손실로 통합한다.
• 비어리드러블 표현(irreps)을 통해 벡터 값 특징인 힘을 자연스럽게 처리할 수 있는 등변성 메시지 전파 네트워크(e.g., Equiformer)를 사용한다.
• 구면 조화 함수를 활용해 힘을 네트워크의 비어리드러블 표현 공간으로 매핑함으로써 등변성 보장을 한다.

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 데이터셋에서는 깊이가 증가한 모델(L_max=3)보다 훈련 시간을 3.1배 단축시키면서도 기준 모델(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