SpookyNet: Learning Force Fields with Electronic Degrees of Freedom and Nonlocal Effects

Oliver T. Unke; Stefan Chmiela; Michael Gastegger; Kristof T. Schütt; Huziel E. Sauceda; Klaus-Robert Müller

Source-linked AI summary

SpookyNet: Learning Force Fields with Electronic Degrees of Freedom and Nonlocal Effects

Oliver T. Unke, Stefan Chmiela, Michael Gastegger, Kristof T. Schütt, Huziel E. Sauceda, Klaus-Robert Müller

arXiv:2105.00304v2May 1, 2021physics.chem-phcs.LG

TL;DR

Existing ML-FFs often cannot distinguish systems with different electronic states and may fail when nonlocal effects matter. SpookyNet addresses both issues with explicit electronic inputs, attention-based nonlocal interactions, and physically motivated corrections, achieving broad generalization and improved benchmark performance while retaining interpretable chemical structure. Its applicability remains challenged by large heterogeneous condensed-phase systems and by guaranteeing coverage of all relevant potential-energy-surface regions.

Problem
Most ML-FFs use nuclear charges and coordinates alone and rely on locality, limiting their treatment of differing electronic states and nonlocal quantum effects.
Method
SpookyNet is an MPNN that directly models charge and spin states, uses attention for nonlocal interactions, and incorporates physically motivated inductive biases and analytical energy corrections.
Results
SpookyNet predicts electronic-state-dependent potential-energy surfaces, models nonlocal dopant effects, generalizes beyond training chemical and conformational space, and improves existing models across quantum-chemical benchmarks.
Takeaways & Limitations
The learned interaction functions resemble atomic orbitals, and the model can extract chemical insight while extending ML-FF applicability beyond purely local, fixed-electronic-state systems.
Takeaways & Limitations
Applying ML-FFs to large heterogeneous condensed-phase systems remains challenging, and training on small molecules does not guarantee accurate coverage of every dynamically visited potential-energy-surface region.

Abstract

from arXiv · show

Machine-learned force fields (ML-FFs) combine the accuracy of ab initio methods with the efficiency of conventional force fields. However, current ML-FFs typically ignore electronic degrees of freedom, such as the total charge or spin state, and assume chemical locality, which is problematic when molecules have inconsistent electronic states, or when nonlocal effects play a significant role. This work introduces SpookyNet, a deep neural network for constructing ML-FFs with explicit treatment of electronic degrees of freedom and quantum nonlocality. Chemically meaningful inductive biases and analytical corrections built into the network architecture allow it to properly model physical limits. SpookyNet improves upon the current state-of-the-art (or achieves similar performance) on popular quantum chemistry data sets. Notably, it is able to generalize across chemical and conformational space and can leverage the learned chemical insights, e.g. by predicting unknown spin states, thus helping to close a further important remaining gap for today's machine learning models in quantum chemistry.

INTRODUCTION

Existing ML force fields are efficient but often omit electronic state and nonlocality, creating ill-defined or physically incomplete learning problems. SpookyNet addresses these gaps with explicit electronic inputs, nonlocal interactions, and physically motivated corrections.

Motivation: Accurate ab initio forces are computationally impractical for large systems or long molecular-dynamics simulations, motivating machine-learned force fields.ML-FFs learn forces from ab initio reference data while retaining efficiency closer to conventional force fields.
Chemical insight: Learned interaction functions closely resemble atomic orbitals, indicating that SpookyNet extracts chemically meaningful structure from data.This behavior is reported for a model trained on QM7-X1.
Electronic degrees of freedom: Atomic numbers and coordinates omit total charge and spin, so systems with inconsistent electronic states become an ill-defined learning problem.This is harmless only when all systems share a consistent electronic state, such as neutral singlets.
Nonlocal effects: Locality assumptions neglect quantum nonlocality, including charge or spin redistribution caused by distant structural changes.Delocalized electrons can make distant atoms influence local properties and potential-energy surfaces.
SpookyNet approach: SpookyNet directly inputs atomic numbers, coordinates, electron number, and spin state, while using implicit angular encoding to retain linear scaling.Its basis functions use Bernstein polynomials and spherical harmonics rather than explicitly computing all neighbor angles.
SpookyNet approach: Distance-independent interactions model nonlocality, while analytical electrostatic, dispersion, and short-range repulsion corrections enforce physically motivated behavior.These corrections reduce the learning burden and improve asymptotic behavior.

