iiinova — where your digital self keeps evolving

It was super fun collaborating with my co-first-author @dft-dutoit.bsky.social, together with the rest of the team across various research groups: @bm-chiheb.bsky.social, Zoé Faure Beaulieu, Bianca Pasça...

100

We hope that you can start using this method to do cool new science!

100

I find our results for the modelling of MAPI (a hybrid perovskite) particularly pleasing: the distribution of cation orientations generated by the teacher and student models during NVT MD are ~identical!

100

We go on to apply this distillation approach to target other chemical domains by distilling different foundation models (Orb, MatterSim @msftresearch.bsky.social) and MACE-OFF), and find that it works well across the board!

100

Beyond error metrics, we extensively validate these models to show they model liquid water well.

100

These student models have relatively few parameters (c. 40k for PaiNN and TensorNet), and so have much lower memory footprint. This lets you scale single GPU experiments very easily!

100

The resulting student models reach impressive accuracy vs DFT while being orders of magnitude faster than the teacher! Note that these student models are of a different architecture to MACE, and in fact ACE is not even NN-based.

100

We start by (i) fine-tuning MACE-MP-0 (@ilyesbatatia.bsky.social) on 25 water structures labelled with an accurate functional, (ii) using this fine-tuned model and structures to generate a large number (10k) new “synthetic” structures, and (iii) training student models on this dataset.

100

Does this distillation approach work? In short, yes! 🤩

100

This approach is very cheap, taking c. 5 calls to the teacher model to generate a new, chemically relevant and uncorrelated structure! We can build large datasets within one hour using this protocol.

100

In this pre-print, we propose a different solution: starting from a (very) small pool of structures, and repeatedly (i) rattling and (ii) crudely relaxing them using the teacher model and a Robbins-Monro procedure.

100

This works well, but has two drawbacks: (1) MD is still quite expensive, and requires many steps to generate uncorrelated structures, and (2) expert knowledge and lots of fiddling is required to get the MD settings right.

100

In previous work, we and others (PFD-kit) have proposed using teacher models to generate "synthetic data" by using them to drive MD, and to sample snapshots along these trajectories as training points.

100

The devil is always in the details however 😈 The main problem we need to solve is how to generate many relevant structures that densely sample the chemical domain we are interested in targeting.

100

At a high level, this builds upon the approach pioneered by Joe Morrow, now extended to the distillation of impressively capable foundation models, and to a range of downstream architectures and chemical domains,

100

Concretely, we train a student to predict the energy and force labels generated by the teacher on a large dataset of structures: this requires no alterations to existing training pipelines, and so is completely agnostic to the architecture of both the teacher and student 😎

100

Both of the above methods try to maximise the amount of information extracted per training structure from the teacher. Our approach is orthogonal to this: we try to maximise the number of structures (that are both sensible and useful) we use to transfer knowledge.

100

Somewhat similarly, @ask1729.bsky.social and others extract additional Hessian information from the teacher. Again, this works well providing you have a training framework that lets you train student models on this data.

100

@gasteigerjo.bsky.social and others attempt to align not only the predictions, but also the internal representations of the teacher and the student. This approach works well for models with similar architectures, but is incompatible with e.g. fast linear models like ACE.

100

At their heart, all model distillation strategies attempt to extract as much information as possible from the teacher model, in a format that is useful for the student. Various existing methods in the literature do this in different ways.

100

In the context of machine learned interatomic potentials, distillation lets you simulate larger systems, for longer times, and with less compute. This lets you explore new science, and democratises access to otherwise expensive simulations/methods and foundation models. 💪

100

Model distillation methods take a large, slow “teacher model”, and try to condense, or “distill”, its knowledge into a smaller, faster “student model” If this can be done well, it is an extremely useful thing!

100

Excited to share the pre-print we’ve been working on for the last ~4 months: “Distillation of atomistic foundation models across architectures and chemical domains” Deep dive thread below! 🤿🧵

4112

graph-pesContentsMenuExpandLight modeDark modeAuto light/dark, in light modeAuto light/dark, in dark mode

If you want to dive straight in with some hands-on examples, you can find several links to Colab notebooks in the graph-pes docs: jla-gardner.github.io/graph-pes/

jla-gardner.github.io

000

GitHub - jla-gardner/graph-pes: train and use graph-based ML models of potential energy surfaces

If graph-pes sounds interesting to you, please do check out the repo and give it a star! Please also reach out via GitHub issues or DM on here if you have any questions or feedback. github.com/jla-gardner/...

github.com

train and use graph-based ML models of potential energy surfaces - jla-gardner/graph-pes

Internally, we've been finding that this ease of use is massively accelerating our research. Watch this space for exciting research using graph-pes coming to your screens soon!

100

I've put a lot of work into making graph-pes easy to use, both for standard MLIP training via the command line interface, and for methods research using the Python API.

100

This lets you fine-tune any of these foundation models on your own datasets, and on any combination of energy, force and stress labels. Crucially, you can easily switch between any of these models by changing a single line in your graph-pes config file.

100

graph-pes now also interfaces to a range of popular foundation models, including MACE-MP, MACE-OFF, MatterSim and Orb v2/3.

I've re-implemented several popular architectures from scratch within graph-pes, including NequIP, PaiNN, TensorNet and MACE. You can take any of these models and train them on your own datasets, before using them to make predictions and run MD in ASE and LAMMPS.

100

graph-pes is a package for building, training and using graph-based machine learning models that predict energies and forces of atomic structures. #ChemSky #CompChemSky

120

John Gardner @jla-gardner.bsky.social · 10/12/2024

I've been working hard behind the scenes, and am excited to announce that graph-pes version 0.1.0 is now out!

180

Thanks @guillemsimeon.bsky.social - I remember first hearing about TensorNet during your talk at CECAM-PsiK Berlin 2023, implementing it was easy thanks to your very detailed paper 😊

120

(and gradients of these functions are just the forces on each atom, and hence are very useful for e.g. running simulations of how molecules/materials will behave/decompose under different conditions)

010

Thanks ☺️ Potential energy *surfaces are functions that map the positions of each atom in a structure to an energy. By implementing these in PyTorch, we can use automatic differentiation to get exact gradients very quickly, rather than approximate gradients slowly via finite difference approaches.

Huge thanks to @vlderinger.bsky.social and the rest of the group here in Oxford for their support and help in testing early versions of graph-pes! Exciting work using this tool will be hitting your feeds very soon…

010

GitHub - jla-gardner/graph-pes: train and use graph-based ML models of potential energy surfaces

Stay tuned for more details in due course, but as a final point: installing graph-pes is easy! pip install graph-pes Any comments, feature requests or bug reports? Find graph-pes on GitHub: github.com/jla-gardner/...

github.com

train and use graph-based ML models of potential energy surfaces - jla-gardner/graph-pes

141

My hope is that by providing this integrated features set, graph-pes can greatly reduce the barrier to entry for MLIP researchers wanting to go beyond “standard” training regimes and model evaluations.

100

By inheriting from our GraphPESModel base class, you can immediately: * get automatic force and stress predictions (if these aren’t already predicted by your model) * train your model just like any other using graph-pes-train * use your model to drive MD via LAMMPS

100

Finally, graph-pes is for research scientists wanting to tinker with existing, and create new, MLIP architectures.

Installing LAMMPS on new systems, and with specialised pair styles, can be tricky and time consuming! To this end, I include a build_lammps.sh script to make installing LAMMPS with graph-pes support easy - this has worked everywhere we’ve tried it so far 😊

100