research-article

NNQS-Transformer: an Efficient and Scalable Neural Network Quantum States Approach for Ab initio Quantum Chemistry

Authors:

Honghui ShangAuthors Info & Claims

SC '23: Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis

Article No.: 42, Pages 1 - 13

https://doi.org/10.1145/3581784.3607061

Published: 11 November 2023 Publication History

Abstract

Neural network quantum state (NNQS) has emerged as a promising candidate for quantum many-body problems, but its practical applications are often hindered by the high cost of sampling and local energy calculation. We develop a high-performance NNQS method for ab initio electronic structure calculations. The major innovations include: (1) A transformer based architecture as the quantum wave function ansatz; (2) A data-centric parallelization scheme for the variational Monte Carlo (VMC) algorithm which preserves data locality and well adapts for different computing architectures; (3) A parallel batch sampling strategy which reduces the sampling cost and achieves good load balance; (4) A parallel local energy evaluation scheme which is both memory and computationally efficient; (5) Study of real chemical systems demonstrates both the superior accuracy of our method compared to state-of-the-art and the strong and weak scalability for large molecular systems with up to 120 spin orbitals.

References

[1]

Konstantinos D. Vogiatzis, Dongxia Ma, Jeppe Olsen, Laura Gagliardi, and Wibe A. de Jong. Pushing configuration-interaction to the limit: Towards massively parallel mcscf calculations. The Journal of Chemical Physics, 147(18):184111, 2017.

[2]

C. David Sherrill and Henry F. Schaefer. The configuration interaction method: Advances in highly correlated approaches. volume 34 of Advances in Quantum Chemistry, pages 143--269. Academic Press, 1999.

[3]

Ireneusz W Bulik, Thomas M Henderson, and Gustavo E Scuseria. Can single-reference coupled cluster theory describe static correlation? Journal of chemical theory and computation, 11(7):3171--3179, 2015.

[4]

Steven R. White. Density matrix formulation for quantum renormalization groups. Physical Review Letters, 69:2863--2866, Nov 1992.

[5]

Steven R. White. Density-matrix algorithms for quantum renormalization groups. Physical Review B, 48:10345--10356, Oct 1993.

[6]

Jiri Brabec, Jan Brandejs, Karol Kowalski, Sotiris Xantheas, Örs Legeza, and Libor Veis. Massively parallel quantum chemical density matrix renormalization group method. Journal of Computational Chemistry, 42(8):534--544, 2021.

[7]

Henrik R Larsson, Huanchen Zhai, Cyrus J Umrigar, and Garnet Kin-Lic Chan. The chromium dimer: closing a chapter of quantum chemistry. Journal of the American Chemical Society, 144(35):15932--15937, 2022.

[8]

D. Perez-Garcia, F. Verstraete, M. M. Wolf, and J. I. Cirac. Matrix product state representations. Quantum Information Computation, 7(5):401--430, jul 2007.

Digital Library

[9]

W. M. C. Foulkes, L. Mitas, R. J. Needs, and G. Rajagopal. Quantum monte carlo simulations of solids. Reviews of Modern Physics, 73:33--83, Jan 2001.

[10]

R J Needs, M D Towler, N D Drummond, and P López Ríos. Continuum variational and diffusion quantum monte carlo calculations. Journal of Physics: Condensed Matter, 22(2):023201, dec 2009.

[11]

Brian M. Austin, Dmitry Yu. Zubarev, and William A. Jr. Lester. Quantum monte carlo and related approaches. Chemical reviews, 112(1):263--288, 2012.

[12]

David Pfau, James S. Spencer, Alexander G. D. G. Matthews, and W. M. C. Foulkes. Ab initio solution of the many-electron schrödinger equation with deep neural networks. Physical review research, 2:033429, Sep 2020.

[13]

Jan Hermann, Zeno Schätzle, and Frank Noé. Deep-neural-network solution of the electronic schrödinger equation. Nature Chemistry, 12(10):891--897, 2020.

