Improving Kohn–Sham Hamiltonian Learning with Grassmannian Supervision

Predicting the Kohn–Sham Hamiltonian in density functional theory has become an important task in AI for science, as it can accelerate electronic structure calculations and enable the prediction of multiple molecular properties from a single model. However, element-wise Hamiltonian regression does not explicitly enforce the accuracy of the occupied orbital subspace, which determines the density matrix from which ground-state observables are evaluated. This subspace therefore provides a natural auxiliary target for improving downstream property predictions. Yet, directly supervising orbital coefficients in an atomic orbital basis presents two challenges: (1) Gauge redundancy. Rotations among occupied orbitals preserve their subspace and density matrix, making the coefficient representation non-unique. (2) Gradient instability. Orbital-based losses require differentiation through an eigensolver, whose eigenvector derivatives can become unstable near degeneracies. To address these challenges, we propose GrassDist, a principled framework that aligns predicted and reference occupied subspaces using a Grassmannian geodesic loss on the Grassmann manifold. GrassDist is invariant to gauge transformations among occupied orbitals by construction, hence resolving the ambiguity in coefficient supervision. We additionally incorporate two gradient-stabilization treatments to ensure stable optimization. Across multiple benchmarks, GrassDist achieves substantial improvements over various downstream properties, while enabling stable training and faster convergence. On QH9, it reduces density matrix and total energy errors by one to two orders of magnitude compared with the Hamiltonian-only baseline and WALoss. GrassDist also exhibits strong data efficiency: on MD17 uracil, using only 10% to 20% of the training data already matches or surpasses the full-data baseline. The resulting model further generalizes to molecules larger than those seen during training. These results establish GrassDist as a significant methodological advance toward reliable and scalable molecular modeling.