Testing the jax implementation of LiTEN

[bhcao]

LiTEN is now included in DIPM (Deep Interatomic Potentials Models in JAX, built on MLIP) and can be accessed from test pypi using

pip install -i https://test.pypi.org/simple/ dipm

Further testing is needed. See https://github.com/bhcao/deep-interatomic-potential-models/ for details.

[bhcao]

Hi @lingcon01, thank you for the contribution of your model to the community. Below is a benchmark test I conducted on MLIPaudit. Your model showed excellent performance on certain metrics, but overall it is slightly worse than the SOTA models:

NOTE: Rank are among Mace, Nequip, VisNet (1 is the best, 4 is the worst) and both trained on SPICE2. See MLIPAudit for details. Testing may not be entirely accurate and is for reference only.

Metric	Score	Rank
bond length distribution	0.88	4
conformer select	0.24	1
dihedral scan	0.68	2
folding stability	0.53	3
noncovalent interactions	0.56	4
reactivity	0.27	1
reference geometry stability	1.0	1
ring planarity	1.0	1(tied)
solvent radial distribution	0.90	4
stability	0.89	4
tautomers	0.13	1

[lingcon01]

Thank you for the evaluation. I am also curious about the LiTEN model setup: since separate checkpoints are used for solvated and non-solvated systems, could you clarify whether the tests were conducted using the corresponding checkpoints for each setting?

[bhcao]

I have implemented your model and fully trained it from scratch on the SPICE2-curated dataset provided by MLIP. The relevant code has been submitted in both my repository and the official MLIP repository.

• My PR to MLIP
• My implementation

The test did not use the checkpoint you provided, but the model I trained on the same dataset as MLIP.

[lingcon01]

OK, I was wondering about the computational efficiency of LiTEN in the JAX framework. Do you have any benchmarks or observations regarding its runtime performance?

[bhcao]

I haven't tested the exact speed, but you can refer to the speeds listed by MLIP. The speed of LiTEN-256 is slower than the listed VisNet, and the speed of LiTEN-128 is faster than the listed VisNet. There is an overall increase of about 4 times in the 300 atom system than torch version. Maybe I can share my model file with you via email so that you can conduct some testing.

[lingcon01]

I am also curious about the hyperparameter settings used for ViSNet and NequIP in your study. The resulting parameter counts appear relatively small—could you clarify how these models were configured?

[lingcon01]

Thank you for sharing!

[bhcao]

I have not trained VisNet/Mace/Nequip by myself. I am comparing the results listed in the MLIPAudit paper. The hyperparameters should be consistent with those in their paper, so, it should indeed be so small.

[lingcon01]

I think I now have a better understanding of the reason. In your evaluations, the number of layers was set to 2 for all models, which is quite interesting. Based on our own benchmarks, LiTEN-256 tends to be somewhat faster than ViSNet, likely because our vector representation has shape [n,3,num_channels], whereas ViSNet uses [n,8,num_channels]. May I ask whether your benchmarks were also conducted on H100 NVIDIA GPU?

[bhcao]

I have sent an email to your Gmail for detailed discussion. If you haven't received it, you can contact me through the email on my homepage.