[14]

Giuseppe Carleo and Matthias Troyer. Solving the quantum many-body problem with artificial neural networks. Science, 355(6325):602--606, 2017.

[15]

Kenny Choo, Titus Neupert, and Giuseppe Carleo. Two-dimensional frustrated J₁-J₂ model studied with neural network quantum states. Physical Review B, 100:125124, Sep 2019.

[16]

Or Sharir, Yoav Levine, Noam Wies, Giuseppe Carleo, and Amnon Shashua. Deep autoregressive models for the efficient variational simulation of many-body quantum systems. Physical Review Letters, 124:020503, Jan 2020.

[17]

Markus Schmitt and Markus Heyl. Quantum many-body dynamics in two dimensions with artificial neural networks. Physical Review Letters, 125:100503, Sep 2020.

[18]

Dong Yuan, He-Ran Wang, Zhong Wang, and Dong-Ling Deng. Solving the liouvillian gap with artificial neural networks. Physical Review Letters, 126:160401, Apr 2021.

[19]

Xuncheng Zhao, Mingfan Li, Qian Xiao, Junshi Chen, Fei Wang, Li Shen, Meijia Zhao, Wenhao Wu, Hong An, Lixin He, and Xiao Liang. Ai for quantum mechanics: High performance quantum many-body simulations via deep learning. In SC22: International Conference for High Performance Computing, Networking, Storage and Analysis, pages 1--15, 2022.

[20]

Javier Robledo Moreno, Giuseppe Carleo, Antoine Georges, and James Stokes. Fermionic wave functions from neural-network constrained hidden states. Proceedings of the National Academy of Sciences, 119(32):e2122059119, 2022.

[21]

Dong-Ling Deng, Xiaopeng Li, and S. Das Sarma. Quantum entanglement in neural network states. Physical Review X, 7:021021, May 2017.

[22]

Ivan Glasser, Nicola Pancotti, Moritz August, Ivan D. Rodriguez, and J. Ignacio Cirac. Neural-network quantum states, string-bond states, and chiral topological states. Physical Review X, 8:011006, Jan 2018.

[23]

Or Sharir, Amnon Shashua, and Giuseppe Carleo. Neural tensor contractions and the expressive power of deep neural quantum states. Physical Review B, 106:205136, Nov 2022.

[24]

Xun Gao and Lu-Ming Duan. Efficient representation of quantum many-body states with deep neural networks. Nature communications, 8(1):662, 2017.

[25]

Kenny Choo, Antonio Mezzacapo, and Giuseppe Carleo. Fermionic neural-network states for ab-initio electronic structure. Nature communications, 11(1):2368, 2020.

[26]

Thomas D Barrett, Aleksei Malyshev, and AI Lvovsky. Autoregressive neural-network wavefunctions for ab initio quantum chemistry. Nature Machine Intelligence, 4(4):351--358, 2022.

[27]

Tianchen Zhao, James Stokes, and Shravan Veerapaneni. Scalable neural quantum states architecture for quantum chemistry. Machine Learning: Science and Technology, 4(2):025034, jun 2023.

[28]

Sandro Sorella and Luca Capriotti. Green function monte carlo with stochastic reconfiguration: An effective remedy for the sign problem. Physical Review B, 61:2599--2612, Jan 2000.

[29]

Sandro Sorella, Michele Casula, and Dario Rocca. Weak binding between two aromatic rings: Feeling the van der waals attraction by quantum monte carlo methods. The Journal of Chemical Physics, 127(1):014105, 2007.

[30]

Honghui Shang, Chu Guo, Yangjun Wu, Zhenyu Li, and Jinlong Yang. Solving schrodinger equation with a language model. arXiv: 2307.09343 [quant-ph].

[31]

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.

Digital Library

[32]

Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al. Improving language understanding by generative pre-training. 2018.

[33]

Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAI blog, 1(8):9, 2019.

