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

Apr 22, 2026

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

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.

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

Ask PDF

Oct 2, 2026

Electronic Density versus Geometry for Machine-Learned Molecular Absorption Spectra

Molecular optical absorption spectroscopy provides a direct probe of electronic structure and is widely used for molecular identification, interpretation of photophysical behaviour, and planning of spectroscopy experiments. Calculating the absorption spectra using first-principle excited-state methods, however, is computationally demanding, at least compared to ground-state calculations, which limits their routine application across large molecular sets. Machine-learning (ML) surrogates can reduce this cost and allow rapid spectral prediction. However, their performance depends strongly on how molecular information is represented. Here, we compare using the ground-state electron density versus the molecular geometry as inputs to a ML model for predicting absorption spectra, for a training set of 6874 molecules selected from the QM7 dataset. For each of these molecules, the density was calculated using density functional theory (DFT) and the absorption spectrum was calculated using linear-response (LR) time-dependent DFT (TDDFT). Utilizing the ground-state density as the input to the ML model is motivated by the Hohenberg-Kohn and Runge-Gross theorems, and the fact that the ground-state density encodes information about bonding, charge localisation, and electronic delocalisation. Hence, it may be a more judicious starting point for the ML model compared to the geometry, as it effectively decouples the chemistry of the ground-state. The question we test is whether the benefits of using the density outweigh the (notprohibitive) penalty of requiring an additional single-point DFT calculation for the density. We find that the density-based convolutional neural network achieves a validation correlation of 0.9926, compared with 0.9795 for the best geometry-based graph model, reducing the residual decorrelation, by approximately 64%.

Siddharth Dhanpal, Peter Elliott, Paolo Emilio Trevisanutto +2

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.

Molecular Property Prediction Molecular Representation Learning Quantum Chemistry Representation Learning

Ask PDF

Oct 2, 2026

Equivariant Flow Matching for Electron Density Prediction

Machine learning surrogates for density functional theory (DFT) have been increasingly used to reduce the cost of first-principles calculations. In this arena, predicting real-space electron densities offers a scalable and transferable initialization for self-consistent field (SCF) procedures. However, current methods face a clear dilemma. That is, grid-based architectures incur a high computational cost, while basis-set methods fail to capture the structural correlations inherent in the coefficient space. Here, we develop OrbFlow, an

\mathrm{SE}(3)

-equivariant generative model that predicts Gaussian-type orbital (GTO) coefficients via flow matching. OrbFlow retains the efficiency of a compact atom-centered basis while replacing pointwise regression with a learned probability path over the full coefficient space. It is trained through a two-phase trajectory curriculum that mitigates discretization drift during numerical integration. OrbFlow achieves state-of-the-art accuracy on QM9, reducing density error by 13.6% relative to the previous best model, and reduces error by 51% to 63% on every molecule of the MD benchmark relative to the strongest prior method sharing its basis. The predicted density also cuts SCF iterations by up to 68% with zero-shot transfer to unseen exchange-correlation functionals and recovers dipole and quadrupole moments to within a few percent of DFT references without any SCF calculation.

Chenxing Liang, Chengdong Wang, Yuchao Lin +2

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

Computational Quantum Chemistry Quantum Chemistry Density Functional Theory Flow-Based Generative Modeling Molecular Property Prediction Flow Matching

Ask PDF

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