Orbital-Free Surrogate Functionals
Yield Transferable Interatomic Potentials and Electron Densities

In OF-DFT, the goal is to construct a functional $E[\rho]$, which for a given electronic system maps an electron density $\rho$ to a corresponding energy $E$. Minimizing $E[\rho]$ over $\rho$ then yields the ground-state energy $E^{*}$ and density $\rho^{*}$ of the system. The density is often expressed as an expansion using atom-centered Gaussian-type basis functions. We follow [40, 31] in applying symmetric Löwdin orthonormalization, yielding an expansion

with density coefficients $p_{\mu}$ and orthonormal basis functions $\omega_{\mu}(\mathbf{r})$.

The energy functional can be split into different physical contributions:

where $T_{\text{S}}$ is the non-interacting kinetic energy, $E_{\text{XC}}$ is the exchange-correlation energy, $E_{\text{H}}$ is the Hartree energy, and $E_{\text{ext}}$ is the interaction energy with an external potential.

In machine-learned OF-DFT [40, 31, 30], the energy functional or parts of it are learned by training on labels obtained from Kohn–Sham reference data. Those labels correspond to electron densities at the electronic ground state as well as off-equilibrium densities extracted from intermediate self-consistent field iterations. During inference, the electron density is initialized using standard methods such as a superposition of atomic densities (SAD). Afterwards, the density is optimized using a gradient-descent scheme where the required gradients of the functional $\gradient_{\mathbf{p}}E_{\text{tot}}$ are obtained by automated differentiation of the neural network.

Weak surrogate functionals [30] lift the requirement of global fidelity to a physical energy functional and instead define a successful functional exclusively through its minimizer in a fixed density optimization procedure. Specifically, a weak surrogate functional $\bar{E}$ must yield the true ground-state density coefficients $\mathbf{p}^{*}$ when minimized from a prescribed initial guess. This allows training using only ground-state density labels, avoiding the need for off-equilibrium densities.

To guarantee convergence during inference, weak surrogates can be trained with the gradient-descent-improvement (GDI) loss. This loss enforces that every optimization step with step size $\lambda$ moves the coefficients closer to the true ground state by at least a predefined contraction factor $0<\beta<1$. With $\mathbf{p}^{\prime}=\mathbf{p}-\lambda\gradient_{\mathbf{p}}\tilde{E}_{\text{tot}}[\mathbf{p};\theta]$ and $[x]_{+}=\max(0,x)$, the loss [30] is

The contraction factor is chosen such that the GDI loss vanishes for most Kohn–Sham labels from off-equilibrium densities as used in [31]. In this sense, the GDI loss is compatible with the physical energy landscape as sampled by these training labels. More details on the choice of $\beta$ and $\lambda$ can be found in Appendix C.3.

To focus model capacity on the regions actually visited during inference, this loss is combined with an adaptive train-time density optimization scheme (see Figure 1(b)). For each molecule $\mathcal{M}$, a cached density coefficient vector $\mathbf{p}_{\mathcal{M}}$ is maintained and updated during training each time the molecule appears in a batch via a gradient descent step:

Note that the step size $\gamma$ may be different from the step size $\lambda$ that is used in the GDI loss. In the original framework, $\gamma$ was kept fixed. Here, we introduce a dynamically computed “ideal” learning rate for the cache updates:

This rate $\gamma^{*}$ corresponds to the nonnegative step size that minimizes the Euclidean distance between the updated density coefficients and the ground-state target $\mathbf{p}^{*}$ for the currently predicted gradient direction (see Appendix A for details). Employing this optimal local step size ensures that train-time trajectories progress as efficiently as possible toward the ground state. At the same time, samples for which the predicted gradient would lead to a step away from the ground state are not updated at all, creating a natural training dynamic where difficult-to-learn densities are kept in cache until the model learns to improve the gradient predictions at these points.

