On Policy Self Distillation for Learning Kohn-Sham Hamiltonian

Neural networks that predict the Kohn–Sham Hamiltonian provide initial guesses that shorten self-consistent field (SCF) calculations. However, every training label is itself a converged SCF calculation. Extending a model to new molecules or functionals therefore requires a new labeled dataset. We observe that a few SCF iterations on the model's own prediction return a corrected Hamiltonian, which indicates how every matrix entry should change. Building on this, we propose on-policy self-distillation with the SCF map (OPSD), which uses the SCF solver as the teacher. A pretrained flow model predicts Hamiltonians, the solver corrects them, and the model learns from the corrections before generating the next round. The model thus amortizes the solver, using only geometries and standard SCF software. In eight transfer settings, including a new functional, implicit solvent, and larger molecules, OPSD has a lower mean error than fine-tuning on converged labels. After a change of functional, for example, the HOMO–LUMO gap error is 2.2 mHa with OPSD and 42 mHa with fine-tuning. On QH9, OPSD lowers the gap error of QHFlow2 models by up to 68%. On the 58-element HELM benchmark, it stabilizes the orbital spectrum of a pretrained QHFlow2 model. OPSD could thus serve as a general method for improving the transferability of Hamiltonian models.