SchNet - a deep learning architecture for molecules and materials

Kristof T. Schütt; Huziel E. Sauceda; Pieter-Jan Kindermans; Alexandre Tkatchenko; Klaus-Robert Müller

Source-linked AI summary

SchNet - a deep learning architecture for molecules and materials

Kristof T. Schütt, Huziel E. Sauceda, Pieter-Jan Kindermans, Alexandre Tkatchenko, Klaus-Robert Müller

arXiv:1712.06113v3Dec 17, 2017physics.chem-phcond-mat.mtrl-sci

TL;DR

Accurate quantum-chemical calculations are too costly for broad exploration of molecular and materials chemical space. SchNet uses continuous-filter convolutions to model atomistic interactions and learns invariant representations, achieving accurate property, potential-energy-surface, and force-field predictions across molecules and materials.

Problem
Accurate quantum-chemical calculations are computationally prohibitive for exploring the vast chemical space of molecules and materials.
Method
SchNet is a deep-learning architecture that models atomistic interactions with continuous-filter convolutions and symmetry-preserving representations.
Results
SchNet accurately predicts diverse molecular and materials properties, potential-energy surfaces, and energy-conserving force fields.
Takeaways & Limitations
SchNet provides fast, accurate predictions while allowing analysis of learned representations through local chemical potentials.

Abstract

from arXiv · show

Deep learning has led to a paradigm shift in artificial intelligence, including web, text and image search, speech recognition, as well as bioinformatics, with growing impact in chemical physics. Machine learning in general and deep learning in particular is ideally suited for representing quantum-mechanical interactions, enabling to model nonlinear potential-energy surfaces or enhancing the exploration of chemical compound space. Here we present the deep learning architecture SchNet that is specifically designed to model atomistic systems by making use of continuous-filter convolutional layers. We demonstrate the capabilities of SchNet by accurately predicting a range of properties across chemical space for \emph{molecules and materials} where our model learns chemically plausible embeddings of atom types across the periodic table. Finally, we employ SchNet to predict potential-energy surfaces and energy-conserving force fields for molecular dynamics simulations of small molecules and perform an exemplary study of the quantum-mechanical properties of C$_{20}$-fullerene that would have been infeasible with regular ab initio molecular dynamics.

I. INTRODUCTION

SchNet addresses the computational cost of exploring chemical space by learning atomistic representations and interactions directly from atomic types and positions. It is designed to predict molecular and material properties, potential-energy surfaces, and force fields while respecting key atomistic symmetries.

Motivation: Accurate quantum-chemical calculations are too costly for exploring the vast chemical space of molecules and materials.Machine-learning approaches reduce the required reference calculations, but earlier methods often rely on handcrafted descriptors or lack atomic positions.
Contribution: SchNet models complex atomic interactions using continuous-filter convolutions and learned representations of atom types and positions.The architecture is a variant of DTNNs and uses filter-generating neural networks to incorporate chemical knowledge and constraints.
Contribution: Rotational, translational, and atom-indexing invariance are built into SchNet, supporting predictions across compositional and configurational chemical space.These symmetries are captured by design rather than learned only from examples.
Contribution: SchNet incorporates periodic boundary conditions through its filter-generating networks, enabling accurate formation-energy predictions for diverse bulk crystals.The filters can encode periodic structure while preserving the required property invariances.
Applications: SchNet enables efficient molecular-dynamics studies by predicting potential-energy surfaces and energy-conserving force fields.The introduction reports nanoseconds of low-temperature PIMD completed within hours instead of years for an exemplary C20-fullerene study.

B. Atom-wise layers

SchNet refines atom-wise representations through shared atom-wise layers and interaction blocks. Continuous-filter convolutions accommodate atoms at arbitrary positions, while shifted-softplus activations support smooth energy and force modeling.

Atom-wise layers: Shared atom-wise layers transform each atom representation independently while preserving scalability with respect to the number of atoms.The layers also process environmental information incorporated through interaction layers.
Interaction blocks: Interaction blocks add refinements based on pair-wise interactions with surrounding atoms using continuous-filter convolutional layers.These convolutions generalize discrete convolutions to atomistic systems whose positions are not arranged on a regular grid.
Interaction blocks: A continuous-filter convolution and three atom-wise layers form the residual mapping of an interaction block.The overall architecture and interaction-block organization are illustrated in Fig. 1.
Activation functions: Shifted-softplus activations satisfy ssp(0) = 0 and have infinite-order continuity, supporting smooth potential-energy surfaces, force fields, and second derivatives.The shift is described as improving convergence.

D. Filter-generating networks

SchNet generates continuous interaction filters from relative atomic positions, incorporating atomistic invariances and allowing resolution to be controlled through Gaussian distance expansions.

