Simulation of quantum many-body dynamics with Tensor Processing Units: Floquet prethermalization

Abstract

Tensor Processing Units (TPUs) are specialized hardware accelerators developed by Google to support large-scale machine-learning tasks, but they can also be leveraged to accelerate and scale other linear-algebra-intensive computations. In this paper we demonstrate the usage of TPUs for massively parallel, classical simulations of quantum many-body dynamics on very long timescales. We apply our methods to study the phenomenon of Floquet prethermalization, i.e., exponentially slow heating in quantum spin chains subject to high-frequency periodic driving. We simulate the dynamics of L = 34 qubits for over 10 5 Floquet periods, corresponding to 4 × 10 6 nearest-neighbor two-qubit gates. This is achieved by distributing the computation over 128 TPU cores. The circuits simulated have no additional symmetries and represent a pure-state evolution in the full 2 Ldimensional Hilbert space. We study the computational cost of the simulations, as a function of both the number of qubits and the number of TPU cores used, up to our maximum capacity of L = 40 qubits which requires 2048 TPU cores. For a 30-qubit benchmark simulation on 128 TPU cores, we find a 230× speedup in wall-clock runtime when compared to a reference multi-core CPU simulation that we take to be representative of the current standard in quantum many-body dynamics research. We also study the accumulation of errors as a function of circuit depth. Our work demonstrates that TPUs can offer significant advantages for state-of-the-art simulations of quantum many-body dynamics.

Large Scale Distributed Linear Algebra With Tensor Processing Units

“…Moreover, a full pod of fourth generation TPUs (8192 cores) is expected to address matrices that double the above linear size in comparable times (work in progress). As shown in subsequent papers, see [9][10][11][12][13][14][15][16], the technology demonstrated here is already significant for a wide range of applications in the context of large-scale simulations and computations of quantum systems, including quantum computation, quantum many-body physics, quantum chemistry, and materials science.…”

“…For instance, anyone with a Google Cloud Platform account can have access to a TPU pod. As a result, a number of large-scale scientific computing tasks such as the ones demonstrated in this paper and in [9][10][11][12][13][14][15][16] are now within reach of any reach group, contributing to democratizing supercomputing throughout the scientific community and beyond.…”

“…The subsequent discussion showcases a selection of distributed dense linear algebra algorithms that function well despite these constraints. We have chosen these specific algorithms because of their widespread use in scientific computing and/or their pivotal role in several applications in [9][10][11][12][13][14][15][16] that were mentioned in the introduction.…”

“…Here we illustrate this with the so-called polar decomposition, which is obtained through applying the sign function to the singular values, where the sign function is approximated by means of a polynomial iteration made of small polynomials of the type in Eq. (12).…”

“…Previous work by others in this vein has concerned discrete [5] and fast [6] distributed Fourier transforms, Monte-Carlo simulation [7], and image processing [8]. Our group's sister papers address quantum circuit simulation [9,10], many-body quantum physics [11][12][13], electronic structure computation via density functional theory (DFT) [14] and coupled cluster (CC) methods [15], and tensor network algorithms such as the density matrix renormalization group (DMRG) [16].…”