[논문 리뷰] Equivariant message passing for the prediction of tensorial properties and molecular spectra

Kristof T. Schütt, Oliver T. Unke외 1명|arXiv (Cornell University)|2021. 02. 05.

Machine Learning in Materials Science인용 수 265

한 줄 요약

논문은 회전 대칭인 메시지 전달 신경망 PaiNN을 molecular graphs에 적용하여 스칼라 및 텐서 특성을 예측하고, 분자 스펙트럼 시뮬레이션을 효율적으로 가능하게 하며, 더 작은 모델과 더 빠른 추론으로 최첨단 성능을 달성합니다.

ABSTRACT

Message passing neural networks have become a method of choice for learning on graphs, in particular the prediction of chemical properties and the acceleration of molecular dynamics studies. While they readily scale to large training data sets, previous approaches have proven to be less data efficient than kernel methods. We identify limitations of invariant representations as a major reason and extend the message passing formulation to rotationally equivariant representations. On this basis, we propose the polarizable atom interaction neural network (PaiNN) and improve on common molecule benchmarks over previous networks, while reducing model size and inference time. We leverage the equivariant atomwise representations obtained by PaiNN for the prediction of tensorial properties. Finally, we apply this to the simulation of molecular spectra, achieving speedups of 4-5 orders of magnitude compared to the electronic structure reference.

연구 동기 및 목표

• 분자 그래프에서 방향 정보를 전달하는 데 회전 불변 표현이 가진 한계를 동기 부여합니다.
• 데이터 효율성과 표현력을 개선하기 위해 회전 대칭 메시지 전달 및 PaiNN 아키텍처를 제안합니다.
• 등가 표현을 사용하여 쌍극 모멘트, 폴라리제이션 등 텐서 특성 예측을 가능하게 합니다.
• 대전자구조 참조 대비 큰 속도향상을 보이며 RPMD 시뮬레이션을 통한 분자 스펙트럼 응용을 입증합니다.

제안 방법

• 스칼라 및 벡터 특징이 회전 대응 함수로 업데이트되는 등가 메시지 전달을 개발합니다.
• 방사형 기저 함수 필터를 사용한 연속 필터 합성( invariants scalar messages )과 벡터 인식 합성으로 회전 대칭 메시지를 사용합니다.
• 잔차 아키텍처를 통해 스칼라 및 벡터 업데이트를 연결하여 방향 정보를 전파합니다.
• 원소별 잠재 특징과 로컬 벡터 특징을 결합한 1계 텐서 분해를 통해 텐서 특성을 예측합니다.
• Adam의 가중치 감소 및 학습률 스케줄링으로 에너지, 힘, 쌍극 모멘트 및 극화도에서 엔드투엔드로 학습합니다.

실험 결과

연구 질문

• RQ1 회전 대칭 메시지 전달이 다양한 데이터 구간에서 불변 GNN보다 표준 분자 특성 벤치마크에서 우수하게 성능을 내는가?
• RQ2 대칭 표현이 쌍극 모멘트나 극화도와 같은 텐서 특성의 정확한 예측을 지나치게 복잡한 모델 없이 가능하게 하는가?
• RQ3 PaiNN이 전자구조 참조 대비 IR/Raman 스펙트럼의 분자 스펙트럼 시뮬레이션을 빠르고 정확한 궤적을 통해 가속시키는가?

주요 결과

• PaiNN은 QM9에서 다수의 스칼라 특성에 대해 최첨단 혹은 경쟁력 있는 결과를 보이며, 매개변수 수가 약 0.6M으로 DimeNet++(약 1.8M)과 같은 기준보다 작습니다.
• PaiNN은 QM9에서 DimeNet++ 대비 추론 시간을 70% 이상 단축시키며(배치당 50분자 각각 13 ms 대 45 ms) 속도 향상을 달성합니다.
• 소거 연구는 등가 아키텍처의 모든 구성요소가 정확도에 기여하며, 벡터 특징 합성 및 스칼라–벡터 결합이 유의하게 이득을 제공합니다.
• 등가 표현은 쌍극 모멘트 및 극화도와 같은 텐서 특성의 정확한 예측을 가능하게 하며, 원자 중심 극화 항 및 전역 기하를 활용합니다.
• PaiNN은 빨라진 RPMD 기반 IR 및 라만 스펙트럼 시뮬레이션을 가능하게 하며, 전자구조 참조 대비 상당한 속도 향상을 보이고 에탄올(RPMD)과 아스피린 데이터셋에 대해 실험과 거의 일치하는 스펙트럼을 생성합니다.
• 등가 특징은 작게 잘리는 컷오프에서도 방향 정보를 포착하는 모델의 능력을 향상시켜, 국소 이웃을 넘은 효율적 전파를 가능하게 합니다.