- Physical Review Journals
All JournalsPhysics Magazine

PRX Quantum

Reuse & Permissions

It is not necessary to obtain permission to reuse this article or its components as it is available under the terms of the Creative Commons Attribution 4.0 International license. This license permits unrestricted use, distribution, and reproduction in any medium, provided attribution to the author(s) and the published article's title, journal citation, and DOI are maintained. Please note that some figures may have been included with permission from other third parties. It is your responsibility to obtain the proper permission from the rights holder directly for these figures.

Export citation

Export citation

Choose format for download:

Download Citation
  • Editors' Suggestion
  • Open Access

Simulation of Quantum Many-Body Dynamics with Tensor Processing Units: Floquet Prethermalization

Alan Morningstar1,2,*, Markus Hauru2, Jackson Beall2, Martin Ganahl2, Adam G.M. Lewis2, Vedika Khemani3, and Guifre Vidal2

  • 1Department of Physics, Princeton University, Princeton, New Jersey 08544, USA
  • 2Sandbox@Alphabet, Mountain View, California 94043, USA
  • 3Department of Physics, Stanford University, Stanford, California 94305, USA

  • *alanmorningstar@princeton.edu
PDF

PRX Quantum 3, 020331 – Published 11 May, 2022

DOI: https://doi.org/10.1103/PRXQuantum.3.020331

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 long time scales. 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 105 Floquet periods, corresponding to circuits with 4×106 nearest-neighbor two-qubit gates. The circuits simulated have no additional symmetries and represent a pure-state evolution in the full 2L-dimensional Hilbert space. This is achieved by distributing the computation over 128 TPU cores. On that size TPU cluster, we find speed-ups in wall-clock run time of 230 times and 15 times when compared to reference CPU and single-graphics-processing-unit (GPU) simulations, respectively, for shorter-time 30-qubit simulations that can be handled by all three platforms. 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 a “full pod” of 2048 TPU cores with tens of terabytes of memory in total. For these simulations, an eight-TPU-core machine is comparable to a single A100 GPU and thus the full TPU pod is comparable to a machine with hundreds of top-of-the-line GPUs. However, the TPU pod is more energy and cost efficient and readily accessible (via Google Cloud), unlike such large many-GPU configurations. We also study the accumulation of numerical error as a function of circuit depth in very deep circuits. Our work demonstrates that TPUs can offer significant advantages for state-of-the-art simulations of quantum many-body dynamics.

View figure in article

Physics Subject Headings (PhySH)

Popular Summary

The simulation of quantum systems on a classical computer is prohibitively costly due to the fact that quantum systems can be in a “highly entangled” superposition of all possible classical states at once. This reality, along with the fact that quantum systems are central to both basic scientific research and technological advancement, was the original motivation to begin the investigation and development of new types of computers that operate quantum mechanically. Despite steady progress, nascent quantum computers are strongly limited by environmental noise and much research in quantum science and engineering continues to rely heavily on large classical-computer simulations. It is, therefore, important to explore state-of-the-art approaches to performing such simulations. In recent years, classical-computing technology has evolved rapidly, in part to meet the demands of large-scale machine-learning tasks. In this work, we repurpose a recent technological advancement in computing to accelerate and scale up classical simulations of quantum dynamics. Tensor processing units (TPUs) are specialized processors built for doing the simple but large-scale algebra involved in training and evaluating neural networks. The same kind of algebra is the core mathematical operation involved in the simulation of quantum systems on a classical computer. We explore the benefits of leveraging TPUs for simulations of quantum dynamics. As an example, we demonstrate state-of-the-art results on the study of how periodically driven quantum systems absorb energy. This effort demonstrates the potential of using TPUs as a platform for the simulation of quantum systems. Further applications of TPU simulation methods may help to enable advances in our understanding of these systems.

Article Text

