[논문 리뷰] UNiTE: Unitary N-body Tensor Equivariant Network with Applications to Quantum Chemistry.

Zhuoran Qiao, Anders S. Christensen외 4명|arXiv (Cornell University)|2021. 05. 31.

Machine Learning in Materials Science참고 문헌 84인용 수 23

한 줄 요약

UNiTE는 대칭 텐서를 다룰 수 있는 선형 시간 복잡도를 가지는 유일한 N체 텐서 등변 신경망을 소개한다. 이는 양자화학 분야에서 확장 가능하고 대칭성을 유지하는 학습을 가능하게 하며, 벤치마크에서 최신 기술 대비 110% 이상의 성능 향상을 달성한다. 또한 전통적인 수치적 방법보다 3개 정도의 지수 빠르게 작동하는 강력한 제로샷 일반화를 실현한다.

ABSTRACT

Equivariant neural networks have been successful in incorporating various types of symmetries, but are mostly limited to vector representations of geometric objects. Despite the prevalence of higher-order tensors in various application domains, e.g. in quantum chemistry, equivariant neural networks for general tensors remain unexplored. Previous strategies for learning equivariant functions on tensors mostly rely on expensive tensor factorization which is not scalable when the dimensionality of the problem becomes large. In this work, we propose unitary $N$-body tensor equivariant neural network (UNiTE), an architecture for a general class of symmetric tensors called $N$-body tensors. The proposed neural network is equivariant with respect to the actions of a unitary group, such as the group of 3D rotations. Furthermore, it has a linear time complexity with respect to the number of non-zero elements in the tensor. We also introduce a normalization method, viz., Equivariant Normalization, to improve generalization of the neural network while preserving symmetry. When applied to quantum chemistry, UNiTE outperforms all state-of-the-art machine learning methods of that domain with over 110% average improvements on multiple benchmarks. Finally, we show that UNiTE achieves a robust zero-shot generalization performance on diverse down stream chemistry tasks, while being three orders of magnitude faster than conventional numerical methods with competitive accuracy.

연구 동기 및 목표

• 양자화학과 같은 분야에서 고차원 텐서에 대해 확장 가능한 등변 신경망의 부족을 해결하기 위해.
• 기존의 텐서 분해 기반 방법의 등변 학습에 대한 확장성 한계를 극복하기 위해.
• 대칭 N체 텐서에 대해 유일한 군 등변성(예: 3차원 회전)을 유지하는 신경망 아키텍처를 개발하기 위해.
• 새로운 대칭 유지 정규화 기법인 등변 정규화를 통해 일반화 성능을 향상시키기 위해.
• 고속이고 정확한 성능을 보장하면서 다양한 응용 화학 작업에 대해 제로샷 일반화를 가능하게 하기 위해.

제안 방법

• 3차원 회전과 같은 유일한 군 작용에 대해 등변성을 유지하는 대칭 N체 텐서를 위한 새로운 신경망 아키텍처인 UNiTE를 제안한다.
• 비용이 많이 들지 않는 인자분해 없이도 등변성을 유지할 수 있도록 체계적으로 텐서 분해 원리를 활용하며, 비용 효율적인 설계로 비례하여 선형 시간 복잡도를 달성한다.
• 대칭성을 유지하면서 모델의 일반화 성능을 향상시키는 등변 정규화라는 정규화 방법을 도입한다.
• 등변성을 유지하기 위해 체계적으로 텐서 분해 원리를 적용하며, 비용이 많이 드는 인자분해를 피한다.
• 특징 변환을 위한 학습 가능한 등변 레이어를 갖춘 N체 상호작용 기반 메시지 전달 프레임워크를 사용한다.
• 양자 시스템의 다체 상호작용을 표현할 수 있는 대칭 텐서 표현 방식을 사용한다.

실험 결과

연구 질문

• RQ1 일반적인 대칭 N체 텐서에 대해 유일한 군 작용(예: 3차원 회전)에 대해 등변성을 유지하는 신경망 아키텍처를 설계할 수 있는가?
• RQ2 비제로 텐서 요소의 수에 대해 선형 시간 복잡도를 달성할 수 있는가? 이를 통해 고차원 문제에 대한 확장성이 보장되는가?
• RQ3 등변 정규화는 대칭성을 유지하면서 일반화 성능을 향상시키는가?
• RQ4 UNiTE는 기존 방법 대비 상당한 성능 향상을 보이며 양자화학 벤치마크에서 최신 기술 수준을 초월할 수 있는가?
• RQ5 UNiTE는 다양한 화학 작업에 대해 제로샷 설정에서도 강력한 일반화 성능을 보이며, 전통적인 수치적 해법보다 빠른 속도로 성능을 뛰어넘을 수 있는가?

주요 결과

• UNiTE는 최신 기계학습 방법 대비 여러 양자화학 벤치마크에서 평균 110% 이상의 성능 향상을 달성한다.
• 다양한 응용 화학 작업에 걸쳐 강력한 제로샷 일반화 성능을 보이며, 대칭성에서 유래한 강력한 인덕티브 바이어스를 확인할 수 있다.
• UNiTE는 전통적인 수치적 방법보다 3개 정도의 지수 빠르게 작동하면서도 경쟁 가능한 정확도를 유지한다.
• 비제로 텐서 요소의 수에 대해 선형 시간 복잡도를 유지함으로써 고차원 시스템으로의 확장 가능성을 확보한다.
• 등변 정규화는 대칭성을 해치지 않으면서 모델의 일반화 성능을 향상시키며, 새로운 작업에 대한 성능 향상에 기여한다.
• UNiTE는 기존 등변 신경망이 다루지 못했던 고차원 텐서에서 등변 함수를 성공적으로 학습한다.