[논문 리뷰] Better, Faster Fermionic Neural Networks

James S. Spencer, David Pfau외 2명|arXiv (Cornell University)|2020. 11. 13.

Machine Learning in Materials Science참고 문헌 15인용 수 7

한 줄 요약

이 논문은 네트워크의 넓이와 행렬식 수를 늘림으로써 페르미온 신경망(FermiNet)을 개선하여 아르곤과 같은 이차 원소에서 화학 정확도를 달성한다. JAX로 이식하고 아키텍처를 단순화함으로써 학습 시간이 10배 감소하여, 바이시클부탄 및 사이클부타디엔의 자가이sovomerization과 같은 도전적인 시스템에서 최신 기술 수준의 성능을 달성한다.

ABSTRACT

The Fermionic Neural Network (FermiNet) is a recently-developed neural network architecture that can be used as a wavefunction Ansatz for many-electron systems, and has already demonstrated high accuracy on small systems. Here we present several improvements to the FermiNet that allow us to set new records for speed and accuracy on challenging systems. We find that increasing the size of the network is sufficient to reach chemical accuracy on atoms as large as argon. Through a combination of implementing FermiNet in JAX and simplifying several parts of the network, we are able to reduce the number of GPU hours needed to train the FermiNet on large systems by an order of magnitude. This enables us to run the FermiNet on the challenging transition of bicyclobutane to butadiene and compare against the PauliNet on the automerization of cyclobutadiene, and we achieve results near the state of the art for both.

연구 동기 및 목표

• 이전 버전이 화학 정확도에 미치지 못했던 이차 원소, 예를 들어 아르곤에서 FermiNet의 정확도를 확장한다.
• 대규모 다전자 시스템의 학습 시간과 계산 비용을 크게 줄인다.
• 아키텍처 개선과 최적화된 학습 프로토콜을 통해 수렴성과 에너지 정확도를 향상시킨다.
• 다중참조 시스템인 사이클부타디엔 자가이sovomerization과 바이시클부탄 전이 상태에서 FermiNet을 PauliNet과 비교하여 벤치마킹한다.
• 신경망 기반 양자화학 파동함수 안사에서 정확도와 학습 속도 사이의 상호 상충 관계를 평가한다.

제안 방법

• 표현력과 정확도를 향상시키기 위해 FermiNet 아키텍처의 일전자 스트림 너비와 행렬식 수를 늘렸다.
• FermiNet의 구현을 텐서플로우에서 JAX로 이식하여 더 빠른 계산과 더 나은 자동 미분를 가능하게 하였다.
• 복합성을 줄이기 위해 전체 공분산 엔벨롭을 제거하고 등방성 가중치 초기화를 사용함으로써 네트워크를 단순화시켰다.
• 무작위 몬테카를로(MCMC) 샘플링 단계를 각 가중치 갱신당 10단계에서 50단계로 늘여 학습 중 평형화를 향상시키고 노이즈를 감소시켰다.
• 큰 시스템에서 수렴을 보장하기 위해 20만 번에서 30만~40만 번의 가중치 갱신을 수행하였다.
• 에너지 기대값을 최소화하기 위해 확률적 경량 최적화 방법과 KFAC 최적화를 사용한 변분 몬테카를로 방법을 사용하였다.

실험 결과

연구 질문

• RQ1 FermiNet 아키텍처의 크기만 증가시켜도 이차 원소, 예를 들어 아르곤에서 화학 정확도를 달성할 수 있는가?
• RQ2 아키텍처 단순화와 JAX 이식로 인해 정확도를 유지하면서 학습 시간을 얼마나 줄일 수 있는가?
• RQ3 바이시클부탄에서 부타디엔으로의 전이 및 사이클부타디엔 자가이sovomerization과 같은 다중참조 시스템에서 FermiNet은 PauliNet에 비해 어떻게 성능을 내는가?
• RQ4 도전적인 유기 반응에서 기초 상태와 전이 상태 간의 상대 에너지를 계산할 때 FermiNet의 상대 정확도는 어떠한가?
• RQ5 왜 FermiNet은 PauliNet보다 학습 속도가 느리며, 이 격차를 정확도를 잃지 않고 줄일 수 있는가?

주요 결과

• 행렬식 수와 일전자 스트림의 너비를 늘림으로써 FermiNet은 황을 제외한 모든 이차 원소에서 화학 정확도(1 kcal/mol 이내)를 달성하였다.
• 아르곤에 대해 FermiNet은 -527.5405(1) $E_h$의 에너지를 기록하여 정확한 값과 화학 정확도 내에서 일치하였다.
• 학습 시간은 텐서플로우 사용 시 11,520 GPU 시간에서 JAX와 단순화된 아키텍처를 사용한 1,104 GPU 시간으로 10배 감소하였다.
• 바이시클부탄에서 부타디엔으로의 전이에서 FermiNet은 최고의 가용 방법과 비교해 화학 정확도 내의 상대 에너지 차이를 달성하였다.
• 사이클부타디엔 자가이sovomerization에서 FermiNet은 PauliNet보다 약 70 m$E_h$ 낮은 에너지를 기록하였으며, 상대 에너지는 10.3 ± 0.1 kcal/mol로 MR-CC 결과와 유사하였다.
• 단위 반복당 계산 비용이 PauliNet보다 약 5배 높지만, 다중참조 시스템에서 절대 에너지 정확도 측면에서 뛰어난 성능을 보였다.