Density Matrix Renormalization Group with Tensor Processing Units

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 therefore be repurposed for other computationally intensive tasks. In this work we demonstrate the use of TPUs accelerating scaling density renormalization group (DMRG), a powerful numerical approach compute ground state local quantum many-body Hamiltonian. The cost DMRG scales with system size $N$ as $O(ND^3)$, where so-called bond dimension $D$ regulates how expressive underlying product (MPS) variational ansatz is. We consider lattice models in two spatial dimensions, square lattices $10\times 10$ (free fermions) $20\times 20$ (transverse field Ising model), which required MPS is known at least $\exp(\sqrt{N})$. Using half TPU v3 pod (namely $1,\!024$ cores) reached an unprecedentedly large $D = 2^{16} 65,\!536$, optimizing single tensor took about 2 minutes.