SchNetPack: A Deep Learning Toolbox For Atomistic Systems

K. T. Schütt; P. Kessel; M. Gastegger; K. Nicoli; A. Tkatchenko; K. -R. Müller

Source-linked AI summary

SchNetPack: A Deep Learning Toolbox For Atomistic Systems

K. T. Schütt, P. Kessel, M. Gastegger, K. Nicoli, A. Tkatchenko, K. -R. Müller

arXiv:1809.01072v1Sep 4, 2018physics.comp-phphysics.chem-ph

TL;DR

Accurate atomistic-property calculations are computationally expensive, motivating tools that make neural-network modeling and evaluation more accessible. SchNetPack unifies representations, training, datasets, and simulation interfaces, and its SchNet architecture consistently outperforms Behler–Parrinello models in the reported results.

Problem
Accurate quantum-chemical calculations are prohibitively costly for large-scale molecular dynamics and exhaustive exploration of chemical space.
Method
SchNetPack provides a PyTorch-based framework combining ACSF, wACSF, and SchNet models with benchmark datasets, multi-GPU training, and an ASE interface.
Results
SchNet consistently outperforms Behler–Parrinello models across the reported experiments, including a RMSE of 0.77 kcal mol−1 on ANI-1 using expanded training splits.
Takeaways & Limitations
The unified toolbox simplifies accessing datasets, training different atomistic architectures, evaluating performance, and combining trained models with ASE functionality.

Abstract

from arXiv · show

SchNetPack is a toolbox for the development and application of deep neural networks to the prediction of potential energy surfaces and other quantum-chemical properties of molecules and materials. It contains basic building blocks of atomistic neural networks, manages their training and provides simple access to common benchmark datasets. This allows for an easy implementation and evaluation of new models. For now, SchNetPack includes implementations of (weighted) atomcentered symmetry functions and the deep tensor neural network SchNet as well as ready-to-use scripts that allow to train these models on molecule and material datasets. Based upon the PyTorch deep learning framework, SchNetPack allows to efficiently apply the neural networks to large datasets with millions of reference calculations as well as parallelize the model across multiple GPUs. Finally, SchNetPack provides an interface to the Atomic Simulation Environment in order to make trained models easily accessible to researchers that are not yet familiar with neural networks.

I. INTRODUCTION

SchNetPack addresses the high cost of accurate atomistic-property calculations by providing a unified framework for descriptor-based and end-to-end neural networks, with training, datasets, and GPU support.

I. INTRODUCTION: Accurate quantum-chemical calculations are too costly for large-scale molecular dynamics and exhaustive chemical-space exploration.Machine learning can reproduce reference calculations at significantly reduced prediction cost after training.
I. INTRODUCTION: Atomistic neural networks broadly comprise descriptor-based models using predefined representations and end-to-end models learning representations from atom types and positions.
I. INTRODUCTION: SchNetPack unifies both categories and currently implements SchNet alongside ACSF- and wACSF-based Behler–Parrinello networks.
I. INTRODUCTION: The toolbox provides benchmark-dataset access, multi-GPU training, extensibility, and PyTorch-based implementation for predicting diverse chemical properties.
I. INTRODUCTION: Models separate atomistic representations from prediction blocks, with representations describing each atom’s chemical environment and prediction blocks estimating system properties.

A. Representations

SchNetPack implements ACSF and wACSF representations that encode local atomic environments through radial and angular symmetry functions, with wACSFs reducing element-dependent descriptor growth.

A. Representations: ACSFs represent each atom’s local chemical environment using radial and angular distribution functions.
A. Representations: Radial ACSFs sum neighbor contributions weighted by element and use Gaussian parameters controlling centers and widths within a spatial cutoff.SchNetPack determines these parameters automatically on an equidistant grid between zero and the cutoff.
A. Representations: Angular ACSFs encode angles between neighboring atoms, with λ shifting the angular maximum and ζ controlling its width.Parameterized angular functions describe local environments across multiple settings.
A. Representations: Standard ACSFs require functions for element pairs or triples, causing the number of descriptors to grow quadratically with the number of chemical species.This can become impractical for systems containing more than four elements, such as QM9.
A. Representations: wACSFs use g(Zj) = Zj radially and g(Zk, Zj) = ZkZj angularly, making the required number of symmetry functions independent of the number of elements.SchNetPack uses wACSFs as the standard descriptor for Behler–Parrinello potentials.

2. SchNet

SchNet is an end-to-end architecture that learns atomistic representations through embeddings and interaction blocks, using continuous-filter convolutions based on interatomic distances.

