Electron-Density-Informed Molecular Property Prediction Using Surrogate Functionals

Manuel V. Klockow; Simon Wagner; Roman Remme; Dominic Plein; Fred A. Hamprecht

physics.chem-phOct 9, 2026

Electron-Density-Informed Molecular Property Prediction Using Surrogate Functionals

Authors: Manuel V. Klockow, Simon Wagner, Roman Remme, Dominic Plein, Fred A. Hamprecht

Organizations: Interdisciplinary Center for Scientific Computing (IWR), Heidelberg University 69120 Heidelberg, Germany

Abstract

Machine learning methods are increasingly used to predict chemical properties directly from the molecular structure to bypass the often prohibitive cost of quantum mechanical calculations. Recent work has shown that including the ground-state electron density as input to the machine learning model can improve accuracy in this setting. This work employs the recent surrogate functional framework originally developed for orbital-free density functional theory and extends it to predicting molecular properties. We show that passing the ground-state density as input to the model, while simultaneously enforcing it to be the minimum of the surrogate functional, surpasses state-of-the-art performance for six out of seven extensive properties of the QM9 dataset. Additionally, the approach shows improved generalization to larger molecules taken from the QM40 dataset. On all extensive properties of both datasets, our approach outperforms models that are informed by the ground-state density but do not enforce the surrogate functional constraints.

Figures & tables

Figure 1: Schematic overview of the three compared frameworks for molecular property prediction. A Geometry-only models predict the property directly from the molecular geometry. B Density-input models predict the property from the geometry in two steps via the ground-state electron density. C Surrogate functional models predict the property from the geometry and the ground-state density, while enforcing that the density is the minimum of a learned surrogate functional. This variational structure makes it possible to reach the ground-state density and property from an initial guess (magenta arrows). At inference time, we use the shortcut of predicting the ground-state density directly with a second network (blue arrow).

Figure 2: Direct comparison of the eSEN and beSEN architectures. The two architectures share the same overall structure, but differ in multiple details of the message passing, node update and readout. The beSEN architecture is used for both the density prediction and the property prediction models, yielding superior results in both of these settings.

	Extensive							Non-extensive
	$U_{0}$	$U$	$H$	$G$	ZPVE	$\alpha$	$C_{v}$	HOMO	LUMO	gap	$\mu$	$\langle R^{2}\rangle$
Unit	meV	meV	meV	meV	meV	$\text{m}a_{0}^{3}$	$\frac{\mathrm{mcal}}{\mathrm{mol\,K}}$	meV	meV	meV	mD	$\text{m}a_{0}^{2}$
Geometry only
DimeNet++ ( Gasteiger et al., 2020a )	6.32	6.28	6.53	7.56	1.21	44	23	24.6	19.5	32.6	29.7	331
SphereNet ( Liu et al., 2022 )	6.26	6.36	6.33	7.78	1.12	46	22	22.8	18.9	31.1	24.5	268
LEFTNet ( Du et al., 2023 )	6	7	6	7	1.33	48	23	24	18	40	12	109

Table 2: Mean absolute error of the predicted properties on the QM9 test set. Values of prior work are reproduced from Aykent & Xia (2025)

Arm	eSEN	beSEN
Geometry only	5.6	5.3
Density input	5.8	4.9
Surrogate	4.2	2.8
gain vs geometry only, density input	0.96 $\times$	1.1 $\times$
gain vs geometry only, surrogate	1.3 $\times$	1.9 $\times$

Table 3: Backbone comparison on

U_{0}

showing the mean absolute error (MAE) in meV on the QM9 test split. beSEN outperforms the eSEN architecture on all three training modes.

	$U_{0}$	$U$	$H$	$G$	ZPVE	$\alpha$	$C_{v}$	HOMO	LUMO	gap	$\mu$	$\langle R^{2}\rangle$
