jVMC: Versatile and performant variational Monte Carlo leveraging automated differentiation and GPU acceleration

Abstract

The introduction of Neural Quantum States (NQS) has recently given a new twist to variational Monte Carlo (VMC). The ability to systematically reduce the bias of the wave function ansatz renders the approach widely applicable. However, performant implementations are crucial to reach the numerical state of the art. Here, we present a Python codebase that supports arbitrary NQS architectures and model Hamiltonians. Additionally leveraging automatic differentiation, just-in-time compilation to accelerators, and distributed computing, it is designed to facilitate the composition of efficient NQS algorithms.

arXiv comment33 pages, 7 figures. Revised version. Code repository: https://github.com/markusschmitt/vmc_jax

References in corpus (43)

• Array Programming with NumPy
• The density-matrix renormalization group in the age of matrix product states
• Electronic Structure Calculations with Dynamical Mean-Field Theory: A Spectral Density Functional Approach
• Solving the Quantum Many-Body Problem with Artificial Neural Networks
• A Practical Introduction to Tensor Networks: Matrix Product States and Projected Entangled Pair States
• Time-dependent variational principle for quantum lattices
• Many-body quantum state tomography with neural networks
• Time-evolution methods for matrix-product states
• Deep neural network solution of the electronic Schrödinger equation
• Ab-Initio Solution of the Many-Electron Schrödinger Equation with Deep Neural Networks
• Real time confinement following a quantum quench to a non-integrable model
• Reconstructing quantum states with generative models
• Neural-Network Approach to Dissipative Quantum Many-Body Dynamics
• Recurrent Neural Network Wave Functions
• Deep autoregressive models for the efficient variational simulation of many-body quantum systems
• Symmetries and many-body excited states with neural-network quantum states
• Fermionic neural-network states for ab-initio electronic structure
• Study of the Two-Dimensional Frustrated J1-J2 Model with Neural Network Quantum States
• Quantum many-body dynamics in two dimensions with artificial neural networks
• Massive Parallel Quantum Computer Simulator
• Neural-Network Quantum States, String-Bond States, and Chiral Topological States
• Neural Quantum States of frustrated magnets: generalization and sign structure
• Massively parallel quantum computer simulator, eleven years later
• NetKet: A Machine Learning Toolkit for Many-Body Quantum Systems
• Green Function Monte Carlo with Stochastic Reconfiguration: an effective remedy for the sign problem disease
• Quantum phase transition dynamics in the two-dimensional transverse-field Ising model
• Chiral Topological Phases from Artificial Neural Networks
• Geometry of variational methods: dynamics of closed quantum systems
• Quenches near Ising quantum criticality as a challenge for artificial neural networks
• Autoregressive Transformer Neural Network for Simulating Open Quantum Systems via a Probabilistic Formulation
• Neural network representation of tensor network and chiral states
• Learning the ground state of a non-stoquastic quantum Hamiltonian in a rugged neural network landscape
• Helping restricted Boltzmann machines with quantum-state representation by restoring symmetry
• Geometry of learning neural quantum states
• Probabilistic Simulation of Quantum Circuits with the Transformer
• Investigating ultrafast quantum magnetism with machine learning
• Neural-network states for the classical simulation of quantum computing
• Role of stochastic noise and generalization error in the time propagation of neural-network quantum states
• Unitary long-time evolution with quantum renormalization groups and artificial neural networks
• Artificial Neural Network Based Computation for Out-of-Time-Ordered Correlators
• Simulating disordered quantum systems via dense and sparse restricted Boltzmann machines
• Neural Autoregressive Distribution Estimation
• Time-dependent variational principle for open quantum systems with artificial neural networks

Cited by in corpus (9)

• Empowering deep neural quantum states through efficient optimization
• Optimizing Design Choices for Neural Quantum States
• Highly resolved spectral functions of two-dimensional systems with neural quantum states
• Wave function network description and Kolmogorov complexity of quantum many-body systems
• Dynamics of correlation spreading in low-dimensional transverse-field Ising models
• Variational Monte Carlo Approach to Partial Differential Equations with Neural Networks
• Zero-temperature Monte Carlo simulations of two-dimensional quantum spin glasses guided by neural network states
• Efficient Optimization of Variational Autoregressive Networks with Natural Gradient
• Simulating dynamics of the two-dimensional transverse-field Ising model: a comparative study of large-scale classical numerics