[論文レビュー] 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ネイティブの行列乗算ユニット（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が機械学習向けに設計されているが、量子多体系の正確なシミュレーションを高速化できるかを調査すること。
• 指数的資源スケーリングにより、通常は実行不可能な量子スピン系における基底状態および時間発展演算をブルートフォース計算で可能にすること。
• 量子モンテカルロ法（符号問題）やテンソルネットワーク法（量子もつれの増大）の制限を、正確な全波動関数シミュレーションによって克服すること。
• 2048コアおよび32TBのHBMを備えたTPU v3ポッドを用いた大規模量子シミュレーションの実現可能性を示すこと。
• 専用ハードウェアプラットフォームを用いて、高パフォーマンスでエラーのない量子シミュレーション手法のベンチマークを提供すること。

提案手法

• 2048コア、1コアあたり16GBのHBM、専用MXUを備えたTPU v3ポッドを、スピン1/2系の正確な量子シミュレーションに再利用すること。
• 2^N次元の量子波動関数をTPUコアの高帯域幅メモリに分散させ、大規模な状態保存を可能にすること。
• TPUのMXUを活用し、高スループットの半精度行列乗算により、ハミルトニアン作用H|Ψ⟩を効率的に計算すること。
• 高速コア間インタコネクト（ICI）を用いて、反復的更新中に波動関数をコア間で動的に再配分すること。
• メモリ効率の良い2段階アプローチにより、ランツォス法を効果的に実装：まずKrylov空間ベクトルを保存せずに三重対角行列T_jを計算し、次にRitzベクトルから基底状態を再構築すること。
• 波動関数更新H|Ψ⟩をプリミティブとして用い、ランツォス法による基底状態計算および指数型作用素分割による時間発展演算を実現すること。

実験結果

リサーチクエスチョン

• RQ1 TPU v3ポッドは、機械学習ワークロードを超えて、正確な量子多体系シミュレーションに効果的に再利用可能か？
• RQ2 ブルートフォースベースのTPU手法でシミュレート可能な最大系サイズ（量子ビット数）はどの程度か？
• RQ3 ハミルトニアン作用および量子系における波動関数発展演算において、TPUのパフォーマンスは一般用途HPCと比べてどうか？
• RQ4 限られたコアあたりメモリを備えたTPUハードウェア上で、ランツォス法を効率的に実行できるか？
• RQ5 TPUベースのシミュレーションは、QMC やテンソルネットワーク法といった近似手法のベンチマークとして、どの程度有効に使えるか？

主な発見

• 2048コアおよび32TBのHBMを備えたTPU v3ポッドは、最大38量子ビットの量子系を正確にシミュレートでき、ブルートフォース法としては顕著なスケールを達成した。
• ネイティブMXUによる半精度行列乗算と、波動関数再配分に向けた高速コア間通信を活用することで、高いパフォーマンスが実現された。
• ランツォス法は、計算を2段階に分割することでTPUの制約に適合させることができ、メモリ使用量を削減しながらも正確性を維持した。
• 1次元スピンチェーンハミルトニアンを対象に、基底状態および時間発展演算のシミュレーションが実施され、相関関数およびもつれ測度による妥当性が検証された。
• このアプローチは、QMC やテンソルネットワーク法が失敗する領域においても、信頼性の高いエラーのないベンチマークを提供する。
• 本研究により、TPUポッドが正確な量子シミュレーションの実現可能性を示し、特に全波動関数へのアクセスが必須となる状況において、高パフォーマンスな代替手段として有効であることが明らかになった。