Simulation of quantum physics with Tensor Processing Units: brute-force computation of ground states and time evolution

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 ⟩$ consisting of $2^N$ complex amplitudes. We demonstrate the usage of TPUs for both (i) computing the ground state $|{\Psi }_{gs}⟩$ of the Hamiltonian $H$, and (ii) simulating the time evolution $|\Psi (t)⟩=e^{-itH}|\Psi (0)⟩$ generated by this Hamiltonian starting from some initial state $|\Psi (0)⟩$. The bottleneck of the above tasks is computing the product $H|\Psi ⟩$, 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 ⟩$ 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.

References in corpus (7)

• A class of quantum many-body states that can be efficiently simulated
• Tensor renormalization group approach to 2D classical lattice models
• Renormalization algorithms for Quantum-Many Body Systems in two and higher dimensions
• A TensorFlow Simulation Framework for Scientific Computing of Fluid Flows on Tensor Processing Units
• Simulation of quantum many-body dynamics with Tensor Processing Units: Floquet prethermalization
• High-resolution imaging on TPUs
• Hardware Acceleration of Explainable Machine Learning using Tensor Processing Units

Cited by in corpus (1)

• Benchmarking highly entangled states on a 60-atom analog quantum simulator