Unit	meV	meV	meV	meV	meV	$\text{m}a_{0}^{3}$	$\frac{\mathrm{mcal}}{\mathrm{mol\,K}}$	meV	meV	meV	mD	$\text{m}a_{0}^{2}$
Density input
at the label density	4.9	4.8	4.8	6.0	1.1	39	22	16	15	34	7.9	140
at the predicted density	4.9	4.8	4.8	6.0	1.1	39	22	16	15	34	7.9	140
Surrogate
at the label density	2.7	2.3	2.6	4.0	0.91	30	17	16	15	32	12	430

Table 4: Mean absolute error of the predicted properties on the QM9 test set, evaluated at the label density and at the density predicted by the direct density prediction model.

	Extensive							Non-extensive
	$U_{0}$	$U$	$H$	$G$	ZPVE	$\alpha$	$C_{v}$	HOMO	LUMO	gap	$\mu$	$\langle R^{2}\rangle$
Unit	meV	meV	meV	meV	meV	$\text{m}a_{0}^{3}$	$\frac{\mathrm{mcal}}{\mathrm{mol\,K}}$	meV	meV	meV	mD	$\text{m}a_{0}^{2}$
Geometry only
GotenNet-S	310	300	300	220	20	7300	4900	6300	540	—	270	$2.7\times 10^{5}$
GotenNet-B	430	460	450	270	22	7400	5400	4700	520	—	240	—
beSEN	180	180	190	120	150	3000	2600	310	490	1400	150	$2.3\times 10^{5}$

Table 5: Mean absolute error of the predicted properties on 600 QM40 molecules from 10 to 40 heavy atoms. The dataset is split into bins sorted by the number of heavy atoms, with width of 5 atoms. From each of these bins 100 molecules with atom types H, C, N, O and F that are neutral and with even electron count are selected. Each evaluated model was trained only on QM9. Especially for the extensive properties, using the surrogate functional approach enables better generalization to larger molecules.

Figure 4: Absolute

U_{0}

error per atom against molecule size. The surrogate and the density input models are evaluated on the density as predicted by the density prediction model. The median is shown as a horizontal line.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

density	stat	bound	cache	GDI	minao	MAE	gain
						5.3	1.00 $\times$
$\bullet$						4.9	1.09 $\times$
$\bullet$	$\bullet$					4.9	1.08 $\times$
$\bullet$		$\bullet$				4.4	1.20 $\times$
$\bullet$	$\bullet$	$\bullet$				4.6	1.16 $\times$
$\bullet$		$\bullet$	$\bullet$			3.9	1.35 $\times$

Appendix

Table 6: Ablations of different pieces of the surrogate machinery, on

U_{0}

, QM9 test split. density determines whether the model receives the predicted ground-state electron density as input. stat is the stationarity constraint

\nabla_{\bm{p}}F({\bm{p}}^{*})=0

at the ground state, bound and GDI are the energy-bound and GDI loss respectively and cache says whether a cache is used. The minao column indicates how the cache is initialized at each reset, if MINAO is marked the MINAO initialization is used. Otherwise the cache entry is the ground-state density displaced along a random, electron-count-conserving direction, with the L2 displacement radius

r

drawn log-uniformly (density

\propto 1/r

) over

[0.002,0.3]

. gain denotes the accuracy improvement relative to the geometry-only model.

Figure 5: Median absolute error per atom of the geometry-only, density-input, and surrogate beSEN models for all twelve target properties against the number of heavy atoms. The shaded leftmost bin corresponds to the QM9 test set and the remaining bins to the 600 QM40 molecules of Table 5 ; all models were trained on QM9 only. Errors are divided by the total number of atoms.

	beSEN	eSEN
Number of layers	4	4
Channels	32	32
Maximum degree $l_{\text{max}}$	4	4
Cutoff radius (Bohr)	6.0	6.0
Gaussian radial basis functions	32	128
Edge channels	32	32

Appendix