RESULTS

SpookyNet combines nuclear and electronic embeddings with local and nonlocal interaction modules, then augments learned atomic energies with physical corrections. It reproduces electronic-state-dependent potential-energy surfaces, while removing charge or spin embeddings causes qualitatively incorrect surfaces and minima.

Architecture: SpookyNet represents nuclear charges, total charge, and spin state as embeddings that initialize atomic features.Total charge encodes electron number, while spin state is represented by the number of unpaired electrons.
Architecture: A chain of interaction modules refines atomic representations through local interactions, nonlocal interactions, and learned transformations.Local functions capture distance and angular information, whereas nonlocal interactions model delocalized electrons.
Energy and forces: The predicted energy combines learned atomic contributions with analytical corrections for nuclear repulsion, electrostatics, and dispersion.Forces are obtained by automatic differentiation of the energy-conserving potential.
Electronic states: SpookyNet faithfully reproduces reference potential-energy surfaces for charge- and spin-distinguished systems, including Ag3+/Ag3− and singlet/triplet CH2.Removing charge/spin embeddings produces qualitatively different surfaces and wrong global minima, whereas the complete model gives minima virtually indistinguishable from reference.
Electronic states: Explicit charge and spin embeddings allow models to distinguish electronic states that can be structurally identical, while implicit electrostatics alone are insufficient for adequate charge-state separation.The charge embedding is especially necessary for systems such as singlet/triplet CH2, which otherwise appear identical to the model.

Nonlocal effects

SpookyNet addresses nonlocal effects that make purely local force fields unreliable, especially for separated atoms and doped materials. It also generalizes across chemical and conformational space, outperforming or matching established models on several benchmarks.

Nonlocal effects: Nonlocal charge redistribution makes separated atoms’ energy contributions depend on distant atoms, which purely local models cannot represent simultaneously across systems.This limitation arises even for simple diatomic molecules because electrons can distribute unevenly according to electronegativity.
Nonlocal effects: Removing nonlocal interactions produces an unphysical energy step at large separations, whereas enabling them faithfully reproduces the reference dissociation curves.The artifact persists despite increasing the local cutoff radius and can be problematic in reaction simulations.
Nonlocal effects: SpookyNet significantly improves Au2–MgO prediction errors over models without nonlocal treatment and achieves lower errors than 4G-BPNNs using charge equilibration.The Au2–MgO system exhibits nonlocal changes induced by dopant atoms in the MgO surface.
Generalization: On QM7-X, SpookyNet achieves lower prediction errors than SchNet and PaiNN for both known-molecule and unknown-molecule generalization tasks.It is only marginally worse when predicting completely unknown molecules, suggesting successful generalization across chemical space.
Generalization: A model trained on molecules with at most seven non-hydrogen atoms also produces correct structures for larger molecules and fullerenes absent from the training data.Reported examples include vitamin B2, cholesterol, deca-alanine, and pure-carbon fullerenes.
Benchmarking: On MD17, SpookyNet reaches lower prediction errors or closely matches other published models for every tested molecule.MD17 contains structures, energies, and forces from ab initio molecular-dynamics simulations at the PBE+TS level of theory.

DISCUSSION

SpookyNet combines attention-based nonlocal interactions, electronic-state inputs, and physically motivated corrections to extend ML force fields beyond purely local, fixed-state settings. The model shows chemically intuitive learned representations, broad generalization, and improved applicability, while large heterogeneous condensed-phase systems remain challenging.

