[論文レビュー] 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

ひとこと要約

本論文は、GoogleのTensor Processing Units（TPUs）上で、密度行列密度行列縮約法（DMRG）アルゴリズムの初の大規模実装を示しており、1,024個のTPU v3コアを用いて、$ D = 2^{16} = 65,536 $にまで達する前例のない結合次元を達成した。TPUの高帯域幅で分散化された行列乗算機能を活用することで、$10\times10$の自由フェルミオン系や$20\times20$の横磁場イジング模型などの2次元量子系に対するDMRGシミュレーションが高速化され、実行時間が数か月から1日未塔に短縮された。

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次元系の収束を達成し、新たなパフォーマンスベンチマークを設定すること。

提案手法

• GoogleのTPU v3クラスタ上でDMRGアルゴリズムを実装し、高帯域幅で分散化された行列乗算機能を活用した。
• 個々のTPUのメモリ容量を超える大きなMPSテンソルを扱うために、分散型で外部記憶装置を活用するテンソル結合戦略を実装した。
• MPSテンソルをTPUコアにデータ並列的に分割することで、1,024コア（TPU v3ポッドの半分）への効率的スケーリングを実現した。
• JAXとXLAを用いてTPUコア間の通信パターンを最適化し、分散テンソル演算における遅延を最小限に抑えた。
• 収束性と安定性を向上させるために、交互にスイープするブロック-DMRGアプローチを適用した。
• TPUコア数を2倍にした際、実行時間は結合次元$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 $という、DMRGシミュレーションで報告された中で最大の結合次元を達成した。
• $10\times10$の自由フェルミオン系や$20\times20$の横磁場イジング模型のシミュレーションが1日未塔で収束した。これは、共有メモリ型CPU上で実行した場合、数か月かかると予想されていた。
• 結合次元$ D = 65,536 $のMPSテンソルを1つ最適化するのに約2分を要した。これは、高い計算スループットを示している。
• TPUコア数を2倍にした際、実行時間は結合次元$D$にほぼ比例して増加した。これは、並列化と強いスケーリングが図られていることを示している。
• 同程度の結合次元において、従来のCPUベースのDMRGに比べて100倍の高速化を達成した。これにより、従来では不可能とされてきたもつれスケールへのアクセスが可能になった。
• 著者らは、1つのTPU v3ポッド（2,048コア）またはTPU v4ポッド（8,192コア）を用いることで、最大達成可能な$D$がそれぞれ$\sqrt{2}$倍から2倍にまで向上すると予測している。