Physically Interpretable Machine‑Learning Compression of Kohn–Sham Wavefunctions

Density functional theory (DFT) is the workhorse for electronic structure modeling, but the Kohn–Sham (KS) wavefunctions represented in the plane wave basis are massive, creating storage and I/O bottlenecks and limiting downstream analyses. We present a variational-autoencoder (VAE) based framework that learns an interpretable and low-dimensional representation for KS wavefunctions. The learned latent space retains physically meaningful content and encode crystal structures and symmetries, enabling prediction of wavefunctions for rotated structures and interpolation from coarse k-grids to fine k-grids. Despite aggressive compression, the representation preserves sufficient information to reconstruct operator matrix elements thereby capturing interactions and correlations among KS states. We further demonstrate transferability across pseudopotentials and exchange–correlation approximations: models trained on PBE data accurately reconstruct results obtained with LDA, DFT+U, and even hybrid functionals. The compressed embeddings also support downstream predictions of quasiparticle band structures within GW and excitonic spectra from the Bethe Salpeter Equation. These results show that our machine learning based compression is not a simple interpolation. Instead, it provides a compact, physically meaningful description that enables faithful reconstruction of KS wavefunctions.