[论文解读] Simulation of quantum physics with Tensor Processing Units: brute-force computation of ground states and time evolution

Markus Hauru, Alan Morningstar+4 more|arXiv (Cornell University)|Nov 19, 2021

Parallel Computing and Optimization Techniques被引用 4

一句话总结

本文展示了如何将原本专为机器学习设计的 Google TPU v3 集群用于量子多体系统的暴力模拟，实现了最多 38 个量子比特的基态和时间演化计算。通过利用 TPU 原生的矩阵乘法单元（MXUs）、高带宽内存（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 是否能够加速量子多体系统的精确模拟。
• 实现量子自旋系统中基态和时间演化的暴力计算，这些系统因资源呈指数级增长而原本难以处理。
• 通过使用精确的全波函数模拟，克服现有方法（如量子蒙特卡洛方法的符号问题）和张量网络方法（纠缠熵增长）的局限性。
• 展示使用配备 2048 个核心和 32 TB HBM 的 TPU v3 集群进行大规模量子模拟的可行性。
• 通过专用硬件平台，为量子模拟方法提供高性能、无错误的基准测试。

提出的方法

• 将 TPU v3 集群（配备 2048 个核心、每个核心 16 GB HBM，以及专用 MXUs）重新用于自旋-1/2 系统的精确量子模拟。
• 将 2^N 维的量子波函数分布到 TPU 核心的高带宽内存中，以实现大规模态的存储。
• 利用 TPU 的 MXUs 通过高吞吐量的半精度矩阵乘法，高效计算哈密顿量作用 H|Ψ⟩。
• 利用快速的核间互连（ICIs）在迭代更新过程中动态重新分配波函数。
• 采用内存高效的两阶段方法实现 Lanczos 算法：首先在不存储 Krylov 向量的情况下计算三对角矩阵 T_j，然后从 Ritz 向量重构基态。
• 将波函数更新 H|Ψ⟩ 作为计算基态（通过 Lanczos）和时间演化（通过指数算子分裂）的基本操作。

实验结果

研究问题

• RQ1 TPU v3 集群是否能被有效重新用于机器学习工作负载之外的精确量子多体模拟？
• RQ2 基于 TPU 的暴力方法可模拟的最大系统规模（以量子比特数计）是多少？
• RQ3 在量子系统的哈密顿量作用和波函数演化方面，TPU 性能与通用 HPC 相比如何？
• RQ4 Lanczos 算法能否在每个核心内存受限的 TPU 硬件上高效运行？
• RQ5 TPU 基于的模拟在多大程度上可作为 QMC 和张量网络等近似方法的基准？

主要发现

• 配备 2048 个核心和 32 TB HBM 的 TPU v3 集群成功模拟了最多 38 个量子比特的量子系统，这是暴力方法实现的显著规模。
• 通过利用原生 MXUs 进行半精度矩阵乘法和快速的核间通信实现波函数再分配，该方法实现了高性能。
• 通过将计算分为两个阶段，成功将 Lanczos 算法适配到 TPU 的约束条件下，减少了内存使用量，同时保持了精度。
• 在 1D 自旋链哈密顿量上展示了基态和时间演化的模拟，结果通过关联函数和纠缠度量得到验证。
• 该方法为近似量子模拟方法（尤其在 QMC 和张量网络失效的区域）提供了可靠且无错误的基准。
• 本研究确立了 TPU 集群作为精确量子模拟的可行高性能替代方案，特别是在需要完整波函数访问的场景中。