Ab InitioWave Function Optimization for Light Atoms A PyTorch Framework for the Psiformer Architecture

Resumen

The application of deep neural networks to Variational Quantum Monte Carlo (VQMC) has significantly advanced the numerical approximation of ground-state energies for quantum many-body systems. Among these advancements, the Psiformer architecture—which leverages self-attention mechanisms to model antisymmetric electron wave functions—has demonstrated state-of-the-art accuracy. However, the original implementation is built exclusively within the JAX ecosystem, presenting a barrier to entry for the broader

research community that predominantly utilizes PyTorch. In the interest of scientific reproducibility and computational accessibility, we present an independent open-source PyTorch implementation of the Psiformer architecture. The model design follows von Glehn, Spencer, and Pfau (2023); the code and numerical experiments reported here are the author’s own. To validate the implementation, we compute ground-state energies for helium through oxygen and compare them to a Hartree–Fock baseline. Our results demonstrate that the PyTorch Psiformer consistently recovers significant electron correlation energy across all tested atoms, confirming that the architecture is correctly implemented and behaves as expected from variational theory. By making the framework openly available, this work lowers the barrier for researchers in computational mathematics, machine learning, and quantum chemistry who wish to experiment with transformer-based wave functions. The full source code is made publicly available to facilitate future experimentation and development.