[논문 리뷰] So3krates: Equivariant attention for interactions on arbitrary length-scales in molecular systems

J. Thorben Frank, Oliver T. Unke외 1명|arXiv (Cornell University)|2022. 05. 28.

Machine Learning in Materials Science인용 수 15

한 줄 요약

So3krates는 구면 조화 좌표(SPHCs)를 동적이고 주목적 기반 표현으로 학습함으로써 분자 시스템 내 비국소 전자 효과를 모델링하는 새로운 등변 주목적 메커니즘을 도입한다. 기하학적 정보를 원자 특징에서 분리하고 SPHC 공간에서 구면 자기주목적을 적용함으로써, 모델은 기존보다 0.25–0.4배 적은 파라미터로 상태최고의 정확도를 달성하며, 학습과 추론 속도가 2–14배 빨라져 임의의 길이 척도에서 효율적인 장거리 상관관계 학습이 가능하다.

ABSTRACT

The application of machine learning methods in quantum chemistry has enabled the study of numerous chemical phenomena, which are computationally intractable with traditional ab-initio methods. However, some quantum mechanical properties of molecules and materials depend on non-local electronic effects, which are often neglected due to the difficulty of modeling them efficiently. This work proposes a modified attention mechanism adapted to the underlying physics, which allows to recover the relevant non-local effects. Namely, we introduce spherical harmonic coordinates (SPHCs) to reflect higher-order geometric information for each atom in a molecule, enabling a non-local formulation of attention in the SPHC space. Our proposed model So3krates - a self-attention based message passing neural network - uncouples geometric information from atomic features, making them independently amenable to attention mechanisms. Thereby we construct spherical filters, which extend the concept of continuous filters in Euclidean space to SPHC space and serve as foundation for a spherical self-attention mechanism. We show that in contrast to other published methods, So3krates is able to describe non-local quantum mechanical effects over arbitrary length scales. Further, we find evidence that the inclusion of higher-order geometric correlations increases data efficiency and improves generalization. So3krates matches or exceeds state-of-the-art performance on popular benchmarks, notably, requiring a significantly lower number of parameters (0.25 - 0.4x) while at the same time giving a substantial speedup (6 - 14x for training and 2 - 11x for inference) compared to other models.

연구 동기 및 목표

• 기존 기계학습력장 모델에서 계산 및 아키텍처적 제약로 인해 자주 忽시되는 분자 시스템 내 비국소 전자 효과를 모델링하는 데 도전한다.
• 국소 이웃 영역을 초월한 장거리 상호작용을 포착하는 고차 기하 상관관계의 미분 가능하고 학습 가능한 표현을 개발한다.
• 3차원 분자 시스템의 SO(3) 대칭성에 적응된 자기주목적 메커니즘을 사용하여 임의의 길이 척도에서 분자 상호작용을 효율적이고 확장 가능한 방식으로 모델링한다.
• 구면 조화 좌표(SPHCs)를 학습 가능한 기능으로 명시적으로 모델링함으로써 기하 상관관계를 반영함으로써 데이터 효율성과 일반화 능력을 향상시킨다.
• 모델 복잡도와 추론 비용을 줄이면서도 분자 성질 예측에서 최신 기술 수준의 성능을 유지하거나 초월한다.

제안 방법

• 모델은 원자 기하 정보의 동적이고 학습 가능한 표현으로서 구면 조화 좌표(SPHCs)를 도입하여 고정되거나 수작업으로 만든 기하 특징을 대체한다.
• 불변 원자 특징과 등변 SPHCs를 분리함으로써 각각에 대해 독립적인 주목적 메커니즘이 작동할 수 있도록 하여 SPHC 공간에서 局부 및 비국소 주목적을 가능하게 한다.
• 구면 자기주목적 메커니즘은 유클리드 공간의 연속 필터를 구면 조화 공간으로 일반화하는 구면 필터를 사용하여 SO(3) 등변성을 유지한다.
• 정보를 장거리로 전파하기 위해 SPHCs에 대한 주목적을 사용한 메시지 전달 아키텍처를 사용함으로써 비국소 전자 효과를 포착할 수 있다.
• 에너지와 힘에 대한 병합 손실을 사용하여 모델을 학습하며, 에너지가 단독으로 감독될 경우 최소제곱 피팅을 통해 에너지 통합 상수를 계산한다.
• SPHCs는 학습 중에 업데이트되므로, 분자 구조 내 고차 상관관계를 반영하는 최적의 기하 표현을 학습할 수 있다.

실험 결과

연구 질문

• RQ1 학습 가능한 주목적 기반의 구면 조화 좌표 표현이 분자 시스템 내 비국소 전자 효과의 모델링을 향상시킬 수 있는가?
• RQ2 SPHCs를 통해 기하 정보를 원자 특징에서 분리함으로써 기계학습력장에서의 데이터 효율성과 일반화 능력이 향상되는가?
• RQ3 구면 자기주목적 메커니즘이 분자 시스템에서 임의의 길이 척도에 걸쳐 장거리 상관관계를 효과적으로 포착할 수 있는가?
• RQ4 고차 기하 상관관계를 포함함으로써 모델 복잡도는 어느 정도 감소하고 예측 정확도는 유지되거나 향상되는가?
• RQ5 기존 최신 기술 수준의 모델들과 비교했을 때 So3krates는 파라미터 효율성, 학습 속도, 추론 속도에서 어떤가?

주요 결과

• So3krates는 QM7-X250 및 MD17를 포함한 여러 벤치마크에서 최신 기술 수준의 성능을 달성하며, 이전 최고 수준의 모델 대비 모델 파라미터 수가 0.25–0.4배 감소했다.
• 다른 주목적 기반 모델들과 비교해 학습 시 6–14배 빠르고 추론 시 2–11배 빠른 성능을 기록했으며, 높은 정확도를 유지함에도 불구하고.
• So3krates는 커불렌의 수소 로터에서 장거리 상관관계를 성공적으로 포착하여, 국소 모델이 실패하는 바늘에서 아비-아이티오 기준값을 충실하게 재현했다.
• SPHCs를 통한 고차 기하 상관관계의 포함은 데이터 효율성과 일반화 능력 향상을 이끌어내었으며, 제한된 데이터에서의 성능 향상으로 이를 입증했다.
• 모델이 아키텍처의 일부로 SPHCs를 학습함으로써 효과적인 비국소 주목적을 실현했으며, 고정되거나 수작업으로 만든 기하 특징을 사용하는 모델들을 능가했다.
• 힘만으로 학습된 경우, 최소제곱 피팅을 통해 에너지 통합 상수를 계산하여 추가 감독 없이도 정확한 에너지 예측을 보장했다.