Unifying machine learning and quantum chemistry -- a deep neural network for molecular wavefunctions

K. T. Schütt; M. Gastegger; A. Tkatchenko; K. -R. Müller; R. J. Maurer

Source-linked AI summary

Unifying machine learning and quantum chemistry -- a deep neural network for molecular wavefunctions

K. T. Schütt, M. Gastegger, A. Tkatchenko, K. -R. Müller, R. J. Maurer

arXiv:1906.10033v1Jun 24, 2019physics.chem-phstat.ML

TL;DR

Existing machine-learning models predict selected molecular properties without explicitly capturing electronic degrees of freedom. This paper develops SchNOrb, which predicts wavefunctions through a local atomic-orbital Hamiltonian representation, enabling derived electronic properties and applications including inverse design and quantum-chemistry acceleration.

Problem
Existing ML models predict selected molecular properties but do not explicitly capture the electronic degrees of freedom underlying molecular quantum chemistry.
Method
SchNOrb predicts molecular electronic structure through the Hamiltonian and wavefunction in a local atomic-orbital representation using symmetry-aware atom and atom-pair features.
Results
The model predicts total energies below 2 meV error, Hamiltonians below 8 meV error, overlap matrices below 1 · 10^-4, and most dipoles below 0.054 D and quadrupoles below 0.058 D Å.
Takeaways & Limitations
The analytically differentiable electronic representation supports electronic-property optimisation, molecular dynamics, derived-property prediction, and enhanced quantum-chemistry calculations.

Abstract

from arXiv · show

Machine learning advances chemistry and materials science by enabling large-scale exploration of chemical space based on quantum chemical calculations. While these models supply fast and accurate predictions of atomistic chemical properties, they do not explicitly capture the electronic degrees of freedom of a molecule, which limits their applicability for reactive chemistry and chemical analysis. Here we present a deep learning framework for the prediction of the quantum mechanical wavefunction in a local basis of atomic orbitals from which all other ground-state properties can be derived. This approach retains full access to the electronic structure via the wavefunction at force field-like efficiency and captures quantum mechanics in an analytically differentiable representation. On several examples, we demonstrate that this opens promising avenues to perform inverse design of molecular structures for target electronic property optimisation and a clear path towards increased synergy of machine learning and quantum chemistry.

I. Introduction

Existing machine-learning models efficiently predict selected molecular properties but do not explicitly represent molecular electronic degrees of freedom. The paper introduces direct wavefunction learning to recover ground-state properties and connect machine learning with quantum chemistry.

Most existing models learn scalar, vector, or tensor molecular properties separately from quantum-chemical data.
Because molecular properties derive from the ground-state wavefunction, directly predicting it could provide all ground-state properties without separate property-specific models.
The framework predicts molecular electronic structure through the Hamiltonian in a local atomic-orbital basis.
The model targets eigenvalue spectra and molecular orbitals near chemical accuracy, approximately 0.04 eV, for organic molecules.
Predicted electronic structure supports charge populations, bond orders, dipole and quadrupole moments, molecular dynamics, inverse design, and quantum-chemistry calculations.
Direct prediction of eigenvalues and wavefunction coefficients is challenged by degeneracies and electronic level crossings, motivating Hamiltonian prediction in a local representation.

B. SchNOrb deep learning framework

SchNOrb extends SchNet to predict molecular electronic structure in a symmetry-aware local atomic-orbital representation. It separately models energy, Hamiltonian, and overlap information using atom and atom-pair environments.

SchNOrb extends SchNet by constructing rotationally invariant atomic-environment representations and symmetry-adapted pairwise features for electronic structure.
Pairwise features combine rotationally invariant and covariant components so atomic-orbital interactions can represent angular momentum and molecular orientation.
The network begins with atom-type embeddings and constructs pair representations from atomic features and interatomic distances.
The architecture uses 2L + 1 SchNOrb interaction blocks to compute pairwise coefficients for angular momenta up to 2L.
Onsite and offsite Hamiltonian blocks are modeled separately and symmetrised, while the overlap matrix is predicted analogously.
The architecture predicts total energy, Hamiltonian, and overlap matrices end-to-end using a combined regression loss.

C. Learning electronic structure and derived properties

SchNOrb predicts molecular electronic structure and derived properties from learned Hamiltonian and overlap matrices, achieving high accuracy across several molecular data sets.