[34]

Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Yang, Zach DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. PyTorch: An Imperative Style, High-Performance Deep Learning Library. Curran Associates Inc., Red Hook, NY, USA, 2019.

[35]

Qiming Sun, Timothy C. Berkelbach, Nick S. Blunt, George H. Booth, Sheng Guo, Zhendong Li, Junzi Liu, James D. McClain, Elvira R. Sayfutyarova, Sandeep Sharma, Sebastian Wouters, and Garnet Kin-Lic Chan. Pyscf: the python-based simulations of chemistry framework. WIREs Computational Molecular Science, 8(1):e1340, 2018.

[36]

Jarrod R McClean, Nicholas C Rubin, Kevin J Sung, Ian D Kivlichan, Xavier Bonet-Monroig, Yudong Cao, Chengyu Dai, E Schuyler Fried, Craig Gidney, Brendan Gimby, Pranav Gokhale, Thomas Häner, Tarini Hardikar, Vojtěch Havlíček, Oscar Higgott, Cupjin Huang, Josh Izaac, Zhang Jiang, Xinle Liu, Sam McArdle, Matthew Neeley, Thomas O'Brien, Bryan O'Gorman, Isil Ozfidan, Maxwell D Radin, Jhonathan Romero, Nicolas P D Sawaya, Bruno Senjean, Kanav Setia, Sukin Sim, Damian S Steiger, Mark Steudtner, Qiming Sun, Wei Sun, Daochen Wang, Fang Zhang, and Ryan Babbush. Openfermion: the electronic structure package for quantum computers. Quantum Science and Technology, 5(3):034014, jun 2020.

[37]

Tianchen Zhao, Saibal De, Brian Chen, James Stokes, and Shravan Veerapaneni. Overcoming barriers to scalability in variational quantum monte carlo. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, SC '21, New York, NY, USA, 2021. Association for Computing Machinery.

Digital Library

[38]

Mathieu Germain, Karol Gregor, Iain Murray, and Hugo Larochelle. Made: Masked autoencoder for distribution estimation. In Francis Bach and David Blei, editors, Proceedings of the 32nd International Conference on Machine Learning, volume 37 of Proceedings of Machine Learning Research, pages 881--889, Lille, France, 07--09 Jul 2015. PMLR.

[39]

Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers), pages 4171--4186, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics.

[40]

John F. Kolen and Stefan C. Kremer. Gradient Flow in Recurrent Nets: The Difficulty of Learning LongTerm Dependencies, pages 237--243. 2001.

[41]

Xiaohui Wang, Yang Wei, Ying Xiong, Guyue Huang, Xian Qian, Yufei Ding, Mingxuan Wang, and Lei Li. Lightseq2: Accelerated training for transformer-based models on gpus. In Proceedings of the International Conference on High Performance Computing, Networking, Storage and Analysis, SC '22. IEEE Press, 2022.

[42]

Xiaohui Wang, Ying Xiong, Yang Wei, Mingxuan Wang, and Lei Li. LightSeq: A high performance inference library for transformers. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies: Industry Papers (NAACL-HLT), pages 113--120. Association for Computational Linguistics, June 2021.

[43]

Jiarui Fang, Yang Yu, Chengduo Zhao, and Jie Zhou. Turbotransformers: An efficient gpu serving system for transformer models. In Proceedings of the 26th ACM SIGPLAN Symposium on Principles and Practice of Parallel Programming, PPoPP '21, page 389--402, New York, NY, USA, 2021. Association for Computing Machinery.

Digital Library

Cited By

Kan BTian YXie DWu YFan YShang H(2024)Solving the Electronic Schrödinger Equation by Pairing Tensor-Network State with Neural Network Quantum StateMathematics10.3390/math1203043312:3(433)Online publication date: 29-Jan-2024
https://doi.org/10.3390/math12030433
Melko RCarrasquilla J(2024)Language models for quantum simulationNature Computational Science10.1038/s43588-023-00578-04:1(11-18)Online publication date: 22-Jan-2024
https://doi.org/10.1038/s43588-023-00578-0
Fu LWu YShang HYang J(2024)Transformer-Based Neural-Network Quantum State Method for Electronic Band Structures of Real SolidsJournal of Chemical Theory and Computation10.1021/acs.jctc.4c0056720:14(6218-6226)Online publication date: 8-Jul-2024
https://doi.org/10.1021/acs.jctc.4c00567

