[논문 리뷰] A simple linear algebra identity to optimize Large-Scale Neural Network Quantum States

Riccardo Rende, Luciano Loris Viteritti외 3명|arXiv (Cornell University)|2023. 10. 09.

Quantum many-body systems참고 문헌 45인용 수 4

한 줄 요약

이 논문은 선형 대수 항등식의 새로운 응용을 제시하여 대규모 신경망 양자 상태(NNQS)에서 정확한 스토하스틱 리컨피그레이션(SR) 최적화를 가능하게 한다. 기존의 $P \times P$ 행렬 역행렬 계산의 성능 저하 문제를 $M \times M$ 문제로 줄여 해결하며, 267,720개의 파라미터와 단 6,000개의 몬테카를로 샘플을 사용하여 $10 \times 10$ 격자에서 $J_2/J_1 = 0.5$인 $J_1$-$J_2$ 헤이젠베르크 모형에 대해 최신 기준 성능을 달성한다.

ABSTRACT

Neural-network architectures have been increasingly used to represent quantum many-body wave functions. These networks require a large number of variational parameters and are challenging to optimize using traditional methods, as gradient descent. Stochastic Reconfiguration (SR) has been effective with a limited number of parameters, but becomes impractical beyond a few thousand parameters. Here, we leverage a simple linear algebra identity to show that SR can be employed even in the deep learning scenario. We demonstrate the effectiveness of our method by optimizing a Deep Transformer architecture with $3 imes 10^5$ parameters, achieving state-of-the-art ground-state energy in the $J_1$-$J_2$ Heisenberg model at $J_2/J_1=0.5$ on the $10 imes10$ square lattice, a challenging benchmark in highly-frustrated magnetism. This work marks a significant step forward in the scalability and efficiency of SR for Neural-Network Quantum States, making them a promising method to investigate unknown quantum phases of matter, where other methods struggle.

연구 동기 및 목표

• 파라미터 수 $P$가 몬테카를로 샘플 수 $M$을 초과하는 대규모 신경망 양자 상태(NNQS)에서 스토하스틱 리컨피그레이션(SR)의 계산 불가능성을 해결한다.
• 수천만 개의 파라미터를 가진 딥 트랜스포머와 같은 깊은 아키텍처에 대해 기존의 $P \times P$ 행렬 역행렬 계산으로는 비가능한 정확한 SR 최적화를 가능하게 한다.
• 강한 불균형을 띠는 양자 스핀 체계, 예를 들어 $J_2/J_1 = 0.5$에서의 $J_1$-$J_2$ 헤이젠베르크 모형과 같이 다른 방법으로는 어려운 고정밀도 기저 상태 에너지 추정을 달성한다.
• NNQS에서 높은 파라미터 수가 큰 샘플 수 $M$이 필요하지 않다는 점을 보여주며, 기존의 $M \gg P$가 필수라는 믿음에 도전한다.
• 파동함수의 부호 패턴에 대한 가정 없이도 SR의 적용 범위를 복잡한 부호-구조 없는 양자 체계로 확장한다.

제안 방법

• 기존 SR에서 발생하는 $P \times P$ 행렬 역행렬 문제를 $M \times M$ 행렬 역행렬 문제로 변환하는 단순한 선형 대수 항등식을 활용한다. 여기서 $M$은 몬테카를로 샘플 수이다.
• SR 프레임워크 내의 피셔 정보 행렬에 이 항등식을 적용하여 근사 없이 정확한 역행렬 계산을 가능하게 한다.
• 변환된 $M \times M$ 행렬을 사용해 신경망 파라미터의 SR 업데이트 방향을 계산함으로써 SR의 기하 최적화 이점을 유지한다.
• JAX를 사용한 미분 가능한 신경망 계산과 mpi4jax를 활용한 분산 병렬 처리를 통한 변분 몬테카를로 프레임워크 내에 이 방법을 구현한다.
• 수정된 SR 업데이트를 실수형 딥 뷰전 트랜스포머(ViT) 아키텍처와 복소수형 완전 연결 출력 레이어를 조합하여 복소수 양자 상태를 표현한다.
• 최적화 과정에서 물리적 대칭성(이동, 회전, 반사, 스핀 페어리티)을 단계적으로 복원하며 일관된 에너지 감소를 관찰한다.

실험 결과

연구 질문

• RQ1 파라미터 수 $P$가 샘플 수 $M$을 크게 초과하는 대규모 NNQS에 대해 $P \times P$ 역행렬 계산 없이 정확한 SR을 적용할 수 있는가?
• RQ2 제안된 선형 대수 항등식이 양자 many-body 체계에서 수십만 개의 파라미터를 가진 깊은 신경망의 효율적이고 안정적인 최적화를 가능하게 하는가?
• RQ3 이 방법이 $J_2/J_1 = 0.5$에서 매우 불균형한 양자 스핀 모형인 $J_1$-$J_2$ 헤이젠베르크 모형에 대해 최신 기준의 기저 상태 에너지 추정치를 달성할 수 있는가?
• RQ4 파라미터 수 $P \approx 267,720$일지라도 단지 $M \approx 6,000$개의 샘플로도 고정밀도 SR을 달성할 수 있는가?
• RQ5 기본 상태의 부호 구조에 대한 가정이 없는 상황에서 물리적 대칭성의 단계적 복원이 변분 에너지에 어떤 영향을 미치는가?

주요 결과

• 제안된 방법은 단지 6,000개의 몬테카를로 샘플만으로도 파라미터 수가 최대 267,720개에 이르는 NNQS에 대해 정확한 SR을 가능하게 하며, $P \times P$ 역행렬 계산의 성능 저하 문제를 회피한다.
• 이 방법은 $10 \times 10$ 격자에서 $J_2/J_1 = 0.5$인 $J_1$-$J_2$ 헤이젠베르크 모형에 대해 단위 당 기저 상태 에너지를 $-0.49575(3)$로 도출하여 새로운 최고 성능 기록을 수립한다.
• 딥 뷰전 트랜스포머 아키텍처에서 이동, 회전, 반사, 스핀 페어리티 대칭성을 단계적으로 복원할수록 에너지가 일관되게 감소함을 확인하여, 이 방법의 안정성과 물리적 일관성을 입증한다.
• 최적화 과정은 초기값의 시드에 민감하지 않으며, 여러 런에 걸쳐 일관된 결과를 보이며 높은 신뢰성을 입증한다.
• 기본 상태의 부호 구조를 가정하지 않아도 성능이 우수하여, 이러한 가정이 불가능한 광범위한 양자 시스템에 적용 가능하다.
• 기존의 $M \gg P$가 정확한 SR에 필수적이라는 전제를 도전하며, $P \approx 267,720$일지라도 $M \approx 6,000$으로도 충분함을 보여준다.