2. SchNet: SchNet constructs atom-wise representations from atom-type embeddings and a series of interaction blocks that incorporate system configuration.
2. SchNet: Unlike ACSF-based networks, SchNet adapts its representation to training data rather than relying on rigid handcrafted features.
2. SchNet: Interaction blocks use continuous-filter convolutions to incorporate neighboring atoms and additively refine the previous representation.
2. SchNet: A filter-generating neural network maps atomic distances to filter values instead of using a fixed parameter tensor.The filter generator receives positions expanded on a grid of radial basis functions related to radial symmetry functions.
2. SchNet: SchNet processes the embedded features through L interaction blocks to produce a latent representation passed to the prediction block.Atom-wise layers reuse weights and biases across atoms, so their parameter count is independent of the number of atoms.

B. Prediction Blocks

SchNetPack prediction blocks transform atom-wise representations into property estimates using aggregation choices suited to the target property. It also provides specialized blocks that encode property-specific structure or element-dependent processing.

Prediction blocks process atom-wise representations through nonlinear layers and aggregate the resulting contributions across atoms.The aggregation is chosen according to the property of interest.
Atomwise blocks express molecular properties as sums of atom-wise contributions.This form is suitable for extensive properties such as energy.
Intensive properties are expressed as averages over atom-wise contributions rather than sums.These properties do not grow with the number of atoms.
DipoleMoment uses latent atomic charges and positions relative to the system center of mass to represent the dipole moment.The charges are modeled by q and the reference point is the center of mass.
ElementalAtomwise applies separate neural networks to different chemical elements, which is particularly useful for (w)ACSF representations.ElementalDipoleMoment is defined analogously for dipole moments.

III. DATA PIPELINE AND TRAINING

SchNetPack provides dataset access and training-management components for atomistic neural networks. Its pipeline supports benchmark data handling, asynchronous loading, monitoring, checkpointing, and multi-GPU training, with compact code for model evaluation.

SchNetPack classes access standard benchmark datasets and manage atomistic neural-network training.Figure 2 summarizes this training setup.
The dataset interface includes QM9, ANI-1, ISO17, MD17, and Materials Project data sources.These cover organic molecules, molecular dynamics, forces, and bulk crystals across broad chemical composition ranges.
AtomsLoader feeds subsets of datasets using multiple threads and computes statistics such as mean and standard deviation.This supports asynchronous preparation of training inputs.
Trainer evaluates validation performance and provides early stopping, while training supports learning-rate schedules, checkpointing, logging, and multiple GPUs.Logging can use CSV files or TensorBoard, and multi-GPU training uses PyTorch’s standard implementation.
The training classes allow atomistic neural networks to be trained and evaluated with a compact amount of code.The paper illustrates this capability in its training example.

IV. IMPLEMENTATION DETAILS

SchNetPack is implemented in Python with PyTorch and NumPy and is integrated with ASE for storing atomistic configurations and using trained models in simulation workflows.

SchNetPack uses Python and PyTorch for neural-network functionality, with NumPy for calculations that do not require automatic differentiation.The implementation specifies PyTorch version 0.4 or later.
ASE stores atomistic configurations and provides a calculator interface for incorporating SchNetPack models into workflows such as molecular dynamics.This integration exposes trained models through ASE workflows.

V. EXAMPLE: TRAINING IN SCHNETPACK

The paper demonstrates SchNetPack through a compact QM9 training script and shows that datasets, representations, and output networks can be changed with small code modifications.

Listing 1 imports SchNetPack modules, dataset utilities, Adam, and PyTorch functional operations for model training.The script loads QM9 and prepares training and validation data.
The example trains a SchNet model to predict total energy U0 on QM9 using 10k training points and 1k validation points.Data loading is asynchronous and uses four worker threads.
Changing the dataset requires only changing the dataset specification in the example script.The paper illustrates this with ANI-1.
Changing the representation from SchNet to wACSF requires replacing the representation line.For wACSF, the example also recommends an ElementalAtomwise output network.
The examples are available in SchNetPack’s source directory under the examples subdirectory.

VI. EXAMPLE: SCHNETPACK FOR CHEMISTS

SchNetPack connects trained atomistic models to ASE calculation tools, enabling optimization, normal-mode analysis, molecular dynamics, and power-spectrum prediction. In the malondialdehyde example, the models reproduce detailed spectral features while accelerating simulation substantially.

