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