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

TL;DR

This paper introduces jVMC, a Python-based variational Monte Carlo framework that uses automatic differentiation and GPU acceleration to efficiently simulate quantum systems with neural network wave functions.

Contribution

It provides a versatile, high-performance codebase supporting arbitrary neural quantum states and Hamiltonians, optimized with modern computational techniques.

Findings

Supports arbitrary NQS architectures and Hamiltonians

Leverages automatic differentiation and GPU acceleration

Facilitates efficient NQS algorithm development

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.