While weak surrogates only ensure convergence to the correct ground-state electron density, strong surrogate functionals are designed to also accurately model the energy at this ground state. From a machine learning perspective, this requires an additional loss term to enforce the true energy at the predicted minimum. However, in the adaptive training scheme described in Section IV.1, the exact ground-state density is never explicitly sampled. The central challenge is therefore how to provide meaningful energy feedback. To address this, we propose a two-step strategy: First, we inject a fraction $0\leq q<1$ of exact ground-state samples into each training batch. For these molecules, instead of using their cached density coefficients, we evaluate the model strictly at the true ground-state density and apply $L_{1}$ penalty losses on both the predicted ground-state energy and its gradients, which should be zero at the ground state: $\mathcal{L}_{E}=\sum_{\mathcal{M}}\left(|\tilde{E}_{\text{tot}}(\mathbf{p}^{*}_{\mathcal{M}})-E_{\text{tot}}(\mathbf{p}^{*}_{\mathcal{M}})|\right)$ and $\mathcal{L}_{\gradient E}=\sum_{\mathcal{M}}\left(|\gradient_{\mathbf{p}}\tilde{E}_{\text{tot}}(\mathbf{p}^{*}_{\mathcal{M}})|\right)$. Incorporating a fraction of exact ground-state labels also provides the architectural flexibility, in principle, to enforce other ground-state properties, such as nuclear gradients.

Second, while exact energy labels at arbitrary off-equilibrium densities are not available, we can exploit physical and empirical energy bounds to construct a robust loss function. A natural lower bound is the ground-state energy itself, since it represents the global minimum of the exact functional. However, penalizing predictions only when they fall below the exact ground-state energy yields a very weak learning signal. Instead, we derive tighter empirical lower and upper bounds from off-equilibrium energies and densities, which we take from [31]. As illustrated in Figure 2, these labels form a set that, empirically, is bounded above and below by two quadratic forms:

where $\mathbf{S}$ is the overlap matrix of the basis functions in the molecule. Based on this observation, we introduce the following energy-bound loss:

The total loss is then given by a weighted sum of all loss components:

The choice of $E_{\text{tot}}$ as energy target is practical, as both the GDI loss and the density cache updates require gradients corresponding to the total energy functional. Nevertheless, there are compelling reasons to use neural networks to model only specific energy components, while computing the remaining contributions using exact physical expressions: (1) To facilitate generalization across diverse chemical systems, it is advantageous to learn only the universal functional components, isolating them from the system-specific external potential. (2) Neural network backbones typically apply a radial spatial cutoff, limiting their ability to model long-range interactions. Explicitly evaluating the Hartree energy enables the model to accurately capture the nonlocal electrostatic contributions. Consequently, learning the sum $E_{\text{TXC}}$ of the non-interacting kinetic energy $T_{\text{S}}$ and the exchange-correlation energy $E_{\text{XC}}$ represents a physically motivated objective, consistent with prior methodologies [40, 31]. In practice this means parameterizing the total energy functional as

where $\tilde{E}_{\text{TXC}}$ is the contribution that is learned by the neural network.

Prior work in OF-DFT [40, 31, 30], including surrogate functionals for OF-DFT, has modeled the density representation using Graphormer-style architectures [39], relying on local canonicalization [23] to express the density coefficients as scalar graph features. SO(3)-equivariant neural networks operating on irreducible tensor features provide a strong alternative for molecular modeling, especially when the target quantities naturally transform under rotations. We therefore evaluate eSEN [12] and EquiformerV2 [22], two modern equivariant architectures. Both architectures build on the efficient eSCN formulation [27], which reduces the cost of equivariant convolutions and tensor-product-like interactions. The main architectural difference is the edge aggregation mechanism: eSEN uses edge-conditioned equivariant message passing with unnormalized sum aggregation, whereas EquiformerV2 uses an SO(3)-equivariant multi-head attention block with softmax-normalized edge weights. To feed the density information into these architectures, a new density embedding is created. The coefficients are first rescaled independently in each dimension and subsequently scattered into a fixed-size “atom-hot” tensor per atom, as in the Graphormer-style architectures. A single equivariant linear layer is then employed to produce the correct structure of irreducible representations that is needed for both eSEN and EquiformerV2. The resulting embedding is added directly onto the respective embeddings employed in the original architectures. Furthermore, an atomic reference module is used to add element-dependent atomic energy biases to the output of the network. In contrast to [40, 31], we do not employ a global energy bias, which we found to degrade the generalization performance to larger molecules.