Recommendations

Quantum correlation swapping

Quantum correlations (QCs), including quantum entanglement and those different, are important quantum resources and have attracted much attention recently. Quantum entanglement swapping as a kernel technique has already been applied to quantum repeaters ...
Quantum discord of ensemble of quantum states

We highlight an information-theoretic meaning of quantum discord as the gap between the accessible information and the Holevo bound in the framework of ensemble of quantum states. This complementary relationship implies that a large amount of ...
Quantum holonomies for displaced Landau---Aharonov---Casher states

We construct displaced Fock states for a Landau---Aharonov---Casher system for neutral particles. Abelian and non-Abelian geometric phases can be obtained in an adiabatic cyclic evolution using this displaced states. Moreover, we show that a possible ...

Comments

Information & Contributors

Information

Published In

cover image ACM Conferences

SC '23: Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis

November 2023

1428 pages

ISBN:9798400701092

DOI:10.1145/3581784

Chair:
Dorian Arnold,
Program Chair:
Rosa M Badia,
Program Co-chair:
Kathryn Mohror

Copyright © 2023 ACM.

Permission to make digital or hard copies of all or part of this work for personal or classroom use is granted without fee provided that copies are not made or distributed for profit or commercial advantage and that copies bear this notice and the full citation on the first page. Copyrights for components of this work owned by others than the author(s) must be honored. Abstracting with credit is permitted. To copy otherwise, or republish, to post on servers or to redistribute to lists, requires prior specific permission and/or a fee. Request permissions from [email protected].

Sponsors

Publisher

Association for Computing Machinery

New York, NY, United States

Publication History

Published: 11 November 2023

Permissions

Request permissions for this article.

Request Permissions

Check for updates

Badges

Author Tags

Qualifiers

Research-article

Funding Sources

National Natural Science Foundation of China

Conference

SC '23

Sponsor:

SIGHPC

SC '23: International Conference for High Performance Computing, Networking, Storage and Analysis

November 12 - 17, 2023

CO, Denver, USA

Acceptance Rates

Overall Acceptance Rate 1,516 of 6,373 submissions, 24%

Contributors

Other Metrics

View Article Metrics

Bibliometrics & Citations

Bibliometrics

Article Metrics

3
Total Citations
View Citations
246
Total Downloads

Downloads (Last 12 months)246
Downloads (Last 6 weeks)29

Reflects downloads up to 10 Aug 2024

Other Metrics

View Author Metrics

Citations

Cited By

Kan BTian YXie DWu YFan YShang H(2024)Solving the Electronic Schrödinger Equation by Pairing Tensor-Network State with Neural Network Quantum StateMathematics10.3390/math1203043312:3(433)Online publication date: 29-Jan-2024
https://doi.org/10.3390/math12030433
Melko RCarrasquilla J(2024)Language models for quantum simulationNature Computational Science10.1038/s43588-023-00578-04:1(11-18)Online publication date: 22-Jan-2024
https://doi.org/10.1038/s43588-023-00578-0
Fu LWu YShang HYang J(2024)Transformer-Based Neural-Network Quantum State Method for Electronic Band Structures of Real SolidsJournal of Chemical Theory and Computation10.1021/acs.jctc.4c0056720:14(6218-6226)Online publication date: 8-Jul-2024
https://doi.org/10.1021/acs.jctc.4c00567

View Options

Get Access

Login options

Check if you have access through your login credentials or your institution to get full access on this article.

Full Access

Get this Publication

View options

PDF

View or Download as a PDF file.

eReader

View online with eReader.

Media

Figures

Other

Tables

View Table of Contents