Add ALIGNN FF (aborted)

[janosh]

All credit for this PR goes to @pbenner. Thanks for this submission! 🙏

This PR copies over a subset of files in pbenner/matbench-discovery#621d4f1a6.

Update 2023-08-09

This model submission was aborted due to ongoing technical challenges. See the model readme for details.

[janosh]

@pbenner Could you check if these package version numbers are correct (i.e. match the ones you were using for this submission).

[pbenner]

Yes, the versions are correct, except for ase -> 3.22.1

[janosh]

@pbenner Just to clarify, you used alignnff_wt10?

[pbenner]

Exactly, it was called best_model.pt before:
alignn/ff/best_model.pt → alignn/ff/alignnff_wt10/best_model.pt

[pbenner]

Not sure if it makes sense to test some of the newly added models. Can give it a try, at least check if the convergence is better

[janosh]

Can we work around this new dependency? I'm not sure how it differs from tqdm? Is tqdm lacking parallel process support?

[pbenner]

Yes, you can just use tqdm. For parallel processing pqdm has a bit nicer output, that's all.

[janosh]

@pbenner I started evaluating the ALIGNN FF predictions and they look bad enough that I think we may have made a mistake. Are we sure ALIGNN FF was trained on the MP legacy correction scheme (MaterialsProjectCompatibility)? Perhaps they trained on raw VASP energies? We'll have to consult the paper or check the code base.

[pbenner]

I think I forgot to remove the energy corrections for the custom ALIGNN model. Uploaded two new files:

• 2023-07-15-custom-alignn-relaxed-wbm-IS2RE.csv.gz: Custom model without energy correction
• 2023-07-15-mp_e_form_alignn-relaxed-wbm-IS2RE.csv.gz: Pre-trained model with correction applied
  But still, I would like to check the ALIGNN-FF model in detail. Will do that next week. Maybe it makes sense to train it from scratch.

[janosh]

Just to clarify, the model you trained from scratch that produced 2023-07-15-custom-alignn-relaxed-wbm-IS2RE.csv.gz was trained on MP ComputedStructureEntries that already include the 2020 MP energy correction scheme. So the predictions it makes are also corrected energies? By "custom model without energy correction", do you mean you subtracted the energy corrections again from every model prediction?

[pbenner]

By "custom model without energy correction" I meant that I didn't apply them. It's this if-statement that I forgot to add to your script:
https://github.com/pbenner/matbench-discovery/blob/25616eeba7e71b0909feccddfef88dfc63c60314/models/alignn/test_alignn_relaxed.py#L118
You can find the changes here in this commit:
pbenner@25616ee
And I re-executed the script twice with the corresponding model_name variable set. Hope this makes sense.

Just also added a few comments to the script.

[pbenner]

Since ALIGNN-FF was pre-trained on the JARVIS data, which is DFT simulation data with OptB88vdW functional (with van der Waals correction), the results are not comparable and I would drop it. Would you consider adding training data to matbench discovery (forces/stresses) that allows the user to train a force field model? I think this would be very valuable and make the benchmark applicable to a wider class of models, where no good pre-trained model exists.

[janosh]

Since ALIGNN-FF was pre-trained on the JARVIS data, which is DFT simulation data with OptB88vdW functional (with van der Waals correction), the results are not comparable and I would drop it.

That's one option. We could also add it anyway with a big disclaimer just to highlight the importance of homogeneous DFT settings between train and test set.

Would you consider adding training data to matbench discovery (forces/stresses) that allows the user to train a force field model?

Yes, that data is definitely part of the MBD training set. I have a local copy of the energies, forces, stresses and magmoms for every ionic step in the MP database. It's ~14.7 GB. I was planning to upload it to figshare at some point.

But considering my local data is identical to the MPtraj dataset used to train CHGNet (~11.3 GB) on Figshare (except for some cleaning which most people will benefit from), I think we might as well just point to that.

[JaGeo]

That's one option. We could also add to the benchmark anyway with a big disclaimer just to highlight how important homogeneous DFT settings between train and test set are.

I wouldn't do this. I think it should be obvious for anyone using DFT that without compatibility scheme you should not do any kind of comparison.

[pbenner]

I converted the MPtraj dataset to ALIGNN compatible format:
https://github.com/pbenner/matbench-discovery/blob/alignn/models/alignn/make_train_data_ff.py
It will probably take too long to train ALIGNN-FF, but I'm just giving it a try.

[janosh]

I wouldn't do this. I think it should be obvious for anyone using DFT that without compatibility scheme you should not do any kind of comparison.

Thanks for your input @JaGeo! Always good to hear other's thoughts. I'm leaning the same way.

It will probably take too long to train ALIGNN-FF, but I'm just giving it a try

@pbenner That's great! Thanks for the update!

Even if that doesn't pan out, I might still merge this PR just to keep a record that we tried to add ALIGNN-FF and show how far we got as well as explain in a readme why we aborted.

[pbenner]

I tried to optimize the ALIGNN-FF training. Unfortunately, the model is quite resource hungry. The 12 GB of CHGNet training data turn into 600 GB of graph data. It still fits into memory, but prevents me to use the accelerator for training on multiple GPUs, which creates a copy for each process. Also, a batch size of 16 already fills up all 80 GB of GPU memory. Training is quite slow with such a small batch size (roughly a day per epoch).

I'm now fine tuning the ALIGNN-FF WT10 model on the CHGNet data, which hopefully takes fewer epochs than training a new model from scratch.

[janosh]

I'm now fine tuning the ALIGNN-FF WT10 model on the CHGNet data, which hopefully takes fewer epochs than training a new model from scratch.

Thanks for trying @pbenner! I think we've done everything that can reasonably be expected here to get ALIGNN FF into MBD. Rather than spend more time on it, how about we call it quits and move on the wrapping up the preprint?

[pbenner]

Yes agreed. I also saw that the WT10 model trained on Jarvis has a worse initial loss on the CHGNet data than the untrained model, which indicates that these two datasets are very much incompatible. I'll let it rest for now...