[논문 리뷰] Equiformer: Equivariant Graph Attention Transformer for 3D Atomistic Graphs

Yi-Lun Liao, Tess Smidt|arXiv (Cornell University)|2022. 06. 23.

Machine Learning in Materials Science인용 수 63

한 줄 요약

Equiformer는 3D 원자 그래프에서 양자 특성을 예측하기 위해 불변 표현(irrep)과 등가 그래프 어텐션 메커니즘을 통합하는 SE(3)/E(3)-equivariant 그래프 트랜스포머로, 경쟁력 있는 정확도와 향상된 효율성을 보인다.

ABSTRACT

Despite their widespread success in various domains, Transformer networks have yet to perform well across datasets in the domain of 3D atomistic graphs such as molecules even when 3D-related inductive biases like translational invariance and rotational equivariance are considered. In this paper, we demonstrate that Transformers can generalize well to 3D atomistic graphs and present Equiformer, a graph neural network leveraging the strength of Transformer architectures and incorporating SE(3)/E(3)-equivariant features based on irreducible representations (irreps). First, we propose a simple and effective architecture by only replacing original operations in Transformers with their equivariant counterparts and including tensor products. Using equivariant operations enables encoding equivariant information in channels of irreps features without complicating graph structures. With minimal modifications to Transformers, this architecture has already achieved strong empirical results. Second, we propose a novel attention mechanism called equivariant graph attention, which improves upon typical attention in Transformers through replacing dot product attention with multi-layer perceptron attention and including non-linear message passing. With these two innovations, Equiformer achieves competitive results to previous models on QM9, MD17 and OC20 datasets.

연구 동기 및 목표

• 3D 기하학적 귀납 편향을 도입하여 Transformer가 3D 원자 그래프에 일반화하도록 동기를 부여하고 가능하게 한다.
• Transformer 연산을 등가 대응으로 대체하고 텐서 곱을 추가하여 Equiformer를 개발한다.
• 내용 및 기하 정보를 결합하기 위한 등가 그래프 어텐션 메커니즘을 도입한다.
• QM9, MD17, OC20에서 평가하여 최첨단 모델에 비해 경쟁력 있는 성능을 보여준다.

제안 방법

• 타입-irrep 피처에서 텐서 곱을 포함한 SE(3)/E(3)-등각 연산으로 Transformer 구성 요소를 대체한다.
• 깊이별 텐서 곱을 사용하여 타입-L 벡터를 효율적으로 결합하고 매개변수 증가를 제어한다.
• 스칼라(타입-0) 피처에서 MLPL 어텐션을 사용하고 상위 차수 피처에서 비선형 메시지 전달을 활용하는 새로운 등가 그래프 어텐션을 채택한다.
• 입력을 원자 및 엣지 차수 임베딩으로 임베딩하고 거리 조건 가중치를 위한 방사 기저 및 방사 함수를 도입한다.
• 스칼라이 아닌 irrep에 게이트 기반 비선형 활성화를 적용하여 등각성을 보존하고, 계산 복잡성을 관리하기 위해 깊이별 접근 방식을 사용한다.

실험 결과

연구 질문

• RQ1 SE(3)/E(3)-equivariant 피처를 장착했을 때 Transformer가 3D 원자 그래프에 효과적으로 일반화할 수 있는가?
• RQ2 제안된 등가 그래프 어텐션이 3D 그래프에 대해 표준 점곱 어텐션보다 더 효과적인가?
• RQ3 Equiformer를 등가 연산으로 최소 수정하여 QM9, MD17, OC20에서 기존의 SE(3)/E(3)-등각 아키텍처를 일치하거나 능가하는가?
• RQ4 Equiformer에서 고차 irrep(Lmax)을 사용했을 때 학습 시간과 확장성의 트레이드오프는 무엇인가?

주요 결과

• Equiformer는 Lmax=2-3에서 QM9, MD17, OC20 데이터셋에서 이전 모델과 비교하여 경쟁력 있는 성능을 달성한다.
• QM9에서 Equiformer는 NequIP 및 TorchMD-NET를 포함한 여러 베이스라인보다 12개 회귀 태스크에서 전반적으로 더 나은 결과를 보인다.
• MD17에서 더 높은 Lmax를 가진 Equiformer가 TorchMD-NET 및 NequIP와 비교하여 여러 분자에서 에너지 및 힘 MAE를 개선한다.
• OC20에서 Equiformer는 IS2RE 데이터로 학습될 때 상태 예측 모델인 SEGNN 및 Graphormer보다 IS2RE 성능이 향상되며 이전 모델과 비교해 학습 시간을 2.3배에서 15.5배까지 줄일 수 있다.
• Equiformer의 등가 그래프 어텐션(비선형 메시지를 가진 MLPL 기반 어텐션)은 점곱 어텐션보다 표현력이 더 크고 등가 트랜스포머의 성능을 향상시킨다.
• 일반적으로 Lmax를 올리면 OC20 및 MD17의 성능이 향상되지만 일부 케이스에서 과적합이 발생할 수 있다.