We train and evaluate our models on molecules from the QM9 dataset [29]. For ground-truth energies and densities, we rely on the same labels utilized in [31], restricted to ground-state labels, which were generated using the PBE functional [28]. Following the methodology established in prior work [40, 31, 30], we test the extrapolation capabilities of our approach by training models on a combined pool of QM9 and QMugs [17] molecules limited to 15 heavy atoms. We then evaluate the models on a distinct test split of significantly larger QMugs configurations containing up to 100 heavy atoms (refer to Appendix C.1 for full details).

We benchmark our accuracy against established machine-learned orbital-free functionals, namely M-OFDFT [40] and STRUCTURES25 [31], as well as weak surrogate functionals from [30]. As additional baselines, we train GotenNet—the current state-of-the-art MLIP on QM9—as well as vanilla eSEN and EquiformerV2 architectures without density embeddings, using only ground-state energy data. Comprehensive model specifications and training hyperparameters are documented in Appendix C.2. Models trained with the strong surrogate pipeline are denoted with the suffix -SF.

Performance metrics on the QM9 dataset are presented in Table 1. Remarkably, all surrogate variants surpass the existing OF-DFT baselines by a wide margin, demonstrating substantial improvements in both energy predictions and density accuracy. Concurrently, models trained with $E_{\text{tot}}$ targets display notably lower energy errors compared to baseline models trained exclusively on molecular geometries without auxiliary electron density information.

Identical performance trends persist when evaluating the models on the larger QMugs molecules, as summarized in Table 2. Notably, the relative performance margin between our strong surrogate functionals and the MLIP baselines widens on this out-of-distribution set.

While the surrogate functionals trained on $E_{\text{TXC}}$ show higher energy and density errors than their $E_{\text{tot}}$ counterparts, they still outperform all previous OF-DFT baselines trained on the same energy target, demonstrating that the train-time density optimization procedure on the auxiliary functional $\tilde{E}_{\text{aux}}$ is sufficient to learn a converging functional.

Compared to the OF-DFT baselines, surrogate functionals show density errors that are lower by approximately one order of magnitude. Furthermore, they do not exhibit large outliers as is the case for STRUCTURES25 [31]. A visual comparison of density errors for example molecules is shown in Figure 3.

It should be noted that all previous machine-learned OF-DFT functionals shown in Tables 1 and 2 are based on a Graphormer architecture [39], which is less expressive than either eSEN or EquiformerV2 and thus explains part of the performance difference to the surrogate functional. However, we found that retraining STRUCTURES25 with either of the alternative architectures instead of a Graphormer yields a functional that does not converge during density optimization. In Figure 4, we visualize the distribution of eigenvalues of the Hessian of an eSEN model trained with the STRUCTURES25 training pipeline and the surrogate pipeline, respectively, evaluated at the ground-state density of a random QM9 molecule. The surrogate functional has only positive eigenvalues that all lie within a narrow interval, confirming that the model has indeed learned a convex functional as enforced by the GDI and energy-bound losses. The eSEN-STR25 model, on the other hand, displays a much broader range of Hessian eigenvalues including both much larger and crucially also negative values, leading to a functional that does not converge during density optimization. The fact that we obtain convergence irrespective of the network architecture thus presents a significant advantage of the surrogate-style training.

In order to isolate the contribution of newly introduced components to the strong surrogate training performance, we conduct targeted ablation experiments. We train an EquiformerV2 surrogate functional under two modified constraints: either disabling the energy-bound loss entirely, or replacing the optimally computed ideal learning rate with a fixed cache update learning rate of $10^{-2}$. The resulting evaluations are detailed in Table 3.

Both ablation scenarios show a degradation in predictive performance, validating the effectiveness of these adaptations within the training pipeline.