tummfm / difftre

Learning neural network potentials from experimental data via Differentiable Trajectory Reweighting
Apache License 2.0
29 stars 9 forks source link

Supplementary Figure 3 a #1

Closed ramess101 closed 2 years ago

ramess101 commented 2 years ago

Greetings. I am very interested in your work. I really enjoyed your paper and this repo has been extremely useful, especially the notebooks!

I have some confusion about the diamond jupyter notebook. The loss curve and time per update seem to be very different from what is reported in supplementary material Figure 3 a.

Were these two figures obtained with different methods? I.e., is the jupyter notebook not the exact same as what you used to generate the plot in supplementary material?

image

image

Thank you

S-Thaler commented 2 years ago

Hi, thanks for reaching out and for the positive feedback!

I assume you're refering to Supplementary Figure 3a of the preprint. Indeed, I had to re-run the diamond example because by the time of the preprint, there was a bug in the Langevin integrator of Jax MD.

The plot in the notebook is consistent with the data shown in Supplementary Figure 5a of the published paper (as well as in the updated arXiv version).

Don't hesitate if something is still unclear!

ramess101 commented 2 years ago

@S-Thaler

Thank you for getting back to me!

I was not aware that this work was published, congratulations! And in Nature, very nice!

You are correct, the image I copied above was from the preprint. The published figure (Supp Figure 5a) is certainly more consistent with what I see in the Jupyter notebook. So that clears up that point. Although I will need to go back and read the published version closely.

However, this does raise one key question/concern:

The old (preprint) figure required far fewer direct MD simulations (as evidenced by the "time per update" curve). Obviously, this means that the preprint version required less real time for the optimization since most update steps relied on reweighting. Do you know why the reweighting was less efficient (i.e., lower Neff) in the published version? Does it make sense that the bug in the Langevin integrator would cause this?

Thanks again

S-Thaler commented 2 years ago

That's certainly a non-trivial question. The bug in the Langevin integrator was, that the atom mass was not considered in the computation of the acceleration, i.e. carbon was simulated, but with the mass of hydrogen. Hence, the kinetic part in the Hamiltonian was far too small and the potential energy component (that DiffTRe is optimizing) artificially dominated.

A strong reduction in N_eff usually means that the optimizer changed the potential energy model a lot. Apparently, with the wrong mass in the simulation, small changes of the model were sufficient such that the simulation sampled states that seamingly generated the correct stiffness.

With the artificially small potential energy in the wrong simulation, maybe the potential was able to more easily drive the atoms towards configurations, that are more consistent with the target stiffness, requiring less changes in the potential. But this is clearly speculation on my side.

Hope this gives some more insights, even though I did not answer your question fully.

ramess101 commented 2 years ago

@S-Thaler

OK, this makes some sense. Thanks for detailing the issue in the Langevin integrator, that was very helpful.

Since you were so kind to answer my many questions, let me introduce myself and our project briefly.

We are hoping to apply this same method for training neural networks to experimental elastic constants for metals, e.g., Aluminum. Perhaps you are familiar with the ANI work Justin Smith did with Aluminium: https://www.nature.com/articles/s41467-021-21376-0 Our plan is to "simply" take that original NN and apply your reweighting method to improve agreement with experiment. Wish us luck!

My background is in training classical potentials to vapor-liquid equilibria experimental data using reweighting approaches (specifically, MBAR): https://pubs.acs.org/doi/abs/10.1021/acs.jced.8b01232 Eventually, we hope to use MBAR to do the reweighting instead.

I guess that raises another question. Did you all ever consider using MBAR to reweight multiple trajectories? I have my doubts that MBAR is really necessary, since most likely only the most recent trajectory would actually contribute significantly...

Thanks again

S-Thaler commented 2 years ago

I did indeed think about reweighting with multiple trajectories, but as you said for most updates only the last trajectory would contribute significantly, so I decided to skip on the implementation.

However, this was with the problems discussed in the current paper in mind, where I started from a potential that is somewhat far away from the target potential, which required a lot of large update steps. However, if the goal is more refining an already reasonable potential, I think the update steps would be rather small so it might be possible to re-use trajectories older than the most recent to improve sampling efficiency.

Good luck with your projects, I'm looking forward to the results!

ramess101 commented 2 years ago

@S-Thaler

Very good point. I am glad that someone else thought about multiple trajectories too, it must not be that crazy of an idea then :-)

Of course, including multiple trajectories introduces some performance questions, especially since you don't know a priori how many update steps will be required. Perhaps we would want a heuristic to determine how many / which trajectories we retain in our reweighting (i.e., just the last 10 trajectories, or based on the contribution to Neff for a given trajectory).

Thanks! It is certainly a very interesting project!

If you don't mind me asking, are you all still working in this space? Are you aware of any other groups trying to employ your method? Just wanting to not "step on anyone's toes".

S-Thaler commented 2 years ago

@ramess101 Yes, there could clearly still be some efficiency gains realized using some smart heuristics. We're currently also working on a more versatile and user-friendly version of DiffTRe, to be open-sourced some months down the road. If you make some progress in this direction, I'd apprechiate a contribution!

Yes our group is still working with DiffTRe and I also know of at least one more group. As far as I'm concerned, I'm of course happy if people are using DiffTRe and I don't think it's very likely that we're all working on the exact same thing as in my view, DiffTRe still offers a lot of open opportunities.

ramess101 commented 2 years ago

@S-Thaler

Cool! I completely agree!

Well, I will certainly be keeping an eye on what you all are doing. As we start to implement this approach, I am sure we will engage in future conversations. I will keep you apprised of any progress we make in the coming months.

Thank you very much and take care!

S-Thaler commented 2 years ago

Looking forward to it, good luck!