Thread by @simonbatzner

🚨 What comes after Neural Message Passing? 🚨

Introducing Allegro - a new approach:

- no message passing or attention
- new SOTA on QM9+MD17
- scales to >100 million atoms
- 1 layer beats all MPNNs+Transformers
- blazing fast
- theory

arxiv.org/abs/2204.05249

How?👇#compchem

First and foremost: this was joint with my co-first author and good friend Albert Musaelian with equal first-author contribution as well as with lab members Anders Johansson, Lixin Sun, Cameron Owen + Mordechai Kornbluth and of course @BKoz / Boris Kozinsky

Message Passing Neural Networks have taken molecular ML by storm and over the past few years, a lot of progress in Machine Learning for molecules and materials has been variations on this theme.

We propose an alternative approach: a new layer based simply on tensor products. We demonstrate it outperforms all existing approaches, including all Message Passing Neural Networks + Transformers. No messages passing, no attention.

The interesting part: it's purely local, using strict cutoffs and no propagation of information. Due to this locality, we can make us of GPU-parallelism and scale it to massive systems. This also challenges the idea that the best models need Message Passing or Attention.

Here's how it works: we start by decomposing the total energy of the system into a set of pairwise contributions E_(i,j), instead of the conventional per-atom decomposition. We then embed these pairs of atoms (i,j) into a weighted spherical harmonics projection.

These E(3)-equivariant features that describe the state of pair (i,j) are updated by computing a tensor product between the current pair feature (i,j) and all other pair features in the environment of atom i, i.e. computing the tensor product between (i,j) and all pairs (i, k).

Naively, this gives exponential scaling. The central mathematical trick is that we can exploit the bilinearity of the tensor product to first compute the sum over the environment and then compute the tensor product. This reduces this operation to a single tensor product.

The weights for the embedding are generated by a MLP. We mix the tensor product outputs by an equivariant linear layer/a MLP and then iteratively apply this procedure. At the output, we sum the pair energies E_(i,j) into the total energy. Autodiff for the forces. That's it.

We benchmark this extremely simple idea and find state-of-the-art performance on revMD17 and QM9, outperforming all other methods. Interestingly, we find that even a single tensor product layer outperforms all message-passing + transformer-based approaches on QM9!

We also test generalization to out-of-distribution data on the 3BPA benchmark and find that Allegro greatly outperforms existing potentials, including also an ANI model pretrained on 8.9 million molecules (we use 500 molecules). The only method that is competitive is NequIP.

This breaks with the notion that Deep Learning Interatomic Potentials don't generalize. They often generalize better, in fact much better than linear models or kernel methods, but you have to use the right ones.

Due to the locality, Allegro is extremely scalable: we show that on as little as 16 GPU nodes (8xA100s each), we can scale it up to 100 million atoms and a speed of 1.5 ns/day, all at excellent accuracy (previous approaches had to use ~27,000 GPUs for sizes like these).

Oh and it is fast: on DFT-sized systems, we can simulate ~90ns/day on 1 NVIDIA DGX A100 GPU, ~10ns/day on 1 million atoms on 1 node, ~50 ns/day on 1 million atoms on 8 nodes!

Everything is integrated with LAMMPS running fully on the GPU (inference + integration) --> no CPU-GPU transfer! We report strong scaling results and see a nice scale-up both intra- and inter-node.

We then show Allegro works well in the wild, i.e. not just MD-17 ;-) We demonstrate Allegro predicts structure + Li-dynamics of a complex amorphous phosphate material with an extremely scalable model.

Code will be made public soon. It's integrated with LAMMPS+ASE. Everything is built on top of our NequIP API to make usage as simple as possible (github.com/mir-group/nequ…). If you're a NequIP user, this is a few-lines change in your input file (if not, welcome to the family)

Tagging some folks who might find this interesting: @ProfvLilienfeld @pfau @jhrmnn @ktschuett @ZhuoranQ @AndersSChristen @jonkhler @jtmargraf @JoshRackers @jchodera @jo_brandstetter @TacoCohen @lab_COSMO @mmbronstein @cwcoley @chaitjo @wellingmax @erikjbekkers @HochreiterSepp