[논문 리뷰] Spherical Channels for Modeling Atomic Interactions

C. Lawrence Zitnick, Abhishek Das외 6명|arXiv (Cornell University)|2022. 06. 29.

Machine Learning in Materials Science인용 수 36

한 줄 요약

구면 채널 네트워크(SCN)를 소개합니다. SCN은 구면 조화를 기반의 채널을 사용하여 원자 에너지와 힘을 모델링하는 GNN으로, 회전적으로 조정 가능한 임베딩 및 variant equivariance를 갖추며 OC20에서 에너지와 힘 예측에서 최첨단 결과를 달성합니다.

ABSTRACT

Modeling the energy and forces of atomic systems is a fundamental problem in computational chemistry with the potential to help address many of the world's most pressing problems, including those related to energy scarcity and climate change. These calculations are traditionally performed using Density Functional Theory, which is computationally very expensive. Machine learning has the potential to dramatically improve the efficiency of these calculations from days or hours to seconds. We propose the Spherical Channel Network (SCN) to model atomic energies and forces. The SCN is a graph neural network where nodes represent atoms and edges their neighboring atoms. The atom embeddings are a set of spherical functions, called spherical channels, represented using spherical harmonics. We demonstrate, that by rotating the embeddings based on the 3D edge orientation, more information may be utilized while maintaining the rotational equivariance of the messages. While equivariance is a desirable property, we find that by relaxing this constraint in both message passing and aggregation, improved accuracy may be achieved. We demonstrate state-of-the-art results on the large-scale Open Catalyst dataset in both energy and force prediction for numerous tasks and metrics.

연구 동기 및 목표

• DFT의 원자 에너지와 힘에 대한 높은 계산 비용에 대한 동기 부여 및 addressed.
• 구면 채널을 통해 각도 정보를 인코딩하는 그래프 신경망을 개발하여 예측 정확도를 향상시키려는 의도.
• 회전 등가성의 강화된 표현력을 개선하고 성능을 향상시키는 완화의 역할을 탐구합니다.
• 에너지 및 힘 태스크에서 OC20(Open Catalyst 2020) 데이터세트에 대한 최첨단 성능을 시演합니다.
• SCN이 기존의 최첨단 모델에 비해 샘플 효율성과 확장성을 어떻게 보이는지 평가합니다.

제안 방법

• 각 원자를 노드로, 이웃을 간선으로 표현하는 그래프 기반 모델을 사용한다.
• 노드 임베딩을 구면의 구면 채널 집합으로 인코딩하고, 차수 L까지의 구면 조화로 확장한다.
• 메시지 전달 중 각 정보를 전달하기 위해 간선 방향과 맞추도록 간선 방향 임베딩을 회전한다.
• 노드 간 메시지를 계산하기 위해 신경망을 사용하고, 전역 프레임으로 회전시키며, 필요 시 다중 z축 회전(taps)을 사용하여 등가성을 촉진한다.
• 구면 격자에 점별 비제약 비선형성을 적용하여 메시지를 집계한 뒤, 업데이트를 위해 다시 구면 조화로 변환한다.
• 방향에 걸친 신경망의 구면 적분을 통해 원자별 에너지를 추정하고, 힘은 에너지 그래디언트 또는 구면 방향에 대한 직접 힘 추정으로 계산한다.
• AdamW로 학습하고 데이터 병렬성 및 혼합 정밀도를 활용하며, 가능한 경우 구면 조화의 다중 해상도(밴드)를 사용하여 메모리와 정확도 간의 균형을 맞춘다.

실험 결과

연구 질문

• RQ1 SCN의 구면 채널 임베딩이 OC20에서 기존 GNN보다 에너지 및 힘 예측을 향상시킬 수 있는가?
• RQ2 엄밀한 회전 등가성(m=0 외의 경우)을 완화하는 것이 모델의 정확도와 학습 효율성에 어떤 영향을 미치는가?
• RQ3 회전된 임베딩과 다중 taps를 사용하는 것이 힘 예측에 어떤 각도 정보의 효과를 가져오는가?
• RQ4 SCN은 OC20 태스크에서 이전 최첨단 모델과 비교하여 샘플 효율성이 얼마나 높은가?

주요 결과

• SCN은 S2EF 태스크에서 여러 기준선보다 힘 MAE 및 힘 코사인에서 최첨단 개선을 달성합니다.
• OC20에서 m이 [-1,1] 또는 [-2,2]인 SCN 변종은 엄밀한 m=0보다 이득을 제공하며, 엄밀한 등가성의 완화를 통해 이익이 나타납니다.
• L을 늘리고 K를 더 깊게 확장하는 SCN의 규모 확장은 IS2RE 에너지 예측과 전체 힘 정확도에서 강한 성능 향상을 보입니다.
• 힘 예측에서 F_e 계산 시 다중 회전(taps)을 샘플링하면 Throughput가 다소 감소하더라도 힘 예측이 크게 개선되며, 두 밴드(L, L-1)를 사용하면 에너지 및 힘 예측이 개선되나 Throughput에 대한 영향은 미미합니다.
• 직접 힘 예측(그래디언트 기반 에너지 힘 대신)이 많은 설정에서 우수한 속도/정확도 트레이드를 제공하면서도 경쟁력 있는 정확도를 유지합니다.
• SCN은 여러 강력한 기준선보다 샘플 효율이 높아, 더 적은 학습 예제로도 비슷한 정확도를 달성할 수 있습니다.