[논문 리뷰] Simulation of quantum physics with Tensor Processing Units: brute-force computation of ground states and time evolution

Markus Hauru, Alan Morningstar외 4명|arXiv (Cornell University)|2021. 11. 19.

Parallel Computing and Optimization Techniques인용 수 4

한 줄 요약

이 논문은 기계 학습을 위해 설계된 Google의 TPU v3 풀을 활용해 양자 다체계의 브루트포스 시뮬레이션을 수행함으로써, 최대 38 큐비트의 기본 상태 및 시간 진화 계산을 달성한다. TPU 네이티브 행렬 곱셈 유닛(MXU), 고대역폭 메모리(HBM), 빠른 코어 간 통신 기능을 활용하여 하미르토니안의 파동함수에 대한 작용을 효율적으로 계산함으로써, QMC 및 텐서 네트워크 방법이 다루기 어려운 영역까지 정확한 시뮬레이션을 가능하게 한다.

ABSTRACT

Tensor Processing Units (TPUs) were developed by Google exclusively to support large-scale machine learning tasks. TPUs can, however, also be used to accelerate and scale up other computationally demanding tasks. In this paper we repurpose TPUs for the challenging problem of simulating quantum spin systems. Consider a lattice model made of $N$ spin-$\\frac{1}{2}$ quantum spins, or qubits, with a Hamiltonian $H = \\sum_i h_i$ that is a sum of local terms $h_i$ and a wavefunction $|\\Psi\ angle$ consisting of $2^N$ complex amplitudes. We demonstrate the usage of TPUs for both (i) computing the ground state $|\\Psi_{gs}\ angle$ of the Hamiltonian $H$, and (ii) simulating the time evolution $|\\Psi(t)\ angle=e^{-itH}|\\Psi(0)\ angle$ generated by this Hamiltonian starting from some initial state $|\\Psi(0)\ angle$. The bottleneck of the above tasks is computing the product $H |\\Psi\ angle$, which can be implemented with remarkable efficiency utilising the native capabilities of TPUs. With a TPU v3 pod, with 2048 cores, we simulate wavefunctions $|\\Psi\ angle$ of up to $N=38$ qubits. The dedicated matrix multiplication units (MXUs), the high bandwidth memory (HBM) on each core, and the fast inter-core interconnects (ICIs) together provide performance far beyond the capabilities of general purpose processors.

연구 동기 및 목표

• 기계 학습을 위해 설계된 TPU가 정확한 양자 다체계 시뮬레이션을 가속화할 수 있는지 조사하기 위해.
• 지수적 자원 증가로 인해 기존에 해결이 어려운 스핀 시스템의 브루트포스 기반 기본 상태 및 시간 진화 계산을 가능하게 하기 위해.
• 기존 방법의 한계를 극복하기 위해, 예를 들어 QMC의 사인 문제 또는 텐서 네트워크의 얽힘 증가 문제를 해결하기 위해 정확하고 전체 파동함수 기반의 시뮬레이션을 사용하기 위해.
• 2048개의 코어와 32TB의 HBM를 갖춘 TPU v3 풀을 사용해 대규모 양자 시뮬레이션의 가능성을 입증하기 위해.
• 특화된 하드웨어 플랫폼을 활용해 고성능이며 오류가 없는 양자 시뮬레이션 방법의 벤치마크를 제공하기 위해.

제안 방법

• 기계 학습을 위해 설계된 TPU v3 풀(2048개의 코어, 코어당 16GB의 HBM, 전용 MXU 포함)을 사용해 스핀-1/2 시스템의 정확한 양자 시뮬레이션을 수행한다.
• 2^N차원의 양자 파동함수를 TPU 코어의 고대역폭 메모리에 분산하여 대규모 상태 저장을 가능하게 한다.
• TPU의 MXU를 활용해 반정밀도 행렬 곱셈의 고처리량 특성을 이용해 하미르토니안 작용 H|Ψ⟩를 효율적으로 계산한다.
• 빠른 코어 간 상호연결(ICIs)을 활용해 반복 업데이트 중에 파동함수를 코어 간에 동적으로 재분배한다.
• 메모리 효율적인 이중 단계 방식으로 랑츠 알고리즘을 구현: 첫 번째 단계에서는 Krylov 벡터를 저장하지 않고 삼중대각 행렬 T_j를 계산하고, 두 번째 단계에서 Ritz 벡터로부터 기본 상태를 재구성한다.
• 파동함수 갱신 H|Ψ⟩를 기본 연산으로 사용해, Lanczos를 통한 기본 상태 계산과 지수 연산 분할을 통한 시간 진화 계산을 수행한다.

실험 결과

연구 질문

• RQ1 TPU v3 풀은 기계 학습 외의 작업을 위해 효과적으로 재사용될 수 있는가?
• RQ2 브루트포스 기반 TPU 기반 방법을 사용할 때 최대 시스템 크기(큐비트 수로 표현)는 얼마인가?
• RQ3 TPU 성능은 하미르토니안 작용 및 양자 시스템의 파동함수 진화에서 일반 목적의 HPC와 비교해 어떻게 되는가?
• RQ4 제한된 코어당 메모리 자원을 가진 TPU 하드웨어에서 랑츠 알고리즘이 효율적으로 작동하도록 적응시킬 수 있는가?
• RQ5 TPU 기반 시뮬레이션은 QMC 및 텐서 네트워크와 같은 근사 방법의 벤치마크로 얼마나 유용한가?

주요 결과

• 2048개의 코어와 32TB의 HBM를 갖춘 TPU v3 풀은 최대 38 큐비트의 양자 시스템을 성공적으로 시뮬레이션하여, 브루트포스 방법의 기존 기준을 크게 초월한다.
• 반정밀도 행렬 곱셈을 위한 네이티브 MXU와 파동함수 재분배를 위한 빠른 코어 간 통신 기능을 활용해 높은 성능을 달성한다.
• 메모리 제약을 고려해 계산을 이중 단계로 나누는 방식으로 랑츠 알고리즘을 성공적으로 TPU 환경에 적응시켰으며, 이로 인해 메모리 사용을 줄이면서도 정확도를 유지했다.
• 1차원 스핀 체인 하미르토니안에서 기본 상태 및 시간 진화의 시뮬레이션을 수행했으며, 상관 함수와 얽힘 측정치를 통해 결과를 검증했다.
• 이 방법은 QMC 및 텐서 네트워크가 실패하는 영역에서 특히 유용한 오류가 없는 신뢰할 수 있는 벤치마크를 제공한다.
• 이 연구는 전체 파동함수 접근이 필요한 경우에 TPU 풀이 정확한 양자 시뮬레이션을 위한 실현 가능하고 고성능 대안임을 입증한다.