[논문 리뷰] TensorNet: Cartesian Tensor Representations for Efficient Learning of Molecular Potentials

Guillem Simeon, Gianni De Fabritiis|arXiv (Cornell University)|2023. 06. 10.

Machine Learning in Materials Science인용 수 17

한 줄 요약

TensorNet은 Cartesian 랭크-2 텐서를 사용하여 O(3)-등각 신경망을 통해 분자 포텐셜을 효율적으로 학습하고, 일부 구면-텐서 모델에 비해 매개변수가 훨씬 적고 계산이 더 낮은 상태에서 최첨단 정확도를 달성합니다. 또한 에너지와 힘과 함께 벡터 및 텐서 분자 양을 예측할 수 있게 합니다.

ABSTRACT

The development of efficient machine learning models for molecular systems representation is becoming crucial in scientific research. We introduce TensorNet, an innovative O(3)-equivariant message-passing neural network architecture that leverages Cartesian tensor representations. By using Cartesian tensor atomic embeddings, feature mixing is simplified through matrix product operations. Furthermore, the cost-effective decomposition of these tensors into rotation group irreducible representations allows for the separate processing of scalars, vectors, and tensors when necessary. Compared to higher-rank spherical tensor models, TensorNet demonstrates state-of-the-art performance with significantly fewer parameters. For small molecule potential energies, this can be achieved even with a single interaction layer. As a result of all these properties, the model's computational cost is substantially decreased. Moreover, the accurate prediction of vector and tensor molecular quantities on top of potential energies and forces is possible. In summary, TensorNet's framework opens up a new space for the design of state-of-the-art equivariant models.

연구 동기 및 목표

• 분자 시뮬레이션과 약물 발견에서 효율적이고 정확한 원자간 포텐셜의 필요성을 제시한다.
• 원자 임베딩에 Cartesian 랭크-2 텐서를 사용하는 O(3)-등각 아키텍처를 제안한다.
• 고차 순의 구면 모델과 비교한 계산 효율성과 매개변수 효율성을 입증한다.
• 에너지, 힘, 텐서 양을 예측하는 모델의 능력을 보여준다.
• 다양한 데이터셋에서 TensorNet을 평가하여 더 적은 매개변수로 최첨단 성능을 입증한다.

제안 방법

• 원자를 전체 3x3 Cartesian 텐서로 표현하고 이것을 스칼라, 벡터, 텐서 구성요소(I, A, S)로 분해한다.
• 상대 위치와 원자 번호로부터 간선 기반 초기화를 사용하여 간선별 텐서 특징을 형성한다.
• 불변 노름을 계산하고 이를 이용해 패리티 보존 행렬 연산으로 원자별 업데이트를 생성한다.
• 레이어 간에 O(3) 등각성과 패리티를 보존하는 매트릭스 곱 스타일 상호작용을 적용한다.
• 최종 텐서 노름에서 유도된 원자별 출력의 합으로 총 에너지를 예측한다.
• 텐서 구성요소에서 추출하여 벡터 및 텐서 출력을 가능하게 하며, 비대칭 부분을 포함한다.

Figure 1: Key steps, from top to bottom, in the embedding and interaction modules for some central atom $i$ and neighbors $j$ and $k$ found within the cutoff radius. a) Relative position vectors are used to initialize edge-wise tensor components, modified using edge-wise invariant functions, and sum

실험 결과

연구 질문

• RQ1 카르테시안 랭크-2 텐서 표현이 더 높은 차수의 구면 모델보다 더 적은 매개변수로 분자 포텐셜에 대해 최첨단 정확도를 달성할 수 있는가?
• RQ2 간단한 매트릭스 곱과 불가약 분해로 된 O(3) 등각 텐서 연산이 정확한 에너지, 힘 및 텐서 특성 예측에 충분한가?
• RQ3 다양한 분자 데이터셋(QM9, rMD17, SPICE/ANI1x/COMP6)에서 TensorNet의 성능은 기존 기준선과 비교해 어떤가?
• RQ4 등각성, 상호작용 설계 및 컷오프가 예측 정확도에 미치는 영향은 무엇인가?
• RQ5 TensorNet이 스칼라, 벡터, 텐서 분자 특성을 동시에 신뢰성 있게 예측할 수 있는가?

주요 결과

• TensorNet은 QM9의 U0, U, H에서 Allegro 및 MACE를 능가하며 3.9(1) meV 스타일의 참조로—Allegro 매개변수의 23%로 최첨단 정확도를 달성한다.
• rMD17에서 1–2 상호작용 레이어를 가진 TensorNet은 에너지 및 힘 MAE에서 경쟁력을 보이며 다른 구면 모델보다 훨씬 적은 매개변수(최저 0.535–0.770M)로 달성한다.
• 가벼운 1L 모델로 ANI1x/COMP6 일반화를 강하게 달성하여 여러 벤치마크에서 여러 기준선보다 더 낮은 에너지 및 힘 MAE를 제공한다.
• 진공 상태의 에탄올에 대해 TensorNet은 FieldSchNet 및 PaiNN보다 약 2배 더 나은 에너지 및 힘 MAE를 달성하고, 또한 정확한 쌍극자 모멘트와 극화도 예측을 제공한다.
• 적소 연구는 전체 O(3) 등각성과 상호작용 곱이 정확도를 크게 개선하고, 0L 대 1L/2L의 트레이드오프는 컷오프와 분자 크기에 따라 달라진다는 것을 보여준다.
• TensorNet은 GPU에서 추론 및 학습 속도가 유리함을 보여주며 수백 원자에 이르는 분자에서도 효율성을 유지하고 더 큰 시스템에서 축소된 등각 변환기들과 경쟁한다.

(b) Interaction and node update