DISCUSSION: SpookyNet is an MPNN that models electronic degrees of freedom and nonlocal interactions using attention.Its architecture is designed to address limitations of models based only on nuclear charges and atomic coordinates.
DISCUSSION: Physically motivated inductive biases and analytical corrections encode electronic configurations, nuclear repulsion, electrostatics, and dispersion.The short-range correction uses the Ziegler-Biersack-Littmark stopping potential, while long-range terms use point-charge electrostatics and empirical dispersion.
DISCUSSION: SpookyNet predicts distinct potential energy surfaces for different electronic states and models nonlocal material-property changes caused by dopants.These capabilities extend ML-FF applicability beyond systems with consistent electronic states and negligible nonlocal effects.
DISCUSSION: The learned interaction functions resemble atomic orbitals, indicating chemically intuitive representations of molecular systems.The paper presents this resemblance as evidence that SpookyNet extracts chemical insight from data.
DISCUSSION: Low test-set error alone cannot establish that an ML model has learned chemically meaningful behavior rather than exploiting data artifacts or Clever Hans effects.Understanding how SpookyNet solves the prediction problem is therefore important for scientific use.
DISCUSSION: Applying ML-FFs to large heterogeneous condensed-phase systems, such as proteins in aqueous solution, remains a challenge because suitable ab initio reference data and complete potential-energy-surface coverage are difficult to guarantee.The authors conjecture that physically motivated inductive biases may help address these problems.

Details on the neural network architecture

SpookyNet combines residual neural-network components with chemically informed atomic embeddings, local equivariant interactions, and explicit nonlocal attention. Its architecture incorporates electronic charge and spin information while preserving smooth, physically meaningful energy predictions.

Core network components: SpookyNet uses residual multilayer perceptrons and smooth generalized SiLU activations to learn nonlinear feature transformations without force discontinuities.The smooth activation is necessary because kinks in predicted potential energies would introduce discontinuities in atomic forces.
Atomic embeddings: Nuclear embeddings combine element-specific parameters with ground-state electronic descriptors, encouraging representations that capture similarities between different elements.The descriptor projection provides an inductive bias for learning alchemical knowledge.
Electronic embeddings: Electronic embeddings map total charge Q and unpaired-electron count S to vectors that distribute electronic information across all atoms.Positive and negative charge inputs use separate parameters, while spin inputs use the nonnegative branch.
Interaction modules: Each interaction module refines atomic features through local and nonlocal blocks, passing x^(t) to the next module and accumulating y^(t) into final descriptors.The local and nonlocal interaction outputs are combined within a chain of interaction modules.
Local interactions: Local interactions use s-, p-, and d-orbital-like basis functions to encode radial and angular information, then form rotationally invariant features from equivariant projections.Bernstein-polynomial radial bases are applied after an exponential distance mapping, while the cutoff smoothly decays them to zero.
Nonlocal interactions: Self-attention models electron delocalization through distance-independent interactions, although standard attention has O(N^2F) time complexity and quadratic atom-count scaling.SpookyNet uses attention to represent nonlocal interactions starting from temporary atomic features.

Training and hyperparameters

SpookyNet is trained with a weighted loss over energies, forces, and optionally dipole moments or partial charges. The implementation uses fixed architectural defaults, physically constrained correction parameters, adaptive optimization, and dataset-specific force weighting.

Default settings: SpookyNet models use T = 6 interaction modules, F = 128 features, and rcut = 10 a0 (approximately 5.29177 Å) unless otherwise specified.These are the stated default settings for the models in this work.
Physical constraints: Nuclear-repulsion coefficients are constrained to positive values and normalized to guarantee correct short-distance asymptotic behavior.Dispersion parameters begin from Hartree–Fock recommendations, and the charge-scaling parameter is initialized at one and kept positive.
Optimization: Parameters are optimized with AMSGrad from a 10^-3 learning rate, while an exponential moving average with smoothing factor 0.999 is evaluated periodically.The learning rate is halved whenever validation performance triggers the stated decay condition.
Objective: The loss combines energy, force, and dipole-moment terms, with relative weights αE = αF = αµ/q = 1 except αF = 100 for MD17 and QM7-X.Increasing the force weight significantly reduces both energy and force prediction errors according to the cited supplementary comparison.
Charge-related supervision: Dipole moments provide reference data for learning partial charges, while energy corrections are omitted when neither dipoles nor reference partial charges are available.Dipole moments are quantum-mechanical observables, whereas alternative charge decompositions may be used when necessary.

Computing and visualizing local chemical potentials

