Skip to main content

Thank you for visiting nature.com. You are using a browser version with limited support for CSS. To obtain the best experience, we recommend you use a more up to date browser (or turn off compatibility mode in Internet Explorer). In the meantime, to ensure continued support, we are displaying the site without styles and JavaScript.

  • Article
  • Published:

A computational framework for neural network-based variational Monte Carlo with Forward Laplacian

A preprint version of the article is available at arXiv.

Abstract

Neural network-based variational Monte Carlo (NN-VMC) has emerged as a promising cutting-edge technique of ab initio quantum chemistry. However, the high computational cost of existing approaches hinders their applications in realistic chemistry problems. Here we report a development of NN-VMC that achieves a remarkable speed-up rate, thereby greatly extending the applicability of NN-VMC to larger systems. Our key design is a computational framework named Forward Laplacian, which computes the Laplacian associated with neural networks, the bottleneck of NN-VMC, through an efficient forward propagation process. We then demonstrate that Forward Laplacian can further facilitate more developments of acceleration methods across various aspects, including optimization for sparse derivative matrix and efficient network design. Empirically, our approach enables NN-VMC to investigate a broader range of systems, providing valuable references to other ab initio methods. The results demonstrate a great potential in applying deep learning methods to solve general quantum mechanical problems.

This is a preview of subscription content, access via your institution

Access options

Buy this article

Prices may be subject to local taxes which are calculated during checkout

Fig. 1: Illustration of the computation process.
Fig. 2: The LapNet architecture.
Fig. 3: Efficiency and performance comparison of different NN-VMC methods.
Fig. 4: Scaling effect with system size.
Fig. 5: Relative energy estimation using LapNet.

Similar content being viewed by others

Data availability

All data supporting the findings of this study are available within the Supplementary Information.

Code availability

