[논문 리뷰] Scalable variational Monte Carlo with graph neural ansatz

Li Yang, Wen-Jun Hu외 1명|arXiv (Cornell University)|2020. 11. 25.

Model Reduction and Neural Networks참고 문헌 42인용 수 11

한 줄 요약

이 논문은 그래프 컬러션 네트워크(GCNs) 기반의 그래프 신경망 앙스타즈(GNA)를 소개하여, 양자 다체계의 유니버설이고 기하학적 비의존성 변분 몬테카를로(VMC) 시뮬레이션을 가능하게 한다. 128개의 TPU 코어에서의 파rameter 공유 및 확장 가능한 분산 학습을 활용하여, 2D 카고메, 삼각형, 무작위로 연결된 격자까지 최대 432개의 스ites에 대해 정확한 기본 상태 에너지 예측을 달성하였으며, 소규모 사전학습된 시스템으로부터의 전이성 덕분에 대규모 시스템에서 빠른 수렴을 이룩하였다.

ABSTRACT

Deep neural networks have been shown as a potentially powerful ansatz in variational Monte Carlo for solving quantum many-body problems. We propose two improvements in this direction. The first is graph neural ansatz (GNA), which is a variational wavefunction universal to arbitrary geometry. GNA results in accurate ground-state energies on 2D Kagome lattices, triangular lattices, and randomly connected graphs. Secondly, we design a distributed workflow on multiple accelerators to scale up the computation. We compute Kagome lattices with sizes up to 432 sites on 128 TPU cores. The parameter sharing nature of the GNA also leads to transferability across different system sizes and geometries.

연구 동기 및 목표

• 임의의 격자 기하학, 특히 비주기적 및 비정규 구조를 포함한 다양한 격자 기하학을 다룰 수 있는 유니버설 변분 앙스타즈를 개발하는 것.
• 이전의 신경망 앙스타즈(예: CNNs, RBMs)의 제한을 극복하는 것. 이들은 격자 대칭성 또는 격자 기반 표현에 의해 제약을 받는다.
• 단일 장치의 메모리 및 계산 병목 현상을 극복하기 위해 분산 가속기(예: TPU)를 사용하여 대규모 양자 시스템에서의 확장 가능한 VMC 시뮬레이션을 가능하게 하는 것.
• 학습된 GNA 모델이 다양한 시스템 크기 및 기하학 간에 전이 가능함을 보여주어 학습 시간을 단축하고 수렴성을 향상시키는 것.
• 강한 상관관계를 가진 양자 시스템을 위한 통합적이고 확장 가능한 프레임워크를 제공하며, 복소수 파동함수 및 페르미온 시스템으로의 확장 가능성도 고려하는 것.

제안 방법

• 각 격자 위치를 노드로 간주하고 메시지 전달을 통해 업데이트되는 학습 가능한 임bedding을 갖는 그래프 컬러션 네트워크(GCNs)를 사용하여 다체계 파동함수를 인코딩하는 그래프 신경망 앙스타즈(GNA)를 제안한다.
• 잔차 연결, ReLU 비선형성, 레이어 정규화를 포함한 GCN 아키텍처를 사용하며, 노드 특징은 정규화된 인접 행렬과 층 간 공유된 가중치 행렬을 통해 업데이트된다.
• 최종 집계된 노드 임베딩을 단일 밀집층을 통해 스칼라 파동함수의 로그 앰플리튜드로 매핑하여 변분 에너지의 미분 가능 최적화를 가능하게 한다.
• 중요도 샘플링 기반 기울기 최적화(ISGO)의 분산 구현을 다수의 가속기(예: TPU)에 걸쳐 적용하여, 대규모 시스템에서 미니배치 샘플링 및 파ram터 업데이트를 가능하게 한다.
• 주기적 경계 조건을 인접 행렬에서 반대편 가장자리를 연결함으로써 구현하여 토러스 기하학 시뮬레이션을 가능하게 한다.
• 소규모 시스템(예: 12개 스ites)에서 GNA를 사전학습하고 대규모 시스템(예: 432개 스ites)에서 미세조정함으로써 전이 학습을 적용하며, 다양한 기하학 간에 공유된 파ram터를 활용한다.

실험 결과

연구 질문

• RQ1 그래프 신경망 기반 앙스타즈가 아키텍처 수정 없이도 카고메, 삼각형, 무작위로 연결된 그래프를 포함한 다양한 격자 기하학에 일반화될 수 있는가?
• RQ2 GNA 내의 파ram터 공유가 소규모에서 대규모 양자 시스템으로의 전이 학습을 얼마나 잘 가능하게 하는가? 이는 학습 수렴성과 안정성 향상에 기여하는가?
• RQ3 TPU와 같은 분산 가속기를 사용하여 GNA 기반 VMC 프레임워크를 대규모 시스템 크기(예: 432개 스ites)로 확장할 수 있는가? 이는 단일 장치의 메모리 제약을 극복할 수 있는가?
• RQ4 다양한 상호작용 강도(t/V)와 격자 유형에서 정확한 대각화 결과와 비교했을 때, GNA-VMC 접근법의 기본 상태 에너지 추정 정확도는 어느 정도인가?
• RQ5 다른 기하학(예: 삼각형 또는 무작위 그래프)에서 사전학습한 모델이 목표로 하는 카고메 격자에서의 수렴을 여전히 가속화하는가?

주요 결과

• GNA는 2D 카고메 및 삼각형 격자에서 정확한 대각화 결과와 0.1% 이내의 기본 상태 에너지 예측을 달성하였으며, 큰 t/V 비율에서도 성능을 유지하였다.
• 432개 스ites를 가진 12×12×3 카고메 격자에 대해, 128개의 TPU 코어를 사용하여 기본 상태 에너지를 성공적으로 계산하였으며, 이는 이전에는 단일 가속기에서는 비가능했던 시스템 크기였다.
• 12개 스ites의 시스템에서 사전학습한 GNA는 수렴 시간을 단지 5개의 샘플링 반복으로 줄였고, 무작위 초기화된 모델은 40회의 반복 후에도 수렴하지 못했다.
• GNA는 전이 가능성을 보였으며, 삼각형 또는 무작위로 연결된 그래프에서 사전학습한 모델도 목표로 하는 카고메 격자에서 훈련을 상당히 가속화하였다. 이는 다양한 기하학 간에 강력한 일반화 능력을 지닌다는 것을 시사한다.
• 확장 가능한 ISGO 구현은 시스템을 분할하고 다수의 가속기 간에 기울기를 집계함으로써 효율적인 분산 학습을 가능하게 하여 대규모 VMC 시뮬레이션을 실현 가능하게 하였다.
• 비정규 및 비정규 기하학을 포함한 다양한 기하학에서 고정밀도를 유지함으로써, GNA가 변분 앙스타즈로서의 유니버설리티를 확인하였다.