A Foundation Model for Simulation-Grade Molecular Electron Densities

Abstract

This paper introduces 3DGrid-VQGAN, a generative framework for the representation and reconstruction of molecular electronic structures as electron charge densities on 3D grid produced by quantum chemical simulations. The model efficiently encodes high-dimensional data into compact latent representations, enabling downstream tasks such as molecular property prediction with enhanced accuracy. Evaluation on the QM9 dataset, which contains quantum chemical properties computed at the density functional theory (DFT) level, demonstrates the model's ability to capture essential features of the electronic structure. The reconstructed charge densities achieve high fidelity, preserving critical details such as electron density cusps at nuclear positions. This is quantified using metrics such as Structural Similarity Index Measure (SSIM), Peak Signal-to-Noise Ratio (PSNR), and Fréchet Inception Distance (FID). These metrics, alongside visual and numerical analyses, demonstrate the model's robustness across diverse molecular structures, including complex geometries and chemical environments. The results suggest that generative approaches like 3DGrid-VQGAN could significantly reduce reliance on computationally intensive quantum chemical simulations, offering simulation-grade data derived directly from learned representations. Future work will focus on extending the model to larger and more complex molecular systems, improving interpretability of latent representations, and integrating the framework into workflows for molecular property prediction and generative design.