Precision

[terryfrankcombe]

I'm fitting a model where the raw energies are O(10⁵). Surprisingly to me, the trained model ran out of precision and produced predictions with significant aliasing... discrete constant calculated energies for a range of geometries, separated from adjacent calculated energies at more distant geometries by significant gaps, O(10^-4) IIRC. For float64 numerics there should still be plenty of precision left to resolve energies between these energies.

Retraining to these energies shifted by a constant so that the training energies were all O(1) gives a model that behaves normally, giving good predictions for all geometries.

Is this expected behavior?

[Linux-cpp-lisp]

Hi @terryfrankcombe ,

Are you using develop, and are you changing any of the dtype settings? What scaling and shifting settings are you using?

[terryfrankcombe]

To clarify a point: I was talking about float64 numerics in my original post as that's what get_potential_energy() returns from the ASE calculator.

[Linux-cpp-lisp]

Aha, yes, on main presently the default is float32 everywhere, which can be changed with default_dtype: float64. In principle our calculator respects dtype (https://github.com/mir-group/nequip/blob/main/nequip/ase/nequip_calculator.py) but maybe ASE is doing some internal up-casting...

In general the solution to this should be to switch to develop (0.6.0), which implements the mixed precision strategy described in section V-B3 of this recent pre-print of ours: https://arxiv.org/abs/2304.10061. This should resolve issues with fidelity of representing potential energies without sacrificing performance. You can get the config used in that manuscript with:

default_dtype: float64
model_dtype: float32
allow_tf32: true   # or `false`

I assume then that I could change that to float64 and everything would be run more precisely? I also guess that the performance penalty for that would depend a lot on the hardware?

Yes and yes. In particular the penalty for non-datacenter GPUs is quite bad, but even for datacenter/cluster class GPUs float64 can be quite slow (see Table IV same preprint).

(For us folks coming from CPU-land that anyone would ever use anything but 64 bit reals comes as a shock!)

Definitely---and entertainingly, the machine learning people often can't see why we would ever use more than 16-bit floats.

[Linux-cpp-lisp]

Side note: I'm not sure if you are running on CPUs, but if you are, the TensorFloat32 flag is irrelevant there.

[terryfrankcombe]

I've seen this develop branch mentioned in a few discussions.

Is it stable enough now for you to recommend switching to 0.6 for production calcs? What are the headline changes?

[Linux-cpp-lisp]

The develop branch is pretty stable and I intend to release it soon; whether you'd like to switch now is a matter of where you are in your project I think. Please note that 0.6.0 is a major release from 0.5.* and involves a number of breaking changes. I don't plan to introduce any further breaking changes but there are a few features and non-critical bugfixes that are still pending.