Feeding electron density to a property model was not enough: make it the minimum of a learned functional
Synopsis
Molecular property prediction gains a new route: instead of only feeding the ground-state electron density in as a feature, the density is required to be the minimum of a learned surrogate functional. Density-as-input already improved accuracy, but the density carried no variational constraint. On six of seven extensive QM9 properties this beats the previous best, and models trained only on small molecules keep lower error on larger QM40 molecules.
Interpretation
The input side of molecular property prediction gains a constraint: the ground-state electron density is not only passed into the model, it must sit at the minimum of the functional the model learns, and the property value is read off at that minimum. Earlier density-as-input work treated the density as a physical descriptor only, with no variational condition on it, so nothing forced a prediction to correspond to a minimum. On six of the seven extensive QM9 properties this surpasses the previous best GotenNet-L, and on every extensive property it beats models that use the same ground-state density without the surrogate functional constraint.
The construction moves the surrogate functional from orbital-free density functional theory onto arbitrary scalar molecular properties: one functional is trained per property, with the ground-state energy replaced by that property's label as the training target. Surrogate functionals were previously used for the electronic energy, whose minimum sits at the ground-state density and whose value there is the ground-state energy, and training needs no physical energies at off-equilibrium densities. At inference the density comes from a second network that predicts it from molecular structure, cutting inference to two forward calls; that density predictor reaches a normalized mean absolute error of 0.002 on the QM9 test set, and swapping predicted for label density changes property results only marginally.
The accompanying equivariant message-passing network beSEN replaces scalar-feature gating in the node update and readout with the Gram matrix of the feature channels. The earlier eSEN gates tensorial features with the scalar part of node features and feeds scalar features to the readout. beSEN beats eSEN in all three training modes, geometry only, density input and surrogate, with the largest gap in the surrogate mode; on one idle A100 at batch size 128 its forward pass is slightly faster despite 6.9M versus 2.7M parameters.
Perspective
The result covers small organic molecules in QM9 (up to 9 heavy atoms) and molecules of 10 to 40 heavy atoms in QM40, across seven extensive and several non-extensive properties; density labels come from Kohn–Sham calculations at the B3LYP/6-31G(2df,p) level and are density-fitted. It is aimed at researchers and engineers working on molecular property machine learning and electronic-structure representations.
The motivation for using surrogate functionals on arbitrary molecular properties is less direct than in orbital-free density functional theory, and the authors state that the exact mechanism behind the improvement is not yet fully understood and should be investigated in future work. The density prediction error is also computed against the density-fitted auxiliary-basis label, which is larger than the error against the original Kohn–Sham density.