search

ACEsuit / mace

MACE - Fast and accurate machine learning interatomic potentials with higher order equivariant message passing.
Other
415 stars 157 forks source link

Adding hessians based on autograd to MACE #367

Closed Nilsgoe closed 1 week ago

Nilsgoe commented 3 months ago

Hello, dear MACE team,

I already answered an Issue post about autograd Hessians and as a result, I got contacted by other people about a autograd Hessian implementation. Therefore, I thought a proper implementation of autograd based Hessians would be very useful for MACE. I tried to make my implementation more clean to hopefully be integrated into a branch of the mace gitHub and here is now a first idea for a possible pull.

To show why autograd Hessians could make sense, I added a comparison in calculation time between a numerical and autograd based Hessian. large_time_vs_size_plot

I have also added a test file for a very simple test, and I am working on more tests. The implementation works for mace_mp, mace_off and self-trained mace potentials. I would welcome a brief discussion about some tests I could run to convince you about the performance and reliability of the implementation. At this point, only the simplest way of getting autograd Hessians is part of the pull request, as the other options either would need larger changes in the mace code or still have some problems (which I would be very happy to discuss at a later point).

ilyes319 commented 3 months ago

Hey , @Nilsgoe thank you very much for your PR, I will have a detailed look, and merge it.

ilyes319 commented 1 week ago

Hey @Nilsgoe, I used your PR and improved a bit the code by swapping from for loops to vmaps. It should be a bit faster, not clear. It is merged now and I made you co-author of the PR so I will close this. Thank you very much.