Radial basis functions parameters

[grsmirnov]

Hi! I have some questions about the choice of radial basis functions parameters.
Am I right that non-trainable normalized basis functions like in your preprint correspond to these values in yaml file?

BesselBasis_trainable: false
normalize_basis: true

I am confused by the fact that the default values are the opposite and the parameter normalize_basis is missed in all yaml examples - I've found it allegro/model/_allegro.py file.
Also, I worry about large values of PolynomialCutoff_p for systems with PBC. I've made some tests for one-component metals with p=48 like you did for Ag and Li3PO4 and always get the equivariance test errors. Tuning other parameters doesn't help, while low p values (like 6 or 8) work fine. Am I missing something? Or this is expected behavior and I should avoid large p with such configurations in the training set?

[simonbatzner]

Hi, thanks for reaching out. Yes, those options are correct for what you want to achieve. The default for normalize_basis should be true. The default for BesselBasis_trainable is true in the example.yaml since it sometimes gives improvement, but we usually see no big effect. Usually safer to just turn it off by default and if you have time/compute try if it helps.

Re the p-value. The default of p=6 works well in most cases. We use p=48 in the Li3PO4 case in the paper since the cutoff is quite small at 4 Angstrom. The envelope function damps interactions at larger radii -- so a smaller p will damp more. We find this helps for extrapolation tests, but if you combine a small p with a short cutoff, you'll overdamp. Therefore we used the p=48 in the paper for the short cutoff. Can you share what your cutoff is and what equivariance errors you're seeing? (i.e. on which symmetries and which targets). Thanks!

[grsmirnov]

@Linux-cpp-lisp @simonbatzner, thanks for your comments!

The default for normalize_basis should be true

Indeed, I overlooked the default value.

A good test to confirm is to run the equivariance tests on default_dtype: float64.

I’ve made some tests with float64 arithmetic and it seems that the equivariance problem was due to the usage of single precision calculations. All errors are gone with default_dtype: float64

I attach the toy problem (10 configuration, bcc titanium, 128 atoms) with the simple yaml file. For example, r_max=4.0 with p=48 (nequip-train ti.yaml --equivariance-test 5) gives in single precision
(parity_k=0, did_translate=True , field=forces ) -> max error=1.019e-03
while in double precision
(parity_k=0, did_translate=True , field=forces ) -> max error=5.130e-13
The complete log of errors in single precision is here:
errors_float32.txt

I will do the additional tests to see the influence of precision on the resultants errors in energy and forces.

[Linux-cpp-lisp]

Whoops I somehow deleted my comment above, but yes I personally wouldn't find that error too alarming. The latest develop branch also contains some developments that might improve this further.

[Linux-cpp-lisp]

Again it's important to note that the equivariance error is of the model with regards to itself, not with regard to the training data. So if it is singificantly smaller than your error with regard to the training data, which it almost certainly is, it shouldn't be an issue.

[simonbatzner]

Thanks for sending this. Yeah for a 4A cutoff, small p-values might hurt you. Can you get away with a larger cutoff? (more expensive but will be more accurate). This overall looks like the smallest config we used in the paper. This will be super scalable and should work, but if you can trade off some speed for accuracy that might be worth it. E.g. I'd definitely use l_max=2, 2 layers, larger MLPs, and more features if I can afford it.

[grsmirnov]

I have made some tests with r_max=4,5,6. Definitely larger values work better in terms of resultant E/N_mae and F_mae. However, I consider r_max> 6 as a problem - atoms start to self-interact through PBC, so we need larger supercells and more demanding DFT. This error is probably small for crystals, but I would like to avoid hard-to-predict things.

Also, so far I have not seen a direct correlation between the r_max and equivariance errors. Some combinations of hyperparameters pass the equivariance test, some don't. Good to know that default tolerance is probably overkill.

[Linux-cpp-lisp]

The only one I'd expect to strongly scale with the number of neighbors is the permutation equivariance, since the sums go over more terms...