Filter generation: The filter-generating network maps the vector from an atom to a neighbor to filter values that determine how their interaction is modeled.This fully connected network can constrain the model and incorporate chemical knowledge.
Distance representation: Rotational invariance is included by computing pairwise distances rather than using relative positions.The distances are then expanded in a Gaussian basis.
Distance representation: Gaussian centers are placed uniformly between zero and the distance cutoff, while the number of Gaussians and γ determine filter resolution.The work uses a grid spacing and scaling parameter γ of 0.1 Å for all models.

2. Periodic boundary conditions

SchNet handles periodic systems by applying periodic boundary conditions directly to the interaction filters, preserving unit-cell invariance and enabling invariant property prediction. Energy-based force prediction additionally yields rotationally equivariant, energy-conserving force fields.

Periodic boundary conditions: Periodic repetitions receive equivalent atom-wise feature vectors, so periodic boundary conditions are applied directly to the filters.This accurately describes atom interactions while keeping predictions invariant to the choice of unit cell.
Periodic filters: The periodic filter sums contributions from all periodic images within the cutoff and can be normalized by the number of nearby atoms.Generated filters reflect the lattice structures of diamond and graphite.
Property prediction: SchNet predicts molecular or material properties from atom-wise contributions by summing extensive properties or averaging intensive properties.This construction retains invariance to atom indexing.
Force prediction: Atomic forces are obtained by differentiating the predicted energy with respect to atomic positions.With a rotationally invariant energy model, this produces rotationally equivariant and energy-conserving force predictions.

F. Training

SchNet is trained by minimizing property-specific squared losses, with a combined energy-and-force loss for molecular-dynamics trajectories. Training uses ADAM with validation-based early stopping, and nearby-atom preprocessing makes training cost linear in system size.

SchNet models each property target by minimizing a squared loss.
Energy and force trajectory models use a combined loss that balances energy and force errors through ρ.ρ controls the trade-off between the two loss components.
Models use mini-batch stochastic gradient descent with ADAM and batches of 32 examples.The learning rate decays exponentially by a ratio of 0.96 every 100,000 steps.
Data are split into training, validation, and test sets, with validation used for early stopping.The remaining data are used to compute test errors.
Precomputing nearby-atom indices makes each training step scale linearly with system size within the cutoff.

III. RESULTS

SchNet is evaluated on molecular and materials property prediction, with strong performance across suitable QM9 targets and improved learning behavior as interaction blocks increase. Its atomic contribution decomposition remains less suitable for some global properties, motivating property-specific output layers.

QM9 experiments train SchNet on 131k small organic molecules with up to nine heavy atoms, using T = 6 and 64-dimensional atomic representations.
SchNet reaches or improves over enn-s2s on 8 out of 12 QM9 properties where atomic decomposition is appropriate.Its predictions of polarizability α and electronic spatial extent ⟨R2⟩ are noticeably less accurate.
Interpretable property-specific output layers, such as for the dipole moment, remain future work.
0.014 eV is the U0 mean absolute error for SchNet with T ∈ {3, 6} interaction blocks and 110k training examples.With T = 2, the error is 0.015 eV; the figure compares learning curves against best-performing DTNN models with T = 3.
T = 6 produces the lowest U0 errors when training on fewer examples, despite having the most parameters.With 110k examples, it also reduces convergence from 2400 epochs for T = 2 to fewer than 750 epochs.

B. Learning formation energies of materials

SchNet predicts formation energies for a broad Materials Project crystal set while learning atom embeddings that reflect periodic-table chemistry. Its learned local chemical potentials also reproduce chemically meaningful molecular and solid-state patterns.

Formation energies: SchNet predicts formation energies for 69,640 Materials Project crystal structures spanning 89 atom types up to Z = 94.The model uses six interaction blocks, 64-dimensional atomic representations, and a 5 Å distance cutoff.
Learned atom embeddings: Figure 4 visualizes the two leading principal components of sp-atom embeddings, with colors denoting periodic-table groups.The projection explains only about 20% of embedding variance, so overlapping points may still be distant in the full space.
Learned atom embeddings: The learned atom embeddings group elements according to periodic-table groups and often order them from lighter to heavier elements.This structure was inferred from geometries and formation-energy targets rather than imposed on the model.
Local chemical potentials: Local chemical potentials are computed by adding a virtual probe atom that acts as a test charge to the SchNet input.The probe enables visualization of the learned representation in its molecular environment.
Local chemical potentials: For QM9 molecules, DTNN and SchNet local potentials capture bond saturation and varying aromaticity, while SchNet produces sharper feature separation.The distinguishing molecular features remain remarkably robust across the two neural-network architectures.
Local chemical potentials: The relationship between inferred local potentials and traditional observables such as electron density, electrostatic potentials, and orbitals remains unresolved.The paper identifies chemical-reactivity trends as a possible future application.
Local chemical potentials: SchNet local-potential cuts for graphite and diamond resemble the periodic structures of the corresponding solids.In solids, these potentials could help investigate defect formation and distribution, including vacancies and interstitials.

