Hyperparameter suggestions considering the speed in lammps

[Hongyu-yu]

To reach a good balance of accuracy and efficiency in MD, is there any advice?
As illustrated in the allegro paper, a very small net is used with one layer to reach high speed in MD.
For general config such as configs/example.yaml, it suggests a large net however way slower than the speed in the paper.
It seems that following hyperparameters are the most important: r_max, num_layers, l_max, env_embed_multiplicity, two_body_latent_mlp_latent_dimensions, latent_mlp_latent_dimensions, edge_eng_mlp_latent_dimensions
I usually use:

num_layers: 2
l_max: 2
parity: o3_full
env_embed_multiplicity: 16
two_body_latent_mlp_latent_dimensions: [32, 64, 128]
latent_mlp_latent_dimensions: [128]
edge_eng_mlp_latent_dimensions: [128, 64]

which is much smaller than configs/example.yaml and end up with 0.368 ns/day with 7680 atoms on 4 GPUs without Kokko as our server doesn't support:(
Is there any advice or idea for a better speed? Thanks!

[simonbatzner]

Hi @Hongyu-yu, thanks for your interest in our code. You're right there is always a trade-off between accuracy and computational efficiency. The parameters you laid out are all important, here is what I usually scan:

• r_max is by far the most important one, Allegro scales linear with # of edges, so reducing that will make a big difference
• make sure you use a large enough latent mlp, often we see that a large, nonlinear latent mlp results in much better networks. Scalar MLPs are also super cheap so you won't pay much for it, I usually set latent_mlp_latent_dimensions: [128, 128, 128] and use a silu nonlinearity in that MLP.
• make sure your 2-body MLP is large enough, so the example.yaml one for example is a good one, you won't pay much for that
• if 1 layers works well for your system (not always true), it will be much faster
• you can usually reduce parity to so3
• once again, r_max is really the main one, if you can deploy a network with large MLPs and a small r_max, that should usually be fast.

Hope this works.

[Hongyu-yu]

Hi @simonbatzner, thanks for your detailed suggestions!

How about env_embed_multiplicity? Do you suggest using larger MLP with smaller env_embed_multiplicity to reach balance?

[Linux-cpp-lisp]

Yes, in general "larger MLP with smaller env_embed_multiplicity" is one of the core ideas of Allegro (the two tracks) 👍

[Hongyu-yu]

Got it! Thanks!

[Linux-cpp-lisp]

(A side point is that we, like pretty much everyone else, see some very large performance differences between NVIDIA GPU generations, so 4 A100 >> 4 V100 >> etc.)

[Hongyu-yu]

Hi, @Linux-cpp-lisp
Thanks for reminding me that!