Skip to main content

New paper: Neural Quantum States in Mixed Precision

This work grew out of a particularly synergistic interdisciplinary collaboration with Massimo Solinas, Agnes Valenti, and Roeland Wiersema at the Center for Computational Quantum Physics, Flatiron Institute.

A basic, but still underappreciated, fact outside machine learning circles is that GPUs are substantially more efficient at lower numerical precision, sometimes by orders of magnitude. Throughput increases, memory traffic decreases, and energy efficiency improves when computations are performed in half or mixed precision, provided the algorithm can tolerate the resulting reduced-precision error.

This caveat is especially important for sampling algorithms that arise, e.g., in scientific machine learning. Markov chain Monte Carlo (MCMC) methods are routinely used to sample high-dimensional distributions, where small numerical errors can accumulate along the Markov chain. While infinite-precision MCMC has been studied extensively, the statistical consequences of finite-precision arithmetic are much less well understood.

In this paper, we ask:

What asymptotic bias does finite-precision arithmetic introduce into MCMC?

Main contributions

On the theoretical side, we derive general analytical bounds on the asymptotic bias induced by reduced-precision Metropolis–Hastings MCMC. These results make explicit how finite numerical precision perturbs the invariant distribution of the Markov chain.

On the computational side, we validate these bounds in the context of neural-network–based Variational Monte Carlo (VMC) for quantum many-body systems. In particular, we show that the ground state approximation can be carried out safely in half precision, with no observable loss of accuracy, while delivering substantial performance gains on GPUs. Depending on the neural-network ansatz, we observe speedups ranging from roughly 4x to 20x.

Broader Implications

From an applications perspective, these results enable faster and more energy-efficient simulations of quantum many-body systems, which are among the most computationally demanding problems out there.

From a mathematical perspective, the work provides a principled framework for reasoning about mixed-precision MCMC, moving beyond heuristic arguments. Rather than treating numerical precision as a fixed background constraint, precision becomes an explicit part of the algorithmic analysis.

More broadly, the framework applies to machine-learning methods that rely on MCMC, where numerical precision plays a role comparable to step sizes, tuning parameters, etc.