Spectrally Constrained Electronic Density Matrix Learning with Subspace Alignment

Learning electronic-structure matrices from density functional theory (DFT) has emerged as an active area of atomistic machine learning, providing electronic information beyond the energies and forces typically targeted by machine learning interatomic potentials (MLIPs). Machine-learned Hamiltonian models (MLHs) and density-matrix models (MLDMs) are two representative approaches. MLDMs avoid the numerically sensitive conversion from predicted Hamiltonians to density matrices, but often still underperform MLHs in energy evaluation. We posit that this performance gap arises in part from how density-matrix learning treats eigenvalues and eigenspaces. We introduce a spectrally constrained framework combining Cayley rotations with an occupied-virtual (OV) loss. Cayley rotations preserve the eigenvalues of a valid initial density in an orthonormal representation, leaving the eigenspaces as the remaining degrees of freedom. The OV loss targets the first-order density errors induced by small rotations away from the target eigenspaces. On QH9, our framework outperforms QHFlow2, a state-of-the-art Hamiltonian model, in energy and real-space density accuracy, achieving a mean one-shot energy error of $0.184$ meV with no errors above $100$ meV. The predicted densities also accelerate self-consistent field (SCF) calculations by reducing effective iteration cost. In our scaling comparison, a 65.9M-parameter Cayley model with OV loss achieves energy accuracy comparable to a 514M-parameter model without it. Our framework generalizes to larger QM40 molecules without fine-tuning, achieving a mean energy error of $0.98$ meV/atom and improved SCF initialization.