Local chemical potentials are computed by introducing a one-way probe atom whose predicted energy contribution measures its interaction with a molecule. Separate evaluations isolate orbital-like and nonlocal contributions to the potential-energy surface.

Local chemical potentials: A probe atom is placed at position r and allowed to interact with the molecule without perturbing the molecule’s prediction.The probe atom’s predicted energy contribution is interpreted as the local chemical potential.
Probe construction: Electronic, local, and nonlocal features for the probe are computed as if it belonged to an N + 1 atom molecule, while the original molecule remains unperturbed.This procedure preserves the one-way probe construction across the different feature contributions.
Orbital decomposition: Orbital-like contributions are isolated by zeroing selected p- and d-feature terms and subtracting the resulting predictions.The s contribution sets both p and d to zero; p and d contributions are obtained through corresponding subtraction procedures.
Nonlocal decomposition: The nonlocal contribution is obtained by subtracting the prediction made with the nonlocal feature n set to zero from the full model prediction.This isolates the component attributed to the nonlocal interaction block.

SchNet and PaiNN training

The QM7-X SchNet and PaiNN baselines use matched feature dimensions but different interaction depths and share Gaussian radial bases, optimizer settings, learning rate, and batch size.

Model settings: SchNet uses F = 128 features and T = 6 interactions, whereas PaiNN uses F = 128 features and T = 3 interactions.Both models use 20 Gaussian radial basis functions with a 5 Å cutoff.
Training settings: Both baselines are trained with Adam at a learning rate of 10^-4 and batch size 10.These settings are specified for the QM7-X experiments.

Data generation

Three new data sets were generated to test SpookyNet on different electronic states and nonlocal interactions. The data include electronic-state recomputations and extensive diatomic bond scans.

Three new data sets test SpookyNet's ability to model different electronic states and nonlocal interactions.
Energies, forces, and dipoles were computed at the semi-empirical GFN2-xTB level of theory.
2,200 structures were generated for each electronic-state data set by sampling 550 structures around both state minima and recomputing them in the other state.
9,000 diatomic structures were generated from nine bond scans spanning 1.5–20 a0 with 1,000 evenly spaced points per molecule.
The diatomic models used all data and an increased cutoff of rcut = 18 a0 to test whether distance-only locality could fit the scans.

DATA AND CODE AVAILABILITY

The authors plan to release the newly generated data sets and a PyTorch reference implementation of SpookyNet upon manuscript acceptance.

The newly generated data sets and a reference PyTorch implementation of SpookyNet will be made available when the manuscript is accepted.
Other data sets used in the work are publicly available through the cited sources, including sGDML, QM7-X, QMspin, and additional references.

Supporting Information for SpookyNet: Learning Force Fields with

The supporting information accompanies the SpookyNet paper and identifies its title, authors, and arXiv version.

The supporting information is titled “Electronic Degrees of Freedom and Nonlocal Effects.”
The document lists Oliver T. Unke, Stefan Chmiela, Michael Gastegger, Kristof T. Schütt, Huziel E. Sauceda, and Klaus-Robert Müller as authors.
The associated preprint is arXiv:2105.00304v2, dated 20 July 2021.

S1. COMPLETENESS OF ATOMIC DESCRIPTORS IN SPOOKYNET

This section examines descriptor completeness and shows how SpookyNet's learned representations distinguish challenging atomic environments and achieve data-efficient energy prediction. It also documents limitations of finite angular descriptors and the role of multiple message-passing steps.

Descriptor completeness: Descriptor completeness requires non-equivalent structures to map to different descriptors under translations, rotations, and permutations of equivalent atoms.
Descriptor completeness: Distance sets alone and finite angular power-spectrum invariants can identify distinct environments as equivalent, while dihedral information may still be insufficient.
Message passing: With T ≥2 message-passing steps, MPNNs distinguish all environments shown in Fig. S1, whereas T = 1 leaves some environments indistinguishable.
Learning efficiency: Chemical accuracy is reached with 1,000 training points, and SpookyNet is about two orders of magnitude more data-efficient than other methods.
Architecture: The generalized SiLU activation reduces to ReLU for α = 1, β = ∞ and to the identity function for α = 2, β = 0.