[论文解读] Density Matrix Renormalization Group with Tensor Processing Units

Martin Ganahl, Jackson Beall+5 more|arXiv (Cornell University)|Apr 12, 2022

Quantum many-body systems被引用 5

一句话总结

本论文首次在谷歌的张量处理单元（TPU）上实现了大规模密度矩阵重整化组（DMRG）算法的部署，仅使用1,024个TPU v3核心，便实现了前所未有的最大纠缠维度 $ D = 2^{16} = 65,536 $。通过利用TPU的高带宽、分布式矩阵乘法能力，该方法显著加速了二维量子系统的DMRG模拟——例如 $10\times10$ 自由费米子系统和 $20\times20$ 横场伊辛模型——将原本需数月的运行时间缩短至一天以内。

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.

研究动机与目标

• 通过利用加速器集群，将DMRG模拟扩展至传统CPU硬件的极限之外。
• 解决在二维量子系统中，矩阵乘积态（MPS）变分假设下高纠缠维度 $D$ 所带来的计算瓶颈。
• 证明将原本专为机器学习设计的TPU用于大规模张量网络模拟在量子多体物理中的可行性与高效性。
• 在以往无法企及的纠缠维度下，实现强纠缠二维系统的收敛，树立新的性能基准。

提出的方法

• 将DMRG算法移植至谷歌TPU v3集群运行，充分利用其高带宽、分布式矩阵乘法能力。
• 实现了一种分布式、外存存储的张量收缩策略，以管理超过单个TPU内存限制的大型MPS张量。
• 采用数据并行方式将MPS张量在TPU核心间分解，实现对1,024个核心（半个TPU v3计算阵列）的高效扩展。
• 利用JAX和XLA优化TPU核心间的通信模式，最大限度减少分布式张量操作的延迟。
• 采用分块DMRG方法并交替进行扫掠，以提升优化过程中的收敛性与稳定性。
• 当TPU核心数量加倍时，运行时间与纠缠维度 $D$ 呈近似线性关系，表明算法与硬件均具备高效性。

实验结果

研究问题

• RQ1 原本专为机器学习设计的TPU能否被有效重用于加速二维量子系统的大型DMRG模拟？
• RQ2 使用TPU集群进行DMRG模拟时，所能达到的最大纠缠维度 $D$ 是多少？与传统的CPU实现相比表现如何？
• RQ3 随着TPU核心数量和纠缠维度 $D$ 的增加，DMRG的性能如何变化？
• RQ4 对称性利用（如U(1)粒子数守恒）能在多大程度上进一步提升性能，并支持更大的 $D$？
• RQ5 基于TPU的DMRG实现能否在 $D=2^{16}$ 下实现对强纠缠二维系统（如临界横场伊辛模型）的收敛？

主要发现

• 作者仅使用1,024个TPU v3核心，便实现了 $ D = 2^{16} = 65,536 $ 的纠缠维度，为目前DMRG模拟中报道的最大值。
• $10\times10$ 自由费米子系统与 $20\times20$ 横场伊辛模型的模拟在一天内完成收敛，而传统共享内存CPU估计需数月时间。
• 在TPU集群上优化单个MPS张量至 $ D = 65,536 $ 仅耗时约2分钟，展现出极高的计算吞吐量。
• 当TPU核心数量加倍时，运行时间与纠缠维度 $D$ 呈近似线性关系，表明并行化效率高且具备强可扩展性。
• 与传统CPU-based DMRG相比，该实现实现了100倍的加速，使原本难以处理的纠缠尺度得以探索。
• 作者预测，若使用完整的TPU v3计算阵列（2,048个核心）或TPU v4计算阵列（8,192个核心），最大可实现的 $D$ 将分别提升 $\sqrt{2}$ 倍与2倍。