References (100)

  1. A. Polkovnikov, K. Sengupta, A. Silva, and M. Vengalattore, Colloquium: Nonequilibrium dynamics of closed interacting quantum systems, Rev. Mod. Phys. 83, 863 (2011).
  2. L. D’Alessio, Y. Kafri, A. Polkovnikov, and M. Rigol, From quantum chaos and eigenstate thermalization to statistical mechanics and thermodynamics, Adv. Phys. 65, 239 (2016).
  3. R. Nandkishore and D. A. Huse, Many-body localization and thermalization in quantum statistical mechanics, Annu. Rev. Condens. Matter Phys. 6, 15 (2015).
  4. D. A. Abanin, E. Altman, I. Bloch, and M. Serbyn, Colloquium: Many-body localization, thermalization, and entanglement, Rev. Mod. Phys. 91, 021001 (2019).
  5. F. Alet and N. Laflorencie, Many-body localization: An introduction and selected topics, C. R. Phys. 19, 498 (2018).
  6. S. Moudgalya, B. A. Bernevig, and N. Regnault, Quantum many-body scars and Hilbert space fragmentation: A review of exact results, ArXiv:2109.00548 (2021).
  7. M. Serbyn, D. A. Abanin, and Z. Papić, Quantum many-body scars and weak breaking of ergodicity, Nat. Phys. 17, 675 (2021).
  8. F. Harper, R. Roy, M. S. Rudner, and S. L. Sondhi, Topology and broken symmetry in Floquet systems, Annu. Rev. Condens. Matter Phys. 11, 345 (2020).
  9. R. P. Feynman, Simulating physics with computers, Int. J. Theor. Phys. 21, 467 (1982).
  10. I. M. Georgescu, S. Ashhab, and F. Nori, Quantum simulation, Rev. Mod. Phys. 86, 153 (2014).
  11. T. D. Ladd, F. Jelezko, R. Laflamme, Y. Nakamura, C. Monroe, and J. L. O’Brien, Quantum computers, Nature 464, 45 (2010).
  12. I. Bloch, J. Dalibard, and S. Nascimbène, Quantum simulations with ultracold quantum gases, Nat. Phys. 8, 267 (2012).
  13. R. Blatt and C. F. Roos, Quantum simulations with trapped ions, Nat. Phys. 8, 277 (2012).
  14. C. Monroe, W. C. Campbell, L.-M. Duan, Z.-X. Gong, A. V. Gorshkov, P. W. Hess, R. Islam, K. Kim, N. M. Linke, and G. Pagano, et al., Programmable quantum simulations of spin systems with trapped ions, Rev. Mod. Phys. 93, 025001 (2021).
  15. A. Aspuru-Guzik and P. Walther, Photonic quantum simulators, Nat. Phys. 8, 285 (2012).
  16. A. A. Houck, H. E. Türeci, and J. Koch, On-chip quantum simulation with superconducting circuits, Nat. Phys. 8, 292 (2012).
  17. J. Preskill, Quantum computing in the NISQ era and beyond, Quantum 2, 79 (2018).
  18. E. Altman, K. R. Brown, G. Carleo, L. D. Carr, E. Demler, C. Chin, B. DeMarco, S. E. Economou, M. A. Eriksson, and K.-M. C. Fu, et al., Quantum Simulators: Architectures and Opportunities, PRX Quantum 2, 017003 (2021).
  19. Y. Alexeev, D. Bacon, K. R. Brown, R. Calderbank, L. D. Carr, F. T. Chong, B. DeMarco, D. Englund, E. Farhi, and B. Fefferman, et al., Quantum Computer Systems for Scientific Discovery, PRX Quantum 2, 017001 (2021).
  20. D. Awschalom, K. K. Berggren, H. Bernien, S. Bhave, L. D. Carr, P. Davids, S. E. Economou, D. Englund, A. Faraon, and M. Fejer, et al., Development of Quantum Interconnects (QuICs) for Next-Generation Information Technologies, PRX Quantum 2, 017002 (2021).
  21. N. Jouppi, D. Yoon, G. Kurian, S. Li, N. Patil, J. Laudon, C. Young, and D. Patterson, A domain-specific supercomputer for training deep neural networks, Commun. ACM 63, 67 (2020).
  22. N. P. Jouppi, C. Young, N. Patil, D. Patterson, G. Agrawal, R. Bajwa, S. Bates, S. Bhatia, N. Boden, and A. Borchers, et al., in Proceedings of the 44th Annual International Symposium on Computer Architecture, ISCA ’17 (Association for Computing Machinery, New York, 2017), p. 1.
  23. J. Bradbury, R. Frostig, P. Hawkins, M. James Johnson, C. Leary, D. Maclaurin, G. Necula, A. Paszke, J. VanderPlas, S. Wanderman-Milne, and Q. Zhang, jax: Composable transformations of python+NumPy programs, 2018, http://github.com/google/jax.
  24. R. Frostig, M. Johnson, and C. Leary, Compiling machine learning programs via high-level tracing, Systems for Machine Learning, 23, 2018.
  25. M. Abadi, P. Barham, J. Chen, Z. Chen, A. Davis, J. Dean, M. Devin, S. Ghemawat, G. Irving, and M. Isard, et al., in 12th USENIX Symposium on Operating Systems Design and Implementation (OSDI 16) (USENIX Association and Savannah, GA, USA, 2016), p. 265.
  26. F. Belletti, D. King, K. Yang, R. Nelet, Y. Shafi, Y.-F. Chen, and J. Anderson, Tensor processing units for financial Monte Carlo, ArXiv:1906.02818 (2020).
  27. Q. Wang, M. Ihme, Y.-F. Chen, and J. Anderson, A TensorFlow simulation framework for scientific computing of fluid flows on tensor processing units, ArXiv:2108.11076 (2021).
  28. Z. Pan and P. Mishra, Hardware acceleration of explainable machine learning using tensor processing units, ArXiv:2103.11927 (2021).
  29. T. Lu, T. Marin, Y. Zhuo, Y.-F. Chen, and C. Ma, in 2021 IEEE 18 h International Symposium on Biomedical (2021), p. 783.
  30. C. Ma, T. Marin, T. J. Lu, Y. F. Chen, and Y. Zhuo, Nonuniform fast Fourier transform on TPUs, 2021.
  31. T. Lu, Y.-F. Chen, B. Hechtman, T. Wang, and J. Anderson, Large-scale discrete Fourier transform on TPUs, ArXiv:2002.03260 (2020).
  32. F. Huot, Y.-F. Chen, R. Clapp, C. Boneti, and J. Anderson, High-resolution imaging on TPUs, ArXiv:1912.08063 (2019).
  33. Single program, multiple data (SPMD): a parallel programming paradigm whereby a single set of instructions (program) is evaluated simultaneously on multiple distinct inputs (data) hosted on a set of parallel processes.
  34. A. Li, O. Subasi, X. Yang, and S. Krishnamoorthy, in SC20: International Conference for High Performance Computing, Networking, Storage and Analysis (2020), p. 1.
  35. T. Vincent, L. J. O’Riordan, M. Andrenkov, J. Brown, N. Killoran, H. Qi, and I. Dhand, jet: Fast quantum circuit simulations with parallel task-based tensor-network contraction, ArXiv:2107.09793 (2021).
  36. S. Efthymiou, S. Ramos-Calderer, C. Bravo-Prieto, A. Pérez-Salinas, D. García-Martín, A. Garcia-Saez, J. I. Latorre, and S. Carrazza, qibo: A framework for quantum simulation with hardware acceleration, Quantum Sci. Technol. 7, 015018 (2021).
  37. Y. Suzuki, Y. Kawase, Y. Masumura, Y. Hiraga, M. Nakadai, J. Chen, K. Nakanishi, K. Mitarai, R. Imai, and S. Tamiya, et al., qulacs: A fast and versatile quantum circuit simulator for research purpose, Quantum 5, 559 (2021).
  38. T. Jones, A. Brown, I. Bush, and S. C. Benjamin, Quest and high performance simulation of quantum computers, Sci. Rep. 9, 10736 (2019).
  39. X.-Z. Luo, J.-G. Liu, P. Zhang, and L. Wang, Yao.jl: Extensible, efficient framework for quantum algorithm design, Quantum 4, 341 (2020).
  40. S. V. Isakov, D. Kafri, O. Martin, C. Vollgraff Heidweiller, W. Mruczkiewicz, M. P. Harrigan, N. C. Rubin, R. Thomson, M. Broughton, and K. Kissell, et al., Simulations of quantum circuits with approximate noise using qsim and cirq, ArXiv:2111.02396 (2021).
  41. F. Pan and P. Zhang, Simulating the Sycamore quantum supremacy circuits, ArXiv:2103.03074 (2021).
  42. S. Mandrà, J. Marshall, E. G. Rieffel, and R. Biswas, in 2021 IEEE/ACM Second International Workshop on Quantum Computing Software (QCS) (IEEE, and St. Louis, MO, USA, 2021), p. 99.
  43. A. Kelly, Simulating quantum computers using opencl, ArXiv:1805.00988 (2018).
  44. C. Huang, F. Zhang, M. Newman, J. Cai, X. Gao, Z. Tian, J. Wu, H. Xu, H. Yu, and B. Yuan, et al., Classical simulation of quantum supremacy circuits, ArXiv:2005.06787 (2020).
  45. J. Gray, quimb: A python library for quantum information and many-body calculations, J. Open Source Softw. 3, 819 (2018).
  46. B. Villalonga, D. Lyakh, S. Boixo, H. Neven, T. S. Humble, R. Biswas, E. G. Rieffel, A. Ho, and S. Mandrà, Establishing the quantum supremacy frontier with a 281 pflop/s simulation, Quantum Sci. Technol. 5, 034003 (2020).
  47. A. G. M. Lewis, J. Beall, M. Ganahl, M/ Hauru, S. B. Mallick, and G. Vidal, Large scale distributed linear algebra with tensor processing units, ArXiv:2112.09017 (2021).
  48. M. Ganahl, et al., Tensor Processing Units for Simulating Quantum Circuits, Sandbox@Alphabet (to be published).
  49. M. Hauru, A. Morningstar, J. Beall, M. Ganahl, A. Lewis, and G. Vidal, Simulation of quantum physics with tensor processing units: Brute-force computation of ground states and time evolution, ArXiv:2111.10466 (2021).
  50. R. Shillito, et al., Classical simulation of superconducting quantum hardware using Tensor Processing Units, Sandbox@Alphabet (to be published).
  51. R. Pederson, et al., Tensor processing units for quantum chemistry, Sandbox@Alphabet (to be published).
  52. R. Song, et al., Simulation of spin light emitting diodes using tensor processing units, Sandbox@Alphabet (to be published).
  53. M. Ganahl, et al., Density matrix renormalization group using tensor processing units, Sandbox@Alphabet (to be published).
  54. E. Gustafson, B. Holzman, J. Kowalkowski, H. Lamm, A. C. Y. Li, G. Perdue, S. Boixo, S. Isakov, O. Martin, and R. Thomson, et al., Large scale multi-node simulations of Z2 gauge theory quantum circuits using google cloud platform, ArXiv:2110.07482 (2021).
  55. T. Oka and S. Kitamura, Floquet engineering of quantum materials, Annu. Rev. Condens. Matter Phys. 10, 387 (2019).
  56. M. Bukov, L. D’Alessio, and A. Polkovnikov, Universal high-frequency behavior of periodically driven systems: From dynamical stabilization to Floquet engineering, Adv. Phys. 64, 139 (2015).
  57. T. Oka and H. Aoki, Photovoltaic Hall effect in graphene, Phys. Rev. B 79, 081406(R) (2009).
  58. M. S. Rudner, N. H. Lindner, E. Berg, and M. Levin, Anomalous Edge States and the Bulk-Edge Correspondence for Periodically Driven Two-Dimensional Systems, Phys. Rev. X 3, 031005 (2013).
  59. P. Titum, E. Berg, M. S. Rudner, G. Refael, and N. H. Lindner, Anomalous Floquet-Anderson Insulator as a Nonadiabatic Quantized Charge Pump, Phys. Rev. X 6, 021013 (2016).
  60. V. Khemani, A. Lazarides, R. Moessner, and S. Sondhi, Phase Structure of Driven Quantum Systems, Phys. Rev. Lett. 116, 250401 (2016).
  61. D. V. Else, B. Bauer, and C. Nayak, Floquet Time Crystals, Phys. Rev. Lett. 117, 090402 (2016).
  62. H. C. Po, L. Fidkowski, T. Morimoto, A. C. Potter, and A. Vishwanath, Chiral Floquet Phases of Many-Body Localized Bosons, Phys. Rev. X 6, 041070 (2016).
  63. M. Ippoliti, K. Kechedzhi, R. Moessner, S. L. Sondhi, and V. Khemani, Many-Body Physics in the NISQ Era: Quantum Programming a Discrete Time Crystal, PRX Quantum 2, 030346 (2021).
  64. X. Mi, M. Ippoliti, C. Quintana, A. Greene, Z. Chen, J. Gross, F. Arute, K. Arya, J. Atalaya, and R. Babbush, et al., Observation of time-crystalline eigenstate order on a quantum processor, ArXiv:2107.13571 (2021).
  65. L. D’Alessio and M. Rigol, Long-time Behavior of Isolated Periodically Driven Interacting Lattice Systems, Phys. Rev. X 4, 041048 (2014).
  66. A. Lazarides, A. Das, and R. Moessner, Equilibrium states of generic quantum systems subject to periodic driving, Phys. Rev. E 90, 012110 (2014).
  67. P. Ponte, Z. Papić, F. Huveneers, and D. A. Abanin, Many-Body Localization in Periodically Driven Systems, Phys. Rev. Lett. 114, 140401 (2015).
  68. A. Lazarides, A. Das, and R. Moessner, Fate of Many-Body Localization Under Periodic Driving, Phys. Rev. Lett. 115, 030402 (2015).
  69. D. A. Abanin, W. De Roeck, and F. Huveneers, Exponentially Slow Heating in Periodically Driven Many-Body Systems, Phys. Rev. Lett. 115, 256803 (2015).
  70. T. Mori, T. Kuwahara, and K. Saito, Rigorous Bound on Energy Absorption and Generic Relaxation in Periodically Driven Quantum Systems, Phys. Rev. Lett. 116, 120401 (2016).
  71. T. Kuwahara, T. Mori, and K. Saito, Floquet-Magnus theory and generic transient dynamics in periodically driven many-body quantum systems, Ann. Phys. (NY) 367, 96 (2016).
  72. D. A. Abanin, W. De Roeck, W. W. Ho, and F. Huveneers, Effective Hamiltonians, prethermalization, and slow energy absorption in periodically driven many-body systems, Phys. Rev. B 95, 014112 (2017).
  73. D. Abanin, W. De Roeck, W. W. Ho, and F. Huveneers, A rigorous theory of many-body prethermalization for periodically driven and closed quantum systems, Commun. Math. Phys. 354, 809 (2017).
  74. D. V. Else, B. Bauer, and C. Nayak, Prethermal Phases of Matter Protected by Time-Translation Symmetry, Phys. Rev. X 7, 011026 (2017).
  75. D. J. Luitz, R. Moessner, S. L. Sondhi, and V. Khemani, Prethermalization Without Temperature, Phys. Rev. X 10, 021046 (2020).
  76. A. Rubio-Abadal, M. Ippoliti, S. Hollerith, D. Wei, J. Rui, S. L. Sondhi, V. Khemani, C. Gross, and I. Bloch, Floquet Prethermalization in a Bose-Hubbard System, Phys. Rev. X 10, 021044 (2020).
  77. P. Peng, C. Yin, X. Huang, C. Ramanathan, and P. Cappellaro, Floquet prethermalization in dipolar spin chains, Nat. Phys. 17, 444 (2021).
  78. F. Machado, G. D. Kahanamoku-Meyer, D. V. Else, C. Nayak, and N. Y. Yao, Exponentially slow heating in short and long-range interacting Floquet systems, Phys. Rev. Res. 1, 033202 (2019).
  79. jax documentation, https://jax.readthedocs.io/en/latest/ (accessed: 2022-01-15).
  80. A. Chan, A. De Luca, and J. T. Chalker, Solution of a Minimal Model for Many-Body Quantum Chaos, Phys. Rev. X 8, 041019 (2018).
  81. C. Sünderhauf, D. Pérez-García, D. A. Huse, N. Schuch, and J. I. Cirac, Localization with random time-periodic quantum circuits, Phys. Rev. B 98, 134204 (2018).
  82. To be sure that we do not choose an atypical sample with, for example, very weak interactions or very strong single-qubit potentials anywhere, we split each Hk into its component noninteracting and interacting pieces and check that the norms of these pieces are ≲1 and approximately 1, respectively.
  83. S. Lloyd, Universal quantum simulators, Science 273, 1073 (1996).
  84. M. Suzuki, General theory of fractal path integrals with applications to many-body theories and statistical physics, J. Math. Phys. 32, 400 (1991).
  85. G. Vidal, Efficient Simulation of One-Dimensional Quantum Many-Body Systems, Phys. Rev. Lett. 93, 040502 (2004).
  86. M. Heyl, P. Hauke, and P. Zoller, Quantum localization bounds Trotter errors in digital quantum simulation, Sci. Adv. 5, eaau8342 (2019).
  87. C. Fleckenstein and M. Bukov, Thermalization and prethermalization in periodically kicked quantum spin chains, Phys. Rev. B 103, 144307 (2021).
  88. Quantum AI team and collaborators, qsim, 2020.
  89. Cirq Developers, cirq, 2021.
  90. P. Weinberg and M. Bukov, QuSpin: A python package for dynamics and exact diagonalisation of quantum many body systems part I: Spin chains, SciPost Phys. 2, 003 (2017).
  91. T. L. M. Lezama and J. H. Bardarson, Apparent slow dynamics in the ergodic phase of a driven many-body localized system without extensive conserved quantities, Phys. Rev. B 99, 161106(R) (2019).
  92. J. Šuntajs, J. Bonča, T. Prosen, and L. Vidmar, Quantum chaos challenges many-body localization, Phys. Rev. E 102, 062144 (2020).
  93. M. Kiefer-Emmanouilidis, R. Unanyan, M. Fleischhauer, and J. Sirker, Slow delocalization of particles in many-body localized phases, Phys. Rev. B 103, 024203 (2021).
  94. D. Sels and A. Polkovnikov, Dynamical obstruction to localization in a disordered spin chain, ArXiv:2009.04501 (2020).
  95. D. A. Abanin, J. H. Bardarson, G. De Tomasi, S. Gopalakrishnan, V. Khemani, S. A. Parameswaran, F. Pollmann, A. C. Potter, M. Serbyn, and R. Vasseur, Distinguishing localization from chaos: Challenges in finite-size systems, Ann. Phys. (NY) 427, 168415 (2021).
  96. P. Sierant, D. Delande, and J. Zakrzewski, Thouless Time Analysis of Anderson and Many-Body Localization Transitions, Phys. Rev. Lett. 124, 186601 (2020).
  97. R. K. Panda, A. Scardicchio, M. Schulz, S. R. Taylor, and M. Žnidarič, Can we study the many-body localisation transition?, EPL (Europhys. Lett.) 128, 67003 (2020).
  98. D. J. Luitz and Y. Bar Lev, Absence of slow particle transport in the many-body localized phase, Phys. Rev. B 102, 100202 (2020).
  99. A. Morningstar, L. Colmenarez, V. Khemani, D. J. Luitz, and D. A. Huse, Avalanches and many-body resonances in many-body localized systems, ArXiv:2107.05642 (2021).
  100. P. Sierant and J. Zakrzewski, Can we observe the many-body localization?, ArXiv:2109.13608 (2021).

Outline

Information

PRX Quantum 3, 020331– Published 11 May, 2022
  • Received 20 November 2021
  • Accepted 20 April 2022
Crossmark - Check for updates button
Creative Commons Logo

Published by the American Physical Society under the terms of the Creative Commons Attribution 4.0 International license. Further distribution of this work must maintain attribution to the author(s) and the published article's title, journal citation, and DOI.

Published by the American Physical Society

Sign In to Your Journals Account

Filter

Filter

Article Lookup

Enter a citation