Table 7: Architecture hyperparameters of beSEN and of the eSEN baseline.

Optimizer	AdamW, $(\beta_{1},\beta_{2})=(0.95,0.99)$
Learning rate	$10^{-3}$ , cosine annealing to 0
Epochs	500
Batch size	128
Loss function	$L_{1}$ (property / density targets)

Appendix

Table 8: Training hyperparameters, shared by all models.

GDI step size $\lambda$	0.04
GDI contraction factor $\beta$	0.86
Energy-bound factors $a_{\text{min}}$ , $a_{\text{max}}$	2, 5
Fraction $q$ of ground-state samples	0.5
Loss weights of $\mathcal{L}_{\text{GDI}}$ , $\mathcal{L}_{\text{EB}}$ , $\mathcal{L}_{E}$ , $\mathcal{L}_{\nabla E}$	20, 1, 1, 1
Cache step size $\gamma$	adaptive, minimum $10^{-4}$
Cache initialization	MINAO initial guess

Appendix

Table 9: Hyperparameters of the surrogate training (Section 3.2 ). The energy bounds refer to the scaled labels (Appendix A.4.2 ).

Density Functional Theory

See all 25 papers

Surrogate Functionals for Machine-Learned Orbital-Free Density Functional Theory
Apr 22, 2026Roman Remme, Fred A. HamprechtSurrogate Modeling Quantum Chemistry
Interdisciplinary Center for Scientific Computing (IWR), Heidelberg University
Electronic Density versus Geometry for Machine-Learned Molecular Absorption Spectra
Oct 2, 2026Siddharth Dhanpal, Peter Elliott, Paolo Emilio Trevisanutto +2Molecular Property Prediction Molecular Representation Learning
Scientific Computing Department, Science and Technology Facilities Council, Rutherford Appleton Laboratory, Didcot OX11 0QX, United Kingdom. · Scientific Computing Department, Science and Technology Facilities Council, Daresbury Laboratory, Warrington, WA4 4AD, United Kingdom.
Equivariant Flow Matching for Electron Density Prediction
Oct 2, 2026Chenxing Liang, Chengdong Wang, Yuchao Lin +2Computational Quantum Chemistry Quantum Chemistry
Department of Computer Science and Engineering, Texas A&M University · Department of Materials Science and Engineering, Texas A&M University · Department of Electrical and Computer Engineering, Texas A&M University +2

Apr 22, 2026J/K move · Enter open · S save

Surrogate Functionals for Machine-Learned Orbital-Free Density Functional Theory

Roman Remme, Fred A. Hamprecht

Interdisciplinary Center for Scientific Computing (IWR), Heidelberg University

Surrogate Modeling Quantum Chemistry Adaptive Sampling Computational Materials Science Density Functional Theory Computational Quantum Chemistry

Open paper AskPDF arXiv

We introduce surrogate functionals: machine-learned energy functionals for orbital-free density functional theory (OF-DFT) which are defined not by universal fidelity to a physical reference, but merely by the requirement that density optimization with a fixed procedure yields the true ground-state density. Helpfully, training surrogate functionals requires only ground-state densities, no energies or gradients away from the ground state. We here propose a gradient-descent-improvement loss that guarantees exponential convergence of the density to the ground state, and combine it with an adaptive sampling scheme that concentrates learning around the optimization trajectories actually visited during inference. On the QM9 and QMugs benchmarks, surrogate functionals achieve density errors competitive with or improving upon the state of the art for fully supervised machine-learned OF-DFT, while eliminating the need for the

O(N^3)

orthononormalization step required by prior work, yielding improved runtime scaling for larger systems.

Abstract

Density Functional Theory

Explore similar work

Surrogate Functionals for Machine-Learned Orbital-Free Density Functional Theory

Electronic Density versus Geometry for Machine-Learned Molecular Absorption Spectra

Equivariant Flow Matching for Electron Density Prediction