Learning Electron Density with Virtual Atoms for Transferable Atomistic Representations

Abstract

Electron density is a fundamental quantity of electronic structure, yet most machine learning density models treat it solely as a reconstruction target and are developed separately from machine learning interatomic potentials. Here, we introduce a virtual-atom formulation within the DPA3 framework, in which density query points are represented as virtual nodes and processed by the same parameterized encoder used for real atoms. An asymmetric neighbor construction allows virtual atoms to receive information from nearby real atoms without altering the real-atom interaction graph. The resulting framework supports density, energy, and force prediction through task-specific readout heads. The proposed model achieves state-of-the-art density reconstruction accuracy on the Materials Project benchmark, matches the best reported performance on nickel-manganese-cobalt battery cathode materials, and is competitive with dedicated density models on QM9. When density labels are used jointly with energy and force supervision, the model improves in-domain accuracy, makes local environment information more accessible under linear probing, and reduces zero-shot force errors on larger QM40 molecules. Density-only checkpoints also provide useful initializations for downstream energy-force learning under the fixed fine-tuning protocol. Together, these results support electron density as a source of electronic-structure supervision for interatomic-potential representations.