We release the source code of Forward Laplacian (https://github.com/YWolfeee/lapjax) and LapNet (https://github.com/bytedance/LapNet) on GitHub. Code for producing the results in this work is available on CodeOcean, together with all the system configuration studied in this work: https://codeocean.com/capsule/1139728/tree/v1 (ref. 60).

References

  1. Needs, R. J., Towler, M. D., Drummond, N. D., López Ríos, P. & Trail, J. R. Variational and diffusion quantum Monte Carlo calculations with the CASINO code. J. Chem. Phys. 152, 154106 (2020).

    Article  ADS  CAS  PubMed  Google Scholar 

  2. Foulkes, W. M. C., Mitas, L., Needs, R. J. & Rajagopal, G. Quantum Monte Carlo simulations of solids. Rev. Modern Phys. 73, 33–83 (2001).

    Article  ADS  CAS  Google Scholar 

  3. Pfau, D., Spencer, J. S., G. Matthews, A. G. & Foulkes, W. M. C. Ab-initio solution of the many-electron Schrödinger equation with deep neural networks. Phys. Rev. Res. 2, 033429 (2020).

    Article  CAS  Google Scholar 

  4. Spencer, J. S., Pfau, D., Botev, A. & Foulkes, W. M. C. Better, faster fermionic neural networks. Preprint at https://arXiv.org/2011.07125 (2020).

  5. Hermann, J., Schätzle, Z. & Noé, F. Deep-neural-network solution of the electronic Schrödinger equation. Nat. Chem. 12, 891–897 (2020).

    Article  CAS  PubMed  Google Scholar 

  6. Glehn, I., Spencer, J.S. & Pfau, D. A self-attention ansatz for ab-initio quantum chemistry. in 11th International Conference on Learning Representations https://openreview.net/forum?id=xveTeHVlF7j (2023.

  7. Schätzle, Z., Szabó, P., Mezera, M., Hermann, J. & Noé, F. DeepQMC: an open-source software suite for variational optimization of deep-learning molecular wave functions. J. Chem. Phys. 159, 094108 (2023).

  8. Bradbury, J. et al. JAX: composable transformations of Python+NumPy programs. GitHub http://github.com/google/jax (2023).

  9. Drummond, N. D., Towler, M. D. & Needs, R. J. Jastrow correlation factor for atoms, molecules, and solids. Phys. Rev. B 70, 235119 (2004).

    Article  ADS  Google Scholar 

  10. Brown, T. et al. Language models are few-shot learners. in Advances in Neural Information Processing Systems Vol. 33 (eds Larochelle, H. et al.) 1877–1901 (Curran Associates, 2020).

  11. Jumper, J. et al. Highly accurate protein structure prediction with AlphaFold. Nature 596, 583–589 (2021).

    Article  ADS  CAS  PubMed  PubMed Central  Google Scholar 

  12. Ren, W., Fu, W., Wu, X. & Chen, J. Towards the ground state of molecules via diffusion Monte Carlo on neural networks. Nat. Commun. 14, 1860 (2023).

    Article  CAS  PubMed  PubMed Central  Google Scholar 

  13. Entwistle, M. T., Schätzle, Z., Erdman, P. A., Hermann, J. & Noé, F. Electronic excited states in deep variational Monte Carlo. Nat. Commun. 14, 274 (2023).

    Article  ADS  CAS  PubMed  PubMed Central  Google Scholar 

  14. Sugar, J. & Corliss, C. Atomic energy levels of the iron-period elements: potassium through nickel. J. Phys. Chem. Ref. Data 14, 1 (1985).

    CAS  Google Scholar 

  15. Page, R. H. & Gudeman, C. S. Completing the iron period: double-resonance, fluorescence-dip rydberg spectroscopy and ionization potentials of titanium, vanadium, iron, cobalt, and nickel. J. Opt. Soc. Am. B 7, 1761–1771 (1990).

    Article  ADS  CAS  Google Scholar 

  16. Sohl, J. E., Zhu, Y. & Knight, R. D. Two-color laser photoionization spectroscopy of Ti i: multichannel quantum defect theory analysis and a new ionization potential. J. Opt. Soc. Am. B 7, 9–14 (1990).

    Article  ADS  CAS  Google Scholar 

  17. James, A. M. et al. Resonant two photon ionization spectroscopy of the molecules V2, VNb, and Nb2. J. Chem. Phys. 101, 4485–4495 (1994).

    Article  ADS  CAS  Google Scholar 

  18. Spiekermann, K., Pattanaik, L. & Green, W. H. High accuracy barrier heights, enthalpies, and rate coefficients for chemical reactions. Sci. Data 9, 417 (2022).

    Article  CAS  PubMed  PubMed Central  Google Scholar 

  19. Goerigk, L. et al. A look at the density functional theory zoo with the advanced GMTKN55 database for general main group thermochemistry, kinetics and noncovalent interactions. Phys. Chem. Chem. Phys. 19, 32184–32215 (2017).

    Article  CAS  PubMed  Google Scholar 

  20. Karton, A. & Martin, J. M. L. Explicitly correlated Wn theory: W1-F12 and W2-F12. J. Chem. Phys. 136, 124114 (2012).

    Article  ADS  PubMed  Google Scholar 

  21. Jurečka, P., Šponer, J., Černy’, J. & Hobza, P. Benchmark database of accurate (MP2 and CCSD(T) complete basis set limit) interaction energies of small model complexes, DNA base pairs, and amino acid pairs. Phys. Chem. Chem. Phys. 8, 1985–1993 (2006).

    Article  PubMed  Google Scholar 

  22. Han, J., Zhang, L. & Weinan, E. Solving many-electron Schrödinger equation using deep neural networks. J. Comput. Phys. 399, 108929 (2019).

    Article  MathSciNet  CAS  Google Scholar 

  23. Lin, J., Goldshlager, G. & Lin, L. Explicitly antisymmetrized neural network layers for variational Monte Carlo simulation. J. Comput. Phys. 474, 111765 (2023).

    Article  MathSciNet  Google Scholar 

  24. Abrahamsen, N. & Lin, L. Taming the sign problem of explicitly antisymmetrized neural networks via rough activation functions. Preprint at https://arXiv.org/2205.12250 (2022).

  25. Gerard, L., Scherbela, M., Marquetand, P. & Grohs, P. Gold-standard solutions to the Schrödinger equation using deep learning: how much physics do we need? in 36th Conference on Neural Information Processing Systems https://openreview.net/forum?id=nX-gReQ0OT (2022).

  26. Li, X., Li, Z. & Chen, J. Ab initio calculation of real solids via neural network ansatz. Nat. Commun. 13, 7895 (2022).

    Article  ADS  CAS  PubMed  PubMed Central  Google Scholar 

  27. Pescia, G., Han, J., Lovato, A., Lu, J. & Carleo, G. Neural-network quantum states for periodic systems in continuous space. Phys. Rev. Res. 4, 023138 (2022).

    Article  CAS  Google Scholar 

  28. Wilson, M. et al. Neural network ansatz for periodic wave functions and the homogeneous electron gas. Phys. Rev. B 107, 235139 (2023).

    Article  ADS  CAS  Google Scholar 

  29. Cassella, G. et al. Discovering quantum phase transitions with fermionic neural networks. Phys. Rev. Lett. 130, 036401 (2023).

    Article  ADS  CAS  PubMed  Google Scholar 

  30. Pescia, G., Nys, J., Kim, J., Lovato, A. & Carleo, G. Message-passing neural quantum states for the homogeneous electron gas. Preprint at https://arXiv.org/2305.07240 (2023).

  31. Qian, Y., Fu, W., Ren, W. & Chen, J. Interatomic force from neural network based variational quantum Monte Carlo. J. Chem. Phys. 157, 164104 (2022).

    Article  ADS  CAS  PubMed  Google Scholar 

  32. Li, X., Qian, Y. & Chen, J. Electric polarization from many-body neural network ansatz. Preprint at https://arXiv.org/2307.02212 (2023).

  33. Li, X., Fan, C., Ren, W. & Chen, J. Fermionic neural network with effective core potential. Phys. Rev. Res. 4, 013021 (2022).

    Article  CAS  Google Scholar 

  34. Wilson, M., Gao, N., Wudarski, F., Rieffel, E. & Tubman, N. M. Simulations of state-of-the-art fermionic neural network wave functions with diffusion Monte Carlo. Preprint at https://arXiv.org/2103.12570 (2021).

  35. Scherbela, M., Reisenhofer, R., Gerard, L., Marquetand, P. & Grohs, P. Solving the electronic schrödinger equation for multiple nuclear geometries with weight-sharing deep neural networks. Nat. Comput. Sci. 2, 331–341 (2022).

    Article  PubMed  Google Scholar 

  36. Gao, N. & Günnemann, S. Ab-initio potential energy surfaces by pairing GNNs with neural wave functions. In 10th International Conference on Learning Representations https://openreview.net/forum?id=apv504XsysP (2022).

  37. Gao, N. & Günnemann, S. Sampling-free inference for ab-initio potential energy surface networks. In 11th International Conference on Learning Representations https://openreview.net/forum?id=Tuk3Pqaizx (2023).

  38. Gao, N. & Günnemann, S. Generalizing neural wave functions. in 40th International Conference on Machine Learning (eds Andreas K. et al.) (2023).

  39. Scherbela, M., Gerard, L. & Grohs, P. Towards a transferable fermionic neural wavefunction for molecules. Nat. Commun. 15, 120 (2024).

  40. Hermann, J., Schätzle, Z., Szabó, P. B., Mezera, M. & DeepQMC contributers. DeepQMC. GitHub https://github.com/deepqmc/deepqmc (2023).

  41. James, S., Spencer, D. P. & FermiNet contributors. FermiNet. GitHub http://github.com/deepmind/ferminet (2023).

  42. Carleo, G. & Troyer, M. Solving the quantum many-body problem with artificial neural networks. Science 355, 602–606 (2017).

    Article  ADS  MathSciNet  CAS  PubMed  Google Scholar 

  43. Adams, C., Carleo, G., Lovato, A. & Rocco, N. Variational monte carlo calculations of A≤4 nuclei with an artificial neural-network correlator ansatz. Phys. Rev. Lett. 127, 022502 (2021).

    Article  ADS  CAS  PubMed  Google Scholar 

  44. Astrakhantsev, N. et al. Broken-symmetry ground states of the heisenberg model on the pyrochlore lattice. Phys. Rev. X 11, 041021 (2021).

    CAS  Google Scholar 

  45. Nomura, Y., Darmawan, A. S., Yamaji, Y. & Imada, M. Restricted Boltzmann machine learning for solving strongly correlated quantum systems. Phys. Rev. B 96, 205152 (2017).

    Article  ADS  Google Scholar 

  46. Saito, H. Solving the Bose-Hubbard model with machine learning. J. Phys. Soc. Japan 86, 093001 (2017).

    Article  ADS  Google Scholar 

  47. Choo, K., Mezzacapo, A. & Carleo, G. Fermionic neural-network states for ab-initio electronic structure. Nat. Commun. 11, 2368 (2020).

    Article  ADS  CAS  PubMed  PubMed Central  Google Scholar 

  48. Yang, P.-J., Sugiyama, M., Tsuda, K. & Yanai, T. Artificial neural networks applied as molecular wave function solvers. J. Chem. Theory Comput. 16, 3513–3529 (2020).

    Article  CAS  PubMed  Google Scholar 

  49. Barrett, T. D., Malyshev, A. & Lvovsky, A. I. Autoregressive neural-network wavefunctions for ab initio quantum chemistry. Nat. Mach. Intell. 4, 351–358 (2022).

    Article  Google Scholar 

  50. Yoshioka, N., Mizukami, W. & Nori, F. Solving quasiparticle band spectra of real solids using neural-network quantum states. Commun. Phys. 4, 106 (2021).

    Article  Google Scholar 

  51. Wang, C., Li, S., He, D. & Wang, L. Is L2 physics informed loss always suitable for training physics informed neural network? in 36th Conference on Neural Information Processing Systems https://openreview.net/forum?id=cy1TKLRAEML (2022).

  52. Raissi, M., Perdikaris, P. & Karniadakis, G. E. Physics-informed neural networks: a deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. J. Comput. Phys. 378, 686–707 (2019).

    Article  ADS  MathSciNet  Google Scholar 

  53. He, D. et al. Learning physics-informed neural networks without stacked back-propagation. in Proc. 26th International Conference on Artificial Intelligence and Statistics (eds Francisco, R.) 3034–3047 (PMLR, 2023).

  54. Mulliken, R. S. Electronic population analysis on LCAO–MO molecular wave functions. i. J. Chem. Phys. 23, 1833–1840 (1955).

    Article  ADS  CAS  Google Scholar 

  55. Paszke, A. et al. PyTorch: An imperative style, high-performance deep learning library. in Advances in Neural Information Processing Systems Vol. 32 (eds Wallach, H. et al.) 8024–8035 (Curran Associates, 2019).

  56. Abadi, M. et al. TensorFlow, large-scale machine learning on heterogeneous systems. in Proc. 12th USENIX Symposium on Operating Systems Design and Implementation (eds Kimberly, K. & Timothy, R.) 265–283 (USENIX, 2016).

  57. Botev, A. The Gauss-Newton Matrix for Deep Learning Models and Its Applications. PhD thesis (Univ. College London, 2020).

  58. Martens, J., Sutskever, I. & Swersky, K. Estimating the Hessian by back-propagating curvature. in 29th International Conference on Machine Learning (2012).

  59. Brown, M. D., Trail, J. R., López Ríos, P. & Needs, R. J. Energies of the first row atoms from quantum Monte Carlo. J. Chem. Phys. 126, 224110 (2007).

    Article  ADS  CAS  PubMed  Google Scholar 

  60. Li, R. et al. Supplementary code of ‘A new computational framework for neural network-based variational Monte Carlo with Forward Laplacian’. Code Ocean https://codeocean.com/capsule/1139728/tree/v1 (2023).

  61. Balabanov, N. B. & Peterson, K. A. Basis set limit electronic excitation energies, ionization potentials, and electron affinities for the 3d transition metal atoms: Coupled cluster and multireference methods. J. Chem. Phys. 125, 074110 (2006).

    Article  ADS  PubMed  Google Scholar 

  62. Marshall, M. S., Burns, L. A. & Sherrill, C. D. Basis set convergence of the coupled-cluster correction, \({\delta }_{{{{\rm{MP2}}}}}^{{{{\rm{CCSD(T)}}}}}\): best practices for benchmarking non-covalent interactions and the attendant revision of the S22, NBC10, HBC6, and HSG databases. J. Chem. Phys. https://doi.org/10.1063/1.3659142 (2011).

Download references

Acknowledgements

We thank H. Li and ByteDance Research for support and inspiration. We thank D. Pfau for his valuable feedback. We thank B. Zhang and H. Wang for their helpful suggestions and discussion. L.W. is supported by the National Key R&D Program of China (grant no. 2022ZD0114900) and the National Science Foundation of China (grant no. NSFC62276005). J.C. acknowledges the National Natural Science Foundation of China for support under grant no. 92165101.

Author information

Authors and Affiliations

Authors

Contributions

R.L. and L.W. conceived the study. H.Y. implemented the main code with contributions from R.L., D.J. and Z.L. X.W., J.C. and W.R. suggested the experiments. R.L., D.J. and H.Y. performed the simulations and analyzed the results. X.W. and X.L. performed the chemical analyses. C.W. and R.L. developed the theoretical formalism. H.Y., R.L., D.H., W.R., D.J. and J.C. performed the figure designing. L.W. and W.R. organized the project. D.H., R.L., H.Y., W.R., J.C., D.J., C.W. and X.W. wrote the paper.

Corresponding authors

Correspondence to Di He, Ji Chen, Weiluo Ren or Liwei Wang.

Ethics declarations

Competing interests

The authors declare no competing interests.

Peer review

Peer review information

Nature Machine Intelligence thanks the anonymous reviewer(s) for their contribution to the peer review of this work.

Additional information

Publisher’s note Springer Nature remains neutral with regard to jurisdictional claims in published maps and institutional affiliations.

Supplementary information

Supplementary Information

Supplementary Figs. 1 and 2, Notes 1–9 and Tables 1–11.

Rights and permissions

Springer Nature or its licensor (e.g. a society or other partner) holds exclusive rights to this article under a publishing agreement with the author(s) or other rightsholder(s); author self-archiving of the accepted manuscript version of this article is solely governed by the terms of such publishing agreement and applicable law.

Reprints and permissions

About this article

Check for updates. Verify currency and authenticity via CrossMark

Cite this article

Li, R., Ye, H., Jiang, D. et al. A computational framework for neural network-based variational Monte Carlo with Forward Laplacian. Nat Mach Intell 6, 209–219 (2024). https://doi.org/10.1038/s42256-024-00794-x

Download citation

  • Received:

  • Accepted:

  • Published:

  • Issue Date:

  • DOI: https://doi.org/10.1038/s42256-024-00794-x

This article is cited by

Search

Quick links

Nature Briefing AI and Robotics

Sign up for the Nature Briefing: AI and Robotics newsletter — what matters in AI and robotics research, free to your inbox weekly.

Get the most important science stories of the day, free in your inbox. Sign up for Nature Briefing: AI and Robotics