[논문 리뷰] Molecular Energy Learning Using Alternative Blackbox Matrix-Matrix Multiplication Algorithm for Exact Gaussian Process

Jiace Sun, Lixue Cheng외 1명|arXiv (Cornell University)|2021. 09. 20.

Machine Learning in Materials Science참고 문헌 25인용 수 5

한 줄 요약

이 논문은 분자 궤도 기반 기계학습(MOB-ML)에서 정확한 가우시안 프로세스(GP) 학습을 가속화하기 위해 보조 블랙박스 행렬-행렬 곱셈(AltBBMM) 알고리즘을 도입한다. 이로 인해 기존 BBMM 대비 4배 빠른 속도로 6,500개 이상의 분자를 포함한 데이터셋에서 스케일러블한 학습이 가능해졌으며, 최신 기준 성능과 이식성은 유지되었다. 이 방법은 QM7b-T 및 GDB-13-T 데이터셋에서 보고된 바 가장 우수한 성능을 기록했으며, 분자 에너지 예측 분야의 새로운 기준을 설정했다.

ABSTRACT

We present an application of the blackbox matrix-matrix multiplication (BBMM) algorithm to scale up the Gaussian Process (GP) training of molecular energies in the molecular-orbital based machine learning (MOB-ML) framework. An alternative implementation of BBMM (AltBBMM) is also proposed to train more efficiently (over four-fold speedup) with the same accuracy and transferability as the original BBMM implementation. The training of MOB-ML was limited to 220 molecules, and BBMM and AltBBMM scale the training of MOB-ML up by over 30 times to 6500 molecules (more than a million pair energies). The accuracy and transferability of both algorithms are examined on the benchmark datasets of organic molecules with 7 and 13 heavy atoms. These lower-scaling implementations of the GP preserve the state-of-the-art learning efficiency in the low-data regime while extending it to the large-data regime with better accuracy than other available machine learning works on molecular energies.

연구 동기 및 목표

• 기존에 약 220개 분자로 제한되었던 MOB-ML에서 정확한 가우시안 프로세스 회귀의 O(N³) 계산 복잡도 문제를 해결하기 위해.
• 정확한 GP 추론과 높은 예측 정확도를 유지하면서도, 대규모 데이터셋—특히 100만 개 이상의 쌍 에너지—로의 확장 가능성을 확보하기 위해.
• 기존 BBMM의 성능을 뛰어넘는 새로운 BBMM 구현체(AltBBMM)를 개발하여 학습 효율성을 향상시키되, 모델 정확도나 이식성에 영향을 주지 않도록 하기 위해.
• 다양한 화학 체계, 특히 더 큰 분자들에서도 분자 에너지 예측에서 최신 기준 성능을 달성할 수 있음을 입증하기 위해.

제안 방법

• 논문은 블랙박스 행렬-행렬 곱셈(BBMM) 알고리즘을 활용하여 GP 학습 복잡도를 O(N³)에서 O(N²)로 감소시켜, 다중 GPU 가속과 정확한 추론을 가능하게 한다.
• 기존 BBMM의 성능을 향상시키기 위해 메모리 액세스 패턴을 최적화하고 계산 오버헤드를 줄이는 새로운 BBMM 구현체(AltBBMM)를 제안한다. 이는 더 빠른 수렴을 달성한다.
• GP 추론의 선형 시스템을 해결하기 위해 공액 그래디언트(CG) 및 블록 공액 그래디언트(BCG) 방법을 사용하며, 수렴 속도를 향상시키기 위해 조건부 조정을 적용한다.
• MOB-ML에서 대각선 및 비대각선 쌍 에너지의 학습을 분리하여 각각 별도의 모델을 사용하고, 학습된 초모수를 갖는 커널 기반 GP 회귀를 적용한다. 이는 로그 우도 최대화를 통해 수행된다.
• 행렬의 특이성 문제를 줄이기 위해 추가 백색 노이즈(σ²_add)를 정규화로 사용하며, 정확도와 학습 속도의 균형을 고려해 σ²_add = 10⁻⁵를 선택하였다.
• 프레임워크는 QM7b-T 및 GDB-13-T 데이터셋에서 평가되었으며, 모델 성능은 평균 절대 오차(MAE) 및 7개의 무거운 원자당 평균 절대 오차(MAE/7HA)로 측정되었다.

실험 결과

연구 질문

• RQ1 BBMM 알고리즘이 기존에 약 220개 분자로 제한되었던 MOB-ML에서의 정확한 GP 학습을 이전의 한계를 초월해 확장하는 데 효과적으로 적용될 수 있는가?
• RQ2 제안된 AltBBMM 구현체가 기존 BBMM 대비 유의미한 속도 향상을 이룰 수 있을 뿐 아니라, 동일한 수준의 예측 정확도와 이식성도 유지하는가?
• RQ3 AltBBMM를 활용한 MOB-ML이 QM7b-T 및 GDB-13-T와 같은 벤치마크 데이터셋에서 특히 대용량 데이터 환경에서 최신 기준 성능을 달성할 수 있는가?
• RQ4 AltBBMM의 성능는 MOB-ML/Full GP 및 DeePHF/확장된 신경망과 같은 기존 최고 수준의 방법들과 비교해 정확도와 분자 크기 범위에 걸친 일반화 능력에서 어떻게 비교되는가?

주요 결과

• 6,500개의 QM7b-T 분자를 대상으로 학습했을 때, AltBBMM는 학습 시간을 4배 빠르게 하였으며(6.24시간 대비 26.52시간), GDB-13-T에서 테스트 MAE가 1% 증가하는 데 그쳤다.
• AltBBMM로 학습한 모델은 QM7b-T에서 테스트 MAE 0.193 kcal/mol, GDB-13-T에서 7개의 무거운 원자당 MAE 0.493 kcal/mol을 기록했으며, 이는 이전 최고 수준의 방법들과 동일하거나 이를 초월하는 성능을 보였다.
• BBMM 및 AltBBMM는 MOB-ML 학습을 1,152,157개의 비대각선 쌍 에너지와 124,973개의 대각선 쌍 에너지로 총 100만 개 이상의 쌍 에너지로 확장하였으며, 이는 이전의 약 220개 분자 기준 대비 30배 증가한 것이다.
• σ²_add = 10⁻⁵를 사용할 경우 최적의 정확도를 확보할 수 있었으며, 더 높은 값(예: 10⁻¹)을 사용할 경우 학습 시간이 반으로 줄어들었음에도 불구하고 예측 MAE가 두 배로 증가하였다.
• AltBBMM는 높은 이식성을 유지하며, GDB-13-T에서 0.493 kcal/mol의 최고로 보고된 MAE/7HA를 기록하여 DeePHF/확장된 신경망 방법을 능가하였다.
• 그림 1의 학습 곡선은 AltBBMM가 저용량 및 대용량 데이터 환경 모두에서 기존 방법보다 더 빨리 화학적 정확도(1 kcal/mol)에 도달하고 더 우수한 일반화 성능을 보임을 보여준다.