Unable to reproduce rMD17 results

[JSLJ23]

Hi @sarpaykent,

I've been trying to reproduce the training on rMD17 with a distilled version of GotenNet without the heavy abstraction from pytorch lightning and hydra (model.py contains all the necessary components for the model and train_rmd17.py contains all the code for the training run) but it seems that the test MAEs are no where near what is being reported in the publication. The energy MAEs are about 4 - 25x higher than what is reported and the force MAEs are a few orders of magnitude larger which is pretty alarming. Any idea what I might be missing here? I used most of the default parameters throughout the code.

Energy MAEs for rMD17

|Molecule     |MACE       |GotenNet publication|GotenNet in-house|
|-------------|-----------|--------------------|-----------------|
|aspirin      |2.298298111|1.552435328         |6.608838104      |
|azobenzene   |0.698162256|0.524705795         |3.499541668      |
|benzene      |0.039027704|0.021682058         |0.596525         |
|ethanol      |0.398949861|0.307885219         |4.703444449      |
|malonaldehyde|0.797899722|0.559397087         |4.542900006      |
|naphthalene  |0.901973598|0.16912005          |2.451166669      |
|paracetamol  |1.400660924|0.919319245         |5.383760001      |
|salicylic    |0.698162256|0.611434026         |5.3423           |
|toluene      |0.299212396|0.190802107         |1.98936          |
|uracil       |0.398949861|0.277530338         |5.382300003      |

model.py
train_rmd17.py

I've attached the full training logs for the 10 molecules from the rMD17 datasets with the 5 test splits each. The validation and test MAEs for the forces can be observed to be about a few hundred meV/A for most of the systems, with a few of them dipping below 100.

training_logs.zip

[sarpaykent]

Hi @JSLJ23,

Thank you for your interest in GotenNet and for taking the time to report this issue with detailed results!

I've added an implementation for rMD17 training in #7 that should address the discrepancies you're seeing. I've verified that this implementation reproduces the published results, and I'll be pushing pretrained weights shortly for reference.

In the meantime, please feel free to test the new implementation. You can train on rMD17 using the following command:

python gotennet/scripts/train.py experiment=rmd17_aspirin.yaml label=aspirin +matmul_precision=highest

Please let me know if this resolves the issue or if you encounter any other problems. I appreciate your patience and thorough testing!