[논문 리뷰] Density Matrix Renormalization Group with Tensor Processing Units

Martin Ganahl, Jackson Beall외 5명|arXiv (Cornell University)|2022. 04. 12.

Quantum many-body systems인용 수 5

한 줄 요약

이 논문은 구글의 텐서 처리 장치(TPU)에서 밀도 행렬 재규격화군(DMRG) 알고리즘의 최초 대규모 구현을 보여주며, 단지 1,024개의 TPU v3 코어를 사용하여 $ D = 2^{16} = 65,536 $에 이르는 예상치 못한 고결합 차원을 달성한다. TPU의 고대역폭 및 분산 행렬 곱셈 기능을 활용함으로써, $10\times10$ 자유 페르미온 및 $20\times20$ 횡방향 자장 이징 모델과 같은 2차원 양자 시스템에 대한 DMRG 시뮬레이션 속도가 월 수준에서 하루 이내로 단축된다.

ABSTRACT

Google's Tensor Processing Units (TPUs) are integrated circuits specifically built to accelerate and scale up machine learning workloads. They can perform fast distributed matrix multiplications and therefore be repurposed for other computationally intensive tasks. In this work we demonstrate the use of TPUs for accelerating and scaling up the density matrix renormalization group (DMRG), a powerful numerical approach to compute the ground state of a local quantum many-body Hamiltonian. The cost of DMRG scales with system size $N$ as $O(ND^3)$, where the so-called bond dimension $D$ regulates how expressive the underlying matrix product state (MPS) variational ansatz is. We consider lattice models in two spatial dimensions, with square lattices of size $10 imes 10$ (free fermions) and $20 imes 20$ (transverse field Ising model), for which the required MPS bond dimension is known to scale at least as $\exp(\sqrt{N})$. Using half of a TPU v3 pod (namely $1,\!024$ TPU v3 cores) we reached an unprecedentedly large bond dimension $D = 2^{16} = 65,\!536$, for which optimizing a single MPS tensor took about 2 minutes.

연구 동기 및 목표

• 기존 CPU 하드웨어의 한계를 넘어서 DMRG 시뮬레이션을 가속기 클러스터를 활용해 확장하는 것.
• 2차원 양자 시스템의 행렬 곱 상태(MPS) 변분 안사에서 고결합 차원 $D$로 인한 계산 블로킹 문제를 해결하는 것.
• 원래 기계 학습을 위해 설계된 TPU를 활용하여 대규모 텐서 네트워크 시뮬레이션을 양자 다체 물리학에서 실현 가능하고 효율적으로 구현하는 것의 타당성과 효율성을 입증하는 것.
• 이전까지는 도달할 수 없었던 고결합 차원에서 강하게 얽힌 2차원 시스템의 수렴을 달성하고, 새로운 성능 기준을 설정하는 것.

제안 방법

• 고성능 TPU v3 클러스터를 활용하여 DMRG 알고리즘을 Google의 TPU v3 클러스터에서 실행하도록 이식하였으며, 고대역폭 및 분산 행렬 곱셈 기능을 최대한 활용하였다.
• 개별 TPU 메모리 용량을 초과하는 큰 MPS 텐서를 처리하기 위해 분산 및 외부 메모리 기반 텐서 수축 전략을 구현하였다.
• TPU 코어 간에 MPS 텐서를 데이터 병렬 방식으로 분할하여, 1,024개 코어(하나의 TPU v3 파드의 절반)로도 효율적인 확장이 가능하도록 하였다.
• JAX와 XLA를 사용하여 TPU 코어 간 통신 패턴을 최적화하여 분산 텐서 연산에서 지연을 최소화하였다.
• 수렴성과 안정성을 향상시키기 위해 교차 스크린을 반복하는 블록-DMRG 방법을 적용하였다.
• TPU 코어 수를 두 배로 늘릴 경우 런타임이 결합 차원 $D$에 비례하여 선형적으로 증가함을 확인하여, 알고리즘적 및 하드웨어적 효율성이 높음을 입증하였다.

실험 결과

연구 질문

• RQ1 원래 기계 학습을 위해 설계된 TPU는 2차원 양자 시스템의 대규모 DMRG 시뮬레이션을 효과적으로 가속화하는 데에 활용될 수 있는가?
• RQ2 TPU 클러스터를 사용할 경우 DMRG 시뮬레이션에서 도달 가능한 최대 결합 차원 $D$는 얼마이며, 기존의 CPU 기반 구현과 비교해 볼 때 어떤가?
• RQ3 TPU 코어 수와 결합 차원 $D$가 증가함에 따라 DMRG의 성능는 어떻게 변화하는가?
• RQ4 예를 들어 U(1) 입자 수 보존 대칭성과 같은 대칭성 활용은 성능 향상에 얼마나 기여할 수 있으며, 더 큰 $D$를 가능하게 할 수 있는가?
• RQ5 TPU 기반 DMRG 구현은 $D=2^{16}$에서 강하게 얽힌 2차원 시스템, 예를 들어 임계 상태의 횡방향 자장 이징 모델에 대해 수렴을 달성할 수 있는가?

주요 결과

• 저자들은 단지 1,024개의 TPU v3 코어를 사용하여, 현재까지 보고된 바 중 가장 큰 결합 차원인 $ D = 2^{16} = 65,536 $을 달성하였다.
• $10\times10$ 자유 페르미온 및 $20\times20$ 횡방향 자장 이징 모델에 대한 시뮬레이션은 하루 이내에 수렴하였으며, 이는 공유 메모리 CPU에서 월 단위로 예상되는 작업을 압도적으로 단축시킨다.
• 결합 차원 $ D = 65,536 $에서 단일 MPS 텐서 최적화에 약 2분이 소요되었으며, 이는 높은 계산 처리 능력을 보여준다.
• TPU 코어 수를 두 배로 늘릴 경우 런타임이 결합 차원 $D$에 비례하여 약선형적으로 증가함을 확인하여, 효율적인 병렬 처리 및 강한 확장성(스케일링)을 입증하였다.
• 유사한 결합 차원에서 전통적인 CPU 기반 DMRG 대비 약 100배의 속도 향상을 달성하여, 이전에는 접근이 어려웠던 얽힘 척도에 도달할 수 있게 되었다.
• 저자들은 전체 TPU v3 파드(2,048개 코어) 또는 TPU v4 파드(8,192개 코어)를 사용할 경우 최대 도달 가능한 $D$가 각각 $\sqrt{2}$배에서 2배까지 증가할 것으로 예측한다.