embedding layer in nequip

[mesonepigreco]

Dear developers,
I'm curious to understand how nequip deals with different atoms. While I know it uses a one-hot vector representation for atomic numbers at the beginning, the paper [Batzner et al, Nat Comm, 13, 2453, 2022] mentions an embedding procedure.
My understanding of embedding is that you use a new trained layer or network to transform this atomic species space into something that has chemical meaning, in which atoms close in this space will have a similar chemical function (if trained on an appropriate dataset). Is this what Nequip is doing? Have you tested if distances in this embedded space have a chemical or physical meaning? Or am I misunderstanding the way nequip works?

P.S.
I would like to experiment with this. What are the parts of the code that deals with this embedding?

[Linux-cpp-lisp]

Hi @mesonepigreco ,

Thanks for your interest in our code!

My understanding of embedding is that you use a new trained layer or network to transform this atomic species space into something that has chemical meaning

Exactly. At a code level, nequip has modules that make up the pieces of the model, and model builders, which put them together. You can see the module that does the one hot embedding as the node features here:
https://github.com/mir-group/nequip/blob/main/nequip/nn/embedding/_one_hot.py#L12. And the model builder for the NequIP model that uses that module here:
https://github.com/mir-group/nequip/blob/main/nequip/model/_eng.py#L103. You can see later that is followed by an AtomwiseLinear module called chemical_embedding (https://github.com/mir-group/nequip/blob/main/nequip/model/_eng.py#L107), which is implemented here: https://github.com/mir-group/nequip/blob/main/nequip/nn/_atomwise.py#L31. You can see that by default it operates on the node features.

https://github.com/mir-group/nequip-example-extension might be helpful in understanding these things as well.

in which atoms close in this space will have a similar chemical function (if trained on an appropriate dataset) [...] Have you tested if distances in this embedded space have a chemical or physical meaning?

Yes, in general they should. However this is a very early part of the network, and these embeddings can be processed further quite a bit throughout the network, so one can't assume that, for example, distances between early embeddings mean something definitive about the difference between later features or final predictions. We have done some work on using the final features (rather than these features at the very beginning) for UQ: https://pubs.aip.org/aip/jcp/article/158/16/164111/2886901.