D. Combined learning of energies and atomic forces

SchNet is evaluated on molecular-dynamics trajectories using energy-only or combined energy-and-force training. Force information substantially improves small-data learning, while the model remains scalable to larger datasets and competitive with GDML on selected molecules.

Benchmark setup: SchNet predicts potential-energy surfaces and force fields on the MD17 benchmark of eight small organic molecules.Experiments use 1,000 or 50,000 reference calculations and compare against GDML and DTNN.
Combined training: On small training sets, adding force information reduces SchNet energy and force errors by 12 orders of magnitude versus energy-only training.With 50,000 examples, the improvement is less apparent because the potential-energy surface is already well sampled.
Comparison with GDML: With 1,000 reference calculations, SchNet outperforms GDML on malonaldehyde and ethanol, whereas GDML achieves lower force errors on trajectories containing aromatic rings.The authors suggest SchNet’s indexing invariance may benefit more flexible molecules.
Accuracy and scalability: SchNet achieves energy and force MAEs below 0.12 kcal/mol and 0.33 kcal/mol/Å, respectively.These results are reported for the tested molecular-dynamics trajectories and prediction setup.
Combined training: SchNet performs better with combined energy-and-force training on 1,000 calculations than with energy-only training on 50,000 examples.The comparison supports the value of force information when reference calculations are limited.

E. Application to molecular dynamics of C20-fullerene

SchNet was trained to reproduce the C20-fullerene potential-energy surface and used for molecular-dynamics simulations including nuclear quantum effects. The resulting model enabled long PIMD simulations with DFT-level accuracy at substantially lower computational cost.

Model and reference data: SchNet was trained on DFT reference data to predict energies and forces for C20-fullerene molecular dynamics.The reference calculations used PBE with Tkatchenko-Scheffler van der Waals interactions.
Model and reference data: The C20 model used 20,000 training references and 4,500 examples for early stopping, with test errors reported on the remaining data.Model selection varied interaction blocks, feature dimensions, and the energy-force trade-off.
Potential-energy surface: The selected model reproduced DFT normal-mode frequencies with a largest error of approximately 1%.The authors conclude that the model successfully reconstructs the C20 potential-energy surface and its symmetries.
Potential-energy surface: The force-only model was selected to drive molecular dynamics because the study focused on the fullerene's mechanical properties.Differentiating the energy model still guarantees an energy-conserving force field.
Nuclear quantum effects: At 300 K, nuclear quantum effects increased nearest-neighbor distances by less than 0.5%, while bond-length delocalization remained considerable.The peaks of the pair distribution function were reported to change little.
Nuclear quantum effects: 1.25 ns of PIMD reduced runtime relative to DFT by 3–4 orders of magnitude, from about 7 years to less than 7 hours.The simulations used SchNet@DFT and included nuclear quantum effects through PIMD.

IV. CONCLUSIONS

The paper presents SchNet as a versatile architecture for quantum chemistry that spans molecular and materials property prediction and accurate potential-energy modeling. Its C20 application demonstrates computationally feasible path-integral molecular dynamics that common DFT approaches could not practically provide.

Conclusions: SchNet learns task-adapted representations from first principles rather than requiring painstaking design of mechanistic force fields or descriptors.The architecture is applied from property prediction across chemical space to force fields for individual molecules.
Conclusions: SchNet models rotational, translational, and permutational invariances by design and can directly represent periodic boundary conditions.The paper also examines learned representations and atomic embeddings.
Conclusions: The C20 PIMD study at the PBE+vdWTS level would not have been computationally feasible with common DFT approaches.The authors identify larger molecules, periodic systems, and interpretable architectures as directions for future work.

Appendix A: Error distributions

The appendix presents logarithmic-scale histograms of prediction errors for QM9 and the Materials Project dataset. It also records the DFT-based C20 simulation setup and the PIMD treatment of nuclear quantum effects.

C20 simulation setup: C20 reference data were generated with classical NVT molecular dynamics at 500 K using PBE plus the Tkatchenko-Scheffler method for van der Waals interactions.The reported timestep was 1 fs.
C20 simulation setup: Nuclear quantum effects were introduced with path-integral molecular dynamics, where the number of beads controls their treatment.The simulations used a 0.5 fs timestep and an NVT path-integral Langevin thermostat.
Error distributions: QM9 error histograms include absolute prediction errors for all predicted properties and use a logarithmic scale to show distribution tails.The histograms aggregate test errors across all three repetitions.
Error distributions: The Materials Project error histogram shows absolute errors for formation energies per atom on a logarithmic scale.The logarithmic presentation is used to visualize the tails of the error distribution.