Model and training: SchNOrb predicts total energies, Hamiltonian matrices, and overlap matrices jointly from molecular electronic-structure data.The model is trained on water, ethanol, malondialdehyde, and uracil using Hartree–Fock and PBE-DFT references.
Prediction accuracy: Below 2 meV mean absolute error is achieved for total energies, while Hamiltonian and overlap-matrix errors remain below 8 meV and 1 · 10−4.
Prediction accuracy: Occupied orbital energies are predicted more accurately than virtual orbitals, with errors below 20 meV and approximately 100 meV, respectively.For water and malondialdehyde, average occupied-orbital errors are below 10 meV; uracil reaches 48 meV.
Prediction accuracy: Both occupied and unoccupied frontier orbital energies and shapes, including HOMO and LUMO, are reproduced with high accuracy.
Derived properties: Wavefunction operators yield electronic dipole errors below 0.054 D and quadrupole errors below 0.058 D Å for most molecules.Uracil is the exception because its delocalised π-system requires more reference data for comparable accuracy.

D. Chemical insights from electronic deep learning

SchNOrb exposes molecular orbitals and their dynamics for chemical interpretation, enables analytic optimization of electronic properties, and interfaces with downstream quantum-chemical calculations.

Chemical interpretation: Molecular orbitals provide direct access to electron distributions, bonding analysis, bond orders, and atomic partial charges relevant to molecular reactivity.
Chemical interpretation: Population analysis of SchNOrb orbitals identifies chemically interpretable charges, carbonyl double bonds, and electron delocalisation in uracil.The carbonyl bond orders are 2.12 and 2.14, respectively.
Reaction dynamics: During malondialdehyde proton transfer, electron density migrates with hydrogen while orbital rearrangement reveals alternating single- and double-bond character.
Reaction dynamics: Orbital forces change from attraction toward the product oxygen to mostly non-bonding character at the intermediate configuration.Forces are projected onto the reaction coordinate, with positive values driving the proton toward the product state.
Inverse design: Analytic derivatives with respect to atomic positions enable HOMO-LUMO-gap optimization from 3.15 eV to 2.68 eV or 3.59 eV.Gradient descent and ascent identify minimum- and maximum-gap structures from a random configuration.
Quantum-chemistry integration: Predicted wavefunctions can initialize further quantum-chemical calculations, reducing SCF iterations by approximately 73% for malondialdehyde without changing the final electronic-structure result.
Quantum-chemistry integration: For ethanol MP2 calculations initialized from predicted wavefunctions, the test-set error is 83 meV and the total error is 93 meV.The relative errors are 0.01% for HF energy and 0.06% for MP2 energy; MP2 correlation energy deviates by 17%.

III. Discussion

SchNOrb provides analytically differentiable electronic wavefunctions in a local atomic-orbital representation, extending machine-learning molecular simulation toward electronic-property optimisation and quantum-chemistry acceleration.

III. Discussion: Analytically differentiable wavefunctions provide atomic derivatives, including molecular-orbital energy and Hamiltonian derivatives for approximating nonadiabatic couplings.Higher-order derivatives describe electron–nuclear response.
III. Discussion: SchNOrb enables inverse molecular design by optimising electronic properties such as the HOMO-LUMO gap.Figure 4a illustrates minimisation and maximisation of the predicted gap for malonaldehyde configurations.
III. Discussion: The framework serves as a proof of principle for direct machine-learning models of electronic structure that enhance further quantum-chemistry calculations.Its predicted wavefunctions reduce the number of DFT-SCF iterations, with Figure 4b reporting an average reduction of 77%.
III. Discussion: Reference calculations used ORCA with the def2-SVP basis set and Pulay-based SCF by default.Integration grid levels 4 and 5 were used for SCF iterations and final property calculations, respectively.
III. Discussion: Malondialdehyde molecular dynamics used a 0.5 fs timestep at 300 K with a Langevin thermostat.Trajectories lasted 50 ps, with the first 10 ps discarded.

1. Basic building blocks and notation

The model begins with element-specific atomic embeddings and SchNet-style interaction layers that iteratively refine local atomic-environment representations before constructing pairwise features for Hamiltonian prediction.

1. Basic building blocks and notation: The architecture uses shifted softplus activations and linear or one-hidden-layer fully connected networks as its basic notation and components.Weights and biases are defined for the corresponding input, hidden, and output dimensions.
1. Basic building blocks and notation: Each atom receives an element-specific embedding based on its nuclear charge, with B = 1000 atom-wise features used for all models.The representations are subsequently refined by SchNet interaction layers.
1. Basic building blocks and notation: SchNet interaction layers use continuous-filter convolutions to incorporate spatial information from interatomic distances.The cutoff radius and radial-basis grid spacing control the distance representation.
1. Basic building blocks and notation: Feature crosstalk is performed atom-wise through fully connected layers, producing refined atom representations at each interaction iteration.Model parameters are shared within layers across atoms and interactions, but not across layers.
1. Basic building blocks and notation: SchNOrb adds a second interaction phase because Hamiltonian prediction requires representations of pair-wise environments rather than only atom-wise environments.These pairwise representations are built from the refined atomic-environment features.

