Modeling electron density grids using a 3D VQ-GAN approach

To convert 3D electron density grids into meaningful latent representations, vector quantized autoencoders have proven effective, particularly in addressing the blurriness typical of traditional variational autoencoders. These models enhance clarity by mapping latent feature vectors at the bottleneck of the autoencoder to a quantized representation derived from a learned codebook. The Vector Quantized Generative Adversarial Network (VQ-GAN), further improves image reconstruction quality by incorporating a discriminator loss. We introduce a foundation model based on a 3D VQ-GAN approach tailored for 3D electron density grids. This model has been trained on approximately 850K grids, totaling 7 TB of data, using 600 GPUs at the Argonne Leadership Computing Facility. The framework processes 3D electron density grids through an encoder that generates a latent code representing the grid's dimensions. The decoder reconstructs the 3D electron density grid using these quantized feature vectors. The results from this generative “no simulation” approach, showcases the ability to produce simulation-grade molecular electronic structures at Hartree-Fock (HF) or density functional theory (DFT) levels. Simulating molecular electronic structures requires solving the electronic Schrödinger equation, parameterized by atomic nuclei positions. Through advanced numerical methods, we can derive the electron density, which indicates the probability of locating an electron in an infinitesimal volume around a point. The electron density exhibits cusps at nuclear positions, encoding critical information about atomic locations. By partitioning the electron density into atomic volumes—using techniques like the atoms-in-molecules formalism—we can reconstruct atomic charges and types. The holographic electron density theorem suggests that a nonzero segment of the ground-state electron density can fully determine the entire molecular electron density. These insights highlight the importance of generative modeling for molecular electron densities, capturing the most informative aspects of quantum chemical simulations. The general architecture of the 3D electron density grids VQ-GAN model.