Neural Quantum States in Mixed Precision

TL;DR

This work tackles the computational bottleneck of sampling in neural-quantum-state variational Monte Carlo by leveraging mixed-precision arithmetic. It develops a rigorous framework that models finite-precision effects as log-density perturbations and derives bounds on the resulting asymptotic bias of MH sampling, with tighter results when the underlying chain mixes rapidly. The authors introduce a mixed-precision VMC workflow that performs sampling in low precision while keeping core computations in high precision, and validate it through toy and realistic quantum models, achieving up to $\$3.5\times$ speedups without sacrificing accuracy in ground-state estimates. These results provide a principled pathway to energy-efficient, scalable MCMC-driven simulations in quantum many-body problems and related machine-learning approaches, with broader applicability to Bayesian learning and energy-based models.

Abstract

Scientific computing has long relied on double precision (64-bit floating point) arithmetic to guarantee accuracy in simulations of real-world phenomena. However, the growing availability of hardware accelerators such as Graphics Processing Units (GPUs) has made low-precision formats attractive due to their superior performance, reduced memory footprint, and improved energy efficiency. In this work, we investigate the role of mixed-precision arithmetic in neural-network based Variational Monte Carlo (VMC), a widely used method for solving computationally otherwise intractable quantum many-body systems. We first derive general analytical bounds on the error introduced by reduced precision on Metropolis-Hastings MCMC, and then empirically validate these bounds on the use-case of VMC. We demonstrate that significant portions of the algorithm, in particular, sampling the quantum state, can be executed in half precision without loss of accuracy. More broadly, this work provides a theoretical framework to assess the applicability of mixed-precision arithmetic in machine-learning approaches that rely on MCMC sampling. In the context of VMC, we additionally demonstrate the practical effectiveness of mixed-precision strategies, enabling more scalable and energy-efficient simulations of quantum many-body systems.

Figures (12)

• Figure 1: Panel (a) displays the acceptance difference $|\Delta \alpha(x,y)|$ as a function of the acceptance ratio $s(x,y)$ and the error difference $\varepsilon(x,y)$. The solid red line corresponds to $s(x,y) = e^{-\varepsilon(x,y)}$. For $\varepsilon(x,y)<0$ the plot is inverted over both the x-axis and y-axis. Panels (b) and (c) present results for the noisy RBM. Panel (b) shows the relative error between expectation values computed with the unperturbed and noisy RBM as a function of noise parameter $\sigma$ compared to the KL bound of Eq. \ref{['eq:tv_gaussian']} and the average bound of Eq. \ref{['eq:avg_bound']}. The expected values of both the Pauli-X operator $\sigma^x$ and the projector onto the even bitstrings $\Pi_{\text{Even}}$ are investigated. Individual points correspond to different random initializations, while shaded regions indicate confidence intervals determined by Monte Carlo sampling noise, calculated as $3\sqrt{\text{Var}[\epsilon_{\text{rel}}]/N_{\text{samples}}}$, where the factor of 3 corresponds to approximately a $99.7\%$ confidence interval assuming a Gaussian distribution of the noise. Results are obtained using $2^{18}$ samples, with the average bound computed for $r=0$. Panel (c) shows the acceptance rate as a function of $\sigma$ for different values of $h/J$ in the TFIM. The dark line represents the bound for $|\Delta\alpha|$, while the dashed lines correspond to the value of the unperturbed acceptance rate $\alpha$.
• Figure 2: Standardized law of $\delta$ for a random state (panel (a)) and for the ground state of the TFIM in the antiferromagnetic phase at $h/J = 0.5$ (panel (b)). The distributions are computed for different data types on a spin chain of length $N = 12$, using all configurations of the Hilbert space as defined in Eq. \ref{['eq:pert_chain']}. Panel (c) shows the standard deviation of each $\delta$ distribution of the states in panels (a) and (b).
• Figure 3: Panel (a) shows the relative energy error with respect to the double-precision reference solution as a function of training step. Optimization is performed using the mixed-precision scheme introduced in this work, with sampling carried out in reduced precision. We use a ResCNN with kernel size $(3,3)$, four residual blocks, and 16 filters to study the two-dimensional TFIM at the critical point with linear system size $L=10$ ($N=L^2$). Panel (b) reports, for the same optimization, the minimum between the KL bound in Eq. \ref{['eq:tv_gaussian']} and the bound in Eq. \ref{['eq:avg_bound']}, evaluated on the absolute difference between gradients from perturbed and unperturbed distributions. Shaded regions indicate the Monte Carlo confidence interval, given by $3\sqrt{2}\sqrt{\mathrm{Var}[\nabla \epsilon_\theta(x)]/N_{\mathrm{samples}}}$, where $\sqrt{2}$ accounts for error propagation in the gradient difference and the factor of 3 corresponds to a $99.7\%$ Gaussian confidence level.
• Figure 4: Sampling speedup obtained on a NVIDIA H100 GPU with RBMs with different data types. Results are shown as a function of both the number of samples $N_s$, with the number of Markov chains set to $N_s/4$, and the density of parameters $\alpha$.
• Figure 5: Speedup over different data types relative to double precision, obtained by multiplying two matrices of shape $(N, N)$ on an NVIDIA H100 GPU. Each result is obtained by averaging over 50 matrix multiplications.

Theorems & Definitions (9)

• Theorem 3.1
• Theorem 3.2
• Theorem 3.3: Local MH with Gaussian Noise
• Lemma : \ref{['lem:piPbound']}
• proof : Proof of Lemma \ref{['lem:piPbound']}
• Lemma : \ref{['lem:varepsbound']}
• proof : Proof of Lemma \ref{['lem:varepsbound']}
• Theorem : \ref{['thm:tv_markov']}
• proof : Proof of Theorem \ref{['thm:tv_markov']}