ASE interface: SchNetPack exposes trained models through an ASE Calculator interface for geometry optimization, normal-mode analysis, and molecular dynamics simulations.The interface is provided through the AseDriver class in the molecular_dynamics module.
ASE interface: Ready-to-use scripts load trained models into the calculator and support molecular dynamics simulations for atomistic systems.The molecular-dynamics script performs simulations out of the box.
Power-spectrum example: The malondialdehyde example computes power spectra at 300 K using SchNets trained on 1000 and 50000 MD17 data points.Harmonic normal-mode vibrations from the electronic-structure reference are shown for comparison.
Power-spectrum example: The models reproduce peak positions accurately even with the smaller training set and resolve the 1700 cm−1 carbonyl-stretching peak structure.The resolved structure corresponds to symmetric and asymmetric stretching vibrations of two carbonyl groups.
Performance: 11 milliseconds per timestep on a Tesla P100 GPU yields almost three orders of magnitude speedup over the electronic-structure reference for malondialdehyde simulations.Behler–Parrinello networks have comparable performance for molecules of this size in the reported GPU setup.

VII. RESULTS

SchNetPack evaluates Behler–Parrinello and SchNet models across QM9, ANI1, MD17, and Materials Project datasets. SchNet consistently outperforms the Behler–Parrinello models, while ACSF and wACSF trade performance across dataset characteristics and Behler–Parrinello models retain a computational-cost advantage.

Evaluation setup: Results average three models trained on different splits and cover QM9, ANI1, MD17, and Materials Project datasets.The reported test-set evaluations were obtained using SchNetPack, with scripts provided for experiments on a Tesla P100 GPU.
SchNet versus Behler–Parrinello: SchNet consistently outperforms Behler–Parrinello networks across the reported experiments.The paper attributes this to data-driven representations, deeper architecture, and advantages for force learning.
Computational trade-offs: Behler–Parrinello models have lower computational cost than SchNet, which is expected to benefit molecular dynamics simulations of large molecules.The reported ACSF and wACSF performance could improve through descriptor fine-tuning, but such tuning is typically tedious.
ACSF versus wACSF: wACSF performs better than ACSF on structurally and chemically diverse QM9 data, whereas ACSF performs better on MD17 tasks resolving small structural variations.The paper links this trade-off to wACSF’s reduced spatial resolution from improved elemental resolution.
SchNet versus Behler–Parrinello: SchNet achieves chemically accurate performance on ANI1, including RMSE 0.89 kcal mol−1 with 10 million reference examples.With 80% of the dataset for training and 10% each for validation and testing, it obtains MAE 0.47 kcal mol−1 and RMSE 0.77 kcal mol−1.

VIII. CONCLUSIONS

SchNetPack unifies dataset access, model training, architecture evaluation, and ASE-based simulations for atomistic neural networks. The authors expect this simplification to help researchers focus on model design and compare architectures more easily.

Conclusions: SchNetPack simplifies access to benchmark datasets, training of different neural-network architectures, and evaluation of their performance.It also provides an ASE interface for molecular-dynamics simulations and related calculation tools.
Conclusions: The framework’s unification is expected to let researchers concentrate on neural-network design and compare different architectures more easily.Planned extensions include more datasets, active sampling, additional quantum-mechanical observables, and further architectures.

Appendix A: Details on Experiments

The experiments report training setups and timing choices for SchNet and Behler–Parrinello models, including learning-rate scheduling, GPU use, and stopping criteria. These choices affect convergence speed and the amount of time spent on final accuracy improvements.

Experimental setup: The experiments average errors and runtimes over three models trained on different data splits using a Tesla P100 GPU.The reported setup includes the experiment scripts and separates model-training configurations across the appendix tables.
Training schedule: A validation-triggered learning-rate decay replaces the fixed 0.96 reduction every 100k iterations used in the original SchNet publications.The authors report comparable long-run results but significantly faster convergence with the new schedule.
Hardware and runtime: SchNetPack supports multi-GPU training, demonstrated for the ANI-1 dataset, to further reduce training time.The cited figure concerns epoch time on 110k QM9 molecules, while the multi-GPU experiment itself was conducted only on ANI-1.
Stopping criteria: Approximately 50% of training time for acetylsalicylic acid is spent fine-tuning the final 0.02 kcal/mol improvement.The example uses 50k reference calculations and motivates stopping once learning progress flattens if that minor improvement is unnecessary.
Configuration details: The appendix organizes separate setup tables for SchNet training, Behler–Parrinello training, and ACSF or wACSF compositions.The wACSF and ACSF symmetry functions were standardized in all experiments.