3. Pair-wise environments and angular momenta

SchNOrb constructs symmetry-aware pairwise features from atom representations and interatomic distances, then uses angular-momentum channels to predict environment-dependent Hamiltonian blocks.

3. Pair-wise environments and angular momenta: The Hamiltonian is organised as matrix blocks H_ij, each depending on atom pair i,j, their chemical environment, and their associated atomic-orbital counts.The full matrix is assembled from blocks spanning all atom pairs.
3. Pair-wise environments and angular momenta: A factorised tensor layer forms raw atom-pair representations from atomic features and interatomic distance.Subsequent interactions refine pair environments through direct pair interactions and interactions with neighbouring atoms.
3. Pair-wise environments and angular momenta: Angular-momentum channels encode rotational covariance, with learnable projections along D = 4 directions.For s-orbital interactions, the special λ = 0 case enforces rotational invariance across directions.
3. Pair-wise environments and angular momenta: The resulting atom-pair features are used to predict coefficients for a basis set associated with angular-momentum channels.These representations provide the inputs for the corresponding Hamiltonian blocks.

4. Predicting the target properties

The target-property module predicts Hamiltonian and overlap matrices from pairwise features, while total energy is predicted separately from atom-wise contributions.

4. Predicting the target properties: Each atom-pair feature Ω^l is mapped to a corresponding Hamiltonian block using linear layers that conserve angular momentum.The output is masked to the actual orbital dimensions for atoms i and j.
4. Predicting the target properties: The predicted Hamiltonian is symmetrised after masking its atom-pair blocks.This produces the final matrix from the predicted off-diagonal and on-site contributions.
4. Predicting the target properties: The overlap matrix is predicted analogously using corresponding pairwise blocks.The same block-based construction is applied to overlap-matrix elements.
4. Predicting the target properties: Total energy is predicted separately as a sum of atom-wise energy contributions, following the conventional SchNet treatment.Energy prediction is therefore distinct from the Hamiltonian and overlap matrix predictions.

C. Data augmentation

The framework augments molecular configurations by randomly rotating them during training, using Wigner D matrices to transform positions, forces, Hamiltonians, and overlap matrices.

C. Data augmentation: Random molecular rotations are applied before each training epoch to reduce the number of reference calculations needed.The rotations are sampled randomly and used as data augmentation to save computing power.
C. Data augmentation: The augmentation learns rotational behavior rather than encoding the full rotational symmetry a priori.This requires training data that reflects enough molecular rotations.
C. Data augmentation: The transformations act on atom positions, atomic forces, the Hamiltonian matrix, and the overlap matrix.

D. Neural network training

Training uses a combined loss over energies, atomic forces, Hamiltonian matrices, and overlap matrices, with ADAM optimization and validation-based learning-rate decay.

D. Neural network training: The model is trained simultaneously on energies E, atomic forces F, Hamiltonian H, and overlap matrices S.A combined loss supports joint prediction of these quantities.
D. Neural network training: Stochastic gradient descent with the ADAM optimizer is used to train the neural networks.
D. Neural network training: The learning rate is multiplied by 0.8 after tpatience epochs without validation improvement and training stops at lr ≤5 · 10−6.The model with the lowest validation error is selected for testing.

Supplementary Material

The supplementary material documents training settings, prediction errors, orbital-level errors, symmetry behavior, derived-property errors, population analysis, and SCF acceleration.

Training settings: Table S1 reports training, validation, and test-set sizes, mini-batch sizes, initial learning rates, and learning-rate patience for each dataset.
Prediction errors: Table S2 reports mean absolute errors for Hamiltonians, overlaps, orbital energies, occupied-MO coefficients, and total energies.
Orbital-level errors: Figures S1 and S2 separate mean absolute orbital-energy errors and mean cosine distances of coefficient vectors by occupied molecular orbital.The cosine distance equals 1−cosine similarity.
Rotational behavior: Figure S3 shows that the learned orbitals are covariant under rotation of the water molecule, with each row corresponding to a molecular orbital.
Derived chemical properties: Figure S4 shows Löwdin partial charges and bond orders predicted by SchNOrb for a randomly selected malondialdehyde configuration.
SCF acceleration: Figure S5 shows that predicted molecular-orbital coefficients used as wavefunction guesses reduce the required SCF iterations by 73% for malondialdehyde.The guesses are used to obtain accurate DFT solutions.