Size extensibility / Periodic Systems

[simonbatzner]

Hi! I'm trying to use SchNet for modeling extended, periodic systems and have two questions regarding the use of SchNetPack in this case:

1. is SchNet size extensible - i.e. if I train on a system of N atoms and want to run Molecular Dynamics on a system of M >> N atoms, is that possible with the current version of the code? If not, is it possible in principle?

2. Is it possible to parallelize SchNetPack over multiple CPUs/GPUs for prediction / Molecular Dynamics (not training, where you can use data parallelism)? This is very helpful for building force-fields of large, extended systems (say 1,000+ atoms). Is that possible in the current version of the code? If not, is it possible in principle or would that not work because the different nodes of the graph need to talk to each other, i.e. it would be hard to parallelize over atomic environments?

Many thanks in advance!

Best, Simon

[mgastegger]

Hi Simon, glad you are interested in SchNetPack. To your questions:

1. SchNet is fully size extensive. You can easily train on smaller systems and then perform the prediction on larger ones. We use the typical atomistic network structure, where the total energy is a sum of atomic energies, which are functions of the atoms environment.
2. Due to the graph like structure, a single prediction procedure is hard to parallelize. It either involves large regions of the molecule or a lot of bookkeeping. The good news is, depending on the RAM of your GPU (~12GB), energy/force predictions should still be possible for systems up to 10k atoms. Constructing the neighborlist might be a bottleneck though.

Hope this helps.

[neojie]

Regarding the scalability or extensibility, I tried using SchNet to train a small ionic silicate system (4 formula unit MgSiO3 melt, MD simulations at P, T using VASP, system size 6.9A6.9A6.9A, cutoff is set as 5 A) and get quite small MAE (energy < 0.01 eV, force <0.05 eV/A). I applied the best fitting model to the a bigger system (16 formula unit MgSiO3 melt, , MD simulations at the same P, T using VASP) and get super big errors for both energy and force, especially for energy. One thing I notice that the energy predicted is always small than that of calculated by VASP, indicating that the small system model under-counting the total energy. I think there are several sources for the bad scalability. The major one may be that Schnet does not explicitly consider the long range Coulomb potential, which is a non-negligibly value especially for the large ionic system. I also tried to train the large system directly. But the I cannot get very good MAE, energy ~ 0.2 eV and 0.05 eV/A no matter how many interaction layer I use (3 or 6) or how large the cutoff I set (5 A or 8 A) or dataset size (2000 snapshots or 5000 snapshots). Do you have any idea on how to improve it? @mgastegger Thank you!

[mgastegger]

Sorry to hear that. Since you also tried fitting the larger system without success and there are ions, my guess would indeed be the electrostatics. There are ways to combine neural network potentials with explicit electrostatics (e.g. https://journals.aps.org/prb/abstract/10.1103/PhysRevB.83.153101 or https://pubs.acs.org/doi/10.1021/acs.jctc.9b00181). There are also plans to include them into SchNetPack, but this might take some time. One would e.g. have to figure out how to do the Ewald summation for periodic systems, etc.

What that might be feasible (quickly) though, is increasing the resolution of the radial functions together with increasing the cutoff. SchNet uses Gaussian basis functions for the radial filters, which are set to a default value. If the cutoff is increased without increasing the number of Gaussians, the resolution suffers. Depending on how you use SchNetPack, this can either be controlled via --num_gaussians from the command line (default should be 50) or the n_gaussians option (default 25) in the spk.representation.SchNet class. In theory, SchNet is able to learn interactions of all kind, especially if they have a high magnitude like the electrostatics in ionic systems. Since you already tried increasing the cutoff without improvement, this might indicate that the filter resolution is not high enough.

Maybe you could try increasing the Gaussians to ~100. Another idea we were playing around with, was to use Cosine basis functions for periodic system, which might be worth a shot and could be implemented quickly.

[simonbatzner]

Thanks a lot for your detailed answer @mgastegger , that was very helpful. That's what I expected regarding size extensiveness and parallelization, so that's good to hear.

The 10k atoms you reference on a single GPU, is that for a molecular structure? I observe that when I use extended, bulk solids with a cutoff of about 7A with SchNetPack, I very quickly have to go to small batch sizes of 5-10 structures on a regular V100 GPU. But bulk systems tend to have many more atoms in their environment then for example a large molecule, so I'm wondering if that 10k estimate is specifically for a molecular system?

[mgastegger]

Ah yes, 10k is for molecules and prediction. Training is more demanding, especially if forces and stress tensor are included.

Out of interest, how large is the unit cell for the bulk / how many atoms does it contain?

[jstoppelman]

@mgastegger I was reading this thread and you mentioned using cosine basis functions for periodic systems, which I'm trying to train using SchNetPack. Would this involve replacing the linear functions in the atom-wise layers with cosine functions?

[mgastegger]

If I understood correctly, what you mean is to modify the last layer of the atomistic output networks. We were thinking more along the lines of replacing the radial basis functions (Gaussians) used to expand the distances in the SchNet convolution by a periodic basis (e.g. Fourier series).