[논문 리뷰] Scaling the leading accuracy of deep equivariant models to biomolecular simulations of realistic size

Albert Musaelian, Anders Johansson외 2명|arXiv (Cornell University)|2023. 04. 20.

Machine Learning in Materials Science인용 수 14

한 줄 요약

Allegro는 확장 가능한 등변 신경 포텐셜과 대규모 GPU 병렬화를 결합하여 양자 정확도에 근접한 확장 가능한 생체분자 시뮬레이션을 제공합니다. 이로써 Perlmutter에서 단나노초 단백질 다이나믹스와 44-million-atom HIV 캡시드 시뮬레이션을 강/약 스케일링으로 가능하게 합니다.

ABSTRACT

This work brings the leading accuracy, sample efficiency, and robustness of deep equivariant neural networks to the extreme computational scale. This is achieved through a combination of innovative model architecture, massive parallelization, and models and implementations optimized for efficient GPU utilization. The resulting Allegro architecture bridges the accuracy-speed tradeoff of atomistic simulations and enables description of dynamics in structures of unprecedented complexity at quantum fidelity. To illustrate the scalability of Allegro, we perform nanoseconds-long stable simulations of protein dynamics and scale up to a 44-million atom structure of a complete, all-atom, explicitly solvated HIV capsid on the Perlmutter supercomputer. We demonstrate excellent strong scaling up to 100 million atoms and 70% weak scaling to 5120 A100 GPUs.

연구 동기 및 목표

• 대규모 생체분자 시스템에 대해 확장 가능하고 정확한 등변 심층 학습 포텐셜을 입증한다.
• 최첨단 등변 모델의 정확도를 엑사스케일 계산 능력과 연결한다.
• 전례 없는 시스템 규모에서의 안정적이고 긴 시뮬레이션(단백질 동역학)을 showcases한다.

제안 방법

• 시스템 에너지를 원자 단위 구성요소와 순서화된 쌍 구성요소로 분해하는 확장 가능한 Allegro 등변 MLIP를 도입한다.
• 스칼라 트랙과 등변 텐서 트랙의 2트랙 구조를 사용해 계산 비용을 관리하면서 용량을 확장한다.
• .strided 메모리 레이아웃과 최적화된 텐서 곱을 이용해 효율적인 GPU 실행과 단일 패스 수축을 가능하게 한다.
• 속도와 수치적 안정성의 균형을 맞추기 위해 혼합 정밀도 계산(TF32와 FP32, 때때로 FP64)을 적용한다.
• 근접 이웃 계산을 정확도를 해치지 않으면서 줄이기 위해 종별 짝 컷오프를 도입한다.
• 기존 MD 공간 분해를 활용하기 위해 Allegro를 GPU 내 거주 파이프라인으로 LAMMPS와 통합한다.

실험 결과

연구 질문

• RQ1 Allegro가 수억 개의 원자 규모에서도 최첨단 등변 정확도를 달성할 수 있는가?
• RQ2 현대 GPU에서 대형 생체분자 시스템에 대한 Allegro의 성능(타임스텝/초, ns/일)은 어떠한가?
• RQ3 현실적인 생체분자 구조에서 긴 시뮬레이션을 다룰 때의 안정성과 정밀도는 어떻게 되는가?
• RQ4 정밀도를 해치지 않으면서 MD를 위한 확장 가능한 등변 학습을 가능하게 하는 건축학적 혁신은 무엇인가?
• RQ5 생체분자 시스템의 정확도와 확장성 측면에서 Allegro가 기존 MLIPs 및 고전적 힘场과 비교하여 어떤 차이가 있는가?

주요 결과

• Allegro는 벤치마크에서 높은 정확도를 달성하고 Perlmutter-규모 하드웨어에서 126.4 million atoms까지 확장되며 100 timesteps/s 성능을 보인다.
• 100 million atoms까지의 강한 스케일링과 5120 A100 GPU(1280 노드)에서의 70% 약한 스케일링을 시연한다.
• 대규모 생체분자 시뮬레이션에 대해 최첨단 ML 인터원자 포텐셜의 최초 적용이다.
• 단백질 동역학의 나노초 길이 안정 시뮬레이션(DHFR, Factor IX 등) 및 44M-atom HIV 캡시드 용매화 시스템을 생성한다.
• 스칼라 트랙과 텐서 트랙을 Strided 메모리 레이아웃 및 최적화된 텐서 곱과 결합해 정확도를 유지하면서 계산 비용을 감소시킨다.
• 종별 짝 컷오프와 혼합 정밀도 기법을 도입해 GPU에서 처리량을 극대화한다.